|
1 | 1 | import re |
2 | 2 | from datetime import datetime, timezone |
3 | | -from typing import Dict, List, Optional, Set, Tuple, cast |
| 3 | +from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Set, Tuple, cast |
4 | 4 |
|
5 | 5 | from fastapi import HTTPException |
6 | 6 | from starlette.datastructures import Headers |
|
30 | 30 | ) |
31 | 31 | from litellm.proxy._types import ( |
32 | 32 | UI_TEAM_ID, |
| 33 | + LiteLLM_ObjectPermissionTable, |
33 | 34 | LiteLLM_TeamTable, |
34 | 35 | ProxyException, |
35 | 36 | SpecialHeaders, |
|
43 | 44 | user_api_key_auth, |
44 | 45 | ) |
45 | 46 | from litellm.proxy.common_utils.http_parsing_utils import _read_request_body |
46 | | -from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl |
| 47 | +from litellm.proxy.common_utils.user_api_key_cache import ( |
| 48 | + USER_NO_MCP_PERMISSION_SENTINEL, |
| 49 | + get_management_object_ttl, |
| 50 | + user_object_permission_id_cache_key, |
| 51 | +) |
47 | 52 | from litellm.repositories.table_repositories import ( |
48 | 53 | AgentsRepository, |
49 | 54 | MCPServerRepository, |
50 | 55 | ) |
| 56 | +from litellm.repositories.user_repository import UserRepository |
51 | 57 | from litellm.types.mcp_server.mcp_server_manager import MCPServer |
52 | 58 |
|
| 59 | +if TYPE_CHECKING: |
| 60 | + from litellm.proxy.utils import PrismaClient |
| 61 | + |
| 62 | + |
| 63 | +def _as_list(values: Sequence[str] | None) -> list[str] | None: # mutable-ok: resolver returns a list |
| 64 | + """Widen a read-only allowlist back to the mutable list the resolver's own contract returns, |
| 65 | + preserving the ``None`` that means "no restriction".""" |
| 66 | + return None if values is None else list(values) |
| 67 | + |
53 | 68 |
|
54 | 69 | def _parse_mcp_server_names_from_path(path: str, mcp_servers_header: Optional[List[str]] = None) -> Optional[List[str]]: |
55 | 70 | """Resolve the single MCP server name a cold-start passthrough bypass may |
@@ -1408,6 +1423,15 @@ async def get_allowed_mcp_servers( |
1408 | 1423 | f"Applied agent intersection filter. Final allowed servers: {allowed_mcp_servers}" |
1409 | 1424 | ) |
1410 | 1425 |
|
| 1426 | + ######################################################### |
| 1427 | + # Apply the internal user's own ceiling (the entitlement attached to the human) |
| 1428 | + ######################################################### |
| 1429 | + capped, user_restricts = await MCPRequestHandler._apply_user_server_ceiling( |
| 1430 | + allowed_mcp_servers, user_api_key_auth, keyless_source=keyless_source |
| 1431 | + ) |
| 1432 | + allowed_mcp_servers = list(capped) |
| 1433 | + has_lower_level_mcp_restrictions = has_lower_level_mcp_restrictions or user_restricts |
| 1434 | + |
1411 | 1435 | ######################################################### |
1412 | 1436 | # Apply org-level ceiling if org_id is set |
1413 | 1437 | ######################################################### |
@@ -1831,6 +1855,12 @@ async def get_allowed_tools_for_server( |
1831 | 1855 | # No team restrictions → use key restrictions |
1832 | 1856 | allowed_tools = cast(List[str], key_tools) |
1833 | 1857 |
|
| 1858 | + allowed_tools = _as_list( |
| 1859 | + await MCPRequestHandler._apply_user_tool_ceiling( |
| 1860 | + allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source |
| 1861 | + ) |
| 1862 | + ) |
| 1863 | + |
1834 | 1864 | return await MCPRequestHandler._apply_agent_and_org_tool_ceilings( |
1835 | 1865 | allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source |
1836 | 1866 | ) |
@@ -2376,6 +2406,203 @@ async def _get_allowed_mcp_servers_for_end_user( |
2376 | 2406 | verbose_logger.warning(f"Failed to get allowed MCP servers for end_user: {str(e)}") |
2377 | 2407 | return [] |
2378 | 2408 |
|
| 2409 | + @staticmethod |
| 2410 | + async def _get_user_object_permission( |
| 2411 | + user_api_key_auth: UserAPIKeyAuth | None = None, |
| 2412 | + ) -> LiteLLM_ObjectPermissionTable | None: |
| 2413 | + """The internal user's OWN object_permission: the entitlement attached to the HUMAN rather |
| 2414 | + than to the credential they authenticated with. |
| 2415 | +
|
| 2416 | + A key's object_permission is the credential's scope and a team's is the group's; this one |
| 2417 | + answers "which MCP servers and tools is this person entitled to", independent of how many keys |
| 2418 | + they hold. Caches the ``user_id -> object_permission_id`` mapping (with a sentinel for "no |
| 2419 | + entitlement") exactly as the agent path does, then reuses the shared ``object_permission_id`` |
| 2420 | + cache, so a warm request reads no rows. |
| 2421 | +
|
| 2422 | + ``None`` means the human places NO ceiling: no user row, or a row naming no permission. The |
| 2423 | + two fault classes are deliberately NOT collapsed into that: a user row we cannot read leaves |
| 2424 | + us unable to say whether they are entitled at all, which is exactly the state before this |
| 2425 | + level existed, so it places no ceiling; a row that NAMES a permission we cannot read is a |
| 2426 | + KNOWN entitlement with unknown contents, so it raises and the caller denies. |
| 2427 | + """ |
| 2428 | + from litellm.proxy.auth.auth_checks import get_object_permission |
| 2429 | + from litellm.proxy.proxy_server import ( |
| 2430 | + prisma_client, |
| 2431 | + proxy_logging_obj, |
| 2432 | + user_api_key_cache, |
| 2433 | + ) |
| 2434 | + |
| 2435 | + if not user_api_key_auth or not user_api_key_auth.user_id: |
| 2436 | + return None |
| 2437 | + |
| 2438 | + if prisma_client is None: |
| 2439 | + verbose_logger.debug("prisma_client is None") |
| 2440 | + return None |
| 2441 | + |
| 2442 | + user_id = user_api_key_auth.user_id |
| 2443 | + object_permission_id = await MCPRequestHandler._user_object_permission_id(user_id, prisma_client) |
| 2444 | + if object_permission_id is None: |
| 2445 | + return None |
| 2446 | + |
| 2447 | + object_permission = await get_object_permission( |
| 2448 | + object_permission_id=object_permission_id, |
| 2449 | + prisma_client=prisma_client, |
| 2450 | + user_api_key_cache=user_api_key_cache, |
| 2451 | + parent_otel_span=user_api_key_auth.parent_otel_span, |
| 2452 | + proxy_logging_obj=proxy_logging_obj, |
| 2453 | + ) |
| 2454 | + if object_permission is None: |
| 2455 | + raise ValueError( |
| 2456 | + f"user {user_id!r} names object_permission_id {object_permission_id!r} which could not be loaded" |
| 2457 | + ) |
| 2458 | + return object_permission |
| 2459 | + |
| 2460 | + @staticmethod |
| 2461 | + async def _user_object_permission_id(user_id: str, prisma_client: "PrismaClient") -> str | None: |
| 2462 | + """The permission row this human's user row links to, or None when they link none. |
| 2463 | +
|
| 2464 | + Caches the link (with a sentinel for "links none") so a human without an entitlement costs no |
| 2465 | + DB read per MCP request. Anything other than an id string is treated as a cache MISS rather |
| 2466 | + than carried into the permission lookup, and a read that fails answers None: not knowing |
| 2467 | + whether someone is entitled is the state that existed before this level, so it places no |
| 2468 | + ceiling. Only a link we DID resolve can make the caller deny. |
| 2469 | + """ |
| 2470 | + from litellm.proxy.proxy_server import user_api_key_cache |
| 2471 | + |
| 2472 | + cache_key = user_object_permission_id_cache_key(user_id) |
| 2473 | + try: |
| 2474 | + cached: object = await user_api_key_cache.async_get_cache(key=cache_key) |
| 2475 | + if cached == USER_NO_MCP_PERMISSION_SENTINEL: |
| 2476 | + return None |
| 2477 | + if isinstance(cached, str) and cached: |
| 2478 | + return cached |
| 2479 | + user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) |
| 2480 | + linked: object = getattr(user_row, "object_permission_id", None) if user_row is not None else None |
| 2481 | + object_permission_id = linked if isinstance(linked, str) and linked else None |
| 2482 | + await user_api_key_cache.async_set_cache( |
| 2483 | + key=cache_key, |
| 2484 | + value=object_permission_id or USER_NO_MCP_PERMISSION_SENTINEL, |
| 2485 | + ttl=get_management_object_ttl(user_api_key_cache), |
| 2486 | + ) |
| 2487 | + return object_permission_id |
| 2488 | + except Exception as e: # noqa: BLE001 # unknown whether entitled at all: no ceiling, as before |
| 2489 | + verbose_logger.warning(f"MCP user entitlement: link for {user_id!r} unresolved, no ceiling: {str(e)}") |
| 2490 | + return None |
| 2491 | + |
| 2492 | + @staticmethod |
| 2493 | + async def _get_allowed_mcp_servers_for_user( |
| 2494 | + user_api_key_auth: UserAPIKeyAuth | None = None, |
| 2495 | + ) -> Sequence[str] | None: |
| 2496 | + """The MCP servers the internal user is entitled to, as server ids. |
| 2497 | +
|
| 2498 | + ``[]`` means this human places no restriction (allow-all from this level); ``None`` means the |
| 2499 | + ceiling is UNRESOLVED, which the caller denies on. Servers named only under |
| 2500 | + ``mcp_tool_permissions`` count as entitled, exactly as they do for a key or a team, so |
| 2501 | + granting one tool never requires naming its server twice. |
| 2502 | + """ |
| 2503 | + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( |
| 2504 | + global_mcp_server_manager, |
| 2505 | + ) |
| 2506 | + |
| 2507 | + try: |
| 2508 | + object_permissions = await MCPRequestHandler._get_user_object_permission(user_api_key_auth) |
| 2509 | + if object_permissions is None: |
| 2510 | + return [] |
| 2511 | + |
| 2512 | + direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or []) |
| 2513 | + access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups( |
| 2514 | + object_permissions.mcp_access_groups or [] |
| 2515 | + ) |
| 2516 | + tool_perm_servers = list( |
| 2517 | + global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys() |
| 2518 | + ) |
| 2519 | + return list(set(direct_mcp_servers + access_group_servers + tool_perm_servers)) |
| 2520 | + except Exception as e: # noqa: BLE001 # any resolution fault is an unresolved ceiling, never "no ceiling" |
| 2521 | + verbose_logger.warning(f"Failed to get allowed MCP servers for user: {str(e)}") |
| 2522 | + return None |
| 2523 | + |
| 2524 | + @staticmethod |
| 2525 | + async def _apply_user_server_ceiling( |
| 2526 | + allowed_mcp_servers: Sequence[str], |
| 2527 | + user_api_key_auth: UserAPIKeyAuth | None = None, |
| 2528 | + *, |
| 2529 | + keyless_source: bool = False, |
| 2530 | + ) -> tuple[tuple[str, ...], bool]: |
| 2531 | + """Narrow a resolved server list by the internal user's own entitlement. |
| 2532 | +
|
| 2533 | + Returns the capped list and whether this human restricted it at all; the caller needs the |
| 2534 | + second value because an org list may only CAP a lower-level restriction, never replace one, so |
| 2535 | + a user ceiling has to be visible to the org step. |
| 2536 | +
|
| 2537 | + RAISES when the entitlement is known but unreadable, which the resolver's own handler turns |
| 2538 | + into deny-all. That is the point of the level: dropping a ceiling we know exists is exactly the |
| 2539 | + silent widening it is there to prevent. |
| 2540 | + """ |
| 2541 | + if keyless_source: |
| 2542 | + return tuple(allowed_mcp_servers), False |
| 2543 | + entitled = await MCPRequestHandler._get_allowed_mcp_servers_for_user(user_api_key_auth) |
| 2544 | + if entitled is None: |
| 2545 | + raise ValueError( |
| 2546 | + f"MCP user ceiling unresolvable for user_id=" |
| 2547 | + f"{user_api_key_auth.user_id if user_api_key_auth else None!r}" |
| 2548 | + ) |
| 2549 | + if not entitled: |
| 2550 | + return tuple(allowed_mcp_servers), False |
| 2551 | + capped = tuple(server for server in allowed_mcp_servers if server in set(entitled)) |
| 2552 | + verbose_logger.debug(f"Applied user ceiling filter. Final allowed servers: {capped}") |
| 2553 | + return capped, True |
| 2554 | + |
| 2555 | + @staticmethod |
| 2556 | + async def _user_places_mcp_ceiling(user_api_key_auth: UserAPIKeyAuth | None = None) -> bool: |
| 2557 | + """Whether this human's own entitlement bounds their MCP access at all. |
| 2558 | +
|
| 2559 | + True when they are entitled to a specific set of servers, and also when that entitlement is |
| 2560 | + UNRESOLVED — a caller uses this to decide whether it may skip the resolver, and skipping it on |
| 2561 | + a transient fault would widen access. |
| 2562 | + """ |
| 2563 | + entitled_servers = await MCPRequestHandler._get_allowed_mcp_servers_for_user(user_api_key_auth) |
| 2564 | + return entitled_servers is None or len(entitled_servers) > 0 |
| 2565 | + |
| 2566 | + @staticmethod |
| 2567 | + async def _apply_user_tool_ceiling( |
| 2568 | + allowed_tools: Sequence[str] | None, |
| 2569 | + server_id: str, |
| 2570 | + user_api_key_auth: UserAPIKeyAuth | None = None, |
| 2571 | + *, |
| 2572 | + keyless_source: bool = False, |
| 2573 | + ) -> Sequence[str] | None: |
| 2574 | + """Narrow a key/team tool allowlist by the internal user's own tool entitlement. |
| 2575 | +
|
| 2576 | + The human's entitlement can only ever narrow: a user naming tools on ``server_id`` intersects |
| 2577 | + (and becomes the allowlist when no lower level restricts), while a user naming none places no |
| 2578 | + restriction. Returns ``[]`` (deny every tool on this server) when the entitlement cannot be |
| 2579 | + resolved, because the caller's own except-handler treats a raise as allow-all for key auth. |
| 2580 | + """ |
| 2581 | + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( |
| 2582 | + global_mcp_server_manager, |
| 2583 | + ) |
| 2584 | + |
| 2585 | + if keyless_source: |
| 2586 | + return allowed_tools |
| 2587 | + |
| 2588 | + try: |
| 2589 | + object_permissions = await MCPRequestHandler._get_user_object_permission(user_api_key_auth) |
| 2590 | + except Exception as e: # noqa: BLE001 # an unresolved human entitlement must deny, not widen |
| 2591 | + verbose_logger.warning(f"MCP user tool ceiling unresolvable, denying tools on {server_id!r}: {str(e)}") |
| 2592 | + return [] |
| 2593 | + |
| 2594 | + if object_permissions is None or not object_permissions.mcp_tool_permissions: |
| 2595 | + return allowed_tools |
| 2596 | + |
| 2597 | + user_tools = global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).get( |
| 2598 | + server_id |
| 2599 | + ) |
| 2600 | + if user_tools is None: |
| 2601 | + return allowed_tools |
| 2602 | + if allowed_tools is None: |
| 2603 | + return list(user_tools) |
| 2604 | + return list(set(allowed_tools) & set(user_tools)) |
| 2605 | + |
2379 | 2606 | # Sentinel stored in cache when an agent has no object_permission, so we |
2380 | 2607 | # don't re-query the DB on every MCP request for that agent. |
2381 | 2608 | _AGENT_NO_PERMISSION_SENTINEL = "__agent_no_mcp_permission__" |
|
0 commit comments