fix(mcp): resolve team-granted toolsets for non-admin keys and dashboard sessions (#43908)

* fix(mcp): expand team and dashboard grants when listing and serving toolsets

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(mcp): honor team toolset grants on the responses gateway path and expose a public team permission lookup

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* chore(mcp): drop the toolset route docstring tweak so the OpenAPI snapshot stays unchanged

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(mcp): scope inherited toolset grants by the key's own MCP ceiling

A key that declares any MCP grant of its own keeps only its own toolsets, one
that declares none inherits its team's, and require_key_mcp_access_defined
stops a virtual key inheriting while dashboard sessions and admitted users
still do. Adds the direct, no-grant, admin and key-ceiling integration cases
and makes the LLM gateway toolset case discriminate a scoped toolset from the
aggregate grant

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(mcp): resolve toolset grants per admitted source and enforce the live team roster

Dashboard sessions and gateway-admitted users now expand into the admitted
subject's per-team sources when resolving toolset grants, so a team-granted
toolset is not capped by the user's own MCP row and is reachable on the
namespaced route. A cached team id no longer grants a toolset unless the live
roster still lists the user, a team lookup fault only drops that team's
inherited grant, and /team/member_add evicts the cached team object so the new
member is authoritative immediately

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(mcp): read a key's named object permission before letting it inherit team toolsets

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(mcp): inject the toolset grant resolver into scope helpers so tests stop patching MCPRequestHandler

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(mcp): drop the duplicate admitted_subject_sources wrapper after merging main and follow its renamed resolvers

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(mcp): honour fresh policy and the session resource scope on pinned toolsets

A pinned toolset scope now reads the toolset through the writer when the admitted session
requires fresh policy, so a tool revoked from the toolset is gone on the next request. A
gateway bearer scoped to one server can only open a toolset that names that server

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(mcp): keep operator-open servers out of toolset gateway urls

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(mcp): audit cells for team-granted toolsets across every surface

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(mcp): bound toolset edit convergence by both cache layers

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(mcp): drop toolset integration cells that test behavior this PR does not change

Repeat-read byte identity, 20 concurrent calls, a stopped peer and a killed worker are covered
generically by test_mcp_resilience.py and test_mcp_user_env_vars.py. The toolset edit cell asserts
pre-existing cache propagation and flaked locally with connection resets while polling

* chore(mcp): drop mutable-ok suppressions that main's LIT013 now flags as unused

* chore(mcp): keep the require_key_mcp_access_defined read from adding an unknown-argument type error

---------

Co-authored-by: ryan <ryan@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-01 14:15:06 -07:00 • committed by GitHub
parent c66c8288c3
commit d7e6154863
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 1704 additions and 134 deletions

View file

@ -1947,10 +1947,41 @@ class MCPRequestHandler:
scope, the tool union and billing attribution are all just different reads of this one
answer — computing it separately per consumer is how they drift (a throttle map scoped by
roster instead of by grant charged unrelated teams' buckets)."""
return [
grants: Final = [
(source, set(await MCPRequestHandler.get_allowed_mcp_servers(source, keyless_source=True)))
for source in await MCPRequestHandler.admitted_subject_sources(auth, allowed_team_ids=allowed_team_ids)
]
scope: Final = await MCPRequestHandler._toolset_scope(auth)
if scope is None:
return grants
return [(source, granted & frozenset(scope)) for source, granted in grants]
@staticmethod
async def _toolset_scope(auth: UserAPIKeyAuth) -> dict[str, list[str]] | None:
"""The ``server_id -> tools`` a namespaced toolset route pinned this subject to via
``mcp_toolset_id``, or None on the aggregate scope. Every source's servers and tools are
intersected with it, so the route narrows a team grant exactly as it narrows the user's own."""
if auth.mcp_toolset_id is None:
return None
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
return await global_mcp_server_manager.resolve_toolset_tool_permissions(
toolset_ids=[auth.mcp_toolset_id], requires_fresh_policy=auth.requires_fresh_policy
)
@staticmethod
async def _narrow_tools_to_toolset(
tools: list[str] | None,
server_id: str,
auth: UserAPIKeyAuth,
) -> list[str] | None:
scope: Final = await MCPRequestHandler._toolset_scope(auth)
if scope is None:
return tools
scoped: Final = frozenset(scope.get(server_id, ()))
return sorted(scoped if tools is None else scoped & frozenset(tools))
@staticmethod
async def resolve_admitted_subject_servers(
@ -2048,9 +2079,9 @@ class MCPRequestHandler:
continue
tools = await MCPRequestHandler.get_allowed_tools_for_server(server_id, source, keyless_source=True)
if tools is None:
return None
return await MCPRequestHandler._narrow_tools_to_toolset(None, server_id, auth)
allowed.update(tools)
return sorted(allowed)
return await MCPRequestHandler._narrow_tools_to_toolset(sorted(allowed), server_id, auth)
@staticmethod
def _get_key_object_permission(
@ -2067,6 +2098,16 @@ class MCPRequestHandler:
return user_api_key_auth.object_permission
@staticmethod
async def team_object_permission(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None:
return await MCPRequestHandler._get_team_object_permission(user_api_key_auth)
@staticmethod
async def key_object_permission_hydrated(
user_api_key_auth: UserAPIKeyAuth,
) -> LiteLLM_ObjectPermissionTable | None:
return await MCPRequestHandler._key_object_permission_hydrated(user_api_key_auth)
@staticmethod
async def _get_team_object_permission(
user_api_key_auth: UserAPIKeyAuth | None = None,

View file

@ -915,11 +915,9 @@ if MCP_AVAILABLE:
) -> UserAPIKeyAuth:
"""The one credential this tools request acts as.
A toolset name narrows the caller's own credential to that toolset; otherwise a dashboard
session is swapped for its admitted subject. The two are mutually exclusive by construction,
which is why they share an owner: the admitted subject resolves per grant source and a team
source deliberately carries none of the caller's ``object_permission``, so a toolset
narrowing layered on top would evaporate on every team-granted server."""
A toolset name pins the acting principal to that toolset through ``_apply_toolset_scope``,
which itself swaps a dashboard session for its admitted subject; otherwise the swap happens
here so both shapes resolve as the same identity."""
if not toolset_name:
return await acting_user_auth(user_api_key_dict)

View file

@ -66,7 +66,13 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import (
get_route_relative_request_path,
well_known_root_suffix,
)
from litellm.proxy._experimental.mcp_server.ui_session_utils import is_ui_session_credential
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
ActingUser,
GrantedToolsetIds,
acting_user_auth,
granted_toolset_ids,
is_ui_session_credential,
)
from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_DESCRIPTION,
LITELLM_MCP_SERVER_NAME,
@ -1515,16 +1521,21 @@ if MCP_AVAILABLE:
async def _apply_toolset_scope(
user_api_key_auth: UserAPIKeyAuth,
toolset_id: str,
acting_user: ActingUser = acting_user_auth,
granted: GrantedToolsetIds = granted_toolset_ids,
) -> UserAPIKeyAuth:
"""
Restrict a key's MCP permissions to a single toolset.
Pin a principal's MCP permissions to a single toolset for /toolset/{name}/mcp.
When a request arrives via /toolset/{name}/mcp we override the key's
object_permission so that only the toolset's tools are visible.
A virtual key (and an admin session) has its object_permission rewritten to
the toolset's servers and tools. A keyless subject resolves per grant source,
so a non-admin dashboard session first becomes its admitted user and the
toolset rides along as ``mcp_toolset_id``, which every source's grant is
intersected with; a team-granted toolset is served without the user's own
row capping it.
Raises HTTPException(403) if the key has an explicit toolset grant list
that does not include toolset_id (i.e. mcp_toolsets is set but empty,
or set to a list that omits this toolset). Admin keys always pass.
Raises HTTPException(403) unless the principal holds toolset_id through one
of its grant sources. Admins always pass.
"""
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
@ -1540,26 +1551,31 @@ if MCP_AVAILABLE:
detail="API key is scoped to no MCP servers; toolset access is denied.",
)
# Access control: non-admin keys must have this toolset in their grant list.
# Use _user_has_admin_view so that PROXY_ADMIN_VIEW_ONLY is also treated as admin.
is_admin: Final = _user_has_admin_view(user_api_key_auth)
if not is_admin:
op: Final = user_api_key_auth.object_permission
granted: Final = getattr(op, "mcp_toolsets", None) if op else None
# granted=None → key has no explicit toolset grants → deny (same semantics as
# fetch_mcp_toolsets which returns [] for non-admin keys with no grants configured).
# granted=[] or list without toolset_id → also deny.
if granted is None or toolset_id not in granted:
acting: Final = await acting_user(user_api_key_auth)
is_admin: Final = _user_has_admin_view(acting)
if not is_admin and toolset_id not in await granted(acting):
raise HTTPException(
status_code=403,
detail=f"API key does not have access to toolset '{toolset_id}'.",
)
if _is_mcp_admitted_user_subject(acting):
resource_server_id: Final = acting.mcp_session_resource_server_id
if resource_server_id is not None and resource_server_id not in (
await operations.global_mcp_server_manager.resolve_toolset_tool_permissions(
toolset_ids=[toolset_id], requires_fresh_policy=acting.requires_fresh_policy
)
):
raise HTTPException(
status_code=403,
detail=f"API key does not have access to toolset '{toolset_id}'.",
)
return acting.model_copy(update={"mcp_toolset_id": toolset_id})
tool_permissions = await operations.global_mcp_server_manager.resolve_toolset_tool_permissions(
toolset_ids=[toolset_id]
)
server_ids: Final = list(tool_permissions.keys())
existing_op: Final = user_api_key_auth.object_permission
existing_op: Final = acting.object_permission
if existing_op is not None:
updated_op = existing_op.model_copy(
update={
@ -1576,7 +1592,12 @@ if MCP_AVAILABLE:
mcp_servers=server_ids,
mcp_tool_permissions=tool_permissions,
)
return user_api_key_auth.model_copy(update={"object_permission": updated_op, "mcp_toolset_id": toolset_id})
return acting.model_copy(update={"object_permission": updated_op, "mcp_toolset_id": toolset_id})
async def _toolset_server_ids(toolset_id: str) -> set[str]:
return set(
await operations.global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=[toolset_id])
)
async def _raise_preemptive_401_for_unauthenticated_servers(
scope: Scope,
@ -2007,8 +2028,7 @@ if MCP_AVAILABLE:
toolset_allowed_server_ids: set[str] | None = None
if active_toolset_id and user_api_key_auth is not None:
user_api_key_auth = await _apply_toolset_scope(user_api_key_auth, active_toolset_id)
op: Final = user_api_key_auth.object_permission
toolset_allowed_server_ids = set(op.mcp_servers or []) if op else set()
toolset_allowed_server_ids = await _toolset_server_ids(active_toolset_id)
# https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response
# Must run after toolset scoping so the challenge set is derived
@ -2355,8 +2375,7 @@ if MCP_AVAILABLE:
toolset_allowed_server_ids: set[str] | None = None
if active_toolset_id and user_api_key_auth is not None:
user_api_key_auth = await _apply_toolset_scope(user_api_key_auth, active_toolset_id)
op: Final = user_api_key_auth.object_permission
toolset_allowed_server_ids = set(op.mcp_servers or []) if op else set()
toolset_allowed_server_ids = await _toolset_server_ids(active_toolset_id)
# https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response
# Must run after toolset scoping so the challenge set is derived

View file

@ -2,14 +2,23 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable
from typing import Final
import asyncio
from collections.abc import Awaitable, Callable, Sequence
from itertools import chain
from typing import Final, TypeAlias
from fastapi import HTTPException
from litellm._logging import verbose_logger
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
EffectiveAuthContexts: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[Sequence[UserAPIKeyAuth]]]
TeamObjectPermission: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_ObjectPermissionTable | None]]
OwnObjectPermission: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_ObjectPermissionTable | None]]
AdmittedContext: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[UserAPIKeyAuth | None]]
ActingUser: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[UserAPIKeyAuth]]
GrantedToolsetIds: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[frozenset[str]]]
def clone_user_api_key_auth_with_team(
@ -109,10 +118,10 @@ async def acting_user_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth:
both surfaces. An admin session keeps its operator view and any caller-passed credential is
returned unchanged, never widened.
Do not combine this with a narrowing that rewrites a single credential's ``object_permission``
(toolset scope): the admitted subject resolves per grant source and a team source deliberately
carries none of the caller's own grants, so the narrowing would silently evaporate on every
team-granted server. A request carrying such a scope keeps the caller's own credential."""
A toolset narrowing is never applied to the admitted subject by rewriting its ``object_permission``:
it resolves per grant source and a team source deliberately carries none of the caller's own grants,
so the rewrite would evaporate on every team-granted server. The route pins ``mcp_toolset_id``
instead, which every source's grant is intersected with."""
if not is_ui_session_credential(user_api_key_auth):
return user_api_key_auth
@ -153,3 +162,135 @@ async def can_access_mcp_server(
if server_id in await allowed_servers(context):
return True
return False
def _restricts_mcp(permission: LiteLLM_ObjectPermissionTable | None) -> bool:
return permission is not None and bool(
permission.mcp_servers
or permission.mcp_toolsets
or permission.mcp_tool_permissions
or permission.mcp_access_groups
)
def is_keyless_mcp_subject(user_api_key_auth: UserAPIKeyAuth) -> bool:
"""A principal with no virtual key to declare MCP access on: the dashboard's own session token or a
gateway-admitted user. Its grants are resolved per source, never through a key row."""
return is_ui_session_credential(user_api_key_auth) or user_api_key_auth.mcp_admitted_user_subject is True
async def toolset_grant_contexts(
user_api_key_auth: UserAPIKeyAuth,
admitted_context: AdmittedContext = admitted_user_context,
admitted_sources: EffectiveAuthContexts | None = None,
) -> Sequence[UserAPIKeyAuth]:
"""The grant sources a toolset is looked up through. A virtual key is its own single source. A keyless
subject fans out exactly as the aggregate /mcp resolution does: its own user row plus every team whose
live roster still lists it, so a membership that only survives in the user's cached team list grants
nothing."""
if not is_keyless_mcp_subject(user_api_key_auth):
return (user_api_key_auth,)
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
load_sources: Final = admitted_sources or MCPRequestHandler.admitted_subject_sources
acting: Final = await admitted_context(user_api_key_auth)
return tuple(await load_sources(acting if acting is not None else user_api_key_auth))
async def _own_toolset_ids(
context: UserAPIKeyAuth,
load_own_permission: OwnObjectPermission,
) -> Sequence[str] | None:
"""The source's own toolsets, or None when it declares no MCP grant of its own. A source that names an
``object_permission_id`` is a known restriction even when the row is unhydrated, unreadable or gone,
so it is loaded rather than read as unrestricted, and grants nothing when it cannot be read."""
if context.object_permission is None and not context.object_permission_id:
return None
try:
own: Final = await load_own_permission(context)
except Exception as exc: # noqa: BLE001 # a named but unreadable own grant must deny, not widen to the team
verbose_logger.warning(
"MCP toolset grants: object permission %s unreadable, granting nothing through it: %s",
context.object_permission_id,
exc,
)
return ()
if own is None:
return ()
if not _restricts_mcp(own):
return None
return own.mcp_toolsets or ()
async def _inherited_toolset_ids(
context: UserAPIKeyAuth,
load_team_permission: TeamObjectPermission,
) -> Sequence[str]:
try:
team: Final = await load_team_permission(context)
except Exception as exc: # noqa: BLE001 # an unreadable team grants nothing through this source and must not fail the caller's other sources
verbose_logger.warning(
"MCP toolset grants: team %s unreadable, inheriting nothing from it: %s",
context.team_id,
exc,
)
return ()
return () if team is None else (team.mcp_toolsets or ())
async def _context_toolset_ids(
context: UserAPIKeyAuth,
inherits_team: bool,
load_team_permission: TeamObjectPermission,
load_own_permission: OwnObjectPermission,
) -> Sequence[str]:
own: Final = await _own_toolset_ids(context, load_own_permission)
if own is not None:
return own
if not inherits_team or not context.team_id:
return ()
return await _inherited_toolset_ids(context, load_team_permission)
async def granted_toolset_ids(
user_api_key_auth: UserAPIKeyAuth,
effective_contexts: EffectiveAuthContexts = toolset_grant_contexts,
team_object_permission: TeamObjectPermission | None = None,
require_key_access: bool | None = None,
own_object_permission: OwnObjectPermission | None = None,
) -> frozenset[str]:
"""Toolset ids the principal holds, resolved per grant source with the key/team rule the aggregate
/mcp listing applies: a source that declares any MCP grant of its own is scoped to its own toolsets and
never reads its team, one that declares none inherits its team's, except a virtual key under
``require_key_mcp_access_defined``, which inherits nothing. A keyless subject's team sources always
inherit. A team that cannot be read contributes nothing while every other source still counts, and an
own grant that is named but cannot be read grants nothing. No grant anywhere yields the empty set."""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
from litellm.proxy.proxy_server import general_settings
require: Final = (
require_key_access
if require_key_access is not None
else bool(
general_settings.get( # pyright: ignore[reportUnknownArgumentType] # general_settings is an untyped dict; truthiness must match the /mcp path's read of this flag
"require_key_mcp_access_defined", False
)
)
)
inherits_team: Final = is_keyless_mcp_subject(user_api_key_auth) or not require
load_team_permission: Final = team_object_permission or MCPRequestHandler.team_object_permission
load_own_permission: Final = own_object_permission or MCPRequestHandler.key_object_permission_hydrated
contexts: Final = await effective_contexts(user_api_key_auth)
per_context: Final = await asyncio.gather(
*(
_context_toolset_ids(context, inherits_team, load_team_permission, load_own_permission)
for context in contexts
)
)
return frozenset(chain.from_iterable(per_context))

View file

@ -184,6 +184,7 @@ if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
admitted_user_context,
build_effective_auth_contexts,
granted_toolset_ids,
is_ui_session_credential,
)
from litellm.proxy._types import (
@ -3332,18 +3333,15 @@ if MCP_AVAILABLE:
):
"""Return toolsets the calling key is allowed to access."""
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
is_admin: Final = _user_has_admin_view(user_api_key_dict)
op: Final = user_api_key_dict.object_permission
# mcp_toolsets=None or [] both mean "not restricted by toolsets".
# For admins: either value → no restriction → return all.
# For non-admins: either value → no toolsets explicitly granted → return nothing.
# (An admin whose DB row has mcp_toolsets=[] should still see all toolsets.)
raw_toolsets: Final = getattr(op, "mcp_toolsets", None) if op else None
if not raw_toolsets:
if is_admin:
if _user_has_admin_view(user_api_key_dict):
op: Final = user_api_key_dict.object_permission
if op is None or not op.mcp_toolsets:
return await list_mcp_toolsets(prisma_client)
return await list_mcp_toolsets(prisma_client, toolset_ids=op.mcp_toolsets)
granted: Final = await granted_toolset_ids(user_api_key_dict)
if not granted:
return []
return await list_mcp_toolsets(prisma_client, toolset_ids=raw_toolsets)
return await list_mcp_toolsets(prisma_client, toolset_ids=sorted(granted))
@router.get(
"/toolset/{toolset_id}",
@ -3355,15 +3353,13 @@ if MCP_AVAILABLE:
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
# Non-admin keys may only fetch toolsets they've been explicitly granted.
if not _user_has_admin_view(user_api_key_dict):
op: Final = user_api_key_dict.object_permission
granted: Final = getattr(op, "mcp_toolsets", None) if op else None
if granted is None or toolset_id not in granted:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={"error": "API key does not have access to this toolset."},
)
if not _user_has_admin_view(user_api_key_dict) and toolset_id not in await granted_toolset_ids(
user_api_key_dict
):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={"error": "API key does not have access to this toolset."},
)
toolset: Final = await get_mcp_toolset(prisma_client, toolset_id)
if toolset is None:
raise HTTPException(

View file

@ -3466,6 +3466,12 @@ async def team_member_add(
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
await delete_cache_team_object(
team_id=data.team_id,
team_alias=complete_team_data.team_alias,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await evict_and_broadcast(
cache_keys=tuple(sorted(user.user_id for user in updated_users)),
user_api_key_cache=user_api_key_cache,

View file

@ -40,6 +40,7 @@ if TYPE_CHECKING:
from mcp.types import CallToolResult
from mcp.types import Tool as MCPTool
from litellm.proxy._experimental.mcp_server.ui_session_utils import GrantedToolsetIds
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
else:
@ -223,7 +224,9 @@ class LiteLLM_Proxy_MCP_Handler:
mcp_servers=all_server_ids,
mcp_tool_permissions=tool_permissions,
)
return user_api_key_auth.model_copy(update={"object_permission": updated_op})
return user_api_key_auth.model_copy(
update={"object_permission": updated_op, "mcp_explicit_grants_only": True}
)
except Exception as _e:
verbose_logger.debug("Could not apply toolset permissions: %s", _e)
return user_api_key_auth
@ -237,6 +240,7 @@ class LiteLLM_Proxy_MCP_Handler:
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
request_tags: list[str] | None = None,
raw_headers: dict[str, str] | None = None,
granted_toolsets: "GrantedToolsetIds | None" = None,
) -> tuple[list[MCPTool], list[str]]:
"""
Get available tools from the MCP server manager.
@ -279,23 +283,19 @@ class LiteLLM_Proxy_MCP_Handler:
if prisma_client is not None:
toolset = await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name)
if toolset is not None:
# Access control: only allow if the key explicitly grants this toolset.
if user_api_key_auth is not None:
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
granted_toolset_ids,
)
from litellm.proxy.management_endpoints.common_utils import (
_user_has_admin_view,
)
is_admin = _user_has_admin_view(user_api_key_auth)
if not is_admin:
op = user_api_key_auth.object_permission
granted = getattr(op, "mcp_toolsets", None) if op else None
# None means no grants configured → deny (consistent with
# fetch_mcp_toolsets which returns [] for unconfigured keys)
if granted is None or toolset.toolset_id not in granted:
verbose_logger.debug(
"Key does not have access to toolset '%s', skipping.", name
)
continue
if not _user_has_admin_view(user_api_key_auth) and toolset.toolset_id not in (
await (granted_toolsets or granted_toolset_ids)(user_api_key_auth)
):
verbose_logger.debug("Key does not have access to toolset '%s', skipping.", name)
continue
resolved_toolset_ids.append(toolset.toolset_id)
# Don't add to resolved_mcp_servers — toolset scope
# restricts via object_permission, not server name filter.

View file

@ -51,12 +51,12 @@ def delete_toolset(gateway: Gateway, identity: str) -> None:
assert response.status_code in (200, 202, 204), response.text
def create_toolset(scenario: Scenario, tools: tuple[tuple[str, str], ...]) -> str:
def create_toolset(scenario: Scenario, tools: tuple[tuple[str, str], ...], toolset_name: str | None = None) -> str:
response: Final = scenario.gateway.request(
"POST",
"/v1/mcp/toolset",
{
"toolset_name": f"integration-{uuid.uuid4().hex[:10]}",
"toolset_name": toolset_name or f"integration-{uuid.uuid4().hex[:10]}",
"tools": [{"server_id": server_id, "tool_name": tool} for server_id, tool in tools],
},
)

View file

@ -9,6 +9,7 @@ import httpx
import pytest
from integration._support.client import Gateway, Scenario
from integration._support.mcp import McpPeer, mcp_peer, register_mcp, tool_calls
from integration._support.mcp_grants import create_toolset
from integration._support.wire import Reply, Request, Wire, wire_server
Surface = Literal["chat", "responses", "messages", "messages_bridge"]
@ -194,9 +195,7 @@ class Rig:
)
def upstream_tools(self) -> tuple[tuple[str, ...], ...]:
return tuple(
_tool_names(json.loads(request.body)) for request in self.wire.drain() if request.method == "POST"
)
return tuple(_tool_names(json.loads(request.body)) for request in self.wire.drain() if request.method == "POST")
def final_text(self, body: Mapping[str, object]) -> str:
if self.surface == "chat":
@ -314,6 +313,25 @@ def test_allowed_tools_narrows_the_tool_list_handed_to_the_model(gateway: Gatewa
assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
def test_toolset_gateway_url_serves_a_team_granted_toolset_to_a_key_without_its_own_grant(
gateway: Gateway, surface: Surface
) -> None:
with _rig(gateway, surface) as rig:
register_mcp(rig.scenario, rig.peer, "open" + uuid.uuid4().hex[:8], allow_all_keys=True)
toolset_name: Final = "ts" + uuid.uuid4().hex[:8]
toolset_id: Final = create_toolset(rig.scenario, ((rig.server_id, "add"),), toolset_name=toolset_name)
sibling_id: Final = create_toolset(rig.scenario, ((rig.server_id, "multiply"),))
team_id: Final = rig.scenario.team(object_permission={"mcp_toolsets": [toolset_id, sibling_id]})
key: Final = rig.scenario.key(team_id=team_id)
response: Final = rig.send(key, [{**AUTO, "server_url": f"litellm_proxy/mcp/{toolset_name}"}])
assert response.status_code == 200, response.text
requests: Final = rig.upstream_tools()
assert requests, "model was never called"
assert all(names == (rig.tool,) for names in requests), requests
assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
def test_server_scoped_gateway_url_exposes_only_that_servers_tools(gateway: Gateway, surface: Surface) -> None:
with _rig(gateway, surface) as rig, mcp_peer() as other_peer:
@ -347,3 +365,41 @@ def test_streaming_chat_executes_the_tool_once_and_streams_the_follow_up(gateway
assert text == ANSWER, response.text
assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
assert len(rig.upstream_tools()) == 2
def test_streaming_chat_through_a_toolset_gateway_url_serves_a_team_key_without_its_own_grant(
gateway: Gateway,
) -> None:
with _rig(gateway, "chat") as rig:
register_mcp(rig.scenario, rig.peer, "open" + uuid.uuid4().hex[:8], allow_all_keys=True)
toolset_name: Final = "ts" + uuid.uuid4().hex[:8]
toolset_id: Final = create_toolset(rig.scenario, ((rig.server_id, "add"),), toolset_name=toolset_name)
key: Final = rig.scenario.key(team_id=rig.scenario.team(object_permission={"mcp_toolsets": [toolset_id]}))
response: Final = rig.send(key, [{**AUTO, "server_url": f"litellm_proxy/mcp/{toolset_name}"}], stream=True)
assert response.status_code == 200, response.text
chunks: Final = tuple(
json.loads(line.removeprefix("data: "))
for line in response.text.splitlines()
if line.startswith("data: ") and line != "data: [DONE]"
)
text: Final = "".join(
str(chunk["choices"][0]["delta"].get("content") or "") for chunk in chunks if chunk.get("choices")
)
assert text == ANSWER, response.text
assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
requests: Final = rig.upstream_tools()
assert len(requests) == 2 and all(names == (rig.tool,) for names in requests), requests
@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
def test_toolset_gateway_url_gives_a_key_of_an_ungranted_team_no_tools_and_never_reaches_the_peer(
gateway: Gateway, surface: Surface
) -> None:
with _rig(gateway, surface) as rig:
toolset_name: Final = "ts" + uuid.uuid4().hex[:8]
create_toolset(rig.scenario, ((rig.server_id, "add"),), toolset_name=toolset_name)
key: Final = rig.scenario.key(team_id=rig.scenario.team())
response: Final = rig.send(key, [{**AUTO, "server_url": f"litellm_proxy/mcp/{toolset_name}"}])
assert _peer_add_calls(rig.peer) == (), "denied caller reached the peer"
assert all(rig.tool not in names for names in rig.upstream_tools()), rig.upstream_tools()
assert response.status_code in (200, 400, 401, 403), response.text

View file

@ -1,25 +1,31 @@
import base64
import hashlib
import secrets
import time
import uuid
from dataclasses import dataclass
from typing import Final
from urllib.parse import parse_qs, urlsplit
import httpx
import jwt
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.mcp import (
ENTRY_POINTS,
INITIALIZE,
EntryPoint,
McpCaller,
McpPeer,
Outcome,
_outcome_from_rpc,
call_tool,
mcp_peer,
register_mcp,
tool_calls,
)
from integration._support.mcp_grants import create_toolset
from integration._support.oauth_server import AuthorizationServer, oauth_server
ADD: Final = {"a": 2, "b": 3}
@ -383,3 +389,117 @@ def test_dcr_bridge_relays_client_registration_and_advertises_gateway_endpoints(
assert issuer.json()["authorization_endpoint"] == f"{_base(gateway)}/{alias}/authorize"
assert issuer.json()["token_endpoint"] == f"{_base(gateway)}/{alias}/token"
assert "S256" in issuer.json()["code_challenge_methods_supported"]
def _ui_session_cookie(gateway: Gateway, user_id: str) -> dict[str, str]:
claims: Final = {"user_id": user_id, "login_method": "username_password", "exp": int(time.time()) + 600}
return {"token": jwt.encode(claims, gateway.key, algorithm="HS256")}
def _gateway_session_bearer(gateway: Gateway, user_id: str, resource: str | None = None) -> str:
registered: Final = gateway.client.post(
"/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"}
)
assert registered.status_code in (200, 201), registered.text
client_id: Final = registered.json()["client_id"]
pkce: Final = _Pkce(secrets.token_urlsafe(48))
cookies: Final = _ui_session_cookie(gateway, user_id)
started: Final = gateway.client.get(
"/authorize",
params={
"client_id": client_id,
"redirect_uri": CLIENT_REDIRECT,
"response_type": "code",
"state": "lit6029",
"code_challenge": pkce.challenge,
"code_challenge_method": "S256",
**({} if resource is None else {"resource": resource}),
},
cookies=cookies,
)
assert started.status_code == 303, started.text
handle: Final = parse_qs(urlsplit(started.headers["location"]).query)["connect_flow"][0]
completed: Final = gateway.client.post(
"/authorize/complete", data={"flow": handle}, cookies={**cookies, **dict(started.cookies)}
)
assert completed.status_code == 303, completed.text
callback: Final = parse_qs(urlsplit(completed.headers["location"]).query)
assert "code" in callback, completed.headers["location"]
issued: Final = gateway.client.post(
"/token",
data={
"grant_type": "authorization_code",
"code": callback["code"][0],
"redirect_uri": CLIENT_REDIRECT,
"client_id": client_id,
"code_verifier": pkce.verifier,
},
)
assert issued.status_code == 200, issued.text
return _issued_token(issued.json())
def _toolset_rpc(gateway: Gateway, bearer: str, name: str, method: str, params: dict[str, object]) -> Outcome:
def post(rpc_method: str, rpc_params: dict[str, object]) -> httpx.Response:
return gateway.client.post(
f"/toolset/{name}/mcp",
json={"jsonrpc": "2.0", "id": 1, "method": rpc_method, "params": rpc_params},
headers={"Authorization": f"Bearer {bearer}", "Accept": "application/json, text/event-stream"},
)
initialized: Final = _outcome_from_rpc(post("initialize", dict(INITIALIZE)))
if not initialized.ok:
return initialized
return _outcome_from_rpc(post(method, params))
def test_gateway_session_bearer_of_a_team_member_is_served_the_team_toolset_on_its_route(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029sess" + uuid.uuid4().hex[:6]
server_id: Final = register_mcp(scenario, peer, alias)
granted_name: Final = "lit6029g" + uuid.uuid4().hex[:8]
withheld_name: Final = "lit6029w" + uuid.uuid4().hex[:8]
granted_id: Final = create_toolset(scenario, ((server_id, "add"),), toolset_name=granted_name)
create_toolset(scenario, ((server_id, "multiply"),), toolset_name=withheld_name)
member: Final = scenario.member(scenario.team(object_permission={"mcp_toolsets": [granted_id]}))
bearer: Final = _gateway_session_bearer(gateway, member)
assert bearer.startswith("llm_session_"), bearer[:16]
listed: Final = _toolset_rpc(gateway, bearer, granted_name, "tools/list", {})
assert listed.tools == (f"{alias}-add",), listed.raw
peer.drain()
called: Final = _toolset_rpc(
gateway, bearer, granted_name, "tools/call", {"name": f"{alias}-add", "arguments": {"a": 4, "b": 5}}
)
assert called.ok and called.text == "9", called.raw
assert len(tool_calls(peer.drain())) == 1
denied: Final = _toolset_rpc(gateway, bearer, withheld_name, "tools/list", {})
assert denied.status == 403, denied.raw
assert tool_calls(peer.drain()) == ()
def test_resource_scoped_session_bearer_opens_a_team_toolset_inside_its_server_and_none_outside(
gateway: Gateway,
) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
inside: Final = "lit6029in" + uuid.uuid4().hex[:6]
outside: Final = "lit6029out" + uuid.uuid4().hex[:6]
inside_server: Final = register_mcp(scenario, peer, inside)
outside_server: Final = register_mcp(scenario, peer, outside)
inside_name: Final = "lit6029i" + uuid.uuid4().hex[:8]
outside_name: Final = "lit6029o" + uuid.uuid4().hex[:8]
inside_id: Final = create_toolset(scenario, ((inside_server, "add"),), toolset_name=inside_name)
outside_id: Final = create_toolset(scenario, ((outside_server, "add"),), toolset_name=outside_name)
member: Final = scenario.member(scenario.team(object_permission={"mcp_toolsets": [inside_id, outside_id]}))
bearer: Final = _gateway_session_bearer(gateway, member, resource=f"{_base(gateway)}/{inside}/mcp")
assert bearer.startswith("llm_session_"), bearer[:16]
listed: Final = _toolset_rpc(gateway, bearer, inside_name, "tools/list", {})
assert listed.tools == (f"{inside}-add",), listed.raw
peer.drain()
called: Final = _toolset_rpc(
gateway, bearer, inside_name, "tools/call", {"name": f"{inside}-add", "arguments": {"a": 4, "b": 5}}
)
assert called.ok and called.text == "9", called.raw
assert len(tool_calls(peer.drain())) == 1
refused: Final = _toolset_rpc(gateway, bearer, outside_name, "tools/list", {})
assert refused.status == 403, refused.raw
assert tool_calls(peer.drain()) == ()

View file

@ -0,0 +1,509 @@
import secrets
import uuid
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, Scenario, object_value
from integration._support.mcp import (
INITIALIZE,
Outcome,
_outcome_from_rest,
_outcome_from_rpc,
mcp_peer,
register_mcp,
tool_calls,
)
from integration._support.mcp_grants import create_toolset
from integration._support.process import owned_proxy
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import LITELLM_SESSION_TOKEN_PREFIX, ExperimentalUIJWTToken
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_bearer_token
ADD: Final = {"a": 4, "b": 5}
def _dashboard_ui_session_token(user_id: str) -> str:
user: Final = LiteLLM_UserTable(user_id=user_id, user_role="internal_user", models=[])
return ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(user)
def _toolset(scenario: Scenario, server_id: str, tool: str) -> tuple[str, str]:
name: Final = "lit6029_" + uuid.uuid4().hex[:10]
return create_toolset(scenario, ((server_id, tool),), toolset_name=name), name
def _toolset_rpc(
gateway: Gateway, headers: dict[str, str], name: str, method: str, params: dict[str, object]
) -> Outcome:
def post(rpc_method: str, rpc_params: dict[str, object]) -> httpx.Response:
return gateway.client.post(
f"/toolset/{name}/mcp",
json={"jsonrpc": "2.0", "id": 1, "method": rpc_method, "params": rpc_params},
headers={**headers, "Accept": "application/json, text/event-stream"},
)
initialized: Final = _outcome_from_rpc(post("initialize", dict(INITIALIZE)))
if not initialized.ok:
return initialized
return _outcome_from_rpc(post(method, params))
def _listed_toolset_ids(gateway: Gateway, headers: dict[str, str]) -> tuple[str, ...]:
response: Final = gateway.client.get("/v1/mcp/toolset", headers=headers)
assert response.status_code == 200, response.text
return tuple(toolset["toolset_id"] for toolset in response.json())
def _assert_team_grants_only(gateway: Gateway, team_id: str, key: str, toolset_id: str) -> None:
team: Final = object_value(gateway.get("/team/info", {"team_id": team_id})["team_info"])
assert object_value(team["object_permission"])["mcp_toolsets"] == [toolset_id], team
key_info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
assert key_info.get("object_permission") is None, f"key must carry no grant of its own: {key_info}"
def test_team_granted_toolset_is_listed_and_served_to_a_team_key(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
server_id: Final = register_mcp(scenario, peer, alias)
granted_id, granted_name = _toolset(scenario, server_id, "add")
withheld_id, withheld_name = _toolset(scenario, server_id, "multiply")
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
key: Final = scenario.key(team_id=team_id)
_assert_team_grants_only(gateway, team_id, key, granted_id)
headers: Final = {"Authorization": f"Bearer {key}"}
assert _listed_toolset_ids(gateway, headers) == (granted_id,)
detail: Final = gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers)
assert detail.status_code == 200, detail.text
assert detail.json()["toolset_name"] == granted_name, detail.text
withheld_detail: Final = gateway.client.get(f"/v1/mcp/toolset/{withheld_id}", headers=headers)
assert withheld_detail.status_code == 403, withheld_detail.text
listed: Final = _toolset_rpc(gateway, headers, granted_name, "tools/list", {})
assert listed.ok, listed.raw
assert listed.tools == (f"{alias}-add",), listed.raw
peer.drain()
called: Final = _toolset_rpc(
gateway, headers, granted_name, "tools/call", {"name": f"{alias}-add", "arguments": ADD}
)
assert called.ok and called.text == "9", called.raw
assert len(tool_calls(peer.drain())) == 1
denied: Final = _toolset_rpc(gateway, headers, withheld_name, "tools/list", {})
assert denied.status == 403, denied.raw
def test_dashboard_session_of_a_team_member_lists_the_team_granted_toolset(
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
server_id: Final = register_mcp(scenario, peer, alias)
granted_id, granted_name = _toolset(scenario, server_id, "add")
_toolset(scenario, server_id, "multiply")
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
user_id: Final = scenario.user(user_role="internal_user", teams=[team_id])
user: Final = object_value(gateway.get("/user/info", {"user_id": user_id})["user_info"])
assert user["teams"] == [team_id], user
headers: Final = {"Authorization": f"Bearer {_dashboard_ui_session_token(user_id)}"}
assert _listed_toolset_ids(gateway, headers) == (granted_id,)
detail: Final = gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers)
assert detail.status_code == 200, detail.text
assert detail.json()["toolset_name"] == granted_name, detail.text
def test_direct_grants_no_grants_and_admin_listing_are_unchanged_by_team_resolution(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
server_id: Final = register_mcp(scenario, peer, alias)
granted_id, granted_name = _toolset(scenario, server_id, "add")
withheld_id, withheld_name = _toolset(scenario, server_id, "multiply")
direct: Final = {"Authorization": f"Bearer {scenario.key(object_permission={'mcp_toolsets': [granted_id]})}"}
ungranted_team: Final = scenario.team()
no_grant: Final = {"Authorization": f"Bearer {scenario.key(team_id=ungranted_team)}"}
admin: Final = {"Authorization": f"Bearer {gateway.key}"}
assert _listed_toolset_ids(gateway, direct) == (granted_id,)
assert _toolset_rpc(gateway, direct, granted_name, "tools/list", {}).tools == (f"{alias}-add",)
assert _toolset_rpc(gateway, direct, withheld_name, "tools/list", {}).status == 403
assert gateway.client.get(f"/v1/mcp/toolset/{withheld_id}", headers=direct).status_code == 403
assert _listed_toolset_ids(gateway, no_grant) == ()
assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=no_grant).status_code == 403
assert _toolset_rpc(gateway, no_grant, granted_name, "tools/list", {}).status == 403
assert {granted_id, withheld_id} <= set(_listed_toolset_ids(gateway, admin))
assert gateway.client.get(f"/v1/mcp/toolset/{withheld_id}", headers=admin).status_code == 200
assert _toolset_rpc(gateway, admin, withheld_name, "tools/list", {}).tools == (f"{alias}-multiply",)
def test_a_key_with_its_own_toolset_grant_does_not_inherit_the_team_toolset(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
server_id: Final = register_mcp(scenario, peer, alias)
own_id, own_name = _toolset(scenario, server_id, "add")
team_only_id, team_only_name = _toolset(scenario, server_id, "multiply")
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [own_id, team_only_id]})
key: Final = scenario.key(team_id=team_id, object_permission={"mcp_toolsets": [own_id]})
headers: Final = {"Authorization": f"Bearer {key}"}
assert _listed_toolset_ids(gateway, headers) == (own_id,)
assert gateway.client.get(f"/v1/mcp/toolset/{team_only_id}", headers=headers).status_code == 403
assert _toolset_rpc(gateway, headers, team_only_name, "tools/list", {}).status == 403
assert _toolset_rpc(gateway, headers, own_name, "tools/list", {}).tools == (f"{alias}-add",)
def _team_member_with_own_grant(scenario: Scenario, team_id: str, own_server_id: str) -> str:
user_id: Final = scenario.user(user_role="internal_user", object_permission={"mcp_servers": [own_server_id]})
scenario.gateway.post("/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}})
return user_id
def test_dashboard_session_serves_the_team_toolset_despite_a_disjoint_grant_on_the_user_row(
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
) -> None:
"""The member's own row grants a different server outright. That grant must not cap the team's toolset
to nothing, and the team's sibling toolset must not leak onto the granted toolset's route."""
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
server_id: Final = register_mcp(scenario, peer, alias)
own_server_id: Final = register_mcp(scenario, peer, "lit6029_own_" + uuid.uuid4().hex[:8])
granted_id, granted_name = _toolset(scenario, server_id, "add")
sibling_id, sibling_name = _toolset(scenario, server_id, "multiply")
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id, sibling_id]})
user_id: Final = _team_member_with_own_grant(scenario, team_id, own_server_id)
headers: Final = {"Authorization": f"Bearer {_dashboard_ui_session_token(user_id)}"}
assert set(_listed_toolset_ids(gateway, headers)) == {granted_id, sibling_id}
listed: Final = _toolset_rpc(gateway, headers, granted_name, "tools/list", {})
assert listed.ok, listed.raw
assert listed.tools == (f"{alias}-add",), listed.raw
assert _toolset_rpc(gateway, headers, sibling_name, "tools/list", {}).tools == (f"{alias}-multiply",)
peer.drain()
called: Final = _toolset_rpc(
gateway, headers, granted_name, "tools/call", {"name": f"{alias}-add", "arguments": ADD}
)
assert called.ok and called.text == "9", called.raw
assert len(tool_calls(peer.drain())) == 1
stranger: Final = {
"Authorization": f"Bearer {_dashboard_ui_session_token(scenario.user(user_role='internal_user'))}"
}
assert _toolset_rpc(gateway, stranger, granted_name, "tools/list", {}).status == 403
def test_a_member_removed_from_the_team_loses_its_toolset_on_the_dashboard_session(
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
server_id: Final = register_mcp(scenario, peer, alias)
granted_id, granted_name = _toolset(scenario, server_id, "add")
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
user_id: Final = scenario.member(team_id)
headers: Final = {"Authorization": f"Bearer {_dashboard_ui_session_token(user_id)}"}
assert _listed_toolset_ids(gateway, headers) == (granted_id,)
assert _toolset_rpc(gateway, headers, granted_name, "tools/list", {}).tools == (f"{alias}-add",)
gateway.post("/team/member_delete", {"team_id": team_id, "user_id": user_id})
assert _listed_toolset_ids(gateway, headers) == ()
assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers).status_code == 403
assert _toolset_rpc(gateway, headers, granted_name, "tools/list", {}).status == 403
def _bearer(token: str) -> dict[str, str]:
return {"Authorization": f"Bearer {token}"}
def _expired_dashboard_token(user_id: str) -> str:
expired: Final = (datetime.now(timezone.utc) - timedelta(minutes=5)).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]
stale: Final = UserAPIKeyAuth(
token="ui-token",
key_name="ui-token",
key_alias="ui-token",
expires=expired + "+00:00",
user_id=user_id,
team_id="litellm-dashboard",
models=[],
user_role=LitellmUserRoles.INTERNAL_USER,
)
return encrypt_bearer_token(stale.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX)
def _rest_list(gateway: Gateway, headers: dict[str, str], params: object) -> httpx.Response:
return gateway.client.get("/mcp-rest/tools/list", headers=headers, params=params)
def _rest_call(gateway: Gateway, headers: dict[str, str], name: str, server_id: str) -> Outcome:
return _outcome_from_rest(
gateway.client.post(
"/mcp-rest/tools/call",
headers=headers,
json={"name": name, "arguments": dict(ADD), "server_id": server_id},
)
)
def _route_tools(gateway: Gateway, headers: dict[str, str], name: str) -> Outcome:
return _toolset_rpc(gateway, headers, name, "tools/list", {})
def _route_call(gateway: Gateway, headers: dict[str, str], name: str, tool: str) -> Outcome:
return _toolset_rpc(gateway, headers, name, "tools/call", {"name": tool, "arguments": dict(ADD)})
def _strict_config(directory: Path) -> Path:
base: Final = yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
strict: Final = {
**base,
"general_settings": {**base.get("general_settings", {}), "require_key_mcp_access_defined": True},
}
path: Final = directory / "require_key_mcp_access.yaml"
path.write_text(yaml.safe_dump(strict))
return path
def test_mcp_rest_toolset_name_narrows_the_list_and_serves_the_call_for_a_team_key_and_a_dashboard_member(
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029rest" + uuid.uuid4().hex[:6]
server_id: Final = register_mcp(scenario, peer, alias)
granted_id, granted_name = _toolset(scenario, server_id, "add")
_, withheld_name = _toolset(scenario, server_id, "multiply")
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
key: Final = scenario.key(team_id=team_id)
member: Final = scenario.member(team_id)
for headers in (_bearer(key), _bearer(_dashboard_ui_session_token(member))):
listed: Final = _rest_list(gateway, headers, {"toolset_name": granted_name})
assert listed.status_code == 200, listed.text
listed_names: Final = tuple(tool["name"] for tool in listed.json()["tools"])
assert len(listed_names) == 1 and listed_names[0].endswith("add"), listed.text
denied: Final = _rest_list(gateway, headers, {"toolset_name": withheld_name})
assert denied.status_code == 200 and denied.json()["tools"] == [], denied.text
assert "does not have access to toolset" in denied.json()["message"], denied.text
peer.drain()
called: Final = _rest_call(gateway, headers, listed_names[0], server_id)
assert called.ok and called.text == "9", called.raw
assert len(tool_calls(peer.drain())) == 1
def test_a_key_restricted_to_its_own_servers_does_not_inherit_the_team_toolset(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029own" + uuid.uuid4().hex[:6]
server_id: Final = register_mcp(scenario, peer, alias)
other_id: Final = register_mcp(scenario, peer, "lit6029other" + uuid.uuid4().hex[:6])
granted_id, granted_name = _toolset(scenario, server_id, "add")
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id], "mcp_servers": [other_id]})
key: Final = scenario.key(team_id=team_id, object_permission={"mcp_servers": [other_id]})
headers: Final = _bearer(key)
assert _listed_toolset_ids(gateway, headers) == ()
assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers).status_code == 403
assert _route_tools(gateway, headers, granted_name).status == 403
assert tool_calls(peer.drain()) == ()
def test_a_team_with_an_empty_or_absent_toolset_grant_gives_its_keys_nothing(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
server_id: Final = register_mcp(scenario, peer, "lit6029empty" + uuid.uuid4().hex[:6])
granted_id, granted_name = _toolset(scenario, server_id, "add")
teams: Final = (
scenario.team(object_permission={"mcp_toolsets": []}),
scenario.team(object_permission={"mcp_toolsets": None}),
scenario.team(),
)
for team_id in teams:
headers: Final = _bearer(scenario.key(team_id=team_id))
assert _listed_toolset_ids(gateway, headers) == (), team_id
assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers).status_code == 403
assert _route_tools(gateway, headers, granted_name).status == 403, team_id
assert tool_calls(peer.drain()) == ()
def test_an_unknown_or_malformed_toolset_name_is_refused_without_peer_traffic_and_the_route_keeps_serving(
gateway: Gateway,
) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029bad" + uuid.uuid4().hex[:6]
server_id: Final = register_mcp(scenario, peer, alias)
granted_id, granted_name = _toolset(scenario, server_id, "add")
_, withheld_name = _toolset(scenario, server_id, "multiply")
headers: Final = _bearer(scenario.key(team_id=scenario.team(object_permission={"mcp_toolsets": [granted_id]})))
unknown: Final = "missing" + uuid.uuid4().hex[:8]
assert _route_tools(gateway, headers, unknown).status == 404
assert _rest_list(gateway, headers, {"toolset_name": unknown}).status_code == 404
malformed: Final = (
[("toolset_name", granted_name), ("toolset_name", withheld_name)],
{"toolset_name": "x" * 5000},
{"toolset_name": ""},
{"toolset_name": granted_name + "\x00"},
)
statuses: Final = tuple(_rest_list(gateway, headers, params).status_code for params in malformed)
assert all(status < 500 for status in statuses), statuses
assert tool_calls(peer.drain()) == ()
assert gateway.client.get("/health/liveliness").status_code == 200
served: Final = _route_tools(gateway, headers, granted_name)
assert served.tools == (f"{alias}-add",), served.raw
def test_garbage_expired_and_tampered_credentials_are_refused_on_every_toolset_surface(
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
with mcp_peer() as peer, gateway.scenario() as scenario:
server_id: Final = register_mcp(scenario, peer, "lit6029cred" + uuid.uuid4().hex[:6])
granted_id, granted_name = _toolset(scenario, server_id, "add")
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
member: Final = scenario.member(team_id)
forged: Final = (
"sk-" + secrets.token_urlsafe(24),
_expired_dashboard_token(member),
"llm_session_" + secrets.token_urlsafe(32),
)
for bearer in forged:
headers: Final = _bearer(bearer)
listed: Final = gateway.client.get("/v1/mcp/toolset", headers=headers)
assert listed.status_code == 401, (bearer[:12], listed.text)
detail: Final = gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers)
assert detail.status_code == 401, (bearer[:12], detail.text)
routed: Final = _route_tools(gateway, headers, granted_name)
assert routed.status == 401, (bearer[:12], routed.raw)
rest: Final = _rest_list(gateway, headers, {"toolset_name": granted_name})
assert rest.status_code == 401, (bearer[:12], rest.text)
assert peer.drain() == ()
def test_a_dashboard_member_of_a_deleted_team_loses_the_toolset_while_a_direct_grant_survives(
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029gone" + uuid.uuid4().hex[:6]
server_id: Final = register_mcp(scenario, peer, alias)
team_id_, team_name = _toolset(scenario, server_id, "add")
own_id, own_name = _toolset(scenario, server_id, "multiply")
doomed: Final = scenario.gateway.post(
"/team/new",
{"team_alias": f"integration-{uuid.uuid4().hex}", "object_permission": {"mcp_toolsets": [team_id_]}},
)
doomed_team: Final = str(doomed["team_id"])
member: Final = scenario.user(user_role="internal_user", teams=[doomed_team])
granted: Final = scenario.user(
user_role="internal_user", teams=[doomed_team], object_permission={"mcp_toolsets": [own_id]}
)
member_headers: Final = _bearer(_dashboard_ui_session_token(member))
granted_headers: Final = _bearer(_dashboard_ui_session_token(granted))
assert _listed_toolset_ids(gateway, member_headers) == (team_id_,)
assert set(_listed_toolset_ids(gateway, granted_headers)) == {team_id_, own_id}
scenario.delete_team(doomed_team)
assert _listed_toolset_ids(gateway, member_headers) == ()
assert gateway.client.get(f"/v1/mcp/toolset/{team_id_}", headers=member_headers).status_code == 403
assert _route_tools(gateway, member_headers, team_name).status == 403
assert _listed_toolset_ids(gateway, granted_headers) == (own_id,)
assert _route_tools(gateway, granted_headers, team_name).status == 403
assert _route_tools(gateway, granted_headers, own_name).tools == (f"{alias}-multiply",)
peer.drain()
kept: Final = _route_call(gateway, granted_headers, own_name, f"{alias}-multiply")
assert kept.ok and kept.text == "20", kept.raw
assert [call["body"]["params"]["name"] for call in tool_calls(peer.drain())] == ["multiply"]
def test_require_key_mcp_access_defined_stops_key_inheritance_but_not_the_dashboard_member(
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
with (
owned_proxy(gateway, tmp_path, {}, config=_strict_config(tmp_path), workers=2) as strict,
mcp_peer() as peer,
strict.scenario() as scenario,
):
alias: Final = "lit6029strict" + uuid.uuid4().hex[:6]
server_id: Final = register_mcp(scenario, peer, alias)
granted_id, granted_name = _toolset(scenario, server_id, "add")
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
inheriting: Final = _bearer(scenario.key(team_id=team_id))
own: Final = _bearer(scenario.key(team_id=team_id, object_permission={"mcp_toolsets": [granted_id]}))
member: Final = _bearer(_dashboard_ui_session_token(scenario.member(team_id)))
assert _listed_toolset_ids(strict, inheriting) == ()
assert _route_tools(strict, inheriting, granted_name).status == 403
assert _listed_toolset_ids(strict, own) == (granted_id,)
assert _route_tools(strict, own, granted_name).tools == (f"{alias}-add",)
assert _listed_toolset_ids(strict, member) == (granted_id,)
assert _route_tools(strict, member, granted_name).tools == (f"{alias}-add",)
peer.drain()
called: Final = _route_call(strict, member, granted_name, f"{alias}-add")
assert called.ok and called.text == "9", called.raw
assert len(tool_calls(peer.drain())) == 1
def test_a_key_with_only_a_vector_store_grant_still_inherits_the_team_toolset(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029vs" + uuid.uuid4().hex[:6]
server_id: Final = register_mcp(scenario, peer, alias)
granted_id, granted_name = _toolset(scenario, server_id, "add")
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
headers: Final = _bearer(
scenario.key(team_id=team_id, object_permission={"vector_stores": ["vs-" + uuid.uuid4().hex[:8]]})
)
assert _listed_toolset_ids(gateway, headers) == (granted_id,)
assert _route_tools(gateway, headers, granted_name).tools == (f"{alias}-add",)
peer.drain()
called: Final = _route_call(gateway, headers, granted_name, f"{alias}-add")
assert called.ok and called.text == "9", called.raw
assert len(tool_calls(peer.drain())) == 1
def test_a_member_added_after_the_team_was_cached_sees_the_toolset_on_every_following_request(
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
with mcp_peer() as peer, gateway.scenario() as scenario:
server_id: Final = register_mcp(scenario, peer, "lit6029cache" + uuid.uuid4().hex[:6])
granted_id, _ = _toolset(scenario, server_id, "add")
team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
user_id: Final = scenario.user(user_role="internal_user")
headers: Final = _bearer(_dashboard_ui_session_token(user_id))
warm: Final = _bearer(scenario.key(team_id=team_id))
assert tuple(_listed_toolset_ids(gateway, warm) for _ in range(4)) == ((granted_id,),) * 4
assert tuple(_listed_toolset_ids(gateway, headers) for _ in range(4)) == ((),) * 4
added: Final = gateway.request(
"POST", "/team/member_add", {"team_id": team_id, "member": {"user_id": user_id, "role": "user"}}
)
assert added.status_code == 200, added.text
listings: Final = tuple(_listed_toolset_ids(gateway, headers) for _ in range(8))
assert listings == ((granted_id,),) * 8, listings
def test_a_member_of_two_teams_sees_the_union_and_each_route_stays_narrowed_to_its_own_toolset(
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
with mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "lit6029two" + uuid.uuid4().hex[:6]
server_id: Final = register_mcp(scenario, peer, alias)
first_id, first_name = _toolset(scenario, server_id, "add")
second_id, second_name = _toolset(scenario, server_id, "multiply")
teams: Final = (
scenario.team(object_permission={"mcp_toolsets": [first_id]}),
scenario.team(object_permission={"mcp_toolsets": [second_id]}),
)
headers: Final = _bearer(
_dashboard_ui_session_token(scenario.user(user_role="internal_user", teams=list(teams)))
)
assert set(_listed_toolset_ids(gateway, headers)) == {first_id, second_id}
assert _route_tools(gateway, headers, first_name).tools == (f"{alias}-add",)
assert _route_tools(gateway, headers, second_name).tools == (f"{alias}-multiply",)
peer.drain()
crossed: Final = _route_call(gateway, headers, first_name, f"{alias}-multiply")
assert not crossed.ok, crossed.raw
assert tool_calls(peer.drain()) == ()

View file

@ -8201,7 +8201,7 @@ class TestGatewaySessionAdmission:
assert not any(k.lower() == "authorization" for k in (raw_headers or {}))
def _make_team(team_id, mcp_servers, *, org_id=None, tool_perms=None, members=("sso-user",)):
def _make_team(team_id, mcp_servers, *, org_id=None, tool_perms=None, members=("sso-user",), toolsets=None):
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, Member
return LiteLLM_TeamTable(
@ -8210,7 +8210,10 @@ def _make_team(team_id, mcp_servers, *, org_id=None, tool_perms=None, members=("
members_with_roles=[Member(user_id=u, role="user") for u in members],
access_group_ids=[],
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id=f"op-{team_id}", mcp_servers=mcp_servers, mcp_tool_permissions=tool_perms
object_permission_id=f"op-{team_id}",
mcp_servers=mcp_servers,
mcp_tool_permissions=tool_perms,
mcp_toolsets=toolsets,
),
)
@ -8271,6 +8274,64 @@ class TestUserSubjectTeamUnion:
result = await MCPRequestHandler.get_allowed_mcp_servers(auth)
assert set(result) == {"srv1", "srv2", "srv3"}
async def test_toolsets_of_a_team_that_dropped_the_user_from_its_roster_are_not_granted(self):
"""The user's cached team list still names team-revoked, but its live roster no longer lists
the user, so its toolset is withheld exactly as its servers are on the aggregate /mcp."""
from litellm.proxy._experimental.mcp_server.ui_session_utils import granted_toolset_ids
teams = {
"team-kept": _make_team("team-kept", [], toolsets=["ts-kept"]),
"team-revoked": _make_team("team-revoked", [], toolsets=["ts-revoked"], members=("someone-else",)),
}
auth = _make_admitted_subject("sso-user")
with self._patch(teams_by_id=teams, user_teams=["team-kept", "team-revoked"]):
granted = await granted_toolset_ids(auth)
assert granted == {"ts-kept"}
async def test_a_pinned_toolset_narrows_every_source_to_the_toolset_servers_and_tools(self):
"""On /toolset/{name}/mcp the admitted subject carries mcp_toolset_id; team-a's grant on srv1 and
srv2 with every tool collapses to the toolset's srv1 and its one tool, and team-b's srv3 drops."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
teams = {"team-a": _make_team("team-a", ["srv1", "srv2"]), "team-b": _make_team("team-b", ["srv3"])}
auth = _make_admitted_subject("sso-user")
pinned = auth.model_copy(update={"mcp_toolset_id": "ts-1"})
resolve = AsyncMock(return_value={"srv1": ["add"]})
with (
self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]),
patch.object(global_mcp_server_manager, "resolve_toolset_tool_permissions", resolve),
):
servers = await MCPRequestHandler.resolve_admitted_subject_servers(pinned)
tools = await MCPRequestHandler.resolve_admitted_subject_tools("srv1", pinned)
unpinned_servers = await MCPRequestHandler.resolve_admitted_subject_servers(auth)
unpinned_tools = await MCPRequestHandler.resolve_admitted_subject_tools("srv1", auth)
assert servers == ["srv1"]
assert tools == ["add"]
assert set(unpinned_servers) == {"srv1", "srv2", "srv3"}
assert unpinned_tools is None
assert {call.kwargs["toolset_ids"][0] for call in resolve.await_args_list} == {"ts-1"}
async def test_a_fresh_policy_pinned_toolset_bypasses_the_toolset_permission_cache(self):
"""A session admitted under requires_fresh_policy reads the pinned toolset from the writer, so a
tool revoked from the toolset is gone on the very next request (Devin Review 4150024092)."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
teams = {"team-a": _make_team("team-a", ["srv1", "srv2"])}
auth = _make_admitted_subject("sso-user")
auth.requires_fresh_policy = True
pinned = auth.model_copy(update={"mcp_toolset_id": "ts-1"})
resolve = AsyncMock(return_value={"srv1": ["add"]})
with (
self._patch(teams_by_id=teams, user_teams=["team-a"]),
patch.object(global_mcp_server_manager, "resolve_toolset_tool_permissions", resolve),
):
servers = await MCPRequestHandler.resolve_admitted_subject_servers(pinned)
tools = await MCPRequestHandler.resolve_admitted_subject_tools("srv1", pinned)
assert servers == ["srv1"]
assert tools == ["add"]
assert resolve.await_args_list
assert all(call.kwargs == {"toolset_ids": ["ts-1"], "requires_fresh_policy": True} for call in resolve.await_args_list)
async def test_key_based_caller_uses_single_team_only(self):
"""A key-based caller (api_key set) with a team_id sees ONLY that team, even though the
same user belongs to other teams: key auth must be byte-identical to before."""

View file

@ -1,10 +1,12 @@
"""Tests for MCP toolset scope enforcement."""
import asyncio
from collections.abc import Awaitable, Callable
from typing import Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from litellm.proxy._types import (
LiteLLM_ObjectPermissionTable,
@ -30,6 +32,19 @@ def _make_auth(
)
def _granted_through_team(*team_toolset_ids: str) -> Callable[[UserAPIKeyAuth], Awaitable[frozenset[str]]]:
"""The real grant resolver over a team that holds ``team_toolset_ids``, with no key access rule."""
from litellm.proxy._experimental.mcp_server.ui_session_utils import granted_toolset_ids
async def team_permission(context: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable:
return LiteLLM_ObjectPermissionTable(object_permission_id="team-op", mcp_toolsets=list(team_toolset_ids))
async def granted(context: UserAPIKeyAuth) -> frozenset[str]:
return await granted_toolset_ids(context, team_object_permission=team_permission, require_key_access=False)
return granted
class TestApplyToolsetScope:
"""Tests for _apply_toolset_scope helper."""
@ -97,6 +112,122 @@ class TestApplyToolsetScope:
assert op.mcp_servers == ["server-a"]
assert op.mcp_tool_permissions == toolset_perms
@pytest.mark.asyncio
async def test_team_granted_toolset_is_served_to_a_key_without_its_own_grant(self):
"""A team key whose own row carries no toolset grant is admitted to the toolset its team
holds (LIT-6029), scoped to that toolset's servers and tools."""
from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope
toolset_perms = {"server-a": ["tool1"]}
auth = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=None)
with patch(
"litellm.proxy._experimental.mcp_server.server."
"global_mcp_server_manager.resolve_toolset_tool_permissions",
new=AsyncMock(return_value=toolset_perms),
):
result = await _apply_toolset_scope(auth, "toolset-123", granted=_granted_through_team("toolset-123"))
assert result.mcp_toolset_id == "toolset-123"
assert result.object_permission is not None
assert result.object_permission.mcp_servers == ["server-a"]
assert result.object_permission.mcp_tool_permissions == toolset_perms
@pytest.mark.asyncio
async def test_a_non_admin_dashboard_session_is_pinned_as_its_admitted_user_instead_of_rewritten(self):
"""The dashboard session acts as its admitted user, whose team grants resolve per source, so a
team-granted toolset is not capped by the user's own row: the row stays intact and the toolset
rides along as mcp_toolset_id (LIT-6029)."""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope
session = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-1")
own_row = LiteLLM_ObjectPermissionTable(object_permission_id="user-op", mcp_servers=["server-own"])
admitted = UserAPIKeyAuth(user_id="user-1", object_permission=own_row)
admitted.mcp_admitted_user_subject = True
granted = AsyncMock(return_value=frozenset({"toolset-123"}))
resolve = AsyncMock(return_value={"server-team": ["tool1"]})
with patch(
"litellm.proxy._experimental.mcp_server.server."
"global_mcp_server_manager.resolve_toolset_tool_permissions",
new=resolve,
):
result = await _apply_toolset_scope(
session, "toolset-123", acting_user=AsyncMock(return_value=admitted), granted=granted
)
assert granted.await_args is not None and granted.await_args.args[0].mcp_admitted_user_subject is True
assert result.mcp_admitted_user_subject is True
assert result.mcp_toolset_id == "toolset-123"
assert result.object_permission == own_row
resolve.assert_not_awaited()
@pytest.mark.asyncio
async def test_a_gateway_admitted_user_without_the_toolset_in_any_source_is_denied(self):
from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope
admitted = UserAPIKeyAuth(user_id="user-1", object_permission=None)
admitted.mcp_admitted_user_subject = True
granted = AsyncMock(return_value=frozenset({"toolset-other"}))
with pytest.raises(HTTPException) as exc_info:
await _apply_toolset_scope(admitted, "toolset-123", granted=granted)
assert exc_info.value.status_code == 403
granted.assert_awaited_once_with(admitted)
@pytest.mark.asyncio
async def test_a_resource_scoped_admitted_user_is_denied_a_team_toolset_on_another_server(self):
"""A gateway bearer scoped to server-own (RFC 8707 resource) cannot open a team toolset whose
servers lie outside that resource, even though the team grants it (Devin Review 4150024267)."""
from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope
admitted = UserAPIKeyAuth(user_id="user-1", object_permission=None)
admitted.mcp_admitted_user_subject = True
admitted.mcp_session_resource_server_id = "server-own"
admitted.requires_fresh_policy = True
granted = AsyncMock(return_value=frozenset({"toolset-123"}))
resolve = AsyncMock(return_value={"server-team": ["tool1"]})
with patch(
"litellm.proxy._experimental.mcp_server.server."
"global_mcp_server_manager.resolve_toolset_tool_permissions",
new=resolve,
):
with pytest.raises(HTTPException) as exc_info:
await _apply_toolset_scope(admitted, "toolset-123", granted=granted)
assert exc_info.value.status_code == 403
resolve.assert_awaited_once_with(toolset_ids=["toolset-123"], requires_fresh_policy=True)
@pytest.mark.asyncio
async def test_a_resource_scoped_admitted_user_opens_a_toolset_inside_its_resource(self):
from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope
admitted = UserAPIKeyAuth(user_id="user-1", object_permission=None)
admitted.mcp_admitted_user_subject = True
admitted.mcp_session_resource_server_id = "server-team"
granted = AsyncMock(return_value=frozenset({"toolset-123"}))
resolve = AsyncMock(return_value={"server-team": ["tool1"], "server-other": ["tool2"]})
with patch(
"litellm.proxy._experimental.mcp_server.server."
"global_mcp_server_manager.resolve_toolset_tool_permissions",
new=resolve,
):
result = await _apply_toolset_scope(admitted, "toolset-123", granted=granted)
assert result.mcp_toolset_id == "toolset-123"
assert result.mcp_session_resource_server_id == "server-team"
resolve.assert_awaited_once_with(toolset_ids=["toolset-123"], requires_fresh_policy=False)
@pytest.mark.asyncio
async def test_team_grant_for_another_toolset_does_not_admit_a_key_to_this_one(self):
from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope
auth = _make_auth(mcp_toolsets=[])
auth.team_id = "team-a"
with pytest.raises(HTTPException) as exc_info:
await _apply_toolset_scope(auth, "toolset-123", granted=_granted_through_team("toolset-other"))
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_non_admin_no_object_permission_raises_403(self):
"""Non-admin key with object_permission=None is denied (no grants configured)."""
@ -250,6 +381,131 @@ class TestFetchMCPToolsetsAccess:
assert len(result) == 2
mock_list.assert_called_once_with(mock_client, toolset_ids=["ts-1", "ts-2"])
@pytest.mark.asyncio
async def test_team_granted_toolsets_are_listed_for_a_key_without_its_own_grant(self):
"""GET /v1/mcp/toolset for a team key lists the team's toolsets (LIT-6029)."""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_mcp_toolsets,
)
auth = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=None)
team_permission = LiteLLM_ObjectPermissionTable(object_permission_id="team-op", mcp_toolsets=["ts-team"])
fake_toolsets = [MagicMock(toolset_id="ts-team")]
mock_client = MagicMock()
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=mock_client,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets",
new=AsyncMock(return_value=fake_toolsets),
) as mock_list,
patch.object(
MCPRequestHandler,
"_get_team_object_permission",
new=AsyncMock(return_value=team_permission),
),
):
result = await fetch_mcp_toolsets(user_api_key_dict=auth)
assert result == fake_toolsets
mock_list.assert_called_once_with(mock_client, toolset_ids=["ts-team"])
@pytest.mark.asyncio
async def test_admin_with_own_grants_is_not_narrowed_by_a_team_lookup(self):
"""An admin's own grant list is the only filter; no team lookup runs for admins."""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_mcp_toolsets,
)
auth = _make_auth(mcp_toolsets=["ts-1"])
auth.user_role = LitellmUserRoles.PROXY_ADMIN
mock_client = MagicMock()
own_toolsets = [{"toolset_id": "ts-1", "toolset_name": "own"}]
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=mock_client,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets",
new=AsyncMock(return_value=own_toolsets),
) as mock_list,
patch.object(
MCPRequestHandler, "_get_team_object_permission", new=AsyncMock(return_value=None)
) as team_lookup,
):
result = await fetch_mcp_toolsets(user_api_key_dict=auth)
assert result == own_toolsets
mock_list.assert_called_once_with(mock_client, toolset_ids=["ts-1"])
team_lookup.assert_not_awaited()
class TestFetchMCPToolsetAccess:
"""Tests for GET /v1/mcp/toolset/{toolset_id} access control."""
@staticmethod
async def _fetch(auth: UserAPIKeyAuth, toolset_id: str, team_toolsets: list[str] | None):
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_mcp_toolset,
)
team_permission = (
LiteLLM_ObjectPermissionTable(object_permission_id="team-op", mcp_toolsets=team_toolsets)
if team_toolsets is not None
else None
)
toolset = MagicMock(toolset_id=toolset_id)
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=MagicMock(),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_toolset",
new=AsyncMock(return_value=toolset),
),
patch.object(
MCPRequestHandler,
"_get_team_object_permission",
new=AsyncMock(return_value=team_permission),
),
):
return await fetch_mcp_toolset(toolset_id=toolset_id, user_api_key_dict=auth)
@pytest.mark.asyncio
async def test_team_granted_toolset_detail_is_served_to_a_key_without_its_own_grant(self):
auth = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=None)
toolset = await self._fetch(auth, "ts-team", team_toolsets=["ts-team"])
assert toolset.toolset_id == "ts-team"
@pytest.mark.asyncio
async def test_toolset_detail_stays_forbidden_when_neither_key_nor_team_holds_it(self):
from fastapi import HTTPException
auth = _make_auth(mcp_toolsets=["ts-own"])
auth.team_id = "team-a"
with pytest.raises(HTTPException) as exc_info:
await self._fetch(auth, "ts-withheld", team_toolsets=["ts-team"])
assert exc_info.value.status_code == 403
class TestToolsetPrefixResolution:
"""Regression for LIT-3419.

View file

@ -6,13 +6,14 @@ import pytest
from fastapi import HTTPException
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
build_effective_auth_contexts,
clone_user_api_key_auth_with_team,
granted_toolset_ids,
toolset_grant_contexts,
resolve_ui_session_team_ids,
)
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
def test_clone_user_api_key_auth_with_team_creates_independent_copy():
@ -258,3 +259,234 @@ async def test_admitted_user_context_carries_the_request_span(monkeypatch):
assert (await acting_user_auth(user_auth)).parent_otel_span is parent_span
assert (await build_effective_auth_contexts(user_auth))[-1].parent_otel_span is parent_span
def _toolset_permission(*toolset_ids: str) -> LiteLLM_ObjectPermissionTable:
return LiteLLM_ObjectPermissionTable(
object_permission_id=f"op-{'-'.join(toolset_ids)}", mcp_toolsets=list(toolset_ids)
)
@pytest.mark.asyncio
async def test_granted_toolset_ids_unions_own_and_team_grants_over_every_effective_context():
"""A dashboard session of a user in two teams holds the toolsets of both teams plus the ones on
the user row itself, exactly the grant sources the aggregate /mcp listing expands."""
session = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-1")
team_a = UserAPIKeyAuth(team_id="team-a", user_id="user-1")
team_b = UserAPIKeyAuth(team_id="team-b", user_id="user-1", object_permission=_toolset_permission())
admitted = UserAPIKeyAuth(user_id="user-1", object_permission=_toolset_permission("ts-user"))
team_grants = {"team-a": _toolset_permission("ts-a", "ts-shared"), "team-b": _toolset_permission("ts-b")}
async def effective_contexts(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]:
assert auth is session
return [team_a, team_b, admitted]
async def team_permission(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None:
return team_grants.get(auth.team_id or "")
granted = await granted_toolset_ids(session, effective_contexts, team_permission)
assert granted == frozenset({"ts-a", "ts-shared", "ts-b", "ts-user"})
@pytest.mark.asyncio
async def test_granted_toolset_ids_is_empty_when_neither_key_nor_team_grants_a_toolset():
key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=_toolset_permission())
async def effective_contexts(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]:
return [auth]
async def no_team_permission(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None:
return None
assert await granted_toolset_ids(key, effective_contexts, no_team_permission) == frozenset()
async def _same_context(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]:
return [auth]
async def _team_grants_ts_team(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None:
return _toolset_permission("ts-team")
@pytest.mark.asyncio
@pytest.mark.parametrize(
"own",
[
LiteLLM_ObjectPermissionTable(object_permission_id="op", mcp_toolsets=["ts-own"]),
LiteLLM_ObjectPermissionTable(object_permission_id="op", mcp_servers=["srv-own"]),
LiteLLM_ObjectPermissionTable(object_permission_id="op", mcp_tool_permissions={"srv-own": ["add"]}),
LiteLLM_ObjectPermissionTable(object_permission_id="op", mcp_access_groups=["group-own"]),
],
)
async def test_a_key_declaring_its_own_mcp_grant_does_not_inherit_the_team_toolsets(own):
"""The key/team rule of the aggregate listing: a key's own MCP grant is a ceiling the team cannot widen."""
key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=own)
granted = await granted_toolset_ids(key, _same_context, _team_grants_ts_team, require_key_access=False)
assert granted == frozenset(own.mcp_toolsets or ())
@pytest.mark.asyncio
async def test_require_key_mcp_access_defined_stops_a_key_inheriting_team_toolsets_but_not_a_session():
key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a")
session = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-1")
assert await granted_toolset_ids(key, _same_context, _team_grants_ts_team, require_key_access=False) == {"ts-team"}
assert await granted_toolset_ids(key, _same_context, _team_grants_ts_team, require_key_access=True) == frozenset()
assert await granted_toolset_ids(session, _same_context, _team_grants_ts_team, require_key_access=True) == {
"ts-team"
}
@pytest.mark.asyncio
async def test_toolset_grant_contexts_of_a_virtual_key_is_the_key_alone():
key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a")
async def never(auth: UserAPIKeyAuth) -> None:
raise AssertionError("a virtual key has no admitted sources")
assert await toolset_grant_contexts(key, admitted_context=never, admitted_sources=never) == (key,)
def _admitted(user_id: str, own: LiteLLM_ObjectPermissionTable | None = None) -> UserAPIKeyAuth:
subject = UserAPIKeyAuth(user_id=user_id, object_permission=own)
subject.mcp_admitted_user_subject = True
return subject
@pytest.mark.asyncio
async def test_toolset_grant_contexts_of_a_dashboard_session_are_its_admitted_users_grant_sources():
"""The dashboard session fans out through the same roster-checked source builder as the aggregate
/mcp resolution, applied to the admitted user it acts as, so a cached membership a team has since
revoked never reaches the toolset check."""
session = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-1")
admitted = _admitted("user-1")
own_source = UserAPIKeyAuth(user_id="user-1")
team_source = UserAPIKeyAuth(user_id="user-1", team_id="team-a")
async def admitted_context(auth: UserAPIKeyAuth) -> UserAPIKeyAuth:
assert auth is session
return admitted
async def admitted_sources(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]:
assert auth is admitted
return [own_source, team_source]
assert await toolset_grant_contexts(session, admitted_context, admitted_sources) == (own_source, team_source)
@pytest.mark.asyncio
async def test_toolset_grant_contexts_of_a_gateway_admitted_user_are_its_own_grant_sources():
admitted = _admitted("user-1")
team_source = UserAPIKeyAuth(user_id="user-1", team_id="team-a")
async def no_dashboard_context(auth: UserAPIKeyAuth) -> None:
return None
async def admitted_sources(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]:
assert auth is admitted
return [team_source]
assert await toolset_grant_contexts(admitted, no_dashboard_context, admitted_sources) == (team_source,)
@pytest.mark.asyncio
async def test_a_source_declaring_its_own_mcp_grant_never_reads_its_team():
key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=_toolset_permission("ts-own"))
team_reads: list[str | None] = [] # mutable-ok: records the lookups the code under test performs
async def team_permission(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None:
team_reads.append(auth.team_id)
return _toolset_permission("ts-team")
granted = await granted_toolset_ids(key, _same_context, team_permission, require_key_access=False)
assert granted == {"ts-own"}
assert team_reads == []
async def _team_a_unreadable(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None:
if auth.team_id == "team-a":
raise RuntimeError("team row unreadable")
return _toolset_permission("ts-b")
@pytest.mark.asyncio
async def test_an_unreadable_team_grants_nothing_while_the_direct_and_other_team_grants_still_count():
"""A dashboard user whose own row grants ts-user and who sits on team-a and team-b keeps ts-user and
ts-b when team-a cannot be read; team-a itself contributes nothing rather than failing the lookup."""
admitted = _admitted("user-1", _toolset_permission("ts-user"))
team_a = UserAPIKeyAuth(user_id="user-1", team_id="team-a")
team_b = UserAPIKeyAuth(user_id="user-1", team_id="team-b")
async def sources(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]:
return [admitted, team_a, team_b]
assert await granted_toolset_ids(admitted, sources, _team_a_unreadable) == {"ts-user", "ts-b"}
@pytest.mark.asyncio
async def test_a_key_whose_only_grant_source_is_an_unreadable_team_is_granted_nothing():
key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a")
assert await granted_toolset_ids(key, _same_context, _team_a_unreadable, require_key_access=False) == frozenset()
async def _hydrates_op_key_to_srv_own(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None:
if auth.object_permission is not None:
return auth.object_permission
if auth.object_permission_id == "op-key":
return LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_servers=["srv-own"])
return None
@pytest.mark.asyncio
async def test_a_key_cached_with_its_own_grant_unhydrated_is_scoped_to_that_grant_not_its_team():
"""The main auth flow can cache a key with object_permission_id set and object_permission None. The
row it names is the key's ceiling, so it is loaded and read as the key's own grant instead of letting the
key inherit its team's toolsets."""
key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission_id="op-key")
granted = await granted_toolset_ids(
key,
_same_context,
_team_grants_ts_team,
require_key_access=False,
own_object_permission=_hydrates_op_key_to_srv_own,
)
assert granted == frozenset()
async def _own_row_unreadable(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None:
raise RuntimeError("object permission row unreadable")
async def _own_row_gone(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None:
return None
@pytest.mark.asyncio
@pytest.mark.parametrize("load_own", [_own_row_unreadable, _own_row_gone])
async def test_a_key_naming_an_own_grant_that_cannot_be_read_is_granted_nothing_rather_than_its_team(load_own):
key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission_id="op-key")
granted = await granted_toolset_ids(
key, _same_context, _team_grants_ts_team, require_key_access=False, own_object_permission=load_own
)
assert granted == frozenset()
@pytest.mark.asyncio
async def test_a_key_naming_no_own_grant_is_not_hydrated_before_inheriting_its_team():
key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a")
granted = await granted_toolset_ids(
key, _same_context, _team_grants_ts_team, require_key_access=False, own_object_permission=_own_row_unreadable
)
assert granted == {"ts-team"}

View file

@ -13596,6 +13596,63 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp
assert mock_audit.call_args.kwargs["team_alias"] == "list-audit"
@pytest.mark.asyncio
async def test_team_member_add_evicts_the_cached_team_roster(monkeypatch):
"""Roster checks read the team through get_team_object, so a cached pre-add roster must be dropped."""
from litellm.proxy._types import TeamMemberAddRequest
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.team_endpoints import team_member_add
team_id = "team-roster-evict"
team_row = LiteLLM_TeamTable(team_id=team_id, team_alias="roster-evict", members_with_roles=[])
cache = UserApiKeyCache()
cache.set_cache(key=f"team_id:{team_id}", value=team_row)
cache.set_cache(key="team_alias:roster-evict", value=team_row)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock())
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id")
joined_user = LiteLLM_UserTable(user_id="joiner", max_budget=None, spend=0.0, models=[])
updated_team = MagicMock()
updated_team.model_dump.return_value = {"team_id": team_id, "members_with_roles": []}
with (
patch(
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
new_callable=AsyncMock,
return_value=team_row,
),
patch(
"litellm.proxy.management_endpoints.team_endpoints._validate_team_member_add_permissions",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_endpoints.team_endpoints._validate_and_populate_member_user_info",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_endpoints.team_endpoints._resolve_existing_member_user_ids",
new_callable=AsyncMock,
return_value=frozenset(),
),
patch(
"litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team",
new_callable=AsyncMock,
return_value=(updated_team, [joined_user], []),
),
patch("litellm.proxy.management_endpoints.team_endpoints._schedule_team_member_add_audit_logs"),
):
await team_member_add(
data=TeamMemberAddRequest(team_id=team_id, member=Member(user_id="joiner", role="user")),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"),
)
assert cache.get_cache(key=f"team_id:{team_id}") is None
assert cache.get_cache(key="team_alias:roster-evict") is None
class _RecordingAuditLogger(CustomLogger):
def __init__(self) -> None:
super().__init__()

View file

@ -8,7 +8,8 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import HTTPException
from mcp.types import CallToolResult, TextContent, Tool as MCPTool
from mcp.types import CallToolResult, TextContent
from mcp.types import Tool as MCPTool
from openai.types.responses.tool_param import Mcp
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing
@ -110,9 +111,7 @@ def test_extract_tool_calls_from_chat_response_handles_tool_calls():
object="chat.completion",
)
tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response(
response
)
tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response(response)
assert len(tool_calls) == 1
assert tool_calls[0]["function"]["name"] == "foo"
@ -182,9 +181,7 @@ def test_transform_mcp_tools_to_openai_uses_chat_format(monkeypatch):
fake_transform_responses,
)
chat_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(
["tool"], target_format="chat"
)
chat_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(["tool"], target_format="chat")
resp_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(["tool"])
assert chat_tools == [{"chat": True}]
@ -304,9 +301,7 @@ async def test_execute_tool_calls_strips_prefix_when_alias_differs_from_server_n
)
from litellm.proxy._experimental.mcp_server import mcp_server_manager as _msm
_msm.global_mcp_server_manager._get_mcp_server_from_tool_name = MagicMock(
return_value=fake_server
)
_msm.global_mcp_server_manager._get_mcp_server_from_tool_name = MagicMock(return_value=fake_server)
tool_name = "my_deepwiki-read_wiki_structure"
tool_calls = [
@ -380,7 +375,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey
fake_manager = types.SimpleNamespace(
get_registry=MagicMock(return_value={}),
call_tool=AsyncMock(side_effect=HTTPException(status_code=500, detail="boom"))
call_tool=AsyncMock(side_effect=HTTPException(status_code=500, detail="boom")),
)
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
@ -388,9 +383,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey
)
tool_name = "deepwiki-read_wiki_structure"
tool_calls = [
{"id": "call-err", "function": {"name": tool_name, "arguments": "{}"}}
]
tool_calls = [{"id": "call-err", "function": {"name": tool_name, "arguments": "{}"}}]
user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user")
@ -408,10 +401,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey
post_call_failure_hook.assert_awaited_once()
assert post_call_failure_hook.await_args is not None
assert (
post_call_failure_hook.await_args.kwargs.get("route")
== "/responses/mcp/call_tool"
)
assert post_call_failure_hook.await_args.kwargs.get("route") == "/responses/mcp/call_tool"
@pytest.mark.asyncio
@ -434,9 +424,7 @@ async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_functio
# NOTE: Don't patch via dotted string path here because `litellm.responses`
# is a function attribute on the `litellm` package (shadowing the submodule),
# which breaks monkeypatch's importpath resolution.
handler_module = importlib.import_module(
"litellm.responses.mcp.litellm_proxy_mcp_handler"
)
handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler")
monkeypatch.setattr(handler_module, "function_setup", fake_function_setup)
tool_name = "deepwiki-read_wiki_structure"
@ -516,7 +504,9 @@ async def test_execute_tool_calls_applies_post_call_hook_content(monkeypatch):
logging_obj = MagicMock()
logging_obj.model_call_details = {}
logging_obj.async_post_mcp_tool_call_hook = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="[REDACTED]")], is_error=True))
logging_obj.async_post_mcp_tool_call_hook = AsyncMock(
return_value=CallToolResult(content=[TextContent(type="text", text="[REDACTED]")], is_error=True)
)
logging_obj.async_success_handler = AsyncMock()
handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler")
monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None))
@ -679,9 +669,7 @@ async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch
user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user")
tools, _server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager(
user_api_key_auth=user_auth,
mcp_tools_with_litellm_proxy=[
{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}
],
mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}],
)
forwarded: Final = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(tools)
@ -700,9 +688,7 @@ async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch
def test_get_parent_request_tags_from_metadata():
tags = LiteLLM_Proxy_MCP_Handler._get_parent_request_tags(
{"metadata": {"tags": ["team-a", "prod"]}}
)
tags = LiteLLM_Proxy_MCP_Handler._get_parent_request_tags({"metadata": {"tags": ["team-a", "prod"]}})
assert tags == ["team-a", "prod"]
@ -739,9 +725,7 @@ async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch):
await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager(
user_api_key_auth=types.SimpleNamespace(api_key="k", user_id="u"),
mcp_tools_with_litellm_proxy=[
{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}
],
mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}],
request_tags=["team-a"],
)
@ -761,9 +745,7 @@ async def test_execute_tool_calls_exposes_sanitized_client_headers_to_logging(mo
captured.update(kwargs)
return None, None
handler_module = importlib.import_module(
"litellm.responses.mcp.litellm_proxy_mcp_handler"
)
handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler")
monkeypatch.setattr(handler_module, "function_setup", fake_function_setup)
tool_name = "deepwiki-read_wiki_structure"
@ -789,9 +771,7 @@ async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monk
captured.update(kwargs)
return None, None
handler_module = importlib.import_module(
"litellm.responses.mcp.litellm_proxy_mcp_handler"
)
handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler")
monkeypatch.setattr(handler_module, "function_setup", fake_function_setup)
tool_name = "deepwiki-read_wiki_structure"
@ -1171,7 +1151,9 @@ def test_create_follow_up_input_keeps_each_reasoning_item_before_its_function_ca
"function_call_output",
"function_call_output",
]
assert [cast(dict[str, Any], item).get("id") or cast(dict[str, Any], item).get("call_id") for item in follow_up[1:5]] == [
assert [
cast(dict[str, Any], item).get("id") or cast(dict[str, Any], item).get("call_id") for item in follow_up[1:5]
] == [
"rs_1",
"call-1",
"rs_2",
@ -1229,16 +1211,20 @@ async def test_mcp_follow_up_call_is_stateless_when_store_is_false(
async def fake_aresponses(**kwargs: Any) -> ResponsesAPIResponse:
captured_calls.append(kwargs)
return first_response if len(captured_calls) == 1 else ResponsesAPIResponse(
id="resp_follow_up",
created_at=1234567891,
model="gpt-5",
object="response",
status="completed",
output=[],
parallel_tool_calls=False,
tool_choice="auto",
tools=[],
return (
first_response
if len(captured_calls) == 1
else ResponsesAPIResponse(
id="resp_follow_up",
created_at=1234567891,
model="gpt-5",
object="response",
status="completed",
output=[],
parallel_tool_calls=False,
tool_choice="auto",
tools=[],
)
)
async def fake_process(**kwargs: Any) -> tuple[list[Any], dict[str, str]]:
@ -1279,12 +1265,14 @@ async def test_mcp_follow_up_call_is_stateless_when_store_is_false(
@pytest.mark.asyncio
async def test_responses_discovery_logs_sanitized_caller_headers(monkeypatch: pytest.MonkeyPatch):
from litellm.proxy._experimental.mcp_server import operations
from litellm.proxy._experimental.mcp_server import mcp_server_manager
from litellm.proxy._experimental.mcp_server import mcp_server_manager, operations
headers: Final = {
"x-app-id": "app-a", "x-nuid": "user-a", "x-user-id": "identity-a",
"x-mcp-deepwiki-authorization": "upstream-sentinel", "authorization": "proxy-sentinel",
"x-app-id": "app-a",
"x-nuid": "user-a",
"x-user-id": "identity-a",
"x-mcp-deepwiki-authorization": "upstream-sentinel",
"authorization": "proxy-sentinel",
}
manager: Final = types.SimpleNamespace(
get_registry=MagicMock(return_value={}),
@ -1298,12 +1286,21 @@ async def test_responses_discovery_logs_sanitized_caller_headers(monkeypatch: py
monkeypatch.setattr(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[]))
monkeypatch.setattr(operations, "function_setup", setup)
response: Final = ResponsesAPIResponse(
id="resp_test", created_at=1234567891, model="test-model", object="response",
status="completed", output=[], parallel_tool_calls=False, tool_choice="auto", tools=[],
id="resp_test",
created_at=1234567891,
model="test-model",
object="response",
status="completed",
output=[],
parallel_tool_calls=False,
tool_choice="auto",
tools=[],
)
monkeypatch.setattr(responses_main, "aresponses", AsyncMock(return_value=response))
result: Final = await responses_main.aresponses_api_with_mcp(
input="hi", model="test-model", tools=[{"type": "mcp", "server_url": "litellm_proxy"}],
input="hi",
model="test-model",
tools=[{"type": "mcp", "server_url": "litellm_proxy"}],
secret_fields={"raw_headers": headers},
)
assert result is response
@ -1311,3 +1308,84 @@ async def test_responses_discovery_logs_sanitized_caller_headers(monkeypatch: py
logged: Final = setup.call_args.kwargs["metadata"]["headers"]
assert logged == {"x-app-id": "app-a", "x-nuid": "user-a", "x-user-id": "identity-a"}
assert headers["x-mcp-deepwiki-authorization"] == "upstream-sentinel"
def _toolset_gateway_manager(toolset_id: str, server_id: str) -> types.SimpleNamespace:
return types.SimpleNamespace(
get_registry=MagicMock(return_value={}),
get_allowed_mcp_servers=AsyncMock(return_value=[]),
get_mcp_servers_from_ids=MagicMock(return_value=[]),
get_mcp_server_by_name=MagicMock(return_value=None),
get_toolset_by_name_cached=AsyncMock(return_value=types.SimpleNamespace(toolset_id=toolset_id)),
resolve_toolset_tool_permissions=AsyncMock(return_value={server_id: ["add"]}),
)
async def _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id: str) -> dict[str, object]:
from litellm.proxy._experimental.mcp_server.ui_session_utils import granted_toolset_ids
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth
mock_get_tools = AsyncMock(return_value=AggregateToolListing(tools=[], outcomes={}))
monkeypatch.setattr("litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", mock_get_tools)
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
_toolset_gateway_manager("ts-granted", "srv-1"),
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
async def team_permission(context: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable:
return LiteLLM_ObjectPermissionTable(object_permission_id="op-team", mcp_toolsets=[team_toolset_id])
async def granted_through_team(context: UserAPIKeyAuth) -> frozenset[str]:
return await granted_toolset_ids(context, team_object_permission=team_permission, require_key_access=False)
team_key: Final = UserAPIKeyAuth(api_key="sk-team", team_id="team-1", user_role=LitellmUserRoles.INTERNAL_USER)
await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager(
user_api_key_auth=team_key,
mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/team-toolset"}],
granted_toolsets=granted_through_team,
)
assert mock_get_tools.await_args is not None
return mock_get_tools.await_args.kwargs
@pytest.mark.asyncio
async def test_toolset_gateway_url_scopes_a_team_granted_toolset_for_a_key_without_its_own_grant(monkeypatch):
kwargs: Final = await _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id="ts-granted")
scoped = kwargs["user_api_key_auth"].object_permission
assert scoped is not None
assert scoped.mcp_servers == ["srv-1"]
assert scoped.mcp_tool_permissions == {"srv-1": ["add"]}
assert kwargs["mcp_servers"] is None
@pytest.mark.asyncio
async def test_toolset_gateway_url_skips_a_toolset_the_team_does_not_grant(monkeypatch):
kwargs: Final = await _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id="ts-other")
assert kwargs["user_api_key_auth"].object_permission is None
assert kwargs["mcp_servers"] is None
@pytest.mark.asyncio
async def test_apply_toolset_permissions_pins_the_auth_to_explicit_grants_only(monkeypatch: pytest.MonkeyPatch):
"""A toolset gateway URL must not widen to operator-open (allow_all_keys) servers."""
from litellm.proxy._types import UserAPIKeyAuth
fake_manager = types.SimpleNamespace(
resolve_toolset_tool_permissions=AsyncMock(return_value={"srv-1": ["add"]}),
)
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
fake_manager,
)
scoped = await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions(
resolved_toolset_ids=["ts-1"],
resolved_mcp_servers=[],
user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", user_id="u1"),
)
assert scoped.mcp_explicit_grants_only is True
assert scoped.object_permission is not None
assert scoped.object_permission.mcp_servers == ["srv-1"]
assert scoped.object_permission.mcp_tool_permissions == {"srv-1": ["add"]}