mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(mcp): resolve config-defined servers in per-user credential and env-var endpoints (#31171)
The per-user BYOK, OAuth (OBO), and env-var management endpoints resolved the target MCP server through a DB-only lookup (get_mcp_server / get_all_mcp_servers_for_user). A server defined in config.yaml lives only in the in-memory registry and never gets a row in LiteLLM_MCPServerTable, so those endpoints raised 404 "MCP Server <id> not found" (or 403 for non-admins) before any credential could be stored, leaving config-server users unable to connect and forced to re-authorize forever. Route all three through a single registry-aware resolver: DB first, then the in-memory registry (built into LiteLLM_MCPServerTable via _build_mcp_server_table, the same fallback fetch_mcp_server already uses), then the canonical get_allowed_mcp_servers authorization the MCP gateway enforces on tool calls. Admins get a 404 for an unknown id; non-admins get 403 for a missing-or-forbidden server so server ids stay non-enumerable. This also closes a gap where the two store endpoints performed no per-server authorization at all.
This commit is contained in:
tin-berri
GitHub
parent
a6b7dcc7d6
commit
360adbe765
2 changed files with 306 additions and 51 deletions
|
|
@ -1930,12 +1930,9 @@ if MCP_AVAILABLE:
|
|||
prisma_client = get_prisma_client_or_throw(
|
||||
"Database not connected. Connect a database to your proxy"
|
||||
)
|
||||
mcp_server = await get_mcp_server(prisma_client, server_id)
|
||||
if mcp_server is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail={"error": f"MCP Server {server_id} not found"},
|
||||
)
|
||||
mcp_server = await _authorize_and_fetch_mcp_server(
|
||||
prisma_client, user_api_key_dict, server_id
|
||||
)
|
||||
if not getattr(mcp_server, "is_byok", False):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
|
|
@ -2010,12 +2007,9 @@ if MCP_AVAILABLE:
|
|||
prisma_client = get_prisma_client_or_throw(
|
||||
"Database not connected. Connect a database to your proxy"
|
||||
)
|
||||
mcp_server = await get_mcp_server(prisma_client, server_id)
|
||||
if mcp_server is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail={"error": f"MCP Server {server_id} not found"},
|
||||
)
|
||||
await _authorize_and_fetch_mcp_server(
|
||||
prisma_client, user_api_key_dict, server_id
|
||||
)
|
||||
user_id = user_api_key_dict.user_id or ""
|
||||
if not user_id:
|
||||
raise HTTPException(
|
||||
|
|
@ -2177,37 +2171,47 @@ if MCP_AVAILABLE:
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
server_id: str,
|
||||
) -> LiteLLM_MCPServerTable:
|
||||
"""Return the MCP server the caller may manage env vars for.
|
||||
"""Resolve the MCP server a caller may manage their own per-user state for.
|
||||
|
||||
Admins look the server up directly. Non-admins reuse the access-scoped
|
||||
listing that already loads every server they can see, so we don't issue
|
||||
a second per-server query just to re-fetch a record the authorization
|
||||
check produced. A non-admin who can't see the server gets 403 (never
|
||||
404) so server ids can't be enumerated.
|
||||
Looks the server up in the DB, then the in-memory registry, so a
|
||||
config-defined server (which never gets a DB row) resolves too. Admins
|
||||
may reach any server and get a 404 for an unknown id. A non-admin may
|
||||
only reach a server in their allowed set and otherwise gets 403 (never
|
||||
404, so server ids can't be enumerated), using the same allowed-server
|
||||
resolution the MCP gateway enforces on tool calls.
|
||||
"""
|
||||
server = await get_mcp_server(prisma_client, server_id)
|
||||
if server is None:
|
||||
registry_server = global_mcp_server_manager.get_mcp_server_by_id(server_id)
|
||||
if registry_server is not None:
|
||||
server = global_mcp_server_manager._build_mcp_server_table(
|
||||
registry_server
|
||||
)
|
||||
|
||||
if _user_has_admin_view(user_api_key_dict):
|
||||
server = await get_mcp_server(prisma_client, server_id)
|
||||
if server is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail={"error": f"MCP Server {server_id} not found"},
|
||||
)
|
||||
return server
|
||||
accessible = await get_all_mcp_servers_for_user(
|
||||
prisma_client, user_api_key_dict
|
||||
)
|
||||
for server in accessible:
|
||||
if server.server_id == server_id:
|
||||
return server
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={
|
||||
"error": (
|
||||
f"User does not have permission to access mcp server with id {server_id}. "
|
||||
"You can only manage env vars for mcp servers that you have access to."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
allowed_server_ids: set[str] = set()
|
||||
for auth_context in await build_effective_auth_contexts(user_api_key_dict):
|
||||
allowed_server_ids.update(
|
||||
await global_mcp_server_manager.get_allowed_mcp_servers(auth_context)
|
||||
)
|
||||
if server is None or server.server_id not in allowed_server_ids:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail={
|
||||
"error": (
|
||||
f"User does not have permission to access mcp server with id {server_id}. "
|
||||
"You can only manage mcp servers that you have access to."
|
||||
)
|
||||
},
|
||||
)
|
||||
return server
|
||||
|
||||
def _compute_user_env_var_status(
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -3468,6 +3468,10 @@ async def test_store_mcp_oauth_user_credential_returns_status():
|
|||
return_value=generate_mock_mcp_server_db_record(server_id=server_id)
|
||||
),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
|
||||
return_value=True,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.store_user_oauth_credential",
|
||||
new=AsyncMock(return_value=None),
|
||||
|
|
@ -4558,8 +4562,13 @@ class TestMCPUserEnvVarsAccessControl:
|
|||
),
|
||||
patch.object(
|
||||
mgmt_endpoints,
|
||||
"get_all_mcp_servers_for_user",
|
||||
AsyncMock(return_value=[_make_env_var_server(server_id="other")]),
|
||||
"build_effective_auth_contexts",
|
||||
AsyncMock(return_value=[object()]),
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
AsyncMock(return_value=["other"]),
|
||||
),
|
||||
patch.object(mgmt_endpoints, "get_user_env_vars", get_user_env_vars),
|
||||
):
|
||||
|
|
@ -4589,7 +4598,12 @@ class TestMCPUserEnvVarsAccessControl:
|
|||
),
|
||||
patch.object(
|
||||
mgmt_endpoints,
|
||||
"get_all_mcp_servers_for_user",
|
||||
"build_effective_auth_contexts",
|
||||
AsyncMock(return_value=[object()]),
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
AsyncMock(return_value=[]),
|
||||
),
|
||||
patch.object(mgmt_endpoints, "merge_user_env_vars", merge_mock),
|
||||
|
|
@ -4623,7 +4637,12 @@ class TestMCPUserEnvVarsAccessControl:
|
|||
),
|
||||
patch.object(
|
||||
mgmt_endpoints,
|
||||
"get_all_mcp_servers_for_user",
|
||||
"build_effective_auth_contexts",
|
||||
AsyncMock(return_value=[object()]),
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
AsyncMock(return_value=[]),
|
||||
),
|
||||
patch.object(mgmt_endpoints, "delete_user_env_vars", delete_mock),
|
||||
|
|
@ -4655,8 +4674,13 @@ class TestMCPUserEnvVarsAccessControl:
|
|||
),
|
||||
patch.object(
|
||||
mgmt_endpoints,
|
||||
"get_all_mcp_servers_for_user",
|
||||
AsyncMock(return_value=[server]),
|
||||
"build_effective_auth_contexts",
|
||||
AsyncMock(return_value=[object()]),
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
AsyncMock(return_value=["srv-1"]),
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints,
|
||||
|
|
@ -4676,13 +4700,13 @@ class TestMCPUserEnvVarsAccessControl:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_bypasses_access_check(self):
|
||||
"""Proxy admins must not be filtered by get_all_mcp_servers_for_user."""
|
||||
"""Proxy admins must not be filtered by the allowed-server check."""
|
||||
server = _make_env_var_server(
|
||||
server_id="srv-1",
|
||||
env_vars=_ENV_VARS_MIXED,
|
||||
static_headers=_STATIC_HEADERS_MIXED,
|
||||
)
|
||||
access_list_mock = AsyncMock(return_value=[])
|
||||
allowed_mock = AsyncMock(return_value=[])
|
||||
with (
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock()
|
||||
|
|
@ -4691,7 +4715,9 @@ class TestMCPUserEnvVarsAccessControl:
|
|||
mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=server)
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_all_mcp_servers_for_user", access_list_mock
|
||||
mgmt_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
allowed_mock,
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_user_env_vars", AsyncMock(return_value={})
|
||||
|
|
@ -4705,22 +4731,34 @@ class TestMCPUserEnvVarsAccessControl:
|
|||
),
|
||||
)
|
||||
assert result.server_id == "srv-1"
|
||||
access_list_mock.assert_not_awaited()
|
||||
allowed_mock.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_admin_gets_403_not_404_for_inaccessible_server(self):
|
||||
"""Authorization must run before the existence check so a non-admin
|
||||
cannot distinguish "server does not exist" (404) from "server exists but
|
||||
you lack access" (403) and enumerate server IDs."""
|
||||
get_mcp_server_mock = AsyncMock(return_value=None)
|
||||
"""A non-admin cannot distinguish "server does not exist" (404) from
|
||||
"server exists but you lack access" (403): both collapse to 403 so server
|
||||
ids stay non-enumerable, even when neither the DB nor the registry has the
|
||||
server."""
|
||||
with (
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock()
|
||||
),
|
||||
patch.object(mgmt_endpoints, "get_mcp_server", get_mcp_server_mock),
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None)
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints,
|
||||
"get_all_mcp_servers_for_user",
|
||||
"build_effective_auth_contexts",
|
||||
AsyncMock(return_value=[object()]),
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints.global_mcp_server_manager,
|
||||
"get_mcp_server_by_id",
|
||||
MagicMock(return_value=None),
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
AsyncMock(return_value=[]),
|
||||
),
|
||||
):
|
||||
|
|
@ -4733,7 +4771,6 @@ class TestMCPUserEnvVarsAccessControl:
|
|||
),
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
get_mcp_server_mock.assert_not_awaited()
|
||||
|
||||
|
||||
def test_oauth2_flow_accepted_on_create_request():
|
||||
|
|
@ -4783,3 +4820,217 @@ def test_oauth2_flow_defaults_to_none_when_omitted():
|
|||
assert (
|
||||
LiteLLM_MCPServerTable(server_id="srv-1", transport="http").oauth2_flow is None
|
||||
)
|
||||
|
||||
|
||||
class TestPerUserCredentialConfigServerResolution:
|
||||
"""Per-user credential and env-var endpoints must resolve config-defined MCP
|
||||
servers, which live only in the in-memory registry and never get a DB row, so
|
||||
a user can store their BYOK key / OAuth token / env vars against them. The
|
||||
same allowed-server authorization the MCP gateway enforces also gates these
|
||||
writes for non-admins.
|
||||
"""
|
||||
|
||||
# 32-char sha256 stable id, the shape a config.yaml server gets.
|
||||
CONFIG_SERVER_ID = "3a6a3f8633340371b49562c8c4682da9"
|
||||
|
||||
def _registry_only_manager(self, *, is_byok: bool = False):
|
||||
"""A manager mock where the server exists only in the registry (DB miss)."""
|
||||
config_server = generate_mock_mcp_server_config_record(
|
||||
server_id=self.CONFIG_SERVER_ID, name="Config Server"
|
||||
)
|
||||
record = generate_mock_mcp_server_db_record(
|
||||
server_id=self.CONFIG_SERVER_ID
|
||||
).model_copy(update={"is_byok": is_byok})
|
||||
manager = MagicMock()
|
||||
manager.get_mcp_server_by_id = MagicMock(
|
||||
side_effect=lambda sid: (
|
||||
config_server if sid == self.CONFIG_SERVER_ID else None
|
||||
)
|
||||
)
|
||||
manager._build_mcp_server_table = MagicMock(return_value=record)
|
||||
manager.get_allowed_mcp_servers = AsyncMock(return_value=[])
|
||||
return manager
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_oauth_credential_resolves_config_server_for_admin(self):
|
||||
"""OBO token persists for a config-defined server (DB miss, registry hit).
|
||||
Before the registry fallback this raised 404 "MCP Server not found"."""
|
||||
manager = self._registry_only_manager()
|
||||
store_mock = AsyncMock(return_value=None)
|
||||
with (
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock()
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None)
|
||||
),
|
||||
patch.object(mgmt_endpoints, "global_mcp_server_manager", manager),
|
||||
patch.object(mgmt_endpoints, "store_user_oauth_credential", store_mock),
|
||||
patch.object(
|
||||
mgmt_endpoints,
|
||||
"get_user_oauth_credential",
|
||||
AsyncMock(return_value={"expires_at": None}),
|
||||
),
|
||||
):
|
||||
result = await mgmt_endpoints.store_mcp_oauth_user_credential(
|
||||
server_id=self.CONFIG_SERVER_ID,
|
||||
payload=mgmt_endpoints.MCPOAuthUserCredentialRequest(
|
||||
access_token="tok", expires_in=3600
|
||||
),
|
||||
user_api_key_dict=generate_mock_user_api_key_auth(user_id="admin"),
|
||||
)
|
||||
assert result.has_credential is True
|
||||
store_mock.assert_awaited_once()
|
||||
manager.get_mcp_server_by_id.assert_called_once_with(self.CONFIG_SERVER_ID)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_byok_credential_resolves_config_server_for_admin(self):
|
||||
"""BYOK key persists for a config-defined BYOK server (DB miss, registry
|
||||
hit). Before the registry fallback this raised 404."""
|
||||
manager = self._registry_only_manager(is_byok=True)
|
||||
store_mock = AsyncMock(return_value=None)
|
||||
with (
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock()
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None)
|
||||
),
|
||||
patch.object(mgmt_endpoints, "global_mcp_server_manager", manager),
|
||||
patch.object(mgmt_endpoints, "store_user_credential", store_mock),
|
||||
):
|
||||
result = await mgmt_endpoints.store_mcp_user_credential(
|
||||
server_id=self.CONFIG_SERVER_ID,
|
||||
payload=mgmt_endpoints.MCPUserCredentialRequest(credential="my-key"),
|
||||
user_api_key_dict=generate_mock_user_api_key_auth(user_id="admin"),
|
||||
)
|
||||
assert result.has_credential is True
|
||||
store_mock.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_oauth_credential_forbidden_for_non_admin_without_access(self):
|
||||
"""A non-admin storing a credential for a server not in their allowed set
|
||||
gets 403 and no row is written (the store endpoints had no authz before)."""
|
||||
manager = self._registry_only_manager()
|
||||
manager.get_allowed_mcp_servers = AsyncMock(return_value=[])
|
||||
store_mock = AsyncMock(return_value=None)
|
||||
with (
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock()
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None)
|
||||
),
|
||||
patch.object(mgmt_endpoints, "global_mcp_server_manager", manager),
|
||||
patch.object(
|
||||
mgmt_endpoints,
|
||||
"build_effective_auth_contexts",
|
||||
AsyncMock(return_value=[object()]),
|
||||
),
|
||||
patch.object(mgmt_endpoints, "store_user_oauth_credential", store_mock),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await mgmt_endpoints.store_mcp_oauth_user_credential(
|
||||
server_id=self.CONFIG_SERVER_ID,
|
||||
payload=mgmt_endpoints.MCPOAuthUserCredentialRequest(
|
||||
access_token="tok", expires_in=3600
|
||||
),
|
||||
user_api_key_dict=generate_mock_user_api_key_auth(
|
||||
user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER
|
||||
),
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
store_mock.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_oauth_credential_allowed_for_non_admin_with_access(self):
|
||||
"""A non-admin with the config server in their allowed set persists the
|
||||
token; proves the non-admin authz uses the registry-aware allowed set."""
|
||||
manager = self._registry_only_manager()
|
||||
manager.get_allowed_mcp_servers = AsyncMock(
|
||||
return_value=[self.CONFIG_SERVER_ID]
|
||||
)
|
||||
store_mock = AsyncMock(return_value=None)
|
||||
with (
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock()
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None)
|
||||
),
|
||||
patch.object(mgmt_endpoints, "global_mcp_server_manager", manager),
|
||||
patch.object(
|
||||
mgmt_endpoints,
|
||||
"build_effective_auth_contexts",
|
||||
AsyncMock(return_value=[object()]),
|
||||
),
|
||||
patch.object(mgmt_endpoints, "store_user_oauth_credential", store_mock),
|
||||
patch.object(
|
||||
mgmt_endpoints,
|
||||
"get_user_oauth_credential",
|
||||
AsyncMock(return_value={"expires_at": None}),
|
||||
),
|
||||
):
|
||||
result = await mgmt_endpoints.store_mcp_oauth_user_credential(
|
||||
server_id=self.CONFIG_SERVER_ID,
|
||||
payload=mgmt_endpoints.MCPOAuthUserCredentialRequest(
|
||||
access_token="tok", expires_in=3600
|
||||
),
|
||||
user_api_key_dict=generate_mock_user_api_key_auth(
|
||||
user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER
|
||||
),
|
||||
)
|
||||
assert result.has_credential is True
|
||||
store_mock.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_env_vars_resolves_config_server_for_non_admin_with_access(
|
||||
self,
|
||||
):
|
||||
"""Per-user env vars persist for a config server a non-admin may access.
|
||||
The non-admin path previously used a DB-only access list that never
|
||||
included config servers, so this 403'd before the fix."""
|
||||
env_var_server = _make_env_var_server(
|
||||
server_id=self.CONFIG_SERVER_ID,
|
||||
env_vars=_ENV_VARS_MIXED,
|
||||
static_headers=_STATIC_HEADERS_MIXED,
|
||||
)
|
||||
manager = MagicMock()
|
||||
manager.get_mcp_server_by_id = MagicMock(
|
||||
return_value=generate_mock_mcp_server_config_record(
|
||||
server_id=self.CONFIG_SERVER_ID
|
||||
)
|
||||
)
|
||||
manager._build_mcp_server_table = MagicMock(return_value=env_var_server)
|
||||
manager.get_allowed_mcp_servers = AsyncMock(
|
||||
return_value=[self.CONFIG_SERVER_ID]
|
||||
)
|
||||
merge_mock = AsyncMock(return_value={"CORP_USERNAME": "alice"})
|
||||
with (
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock()
|
||||
),
|
||||
patch.object(
|
||||
mgmt_endpoints, "get_mcp_server", AsyncMock(return_value=None)
|
||||
),
|
||||
patch.object(mgmt_endpoints, "global_mcp_server_manager", manager),
|
||||
patch.object(
|
||||
mgmt_endpoints,
|
||||
"build_effective_auth_contexts",
|
||||
AsyncMock(return_value=[object()]),
|
||||
),
|
||||
patch.object(mgmt_endpoints, "merge_user_env_vars", merge_mock),
|
||||
):
|
||||
result = await mgmt_endpoints.store_mcp_user_env_vars(
|
||||
server_id=self.CONFIG_SERVER_ID,
|
||||
payload=mgmt_endpoints.MCPUserEnvVarsRequest(
|
||||
values={"CORP_USERNAME": "alice"}
|
||||
),
|
||||
user_api_key_dict=generate_mock_user_api_key_auth(
|
||||
user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER
|
||||
),
|
||||
)
|
||||
merge_mock.assert_awaited_once()
|
||||
_, _, _, updates, _ = merge_mock.await_args.args
|
||||
assert updates == {"CORP_USERNAME": "alice"}
|
||||
assert result.server_id == self.CONFIG_SERVER_ID
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue