From 68d6db45ef9ae7ae1fe11f1ff61a1d3353959bb3 Mon Sep 17 00:00:00 2001 From: shrey Date: Fri, 13 Nov 2026 16:36:49 -0800 Subject: [PATCH] fix(mcp): give the MCP client exactly the servers left on in the connect page (#50973) --- litellm/proxy/_lazy_openapi_snapshot.json | 10 +- .../mcp_server/discoverable_endpoints.py | 12 +- litellm/proxy/mcp_server/gateway_dcr_flow.py | 81 ++- .../chat/ConnectFlowBanner.test.tsx | 30 +- .../src/components/chat/ConnectFlowBanner.tsx | 37 +- .../ConnectFlowSurface.integration.test.tsx | 397 ++++++++++++ .../chat/ConnectFlowSurface.test.tsx | 220 +++++-- .../components/chat/ConnectFlowSurface.tsx | 59 +- .../chat/MCPAppsPanel.integration.test.tsx | 587 ++++++++++++++++++ .../src/components/chat/MCPAppsPanel.test.tsx | 302 --------- .../src/components/chat/MCPAppsPanel.tsx | 204 ++++-- .../src/components/networking.tsx | 2 + .../src/hooks/useUserMcpOAuthFlow.test.tsx | 40 +- .../src/hooks/useUserMcpOAuthFlow.tsx | 11 +- .../src/i18n/locales/en.json | 17 +- .../src/lib/http/schema.d.ts | 6 + .../mcp_server/test_discoverable_endpoints.py | 84 +++ .../proxy/mcp_server/test_gateway_dcr_flow.py | 336 ++++++++++ 18 files changed, 2011 insertions(+), 424 deletions(-) create mode 100644 litellm/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.integration.test.tsx create mode 100644 litellm/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.integration.test.tsx delete mode 100644 litellm/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5cba49bf884..eee6beb83c7 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -39630,6 +39630,14 @@ "title": "Flow", "type": "string" }, + "server": { + "default": [], + "items": { + "type": "string" + }, + "title": "Server", + "type": "array" + }, "team_id": { "anyOf": [ { @@ -42205,7 +42213,7 @@ }, "/authorize/complete": { "post": { - "description": "Finish an aggregate connect flow: mint the gateway authorization code for the\nsigned-in user and hand it back to the DCR client, by 303 redirect (default) or, for\na loopback client on a different machine, as a copyable callback URL\n(``delivery=manual``). POST plus the per-flow HttpOnly cookie set at /authorize; an\nanonymous or bad-flow request just 400s. The native-client consent page adds\n``decision`` (approve or deny) and the ``team_id`` the credential is attributed to.", + "description": "Finish an aggregate connect flow: mint the gateway authorization code for the\nsigned-in user and hand it back to the DCR client, by 303 redirect (default) or, for\na loopback client on a different machine, as a copyable callback URL\n(``delivery=manual``). POST plus the per-flow HttpOnly cookie set at /authorize; an\nanonymous or bad-flow request just 400s. The native-client consent page adds\n``decision`` (approve or deny) and the ``team_id`` the credential is attributed to.\nEach ``server`` field is the id of a server the connect page shows as on.", "operationId": "authorize_complete_authorize_complete_post", "requestBody": { "content": { diff --git a/litellm/proxy/mcp_server/discoverable_endpoints.py b/litellm/proxy/mcp_server/discoverable_endpoints.py index 36a312d7c1f..028855666b2 100644 --- a/litellm/proxy/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/mcp_server/discoverable_endpoints.py @@ -8,6 +8,7 @@ from contextlib import asynccontextmanager from datetime import datetime, timezone from typing import ( TYPE_CHECKING, + Annotated, Final, Literal, Optional, @@ -2614,14 +2615,19 @@ async def authorize_complete( delivery: str | None = Form(None), team_id: str | None = Form(None), decision: str | None = Form(None), + server: Annotated[tuple[str, ...], Form()] = (), ) -> Response: """Finish an aggregate connect flow: mint the gateway authorization code for the signed-in user and hand it back to the DCR client, by 303 redirect (default) or, for a loopback client on a different machine, as a copyable callback URL (``delivery=manual``). POST plus the per-flow HttpOnly cookie set at /authorize; an anonymous or bad-flow request just 400s. The native-client consent page adds - ``decision`` (approve or deny) and the ``team_id`` the credential is attributed to.""" - from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 # circular import at module load + ``decision`` (approve or deny) and the ``team_id`` the credential is attributed to. + Each ``server`` field is the id of a server the connect page shows as on.""" + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load + master_key, + user_api_key_cache, + ) return await complete_connect_flow( request=request, @@ -2633,6 +2639,8 @@ async def authorize_complete( decision=decision, lookup_vendor_credential=_vendor_credential_state, lookup_server_reachability=_user_can_reach_mcp_server, + selected_server_ids=server, + master_key=master_key, ) diff --git a/litellm/proxy/mcp_server/gateway_dcr_flow.py b/litellm/proxy/mcp_server/gateway_dcr_flow.py index ae74d878b36..2ed08da4bc7 100644 --- a/litellm/proxy/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/mcp_server/gateway_dcr_flow.py @@ -411,6 +411,7 @@ class _ConnectFlow(_ServerScopedSealedModel): audience: SessionAudience | None = None requested_scopes: tuple[str, ...] = () upstream_scopes: tuple[UpstreamScope, ...] = () + unresolved_resource: Literal[True] | None = None class _GatewayAuthCode(_ServerScopedSealedModel): @@ -756,6 +757,9 @@ def aggregate_authorize( if resource_scoped_server is not None and not upstream_scopes else (), upstream_scopes=upstream_scopes, + unresolved_resource=True + if resource is not None and resource_scoped_server is None and not _is_aggregate_resource(request, resource) + else None, ) connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),)) response: Final = RedirectResponse(connect_url, status_code=303) @@ -919,6 +923,7 @@ def _new_connect_flow( audience: SessionAudience | None, requested_scopes: tuple[str, ...] = (), upstream_scopes: tuple[UpstreamScope, ...] = (), + unresolved_resource: Literal[True] | None = None, ) -> _ConnectFlow: now: Final = datetime.now(timezone.utc) return _ConnectFlow( @@ -933,6 +938,7 @@ def _new_connect_flow( audience=audience, requested_scopes=requested_scopes, upstream_scopes=upstream_scopes, + unresolved_resource=unresolved_resource, ) @@ -1028,6 +1034,8 @@ class ConnectFlowDescription(TypedDict): connected: ReadOnly[bool | None] servers: ReadOnly[Sequence[ConnectFlowServer] | None] requested_scopes: NotRequired[ReadOnly[tuple[str, ...]]] + selectable: NotRequired[ReadOnly[Literal[False]]] + required_server_ids: NotRequired[ReadOnly[tuple[str, ...]]] async def _describe_connect_flow_server( @@ -1174,6 +1182,17 @@ async def _describe_upstream_scope_flow( return descriptions[0] +def _with_selection_limits(flow: _ConnectFlow, described: ConnectFlowDescription) -> ConnectFlowDescription: + if flow.unresolved_resource: + unselectable: Final[ConnectFlowDescription] = {**described, "selectable": False} + return unselectable + if not flow.upstream_scopes: + return described + required: Final = tuple(dict.fromkeys(item.server_id for item in flow.upstream_scopes)) + with_required: Final[ConnectFlowDescription] = {**described, "required_server_ids": required} + return with_required + + async def describe_connect_flow( request: Request, flow_handle: str, @@ -1188,7 +1207,7 @@ async def describe_connect_flow( return ( described if isinstance(described, Response) - else JSONResponse(content=described, headers=TOKEN_NO_CACHE_HEADERS) + else JSONResponse(content=_with_selection_limits(opened, described), headers=TOKEN_NO_CACHE_HEADERS) ) @@ -1202,6 +1221,8 @@ async def complete_connect_flow( decision: str | None = None, lookup_vendor_credential: LookupVendorCredential = _unavailable_vendor_credential, lookup_server_reachability: LookupServerReachability = _unreachable_server, + selected_server_ids: Sequence[str] = (), + master_key: str | None = None, ) -> Response: """Mint the code only after a deliberate POST by the sealed user. @@ -1227,6 +1248,13 @@ async def complete_connect_flow( return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing") if described["connected"] is False: return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing") + narrowed: Final = ( + opened + if decision == "deny" + else await _narrowed_to_selection(opened, selected_server_ids, lookup_server_reachability, master_key) + ) + if isinstance(narrowed, Response): + return narrowed flow_refusal: Final = _claim_refusal( await _SingleUseGuard(cache).claim( f"{_USED_FLOW_CACHE_PREFIX}{opened.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS @@ -1238,13 +1266,62 @@ async def complete_connect_flow( if flow_refusal is not None: return flow_refusal response: Final = ( - _denied_flow_response(opened) if decision == "deny" else _approved_flow_response(opened, delivery, team_id, now) + _denied_flow_response(narrowed) + if decision == "deny" + else _approved_flow_response(narrowed, delivery, team_id, now) ) path, secure = _cookie_path_and_secure(request) response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax") return response +async def _narrowed_to_selection( + flow: _ConnectFlow, + selected_server_ids: Sequence[str], + lookup_server_reachability: LookupServerReachability, + master_key: str | None, +) -> _ConnectFlow | Response: + if ( + not any(server_id.strip() for server_id in selected_server_ids) + or flow.resource_server_ids is not None + or flow.audience is not None + or flow.unresolved_resource + ): + return flow + from litellm.proxy.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # import cycle + global_mcp_server_manager, + ) + + listed_ids: Final = frozenset(server.server_id for server in global_mcp_server_manager.get_registry().values()) + known_ids: Final = tuple(dict.fromkeys(server_id for server_id in selected_server_ids if server_id in listed_ids)) + if not known_ids: + return _oauth_error(400, "invalid_request", "the requested MCP server is no longer available") + reachable_ids: Final = tuple( + [server_id for server_id in known_ids if await lookup_server_reachability(flow.user_id, server_id)] + ) + if not reachable_ids: + return _oauth_error(400, "invalid_request", "you do not have access to the requested MCP server") + upstream_ids: Final = tuple(item.server_id for item in flow.upstream_scopes) + granted: Final = tuple(dict.fromkeys((*upstream_ids, *reachable_ids))) + if not _session_tokens_fit(flow, granted, master_key): + return _oauth_error( + 400, "invalid_request", "too many MCP servers are turned on for one connection, turn some off and try again" + ) + return flow.model_copy(update={"resource_server_ids": granted}) + + +def _session_tokens_fit(flow: _ConnectFlow, server_ids: tuple[str, ...], master_key: str | None) -> bool: + """A grant that cannot be minted would fail at /token after the code is spent.""" + keys: Final = None if master_key is None else active_session_signing_keys(master_key) + if not isinstance(keys, SessionSigningKeys): + return True + principal: Final = SessionPrincipal(user_id=flow.user_id, client_id=flow.client_id, resource_server_ids=server_ids) + now: Final = datetime.now(timezone.utc) + access: Final = mint_session_token(principal, keys, now) + rotated_refresh: Final = mint_session_refresh_token(principal, keys, now, family=secrets.token_urlsafe(16)) + return isinstance(access, MintedSessionToken) and isinstance(rotated_refresh, MintedSessionToken) + + def _state_param(state: str) -> tuple[tuple[str, str], ...]: return (("state", state),) if state else () diff --git a/litellm/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx b/litellm/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx index d3a24d7b9a5..79f5ef64d2d 100644 --- a/litellm/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx +++ b/litellm/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx @@ -1,7 +1,7 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { render, screen } from "@testing-library/react"; import type { ConnectFlowStatus } from "@/components/networking"; -import ConnectFlowBanner, { isLoopbackOrigin } from "./ConnectFlowBanner"; +import ConnectFlowBanner, { isLoopbackOrigin, type FinishSelection } from "./ConnectFlowBanner"; const { startOAuthFlow } = vi.hoisted(() => ({ startOAuthFlow: vi.fn() })); @@ -27,7 +27,7 @@ const unscoped = (client_origin: string): ConnectFlowStatus => ({ servers: null, }); -const renderBanner = (clientOrigin: string) => +const renderBanner = (clientOrigin: string, finish?: FinishSelection) => render( accessToken="tok" onConnected={vi.fn()} failed={false} + finish={finish} />, ); @@ -51,6 +52,31 @@ describe("ConnectFlowBanner", () => { expect(screen.queryByRole("button", { name: "Cancel" })).not.toBeInTheDocument(); }); + it("posts the id of each server that is on with Finish", () => { + renderBanner("https://claude.ai", { status: "ready", serverIds: ["catalog-id", "orders-id"] }); + + expect(screen.getByDisplayValue("catalog-id")).toHaveAttribute("name", "server"); + expect(screen.getByDisplayValue("orders-id")).toHaveAttribute("name", "server"); + expect(screen.getByRole("button", { name: /finish connecting/i })).toBeEnabled(); + }); + + it("holds Finish until the page has loaded its servers", () => { + renderBanner("https://claude.ai", { status: "waiting" }); + + expect(screen.getByRole("button", { name: /finish connecting/i })).toBeDisabled(); + expect(screen.queryByText("Turn on at least one server to finish connecting.")).not.toBeInTheDocument(); + }); + + it("holds Finish and says why when every server is off", () => { + renderBanner("https://claude.ai", { status: "nothing-on" }); + + expect(screen.getByRole("button", { name: /finish connecting/i })).toBeDisabled(); + expect(screen.getByRole("button", { name: /finish connecting/i })).toHaveAccessibleDescription( + "Turn on at least one server to finish connecting.", + ); + expect(screen.queryByDisplayValue(/./, { exact: false })).toHaveAttribute("name", "flow"); + }); + it("lists multi-server consent and requires every server before finishing", () => { const flow: ConnectFlowStatus = { state: "multi", diff --git a/litellm/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx b/litellm/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx index 36cf048328b..f29e2411702 100644 --- a/litellm/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx +++ b/litellm/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx @@ -6,6 +6,7 @@ import type { TFunction } from "i18next"; import { useTranslation } from "react-i18next"; import { getProxyBaseUrl, ConnectFlowStatus } from "@/components/networking"; import { OAuth2ConnectButton } from "@/components/chat/MCPAppsPanel"; +import { Button } from "@/components/ui/button"; import i18n from "@/i18n/i18n"; interface Props { @@ -14,8 +15,16 @@ interface Props { accessToken: string; onConnected: () => void; failed: boolean; + finish?: FinishSelection; } +export type FinishSelection = + | { status: "waiting" } + | { status: "nothing-on" } + | { status: "ready"; serverIds: readonly string[] }; + +const FINISH_WITHOUT_SERVERS: FinishSelection = { status: "ready", serverIds: [] }; + /** Finish remains an explicit POST because a cross-site navigation must never mint a code. */ export function isLoopbackOrigin(origin: string | null): boolean { if (!origin) return false; @@ -44,7 +53,9 @@ const copyFor = ( if (flow.state === "unscoped") { return [ t("chat.connectFlowBanner.unscoped.title", { client: clientLabel }), - t("chat.connectFlowBanner.unscoped.description", { client: clientLabel }), + flow.selectable === false + ? t("chat.connectFlowBanner.unscoped.fixed", { client: clientLabel }) + : t("chat.connectFlowBanner.unscoped.description", { client: clientLabel }), ]; } if (flow.state === "multi") { @@ -89,7 +100,14 @@ const vendorServersFor = ( return [{ server_id: flow.server_id, server_name: flow.server_name }]; }; -const ConnectFlowBanner: React.FC = ({ flowHandle, flow, accessToken, onConnected, failed }) => { +const ConnectFlowBanner: React.FC = ({ + flowHandle, + flow, + accessToken, + onConnected, + failed, + finish = FINISH_WITHOUT_SERVERS, +}) => { const { t } = useTranslation(); const action = `${getProxyBaseUrl()}/authorize/complete`; const state = failed || flow === undefined ? "stale" : flow.state; @@ -130,13 +148,17 @@ const ConnectFlowBanner: React.FC = ({ flowHandle, flow, accessToken, onC ))}
+ {finish.status === "ready" && + finish.serverIds.map((serverId) => )} {canFinish && ( - + )} {canCancel && (