mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(mcp): give the MCP client exactly the servers left on in the connect page (#50973)
This commit is contained in:
shrey
GitHub
parent
c4a72b8fac
commit
68d6db45ef
18 changed files with 2011 additions and 424 deletions
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 ()
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
<ConnectFlowBanner
|
||||
flowHandle="flow-handle-123"
|
||||
|
|
@ -35,6 +35,7 @@ const renderBanner = (clientOrigin: string) =>
|
|||
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",
|
||||
|
|
|
|||
|
|
@ -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<Props> = ({ flowHandle, flow, accessToken, onConnected, failed }) => {
|
||||
const ConnectFlowBanner: React.FC<Props> = ({
|
||||
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<Props> = ({ flowHandle, flow, accessToken, onC
|
|||
))}
|
||||
<form method="POST" action={action}>
|
||||
<input type="hidden" name="flow" value={flowHandle} />
|
||||
{finish.status === "ready" &&
|
||||
finish.serverIds.map((serverId) => <input key={serverId} type="hidden" name="server" value={serverId} />)}
|
||||
{canFinish && (
|
||||
<button
|
||||
<Button
|
||||
type="submit"
|
||||
className="h-[38px] rounded-md bg-primary px-4 text-sm font-semibold text-primary-foreground hover:bg-primary/90"
|
||||
disabled={finish.status !== "ready"}
|
||||
aria-describedby={finish.status === "nothing-on" ? "connect-flow-nothing-on" : undefined}
|
||||
className="h-[38px] px-4 font-semibold"
|
||||
>
|
||||
{t("chat.connectFlowBanner.finish")}
|
||||
</button>
|
||||
</Button>
|
||||
)}
|
||||
{canCancel && (
|
||||
<button
|
||||
|
|
@ -154,6 +176,11 @@ const ConnectFlowBanner: React.FC<Props> = ({ flowHandle, flow, accessToken, onC
|
|||
{t("chat.connectFlowBanner.manualDelivery")}
|
||||
</label>
|
||||
)}
|
||||
{finish.status === "nothing-on" && (
|
||||
<p id="connect-flow-nothing-on" className="mt-2 text-[13px] text-muted-foreground">
|
||||
{t("chat.connectFlowBanner.nothingOn")}
|
||||
</p>
|
||||
)}
|
||||
</form>
|
||||
</div>
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,397 @@
|
|||
import { webcrypto } from "node:crypto";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import ConnectFlowSurface from "./ConnectFlowSurface";
|
||||
import McpOAuthCallbackPage from "@/app/mcp/oauth/callback/page";
|
||||
import {
|
||||
exchangeMcpOAuthToken,
|
||||
fetchConnectFlow,
|
||||
fetchMCPServers,
|
||||
getMCPOAuthUserCredentialStatus,
|
||||
listMCPTools,
|
||||
registerMcpOAuthClient,
|
||||
storeMCPOAuthUserCredential,
|
||||
type ConnectFlowStatus,
|
||||
} from "@/components/networking";
|
||||
import type { MCPServer } from "@/components/mcp_tools/types";
|
||||
|
||||
const query = vi.hoisted(() => ({ params: {} as Record<string, string | null> }));
|
||||
|
||||
vi.mock("next/navigation", () => ({
|
||||
useRouter: () => ({ replace: vi.fn() }),
|
||||
useSearchParams: () => ({ get: (key: string) => query.params[key] ?? null }),
|
||||
}));
|
||||
vi.mock("@/components/networking", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@/components/networking")>()),
|
||||
fetchConnectFlow: vi.fn(),
|
||||
fetchMCPServers: vi.fn(),
|
||||
listMCPTools: vi.fn(),
|
||||
getMCPOAuthUserCredentialStatus: vi.fn(),
|
||||
registerMcpOAuthClient: vi.fn(),
|
||||
exchangeMcpOAuthToken: vi.fn(),
|
||||
storeMCPOAuthUserCredential: vi.fn(),
|
||||
}));
|
||||
|
||||
const CONNECT_PAGE = "https://gateway.example.com/ui/connect?connect_flow=flow-handle-123";
|
||||
const realLocation = Object.getOwnPropertyDescriptor(window, "location");
|
||||
|
||||
const browser = { visits: [CONNECT_PAGE] };
|
||||
|
||||
const stubBrowserLocation = () => {
|
||||
browser.visits = [CONNECT_PAGE];
|
||||
const location = {
|
||||
origin: "https://gateway.example.com",
|
||||
pathname: "/ui/connect",
|
||||
get href() {
|
||||
return browser.visits[browser.visits.length - 1];
|
||||
},
|
||||
set href(url: string) {
|
||||
browser.visits.push(url);
|
||||
},
|
||||
replace: (url: string) => {
|
||||
browser.visits.push(url);
|
||||
},
|
||||
};
|
||||
Object.defineProperty(window, "location", { configurable: true, value: location });
|
||||
};
|
||||
|
||||
const vendorVisits = () => browser.visits.filter((url) => url.includes("/v1/mcp/server/oauth/"));
|
||||
|
||||
const servers = [
|
||||
{ server_id: "catalog-id", server_name: "catalog", auth_type: "none", connected_app_reachable: true },
|
||||
{ server_id: "orders-id", server_name: "orders", auth_type: "none", connected_app_reachable: true },
|
||||
{ server_id: "tracker-id", server_name: "tracker", auth_type: "none", connected_app_reachable: true },
|
||||
] as MCPServer[];
|
||||
|
||||
const vault = {
|
||||
server_id: "vault-id",
|
||||
server_name: "vault",
|
||||
auth_type: "oauth2",
|
||||
oauth2_flow: "authorization_code",
|
||||
connected_app_reachable: true,
|
||||
} as MCPServer;
|
||||
|
||||
const flow = (kind: ConnectFlowStatus["state"], connected: boolean | null = null): ConnectFlowStatus => ({
|
||||
state: kind,
|
||||
client_origin: "https://claude.ai",
|
||||
server_id: kind === "interactive" || kind === "m2m" ? "s-design" : null,
|
||||
server_name: kind === "interactive" || kind === "m2m" ? "design_tool" : null,
|
||||
connected,
|
||||
servers: null,
|
||||
});
|
||||
|
||||
const deferred = <T,>() => {
|
||||
let resolve: (value: T) => void = () => {};
|
||||
let reject: (reason: Error) => void = () => {};
|
||||
const promise = new Promise<T>((res, rej) => {
|
||||
resolve = res;
|
||||
reject = rej;
|
||||
});
|
||||
return { promise, resolve: (value: T) => resolve(value), reject: (reason: Error) => reject(reason) };
|
||||
};
|
||||
|
||||
const renderSurface = () => {
|
||||
const client = new QueryClient({ defaultOptions: { queries: { retry: false } } });
|
||||
const tree = (
|
||||
<QueryClientProvider client={client}>
|
||||
<ConnectFlowSurface accessToken="token-123" selectedServers={[]} onChange={vi.fn()} />
|
||||
</QueryClientProvider>
|
||||
);
|
||||
const view = render(tree);
|
||||
return { unmount: view.unmount, flushPendingRenders: () => view.rerender(tree) };
|
||||
};
|
||||
|
||||
const openConnectPage = (result: ConnectFlowStatus = flow("unscoped")) => {
|
||||
vi.mocked(fetchConnectFlow).mockResolvedValue(result);
|
||||
return renderSurface();
|
||||
};
|
||||
|
||||
const returnFromVendor = async (page: { unmount: () => void }) => {
|
||||
await waitFor(() => expect(window.location.href).toContain("/authorize"));
|
||||
const vendorState = new URL(window.location.href).searchParams.get("state");
|
||||
page.unmount();
|
||||
query.params = { code: "vendor-code", state: vendorState };
|
||||
const callback = render(<McpOAuthCallbackPage />);
|
||||
await waitFor(() => expect(window.location.href).toContain("mcpOauthReturn=apps"));
|
||||
callback.unmount();
|
||||
query.params = { connect_flow: "flow-handle-123", mcpOauthReturn: "apps" };
|
||||
return renderSurface();
|
||||
};
|
||||
|
||||
const serverSwitch = (name: string) => screen.findByRole("switch", { name: `Include ${name} in this connection` });
|
||||
const finishButton = () => screen.getByRole("button", { name: /finish connecting/i });
|
||||
const serverIdsFinishSends = () =>
|
||||
screen
|
||||
.queryAllByDisplayValue(/./)
|
||||
.filter((field) => field.getAttribute("name") === "server")
|
||||
.map((field) => field.getAttribute("value"));
|
||||
|
||||
beforeEach(() => {
|
||||
query.params = { connect_flow: "flow-handle-123" };
|
||||
vi.stubGlobal("crypto", webcrypto);
|
||||
stubBrowserLocation();
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(servers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
vi.mocked(registerMcpOAuthClient).mockResolvedValue({ client_id: "registered-client" });
|
||||
vi.mocked(exchangeMcpOAuthToken).mockResolvedValue({ access_token: "vendor-token" });
|
||||
vi.mocked(getMCPOAuthUserCredentialStatus).mockImplementation(async (_token, serverId) => ({
|
||||
server_id: serverId,
|
||||
has_credential: false,
|
||||
is_expired: false,
|
||||
}));
|
||||
vi.mocked(storeMCPOAuthUserCredential).mockImplementation(async (_token, serverId) => ({
|
||||
server_id: serverId,
|
||||
has_credential: true,
|
||||
is_expired: false,
|
||||
}));
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
query.params = {};
|
||||
sessionStorage.clear();
|
||||
vi.unstubAllGlobals();
|
||||
if (realLocation) Object.defineProperty(window, "location", realLocation);
|
||||
vi.resetAllMocks();
|
||||
});
|
||||
|
||||
describe("the connect page as one tree", () => {
|
||||
it.each([
|
||||
{ result: flow("unscoped"), grid: true, finish: true, cancel: false },
|
||||
{ result: flow("interactive", true), grid: false, finish: true, cancel: true },
|
||||
{ result: flow("m2m", true), grid: false, finish: true, cancel: true },
|
||||
{ result: flow("stale"), grid: false, finish: false, cancel: true },
|
||||
])("renders $result.state without widening its action surface", async ({ result, grid, finish, cancel }) => {
|
||||
openConnectPage(result);
|
||||
|
||||
await screen.findAllByText(/https:\/\/claude\.ai/);
|
||||
expect(screen.queryByRole("tablist") !== null).toBe(grid);
|
||||
expect(screen.queryByRole("button", { name: /finish connecting/i }) !== null).toBe(finish);
|
||||
expect(screen.queryByRole("button", { name: "Cancel" }) !== null).toBe(cancel);
|
||||
expect(vendorVisits()).toEqual([]);
|
||||
});
|
||||
|
||||
it.each([flow("interactive", true), flow("m2m", true)])(
|
||||
"leaves Finish open on a $state flow, which has no panel to wait for",
|
||||
async (result) => {
|
||||
openConnectPage(result);
|
||||
|
||||
expect(await screen.findByRole("button", { name: /finish connecting/i })).toBeEnabled();
|
||||
},
|
||||
);
|
||||
|
||||
it("keeps the grid and Finish hidden until the gateway accepts a handle", () => {
|
||||
vi.mocked(fetchConnectFlow).mockReturnValue(new Promise(() => {}));
|
||||
renderSurface();
|
||||
|
||||
expect(screen.queryByRole("tablist")).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Cancel" })).toHaveAttribute("value", "deny");
|
||||
});
|
||||
|
||||
it("keeps the grid and Finish hidden when flow validation fails", async () => {
|
||||
query.params = { connect_flow: "invalid-handle" };
|
||||
vi.mocked(fetchConnectFlow).mockRejectedValue(new Error("invalid flow"));
|
||||
renderSurface();
|
||||
|
||||
await screen.findByRole("button", { name: "Cancel" });
|
||||
expect(screen.queryByRole("tablist")).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("sends the user to the vendor once on an interactive flow that is not signed in", async () => {
|
||||
openConnectPage(flow("interactive", false));
|
||||
|
||||
await waitFor(() => expect(window.location.href).toContain("/v1/mcp/server/oauth/s-design/authorize"));
|
||||
expect(vendorVisits()).toHaveLength(1);
|
||||
expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument();
|
||||
expect(screen.queryByRole("tablist")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows Finish once a login returning from the vendor is saved", async () => {
|
||||
let saved = false;
|
||||
vi.mocked(fetchConnectFlow).mockImplementation(async () => flow("interactive", saved));
|
||||
vi.mocked(storeMCPOAuthUserCredential).mockImplementation(async (_token, serverId) => {
|
||||
saved = true;
|
||||
return { server_id: serverId, has_credential: true, is_expired: false };
|
||||
});
|
||||
|
||||
await returnFromVendor(renderSurface());
|
||||
|
||||
expect(await screen.findByRole("button", { name: /finish connecting/i })).toBeEnabled();
|
||||
expect(exchangeMcpOAuthToken).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ serverId: "s-design", code: "vendor-code" }),
|
||||
);
|
||||
expect(storeMCPOAuthUserCredential).toHaveBeenCalledWith(
|
||||
"token-123",
|
||||
"s-design",
|
||||
expect.objectContaining({ access_token: "vendor-token" }),
|
||||
);
|
||||
});
|
||||
|
||||
it("offers Connect again, without Finish, when a login returning from the vendor cannot be saved", async () => {
|
||||
vi.mocked(fetchConnectFlow).mockResolvedValue(flow("interactive", false));
|
||||
const save = deferred<never>();
|
||||
vi.mocked(storeMCPOAuthUserCredential).mockReturnValue(save.promise);
|
||||
|
||||
await returnFromVendor(renderSurface());
|
||||
await screen.findByRole("button", { name: "Connecting…" });
|
||||
await act(async () => {
|
||||
save.reject(new Error("save failed"));
|
||||
});
|
||||
|
||||
expect(await screen.findByRole("button", { name: "Connect" })).toBeEnabled();
|
||||
expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("holds Finish while a login returning from the vendor is being saved", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue([...servers, vault]);
|
||||
const exchange = deferred<{ access_token: string }>();
|
||||
vi.mocked(exchangeMcpOAuthToken).mockReturnValue(exchange.promise);
|
||||
const page = openConnectPage();
|
||||
fireEvent.click(await screen.findByText("Connect", { exact: true }));
|
||||
|
||||
await returnFromVendor(page);
|
||||
|
||||
expect(await screen.findByRole("button", { name: /finish connecting/i })).toBeDisabled();
|
||||
expect(serverIdsFinishSends()).toEqual([]);
|
||||
|
||||
await act(async () => {
|
||||
exchange.resolve({ access_token: "vendor-token" });
|
||||
});
|
||||
|
||||
await waitFor(() => expect(finishButton()).toBeEnabled());
|
||||
expect(storeMCPOAuthUserCredential).toHaveBeenCalledWith(
|
||||
"token-123",
|
||||
"vault-id",
|
||||
expect.objectContaining({ access_token: "vendor-token" }),
|
||||
);
|
||||
expect(serverIdsFinishSends()).toEqual([]);
|
||||
});
|
||||
|
||||
it("sends only the flow while every switch is on, so the session stays as on main", async () => {
|
||||
openConnectPage();
|
||||
|
||||
expect(await serverSwitch("catalog")).toBeChecked();
|
||||
expect(await serverSwitch("orders")).toBeChecked();
|
||||
expect(await serverSwitch("tracker")).toBeChecked();
|
||||
expect(finishButton()).toBeEnabled();
|
||||
expect(screen.getByDisplayValue("flow-handle-123")).toHaveAttribute("name", "flow");
|
||||
expect(serverIdsFinishSends()).toEqual([]);
|
||||
});
|
||||
|
||||
it("sends the servers still on once one is switched off, and only the flow once it is back on", async () => {
|
||||
openConnectPage();
|
||||
const orders = await serverSwitch("orders");
|
||||
|
||||
fireEvent.click(orders);
|
||||
|
||||
expect(orders).not.toBeChecked();
|
||||
expect(serverIdsFinishSends()).toEqual(["catalog-id", "tracker-id"]);
|
||||
expect(finishButton()).toBeEnabled();
|
||||
|
||||
fireEvent.click(orders);
|
||||
|
||||
expect(serverIdsFinishSends()).toEqual([]);
|
||||
});
|
||||
|
||||
it("holds Finish and says why once every switch is off", async () => {
|
||||
openConnectPage();
|
||||
|
||||
fireEvent.click(await serverSwitch("catalog"));
|
||||
fireEvent.click(await serverSwitch("orders"));
|
||||
fireEvent.click(await serverSwitch("tracker"));
|
||||
|
||||
expect(finishButton()).toBeDisabled();
|
||||
expect(finishButton()).toHaveAccessibleDescription("Turn on at least one server to finish connecting.");
|
||||
expect(serverIdsFinishSends()).toEqual([]);
|
||||
});
|
||||
|
||||
it("keeps a server the client requires on, locked and sent", async () => {
|
||||
openConnectPage({ ...flow("unscoped"), required_server_ids: ["tracker-id"] });
|
||||
const tracker = await serverSwitch("tracker");
|
||||
|
||||
fireEvent.click(tracker);
|
||||
fireEvent.click(await serverSwitch("orders"));
|
||||
|
||||
expect(tracker).toBeChecked();
|
||||
expect(screen.getByText("Required by this client")).toBeInTheDocument();
|
||||
expect(serverIdsFinishSends()).toEqual(["catalog-id", "tracker-id"]);
|
||||
});
|
||||
|
||||
it("keeps every switch on and sends only the flow when the flow cannot be narrowed", async () => {
|
||||
openConnectPage({ ...flow("unscoped"), selectable: false });
|
||||
const orders = await serverSwitch("orders");
|
||||
|
||||
fireEvent.click(orders);
|
||||
|
||||
expect(orders).toBeChecked();
|
||||
expect(screen.getByText(/https:\/\/claude\.ai will get every server you can reach/)).toBeInTheDocument();
|
||||
expect(finishButton()).toBeEnabled();
|
||||
expect(serverIdsFinishSends()).toEqual([]);
|
||||
});
|
||||
|
||||
it("keeps a switch off across a reload and holds Finish until the server list is back", async () => {
|
||||
const { unmount } = openConnectPage();
|
||||
fireEvent.click(await serverSwitch("orders"));
|
||||
unmount();
|
||||
const list = deferred<MCPServer[]>();
|
||||
vi.mocked(fetchMCPServers).mockReturnValue(list.promise);
|
||||
|
||||
openConnectPage();
|
||||
|
||||
expect(await screen.findByRole("button", { name: /finish connecting/i })).toBeDisabled();
|
||||
expect(serverIdsFinishSends()).toEqual([]);
|
||||
|
||||
await act(async () => {
|
||||
list.resolve(servers);
|
||||
});
|
||||
|
||||
expect(await serverSwitch("orders")).not.toBeChecked();
|
||||
await waitFor(() => expect(serverIdsFinishSends()).toEqual(["catalog-id", "tracker-id"]));
|
||||
});
|
||||
|
||||
it("holds Finish after a reload with a switch off when the server list fails to load", async () => {
|
||||
const { unmount } = openConnectPage();
|
||||
fireEvent.click(await serverSwitch("orders"));
|
||||
unmount();
|
||||
vi.mocked(fetchMCPServers).mockRejectedValue(new Error("list failed"));
|
||||
|
||||
const page = openConnectPage();
|
||||
|
||||
expect(await screen.findByText(/Could not load your MCP servers/)).toBeInTheDocument();
|
||||
page.flushPendingRenders();
|
||||
expect(finishButton()).toBeDisabled();
|
||||
expect(serverIdsFinishSends()).toEqual([]);
|
||||
});
|
||||
|
||||
it("lets the user finish as on main while the server list is still loading", async () => {
|
||||
vi.mocked(fetchMCPServers).mockReturnValue(new Promise(() => {}));
|
||||
openConnectPage();
|
||||
|
||||
expect(await screen.findByRole("button", { name: /finish connecting/i })).toBeEnabled();
|
||||
expect(serverIdsFinishSends()).toEqual([]);
|
||||
});
|
||||
|
||||
it.each([
|
||||
{
|
||||
when: "fails to load",
|
||||
list: () => Promise.reject<MCPServer[]>(new Error("list failed")),
|
||||
shown: /Could not load your MCP servers/,
|
||||
},
|
||||
{
|
||||
when: "comes back empty",
|
||||
list: () => Promise.resolve<MCPServer[]>([]),
|
||||
shown: /No MCP servers are available to this connection yet/,
|
||||
},
|
||||
])("lets the user finish as on main when the server list $when", async ({ list, shown }) => {
|
||||
vi.mocked(fetchMCPServers).mockImplementation(list);
|
||||
const page = openConnectPage();
|
||||
|
||||
expect(await screen.findByText(shown)).toBeInTheDocument();
|
||||
page.flushPendingRenders();
|
||||
expect(finishButton()).toBeEnabled();
|
||||
expect(serverIdsFinishSends()).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
|
@ -2,11 +2,23 @@ import { afterEach, describe, expect, it, vi } from "vitest";
|
|||
import { act, render, screen, waitFor } from "@testing-library/react";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import ConnectFlowSurface from "./ConnectFlowSurface";
|
||||
import { fetchConnectFlow } from "@/components/networking";
|
||||
import type { FinishSelection } from "./ConnectFlowBanner";
|
||||
import type { LockedOn, ServerListReport } from "./MCPAppsPanel";
|
||||
import { fetchConnectFlow, type ConnectFlowStatus } from "@/components/networking";
|
||||
import { pendingUserMcpOAuthServerId } from "@/hooks/useUserMcpOAuthFlow";
|
||||
|
||||
const { startOAuthFlow, state, onSuccess } = vi.hoisted(() => ({
|
||||
startOAuthFlow: vi.fn(),
|
||||
onSuccess: { current: undefined as (() => void) | undefined },
|
||||
const { state, panel, banner } = vi.hoisted(() => ({
|
||||
panel: {
|
||||
offServerIds: undefined as readonly string[] | undefined,
|
||||
lockedOn: undefined as LockedOn,
|
||||
onOffServerIdsChange: undefined as ((serverIds: string[]) => void) | undefined,
|
||||
onConnectedServers: undefined as ((servers: ServerListReport) => void) | undefined,
|
||||
},
|
||||
banner: {
|
||||
flow: undefined as ConnectFlowStatus | undefined,
|
||||
failed: undefined as boolean | undefined,
|
||||
finish: undefined as FinishSelection | undefined,
|
||||
},
|
||||
state: { oauthReturn: null as string | null, connectFlow: null as string | null },
|
||||
}));
|
||||
|
||||
|
|
@ -16,110 +28,190 @@ vi.mock("next/navigation", () => ({
|
|||
get: (key: string) => ({ mcpOauthReturn: state.oauthReturn, connect_flow: state.connectFlow })[key] ?? null,
|
||||
}),
|
||||
}));
|
||||
vi.mock("@/components/networking", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@/components/networking")>()),
|
||||
fetchConnectFlow: vi.fn(),
|
||||
getProxyBaseUrl: () => "https://gateway.example.com",
|
||||
vi.mock("@/components/networking", () => ({ fetchConnectFlow: vi.fn() }));
|
||||
vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({ pendingUserMcpOAuthServerId: vi.fn(() => null) }));
|
||||
vi.mock("@/components/chat/MCPAppsPanel", () => ({
|
||||
default: ({ offServerIds, lockedOn, onOffServerIdsChange, onConnectedServers }: typeof panel) => {
|
||||
panel.offServerIds = offServerIds;
|
||||
panel.lockedOn = lockedOn;
|
||||
panel.onOffServerIdsChange = onOffServerIdsChange;
|
||||
panel.onConnectedServers = onConnectedServers;
|
||||
return <div data-testid="mcp-apps-panel" />;
|
||||
},
|
||||
}));
|
||||
vi.mock("@/components/chat/MCPAppsPanel", async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import("@/components/chat/MCPAppsPanel")>()),
|
||||
default: () => <div data-testid="mcp-apps-panel" />,
|
||||
}));
|
||||
vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({
|
||||
useUserMcpOAuthFlow: ({ onSuccess: success }: { onSuccess: () => void }) => {
|
||||
onSuccess.current = success;
|
||||
return { startOAuthFlow, status: "idle" };
|
||||
vi.mock("@/components/chat/ConnectFlowBanner", () => ({
|
||||
default: ({ flow, failed, finish }: typeof banner) => {
|
||||
banner.flow = flow;
|
||||
banner.failed = failed;
|
||||
banner.finish = finish;
|
||||
return <div data-testid="connect-flow-banner" />;
|
||||
},
|
||||
}));
|
||||
|
||||
const flow = (state: "unscoped" | "interactive" | "m2m" | "multi" | "stale", connected: boolean | null = null) => ({
|
||||
const flow = (state: ConnectFlowStatus["state"], connected: boolean | null = null): ConnectFlowStatus => ({
|
||||
state,
|
||||
client_origin: "https://claude.ai",
|
||||
server_id: state === "interactive" || state === "m2m" ? "s-design" : null,
|
||||
server_name: state === "interactive" || state === "m2m" ? "design_tool" : null,
|
||||
connected,
|
||||
servers:
|
||||
state === "multi"
|
||||
? [
|
||||
{ server_id: "s-alpha", server_name: "alpha", connected: false },
|
||||
{ server_id: "s-beta", server_name: "beta", connected: false },
|
||||
]
|
||||
: null,
|
||||
servers: null,
|
||||
});
|
||||
|
||||
const renderSurface = () =>
|
||||
const everyServer: FinishSelection = { status: "ready", serverIds: [] };
|
||||
const waiting: FinishSelection = { status: "waiting" };
|
||||
const finishHanded = () => banner.finish ?? everyServer;
|
||||
|
||||
const renderSurface = (onChange = vi.fn()) =>
|
||||
render(
|
||||
<QueryClientProvider client={new QueryClient({ defaultOptions: { queries: { retry: false } } })}>
|
||||
<ConnectFlowSurface accessToken="token-123" selectedServers={[]} onChange={vi.fn()} />
|
||||
<ConnectFlowSurface accessToken="token-123" selectedServers={[]} onChange={onChange} />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
|
||||
const openConnectPage = async (result: ConnectFlowStatus = flow("unscoped"), handle = "flow-handle-123") => {
|
||||
state.connectFlow = handle;
|
||||
vi.mocked(fetchConnectFlow).mockResolvedValue(result);
|
||||
const page = renderSurface();
|
||||
await screen.findByTestId("mcp-apps-panel");
|
||||
return page;
|
||||
};
|
||||
|
||||
afterEach(() => {
|
||||
state.oauthReturn = null;
|
||||
state.connectFlow = null;
|
||||
onSuccess.current = undefined;
|
||||
banner.flow = undefined;
|
||||
banner.failed = undefined;
|
||||
banner.finish = undefined;
|
||||
sessionStorage.clear();
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe("ConnectFlowSurface", () => {
|
||||
it.each([
|
||||
{ result: flow("unscoped"), grid: true, finish: true, cancel: false, oauthStarts: 0 },
|
||||
{ result: flow("interactive", false), grid: false, finish: false, cancel: true, oauthStarts: 1 },
|
||||
{ result: flow("interactive", true), grid: false, finish: true, cancel: true, oauthStarts: 0 },
|
||||
{ result: flow("m2m", true), grid: false, finish: true, cancel: true, oauthStarts: 0 },
|
||||
{ result: flow("stale"), grid: false, finish: false, cancel: true, oauthStarts: 0 },
|
||||
])(
|
||||
"renders $result.state without widening its action surface",
|
||||
async ({ result, grid, finish, cancel, oauthStarts }) => {
|
||||
state.connectFlow = "flow-handle-123";
|
||||
vi.mocked(fetchConnectFlow).mockResolvedValue(result);
|
||||
renderSurface();
|
||||
{ result: flow("unscoped"), grid: true },
|
||||
{ result: flow("interactive", false), grid: false },
|
||||
{ result: flow("interactive", true), grid: false },
|
||||
{ result: flow("m2m", true), grid: false },
|
||||
{ result: flow("stale"), grid: false },
|
||||
])("hands the banner a $result.state flow and shows the grid only when it is unscoped", async ({ result, grid }) => {
|
||||
state.connectFlow = "flow-handle-123";
|
||||
vi.mocked(fetchConnectFlow).mockResolvedValue(result);
|
||||
renderSurface();
|
||||
|
||||
await screen.findByRole("button", { name: /finish connecting|cancel|connect/i });
|
||||
await waitFor(() => expect(startOAuthFlow).toHaveBeenCalledTimes(oauthStarts));
|
||||
expect(screen.queryByTestId("mcp-apps-panel") !== null).toBe(grid);
|
||||
expect(screen.queryByRole("button", { name: /finish connecting/i }) !== null).toBe(finish);
|
||||
expect(screen.queryByRole("button", { name: "Cancel" }) !== null).toBe(cancel);
|
||||
},
|
||||
);
|
||||
await waitFor(() => expect(banner.flow).toEqual(result));
|
||||
expect(screen.queryByTestId("mcp-apps-panel") !== null).toBe(grid);
|
||||
});
|
||||
|
||||
it("keeps the grid and Finish hidden until the gateway accepts a handle", () => {
|
||||
it("keeps the grid hidden until the gateway accepts a handle", () => {
|
||||
state.connectFlow = "flow-handle-123";
|
||||
vi.mocked(fetchConnectFlow).mockReturnValue(new Promise(() => {}));
|
||||
renderSurface();
|
||||
|
||||
expect(screen.queryByTestId("mcp-apps-panel")).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Cancel" })).toHaveAttribute("value", "deny");
|
||||
expect(banner.flow).toBeUndefined();
|
||||
expect(banner.failed).toBe(false);
|
||||
});
|
||||
|
||||
it("keeps the grid and Finish hidden when flow validation fails", async () => {
|
||||
it("keeps the grid hidden and tells the banner when flow validation fails", async () => {
|
||||
state.connectFlow = "invalid-handle";
|
||||
vi.mocked(fetchConnectFlow).mockRejectedValue(new Error("invalid flow"));
|
||||
renderSurface();
|
||||
|
||||
await screen.findByRole("button", { name: "Cancel" });
|
||||
await waitFor(() => expect(banner.failed).toBe(true));
|
||||
expect(screen.queryByTestId("mcp-apps-panel")).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("refetches the sealed flow after the vendor connection completes", async () => {
|
||||
state.connectFlow = "flow-handle-123";
|
||||
vi.mocked(fetchConnectFlow)
|
||||
.mockResolvedValueOnce(flow("interactive", false))
|
||||
.mockResolvedValueOnce(flow("interactive", true));
|
||||
renderSurface();
|
||||
|
||||
await waitFor(() => expect(startOAuthFlow).toHaveBeenCalledOnce());
|
||||
await act(async () => onSuccess.current?.());
|
||||
|
||||
await screen.findByRole("button", { name: /finish connecting/i });
|
||||
});
|
||||
|
||||
it("renders the ordinary panel without a flow handle", () => {
|
||||
renderSurface();
|
||||
|
||||
expect(fetchConnectFlow).not.toHaveBeenCalled();
|
||||
expect(screen.getByTestId("mcp-apps-panel")).toBeInTheDocument();
|
||||
expect(screen.queryByTestId("connect-flow-banner")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps the servers turned off across the OAuth round trip", async () => {
|
||||
const { unmount } = await openConnectPage();
|
||||
act(() => panel.onOffServerIdsChange?.(["orders-id"]));
|
||||
unmount();
|
||||
|
||||
await openConnectPage();
|
||||
|
||||
expect(panel.offServerIds).toEqual(["orders-id"]);
|
||||
});
|
||||
|
||||
it("starts a new connect flow with every server on", async () => {
|
||||
const { unmount } = await openConnectPage();
|
||||
act(() => panel.onOffServerIdsChange?.(["orders-id"]));
|
||||
unmount();
|
||||
|
||||
await openConnectPage(flow("unscoped"), "flow-handle-456");
|
||||
|
||||
expect(panel.offServerIds).toEqual([]);
|
||||
expect(panel.lockedOn).toBeUndefined();
|
||||
});
|
||||
|
||||
it("locks every switch and never narrows a flow that cannot be narrowed", async () => {
|
||||
await openConnectPage({ ...flow("unscoped"), selectable: false });
|
||||
expect(panel.lockedOn).toBe("all");
|
||||
|
||||
act(() => panel.onConnectedServers?.({ on: ["catalog-id"], all: ["catalog-id", "orders-id"] }));
|
||||
|
||||
expect(finishHanded()).toEqual(everyServer);
|
||||
});
|
||||
|
||||
it("locks only the servers the client requires and still narrows to the servers that are on", async () => {
|
||||
await openConnectPage({ ...flow("unscoped"), required_server_ids: ["tracker-id"] });
|
||||
expect(panel.lockedOn).toEqual(["tracker-id"]);
|
||||
|
||||
act(() => panel.onConnectedServers?.({ on: ["tracker-id"], all: ["tracker-id", "orders-id"] }));
|
||||
|
||||
expect(banner.finish).toEqual({ status: "ready", serverIds: ["tracker-id"] });
|
||||
});
|
||||
|
||||
it("holds Finish on return from a vendor sign-in before the panel has said anything", async () => {
|
||||
vi.mocked(pendingUserMcpOAuthServerId).mockReturnValueOnce("vault-id");
|
||||
|
||||
await openConnectPage();
|
||||
|
||||
expect(banner.finish).toEqual(waiting);
|
||||
});
|
||||
|
||||
it.each<{ when: string; off: string[]; report: ServerListReport; finish: FinishSelection }>([
|
||||
{ when: "the list is loading and nothing is off", off: [], report: "loading", finish: everyServer },
|
||||
{
|
||||
when: "the list is loading after a server was turned off",
|
||||
off: ["orders-id"],
|
||||
report: "loading",
|
||||
finish: waiting,
|
||||
},
|
||||
{ when: "the list failed to load and nothing is off", off: [], report: "failed", finish: everyServer },
|
||||
{ when: "the list failed after a server was turned off", off: ["orders-id"], report: "failed", finish: waiting },
|
||||
{ when: "a login returning from OAuth is being saved", off: [], report: "saving-login", finish: waiting },
|
||||
{
|
||||
when: "every server is on",
|
||||
off: [],
|
||||
report: { on: ["catalog-id", "orders-id"], all: ["catalog-id", "orders-id"] },
|
||||
finish: everyServer,
|
||||
},
|
||||
{
|
||||
when: "one server is off",
|
||||
off: ["orders-id"],
|
||||
report: { on: ["catalog-id"], all: ["catalog-id", "orders-id"] },
|
||||
finish: { status: "ready", serverIds: ["catalog-id"] },
|
||||
},
|
||||
{
|
||||
when: "every server is off",
|
||||
off: ["catalog-id", "orders-id"],
|
||||
report: { on: [], all: ["catalog-id", "orders-id"] },
|
||||
finish: { status: "nothing-on" },
|
||||
},
|
||||
{ when: "no server is reachable", off: [], report: { on: [], all: [] }, finish: everyServer },
|
||||
])("tells the banner Finish is $finish.status when $when", async ({ off, report, finish }) => {
|
||||
await openConnectPage();
|
||||
|
||||
act(() => panel.onOffServerIdsChange?.(off));
|
||||
act(() => panel.onConnectedServers?.(report));
|
||||
|
||||
expect(banner.finish).toEqual(finish);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,11 +1,12 @@
|
|||
"use client";
|
||||
|
||||
import React, { useEffect } from "react";
|
||||
import React, { useEffect, useState } from "react";
|
||||
import { useRouter, useSearchParams } from "next/navigation";
|
||||
import { useQuery } from "@tanstack/react-query";
|
||||
import MCPAppsPanel from "@/components/chat/MCPAppsPanel";
|
||||
import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner";
|
||||
import MCPAppsPanel, { type ServerListReport } from "@/components/chat/MCPAppsPanel";
|
||||
import ConnectFlowBanner, { type FinishSelection } from "@/components/chat/ConnectFlowBanner";
|
||||
import { fetchConnectFlow } from "@/components/networking";
|
||||
import { pendingUserMcpOAuthServerId } from "@/hooks/useUserMcpOAuthFlow";
|
||||
|
||||
interface Props {
|
||||
accessToken: string;
|
||||
|
|
@ -13,6 +14,33 @@ interface Props {
|
|||
onChange: (servers: string[]) => void;
|
||||
}
|
||||
|
||||
const offServersKey = (flow: string) => `litellm-connect-off:${flow}`;
|
||||
|
||||
const storedServerIds = (key: string): string[] => {
|
||||
try {
|
||||
const value: unknown = JSON.parse(sessionStorage.getItem(key) ?? "null");
|
||||
return Array.isArray(value) ? value.filter((item): item is string => typeof item === "string") : [];
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
};
|
||||
|
||||
const rememberServerIds = (key: string, serverIds: string[]) => {
|
||||
try {
|
||||
sessionStorage.setItem(key, JSON.stringify(serverIds));
|
||||
} catch {}
|
||||
};
|
||||
|
||||
const finishSelectionFor = (connected: ServerListReport, offServerIds: readonly string[]): FinishSelection => {
|
||||
if (connected === "saving-login") return { status: "waiting" };
|
||||
if (typeof connected === "string") {
|
||||
return offServerIds.length === 0 ? { status: "ready", serverIds: [] } : { status: "waiting" };
|
||||
}
|
||||
if (connected.on.length === connected.all.length) return { status: "ready", serverIds: [] };
|
||||
if (connected.on.length === 0) return { status: "nothing-on" };
|
||||
return { status: "ready", serverIds: connected.on };
|
||||
};
|
||||
|
||||
/** Renders the sealed gateway connect flow without trusting URL context. */
|
||||
const ConnectFlowSurface: React.FC<Props> = ({ accessToken, selectedServers, onChange }) => {
|
||||
const router = useRouter();
|
||||
|
|
@ -35,11 +63,24 @@ const ConnectFlowSurface: React.FC<Props> = ({ accessToken, selectedServers, onC
|
|||
retry: false,
|
||||
};
|
||||
const { data: flow, isError, refetch } = useQuery(flowQuery);
|
||||
const [connected, setConnected] = useState<ServerListReport>(() =>
|
||||
pendingUserMcpOAuthServerId() === null ? "loading" : "saving-login",
|
||||
);
|
||||
const [offServerIds, setOffServerIds] = useState(() =>
|
||||
connectFlow ? storedServerIds(offServersKey(connectFlow)) : [],
|
||||
);
|
||||
|
||||
const changeOffServerIds = (serverIds: string[]) => {
|
||||
if (connectFlow) rememberServerIds(offServersKey(connectFlow), serverIds);
|
||||
setOffServerIds(serverIds);
|
||||
};
|
||||
|
||||
if (connectFlow === null) {
|
||||
return <MCPAppsPanel accessToken={accessToken} selectedServers={selectedServers} onChange={onChange} />;
|
||||
}
|
||||
|
||||
const narrowable = flow?.state === "unscoped" && flow.selectable !== false;
|
||||
|
||||
return (
|
||||
<>
|
||||
<ConnectFlowBanner
|
||||
|
|
@ -48,9 +89,19 @@ const ConnectFlowSurface: React.FC<Props> = ({ accessToken, selectedServers, onC
|
|||
accessToken={accessToken}
|
||||
onConnected={refetch}
|
||||
failed={isError}
|
||||
finish={narrowable ? finishSelectionFor(connected, offServerIds) : undefined}
|
||||
/>
|
||||
{flow?.state === "unscoped" && (
|
||||
<MCPAppsPanel accessToken={accessToken} selectedServers={selectedServers} onChange={onChange} connectMode />
|
||||
<MCPAppsPanel
|
||||
accessToken={accessToken}
|
||||
selectedServers={selectedServers}
|
||||
onChange={onChange}
|
||||
connectMode
|
||||
offServerIds={offServerIds}
|
||||
onOffServerIdsChange={changeOffServerIds}
|
||||
lockedOn={narrowable ? flow.required_server_ids : "all"}
|
||||
onConnectedServers={setConnected}
|
||||
/>
|
||||
)}
|
||||
</>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -0,0 +1,587 @@
|
|||
import React from "react";
|
||||
import { render, screen, fireEvent, waitFor, act } from "@testing-library/react";
|
||||
import { describe, it, expect, vi, afterEach, beforeEach } from "vitest";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import MCPAppsPanel, { type LockedOn, type ServerListReport } from "./MCPAppsPanel";
|
||||
import {
|
||||
deleteMCPOAuthUserCredential,
|
||||
fetchMCPServers,
|
||||
getMCPOAuthUserCredentialStatus,
|
||||
listMCPTools,
|
||||
} from "../networking";
|
||||
import type { MCPOAuthUserCredentialStatus } from "../networking";
|
||||
import type { MCPServer } from "../mcp_tools/types";
|
||||
import { setServerRootPath } from "@/lib/serverRootPath";
|
||||
|
||||
vi.mock("../networking", () => ({
|
||||
fetchMCPServers: vi.fn(),
|
||||
getMCPOAuthUserCredentialStatus: vi.fn(),
|
||||
listMCPTools: vi.fn(),
|
||||
deleteMCPOAuthUserCredential: vi.fn(),
|
||||
}));
|
||||
|
||||
const oauthReturn = vi.hoisted(() => ({ serverId: null as string | null, status: "idle", onSuccess: () => {} }));
|
||||
vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({
|
||||
pendingUserMcpOAuthServerId: () => oauthReturn.serverId,
|
||||
useUserMcpOAuthFlow: ({ serverId, onSuccess }: { serverId: string; onSuccess: () => void }) => {
|
||||
if (serverId !== oauthReturn.serverId) return { startOAuthFlow: vi.fn(), status: "idle" };
|
||||
oauthReturn.onSuccess = onSuccess;
|
||||
return { startOAuthFlow: vi.fn(), status: oauthReturn.status };
|
||||
},
|
||||
}));
|
||||
|
||||
const servers = [
|
||||
{
|
||||
server_id: "s-ext",
|
||||
server_name: "external_logo",
|
||||
auth_type: "none",
|
||||
mcp_info: { server_name: "external_logo", logo_url: "https://cdn.example.com/ext.png" },
|
||||
},
|
||||
{
|
||||
server_id: "s-local",
|
||||
server_name: "local_logo",
|
||||
auth_type: "none",
|
||||
mcp_info: { server_name: "local_logo", logo_url: "/ui/assets/logos/github.svg" },
|
||||
},
|
||||
{
|
||||
server_id: "s-none",
|
||||
server_name: "no_logo",
|
||||
auth_type: "none",
|
||||
},
|
||||
] as MCPServer[];
|
||||
|
||||
const renderPanel = () =>
|
||||
render(
|
||||
<QueryClientProvider client={new QueryClient({ defaultOptions: { queries: { retry: false } } })}>
|
||||
<MCPAppsPanel accessToken="tok" selectedServers={[]} onChange={vi.fn()} />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
|
||||
describe("MCPAppsPanel logos", () => {
|
||||
afterEach(() => {
|
||||
setServerRootPath("/");
|
||||
});
|
||||
|
||||
it("resolves backend logo_url values in the server grid", async () => {
|
||||
setServerRootPath("/litellm");
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(servers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderPanel();
|
||||
|
||||
expect(await screen.findByText("external_logo")).toBeInTheDocument();
|
||||
expect(screen.getByAltText("external_logo logo")).toHaveAttribute("src", "https://cdn.example.com/ext.png");
|
||||
expect(screen.getByAltText("local_logo logo")).toHaveAttribute("src", "/litellm/ui/assets/logos/github.svg");
|
||||
});
|
||||
|
||||
it("renders a colored letter avatar for servers without logo_url", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(servers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderPanel();
|
||||
|
||||
expect(await screen.findByText("no_logo")).toBeInTheDocument();
|
||||
expect(screen.queryByAltText("no_logo logo")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("N")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("resolves the logo_url in the detail header", async () => {
|
||||
setServerRootPath("/litellm");
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(servers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderPanel();
|
||||
|
||||
fireEvent.click(await screen.findByText("local_logo"));
|
||||
|
||||
expect(await screen.findByRole("heading", { name: "local_logo" })).toBeInTheDocument();
|
||||
expect(screen.getByAltText("local_logo logo")).toHaveAttribute("src", "/litellm/ui/assets/logos/github.svg");
|
||||
});
|
||||
});
|
||||
|
||||
const connectServers = [
|
||||
{
|
||||
server_id: "s-m2m",
|
||||
server_name: "service_tool",
|
||||
auth_type: "oauth2",
|
||||
oauth2_flow: "client_credentials",
|
||||
connected_app_reachable: true,
|
||||
},
|
||||
{
|
||||
server_id: "s-reach",
|
||||
server_name: "reachable_srv",
|
||||
auth_type: "none",
|
||||
connected_app_reachable: true,
|
||||
},
|
||||
{
|
||||
server_id: "s-unreach",
|
||||
server_name: "unreachable_srv",
|
||||
auth_type: "none",
|
||||
connected_app_reachable: false,
|
||||
},
|
||||
] as MCPServer[];
|
||||
|
||||
const renderConnectPanel = (connectMode: boolean, selectedServers: string[] = []) =>
|
||||
render(
|
||||
<QueryClientProvider client={new QueryClient({ defaultOptions: { queries: { retry: false } } })}>
|
||||
<MCPAppsPanel accessToken="tok" selectedServers={selectedServers} onChange={vi.fn()} connectMode={connectMode} />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
|
||||
const oauthServer = {
|
||||
server_id: "s-oauth",
|
||||
server_name: "oauth_srv",
|
||||
auth_type: "oauth2",
|
||||
oauth2_flow: "authorization_code",
|
||||
connected_app_reachable: true,
|
||||
} as MCPServer;
|
||||
|
||||
const SwitchPage = ({
|
||||
onConnectedServers,
|
||||
initialOff = [],
|
||||
lockedOn,
|
||||
}: {
|
||||
onConnectedServers: (servers: ServerListReport) => void;
|
||||
initialOff?: string[];
|
||||
lockedOn?: LockedOn;
|
||||
}) => {
|
||||
const [off, setOff] = React.useState<string[]>(initialOff);
|
||||
return (
|
||||
<QueryClientProvider client={new QueryClient({ defaultOptions: { queries: { retry: false } } })}>
|
||||
<MCPAppsPanel
|
||||
accessToken="tok"
|
||||
selectedServers={[]}
|
||||
onChange={vi.fn()}
|
||||
connectMode
|
||||
offServerIds={off}
|
||||
onOffServerIdsChange={setOff}
|
||||
lockedOn={lockedOn}
|
||||
onConnectedServers={onConnectedServers}
|
||||
/>
|
||||
</QueryClientProvider>
|
||||
);
|
||||
};
|
||||
|
||||
const lastReport = (onConnectedServers: ReturnType<typeof vi.fn>) => onConnectedServers.mock.lastCall?.[0];
|
||||
|
||||
describe("MCPAppsPanel switches in connect mode", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
oauthReturn.serverId = null;
|
||||
oauthReturn.status = "idle";
|
||||
oauthReturn.onSuccess = () => {};
|
||||
});
|
||||
|
||||
it("reports every reachable server as on before any login is checked", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue([
|
||||
oauthServer,
|
||||
...connectServers,
|
||||
{ server_id: "s-unsup", server_name: "unsupported_srv", auth_type: "oauth2_token_exchange" },
|
||||
] as MCPServer[]);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
vi.mocked(getMCPOAuthUserCredentialStatus).mockReturnValue(new Promise(() => {}));
|
||||
const onConnectedServers = vi.fn();
|
||||
|
||||
render(<SwitchPage onConnectedServers={onConnectedServers} />);
|
||||
|
||||
await waitFor(() =>
|
||||
expect(lastReport(onConnectedServers)).toEqual({
|
||||
on: ["s-oauth", "s-m2m", "s-reach"],
|
||||
all: ["s-oauth", "s-m2m", "s-reach"],
|
||||
}),
|
||||
);
|
||||
expect(screen.getByRole("switch", { name: "Include oauth_srv in this connection" })).toBeChecked();
|
||||
expect(screen.getByRole("switch", { name: "Include service_tool in this connection" })).toBeChecked();
|
||||
expect(
|
||||
screen.queryByRole("switch", { name: "Include unsupported_srv in this connection" }),
|
||||
).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("leaves a server out once it is turned off, and back in once it is turned on", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
const onConnectedServers = vi.fn();
|
||||
render(<SwitchPage onConnectedServers={onConnectedServers} />);
|
||||
const reachable = await screen.findByRole("switch", { name: "Include reachable_srv in this connection" });
|
||||
|
||||
fireEvent.click(reachable);
|
||||
|
||||
expect(lastReport(onConnectedServers)).toEqual({ on: ["s-m2m"], all: ["s-m2m", "s-reach"] });
|
||||
expect(reachable).not.toBeChecked();
|
||||
expect(screen.queryByRole("heading", { name: "reachable_srv" })).not.toBeInTheDocument();
|
||||
|
||||
fireEvent.click(reachable);
|
||||
|
||||
expect(lastReport(onConnectedServers)).toEqual({ on: ["s-m2m", "s-reach"], all: ["s-m2m", "s-reach"] });
|
||||
});
|
||||
|
||||
it("lists only the servers that are on under the Connected tab", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
render(<SwitchPage onConnectedServers={vi.fn()} />);
|
||||
fireEvent.click(await screen.findByRole("switch", { name: "Include reachable_srv in this connection" }));
|
||||
|
||||
fireEvent.click(screen.getByRole("tab", { name: "On (1)" }));
|
||||
|
||||
expect(await screen.findByText("service_tool")).toBeInTheDocument();
|
||||
expect(screen.queryByText("reachable_srv")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("turns a server off from its detail view too", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
const onConnectedServers = vi.fn();
|
||||
render(<SwitchPage onConnectedServers={onConnectedServers} />);
|
||||
fireEvent.click(await screen.findByText("service_tool"));
|
||||
await screen.findByRole("heading", { name: "service_tool" });
|
||||
|
||||
fireEvent.click(screen.getByRole("switch", { name: "Include service_tool in this connection" }));
|
||||
|
||||
expect(lastReport(onConnectedServers)).toEqual({ on: ["s-reach"], all: ["s-m2m", "s-reach"] });
|
||||
expect(screen.getByText("Off")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps an OAuth server with no saved login on and offers a sign in", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue([oauthServer]);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
vi.mocked(getMCPOAuthUserCredentialStatus).mockResolvedValue({
|
||||
server_id: "s-oauth",
|
||||
has_credential: false,
|
||||
is_expired: false,
|
||||
} as MCPOAuthUserCredentialStatus);
|
||||
|
||||
render(<SwitchPage onConnectedServers={vi.fn()} />);
|
||||
|
||||
expect(await screen.findByText("Sign in to use")).toBeInTheDocument();
|
||||
expect(screen.getByText("Connect", { exact: true })).toBeInTheDocument();
|
||||
expect(screen.getByRole("switch", { name: "Include oauth_srv in this connection" })).toBeChecked();
|
||||
});
|
||||
|
||||
it("keeps the saved login when an OAuth server is turned off", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue([oauthServer, ...connectServers]);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
vi.mocked(getMCPOAuthUserCredentialStatus).mockResolvedValue({
|
||||
server_id: "s-oauth",
|
||||
has_credential: true,
|
||||
is_expired: false,
|
||||
} as MCPOAuthUserCredentialStatus);
|
||||
const onConnectedServers = vi.fn();
|
||||
render(<SwitchPage onConnectedServers={onConnectedServers} />);
|
||||
await waitFor(() => expect(getMCPOAuthUserCredentialStatus).toHaveBeenCalled());
|
||||
|
||||
fireEvent.click(await screen.findByRole("switch", { name: "Include oauth_srv in this connection" }));
|
||||
|
||||
expect(lastReport(onConnectedServers)).toEqual({ on: ["s-m2m", "s-reach"], all: ["s-oauth", "s-m2m", "s-reach"] });
|
||||
expect(deleteMCPOAuthUserCredential).not.toHaveBeenCalled();
|
||||
expect(screen.queryByText("Sign in to use")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps the switch where it was when a login returning from OAuth is saved", async () => {
|
||||
oauthReturn.serverId = "s-oauth";
|
||||
oauthReturn.status = "exchanging";
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue([oauthServer, ...connectServers]);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
vi.mocked(getMCPOAuthUserCredentialStatus).mockResolvedValue({
|
||||
server_id: "s-oauth",
|
||||
has_credential: false,
|
||||
is_expired: false,
|
||||
} as MCPOAuthUserCredentialStatus);
|
||||
const onConnectedServers = vi.fn();
|
||||
render(<SwitchPage onConnectedServers={onConnectedServers} initialOff={["s-oauth"]} />);
|
||||
await screen.findByText("oauth_srv");
|
||||
|
||||
expect(lastReport(onConnectedServers)).toBe("saving-login");
|
||||
|
||||
await act(async () => {
|
||||
oauthReturn.status = "success";
|
||||
oauthReturn.onSuccess();
|
||||
});
|
||||
|
||||
expect(screen.getByRole("switch", { name: "Include oauth_srv in this connection" })).not.toBeChecked();
|
||||
expect(lastReport(onConnectedServers)).toEqual({ on: ["s-m2m", "s-reach"], all: ["s-oauth", "s-m2m", "s-reach"] });
|
||||
});
|
||||
|
||||
it("reports a returning login as being saved from the first report, before the hook starts", async () => {
|
||||
oauthReturn.serverId = "s-oauth";
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue([oauthServer, ...connectServers]);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
vi.mocked(getMCPOAuthUserCredentialStatus).mockReturnValue(new Promise(() => {}));
|
||||
const onConnectedServers = vi.fn();
|
||||
|
||||
render(<SwitchPage onConnectedServers={onConnectedServers} />);
|
||||
await screen.findByText("oauth_srv");
|
||||
|
||||
expect(new Set(onConnectedServers.mock.calls.map(([report]) => report))).toEqual(new Set(["saving-login"]));
|
||||
});
|
||||
|
||||
it("shows every switch on and locked when the flow cannot be narrowed", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
render(<SwitchPage onConnectedServers={vi.fn()} initialOff={["s-reach"]} lockedOn="all" />);
|
||||
|
||||
const reachable = await screen.findByRole("switch", { name: "Include reachable_srv in this connection" });
|
||||
expect(reachable).toBeChecked();
|
||||
expect(reachable).toHaveAttribute("data-disabled");
|
||||
expect(screen.queryByText("Required by this client")).not.toBeInTheDocument();
|
||||
|
||||
fireEvent.click(reachable);
|
||||
|
||||
expect(reachable).toBeChecked();
|
||||
});
|
||||
|
||||
it("keeps a server the client requires on and locked, and leaves the others switchable", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
const onConnectedServers = vi.fn();
|
||||
render(<SwitchPage onConnectedServers={onConnectedServers} initialOff={["s-reach"]} lockedOn={["s-reach"]} />);
|
||||
const required = await screen.findByRole("switch", { name: "Include reachable_srv in this connection" });
|
||||
expect(required).toBeChecked();
|
||||
expect(required).toHaveAttribute("data-disabled");
|
||||
expect(screen.getByText("Required by this client")).toBeInTheDocument();
|
||||
|
||||
fireEvent.click(required);
|
||||
fireEvent.click(screen.getByRole("switch", { name: "Include service_tool in this connection" }));
|
||||
|
||||
expect(required).toBeChecked();
|
||||
expect(lastReport(onConnectedServers)).toEqual({ on: ["s-reach"], all: ["s-m2m", "s-reach"] });
|
||||
});
|
||||
|
||||
it("ignores a stored off switch for a server that is no longer listed", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
const onConnectedServers = vi.fn();
|
||||
|
||||
render(<SwitchPage onConnectedServers={onConnectedServers} initialOff={["s-deleted"]} />);
|
||||
|
||||
await waitFor(() =>
|
||||
expect(lastReport(onConnectedServers)).toEqual({ on: ["s-m2m", "s-reach"], all: ["s-m2m", "s-reach"] }),
|
||||
);
|
||||
});
|
||||
|
||||
it("reports nothing while the list is still loading", async () => {
|
||||
vi.mocked(fetchMCPServers).mockReturnValue(new Promise(() => {}));
|
||||
const onConnectedServers = vi.fn();
|
||||
|
||||
render(<SwitchPage onConnectedServers={onConnectedServers} />);
|
||||
|
||||
await waitFor(() => expect(fetchMCPServers).toHaveBeenCalled());
|
||||
expect(onConnectedServers.mock.calls.map(([servers]) => servers)).toEqual(["loading"]);
|
||||
});
|
||||
|
||||
it("reports nothing and says so when the list fails to load", async () => {
|
||||
vi.mocked(fetchMCPServers).mockRejectedValue(new Error("list failed"));
|
||||
const onConnectedServers = vi.fn();
|
||||
|
||||
render(<SwitchPage onConnectedServers={onConnectedServers} />);
|
||||
|
||||
expect(
|
||||
await screen.findByText("Could not load your MCP servers. Reload the page to try again."),
|
||||
).toBeInTheDocument();
|
||||
expect(onConnectedServers.mock.calls.map(([servers]) => servers)).toEqual(["loading", "failed"]);
|
||||
});
|
||||
|
||||
it("reports again once a reload succeeds after the list failed to load", async () => {
|
||||
vi.mocked(fetchMCPServers).mockRejectedValueOnce(new Error("list failed")).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
const onConnectedServers = vi.fn();
|
||||
const panel = (token: string) => (
|
||||
<QueryClientProvider client={new QueryClient({ defaultOptions: { queries: { retry: false } } })}>
|
||||
<MCPAppsPanel
|
||||
accessToken={token}
|
||||
selectedServers={[]}
|
||||
onChange={vi.fn()}
|
||||
connectMode
|
||||
onConnectedServers={onConnectedServers}
|
||||
/>
|
||||
</QueryClientProvider>
|
||||
);
|
||||
const { rerender } = render(panel("tok"));
|
||||
expect(await screen.findByText(/could not load your mcp servers/i)).toBeInTheDocument();
|
||||
|
||||
rerender(panel("tok-refreshed"));
|
||||
|
||||
await waitFor(() =>
|
||||
expect(lastReport(onConnectedServers)).toEqual({ on: ["s-m2m", "s-reach"], all: ["s-m2m", "s-reach"] }),
|
||||
);
|
||||
});
|
||||
|
||||
it("leaves a login returning from OAuth to the server's own badge outside connect mode", async () => {
|
||||
oauthReturn.serverId = "s-oauth";
|
||||
oauthReturn.status = "exchanging";
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue([
|
||||
{ server_id: "s-oauth", server_name: "oauth_srv", auth_type: "oauth2", oauth2_flow: "authorization_code" },
|
||||
] as MCPServer[]);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
vi.mocked(getMCPOAuthUserCredentialStatus).mockResolvedValue({
|
||||
server_id: "s-oauth",
|
||||
has_credential: false,
|
||||
is_expired: false,
|
||||
} as MCPOAuthUserCredentialStatus);
|
||||
|
||||
renderConnectPanel(false);
|
||||
|
||||
expect(await screen.findByText("Connecting…")).toBeInTheDocument();
|
||||
expect(screen.queryByRole("switch")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps the plain panel's empty text when the list fails to load outside connect mode", async () => {
|
||||
vi.mocked(fetchMCPServers).mockRejectedValue(new Error("list failed"));
|
||||
|
||||
renderConnectPanel(false);
|
||||
|
||||
expect(await screen.findByText(/no mcp servers configured/i)).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
describe("MCPAppsPanel connected-app reachability (LIT-4861)", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
oauthReturn.serverId = null;
|
||||
oauthReturn.status = "idle";
|
||||
});
|
||||
|
||||
it("requests the connected-app view and hides unreachable servers in connect mode", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(true, ["reachable_srv", "unreachable_srv"]);
|
||||
|
||||
expect(await screen.findByText("reachable_srv")).toBeInTheDocument();
|
||||
expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, true);
|
||||
expect(screen.queryByText("unreachable_srv")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("On (2)")).toBeInTheDocument();
|
||||
expect(screen.getByText("service_tool")).toBeInTheDocument();
|
||||
expect(screen.queryByText("Connect", { exact: true })).not.toBeInTheDocument();
|
||||
const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]);
|
||||
expect(toolCountFetchedIds).toContain("s-reach");
|
||||
expect(toolCountFetchedIds).not.toContain("s-unreach");
|
||||
});
|
||||
|
||||
it("blocks connecting an unsupported server from the detail view in connect mode", async () => {
|
||||
const detailServers = [
|
||||
...connectServers,
|
||||
{
|
||||
server_id: "s-unsup",
|
||||
server_name: "unsupported_srv",
|
||||
auth_type: "oauth2_token_exchange",
|
||||
connected_app_reachable: true,
|
||||
},
|
||||
] as MCPServer[];
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(detailServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(true);
|
||||
|
||||
fireEvent.click(await screen.findByText("unsupported_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "unsupported_srv" })).toBeInTheDocument();
|
||||
expect(screen.queryByRole("button", { name: "Connect" })).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("switch")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Not supported on this connection")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps the detail-view Connect action outside connect mode", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(false);
|
||||
|
||||
fireEvent.click(await screen.findByText("unreachable_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "unreachable_srv" })).toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Connect" })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("ignores the flag and skips no server outside connect mode", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(false, ["reachable_srv", "unreachable_srv"]);
|
||||
|
||||
expect(await screen.findByText("unreachable_srv")).toBeInTheDocument();
|
||||
expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, false);
|
||||
expect(screen.queryByText("Not available to connected apps")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Connected (2)")).toBeInTheDocument();
|
||||
const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]);
|
||||
expect(toolCountFetchedIds).toContain("s-unreach");
|
||||
});
|
||||
|
||||
const revocable = (reachable: boolean) =>
|
||||
[
|
||||
{ server_id: "s-reach", server_name: "reachable_srv", auth_type: "none", connected_app_reachable: true },
|
||||
{ server_id: "s-drop", server_name: "revoked_srv", auth_type: "none", connected_app_reachable: reachable },
|
||||
] as MCPServer[];
|
||||
|
||||
const ConnectPanel = ({
|
||||
token,
|
||||
client,
|
||||
onConnectedServers = vi.fn(),
|
||||
}: {
|
||||
token: string;
|
||||
client: QueryClient;
|
||||
onConnectedServers?: (servers: ServerListReport) => void;
|
||||
}) => (
|
||||
<QueryClientProvider client={client}>
|
||||
<MCPAppsPanel
|
||||
accessToken={token}
|
||||
selectedServers={[]}
|
||||
onChange={vi.fn()}
|
||||
connectMode
|
||||
onConnectedServers={onConnectedServers}
|
||||
/>
|
||||
</QueryClientProvider>
|
||||
);
|
||||
|
||||
const newClient = () => new QueryClient({ defaultOptions: { queries: { retry: false } } });
|
||||
|
||||
it("drops an open detail view when a refetch removes that server from the reachable set", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValueOnce(revocable(true)).mockResolvedValueOnce(revocable(false));
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
const client = newClient();
|
||||
const { rerender } = render(<ConnectPanel token="tok" client={client} />);
|
||||
|
||||
fireEvent.click(await screen.findByText("revoked_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument();
|
||||
|
||||
rerender(<ConnectPanel token="tok-refreshed" client={client} />);
|
||||
|
||||
await waitFor(() => expect(screen.queryByRole("heading", { name: "revoked_srv" })).not.toBeInTheDocument());
|
||||
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("reachable_srv")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("stops reporting a server once a refetch removes it from the reachable set", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValueOnce(revocable(true)).mockResolvedValueOnce(revocable(false));
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
const onConnectedServers = vi.fn();
|
||||
const client = newClient();
|
||||
const { rerender } = render(<ConnectPanel token="tok" client={client} onConnectedServers={onConnectedServers} />);
|
||||
await waitFor(() =>
|
||||
expect(lastReport(onConnectedServers)).toEqual({ on: ["s-reach", "s-drop"], all: ["s-reach", "s-drop"] }),
|
||||
);
|
||||
|
||||
rerender(<ConnectPanel token="tok-refreshed" client={client} onConnectedServers={onConnectedServers} />);
|
||||
|
||||
await waitFor(() => expect(lastReport(onConnectedServers)).toEqual({ on: ["s-reach"], all: ["s-reach"] }));
|
||||
});
|
||||
|
||||
it("does not let a superseded list load overwrite the current reachable set", async () => {
|
||||
let finishStaleLoad: (servers: MCPServer[]) => void = () => {};
|
||||
vi.mocked(fetchMCPServers)
|
||||
.mockImplementationOnce(() => new Promise((resolve) => (finishStaleLoad = resolve)))
|
||||
.mockResolvedValueOnce(revocable(false));
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
const client = newClient();
|
||||
const { rerender } = render(<ConnectPanel token="tok" client={client} />);
|
||||
rerender(<ConnectPanel token="tok-refreshed" client={client} />);
|
||||
|
||||
expect(await screen.findByText("reachable_srv")).toBeInTheDocument();
|
||||
|
||||
await act(async () => {
|
||||
finishStaleLoad(revocable(true));
|
||||
});
|
||||
|
||||
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
@ -1,302 +0,0 @@
|
|||
import React from "react";
|
||||
import { render, screen, fireEvent, waitFor, act } from "@testing-library/react";
|
||||
import { describe, it, expect, vi, afterEach, beforeEach } from "vitest";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import MCPAppsPanel from "./MCPAppsPanel";
|
||||
import { fetchMCPServers, listMCPTools } from "../networking";
|
||||
import type { MCPServer } from "../mcp_tools/types";
|
||||
import { setServerRootPath } from "@/lib/serverRootPath";
|
||||
|
||||
vi.mock("../networking", () => ({
|
||||
fetchMCPServers: vi.fn(),
|
||||
getMCPOAuthUserCredentialStatus: vi.fn(),
|
||||
listMCPTools: vi.fn(),
|
||||
deleteMCPOAuthUserCredential: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({
|
||||
useUserMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle" }),
|
||||
}));
|
||||
|
||||
const servers = [
|
||||
{
|
||||
server_id: "s-ext",
|
||||
server_name: "external_logo",
|
||||
auth_type: "none",
|
||||
mcp_info: { server_name: "external_logo", logo_url: "https://cdn.example.com/ext.png" },
|
||||
},
|
||||
{
|
||||
server_id: "s-local",
|
||||
server_name: "local_logo",
|
||||
auth_type: "none",
|
||||
mcp_info: { server_name: "local_logo", logo_url: "/ui/assets/logos/github.svg" },
|
||||
},
|
||||
{
|
||||
server_id: "s-none",
|
||||
server_name: "no_logo",
|
||||
auth_type: "none",
|
||||
},
|
||||
] as MCPServer[];
|
||||
|
||||
const renderPanel = () =>
|
||||
render(
|
||||
<QueryClientProvider client={new QueryClient({ defaultOptions: { queries: { retry: false } } })}>
|
||||
<MCPAppsPanel accessToken="tok" selectedServers={[]} onChange={vi.fn()} />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
|
||||
describe("MCPAppsPanel logos", () => {
|
||||
afterEach(() => {
|
||||
setServerRootPath("/");
|
||||
});
|
||||
|
||||
it("resolves backend logo_url values in the server grid", async () => {
|
||||
setServerRootPath("/litellm");
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(servers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderPanel();
|
||||
|
||||
expect(await screen.findByText("external_logo")).toBeInTheDocument();
|
||||
expect(screen.getByAltText("external_logo logo")).toHaveAttribute("src", "https://cdn.example.com/ext.png");
|
||||
expect(screen.getByAltText("local_logo logo")).toHaveAttribute("src", "/litellm/ui/assets/logos/github.svg");
|
||||
});
|
||||
|
||||
it("renders a colored letter avatar for servers without logo_url", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(servers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderPanel();
|
||||
|
||||
expect(await screen.findByText("no_logo")).toBeInTheDocument();
|
||||
expect(screen.queryByAltText("no_logo logo")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("N")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("resolves the logo_url in the detail header", async () => {
|
||||
setServerRootPath("/litellm");
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(servers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderPanel();
|
||||
|
||||
fireEvent.click(await screen.findByText("local_logo"));
|
||||
|
||||
expect(await screen.findByRole("heading", { name: "local_logo" })).toBeInTheDocument();
|
||||
expect(screen.getByAltText("local_logo logo")).toHaveAttribute("src", "/litellm/ui/assets/logos/github.svg");
|
||||
});
|
||||
});
|
||||
|
||||
const connectServers = [
|
||||
{
|
||||
server_id: "s-m2m",
|
||||
server_name: "service_tool",
|
||||
auth_type: "oauth2",
|
||||
oauth2_flow: "client_credentials",
|
||||
connected_app_reachable: true,
|
||||
},
|
||||
{
|
||||
server_id: "s-reach",
|
||||
server_name: "reachable_srv",
|
||||
auth_type: "none",
|
||||
connected_app_reachable: true,
|
||||
},
|
||||
{
|
||||
server_id: "s-unreach",
|
||||
server_name: "unreachable_srv",
|
||||
auth_type: "none",
|
||||
connected_app_reachable: false,
|
||||
},
|
||||
] as MCPServer[];
|
||||
|
||||
const renderConnectPanel = (connectMode: boolean, selectedServers: string[] = []) =>
|
||||
render(
|
||||
<QueryClientProvider client={new QueryClient({ defaultOptions: { queries: { retry: false } } })}>
|
||||
<MCPAppsPanel accessToken="tok" selectedServers={selectedServers} onChange={vi.fn()} connectMode={connectMode} />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
|
||||
describe("MCPAppsPanel connected-app reachability (LIT-4861)", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("requests the connected-app view and hides unreachable servers in connect mode", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(true, ["reachable_srv", "unreachable_srv"]);
|
||||
|
||||
expect(await screen.findByText("reachable_srv")).toBeInTheDocument();
|
||||
expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, true);
|
||||
expect(screen.queryByText("unreachable_srv")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Connected (1)")).toBeInTheDocument();
|
||||
expect(screen.getByText("service_tool")).toBeInTheDocument();
|
||||
expect(screen.queryByText("Connect", { exact: true })).not.toBeInTheDocument();
|
||||
const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]);
|
||||
expect(toolCountFetchedIds).toContain("s-reach");
|
||||
expect(toolCountFetchedIds).not.toContain("s-unreach");
|
||||
});
|
||||
|
||||
it("blocks connecting an unsupported server from the detail view in connect mode", async () => {
|
||||
const detailServers = [
|
||||
...connectServers,
|
||||
{
|
||||
server_id: "s-unsup",
|
||||
server_name: "unsupported_srv",
|
||||
auth_type: "oauth2_token_exchange",
|
||||
connected_app_reachable: true,
|
||||
},
|
||||
] as MCPServer[];
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(detailServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(true);
|
||||
|
||||
fireEvent.click(await screen.findByText("unsupported_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "unsupported_srv" })).toBeInTheDocument();
|
||||
expect(screen.queryByRole("button", { name: "Connect" })).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Not supported on this connection")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps the detail-view Connect action outside connect mode", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(false);
|
||||
|
||||
fireEvent.click(await screen.findByText("unreachable_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "unreachable_srv" })).toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: "Connect" })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("ignores the flag and skips no server outside connect mode", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValue(connectServers);
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
renderConnectPanel(false, ["reachable_srv", "unreachable_srv"]);
|
||||
|
||||
expect(await screen.findByText("unreachable_srv")).toBeInTheDocument();
|
||||
expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, false);
|
||||
expect(screen.queryByText("Not available to connected apps")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Connected (2)")).toBeInTheDocument();
|
||||
const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]);
|
||||
expect(toolCountFetchedIds).toContain("s-unreach");
|
||||
});
|
||||
|
||||
const revocable = (reachable: boolean) =>
|
||||
[
|
||||
{ server_id: "s-reach", server_name: "reachable_srv", auth_type: "none", connected_app_reachable: true },
|
||||
{ server_id: "s-drop", server_name: "revoked_srv", auth_type: "none", connected_app_reachable: reachable },
|
||||
] as MCPServer[];
|
||||
|
||||
const ConnectPanel = ({
|
||||
token,
|
||||
onChange,
|
||||
client,
|
||||
}: {
|
||||
token: string;
|
||||
onChange: (servers: string[]) => void;
|
||||
client: QueryClient;
|
||||
}) => (
|
||||
<QueryClientProvider client={client}>
|
||||
<MCPAppsPanel accessToken={token} selectedServers={[]} onChange={onChange} connectMode />
|
||||
</QueryClientProvider>
|
||||
);
|
||||
|
||||
const newClient = () => new QueryClient({ defaultOptions: { queries: { retry: false } } });
|
||||
|
||||
it("drops an open detail view when a refetch removes that server from the reachable set", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValueOnce(revocable(true)).mockResolvedValueOnce(revocable(false));
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
const client = newClient();
|
||||
const { rerender } = render(<ConnectPanel token="tok" onChange={vi.fn()} client={client} />);
|
||||
|
||||
fireEvent.click(await screen.findByText("revoked_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument();
|
||||
|
||||
rerender(<ConnectPanel token="tok-refreshed" onChange={vi.fn()} client={client} />);
|
||||
|
||||
await waitFor(() => expect(screen.queryByRole("heading", { name: "revoked_srv" })).not.toBeInTheDocument());
|
||||
expect(screen.queryByRole("button", { name: "Connect" })).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("reachable_srv")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does not select a server whose Connect finishes after a refetch removed it", async () => {
|
||||
vi.mocked(fetchMCPServers).mockResolvedValueOnce(revocable(true)).mockResolvedValueOnce(revocable(false));
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
const onChange = vi.fn();
|
||||
const client = newClient();
|
||||
const { rerender } = render(<ConnectPanel token="tok" onChange={onChange} client={client} />);
|
||||
|
||||
fireEvent.click(await screen.findByText("revoked_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument();
|
||||
|
||||
let finishConnect: (result: { tools: never[] }) => void = () => {};
|
||||
vi.mocked(listMCPTools).mockImplementationOnce(() => new Promise((resolve) => (finishConnect = resolve)));
|
||||
fireEvent.click(screen.getByRole("button", { name: "Connect" }));
|
||||
|
||||
rerender(<ConnectPanel token="tok-refreshed" onChange={onChange} client={client} />);
|
||||
await waitFor(() => expect(screen.queryByRole("heading", { name: "revoked_srv" })).not.toBeInTheDocument());
|
||||
|
||||
await act(async () => {
|
||||
finishConnect({ tools: [] });
|
||||
});
|
||||
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
|
||||
expect(screen.getByText("Connected", { exact: false })).toHaveTextContent("Connected");
|
||||
});
|
||||
|
||||
it("does not select a server when Connect resolves in the same tick the refetch drops it", async () => {
|
||||
let finishRefetch: (servers: MCPServer[]) => void = () => {};
|
||||
vi.mocked(fetchMCPServers)
|
||||
.mockResolvedValueOnce(revocable(true))
|
||||
.mockImplementationOnce(() => new Promise((resolve) => (finishRefetch = resolve)));
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
const onChange = vi.fn();
|
||||
const client = newClient();
|
||||
const { rerender } = render(<ConnectPanel token="tok" onChange={onChange} client={client} />);
|
||||
|
||||
fireEvent.click(await screen.findByText("revoked_srv"));
|
||||
expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument();
|
||||
|
||||
let finishConnect: (result: { tools: never[] }) => void = () => {};
|
||||
vi.mocked(listMCPTools).mockImplementationOnce(() => new Promise((resolve) => (finishConnect = resolve)));
|
||||
fireEvent.click(screen.getByRole("button", { name: "Connect" }));
|
||||
|
||||
rerender(<ConnectPanel token="tok-refreshed" onChange={onChange} client={client} />);
|
||||
|
||||
await act(async () => {
|
||||
finishRefetch(revocable(false));
|
||||
finishConnect({ tools: [] });
|
||||
});
|
||||
|
||||
expect(onChange).not.toHaveBeenCalled();
|
||||
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does not let a superseded list load overwrite the current reachable set", async () => {
|
||||
let finishStaleLoad: (servers: MCPServer[]) => void = () => {};
|
||||
vi.mocked(fetchMCPServers)
|
||||
.mockImplementationOnce(() => new Promise((resolve) => (finishStaleLoad = resolve)))
|
||||
.mockResolvedValueOnce(revocable(false));
|
||||
vi.mocked(listMCPTools).mockResolvedValue({ tools: [] });
|
||||
|
||||
const client = newClient();
|
||||
const { rerender } = render(<ConnectPanel token="tok" onChange={vi.fn()} client={client} />);
|
||||
rerender(<ConnectPanel token="tok-refreshed" onChange={vi.fn()} client={client} />);
|
||||
|
||||
expect(await screen.findByText("reachable_srv")).toBeInTheDocument();
|
||||
|
||||
await act(async () => {
|
||||
finishStaleLoad(revocable(true));
|
||||
});
|
||||
|
||||
expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
@ -1,12 +1,13 @@
|
|||
"use client";
|
||||
|
||||
import React, { useCallback, useEffect, useRef, useState } from "react";
|
||||
import React, { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { useQuery } from "@tanstack/react-query";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { Search, ArrowLeft, ChevronRight, Wrench, CheckCircle, Loader2 } from "lucide-react";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { Skeleton } from "@/components/ui/skeleton";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
import {
|
||||
deleteMCPOAuthUserCredential,
|
||||
|
|
@ -24,7 +25,7 @@ import {
|
|||
import { Logo } from "@/components/molecules/logo/Logo";
|
||||
import i18n from "@/i18n/i18n";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow";
|
||||
import { pendingUserMcpOAuthServerId, useUserMcpOAuthFlow, type UserMcpOAuthStatus } from "@/hooks/useUserMcpOAuthFlow";
|
||||
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
|
||||
|
||||
interface OAuth2ConnectButtonProps {
|
||||
|
|
@ -91,13 +92,58 @@ export const OAuth2ConnectButton: React.FC<OAuth2ConnectButtonProps> = ({
|
|||
);
|
||||
};
|
||||
|
||||
export interface ConnectedServers {
|
||||
on: string[];
|
||||
all: string[];
|
||||
}
|
||||
|
||||
export type ServerListReport = ConnectedServers | "loading" | "failed" | "saving-login";
|
||||
|
||||
const savingReturnedLogin = (serverId: string, status: UserMcpOAuthStatus): boolean =>
|
||||
serverId !== "" && (status === "idle" || status === "exchanging");
|
||||
|
||||
const listReport = (
|
||||
loading: boolean,
|
||||
loadFailed: boolean,
|
||||
savingLogin: boolean,
|
||||
servers: ConnectedServers,
|
||||
): ServerListReport => {
|
||||
if (savingLogin) return "saving-login";
|
||||
if (loading) return "loading";
|
||||
if (loadFailed) return "failed";
|
||||
return servers;
|
||||
};
|
||||
|
||||
interface Props {
|
||||
accessToken: string;
|
||||
selectedServers: string[];
|
||||
onChange: (servers: string[]) => void;
|
||||
connectMode?: boolean;
|
||||
offServerIds?: readonly string[];
|
||||
onOffServerIdsChange?: (serverIds: string[]) => void;
|
||||
lockedOn?: LockedOn;
|
||||
onConnectedServers?: (servers: ServerListReport) => void;
|
||||
}
|
||||
|
||||
const NO_SERVER_IDS: readonly string[] = [];
|
||||
|
||||
export type LockedOn = readonly string[] | "all" | undefined;
|
||||
|
||||
const isLockedOn = (lockedOn: LockedOn, serverId: string): boolean =>
|
||||
lockedOn === "all" || (lockedOn?.includes(serverId) ?? false);
|
||||
|
||||
const CHAT_LABELS = {
|
||||
tab: "chat.mcpAppsPanel.tabs.connected",
|
||||
on: "chat.mcpAppsPanel.info.connected",
|
||||
off: "chat.mcpAppsPanel.info.notConnected",
|
||||
} as const;
|
||||
const CONNECT_LABELS = {
|
||||
tab: "chat.mcpAppsPanel.tabs.on",
|
||||
on: "chat.mcpAppsPanel.info.on",
|
||||
off: "chat.mcpAppsPanel.info.off",
|
||||
} as const;
|
||||
const labelsFor = (connectMode: boolean | undefined) => (connectMode ? CONNECT_LABELS : CHAT_LABELS);
|
||||
|
||||
const AVATAR_COLORS = [
|
||||
"#1677ff",
|
||||
"#52c41a",
|
||||
|
|
@ -121,10 +167,20 @@ type TabKey = "all" | "connected";
|
|||
|
||||
const TOOLS_FETCH_CONCURRENCY = 5;
|
||||
|
||||
const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange, connectMode }) => {
|
||||
const MCPAppsPanel: React.FC<Props> = ({
|
||||
accessToken,
|
||||
selectedServers,
|
||||
onChange,
|
||||
connectMode,
|
||||
offServerIds = NO_SERVER_IDS,
|
||||
onOffServerIdsChange,
|
||||
lockedOn,
|
||||
onConnectedServers,
|
||||
}) => {
|
||||
const { t } = useTranslation();
|
||||
const [servers, setServers] = useState<MCPServer[]>([]);
|
||||
const [loading, setLoading] = useState(true);
|
||||
const [loadFailed, setLoadFailed] = useState(false);
|
||||
const [query, setQuery] = useState("");
|
||||
const [activeTab, setActiveTab] = useState<TabKey>("all");
|
||||
const [togglingOn, setTogglingOn] = useState<Set<string>>(new Set());
|
||||
|
|
@ -169,6 +225,15 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
[connectUnavailabilityLabel],
|
||||
);
|
||||
|
||||
const [returningServerId] = useState(() => (connectMode ? pendingUserMcpOAuthServerId() : null) ?? "");
|
||||
const returningOAuth = useUserMcpOAuthFlow({
|
||||
accessToken,
|
||||
serverId: returningServerId,
|
||||
onSuccess: useCallback(() => {
|
||||
setOauthConnected((prev) => new Set(prev).add(returningServerId));
|
||||
}, [returningServerId]),
|
||||
});
|
||||
|
||||
const fetchToolCount = useCallback(
|
||||
async (server: MCPServer, isCurrentLoad: () => boolean) => {
|
||||
try {
|
||||
|
|
@ -218,6 +283,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
const oauthServers = reachable.filter((s) => getMcpOAuthMode(s) === "authorization_code");
|
||||
commitServers(reachable);
|
||||
setOauthChecking(new Set(oauthServers.map((s) => s.server_id)));
|
||||
setLoadFailed(false);
|
||||
setLoading(false);
|
||||
|
||||
oauthServers.forEach((s) => checkOauthCredential(s, isCurrentLoad));
|
||||
|
|
@ -235,6 +301,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
.catch(() => {
|
||||
if (isCurrentLoad()) {
|
||||
commitServers([]);
|
||||
setLoadFailed(true);
|
||||
setLoading(false);
|
||||
}
|
||||
});
|
||||
|
|
@ -292,6 +359,63 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
}
|
||||
};
|
||||
|
||||
const isOn = (server: MCPServer) => {
|
||||
if (connectUnavailabilityLabel(server) !== null) return false;
|
||||
if (!connectMode) return selectedServers.includes(nameOf(server));
|
||||
return isLockedOn(lockedOn, server.server_id) || !offServerIds.includes(server.server_id);
|
||||
};
|
||||
|
||||
const setServerOn = (server: MCPServer, on: boolean) =>
|
||||
onOffServerIdsChange?.(
|
||||
on ? offServerIds.filter((id) => id !== server.server_id) : [...offServerIds, server.server_id],
|
||||
);
|
||||
|
||||
const renderServerSwitch = (server: MCPServer) => (
|
||||
<span className="ml-2 inline-flex items-center gap-2" onClick={(e) => e.stopPropagation()}>
|
||||
{lockedOn !== "all" && isLockedOn(lockedOn, server.server_id) && (
|
||||
<span className="text-[11px] text-muted-foreground whitespace-nowrap">
|
||||
{t("chat.mcpAppsPanel.requiredByClient")}
|
||||
</span>
|
||||
)}
|
||||
<Switch
|
||||
checked={isOn(server)}
|
||||
disabled={isLockedOn(lockedOn, server.server_id)}
|
||||
onCheckedChange={(checked) => setServerOn(server, checked)}
|
||||
aria-label={t("chat.mcpAppsPanel.useServer", { serverName: nameOf(server) })}
|
||||
/>
|
||||
</span>
|
||||
);
|
||||
|
||||
const renderLoginState = (server: MCPServer) => {
|
||||
const mode = getMcpOAuthMode(server);
|
||||
if (mode === "m2m") return <CheckCircle className="h-3.5 w-3.5 text-success shrink-0" />;
|
||||
if (mode !== "authorization_code") return null;
|
||||
if (oauthConnected.has(server.server_id)) {
|
||||
return <CheckCircle className="h-3.5 w-3.5 text-success shrink-0" />;
|
||||
}
|
||||
if (
|
||||
oauthChecking.has(server.server_id) ||
|
||||
(returningOAuth.status === "exchanging" && server.server_id === returningServerId)
|
||||
) {
|
||||
return <Skeleton className="h-6 w-16 shrink-0 rounded-md" />;
|
||||
}
|
||||
return (
|
||||
<>
|
||||
{connectMode && (
|
||||
<span className="text-[11px] text-muted-foreground shrink-0 whitespace-nowrap">
|
||||
{t("chat.mcpAppsPanel.signInToUse")}
|
||||
</span>
|
||||
)}
|
||||
<OAuth2ConnectButton
|
||||
server={server}
|
||||
accessToken={accessToken}
|
||||
onConnect={(id) => setOauthConnected((prev) => new Set(prev).add(id))}
|
||||
variant="badge"
|
||||
/>
|
||||
</>
|
||||
);
|
||||
};
|
||||
|
||||
const renderConnectionIndicator = (server: MCPServer) => {
|
||||
const unavailabilityLabel = connectUnavailabilityLabel(server);
|
||||
if (unavailabilityLabel !== null) {
|
||||
|
|
@ -299,25 +423,16 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
<span className="text-[11px] text-muted-foreground shrink-0 whitespace-nowrap">{unavailabilityLabel}</span>
|
||||
);
|
||||
}
|
||||
if (getMcpOAuthMode(server) === "m2m") {
|
||||
return <CheckCircle className="h-3.5 w-3.5 text-success shrink-0" />;
|
||||
}
|
||||
if (getMcpOAuthMode(server) === "authorization_code") {
|
||||
if (oauthConnected.has(server.server_id)) {
|
||||
return <CheckCircle className="h-3.5 w-3.5 text-success shrink-0" />;
|
||||
}
|
||||
if (oauthChecking.has(server.server_id)) {
|
||||
return <Skeleton className="h-6 w-16 shrink-0 rounded-md" />;
|
||||
}
|
||||
const loginState = renderLoginState(server);
|
||||
if (connectMode) {
|
||||
return (
|
||||
<OAuth2ConnectButton
|
||||
server={server}
|
||||
accessToken={accessToken}
|
||||
onConnect={(id) => setOauthConnected((prev) => new Set(prev).add(id))}
|
||||
variant="badge"
|
||||
/>
|
||||
<div className="flex items-center gap-2 shrink-0">
|
||||
{loginState}
|
||||
{renderServerSwitch(server)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
if (loginState !== null) return loginState;
|
||||
if (selectedServers.includes(nameOf(server))) {
|
||||
return <span className="w-[7px] h-[7px] rounded-full bg-success shrink-0" />;
|
||||
}
|
||||
|
|
@ -337,16 +452,25 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
!query.trim() ||
|
||||
name.toLowerCase().includes(query.toLowerCase()) ||
|
||||
(s.description ?? "").toLowerCase().includes(query.toLowerCase());
|
||||
const matchesTab =
|
||||
activeTab === "all" || (selectedServers.includes(name) && connectUnavailabilityLabel(s) === null);
|
||||
const matchesTab = activeTab === "all" || isOn(s);
|
||||
return matchesQuery && matchesTab;
|
||||
});
|
||||
|
||||
const connectedCount = servers.filter(
|
||||
(s) => selectedServers.includes(nameOf(s)) && connectUnavailabilityLabel(s) === null,
|
||||
).length;
|
||||
const connectedCount = servers.filter(isOn).length;
|
||||
const labels = labelsFor(connectMode);
|
||||
|
||||
const connectedServers = useMemo(() => {
|
||||
const all = servers.filter((s) => connectUnavailabilityLabel(s) === null).map((s) => s.server_id);
|
||||
return { on: all.filter((id) => isLockedOn(lockedOn, id) || !offServerIds.includes(id)), all };
|
||||
}, [servers, offServerIds, lockedOn, connectUnavailabilityLabel]);
|
||||
useEffect(() => {
|
||||
onConnectedServers?.(
|
||||
listReport(loading, loadFailed, savingReturnedLogin(returningServerId, returningOAuth.status), connectedServers),
|
||||
);
|
||||
}, [loading, loadFailed, returningServerId, returningOAuth.status, connectedServers, onConnectedServers]);
|
||||
|
||||
const emptyStateText = () => {
|
||||
if (loadFailed && connectMode) return t("chat.mcpAppsPanel.empty.loadFailed");
|
||||
if (servers.length === 0) {
|
||||
return connectMode ? t("chat.mcpAppsPanel.empty.connectMode") : t("chat.mcpAppsPanel.empty.noServers");
|
||||
}
|
||||
|
|
@ -358,19 +482,16 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
|
||||
if (detailServer) {
|
||||
const name = nameOf(detailServer);
|
||||
const isConnected = selectedServers.includes(name);
|
||||
const isConnected = isOn(detailServer);
|
||||
const isTogglingOn = togglingOn.has(name);
|
||||
const color = getAvatarColor(name);
|
||||
|
||||
const renderDetailAction = () => {
|
||||
const unavailabilityLabel = connectUnavailabilityLabel(detailServer);
|
||||
if (unavailabilityLabel !== null) {
|
||||
return <span className="text-[13px] text-muted-foreground py-2.5 shrink-0">{unavailabilityLabel}</span>;
|
||||
}
|
||||
const renderDetailLogin = () => {
|
||||
if (getMcpOAuthMode(detailServer) === "m2m") {
|
||||
return <span className="text-[13px] text-muted-foreground">{t("chat.mcpAppsPanel.authorized")}</span>;
|
||||
}
|
||||
if (getMcpOAuthMode(detailServer) !== "authorization_code") {
|
||||
if (connectMode) return null;
|
||||
return (
|
||||
<Button
|
||||
variant={isConnected ? "outline" : "default"}
|
||||
|
|
@ -398,7 +519,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
n.delete(detailServer.server_id);
|
||||
return n;
|
||||
});
|
||||
onChangeRef.current(selectedServersRef.current.filter((s) => s !== name));
|
||||
if (!connectMode) onChangeRef.current(selectedServersRef.current.filter((s) => s !== name));
|
||||
}}
|
||||
className="font-semibold h-[38px] min-w-[110px]"
|
||||
>
|
||||
|
|
@ -418,6 +539,20 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
);
|
||||
};
|
||||
|
||||
const renderDetailAction = () => {
|
||||
const unavailabilityLabel = connectUnavailabilityLabel(detailServer);
|
||||
if (unavailabilityLabel !== null) {
|
||||
return <span className="text-[13px] text-muted-foreground py-2.5 shrink-0">{unavailabilityLabel}</span>;
|
||||
}
|
||||
if (!connectMode) return renderDetailLogin();
|
||||
return (
|
||||
<div className="flex items-center gap-3 shrink-0">
|
||||
{renderDetailLogin()}
|
||||
{renderServerSwitch(detailServer)}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="w-full">
|
||||
<Button
|
||||
|
|
@ -459,10 +594,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
{[
|
||||
[t("chat.mcpAppsPanel.info.serverId"), detailServer.server_id],
|
||||
[t("chat.mcpAppsPanel.info.transport"), handleTransport(detailServer.transport, detailServer.spec_path)],
|
||||
[
|
||||
t("chat.mcpAppsPanel.info.status"),
|
||||
isConnected ? t("chat.mcpAppsPanel.info.connected") : t("chat.mcpAppsPanel.info.notConnected"),
|
||||
],
|
||||
[t("chat.mcpAppsPanel.info.status"), isConnected ? t(labels.on) : t(labels.off)],
|
||||
]
|
||||
.filter(([, v]) => v)
|
||||
.map(([label, value], i, arr) => (
|
||||
|
|
@ -557,7 +689,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange,
|
|||
{t("chat.mcpAppsPanel.tabs.all")}
|
||||
</TabsTrigger>
|
||||
<TabsTrigger value="connected" className="rounded-none px-4 py-2 text-[13px]">
|
||||
{t("chat.mcpAppsPanel.tabs.connected")}
|
||||
{t(labels.tab)}
|
||||
{connectedCount > 0 ? ` (${connectedCount})` : ""}
|
||||
</TabsTrigger>
|
||||
</TabsList>
|
||||
|
|
|
|||
|
|
@ -4456,6 +4456,8 @@ export interface ConnectFlowStatus {
|
|||
connected: boolean | null;
|
||||
requested_scopes?: string[];
|
||||
servers: { server_id: string; server_name: string; connected: boolean }[] | null;
|
||||
selectable?: false;
|
||||
required_server_ids?: string[];
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import { renderHook, waitFor } from "@testing-library/react";
|
|||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import * as networking from "@/components/networking";
|
||||
import { setSecureItem } from "@/utils/secureStorage";
|
||||
import { useUserMcpOAuthFlow } from "./useUserMcpOAuthFlow";
|
||||
import { pendingUserMcpOAuthServerId, useUserMcpOAuthFlow } from "./useUserMcpOAuthFlow";
|
||||
|
||||
vi.mock("@/components/networking", () => ({
|
||||
exchangeMcpOAuthToken: vi.fn(),
|
||||
|
|
@ -68,4 +68,42 @@ describe("user OAuth granted scopes", () => {
|
|||
);
|
||||
expect(onSuccess).toHaveBeenCalledOnce();
|
||||
});
|
||||
|
||||
it("keeps the returned login until there is a token to save it with", async () => {
|
||||
vi.mocked(networking.exchangeMcpOAuthToken).mockResolvedValue({
|
||||
access_token: "vendor-token",
|
||||
refresh_token: "vendor-refresh",
|
||||
});
|
||||
const onSuccess = vi.fn();
|
||||
const { result, rerender } = renderHook(
|
||||
({ accessToken }) => useUserMcpOAuthFlow({ accessToken, serverId: "server-1", onSuccess }),
|
||||
{ initialProps: { accessToken: "" } },
|
||||
);
|
||||
|
||||
expect(networking.exchangeMcpOAuthToken).not.toHaveBeenCalled();
|
||||
expect(pendingUserMcpOAuthServerId()).toBe("server-1");
|
||||
|
||||
rerender({ accessToken: "caller" });
|
||||
|
||||
await waitFor(() => expect(result.current.status).toBe("success"));
|
||||
expect(networking.storeMCPOAuthUserCredential).toHaveBeenCalledWith("caller", "server-1", expect.anything());
|
||||
expect(onSuccess).toHaveBeenCalledOnce();
|
||||
});
|
||||
});
|
||||
|
||||
describe("pending user OAuth return", () => {
|
||||
beforeEach(() => {
|
||||
window.sessionStorage.clear();
|
||||
setSecureItem("litellm-user-mcp-oauth-flow-state", JSON.stringify({ serverId: "server-1" }));
|
||||
});
|
||||
|
||||
it("names the server whose login came back", () => {
|
||||
setSecureItem("litellm-user-mcp-oauth-result", JSON.stringify({ code: "provider-code" }));
|
||||
|
||||
expect(pendingUserMcpOAuthServerId()).toBe("server-1");
|
||||
});
|
||||
|
||||
it("names nothing while the login has not come back", () => {
|
||||
expect(pendingUserMcpOAuthServerId()).toBeNull();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -69,6 +69,15 @@ const getStorage = (key: string): string | null => {
|
|||
return getSecureItem(key);
|
||||
};
|
||||
|
||||
export const pendingUserMcpOAuthServerId = (): string | null => {
|
||||
if (typeof window === "undefined" || getStorage(RESULT_KEY) === null) return null;
|
||||
try {
|
||||
return (JSON.parse(getStorage(FLOW_STATE_KEY) ?? "null") as StoredFlowState | null)?.serverId ?? null;
|
||||
} catch (_) {
|
||||
return null;
|
||||
}
|
||||
};
|
||||
|
||||
const grantedScopes = (scope: unknown, requested: string[] | undefined): string[] | undefined =>
|
||||
typeof scope === "string" ? scope.split(" ").filter(Boolean) : requested;
|
||||
|
||||
|
|
@ -150,7 +159,7 @@ export const useUserMcpOAuthFlow = ({
|
|||
}, [accessToken, serverId, serverAlias, scopes, preClientId]);
|
||||
|
||||
const resumeOAuthFlow = useCallback(async () => {
|
||||
if (typeof window === "undefined" || processingRef.current) return;
|
||||
if (typeof window === "undefined" || processingRef.current || !accessToken) return;
|
||||
|
||||
const storedResult = getStorage(RESULT_KEY);
|
||||
if (!storedResult) return;
|
||||
|
|
|
|||
|
|
@ -9383,11 +9383,15 @@
|
|||
"toolsLoadFailed": "Could not load tools for {{serverName}}",
|
||||
"empty": {
|
||||
"connectMode": "No MCP servers are available to this connection yet. Ask an admin to grant your user or team access.",
|
||||
"loadFailed": "Could not load your MCP servers. Reload the page to try again.",
|
||||
"noServers": "No MCP servers configured. Add servers in Tools -> MCP Servers.",
|
||||
"noneConnected": "No servers connected yet.",
|
||||
"noMatch": "No servers match your search."
|
||||
},
|
||||
"authorized": "Authorized",
|
||||
"signInToUse": "Sign in to use",
|
||||
"requiredByClient": "Required by this client",
|
||||
"useServer": "Include {{serverName}} in this connection",
|
||||
"disconnect": "Disconnect",
|
||||
"back": "Back",
|
||||
"defaultDescription": "MCP server",
|
||||
|
|
@ -9397,13 +9401,15 @@
|
|||
"transport": "Transport",
|
||||
"status": "Status",
|
||||
"connected": "Connected",
|
||||
"notConnected": "Not connected"
|
||||
"notConnected": "Not connected",
|
||||
"on": "On",
|
||||
"off": "Off"
|
||||
},
|
||||
"availableTools": "Available Tools",
|
||||
"noTools": "No tools available",
|
||||
"title": "MCP Servers",
|
||||
"beta": "Beta",
|
||||
"connectModeSubtitle": "Click a server to see its tools and connect",
|
||||
"connectModeSubtitle": "Turn off any server this connection should not use. Click a server to see its tools or sign in",
|
||||
"subtitle": "Browse tools, authenticate once, use in chat",
|
||||
"loadingTools": "Loading tools...",
|
||||
"toolsAvailable_one": "{{count}} tool available",
|
||||
|
|
@ -9411,7 +9417,8 @@
|
|||
"searchPlaceholder": "Search servers...",
|
||||
"tabs": {
|
||||
"all": "All",
|
||||
"connected": "Connected"
|
||||
"connected": "Connected",
|
||||
"on": "On"
|
||||
}
|
||||
},
|
||||
"connectFlowBanner": {
|
||||
|
|
@ -9423,7 +9430,8 @@
|
|||
},
|
||||
"unscoped": {
|
||||
"title": "Connect your MCP servers to {{client}}",
|
||||
"description": "Authorize the servers you want to use below, then click Finish connecting to return to {{client}}."
|
||||
"description": "{{client}} gets every server that is on. Turn off the ones it should not use, sign in where needed, then click Finish connecting.",
|
||||
"fixed": "{{client}} will get every server you can reach. Sign in to the ones that need it below, then click Finish connecting to return to {{client}}."
|
||||
},
|
||||
"multi": {
|
||||
"title": "Allow {{client}} to use {{servers}}",
|
||||
|
|
@ -9437,6 +9445,7 @@
|
|||
},
|
||||
"connectServer": "Connect {{server}}",
|
||||
"finish": "Finish connecting",
|
||||
"nothingOn": "Turn on at least one server to finish connecting.",
|
||||
"cancel": "Cancel",
|
||||
"manualDelivery": "My client is on a remote or SSH machine"
|
||||
},
|
||||
|
|
|
|||
6
litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
6
litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -1408,6 +1408,7 @@ export interface paths {
|
|||
* (``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.
|
||||
* Each ``server`` field is the id of a server the connect page shows as on.
|
||||
*/
|
||||
post: operations["authorize_complete_authorize_complete_post"];
|
||||
delete?: never;
|
||||
|
|
@ -29236,6 +29237,11 @@ export interface components {
|
|||
delivery?: string | null;
|
||||
/** Flow */
|
||||
flow: string;
|
||||
/**
|
||||
* Server
|
||||
* @default []
|
||||
*/
|
||||
server: string[];
|
||||
/** Team Id */
|
||||
team_id?: string | null;
|
||||
};
|
||||
|
|
|
|||
|
|
@ -12211,6 +12211,90 @@ def test_gateway_connect_rechecks_session_after_authorize(
|
|||
assert parse_qs(urlparse(response.headers["location"]).query)["code"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("resource", "sealed_ids", "selectable"),
|
||||
[
|
||||
(None, ("alpha-id", "beta-id"), "absent"),
|
||||
("http://testserver", ("alpha-id", "beta-id"), "absent"),
|
||||
("http://testserver/mcp", ("alpha-id", "beta-id"), "absent"),
|
||||
("http://testserver/mcp/alpha,beta", None, False),
|
||||
],
|
||||
)
|
||||
def test_gateway_connect_complete_reads_each_picked_server(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
browser_session_database: _BrowserSessionDatabase,
|
||||
resource: str | None,
|
||||
sealed_ids: tuple[str, ...] | None,
|
||||
selectable: Literal["absent", False],
|
||||
) -> None:
|
||||
import json
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
|
||||
from litellm.models.user import LiteLLM_UserTable
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
|
||||
from litellm.proxy.mcp_server.gateway_dcr_flow import GATEWAY_AUTH_CODE_PREFIX
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
monkeypatch.setattr(global_mcp_server_manager, "registry", global_mcp_server_manager.registry.copy())
|
||||
client, _, _ = _native_client_app(monkeypatch)
|
||||
user_api_key_cache.set_cache(
|
||||
"u1",
|
||||
LiteLLM_UserTable(
|
||||
user_id="u1",
|
||||
object_permission=LiteLLM_ObjectPermissionTable(
|
||||
object_permission_id="u1-mcp", mcp_servers=["alpha-id", "beta-id"]
|
||||
),
|
||||
),
|
||||
)
|
||||
for name in ("alpha", "beta"):
|
||||
server = _create_oauth2_server(server_id=f"{name}-id", name=name, server_name=name, alias=name)
|
||||
monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server)
|
||||
client.cookies.set("token", _browser_session_cookie(monkeypatch, "u1", session_key="sk-browser-session"))
|
||||
registered: Final = client.post("/register", json={"redirect_uris": ["http://127.0.0.1/callback"]})
|
||||
client_id: Final = registered.json()["client_id"]
|
||||
resource_param: Final = {"resource": resource} if resource else {}
|
||||
authorized: Final = client.get(
|
||||
"/authorize/mcp-session",
|
||||
params={
|
||||
"client_id": client_id,
|
||||
"redirect_uri": "http://127.0.0.1/callback",
|
||||
"response_type": "code",
|
||||
"code_challenge": _s256("v" * 43),
|
||||
"code_challenge_method": "S256",
|
||||
**resource_param,
|
||||
},
|
||||
follow_redirects=False,
|
||||
)
|
||||
flow: Final = parse_qs(urlparse(authorized.headers["location"]).query)["connect_flow"][0]
|
||||
described: Final = client.get("/authorize/flow", params={"flow": flow})
|
||||
assert described.json().get("selectable", "absent") == selectable
|
||||
response: Final = client.post(
|
||||
"/authorize/complete", data={"flow": flow, "server": ["alpha-id", "beta-id"]}, follow_redirects=False
|
||||
)
|
||||
assert response.status_code == 303, response.text
|
||||
code: Final = parse_qs(urlparse(response.headers["location"]).query)["code"][0]
|
||||
sealed: Final = decrypt_value_helper(
|
||||
code.removeprefix(GATEWAY_AUTH_CODE_PREFIX), "gateway_authorization_code", return_original_value=False
|
||||
)
|
||||
assert isinstance(sealed, str)
|
||||
granted: Final = json.loads(sealed).get("resource_server_ids")
|
||||
assert (None if granted is None else tuple(granted)) == sealed_ids
|
||||
token: Final = client.post(
|
||||
"/token",
|
||||
data={
|
||||
"grant_type": "authorization_code",
|
||||
"client_id": client_id,
|
||||
"code": code,
|
||||
"redirect_uri": "http://127.0.0.1/callback",
|
||||
"code_verifier": "v" * 43,
|
||||
**resource_param,
|
||||
},
|
||||
)
|
||||
assert token.status_code == 200, token.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("authorize_path", ["/authorize/mcp-session", "/authorize"])
|
||||
def test_gateway_authorize_routes_forward_server_scope(monkeypatch: pytest.MonkeyPatch, authorize_path: str) -> None:
|
||||
from http.cookies import SimpleCookie
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import hashlib
|
|||
import json
|
||||
import re
|
||||
from base64 import urlsafe_b64encode
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from http.cookies import SimpleCookie
|
||||
from typing import Final
|
||||
|
|
@ -28,6 +29,7 @@ from litellm.proxy.mcp_server.gateway_dcr_flow import (
|
|||
TOKEN_EXCHANGE_GRANT_TYPE,
|
||||
AssertionIdentity,
|
||||
ConsentTeam,
|
||||
LookupServerReachability,
|
||||
MintedProxyCredential,
|
||||
SubjectIdentity,
|
||||
SubjectTokenRefusal,
|
||||
|
|
@ -63,6 +65,7 @@ from litellm.proxy.mcp_server.outbound_credentials.session_token import (
|
|||
mint_session_refresh_token,
|
||||
mint_session_token,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
MASTER_KEY = "sk-gateway-dcr-flow-tests"
|
||||
REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback"
|
||||
|
|
@ -3499,3 +3502,336 @@ def test_aggregate_metadata_advertises_basic_client_auth_only_with_the_grant_tha
|
|||
assert isinstance(methods, list)
|
||||
assert ("client_secret_basic" in methods) is jwt_bearer_available
|
||||
assert "none" in methods and "client_secret_post" in methods
|
||||
|
||||
|
||||
async def _authorized_flow(servers: Sequence[MCPServer], scope: str | None = None) -> tuple[str, Response]:
|
||||
from unittest.mock import patch
|
||||
|
||||
client_id: Final = (await _register([REDIRECT_URI]))["client_id"]
|
||||
with patch(_MANAGER_PATCH) as manager:
|
||||
manager.get_mcp_server_by_identifier.side_effect = {server.name: server for server in servers}.get
|
||||
return client_id, _scoped_authorize(client_id, "https://llm.example.com/mcp", scope=scope)
|
||||
|
||||
|
||||
async def _finish(
|
||||
authorized: Response,
|
||||
selected: Sequence[str],
|
||||
servers: Sequence[MCPServer],
|
||||
cache: DualCache,
|
||||
decision: str | None = None,
|
||||
reachable: LookupServerReachability | None = None,
|
||||
) -> Response:
|
||||
from unittest.mock import patch
|
||||
|
||||
handle, cookies = _flow_cookie_from(authorized)
|
||||
|
||||
async def vendor_credential(user_id: str, server_id: str, requested_scopes: tuple[str, ...] = ()) -> str:
|
||||
return "present"
|
||||
|
||||
with patch(_MANAGER_PATCH) as manager:
|
||||
manager.get_mcp_server_by_identifier.side_effect = {server.name: server for server in servers}.get
|
||||
manager.get_mcp_server_by_id.side_effect = {server.server_id: server for server in servers}.get
|
||||
manager.get_registry.return_value = {server.server_id: server for server in servers}
|
||||
return await complete_connect_flow(
|
||||
request=_request("/authorize/complete", cookies=cookies, method="POST"),
|
||||
flow_handle=handle,
|
||||
session_user_id="u1",
|
||||
cache=cache,
|
||||
decision=decision,
|
||||
lookup_vendor_credential=vendor_credential,
|
||||
lookup_server_reachability=reachable or _ServerReachability(),
|
||||
selected_server_ids=selected,
|
||||
master_key=MASTER_KEY,
|
||||
)
|
||||
|
||||
|
||||
async def _complete_with_selection(
|
||||
selected: Sequence[str],
|
||||
servers: Sequence[MCPServer],
|
||||
decision: str | None = None,
|
||||
scope: str | None = None,
|
||||
reachable: LookupServerReachability | None = None,
|
||||
) -> tuple[str, Response]:
|
||||
client_id, authorized = await _authorized_flow(servers, scope)
|
||||
return client_id, await _finish(authorized, selected, servers, DualCache(), decision, reachable)
|
||||
|
||||
|
||||
def _sealed_code_from(completed: Response) -> tuple[str, Mapping[str, object]]:
|
||||
code: Final = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
|
||||
return code, _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")
|
||||
|
||||
|
||||
async def _session_servers(code: str, client_id: str) -> tuple[str, ...] | None:
|
||||
token_response: Final = await _redeem(code, client_id, resource="https://llm.example.com/mcp")
|
||||
assert token_response.status_code == 200
|
||||
return _opened_principal(json.loads(token_response.body)).resource_server_ids
|
||||
|
||||
|
||||
CATALOG: Final = _scoped_mcp_server(name="catalog", auth_type="api_key")
|
||||
WAREHOUSE: Final = _scoped_mcp_server(name="warehouse", oauth2_flow="client_credentials")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"picked",
|
||||
[
|
||||
_scoped_mcp_server(name="tracker"),
|
||||
_scoped_mcp_server(name="delegated", delegate_auth_to_upstream=True),
|
||||
_scoped_mcp_server(name="relayed", per_server_oauth_discovery=True),
|
||||
_scoped_mcp_server(name="headers", auth_type="api_key", extra_headers=["Authorization"]),
|
||||
WAREHOUSE,
|
||||
],
|
||||
)
|
||||
async def test_picked_servers_scope_the_session_token(picked: MCPServer) -> None:
|
||||
inventory: Final = _scoped_mcp_server(name="inventory", auth_type="api_key")
|
||||
client_id, completed = await _complete_with_selection(
|
||||
("catalog-id", picked.server_id, "catalog-id"), (CATALOG, picked, inventory)
|
||||
)
|
||||
assert completed.status_code == 303
|
||||
code, sealed = _sealed_code_from(completed)
|
||||
assert sealed["resource_server_ids"] == ["catalog-id", picked.server_id]
|
||||
assert await _session_servers(code, client_id) == ("catalog-id", picked.server_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_servers_that_were_not_picked_stay_out() -> None:
|
||||
unpicked: Final = _scoped_mcp_server(name="tracker")
|
||||
client_id, completed = await _complete_with_selection(("catalog-id",), (CATALOG, WAREHOUSE, unpicked))
|
||||
assert completed.status_code == 303
|
||||
code, _ = _sealed_code_from(completed)
|
||||
assert await _session_servers(code, client_id) == ("catalog-id",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("selected", [(), ("",), (" ", "\t")])
|
||||
async def test_without_a_pick_the_unscoped_flow_mints_an_unscoped_code(selected: tuple[str, ...]) -> None:
|
||||
_, completed = await _complete_with_selection(selected, (CATALOG, WAREHOUSE))
|
||||
assert completed.status_code == 303
|
||||
_, sealed = _sealed_code_from(completed)
|
||||
assert "resource_server_ids" not in sealed
|
||||
assert "resource_server_id" not in sealed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_picks_of_deleted_servers_are_dropped() -> None:
|
||||
client_id, completed = await _complete_with_selection(("deleted-id", "catalog-id"), (CATALOG,))
|
||||
assert completed.status_code == 303
|
||||
code, _ = _sealed_code_from(completed)
|
||||
assert await _session_servers(code, client_id) == ("catalog-id",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finish_refuses_when_every_pick_is_gone() -> None:
|
||||
_, completed = await _complete_with_selection(("deleted-id",), (CATALOG, WAREHOUSE))
|
||||
assert completed.status_code == 400
|
||||
assert json.loads(completed.body) == {
|
||||
"error": "invalid_request",
|
||||
"error_description": "the requested MCP server is no longer available",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_refused_finish_can_be_retried_with_picks_that_can_be_granted() -> None:
|
||||
client_id, authorized = await _authorized_flow((CATALOG,))
|
||||
cache: Final = DualCache()
|
||||
refused: Final = await _finish(authorized, ("deleted-id",), (CATALOG,), cache)
|
||||
assert refused.status_code == 400
|
||||
|
||||
retried: Final = await _finish(authorized, ("catalog-id",), (CATALOG,), cache)
|
||||
|
||||
assert retried.status_code == 303
|
||||
code, _ = _sealed_code_from(retried)
|
||||
assert await _session_servers(code, client_id) == ("catalog-id",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("selected", [("catalog-id",), ("deleted-id", "catalog-id")])
|
||||
async def test_finish_refuses_when_every_known_pick_is_inaccessible(selected: tuple[str, ...]) -> None:
|
||||
_, completed = await _complete_with_selection(selected, (CATALOG, WAREHOUSE), reachable=_ServerReachability(False))
|
||||
assert completed.status_code == 400
|
||||
assert json.loads(completed.body) == {
|
||||
"error": "invalid_request",
|
||||
"error_description": "you do not have access to the requested MCP server",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inaccessible_picks_are_dropped() -> None:
|
||||
orders: Final = _scoped_mcp_server(name="orders", auth_type="api_key")
|
||||
|
||||
async def reachable(user_id: str, server_id: str) -> bool:
|
||||
return user_id == "u1" and server_id == CATALOG.server_id
|
||||
|
||||
client_id, completed = await _complete_with_selection(
|
||||
("catalog-id", "orders-id", "deleted-id"), (CATALOG, orders, WAREHOUSE), reachable=reachable
|
||||
)
|
||||
assert completed.status_code == 303
|
||||
code, _ = _sealed_code_from(completed)
|
||||
assert await _session_servers(code, client_id) == (CATALOG.server_id,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_server_picked_twice_is_checked_once() -> None:
|
||||
reachable: Final = _ServerReachability()
|
||||
client_id, completed = await _complete_with_selection(
|
||||
("catalog-id", "catalog-id"), (CATALOG, WAREHOUSE), reachable=reachable
|
||||
)
|
||||
assert completed.status_code == 303
|
||||
assert reachable.calls == [("u1", "catalog-id")]
|
||||
code, _ = _sealed_code_from(completed)
|
||||
assert await _session_servers(code, client_id) == ("catalog-id",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_pick_grants_only_that_server_when_another_shares_its_name() -> None:
|
||||
twin: Final = CATALOG.model_copy(update={"server_id": "catalog-db-id"})
|
||||
client_id, completed = await _complete_with_selection((twin.server_id,), (CATALOG, twin))
|
||||
assert completed.status_code == 303
|
||||
code, _ = _sealed_code_from(completed)
|
||||
assert await _session_servers(code, client_id) == (twin.server_id,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_server_name_is_not_taken_as_a_pick() -> None:
|
||||
_, completed = await _complete_with_selection((CATALOG.name,), (CATALOG,))
|
||||
assert completed.status_code == 400
|
||||
assert json.loads(completed.body)["error_description"] == "the requested MCP server is no longer available"
|
||||
|
||||
|
||||
def _fleet_server(index: int) -> MCPServer:
|
||||
return _scoped_mcp_server(name=f"tool{index}", auth_type="api_key").model_copy(
|
||||
update={"server_id": f"00000000-0000-4000-8000-{index:012d}"}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_picks_too_many_for_a_session_token_are_refused() -> None:
|
||||
picked: Final = tuple(_fleet_server(index) for index in range(80))
|
||||
_, completed = await _complete_with_selection(tuple(server.server_id for server in picked), picked)
|
||||
assert completed.status_code == 400
|
||||
assert json.loads(completed.body)["error_description"] == (
|
||||
"too many MCP servers are turned on for one connection, turn some off and try again"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_grant_is_refused_exactly_when_its_refresh_tokens_would_not_fit() -> None:
|
||||
fleet: Final = tuple(_fleet_server(index) for index in range(69))
|
||||
finishes: Final = tuple(
|
||||
[
|
||||
await _complete_with_selection(tuple(server.server_id for server in fleet[:count]), fleet[:count])
|
||||
for count in (66, 67, 68, 69)
|
||||
]
|
||||
)
|
||||
assert [completed.status_code for _, completed in finishes] == [303, 303, 400, 400]
|
||||
client_id, largest_fit = finishes[1]
|
||||
code, _ = _sealed_code_from(largest_fit)
|
||||
cache: Final = DualCache()
|
||||
redeemed: Final = await _redeem(code, client_id, cache=cache, resource="https://llm.example.com/mcp")
|
||||
assert redeemed.status_code == 200
|
||||
rotated: Final = await _redeem(
|
||||
None,
|
||||
client_id,
|
||||
cache=cache,
|
||||
grant_type="refresh_token",
|
||||
refresh_token=json.loads(redeemed.body)["refresh_token"],
|
||||
resource="https://llm.example.com/mcp",
|
||||
)
|
||||
assert rotated.status_code == 200
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_already_scoped_flow_ignores_picks() -> None:
|
||||
alpha: Final = _scoped_mcp_server(name="alpha", auth_type="api_key")
|
||||
client_id, completed = await _complete_with_selection(
|
||||
("catalog-id",), (alpha, CATALOG), scope="litellm:mcp_server:alpha"
|
||||
)
|
||||
assert completed.status_code == 303
|
||||
code, sealed = _sealed_code_from(completed)
|
||||
assert sealed["resource_server_id"] == "alpha-id"
|
||||
assert await _session_servers(code, client_id) == ("alpha-id",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_native_client_consent_flow_ignores_picks() -> None:
|
||||
client_id: Final = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
|
||||
consent: Final = await _native_authorize(client_id)
|
||||
completed: Final = await _complete_consent(consent, team_id="team-b", selected_server_ids=("catalog-id",))
|
||||
assert completed.status_code == 303
|
||||
sealed: Final = _sealed_wire_json(_code_from(completed), GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")
|
||||
assert "resource_server_ids" not in sealed
|
||||
assert "resource_server_id" not in sealed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_selection_on_an_upstream_step_up_keeps_the_step_up_server() -> None:
|
||||
from litellm.proxy.mcp_server.oauth_utils import encode_upstream_scope
|
||||
|
||||
tracker: Final = _scoped_mcp_server(name="tracker")
|
||||
client_id, completed = await _complete_with_selection(
|
||||
("catalog-id",), (tracker, CATALOG), scope=encode_upstream_scope(tracker.server_id, "write")
|
||||
)
|
||||
assert completed.status_code == 303
|
||||
code, sealed = _sealed_code_from(completed)
|
||||
assert sealed["resource_server_ids"] == ["tracker-id", "catalog-id"]
|
||||
assert await _session_servers(code, client_id) == ("tracker-id", "catalog-id")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("resource", "step_up", "selectable", "required"),
|
||||
[
|
||||
("https://llm.example.com/mcp/alpha,beta", False, False, "absent"),
|
||||
("https://llm.example.com/mcp/alpha,beta", True, False, "absent"),
|
||||
("https://llm.example.com/mcp", True, "absent", ["tracker-id"]),
|
||||
("https://llm.example.com/mcp", False, "absent", "absent"),
|
||||
],
|
||||
)
|
||||
async def test_the_page_is_told_which_switches_it_cannot_change(
|
||||
resource: str, step_up: bool, selectable: bool | str, required: list[str] | str
|
||||
) -> None:
|
||||
from unittest.mock import patch
|
||||
|
||||
from litellm.proxy.mcp_server.oauth_utils import encode_upstream_scope
|
||||
|
||||
tracker: Final = _scoped_mcp_server(name="tracker")
|
||||
client_id: Final = (await _register([REDIRECT_URI]))["client_id"]
|
||||
with patch(_MANAGER_PATCH) as manager:
|
||||
manager.get_mcp_server_by_identifier.return_value = None
|
||||
manager.get_mcp_server_by_id.return_value = tracker
|
||||
authorized: Final = _scoped_authorize(
|
||||
client_id, resource, scope=encode_upstream_scope(tracker.server_id, "write") if step_up else None
|
||||
)
|
||||
|
||||
async def vendor_credential(user_id: str, server_id: str, requested_scopes: tuple[str, ...] = ()) -> str:
|
||||
return "present"
|
||||
|
||||
described: Final = await _describe_page(authorized, scoped_server=tracker, vendor=vendor_credential)
|
||||
|
||||
body: Final = json.loads(described.body)
|
||||
assert (body["state"], body.get("selectable", "absent")) == ("unscoped", selectable)
|
||||
assert body.get("required_server_ids", "absent") == required
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deny_ignores_the_selection() -> None:
|
||||
_, completed = await _complete_with_selection(("deleted-id",), (), decision="deny")
|
||||
assert completed.status_code == 303
|
||||
assert parse_qs(urlparse(completed.headers["location"]).query)["error"] == ["access_denied"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("resource", "error"),
|
||||
[("https://llm.example.com/mcp", None), (None, None), ("https://other.example.com/x", "invalid_target")],
|
||||
)
|
||||
async def test_a_picked_code_redeems_only_for_the_resource_it_was_authorized_for(
|
||||
resource: str | None, error: str | None
|
||||
) -> None:
|
||||
client_id, completed = await _complete_with_selection(("catalog-id",), (CATALOG,))
|
||||
assert completed.status_code == 303
|
||||
code, _ = _sealed_code_from(completed)
|
||||
redeemed: Final = await _redeem(code, client_id, resource=resource)
|
||||
assert (redeemed.status_code, json.loads(redeemed.body).get("error")) == (
|
||||
(200, None) if error is None else (400, error)
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue