diff --git a/package-lock.json b/package-lock.json index 9615029858863..f46800f13364a 100644 --- a/package-lock.json +++ b/package-lock.json @@ -73,6 +73,7 @@ }, "devDependencies": { "@anthropic-ai/claude-agent-sdk": "0.2.128", + "@modelcontextprotocol/sdk": "^1.29.0", "@playwright/cli": "^0.1.9", "@playwright/test": "^1.56.1", "@stylistic/eslint-plugin-ts": "^2.8.0", diff --git a/package.json b/package.json index c3e804079299b..cb7c6470c0453 100644 --- a/package.json +++ b/package.json @@ -150,6 +150,7 @@ }, "devDependencies": { "@anthropic-ai/claude-agent-sdk": "0.2.128", + "@modelcontextprotocol/sdk": "^1.29.0", "@playwright/cli": "^0.1.9", "@playwright/test": "^1.56.1", "@stylistic/eslint-plugin-ts": "^2.8.0", diff --git a/src/vs/base/common/jsonRpcProtocol.ts b/src/vs/base/common/jsonRpcProtocol.ts index 35d7144ba82bf..ce47db4c4650d 100644 --- a/src/vs/base/common/jsonRpcProtocol.ts +++ b/src/vs/base/common/jsonRpcProtocol.ts @@ -290,7 +290,15 @@ export function isJsonRpcRequest(message: JsonRpcMessage): message is IJsonRpcRe } export function isJsonRpcResponse(message: JsonRpcMessage): message is IJsonRpcSuccessResponse | IJsonRpcErrorResponse { - return hasKey(message, { id: true, result: true }) || hasKey(message, { id: true, error: true }); + return isJsonRpcSuccessResponse(message) || isJsonRpcErrorResponse(message); +} + +export function isJsonRpcSuccessResponse(message: JsonRpcMessage): message is IJsonRpcSuccessResponse { + return hasKey(message, { id: true, result: true }); +} + +export function isJsonRpcErrorResponse(message: JsonRpcMessage): message is IJsonRpcErrorResponse { + return hasKey(message, { id: true, error: true }); } export function isJsonRpcNotification(message: JsonRpcMessage): message is IJsonRpcNotification { diff --git a/src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts b/src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts index f152c2e221a10..c834d3282a679 100644 --- a/src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts +++ b/src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts @@ -268,7 +268,11 @@ export class RemoteAgentHostProtocolClient extends Disposable implements IAgentC * Authenticate with the remote agent host using a specific scheme. */ async authenticate(params: AuthenticateParams): Promise { - await this._sendRequest('authenticate', params); + await this._sendRequest('authenticate', { + resource: params.resource, + token: params.token, + ...(params.server ? { server: params.server.toString() } : {}), + }); return { authenticated: true }; } diff --git a/src/vs/platform/agentHost/common/agentService.ts b/src/vs/platform/agentHost/common/agentService.ts index 21b89b6eba78b..ac5a606564b92 100644 --- a/src/vs/platform/agentHost/common/agentService.ts +++ b/src/vs/platform/agentHost/common/agentService.ts @@ -180,6 +180,13 @@ export interface AuthenticateParams { /** The bearer token value (RFC 6750). */ readonly token: string; + + /** + * Optional: if set, scopes the token to a specific MCP server's + * proxy. The host routes the call to {@link IMcpHostService.getServer} + * instead of fanning out to providers. + */ + readonly server?: URI; } /** diff --git a/src/vs/platform/agentHost/common/mcpHost/mcpHostService.ts b/src/vs/platform/agentHost/common/mcpHost/mcpHostService.ts new file mode 100644 index 0000000000000..963a1b05de5fd --- /dev/null +++ b/src/vs/platform/agentHost/common/mcpHost/mcpHostService.ts @@ -0,0 +1,189 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { IDisposable } from '../../../../base/common/lifecycle.js'; +import { IObservable } from '../../../../base/common/observable.js'; +import { URI } from '../../../../base/common/uri.js'; +import { IMcpServerDefinition } from '../../../agentPlugins/common/pluginParsers.js'; +import { createDecorator } from '../../../instantiation/common/instantiation.js'; +import { McpMethodCallParams, McpMethodCallResult, McpNotificationParams } from '../state/protocol/commands.js'; +import { AhpMcpUiHostCapabilities, McpServerSummary } from '../state/protocol/state.js'; + +/** + * The MCP Apps `_meta.ui` payload carried by a `Tool` definition. + * + * Mirrors the + * [MCP Apps spec](https://github.com/modelcontextprotocol/ext-apps/blob/main/specification/2026-01-26/apps.mdx). + * The host captures these payloads opportunistically from `tools/list` + * responses; consumers (typically a Copilot/Claude agent's tool-call + * mapper) look them up by tool name when surfacing a tool call to AHP + * so the resulting `ToolCallBase._meta` carries `ui` / + * `uiHostCapabilities` for MCP-App-aware clients. + */ +export interface IMcpUiToolMeta { + /** `ui://…` URI of the UI resource for rendering the tool result. */ + readonly resourceUri?: string; + /** + * Who can access this tool. Default `["model", "app"]`. `"model"` makes + * the tool visible to the agent; `"app"` makes it callable by the UI + * view from the same MCP server. + */ + readonly visibility?: readonly ('model' | 'app')[]; +} + +/** + * An upstream-originated `mcpMethodCall` request. The host invokes the + * configured {@link IMcpHostUpstreamDelegate.handleUpstreamRequest} once + * per inbound JSON-RPC request from an MCP server and writes the + * returned outcome back as a JSON-RPC response. + */ +export interface IUpstreamMcpRequest { + /** MCP server URI; matches {@link McpServerSummary.resource}. */ + readonly server: URI; + /** JSON-RPC method (e.g. `sampling/createMessage`, `ui/open-link`). */ + readonly method: string; + /** Method params; opaque to AHP. */ + readonly params?: unknown; +} + +/** + * Outcome the delegate produces for an {@link IUpstreamMcpRequest}. + * + * Either {@link result} or {@link error} is set. The host forwards either + * one to the upstream MCP server as a JSON-RPC success or error response. + */ +export interface IUpstreamMcpResponse { + readonly result?: unknown; + readonly error?: { code: number; message: string; data?: unknown }; +} + +/** + * Upstream-originated `mcpNotification` payload. Fire-and-forget; the + * host calls {@link IMcpHostUpstreamDelegate.handleUpstreamNotification} + * whenever the upstream MCP server publishes a JSON-RPC notification + * (e.g. `notifications/tools/list_changed`). + */ +export interface IUpstreamMcpNotification { + readonly server: URI; + readonly method: string; + readonly params?: unknown; +} + +/** + * Bridge for forwarding upstream-originated MCP traffic out to an AHP + * client (typically wired up by `ProtocolServerHandler`). At most one + * delegate is installed at a time; calls when none is installed yield + * `MethodNotFound`. + */ +export interface IMcpHostUpstreamDelegate { + handleUpstreamRequest(request: IUpstreamMcpRequest): Promise; + handleUpstreamNotification(notification: IUpstreamMcpNotification): void; +} + +/** + * Per-MCP-server handle owned by {@link IMcpHostService}. Lifetime is bounded + * by {@link IMcpHostService.setSessionServers} — once the parent session no + * longer lists the server, the handle is disposed. + */ +export interface IMcpServerHandle extends IDisposable { + /** The `mcp://` URI; matches {@link McpServerSummary.resource}. */ + readonly resource: URI; + + /** Latest summary observed for this server. */ + readonly summary: IObservable; + + /** + * Endpoint the upstream SDK should connect to. `undefined` until the + * proxy is up. + */ + readonly endpoint: IObservable; + + /** + * Push a bearer token. Returns `true` when accepted by the upstream; + * `false` when the resource doesn't match this server. + */ + authenticate(resource: string, token: string): Promise; + + /** + * Forward an AHP-client → MCP server JSON-RPC **request** and resolve + * with the upstream's result. Rejects with a `ProtocolError` if the + * upstream returns a JSON-RPC error. + */ + callMethod(params: McpMethodCallParams): Promise; + + /** + * Forward an AHP-client → MCP server JSON-RPC **notification**. + * Fire-and-forget. + */ + notify(params: McpNotificationParams): void; + + /** + * Most recent `_meta.ui` payload the upstream MCP server advertised + * for `toolName` via `tools/list`. Returns `undefined` when no + * MCP-Apps-enabled tool with that name has been observed (for + * example, before the SDK has listed tools, or for non-app tools). + * + * Captured opportunistically as `tools/list` responses flow through + * the proxy from either the upstream SDK or AHP-client + * `mcpMethodCall` traffic. + */ + getToolUiMeta(toolName: string): IMcpUiToolMeta | undefined; + + /** + * The MCP-Apps host-capability set the AHP host satisfies for a View + * backed by this server, derived from the upstream's `initialize` + * response. Producers attach this to `_meta.uiHostCapabilities` on + * every tool-call state surfacing an MCP App so consumers can + * forward the same set into the view's `ui/initialize` response. + * + * Empty until the SDK's `initialize` handshake has completed through + * the proxy. + */ + getUiHostCapabilities(): AhpMcpUiHostCapabilities; +} + +/** + * Owns per-(session, MCP server) state and bridges AHP traffic to the + * underlying MCP transports. Platform-agnostic; the node-only proxy + * mechanics live in `IMcpProxyFactory`. + */ +export interface IMcpHostService { + readonly _serviceBrand: undefined; + + /** + * Replace the set of MCP servers registered for `session`. Diffing against + * the previous set, the service dispatches `session/mcpServerAdded` and + * `session/mcpServerRemoved` actions through the state manager. + */ + setSessionServers(session: URI, servers: readonly IMcpServerDefinition[]): readonly IMcpServerHandle[]; + + /** Return the current lightweight summaries for all MCP servers in `session`. */ + getServerSummaries(session: URI): readonly McpServerSummary[]; + + /** Look up a handle by its resource URI. */ + getServer(resource: URI): IMcpServerHandle | undefined; + + /** + * Pure router for client `mcpMethodCall` requests. Resolves the handle + * from `params.server` and forwards. Used by `ProtocolServerHandler`. + */ + callMethod(params: McpMethodCallParams): Promise; + + /** + * Pure router for client `mcpNotification` notifications. + */ + notify(params: McpNotificationParams): void; + + /** + * Install a delegate for upstream-originated traffic (the only direction + * the host service cannot satisfy locally). The disposable removes the + * delegate, after which upstream requests fail with `MethodNotFound`. + */ + setUpstreamDelegate(delegate: IMcpHostUpstreamDelegate): IDisposable; +} + +export const IMcpHostService = createDecorator('mcpHostService'); + + diff --git a/src/vs/platform/agentHost/common/mcpHost/nullMcpHostService.ts b/src/vs/platform/agentHost/common/mcpHost/nullMcpHostService.ts new file mode 100644 index 0000000000000..40ad011e13b43 --- /dev/null +++ b/src/vs/platform/agentHost/common/mcpHost/nullMcpHostService.ts @@ -0,0 +1,47 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { IDisposable } from '../../../../base/common/lifecycle.js'; +import { URI } from '../../../../base/common/uri.js'; +import { IMcpServerDefinition } from '../../../agentPlugins/common/pluginParsers.js'; +import { McpMethodCallParams, McpMethodCallResult, McpNotificationParams } from '../state/protocol/commands.js'; +import type { McpServerSummary } from '../state/protocol/state.js'; +import { JsonRpcErrorCodes, ProtocolError } from '../state/sessionProtocol.js'; +import { IMcpHostService, IMcpHostUpstreamDelegate, IMcpServerHandle } from './mcpHostService.js'; + +/** + * No-op {@link IMcpHostService}. Returns no servers, refuses `mcpMethodCall` + * requests with `MethodNotFound`, and silently drops `mcpNotification`s. + * Used by surfaces that don't host MCP traffic (browser-side stubs and tests). + */ +export class NullMcpHostService implements IMcpHostService { + declare readonly _serviceBrand: undefined; + + setSessionServers(_session: URI, _servers: readonly IMcpServerDefinition[]): readonly IMcpServerHandle[] { + return []; + } + + getServerSummaries(_session: URI): readonly McpServerSummary[] { + return []; + } + + getServer(_resource: URI): IMcpServerHandle | undefined { + return undefined; + } + + async callMethod(_params: McpMethodCallParams): Promise { + throw new ProtocolError(JsonRpcErrorCodes.MethodNotFound, 'mcpMethodCall is not supported by this host'); + } + + notify(_params: McpNotificationParams): void { + // no-op + } + + setUpstreamDelegate(_delegate: IMcpHostUpstreamDelegate): IDisposable { + return { dispose: () => { /* no-op */ } }; + } +} + + diff --git a/src/vs/platform/agentHost/common/state/mcpServerUri.ts b/src/vs/platform/agentHost/common/state/mcpServerUri.ts new file mode 100644 index 0000000000000..46791c8241e12 --- /dev/null +++ b/src/vs/platform/agentHost/common/state/mcpServerUri.ts @@ -0,0 +1,44 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { URI } from '../../../../base/common/uri.js'; + +/** URI scheme for MCP server resources surfaced as `McpServerSummary.resource`. */ +export const MCP_SERVER_SCHEME = 'mcp'; + +/** + * Build an `mcp://` URI. + * + * The path encodes the session URI's path component (typically a UUID) so + * the server URI is **self-identifying to a session** per + * `McpMessageParams` — the host service can derive the owning session by + * parsing this URI in O(1) without a side map. + */ +export function buildMcpServerUri(session: URI, serverId: string): URI { + const sessionPath = session.path.replace(/^\/+/, ''); + return URI.from({ + scheme: MCP_SERVER_SCHEME, + path: `/${sessionPath}/${encodeURIComponent(serverId)}`, + }); +} + +/** + * Parse an `mcp://` URI back into its parts. Returns + * `undefined` if the URI is not in the expected scheme/shape. + */ +export function parseMcpServerUri(uri: URI): { sessionPath: string; serverId: string } | undefined { + if (uri.scheme !== MCP_SERVER_SCHEME) { + return undefined; + } + const path = uri.path.replace(/^\/+/, ''); + const lastSlash = path.lastIndexOf('/'); + if (lastSlash < 1) { + return undefined; + } + return { + sessionPath: path.slice(0, lastSlash), + serverId: decodeURIComponent(path.slice(lastSlash + 1)), + }; +} diff --git a/src/vs/platform/agentHost/common/state/protocol/.ahp-version b/src/vs/platform/agentHost/common/state/protocol/.ahp-version index c695b4dca63ea..bb9ee3ada660d 100644 --- a/src/vs/platform/agentHost/common/state/protocol/.ahp-version +++ b/src/vs/platform/agentHost/common/state/protocol/.ahp-version @@ -1 +1 @@ -5f79fe4 +a3ea9b4 diff --git a/src/vs/platform/agentHost/common/state/protocol/action-origin.generated.ts b/src/vs/platform/agentHost/common/state/protocol/action-origin.generated.ts index 5cd909e906317..1029720f755a9 100644 --- a/src/vs/platform/agentHost/common/state/protocol/action-origin.generated.ts +++ b/src/vs/platform/agentHost/common/state/protocol/action-origin.generated.ts @@ -9,10 +9,10 @@ // Generated from types/actions.ts — do not edit // Run `npm run generate` to regenerate. -import { ActionType, type StateAction, type RootAgentsChangedAction, type RootActiveSessionsChangedAction, type RootTerminalsChangedAction, type RootConfigChangedAction, type SessionReadyAction, type SessionCreationFailedAction, type SessionTurnStartedAction, type SessionDeltaAction, type SessionResponsePartAction, type SessionToolCallStartAction, type SessionToolCallDeltaAction, type SessionToolCallReadyAction, type SessionToolCallConfirmedAction, type SessionToolCallCompleteAction, type SessionToolCallResultConfirmedAction, type SessionToolCallContentChangedAction, type SessionTurnCompleteAction, type SessionTurnCancelledAction, type SessionErrorAction, type SessionTitleChangedAction, type SessionUsageAction, type SessionReasoningAction, type SessionModelChangedAction, type SessionServerToolsChangedAction, type SessionActiveClientChangedAction, type SessionActiveClientToolsChangedAction, type SessionPendingMessageSetAction, type SessionPendingMessageRemovedAction, type SessionQueuedMessagesReorderedAction, type SessionInputRequestedAction, type SessionInputAnswerChangedAction, type SessionInputCompletedAction, type SessionCustomizationsChangedAction, type SessionCustomizationToggledAction, type SessionTruncatedAction, type SessionIsReadChangedAction, type SessionIsArchivedChangedAction, type SessionActivityChangedAction, type SessionDiffsChangedAction, type SessionConfigChangedAction, type SessionMetaChangedAction, type TerminalDataAction, type TerminalInputAction, type TerminalResizedAction, type TerminalClaimedAction, type TerminalTitleChangedAction, type TerminalCwdChangedAction, type TerminalExitedAction, type TerminalClearedAction, type TerminalCommandDetectionAvailableAction, type TerminalCommandExecutedAction, type TerminalCommandFinishedAction } from './actions.js'; +import { ActionType, type StateAction, type RootAgentsChangedAction, type RootActiveSessionsChangedAction, type RootTerminalsChangedAction, type RootConfigChangedAction, type SessionReadyAction, type SessionCreationFailedAction, type SessionTurnStartedAction, type SessionDeltaAction, type SessionResponsePartAction, type SessionToolCallStartAction, type SessionToolCallDeltaAction, type SessionToolCallReadyAction, type SessionToolCallConfirmedAction, type SessionToolCallCompleteAction, type SessionToolCallResultConfirmedAction, type SessionToolCallContentChangedAction, type SessionTurnCompleteAction, type SessionTurnCancelledAction, type SessionErrorAction, type SessionTitleChangedAction, type SessionUsageAction, type SessionReasoningAction, type SessionModelChangedAction, type SessionServerToolsChangedAction, type SessionActiveClientChangedAction, type SessionActiveClientToolsChangedAction, type SessionPendingMessageSetAction, type SessionPendingMessageRemovedAction, type SessionQueuedMessagesReorderedAction, type SessionInputRequestedAction, type SessionInputAnswerChangedAction, type SessionInputCompletedAction, type SessionCustomizationsChangedAction, type SessionCustomizationToggledAction, type SessionTruncatedAction, type SessionIsReadChangedAction, type SessionIsArchivedChangedAction, type SessionActivityChangedAction, type SessionDiffsChangedAction, type SessionConfigChangedAction, type SessionMetaChangedAction, type TerminalDataAction, type TerminalInputAction, type TerminalResizedAction, type TerminalClaimedAction, type TerminalTitleChangedAction, type TerminalCwdChangedAction, type TerminalExitedAction, type TerminalClearedAction, type TerminalCommandDetectionAvailableAction, type TerminalCommandExecutedAction, type TerminalCommandFinishedAction, type McpServerAddedAction, type McpServerRemovedAction, type McpServerStatusChangedAction } from './actions.js'; -// ─── Root vs Session vs Terminal Action Unions ─────────────────────────────── +// ─── Root vs Session vs Terminal Action Unions ─────────────────────────── /** Union of all root-scoped actions. */ export type RootAction = @@ -73,6 +73,9 @@ export type SessionAction = | SessionDiffsChangedAction | SessionConfigChangedAction | SessionMetaChangedAction + | McpServerAddedAction + | McpServerRemovedAction + | McpServerStatusChangedAction ; /** Union of session actions that clients may dispatch. */ @@ -118,6 +121,9 @@ export type ServerSessionAction = | SessionActivityChangedAction | SessionDiffsChangedAction | SessionMetaChangedAction + | McpServerAddedAction + | McpServerRemovedAction + | McpServerStatusChangedAction ; /** Union of all terminal-scoped actions. */ @@ -213,4 +219,7 @@ export const IS_CLIENT_DISPATCHABLE: { readonly [K in StateAction['type']]: bool [ActionType.TerminalCommandDetectionAvailable]: false, [ActionType.TerminalCommandExecuted]: false, [ActionType.TerminalCommandFinished]: false, + [ActionType.McpServerAdded]: false, + [ActionType.McpServerRemoved]: false, + [ActionType.McpServerStatusChanged]: false, }; diff --git a/src/vs/platform/agentHost/common/state/protocol/actions.ts b/src/vs/platform/agentHost/common/state/protocol/actions.ts index 1d3510a353cea..ab8d698f48a64 100644 --- a/src/vs/platform/agentHost/common/state/protocol/actions.ts +++ b/src/vs/platform/agentHost/common/state/protocol/actions.ts @@ -6,7 +6,7 @@ // allow-any-unicode-comment-file // DO NOT EDIT -- auto-generated by scripts/sync-agent-host-protocol.ts -import { ToolCallConfirmationReason, ToolCallCancellationReason, PendingMessageKind, type URI, type StringOrMarkdown, type AgentInfo, type ErrorInfo, type ModelSelection, type UserMessage, type ResponsePart, type ToolCallResult, type ToolResultContent, type ToolDefinition, type SessionActiveClient, type UsageInfo, type SessionCustomization, type FileEdit, type SessionInputAnswer, type SessionInputRequest, type TerminalInfo, type TerminalClaim, type SessionInputResponseKind, type ConfirmationOption } from './state.js'; +import { ToolCallConfirmationReason, ToolCallCancellationReason, PendingMessageKind, type URI, type StringOrMarkdown, type AgentInfo, type ErrorInfo, type ModelSelection, type UserMessage, type ResponsePart, type ToolCallResult, type ToolResultContent, type ToolDefinition, type SessionActiveClient, type UsageInfo, type SessionCustomization, type FileEdit, type SessionInputAnswer, type SessionInputRequest, type TerminalInfo, type TerminalClaim, type SessionInputResponseKind, type ConfirmationOption, type McpServerSummary, type McpServerStatus } from './state.js'; // ─── Action Type Enum ──────────────────────────────────────────────────────── @@ -69,6 +69,9 @@ export const enum ActionType { TerminalCommandDetectionAvailable = 'terminal/commandDetectionAvailable', TerminalCommandExecuted = 'terminal/commandExecuted', TerminalCommandFinished = 'terminal/commandFinished', + McpServerAdded = 'session/mcpServerAdded', + McpServerRemoved = 'session/mcpServerRemoved', + McpServerStatusChanged = 'session/mcpServerStatusChanged', } // ─── Action Envelope ───────────────────────────────────────────────────────── @@ -1135,6 +1138,58 @@ export interface TerminalCommandFinishedAction { durationMs?: number; } +// ─── Mcp Actions ───────────────────────────────────────────────────────────── + +/** + * An MCP server has been registered with a session. Pushes a new + * {@link McpServerSummary} onto {@link SessionState.mcpServers}. + * + * @category Session Actions + * @version 1 + */ +export interface McpServerAddedAction { + type: ActionType.McpServerAdded; + /** Session URI this server belongs to. */ + session: URI; + /** The summary of the newly-added server. */ + server: McpServerSummary; +} + +/** + * An MCP server has been removed from a session. Removes the matching + * {@link McpServerSummary} from {@link SessionState.mcpServers}. + * + * @category Session Actions + * @version 1 + */ +export interface McpServerRemovedAction { + type: ActionType.McpServerRemoved; + /** Session URI this server belonged to. */ + session: URI; + /** {@link McpServerSummary.resource} of the removed server. */ + mcpServer: URI; +} + +/** + * The status of an MCP server changed. Updates the matching summary on + * {@link SessionState.mcpServers}. + * + * Authentication transitions arrive via this action only — there is no + * `notify/authRequired` notification for MCP servers. + * + * @category Session Actions + * @version 1 + */ +export interface McpServerStatusChangedAction { + type: ActionType.McpServerStatusChanged; + /** Session URI this server belongs to. */ + session: URI; + /** {@link McpServerSummary.resource} of the affected server. */ + mcpServer: URI; + /** New status. */ + status: McpServerStatus; +} + // ─── Discriminated Union ───────────────────────────────────────────────────── /** @@ -1192,4 +1247,7 @@ export type StateAction = | TerminalClearedAction | TerminalCommandDetectionAvailableAction | TerminalCommandExecutedAction - | TerminalCommandFinishedAction; + | TerminalCommandFinishedAction + | McpServerAddedAction + | McpServerRemovedAction + | McpServerStatusChangedAction; diff --git a/src/vs/platform/agentHost/common/state/protocol/commands.ts b/src/vs/platform/agentHost/common/state/protocol/commands.ts index 9e52e0fd1d74e..b59ddd40e3008 100644 --- a/src/vs/platform/agentHost/common/state/protocol/commands.ts +++ b/src/vs/platform/agentHost/common/state/protocol/commands.ts @@ -13,6 +13,32 @@ export type { ConfigPropertySchema, ConfigSchema, SessionConfigPropertySchema, S // ─── initialize ────────────────────────────────────────────────────────────── +/** + * Capabilities advertised by the **client** during `initialize`. All fields + * are optional and additive; the server treats absence the same as `false` / + * "not supported". + * + * Currently empty — MCP App support is negotiated per tool call via + * `_meta.uiHostCapabilities` on the tool call state. + * + * @category Commands + */ +export interface ClientCapabilities { +} + +/** + * Capabilities advertised by the **server** in the `initialize` result. All + * fields are optional and additive; the client treats absence the same as + * `false` / "not supported". + * + * Currently empty — MCP App support is negotiated per tool call via + * `_meta.uiHostCapabilities` on the tool call state. + * + * @category Commands + */ +export interface ServerCapabilities { +} + /** * Establishes a new connection and negotiates the protocol version. * This MUST be the first message sent by the client. @@ -45,6 +71,12 @@ export interface InitializeParams { * user-facing strings such as confirmation option labels. */ locale?: string; + /** + * Optional capabilities the client supports. Servers MUST treat absent + * fields as "not supported" and MUST NOT initiate traffic that depends + * on a capability the client did not advertise. + */ + capabilities?: ClientCapabilities; } /** @@ -76,6 +108,12 @@ export interface InitializeResult { * `'@'` or `'/'`. */ completionTriggerCharacters?: string[]; + /** + * Optional capabilities the server supports. Clients MUST treat absent + * fields as "not supported" and MUST NOT issue commands that depend on + * a capability the server did not advertise. + */ + capabilities?: ServerCapabilities; } // ─── reconnect ─────────────────────────────────────────────────────────────── @@ -745,8 +783,13 @@ export interface ResourceMoveResult { /** * Pushes a Bearer token for a protected resource. The `resource` field MUST - * match a `ProtectedResourceMetadata.resource` value declared by an agent - * in `AgentInfo.protectedResources`. + * match either: + * + * - a `resource` value from {@link ProtectedResourceMetadata} declared by an + * agent in {@link AgentInfo.protectedResources}, **or** + * - the `resource` field of a current + * {@link McpServerStatusAuthRequired} on a subscribed session (when + * `server` is set). * * Tokens are delivered using [RFC 6750](https://datatracker.ietf.org/doc/html/rfc6750) * (Bearer Token Usage) semantics. The client obtains the token from the @@ -761,10 +804,16 @@ export interface ResourceMoveResult { * @see {@link /specification/authentication | Authentication} * @example * ```jsonc - * // Client → Server + * // Client → Server (agent-level resource) * { "jsonrpc": "2.0", "id": 3, "method": "authenticate", * "params": { "resource": "https://api.github.com", "token": "gho_xxxx" } } * + * // Client → Server (per-MCP-server scoping) + * { "jsonrpc": "2.0", "id": 4, "method": "authenticate", + * "params": { "resource": "https://api.github.com", + * "token": "gho_xxxx", + * "server": "mcp://github" } } + * * // Server → Client (success) * { "jsonrpc": "2.0", "id": 3, "result": {} } * @@ -774,12 +823,29 @@ export interface ResourceMoveResult { */ export interface AuthenticateParams { /** - * The protected resource identifier. MUST match a `resource` value from - * `ProtectedResourceMetadata` declared in `AgentInfo.protectedResources`. + * The protected resource identifier. MUST match either a `resource` value + * from {@link ProtectedResourceMetadata} declared in + * {@link AgentInfo.protectedResources}, **or** the `resource` field of a + * current {@link McpServerStatusAuthRequired} on a subscribed session + * (when `server` is set). */ resource: string; - /** Bearer token obtained from the resource's authorization server */ + /** Bearer token obtained from the resource's authorization server. */ token: string; + /** + * Optional MCP server URI (`mcp://`) this token is + * being presented for. When set, the token is scoped to *only* that + * server — other MCP servers that share the same `resource` URL remain + * in {@link McpServerStatusKind.AuthRequired}. This lets a + * security-conscious client authorize only the MCP servers it trusts and + * leave others in an `AuthRequired` state. + * + * When omitted, the token is bound to the agent-level resource as + * before. Hosts that do not yet implement per-server scoping MAY ignore + * `server` and treat the token as resource-wide; clients SHOULD treat + * that as the worst-case fallback. + */ + server?: URI; } /** @@ -935,6 +1001,120 @@ export interface SessionConfigCompletionsResult { items: SessionConfigValueItem[]; } +// ─── mcpMethodCall ────────────────────────────────────────────────────────── + +/** + * Forwards a JSON-RPC **method call** between an AHP peer and a specific + * MCP server managed by the host. Bidirectional: both the client and the + * server may issue `mcpMethodCall` requests. + * + * - Client → Server: the client invokes an MCP method exposed by the + * underlying MCP server (for example `tools/list`, `resources/read`, + * or `ui/*` traffic when the tool surfaces an MCP App). + * - Server → Client: the host invokes a method that the client is expected + * to satisfy (for example `sampling/createMessage`, `ui/open-link`, or + * `resources/read` against a client-owned resource). + * + * Notifications travel separately via `mcpNotification` (see + * {@link McpNotificationParams}). No AHP-level correlation id is required: + * each `mcpMethodCall` is a normal JSON-RPC request whose response *is* + * the response. + * + * Peers MUST refuse method calls that fall outside the capabilities they + * have advertised \u2014 today, that means the per-tool-call + * `_meta.uiHostCapabilities` declared by the producer of an MCP App tool + * call. Refusal is signalled by returning a JSON-RPC error. + * + * @category Commands + * @method mcpMethodCall + * @direction Bidirectional + * @messageType Request + * @version 1 + * @example + * ```jsonc + * // Client \u2192 Server (a UI app asks the host to open a link) + * { "jsonrpc": "2.0", "id": 30, "method": "mcpMethodCall", + * "params": { + * "server": "mcp://github", + * "method": "ui/open-link", + * "params": { "url": "https://github.com/owner/repo" } + * } } + * + * // Server \u2192 Client + * { "jsonrpc": "2.0", "id": 30, "result": { "result": {} } } + * ``` + */ +export interface McpMethodCallParams { + /** + * Stable identifier (typically `mcp://`) of the + * MCP server this call is scoped to. AHP treats the value as opaque \u2014 + * the host defines the scheme. Matches the + * {@link McpServerSummary.resource} value surfaced on the session. + */ + server: URI; + /** JSON-RPC method (e.g. `tools/list`, `ui/open-link`). */ + method: string; + /** Method params; opaque to AHP \u2014 typed by the MCP method. */ + params?: unknown; +} + +/** + * Result of `mcpMethodCall`. + */ +export interface McpMethodCallResult { + /** + * The result payload returned by the underlying MCP method, opaque to + * AHP. Receivers that wish to surface MCP-level errors SHOULD instead + * return a JSON-RPC error response. + */ + result: unknown; +} + +// ─── mcpNotification ──────────────────────────────────────────────────────── + +/** + * Forwards a JSON-RPC **notification** between an AHP peer and a specific + * MCP server managed by the host. Bidirectional and fire-and-forget; no + * response is expected. + * + * Examples: + * + * - Server \u2192 Client: `notifications/tools/list_changed`, + * `notifications/resources/updated`, `ui/notifications/host-context-changed`. + * - Client \u2192 Server: `notifications/message` log entries emitted by the + * client on behalf of an MCP App. + * + * Peers MUST only emit notifications consistent with the capabilities + * they have advertised for the relevant tool call (see + * {@link AhpMcpUiHostCapabilities}). + * + * @category Notifications + * @method mcpNotification + * @direction Bidirectional + * @messageType Notification + * @version 1 + * @example + * ```jsonc + * // Server \u2192 Client (the MCP server reported a tool list change) + * { "jsonrpc": "2.0", "method": "mcpNotification", + * "params": { + * "server": "mcp://github", + * "method": "notifications/tools/list_changed" + * } } + * ``` + */ +export interface McpNotificationParams { + /** + * Stable identifier of the MCP server this notification is scoped to. + * Matches {@link McpMethodCallParams.server}. + */ + server: URI; + /** JSON-RPC method name. */ + method: string; + /** Method params; opaque to AHP \u2014 typed by the MCP method. */ + params?: unknown; +} + // ─── completions ───────────────────────────────────────────────────────────── /** diff --git a/src/vs/platform/agentHost/common/state/protocol/messages.ts b/src/vs/platform/agentHost/common/state/protocol/messages.ts index 6fe656b48d3a3..008250a435e4e 100644 --- a/src/vs/platform/agentHost/common/state/protocol/messages.ts +++ b/src/vs/platform/agentHost/common/state/protocol/messages.ts @@ -6,7 +6,7 @@ // allow-any-unicode-comment-file // DO NOT EDIT -- auto-generated by scripts/sync-agent-host-protocol.ts -import type { InitializeParams, InitializeResult, ReconnectParams, ReconnectResult, SubscribeParams, SubscribeResult, CreateSessionParams, DisposeSessionParams, CreateTerminalParams, DisposeTerminalParams, ListSessionsParams, ListSessionsResult, ResourceReadParams, ResourceReadResult, ResourceWriteParams, ResourceWriteResult, ResourceListParams, ResourceListResult, ResourceCopyParams, ResourceCopyResult, ResourceDeleteParams, ResourceDeleteResult, ResourceMoveParams, ResourceMoveResult, ResourceRequestParams, ResourceRequestResult, FetchTurnsParams, FetchTurnsResult, UnsubscribeParams, DispatchActionParams, AuthenticateParams, AuthenticateResult, ResolveSessionConfigParams, ResolveSessionConfigResult, SessionConfigCompletionsParams, SessionConfigCompletionsResult, CompletionsParams, CompletionsResult } from './commands.js'; +import type { InitializeParams, InitializeResult, ReconnectParams, ReconnectResult, SubscribeParams, SubscribeResult, CreateSessionParams, DisposeSessionParams, CreateTerminalParams, DisposeTerminalParams, ListSessionsParams, ListSessionsResult, ResourceReadParams, ResourceReadResult, ResourceWriteParams, ResourceWriteResult, ResourceListParams, ResourceListResult, ResourceCopyParams, ResourceCopyResult, ResourceDeleteParams, ResourceDeleteResult, ResourceMoveParams, ResourceMoveResult, ResourceRequestParams, ResourceRequestResult, FetchTurnsParams, FetchTurnsResult, UnsubscribeParams, DispatchActionParams, AuthenticateParams, AuthenticateResult, ResolveSessionConfigParams, ResolveSessionConfigResult, SessionConfigCompletionsParams, SessionConfigCompletionsResult, CompletionsParams, CompletionsResult, McpMethodCallParams, McpMethodCallResult, McpNotificationParams } from './commands.js'; import type { ActionEnvelope } from './actions.js'; import type { ProtocolNotification } from './notifications.js'; @@ -93,6 +93,7 @@ export interface CommandMap { 'resolveSessionConfig': { params: ResolveSessionConfigParams; result: ResolveSessionConfigResult }; 'sessionConfigCompletions': { params: SessionConfigCompletionsParams; result: SessionConfigCompletionsResult }; 'completions': { params: CompletionsParams; result: CompletionsResult }; + 'mcpMethodCall': { params: McpMethodCallParams; result: McpMethodCallResult }; } /** @@ -108,6 +109,7 @@ export interface CommandMap { */ export interface ServerCommandMap { 'resourceRequest': { params: ResourceRequestParams; result: ResourceRequestResult }; + 'mcpMethodCall': { params: McpMethodCallParams; result: McpMethodCallResult }; } // ─── Notification Maps ─────────────────────────────────────────────────────── @@ -125,6 +127,7 @@ export interface NotificationMethodParams { export interface ClientNotificationMap { 'unsubscribe': { params: UnsubscribeParams }; 'dispatchAction': { params: DispatchActionParams }; + 'mcpNotification': { params: McpNotificationParams }; } /** @@ -135,6 +138,7 @@ export interface ClientNotificationMap { export interface ServerNotificationMap { 'action': { params: ActionEnvelope }; 'notification': { params: NotificationMethodParams }; + 'mcpNotification': { params: McpNotificationParams }; } /** Combined notification map for all directions. */ diff --git a/src/vs/platform/agentHost/common/state/protocol/reducers.ts b/src/vs/platform/agentHost/common/state/protocol/reducers.ts index a70f5b9edcbe0..4df897bfcc737 100644 --- a/src/vs/platform/agentHost/common/state/protocol/reducers.ts +++ b/src/vs/platform/agentHost/common/state/protocol/reducers.ts @@ -7,7 +7,7 @@ // DO NOT EDIT -- auto-generated by scripts/sync-agent-host-protocol.ts import { ActionType } from './actions.js'; -import { SessionLifecycle, SessionStatus, TurnState, ToolCallStatus, ToolCallConfirmationReason, ToolCallCancellationReason, ResponsePartKind, PendingMessageKind, type RootState, type SessionInputRequest, type SessionState, type TerminalState, type TerminalContentPart, type ToolCallState, type ResponsePart, type ToolCallResponsePart, type Turn, type PendingMessage, type ConfirmationOption } from './state.js'; +import { SessionLifecycle, SessionStatus, TurnState, ToolCallStatus, ToolCallConfirmationReason, ToolCallCancellationReason, ResponsePartKind, PendingMessageKind, type RootState, type SessionInputRequest, type SessionState, type TerminalState, type TerminalContentPart, type ToolCallState, type ResponsePart, type ToolCallResponsePart, type Turn, type PendingMessage, type ConfirmationOption, type McpServerSummary } from './state.js'; import { IS_CLIENT_DISPATCHABLE, type RootAction, type ClientRootAction, type SessionAction, type ClientSessionAction, type TerminalAction, type ClientTerminalAction } from './action-origin.generated.js'; // ─── Helpers ───────────────────────────────────────────────────────────────── @@ -627,6 +627,44 @@ export function sessionReducer(state: SessionState, action: SessionAction, log?: activeClient: { ...state.activeClient, tools: action.tools }, }; + // ── MCP Servers (session-scoped summary) ───────────────────────────── + + case ActionType.McpServerAdded: { + const existing = state.mcpServers ?? []; + // Replace any prior entry with the same resource URI (idempotent on re-adds). + const filtered = existing.filter(s => s.resource !== action.server.resource); + return { ...state, mcpServers: [...filtered, action.server] }; + } + + case ActionType.McpServerRemoved: { + const existing = state.mcpServers; + if (!existing || existing.length === 0) { + return state; + } + const filtered = existing.filter(s => s.resource !== action.mcpServer); + if (filtered.length === existing.length) { + return state; + } + // Drop the array entirely when empty so the field stays optional-shaped. + return filtered.length === 0 + ? (() => { const { mcpServers: _, ...rest } = state; return rest; })() + : { ...state, mcpServers: filtered }; + } + + case ActionType.McpServerStatusChanged: { + const existing = state.mcpServers; + if (!existing) { + return state; + } + const idx = existing.findIndex(s => s.resource === action.mcpServer); + if (idx < 0) { + return state; + } + const updated: McpServerSummary[] = [...existing]; + updated[idx] = { ...existing[idx], status: action.status }; + return { ...state, mcpServers: updated }; + } + // ── Customizations ────────────────────────────────────────────────── case ActionType.SessionCustomizationsChanged: diff --git a/src/vs/platform/agentHost/common/state/protocol/state.ts b/src/vs/platform/agentHost/common/state/protocol/state.ts index 5a95ad6535abe..12d4cd62fc86a 100644 --- a/src/vs/platform/agentHost/common/state/protocol/state.ts +++ b/src/vs/platform/agentHost/common/state/protocol/state.ts @@ -346,6 +346,15 @@ export interface SessionState { * workingDirectory. */ _meta?: Record; + /** + * Lightweight summary of MCP servers running for this session. + * + * Each entry carries just enough state for clients to render an "MCP + * servers" affordance and react to authentication challenges. + * + * @see {@link McpServerSummary} + */ + mcpServers?: McpServerSummary[]; } /** @@ -1220,9 +1229,20 @@ interface ToolCallBase { * Additional provider-specific metadata for this tool call. * * Clients MAY look for well-known keys here to provide enhanced UI. - * For example, a `ptyTerminal` key with `{ input: string; output: string }` - * indicates the tool operated on a terminal (both `input` and `output` may - * contain escape sequences). + * Recognised well-known keys: + * + * - `ptyTerminal: { input: string; output: string }` — the tool + * operated on a terminal (both `input` and `output` may contain + * escape sequences). + * - `ui` — the tool surfaces an [MCP App](https://github.com/modelcontextprotocol/ext-apps). + * The value mirrors the MCP `_meta` `"io.modelcontextprotocol/ui"` + * payload exactly (opaque to AHP). When this key is present, + * producers SHOULD also set `uiHostCapabilities`. + * - `uiHostCapabilities: {@link AhpMcpUiHostCapabilities}` — the set + * of MCP host capabilities the producer is willing to satisfy for + * this app via `mcpMethodCall` and `mcpNotification`. Each advertised + * capability MUST correspond to methods the producer actually + * accepts; absent capabilities mean "not supported". */ _meta?: Record; } @@ -1823,6 +1843,216 @@ export interface TerminalCommandPart { durationMs?: number; } +// ─── MCP Server State ──────────────────────────────────────────────────────── + +/** + * Discriminant for the {@link McpServerStatus} union. + * + * @category MCP Server State + */ +export const enum McpServerStatusKind { + /** Server has been registered but is not yet running. */ + Starting = 'starting', + /** Server is running and serving requests. */ + Ready = 'ready', + /** Server is reachable but cannot serve requests until the client authenticates. */ + AuthRequired = 'authRequired', + /** Server failed to start, crashed, or otherwise transitioned to a fatal error. */ + Error = 'error', + /** Server has been shut down. */ + Stopped = 'stopped', +} + +/** + * Why an MCP server is currently in the {@link McpServerStatusKind.AuthRequired} + * state. Mirrors the three failure modes defined by the + * [MCP authorization spec](https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization.md). + * + * @category MCP Server State + */ +export const enum McpAuthRequiredReason { + /** No token has been provided yet (HTTP 401, no prior token). */ + Required = 'required', + /** A previously valid token expired or was revoked (HTTP 401). */ + Expired = 'expired', + /** + * Step-up auth: a token is present but its scopes are insufficient for + * the requested operation (HTTP 403 with + * `WWW-Authenticate: Bearer error="insufficient_scope"`). + */ + InsufficientScope = 'insufficientScope', +} + +/** + * Server is registered with the host but has not yet started. + * + * @category MCP Server State + */ +export interface McpServerStatusStarting { + kind: McpServerStatusKind.Starting; +} + +/** + * Server is running and serving requests. + * + * @category MCP Server State + */ +export interface McpServerStatusReady { + kind: McpServerStatusKind.Ready; +} + +/** + * Server is reachable but cannot serve requests until the client + * authenticates. Mirrors the discovery flow defined by + * [RFC 9728](https://datatracker.ietf.org/doc/html/rfc9728) + * (Protected Resource Metadata) and the OAuth 2.1 / RFC 6750 challenge + * semantics required by the MCP authorization spec. + * + * Clients react to this state by calling the existing `authenticate` + * command with the {@link ProtectedResourceMetadata.resource | resource} + * carried here. There is **no** `notify/authRequired` notification for + * MCP servers — the action stream is the single source of truth. + * + * @category MCP Server State + */ +export interface McpServerStatusAuthRequired { + kind: McpServerStatusKind.AuthRequired; + /** Why authentication is required. */ + reason: McpAuthRequiredReason; + /** + * RFC 9728 Protected Resource Metadata. The `resource` field is the + * canonical MCP server URI per RFC 8707, used as the OAuth `resource` + * indicator. `authorization_servers` is REQUIRED by the MCP + * authorization spec. + */ + resource: ProtectedResourceMetadata; + /** + * Scopes required for the current challenge, parsed from the + * `WWW-Authenticate: Bearer scope="…"` header (or `scopes_supported` + * fallback). Authoritative for the next authorization request — clients + * MUST NOT assume any subset/superset relationship to + * `resource.scopes_supported`. + */ + requiredScopes?: string[]; + /** Human-readable hint, typically from the OAuth `error_description`. */ + description?: string; +} + +/** + * Server failed to start, crashed, or otherwise transitioned to a + * non-recoverable error. Use {@link McpServerStatusKind.AuthRequired} + * for authentication failures. + * + * @category MCP Server State + */ +export interface McpServerStatusError { + kind: McpServerStatusKind.Error; + /** Error details. */ + error: ErrorInfo; +} + +/** + * Server has been shut down. The host MAY remove the server from the + * session entirely shortly after this state. + * + * @category MCP Server State + */ +export interface McpServerStatusStopped { + kind: McpServerStatusKind.Stopped; +} + +/** + * Discriminated union of all MCP server statuses. Discriminated by `kind`. + * + * @category MCP Server State + */ +export type McpServerStatus = + | McpServerStatusStarting + | McpServerStatusReady + | McpServerStatusAuthRequired + | McpServerStatusError + | McpServerStatusStopped; + +/** + * Lightweight summary of an MCP server running for a session. Lives on + * {@link SessionState.mcpServers}. + * + * @category MCP Server State + */ +export interface McpServerSummary { + /** + * Stable, host-defined identifier for this MCP server (typically a + * `mcp://` URI). Used as the `server` field on + * `mcpMethodCall` requests and `mcpNotification` notifications to address + * traffic to or from this server. + */ + resource: URI; + /** Human-readable display label. */ + label: string; + /** Current status. */ + status: McpServerStatus; +} + +/** + * The subset of MCP App + * [`HostCapabilities`](https://github.com/modelcontextprotocol/ext-apps/blob/main/specification/draft/apps.mdx) + * an AHP host can derive from the upstream MCP server (and from AHP's own + * forwarding plumbing). Used as the value of `_meta.uiHostCapabilities` on + * tool call states that surface an MCP App (see {@link ToolCallState}). + * + * Field names mirror the MCP Apps spec exactly so the AHP-side producer + * can pass them straight through into the `hostCapabilities` of the + * `ui/initialize` response delivered to the View. + * + * Capabilities outside this set (`openLinks`, `downloadFile`, `sandbox`, + * `experimental`) are decided locally by whichever AHP client renders the + * View and are NOT part of this AHP-level advertisement — only the + * server-derived subset is. + * + * A producer MUST only advertise a capability when it actually accepts the + * corresponding methods/notifications via `mcpMethodCall` / `mcpNotification`: + * + * - {@link serverTools}: producer proxies `tools/list` and `tools/call` to + * the MCP server. When `listChanged` is `true`, the producer also forwards + * `notifications/tools/list_changed`. + * - {@link serverResources}: producer proxies `resources/read`, + * `resources/list`, and `resources/templates/list` to the MCP server. + * When `listChanged` is `true`, the producer also forwards + * `notifications/resources/list_changed`. + * - {@link logging}: producer accepts `notifications/message` log entries + * from the App and forwards them via `mcpNotification` (and forwards + * `logging/setLevel` calls to the server). + * - {@link sampling}: producer serves `sampling/createMessage` via + * `mcpMethodCall`. When `sampling.tools` is present, the producer also + * accepts SEP-1577 `tools` / `toolChoice` / `tool_use` content blocks + * inside `CreateMessageRequest`. + * + * @category MCP Server State + * @see {@link https://github.com/modelcontextprotocol/ext-apps/blob/main/specification/draft/apps.mdx | MCP Apps spec (SEP-1865)} + */ +export interface AhpMcpUiHostCapabilities { + /** Producer proxies the MCP `tools/*` methods to the upstream server. */ + serverTools?: { + /** Producer forwards `notifications/tools/list_changed` from the server. */ + listChanged?: boolean; + }; + /** Producer proxies the MCP `resources/*` methods to the upstream server. */ + serverResources?: { + /** Producer forwards `notifications/resources/list_changed` from the server. */ + listChanged?: boolean; + }; + /** Producer accepts `notifications/message` log entries from the App via `mcpNotification`. */ + logging?: Record; + /** Producer serves `sampling/createMessage` via `mcpMethodCall`. */ + sampling?: { + /** + * Producer accepts SEP-1577 `tools` / `toolChoice` / `tool_use` content + * blocks inside `CreateMessageRequest`. + */ + tools?: Record; + }; +} + // ─── Common Types ──────────────────────────────────────────────────────────── /** @@ -1837,6 +2067,11 @@ export interface UsageInfo { model?: string; /** Tokens read from cache */ cacheReadTokens?: number; + /** + * Additional provider-specific metadata for this usage report. + * Clients MAY look for well-known optional keys here to provide enhanced UI. + */ + _meta?: Record; } /** diff --git a/src/vs/platform/agentHost/common/state/protocol/version/registry.ts b/src/vs/platform/agentHost/common/state/protocol/version/registry.ts index 74691c764bce8..7ad315ab78b23 100644 --- a/src/vs/platform/agentHost/common/state/protocol/version/registry.ts +++ b/src/vs/platform/agentHost/common/state/protocol/version/registry.ts @@ -108,6 +108,9 @@ export const ACTION_INTRODUCED_IN: { readonly [K in StateAction['type']]: string [ActionType.TerminalCommandDetectionAvailable]: '0.1.0', [ActionType.TerminalCommandExecuted]: '0.1.0', [ActionType.TerminalCommandFinished]: '0.1.0', + [ActionType.McpServerAdded]: '0.2.0', + [ActionType.McpServerRemoved]: '0.2.0', + [ActionType.McpServerStatusChanged]: '0.2.0', }; /** diff --git a/src/vs/platform/agentHost/common/state/sessionActions.ts b/src/vs/platform/agentHost/common/state/sessionActions.ts index d626af497249d..6214dca80a8c7 100644 --- a/src/vs/platform/agentHost/common/state/sessionActions.ts +++ b/src/vs/platform/agentHost/common/state/sessionActions.ts @@ -53,6 +53,9 @@ export { type SessionIsReadChangedAction, type SessionIsArchivedChangedAction, type SessionToolCallContentChangedAction, + type McpServerAddedAction, + type McpServerRemovedAction, + type McpServerStatusChangedAction, type StateAction, } from './protocol/actions.js'; diff --git a/src/vs/platform/agentHost/node/agentHostMain.ts b/src/vs/platform/agentHost/node/agentHostMain.ts index e2eb11b0f1b99..d23ca55c7afc6 100644 --- a/src/vs/platform/agentHost/node/agentHostMain.ts +++ b/src/vs/platform/agentHost/node/agentHostMain.ts @@ -51,10 +51,14 @@ import { AgentHostClientFileSystemProvider } from '../common/agentHostClientFile import { AGENT_CLIENT_SCHEME } from '../common/agentClientUri.js'; import { AGENT_HOST_CLIENT_RESOURCE_CHANNEL, createAgentHostClientResourceConnection } from '../common/agentHostClientResourceChannel.js'; import { IAgentPluginManager } from '../common/agentPluginManager.js'; +import { IMcpHostService } from '../common/mcpHost/mcpHostService.js'; import { AgentPluginManager } from './agentPluginManager.js'; import { AgentHostGitService, IAgentHostGitService } from './agentHostGitService.js'; +import { McpHostServiceImpl } from './mcpHost/mcpHostServiceImpl.js'; +import { McpProxyFactory } from './mcpHost/mcpProxy.js'; import { registerPendingEditContentProvider } from './copilot/pendingEditContentStore.js'; import { join } from '../../../base/common/path.js'; +import { NullMcpHostService } from '../common/mcpHost/nullMcpHostService.js'; // Entry point for the agent host utility process. // Sets up IPC, logging, and registers agent providers (Copilot). @@ -97,6 +101,12 @@ function startAgentHost(): void { const sessionDataService = new SessionDataService(URI.file(environmentService.userDataPath), fileService, logService); const rootConfigResource = joinPath(environmentService.appSettingsHome, 'globalStorage', 'agent-host-config.json'); + // Single shared MCP host service backed by the per-server proxy factory. + // The host service requires `agentService.stateManager` (constructed inside + // `AgentService`) so the real wiring happens via {@link AgentService.setMcpHostService} + // once both have been built. + let mcpHostService = new NullMcpHostService(); + // Create the real service implementation that lives in this process let agentService: AgentService; try { @@ -118,6 +128,10 @@ function startAgentHost(): void { const claudeAgentSdkService = instantiationService.createInstance(ClaudeAgentSdkService); diServices.set(IClaudeAgentSdkService, claudeAgentSdkService); agentService = new AgentService(logService, fileService, sessionDataService, productService, gitService, rootConfigResource); + const proxyFactory = disposables.add(new McpProxyFactory(logService)); + mcpHostService = disposables.add(new McpHostServiceImpl(agentService.stateManager, proxyFactory, logService)); + agentService.setMcpHostService(mcpHostService); + diServices.set(IMcpHostService, mcpHostService); const pluginManager = new AgentPluginManager(URI.file(environmentService.userDataPath), fileService, logService); diServices.set(IAgentPluginManager, pluginManager); const diffComputeService = disposables.add(new NodeWorkerDiffComputeService(logService)); @@ -201,6 +215,7 @@ function startAgentHost(): void { completionTriggerCharacters: agentService.completionTriggerCharacters, }, clientFileSystemProvider, + mcpHostService, logService, )); disposables.add(protocolHandler.onDidChangeConnectionCount(count => connectionCountEmitter.fire(count))); @@ -260,7 +275,7 @@ function startAgentHost(): void { server.registerChannel(AgentHostIpcChannels.ConnectionTracker, connectionTrackerChannel); // Start WebSocket server for external clients if configured (env-var flow for CLI/server) - startWebSocketServer(agentService, clientFileSystemProvider, logService, disposables, count => connectionCountEmitter.fire(count)).catch(err => { + startWebSocketServer(agentService, clientFileSystemProvider, mcpHostService, logService, disposables, count => connectionCountEmitter.fire(count)).catch(err => { logService.error('Failed to start WebSocket server', err); }); @@ -277,7 +292,7 @@ function startAgentHost(): void { * This reuses the same {@link AgentService} and {@link AgentHostStateManager} * that the IPC channel uses, so both IPC and WebSocket clients share state. */ -async function startWebSocketServer(agentService: AgentService, clientFileSystemProvider: AgentHostClientFileSystemProvider, logService: ILogService, disposables: DisposableStore, onConnectionCountChanged: (count: number) => void): Promise { +async function startWebSocketServer(agentService: AgentService, clientFileSystemProvider: AgentHostClientFileSystemProvider, mcpHostService: IMcpHostService, logService: ILogService, disposables: DisposableStore, onConnectionCountChanged: (count: number) => void): Promise { const port = process.env['VSCODE_AGENT_HOST_PORT']; const socketPath = process.env['VSCODE_AGENT_HOST_SOCKET_PATH']; @@ -315,6 +330,7 @@ async function startWebSocketServer(agentService: AgentService, clientFileSystem completionTriggerCharacters: agentService.completionTriggerCharacters, }, clientFileSystemProvider, + mcpHostService, logService, )); disposables.add(protocolHandler.onDidChangeConnectionCount(onConnectionCountChanged)); diff --git a/src/vs/platform/agentHost/node/agentHostServerMain.ts b/src/vs/platform/agentHost/node/agentHostServerMain.ts index 9944ed577b362..9b90d597f90be 100644 --- a/src/vs/platform/agentHost/node/agentHostServerMain.ts +++ b/src/vs/platform/agentHost/node/agentHostServerMain.ts @@ -55,6 +55,9 @@ import { AGENT_CLIENT_SCHEME } from '../common/agentClientUri.js'; import { resolveServerUrls } from './serverUrls.js'; import { AgentPluginManager } from './agentPluginManager.js'; import { IAgentPluginManager } from '../common/agentPluginManager.js'; +import { IMcpHostService } from '../common/mcpHost/mcpHostService.js'; +import { McpHostServiceImpl } from './mcpHost/mcpHostServiceImpl.js'; +import { McpProxyFactory } from './mcpHost/mcpProxy.js'; import { registerPendingEditContentProvider } from './copilot/pendingEditContentStore.js'; import { AgentHostGitService, IAgentHostGitService } from './agentHostGitService.js'; @@ -196,6 +199,14 @@ async function main(): Promise { const agentService = new AgentService(logService, fileService, sessionDataService, productService, gitService, rootConfigResource); disposables.add(agentService); + // Single shared MCP host service backed by the per-server proxy factory. + // Constructed after `agentService` because it consumes the agent service's + // `AgentHostStateManager`. + const proxyFactory = disposables.add(new McpProxyFactory(logService)); + const mcpHostService: IMcpHostService = disposables.add(new McpHostServiceImpl(agentService.stateManager, proxyFactory, logService)); + agentService.setMcpHostService(mcpHostService); + diServices.set(IMcpHostService, mcpHostService); + // Register agents if (!options.quiet) { // Production agents (require DI) @@ -259,6 +270,7 @@ async function main(): Promise { completionTriggerCharacters: agentService.completionTriggerCharacters, }, clientFileSystemProvider, + mcpHostService, logService, )); diff --git a/src/vs/platform/agentHost/node/agentHostStateManager.ts b/src/vs/platform/agentHost/node/agentHostStateManager.ts index c3ac5758ca6cf..f8ae8b08c21b4 100644 --- a/src/vs/platform/agentHost/node/agentHostStateManager.ts +++ b/src/vs/platform/agentHost/node/agentHostStateManager.ts @@ -12,6 +12,7 @@ import { ActionType, NotificationType, ActionEnvelope, ActionOrigin, INotificati import type { IStateSnapshot } from '../common/state/sessionProtocol.js'; import { rootReducer, sessionReducer } from '../common/state/sessionReducers.js'; import { createRootState, createSessionState, SessionLifecycle, type RootState, type SessionMeta, type SessionState, type SessionSummary, type Turn, type URI, ROOT_STATE_URI } from '../common/state/sessionState.js'; +import type { McpServerStatus, McpServerSummary } from '../common/state/protocol/state.js'; import { IPermissionsValue, platformRootSchema } from '../common/agentHostSchema.js'; import { SessionConfigKey } from '../common/sessionConfigKeys.js'; @@ -145,7 +146,7 @@ export class AgentHostStateManager extends Disposable { * and writes its on-disk metadata). Call {@link markSessionPersisted} * afterwards to fire the deferred notification. */ - createSession(summary: SessionSummary, options?: { readonly emitNotification?: boolean }): SessionState { + createSession(summary: SessionSummary, options?: { readonly emitNotification?: boolean; readonly mcpServers?: readonly McpServerSummary[] }): SessionState { const key = summary.resource; if (this._sessionStates.has(key)) { this._logService.warn(`[AgentHostStateManager] Session already exists: ${key}`); @@ -153,6 +154,9 @@ export class AgentHostStateManager extends Disposable { } const state = createSessionState(summary); + if (options?.mcpServers && options.mcpServers.length > 0) { + state.mcpServers = [...options.mcpServers]; + } this._sessionStates.set(key, state); this._logService.trace(`[AgentHostStateManager] Created session: ${key}`); @@ -314,6 +318,26 @@ export class AgentHostStateManager extends Disposable { this.dispatchServerAction({ type: ActionType.SessionMetaChanged, session, _meta: meta }); } + // ---- MCP servers -------------------------------------------------------- + + /** Registers a new MCP server with a session. */ + createMcpServer(session: URI, summary: McpServerSummary): void { + this._logService.trace(`[AgentHostStateManager] createMcpServer: ${summary.resource} on ${session}`); + this.dispatchServerAction({ type: ActionType.McpServerAdded, session, server: summary }); + } + + /** Removes an MCP server from a session. */ + removeMcpServer(session: URI, mcpServer: URI): void { + this._logService.trace(`[AgentHostStateManager] removeMcpServer: ${mcpServer} from ${session}`); + this.dispatchServerAction({ type: ActionType.McpServerRemoved, session, mcpServer }); + } + + /** Updates the status of an MCP server. */ + setMcpServerStatus(session: URI, mcpServer: URI, status: McpServerStatus): void { + this._logService.trace(`[AgentHostStateManager] setMcpServerStatus: ${mcpServer} -> ${status.kind}`); + this.dispatchServerAction({ type: ActionType.McpServerStatusChanged, session, mcpServer, status }); + } + // ---- Turn tracking ------------------------------------------------------ /** @@ -405,6 +429,9 @@ export class AgentHostStateManager extends Disposable { } } + // MCP server lifecycle actions are session-scoped and handled by the + // session reducer above; there is no longer a per-server state slice. + // Emit envelope const envelope: ActionEnvelope = { action, diff --git a/src/vs/platform/agentHost/node/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts index 8ae86c7ad744a..87a0bf2c95f2d 100644 --- a/src/vs/platform/agentHost/node/agentService.ts +++ b/src/vs/platform/agentHost/node/agentService.ts @@ -37,6 +37,8 @@ import { AgentHostStateManager } from './agentHostStateManager.js'; import { IAgentHostGitService } from './agentHostGitService.js'; import { AgentHostCompletions, IAgentHostCompletions } from './agentHostCompletions.js'; import { AgentHostFileCompletionProvider } from './agentHostFileCompletionProvider.js'; +import { IMcpHostService } from '../common/mcpHost/mcpHostService.js'; +import { NullMcpHostService } from '../common/mcpHost/nullMcpHostService.js'; import { AgentHostWorkspaceFiles } from './agentHostWorkspaceFiles.js'; import { toAgentClientUri } from '../common/agentClientUri.js'; @@ -91,6 +93,9 @@ export class AgentService extends Disposable implements IAgentService { /** Pluggable completion item providers (e.g. workspace file completions, agent-specific @-mentions). */ private readonly _completions: IAgentHostCompletions; + /** Routes MCP traffic to the per-server proxies. Defaults to {@link NullMcpHostService} for surfaces (e.g. tests) that don't host MCP. */ + private _mcpHostService: IMcpHostService; + /** * Authoritative server-side per-resource subscription refcount, keyed by * resource URI string and valued by the set of subscribed protocol @@ -126,8 +131,10 @@ export class AgentService extends Disposable implements IAgentService { private readonly _productService: IProductService, private readonly _gitService: IAgentHostGitService, private readonly _rootConfigResource?: URI, + mcpHostService?: IMcpHostService, ) { super(); + this._mcpHostService = mcpHostService ?? new NullMcpHostService(); this._logService.info('AgentService initialized'); this._stateManager = this._register(new AgentHostStateManager(_logService)); this._register(this._stateManager.onDidEmitEnvelope(e => this._onDidAction.fire(e))); @@ -168,6 +175,19 @@ export class AgentService extends Disposable implements IAgentService { this._terminalManager = this._register(instantiationService.createInstance(AgentHostTerminalManager, this._stateManager)); } + /** + * Replace the MCP host service used to route MCP traffic. + * + * Production wires the real {@link McpHostServiceImpl} via this setter + * after both the agent service and the proxy factory have been + * constructed (the host service requires the agent service's + * {@link AgentHostStateManager}, so circular construction is avoided + * via setter-injection rather than constructor-injection). + */ + setMcpHostService(svc: IMcpHostService): void { + this._mcpHostService = svc; + } + // ---- provider registration ---------------------------------------------- registerProvider(provider: IAgent): void { @@ -191,7 +211,26 @@ export class AgentService extends Disposable implements IAgentService { // ---- auth --------------------------------------------------------------- async authenticate(params: AuthenticateParams): Promise { - this._logService.trace(`[AgentService] authenticate called: resource=${params.resource}`); + this._logService.trace(`[AgentService] authenticate called: resource=${params.resource}, server=${params.server?.toString() ?? '(agent-level)'}`); + + // Per-server auth path: token is scoped to one MCP server's proxy + // only. The host service's token vault is per-handle, so + // adversarial servers don't see each others' tokens. + if (params.server) { + const handle = this._mcpHostService.getServer(params.server); + if (!handle) { + this._logService.warn(`[AgentService] authenticate: no MCP server registered for ${params.server.toString()}`); + return { authenticated: false }; + } + try { + const ok = await handle.authenticate(params.resource, params.token); + return { authenticated: ok }; + } catch (err) { + this._logService.error(err, `[AgentService] authenticate failed for MCP server ${params.server.toString()}`); + return { authenticated: false }; + } + } + // Multiple providers may share the same protected resource (e.g. // both Copilot CLI and Claude consume the GitHub Copilot token). // Fan out to every matching provider in parallel; the request is @@ -338,6 +377,7 @@ export class AgentService extends Disposable implements IAgentService { this._logService.trace(`[AgentService] createSession returned: ${session.toString()}`); const sessionConfig = await this._resolveCreatedSessionConfig(provider, config); + const mcpServers = this._mcpHostService.getServerSummaries(session); // When forking, populate the new session's protocol state with // the source session's turns so the client sees the forked history. @@ -350,7 +390,7 @@ export class AgentService extends Disposable implements IAgentService { } const summary = this._buildInitialSummary(provider, session, config, created, sourceState?.summary.title ?? 'Forked Session'); - const state = this._stateManager.createSession(summary); + const state = this._stateManager.createSession(summary, { mcpServers }); state.config = sessionConfig; state.turns = sourceTurns; state.activeClient = config.activeClient; @@ -362,7 +402,7 @@ export class AgentService extends Disposable implements IAgentService { // clients can subscribe and stream config / model changes that // the agent will pick up at materialization time. const summary = this._buildInitialSummary(provider, session, config, created, ''); - const state = this._stateManager.createSession(summary, { emitNotification: !created.provisional }); + const state = this._stateManager.createSession(summary, { emitNotification: !created.provisional, mcpServers }); state.config = sessionConfig; state.activeClient = config?.activeClient; } @@ -576,6 +616,12 @@ export class AgentService extends Disposable implements IAgentService { let snapshot = this._stateManager.getSnapshot(resourceStr); if (!snapshot) { + // `mcp:/` resources have no restore semantics: the host service + // owns their lifetime. If no entry exists, the server isn't + // registered — there is nothing to rehydrate. + if (resourceStr.startsWith('mcp:/')) { + throw new Error(`Cannot subscribe to unknown MCP server: ${resourceStr}`); + } // Try subagent restore before regular session restore const parsed = parseSubagentSessionUri(resource); if (parsed) { diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts index 8ded0d56227fb..e663a12211538 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts @@ -6,19 +6,19 @@ import { CopilotClient, ResumeSessionConfig, type CopilotClientOptions, type SessionConfig } from '@github/copilot-sdk'; import { rgPath } from '@vscode/ripgrep'; import * as fs from 'fs/promises'; -import { Limiter, SequencerByKey } from '../../../../base/common/async.js'; +import { Limiter, raceTimeout, SequencerByKey } from '../../../../base/common/async.js'; +import { CancellationToken } from '../../../../base/common/cancellation.js'; import { Emitter, Event } from '../../../../base/common/event.js'; import { appendEscapedMarkdownInlineCode } from '../../../../base/common/htmlContent.js'; -import { Disposable, DisposableMap, toDisposable } from '../../../../base/common/lifecycle.js'; -import { ResourceMap } from '../../../../base/common/map.js'; +import { Disposable, DisposableMap, DisposableResourceMap, toDisposable } from '../../../../base/common/lifecycle.js'; import { FileAccess } from '../../../../base/common/network.js'; import { equals } from '../../../../base/common/objects.js'; -import { observableValue } from '../../../../base/common/observable.js'; +import { observableValue, waitForState } from '../../../../base/common/observable.js'; import { basename, delimiter, dirname } from '../../../../base/common/path.js'; import { URI } from '../../../../base/common/uri.js'; import { generateUuid } from '../../../../base/common/uuid.js'; import { localize } from '../../../../nls.js'; -import { IParsedPlugin, parsePlugin } from '../../../agentPlugins/common/pluginParsers.js'; +import { IMcpServerDefinition, IParsedPlugin, parsePlugin } from '../../../agentPlugins/common/pluginParsers.js'; import { IFileService } from '../../../files/common/files.js'; import { IInstantiationService } from '../../../instantiation/common/instantiation.js'; import { ILogService } from '../../../log/common/log.js'; @@ -26,10 +26,12 @@ import { AgentHostConfigKey, agentHostCustomizationConfigSchema } from '../../co import { AutoApproveLevel, ISchemaProperty, SessionMode, createSchema, platformSessionSchema, schemaProperty } from '../../common/agentHostSchema.js'; import { IAgentPluginManager, ISyncedCustomization } from '../../common/agentPluginManager.js'; import { AgentSession, AgentSignal, GITHUB_COPILOT_PROTECTED_RESOURCE, IAgent, IAgentCreateSessionConfig, IAgentCreateSessionResult, IAgentDescriptor, IAgentMaterializeSessionEvent, IAgentModelInfo, IAgentResolveSessionConfigParams, IAgentSessionConfigCompletionsParams, IAgentSessionMetadata, IAgentSessionProjectInfo } from '../../common/agentService.js'; +import { IMcpHostService } from '../../common/mcpHost/mcpHostService.js'; import { SessionConfigKey } from '../../common/sessionConfigKeys.js'; import { ISessionDataService, SESSION_DB_FILENAME } from '../../common/sessionDataService.js'; +import { buildMcpServerUri } from '../../common/state/mcpServerUri.js'; import type { ResolveSessionConfigResult, SessionConfigCompletionsResult } from '../../common/state/protocol/commands.js'; -import { ProtectedResourceMetadata, type ConfigSchema, type ModelSelection, type SessionCustomization, type ToolDefinition } from '../../common/state/protocol/state.js'; +import { ProtectedResourceMetadata, McpServerStatusKind, type ConfigSchema, type ModelSelection, type SessionCustomization, type ToolDefinition } from '../../common/state/protocol/state.js'; import { AHP_AUTH_REQUIRED, ProtocolError } from '../../common/state/sessionProtocol.js'; import { CustomizationRef, CustomizationStatus, ResponsePartKind, SessionInputResponseKind, parseSubagentSessionUri, type MessageAttachment, type PendingMessage, type PolicyState, type ResponsePart, type SessionInputAnswer, type ToolCallResult, type Turn } from '../../common/state/sessionState.js'; import { IAgentConfigurationService } from '../agentConfigurationService.js'; @@ -228,6 +230,13 @@ function prependAnnouncementToFirstTurn( export class CopilotAgent extends Disposable implements IAgent { readonly id = 'copilotcli' as const; private static readonly _BRANCH_COMPLETION_LIMIT = 25; + /** + * How long {@link _resolveMcpServersForSdk} waits for each MCP server's + * proxy endpoint to bind before dropping it from the SDK config. The + * server will be re-attached on the next session refresh once the proxy + * comes online. + */ + private static readonly _MCP_SERVER_READY_TIMEOUT_MS = 10_000; private readonly _onDidSessionProgress = this._register(new Emitter()); readonly onDidSessionProgress = this._onDidSessionProgress.event; @@ -263,7 +272,7 @@ export class CopilotAgent extends Disposable implements IAgent { private readonly _plugins: PluginController; readonly onDidCustomizationsChange: Event; /** Per-session active client state for tools + plugin snapshot tracking. */ - private readonly _activeClients = new ResourceMap(); + private readonly _activeClients = this._register(new DisposableResourceMap()); constructor( @ILogService private readonly _logService: ILogService, @@ -273,10 +282,24 @@ export class CopilotAgent extends Disposable implements IAgent { @IAgentHostGitService private readonly _gitService: IAgentHostGitService, @IAgentHostTerminalManager private readonly _terminalManager: IAgentHostTerminalManager, @IAgentConfigurationService private readonly _configurationService: IAgentConfigurationService, + @IMcpHostService private readonly _mcpHostService: IMcpHostService, ) { super(); this._plugins = this._register(this._instantiationService.createInstance(PluginController)); this.onDidCustomizationsChange = this._plugins.onDidChange; + // Republish MCP servers for every session that has an active client + // whenever the parsed plugin set changes. Each ActiveClient dedups + // against its last-published set so unchanged sessions are no-ops. + // Iterating active clients (rather than only live SDK sessions) is + // what lets pre-materialization (provisional) sessions warm up + // their MCP proxies before the user sends their first message. + this._register(this._plugins.onDidChange(() => { + for (const ac of this._activeClients.values()) { + void ac.republish().catch(err => { + this._logService.warn(`[Copilot] MCP republish on plugin change failed: ${err instanceof Error ? err.message : String(err)}`); + }); + } + })); } protected _createCopilotClient(options: CopilotClientOptions): CopilotClient { @@ -752,6 +775,7 @@ export class CopilotAgent extends Disposable implements IAgent { if (config.activeClient.customizations !== undefined) { await this._plugins.sync(config.activeClient.clientId, config.activeClient.customizations); } + await ac.snapshot(config.workingDirectory); } // Compute project metadata cheaply from the original working dir. @@ -770,7 +794,12 @@ export class CopilotAgent extends Disposable implements IAgent { } this._logService.info(`[Copilot] Session created (provisional): ${sessionUri.toString()}`); - return { session: sessionUri, workingDirectory: config.workingDirectory, provisional: true, ...(project ? { project } : {}) }; + return { + session: sessionUri, + workingDirectory: config.workingDirectory, + provisional: true, + ...(project ? { project } : {}), + }; } /** @@ -815,7 +844,7 @@ export class CopilotAgent extends Disposable implements IAgent { const snapshot = activeClient ? await activeClient.snapshot(customizationDirectory) : undefined; const workingDirectory = await this._resolveSessionWorkingDirectory(materializedConfig, sessionId, prompt); const shellManager = this._instantiationService.createInstance(ShellManager, sessionUri, workingDirectory); - const sessionConfigBuilder = this._buildSessionConfig(snapshot, shellManager); + const sessionConfigBuilder = this._buildSessionConfig(sessionUri, snapshot, shellManager); const factory: SessionWrapperFactory = async callbacks => { const raw = await client.createSession({ @@ -1300,7 +1329,7 @@ export class CopilotAgent extends Disposable implements IAgent { private _getOrCreateActiveClient(session: URI): ActiveClient { let client = this._activeClients.get(session); if (!client) { - client = new ActiveClient(directory => this._plugins.getAppliedPlugins(directory)); + client = this._instantiationService.createInstance(ActiveClient, session, directory => this._plugins.getAppliedPlugins(directory)); this._activeClients.set(session, client); } return client; @@ -1333,14 +1362,20 @@ export class CopilotAgent extends Disposable implements IAgent { } private async _destroyAndDisposeSession(sessionId: string): Promise { + const sessionUri = AgentSession.uri(this.id, sessionId); + // Disposing the active client clears its published MCP servers via + // `setSessionServers([])`. Idempotent when the session never had + // any registered (e.g. provisional sessions disposed before + // materialization). + this._activeClients.deleteAndDispose(sessionUri); + // Provisional sessions have no SDK session, no worktree, and no - // on-disk metadata — drop the in-memory record and clean up the - // active-client snapshot. The state-manager entry is removed by the - // caller via {@link IAgentService.disposeSession}. + // on-disk metadata — drop the in-memory record. The state-manager + // entry is removed by the caller via + // {@link IAgentService.disposeSession}. const provisional = this._provisionalSessions.get(sessionId); if (provisional) { this._provisionalSessions.delete(sessionId); - this._activeClients.delete(provisional.sessionUri); return; } const entry = this._sessions.get(sessionId); @@ -1362,18 +1397,27 @@ export class CopilotAgent extends Disposable implements IAgent { * Returns an async function that resolves the final config given the * session's permission/hook callbacks, so it can be called lazily * inside the {@link SessionWrapperFactory}. + * + * MCP servers are routed through {@link IMcpHostService}'s per-server + * proxies. Servers whose proxy endpoint is not yet bound are skipped + * from the SDK config; once the endpoint becomes available, the next + * customization-change-driven session refresh (via + * {@link ActiveClient.isOutdated}) picks them up. This means the SDK's + * very first connect after materialization may see an empty + * `mcpServers` set. */ - private _buildSessionConfig(snapshot: IActiveClientSnapshot | undefined, shellManager: ShellManager): (args: Parameters[0]) => Promise { + private _buildSessionConfig(sessionUri: URI, snapshot: IActiveClientSnapshot | undefined, shellManager: ShellManager): (args: Parameters[0]) => Promise { const plugins = snapshot?.plugins ?? []; return async (callbacks: Parameters[0]) => { const shellTools = await createShellTools(shellManager, this._terminalManager, this._logService); const customAgents = await toSdkCustomAgents(plugins.flatMap(p => p.agents), this._fileService); + const mcpServers = await this._resolveMcpServersForSdk(sessionUri, plugins.flatMap(p => p.mcpServers)); return { onPermissionRequest: callbacks.onPermissionRequest, onUserInputRequest: callbacks.onUserInputRequest, hooks: toSdkHooks(plugins.flatMap(p => p.hooks), callbacks.hooks), - mcpServers: toSdkMcpServers(plugins.flatMap(p => p.mcpServers)), + mcpServers, customAgents, skillDirectories: toSdkSkillDirectories(plugins.flatMap(p => p.skills)), systemMessage: COPILOT_AGENT_HOST_SYSTEM_MESSAGE, @@ -1388,6 +1432,63 @@ export class CopilotAgent extends Disposable implements IAgent { }; } + /** + * Resolves each MCP server's local proxy endpoint and produces the SDK + * `mcpServers` config. Waits up to + * {@link CopilotAgent._MCP_SERVER_READY_TIMEOUT_MS} for each server to + * reach {@link McpServerStatusKind.Ready}. Servers in any other state + * (Starting after timeout, AuthRequired, Error, Stopped) are dropped + * from this round of SDK config — we never advertise an unhealthy or + * pending-auth server to the SDK because the SDK will eagerly try to + * `initialize` against it and fail. Skipped servers are picked up on + * the next session refresh, which fires when the active client's + * customizations change or when the user authenticates and the + * server transitions to Ready. + */ + private async _resolveMcpServersForSdk(sessionUri: URI, defs: readonly IMcpServerDefinition[]): Promise> { + const resolved = new Map(); + await Promise.all(defs.map(async def => { + const handle = this._mcpHostService.getServer(buildMcpServerUri(sessionUri, def.name)); + if (!handle) { + return; + } + const currentKind = handle.summary.get().status.kind; + if (currentKind === McpServerStatusKind.Ready) { + const endpoint = handle.endpoint.get(); + if (endpoint) { + resolved.set(def, endpoint); + } + return; + } + if (currentKind !== McpServerStatusKind.Starting) { + // AuthRequired / Error / Stopped — nothing to wait for. + this._logService.info(`[Copilot] MCP server '${def.name}' is in state '${currentKind}'; not advertising to SDK`); + return; + } + try { + const readySummary = await raceTimeout( + waitForState(handle.summary, value => value.status.kind !== McpServerStatusKind.Starting, undefined, CancellationToken.None), + CopilotAgent._MCP_SERVER_READY_TIMEOUT_MS, + ); + if (!readySummary) { + this._logService.warn(`[Copilot] MCP server '${def.name}' did not reach Ready within ${CopilotAgent._MCP_SERVER_READY_TIMEOUT_MS}ms; skipping`); + return; + } + if (readySummary.status.kind !== McpServerStatusKind.Ready) { + this._logService.info(`[Copilot] MCP server '${def.name}' settled in '${readySummary.status.kind}'; not advertising to SDK`); + return; + } + const endpoint = handle.endpoint.get(); + if (endpoint) { + resolved.set(def, endpoint); + } + } catch (err) { + this._logService.warn(`[Copilot] MCP server '${def.name}' readiness wait failed: ${err instanceof Error ? err.message : String(err)}`); + } + })); + return toSdkMcpServers(defs, def => resolved.get(def)); + } + protected async _resumeSession(sessionId: string): Promise { this._logService.info(`[Copilot:${sessionId}] _resumeSession called — session not in memory, resuming...`); const client = await this._ensureClient(); @@ -1407,7 +1508,7 @@ export class CopilotAgent extends Disposable implements IAgent { } const shellManager = this._instantiationService.createInstance(ShellManager, sessionUri, workingDirectory); - const sessionConfig = this._buildSessionConfig(snapshot, shellManager); + const sessionConfig = this._buildSessionConfig(sessionUri, snapshot, shellManager); const factory: SessionWrapperFactory = async callbacks => { const config = await sessionConfig(callbacks); @@ -1910,18 +2011,31 @@ class PluginController extends Disposable { } /** - * Tracks per-session active client contributions (tools and plugins). - * The {@link snapshot} captures the state at session creation time, and - * {@link isOutdated} detects when the session needs to be refreshed. + * Tracks per-session active client contributions (tools and plugins) AND + * owns the publishing of resolved MCP servers into {@link IMcpHostService}. + * + * Every {@link snapshot} call refreshes the resolved plugin set and + * forwards the resulting MCP server definitions to the host service, + * deduplicating against the previously published set. {@link republish} + * does the same without changing the captured tool/clientId state — used + * by plugin-change subscribers that don't have a directory in hand. The + * disposable lifecycle is bound to the session: disposing this instance + * clears the session's MCP servers via `setSessionServers([])`. */ -class ActiveClient { +export class ActiveClient extends Disposable { private _tools: readonly ToolDefinition[] = []; private _clientId = ''; + private _lastDirectory: URI | undefined; + private _publishedMcpServers: readonly IMcpServerDefinition[] | undefined; constructor( - /** Resolves the current set of applied plugins. May block while a sync is in progress. */ + private readonly _sessionUri: URI, private readonly _resolvePlugins: (directory: URI | undefined) => Promise, - ) { } + @IMcpHostService private readonly _mcpHostService: IMcpHostService, + @ILogService private readonly _logger: ILogService, + ) { + super(); + } updateTools(clientId: string, tools: readonly ToolDefinition[]): void { this._clientId = clientId; @@ -1929,7 +2043,24 @@ class ActiveClient { } async snapshot(directory: URI | undefined): Promise { - return { clientId: this._clientId, tools: this._tools, plugins: await this._resolvePlugins(directory) }; + this._lastDirectory = directory; + const plugins = await this._resolvePlugins(directory); + this._republishMcpServers(plugins); + return { + clientId: this._clientId, + tools: this._tools, + plugins, + mcpReadiness: this._captureMcpReadiness(plugins), + }; + } + + /** + * Re-resolves plugins using the most recent {@link snapshot} directory + * and republishes the MCP server set. Idempotent via the dedup check. + */ + async republish(): Promise { + const plugins = await this._resolvePlugins(this._lastDirectory); + this._republishMcpServers(plugins); } async isOutdated(snap: IActiveClientSnapshot, directory: URI | undefined): Promise { @@ -1951,6 +2082,63 @@ class ActiveClient { return true; } } + // Detect MCP server readiness flips: when a previously skipped + // server (AuthRequired / Error / Stopped) reaches Ready (typically + // after `authenticate`), the cached SDK session — built without + // that server's tools — is stale and must be rebuilt. + const currentReadiness = this._captureMcpReadiness(plugins); + for (const [resource, ready] of Object.entries(currentReadiness)) { + if (snap.mcpReadiness[resource] !== ready) { + return true; + } + } + for (const resource of Object.keys(snap.mcpReadiness)) { + if (!Object.prototype.hasOwnProperty.call(currentReadiness, resource)) { + return true; + } + } return false; } + + override dispose(): void { + if (this._publishedMcpServers !== undefined) { + try { + this._mcpHostService.setSessionServers(this._sessionUri, []); + } catch (err) { + this._logger.warn(`[Copilot:ActiveClient] dispose: failed to clear MCP servers for ${this._sessionUri.toString()}: ${err instanceof Error ? err.message : String(err)}`); + } + this._publishedMcpServers = undefined; + } + super.dispose(); + } + + private _republishMcpServers(plugins: readonly IParsedPlugin[]): void { + const servers = plugins.flatMap(p => p.mcpServers); + if (this._publishedMcpServers !== undefined && equals(this._publishedMcpServers, servers)) { + return; + } + this._publishedMcpServers = servers; + this._mcpHostService.setSessionServers(this._sessionUri, servers); + } + + /** + * Snapshot the current `Ready` flag for every MCP server in `plugins`, + * keyed by the same `mcp:/...` resource URI the SDK config uses. Missing + * handles or any non-`Ready` status (including `AuthRequired`) map to + * `false`; only `Ready` counts as "advertised to the SDK". + * + * The snapshot is compared inside {@link isOutdated} to detect Ready + * transitions that change what `_resolveMcpServersForSdk` would + * advertise, so the agent can rebuild the cached SDK session and + * re-advertise the freshly-authenticated tool set. + */ + private _captureMcpReadiness(plugins: readonly IParsedPlugin[]): Record { + const readiness: Record = {}; + for (const def of plugins.flatMap(p => p.mcpServers)) { + const resource = buildMcpServerUri(this._sessionUri, def.name).toString(); + const handle = this._mcpHostService.getServer(URI.parse(resource)); + readiness[resource] = handle?.summary.get().status.kind === McpServerStatusKind.Ready; + } + return readiness; + } } diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts index 5d643a180b594..c71601e41c53c 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts @@ -27,6 +27,8 @@ import { ISessionDatabase, ISessionDataService, SESSION_ATTACHMENTS_DIRNAME } fr import { MessageAttachmentKind, type FileEdit, type MessageAttachment, type ToolDefinition } from '../../common/state/protocol/state.js'; import { ActionType, type SessionAction } from '../../common/state/sessionActions.js'; import { ResponsePartKind, SessionInputAnswerState, SessionInputAnswerValueKind, SessionInputQuestionKind, SessionInputResponseKind, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType, type PendingMessage, type URI as ProtocolURI, type SessionInputAnswer, type SessionInputRequest, type ToolCallResult, type ToolResultContent, type Turn } from '../../common/state/sessionState.js'; +import { IMcpHostService } from '../../common/mcpHost/mcpHostService.js'; +import { buildMcpServerUri } from '../../common/state/mcpServerUri.js'; import { IAgentConfigurationService } from '../agentConfigurationService.js'; import type { IExitPlanModeRequestParams, IExitPlanModeResponse } from './copilotAgent.js'; import { CopilotSessionWrapper } from './copilotSessionWrapper.js'; @@ -127,6 +129,15 @@ export interface IActiveClientSnapshot { readonly clientId: string; readonly tools: readonly ToolDefinition[]; readonly plugins: readonly IParsedPlugin[]; + /** + * `Ready` state of each MCP server (keyed by `McpServerSummary.resource`) + * at snapshot time. Captured because the SDK config is built from this + * snapshot and only advertises MCP servers that were `Ready` when + * `_resolveMcpServersForSdk` ran — if any server later transitions + * to/from `Ready` (e.g. after `authenticate` succeeds), the cached SDK + * session must be rebuilt so the new tool list is exposed to the model. + */ + readonly mcpReadiness: Readonly>; } /** @@ -248,6 +259,7 @@ export class CopilotAgentSession extends Disposable { @IFileService private readonly _fileService: IFileService, @INativeEnvironmentService private readonly _environmentService: INativeEnvironmentService, @IAgentConfigurationService private readonly _configurationService: IAgentConfigurationService, + @IMcpHostService private readonly _mcpHostService: IMcpHostService, ) { super(); this.sessionId = options.rawSessionId; @@ -258,7 +270,7 @@ export class CopilotAgentSession extends Disposable { this._workingDirectory = options.workingDirectory; this._customizationDirectory = options.customizationDirectory; - this._appliedSnapshot = options.clientSnapshot ?? { clientId: '', tools: [], plugins: [] }; + this._appliedSnapshot = options.clientSnapshot ?? { clientId: '', tools: [], plugins: [], mcpReadiness: {} }; this._clientToolNames = new Set(this._appliedSnapshot.tools.map(t => t.name)); this._databaseRef = sessionDataService.openDatabase(options.sessionUri); @@ -1220,6 +1232,22 @@ export class CopilotAgentSession extends Disposable { if (e.data.mcpToolName) { meta.mcpToolName = e.data.mcpToolName; } + // MCP Apps: when the upstream MCP server advertised a UI + // resource for this tool via `tools/list`, attach the + // `_meta.ui` payload from the spec and the per-tool-call + // host-capability set derived from the upstream's + // `initialize` response. The proxy captures both + // opportunistically; non-app tools and pre-handshake races + // leave `meta` alone. + if (e.data.mcpServerName && e.data.mcpToolName) { + const serverUri = buildMcpServerUri(this.sessionUri, e.data.mcpServerName); + const handle = this._mcpHostService.getServer(serverUri); + const uiMeta = handle?.getToolUiMeta(e.data.mcpToolName); + if (uiMeta) { + meta.ui = uiMeta; + meta.uiHostCapabilities = handle!.getUiHostCapabilities(); + } + } const protocolSession = this._protocolSession(); this._emitAction({ diff --git a/src/vs/platform/agentHost/node/copilot/copilotPluginConverters.ts b/src/vs/platform/agentHost/node/copilot/copilotPluginConverters.ts index 61e93e5a6c1e3..a9b086beda541 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotPluginConverters.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotPluginConverters.ts @@ -6,8 +6,8 @@ import { spawn } from 'child_process'; import type { CustomAgentConfig, MCPServerConfig, SessionConfig } from '@github/copilot-sdk'; import { OperatingSystem, OS } from '../../../../base/common/platform.js'; +import { URI } from '../../../../base/common/uri.js'; import { IFileService } from '../../../files/common/files.js'; -import { McpServerType } from '../../../mcp/common/mcpPlatformTypes.js'; import type { IMcpServerDefinition, INamedPluginResource, IParsedHookCommand, IParsedHookGroup, IParsedPlugin } from '../../../agentPlugins/common/pluginParsers.js'; import { dirname } from '../../../../base/common/path.js'; @@ -24,42 +24,38 @@ type ErrorOccurredHookInput = Parameters { - const result: Record = {}; - for (const def of defs) { - const config = def.configuration; - if (config.type === McpServerType.LOCAL) { - result[def.name] = { - type: 'local', - command: config.command, - args: config.args ? [...config.args] : [], - tools: ['*'], - ...(config.env && { env: toStringEnv(config.env) }), - ...(config.cwd && { cwd: config.cwd }), - }; - } else { - result[def.name] = { - type: 'http', - url: config.url, - tools: ['*'], - ...(config.headers && { headers: { ...config.headers } }), - }; - } - } - return result; -} +export type McpEndpointResolver = (def: IMcpServerDefinition) => URI | undefined; /** - * Ensures all env values are strings (the SDK requires `Record`). + * Converts parsed MCP server definitions into the SDK's `mcpServers` config. + * + * Every server is fronted by a local HTTP proxy owned by `IMcpHostService`; + * `resolveEndpoint` returns the proxy URI for a given definition (or + * `undefined` while the proxy is still booting). Servers without a bound + * endpoint are skipped — they will be re-added on the next session + * refresh once their proxy comes up. */ -function toStringEnv(env: Record): Record { - const result: Record = {}; - for (const [key, value] of Object.entries(env)) { - if (value !== null) { - result[key] = String(value); +export function toSdkMcpServers( + defs: readonly IMcpServerDefinition[], + resolveEndpoint: McpEndpointResolver, +): Record { + const result: Record = {}; + for (const def of defs) { + const endpoint = resolveEndpoint(def); + if (!endpoint) { + continue; } + result[def.name] = { + type: 'http', + url: endpoint.toString(), + tools: ['*'], + }; } return result; } diff --git a/src/vs/platform/agentHost/node/mcpHost/AGENTS.md b/src/vs/platform/agentHost/node/mcpHost/AGENTS.md new file mode 100644 index 0000000000000..e15db4ac941ad --- /dev/null +++ b/src/vs/platform/agentHost/node/mcpHost/AGENTS.md @@ -0,0 +1,123 @@ +# MCP host (`src/vs/platform/agentHost/node/mcpHost/`) + +This folder implements the agent-host's MCP gateway. It owns one +lightweight HTTP proxy per MCP server that the host exposes to the +agent SDK (Copilot CLI today; Claude in the future). Every byte the +SDK exchanges with an MCP server flows through code in this folder. + +## What lives here + +- **`mcpProxy.ts`** — public façade and `IMcpProxyFactory`. Owns the + shared HTTP listener; creates one `McpProxyRoute` per advertised + server. +- **`mcpProxyHttpListener.ts`** — single shared `127.0.0.1:0` + HTTP server. Routes are registered at randomized + `/mcp//message` paths. Loopback-only by design. +- **`mcpProxyRoute.ts`** — per-server JSON-RPC bridge. Blind + pass-through with three taps: + 1. `IInitializeInjector` — adds extension capabilities to + client-→upstream `initialize` requests (e.g. MCP Apps). + 2. `onUpstreamRequest(method, params)` — invoked for every + JSON-RPC **request** the upstream emits. Returns a + `Promise` whose result or error is + written back to the upstream as a JSON-RPC response using the + original request id. + 3. `onUpstreamNotification(method, params)` — invoked for every + JSON-RPC **notification** the upstream emits. Fire-and-forget. +- **`mcpUpstream.ts`** — `IMcpUpstream` abstraction. Implementations: + - `mcpStdioUpstream.ts` — spawns a child process on demand + (lazy spawn — see plan §3.4). + - `mcpHttpUpstream.ts` — HTTP discovery handshake; parses + `WWW-Authenticate` for RFC 9728 / RFC 6750 challenges. +- **`mcpStdioStateHandler.ts`** — graceful shutdown for stdio + children (stdin.end → SIGTERM → SIGKILL). Inlined copy of the + workbench helper to avoid a layer-violating import. +- **`mcpAuthChallengeParser.ts`** — defensive `WWW-Authenticate` + parser; turns 401/403 + RFC 9728 metadata into an + `McpServerStatusAuthRequired`. +- **`mcpInitializeInjector.ts`** — `McpAppsInitializeInjector` adds + the MCP Apps extension capability to the SDK→upstream `initialize` + request while preserving everything else. Independent of AHP + capability negotiation; the upstream MCP server decides whether + it speaks Apps based on this extension. +- **`mcpHostServiceImpl.ts`** — node-side `IMcpHostService`. Owns + the per-(session, server) registry, drives `IMcpProxy` lifetimes, + dispatches `session/mcpServer*` actions through + `AgentHostStateManager`, routes `mcpMethodCall` and + `mcpNotification` traffic, and forwards upstream-originated MCP + requests and notifications to whichever AHP client owns the + session via `IMcpHostUpstreamDelegate`. + +## How it fits into AHP + +1. A provider (e.g. `CopilotAgent`) advertises its MCP servers via + `IAgent.getMcpServersForSession()` and fires + `onDidMcpServersChange` whenever they change. +2. `AgentService` watches the provider event and calls + `IMcpHostService.setSessionServers(session, defs)`. +3. The host service diffs against the live registry. For added + servers it dispatches `session/mcpServerAdded` (status `Starting`) + immediately and kicks off async proxy creation. For removed ones + it dispatches `session/mcpServerRemoved` and disposes the proxy. +4. Once the proxy is bound, the host service hands its endpoint URI + to the SDK via `toSdkMcpServers(...)` (in + [../copilot/copilotPluginConverters.ts](../copilot/copilotPluginConverters.ts)). + The SDK now connects over loopback HTTP. +5. JSON-RPC traffic now flows directly through the bidirectional + `mcpMethodCall` / `mcpNotification` methods on the AHP wire — there + is no per-server state mailbox. Upstream-originated requests are + round-tripped via the host service's installed + `IMcpHostUpstreamDelegate` (typically `ProtocolServerHandler`), + which issues a reverse `mcpMethodCall` to the AHP client and + writes the client's result back as the upstream JSON-RPC response. + Upstream-originated notifications are likewise forwarded via the + delegate's `handleUpstreamNotification`, which sends an outbound + `mcpNotification` to the AHP client. + +## Authentication + +- HTTP MCP servers may return 401/403 with a `WWW-Authenticate` + Bearer challenge. `mcpHttpUpstream.ts` handles the discovery. +- The proxy publishes `McpServerStatusAuthRequired` to AHP; clients + push tokens via `authenticate({ resource, token, server })` (see + protocol §4.5 + Phase 4d). +- Per-server scoping: tokens are stored only on the named handle; + other servers sharing the same `resource` URL keep separate state. + +## Lifecycle constraints + +- **Lazy stdio spawn.** Children are not spawned until the SDK + POSTs to the proxy endpoint. Spawn failure surfaces as + `McpServerStatusKind.Error`. +- **HTTP probe is real.** `McpHttpUpstream.start()` sends a real + `initialize` to the upstream during discovery. Server-bearing + servers will see the probe; this is acceptable per MCP. +- **First-message race.** The first SDK config after session + materialization may have an empty `mcpServers` list because proxy + endpoints bind asynchronously. The next plugins-change cycle + picks up the endpoints. See [MCP_AHP_PLAN](../../../../../../MCP_AHP_PLAN.md) §4c. +- **No SSE today.** Upstream-originated messages on HTTP are + recorded in AHP state but not pushed back to the SDK. Real + bidirectional flow is a Phase 6 follow-up. + +## Known gaps tracked for follow-up + +- **Telemetry.** `ITelemetryService` is not wired into the agent + host process. Plan §6.2 sketches events; implementation requires + a separate-process telemetry channel and is deferred. +- **Sandboxing parity.** Stdio children spawn unsandboxed; matches + the existing workbench gateway. See the TODO comment in + [mcpStdioUpstream.ts](mcpStdioUpstream.ts) and plan §6.3. +- **Per-tool MCP App capability negotiation.** MCP App support is + now negotiated per tool call via `_meta.uiHostCapabilities` on tool + call states (see `AhpMcpUiHostCapabilities` in the protocol). + CopilotAgent does not yet surface MCP App tool calls; when it + does, the producer must populate `_meta.ui` and + `_meta.uiHostCapabilities` with the subset of capabilities the + host actually proxies. + +## Tests + +- [../../test/node/mcpHost/](../../test/node/mcpHost/) — unit tests for each module. +- [../../test/node/mcpHost/mcpProxy.integrationTest.ts](../../test/node/mcpHost/mcpProxy.integrationTest.ts) — end-to-end through the real listener. +- [../../test/node/mcpHost/mcpAppsRoundTrip.integrationTest.ts](../../test/node/mcpHost/mcpAppsRoundTrip.integrationTest.ts) — full MCP Apps lifecycle (initialize → notifications → upstream request → client response). diff --git a/src/vs/platform/agentHost/node/mcpHost/mcpAuthChallengeParser.ts b/src/vs/platform/agentHost/node/mcpHost/mcpAuthChallengeParser.ts new file mode 100644 index 0000000000000..bff84fe86bd52 --- /dev/null +++ b/src/vs/platform/agentHost/node/mcpHost/mcpAuthChallengeParser.ts @@ -0,0 +1,129 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { McpAuthRequiredReason, McpServerStatusAuthRequired, McpServerStatusKind, ProtectedResourceMetadata } from '../../common/state/protocol/state.js'; + +/** + * Parsed parts of an HTTP 401/403 auth challenge. + */ +export interface IMcpAuthChallenge { + /** Required scopes from the `scope=` directive on `WWW-Authenticate`. */ + readonly scopes: readonly string[] | undefined; + /** OAuth `error` directive (`invalid_token`, `insufficient_scope`, etc.) */ + readonly error: string | undefined; + /** OAuth `error_description` directive — human-readable. */ + readonly errorDescription: string | undefined; + /** URL of the RFC 9728 protected resource metadata document. */ + readonly resourceMetadataUrl: string | undefined; +} + +const EMPTY_CHALLENGE: IMcpAuthChallenge = { + scopes: undefined, + error: undefined, + errorDescription: undefined, + resourceMetadataUrl: undefined, +}; + +/** + * Parse a `WWW-Authenticate` header value. Tolerates missing fields and + * returns all-undefined when the header is empty or doesn't start with + * `Bearer`. Whitespace-tolerant; quoted values are unquoted. + * + * Caveat: this is NOT a fully RFC-7235-compliant parser — it handles + * the subset of `Bearer ...` challenges MCP is documented to use. + */ +export function parseWwwAuthenticate(header: string | undefined): IMcpAuthChallenge { + if (!header) { + return EMPTY_CHALLENGE; + } + + const trimmed = header.trim(); + const bearerMatch = /^Bearer(?:\s+(.*))?$/i.exec(trimmed); + if (!bearerMatch) { + return EMPTY_CHALLENGE; + } + + const params = parseChallengeParams(bearerMatch[1] ?? ''); + + const scopeRaw = params.get('scope'); + const scopes = scopeRaw === undefined + ? undefined + : scopeRaw.split(/\s+/).filter(s => s.length > 0); + + return { + scopes: scopes && scopes.length > 0 ? scopes : undefined, + error: params.get('error'), + errorDescription: params.get('error_description'), + resourceMetadataUrl: params.get('resource_metadata'), + }; +} + +/** + * Parse a comma-separated list of `key=value` or `key="quoted value"` + * pairs from the auth-param portion of a Bearer challenge. Defensive: + * malformed input does not throw — keys without recognizable values are + * skipped. + */ +function parseChallengeParams(input: string): Map { + const result = new Map(); + // Matches: key=value where value is either a quoted string (with + // optional escaped chars) or a token of non-whitespace, non-comma chars. + const re = /([A-Za-z0-9_-]+)\s*=\s*(?:"((?:\\.|[^"\\])*)"|([^,\s]*))/g; + let m: RegExpExecArray | null; + while ((m = re.exec(input)) !== null) { + const key = m[1].toLowerCase(); + const quoted = m[2]; + const bare = m[3]; + const value = quoted !== undefined + ? quoted.replace(/\\(.)/g, '$1') + : (bare ?? ''); + result.set(key, value); + } + return result; +} + +/** + * Compose an `McpServerStatusAuthRequired` from an HTTP status, the + * parsed challenge, the resource metadata fetched from + * `resource_metadata`, and whether a prior token had been used. + */ +export function buildAuthRequiredStatus(opts: { + httpStatus: 401 | 403; + challenge: IMcpAuthChallenge; + /** RFC 9728 metadata fetched out-of-band. */ + resource: ProtectedResourceMetadata; + /** True if the upstream had been previously authenticated. */ + hadPriorToken: boolean; +}): McpServerStatusAuthRequired { + const { httpStatus, challenge, resource, hadPriorToken } = opts; + + let reason: McpAuthRequiredReason; + if (httpStatus === 403 && challenge.error === 'insufficient_scope') { + reason = McpAuthRequiredReason.InsufficientScope; + } else if (httpStatus === 401 && hadPriorToken) { + reason = McpAuthRequiredReason.Expired; + } else { + reason = McpAuthRequiredReason.Required; + } + + const requiredScopes = challenge.scopes && challenge.scopes.length > 0 + ? [...challenge.scopes] + : (resource.scopes_supported && resource.scopes_supported.length > 0 + ? [...resource.scopes_supported] + : undefined); + + const status: McpServerStatusAuthRequired = { + kind: McpServerStatusKind.AuthRequired, + reason, + resource, + }; + if (requiredScopes) { + status.requiredScopes = requiredScopes; + } + if (challenge.errorDescription) { + status.description = challenge.errorDescription; + } + return status; +} diff --git a/src/vs/platform/agentHost/node/mcpHost/mcpHostServiceImpl.ts b/src/vs/platform/agentHost/node/mcpHost/mcpHostServiceImpl.ts new file mode 100644 index 0000000000000..69b36ea66c13b --- /dev/null +++ b/src/vs/platform/agentHost/node/mcpHost/mcpHostServiceImpl.ts @@ -0,0 +1,478 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { + IJsonRpcNotification, + IJsonRpcRequest, + JsonRpcMessage, + isJsonRpcErrorResponse, + isJsonRpcSuccessResponse, +} from '../../../../base/common/jsonRpcProtocol.js'; +import { Disposable, IDisposable } from '../../../../base/common/lifecycle.js'; +import { ResourceMap } from '../../../../base/common/map.js'; +import { equals } from '../../../../base/common/objects.js'; +import { IObservable, observableValue } from '../../../../base/common/observable.js'; +import { URI } from '../../../../base/common/uri.js'; +import { generateUuid } from '../../../../base/common/uuid.js'; +import { IMcpServerDefinition } from '../../../agentPlugins/common/pluginParsers.js'; +import { ILogService, ILogger } from '../../../log/common/log.js'; +import { McpServerType } from '../../../mcp/common/mcpPlatformTypes.js'; +import { + IMcpHostService, + IMcpHostUpstreamDelegate, + IMcpServerHandle, + IMcpUiToolMeta, +} from '../../common/mcpHost/mcpHostService.js'; +import { buildMcpServerUri } from '../../common/state/mcpServerUri.js'; +import { JsonRpcErrorCodes } from '../../common/state/protocol/errors.js'; +import { McpMethodCallParams, McpMethodCallResult, McpNotificationParams } from '../../common/state/protocol/commands.js'; +import { + AhpMcpUiHostCapabilities, + McpServerStatus, + McpServerStatusKind, + McpServerSummary, +} from '../../common/state/protocol/state.js'; +import { ProtocolError } from '../../common/state/sessionProtocol.js'; +import { AgentHostStateManager } from '../agentHostStateManager.js'; +import { McpHttpUpstream } from './mcpHttpUpstream.js'; +import { McpAppsInitializeInjector } from './mcpInitializeInjector.js'; +import { McpStdioUpstream } from './mcpStdioUpstream.js'; +import { IMcpProxy, IMcpProxyFactory, IMcpProxyOptions } from './mcpProxy.js'; +import type { IUpstreamRequestOutcome } from './mcpProxyRoute.js'; +import { IMcpUpstream } from './mcpUpstream.js'; + +/** Internal state of an {@link Entry} — disposed entries refuse new mutations. */ +const enum EntryLifecycle { + Live = 0, + Disposed = 1, +} + +/** + * Per-`(session, server)` entry. Owns the upstream + proxy and exposes + * the {@link IMcpServerHandle} surface back to {@link McpHostServiceImpl}. + */ +class Entry extends Disposable implements IMcpServerHandle { + + private readonly _summary; + public readonly summary: IObservable; + + private readonly _endpoint; + public readonly endpoint: IObservable; + + private _proxy: IMcpProxy | undefined; + private _lifecycle: EntryLifecycle = EntryLifecycle.Live; + + constructor( + public readonly session: URI, + public readonly definition: IMcpServerDefinition, + public readonly resource: URI, + initialStatus: McpServerStatus, + private readonly _logService: ILogService, + ) { + super(); + this._summary = observableValue(this, { + resource: resource.toString(), + label: definition.name, + status: initialStatus, + }); + this.summary = this._summary; + this._endpoint = observableValue(this, undefined); + this.endpoint = this._endpoint; + } + + public get isDisposed(): boolean { + return this._lifecycle === EntryLifecycle.Disposed; + } + + public override dispose(): void { + this._lifecycle = EntryLifecycle.Disposed; + super.dispose(); + } + + public registerUpstream(upstream: IMcpUpstream): void { + this._register(upstream); + } + + public setProxy(proxy: IMcpProxy): void { + this._proxy = this._register(proxy); + this._endpoint.set(proxy.endpoint, undefined); + this._logService.info(`[McpHostService] proxy ready for '${this.resource.toString()}' → ${proxy.endpoint?.toString() ?? ''}`); + } + + public setStatus(status: McpServerStatus): void { + const current = this._summary.get(); + if (current.status === status) { + return; + } + this._summary.set({ ...current, status }, undefined); + } + + public async authenticate(resource: string, token: string): Promise { + if (!this._proxy) { + return false; + } + return this._proxy.authenticate(resource, token); + } + + public async callMethod(params: McpMethodCallParams): Promise { + const proxy = this._proxy; + if (!proxy) { + throw new ProtocolError(JsonRpcErrorCodes.InternalError, 'MCP server is not yet ready'); + } + const request: IJsonRpcRequest = { + jsonrpc: '2.0', + id: generateUuid(), + method: params.method, + params: params.params, + }; + const response: JsonRpcMessage | undefined = await proxy.sendClientMessage(request); + if (!response) { + throw new ProtocolError(JsonRpcErrorCodes.InternalError, 'Upstream MCP server returned no response'); + } + if (isJsonRpcSuccessResponse(response)) { + return { result: response.result }; + } + if (isJsonRpcErrorResponse(response)) { + throw new ProtocolError(response.error.code, response.error.message, response.error.data); + } + throw new ProtocolError(JsonRpcErrorCodes.InternalError, 'Upstream MCP server returned an invalid JSON-RPC message'); + } + + public notify(params: McpNotificationParams): void { + const proxy = this._proxy; + if (!proxy) { + this._logService.warn(`[McpHostService] notify dropped for '${this.resource.toString()}': proxy not ready`); + return; + } + const notification: IJsonRpcNotification = { + jsonrpc: '2.0', + method: params.method, + params: params.params, + }; + void proxy.sendClientMessage(notification); + } + + public getToolUiMeta(toolName: string): IMcpUiToolMeta | undefined { + return this._proxy?.getToolUiMeta(toolName); + } + + public getUiHostCapabilities(): AhpMcpUiHostCapabilities { + return this._proxy?.getUiHostCapabilities() ?? {}; + } +} + +/** + * Real {@link IMcpHostService} implementation backed by + * {@link IMcpProxyFactory}. + * + * Holds one {@link Entry} per registered `(session, server)` pair. + * Entries are keyed by their `mcp:/...` resource URI in a + * {@link ResourceMap}; that mapping is the ground truth for + * {@link getServer}, {@link callMethod}, {@link notify}, and + * {@link setSessionServers} diffing. + */ +export class McpHostServiceImpl extends Disposable implements IMcpHostService { + + public readonly _serviceBrand: undefined; + + private _upstreamDelegate: IMcpHostUpstreamDelegate | undefined; + + /** All live entries, keyed by resource URI. */ + private readonly _entries = new ResourceMap(); + + /** Sessions → set of resource URI strings registered for that session. */ + private readonly _bySession = new ResourceMap>(); + + constructor( + private readonly _stateManager: AgentHostStateManager, + private readonly _proxyFactory: IMcpProxyFactory, + @ILogService private readonly _logService: ILogService, + ) { + super(); + } + + public override dispose(): void { + for (const entry of this._entries.values()) { + entry.dispose(); + } + this._entries.clear(); + this._bySession.clear(); + super.dispose(); + } + + public setSessionServers(session: URI, servers: readonly IMcpServerDefinition[]): readonly IMcpServerHandle[] { + const desired = new Map(); + for (const def of servers) { + desired.set(def.name, def); + } + + const previousResources = this._bySession.get(session) ?? new Set(); + const nextResources = new Set(); + const result: IMcpServerHandle[] = []; + + // Compute: removed = previous \ desired (by serverId), reconfigured = previous & desired with different config. + // We compare by serverId derived from the existing entry's definition. + const previousByServerId = new Map(); + for (const resourceString of previousResources) { + const entry = this._entries.get(URI.parse(resourceString)); + if (entry) { + previousByServerId.set(entry.definition.name, entry); + } + } + + // Process additions, reconfigurations, and unchanged. + for (const def of servers) { + const existing = previousByServerId.get(def.name); + if (existing && equals(existing.definition.configuration, def.configuration)) { + // Unchanged — reuse. + nextResources.add(existing.resource.toString()); + result.push(existing); + continue; + } + if (existing) { + // Reconfigured — remove the old entry first. + this._removeEntry(session, existing); + } + // Added or reconfigured — mint a fresh entry. + const entry = this._addEntry(session, def); + nextResources.add(entry.resource.toString()); + result.push(entry); + } + + // Process removals (servers in previous but not in desired). + for (const [serverId, entry] of previousByServerId) { + if (!desired.has(serverId) && !nextResources.has(entry.resource.toString())) { + this._removeEntry(session, entry); + } + } + + if (nextResources.size === 0) { + this._bySession.delete(session); + } else { + this._bySession.set(session, nextResources); + } + + return result; + } + + public getServer(resource: URI): IMcpServerHandle | undefined { + return this._entries.get(resource); + } + + public getServerSummaries(session: URI): readonly McpServerSummary[] { + const resources = this._bySession.get(session); + if (!resources) { + return []; + } + const result: McpServerSummary[] = []; + for (const resourceString of resources) { + const entry = this._entries.get(URI.parse(resourceString)); + if (entry) { + result.push(entry.summary.get()); + } + } + return result; + } + + public async callMethod(params: McpMethodCallParams): Promise { + const resource = URI.parse(params.server); + const entry = this._entries.get(resource); + if (!entry) { + throw new ProtocolError(JsonRpcErrorCodes.InvalidParams, `Unknown MCP server: ${params.server}`); + } + return entry.callMethod(params); + } + + public notify(params: McpNotificationParams): void { + const resource = URI.parse(params.server); + const entry = this._entries.get(resource); + if (!entry) { + this._logService.warn(`[McpHostService] notify dropped for unknown server '${params.server}'`); + return; + } + entry.notify(params); + } + + public setUpstreamDelegate(delegate: IMcpHostUpstreamDelegate): IDisposable { + if (this._upstreamDelegate) { + this._logService.warn('[McpHostService] setUpstreamDelegate replacing existing delegate'); + } + this._upstreamDelegate = delegate; + return { + dispose: () => { + if (this._upstreamDelegate === delegate) { + this._upstreamDelegate = undefined; + } + }, + }; + } + + // ---- Internals --------------------------------------------------------- + + private _addEntry(session: URI, def: IMcpServerDefinition): Entry { + const resource = buildMcpServerUri(session, def.name); + const initialStatus: McpServerStatus = { kind: McpServerStatusKind.Starting }; + const entry = new Entry(session, def, resource, initialStatus, this._logService); + this._entries.set(resource, entry); + + if (this._hasSessionState(session)) { + // Dispatch `mcp/serverAdded` synchronously BEFORE kicking off the + // async proxy creation so the AHP client sees the `Starting` server + // immediately. + this._stateManager.createMcpServer(session.toString(), { + resource: resource.toString(), + label: def.name, + status: initialStatus, + }); + } + + // Async proxy creation — don't await; the caller of + // setSessionServers should not block on transport startup. + void this._createProxyFor(entry); + + return entry; + } + + private _removeEntry(session: URI, entry: Entry): void { + this._entries.delete(entry.resource); + const resources = this._bySession.get(session); + if (resources) { + resources.delete(entry.resource.toString()); + } + // Dispatch `mcp/serverRemoved` BEFORE disposing — entry.dispose() is silent. + if (this._hasSessionState(session)) { + this._stateManager.removeMcpServer(session.toString(), entry.resource.toString()); + } + entry.dispose(); + } + + private async _createProxyFor(entry: Entry): Promise { + let upstream: IMcpUpstream; + try { + upstream = this._createUpstream(entry.definition, this._logService); + } catch (err) { + this._reportProxyCreateError(entry, err); + return; + } + if (entry.isDisposed) { + upstream.dispose(); + return; + } + entry.registerUpstream(upstream); + + const options: IMcpProxyOptions = { + resource: entry.resource, + upstream, + // Always inject the MCP Apps capability when the SDK initializes + // the upstream. Servers that don't speak Apps simply ignore the + // extension entry. Per-tool-call gating now happens via + // `_meta.uiHostCapabilities` on tool call states; the AHP host + // remains a transparent forwarder for `mcpMethodCall` / + // `mcpNotification` traffic. + initializeInjector: new McpAppsInitializeInjector(), + onUpstreamRequest: (method, params) => this._onUpstreamRequest(entry, method, params), + onUpstreamNotification: (method, params) => this._onUpstreamNotification(entry, method, params), + onAuthRequired: status => this._onProxyStatusChange(entry, status), + onStateChange: status => this._onProxyStatusChange(entry, status), + logger: this._logService, + }; + + let proxy: IMcpProxy; + try { + proxy = await this._proxyFactory.create(options); + } catch (err) { + this._reportProxyCreateError(entry, err); + return; + } + + // The entry could have been disposed while we were awaiting the + // proxy. If so, dispose the proxy we just created and bail. + if (entry.isDisposed) { + proxy.dispose(); + return; + } + + entry.setProxy(proxy); + + // Drive initial discovery. Status emissions from the proxy's + // autorun handle the rest. + try { + await upstream.start(); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + this._logService.warn(`[McpHostService] upstream.start() failed for '${entry.resource.toString()}': ${message}`); + } + } + + /** + * Override seam for tests — swaps the real stdio/HTTP upstream + * factories for a stub. Production callers go through the default + * implementation below. + */ + protected _createUpstream(def: IMcpServerDefinition, logger: ILogger): IMcpUpstream { + const config = def.configuration; + if (config.type === McpServerType.LOCAL) { + return new McpStdioUpstream({ config, logger }); + } + return new McpHttpUpstream({ config, logger }); + } + + private async _onUpstreamRequest(entry: Entry, method: string, params: unknown): Promise { + const delegate = this._upstreamDelegate; + if (!delegate) { + return { + error: { + code: JsonRpcErrorCodes.MethodNotFound, + message: `No AHP client is listening for upstream MCP requests on '${entry.resource.toString()}'`, + }, + }; + } + try { + const response = await delegate.handleUpstreamRequest({ server: entry.resource, method, params }); + return response.error ? { error: response.error } : { result: response.result }; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + this._logService.warn(`[McpHostService] upstream request handler threw for method '${method}' on '${entry.resource.toString()}': ${message}`); + return { error: { code: JsonRpcErrorCodes.InternalError, message } }; + } + } + + private _onUpstreamNotification(entry: Entry, method: string, params: unknown): void { + const delegate = this._upstreamDelegate; + if (!delegate) { + return; + } + try { + delegate.handleUpstreamNotification({ server: entry.resource, method, params }); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + this._logService.warn(`[McpHostService] upstream notification handler threw for method '${method}' on '${entry.resource.toString()}': ${message}`); + } + } + + private _onProxyStatusChange(entry: Entry, status: McpServerStatus): void { + entry.setStatus(status); + if (this._hasSessionState(entry.session)) { + this._stateManager.setMcpServerStatus(entry.session.toString(), entry.resource.toString(), status); + } + } + + private _reportProxyCreateError(entry: Entry, err: unknown): void { + const message = err instanceof Error ? err.message : String(err); + this._logService.warn(`[McpHostService] failed to create proxy for '${entry.resource.toString()}': ${message}`); + const status: McpServerStatus = { + kind: McpServerStatusKind.Error, + error: { errorType: 'proxyCreateFailed', message }, + }; + entry.setStatus(status); + if (this._hasSessionState(entry.session)) { + this._stateManager.setMcpServerStatus(entry.session.toString(), entry.resource.toString(), status); + } + } + + private _hasSessionState(session: URI): boolean { + return this._stateManager.getSessionState(session.toString()) !== undefined; + } +} diff --git a/src/vs/platform/agentHost/node/mcpHost/mcpHttpUpstream.ts b/src/vs/platform/agentHost/node/mcpHost/mcpHttpUpstream.ts new file mode 100644 index 0000000000000..cfc75347ef2b8 --- /dev/null +++ b/src/vs/platform/agentHost/node/mcpHost/mcpHttpUpstream.ts @@ -0,0 +1,478 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { CancellationError } from '../../../../base/common/errors.js'; +import { Emitter, type Event } from '../../../../base/common/event.js'; +import { Disposable } from '../../../../base/common/lifecycle.js'; +import { observableValue, type IObservable } from '../../../../base/common/observable.js'; +import type { JsonRpcMessage } from '../../../../base/common/jsonRpcProtocol.js'; +import { SSEParser } from '../../../../base/common/sseParser.js'; +import { generateUuid } from '../../../../base/common/uuid.js'; +import type { ILogger } from '../../../log/common/log.js'; +import type { IMcpRemoteServerConfiguration } from '../../../mcp/common/mcpPlatformTypes.js'; +import { + McpServerStatusKind, + type McpServerStatus, + type ProtectedResourceMetadata, +} from '../../common/state/protocol/state.js'; +import { buildAuthRequiredStatus, parseWwwAuthenticate } from './mcpAuthChallengeParser.js'; +import type { IMcpUpstream, IMcpUpstreamCapabilities } from './mcpUpstream.js'; + +/** HTTP-fetcher signature. Test seam — defaults to global `fetch`. */ +export type HttpFetch = (url: string, init: { + method: 'GET' | 'POST'; + headers: Record; + body?: string; + signal?: AbortSignal; +}) => Promise; + +export interface IHttpResponse { + status: number; + headers: { get(name: string): string | null }; + text(): Promise; + /** + * Optional response body stream. Required for `text/event-stream` + * responses. When `undefined`/`null`, the upstream falls back to + * {@link text} and skips SSE handling. + */ + body?: ReadableStream | null; +} + +export interface IMcpHttpUpstreamOptions { + readonly config: IMcpRemoteServerConfiguration; + readonly logger: ILogger; + /** Test seam — defaults to global `fetch`. */ + readonly fetch?: HttpFetch; +} + +const PROBE_BODY = JSON.stringify({ + jsonrpc: '2.0', + id: 0, + method: 'initialize', + params: { + protocolVersion: '2024-11-05', + capabilities: {}, + clientInfo: { name: 'ahp-proxy', version: '0.0.0' }, + }, +}); + +/** + * MCP upstream backed by a remote HTTP endpoint. The class performs a + * discovery handshake on {@link start} by sending a method-level + * `initialize` probe to the configured URL: the result is discarded + * but the HTTP status drives the state machine (2xx → Ready, 401/403 → + * AuthRequired, otherwise Error). + * + * Supports the streamable HTTP transport: a POST may receive either a + * one-shot `application/json` reply or an SSE (`text/event-stream`) + * stream that delivers one or more JSON-RPC messages. Both are routed + * through {@link IMcpUpstream.onMessage}. + * + * Caveats: + * - The probe may briefly initialize the upstream server before the + * SDK does. For session-bearing servers this is still acceptable + * per MCP, but it does mean the handshake is observable upstream. + * - A long-lived `GET` SSE stream for purely server-initiated + * messages (the second half of the streamable HTTP spec) is not + * yet plumbed; only POST-response streams are consumed. + */ +export class McpHttpUpstream extends Disposable implements IMcpUpstream { + + private readonly _status = observableValue(this, { kind: McpServerStatusKind.Stopped }); + public readonly status: IObservable = this._status; + + private readonly _onMessage = this._register(new Emitter()); + public readonly onMessage: Event = this._onMessage.event; + + private readonly _upstreamCapabilities = observableValue(this, undefined); + public readonly upstreamCapabilities: IObservable = this._upstreamCapabilities; + + private readonly _config: IMcpRemoteServerConfiguration; + private readonly _logger: ILogger; + private readonly _fetch: HttpFetch; + + private _bearerToken: string | undefined; + private _hadPriorToken = false; + private readonly _pendingRequests = new Set(); + private _disposed = false; + private _startInFlight: Promise | undefined; + + constructor(options: IMcpHttpUpstreamOptions) { + super(); + this._config = options.config; + this._logger = options.logger; + this._fetch = options.fetch ?? ((url, init) => globalThis.fetch(url, init) as Promise); + } + + public async start(): Promise { + if (this._disposed) { + return this._status.get(); + } + const current = this._status.get(); + if (current.kind === McpServerStatusKind.Ready) { + return current; + } + if (this._startInFlight) { + return this._startInFlight; + } + const promise = this._doStart(); + this._startInFlight = promise; + try { + return await promise; + } finally { + this._startInFlight = undefined; + } + } + + private async _doStart(): Promise { + this._status.set({ kind: McpServerStatusKind.Starting }, undefined); + + const ac = new AbortController(); + + return this._trackRequest(ac, async () => { + let response: IHttpResponse; + try { + response = await this._traceFetch('probe', this._config.url, { + method: 'POST', + headers: this._buildHeaders(), + body: PROBE_BODY, + signal: ac.signal, + }); + } catch (err) { + if (this._disposed || ac.signal.aborted) { + return this._status.get(); + } + const message = err instanceof Error ? err.message : String(err); + const status: McpServerStatus = { + kind: McpServerStatusKind.Error, + error: { errorType: 'httpError', message }, + }; + this._status.set(status, undefined); + return status; + } + + if (this._disposed || ac.signal.aborted) { + return this._status.get(); + } + + if (response.status >= 200 && response.status < 300) { + this._status.set({ kind: McpServerStatusKind.Ready }, undefined); + // Drain or close the probe response body. The probe payload is + // not consumed (the SDK runs its own initialize), but leaving an + // SSE stream open would hold the underlying connection. + response.body?.cancel().catch(() => { /* best-effort */ }); + return this._status.get(); + } + + if (response.status === 401 || response.status === 403) { + const status = await this._buildAuthRequired(response, response.status); + if (this._disposed) { + return this._status.get(); + } + this._status.set(status, undefined); + return status; + } + + let bodyText = ''; + try { + bodyText = await response.text(); + } catch { + // best-effort capture only + } + const message = `HTTP ${response.status}${bodyText ? `: ${truncate(bodyText, 256)}` : ''}`; + const errStatus: McpServerStatus = { + kind: McpServerStatusKind.Error, + error: { errorType: 'httpError', message }, + }; + this._status.set(errStatus, undefined); + return errStatus; + }); + } + + private _trackRequest(ac: AbortController, work: () => Promise): Promise { + this._pendingRequests.add(ac); + return work().finally(() => this._pendingRequests.delete(ac)); + } + + /** + * Wraps {@link _fetch} with trace-level logging. Each call produces a + * short correlation id so the request and response lines pair up in + * logs. Bodies are truncated; bearer tokens and configured headers + * are redacted. + */ + private async _traceFetch(label: string, url: string, init: { + readonly method: 'GET' | 'POST'; + readonly headers: Record; + readonly body?: string; + readonly signal?: AbortSignal; + }): Promise { + const id = generateUuid().slice(0, 8); + const bodyPreview = init.body ? `, body=${truncate(init.body, 256)}` : ''; + this._logger.trace(`McpHttpUpstream[${label}] -> ${id}: ${init.method} ${url}${bodyPreview}`); + try { + const response = await this._fetch(url, init); + const contentType = response.headers.get('content-type') ?? ''; + this._logger.trace(`McpHttpUpstream[${label}] <- ${id}: HTTP ${response.status}${contentType ? ` content-type=${contentType}` : ''}`); + return response; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + this._logger.trace(`McpHttpUpstream[${label}] <- ${id}: error: ${message}`); + throw err; + } + } + + public async send(message: JsonRpcMessage): Promise { + const current = this._status.get(); + if (current.kind !== McpServerStatusKind.Ready) { + throw new Error(`McpHttpUpstream: cannot send while in state '${current.kind}'`); + } + + const ac = new AbortController(); + + return this._trackRequest(ac, async () => { + let response: IHttpResponse; + try { + response = await this._traceFetch('send', this._config.url, { + method: 'POST', + headers: this._buildHeaders(), + body: JSON.stringify(message), + signal: ac.signal, + }); + } catch (err) { + if (ac.signal.aborted) { + throw new CancellationError(); + } + throw err; + } + + if (response.status === 401 || response.status === 403) { + const authStatus = await this._buildAuthRequired(response, response.status); + if (this._disposed) { + // Race: a concurrent dispose dropped us into Stopped — leave alone. + throw new CancellationError(); + } + this._status.set(authStatus, undefined); + throw new Error('McpHttpUpstream: upstream is now AuthRequired'); + } + + await this._consumeResponseBody(response, ac, 'send'); + }); + } + + /** + * Reads a POST response body and emits any contained JSON-RPC + * messages via {@link onMessage}. Branches on `Content-Type`: + * - `text/event-stream` → stream-parsed via {@link SSEParser}; + * each `data` event whose payload is JSON-RPC fires `onMessage`. + * Resolves once the stream closes or the upstream is disposed. + * - anything else → reads the body as a one-shot JSON message. + */ + private async _consumeResponseBody(response: IHttpResponse, ac: AbortController, label: string): Promise { + const contentType = (response.headers.get('content-type') ?? '').toLowerCase(); + if (contentType.includes('text/event-stream')) { + await this._consumeSseStream(response, ac, label); + return; + } + + const text = await response.text(); + if (text) { + this._logger.trace(`McpHttpUpstream[${label}] body: ${truncate(text, 256)}`); + } + if (!text) { + return; + } + let parsed: JsonRpcMessage; + try { + parsed = JSON.parse(text) as JsonRpcMessage; + } catch (err) { + const errMessage = err instanceof Error ? err.message : String(err); + this._logger.error(`McpHttpUpstream: failed to parse response body: ${errMessage}`); + return; + } + this._onMessage.fire(parsed); + } + + /** + * Streams an SSE response body through {@link SSEParser}. Each + * `message` event whose `data` parses as JSON-RPC is emitted via + * {@link onMessage}. Custom event types are logged and ignored. + * Resolves when the stream closes (server side) or the upstream is + * disposed/aborted. + */ + private async _consumeSseStream(response: IHttpResponse, ac: AbortController, label: string): Promise { + const body = response.body; + if (!body) { + this._logger.warn(`McpHttpUpstream[${label}]: text/event-stream response has no body`); + return; + } + const parser = new SSEParser(event => { + // Per the MCP streamable HTTP spec, JSON-RPC payloads are carried + // on the default `message` event. Other event names are reserved + // for future protocol extensions — trace and ignore. + if (event.type && event.type !== 'message') { + this._logger.trace(`McpHttpUpstream[${label}] sse: ignoring event type=${event.type}`); + return; + } + if (!event.data) { + return; + } + this._logger.trace(`McpHttpUpstream[${label}] sse: ${truncate(event.data, 256)}`); + let parsed: JsonRpcMessage; + try { + parsed = JSON.parse(event.data) as JsonRpcMessage; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + this._logger.error(`McpHttpUpstream: failed to parse SSE event data: ${message}`); + return; + } + this._onMessage.fire(parsed); + }); + + const reader = body.getReader(); + try { + for (; ;) { + if (ac.signal.aborted || this._disposed) { + return; + } + const { done, value } = await reader.read(); + if (done) { + return; + } + if (value && value.length > 0) { + parser.feed(value); + } + } + } catch (err) { + if (ac.signal.aborted || this._disposed) { + return; + } + const message = err instanceof Error ? err.message : String(err); + this._logger.warn(`McpHttpUpstream[${label}]: SSE read error: ${message}`); + } finally { + try { + await reader.cancel(); + } catch { + // best-effort + } + } + } + + public setBearerToken(token: string | undefined): void { + this._bearerToken = token; + this._hadPriorToken = !!token; + } + + public setUpstreamCapabilities(caps: IMcpUpstreamCapabilities | undefined): void { + this._upstreamCapabilities.set(caps, undefined); + } + + private _buildHeaders(): Record { + const headers: Record = { + 'Content-Type': 'application/json', + 'Accept': 'application/json, text/event-stream', + }; + if (this._config.headers) { + for (const [key, value] of Object.entries(this._config.headers)) { + headers[key] = value; + } + } + if (this._bearerToken) { + headers['Authorization'] = `Bearer ${this._bearerToken}`; + } + return headers; + } + + private async _buildAuthRequired(response: IHttpResponse, httpStatus: 401 | 403): Promise { + const challenge = parseWwwAuthenticate(response.headers.get('WWW-Authenticate') ?? undefined); + let resource: ProtectedResourceMetadata | undefined; + if (challenge.resourceMetadataUrl) { + if (this._isSafeMetadataUrl(challenge.resourceMetadataUrl)) { + resource = await this._fetchResourceMetadata(challenge.resourceMetadataUrl); + } else { + this._logger.warn(`McpHttpUpstream: ignoring resource_metadata URL '${challenge.resourceMetadataUrl}' — not at the same origin/scheme as ${this._config.url}`); + } + } + if (!resource) { + this._logger.warn(`McpHttpUpstream: server returned ${httpStatus} without usable resource_metadata; synthesizing minimal metadata for ${this._config.url}`); + resource = { resource: this._config.url }; + } + return buildAuthRequiredStatus({ + httpStatus, + challenge, + resource, + hadPriorToken: this._hadPriorToken, + }); + } + + /** + * RFC 9728 expects the protected-resource metadata to live at the same + * authority as the protected resource. We enforce that strictly to + * prevent a malicious or misconfigured server from steering us into + * fetching arbitrary URLs (e.g. cloud-metadata endpoints, intranet + * services, `file:`/`javascript:`/`data:` URIs). + */ + private _isSafeMetadataUrl(metadataUrl: string): boolean { + let parsed: URL; + let configUrl: URL; + try { + parsed = new URL(metadataUrl, this._config.url); + configUrl = new URL(this._config.url); + } catch { + return false; + } + if (parsed.protocol !== 'https:' && parsed.protocol !== 'http:') { + return false; + } + // `host` includes hostname + port, so this enforces same authority. + return parsed.protocol === configUrl.protocol && parsed.host === configUrl.host; + } + + private async _fetchResourceMetadata(url: string): Promise { + const ac = new AbortController(); + return this._trackRequest(ac, async () => { + try { + const response = await this._traceFetch('resource_metadata', url, { + method: 'GET', + headers: { 'Accept': 'application/json' }, + signal: ac.signal, + }); + if (this._disposed || ac.signal.aborted) { + return undefined; + } + if (response.status < 200 || response.status >= 300) { + this._logger.warn(`McpHttpUpstream: resource_metadata fetch returned HTTP ${response.status}`); + return undefined; + } + const text = await response.text(); + return JSON.parse(text) as ProtectedResourceMetadata; + } catch (err) { + if (ac.signal.aborted) { + return undefined; + } + const message = err instanceof Error ? err.message : String(err); + this._logger.warn(`McpHttpUpstream: failed to fetch or parse resource_metadata: ${message}`); + return undefined; + } + }); + } + + public override dispose(): void { + if (this._disposed) { + return; + } + this._disposed = true; + for (const ac of this._pendingRequests) { + ac.abort(); + } + this._pendingRequests.clear(); + this._upstreamCapabilities.set(undefined, undefined); + this._status.set({ kind: McpServerStatusKind.Stopped }, undefined); + super.dispose(); + } +} + +function truncate(s: string, max: number): string { + return s.length <= max ? s : s.slice(0, max) + '…'; +} diff --git a/src/vs/platform/agentHost/node/mcpHost/mcpInitializeInjector.ts b/src/vs/platform/agentHost/node/mcpHost/mcpInitializeInjector.ts new file mode 100644 index 0000000000000..393ae44f94e93 --- /dev/null +++ b/src/vs/platform/agentHost/node/mcpHost/mcpInitializeInjector.ts @@ -0,0 +1,83 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import type { IJsonRpcRequest } from '../../../../base/common/jsonRpcProtocol.js'; + +/** + * MIME type the MCP Apps extension advertises support for. + */ +const MCP_APPS_MIME_TYPE = 'text/html;profile=mcp-app'; + +/** + * Capability key for the MCP Apps extension under + * `params.capabilities.extensions`. + */ +const MCP_APPS_EXTENSION_KEY = 'io.modelcontextprotocol/ui'; + +/** + * Mutates an MCP `initialize` request payload from the upstream SDK + * to advertise additional client capabilities (e.g. MCP Apps) when + * the AHP client supports them. + * + * The proxy is otherwise blind. This is the only inbound (client → + * upstream) rewrite it performs. + */ +export interface IInitializeInjector { + /** + * Apply mutations to the `params` object of an `initialize` request. + * Idempotent: re-running on already-injected params is a no-op. + * Caller-provided client capabilities are PRESERVED — never + * overwritten — only merged. Unknown extension keys remain. + */ + inject(request: IJsonRpcRequest): IJsonRpcRequest; +} + +interface IMutableInitializeParams { + capabilities?: IMutableCapabilities; + [k: string]: unknown; +} + +interface IMutableCapabilities { + extensions?: Record; + [k: string]: unknown; +} + +/** + * Injector that adds the MCP Apps extension capability under + * `params.capabilities.extensions['io.modelcontextprotocol/ui']`. + * Other fields on `capabilities` (sampling, elicitation, roots, + * tasks, etc.) are preserved. + */ +export class McpAppsInitializeInjector implements IInitializeInjector { + inject(request: IJsonRpcRequest): IJsonRpcRequest { + const originalParams = (request.params ?? {}) as IMutableInitializeParams; + const originalCaps = (originalParams.capabilities ?? {}) as IMutableCapabilities; + const originalExtensions = (originalCaps.extensions ?? {}) as Record; + + const newExtensions: Record = { + ...originalExtensions, + [MCP_APPS_EXTENSION_KEY]: { + mimeTypes: [MCP_APPS_MIME_TYPE], + }, + }; + + const newCapabilities: IMutableCapabilities = { + ...originalCaps, + extensions: newExtensions, + }; + + const newParams: IMutableInitializeParams = { + ...originalParams, + capabilities: newCapabilities, + }; + + return { + jsonrpc: request.jsonrpc, + id: request.id, + method: request.method, + params: newParams, + }; + } +} diff --git a/src/vs/platform/agentHost/node/mcpHost/mcpProxy.ts b/src/vs/platform/agentHost/node/mcpHost/mcpProxy.ts new file mode 100644 index 0000000000000..d674e5a20dd83 --- /dev/null +++ b/src/vs/platform/agentHost/node/mcpHost/mcpProxy.ts @@ -0,0 +1,185 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import type { JsonRpcMessage } from '../../../../base/common/jsonRpcProtocol.js'; +import { Disposable, type IDisposable } from '../../../../base/common/lifecycle.js'; +import { autorun } from '../../../../base/common/observable.js'; +import type { URI } from '../../../../base/common/uri.js'; +import { createDecorator } from '../../../instantiation/common/instantiation.js'; +import type { ILogger } from '../../../log/common/log.js'; +import { + McpServerStatusKind, + type AhpMcpUiHostCapabilities, + type McpServerStatus, + type McpServerStatusAuthRequired, +} from '../../common/state/protocol/state.js'; +import type { IInitializeInjector } from './mcpInitializeInjector.js'; +import { McpProxyHttpListener, type IRouteRegistration } from './mcpProxyHttpListener.js'; +import { McpProxyRoute, type IMcpUiToolMeta, type IUpstreamRequestOutcome } from './mcpProxyRoute.js'; +import type { IMcpUpstream } from './mcpUpstream.js'; + +export type { IMcpUiToolMeta } from './mcpProxyRoute.js'; + +export interface IMcpProxyOptions { + /** mcp:// URI; matches the AHP McpServerSummary.resource. */ + readonly resource: URI; + readonly upstream: IMcpUpstream; + readonly initializeInjector?: IInitializeInjector; + /** + * Tap fired when the upstream emits a JSON-RPC **request**. The + * implementation forwards the call to the AHP client as a reverse + * `mcpMethodCall` and resolves with the client's response. + */ + readonly onUpstreamRequest: (method: string, params: unknown) => Promise; + /** + * Tap fired when the upstream emits a JSON-RPC **notification**. + * Fire-and-forget. + */ + readonly onUpstreamNotification: (method: string, params: unknown) => void; + readonly onAuthRequired: (status: McpServerStatusAuthRequired) => void; + readonly onStateChange: (status: McpServerStatus) => void; + readonly logger: ILogger; +} + +export interface IMcpProxy extends IDisposable { + /** mcp:// URI; matches the AHP McpServerSummary.resource. */ + readonly resource: URI; + /** HTTP endpoint the upstream-facing SDK should connect to. */ + readonly endpoint: URI; + /** + * Push a bearer token. Returns true if the upstream accepted the + * token (state transitioned to {@link McpServerStatusKind.Ready}). + * The `resource` argument is the protected-resource identifier the + * caller obtained via the most recent `AuthRequired` status; it is + * cross-checked against the proxy's last challenge. + */ + authenticate(resource: string, token: string): Promise; + /** Forward a message from the AHP client to the upstream. */ + sendClientMessage(message: JsonRpcMessage): Promise; + /** + * Latest `_meta.ui` payload the upstream advertised for `toolName` via + * `tools/list`, captured opportunistically as the proxy sniffs JSON-RPC + * traffic. Returns `undefined` for tools that don't surface an MCP App + * or for tool names that have never been listed. + */ + getToolUiMeta(toolName: string): IMcpUiToolMeta | undefined; + + /** + * The set of MCP App host capabilities the AHP proxy can satisfy for + * a View backed by this upstream server, derived from the upstream's + * `initialize` response. Empty until the SDK's `initialize` + * handshake has completed through the proxy. + */ + getUiHostCapabilities(): AhpMcpUiHostCapabilities; +} + +export const IMcpProxyFactory = createDecorator('mcpProxyFactory'); + +export interface IMcpProxyFactory { + readonly _serviceBrand: undefined; + /** + * Create a proxy AND register its HTTP route. Resolves once the + * shared HTTP listener has bound and the route is reachable. + */ + create(options: IMcpProxyOptions): Promise; +} + +class McpProxy extends Disposable implements IMcpProxy { + + public readonly resource: URI; + public readonly endpoint: URI; + private readonly _route: McpProxyRoute; + private readonly _options: IMcpProxyOptions; + private _lastAuthChallenge: McpServerStatusAuthRequired | undefined; + + constructor(options: IMcpProxyOptions, route: McpProxyRoute, registration: IRouteRegistration) { + super(); + this._options = options; + this.resource = options.resource; + this.endpoint = registration.endpoint; + this._route = this._register(route); + this._register({ dispose: () => registration.dispose() }); + + this._register(autorun(reader => { + const status = options.upstream.status.read(reader); + if (status.kind === McpServerStatusKind.AuthRequired) { + this._lastAuthChallenge = status; + options.onAuthRequired(status); + } else { + options.onStateChange(status); + } + })); + } + + public async authenticate(resource: string, token: string): Promise { + const challenge = this._lastAuthChallenge; + if (challenge && challenge.resource.resource !== resource) { + this._options.logger.warn(`McpProxy: authenticate called with resource '${resource}' but the most recent challenge was for '${challenge.resource.resource}'`); + return false; + } + this._options.upstream.setBearerToken(token); + try { + const status = await this._options.upstream.start(); + return status.kind === McpServerStatusKind.Ready; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + this._options.logger.error(`McpProxy: authenticate failed: ${message}`); + return false; + } + } + + public sendClientMessage(message: JsonRpcMessage): Promise { + return this._route.sendClientMessage(message); + } + + public getToolUiMeta(toolName: string): IMcpUiToolMeta | undefined { + return this._route.getToolUiMeta(toolName); + } + + public getUiHostCapabilities(): AhpMcpUiHostCapabilities { + return this._route.getUiHostCapabilities(); + } +} + +/** + * Default implementation of {@link IMcpProxyFactory}. Owns a single + * shared {@link McpProxyHttpListener}; the listener is bound on first + * `create()` and shut down when the last route is removed. + */ +export class McpProxyFactory extends Disposable implements IMcpProxyFactory { + + public readonly _serviceBrand: undefined; + + private readonly _listener: McpProxyHttpListener; + + constructor(logger: ILogger) { + super(); + this._listener = this._register(new McpProxyHttpListener(logger)); + } + + public async create(options: IMcpProxyOptions): Promise { + const route = new McpProxyRoute({ + upstream: options.upstream, + logger: options.logger, + initializeInjector: options.initializeInjector, + onUpstreamRequest: options.onUpstreamRequest, + onUpstreamNotification: options.onUpstreamNotification, + }); + let registration: IRouteRegistration; + try { + registration = await this._listener.registerRoute(body => route.handleSdkBody(body)); + } catch (err) { + route.dispose(); + throw err; + } + try { + return new McpProxy(options, route, registration); + } catch (err) { + registration.dispose(); + route.dispose(); + throw err; + } + } +} diff --git a/src/vs/platform/agentHost/node/mcpHost/mcpProxyHttpListener.ts b/src/vs/platform/agentHost/node/mcpHost/mcpProxyHttpListener.ts new file mode 100644 index 0000000000000..f3fa6ab1bc0ec --- /dev/null +++ b/src/vs/platform/agentHost/node/mcpHost/mcpProxyHttpListener.ts @@ -0,0 +1,231 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import type { IncomingMessage, Server, ServerResponse } from 'http'; +import type { AddressInfo } from 'net'; +import { Disposable, toDisposable } from '../../../../base/common/lifecycle.js'; +import { URI } from '../../../../base/common/uri.js'; +import { generateUuid } from '../../../../base/common/uuid.js'; +import type { ILogger } from '../../../log/common/log.js'; + +function truncate(s: string, max: number): string { + return s.length <= max ? s : s.slice(0, max) + '…'; +} + +/** + * Handles a single inbound HTTP request body for a registered route. + * Each route owns its own JSON-RPC dispatch. + * + * @param body The request body as a UTF-8 string. + * @returns The response body to send back. The listener writes + * `200 OK` with `Content-Type: application/json` for non-empty bodies + * and `204 No Content` for empty strings. + */ +export type RouteHandler = (body: string) => Promise; + +export interface IRouteRegistration { + /** The URL the upstream-facing SDK should POST to. */ + readonly endpoint: URI; + /** Removes the route. Shuts down the HTTP server when the last route is removed. */ + dispose(): void; +} + +/** Maximum accepted request-body size (1MB). */ +const MAX_REQUEST_BODY_BYTES = 1024 * 1024; + +/** Path prefix every registered route lives under. */ +const ROUTE_PATH_PREFIX = '/mcp/'; + +/** Path suffix every registered route lives under. */ +const ROUTE_PATH_SUFFIX = '/message'; + +/** + * A shared HTTP listener bound on `127.0.0.1` that hosts per-MCP-server + * routes at randomized paths. + * + * The first registration starts the server; disposing the last route + * shuts it down. Loopback-only — never binds to a non-loopback + * interface. The randomized path is the access control mechanism; + * clients must know the URL exactly. URLs are passed only to the local + * SDK process spawned by the proxy. + */ +export class McpProxyHttpListener extends Disposable { + + private _server: Server | undefined; + private _origin: string | undefined; + private readonly _routes = new Map(); + private _disposed = false; + + constructor(private readonly _logger: ILogger) { + super(); + this._register(toDisposable(() => this._shutdown())); + } + + /** + * Register a route. Returns the URI the SDK should connect to and a + * disposable that removes the route. When the last route is removed, + * the server is shut down. + */ + public async registerRoute(handler: RouteHandler): Promise { + if (this._disposed) { + throw new Error('McpProxyHttpListener: cannot register route after dispose'); + } + const origin = await this._ensureServer(); + const routeId = generateUuid(); + this._routes.set(routeId, handler); + const path = `${ROUTE_PATH_PREFIX}${routeId}${ROUTE_PATH_SUFFIX}`; + const endpoint = URI.parse(`${origin}${path}`); + let removed = false; + const dispose = () => { + if (removed) { + return; + } + removed = true; + this._routes.delete(routeId); + if (this._routes.size === 0) { + this._closeServer(); + } + }; + return { endpoint, dispose }; + } + + private async _ensureServer(): Promise { + if (this._server && this._origin) { + return this._origin; + } + const { createServer } = await import('http'); + return new Promise((resolve, reject) => { + const server = createServer((req, res) => this._handleRequest(req, res)); + server.on('error', err => { + this._logger.error(`McpProxyHttpListener: server error: ${err instanceof Error ? err.message : String(err)}`); + if (!this._origin) { + reject(err); + } + }); + server.listen(0, '127.0.0.1', () => { + const address = server.address() as AddressInfo | null; + if (!address || typeof address !== 'object') { + reject(new Error('McpProxyHttpListener: failed to determine listening port')); + return; + } + this._server = server; + this._origin = `http://127.0.0.1:${address.port}`; + resolve(this._origin); + }); + }); + } + + private _handleRequest(req: IncomingMessage, res: ServerResponse): void { + if (req.method !== 'POST') { + res.statusCode = 405; + res.setHeader('Allow', 'POST'); + res.end(); + return; + } + + const url = req.url ?? ''; + if (!url.startsWith(ROUTE_PATH_PREFIX)) { + res.statusCode = 404; + res.end(); + return; + } + + const segments = url.split('/').filter(s => s.length > 0); + // Expect ['mcp', '', 'message']. + if (segments.length < 3 || segments[0] !== 'mcp' || segments[2] !== 'message') { + res.statusCode = 404; + res.end(); + return; + } + const routeId = segments[1]; + const handler = this._routes.get(routeId); + if (!handler) { + res.statusCode = 404; + res.end(); + return; + } + + const declaredLength = Number(req.headers['content-length']); + if (Number.isFinite(declaredLength) && declaredLength > MAX_REQUEST_BODY_BYTES) { + res.statusCode = 413; + res.end(); + return; + } + + const chunks: Buffer[] = []; + let received = 0; + let aborted = false; + req.on('data', (chunk: Buffer) => { + if (aborted) { + return; + } + received += chunk.length; + if (received > MAX_REQUEST_BODY_BYTES) { + aborted = true; + res.statusCode = 413; + res.end(); + req.destroy(); + return; + } + chunks.push(chunk); + }); + req.on('error', err => { + this._logger.warn(`McpProxyHttpListener: request error: ${err instanceof Error ? err.message : String(err)}`); + }); + req.on('end', () => { + if (aborted) { + return; + } + const body = Buffer.concat(chunks).toString('utf8'); + const traceId = generateUuid().slice(0, 8); + this._logger.trace(`McpProxyHttpListener[${routeId.slice(0, 8)}] -> ${traceId}: POST ${url}${body ? `, body=${truncate(body, 256)}` : ''}`); + handler(body).then( + responseBody => { + if (responseBody.length === 0) { + this._logger.trace(`McpProxyHttpListener[${routeId.slice(0, 8)}] <- ${traceId}: 204`); + res.statusCode = 204; + res.end(); + return; + } + this._logger.trace(`McpProxyHttpListener[${routeId.slice(0, 8)}] <- ${traceId}: 200, body=${truncate(responseBody, 256)}`); + res.statusCode = 200; + res.setHeader('Content-Type', 'application/json'); + res.end(responseBody); + }, + err => { + this._logger.trace(`McpProxyHttpListener[${routeId.slice(0, 8)}] <- ${traceId}: 500: ${err instanceof Error ? err.message : String(err)}`); + this._logger.error(`McpProxyHttpListener: route handler error: ${err instanceof Error ? err.message : String(err)}`); + if (!res.headersSent) { + res.statusCode = 500; + } + res.end(); + }, + ); + }); + } + + private _closeServer(): void { + const server = this._server; + if (!server) { + return; + } + this._server = undefined; + this._origin = undefined; + server.close(err => { + if (err) { + this._logger.warn(`McpProxyHttpListener: server.close error: ${err instanceof Error ? err.message : String(err)}`); + } + }); + } + + private _shutdown(): void { + if (this._disposed) { + return; + } + this._disposed = true; + this._routes.clear(); + this._closeServer(); + } +} diff --git a/src/vs/platform/agentHost/node/mcpHost/mcpProxyRoute.ts b/src/vs/platform/agentHost/node/mcpHost/mcpProxyRoute.ts new file mode 100644 index 0000000000000..e59f6d660c1c4 --- /dev/null +++ b/src/vs/platform/agentHost/node/mcpHost/mcpProxyRoute.ts @@ -0,0 +1,460 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { DeferredPromise, disposableTimeout } from '../../../../base/common/async.js'; +import { + isJsonRpcNotification, + isJsonRpcRequest, + isJsonRpcResponse, + isJsonRpcSuccessResponse, + type IJsonRpcErrorResponse, + type IJsonRpcNotification, + type IJsonRpcRequest, + type IJsonRpcSuccessResponse, + type JsonRpcId, + type JsonRpcMessage, +} from '../../../../base/common/jsonRpcProtocol.js'; +import { Disposable } from '../../../../base/common/lifecycle.js'; +import { generateUuid } from '../../../../base/common/uuid.js'; +import type { ILogger } from '../../../log/common/log.js'; +import { MCP } from '../../../mcp/common/modelContextProtocol.js'; +import type { AhpMcpUiHostCapabilities } from '../../common/state/protocol/state.js'; +import type { IInitializeInjector } from './mcpInitializeInjector.js'; +import type { IMcpUpstream, IMcpUpstreamCapabilities } from './mcpUpstream.js'; + +/** Default timeout (ms) for an SDK→upstream request awaiting a response. */ +const DEFAULT_REQUEST_TIMEOUT_MS = 30_000; + +/** JSON-RPC error code used when the request body cannot be parsed. */ +const RPC_PARSE_ERROR = -32700; + +/** JSON-RPC error code used when the upstream times out. */ +const RPC_INTERNAL_ERROR = -32603; + +/** Build a JSON-RPC error response. */ +function jsonRpcError(id: JsonRpcId, code: number, message: string, data?: unknown): IJsonRpcErrorResponse { + const error: IJsonRpcErrorResponse['error'] = { code, message }; + if (data !== undefined) { + error.data = data; + } + return { jsonrpc: '2.0', id, error }; +} + +/** Build a JSON-RPC success response. */ +function jsonRpcSuccess(id: JsonRpcId, result: unknown): IJsonRpcSuccessResponse { + return { jsonrpc: '2.0', id, result }; +} + +/** + * Outcome of an upstream → AHP-client request: either a JSON-RPC `result` + * or a JSON-RPC `error`. + */ +export interface IUpstreamRequestOutcome { + readonly result?: unknown; + readonly error?: { code: number; message: string; data?: unknown }; +} + +/** + * The MCP Apps `_meta.ui` payload carried by a `Tool` definition. + * + * Mirrors the + * [MCP Apps spec](https://github.com/modelcontextprotocol/ext-apps/blob/main/specification/2026-01-26/apps.mdx). + * The proxy captures these payloads from `tools/list` responses so the + * agent can decorate AHP `ToolCallBase._meta` with `ui` / `uiHostCapabilities` + * when the Copilot SDK invokes the matching tool. + */ +export interface IMcpUiToolMeta { + /** `ui://…` URI of the UI resource for rendering the tool result. */ + readonly resourceUri?: string; + /** + * Who can access this tool. Default `["model", "app"]`. `"model"` makes + * the tool visible to the agent; `"app"` makes it callable by the UI + * view from the same MCP server. + */ + readonly visibility?: readonly ('model' | 'app')[]; +} + +/** + * Project the upstream MCP server's `InitializeResult.capabilities` into + * the subset of MCP-Apps host capabilities the AHP proxy is willing to + * satisfy for a View backed by that server. + * + * - `serverTools` / `serverResources` — advertised iff the upstream + * advertised the matching `tools` / `resources` capability. The + * `listChanged` flag mirrors what the upstream said; the proxy + * forwards those notifications transparently when present. + * - `logging` — advertised iff the upstream said it accepts log + * notifications. The proxy forwards `notifications/message` from the + * View and `logging/setLevel` from the View back to the server. + * - `sampling` — never advertised today: even when the upstream + * supports it, the AHP host has no LLM bridge to serve + * `sampling/createMessage` from the View. + * + * Tolerates malformed upstream payloads — bad data simply produces an + * empty capability set rather than throwing. + */ +function deriveUiHostCapabilities(upstream: IMcpUpstreamCapabilities | undefined): AhpMcpUiHostCapabilities { + if (!upstream) { + return {}; + } + try { + const caps = upstream as MCP.ServerCapabilities; + return { + ...(caps.tools && { serverTools: { listChanged: caps.tools.listChanged === true } }), + ...(caps.resources && { serverResources: { listChanged: caps.resources.listChanged === true } }), + ...(caps.logging && { logging: {} }), + }; + } catch { + return {}; + } +} + +export interface IMcpProxyRouteOptions { + readonly upstream: IMcpUpstream; + readonly logger: ILogger; + /** + * Optional: rewrites client→upstream `initialize` requests to add + * extension capabilities. Caller chooses based on the AHP client's + * advertised support. + */ + readonly initializeInjector?: IInitializeInjector; + /** + * Tap fired when the upstream emits a JSON-RPC **request**. The route + * awaits the returned promise and writes the JSON-RPC response back + * to the upstream transport using the original request id. + */ + readonly onUpstreamRequest: (method: string, params: unknown) => Promise; + /** + * Tap fired when the upstream emits a JSON-RPC **notification**. + * Fire-and-forget; no response is expected. + */ + readonly onUpstreamNotification: (method: string, params: unknown) => void; + /** Override request timeout (ms). Defaults to 30s. */ + readonly requestTimeoutMs?: number; +} + +/** + * One per advertised MCP server. Bridges: + * SDK ⟷ HTTP ⟷ McpProxyRoute ⟷ IMcpUpstream + * + * Pass-through is blind in both directions, with three hooks. + * 1. client→upstream `initialize` requests are rewritten via + * {@link IInitializeInjector} when one is configured. + * 2. upstream-originated notifications are tapped and forwarded to the + * AHP client via {@link IMcpProxyRouteOptions.onUpstreamNotification} + * (which routes them out as `mcpNotification`). + * 3. upstream-originated requests are tapped and forwarded to the AHP + * client via {@link IMcpProxyRouteOptions.onUpstreamRequest} (which + * routes them out as `mcpMethodCall`). The route awaits the + * returned outcome and writes the JSON-RPC response back to the + * upstream transport using the original request id. + */ +export class McpProxyRoute extends Disposable { + + /** + * Maps SDK JSON-RPC id (as a string for safe Map keys) → entry. The + * route forwards the SDK's request to the upstream and resolves the + * deferred once the upstream replies. The original {@link JsonRpcId} + * is retained so disposal can complete the deferred with an error + * response carrying the id the SDK actually sent (not `0`). + */ + private readonly _pendingSdkRequests = new Map }>(); + + /** + * Set of SDK JSON-RPC ids (string-keyed) that the route has forwarded + * to the upstream as `initialize` requests. When a matching response + * arrives we extract `result.capabilities` and push it to + * {@link IMcpUpstream.setUpstreamCapabilities}. Cleared on dispose. + */ + private readonly _pendingInitializeIds = new Set(); + + /** + * Set of SDK JSON-RPC ids (string-keyed) that the route has forwarded + * to the upstream as `tools/list` requests. When a matching response + * arrives the route parses `result.tools[]._meta.ui` and refreshes + * {@link _toolUiMeta} so the agent can decorate MCP tool calls with + * the MCP Apps payload. + */ + private readonly _pendingToolsListIds = new Set(); + + /** + * Latest `_meta.ui` payload captured per upstream tool name. Populated + * lazily as `tools/list` responses flow through the proxy (from either + * the SDK or AHP-client `mcpMethodCall` traffic) and exposed via + * {@link getToolUiMeta} for the agent's tool-call decoration path. + */ + private readonly _toolUiMeta = new Map(); + + private readonly _requestTimeoutMs: number; + private _disposed = false; + + constructor(private readonly _options: IMcpProxyRouteOptions) { + super(); + this._requestTimeoutMs = _options.requestTimeoutMs ?? DEFAULT_REQUEST_TIMEOUT_MS; + this._register(this._options.upstream.onMessage(msg => this._onUpstreamMessage(msg))); + } + + /** + * Handle an inbound HTTP body from the SDK. The body is a single + * JSON-RPC message: request, notification, or response (responses + * arrive when the upstream had previously sent the SDK a request). + * + * Returns the response body to send back as HTTP 200, or empty + * string for a notification (HTTP 204). + */ + public async handleSdkBody(body: string): Promise { + let parsed: JsonRpcMessage; + try { + parsed = JSON.parse(body) as JsonRpcMessage; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + this._options.logger.warn(`McpProxyRoute: failed to parse SDK body: ${message}`); + return JSON.stringify(jsonRpcError(0, RPC_PARSE_ERROR, `Parse error: ${message}`)); + } + + if (isJsonRpcRequest(parsed)) { + return this._handleSdkRequest(parsed); + } + if (isJsonRpcNotification(parsed)) { + await this._sendToUpstream(parsed); + return ''; + } + if (isJsonRpcResponse(parsed)) { + // SDK is answering a previously-tapped upstream request. Forward + // it through; the host service tracks the messageId mapping + // separately, so all we need to do is relay. + await this._sendToUpstream(parsed); + return ''; + } + this._options.logger.warn('McpProxyRoute: SDK body is not a valid JSON-RPC message'); + return JSON.stringify(jsonRpcError(0, RPC_PARSE_ERROR, 'Invalid JSON-RPC message')); + } + + /** + * Forward a client-initiated message (from `mcpMethodCall` / + * `mcpNotification`) to the upstream. Notifications return + * immediately; requests await the response. The caller is expected + * to mint the JSON-RPC id when the message is a request (the proxy + * does not rewrite ids here). + */ + public async sendClientMessage(message: JsonRpcMessage): Promise { + if (isJsonRpcRequest(message)) { + // If the caller did not provide an id, mint one. + const id: JsonRpcId = message.id ?? generateUuid(); + const patched: IJsonRpcRequest = { ...message, id }; + // Both SDK→upstream and AHP-client→upstream `tools/list` traffic + // goes through `_dispatchSdkRequest`; tracking the id here keeps + // {@link _toolUiMeta} fresh regardless of which peer asked. + if (patched.method === 'tools/list') { + this._pendingToolsListIds.add(String(patched.id)); + } + return this._dispatchSdkRequest(patched); + } + if (isJsonRpcNotification(message) || isJsonRpcResponse(message)) { + await this._sendToUpstream(message); + return undefined; + } + this._options.logger.warn('McpProxyRoute: sendClientMessage received an invalid JSON-RPC message'); + return undefined; + } + + private async _handleSdkRequest(request: IJsonRpcRequest): Promise { + let outbound = request; + if (request.method === 'initialize' && this._options.initializeInjector) { + outbound = this._options.initializeInjector.inject(request); + } + if (request.method === 'initialize') { + this._pendingInitializeIds.add(String(outbound.id)); + } + if (request.method === 'tools/list') { + this._pendingToolsListIds.add(String(outbound.id)); + } + const reply = await this._dispatchSdkRequest(outbound); + return JSON.stringify(reply); + } + + private async _dispatchSdkRequest(request: IJsonRpcRequest): Promise { + const key = String(request.id); + const deferred = new DeferredPromise(); + this._pendingSdkRequests.set(key, { id: request.id, deferred }); + + const timeoutHandle = disposableTimeout(() => { + this._pendingSdkRequests.delete(key); + this._options.logger.warn(`McpProxyRoute: upstream timed out for id '${key}'`); + deferred.complete(jsonRpcError(request.id, RPC_INTERNAL_ERROR, 'Upstream MCP server did not respond within timeout')); + }, this._requestTimeoutMs); + + try { + await this._sendToUpstream(request); + } catch (err) { + timeoutHandle.dispose(); + this._pendingSdkRequests.delete(key); + const message = err instanceof Error ? err.message : String(err); + this._options.logger.warn(`McpProxyRoute: upstream send failed for id '${key}': ${message}`); + return jsonRpcError(request.id, RPC_INTERNAL_ERROR, message); + } + + try { + return await deferred.p; + } finally { + timeoutHandle.dispose(); + } + } + + private async _sendToUpstream(message: JsonRpcMessage): Promise { + await this._options.upstream.send(message); + } + + private _captureUpstreamCapabilities(result: unknown): void { + if (!result || typeof result !== 'object') { + return; + } + const caps = (result as { capabilities?: unknown }).capabilities; + if (caps && typeof caps === 'object') { + this._options.upstream.setUpstreamCapabilities(caps as IMcpUpstreamCapabilities); + } + } + + /** + * Parse a `tools/list` response body and refresh {@link _toolUiMeta} + * with whatever `_meta.ui` payloads the upstream advertised. Each + * `tools/list` call replaces the entire map so stale entries don't + * linger when the server removes a tool. Tolerates malformed payloads + * — bad data simply yields an empty map rather than throwing. + */ + private _captureToolUiMeta(result: unknown): void { + try { + const { tools } = result as MCP.ListToolsResult; + this._toolUiMeta.clear(); + for (const tool of tools) { + const ui = tool._meta?.ui as IMcpUiToolMeta | undefined; + if (ui && tool.name) { + this._toolUiMeta.set(tool.name, ui); + } + } + } catch (err) { + this._options.logger.warn(`McpProxyRoute: failed to capture tool _meta.ui: ${err instanceof Error ? err.message : String(err)}`); + } + } + + /** + * Look up the most recently captured MCP Apps `_meta.ui` payload for a + * tool advertised by this upstream, or `undefined` if no + * MCP-Apps-enabled tool with this name has been observed. + */ + public getToolUiMeta(toolName: string): IMcpUiToolMeta | undefined { + return this._toolUiMeta.get(toolName); + } + + /** + * The set of MCP App host capabilities the AHP proxy can satisfy for + * a View backed by THIS upstream server, derived from the upstream's + * `initialize` response. + * + * For each capability in {@link AhpMcpUiHostCapabilities} the proxy + * advertises it only when both + * (a) the upstream advertised the corresponding server capability, + * and + * (b) the proxy actually forwards the matching MCP method/notification + * traffic. + * + * `(b)` always holds for `serverTools`, `serverResources`, and + * `logging` — the proxy is a blind passthrough for those. `sampling` + * is intentionally omitted today: even when the upstream supports it, + * the AHP host has no LLM bridge to satisfy `sampling/createMessage` + * calls from the View. + * + * Returns an empty object when no `initialize` has been observed yet + * (the upstream hasn't told us what it supports, so we can't honestly + * advertise anything). + */ + public getUiHostCapabilities(): AhpMcpUiHostCapabilities { + const upstream = this._options.upstream.upstreamCapabilities.get(); + return deriveUiHostCapabilities(upstream); + } + + private _onUpstreamMessage(message: JsonRpcMessage): void { + if (isJsonRpcRequest(message)) { + this._handleUpstreamRequest(message); + return; + } + if (isJsonRpcResponse(message)) { + const id = message.id; + if (id === undefined) { + this._options.logger.warn('McpProxyRoute: upstream response missing id'); + return; + } + const key = String(id); + if (this._pendingInitializeIds.delete(key) && isJsonRpcSuccessResponse(message)) { + this._captureUpstreamCapabilities(message.result); + } + if (this._pendingToolsListIds.delete(key) && isJsonRpcSuccessResponse(message)) { + this._captureToolUiMeta(message.result); + } + const entry = this._pendingSdkRequests.get(key); + if (!entry) { + this._options.logger.warn(`McpProxyRoute: upstream response for unknown id '${key}'`); + return; + } + this._pendingSdkRequests.delete(key); + entry.deferred.complete(message); + return; + } + if (isJsonRpcNotification(message)) { + this._handleUpstreamNotification(message); + return; + } + this._options.logger.warn('McpProxyRoute: upstream emitted an unrecognized JSON-RPC message'); + } + + private _handleUpstreamRequest(request: IJsonRpcRequest): void { + const method = request.method; + const params = request.params; + this._options.onUpstreamRequest(method, params).then(outcome => { + if (this._disposed) { + return; + } + const reply: IJsonRpcSuccessResponse | IJsonRpcErrorResponse = outcome.error + ? jsonRpcError(request.id, outcome.error.code, outcome.error.message, outcome.error.data) + : jsonRpcSuccess(request.id, outcome.result); + void this._sendToUpstream(reply); + }, err => { + if (this._disposed) { + return; + } + const message = err instanceof Error ? err.message : String(err); + this._options.logger.warn(`McpProxyRoute: upstream request handler threw for method '${method}': ${message}`); + void this._sendToUpstream(jsonRpcError(request.id, RPC_INTERNAL_ERROR, message)); + }); + } + + private _handleUpstreamNotification(notification: IJsonRpcNotification): void { + if (typeof notification.method !== 'string' || notification.method.length === 0) { + this._options.logger.warn('McpProxyRoute: upstream notification was malformed'); + return; + } + try { + this._options.onUpstreamNotification(notification.method, notification.params); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + this._options.logger.warn(`McpProxyRoute: upstream notification handler threw for method '${notification.method}': ${message}`); + } + } + + public override dispose(): void { + if (this._disposed) { + return; + } + this._disposed = true; + for (const { id, deferred } of this._pendingSdkRequests.values()) { + deferred.complete(jsonRpcError(id, RPC_INTERNAL_ERROR, 'Proxy route disposed')); + } + this._pendingSdkRequests.clear(); + this._pendingInitializeIds.clear(); + this._pendingToolsListIds.clear(); + this._toolUiMeta.clear(); + super.dispose(); + } +} diff --git a/src/vs/platform/agentHost/node/mcpHost/mcpStdioStateHandler.ts b/src/vs/platform/agentHost/node/mcpHost/mcpStdioStateHandler.ts new file mode 100644 index 0000000000000..6ca3c26e0e630 --- /dev/null +++ b/src/vs/platform/agentHost/node/mcpHost/mcpStdioStateHandler.ts @@ -0,0 +1,106 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { ChildProcessWithoutNullStreams } from 'child_process'; +import { TimeoutTimer } from '../../../../base/common/async.js'; +import { IDisposable } from '../../../../base/common/lifecycle.js'; +import { isWindows } from '../../../../base/common/platform.js'; +import { killTree } from '../../../../base/node/processes.js'; + +/** + * Inlined copy of {@link `src/vs/workbench/contrib/mcp/node/mcpStdioStateHandler.ts`}. + * `vs/platform` cannot import from `vs/workbench` per the layering rules + * enforced by `eslint.config.js`. + */ +const enum McpProcessState { + Running, + StdinEnded, + KilledPolite, + KilledForceful, +} + +/** + * Manages graceful shutdown of MCP stdio connections following the MCP specification. + * + * Per spec, shutdown should: + * 1. Close the input stream to the child process + * 2. Wait for the server to exit, or send SIGTERM if it doesn't exit within 10 seconds + * 3. Send SIGKILL if the server doesn't exit within 10 seconds after SIGTERM + * 4. Allow forceful killing if called twice + */ +export class McpStdioStateHandler implements IDisposable { + private static readonly GRACE_TIME_MS = 10_000; + + private _procState = McpProcessState.Running; + private _nextTimeout?: IDisposable; + + public get stopped() { + return this._procState !== McpProcessState.Running; + } + + constructor( + private readonly _child: ChildProcessWithoutNullStreams, + private readonly _graceTimeMs: number = McpStdioStateHandler.GRACE_TIME_MS + ) { } + + /** + * Initiates graceful shutdown. If called while shutdown is already in progress, + * forces immediate termination. + */ + public stop(): void { + if (this._procState === McpProcessState.Running) { + let graceTime = this._graceTimeMs; + try { + this._child.stdin.end(); + } catch (error) { + // If stdin.end() fails, continue with termination sequence + // This can happen if the stream is already in an error state + graceTime = 1; + } + this._procState = McpProcessState.StdinEnded; + this._nextTimeout = new TimeoutTimer(() => this.killPolite(), graceTime); + } else { + this._nextTimeout?.dispose(); + this.killForceful(); + } + } + + private async killPolite() { + this._procState = McpProcessState.KilledPolite; + this._nextTimeout = new TimeoutTimer(() => this.killForceful(), this._graceTimeMs); + + if (this._child.pid) { + if (!isWindows) { + await killTree(this._child.pid, false).catch(() => { + this._child.kill('SIGTERM'); + }); + } + } else { + this._child.kill('SIGTERM'); + } + } + + private async killForceful() { + this._procState = McpProcessState.KilledForceful; + + if (this._child.pid) { + await killTree(this._child.pid, true).catch(() => { + this._child.kill('SIGKILL'); + }); + } else { + this._child.kill(); + } + } + + public write(message: string): void { + if (!this.stopped) { + this._child.stdin.write(message + '\n'); + } + } + + public dispose() { + this._nextTimeout?.dispose(); + } +} diff --git a/src/vs/platform/agentHost/node/mcpHost/mcpStdioUpstream.ts b/src/vs/platform/agentHost/node/mcpHost/mcpStdioUpstream.ts new file mode 100644 index 0000000000000..f22a653ce69ca --- /dev/null +++ b/src/vs/platform/agentHost/node/mcpHost/mcpStdioUpstream.ts @@ -0,0 +1,225 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { ChildProcessWithoutNullStreams, spawn as defaultSpawn } from 'child_process'; +import { Emitter, type Event } from '../../../../base/common/event.js'; +import { Disposable, toDisposable } from '../../../../base/common/lifecycle.js'; +import { observableValue, type IObservable } from '../../../../base/common/observable.js'; +import type { JsonRpcMessage } from '../../../../base/common/jsonRpcProtocol.js'; +import type { ILogger } from '../../../log/common/log.js'; +import { StreamSplitter } from '../../../../base/node/nodeStreams.js'; +import type { IMcpStdioServerConfiguration } from '../../../mcp/common/mcpPlatformTypes.js'; +import { McpServerStatusKind, type McpServerStatus } from '../../common/state/protocol/state.js'; +import { McpStdioStateHandler } from './mcpStdioStateHandler.js'; +import type { IMcpUpstream, IMcpUpstreamCapabilities } from './mcpUpstream.js'; + +/** Spawn function (injectable for tests). */ +export type StdioSpawn = (command: string, args: readonly string[], options: { + readonly cwd?: string; + readonly env?: NodeJS.ProcessEnv; +}) => ChildProcessWithoutNullStreams; + +export interface IMcpStdioUpstreamOptions { + readonly config: IMcpStdioServerConfiguration; + readonly logger: ILogger; + /** Test seam — defaults to Node `child_process.spawn`. */ + readonly spawn?: StdioSpawn; +} + +/** + * MCP upstream backed by a child process speaking JSON-RPC over + * NDJSON on stdin/stdout (the canonical "stdio" transport). + * + * The child is spawned lazily on the first call to {@link start}; the + * constructor performs no I/O. Bearer tokens are not plumbed through + * stdio transports — see {@link setBearerToken}. + */ +export class McpStdioUpstream extends Disposable implements IMcpUpstream { + + private readonly _status = observableValue(this, { kind: McpServerStatusKind.Stopped }); + public readonly status: IObservable = this._status; + + private readonly _onMessage = this._register(new Emitter()); + public readonly onMessage: Event = this._onMessage.event; + + private readonly _upstreamCapabilities = observableValue(this, undefined); + public readonly upstreamCapabilities: IObservable = this._upstreamCapabilities; + + private readonly _config: IMcpStdioServerConfiguration; + private readonly _logger: ILogger; + private readonly _spawn: StdioSpawn; + + private _stateHandler: McpStdioStateHandler | undefined; + private _disposed = false; + private _stopRequested = false; + + constructor(options: IMcpStdioUpstreamOptions) { + super(); + this._config = options.config; + this._logger = options.logger; + this._spawn = options.spawn ?? defaultSpawn; + this._register(toDisposable(() => this._stop())); + } + + public async start(): Promise { + if (this._disposed) { + return this._status.get(); + } + const current = this._status.get(); + if (current.kind === McpServerStatusKind.Ready) { + return current; + } + if (current.kind === McpServerStatusKind.Starting) { + return current; + } + + this._status.set({ kind: McpServerStatusKind.Starting }, undefined); + + try { + // TODO(mcp/sandbox): the workbench-side McpServerLaunch path consults + // `IMcpSandboxConfiguration` for stdio servers (see + // `src/vs/platform/mcp/common/mcpPlatformTypes.ts` and the launcher in + // `src/vs/workbench/contrib/mcp/`). The agent-host proxy currently + // spawns unsandboxed; this matches the existing workbench gateway as + // of v1. + const child = this._spawn(this._config.command, this._config.args ?? [], { + cwd: this._config.cwd, + env: this._buildEnv(), + }); + this._stateHandler = new McpStdioStateHandler(child); + this._wireChild(child); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + const status: McpServerStatus = { + kind: McpServerStatusKind.Error, + error: { errorType: 'spawnFailed', message }, + }; + this._status.set(status, undefined); + return status; + } + + this._status.set({ kind: McpServerStatusKind.Ready }, undefined); + return this._status.get(); + } + + public async send(message: JsonRpcMessage): Promise { + const current = this._status.get(); + if (current.kind !== McpServerStatusKind.Ready) { + throw new Error(`McpStdioUpstream: cannot send while in state '${current.kind}'`); + } + this._stateHandler!.write(JSON.stringify(message)); + } + + public setBearerToken(_token: string | undefined): void { + // no-op for stdio; tokens are not yet plumbed through stdio MCP transports. + } + + public setUpstreamCapabilities(caps: IMcpUpstreamCapabilities | undefined): void { + this._upstreamCapabilities.set(caps, undefined); + } + + private _buildEnv(): NodeJS.ProcessEnv | undefined { + if (!this._config.env) { + return undefined; + } + const env: NodeJS.ProcessEnv = { ...process.env }; + for (const [key, value] of Object.entries(this._config.env)) { + if (value === null) { + delete env[key]; + } else { + env[key] = String(value); + } + } + return env; + } + + private _wireChild(child: ChildProcessWithoutNullStreams): void { + // `StreamSplitter` emits one chunk per `\n`-delimited line and also + // flushes any unterminated trailing chunk on stream end, so we don't + // have to keep our own line buffer or hand-roll an exit-time flush. + child.stdout.pipe(new StreamSplitter('\n')).on('data', (chunk: Buffer) => this._onStdoutLine(chunk)).resume(); + child.stderr.pipe(new StreamSplitter('\n')).on('data', (chunk: Buffer) => this._onStderrLine(chunk)).resume(); + + child.on('error', (err: Error) => { + this._logger.error(`McpStdioUpstream: child error: ${err.message}`); + if (this._disposed || this._stopRequested) { + return; + } + const current = this._status.get(); + if (current.kind === McpServerStatusKind.Error || current.kind === McpServerStatusKind.Stopped) { + return; + } + this._status.set({ + kind: McpServerStatusKind.Error, + error: { errorType: 'spawnFailed', message: err.message }, + }, undefined); + }); + + child.on('exit', (code: number | null, signal: NodeJS.Signals | null) => { + this._onChildExit(code, signal); + }); + } + + private _onStdoutLine(chunk: Buffer): void { + const trimmed = chunk.toString('utf8').trimEnd(); + if (!trimmed) { + return; + } + let parsed: JsonRpcMessage; + try { + parsed = JSON.parse(trimmed) as JsonRpcMessage; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + this._logger.error(`McpStdioUpstream: failed to parse stdout line: ${message}`); + return; + } + this._onMessage.fire(parsed); + } + + private _onStderrLine(chunk: Buffer): void { + const text = chunk.toString('utf8').trimEnd(); + if (text.length > 0) { + this._logger.info(`McpStdioUpstream[stderr]: ${text}`); + } + } + + private _onChildExit(code: number | null, signal: NodeJS.Signals | null): void { + if (this._disposed || this._stopRequested) { + this._status.set({ kind: McpServerStatusKind.Stopped }, undefined); + return; + } + if (code !== null && code !== 0) { + this._status.set({ + kind: McpServerStatusKind.Error, + error: { + errorType: 'childExited', + message: `MCP stdio child exited with code ${code}${signal ? ` (signal ${signal})` : ''}`, + }, + }, undefined); + } else { + this._status.set({ kind: McpServerStatusKind.Stopped }, undefined); + } + } + + private _stop(): void { + if (this._disposed) { + return; + } + this._disposed = true; + this._stopRequested = true; + this._upstreamCapabilities.set(undefined, undefined); + const handler = this._stateHandler; + this._stateHandler = undefined; + if (handler) { + handler.stop(); + handler.dispose(); + } + } + + public override dispose(): void { + this._stop(); + super.dispose(); + } +} diff --git a/src/vs/platform/agentHost/node/mcpHost/mcpUpstream.ts b/src/vs/platform/agentHost/node/mcpHost/mcpUpstream.ts new file mode 100644 index 0000000000000..d72f840ff8a9b --- /dev/null +++ b/src/vs/platform/agentHost/node/mcpHost/mcpUpstream.ts @@ -0,0 +1,78 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import type { Event } from '../../../../base/common/event.js'; +import type { IDisposable } from '../../../../base/common/lifecycle.js'; +import type { IObservable } from '../../../../base/common/observable.js'; +import type { JsonRpcMessage } from '../../../../base/common/jsonRpcProtocol.js'; +import type { McpServerStatus } from '../../common/state/protocol/state.js'; + +/** + * Server-side MCP capabilities advertised by the upstream in its + * `initialize` result. Loosely typed — the proxy treats the value + * opaquely except when checking for specific extension keys (e.g. + * `extensions['io.modelcontextprotocol/ui']`). The shape mirrors + * MCP's `InitializeResult.capabilities`. + */ +export interface IMcpUpstreamCapabilities { + readonly extensions?: { readonly [key: string]: unknown }; + readonly [key: string]: unknown; +} + +/** + * Transport-agnostic upstream MCP connection. Wraps either a stdio + * child process or a remote HTTP endpoint behind a uniform JSON-RPC + * message bus. + * + * Lifecycle: + * - constructed in `Stopped` state + * - `start()` transitions to `Starting`; on success → `Ready`; on + * auth failure → `AuthRequired`; on hard failure → `Error`. + * - `send()` is rejected unless state is `Ready`. + * - `dispose()` initiates graceful shutdown and resolves to `Stopped`. + */ +export interface IMcpUpstream extends IDisposable { + /** Current upstream status (mirrors what the proxy publishes to AHP). */ + readonly status: IObservable; + + /** Fires when a message arrives from the upstream MCP server. */ + readonly onMessage: Event; + + /** + * Capabilities the upstream advertised in its `initialize` result. + * `undefined` until the SDK's `initialize` round-trip has completed + * through the proxy. Updated each time a fresh `initialize` is + * observed (e.g. on session refresh). + */ + readonly upstreamCapabilities: IObservable; + + /** + * Begin connecting. Idempotent: a second call after `Ready` is a no-op, + * after `AuthRequired` retries the handshake with the current token. + * Resolves once the upstream reaches a terminal handshake state + * (`Ready`, `AuthRequired`, or `Error`). + */ + start(): Promise; + + /** + * Send a JSON-RPC message to the upstream. Rejects if state is not + * `Ready`. + */ + send(message: JsonRpcMessage): Promise; + + /** + * Set the bearer token used for HTTP requests. Stdio implementations + * MAY ignore this. After updating, call `start()` to retry. + */ + setBearerToken(token: string | undefined): void; + + /** + * Called by the proxy when it observes a successful `initialize` + * response from the upstream. Exposed because the route, not the + * upstream, is the layer that pairs SDK requests with their replies. + * Pass `undefined` to clear (e.g. on a re-handshake). + */ + setUpstreamCapabilities(caps: IMcpUpstreamCapabilities | undefined): void; +} diff --git a/src/vs/platform/agentHost/node/protocolServerHandler.ts b/src/vs/platform/agentHost/node/protocolServerHandler.ts index f7795bb00185e..7c2a61116b4bd 100644 --- a/src/vs/platform/agentHost/node/protocolServerHandler.ts +++ b/src/vs/platform/agentHost/node/protocolServerHandler.ts @@ -11,6 +11,8 @@ import { URI } from '../../../base/common/uri.js'; import { ILogService } from '../../log/common/log.js'; import { AHPFileSystemProvider } from '../common/agentHostFileSystemProvider.js'; import { AgentSession, type IAgentService } from '../common/agentService.js'; +import { IMcpHostService, IMcpHostUpstreamDelegate, IUpstreamMcpNotification, IUpstreamMcpRequest, IUpstreamMcpResponse } from '../common/mcpHost/mcpHostService.js'; +import type { ClientCapabilities, ServerCapabilities } from '../common/state/protocol/commands.js'; import type { CommandMap } from '../common/state/protocol/messages.js'; import type { UnsupportedProtocolVersionErrorData } from '../common/state/protocol/errors.js'; import { ActionEnvelope, ActionType, INotification, isSessionAction, isTerminalAction, type SessionAction, type TerminalAction, type IRootConfigChangedAction } from '../common/state/sessionActions.js'; @@ -84,6 +86,8 @@ interface IConnectedClient { readonly transport: IProtocolTransport; readonly subscriptions: Set; readonly disposables: DisposableStore; + /** Capabilities the client advertised during initialize. Currently empty in the spec but kept for future extension points. */ + readonly capabilities: ClientCapabilities | undefined; } /** @@ -122,6 +126,7 @@ export class ProtocolServerHandler extends Disposable { private readonly _server: IProtocolServer, private readonly _config: IProtocolServerConfig, private readonly _clientFileSystemProvider: AHPFileSystemProvider, + private readonly _mcpHostService: IMcpHostService, @ILogService private readonly _logService: ILogService, ) { super(); @@ -141,6 +146,14 @@ export class ProtocolServerHandler extends Disposable { this._register(this._stateManager.onDidEmitNotification(notification => { this._broadcastNotification(notification); })); + + // Install the host-side delegate that forwards upstream MCP traffic to + // a connected AHP client via reverse `mcpMethodCall` / `mcpNotification`. + const delegate: IMcpHostUpstreamDelegate = { + handleUpstreamRequest: (request) => this._handleUpstreamMcpRequest(request), + handleUpstreamNotification: (notification) => this._handleUpstreamMcpNotification(notification), + }; + this._register(this._mcpHostService.setUpstreamDelegate(delegate)); } // ---- Connection handling ------------------------------------------------- @@ -206,6 +219,15 @@ export class ProtocolServerHandler extends Disposable { } } break; + case 'mcpNotification': + if (client) { + try { + this._mcpHostService.notify(msg.params); + } catch (err) { + this._logService.warn(`[ProtocolServer] mcpNotification handler threw: ${err instanceof Error ? err.message : String(err)}`); + } + } + break; } } else if (isJsonRpcResponse(msg)) { const pending = this._pendingReverseRequests.get(msg.id); @@ -270,6 +292,7 @@ export class ProtocolServerHandler extends Disposable { transport, subscriptions: new Set(), disposables, + capabilities: params.capabilities, }; this._clients.set(params.clientId, client); this._onDidChangeConnectionCount.fire(this._clients.size); @@ -298,6 +321,8 @@ export class ProtocolServerHandler extends Disposable { } } + const serverCapabilities: ServerCapabilities = {}; + return { client, response: { @@ -305,6 +330,7 @@ export class ProtocolServerHandler extends Disposable { serverSeq: this._stateManager.serverSeq, snapshots, defaultDirectory: this._config.defaultDirectory, + capabilities: serverCapabilities, completionTriggerCharacters: this._config.completionTriggerCharacters, }, }; @@ -320,12 +346,16 @@ export class ProtocolServerHandler extends Disposable { // Synchronously install the client so messages arriving on this transport // while we restore subscriptions can find a valid client object. The // reconnect response is only sent once `responsePromise` resolves below. + // Reconnect doesn't carry capabilities; preserve them from the prior + // client record if present, otherwise treat as absent. + const previousClient = this._clients.get(params.clientId); const client: IConnectedClient = { clientId: params.clientId, protocolVersion: PROTOCOL_VERSION, transport, subscriptions: new Set(), disposables, + capabilities: previousClient?.capabilities, }; this._clients.set(params.clientId, client); this._onDidChangeConnectionCount.fire(this._clients.size); @@ -638,7 +668,11 @@ export class ProtocolServerHandler extends Disposable { return {}; }, authenticate: async (_client, params) => { - const result = await this._agentService.authenticate(params); + const result = await this._agentService.authenticate({ + resource: params.resource, + token: params.token, + server: params.server ? URI.parse(params.server) : undefined, + }); if (!result.authenticated) { throw new ProtocolError(AHP_AUTH_REQUIRED, 'Authentication failed for resource: ' + params.resource); } @@ -652,6 +686,9 @@ export class ProtocolServerHandler extends Disposable { await this._agentService.disposeTerminal(URI.parse(params.terminal)); return null; }, + mcpMethodCall: async (_client, params) => { + return this._mcpHostService.callMethod(params); + }, }; @@ -751,6 +788,81 @@ export class ProtocolServerHandler extends Disposable { } } + // ---- Upstream MCP forwarding ------------------------------------------- + + /** + * Pick a connected client to forward an upstream-originated MCP request + * to. Today we pick the first client subscribed to the session that owns + * the server (looked up via `mcp://` segment); if + * none are subscribed, we fall back to the first connected client. If + * there are no clients at all, returns `undefined`. + */ + private _pickUpstreamMcpTarget(serverResource: URI): IConnectedClient | undefined { + const segments = serverResource.path.split('/').filter(s => s.length > 0); + if (segments.length >= 1) { + // Reconstruct the session URI. The convention is + // `mcp://` where `` is the + // **path component** of the parent session URI. We don't know the + // session's scheme so we match by suffix. + const sessionPathTail = segments[0]; + for (const client of this._clients.values()) { + for (const sub of client.subscriptions) { + try { + const parsed = URI.parse(sub); + if (parsed.scheme !== 'mcp' && parsed.path.endsWith('/' + sessionPathTail)) { + return client; + } + } catch { + // ignore + } + } + } + } + return this._clients.values().next().value; + } + + private async _handleUpstreamMcpRequest(request: IUpstreamMcpRequest): Promise { + const client = this._pickUpstreamMcpTarget(request.server); + if (!client) { + return { + error: { + code: JsonRpcErrorCodes.MethodNotFound, + message: `No AHP client is connected to forward MCP request '${request.method}' to '${request.server.toString()}'`, + }, + }; + } + try { + const result = await this._sendReverseRequest(client.clientId, 'mcpMethodCall', { + server: request.server.toString(), + method: request.method, + params: request.params, + }); + return { result }; + } catch (err) { + if (err instanceof ProtocolError) { + return { error: { code: err.code, message: err.message, data: err.data } }; + } + const message = err instanceof Error ? err.message : String(err); + return { error: { code: JsonRpcErrorCodes.InternalError, message } }; + } + } + + private _handleUpstreamMcpNotification(notification: IUpstreamMcpNotification): void { + const msg = { + jsonrpc: '2.0' as const, + method: 'mcpNotification' as const, + params: { + server: notification.server.toString(), + method: notification.method, + params: notification.params, + }, + }; + const target = this._pickUpstreamMcpTarget(notification.server); + if (target) { + target.transport.send(msg); + } + } + private _isRelevantToClient(client: IConnectedClient, envelope: ActionEnvelope): boolean { const action = envelope.action; if (action.type.startsWith('root/')) { @@ -765,6 +877,11 @@ export class ProtocolServerHandler extends Disposable { return false; } + /** @internal Test accessor — do not call from production. */ + getClientCapabilitiesForTest(clientId: string): ClientCapabilities | undefined { + return this._clients.get(clientId)?.capabilities; + } + override dispose(): void { for (const client of this._clients.values()) { client.disposables.dispose(); diff --git a/src/vs/platform/agentHost/test/common/mcpServerUri.test.ts b/src/vs/platform/agentHost/test/common/mcpServerUri.test.ts new file mode 100644 index 0000000000000..e035140024447 --- /dev/null +++ b/src/vs/platform/agentHost/test/common/mcpServerUri.test.ts @@ -0,0 +1,51 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as assert from 'assert'; +import { URI } from '../../../../base/common/uri.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import { buildMcpServerUri, MCP_SERVER_SCHEME, parseMcpServerUri } from '../../common/state/mcpServerUri.js'; + +suite('mcpServerUri', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + test('round-trips a simple session/server pair', () => { + const session = URI.parse('copilot:/abcd-1234'); + const built = buildMcpServerUri(session, 'github-mcp'); + assert.deepStrictEqual({ + scheme: built.scheme, + path: built.path, + parsed: parseMcpServerUri(built), + }, { + scheme: MCP_SERVER_SCHEME, + path: '/abcd-1234/github-mcp', + parsed: { sessionPath: 'abcd-1234', serverId: 'github-mcp' }, + }); + }); + + test('rejects URIs with a non-mcp scheme', () => { + const other = URI.parse('copilot:/abcd-1234/some-server'); + assert.strictEqual(parseMcpServerUri(other), undefined); + }); + + test('encodes and decodes server names with special characters', () => { + const session = URI.parse('copilot:/sess-1'); + const serverId = 'a/b c@github.com'; + const built = buildMcpServerUri(session, serverId); + const parsed = parseMcpServerUri(built); + assert.deepStrictEqual({ + containsRawSlash: built.path.lastIndexOf('/') === '/sess-1'.length, + parsed, + }, { + containsRawSlash: true, + parsed: { sessionPath: 'sess-1', serverId }, + }); + }); + + test('rejects malformed mcp URIs', () => { + const noServer = URI.from({ scheme: MCP_SERVER_SCHEME, path: '/sess-only' }); + assert.strictEqual(parseMcpServerUri(noServer), undefined); + }); +}); diff --git a/src/vs/platform/agentHost/test/common/nullMcpHostService.test.ts b/src/vs/platform/agentHost/test/common/nullMcpHostService.test.ts new file mode 100644 index 0000000000000..7105a5ef305e9 --- /dev/null +++ b/src/vs/platform/agentHost/test/common/nullMcpHostService.test.ts @@ -0,0 +1,64 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { URI } from '../../../../base/common/uri.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import { McpServerType } from '../../../mcp/common/mcpPlatformTypes.js'; +import type { IMcpServerDefinition } from '../../../agentPlugins/common/pluginParsers.js'; +import { NullMcpHostService } from '../../common/mcpHost/nullMcpHostService.js'; +import { JsonRpcErrorCodes, ProtocolError } from '../../common/state/sessionProtocol.js'; + +suite('NullMcpHostService', () => { + + ensureNoDisposablesAreLeakedInTestSuite(); + + test('setSessionServers returns an empty array', () => { + const service = new NullMcpHostService(); + const session = URI.parse('agent-session:/session-1'); + const def: IMcpServerDefinition = { + name: 'test-server', + uri: URI.file('/plugin'), + configuration: { + type: McpServerType.LOCAL, + command: 'node', + }, + }; + + const handles = service.setSessionServers(session, [def]); + assert.deepStrictEqual(handles, []); + }); + + test('getServer returns undefined', () => { + const service = new NullMcpHostService(); + const resource = URI.parse('mcp:/session-1/test-server'); + assert.strictEqual(service.getServer(resource), undefined); + }); + + test('callMethod rejects with MethodNotFound ProtocolError', async () => { + const service = new NullMcpHostService(); + const error = await service.callMethod( + { server: 'mcp:/session-1/test-server', method: 'tools/list', params: {} }, + ).then(() => undefined, (err: unknown) => err); + + assert.ok(error instanceof ProtocolError); + assert.strictEqual(error.code, JsonRpcErrorCodes.MethodNotFound); + }); + + test('notify is a no-op', () => { + const service = new NullMcpHostService(); + service.notify({ server: 'mcp:/session-1/test-server', method: 'notifications/message', params: {} }); + }); + + test('setUpstreamDelegate returns a disposable that does not throw', () => { + const service = new NullMcpHostService(); + const reg = service.setUpstreamDelegate({ + handleUpstreamRequest: async () => ({ result: {} }), + handleUpstreamNotification: () => { /* no-op */ }, + }); + reg.dispose(); + }); +}); + diff --git a/src/vs/platform/agentHost/test/node/agentHostStateManager.test.ts b/src/vs/platform/agentHost/test/node/agentHostStateManager.test.ts index be2d65f0921fa..c17f2a122c748 100644 --- a/src/vs/platform/agentHost/test/node/agentHostStateManager.test.ts +++ b/src/vs/platform/agentHost/test/node/agentHostStateManager.test.ts @@ -12,6 +12,7 @@ import { NullLogService } from '../../../log/common/log.js'; import { ActionType, NotificationType, type ActionEnvelope, type INotification } from '../../common/state/sessionActions.js'; import { SessionSummary, ResponsePartKind, ROOT_STATE_URI, SessionLifecycle, SessionStatus, TurnState, buildSubagentSessionUri, buildSubagentSessionUriPrefix, isSubagentSession, parseSubagentSessionUri, type MarkdownResponsePart, type SessionState } from '../../common/state/sessionState.js'; import { type SessionSummaryChangedNotification } from '../../common/state/protocol/notifications.js'; +import { McpServerStatusKind, type McpServerSummary } from '../../common/state/protocol/state.js'; import { AgentHostStateManager } from '../../node/agentHostStateManager.js'; suite('AgentHostStateManager', () => { @@ -471,6 +472,95 @@ suite('AgentHostStateManager', () => { assert.strictEqual(changed[0].changes.status, SessionStatus.Idle, 'status should be Idle so the spinner clears'); }); }); + + suite('mcp servers', () => { + + const mcpResource = 'mcp:/test-session/server-1'; + + function makeServerSummary(overrides?: Partial): McpServerSummary { + return { + resource: mcpResource, + label: 'Server 1', + status: { kind: McpServerStatusKind.Starting }, + ...overrides, + }; + } + + test('createMcpServer registers the server and emits a McpServerAdded envelope', () => { + manager.createSession(makeSessionSummary()); + + const envelopes: ActionEnvelope[] = []; + disposables.add(manager.onDidEmitEnvelope(e => envelopes.push(e))); + + const summary = makeServerSummary(); + manager.createMcpServer(sessionUri, summary); + + assert.deepStrictEqual({ + envelopeTypes: envelopes.map(e => e.action.type), + sessionMcpServers: manager.getSessionState(sessionUri)?.mcpServers, + }, { + envelopeTypes: [ActionType.McpServerAdded], + sessionMcpServers: [summary], + }); + }); + + test('setMcpServerStatus updates session summary', () => { + manager.createSession(makeSessionSummary()); + manager.createMcpServer(sessionUri, makeServerSummary()); + + const envelopes: ActionEnvelope[] = []; + disposables.add(manager.onDidEmitEnvelope(e => envelopes.push(e))); + + const newStatus = { kind: McpServerStatusKind.Ready } as const; + manager.setMcpServerStatus(sessionUri, mcpResource, newStatus); + + assert.deepStrictEqual({ + envelopeTypes: envelopes.map(e => e.action.type), + sessionMcpServers: manager.getSessionState(sessionUri)?.mcpServers, + }, { + envelopeTypes: [ActionType.McpServerStatusChanged], + sessionMcpServers: [{ resource: mcpResource, label: 'Server 1', status: newStatus }], + }); + }); + + test('removeMcpServer drops session summary entry', () => { + manager.createSession(makeSessionSummary()); + manager.createMcpServer(sessionUri, makeServerSummary()); + + const envelopes: ActionEnvelope[] = []; + disposables.add(manager.onDidEmitEnvelope(e => envelopes.push(e))); + + manager.removeMcpServer(sessionUri, mcpResource); + + assert.deepStrictEqual({ + envelopeTypes: envelopes.map(e => e.action.type), + sessionMcpServers: manager.getSessionState(sessionUri)?.mcpServers, + snapshot: manager.getSnapshot(mcpResource), + }, { + envelopeTypes: [ActionType.McpServerRemoved], + sessionMcpServers: undefined, + snapshot: undefined, + }); + }); + + test('createSession seeds initial MCP server summaries', () => { + const envelopes: ActionEnvelope[] = []; + disposables.add(manager.onDidEmitEnvelope(e => envelopes.push(e))); + + const summary = makeServerSummary(); + const state = manager.createSession(makeSessionSummary(), { mcpServers: [summary] }); + + assert.deepStrictEqual({ + envelopeTypes: envelopes.map(e => e.action.type), + stateMcpServers: state.mcpServers, + snapshotMcpServers: (manager.getSnapshot(sessionUri)?.state as SessionState | undefined)?.mcpServers, + }, { + envelopeTypes: [], + stateMcpServers: [summary], + snapshotMcpServers: [summary], + }); + }); + }); }); suite('Subagent URI helpers', () => { diff --git a/src/vs/platform/agentHost/test/node/agentService.test.ts b/src/vs/platform/agentHost/test/node/agentService.test.ts index 620caf54d7609..70723a8f93dec 100644 --- a/src/vs/platform/agentHost/test/node/agentService.test.ts +++ b/src/vs/platform/agentHost/test/node/agentService.test.ts @@ -8,28 +8,31 @@ import { mkdtempSync, readFileSync, rmSync } from 'fs'; import { tmpdir } from 'os'; import { fileURLToPath } from 'url'; import { encodeBase64, VSBuffer } from '../../../../base/common/buffer.js'; -import { DisposableStore, IReference, toDisposable } from '../../../../base/common/lifecycle.js'; +import { DisposableStore, IDisposable, IReference, toDisposable } from '../../../../base/common/lifecycle.js'; import { Schemas } from '../../../../base/common/network.js'; +import { observableValue, type IObservable } from '../../../../base/common/observable.js'; import { joinPath } from '../../../../base/common/resources.js'; +import { hasKey } from '../../../../base/common/types.js'; import { URI } from '../../../../base/common/uri.js'; -import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import { runWithFakedTimers } from '../../../../base/test/common/timeTravelScheduler.js'; -import { hasKey } from '../../../../base/common/types.js'; -import { NullLogService } from '../../../log/common/log.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import type { IMcpServerDefinition } from '../../../agentPlugins/common/pluginParsers.js'; import { FileService } from '../../../files/common/fileService.js'; import { InMemoryFileSystemProvider } from '../../../files/common/inMemoryFilesystemProvider.js'; +import { NullLogService } from '../../../log/common/log.js'; +import { IProductService } from '../../../product/common/productService.js'; import { AgentSession } from '../../common/agentService.js'; +import { IMcpHostService, IMcpServerHandle } from '../../common/mcpHost/mcpHostService.js'; import { ISessionDatabase, ISessionDataService } from '../../common/sessionDataService.js'; -import { SessionDatabase } from '../../node/sessionDatabase.js'; -import { ActionType, ActionEnvelope } from '../../common/state/sessionActions.js'; -import { MessageAttachmentKind, SessionActiveClient, ResponsePartKind, SessionLifecycle, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType, TurnState, buildSubagentSessionUri, type MarkdownResponsePart, type ToolCallCompletedState, type ToolCallResponsePart } from '../../common/state/sessionState.js'; -import { type MessageResourceAttachment } from '../../common/state/protocol/state.js'; -import { IProductService } from '../../../product/common/productService.js'; +import { McpServerStatusKind, MessageResourceAttachment, type McpServerSummary } from '../../common/state/protocol/state.js'; +import { ActionEnvelope, ActionType } from '../../common/state/sessionActions.js'; +import { buildSubagentSessionUri, MessageAttachmentKind, ResponsePartKind, SessionActiveClient, SessionLifecycle, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType, TurnState, type MarkdownResponsePart, type ToolCallCompletedState, type ToolCallResponsePart } from '../../common/state/sessionState.js'; import { AgentService } from '../../node/agentService.js'; -import { MockAgent, ScriptedMockAgent } from './mockAgent.js'; -import { mapSessionEventsToHistoryRecords } from './historyRecordFixtures.js'; import { type ISessionEvent } from '../../node/copilot/mapSessionEvents.js'; +import { SessionDatabase } from '../../node/sessionDatabase.js'; import { createNoopGitService, createSessionDataService } from '../common/sessionTestHelpers.js'; +import { mapSessionEventsToHistoryRecords } from './historyRecordFixtures.js'; +import { MockAgent, ScriptedMockAgent } from './mockAgent.js'; /** * Loads a JSONL fixture of raw Copilot SDK events, runs them through @@ -172,6 +175,26 @@ suite('AgentService (node dispatcher)', () => { }); }); + suite('subscribe to mcp:/ resources', () => { + + test('rejects unknown mcp:/ server without retaining subscription', async () => { + const mcpResource = URI.parse('mcp:/unknown/server'); + + await assert.rejects( + () => service.subscribe(mcpResource, 'client-1'), + (err: Error) => err.message.includes('mcp:/unknown/server'), + ); + + // The catch block must roll the addSubscriber registration back — + // a second subscribe attempt should reject the same way (not + // e.g. fail because the subscription already exists). + await assert.rejects( + () => service.subscribe(mcpResource, 'client-1'), + (err: Error) => err.message.includes('mcp:/unknown/server'), + ); + }); + }); + // ---- attachment rewriting ------------------------------------------ suite('user-message attachment rewriting', () => { @@ -691,6 +714,39 @@ suite('AgentService (node dispatcher)', () => { }); }); + test('seeds initial MCP servers from the MCP host summaries', async () => { + const session = AgentSession.uri('copilot', 'mcp-seeded'); + const mcpServer: McpServerSummary = { + resource: `mcp:/${AgentSession.id(session)}/github`, + label: 'github', + status: { kind: McpServerStatusKind.Starting }, + }; + class McpAgent extends MockAgent { + override async createSession() { + return { + session, + project: { uri: URI.file('/test'), displayName: 'Test' }, + }; + } + } + const mcpHostService = new SpyMcpHostService(); + mcpHostService.sessionSummaries.set(session.toString(), [mcpServer]); + const localService = disposables.add(new AgentService(new NullLogService(), fileService, nullSessionDataService, { _serviceBrand: undefined } as IProductService, createNoopGitService(), undefined, mcpHostService)); + const agent = new McpAgent('copilot'); + disposables.add(toDisposable(() => agent.dispose())); + localService.registerProvider(agent); + + await localService.createSession({ provider: 'copilot' }); + + assert.deepStrictEqual({ + mcpServers: localService.stateManager.getSessionState(session.toString())?.mcpServers, + mcpHostCalls: mcpHostService.setSessionServersCalls, + }, { + mcpServers: [mcpServer], + mcpHostCalls: [], + }); + }); + test('omits activeClient from the initial session state when not provided', async () => { service.registerProvider(copilotAgent); @@ -786,6 +842,111 @@ suite('AgentService (node dispatcher)', () => { assert.deepStrictEqual(result, { authenticated: false }); }); + + // ---- per-server (MCP) --------------------------------------------- + + function newServiceWithSpy(spy: SpyMcpHostService): AgentService { + return disposables.add(new AgentService( + new NullLogService(), fileService, nullSessionDataService, + { _serviceBrand: undefined } as IProductService, + createNoopGitService(), undefined, spy, + )); + } + + test('routes per-server token to the registered MCP handle', async () => { + const spy = new SpyMcpHostService(); + const serverUri = URI.parse('mcp:/sess1/srv1'); + spy.registerStubServer(serverUri, true); + const svc = newServiceWithSpy(spy); + + const result = await svc.authenticate({ + resource: 'https://api.github.com', + token: 'tok', + server: serverUri, + }); + + assert.deepStrictEqual({ + result, + calls: spy.authenticateCalls, + }, { + result: { authenticated: true }, + calls: [{ resource: 'https://api.github.com', token: 'tok', server: serverUri }], + }); + }); + + test('per-server authenticate returns false when the server is unregistered (no throw)', async () => { + const spy = new SpyMcpHostService(); + const svc = newServiceWithSpy(spy); + + const result = await svc.authenticate({ + resource: 'https://api.github.com', + token: 'tok', + server: URI.parse('mcp:/sess1/missing'), + }); + + assert.deepStrictEqual({ + result, + calls: spy.authenticateCalls, + }, { + result: { authenticated: false }, + calls: [], + }); + }); + + test('per-server authenticate does NOT fan out to providers', async () => { + const spy = new SpyMcpHostService(); + const serverUri = URI.parse('mcp:/sess1/srv1'); + spy.registerStubServer(serverUri, true); + const svc = newServiceWithSpy(spy); + svc.registerProvider(copilotAgent); // owns https://api.github.com + + await svc.authenticate({ + resource: 'https://api.github.com', + token: 'tok', + server: serverUri, + }); + + assert.deepStrictEqual({ + providerCalls: copilotAgent.authenticateCalls, + handleCalls: spy.authenticateCalls, + }, { + providerCalls: [], + handleCalls: [{ resource: 'https://api.github.com', token: 'tok', server: serverUri }], + }); + }); + + test('agent-level authenticate (no server) still fans out to providers', async () => { + const spy = new SpyMcpHostService(); + const svc = newServiceWithSpy(spy); + svc.registerProvider(copilotAgent); + + const result = await svc.authenticate({ resource: 'https://api.github.com', token: 'tok' }); + + assert.deepStrictEqual({ + result, + providerCalls: copilotAgent.authenticateCalls, + handleCalls: spy.authenticateCalls, + }, { + result: { authenticated: true }, + providerCalls: [{ resource: 'https://api.github.com', token: 'tok' }], + handleCalls: [], + }); + }); + + test('per-server authenticate returns false when the handle throws (no throw)', async () => { + const spy = new SpyMcpHostService(); + const serverUri = URI.parse('mcp:/sess1/srv1'); + spy.registerStubServer(serverUri, new Error('boom')); + const svc = newServiceWithSpy(spy); + + const result = await svc.authenticate({ + resource: 'https://api.github.com', + token: 'tok', + server: serverUri, + }); + + assert.deepStrictEqual(result, { authenticated: false }); + }); }); // ---- shutdown ------------------------------------------------------- @@ -1548,4 +1709,65 @@ suite('AgentService (node dispatcher)', () => { assert.strictEqual(state?.summary.workingDirectory, worktreeDir.toString()); }); }); + + class SpyMcpHostService implements IMcpHostService { + declare readonly _serviceBrand: undefined; + readonly setSessionServersCalls: { session: URI; servers: readonly IMcpServerDefinition[] }[] = []; + readonly authenticateCalls: { resource: string; token: string; server: URI }[] = []; + readonly sessionSummaries = new Map(); + + private readonly _stubs = new Map(); + + setSessionServers(session: URI, servers: readonly IMcpServerDefinition[]): readonly IMcpServerHandle[] { + this.setSessionServersCalls.push({ session, servers }); + return []; + } + getServerSummaries(session: URI): readonly McpServerSummary[] { + return this.sessionSummaries.get(session.toString()) ?? []; + } + getServer(resource: URI): IMcpServerHandle | undefined { + return this._stubs.get(resource.toString()); + } + async callMethod(): Promise<{ result: unknown }> { throw new Error('not used in test'); } + notify(): void { /* no-op */ } + setUpstreamDelegate(): IDisposable { return { dispose: () => { /* no-op */ } }; } + + /** + * Register a stub {@link IMcpServerHandle} returned from + * {@link getServer} for `resource`. The handle's + * {@link IMcpServerHandle.authenticate} records calls into + * {@link authenticateCalls} and resolves with `accept`. If + * `accept` is an `Error`, the handle rejects with that error + * — used to exercise the catch branch in + * {@link AgentService.authenticate}. + */ + registerStubServer(resource: URI, accept: boolean | Error): IMcpServerHandle { + const summary: McpServerSummary = { + resource: resource.toString(), + label: 'stub', + status: { kind: McpServerStatusKind.Ready }, + }; + const summaryObs: IObservable = observableValue('summary', summary); + const endpointObs: IObservable = observableValue('endpoint', undefined); + const handle: IMcpServerHandle = { + resource, + summary: summaryObs, + endpoint: endpointObs, + authenticate: async (res, tok) => { + this.authenticateCalls.push({ resource: res, token: tok, server: resource }); + if (accept instanceof Error) { + throw accept; + } + return accept; + }, + callMethod: async () => { throw new Error('not used in test'); }, + notify: () => { /* no-op */ }, + getToolUiMeta: () => undefined, + getUiHostCapabilities: () => ({}), + dispose: () => { /* no-op */ }, + }; + this._stubs.set(resource.toString(), handle); + return handle; + } + } }); diff --git a/src/vs/platform/agentHost/test/node/copilotAgent.test.ts b/src/vs/platform/agentHost/test/node/copilotAgent.test.ts index aab20eef2a409..52292628fc1a9 100644 --- a/src/vs/platform/agentHost/test/node/copilotAgent.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotAgent.test.ts @@ -8,7 +8,7 @@ import assert from 'assert'; import * as fs from 'fs/promises'; import * as os from 'os'; import { Disposable, type DisposableStore, type IDisposable, type IReference } from '../../../../base/common/lifecycle.js'; -import { waitForState } from '../../../../base/common/observable.js'; +import { observableValue, waitForState, type IObservable } from '../../../../base/common/observable.js'; import { URI } from '../../../../base/common/uri.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import { INativeEnvironmentService } from '../../../environment/common/environment.js'; @@ -20,16 +20,22 @@ import { ServiceCollection } from '../../../instantiation/common/serviceCollecti import { ILogService, NullLogService } from '../../../log/common/log.js'; import { IAgentPluginManager, ISyncedCustomization } from '../../common/agentPluginManager.js'; import { AgentSession, type AgentSignal, type IAgentActionSignal, type IAgentSessionMetadata } from '../../common/agentService.js'; +import { IMcpHostService, type IMcpServerHandle } from '../../common/mcpHost/mcpHostService.js'; +import { NullMcpHostService } from '../../common/mcpHost/nullMcpHostService.js'; import { ISessionDataService } from '../../common/sessionDataService.js'; import { AHP_AUTH_REQUIRED, ProtocolError } from '../../common/state/sessionProtocol.js'; import { buildSubagentSessionUri, ResponsePartKind, SessionCustomization, TurnState, type CustomizationRef, type MarkdownResponsePart, type ToolCallResult, type Turn } from '../../common/state/sessionState.js'; import { ActionType, type IDeltaAction } from '../../common/state/sessionActions.js'; +import type { IMcpServerDefinition, IParsedPlugin } from '../../../agentPlugins/common/pluginParsers.js'; +import { McpServerType } from '../../../mcp/common/mcpPlatformTypes.js'; +import { McpAuthRequiredReason, McpServerStatusKind, type McpServerSummary } from '../../common/state/protocol/state.js'; +import { buildMcpServerUri } from '../../common/state/mcpServerUri.js'; import { AgentConfigurationService, IAgentConfigurationService } from '../../node/agentConfigurationService.js'; import { AgentHostStateManager } from '../../node/agentHostStateManager.js'; import { IAgentHostGitService } from '../../node/agentHostGitService.js'; import { IAgentHostTerminalManager } from '../../node/agentHostTerminalManager.js'; -import { COPILOT_AGENT_HOST_SYSTEM_MESSAGE, CopilotAgent, getCopilotBranchNameHintFromMessage, getCopilotWorktreeBranchName, getCopilotWorktreeName, getCopilotWorktreesRoot } from '../../node/copilot/copilotAgent.js'; +import { COPILOT_AGENT_HOST_SYSTEM_MESSAGE, ActiveClient, CopilotAgent, getCopilotBranchNameHintFromMessage, getCopilotWorktreeBranchName, getCopilotWorktreeName, getCopilotWorktreesRoot } from '../../node/copilot/copilotAgent.js'; import { CopilotAgentSession, type SessionWrapperFactory } from '../../node/copilot/copilotAgentSession.js'; import { CopilotSessionWrapper } from '../../node/copilot/copilotSessionWrapper.js'; import { ShellManager } from '../../node/copilot/copilotShellTools.js'; @@ -226,8 +232,9 @@ class TestableCopilotAgent extends CopilotAgent { @IAgentHostGitService gitService: IAgentHostGitService, @IAgentHostTerminalManager terminalManager: IAgentHostTerminalManager, @IAgentConfigurationService configurationService: IAgentConfigurationService, + @IMcpHostService mcpHostService: IMcpHostService, ) { - super(logService, instantiationService, fileService, sessionDataService, gitService, terminalManager, configurationService); + super(logService, instantiationService, fileService, sessionDataService, gitService, terminalManager, configurationService, mcpHostService); this._enablePlanModeOnClient(this._copilotClient as CopilotClient); } @@ -278,7 +285,7 @@ class TestableCopilotAgent extends CopilotAgent { } } -function createTestAgentContext(disposables: Pick, options?: { sessionDataService?: ISessionDataService; copilotClient?: ITestCopilotClient; gitService?: TestAgentHostGitService; environmentServiceRegistration?: 'native' | 'none'; pluginManager?: IAgentPluginManager }): { agent: CopilotAgent; instantiationService: IInstantiationService } { +function createTestAgentContext(disposables: Pick, options?: { sessionDataService?: ISessionDataService; copilotClient?: ITestCopilotClient; gitService?: TestAgentHostGitService; environmentServiceRegistration?: 'native' | 'none'; pluginManager?: IAgentPluginManager; mcpHostService?: IMcpHostService }): { agent: CopilotAgent; instantiationService: IInstantiationService } { const services = new ServiceCollection(); const logService = new NullLogService(); const fileService = disposables.add(new FileService(logService)); @@ -291,6 +298,7 @@ function createTestAgentContext(disposables: Pick, optio services.set(IAgentPluginManager, options?.pluginManager ?? new TestAgentPluginManager()); services.set(IAgentHostGitService, options?.gitService ?? new TestAgentHostGitService()); services.set(IAgentHostTerminalManager, new TestAgentHostTerminalManager()); + services.set(IMcpHostService, options?.mcpHostService ?? new NullMcpHostService()); if (options?.environmentServiceRegistration !== 'none') { const environmentService = { _serviceBrand: undefined, @@ -306,7 +314,7 @@ function createTestAgentContext(disposables: Pick, optio return { agent, instantiationService }; } -function createTestAgent(disposables: Pick, options?: { sessionDataService?: ISessionDataService; copilotClient?: ITestCopilotClient; gitService?: TestAgentHostGitService; environmentServiceRegistration?: 'native' | 'none'; pluginManager?: IAgentPluginManager }): CopilotAgent { +function createTestAgent(disposables: Pick, options?: { sessionDataService?: ISessionDataService; copilotClient?: ITestCopilotClient; gitService?: TestAgentHostGitService; environmentServiceRegistration?: 'native' | 'none'; pluginManager?: IAgentPluginManager; mcpHostService?: IMcpHostService }): CopilotAgent { return createTestAgentContext(disposables, options).agent; } @@ -1142,4 +1150,236 @@ suite('CopilotAgent', () => { }); }); + + suite('mcp servers (direct publish)', () => { + + function makeServerDef(name: string, command: string): IMcpServerDefinition { + return { name, uri: URI.file(`/plugin/${name}`), configuration: { type: McpServerType.LOCAL, command } }; + } + + function makePluginWith(servers: readonly IMcpServerDefinition[]): IParsedPlugin { + return { hooks: [], skills: [], agents: [], mcpServers: servers }; + } + + interface ICopilotAgentInternals { + _resolveMcpServersForSdk(sessionUri: URI, defs: readonly IMcpServerDefinition[]): Promise>; + } + + test('ActiveClient.snapshot republishes MCP servers only when the resolved set changed', async () => { + const spy = new SpyMcpHostService(); + const sessionUri = AgentSession.uri('copilotcli', 'pub-1'); + const a = makeServerDef('a', 'cmdA'); + const b = makeServerDef('b', 'cmdB'); + + // The plugin resolver returns whatever the test sets here; this + // avoids the PluginController plumbing and tests the dedup behavior + // of ActiveClient in isolation. + let next: readonly IParsedPlugin[] = []; + const client = disposables.add(new ActiveClient( + sessionUri, + async () => next, + spy, + new NullLogService(), + )); + + next = [makePluginWith([a])]; + await client.snapshot(undefined); + next = [makePluginWith([a])]; // unchanged + await client.snapshot(undefined); + next = [makePluginWith([a, b])]; // changed + await client.snapshot(undefined); + + assert.deepStrictEqual( + spy.setSessionServersCalls.map(c => ({ session: c.session.toString(), serverNames: c.servers.map(s => s.name) })), + [ + { session: sessionUri.toString(), serverNames: ['a'] }, + { session: sessionUri.toString(), serverNames: ['a', 'b'] }, + ], + ); + }); + + test('createSession publishes MCP servers and returns host summaries for initial state', async () => { + const spy = new SpyMcpHostService(); + const agent = createTestAgent(disposables, { + copilotClient: new TestCopilotClient([]), + mcpHostService: spy, + }); + try { + await agent.authenticate('https://api.github.com', 'token'); + + const session = AgentSession.uri('copilotcli', 'eager-1'); + const serverSummary: McpServerSummary = { + resource: buildMcpServerUri(session, 'github').toString(), + label: 'github', + status: { kind: McpServerStatusKind.Starting }, + }; + spy.sessionSummaries.set(session.toString(), [serverSummary]); + + await agent.createSession({ + session, + workingDirectory: URI.file('/workspace'), + activeClient: { + clientId: 'client-1', + tools: [], + customizations: [], + }, + }); + + assert.deepStrictEqual({ + publishedSessions: spy.setSessionServersCalls.map(call => call.session.toString()), + }, { + publishedSessions: [session.toString()], + }); + } finally { + await disposeAgent(agent); + } + }); + + test('disposeSession clears the session\'s MCP servers', async () => { + const spy = new SpyMcpHostService(); + const agent = createTestAgent(disposables, { + copilotClient: new TestCopilotClient([]), + mcpHostService: spy, + }); + try { + await agent.authenticate('https://api.github.com', 'token'); + + // Create a provisional session WITH an active client so the + // active-client lifecycle publishes a server set that dispose + // can later clear. + const session = AgentSession.uri('copilotcli', 'dispose-1'); + const result = await agent.createSession({ + session, + workingDirectory: URI.file('/workspace'), + activeClient: { + clientId: 'client-1', + tools: [], + customizations: [], + }, + }); + + const beforeDispose = spy.setSessionServersCalls.length; + await agent.disposeSession(result.session); + + const teardown = spy.setSessionServersCalls.slice(beforeDispose); + assert.deepStrictEqual( + teardown.map(c => ({ session: c.session.toString(), serverCount: c.servers.length })), + [{ session: result.session.toString(), serverCount: 0 }], + ); + } finally { + await disposeAgent(agent); + } + }); + + test('_resolveMcpServersForSdk waits for the server to reach Ready before resolving the endpoint', async () => { + const spy = new SpyMcpHostService(); + const agent = createTestAgent(disposables, { + copilotClient: new TestCopilotClient([]), + mcpHostService: spy, + }); + try { + const sessionUri = AgentSession.uri('copilotcli', 'wait-1'); + const def = makeServerDef('srv', 'cmd'); + + const endpointObs = observableValue('endpoint', undefined); + const summaryObs = observableValue('summary', { + resource: buildMcpServerUri(sessionUri, def.name).toString(), + label: 'stub', + status: { kind: McpServerStatusKind.Starting }, + }); + spy.registerStubServer(buildMcpServerUri(sessionUri, def.name), endpointObs, summaryObs); + + const internals = agent as unknown as ICopilotAgentInternals; + const promise = internals._resolveMcpServersForSdk(sessionUri, [def]); + + // The transport binds asynchronously: the endpoint is set first, + // then the upstream transitions to Ready. Until Ready, the SDK + // config must NOT include this server. + queueMicrotask(() => { + endpointObs.set(URI.parse('http://127.0.0.1:9999/mcp'), undefined); + summaryObs.set({ + resource: buildMcpServerUri(sessionUri, def.name).toString(), + label: 'stub', + status: { kind: McpServerStatusKind.Ready }, + }, undefined); + }); + + const result = await promise; + assert.deepStrictEqual(result, { + srv: { type: 'http', url: 'http://127.0.0.1:9999/mcp', tools: ['*'] }, + }); + } finally { + await disposeAgent(agent); + } + }); + + test('_resolveMcpServersForSdk drops servers stuck in non-Ready states', async () => { + const spy = new SpyMcpHostService(); + const agent = createTestAgent(disposables, { + copilotClient: new TestCopilotClient([]), + mcpHostService: spy, + }); + try { + const sessionUri = AgentSession.uri('copilotcli', 'wait-2'); + const def = makeServerDef('srv', 'cmd'); + + const endpointObs = observableValue('endpoint', URI.parse('http://127.0.0.1:9999/mcp')); + const summaryObs = observableValue('summary', { + resource: buildMcpServerUri(sessionUri, def.name).toString(), + label: 'stub', + status: { kind: McpServerStatusKind.AuthRequired, reason: McpAuthRequiredReason.Required, resource: { resource: 'https://x' } }, + }); + spy.registerStubServer(buildMcpServerUri(sessionUri, def.name), endpointObs, summaryObs); + + const internals = agent as unknown as ICopilotAgentInternals; + const result = await internals._resolveMcpServersForSdk(sessionUri, [def]); + assert.deepStrictEqual(result, {}); + } finally { + await disposeAgent(agent); + } + }); + }); }); + +class SpyMcpHostService implements IMcpHostService { + declare readonly _serviceBrand: undefined; + readonly setSessionServersCalls: { session: URI; servers: readonly IMcpServerDefinition[] }[] = []; + readonly sessionSummaries = new Map(); + + private readonly _stubs = new Map(); + + setSessionServers(session: URI, servers: readonly IMcpServerDefinition[]): readonly IMcpServerHandle[] { + this.setSessionServersCalls.push({ session, servers }); + return []; + } + getServer(resource: URI): IMcpServerHandle | undefined { + return this._stubs.get(resource.toString()); + } + getServerSummaries(session: URI): readonly McpServerSummary[] { + return this.sessionSummaries.get(session.toString()) ?? []; + } + async callMethod(): Promise<{ result: unknown }> { throw new Error('not used in test'); } + notify(): void { /* no-op */ } + setUpstreamDelegate(): IDisposable { return { dispose: () => { /* no-op */ } }; } + + registerStubServer(resource: URI, endpoint: IObservable, summary?: IObservable): IMcpServerHandle { + const summaryObs = summary ?? observableValue('summary', { + resource: resource.toString(), + label: 'stub', + status: { kind: McpServerStatusKind.Ready }, + }); + const handle: IMcpServerHandle = { + resource, + summary: summaryObs, + endpoint, + authenticate: async () => true, + callMethod: async () => { throw new Error('not used in test'); }, + notify: () => { /* no-op */ }, + getToolUiMeta: () => undefined, + getUiHostCapabilities: () => ({}), + dispose: () => { /* no-op */ }, + }; + this._stubs.set(resource.toString(), handle); + return handle; + } +} diff --git a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts index 40389a469357d..52de49b95e4e6 100644 --- a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts @@ -1456,6 +1456,7 @@ suite('CopilotAgentSession', () => { inputSchema: { type: 'object', properties: {} }, }], plugins: [], + mcpReadiness: {}, }; test('client tool handler waits for completion without emitting tool_ready', async () => { diff --git a/src/vs/platform/agentHost/test/node/copilotPluginConverters.test.ts b/src/vs/platform/agentHost/test/node/copilotPluginConverters.test.ts index 990a15bbe52bc..634f63a9baca9 100644 --- a/src/vs/platform/agentHost/test/node/copilotPluginConverters.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotPluginConverters.test.ts @@ -35,7 +35,10 @@ suite('copilotPluginConverters', () => { suite('toSdkMcpServers', () => { - test('converts local server definitions', () => { + const proxyEndpoint = (name: string) => URI.parse(`http://127.0.0.1:1/mcp/${name}`); + const resolveAll = (def: IMcpServerDefinition) => proxyEndpoint(def.name); + + test('emits an http entry pointing at the proxy endpoint for each resolved server', () => { const defs: IMcpServerDefinition[] = [{ name: 'test-server', uri: URI.file('/plugin'), @@ -48,20 +51,17 @@ suite('copilotPluginConverters', () => { }, }]; - const result = toSdkMcpServers(defs); + const result = toSdkMcpServers(defs, resolveAll); assert.deepStrictEqual(result, { 'test-server': { - type: 'local', - command: 'node', - args: ['server.js', '--port', '3000'], + type: 'http', + url: 'http://127.0.0.1:1/mcp/test-server', tools: ['*'], - env: { NODE_ENV: 'production', PORT: '3000' }, - cwd: '/workspace', }, }); }); - test('converts remote/http server definitions', () => { + test('emits the same http shape regardless of upstream transport', () => { const defs: IMcpServerDefinition[] = [{ name: 'remote-server', uri: URI.file('/plugin'), @@ -72,53 +72,44 @@ suite('copilotPluginConverters', () => { }, }]; - const result = toSdkMcpServers(defs); + const result = toSdkMcpServers(defs, resolveAll); assert.deepStrictEqual(result, { 'remote-server': { type: 'http', - url: 'https://example.com/mcp', + url: 'http://127.0.0.1:1/mcp/remote-server', tools: ['*'], - headers: { 'Authorization': 'Bearer token' }, }, }); }); test('handles empty definitions', () => { - const result = toSdkMcpServers([]); + const result = toSdkMcpServers([], resolveAll); assert.deepStrictEqual(result, {}); }); - test('omits optional fields when undefined', () => { + test('skips servers whose proxy endpoint is unresolved', () => { const defs: IMcpServerDefinition[] = [{ - name: 'minimal', + name: 'pending', uri: URI.file('/plugin'), - configuration: { - type: McpServerType.LOCAL, - command: 'echo', - }, + configuration: { type: McpServerType.LOCAL, command: 'echo' }, }]; - const result = toSdkMcpServers(defs); - assert.strictEqual(result['minimal'].type, 'local'); - assert.deepStrictEqual((result['minimal'] as { args?: string[] }).args, []); - assert.strictEqual(Object.hasOwn(result['minimal'], 'env'), false); - assert.strictEqual(Object.hasOwn(result['minimal'], 'cwd'), false); + const result = toSdkMcpServers(defs, () => undefined); + assert.deepStrictEqual(result, {}); }); - test('filters null values from env', () => { - const defs: IMcpServerDefinition[] = [{ - name: 'with-null-env', - uri: URI.file('/plugin'), - configuration: { - type: McpServerType.LOCAL, - command: 'test', - env: { KEEP: 'value', DROP: null as unknown as string }, - }, - }]; + test('only includes resolved entries when some endpoints are pending', () => { + const defs: IMcpServerDefinition[] = [ + { name: 'ready', uri: URI.file('/p1'), configuration: { type: McpServerType.LOCAL, command: 'a' } }, + { name: 'pending', uri: URI.file('/p2'), configuration: { type: McpServerType.LOCAL, command: 'b' } }, + { name: 'also-ready', uri: URI.file('/p3'), configuration: { type: McpServerType.REMOTE, url: 'https://x' } }, + ]; - const result = toSdkMcpServers(defs); - const env = (result['with-null-env'] as { env?: Record }).env; - assert.deepStrictEqual(env, { KEEP: 'value' }); + const result = toSdkMcpServers(defs, def => def.name === 'pending' ? undefined : proxyEndpoint(def.name)); + assert.deepStrictEqual(result, { + 'ready': { type: 'http', url: 'http://127.0.0.1:1/mcp/ready', tools: ['*'] }, + 'also-ready': { type: 'http', url: 'http://127.0.0.1:1/mcp/also-ready', tools: ['*'] }, + }); }); }); diff --git a/src/vs/platform/agentHost/test/node/mcpHost/mcpAppsRoundTrip.integrationTest.ts b/src/vs/platform/agentHost/test/node/mcpHost/mcpAppsRoundTrip.integrationTest.ts new file mode 100644 index 0000000000000..f81f1c8a28196 --- /dev/null +++ b/src/vs/platform/agentHost/test/node/mcpHost/mcpAppsRoundTrip.integrationTest.ts @@ -0,0 +1,389 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * End-to-end MCP Apps round-trip test. Wires together {@link McpHostServiceImpl}, + * a real {@link McpProxyFactory} (which binds a real localhost HTTP listener), + * a real {@link AgentHostStateManager}, and a fake in-memory MCP App upstream + * that simulates the wire behaviour of an MCP server speaking the + * `io.modelcontextprotocol/ui` extension. + * + * The single test in this suite exercises the full lifecycle the MCP Apps + * extension relies on, in the order it occurs at runtime: + * + * 1. `setSessionServers` registers the MCP server, the host kicks off + * proxy creation, and the per-(session, server) entry exposes an HTTP + * endpoint URI back to the SDK. + * + * 2. The SDK (simulated via `fetch`) issues `initialize` to the proxy + * endpoint. {@link McpAppsInitializeInjector} rewrites the outbound + * params to advertise `mcp.apps`. The fake upstream replies with its + * own Apps capability; {@link McpProxyRoute} captures it on the + * upstream's `upstreamCapabilities` observable, which is what + * {@link McpHostServiceImpl.sendMessage} consults when gating `ui/*` + * methods. + * + * 3. The fake upstream pushes a `ui/notifications/host-context-changed` + * notification. The proxy taps it and fires `onUpstreamMessage`, + * which causes the host service to dispatch `mcp/messageReceived` + * followed by an immediate `mcp/messageRemoved` (notifications have + * no response phase). + * + * 4. An AHP client invokes `mcpMessage` for `ui/some-method`. The host + * service's allowlist permits the call (client advertises `mcp.apps`, + * upstream advertises Apps), the request reaches the fake which + * auto-replies, and the result is propagated back as + * {@link McpMessageResult}. + * + * 5. The fake upstream pushes a server→client `ui/open-link` request. + * The proxy taps it; the host service mints a fresh messageId and + * dispatches `mcp/messageReceived` with the call. + * + * 6. The AHP client supplies a response via `deliverResponse` (the + * handler for the client-dispatchable `mcp/messageResponded` action). + * The proxy resolves the messageId back to the upstream's original + * JSON-RPC id and writes the response onto the upstream transport. + * The host service then dispatches `mcp/messageRemoved`. + * + * Run via `scripts/test-integration.sh` / `scripts\test-integration.bat`. + */ + +import * as assert from 'assert'; +import { timeout } from '../../../../../base/common/async.js'; +import { Emitter, type Event } from '../../../../../base/common/event.js'; +import { + isJsonRpcRequest, + isJsonRpcResponse, + type IJsonRpcRequest, + type IJsonRpcSuccessResponse, + type JsonRpcId, + type JsonRpcMessage, +} from '../../../../../base/common/jsonRpcProtocol.js'; +import { Disposable } from '../../../../../base/common/lifecycle.js'; +import { + observableValue, + type IObservable, + type ISettableObservable, +} from '../../../../../base/common/observable.js'; +import { URI } from '../../../../../base/common/uri.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; +import { IMcpServerDefinition } from '../../../../agentPlugins/common/pluginParsers.js'; +import { ILogger, NullLogService } from '../../../../log/common/log.js'; +import { McpServerType } from '../../../../mcp/common/mcpPlatformTypes.js'; +import { type ActionEnvelope } from '../../../common/state/sessionActions.js'; +import { + McpServerStatusKind, + type McpServerStatus, +} from '../../../common/state/protocol/state.js'; +import { AgentHostStateManager } from '../../../node/agentHostStateManager.js'; +import type { IMcpServerHandle, IUpstreamMcpNotification, IUpstreamMcpRequest, IUpstreamMcpResponse } from '../../../common/mcpHost/mcpHostService.js'; +import { McpHostServiceImpl } from '../../../node/mcpHost/mcpHostServiceImpl.js'; +import { McpProxyFactory } from '../../../node/mcpHost/mcpProxy.js'; +import type { IMcpUpstream, IMcpUpstreamCapabilities } from '../../../node/mcpHost/mcpUpstream.js'; + +// ----- Fake MCP App upstream ---------------------------------------------- + +const APPS_EXTENSION_KEY = 'io.modelcontextprotocol/ui'; + +const APPS_INITIALIZE_RESULT = { + protocolVersion: '2024-11-05', + capabilities: { + extensions: { + [APPS_EXTENSION_KEY]: { mimeTypes: ['text/html;profile=mcp-app'] }, + }, + }, + serverInfo: { name: 'fake-apps-server', version: '0.0.1' }, +} as const; + +/** + * In-memory {@link IMcpUpstream} that simulates an MCP server speaking the + * Apps extension. Auto-replies to `initialize` with the Apps capability, + * captures everything sent by the SDK/proxy, and exposes + * {@link FakeMcpAppsUpstream.pushNotification} and + * {@link FakeMcpAppsUpstream.pushRequest} so the test can drive + * server→client traffic. + */ +class FakeMcpAppsUpstream extends Disposable implements IMcpUpstream { + + private readonly _status: ISettableObservable = + observableValue('fake-upstream', { kind: McpServerStatusKind.Stopped }); + public readonly status: IObservable = this._status; + + private readonly _onMessage = this._register(new Emitter()); + public readonly onMessage: Event = this._onMessage.event; + + private readonly _upstreamCapabilities: ISettableObservable = + observableValue('fake-upstream-caps', undefined); + public readonly upstreamCapabilities: IObservable = this._upstreamCapabilities; + + /** Every JSON-RPC message the proxy/SDK sent to the upstream. */ + public readonly received: JsonRpcMessage[] = []; + + /** Auto-reply policy keyed by request method (used for `ui/*` calls). */ + private readonly _replies = new Map(); + + private _upstreamReqCounter = 0; + + public override dispose(): void { + super.dispose(); + } + + public async start(): Promise { + this._status.set({ kind: McpServerStatusKind.Starting }, undefined); + // Yield once so any observer that latched onto Starting can run + // before we transition to Ready — mirrors the real stdio upstream. + await Promise.resolve(); + const ready: McpServerStatus = { kind: McpServerStatusKind.Ready }; + this._status.set(ready, undefined); + return ready; + } + + public async send(message: JsonRpcMessage): Promise { + this.received.push(message); + + if (isJsonRpcRequest(message)) { + if (message.method === 'initialize') { + this._onMessage.fire({ + jsonrpc: '2.0', + id: message.id, + result: APPS_INITIALIZE_RESULT, + } satisfies IJsonRpcSuccessResponse); + return; + } + if (this._replies.has(message.method)) { + const result = this._replies.get(message.method); + this._onMessage.fire({ + jsonrpc: '2.0', + id: message.id, + result, + } satisfies IJsonRpcSuccessResponse); + } + } + } + + public setBearerToken(_token: string | undefined): void { /* noop */ } + + public setUpstreamCapabilities(caps: IMcpUpstreamCapabilities | undefined): void { + this._upstreamCapabilities.set(caps, undefined); + } + + // ---- Test driver helpers ---- + + public replyTo(method: string, result: unknown): void { + this._replies.set(method, result); + } + + public pushNotification(method: string, params?: unknown): void { + this._onMessage.fire({ jsonrpc: '2.0', method, params }); + } + + public pushRequest(method: string, params?: unknown): JsonRpcId { + const id: JsonRpcId = ++this._upstreamReqCounter; + this._onMessage.fire({ jsonrpc: '2.0', id, method, params } satisfies IJsonRpcRequest); + return id; + } + + public lastReceivedRequest(method: string): IJsonRpcRequest | undefined { + for (let i = this.received.length - 1; i >= 0; i--) { + const m = this.received[i]; + if (isJsonRpcRequest(m) && m.method === method) { + return m; + } + } + return undefined; + } + + public lastResponseTo(id: JsonRpcId): JsonRpcMessage | undefined { + for (let i = this.received.length - 1; i >= 0; i--) { + const m = this.received[i]; + if (isJsonRpcResponse(m) && m.id === id) { + return m; + } + } + return undefined; + } +} + +// ----- Test seam ---------------------------------------------------------- + +/** + * Subclass that swaps the real stdio/HTTP upstreams for a caller-supplied + * factory. The test uses this to inject {@link FakeMcpAppsUpstream} instead + * of the production transports. + */ +class TestMcpHostService extends McpHostServiceImpl { + constructor( + stateManager: AgentHostStateManager, + proxyFactory: McpProxyFactory, + logService: NullLogService, + private readonly _upstreamFactory: () => IMcpUpstream, + ) { + super(stateManager, proxyFactory, logService); + } + protected override _createUpstream(_def: IMcpServerDefinition, _logger: ILogger): IMcpUpstream { + return this._upstreamFactory(); + } +} + +// ----- Helpers ------------------------------------------------------------ + +async function fetchPostJson(endpoint: URI, body: unknown): Promise<{ status: number; json: unknown }> { + const response = await fetch(endpoint.toString(true), { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + }); + const text = await response.text(); + return { status: response.status, json: text.length > 0 ? JSON.parse(text) : undefined }; +} + +async function waitForEndpoint(handle: IMcpServerHandle): Promise { + for (let i = 0; i < 200; i++) { + const ep = handle.endpoint.get(); + if (ep) { + return ep; + } + await timeout(10); + } + throw new Error('Timed out waiting for proxy endpoint'); +} + +async function flushMicrotasks(): Promise { + // Three turns is enough for autorun + the route's onUpstreamMessage tap + // to settle through the host service's dispatch path. + await Promise.resolve(); + await Promise.resolve(); + await Promise.resolve(); +} + +// ----- Test --------------------------------------------------------------- + +suite('MCP Apps round-trip (integration)', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + test('full ui/* lifecycle: initialize, notification, client→server call, server→client request', async () => { + const logService = new NullLogService(); + const stateManager = new AgentHostStateManager(logService); + const proxyFactory = new McpProxyFactory(logService); + const fakeUpstream = new FakeMcpAppsUpstream(); + fakeUpstream.replyTo('ui/some-method', { ok: true, payload: 42 }); + + const hostService = new TestMcpHostService(stateManager, proxyFactory, logService, () => fakeUpstream); + + const envelopes: ActionEnvelope[] = []; + const envelopeSubscription = stateManager.onDidEmitEnvelope(e => envelopes.push(e)); + + // Install an upstream delegate that records every upstream-originated + // request/notification for later assertions. + const receivedNotifications: IUpstreamMcpNotification[] = []; + const receivedRequests: IUpstreamMcpRequest[] = []; + let nextRequestOutcome: IUpstreamMcpResponse = { result: { handled: true } }; + const delegateRegistration = hostService.setUpstreamDelegate({ + handleUpstreamRequest: async (request) => { + receivedRequests.push(request); + return nextRequestOutcome; + }, + handleUpstreamNotification: (notification) => { + receivedNotifications.push(notification); + }, + }); + + try { + // ---- 1. Register the MCP server. ---------------------------- + const session = URI.parse('copilot:/00000000-0000-0000-0000-000000000001'); + const def: IMcpServerDefinition = { + name: 'apps-server', + uri: URI.parse('inmemory:/apps-server'), + configuration: { type: McpServerType.LOCAL, command: 'unused', args: [] }, + }; + const [handle] = hostService.setSessionServers(session, [def]); + const endpoint = await waitForEndpoint(handle); + + // ---- 2. SDK posts `initialize` through the proxy. ----------- + const initResp = await fetchPostJson(endpoint, { + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params: { + protocolVersion: '2024-11-05', + capabilities: { sampling: {} }, + clientInfo: { name: 'test-sdk', version: '0.0.0' }, + }, + }); + + const initResult = (initResp.json as { result?: { capabilities?: { extensions?: Record } } } | undefined)?.result; + const sentInitParams = (fakeUpstream.lastReceivedRequest('initialize')?.params ?? {}) as { + capabilities?: { extensions?: Record; sampling?: unknown }; + }; + + const phase2Summary = { + httpStatus: initResp.status, + upstreamSeesAppsExtension: !!sentInitParams.capabilities?.extensions?.[APPS_EXTENSION_KEY], + upstreamSeesPreservedSampling: sentInitParams.capabilities?.sampling, + sdkSeesAppsExtension: !!initResult?.capabilities?.extensions?.[APPS_EXTENSION_KEY], + upstreamCapabilitiesObservable: !!fakeUpstream.upstreamCapabilities.get()?.extensions?.[APPS_EXTENSION_KEY], + }; + assert.deepStrictEqual(phase2Summary, { + httpStatus: 200, + upstreamSeesAppsExtension: true, + upstreamSeesPreservedSampling: {}, + sdkSeesAppsExtension: true, + upstreamCapabilitiesObservable: true, + }); + + // ---- 3. Server→client notification. ------------------------- + receivedNotifications.length = 0; + envelopes.length = 0; + fakeUpstream.pushNotification('ui/notifications/host-context-changed', { theme: 'dark' }); + await flushMicrotasks(); + + assert.deepStrictEqual({ + notifications: receivedNotifications.map(n => ({ method: n.method, params: n.params })), + envelopeKinds: envelopes.map(e => e.action.type), + }, { + notifications: [{ method: 'ui/notifications/host-context-changed', params: { theme: 'dark' } }], + envelopeKinds: [], + }); + + // ---- 4. Client→server `ui/*` call via mcpMethodCall. -------- + const clientResult = await hostService.callMethod({ + server: handle.resource.toString(), + method: 'ui/some-method', + params: { foo: 1 }, + }); + + const sentUiRequest = fakeUpstream.lastReceivedRequest('ui/some-method'); + assert.deepStrictEqual({ + clientResult, + sentMethod: sentUiRequest?.method, + sentParams: sentUiRequest?.params, + }, { + clientResult: { result: { ok: true, payload: 42 } }, + sentMethod: 'ui/some-method', + sentParams: { foo: 1 }, + }); + + // ---- 5. Server→client request (`ui/open-link`). ------------- + receivedRequests.length = 0; + nextRequestOutcome = { result: { handled: true } }; + const upstreamReqId = fakeUpstream.pushRequest('ui/open-link', { url: 'https://example.com' }); + await flushMicrotasks(); + + assert.deepStrictEqual({ + requests: receivedRequests.map(r => ({ method: r.method, params: r.params })), + upstreamResponse: fakeUpstream.lastResponseTo(upstreamReqId), + }, { + requests: [{ method: 'ui/open-link', params: { url: 'https://example.com' } }], + upstreamResponse: { jsonrpc: '2.0', id: upstreamReqId, result: { handled: true } }, + }); + } finally { + delegateRegistration.dispose(); + envelopeSubscription.dispose(); + hostService.dispose(); + proxyFactory.dispose(); + stateManager.dispose(); + } + }); +}); diff --git a/src/vs/platform/agentHost/test/node/mcpHost/mcpAuthChallengeParser.test.ts b/src/vs/platform/agentHost/test/node/mcpHost/mcpAuthChallengeParser.test.ts new file mode 100644 index 0000000000000..a81d6588e0c36 --- /dev/null +++ b/src/vs/platform/agentHost/test/node/mcpHost/mcpAuthChallengeParser.test.ts @@ -0,0 +1,140 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as assert from 'assert'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; +import { McpAuthRequiredReason, McpServerStatusKind, type ProtectedResourceMetadata } from '../../../common/state/protocol/state.js'; +import { buildAuthRequiredStatus, parseWwwAuthenticate } from '../../../node/mcpHost/mcpAuthChallengeParser.js'; + +suite('mcpAuthChallengeParser', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + const emptyChallenge = { + scopes: undefined, + error: undefined, + errorDescription: undefined, + resourceMetadataUrl: undefined, + }; + + test('parseWwwAuthenticate returns all-undefined for missing or empty header', () => { + assert.deepStrictEqual(parseWwwAuthenticate(undefined), emptyChallenge); + assert.deepStrictEqual(parseWwwAuthenticate(''), emptyChallenge); + }); + + test('parseWwwAuthenticate returns all-undefined for non-Bearer scheme', () => { + assert.deepStrictEqual(parseWwwAuthenticate('Basic realm="api"'), emptyChallenge); + }); + + test('parseWwwAuthenticate accepts Bearer with only realm', () => { + assert.deepStrictEqual(parseWwwAuthenticate('Bearer realm="api"'), emptyChallenge); + }); + + test('parseWwwAuthenticate parses full challenge with quoted values', () => { + const header = 'Bearer error="invalid_token", error_description="token expired", scope="read:user user:email", resource_metadata="https://example.com/.well-known/oauth-protected-resource"'; + assert.deepStrictEqual(parseWwwAuthenticate(header), { + scopes: ['read:user', 'user:email'], + error: 'invalid_token', + errorDescription: 'token expired', + resourceMetadataUrl: 'https://example.com/.well-known/oauth-protected-resource', + }); + }); + + test('parseWwwAuthenticate parses unquoted token values', () => { + assert.deepStrictEqual(parseWwwAuthenticate('Bearer error=insufficient_scope'), { + scopes: undefined, + error: 'insufficient_scope', + errorDescription: undefined, + resourceMetadataUrl: undefined, + }); + }); + + const resource: ProtectedResourceMetadata = { + resource: 'https://api.example.com', + authorization_servers: ['https://auth.example.com'], + scopes_supported: ['read', 'write'], + }; + + test('buildAuthRequiredStatus 401 without prior token → Required', () => { + assert.deepStrictEqual( + buildAuthRequiredStatus({ + httpStatus: 401, + challenge: emptyChallenge, + resource, + hadPriorToken: false, + }), + { + kind: McpServerStatusKind.AuthRequired, + reason: McpAuthRequiredReason.Required, + resource, + requiredScopes: ['read', 'write'], + }, + ); + }); + + test('buildAuthRequiredStatus 401 with prior token → Expired', () => { + const status = buildAuthRequiredStatus({ + httpStatus: 401, + challenge: { ...emptyChallenge, scopes: ['only-this'] }, + resource, + hadPriorToken: true, + }); + assert.deepStrictEqual(status, { + kind: McpServerStatusKind.AuthRequired, + reason: McpAuthRequiredReason.Expired, + resource, + requiredScopes: ['only-this'], + }); + }); + + test('buildAuthRequiredStatus 403 + insufficient_scope → InsufficientScope', () => { + const status = buildAuthRequiredStatus({ + httpStatus: 403, + challenge: { + scopes: ['admin'], + error: 'insufficient_scope', + errorDescription: 'need admin', + resourceMetadataUrl: undefined, + }, + resource, + hadPriorToken: true, + }); + assert.deepStrictEqual(status, { + kind: McpServerStatusKind.AuthRequired, + reason: McpAuthRequiredReason.InsufficientScope, + resource, + requiredScopes: ['admin'], + description: 'need admin', + }); + }); + + test('buildAuthRequiredStatus 403 + other error → Required (conservative)', () => { + const status = buildAuthRequiredStatus({ + httpStatus: 403, + challenge: { ...emptyChallenge, error: 'invalid_token' }, + resource, + hadPriorToken: true, + }); + assert.deepStrictEqual(status, { + kind: McpServerStatusKind.AuthRequired, + reason: McpAuthRequiredReason.Required, + resource, + requiredScopes: ['read', 'write'], + }); + }); + + test('buildAuthRequiredStatus omits requiredScopes when absent everywhere', () => { + const status = buildAuthRequiredStatus({ + httpStatus: 401, + challenge: emptyChallenge, + resource: { resource: 'https://api.example.com' }, + hadPriorToken: false, + }); + assert.deepStrictEqual(status, { + kind: McpServerStatusKind.AuthRequired, + reason: McpAuthRequiredReason.Required, + resource: { resource: 'https://api.example.com' }, + }); + }); +}); diff --git a/src/vs/platform/agentHost/test/node/mcpHost/mcpHostServiceImpl.test.ts b/src/vs/platform/agentHost/test/node/mcpHost/mcpHostServiceImpl.test.ts new file mode 100644 index 0000000000000..15ed78a4cfa4e --- /dev/null +++ b/src/vs/platform/agentHost/test/node/mcpHost/mcpHostServiceImpl.test.ts @@ -0,0 +1,487 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as assert from 'assert'; +import { DeferredPromise, timeout } from '../../../../../base/common/async.js'; +import { JsonRpcMessage, isJsonRpcRequest } from '../../../../../base/common/jsonRpcProtocol.js'; +import { DisposableStore } from '../../../../../base/common/lifecycle.js'; +import { observableValue, type ISettableObservable } from '../../../../../base/common/observable.js'; +import { URI } from '../../../../../base/common/uri.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; +import { NullLogService, ILogger } from '../../../../log/common/log.js'; +import { IMcpServerDefinition } from '../../../../agentPlugins/common/pluginParsers.js'; +import { + IMcpRemoteServerConfiguration, + IMcpStdioServerConfiguration, + McpServerType, +} from '../../../../mcp/common/mcpPlatformTypes.js'; +import { ActionType, type ActionEnvelope } from '../../../common/state/sessionActions.js'; +import { SessionStatus } from '../../../common/state/sessionState.js'; +import { + McpServerStatusKind, + type McpServerStatus, +} from '../../../common/state/protocol/state.js'; +import { JsonRpcErrorCodes } from '../../../common/state/protocol/errors.js'; +import { ProtocolError } from '../../../common/state/sessionProtocol.js'; +import { AgentHostStateManager } from '../../../node/agentHostStateManager.js'; +import { McpAppsInitializeInjector } from '../../../node/mcpHost/mcpInitializeInjector.js'; +import { McpHostServiceImpl } from '../../../node/mcpHost/mcpHostServiceImpl.js'; +import { buildMcpServerUri } from '../../../common/state/mcpServerUri.js'; +import type { IMcpUpstream, IMcpUpstreamCapabilities } from '../../../node/mcpHost/mcpUpstream.js'; +import type { IMcpProxy, IMcpProxyFactory, IMcpProxyOptions } from '../../../node/mcpHost/mcpProxy.js'; +import type { IUpstreamRequestOutcome } from '../../../node/mcpHost/mcpProxyRoute.js'; + +// ----- Mocks --------------------------------------------------------------- + +class StubUpstream implements IMcpUpstream { + public readonly status = { read: () => ({ kind: McpServerStatusKind.Stopped }) } as unknown as IMcpUpstream['status']; + public readonly onMessage = (() => () => { /* unsubscribe */ }) as unknown as IMcpUpstream['onMessage']; + private readonly _upstreamCapabilities: ISettableObservable = + observableValue('stub-caps', undefined); + public readonly upstreamCapabilities = this._upstreamCapabilities; + public startCalls = 0; + public disposed = false; + + public async start(): Promise { + this.startCalls++; + return { kind: McpServerStatusKind.Ready }; + } + public async send(_message: JsonRpcMessage): Promise { /* noop */ } + public setBearerToken(_token: string | undefined): void { /* noop */ } + public setUpstreamCapabilities(caps: IMcpUpstreamCapabilities | undefined): void { + this._upstreamCapabilities.set(caps, undefined); + } + public dispose(): void { this.disposed = true; } +} + +class StubProxy implements IMcpProxy { + public readonly resource: URI; + public readonly endpoint: URI; + public disposed = false; + public readonly sentMessages: JsonRpcMessage[] = []; + + /** + * Synchronous reply for the next `sendClientMessage` call. If unset, + * `sendClientMessage` returns a pending DeferredPromise the test must + * resolve manually. + */ + public nextResponse: JsonRpcMessage | undefined; + + constructor(public readonly options: IMcpProxyOptions) { + this.resource = options.resource; + this.endpoint = URI.parse(`http://127.0.0.1:1/mcp/${encodeURIComponent(options.resource.toString())}`); + } + + public async authenticate(_resource: string, _token: string): Promise { + return true; + } + + public sendClientMessage(message: JsonRpcMessage): Promise { + this.sentMessages.push(message); + if (this.nextResponse !== undefined) { + const response = this.nextResponse; + this.nextResponse = undefined; + if (isJsonRpcRequest(message)) { + return Promise.resolve({ ...(response as { id?: unknown }), id: message.id } as JsonRpcMessage); + } + return Promise.resolve(response); + } + return Promise.resolve(undefined); + } + + public getToolUiMeta(): undefined { + return undefined; + } + + public getUiHostCapabilities() { + return {}; + } + + public dispose(): void { this.disposed = true; } +} + +interface IRecordedCreate { + readonly options: IMcpProxyOptions; + readonly upstream: StubUpstream; +} + +class StubProxyFactory implements IMcpProxyFactory { + public readonly _serviceBrand: undefined; + public readonly created: IRecordedCreate[] = []; + public readonly proxies: StubProxy[] = []; + public failNextCreate: Error | undefined; + /** When set, `create()` blocks until this is resolved. */ + public pendingCreate: DeferredPromise | undefined; + + public async create(options: IMcpProxyOptions): Promise { + if (this.pendingCreate) { + await this.pendingCreate.p; + } + if (this.failNextCreate) { + const err = this.failNextCreate; + this.failNextCreate = undefined; + throw err; + } + const upstream = options.upstream as StubUpstream; + this.created.push({ options, upstream }); + const proxy = new StubProxy(options); + this.proxies.push(proxy); + return proxy; + } +} + +/** + * Subclass that swaps the real stdio/HTTP upstreams for {@link StubUpstream}. + */ +class TestMcpHostService extends McpHostServiceImpl { + public readonly createdUpstreams: StubUpstream[] = []; + protected override _createUpstream(_def: IMcpServerDefinition, _logger: ILogger): IMcpUpstream { + const upstream = new StubUpstream(); + this.createdUpstreams.push(upstream); + return upstream; + } +} + +// ----- Helpers -------------------------------------------------------------- + +function stdioDef(name: string, command = 'node', args: readonly string[] = []): IMcpServerDefinition { + const configuration: IMcpStdioServerConfiguration = { + type: McpServerType.LOCAL, + command, + args, + }; + return { name, configuration, uri: URI.parse('file:///plugins/test') }; +} + +function remoteDef(name: string, url: string): IMcpServerDefinition { + const configuration: IMcpRemoteServerConfiguration = { + type: McpServerType.REMOTE, + url, + }; + return { name, configuration, uri: URI.parse('file:///plugins/test') }; +} + +interface IHarness { + readonly service: TestMcpHostService; + readonly factory: StubProxyFactory; + readonly stateManager: AgentHostStateManager; + readonly envelopes: ActionEnvelope[]; + readonly session: URI; +} + +function setupHarness(disposables: DisposableStore): IHarness { + const logService = new NullLogService(); + const stateManager = disposables.add(new AgentHostStateManager(logService)); + const factory = new StubProxyFactory(); + const service = disposables.add(new TestMcpHostService(stateManager, factory, logService)); + const envelopes: ActionEnvelope[] = []; + const session = URI.parse('copilot:/test-session'); + stateManager.createSession({ + resource: session.toString(), + provider: 'copilot', + title: 'Test', + status: SessionStatus.Idle, + createdAt: Date.now(), + modifiedAt: Date.now(), + }); + disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e))); + return { + service, + factory, + stateManager, + envelopes, + session, + }; +} + +/** Wait for the proxy factory to record at least `count` creates. */ +async function waitForCreate(factory: StubProxyFactory, count = 1): Promise { + for (let i = 0; i < 50; i++) { + if (factory.created.length >= count) { + return; + } + await timeout(0); + } + throw new Error(`StubProxyFactory.create not called ${count} times after timeout`); +} + +// ----- Tests --------------------------------------------------------------- + +suite('McpHostServiceImpl', () => { + + let disposables: DisposableStore; + + setup(() => { + disposables = new DisposableStore(); + }); + + teardown(() => { + disposables.dispose(); + }); + + ensureNoDisposablesAreLeakedInTestSuite(); + + test('setSessionServers emits McpServerAdded with Starting status', async () => { + const h = setupHarness(disposables); + const def = stdioDef('foo'); + + const handles = h.service.setSessionServers(h.session, [def]); + await waitForCreate(h.factory); + + const expectedResource = buildMcpServerUri(h.session, 'foo').toString(); + const addedEnvelopes = h.envelopes.filter(e => e.action.type === ActionType.McpServerAdded); + + assert.deepStrictEqual({ + handleCount: handles.length, + handleResource: handles[0]?.resource.toString(), + added: addedEnvelopes.map(e => e.action.type === ActionType.McpServerAdded ? e.action.server : null), + }, { + handleCount: 1, + handleResource: expectedResource, + added: [{ + resource: expectedResource, + label: 'foo', + status: { kind: McpServerStatusKind.Starting }, + }], + }); + }); + + test('setSessionServers records summaries without emitting actions before session state exists', async () => { + const logService = new NullLogService(); + const stateManager = disposables.add(new AgentHostStateManager(logService)); + const factory = new StubProxyFactory(); + const service = disposables.add(new TestMcpHostService(stateManager, factory, logService)); + const envelopes: ActionEnvelope[] = []; + disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e))); + const session = URI.parse('copilot:/pre-state'); + + service.setSessionServers(session, [stdioDef('foo')]); + await waitForCreate(factory); + factory.created[0].options.onStateChange({ kind: McpServerStatusKind.Ready }); + + assert.deepStrictEqual({ + envelopes: envelopes.map(e => e.action.type), + summaries: service.getServerSummaries(session), + }, { + envelopes: [], + summaries: [{ + resource: buildMcpServerUri(session, 'foo').toString(), + label: 'foo', + status: { kind: McpServerStatusKind.Ready }, + }], + }); + }); + + test('successful proxy.create flips status to Ready via onStateChange', async () => { + const h = setupHarness(disposables); + h.service.setSessionServers(h.session, [stdioDef('foo')]); + + await waitForCreate(h.factory); + // Simulate the proxy emitting Ready (real proxy autorun does this + // when upstream observable transitions). + h.factory.created[0].options.onStateChange({ kind: McpServerStatusKind.Ready }); + + const statusChanges = h.envelopes + .filter(e => e.action.type === ActionType.McpServerStatusChanged) + .map(e => e.action.type === ActionType.McpServerStatusChanged ? e.action.status : null); + assert.deepStrictEqual(statusChanges.at(-1), { kind: McpServerStatusKind.Ready }); + }); + + test('failed proxy.create flips status to Error', async () => { + const h = setupHarness(disposables); + h.factory.failNextCreate = new Error('listener bind failed'); + h.service.setSessionServers(h.session, [stdioDef('foo')]); + + // Wait for the rejection to propagate. + for (let i = 0; i < 50; i++) { + if (h.envelopes.some(e => e.action.type === ActionType.McpServerStatusChanged)) { break; } + await timeout(0); + } + + const errorStatus = h.envelopes + .filter(e => e.action.type === ActionType.McpServerStatusChanged) + .map(e => e.action.type === ActionType.McpServerStatusChanged ? e.action.status : null) + .at(-1); + assert.deepStrictEqual(errorStatus, { + kind: McpServerStatusKind.Error, + error: { errorType: 'proxyCreateFailed', message: 'listener bind failed' }, + }); + }); + + test('setSessionServers([]) removes the entry and disposes the proxy', async () => { + const h = setupHarness(disposables); + h.service.setSessionServers(h.session, [stdioDef('foo')]); + await waitForCreate(h.factory); + + h.envelopes.length = 0; + h.service.setSessionServers(h.session, []); + + assert.deepStrictEqual({ + removedCount: h.envelopes.filter(e => e.action.type === ActionType.McpServerRemoved).length, + proxyDisposed: h.factory.proxies[0].disposed, + }, { + removedCount: 1, + proxyDisposed: true, + }); + }); + + test('setSessionServers reconfigures (same id, different config) → remove + add', async () => { + const h = setupHarness(disposables); + h.service.setSessionServers(h.session, [stdioDef('foo', 'node', ['v1'])]); + await waitForCreate(h.factory); + + h.envelopes.length = 0; + h.service.setSessionServers(h.session, [stdioDef('foo', 'node', ['v2'])]); + await waitForCreate(h.factory, 2); + + assert.deepStrictEqual( + h.envelopes.map(e => e.action.type), + [ActionType.McpServerRemoved, ActionType.McpServerAdded], + ); + }); + + test('setSessionServers no-op (same id, same config) emits nothing', async () => { + const h = setupHarness(disposables); + const def = stdioDef('foo'); + h.service.setSessionServers(h.session, [def]); + await waitForCreate(h.factory); + + h.envelopes.length = 0; + const handles = h.service.setSessionServers(h.session, [stdioDef('foo')]); + + assert.deepStrictEqual({ + envelopes: h.envelopes, + handleCount: handles.length, + createCount: h.factory.created.length, + }, { + envelopes: [], + handleCount: 1, + createCount: 1, + }); + }); + + test('getServer returns undefined for unknown resources', () => { + const h = setupHarness(disposables); + assert.strictEqual(h.service.getServer(URI.parse('mcp:/no-such/server')), undefined); + }); + + test('callMethod for unknown server throws ProtocolError(InvalidParams)', async () => { + const h = setupHarness(disposables); + const unknown = buildMcpServerUri(h.session, 'missing'); + await assert.rejects( + () => h.service.callMethod( + { server: unknown.toString(), method: 'tools/list' }, + ), + (err: Error) => err instanceof ProtocolError && err.code === JsonRpcErrorCodes.InvalidParams, + ); + }); + + test('callMethod for known server forwards to proxy.sendClientMessage', async () => { + const h = setupHarness(disposables); + h.service.setSessionServers(h.session, [stdioDef('foo')]); + await waitForCreate(h.factory); + + const proxy = h.factory.proxies[0]; + proxy.nextResponse = { jsonrpc: '2.0', id: 'replaced', result: { ok: true } } as JsonRpcMessage; + const result = await h.service.callMethod( + { server: buildMcpServerUri(h.session, 'foo').toString(), method: 'tools/list' }, + ); + + assert.deepStrictEqual({ + result, + sentCount: proxy.sentMessages.length, + method: (proxy.sentMessages[0] as { method?: string }).method, + }, { + result: { result: { ok: true } }, + sentCount: 1, + method: 'tools/list', + }); + }); + + test('notify for known server sends a JSON-RPC notification through the proxy', async () => { + const h = setupHarness(disposables); + h.service.setSessionServers(h.session, [stdioDef('foo')]); + await waitForCreate(h.factory); + + const proxy = h.factory.proxies[0]; + h.service.notify({ + server: buildMcpServerUri(h.session, 'foo').toString(), + method: 'notifications/message', + params: { level: 'info', data: 'hi' }, + }); + + assert.strictEqual(proxy.sentMessages.length, 1); + assert.deepStrictEqual(proxy.sentMessages[0], { + jsonrpc: '2.0', + method: 'notifications/message', + params: { level: 'info', data: 'hi' }, + }); + }); + + test('proxy is created with an McpAppsInitializeInjector', async () => { + const h = setupHarness(disposables); + h.service.setSessionServers(h.session, [stdioDef('foo')]); + await waitForCreate(h.factory); + + assert.ok( + h.factory.created[0].options.initializeInjector instanceof McpAppsInitializeInjector, + 'expected initializeInjector to be McpAppsInitializeInjector', + ); + }); + + test('upstream request is routed to the installed upstream delegate', async () => { + const h = setupHarness(disposables); + h.service.setSessionServers(h.session, [stdioDef('foo')]); + await waitForCreate(h.factory); + + let receivedMethod: string | undefined; + h.service.setUpstreamDelegate({ + handleUpstreamRequest: async (request) => { + receivedMethod = request.method; + return { result: { ok: true } }; + }, + handleUpstreamNotification: () => { /* unused */ }, + }); + + const outcome: IUpstreamRequestOutcome = await h.factory.created[0].options.onUpstreamRequest('sampling/createMessage', { messages: [] }); + + assert.deepStrictEqual({ method: receivedMethod, outcome }, { + method: 'sampling/createMessage', + outcome: { result: { ok: true } }, + }); + }); + + test('upstream request without an installed delegate yields MethodNotFound', async () => { + const h = setupHarness(disposables); + h.service.setSessionServers(h.session, [stdioDef('foo')]); + await waitForCreate(h.factory); + + const outcome = await h.factory.created[0].options.onUpstreamRequest('sampling/createMessage', { messages: [] }); + assert.strictEqual(outcome.error?.code, JsonRpcErrorCodes.MethodNotFound); + }); + + test('upstream notification is forwarded to the installed delegate', async () => { + const h = setupHarness(disposables); + h.service.setSessionServers(h.session, [stdioDef('foo')]); + await waitForCreate(h.factory); + + let receivedMethod: string | undefined; + h.service.setUpstreamDelegate({ + handleUpstreamRequest: async () => ({ result: {} }), + handleUpstreamNotification: (notification) => { receivedMethod = notification.method; }, + }); + + h.factory.created[0].options.onUpstreamNotification('notifications/tools/list_changed', undefined); + assert.strictEqual(receivedMethod, 'notifications/tools/list_changed'); + }); + + test('remote MCP definition routes through HTTP upstream', async () => { + const h = setupHarness(disposables); + h.service.setSessionServers(h.session, [remoteDef('rem', 'https://example.com/mcp')]); + await waitForCreate(h.factory); + + assert.strictEqual(h.factory.created.length, 1); + }); +}); diff --git a/src/vs/platform/agentHost/test/node/mcpHost/mcpHttpUpstream.test.ts b/src/vs/platform/agentHost/test/node/mcpHost/mcpHttpUpstream.test.ts new file mode 100644 index 0000000000000..abc6be18587d2 --- /dev/null +++ b/src/vs/platform/agentHost/test/node/mcpHost/mcpHttpUpstream.test.ts @@ -0,0 +1,591 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as assert from 'assert'; +import type { JsonRpcMessage } from '../../../../../base/common/jsonRpcProtocol.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; +import { NullLogger } from '../../../../log/common/log.js'; +import { McpServerType, type IMcpRemoteServerConfiguration } from '../../../../mcp/common/mcpPlatformTypes.js'; +import { McpAuthRequiredReason, McpServerStatusKind } from '../../../common/state/protocol/state.js'; +import { McpHttpUpstream, type HttpFetch, type IHttpResponse } from '../../../node/mcpHost/mcpHttpUpstream.js'; + +interface IFetchCall { + url: string; + method: string; + headers: Record; + body?: string; +} + +interface IFetchResponseSpec { + status: number; + headers?: Record; + body?: string; + throw?: Error; +} + +function makeResponse(spec: IFetchResponseSpec): IHttpResponse { + const headers = spec.headers ?? {}; + return { + status: spec.status, + headers: { + get(name: string): string | null { + const lower = name.toLowerCase(); + for (const [k, v] of Object.entries(headers)) { + if (k.toLowerCase() === lower) { + return v; + } + } + return null; + }, + }, + text: async () => spec.body ?? '', + }; +} + +function makeFetch(responses: (IFetchResponseSpec | ((url: string) => IFetchResponseSpec))[]): { fetch: HttpFetch; calls: IFetchCall[] } { + const calls: IFetchCall[] = []; + let i = 0; + const fetch: HttpFetch = async (url, init) => { + calls.push({ url, method: init.method, headers: { ...init.headers }, body: init.body }); + const next = responses[i++]; + if (!next) { + throw new Error(`unexpected fetch call to ${url}`); + } + const spec = typeof next === 'function' ? next(url) : next; + if (spec.throw) { + throw spec.throw; + } + return makeResponse(spec); + }; + return { fetch, calls }; +} + +interface IRecordedLog { + level: 'info' | 'warn' | 'error'; + message: string; +} + +class RecordingLogger extends NullLogger { + public readonly records: IRecordedLog[] = []; + override info(message: string): void { this.records.push({ level: 'info', message }); } + override warn(message: string): void { this.records.push({ level: 'warn', message }); } + override error(message: string | Error): void { + this.records.push({ level: 'error', message: message instanceof Error ? message.message : message }); + } +} + +const baseConfig: IMcpRemoteServerConfiguration = { + type: McpServerType.REMOTE, + url: 'https://mcp.example.com/v1', +}; + +suite('McpHttpUpstream', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + test('start() posts initialize probe with content-type and config headers, transitions to Ready on 2xx', async () => { + const { fetch, calls } = makeFetch([{ status: 200 }]); + const config: IMcpRemoteServerConfiguration = { ...baseConfig, headers: { 'X-Custom': 'val' } }; + const upstream = new McpHttpUpstream({ config, logger: new NullLogger(), fetch }); + try { + const status = await upstream.start(); + const body = calls[0].body ? JSON.parse(calls[0].body) : undefined; + assert.deepStrictEqual({ + status, + callCount: calls.length, + url: calls[0].url, + method: calls[0].method, + contentType: calls[0].headers['Content-Type'], + custom: calls[0].headers['X-Custom'], + bodyMethod: body?.method, + bodyId: body?.id, + hasAuth: calls[0].headers.hasOwnProperty('Authorization'), + }, { + status: { kind: McpServerStatusKind.Ready }, + callCount: 1, + url: 'https://mcp.example.com/v1', + method: 'POST', + contentType: 'application/json', + custom: 'val', + bodyMethod: 'initialize', + bodyId: 0, + hasAuth: false, + }); + } finally { + upstream.dispose(); + } + }); + + test('401 with resource_metadata fetches it and emits AuthRequired (Required, no prior token)', async () => { + const metadata = { resource: 'https://mcp.example.com/v1', authorization_servers: ['https://auth.example.com'] }; + const { fetch, calls } = makeFetch([ + { + status: 401, + headers: { 'WWW-Authenticate': 'Bearer error="invalid_token", resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource"' }, + }, + { + status: 200, + body: JSON.stringify(metadata), + }, + ]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + const status = await upstream.start(); + assert.deepStrictEqual({ + status, + secondCallUrl: calls[1]?.url, + secondCallMethod: calls[1]?.method, + }, { + status: { + kind: McpServerStatusKind.AuthRequired, + reason: McpAuthRequiredReason.Required, + resource: metadata, + }, + secondCallUrl: 'https://mcp.example.com/.well-known/oauth-protected-resource', + secondCallMethod: 'GET', + }); + } finally { + upstream.dispose(); + } + }); + + test('401 after a prior setBearerToken yields AuthRequired (Expired)', async () => { + const metadata = { resource: 'https://mcp.example.com/v1', authorization_servers: ['https://auth.example.com'] }; + const { fetch } = makeFetch([ + { + status: 401, + headers: { 'WWW-Authenticate': 'Bearer resource_metadata="https://mcp.example.com/.well-known/m"' }, + }, + { status: 200, body: JSON.stringify(metadata) }, + ]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + upstream.setBearerToken('prior-token'); + const status = await upstream.start(); + assert.strictEqual(status.kind, McpServerStatusKind.AuthRequired); + assert.strictEqual(status.kind === McpServerStatusKind.AuthRequired && status.reason, McpAuthRequiredReason.Expired); + } finally { + upstream.dispose(); + } + }); + + test('reclassifies as Required after token is cleared', async () => { + const metadata = { resource: 'https://mcp.example.com/v1', authorization_servers: ['https://auth.example.com'] }; + const { fetch } = makeFetch([ + // First start: token 't1' present, server replies 401 → Expired + { + status: 401, + headers: { 'WWW-Authenticate': 'Bearer resource_metadata="https://mcp.example.com/.well-known/m"' }, + }, + { status: 200, body: JSON.stringify(metadata) }, + // Second start (after token cleared): another 401 → Required + { + status: 401, + headers: { 'WWW-Authenticate': 'Bearer resource_metadata="https://mcp.example.com/.well-known/m"' }, + }, + { status: 200, body: JSON.stringify(metadata) }, + ]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + upstream.setBearerToken('t1'); + const first = await upstream.start(); + assert.strictEqual(first.kind === McpServerStatusKind.AuthRequired && first.reason, McpAuthRequiredReason.Expired); + + upstream.setBearerToken(undefined); + const second = await upstream.start(); + assert.strictEqual(second.kind === McpServerStatusKind.AuthRequired && second.reason, McpAuthRequiredReason.Required); + } finally { + upstream.dispose(); + } + }); + + test('rejects resource_metadata with a different origin', async () => { + const { fetch, calls } = makeFetch([ + { + status: 401, + headers: { 'WWW-Authenticate': 'Bearer resource_metadata="http://169.254.169.254/latest/meta-data/"' }, + }, + ]); + const logger = new RecordingLogger(); + const upstream = new McpHttpUpstream({ config: baseConfig, logger, fetch }); + try { + const status = await upstream.start(); + assert.deepStrictEqual({ + status, + callCount: calls.length, + warned: logger.records.some(r => r.level === 'warn' && /resource_metadata/i.test(r.message)), + }, { + status: { + kind: McpServerStatusKind.AuthRequired, + reason: McpAuthRequiredReason.Required, + resource: { resource: 'https://mcp.example.com/v1' }, + }, + callCount: 1, + warned: true, + }); + } finally { + upstream.dispose(); + } + }); + + test('rejects resource_metadata with a different scheme', async () => { + const { fetch, calls } = makeFetch([ + { + status: 401, + headers: { 'WWW-Authenticate': 'Bearer resource_metadata="http://mcp.example.com/.well-known/oauth-protected-resource"' }, + }, + ]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + const status = await upstream.start(); + assert.deepStrictEqual({ + status, + callCount: calls.length, + }, { + status: { + kind: McpServerStatusKind.AuthRequired, + reason: McpAuthRequiredReason.Required, + resource: { resource: 'https://mcp.example.com/v1' }, + }, + callCount: 1, + }); + } finally { + upstream.dispose(); + } + }); + + test('rejects file:/javascript:/data: schemes outright', async () => { + for (const url of ['file:///etc/passwd', 'javascript:alert(1)', 'data:application/json,{}']) { + const { fetch, calls } = makeFetch([ + { + status: 401, + headers: { 'WWW-Authenticate': `Bearer resource_metadata="${url}"` }, + }, + ]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + const status = await upstream.start(); + assert.deepStrictEqual({ + url, + kind: status.kind, + resource: status.kind === McpServerStatusKind.AuthRequired ? status.resource : undefined, + callCount: calls.length, + }, { + url, + kind: McpServerStatusKind.AuthRequired, + resource: { resource: 'https://mcp.example.com/v1' }, + callCount: 1, + }); + } finally { + upstream.dispose(); + } + } + }); + + test('accepts resource_metadata at the same origin', async () => { + const metadata = { resource: 'https://mcp.example.com/v1', authorization_servers: ['https://auth.example.com'] }; + const { fetch, calls } = makeFetch([ + { + status: 401, + headers: { 'WWW-Authenticate': 'Bearer resource_metadata="https://mcp.example.com/.well-known/oauth-protected-resource"' }, + }, + { status: 200, body: JSON.stringify(metadata) }, + ]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + const status = await upstream.start(); + assert.deepStrictEqual({ + kind: status.kind, + resource: status.kind === McpServerStatusKind.AuthRequired ? status.resource : undefined, + secondCallUrl: calls[1]?.url, + }, { + kind: McpServerStatusKind.AuthRequired, + resource: metadata, + secondCallUrl: 'https://mcp.example.com/.well-known/oauth-protected-resource', + }); + } finally { + upstream.dispose(); + } + }); + + test('403 with insufficient_scope yields AuthRequired (InsufficientScope)', async () => { + const metadata = { resource: 'https://mcp.example.com/v1' }; + const { fetch } = makeFetch([ + { + status: 403, + headers: { 'WWW-Authenticate': 'Bearer error="insufficient_scope", scope="admin", resource_metadata="https://mcp.example.com/.well-known/m"' }, + }, + { status: 200, body: JSON.stringify(metadata) }, + ]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + const status = await upstream.start(); + assert.strictEqual(status.kind, McpServerStatusKind.AuthRequired); + if (status.kind === McpServerStatusKind.AuthRequired) { + assert.strictEqual(status.reason, McpAuthRequiredReason.InsufficientScope); + assert.deepStrictEqual(status.requiredScopes, ['admin']); + } + } finally { + upstream.dispose(); + } + }); + + test('after setBearerToken, retried start() sends Authorization header', async () => { + const { fetch, calls } = makeFetch([ + { status: 401, headers: { 'WWW-Authenticate': 'Bearer' } }, + { status: 200 }, + ]); + const logger = new RecordingLogger(); + const upstream = new McpHttpUpstream({ config: baseConfig, logger, fetch }); + try { + await upstream.start(); + upstream.setBearerToken('shiny-token'); + const status = await upstream.start(); + assert.deepStrictEqual({ + firstAuth: calls[0].headers.hasOwnProperty('Authorization'), + secondAuth: calls[1].headers['Authorization'], + status, + }, { + firstAuth: false, + secondAuth: 'Bearer shiny-token', + status: { kind: McpServerStatusKind.Ready }, + }); + } finally { + upstream.dispose(); + } + }); + + test('send() rejects when not Ready', async () => { + const { fetch } = makeFetch([]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + await assert.rejects( + upstream.send({ jsonrpc: '2.0', id: 1, method: 'ping' }), + /cannot send while in state 'stopped'/, + ); + } finally { + upstream.dispose(); + } + }); + + test('send() consumes text/event-stream responses and emits one onMessage per JSON-RPC event', async () => { + // Encode two SSE `message` events whose data fields are JSON-RPC payloads. + const sseBytes = new TextEncoder().encode([ + 'event: message', + `data: ${JSON.stringify({ jsonrpc: '2.0', id: 1, result: { ok: true } })}`, + '', + 'event: message', + `data: ${JSON.stringify({ jsonrpc: '2.0', method: 'notifications/progress', params: { value: 42 } })}`, + '', + '', + ].join('\n')); + + // Probe response: one-shot JSON, transitions to Ready. + // Send response: SSE stream with the two events above. + const responses: ((url: string) => IFetchResponseSpec)[] = [ + () => ({ status: 200 }), + () => ({ status: 200, headers: { 'Content-Type': 'text/event-stream' } }), + ]; + let i = 0; + const calls: IFetchCall[] = []; + const fetch: HttpFetch = async (url, init) => { + calls.push({ url, method: init.method, headers: { ...init.headers }, body: init.body }); + const spec = responses[i++](url); + const base = makeResponse(spec); + if (i === 2) { + // Attach a real ReadableStream as the response body for the SSE branch. + return { + ...base, + body: new ReadableStream({ + start(controller) { + controller.enqueue(sseBytes); + controller.close(); + }, + }), + }; + } + return base; + }; + + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + const received: JsonRpcMessage[] = []; + const sub = upstream.onMessage(m => received.push(m)); + try { + const startStatus = await upstream.start(); + assert.strictEqual(startStatus.kind, McpServerStatusKind.Ready); + await upstream.send({ jsonrpc: '2.0', id: 1, method: 'tools/list' }); + assert.deepStrictEqual(received, [ + { jsonrpc: '2.0', id: 1, result: { ok: true } }, + { jsonrpc: '2.0', method: 'notifications/progress', params: { value: 42 } }, + ]); + } finally { + sub.dispose(); + upstream.dispose(); + } + }); + + test('network error transitions to Error', async () => { + const { fetch } = makeFetch([{ status: 0, throw: new Error('ECONNREFUSED') }]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + const status = await upstream.start(); + assert.strictEqual(status.kind, McpServerStatusKind.Error); + if (status.kind === McpServerStatusKind.Error) { + assert.strictEqual(status.error.errorType, 'httpError'); + assert.match(status.error.message, /ECONNREFUSED/); + } + } finally { + upstream.dispose(); + } + }); + + test('401 without resource_metadata synthesizes minimal metadata and warns', async () => { + const { fetch, calls } = makeFetch([ + { status: 401, headers: { 'WWW-Authenticate': 'Bearer error="invalid_token"' } }, + ]); + const logger = new RecordingLogger(); + const upstream = new McpHttpUpstream({ config: baseConfig, logger, fetch }); + try { + const status = await upstream.start(); + assert.deepStrictEqual({ + status, + callCount: calls.length, + warned: logger.records.some(r => r.level === 'warn' && /resource_metadata/.test(r.message)), + }, { + status: { + kind: McpServerStatusKind.AuthRequired, + reason: McpAuthRequiredReason.Required, + resource: { resource: 'https://mcp.example.com/v1' }, + }, + callCount: 1, + warned: true, + }); + } finally { + upstream.dispose(); + } + }); + + test('send() on 401 mid-session transitions to AuthRequired and throws', async () => { + const metadata = { resource: 'https://mcp.example.com/v1', authorization_servers: ['https://auth.example.com'] }; + const { fetch } = makeFetch([ + { status: 200 }, + { + status: 401, + headers: { 'WWW-Authenticate': 'Bearer error="invalid_token", resource_metadata="https://mcp.example.com/.well-known/m"' }, + }, + { status: 200, body: JSON.stringify(metadata) }, + ]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + await upstream.start(); + await assert.rejects( + upstream.send({ jsonrpc: '2.0', id: 1, method: 'tools/list' }), + /AuthRequired/, + ); + assert.deepStrictEqual(upstream.status.get(), { + kind: McpServerStatusKind.AuthRequired, + reason: McpAuthRequiredReason.Required, + resource: metadata, + }); + } finally { + upstream.dispose(); + } + }); + + test('send() on 401 after a prior token transitions to AuthRequired (Expired)', async () => { + const metadata = { resource: 'https://mcp.example.com/v1', authorization_servers: ['https://auth.example.com'] }; + const { fetch } = makeFetch([ + { status: 200 }, + { + status: 401, + headers: { 'WWW-Authenticate': 'Bearer error="invalid_token", resource_metadata="https://mcp.example.com/.well-known/m"' }, + }, + { status: 200, body: JSON.stringify(metadata) }, + ]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + await upstream.start(); + upstream.setBearerToken('tok'); + await assert.rejects(upstream.send({ jsonrpc: '2.0', id: 1, method: 'tools/list' })); + const status = upstream.status.get(); + assert.strictEqual(status.kind, McpServerStatusKind.AuthRequired); + if (status.kind === McpServerStatusKind.AuthRequired) { + assert.strictEqual(status.reason, McpAuthRequiredReason.Expired); + } + } finally { + upstream.dispose(); + } + }); + + test('send() on 403 insufficient_scope transitions to AuthRequired (InsufficientScope)', async () => { + const { fetch } = makeFetch([ + { status: 200 }, + { + status: 403, + headers: { 'WWW-Authenticate': 'Bearer error="insufficient_scope", scope="admin"' }, + }, + ]); + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + await upstream.start(); + await assert.rejects(upstream.send({ jsonrpc: '2.0', id: 1, method: 'tools/list' })); + const status = upstream.status.get(); + assert.strictEqual(status.kind, McpServerStatusKind.AuthRequired); + if (status.kind === McpServerStatusKind.AuthRequired) { + assert.strictEqual(status.reason, McpAuthRequiredReason.InsufficientScope); + assert.deepStrictEqual(status.requiredScopes, ['admin']); + } + } finally { + upstream.dispose(); + } + }); + + test('dispose aborts all in-flight requests, not just the most recent', async () => { + // First fetch resolves immediately to bring us to Ready; subsequent + // send() fetches are deferred and never resolve until released. + const signals: AbortSignal[] = []; + let pendingResolve: (() => void) | undefined; + const fetch: HttpFetch = async (url, init) => { + if (init.method === 'POST' && init.body && JSON.parse(init.body).method === 'initialize') { + return makeResponse({ status: 200 }); + } + if (init.signal) { + signals.push(init.signal); + } + // Never resolve; the AbortController abort will reject via abort listener. + return new Promise((_, reject) => { + const onAbort = () => reject(new Error('aborted')); + if (init.signal?.aborted) { + onAbort(); + } else { + init.signal?.addEventListener('abort', onAbort, { once: true }); + } + pendingResolve = () => reject(new Error('test cleanup')); + }); + }; + const upstream = new McpHttpUpstream({ config: baseConfig, logger: new NullLogger(), fetch }); + try { + await upstream.start(); + const sends = [ + upstream.send({ jsonrpc: '2.0', id: 1, method: 'tools/list' }).catch(() => 'rejected'), + upstream.send({ jsonrpc: '2.0', id: 2, method: 'tools/list' }).catch(() => 'rejected'), + upstream.send({ jsonrpc: '2.0', id: 3, method: 'tools/list' }).catch(() => 'rejected'), + ]; + // Yield so the fetches register their AbortSignals. + await new Promise(resolve => setTimeout(resolve, 0)); + upstream.dispose(); + await Promise.all(sends); + assert.deepStrictEqual({ + count: signals.length, + aborted: signals.map(s => s.aborted), + }, { + count: 3, + aborted: [true, true, true], + }); + } finally { + pendingResolve?.(); + upstream.dispose(); + } + }); +}); diff --git a/src/vs/platform/agentHost/test/node/mcpHost/mcpInitializeInjector.test.ts b/src/vs/platform/agentHost/test/node/mcpHost/mcpInitializeInjector.test.ts new file mode 100644 index 0000000000000..bf483c62ca07e --- /dev/null +++ b/src/vs/platform/agentHost/test/node/mcpHost/mcpInitializeInjector.test.ts @@ -0,0 +1,94 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as assert from 'assert'; +import type { IJsonRpcRequest } from '../../../../../base/common/jsonRpcProtocol.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; +import { McpAppsInitializeInjector } from '../../../node/mcpHost/mcpInitializeInjector.js'; + +suite('McpAppsInitializeInjector', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + const baseRequest = (params: unknown): IJsonRpcRequest => ({ + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params, + }); + + const expectedExtension = { + 'io.modelcontextprotocol/ui': { + mimeTypes: ['text/html;profile=mcp-app'], + }, + }; + + test('injects extensions into empty capabilities', () => { + const injector = new McpAppsInitializeInjector(); + const result = injector.inject(baseRequest({ capabilities: {} })); + + assert.deepStrictEqual(result, { + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params: { + capabilities: { + extensions: expectedExtension, + }, + }, + }); + }); + + test('preserves caller-provided capabilities while merging extensions', () => { + const injector = new McpAppsInitializeInjector(); + const result = injector.inject(baseRequest({ + protocolVersion: '2025-11-25', + capabilities: { + sampling: {}, + roots: { listChanged: true }, + extensions: { + foo: { bar: 1 }, + }, + }, + clientInfo: { name: 'sdk', version: '1.0.0' }, + })); + + assert.deepStrictEqual(result, { + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params: { + protocolVersion: '2025-11-25', + capabilities: { + sampling: {}, + roots: { listChanged: true }, + extensions: { + foo: { bar: 1 }, + ...expectedExtension, + }, + }, + clientInfo: { name: 'sdk', version: '1.0.0' }, + }, + }); + }); + + test('is idempotent', () => { + const injector = new McpAppsInitializeInjector(); + const first = injector.inject(baseRequest({ capabilities: { sampling: {} } })); + const second = injector.inject(first); + assert.deepStrictEqual(second, first); + }); + + test('does not mutate the original request or its params', () => { + const injector = new McpAppsInitializeInjector(); + const originalParams = { capabilities: { sampling: {}, extensions: { foo: 1 } } }; + const original = baseRequest(originalParams); + const snapshot = JSON.parse(JSON.stringify(original)); + + injector.inject(original); + + assert.deepStrictEqual(original, snapshot); + assert.strictEqual(original.params, originalParams); + }); +}); diff --git a/src/vs/platform/agentHost/test/node/mcpHost/mcpProxy.integrationTest.ts b/src/vs/platform/agentHost/test/node/mcpHost/mcpProxy.integrationTest.ts new file mode 100644 index 0000000000000..b6057ad5dbefe --- /dev/null +++ b/src/vs/platform/agentHost/test/node/mcpHost/mcpProxy.integrationTest.ts @@ -0,0 +1,449 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * Integration tests for {@link McpProxy}/{@link McpProxyFactory} that bind + * a real localhost HTTP listener and exercise the full SDK ⟷ HTTP ⟷ + * route ⟷ upstream pipeline. Kept out of the unit suite because the + * listener does real socket work; same convention as + * `agentHostGitService.integrationTest.ts`. + * + * Run via `scripts/test-integration.sh`. + */ + +import * as assert from 'assert'; +import { DeferredPromise } from '../../../../../base/common/async.js'; +import { Emitter, type Event } from '../../../../../base/common/event.js'; +import { isJsonRpcRequest, type JsonRpcMessage } from '../../../../../base/common/jsonRpcProtocol.js'; +import { observableValue, type IObservable, type ISettableObservable } from '../../../../../base/common/observable.js'; +import { URI } from '../../../../../base/common/uri.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; +import { NullLogger } from '../../../../log/common/log.js'; +import { + McpAuthRequiredReason, + McpServerStatusKind, + type McpServerStatus, + type McpServerStatusAuthRequired, +} from '../../../common/state/protocol/state.js'; +import { McpAppsInitializeInjector } from '../../../node/mcpHost/mcpInitializeInjector.js'; +import { McpProxyFactory, type IMcpProxyOptions } from '../../../node/mcpHost/mcpProxy.js'; +import type { IUpstreamRequestOutcome } from '../../../node/mcpHost/mcpProxyRoute.js'; +import type { IMcpUpstream, IMcpUpstreamCapabilities } from '../../../node/mcpHost/mcpUpstream.js'; + +class StubUpstream implements IMcpUpstream { + private readonly _status: ISettableObservable = observableValue('stub-upstream', { kind: McpServerStatusKind.Stopped }); + public readonly status: IObservable = this._status; + private readonly _onMessage = new Emitter(); + public readonly onMessage: Event = this._onMessage.event; + private readonly _upstreamCapabilities: ISettableObservable = + observableValue('stub-upstream-caps', undefined); + public readonly upstreamCapabilities: IObservable = this._upstreamCapabilities; + + public readonly sent: JsonRpcMessage[] = []; + public readonly tokens: (string | undefined)[] = []; + public startCalls = 0; + public disposeCalls = 0; + + public startResult: McpServerStatus = { kind: McpServerStatusKind.Ready }; + public sendThrows: Error | undefined; + + /** Optional reaction invoked synchronously after each `send`. */ + public onSend: ((msg: JsonRpcMessage) => void) | undefined; + + public async start(): Promise { + this.startCalls++; + this._status.set(this.startResult, undefined); + return this.startResult; + } + + public async send(message: JsonRpcMessage): Promise { + if (this.sendThrows) { + throw this.sendThrows; + } + this.sent.push(message); + this.onSend?.(message); + } + + public setBearerToken(token: string | undefined): void { + this.tokens.push(token); + } + + public setUpstreamCapabilities(caps: IMcpUpstreamCapabilities | undefined): void { + this._upstreamCapabilities.set(caps, undefined); + } + + public emit(message: JsonRpcMessage): void { + this._onMessage.fire(message); + } + + public setStatus(status: McpServerStatus): void { + this._status.set(status, undefined); + } + + public dispose(): void { + this.disposeCalls++; + this._onMessage.dispose(); + } +} + +interface IRecordedRequest { + readonly method: string; + readonly params: unknown; + readonly outcome: DeferredPromise; +} + +interface IRecordedNotification { + readonly method: string; + readonly params: unknown; +} + +interface ITestHarness { + readonly factory: McpProxyFactory; + readonly upstream: StubUpstream; + readonly recordedRequests: IRecordedRequest[]; + readonly recordedNotifications: IRecordedNotification[]; + readonly authChallenges: McpServerStatusAuthRequired[]; + readonly stateChanges: McpServerStatus[]; + makeOptions(overrides?: Partial): IMcpProxyOptions; +} + +function createHarness(): ITestHarness { + const factory = new McpProxyFactory(new NullLogger()); + const upstream = new StubUpstream(); + const recordedRequests: IRecordedRequest[] = []; + const recordedNotifications: IRecordedNotification[] = []; + const authChallenges: McpServerStatusAuthRequired[] = []; + const stateChanges: McpServerStatus[] = []; + return { + factory, + upstream, + recordedRequests, + recordedNotifications, + authChallenges, + stateChanges, + makeOptions(overrides) { + const base: IMcpProxyOptions = { + resource: URI.parse('mcp:/session-1/server-1'), + upstream, + logger: new NullLogger(), + onUpstreamRequest: (method, params) => { + const outcome = new DeferredPromise(); + recordedRequests.push({ method, params, outcome }); + return outcome.p; + }, + onUpstreamNotification: (method, params) => { + recordedNotifications.push({ method, params }); + }, + onAuthRequired: status => { + authChallenges.push(status); + }, + onStateChange: status => { + stateChanges.push(status); + }, + }; + return { ...base, ...overrides }; + }, + }; +} + +async function postJson(endpoint: URI, body: object | string): Promise<{ status: number; text: string }> { + const payload = typeof body === 'string' ? body : JSON.stringify(body); + const response = await fetch(endpoint.toString(true), { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: payload, + }); + const text = await response.text(); + return { status: response.status, text }; +} + +suite('McpProxy (integration)', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + test('initialize injection rewrites params before forwarding to upstream', async () => { + const harness = createHarness(); + let proxy; + try { + harness.upstream.onSend = msg => { + if (isJsonRpcRequest(msg) && msg.method === 'initialize') { + harness.upstream.emit({ jsonrpc: '2.0', id: msg.id, result: { ok: true } }); + } + }; + proxy = await harness.factory.create(harness.makeOptions({ initializeInjector: new McpAppsInitializeInjector() })); + await harness.upstream.start(); + + const { status, text } = await postJson(proxy.endpoint, { + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params: { capabilities: { sampling: {} }, clientInfo: { name: 'sdk', version: '0.1' } }, + }); + + const sent = harness.upstream.sent[0] as { method?: string; params?: { capabilities?: { sampling?: unknown; extensions?: Record } } }; + assert.deepStrictEqual({ + status, + text: JSON.parse(text), + upstreamMethod: sent?.method, + preservedSampling: sent?.params?.capabilities?.sampling, + injectedExtension: sent?.params?.capabilities?.extensions?.['io.modelcontextprotocol/ui'], + }, { + status: 200, + text: { jsonrpc: '2.0', id: 1, result: { ok: true } }, + upstreamMethod: 'initialize', + preservedSampling: {}, + injectedExtension: { mimeTypes: ['text/html;profile=mcp-app'] }, + }); + } finally { + proxy?.dispose(); + harness.factory.dispose(); + harness.upstream.dispose(); + } + }); + + test('SDK request is forwarded and upstream response routed back as HTTP 200', async () => { + const harness = createHarness(); + let proxy; + try { + harness.upstream.onSend = msg => { + if (isJsonRpcRequest(msg) && msg.method === 'tools/list') { + harness.upstream.emit({ jsonrpc: '2.0', id: msg.id, result: { tools: [{ name: 'echo' }] } }); + } + }; + proxy = await harness.factory.create(harness.makeOptions()); + await harness.upstream.start(); + + const { status, text } = await postJson(proxy.endpoint, { jsonrpc: '2.0', id: 7, method: 'tools/list' }); + + assert.deepStrictEqual({ status, body: JSON.parse(text) }, { + status: 200, + body: { jsonrpc: '2.0', id: 7, result: { tools: [{ name: 'echo' }] } }, + }); + } finally { + proxy?.dispose(); + harness.factory.dispose(); + harness.upstream.dispose(); + } + }); + + test('SDK notification is forwarded to upstream and yields HTTP 204', async () => { + const harness = createHarness(); + let proxy; + try { + proxy = await harness.factory.create(harness.makeOptions()); + await harness.upstream.start(); + + const { status, text } = await postJson(proxy.endpoint, { + jsonrpc: '2.0', + method: 'notifications/cancelled', + params: { requestId: 7 }, + }); + + assert.deepStrictEqual({ + status, + text, + sent: harness.upstream.sent, + }, { + status: 204, + text: '', + sent: [{ jsonrpc: '2.0', method: 'notifications/cancelled', params: { requestId: 7 } }], + }); + } finally { + proxy?.dispose(); + harness.factory.dispose(); + harness.upstream.dispose(); + } + }); + + test('upstream notification fires onUpstreamNotification', async () => { + const harness = createHarness(); + let proxy; + try { + proxy = await harness.factory.create(harness.makeOptions()); + await harness.upstream.start(); + + harness.upstream.emit({ jsonrpc: '2.0', method: 'notifications/tools/list_changed' }); + + assert.deepStrictEqual(harness.recordedNotifications, [{ + method: 'notifications/tools/list_changed', + params: undefined, + }]); + } finally { + proxy?.dispose(); + harness.factory.dispose(); + harness.upstream.dispose(); + } + }); + + test('upstream request is tapped and the resolved outcome is written back with the original id', async () => { + const harness = createHarness(); + let proxy; + try { + proxy = await harness.factory.create(harness.makeOptions()); + await harness.upstream.start(); + + harness.upstream.emit({ jsonrpc: '2.0', id: 42, method: 'sampling/createMessage', params: { foo: 'bar' } }); + + // Wait for the request to be recorded, then resolve the outcome. + assert.strictEqual(harness.recordedRequests.length, 1); + assert.deepStrictEqual({ + method: harness.recordedRequests[0].method, + params: harness.recordedRequests[0].params, + }, { + method: 'sampling/createMessage', + params: { foo: 'bar' }, + }); + + harness.recordedRequests[0].outcome.complete({ result: { ok: true } }); + await new Promise(resolve => setImmediate(resolve)); + + assert.deepStrictEqual(harness.upstream.sent, [{ jsonrpc: '2.0', id: 42, result: { ok: true } }]); + } finally { + proxy?.dispose(); + harness.factory.dispose(); + harness.upstream.dispose(); + } + }); + + test('AuthRequired status invokes onAuthRequired; authenticate() retries with bearer token', async () => { + const harness = createHarness(); + let proxy; + try { + const challenge: McpServerStatusAuthRequired = { + kind: McpServerStatusKind.AuthRequired, + reason: McpAuthRequiredReason.Required, + resource: { resource: 'https://example/' }, + }; + harness.upstream.startResult = challenge; + proxy = await harness.factory.create(harness.makeOptions()); + await harness.upstream.start(); + + // Now flip the upstream so the next start() returns Ready. + harness.upstream.startResult = { kind: McpServerStatusKind.Ready }; + const ok = await proxy.authenticate('https://example/', 'token-xyz'); + + assert.deepStrictEqual({ + ok, + challenges: harness.authChallenges, + tokens: harness.upstream.tokens, + startCalls: harness.upstream.startCalls, + }, { + ok: true, + challenges: [challenge], + tokens: ['token-xyz'], + startCalls: 2, + }); + } finally { + proxy?.dispose(); + harness.factory.dispose(); + harness.upstream.dispose(); + } + }); + + test('factory.dispose() shuts the listener down so subsequent fetch fails', async () => { + const harness = createHarness(); + const proxy = await harness.factory.create(harness.makeOptions()); + await harness.upstream.start(); + proxy.dispose(); + harness.factory.dispose(); + harness.upstream.dispose(); + + // Allow the close to propagate. + const closed = new DeferredPromise<{ ok: boolean }>(); + fetch(proxy.endpoint.toString(true), { method: 'POST', body: '{}' }) + .then(() => closed.complete({ ok: true })) + .catch(() => closed.complete({ ok: false })); + + const { ok } = await closed.p; + assert.strictEqual(ok, false); + }); + + test('disposing route while SDK request is in-flight returns error with the original id', async () => { + const harness = createHarness(); + const proxy = await harness.factory.create(harness.makeOptions()); + try { + await harness.upstream.start(); + // upstream.onSend is unset → upstream never emits a reply, so + // the route's pending-deferred is still in the map when we + // dispose below. + + const requestId = 'req-x'; + const fetchPromise = postJson(proxy.endpoint, { + jsonrpc: '2.0', + id: requestId, + method: 'tools/list', + }); + + // Wait for the SDK's request to be observed by the upstream + // (i.e. registered as a pending SDK request on the route). + for (let i = 0; i < 50 && harness.upstream.sent.length === 0; i++) { + await new Promise(resolve => setTimeout(resolve, 0)); + } + assert.strictEqual(harness.upstream.sent.length, 1, 'expected upstream to have observed the SDK request'); + + proxy.dispose(); + const { status, text } = await fetchPromise; + const body = JSON.parse(text); + assert.deepStrictEqual({ + status, + id: body.id, + hasError: !!body.error, + }, { + status: 200, + id: requestId, + hasError: true, + }); + } finally { + harness.factory.dispose(); + harness.upstream.dispose(); + } + }); + + test('initialize response capabilities are captured on the upstream', async () => { + const harness = createHarness(); + let proxy; + try { + harness.upstream.onSend = msg => { + if (isJsonRpcRequest(msg) && msg.method === 'initialize') { + harness.upstream.emit({ + jsonrpc: '2.0', + id: msg.id, + result: { + capabilities: { + extensions: { + 'io.modelcontextprotocol/ui': { mimeTypes: ['text/html;profile=mcp-app'] }, + }, + }, + }, + }); + } + }; + proxy = await harness.factory.create(harness.makeOptions()); + await harness.upstream.start(); + + const { status } = await postJson(proxy.endpoint, { + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params: { capabilities: {} }, + }); + + assert.deepStrictEqual({ + status, + caps: harness.upstream.upstreamCapabilities.get(), + }, { + status: 200, + caps: { + extensions: { + 'io.modelcontextprotocol/ui': { mimeTypes: ['text/html;profile=mcp-app'] }, + }, + }, + }); + } finally { + proxy?.dispose(); + harness.factory.dispose(); + harness.upstream.dispose(); + } + }); +}); diff --git a/src/vs/platform/agentHost/test/node/mcpHost/mcpStdioUpstream.test.ts b/src/vs/platform/agentHost/test/node/mcpHost/mcpStdioUpstream.test.ts new file mode 100644 index 0000000000000..f10c6637cf6f0 --- /dev/null +++ b/src/vs/platform/agentHost/test/node/mcpHost/mcpStdioUpstream.test.ts @@ -0,0 +1,325 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as assert from 'assert'; +import { ChildProcessWithoutNullStreams } from 'child_process'; +import { EventEmitter } from 'events'; +import { PassThrough } from 'stream'; +import { DeferredPromise, timeout } from '../../../../../base/common/async.js'; +import { DisposableStore } from '../../../../../base/common/lifecycle.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; +import { NullLogger } from '../../../../log/common/log.js'; +import { McpServerType, type IMcpStdioServerConfiguration } from '../../../../mcp/common/mcpPlatformTypes.js'; +import { McpServerStatusKind } from '../../../common/state/protocol/state.js'; +import { McpStdioUpstream, type StdioSpawn } from '../../../node/mcpHost/mcpStdioUpstream.js'; + +interface IFakeChild extends ChildProcessWithoutNullStreams { + _emitStdout(line: string): void; + _emitStderr(line: string): void; + _emitStdoutRaw(chunk: string): void; + _emitStderrRaw(chunk: string): void; + _emitError(err: Error): void; + _exit(code: number | null, signal?: NodeJS.Signals | null): void; + _killCalls: number; +} + +interface ISpawnRecord { + command: string; + args: readonly string[]; + options: { cwd?: string; env?: NodeJS.ProcessEnv }; +} + +function createFakeChild(): IFakeChild { + const stdin = new PassThrough(); + const stdout = new PassThrough(); + const stderr = new PassThrough(); + const ee = new EventEmitter(); + const fake = ee as unknown as IFakeChild; + (fake as unknown as { stdin: PassThrough }).stdin = stdin; + (fake as unknown as { stdout: PassThrough }).stdout = stdout; + (fake as unknown as { stderr: PassThrough }).stderr = stderr; + (fake as unknown as { pid: number }).pid = 0; // 0 avoids the killTree path on dispose + fake._killCalls = 0; + (fake as unknown as { kill: (s?: NodeJS.Signals) => boolean }).kill = () => { + fake._killCalls++; + return true; + }; + fake._emitStdout = line => stdout.write(line + '\n'); + fake._emitStderr = line => stderr.write(line + '\n'); + fake._emitStdoutRaw = chunk => stdout.write(chunk); + fake._emitStderrRaw = chunk => stderr.write(chunk); + fake._emitError = err => ee.emit('error', err); + fake._exit = (code, signal) => { + // Real child processes close stdio on exit; the splitter pipeline + // only flushes a trailing partial line once its source stream ends. + stdout.end(); + stderr.end(); + ee.emit('exit', code, signal ?? null); + }; + return fake; +} + +interface IRecordedLog { + level: 'info' | 'warn' | 'error'; + message: string; +} + +class RecordingLogger extends NullLogger { + public readonly records: IRecordedLog[] = []; + override info(message: string): void { this.records.push({ level: 'info', message }); } + override warn(message: string): void { this.records.push({ level: 'warn', message }); } + override error(message: string | Error): void { + this.records.push({ level: 'error', message: message instanceof Error ? message.message : message }); + } +} + +const baseConfig: IMcpStdioServerConfiguration = { + type: McpServerType.LOCAL, + command: 'mcp-server', + args: ['--flag'], +}; + +suite('McpStdioUpstream', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + function makeUpstream(overrides: { spawn?: StdioSpawn; logger?: RecordingLogger; spawnRecords?: ISpawnRecord[]; childRef?: { current: IFakeChild | undefined }; throwOnSpawn?: Error } = {}) { + const logger = overrides.logger ?? new RecordingLogger(); + const records: ISpawnRecord[] = overrides.spawnRecords ?? []; + const childRef = overrides.childRef ?? { current: undefined }; + const spawn: StdioSpawn = overrides.spawn ?? ((command, args, options) => { + records.push({ command, args, options }); + if (overrides.throwOnSpawn) { + throw overrides.throwOnSpawn; + } + const child = createFakeChild(); + childRef.current = child; + return child; + }); + const upstream = new McpStdioUpstream({ config: baseConfig, logger, spawn }); + return { upstream, logger, records, childRef }; + } + + test('constructor does not spawn', () => { + let spawnCalls = 0; + const spawn: StdioSpawn = () => { + spawnCalls++; + return createFakeChild(); + }; + const upstream = new McpStdioUpstream({ config: baseConfig, logger: new NullLogger(), spawn }); + try { + assert.strictEqual(spawnCalls, 0); + assert.deepStrictEqual(upstream.status.get(), { kind: McpServerStatusKind.Stopped }); + } finally { + upstream.dispose(); + } + }); + + test('start() spawns and transitions Stopped → Ready', async () => { + const { upstream, records, childRef } = makeUpstream(); + try { + assert.deepStrictEqual(upstream.status.get(), { kind: McpServerStatusKind.Stopped }); + const result = await upstream.start(); + assert.deepStrictEqual({ + result, + records, + hasChild: !!childRef.current, + finalStatus: upstream.status.get(), + }, { + result: { kind: McpServerStatusKind.Ready }, + records: [{ + command: 'mcp-server', + args: ['--flag'], + options: { cwd: undefined, env: undefined }, + }], + hasChild: true, + finalStatus: { kind: McpServerStatusKind.Ready }, + }); + } finally { + upstream.dispose(); + } + }); + + test('parses stdout NDJSON and emits via onMessage', async () => { + const store = new DisposableStore(); + const { upstream, childRef } = makeUpstream(); + try { + const got: unknown[] = []; + store.add(upstream.onMessage(m => got.push(m))); + await upstream.start(); + childRef.current!._emitStdout(JSON.stringify({ jsonrpc: '2.0', id: 1, result: { ok: true } })); + await timeout(0); + assert.deepStrictEqual(got, [{ jsonrpc: '2.0', id: 1, result: { ok: true } }]); + } finally { + store.dispose(); + upstream.dispose(); + } + }); + + test('drops malformed stdout lines and logs error', async () => { + const store = new DisposableStore(); + const { upstream, logger, childRef } = makeUpstream(); + try { + const got: unknown[] = []; + store.add(upstream.onMessage(m => got.push(m))); + await upstream.start(); + childRef.current!._emitStdout('not-json'); + await timeout(0); + assert.deepStrictEqual(got, []); + assert.ok(logger.records.some(r => r.level === 'error' && /failed to parse stdout line/.test(r.message))); + } finally { + store.dispose(); + upstream.dispose(); + } + }); + + test('forwards stderr lines as logger.info', async () => { + const { upstream, logger, childRef } = makeUpstream(); + try { + await upstream.start(); + childRef.current!._emitStderr('hello world'); + await timeout(0); + assert.ok(logger.records.some(r => r.level === 'info' && /hello world/.test(r.message))); + } finally { + upstream.dispose(); + } + }); + + test('send() writes \\n-terminated JSON to stdin', async () => { + const { upstream, childRef } = makeUpstream(); + try { + await upstream.start(); + const written = new DeferredPromise(); + let buf = ''; + childRef.current!.stdin.on('data', (chunk: Buffer | string) => { + buf += typeof chunk === 'string' ? chunk : chunk.toString('utf8'); + if (buf.endsWith('\n')) { + written.complete(buf); + } + }); + await upstream.send({ jsonrpc: '2.0', id: 7, method: 'ping' }); + const got = await written.p; + assert.strictEqual(got, JSON.stringify({ jsonrpc: '2.0', id: 7, method: 'ping' }) + '\n'); + } finally { + upstream.dispose(); + } + }); + + test('send() rejects when not Ready', async () => { + const { upstream } = makeUpstream(); + try { + await assert.rejects( + upstream.send({ jsonrpc: '2.0', id: 1, method: 'ping' }), + /cannot send while in state 'stopped'/, + ); + } finally { + upstream.dispose(); + } + }); + + test('child exit (clean) transitions to Stopped', async () => { + const { upstream, childRef } = makeUpstream(); + try { + await upstream.start(); + childRef.current!._exit(0); + await timeout(0); + assert.deepStrictEqual(upstream.status.get(), { kind: McpServerStatusKind.Stopped }); + } finally { + upstream.dispose(); + } + }); + + test('child exit (non-zero) transitions to Error', async () => { + const { upstream, childRef } = makeUpstream(); + try { + await upstream.start(); + childRef.current!._exit(2); + await timeout(0); + const status = upstream.status.get(); + assert.strictEqual(status.kind, McpServerStatusKind.Error); + assert.strictEqual(status.kind === McpServerStatusKind.Error && status.error.errorType, 'childExited'); + } finally { + upstream.dispose(); + } + }); + + test('spawn failure resolves with Error', async () => { + const { upstream } = makeUpstream({ throwOnSpawn: new Error('ENOENT: mcp-server') }); + try { + const result = await upstream.start(); + assert.strictEqual(result.kind, McpServerStatusKind.Error); + assert.strictEqual(result.kind === McpServerStatusKind.Error && result.error.errorType, 'spawnFailed'); + assert.ok(result.kind === McpServerStatusKind.Error && /ENOENT/.test(result.error.message)); + } finally { + upstream.dispose(); + } + }); + + test('asynchronous child error transitions status to Error', async () => { + const { upstream, childRef } = makeUpstream(); + try { + const synchronousStart = await upstream.start(); + assert.strictEqual(synchronousStart.kind, McpServerStatusKind.Ready); + + // Fire the error asynchronously, after start() returned. + await new Promise(resolve => setImmediate(() => { + childRef.current!._emitError(new Error('spawn ENOENT: bogus-mcp-server')); + resolve(); + })); + await timeout(0); + + const status = upstream.status.get(); + assert.strictEqual(status.kind, McpServerStatusKind.Error); + if (status.kind === McpServerStatusKind.Error) { + assert.strictEqual(status.error.errorType, 'spawnFailed'); + assert.match(status.error.message, /ENOENT/); + } + await assert.rejects( + upstream.send({ jsonrpc: '2.0', id: 1, method: 'ping' }), + /cannot send while in state 'error'/, + ); + } finally { + upstream.dispose(); + } + }); + + test('delivers a final unterminated stdout line on child exit', async () => { + const store = new DisposableStore(); + const { upstream, childRef } = makeUpstream(); + try { + const got: unknown[] = []; + store.add(upstream.onMessage(m => got.push(m))); + await upstream.start(); + // Partial line without a trailing newline. `StreamSplitter._flush` + // emits it on stream end, so the consumer still sees the message + // even when the child crashes mid-line. + childRef.current!._emitStdoutRaw('{"jsonrpc":"2.0","id":1,"result":{}}'); + childRef.current!._exit(0); + // Stream-end → `StreamSplitter._flush` → final `data` event takes + // more than one event-loop turn through the PassThrough chain. + await timeout(10); + assert.deepStrictEqual(got, [{ jsonrpc: '2.0', id: 1, result: {} }]); + } finally { + store.dispose(); + upstream.dispose(); + } + }); + + test('dispose is idempotent', async () => { + const { upstream } = makeUpstream(); + await upstream.start(); + upstream.dispose(); + assert.doesNotThrow(() => upstream.dispose()); + }); + + test('setBearerToken is a no-op for stdio', () => { + const { upstream } = makeUpstream(); + try { + assert.doesNotThrow(() => upstream.setBearerToken('abc')); + assert.doesNotThrow(() => upstream.setBearerToken(undefined)); + } finally { + upstream.dispose(); + } + }); +}); diff --git a/src/vs/platform/agentHost/test/node/protocol/mcpRealSdk.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/mcpRealSdk.integrationTest.ts new file mode 100644 index 0000000000000..2263caa43c345 --- /dev/null +++ b/src/vs/platform/agentHost/test/node/protocol/mcpRealSdk.integrationTest.ts @@ -0,0 +1,1006 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * Integration tests for the agent host's MCP plugin pipeline using real + * `@modelcontextprotocol/sdk` servers and the real Copilot SDK. + * + * Each test: + * 1. Builds an in-memory Open Plugin filesystem containing a `.mcp.json` + * plus any server source the upstream needs. + * 2. Starts a real agent host server (no mock agent). + * 3. Hooks `resourceList` / `resourceRead` reverse-RPC handlers on the + * test client so the host's plugin manager can copy the plugin via + * `vscode-agent-client://` URIs without anything ever touching disk. + * 4. Creates a `copilotcli` session whose `activeClient.customizations` + * references the in-memory plugin. The agent eagerly snapshots and + * publishes the MCP server immediately — no warmup turn required. + * 5. Waits for the MCP server to reach `Ready` (handling + * `AuthRequired` for the HTTP test). + * 6. Dispatches a real LLM turn that asks the model to invoke + * `say_hello`, auto-approves any tool-confirmation prompts, and + * asserts the resulting tool output text — proving the upstream + * MCP server was actually exercised end-to-end through the SDK. + * + * These tests are **disabled by default**. Enable them with + * `AGENT_HOST_REAL_SDK=1` (matching `toolApprovalRealSdk.integrationTest.ts`). + */ + +import assert from 'assert'; +import { execSync } from 'child_process'; +import { mkdtempSync, rmSync } from 'fs'; +import type { IncomingMessage, Server as HttpServer, ServerResponse } from 'http'; +import type { AddressInfo } from 'net'; +import { tmpdir } from 'os'; +import { fileURLToPath } from 'url'; +import { encodeBase64, VSBuffer } from '../../../../../base/common/buffer.js'; +import { dirname, join, resolve as resolvePath } from '../../../../../base/common/path.js'; +import { URI } from '../../../../../base/common/uri.js'; +import { ContentEncoding, SubscribeResult, type DirectoryEntry, type McpMethodCallParams, type McpMethodCallResult, type ResourceListParams, type ResourceListResult, type ResourceReadParams, type ResourceReadResult } from '../../../common/state/protocol/commands.js'; +import type { SessionToolCallCompleteAction, SessionToolCallReadyAction, SessionToolCallStartAction } from '../../../common/state/protocol/actions.js'; +import { McpServerStatusKind, ToolResultContentType, type AhpMcpUiHostCapabilities, type McpServerStatus, type SessionState, type ToolCallResult, type ToolResultContent, type ToolResultTextContent } from '../../../common/state/protocol/state.js'; +import { PROTOCOL_VERSION } from '../../../common/state/protocol/version/registry.js'; +import { ProtocolError } from '../../../common/state/sessionProtocol.js'; +import { + getActionEnvelope, + isActionNotification, + IServerHandle, + startRealServer, + TestProtocolClient, +} from './testHelpers.js'; + +/** Set `AGENT_HOST_REAL_SDK=1` to run these tests. */ +const REAL_SDK_ENABLED = process.env['AGENT_HOST_REAL_SDK'] === '1'; + +/** AHP `NotFound` (-32008) — surfaced for missing files/directories. */ +const RPC_RESOURCE_NOT_FOUND = -32008; + +/** + * Resolve the GitHub token used to authenticate the Copilot SDK. Falls + * back to `gh auth token` when `GITHUB_TOKEN` is not set, mirroring + * `toolApprovalRealSdk.integrationTest.ts`. + */ +function resolveGitHubToken(): string { + const envToken = process.env['GITHUB_TOKEN']; + if (envToken) { + return envToken; + } + try { + return execSync('gh auth token', { encoding: 'utf-8' }).trim(); + } catch { + throw new Error('No GITHUB_TOKEN set and `gh auth token` failed. Run `gh auth login` first.'); + } +} + +// --------------------------------------------------------------------------- +// In-memory plugin filesystem +// --------------------------------------------------------------------------- + +/** + * A tiny in-memory plugin filesystem keyed by relative paths under a single + * root URI. Serves the two reverse-RPC methods the agent host issues during + * plugin sync (`resourceList` / `resourceRead`). + * + * Also exposes every ancestor of {@link rootUri} as a virtual directory + * containing only the next-deeper segment, so that `stat`-style probes the + * agent host fires (which list the parent to identify a child's type) + * succeed all the way up to the URI root. + */ +class InMemoryPluginFilesystem { + + /** Map from absolute URI string → file contents. */ + private readonly _files = new Map(); + /** + * Map from absolute URI string → set of immediate child names (with + * their types). Pre-populated for the root + every ancestor of the + * root + every directory implied by a file path. + */ + private readonly _dirs = new Map>(); + + constructor(public readonly rootUri: URI, files: Record) { + this._registerAncestors(rootUri); + + for (const [relPath, content] of Object.entries(files)) { + const absolute = this._resolveRelative(relPath); + this._files.set(absolute.toString(), content); + // Parent of the file is a directory containing it; bubble up + // to the root, registering each intermediate directory. + let cursor = absolute; + let childName = this._basename(cursor.path); + let childType: 'file' | 'directory' = 'file'; + while (true) { + const parent = this._parent(cursor); + this._addEntry(parent, childName, childType); + if (parent.toString() === this.rootUri.toString()) { + break; + } + cursor = parent; + childName = this._basename(cursor.path); + childType = 'directory'; + } + } + } + + private _registerAncestors(uri: URI): void { + // Ensure rootUri itself is a (possibly empty) directory. + if (!this._dirs.has(uri.toString())) { + this._dirs.set(uri.toString(), new Map()); + } + let cursor = uri; + while (cursor.path !== '/' && cursor.path !== '') { + const parent = this._parent(cursor); + const name = this._basename(cursor.path); + this._addEntry(parent, name, 'directory'); + cursor = parent; + } + } + + private _parent(uri: URI): URI { + const trimmed = uri.path.replace(/\/+$/, ''); + const idx = trimmed.lastIndexOf('/'); + const parentPath = idx <= 0 ? '/' : trimmed.substring(0, idx); + return uri.with({ path: parentPath }); + } + + private _basename(path: string): string { + const trimmed = path.replace(/\/+$/, ''); + const idx = trimmed.lastIndexOf('/'); + return idx === -1 ? trimmed : trimmed.substring(idx + 1); + } + + private _addEntry(dir: URI, name: string, type: 'file' | 'directory'): void { + const key = dir.toString(); + let entries = this._dirs.get(key); + if (!entries) { + entries = new Map(); + this._dirs.set(key, entries); + } + const existing = entries.get(name); + // Files take precedence over auto-promoted directory entries; once a + // child is registered as a file, never demote it back to a directory. + if (existing === 'file' && type === 'directory') { + return; + } + entries.set(name, type); + } + + private _resolveRelative(relPath: string): URI { + const normalized = relPath.replace(/\\/g, '/').replace(/^\/+/, ''); + const root = this.rootUri.toString().replace(/\/$/, ''); + return URI.parse(`${root}/${normalized}`); + } + + list(uri: string): ResourceListResult { + if (this._files.has(uri)) { + throw new ProtocolError(RPC_RESOURCE_NOT_FOUND, `Not a directory: ${uri}`); + } + const entries = this._dirs.get(uri); + if (!entries) { + throw new ProtocolError(RPC_RESOURCE_NOT_FOUND, `Directory not found: ${uri}`); + } + const list: DirectoryEntry[] = []; + for (const [name, type] of entries) { + list.push({ name, type }); + } + list.sort((a, b) => a.name.localeCompare(b.name)); + return { entries: list }; + } + + read(uri: string): ResourceReadResult { + const content = this._files.get(uri); + if (content === undefined) { + throw new ProtocolError(RPC_RESOURCE_NOT_FOUND, `File not found: ${uri}`); + } + if (typeof content === 'string') { + return { data: content, encoding: ContentEncoding.Utf8, contentType: 'text/plain' }; + } + return { + data: encodeBase64(VSBuffer.wrap(content)), + encoding: ContentEncoding.Base64, + contentType: 'application/octet-stream', + }; + } +} + +// --------------------------------------------------------------------------- +// Hello-world MCP server source +// --------------------------------------------------------------------------- + +/** + * Source of a minimal stdio MCP server built on `@modelcontextprotocol/sdk`. + * Registers a single `say_hello` tool that returns `"Hello, !"`. + * + * Resolved against the workspace's `node_modules` via `NODE_PATH` so the + * spawned child can `require()` the SDK from any plugin directory. + */ +/** URI the upstream MCP server advertises for the `say_hello` tool's UI view. */ +const HELLO_MCP_UI_RESOURCE_URI = 'ui://hello-mcp/main'; +/** HTML body the MCP server serves at {@link HELLO_MCP_UI_RESOURCE_URI}. */ +const HELLO_MCP_UI_HTML = 'hello-mcp

Hello, MCP Apps!

'; + +const HELLO_MCP_SERVER_JS = ` +'use strict'; +const { Server } = require('@modelcontextprotocol/sdk/server/index.js'); +const { StdioServerTransport } = require('@modelcontextprotocol/sdk/server/stdio.js'); +const { + CallToolRequestSchema, + ListResourcesRequestSchema, + ListToolsRequestSchema, + ReadResourceRequestSchema, +} = require('@modelcontextprotocol/sdk/types.js'); + +const UI_RESOURCE_URI = ${JSON.stringify(HELLO_MCP_UI_RESOURCE_URI)}; +const UI_HTML = ${JSON.stringify(HELLO_MCP_UI_HTML)}; + +const server = new Server( + { name: 'hello-mcp', version: '1.0.0' }, + { + capabilities: { + // Explicitly advertise the MCP-Apps host capabilities the AHP + // proxy should mirror back through \`_meta.uiHostCapabilities\` + // on every tool call: tool / resource discovery with change + // notifications, plus log forwarding. + tools: { listChanged: true }, + resources: { listChanged: true }, + logging: {}, + }, + }, +); + +server.setRequestHandler(ListToolsRequestSchema, async () => ({ + tools: [{ + name: 'say_hello', + description: 'Returns a friendly greeting for the given name.', + inputSchema: { + type: 'object', + properties: { name: { type: 'string', description: 'Name to greet.' } }, + required: ['name'], + }, + _meta: { + // MCP Apps spec — tools advertise a UI resource that the host + // renders alongside the tool result. + ui: { + resourceUri: UI_RESOURCE_URI, + visibility: ['model', 'app'], + }, + }, + }], +})); + +server.setRequestHandler(CallToolRequestSchema, async (request) => { + if (request.params.name !== 'say_hello') { + throw new Error('Unknown tool: ' + request.params.name); + } + const target = (request.params.arguments && request.params.arguments.name) || 'world'; + return { content: [{ type: 'text', text: 'Hello, ' + target + '!' }] }; +}); + +server.setRequestHandler(ListResourcesRequestSchema, async () => ({ + resources: [{ + uri: UI_RESOURCE_URI, + name: 'hello-app', + mimeType: 'text/html;profile=mcp-app', + }], +})); + +server.setRequestHandler(ReadResourceRequestSchema, async (request) => { + if (request.params.uri !== UI_RESOURCE_URI) { + throw new Error('Unknown resource: ' + request.params.uri); + } + return { + contents: [{ + uri: UI_RESOURCE_URI, + mimeType: 'text/html;profile=mcp-app', + text: UI_HTML, + }], + }; +}); + +(async () => { + const transport = new StdioServerTransport(); + await server.connect(transport); +})().catch(err => { + console.error('hello-mcp server failed:', err); + process.exit(1); +}); +`; + +/** + * Resolve the absolute path to the workspace's `node_modules` directory so + * the spawned MCP server can locate `@modelcontextprotocol/sdk` (and its + * peer `zod`) regardless of where the plugin gets synced on disk. + */ +function getWorkspaceNodeModules(): string { + const here = fileURLToPath(import.meta.url); + // out/vs/platform/agentHost/test/node/protocol/.js — 7 segments + // from the file directory up to the vscode repository root. + const repoRoot = resolvePath(dirname(here), '..', '..', '..', '..', '..', '..', '..'); + return join(repoRoot, 'node_modules'); +} + +// --------------------------------------------------------------------------- +// Shared helpers +// --------------------------------------------------------------------------- + +/** + * Wire reverse-RPC handlers for `resourceList` / `resourceRead` so the + * agent host's `AHPFileSystemProvider` can sync the in-memory plugin via + * `vscode-agent-client://` URIs. Must be installed BEFORE `createSession` + * is called because the agent's plugin sync runs inline with that call. + */ +function installPluginFsHandlers(client: TestProtocolClient, plugin: InMemoryPluginFilesystem): void { + client.setRequestHandler('resourceList', async params => { + return plugin.list((params as ResourceListParams).uri.toString()); + }); + client.setRequestHandler('resourceRead', async params => { + return plugin.read((params as ResourceReadParams).uri); + }); +} + +/** + * Initialize the AHP client, authenticate against the Copilot API, create + * a `copilotcli` session, subscribe to it, and then publish the + * customization via `session/activeClientChanged`. We deliberately split + * `createSession` and `activeClientChanged` rather than passing + * `activeClient` on `createSession`: + * + * - Historically, publishing MCP servers during `createSession` raced + * the state-manager session entry. The MCP host now keeps those + * summaries as its own state while suppressing pre-session action + * emission; this split flow still exercises the live + * `activeClientChanged` republish path. + * + * - Dispatching `activeClientChanged` after `subscribe` flows through + * `agentSideEffects` → `setClientCustomizations` → plugin sync → + * `PluginController.onDidChange` → `ActiveClient.republish()` → + * `setSessionServers(...)`. By that point the session entry exists + * and the client is subscribed, so every MCP lifecycle action is + * routed to us. + */ +async function createSessionWithPlugin( + client: TestProtocolClient, + clientId: string, + pluginUri: string, + sessionScheme: string, + sessionLabel: string, +): Promise { + await client.call('initialize', { protocolVersions: [PROTOCOL_VERSION], clientId }, 30_000); + await client.call('authenticate', { resource: 'https://api.github.com', token: resolveGitHubToken() }, 30_000); + + const sessionUri = URI.from({ scheme: sessionScheme, path: `/${sessionLabel}-${Date.now()}` }).toString(); + await client.call('createSession', { + session: sessionUri, + provider: 'copilotcli', + workingDirectory: URI.file(tmpdir()).toString(), + }, 30_000); + await client.call('subscribe', { resource: sessionUri }, 30_000); + + client.notify('dispatchAction', { + clientSeq: 1, + action: { + type: 'session/activeClientChanged', + session: sessionUri, + activeClient: { + clientId, + displayName: 'MCP Test Client', + tools: [], + customizations: [{ + uri: pluginUri, + displayName: 'MCP Test Plugin', + nonce: 'v1', + }], + }, + }, + }); + + return sessionUri; +} + +/** + * Wait for an `session/mcpServerStatusChanged` action whose status + * matches `predicate`, returning the matched `McpServerStatus` payload + * and the affected server's `mcp:/...` URI. + */ +async function waitForMcpServerStatus( + client: TestProtocolClient, + predicate: (status: McpServerStatus) => boolean, + timeoutMs: number, +): Promise<{ readonly status: McpServerStatus; readonly mcpServer: string }> { + const notification = await client.waitForNotification(n => { + if (!isActionNotification(n, 'session/mcpServerStatusChanged')) { + return false; + } + const action = getActionEnvelope(n).action as { status: McpServerStatus }; + return predicate(action.status); + }, timeoutMs); + const action = getActionEnvelope(notification).action as { status: McpServerStatus; mcpServer: string }; + return { status: action.status, mcpServer: action.mcpServer }; +} + +/** + * Re-subscribe to `sessionUri` and return the named MCP server's + * latest summary (or `undefined` if no such server is registered). + */ +async function readMcpServerSummary( + client: TestProtocolClient, + sessionUri: string, + label: string, +) { + const subscribe = await client.call('subscribe', { resource: sessionUri }, 30_000); + const state = subscribe.snapshot.state as SessionState; + return state.mcpServers?.find(s => s.label === label); +} + +/** + * Aggregate the `text` payloads from a `ToolCallResult.content` array + * into a single string. Returns the empty string when the result has no + * text content (e.g. binary-only output, which `say_hello` never produces). + */ +function toolResultText(result: ToolCallResult): string { + const content: readonly ToolResultContent[] = result.content ?? []; + return content + .filter((c): c is ToolResultTextContent => c.type === ToolResultContentType.Text) + .map(c => c.text) + .join(''); +} + +interface IObservedToolCall { + readonly toolName: string; + readonly result: ToolCallResult; + /** The `session/toolCallStart` action that opened the matched tool call. */ + readonly start: SessionToolCallStartAction; +} + +/** + * Dispatch a real chat turn and drive it until either the model invokes + * a tool matching `predicate` (returned), or the turn finishes/errors + * without matching (throws). + * + * Auto-approves every `session/toolCallReady` that arrives without a + * `confirmed` flag — Copilot's MCP integration surfaces tool calls as + * permission requests, so we approve them on the AHP client's behalf. + */ +async function driveTurnUntilTool( + client: TestProtocolClient, + sessionUri: string, + turnId: string, + prompt: string, + startSeq: number, + predicate: (toolName: string) => boolean, + timeoutMs: number, +): Promise { + client.clearReceived(); + client.notify('dispatchAction', { + clientSeq: startSeq, + action: { + type: 'session/turnStarted', + session: sessionUri, + turnId, + userMessage: { text: prompt }, + }, + }); + + let nextSeq = startSeq + 1; + const seen = new Set(); + const startsById = new Map(); + + const deadline = Date.now() + timeoutMs; + while (true) { + const remaining = deadline - Date.now(); + if (remaining <= 0) { + throw new Error(`driveTurnUntilTool: timed out after ${timeoutMs}ms waiting for a tool call match in turn '${turnId}'`); + } + const notification = await client.waitForNotification(n => { + if (seen.has(n as object)) { + return false; + } + return isActionNotification(n, 'session/toolCallStart') + || isActionNotification(n, 'session/toolCallReady') + || isActionNotification(n, 'session/toolCallComplete') + || isActionNotification(n, 'session/turnComplete') + || isActionNotification(n, 'session/turnCancelled') + || isActionNotification(n, 'session/error'); + }, remaining); + seen.add(notification as object); + + if (isActionNotification(notification, 'session/error')) { + throw new Error(`Session error while driving turn '${turnId}'`); + } + + if (isActionNotification(notification, 'session/turnComplete') || isActionNotification(notification, 'session/turnCancelled')) { + throw new Error(`Turn '${turnId}' ended before the model invoked a matching tool. Recorded tool calls: ${[...startsById.values()].map(a => a.toolName).join(', ') || ''}`); + } + + if (isActionNotification(notification, 'session/toolCallStart')) { + const action = getActionEnvelope(notification).action as SessionToolCallStartAction; + startsById.set(action.toolCallId, action); + continue; + } + + if (isActionNotification(notification, 'session/toolCallReady')) { + const action = getActionEnvelope(notification).action as SessionToolCallReadyAction; + if (!action.confirmed) { + client.notify('dispatchAction', { + clientSeq: nextSeq++, + action: { + type: 'session/toolCallConfirmed', + session: sessionUri, + turnId, + toolCallId: action.toolCallId, + approved: true, + }, + }); + } + continue; + } + + // toolCallComplete + const action = getActionEnvelope(notification).action as SessionToolCallCompleteAction; + const start = startsById.get(action.toolCallId); + const toolName = start?.toolName ?? ''; + if (start && predicate(toolName)) { + return { toolName, result: action.result, start }; + } + } +} + +// --------------------------------------------------------------------------- +// Suite +// --------------------------------------------------------------------------- + +(REAL_SDK_ENABLED ? suite : suite.skip)('Protocol WebSocket — Real MCP SDK', function () { + + let server: IServerHandle; + let client: TestProtocolClient; + let userDataDir: string; + const createdSessions: string[] = []; + const tempDirs: string[] = []; + + suiteSetup(async function () { + this.timeout(60_000); + // Use a dedicated user-data dir so plugin syncs land in a temp location + // that is wiped between runs. + userDataDir = mkdtempSync(join(tmpdir(), 'ahp-mcp-real-')); + tempDirs.push(userDataDir); + server = await startRealServer({ userDataDir }); + }); + + suiteTeardown(function () { + server?.process.kill(); + if (process.env['AGENT_HOST_REAL_SDK_KEEP_TEMP'] === '1') { + console.log('[Real MCP SDK] kept temp dirs:', tempDirs); + return; + } + for (const dir of tempDirs) { + try { + rmSync(dir, { recursive: true, force: true, maxRetries: 5, retryDelay: 200 }); + } catch { + // best-effort cleanup — Windows may still hold transient file handles + // from the just-killed server. + } + } + }); + + setup(async function () { + this.timeout(30_000); + client = new TestProtocolClient(server.port); + await client.connect(); + }); + + teardown(async function () { + this.timeout(30_000); + for (const session of createdSessions) { + try { + await client.call('disposeSession', { session }, 10_000); + } catch { + // best-effort + } + } + createdSessions.length = 0; + client.close(); + }); + + test('stdio MCP plugin: model calls say_hello end-to-end via Copilot SDK', async function () { + this.timeout(120_000); + + const clientId = 'real-mcp-stdio-client'; + const pluginUri = 'file:///inmem/hello-mcp'; + + // `${PLUGIN_ROOT}` is expanded by the agent host at parse time to the + // synced plugin's local fsPath, so the spawned child reads `server.js` + // from the on-disk copy. + const mcpJson = { + mcpServers: { + hello: { + command: 'node', + args: ['${PLUGIN_ROOT}/server.js'], + env: { NODE_PATH: getWorkspaceNodeModules() }, + }, + }, + }; + const plugin = new InMemoryPluginFilesystem(URI.parse(pluginUri), { + '.plugin/plugin.json': JSON.stringify({ name: 'hello-mcp', version: '1.0.0' }), + '.mcp.json': JSON.stringify(mcpJson), + 'server.js': HELLO_MCP_SERVER_JS, + }); + + // Reverse-RPC handlers MUST be installed before `activeClientChanged` + // because the agent's plugin sync issues `resourceList`/`resourceRead` + // callbacks while the client is being applied. + installPluginFsHandlers(client, plugin); + + const sessionUri = await createSessionWithPlugin(client, clientId, pluginUri, 'copilotcli', 'mcp-real-stdio'); + createdSessions.push(sessionUri); + + // End-to-end single turn: dispatching `turnStarted` here triggers + // `_materializeProvisional`, which calls `ActiveClient.snapshot()` + // → publishes MCP servers → spawns the stdio child → + // `_resolveMcpServersForSdk` blocks until Ready → SDK gets the + // tool list including `say_hello`. The model is asked to call the + // tool; the model's tool call traverses SDK → loopback proxy → + // stdio child → say_hello → back through the proxy → SDK → + // `toolCallComplete` on the wire. We auto-approve any + // `toolCallReady` that needs confirmation along the way. + const observed = await driveTurnUntilTool( + client, + sessionUri, + 'turn-mcp-stdio-real', + 'Please call the `say_hello` MCP tool with the argument `{"name": "integration"}`. Do not explain — just invoke the tool.', + 2, + toolName => /say_hello/.test(toolName), + 90_000, + ); + + const text = toolResultText(observed.result); + assert.match( + text, + /Hello, integration!/, + `expected say_hello tool result to contain the greeting; got ${JSON.stringify(observed.result)}`, + ); + + const helloServer = await readMcpServerSummary(client, sessionUri, 'hello'); + assert.ok(helloServer, 'expected hello MCP server on the session'); + assert.strictEqual(helloServer!.status.kind, McpServerStatusKind.Ready); + + // MCP Apps: the proxy sniffed `tools/list` and the agent decorated + // the `toolCallStart` action with `_meta.ui` (mirroring the MCP + // `_meta."io.modelcontextprotocol/ui"` payload) plus the per-tool + // `_meta.uiHostCapabilities` set the AHP host can satisfy. + const startMeta = (observed.start._meta ?? {}) as { + ui?: { resourceUri?: string; visibility?: readonly string[] }; + uiHostCapabilities?: AhpMcpUiHostCapabilities; + }; + assert.deepStrictEqual( + startMeta.ui, + { resourceUri: HELLO_MCP_UI_RESOURCE_URI, visibility: ['model', 'app'] }, + `expected toolCallStart._meta.ui to mirror the MCP server payload; got ${JSON.stringify(startMeta)}`, + ); + assert.ok( + startMeta.uiHostCapabilities?.serverTools?.listChanged, + `expected toolCallStart._meta.uiHostCapabilities.serverTools.listChanged; got ${JSON.stringify(startMeta.uiHostCapabilities)}`, + ); + assert.ok( + startMeta.uiHostCapabilities?.serverResources?.listChanged, + `expected toolCallStart._meta.uiHostCapabilities.serverResources.listChanged; got ${JSON.stringify(startMeta.uiHostCapabilities)}`, + ); + assert.ok( + startMeta.uiHostCapabilities?.logging !== undefined, + `expected toolCallStart._meta.uiHostCapabilities.logging; got ${JSON.stringify(startMeta.uiHostCapabilities)}`, + ); + + // The AHP client can fetch the app's HTML body by invoking + // `mcpMethodCall` with the standard MCP `resources/read` method — + // no special envelope is needed beyond pointing at the server. + const readResult = await client.call('mcpMethodCall', { + server: helloServer!.resource, + method: 'resources/read', + params: { uri: HELLO_MCP_UI_RESOURCE_URI }, + } satisfies McpMethodCallParams, 15_000); + const contents = (readResult.result as { contents?: { uri: string; mimeType?: string; text?: string }[] } | undefined)?.contents; + assert.ok(contents && contents.length === 1, `expected one resource content block; got ${JSON.stringify(readResult)}`); + assert.strictEqual(contents[0].uri, HELLO_MCP_UI_RESOURCE_URI); + assert.strictEqual(contents[0].mimeType, 'text/html;profile=mcp-app'); + assert.strictEqual(contents[0].text, HELLO_MCP_UI_HTML); + }); + + test('HTTP MCP plugin: AuthRequired → authenticate → model calls say_hello', async function () { + this.timeout(120_000); + + const clientId = 'real-mcp-http-client'; + const pluginUri = 'file:///inmem/secure-mcp'; + const expectedToken = 'test-bearer-token-' + Date.now(); + const oauthResource = `local-secure-mcp-${Date.now()}`; + + const httpServer = await startSecureHttpMcpServer({ expectedToken, oauthResource }); + try { + const mcpJson = { + mcpServers: { + 'secure-hello': { + type: 'http', + url: httpServer.url, + }, + }, + }; + const plugin = new InMemoryPluginFilesystem(URI.parse(pluginUri), { + '.plugin/plugin.json': JSON.stringify({ name: 'secure-mcp', version: '1.0.0' }), + '.mcp.json': JSON.stringify(mcpJson), + }); + + installPluginFsHandlers(client, plugin); + + const sessionUri = await createSessionWithPlugin(client, clientId, pluginUri, 'copilotcli', 'mcp-real-http'); + createdSessions.push(sessionUri); + + // Phase 1: dispatch a throwaway turn whose only purpose is to + // trigger `_materializeProvisional`. Inside materialization, + // `ActiveClient.snapshot()` publishes the MCP server, which + // performs the HTTP probe, gets a 401, and surfaces + // `AuthRequired` with metadata fetched from + // `/.well-known/oauth-protected-resource`. The turn itself is + // allowed to run (the SDK skips the auth-required server, so + // the model has no MCP tools yet); we cancel it once we have + // the auth challenge. + client.notify('dispatchAction', { + clientSeq: 2, + action: { + type: 'session/turnStarted', + session: sessionUri, + turnId: 'turn-mcp-http-warmup', + userMessage: { text: 'Reply with the single word OK.' }, + }, + }); + + const authReqEvent = await waitForMcpServerStatus( + client, + s => s.kind === McpServerStatusKind.AuthRequired, + 60_000, + ); + const authStatus = authReqEvent.status; + assert.strictEqual(authStatus.kind, McpServerStatusKind.AuthRequired); + if (authStatus.kind !== McpServerStatusKind.AuthRequired) { + throw new Error('unreachable'); + } + assert.strictEqual( + authStatus.resource.resource, + oauthResource, + `expected fetched resource_metadata to carry the synthetic resource id; got ${JSON.stringify(authStatus.resource)}`, + ); + + client.notify('dispatchAction', { + clientSeq: 3, + action: { + type: 'session/turnCancelled', + session: sessionUri, + turnId: 'turn-mcp-http-warmup', + }, + }); + + // Phase 2: push a token. The host re-runs the upstream `start()` + // with the bearer header and transitions the server to Ready. + // `authenticate` resolves with `{}` on success and throws on + // failure (per protocol — the empty result IS the success + // signal), so reaching the next line means the host accepted + // the token. + await client.call<{}>('authenticate', { + resource: authStatus.resource.resource, + token: expectedToken, + server: authReqEvent.mcpServer, + }, 15_000); + + await waitForMcpServerStatus(client, s => s.kind === McpServerStatusKind.Ready, 30_000); + + const secureServer = await readMcpServerSummary(client, sessionUri, 'secure-hello'); + assert.ok(secureServer, 'expected an MCP server named secure-hello on the session'); + assert.strictEqual(secureServer!.status.kind, McpServerStatusKind.Ready); + + // Phase 3: ask the model to invoke `say_hello`. The SDK's + // cached session was built when `secure-hello` was still + // `AuthRequired`, so it doesn't yet know about `say_hello`. + // `ActiveClient.isOutdated()` compares the snapshot's per-MCP + // `mcpReadiness` map against the current Ready set, sees that + // `secure-hello` flipped from false → true, and triggers a + // full SDK-session rebuild on the next `sendMessage`. The new + // `_resolveMcpServersForSdk` sees `Ready` and advertises the + // tool to the model. + const observed = await driveTurnUntilTool( + client, + sessionUri, + 'turn-mcp-http-real', + 'Please call the `say_hello` MCP tool with the argument `{"name": "authed"}`. Do not explain — just invoke the tool.', + 4, + toolName => /say_hello/.test(toolName), + 90_000, + ); + + const text = toolResultText(observed.result); + assert.match( + text, + /Hello, authed!/, + `expected say_hello tool result to contain the greeting; got ${JSON.stringify(observed.result)}`, + ); + + // Sanity: the upstream HTTP server must have observed at least + // one unauth'd probe (the discovery handshake) and at least one + // authenticated request thereafter — proving the host service + // propagated the bearer token through `setBearerToken` and the + // SDK actually invoked the tool against the secured endpoint. + assert.ok( + httpServer.authedRequestCount > 0, + 'expected at least one authenticated HTTP request to the upstream', + ); + assert.ok( + httpServer.unauthedRequestCount >= 1, + `expected at least one unauthenticated probe; got ${httpServer.unauthedRequestCount}`, + ); + } finally { + await httpServer.close(); + } + }); +}); + +// --------------------------------------------------------------------------- +// Tiny HTTP MCP server with bearer-token auth +// --------------------------------------------------------------------------- + +interface ISecureHttpMcpServer { + readonly url: string; + readonly close: () => Promise; + readonly authedRequestCount: number; + readonly unauthedRequestCount: number; +} + +interface ISecureHttpMcpServerOptions { + readonly expectedToken: string; + /** + * Synthetic OAuth resource identifier returned in + * `/.well-known/oauth-protected-resource`. The test uses this to + * assert the host actually fetched the metadata document rather than + * synthesizing fallback metadata from the request URL. + */ + readonly oauthResource: string; +} + +/** + * Spin up a localhost HTTP server that mimics a remote MCP endpoint with + * RFC 9728-style bearer auth. Without a token, the `/mcp` POST returns + * 401 + a `WWW-Authenticate: Bearer resource_metadata="…"` header that + * points at `/.well-known/oauth-protected-resource` on the same origin. + * With the configured token it answers JSON-RPC `initialize`, + * `tools/list`, and `tools/call`. + * + * Listening on `127.0.0.1:0` keeps the test isolated from any real + * service and matches the loopback-only convention the agent host uses + * for its own proxy listener. + */ +async function startSecureHttpMcpServer(options: ISecureHttpMcpServerOptions): Promise { + const counters = { authed: 0, unauthed: 0 }; + const { createServer } = await import('http'); + + const handler = async (req: IncomingMessage, res: ServerResponse): Promise => { + const baseUrl = `http://${req.headers.host ?? '127.0.0.1'}`; + const url = new URL(req.url ?? '/', baseUrl); + + if (req.method === 'GET' && url.pathname === '/.well-known/oauth-protected-resource') { + res.writeHead(200, { 'Content-Type': 'application/json' }); + res.end(JSON.stringify({ + resource: options.oauthResource, + authorization_servers: [`${baseUrl}/auth`], + scopes_supported: ['mcp:read'], + })); + return; + } + + if (req.method !== 'POST' || url.pathname !== '/mcp') { + res.writeHead(404, { 'Content-Type': 'text/plain' }); + res.end('not found'); + return; + } + + const auth = req.headers['authorization']; + if (auth !== `Bearer ${options.expectedToken}`) { + counters.unauthed++; + res.writeHead(401, { + 'Content-Type': 'application/json', + 'WWW-Authenticate': `Bearer resource_metadata="${baseUrl}/.well-known/oauth-protected-resource"`, + }); + res.end(JSON.stringify({ error: 'unauthorized' })); + return; + } + + counters.authed++; + const body = await readRequestBody(req); + let parsed: { id?: number | string; method: string; params?: unknown }; + try { + parsed = JSON.parse(body); + } catch { + res.writeHead(400, { 'Content-Type': 'application/json' }); + res.end(JSON.stringify({ error: 'invalid json' })); + return; + } + + const reply = handleMcpMethod(parsed); + if (reply === undefined) { + // Notification — no response body, HTTP 204. + res.writeHead(204); + res.end(); + return; + } + res.writeHead(200, { 'Content-Type': 'application/json' }); + res.end(JSON.stringify({ jsonrpc: '2.0', id: parsed.id, ...reply })); + }; + + const server: HttpServer = createServer((req, res) => { + Promise.resolve(handler(req, res)).catch(err => { + res.writeHead(500, { 'Content-Type': 'text/plain' }); + res.end(err instanceof Error ? err.message : String(err)); + }); + }); + + await new Promise((resolve, reject) => { + server.once('error', reject); + server.listen(0, '127.0.0.1', () => { + server.removeListener('error', reject); + resolve(); + }); + }); + + const address = server.address() as AddressInfo; + const url = `http://127.0.0.1:${address.port}/mcp`; + + return { + url, + close: () => new Promise(resolve => server.close(() => resolve())), + get authedRequestCount() { return counters.authed; }, + get unauthedRequestCount() { return counters.unauthed; }, + }; +} + +function readRequestBody(req: IncomingMessage): Promise { + return new Promise((resolve, reject) => { + const chunks: Buffer[] = []; + req.on('data', (chunk: Buffer) => chunks.push(chunk)); + req.on('end', () => resolve(Buffer.concat(chunks).toString('utf-8'))); + req.on('error', reject); + }); +} + +function handleMcpMethod(message: { id?: number | string; method: string; params?: unknown }): { result: unknown } | { error: { code: number; message: string } } | undefined { + switch (message.method) { + case 'initialize': + return { + result: { + protocolVersion: '2024-11-05', + capabilities: { tools: {} }, + serverInfo: { name: 'secure-hello', version: '1.0.0' }, + }, + }; + case 'notifications/initialized': + return undefined; + case 'tools/list': + return { + result: { + tools: [{ + name: 'say_hello', + description: 'Returns a friendly greeting for the given name.', + inputSchema: { + type: 'object', + properties: { name: { type: 'string' } }, + required: ['name'], + }, + }], + }, + }; + case 'tools/call': { + const params = (message.params ?? {}) as { name?: string; arguments?: { name?: string } }; + if (params.name !== 'say_hello') { + return { error: { code: -32601, message: `Unknown tool: ${params.name}` } }; + } + const target = params.arguments?.name ?? 'world'; + return { result: { content: [{ type: 'text', text: `Hello, ${target}!` }] } }; + } + default: + return { error: { code: -32601, message: `Method not found: ${message.method}` } }; + } +} diff --git a/src/vs/platform/agentHost/test/node/protocol/testHelpers.ts b/src/vs/platform/agentHost/test/node/protocol/testHelpers.ts index 6ecadc8c1c1cc..901ba54240997 100644 --- a/src/vs/platform/agentHost/test/node/protocol/testHelpers.ts +++ b/src/vs/platform/agentHost/test/node/protocol/testHelpers.ts @@ -12,8 +12,11 @@ import type { ActionEnvelope, SessionAddedNotification } from '../../../common/s import { PROTOCOL_VERSION } from '../../../common/state/protocol/version/registry.js'; import { isJsonRpcNotification, + isJsonRpcRequest, isJsonRpcResponse, + ProtocolError, type AhpNotification, + type AhpRequest, type JsonRpcErrorResponse, type JsonRpcSuccessResponse, type INotificationBroadcastParams, @@ -27,12 +30,22 @@ interface IPendingCall { reject: (err: Error) => void; } +/** + * Handler for an incoming reverse-RPC request from the server. + * Throw a {@link ProtocolError} to send a structured JSON-RPC error + * response back to the server; any other error is sent as + * `JSON_RPC_INTERNAL_ERROR` (`-32603`). + */ +export type ReverseRequestHandler = (params: unknown) => Promise | unknown; + export class TestProtocolClient { private readonly _ws: WebSocket; private _nextId = 1; private readonly _pendingCalls = new Map(); private readonly _notifications: AhpNotification[] = []; private readonly _notifWaiters: { predicate: (n: AhpNotification) => boolean; resolve: (n: AhpNotification) => void; reject: (err: Error) => void }[] = []; + /** Reverse-RPC handlers — the server sends requests like `resourceRead`/`resourceList`. */ + private readonly _requestHandlers = new Map(); constructor(port: number) { this._ws = new WebSocket(`ws://127.0.0.1:${port}`); @@ -64,6 +77,8 @@ export class TestProtocolClient { pending.resolve((msg as JsonRpcSuccessResponse).result); } } + } else if (isJsonRpcRequest(msg)) { + this._handleIncomingRequest(msg); } else if (isJsonRpcNotification(msg)) { const notif = msg; for (let i = this._notifWaiters.length - 1; i >= 0; i--) { @@ -76,6 +91,42 @@ export class TestProtocolClient { } } + private _handleIncomingRequest(msg: AhpRequest): void { + const handler = this._requestHandlers.get(msg.method); + if (!handler) { + this._ws.send(JSON.stringify({ + jsonrpc: '2.0', + id: msg.id, + error: { code: -32601, message: `Unknown reverse-RPC method: ${msg.method}` }, + })); + return; + } + Promise.resolve() + .then(() => handler(msg.params)) + .then(result => { + this._ws.send(JSON.stringify({ jsonrpc: '2.0', id: msg.id, result: result ?? null })); + }, err => { + const error = err instanceof ProtocolError + ? { code: err.code, message: err.message, ...(err.data !== undefined ? { data: err.data } : {}) } + : { code: -32603, message: err instanceof Error ? err.message : String(err) }; + this._ws.send(JSON.stringify({ jsonrpc: '2.0', id: msg.id, error })); + }); + } + + /** + * Register a handler for an incoming reverse-RPC request method. Replaces + * any prior handler for the same method. Returns a disposer that removes + * the handler. + */ + setRequestHandler(method: string, handler: ReverseRequestHandler): () => void { + this._requestHandlers.set(method, handler); + return () => { + if (this._requestHandlers.get(method) === handler) { + this._requestHandlers.delete(method); + } + }; + } + /** Send a JSON-RPC notification (fire-and-forget). */ notify(method: string, params?: unknown): void { this._ws.send(JSON.stringify({ jsonrpc: '2.0', method, params })); @@ -225,12 +276,16 @@ export async function startServer(options?: { readonly quiet?: boolean; readonly * Start the agent host server with the real Copilot SDK agent (no mock agent). * The server is started with logging enabled so the CopilotAgent is registered. */ -export async function startRealServer(): Promise { +export async function startRealServer(options?: { readonly userDataDir?: string; readonly env?: NodeJS.ProcessEnv }): Promise { return new Promise((resolve, reject) => { const serverPath = fileURLToPath(new URL('../../../node/agentHostServerMain.js', import.meta.url)); const args = ['--port', '0', '--without-connection-token']; + if (options?.userDataDir) { + args.push('--user-data-dir', options.userDataDir); + } const child = fork(serverPath, args, { stdio: ['pipe', 'pipe', 'pipe', 'ipc'], + env: options?.env ? { ...process.env, ...options.env } : process.env, }); const timer = setTimeout(() => { @@ -249,6 +304,9 @@ export async function startRealServer(): Promise { child.stderr!.on('data', () => { // Intentionally swallowed - the test runner fails if console.error is used. + // Set `AGENT_HOST_REAL_SDK_KEEP_TEMP=1` on the test to keep the + // agent host's user-data-dir; the structured `agenthost-server.log` + // inside it is much easier to read than raw stderr. }); child.on('error', err => { diff --git a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts index a566d700a5094..68f6eb797da84 100644 --- a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts +++ b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts @@ -9,17 +9,20 @@ import { DisposableStore } from '../../../../base/common/lifecycle.js'; import { URI } from '../../../../base/common/uri.js'; import { runWithFakedTimers } from '../../../../base/test/common/timeTravelScheduler.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import { IMcpServerDefinition } from '../../../agentPlugins/common/pluginParsers.js'; import { NullLogService } from '../../../log/common/log.js'; -import { type IAgentCreateSessionConfig, type IAgentResolveSessionConfigParams, type IAgentService, type IAgentSessionConfigCompletionsParams, type IAgentSessionMetadata, type AuthenticateParams, type AuthenticateResult } from '../../common/agentService.js'; -import { CompletionsParams, CompletionsResult, ListSessionsResult, ResourceReadResult, ResolveSessionConfigResult, SessionConfigCompletionsResult } from '../../common/state/protocol/commands.js'; -import { ActionType, type IRootConfigChangedAction, type SessionAction, type TerminalAction } from '../../common/state/sessionActions.js'; +import { AgentHostFileSystemProvider } from '../../common/agentHostFileSystemProvider.js'; +import { type AuthenticateParams, type AuthenticateResult, type IAgentCreateSessionConfig, type IAgentResolveSessionConfigParams, type IAgentService, type IAgentSessionConfigCompletionsParams, type IAgentSessionMetadata } from '../../common/agentService.js'; +import { IMcpHostService, IMcpHostUpstreamDelegate, IMcpServerHandle } from '../../common/mcpHost/mcpHostService.js'; +import { CompletionsParams, CompletionsResult, ListSessionsResult, McpMethodCallParams, McpMethodCallResult, McpNotificationParams, ResolveSessionConfigResult, ResourceReadResult, SessionConfigCompletionsResult, type ClientCapabilities } from '../../common/state/protocol/commands.js'; +import { McpServerStatusKind, type McpServerSummary } from '../../common/state/protocol/state.js'; import { PROTOCOL_VERSION } from '../../common/state/protocol/version/registry.js'; -import { isJsonRpcNotification, isJsonRpcResponse, JSON_RPC_INTERNAL_ERROR, ProtocolError, AHP_UNSUPPORTED_PROTOCOL_VERSION, type AhpNotification, type InitializeResult, type ProtocolMessage, type ReconnectResult, type ResourceListResult, type ResourceWriteParams, type ResourceWriteResult, type IStateSnapshot } from '../../common/state/sessionProtocol.js'; +import { ActionType, type IRootConfigChangedAction, type SessionAction, type TerminalAction } from '../../common/state/sessionActions.js'; +import { AHP_UNSUPPORTED_PROTOCOL_VERSION, isJsonRpcNotification, isJsonRpcRequest, isJsonRpcResponse, JSON_RPC_INTERNAL_ERROR, JsonRpcErrorCodes, ProtocolError, type AhpNotification, type InitializeResult, type IStateSnapshot, type ProtocolMessage, type ReconnectResult, type ResourceListResult, type ResourceWriteParams, type ResourceWriteResult } from '../../common/state/sessionProtocol.js'; import { ResponsePartKind, SessionStatus, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType, type SessionSummary } from '../../common/state/sessionState.js'; import type { IProtocolServer, IProtocolTransport } from '../../common/state/sessionTransport.js'; -import { ProtocolServerHandler } from '../../node/protocolServerHandler.js'; import { AgentHostStateManager } from '../../node/agentHostStateManager.js'; -import { AgentHostFileSystemProvider } from '../../common/agentHostFileSystemProvider.js'; +import { ProtocolServerHandler } from '../../node/protocolServerHandler.js'; // ---- Mock helpers ----------------------------------------------------------- @@ -156,6 +159,48 @@ class MockAgentService implements IAgentService { // ---- Helpers ---------------------------------------------------------------- +/** + * Test double for {@link IMcpHostService} that records every `callMethod`, + * `notify`, and upstream delegate registration. Tests can swap in a custom + * `callMethodImpl` (e.g. throw) by reassigning the field. + */ +class RecordingMcpHostService implements IMcpHostService { + declare readonly _serviceBrand: undefined; + + readonly callMethodCalls: { params: McpMethodCallParams }[] = []; + readonly notifyCalls: { params: McpNotificationParams }[] = []; + delegate: IMcpHostUpstreamDelegate | undefined; + + callMethodImpl: (params: McpMethodCallParams) => Promise = + async () => ({ result: 'ok' }); + + setSessionServers(_session: URI, _servers: readonly IMcpServerDefinition[]): readonly IMcpServerHandle[] { + return []; + } + + getServerSummaries(_session: URI): readonly McpServerSummary[] { + return []; + } + + getServer(_resource: URI): IMcpServerHandle | undefined { + return undefined; + } + + async callMethod(params: McpMethodCallParams): Promise { + this.callMethodCalls.push({ params }); + return this.callMethodImpl(params); + } + + notify(params: McpNotificationParams): void { + this.notifyCalls.push({ params }); + } + + setUpstreamDelegate(delegate: IMcpHostUpstreamDelegate) { + this.delegate = delegate; + return { dispose: () => { this.delegate = undefined; } }; + } +} + function notification(method: string, params?: unknown): ProtocolMessage { return { jsonrpc: '2.0', method, params } as ProtocolMessage; } @@ -185,6 +230,7 @@ suite('ProtocolServerHandler', () => { let server: MockProtocolServer; let agentService: MockAgentService; let handler: ProtocolServerHandler; + let recordingMcpHostService: RecordingMcpHostService; const sessionUri = URI.from({ scheme: 'copilot', path: '/test-session' }).toString(); @@ -218,12 +264,14 @@ suite('ProtocolServerHandler', () => { agentService = new MockAgentService(); agentService.setStateManager(stateManager); disposables.add(agentService); + recordingMcpHostService = new RecordingMcpHostService(); disposables.add(handler = new ProtocolServerHandler( agentService, stateManager, server, { defaultDirectory: URI.file('/home/testuser').toString() }, disposables.add(new AgentHostFileSystemProvider()), + recordingMcpHostService, new NullLogService(), )); }); @@ -244,6 +292,33 @@ suite('ProtocolServerHandler', () => { assert.strictEqual(result.serverSeq, stateManager.serverSeq); }); + test('initialize response includes empty server capabilities', () => { + const transport = connectClient('client-caps-1'); + + const resp = findResponse(transport.sent, 1); + assert.ok(resp, 'should have sent initialize response'); + const result = (resp as { result: InitializeResult }).result; + assert.deepStrictEqual(result.capabilities, {}); + }); + + test('initialize captures client capabilities for later use', () => { + const transport = new MockProtocolTransport(); + server.simulateConnection(transport); + // ClientCapabilities is now an empty extension point; passing an empty + // object exercises the capture path without depending on any specific + // field shape. + const clientCaps: ClientCapabilities = {}; + transport.simulateMessage(request(1, 'initialize', { + protocolVersions: [PROTOCOL_VERSION], + clientId: 'client-caps-2', + capabilities: clientCaps, + })); + + const resp = findResponse(transport.sent, 1); + assert.ok(resp, 'should have sent initialize response'); + assert.deepStrictEqual(handler.getClientCapabilitiesForTest('client-caps-2'), clientCaps); + }); + test('handshake rejects unsupported protocol versions', () => { const transport = new MockProtocolTransport(); server.simulateConnection(transport); @@ -945,6 +1020,41 @@ suite('ProtocolServerHandler', () => { agentService.authenticate = origHandler; }); + test('authenticate forwards params.server as a parsed URI to the agent service', async () => { + const origHandler = agentService.authenticate; + const seen: AuthenticateParams[] = []; + agentService.authenticate = async (params: AuthenticateParams): Promise => { + seen.push(params); + return { authenticated: true }; + }; + + const transport = connectClient('client-auth-mcp'); + transport.sent.length = 0; + + const responsePromise = waitForResponse(transport, 2); + transport.simulateMessage(request(2, 'authenticate', { + resource: 'https://api.github.com', + token: 'tok', + server: 'mcp:/sess1/srv1', + })); + const resp = await responsePromise as { result?: Record; error?: { code: number; message: string } }; + + assert.ok(!resp.error, `unexpected error: ${resp.error?.message}`); + assert.strictEqual(seen.length, 1); + assert.ok(URI.isUri(seen[0].server), 'server must be a URI instance'); + assert.deepStrictEqual({ + resource: seen[0].resource, + token: seen[0].token, + server: seen[0].server!.toString(), + }, { + resource: 'https://api.github.com', + token: 'tok', + server: 'mcp:/sess1/srv1', + }); + + agentService.authenticate = origHandler; + }); + // ---- Connection count event ----------------------------------------- test('onDidChangeConnectionCount fires on connect and disconnect', () => { @@ -986,6 +1096,136 @@ suite('ProtocolServerHandler', () => { assert.deepStrictEqual(counts, [1, 1, 0]); }); + // ---- MCP routing ----------------------------------------------------- + + suite('MCP', () => { + + const mcpServer1 = 'mcp:/sess1/srv1'; + const mcpServer2 = 'mcp:/sess1/srv2'; + + function createMcpSessionWithServers(): void { + const summary = makeSessionSummary(); + stateManager.createSession(summary); + stateManager.dispatchServerAction({ type: ActionType.SessionReady, session: sessionUri }); + stateManager.createMcpServer(sessionUri, { + resource: mcpServer1, + label: 'srv1', + status: { kind: McpServerStatusKind.Ready }, + }); + stateManager.createMcpServer(sessionUri, { + resource: mcpServer2, + label: 'srv2', + status: { kind: McpServerStatusKind.Ready }, + }); + } + + test('mcpMethodCall forwards params to IMcpHostService and returns result', async () => { + const transport = new MockProtocolTransport(); + server.simulateConnection(transport); + transport.simulateMessage(request(1, 'initialize', { + protocolVersions: [PROTOCOL_VERSION], + clientId: 'client-mcp-msg', + })); + transport.sent.length = 0; + + const params: McpMethodCallParams = { + server: mcpServer1, + method: 'tools/call', + params: { foo: 1 }, + }; + + const responsePromise = waitForResponse(transport, 2); + transport.simulateMessage(request(2, 'mcpMethodCall', params)); + const resp = await responsePromise as { result?: { result: unknown }; error?: { code: number } }; + + assert.deepStrictEqual({ + calls: recordingMcpHostService.callMethodCalls, + response: resp, + }, { + calls: [{ params }], + response: { jsonrpc: '2.0', id: 2, result: { result: 'ok' } }, + }); + + transport.simulateClose(); + transport.dispose(); + }); + + test('mcpMethodCall propagates ProtocolError as JSON-RPC error response', async () => { + recordingMcpHostService.callMethodImpl = async () => { + throw new ProtocolError(JsonRpcErrorCodes.MethodNotFound, 'no such method'); + }; + const transport = connectClient('client-mcp-err'); + transport.sent.length = 0; + const responsePromise = waitForResponse(transport, 2); + + transport.simulateMessage(request(2, 'mcpMethodCall', { + server: mcpServer1, + method: 'tools/call', + })); + const resp = await responsePromise as { error?: { code: number; message: string } }; + + assert.deepStrictEqual( + { code: resp.error?.code, message: resp.error?.message }, + { code: JsonRpcErrorCodes.MethodNotFound, message: 'no such method' }, + ); + }); + + test('mcpNotification notification is forwarded to IMcpHostService.notify', () => { + const transport = connectClient('client-mcp-notify'); + transport.sent.length = 0; + + const params: McpNotificationParams = { + server: mcpServer1, + method: 'notifications/message', + params: { level: 'info', data: 'hello' }, + }; + transport.simulateMessage(notification('mcpNotification', params)); + + assert.deepStrictEqual(recordingMcpHostService.notifyCalls, [{ params }]); + }); + + test('upstream MCP request is round-tripped via reverse mcpMethodCall to a connected client', async () => { + createMcpSessionWithServers(); + const transport = connectClient('client-mcp', [sessionUri]); + transport.sent.length = 0; + + const upstreamPromise = recordingMcpHostService.delegate!.handleUpstreamRequest({ + server: URI.parse(mcpServer1), + method: 'sampling/createMessage', + params: { foo: 1 }, + }); + + // `_sendReverseRequest` synchronously dispatches the outbound JSON-RPC request, + // so the mcpMethodCall message is already in `transport.sent`. + const sent = transport.sent.find(m => isJsonRpcRequest(m) && m.method === 'mcpMethodCall'); + assert.ok(sent && isJsonRpcRequest(sent)); + transport.simulateMessage({ jsonrpc: '2.0', id: sent.id, result: { ok: true } } as ProtocolMessage); + + const outcome = await upstreamPromise; + assert.deepStrictEqual(outcome, { result: { ok: true } }); + }); + + test('upstream MCP notification is forwarded as an outbound mcpNotification', () => { + createMcpSessionWithServers(); + const transport = connectClient('client-mcp', [sessionUri]); + transport.sent.length = 0; + + recordingMcpHostService.delegate!.handleUpstreamNotification({ + server: URI.parse(mcpServer1), + method: 'notifications/tools/list_changed', + params: undefined, + }); + + const sent = transport.sent.find(m => isJsonRpcNotification(m) && m.method === 'mcpNotification'); + assert.ok(sent && isJsonRpcNotification(sent)); + assert.deepStrictEqual(sent.params, { + server: mcpServer1, + method: 'notifications/tools/list_changed', + params: undefined, + }); + }); + }); + // ---- createSession activeClient ------------------------------------- suite('createSession activeClient', () => { diff --git a/src/vs/platform/agentHost/test/node/reducers.test.ts b/src/vs/platform/agentHost/test/node/reducers.test.ts index 0c333e38fc71f..b60a1e0c7e24e 100644 --- a/src/vs/platform/agentHost/test/node/reducers.test.ts +++ b/src/vs/platform/agentHost/test/node/reducers.test.ts @@ -7,6 +7,7 @@ import assert from 'assert'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import { sessionReducer } from '../../common/state/protocol/reducers.js'; import { ActionType } from '../../common/state/sessionActions.js'; +import { McpServerStatusKind, type McpServerStatus, type McpServerSummary } from '../../common/state/protocol/state.js'; import { SessionInputAnswerState, SessionInputAnswerValueKind, SessionInputQuestionKind, SessionInputResponseKind, SessionLifecycle, SessionStatus, ToolCallConfirmationReason, type SessionState } from '../../common/state/sessionState.js'; function makeSession(): SessionState { @@ -192,3 +193,89 @@ suite('sessionReducer – summaryStatus with tool call confirmations and input r assert.strictEqual(state.summary.status, SessionStatus.InputNeeded); }); }); + +suite('sessionReducer – mcp lifecycle', () => { + + ensureNoDisposablesAreLeakedInTestSuite(); + + const SESSION_URI = 'copilot:/test'; + + function makeServer(overrides?: Partial): McpServerSummary { + return { + resource: 'mcp:/sess1/server1', + label: 'server1', + status: { kind: McpServerStatusKind.Starting }, + ...overrides, + }; + } + + test('McpServerAdded inserts a single summary into mcpServers', () => { + const server = makeServer(); + const state = sessionReducer(makeSession(), { + type: ActionType.McpServerAdded, + session: SESSION_URI, + server, + }); + + assert.deepStrictEqual(state.mcpServers, [server]); + }); + + test('McpServerAdded with a duplicate resource replaces the prior entry', () => { + let state = sessionReducer(makeSession(), { + type: ActionType.McpServerAdded, + session: SESSION_URI, + server: makeServer(), + }); + + const replaced = makeServer({ + label: 'server1-renamed', + status: { kind: McpServerStatusKind.Ready }, + }); + state = sessionReducer(state, { + type: ActionType.McpServerAdded, + session: SESSION_URI, + server: replaced, + }); + + assert.deepStrictEqual(state.mcpServers, [replaced]); + }); + + test('McpServerRemoved drops the matching entry', () => { + const server = makeServer(); + let state = sessionReducer(makeSession(), { + type: ActionType.McpServerAdded, + session: SESSION_URI, + server, + }); + state = sessionReducer(state, { + type: ActionType.McpServerRemoved, + session: SESSION_URI, + mcpServer: server.resource, + }); + + assert.strictEqual(state.mcpServers, undefined); + }); + + test('McpServerStatusChanged updates only the status of the matching entry', () => { + const server = makeServer(); + let state = sessionReducer(makeSession(), { + type: ActionType.McpServerAdded, + session: SESSION_URI, + server, + }); + + const newStatus: McpServerStatus = { kind: McpServerStatusKind.Ready }; + state = sessionReducer(state, { + type: ActionType.McpServerStatusChanged, + session: SESSION_URI, + mcpServer: server.resource, + status: newStatus, + }); + + assert.deepStrictEqual(state.mcpServers, [{ + resource: server.resource, + label: server.label, + status: newStatus, + }]); + }); +}); diff --git a/src/vs/sessions/contrib/agentHost/browser/baseAgentHostSessionsProvider.ts b/src/vs/sessions/contrib/agentHost/browser/baseAgentHostSessionsProvider.ts index 56cdabda52c47..07415f187ad68 100644 --- a/src/vs/sessions/contrib/agentHost/browser/baseAgentHostSessionsProvider.ts +++ b/src/vs/sessions/contrib/agentHost/browser/baseAgentHostSessionsProvider.ts @@ -20,16 +20,20 @@ import { AgentSession, IAgentConnection, IAgentSessionMetadata } from '../../../ import { KNOWN_AUTO_APPROVE_VALUES, SessionConfigKey } from '../../../../platform/agentHost/common/sessionConfigKeys.js'; import { ResolveSessionConfigResult } from '../../../../platform/agentHost/common/state/protocol/commands.js'; import { NotificationType } from '../../../../platform/agentHost/common/state/protocol/notifications.js'; -import { FileEdit, ModelSelection, SessionStatus as ProtocolSessionStatus, RootConfigState, RootState, SessionState, SessionSummary } from '../../../../platform/agentHost/common/state/protocol/state.js'; +import { FileEdit, ModelSelection, SessionStatus as ProtocolSessionStatus, RootConfigState, RootState, SessionState, SessionSummary, type McpServerSummary } from '../../../../platform/agentHost/common/state/protocol/state.js'; import { ActionType, isSessionAction } from '../../../../platform/agentHost/common/state/sessionActions.js'; import { readSessionGitState, SessionMeta, StateComponents, type ISessionGitState } from '../../../../platform/agentHost/common/state/sessionState.js'; import { IConfigurationService } from '../../../../platform/configuration/common/configuration.js'; import { ILogService } from '../../../../platform/log/common/log.js'; import { ChatViewPaneTarget, IChatWidgetService } from '../../../../workbench/contrib/chat/browser/chat.js'; +import { authenticateMcpServerCandidates } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.js'; +import { IAgentHostActiveClientRegistry } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostActiveClientRegistry.js'; +import { IAgentHostMcpAuthRegistry, type IAgentHostMcpAuthSessionEntry } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthRegistry.js'; import { IChatSendRequestOptions, IChatService } from '../../../../workbench/contrib/chat/common/chatService/chatService.js'; import { IChatSessionFileChange, IChatSessionFileChange2, IChatSessionsService } from '../../../../workbench/contrib/chat/common/chatSessionsService.js'; import { ChatAgentLocation, ChatConfiguration, ChatModeKind } from '../../../../workbench/contrib/chat/common/constants.js'; import { ILanguageModelsService } from '../../../../workbench/contrib/chat/common/languageModels.js'; +import { IAuthenticationService } from '../../../../workbench/services/authentication/common/authentication.js'; import { buildMutableConfigSchema, IAgentHostSessionsProvider, resolvedConfigsEqual } from '../../../common/agentHostSessionsProvider.js'; import { agentHostSessionWorkspaceKey } from '../../../common/agentHostSessionWorkspace.js'; import { isSessionConfigComplete } from '../../../common/sessionConfig.js'; @@ -363,6 +367,30 @@ export class AgentHostSessionAdapter implements ISession { // NewSession — bundles the in-flight new-session state // ============================================================================ +/** + * Per-NewSession dependencies needed to surface MCP-auth state and + * drive interactive auth before the session has been graduated into a + * real chat session (i.e. before `provideChatSessionContent` runs). + * + * Mirrors the auth machinery `AgentHostSessionHandler` uses for active + * sessions; both surfaces share `agentHostAuth.ts` helpers so the + * silent / remembered-scopes semantics stay identical across the two + * code paths. + */ +export interface IProvisionalMcpAuthBinding { + readonly registry: IAgentHostMcpAuthRegistry; + readonly authenticationService: IAuthenticationService; + /** + * Stable agent-host identifier (`'local'` for the in-process host, + * the remote address for remote hosts). Combined with the OAuth + * resource URI to scope the registry's persistent scope memory so + * that approving the same OAuth resource on a different agent host + * prompts the user again instead of silently reusing a token from + * another host. + */ + readonly mcpAuthHostKey: string; +} + /** * Inputs needed to construct a {@link NewSession}. */ @@ -381,6 +409,22 @@ interface INewSessionConstructionContext { * the picker reflects the user's preference immediately. */ readonly initialConfigValues?: Record; + /** + * Optional MCP-auth binding. When provided, the {@link NewSession} + * registers itself with the workbench MCP-auth registry once its + * eager backend session opens, so the chat-input MCP-auth indicator + * appears for provisional sessions on the new-chat view. + */ + readonly mcpAuthBinding?: IProvisionalMcpAuthBinding; + /** + * Optional registry for retrieving the chat-session handler's + * active-client snapshot (`clientId`, client tools, customizations) + * by `sessionType`. Passed straight through to + * `connection.createSession`'s `activeClient` parameter so the agent + * host establishes the active client at creation time without a + * follow-up `session/activeClientChanged` dispatch. + */ + readonly activeClientRegistry?: IAgentHostActiveClientRegistry; } /** @@ -440,6 +484,23 @@ class NewSession extends Disposable { private readonly _logService: ILogService; private readonly _providerId: string; + /** + * Chat-session type identifier used to look up the active-client + * snapshot. Equals `resourceScheme` (matches the value the chat + * session contribution registers under) — the agent provider id + * stored on {@link agentProvider} differs from this for remote + * hosts, so we capture the right key explicitly. + */ + private readonly _chatSessionType: string; + private readonly _mcpAuthBinding: IProvisionalMcpAuthBinding | undefined; + private readonly _activeClientRegistry: IAgentHostActiveClientRegistry | undefined; + /** + * MCP-auth registration cleanup. Lazily filled in when + * {@link eagerCreate}'s subscription opens; cleared on + * {@link graduate} or dispose so the registry entry doesn't outlive + * the eager backend session. + */ + private readonly _mcpAuthRegistration = this._register(new MutableDisposable()); constructor(ctx: INewSessionConstructionContext) { super(); @@ -450,7 +511,10 @@ class NewSession extends Disposable { this.workspaceUri = workspaceUri; this.agentProvider = ctx.sessionType.id; this._providerId = ctx.providerId; + this._chatSessionType = ctx.resourceScheme; this._logService = ctx.logService; + this._mcpAuthBinding = ctx.mcpAuthBinding; + this._activeClientRegistry = ctx.activeClientRegistry; const resource = URI.from({ scheme: ctx.resourceScheme, path: `/${generateUuid()}` }); this._status = observableValue(this, SessionStatus.Untitled); @@ -606,10 +670,18 @@ class NewSession extends Disposable { void (async () => { try { + // Look up the chat-session handler's active-client snapshot + // (`clientId`, tools, customizations) for this session type + // and pass it through. The agent host treats this as + // equivalent to dispatching `session/activeClientChanged` + // immediately after creation, so customizations and client + // tools are in place for the very first turn. + const activeClient = this._activeClientRegistry?.get(this._chatSessionType); await connection.createSession({ provider: this.agentProvider, session: backendUri, workingDirectory: this.workspaceUri, + activeClient, }); } catch (err) { this._logService.warn(`[${this._providerId}] Eager createSession failed for ${backendUri.toString()}: ${err}`); @@ -637,9 +709,68 @@ class NewSession extends Disposable { // when chat content opens, so when we release this ref on // graduation the wire-level refcount stays positive. this._subscription = connection.getSubscription(StateComponents.Session, backendUri); + + // Register an MCP-auth entry against the chat session resource + // so the chat-input indicator (and the new-chat-input indicator + // in particular) can surface MCP servers in `AuthRequired` + // before the user sends their first message — the + // `provideChatSessionContent` registration only runs at + // graduation. Skipped silently when no binding was provided + // (tests, surfaces that don't surface the indicator). + this._registerMcpAuth(connection); })(); } + /** + * Wires `_subscription`'s `mcpServers` slice into the workbench-scoped + * MCP-auth registry. Called from {@link eagerCreate} after the + * subscription opens; the resulting registration lives until + * {@link graduate} or dispose, at which point the chat-session-bound + * registration takes over (or the session goes away). + */ + private _registerMcpAuth(connection: IAgentConnection): void { + const binding = this._mcpAuthBinding; + if (!binding) { + return; + } + const sub = this._subscription; + if (!sub) { + return; + } + const store = new DisposableStore(); + const typedSub = sub.object as { value: SessionState | Error | undefined; onDidChange: Event }; + const mcpServersObs: ISettableObservable = observableValue(this, []); + const update = () => { + const v = typedSub.value; + mcpServersObs.set((v && !(v instanceof Error)) ? (v.mcpServers ?? []) : [], undefined); + }; + update(); + store.add(typedSub.onDidChange(update)); + + const logPrefix = `[NewSession-MCP/${this._providerId}]`; + const entry: IAgentHostMcpAuthSessionEntry = { + mcpServers: mcpServersObs, + authenticate: target => authenticateMcpServerCandidates( + mcpServersObs.get(), + target, + { + authenticate: req => connection.authenticate({ + resource: req.resource, + token: req.token, + server: req.server, + }), + authenticationService: binding.authenticationService, + logService: this._logService, + logPrefix, + mcpAuthMemory: binding.registry, + mcpAuthHostKey: binding.mcpAuthHostKey, + }, + ), + }; + store.add(binding.registry.registerSession(this.session.resource, entry)); + this._mcpAuthRegistration.value = store; + } + /** * Release the backend subscription without firing `disposeSession`. * Used on the success path in `sendAndCreateChat` when the session has @@ -651,6 +782,10 @@ class NewSession extends Disposable { this._backendUri = undefined; this._connection = undefined; this._configRequestSeq++; + // Drop our MCP-auth registration; the chat-session handler will + // re-register an equivalent entry against the same resource via + // `provideChatSessionContent`. + this._mcpAuthRegistration.clear(); } override dispose(): void { @@ -778,6 +913,9 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement @IConfigurationService protected readonly _baseConfigurationService: IConfigurationService, @ILogService protected readonly _logService: ILogService, @IGitHubService protected readonly _gitHubService: IGitHubService, + @IAuthenticationService protected readonly _authenticationService: IAuthenticationService, + @IAgentHostMcpAuthRegistry protected readonly _mcpAuthRegistry: IAgentHostMcpAuthRegistry, + @IAgentHostActiveClientRegistry protected readonly _activeClientRegistry: IAgentHostActiveClientRegistry, ) { super(); } @@ -790,6 +928,16 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement /** Provider-level authentication-pending observable used to derive `loading` for sessions. */ protected abstract get authenticationPending(): IObservable; + /** + * Stable identifier for the agent host this provider talks to, + * used as the `host` axis of {@link IAgentHostMcpAuthRegistry}'s + * persistent scope memory. `'local'` for the in-process host, + * the remote address for remote hosts. The host axis prevents + * authorizing an OAuth resource on one host from silently + * authenticating another. + */ + protected abstract get _mcpAuthHostKey(): string; + /** * Subclass-specific portion of the adapter options. Base fills in * the bits that are uniform across hosts (`icon`, `loading`, @@ -975,6 +1123,12 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement authenticationPending: this.authenticationPending, logService: this._logService, initialConfigValues: this._initialNewSessionConfig(), + mcpAuthBinding: { + registry: this._mcpAuthRegistry, + authenticationService: this._authenticationService, + mcpAuthHostKey: this._mcpAuthHostKey, + }, + activeClientRegistry: this._activeClientRegistry, }); this._newSession = newSession; this._onDidChangeSessionConfig.fire(newSession.sessionId); diff --git a/src/vs/sessions/contrib/agentHost/browser/localAgentHostSessionsProvider.ts b/src/vs/sessions/contrib/agentHost/browser/localAgentHostSessionsProvider.ts index f307eb2d35cab..764914d4fa8cd 100644 --- a/src/vs/sessions/contrib/agentHost/browser/localAgentHostSessionsProvider.ts +++ b/src/vs/sessions/contrib/agentHost/browser/localAgentHostSessionsProvider.ts @@ -17,9 +17,12 @@ import { IConfigurationService } from '../../../../platform/configuration/common import { ILabelService } from '../../../../platform/label/common/label.js'; import { ILogService } from '../../../../platform/log/common/log.js'; import { IChatWidgetService } from '../../../../workbench/contrib/chat/browser/chat.js'; +import { IAgentHostActiveClientRegistry } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostActiveClientRegistry.js'; +import { IAgentHostMcpAuthRegistry } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthRegistry.js'; import { IChatService } from '../../../../workbench/contrib/chat/common/chatService/chatService.js'; import { IChatSessionsService } from '../../../../workbench/contrib/chat/common/chatSessionsService.js'; import { ILanguageModelsService } from '../../../../workbench/contrib/chat/common/languageModels.js'; +import { IAuthenticationService } from '../../../../workbench/services/authentication/common/authentication.js'; import { BaseAgentHostSessionsProvider } from './baseAgentHostSessionsProvider.js'; import { buildAgentHostSessionWorkspace, readBranchProtectionPatterns } from '../../../common/agentHostSessionWorkspace.js'; import { ISessionWorkspace, ISessionWorkspaceBrowseAction, SESSION_WORKSPACE_GROUP_LOCAL } from '../../../services/sessions/common/session.js'; @@ -58,8 +61,11 @@ export class LocalAgentHostSessionsProvider extends BaseAgentHostSessionsProvide @IConfigurationService private readonly _configurationService: IConfigurationService, @ILogService logService: ILogService, @IGitHubService gitHubService: IGitHubService, + @IAuthenticationService authenticationService: IAuthenticationService, + @IAgentHostMcpAuthRegistry mcpAuthRegistry: IAgentHostMcpAuthRegistry, + @IAgentHostActiveClientRegistry activeClientRegistry: IAgentHostActiveClientRegistry, ) { - super(chatSessionsService, chatService, chatWidgetService, languageModelsService, _configurationService, logService, gitHubService); + super(chatSessionsService, chatService, chatWidgetService, languageModelsService, _configurationService, logService, gitHubService, authenticationService, mcpAuthRegistry, activeClientRegistry); this.label = localize('localAgentHostLabel', "Local Agent Host"); @@ -101,6 +107,8 @@ export class LocalAgentHostSessionsProvider extends BaseAgentHostSessionsProvide protected get authenticationPending(): IObservable { return this._agentHostService.authenticationPending; } + protected get _mcpAuthHostKey(): string { return 'local'; } + /** * Local resource scheme: `agent-host-${provider}`. Must match the type * string registered by AgentHostContribution. Distinct from the logical diff --git a/src/vs/sessions/contrib/agentHost/test/browser/localAgentHostSessionsProvider.test.ts b/src/vs/sessions/contrib/agentHost/test/browser/localAgentHostSessionsProvider.test.ts index 1497e4c857bc0..06ca89c3bfba8 100644 --- a/src/vs/sessions/contrib/agentHost/test/browser/localAgentHostSessionsProvider.test.ts +++ b/src/vs/sessions/contrib/agentHost/test/browser/localAgentHostSessionsProvider.test.ts @@ -24,9 +24,12 @@ import { TestConfigurationService } from '../../../../../platform/configuration/ import { IFileDialogService } from '../../../../../platform/dialogs/common/dialogs.js'; import { TestInstantiationService } from '../../../../../platform/instantiation/test/common/instantiationServiceMock.js'; import { IChatWidget, IChatWidgetService } from '../../../../../workbench/contrib/chat/browser/chat.js'; +import { IAgentHostActiveClientRegistry } from '../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostActiveClientRegistry.js'; +import { IAgentHostMcpAuthRegistry } from '../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthRegistry.js'; import { IChatService, type ChatSendResult, type IChatSendRequestOptions } from '../../../../../workbench/contrib/chat/common/chatService/chatService.js'; import { IChatSessionsService } from '../../../../../workbench/contrib/chat/common/chatSessionsService.js'; import { ILanguageModelsService } from '../../../../../workbench/contrib/chat/common/languageModels.js'; +import { IAuthenticationService } from '../../../../../workbench/services/authentication/common/authentication.js'; import { ISessionChangeEvent } from '../../../../services/sessions/common/sessionsProvider.js'; import { SessionStatus } from '../../../../services/sessions/common/session.js'; import { LocalAgentHostSessionsProvider } from '../../browser/localAgentHostSessionsProvider.js'; @@ -251,6 +254,13 @@ function createProvider(disposables: DisposableStore, agentHostService: MockAgen instantiationService.stub(IGitHubService, new class extends mock() { override findPullRequestNumberByHeadBranch = async () => undefined; }()); + instantiationService.stub(IAuthenticationService, new class extends mock() { }()); + instantiationService.stub(IAgentHostMcpAuthRegistry, new class extends mock() { + override registerSession = () => ({ dispose() { } }); + }()); + instantiationService.stub(IAgentHostActiveClientRegistry, new class extends mock() { + override get = () => undefined; + }()); return disposables.add(instantiationService.createInstance(LocalAgentHostSessionsProvider)); } diff --git a/src/vs/sessions/contrib/chat/browser/agentHost/newChatMcpAuthIndicator.ts b/src/vs/sessions/contrib/chat/browser/agentHost/newChatMcpAuthIndicator.ts new file mode 100644 index 0000000000000..865eb247671c6 --- /dev/null +++ b/src/vs/sessions/contrib/chat/browser/agentHost/newChatMcpAuthIndicator.ts @@ -0,0 +1,92 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import * as dom from '../../../../../base/browser/dom.js'; +import { Button } from '../../../../../base/browser/ui/button/button.js'; +import { Codicon } from '../../../../../base/common/codicons.js'; +import { Disposable } from '../../../../../base/common/lifecycle.js'; +import { autorun } from '../../../../../base/common/observable.js'; +import { localize } from '../../../../../nls.js'; +import { IContextMenuService } from '../../../../../platform/contextview/browser/contextView.js'; +import { McpServerStatusKind } from '../../../../../platform/agentHost/common/state/protocol/state.js'; +import { IAgentHostMcpAuthRegistry } from '../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthRegistry.js'; +import { showMcpAuthContextMenu } from '../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/chatMcpAuthAction.js'; +import { ISessionsManagementService } from '../../../../services/sessions/common/sessionsManagement.js'; +import { ThemeIcon } from '../../../../../base/common/themables.js'; + +/** + * Toolbar indicator for the new-chat-input that surfaces MCP servers + * needing authentication on the currently active session. Mirrors the + * chat-input toolbar's {@link OpenMcpAuthAction} but renders directly + * because the new-chat-input doesn't host `MenuId.ChatExecute`. + * + * Visibility tracks + * {@link ISessionsManagementService.activeSession}: when the active + * session has at least one MCP server in `AuthRequired` state (looked + * up via {@link IAgentHostMcpAuthRegistry}), the indicator appears; + * otherwise it's hidden. Clicking opens the same context menu used by + * the chat-input action. + */ +export class NewChatMcpAuthIndicator extends Disposable { + + private readonly _button: Button; + private readonly _container: HTMLElement; + + constructor( + container: HTMLElement, + @ISessionsManagementService private readonly _sessionsManagementService: ISessionsManagementService, + @IAgentHostMcpAuthRegistry private readonly _registry: IAgentHostMcpAuthRegistry, + @IContextMenuService private readonly _contextMenuService: IContextMenuService, + ) { + super(); + + this._container = dom.append(container, dom.$('.sessions-chat-mcp-auth-indicator')); + this._container.classList.add('hidden'); + + this._button = this._register(new Button(this._container, { + secondary: true, + title: localize('newChatMcpAuth.title', "Authenticate MCP Servers"), + ariaLabel: localize('newChatMcpAuth.ariaLabel', "One or more MCP servers require authentication"), + })); + // Match the codicon used by `OpenMcpAuthAction` so users see the + // same affordance on both surfaces. + this._button.element.classList.add(...ThemeIcon.asClassNameArray(Codicon.mcp)); + + this._register(this._button.onDidClick(() => this._onClick())); + + // Single autorun: re-evaluates when the active session changes, + // when a registry entry is registered/unregistered for the + // current session resource (race: the session handler often + // registers AFTER this contribution first reads the active + // session), or when the entry's `mcpServers` observable + // changes. + this._register(autorun(reader => { + const activeSession = this._sessionsManagementService.activeSession.read(reader); + const entry = activeSession + ? this._registry.getEntry(activeSession.resource, reader) + : undefined; + if (!entry) { + this._setHidden(true); + return; + } + const summaries = entry.mcpServers.read(reader); + const hasAuthRequired = summaries.some(s => s.status.kind === McpServerStatusKind.AuthRequired); + this._setHidden(!hasAuthRequired); + })); + } + + private _setHidden(hidden: boolean): void { + this._container.classList.toggle('hidden', hidden); + } + + private _onClick(): void { + const activeSession = this._sessionsManagementService.activeSession.get(); + const entry = activeSession ? this._registry.getEntry(activeSession.resource) : undefined; + if (!entry) { + return; + } + showMcpAuthContextMenu(entry, this._button.element, this._contextMenuService); + } +} diff --git a/src/vs/sessions/contrib/chat/browser/media/chatInput.css b/src/vs/sessions/contrib/chat/browser/media/chatInput.css index dba8994ad6033..8871b152f8298 100644 --- a/src/vs/sessions/contrib/chat/browser/media/chatInput.css +++ b/src/vs/sessions/contrib/chat/browser/media/chatInput.css @@ -86,6 +86,55 @@ flex: 1; } +/* MCP-auth indicator (see `NewChatMcpAuthIndicator`). Hidden via the + `.hidden` class while the active session has no servers in + `AuthRequired`. Color and pulsing animation mirror + `chatMcpAuthAction.css` so the indicator looks the same on both + surfaces. */ +.sessions-chat-mcp-auth-indicator { + display: flex; + align-items: center; +} + +.sessions-chat-mcp-auth-indicator.hidden { + display: none; +} + +.sessions-chat-mcp-auth-indicator .monaco-button { + background: transparent; + border: none; + padding: 2px 4px; + color: var(--vscode-notificationsWarningIcon-foreground); + cursor: pointer; + display: flex; + align-items: center; + justify-content: center; + animation: sessions-chat-mcp-auth-pulse 2s ease-in-out infinite; +} + +.sessions-chat-mcp-auth-indicator .monaco-button:hover, +.sessions-chat-mcp-auth-indicator .monaco-button:focus-visible { + background: var(--vscode-toolbar-hoverBackground); + animation: none; +} + +@keyframes sessions-chat-mcp-auth-pulse { + 0%, 100% { + opacity: 1; + transform: scale(1); + } + 50% { + opacity: 0.65; + transform: scale(1.08); + } +} + +@media (prefers-reduced-motion: reduce) { + .sessions-chat-mcp-auth-indicator .monaco-button { + animation: none; + } +} + /* Model picker - uses workbench ModelPickerActionItem */ /* Session config toolbar (mode, model pickers via MenuWorkbenchToolBar) */ .sessions-chat-config-toolbar { diff --git a/src/vs/sessions/contrib/chat/browser/newChatInput.ts b/src/vs/sessions/contrib/chat/browser/newChatInput.ts index 163de415abb02..79371e8c704e9 100644 --- a/src/vs/sessions/contrib/chat/browser/newChatInput.ts +++ b/src/vs/sessions/contrib/chat/browser/newChatInput.ts @@ -39,6 +39,7 @@ import { getSimpleEditorOptions } from '../../../../workbench/contrib/codeEditor import { NewChatContextAttachments } from './newChatContextAttachments.js'; import { SessionTypePicker } from './sessionTypePicker.js'; import { MobileSessionTypePicker } from './mobile/mobileSessionTypePicker.js'; +import { NewChatMcpAuthIndicator } from './agentHost/newChatMcpAuthIndicator.js'; import { installMobileChipLaneScroll } from '../../../browser/parts/mobile/mobileChipLaneScroll.js'; import { IWorkbenchLayoutService } from '../../../../workbench/services/layout/browser/layoutService.js'; import { Menus } from '../../../browser/menus.js'; @@ -441,6 +442,12 @@ export class NewChatInputWidget extends Disposable implements IHistoryNavigation dom.append(toolbar, dom.$('.sessions-chat-toolbar-spacer')); + // MCP-auth indicator: appears when the active session has any + // MCP server in `AuthRequired`. Mirrors the chat-input toolbar + // indicator (see `OpenMcpAuthAction`) but renders here directly + // because the new-chat-input doesn't host `MenuId.ChatExecute`. + this._register(this.instantiationService.createInstance(NewChatMcpAuthIndicator, toolbar)); + this._loadingSpinner = dom.append(toolbar, dom.$('.sessions-chat-loading-spinner')); this._register(this.hoverService.setupManagedHover(getDefaultHoverDelegate('mouse'), this._loadingSpinner, localize('loading', "Loading..."))); diff --git a/src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHost.contribution.ts b/src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHost.contribution.ts index bc9e345e7bc92..764ad7d534122 100644 --- a/src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHost.contribution.ts +++ b/src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHost.contribution.ts @@ -5,7 +5,7 @@ import { Disposable, DisposableMap, DisposableStore, toDisposable } from '../../../../base/common/lifecycle.js'; import { Event } from '../../../../base/common/event.js'; -import { observableValue } from '../../../../base/common/observable.js'; +import { autorun, observableValueOpts } from '../../../../base/common/observable.js'; import { URI } from '../../../../base/common/uri.js'; import * as nls from '../../../../nls.js'; import { agentHostAuthority } from '../../../../platform/agentHost/common/agentHostUri.js'; @@ -31,7 +31,8 @@ import { IWorkbenchContribution, registerWorkbenchContribution2, WorkbenchPhase import { registerAction2 } from '../../../../platform/actions/common/actions.js'; import { OpenSessionEventsFileAction } from '../../agentHost/browser/openSessionEventsFileActions.js'; import { AgentCustomizationSyncProvider } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentCustomizationSyncProvider.js'; -import { authenticateProtectedResources, AgentHostAuthTokenCache, resolveAuthenticationInteractively } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.js'; +import { authenticateProtectedResources, AgentHostAuthTokenCache, resolveAuthenticationInteractively, resolveMcpServerAuthenticationInteractively, resolveMcpServerAuthenticationSilently } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.js'; +import { IAgentHostMcpAuthRegistry } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthRegistry.js'; import { AgentHostLanguageModelProvider } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLanguageModelProvider.js'; import { AgentHostSessionHandler } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.js'; import { LoggingAgentConnection } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/loggingAgentConnection.js'; @@ -119,6 +120,7 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc @IAgentHostTerminalService private readonly _agentHostTerminalService: IAgentHostTerminalService, @ITelemetryService private readonly _telemetryService: ITelemetryService, @IPromptsService private readonly _promptsService: IPromptsService, + @IAgentHostMcpAuthRegistry private readonly _mcpAuthRegistry: IAgentHostMcpAuthRegistry, ) { super(); @@ -481,7 +483,7 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc const bundler = agentStore.add(this._instantiationService.createInstance(SyncedCustomizationBundler, sessionType)); // Agent-level customizations observable - const customizations = observableValue('agentCustomizations', []); + const customizations = observableValueOpts({ debugName: 'agentCustomizations', equalsFn: deepEquals }, []); const updateCustomizations = async () => { const refs = await resolveCustomizationRefs(this._promptsService, syncProvider, this._agentPluginService, bundler, sessionType); customizations.set(refs, undefined); @@ -493,7 +495,18 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc this._promptsService.onDidChangeSkills, this._promptsService.onDidChangeInstructions, )(() => updateCustomizations())); - updateCustomizations(); // resolve initial state + // Refresh when the installed plugin set or any plugin's MCP server + // definitions / hooks change — these contributions are not surfaced + // via the prompts service events above, so MCP-only plugins would + // otherwise never trigger an update. + agentStore.add(autorun(reader => { + const plugins = this._agentPluginService.plugins.read(reader); + for (const plugin of plugins) { + plugin.mcpServerDefinitions.read(reader); + plugin.hooks.read(reader); + } + updateCustomizations(); + })); // Session handler (unified) const sessionHandler = agentStore.add(this._instantiationService.createInstance( @@ -510,6 +523,8 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc resolveWorkingDirectory, isNewSession, resolveAuthentication: (resources) => this._resolveAuthenticationInteractively(address, loggedConnection, resources), + resolveMcpAuthentication: (mcpServer, resource, scopes) => this._resolveMcpAuthenticationInteractively(address, loggedConnection, mcpServer, resource, scopes), + resolveMcpAuthenticationSilently: (mcpServer, resource, scopes) => this._resolveMcpAuthenticationSilently(address, loggedConnection, mcpServer, resource, scopes), customizations, })); agentStore.add(this._chatSessionsService.registerChatSessionContentProvider(sessionType, sessionHandler)); @@ -586,6 +601,49 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc } return false; } + + /** + * Interactively authenticate a single MCP server's protected resource + * and forward the token scoped to that server URI. Used by the + * chat-input MCP-auth indicator. + */ + private async _resolveMcpAuthenticationInteractively(address: string, loggedConnection: LoggingAgentConnection, mcpServer: URI, resource: ProtectedResourceMetadata, scopes: readonly string[] | undefined): Promise { + try { + return await resolveMcpServerAuthenticationInteractively(mcpServer, resource, scopes, { + mcpAuthMemory: this._mcpAuthRegistry, + mcpAuthHostKey: address, + authenticationService: this._authenticationService, + logPrefix: '[RemoteAgentHost]', + logService: this._logService, + authenticate: request => loggedConnection.authenticate(request), + }); + } catch (err) { + this._logService.error('[RemoteAgentHost] Interactive MCP authentication failed', err); + loggedConnection.logError('resolveMcpAuthenticationInteractively', err); + } + return false; + } + + /** + * Silently re-authenticate a single MCP server using an existing + * OAuth session, without prompting. Used as the auto-recovery path + * when an MCP server transitions to `AuthRequired`. + */ + private async _resolveMcpAuthenticationSilently(address: string, loggedConnection: LoggingAgentConnection, mcpServer: URI, resource: ProtectedResourceMetadata, scopes: readonly string[] | undefined): Promise { + try { + return await resolveMcpServerAuthenticationSilently(mcpServer, resource, scopes, { + mcpAuthMemory: this._mcpAuthRegistry, + mcpAuthHostKey: address, + authenticationService: this._authenticationService, + logPrefix: '[RemoteAgentHost]', + logService: this._logService, + authenticate: request => loggedConnection.authenticate(request), + }); + } catch (err) { + this._logService.warn('[RemoteAgentHost] Silent MCP authentication failed', err); + } + return false; + } } registerWorkbenchContribution2(RemoteAgentHostContribution.ID, RemoteAgentHostContribution, WorkbenchPhase.AfterRestored); @@ -672,3 +730,4 @@ Registry.as(ConfigurationExtensions.Configuration).regis import './remoteAgentHostActions.js'; import './manageRemoteAgentHosts.js'; import '../../chat/browser/agentHost/agentHostModelPicker.js'; +import { equals as deepEquals } from '../../../../base/common/objects.js'; diff --git a/src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHostSessionsProvider.ts b/src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHostSessionsProvider.ts index f79bc4fb31407..30938cb545737 100644 --- a/src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHostSessionsProvider.ts +++ b/src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHostSessionsProvider.ts @@ -27,9 +27,12 @@ import { ILogService } from '../../../../platform/log/common/log.js'; import { INotificationService } from '../../../../platform/notification/common/notification.js'; import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js'; import { IChatWidgetService } from '../../../../workbench/contrib/chat/browser/chat.js'; +import { IAgentHostActiveClientRegistry } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostActiveClientRegistry.js'; +import { IAgentHostMcpAuthRegistry } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthRegistry.js'; import { IChatService } from '../../../../workbench/contrib/chat/common/chatService/chatService.js'; import { IChatSessionsService } from '../../../../workbench/contrib/chat/common/chatSessionsService.js'; import { ILanguageModelsService } from '../../../../workbench/contrib/chat/common/languageModels.js'; +import { IAuthenticationService } from '../../../../workbench/services/authentication/common/authentication.js'; import { AgentHostSessionAdapter, BaseAgentHostSessionsProvider } from '../../agentHost/browser/baseAgentHostSessionsProvider.js'; import { IGitHubService } from '../../github/browser/githubService.js'; import { buildAgentHostSessionWorkspace, readBranchProtectionPatterns } from '../../../common/agentHostSessionWorkspace.js'; @@ -204,8 +207,11 @@ export class RemoteAgentHostSessionsProvider extends BaseAgentHostSessionsProvid @IConfigurationService private readonly _configurationService: IConfigurationService, @ILogService logService: ILogService, @IGitHubService gitHubService: IGitHubService, + @IAuthenticationService authenticationService: IAuthenticationService, + @IAgentHostMcpAuthRegistry mcpAuthRegistry: IAgentHostMcpAuthRegistry, + @IAgentHostActiveClientRegistry activeClientRegistry: IAgentHostActiveClientRegistry, ) { - super(chatSessionsService, chatService, chatWidgetService, languageModelsService, _configurationService, logService, gitHubService); + super(chatSessionsService, chatService, chatWidgetService, languageModelsService, _configurationService, logService, gitHubService, authenticationService, mcpAuthRegistry, activeClientRegistry); this._connectionAuthority = agentHostAuthority(config.address); this._connectOnDemand = config.connectOnDemand; @@ -258,6 +264,8 @@ export class RemoteAgentHostSessionsProvider extends BaseAgentHostSessionsProvid protected get authenticationPending(): IObservable { return this._authenticationPending; } + protected get _mcpAuthHostKey(): string { return this.remoteAddress; } + protected override createAdapter(meta: IAgentSessionMetadata): AgentHostSessionAdapter { this._metaByRawId.set(AgentSession.id(meta.session), meta); return super.createAdapter(meta); diff --git a/src/vs/sessions/contrib/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts b/src/vs/sessions/contrib/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts index 37780d308f801..58b23a9f3c7e9 100644 --- a/src/vs/sessions/contrib/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts +++ b/src/vs/sessions/contrib/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts @@ -25,9 +25,12 @@ import { TestInstantiationService } from '../../../../../platform/instantiation/ import { INotificationService } from '../../../../../platform/notification/common/notification.js'; import { InMemoryStorageService, IStorageService } from '../../../../../platform/storage/common/storage.js'; import { IChatWidget, IChatWidgetService } from '../../../../../workbench/contrib/chat/browser/chat.js'; +import { IAgentHostActiveClientRegistry } from '../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostActiveClientRegistry.js'; +import { IAgentHostMcpAuthRegistry } from '../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthRegistry.js'; import { IChatService, type ChatSendResult, type IChatSendRequestOptions } from '../../../../../workbench/contrib/chat/common/chatService/chatService.js'; import { IChatSessionsService } from '../../../../../workbench/contrib/chat/common/chatSessionsService.js'; import { ILanguageModelsService } from '../../../../../workbench/contrib/chat/common/languageModels.js'; +import { IAuthenticationService } from '../../../../../workbench/services/authentication/common/authentication.js'; import { ISessionChangeEvent } from '../../../../services/sessions/common/sessionsProvider.js'; import { SessionStatus, COPILOT_CLI_SESSION_TYPE } from '../../../../services/sessions/common/session.js'; import { RemoteAgentHostSessionsProvider, type IRemoteAgentHostSessionsProviderConfig } from '../../browser/remoteAgentHostSessionsProvider.js'; @@ -215,6 +218,13 @@ function createProvider(disposables: DisposableStore, connection: MockAgentConne instantiationService.stub(IGitHubService, new class extends mock() { override findPullRequestNumberByHeadBranch = async () => undefined; }()); + instantiationService.stub(IAuthenticationService, new class extends mock() { }()); + instantiationService.stub(IAgentHostMcpAuthRegistry, new class extends mock() { + override registerSession = () => ({ dispose() { } }); + }()); + instantiationService.stub(IAgentHostActiveClientRegistry, new class extends mock() { + override get = () => undefined; + }()); const config: IRemoteAgentHostSessionsProviderConfig = { address: overrides?.address ?? 'localhost:4321', diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostActiveClientRegistry.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostActiveClientRegistry.ts new file mode 100644 index 0000000000000..f1f312e556a31 --- /dev/null +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostActiveClientRegistry.ts @@ -0,0 +1,79 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { Disposable, IDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js'; +import { type SessionActiveClient } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import { InstantiationType, registerSingleton } from '../../../../../../platform/instantiation/common/extensions.js'; +import { createDecorator } from '../../../../../../platform/instantiation/common/instantiation.js'; + +/** + * Snapshot the active-client info this client would dispatch via + * `session/activeClientChanged` for a session it owns. Recomputed on + * every read so the caller sees the latest customizations / tools. + */ +export type AgentHostActiveClientSnapshot = () => SessionActiveClient; + +export const IAgentHostActiveClientRegistry = createDecorator('agentHostActiveClientRegistry'); + +/** + * Workbench-scoped registry that exposes each + * {@link AgentHostSessionHandler}'s active-client snapshot — clientId, + * client-side tools, customizations — to the sessions provider that + * eagerly creates backend sessions on the new-chat view. + * + * The provider's `NewSession.eagerCreate` reads the snapshot here and + * passes it to `connection.createSession`'s `activeClient` parameter, + * letting the agent host establish the active client at the moment of + * creation. Without this bridge, the handler would have to re-dispatch + * `session/activeClientChanged` for every provisional session it + * discovers post-hoc — which both adds plumbing and leaves a window + * where the host has no active client for sessions the user is + * actively configuring. + * + * Keyed by chat `sessionType` (the same identifier the chat sessions + * service uses to route to a content provider). One registration per + * session type at a time; the contribution that owns the handler also + * owns the registration. + */ +export interface IAgentHostActiveClientRegistry { + readonly _serviceBrand: undefined; + + /** + * Registers a snapshot getter for `sessionType`. Replaces any prior + * registration. Returns a disposable that removes the registration + * if it's still the current one for `sessionType`. + */ + register(sessionType: string, snapshot: AgentHostActiveClientSnapshot): IDisposable; + + /** + * Returns the active-client snapshot for `sessionType`, or + * `undefined` if no handler has registered one (e.g. the agent host + * isn't connected for this session type, or the contribution + * unregistered). + */ + get(sessionType: string): SessionActiveClient | undefined; +} + +/** Exported for tests. Production code MUST use {@link IAgentHostActiveClientRegistry}. */ +export class AgentHostActiveClientRegistry extends Disposable implements IAgentHostActiveClientRegistry { + declare readonly _serviceBrand: undefined; + + private readonly _snapshots = new Map(); + + register(sessionType: string, snapshot: AgentHostActiveClientSnapshot): IDisposable { + this._snapshots.set(sessionType, snapshot); + return toDisposable(() => { + if (this._snapshots.get(sessionType) === snapshot) { + this._snapshots.delete(sessionType); + } + }); + } + + get(sessionType: string): SessionActiveClient | undefined { + return this._snapshots.get(sessionType)?.(); + } +} + +registerSingleton(IAgentHostActiveClientRegistry, AgentHostActiveClientRegistry, InstantiationType.Delayed); diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.ts index fafe9633d851e..d34aeab046d99 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.ts @@ -4,10 +4,11 @@ *--------------------------------------------------------------------------------------------*/ import { URI } from '../../../../../../base/common/uri.js'; -import { type ProtectedResourceMetadata } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import { McpServerStatusKind, type McpServerSummary, type ProtectedResourceMetadata } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; import { type AgentInfo } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { ILogService } from '../../../../../../platform/log/common/log.js'; import { IAuthenticationService } from '../../../../../services/authentication/common/authentication.js'; +import { type IAgentHostMcpAuthRegistry } from './agentHostMcpAuthRegistry.js'; /** * Tracks the last bearer token pushed to a given agent host connection @@ -110,10 +111,36 @@ export async function resolveTokenForResource( export interface IAgentHostAuthenticateRequest { readonly resource: string; readonly token: string; + /** + * Optional MCP server URI to scope the token to. Forwarded as + * `AuthenticateParams.server`; the host treats the token as + * server-scoped when set, otherwise as resource-wide. + */ + readonly server?: URI; } export interface IAgentHostAuthenticationOptions { readonly authTokenCache?: AgentHostAuthTokenCache; + /** + * Persistent memory of previously approved MCP scopes per OAuth + * resource. Used by the per-MCP-server helpers to silent-auth with + * the scopes the user originally consented to, even when the + * server's current `requiredScopes` are narrower or absent. The + * helpers only call `recall`/`remember`; the caller may pass the + * full {@link IAgentHostMcpAuthRegistry} since it implements both. + * + * Memory is consulted only when {@link mcpAuthHostKey} is also + * provided — the host dimension prevents a token approved against + * one agent host from silently authenticating another. + */ + readonly mcpAuthMemory?: Pick; + /** + * Stable identifier for the agent host these credentials live under + * (`'local'` for the in-process agent host, an address string for + * remote hosts). Combined with the OAuth resource URI to scope + * {@link mcpAuthMemory} reads/writes. + */ + readonly mcpAuthHostKey?: string; readonly authenticationService: IAuthenticationService; readonly logPrefix: string; readonly logService: ILogService; @@ -207,3 +234,181 @@ export async function resolveAuthenticationInteractively( return false; } + +/** + * Try to silently push a token to the agent host for a single MCP + * server. Resolves an existing OAuth session matching the previously + * approved scopes (no `createSession` fallback) and forwards the + * bearer token scoped to {@link mcpServer} via + * {@link IAgentHostAuthenticateRequest.server}. + * + * Used to auto-recover from `AuthRequired` transitions without + * surfacing UI when the user has already authenticated the underlying + * OAuth resource — for example, after the agent host re-emits + * `AuthRequired` for a server we had previously authenticated. + * + * Silent authentication only proceeds when + * {@link IAgentHostAuthenticationOptions.mcpAuthMemory} has a + * recorded entry for this `(host, resource)` AND the remembered + * scopes cover the currently-requested scopes. Without a recorded + * approval that subsumes the new request, we'd be quietly granting + * the agent host access to scopes the user never agreed to — fall + * back to the interactive flow instead. + * + * Like {@link resolveMcpServerAuthenticationInteractively}, this does + * not consult or update the resource-wide + * {@link AgentHostAuthTokenCache}: per-server tokens are independent. + */ +export async function resolveMcpServerAuthenticationSilently( + mcpServer: URI, + resource: ProtectedResourceMetadata, + scopes: readonly string[] | undefined, + options: IAgentHostAuthenticationOptions, +): Promise { + if (!options.mcpAuthMemory || !options.mcpAuthHostKey) { + return false; + } + const remembered = options.mcpAuthMemory.recall(options.mcpAuthHostKey, resource.resource); + if (!remembered) { + return false; + } + const requestedScopes = scopes ?? resource.scopes_supported ?? []; + const rememberedSet = new Set(remembered); + for (const scope of requestedScopes) { + if (!rememberedSet.has(scope)) { + return false; + } + } + const resourceUri = URI.parse(resource.resource); + const token = await resolveTokenForResource( + resourceUri, + resource.authorization_servers ?? [], + remembered, + options.authenticationService, + options.logService, + options.logPrefix, + ); + if (!token) { + return false; + } + await options.authenticate({ resource: resource.resource, token, server: mcpServer }); + options.logService.info(`${options.logPrefix} Silent MCP authentication succeeded for ${resource.resource} (${mcpServer.toString()})`); + return true; +} + +/** + * Authenticate a single MCP server's protected resource and forward the + * resulting token to the agent host scoped to that server URI. Mirrors + * {@link resolveAuthenticationInteractively} but always scopes to a + * specific MCP server proxy via {@link IAgentHostAuthenticateRequest.server}. + * + * Tries {@link resolveMcpServerAuthenticationSilently} first; if memory + * doesn't authorize a silent reuse, falls back to looking up an + * existing session matching the requested scopes (the user is actively + * authorizing this server, so reusing a matching session is OK even + * without a memory record). Only triggers `createSession` (which + * prompts) when no matching session exists at all. Successful + * authentications are recorded into + * {@link IAgentHostAuthenticationOptions.mcpAuthMemory} so future + * reloads can silent-auth without prompting. + * + * Unlike the agent-level helper, this does NOT consult or update the + * resource-wide token cache: per-server tokens may differ from the + * agent-scoped token for the same OAuth resource, and writing to the + * shared cache could mask a later resource-wide call. + */ +export async function resolveMcpServerAuthenticationInteractively( + mcpServer: URI, + resource: ProtectedResourceMetadata, + scopes: readonly string[] | undefined, + options: IAgentHostAuthenticationOptions, +): Promise { + if (await resolveMcpServerAuthenticationSilently(mcpServer, resource, scopes, options)) { + return true; + } + + const effectiveScopes = scopes ?? resource.scopes_supported ?? []; + const resourceUri = URI.parse(resource.resource); + + // Reuse an existing OAuth session that already covers the requested + // scopes before prompting, even when memory has no record. The user + // is actively authorizing this server, so binding an already-granted + // session to it doesn't escalate consent. + const existingToken = await resolveTokenForResource( + resourceUri, + resource.authorization_servers ?? [], + effectiveScopes, + options.authenticationService, + options.logService, + options.logPrefix, + ); + if (existingToken) { + await options.authenticate({ resource: resource.resource, token: existingToken, server: mcpServer }); + if (options.mcpAuthMemory && options.mcpAuthHostKey) { + options.mcpAuthMemory.remember(options.mcpAuthHostKey, resource.resource, effectiveScopes); + } + options.logService.info(`${options.logPrefix} Reused existing session for ${resource.resource} (${mcpServer.toString()})`); + return true; + } + + for (const server of resource.authorization_servers ?? []) { + const serverUri = URI.parse(server); + const providerId = await options.authenticationService.getOrActivateProviderIdForServer(serverUri, resourceUri); + if (!providerId) { + continue; + } + + const session = await options.authenticationService.createSession(providerId, [...effectiveScopes], { + activateImmediate: true, + authorizationServer: serverUri, + }); + + await options.authenticate({ resource: resource.resource, token: session.accessToken, server: mcpServer }); + if (options.mcpAuthMemory && options.mcpAuthHostKey) { + options.mcpAuthMemory.remember(options.mcpAuthHostKey, resource.resource, effectiveScopes); + } + options.logService.info(`${options.logPrefix} Interactive MCP authentication succeeded for ${resource.resource} (${mcpServer.toString()})`); + return true; + } + + return false; +} + +/** + * Drive interactive authentication for one MCP server (`target`) or for + * every server in the supplied snapshot that's currently in + * `AuthRequired`. Per-server failures are logged and skipped so a + * single rejection doesn't abort the rest of the batch. + * + * Used by both the chat-session-bound auth path and the provisional + * (new-chat) path so the iteration semantics match across the two + * surfaces. + * + * @returns `true` when at least one of the attempted servers + * authenticated successfully. + */ +export async function authenticateMcpServerCandidates( + summaries: readonly McpServerSummary[], + target: McpServerSummary | undefined, + options: IAgentHostAuthenticationOptions, +): Promise { + const candidates = target ? [target] : summaries.filter(s => s.status.kind === McpServerStatusKind.AuthRequired); + let anySucceeded = false; + for (const summary of candidates) { + if (summary.status.kind !== McpServerStatusKind.AuthRequired) { + continue; + } + try { + const ok = await resolveMcpServerAuthenticationInteractively( + URI.parse(summary.resource), + summary.status.resource, + summary.status.requiredScopes, + options, + ); + anySucceeded ||= ok; + } catch (err) { + options.logService.error(`${options.logPrefix} MCP authentication failed for ${summary.resource}`, err); + } + } + return anySucceeded; +} diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostChatContribution.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostChatContribution.ts index 009d27099fe89..92824244cd40b 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostChatContribution.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostChatContribution.ts @@ -6,8 +6,9 @@ import { Codicon } from '../../../../../../base/common/codicons.js'; import { Disposable, DisposableMap, DisposableStore, toDisposable } from '../../../../../../base/common/lifecycle.js'; import { Event } from '../../../../../../base/common/event.js'; -import { observableValue } from '../../../../../../base/common/observable.js'; +import { autorun, observableValueOpts } from '../../../../../../base/common/observable.js'; import { ThemeIcon } from '../../../../../../base/common/themables.js'; +import { URI } from '../../../../../../base/common/uri.js'; import { localize } from '../../../../../../nls.js'; import { AgentHostEnabledSettingId, IAgentHostService, type AgentProvider } from '../../../../../../platform/agentHost/common/agentService.js'; import { type ProtectedResourceMetadata } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; @@ -28,12 +29,14 @@ import { IAgentPluginService } from '../../../common/plugins/agentPluginService. import { IPromptsService, PromptsStorage } from '../../../common/promptSyntax/service/promptsService.js'; import { AgentCustomizationSyncProvider } from './agentCustomizationSyncProvider.js'; import { LocalAgentHostCustomizationItemProvider, resolveCustomizationRefs } from './agentHostLocalCustomizations.js'; -import { authenticateProtectedResources, AgentHostAuthTokenCache, resolveAuthenticationInteractively } from './agentHostAuth.js'; +import { authenticateProtectedResources, AgentHostAuthTokenCache, resolveAuthenticationInteractively, resolveMcpServerAuthenticationInteractively, resolveMcpServerAuthenticationSilently } from './agentHostAuth.js'; +import { IAgentHostMcpAuthRegistry } from './agentHostMcpAuthRegistry.js'; import { AgentHostLanguageModelProvider } from './agentHostLanguageModelProvider.js'; import { AgentHostSessionHandler } from './agentHostSessionHandler.js'; import { AgentHostSessionListController } from './agentHostSessionListController.js'; import { LoggingAgentConnection } from './loggingAgentConnection.js'; import { SyncedCustomizationBundler } from './syncedCustomizationBundler.js'; +import { equals as deepEquals } from '../../../../../../base/common/objects.js'; export { AgentHostSessionHandler } from './agentHostSessionHandler.js'; export { AgentHostSessionListController } from './agentHostSessionListController.js'; @@ -77,6 +80,7 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr @IAgentPluginService private readonly _agentPluginService: IAgentPluginService, @IPromptsService private readonly _promptsService: IPromptsService, @IWorkbenchEnvironmentService environmentService: IWorkbenchEnvironmentService, + @IAgentHostMcpAuthRegistry private readonly _mcpAuthRegistry: IAgentHostMcpAuthRegistry, ) { super(); @@ -206,7 +210,7 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr itemProvider, })); - const customizations = observableValue('agentCustomizations', []); + const customizations = observableValueOpts({ debugName: 'agentCustomizations', equalsFn: deepEquals }, []); const updateCustomizations = async () => { const refs = await resolveCustomizationRefs(this._promptsService, syncProvider, this._agentPluginService, bundler, sessionType); customizations.set(refs, undefined); @@ -218,7 +222,18 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr this._promptsService.onDidChangeSkills, this._promptsService.onDidChangeInstructions, )(() => updateCustomizations())); - updateCustomizations(); // resolve initial state + // Refresh when the installed plugin set or any plugin's MCP server + // definitions / hooks change — these contributions are not surfaced + // via the prompts service events above, so MCP-only plugins would + // otherwise never trigger an update. + store.add(autorun(reader => { + const plugins = this._agentPluginService.plugins.read(reader); + for (const plugin of plugins) { + plugin.mcpServerDefinitions.read(reader); + plugin.hooks.read(reader); + } + updateCustomizations(); + })); // Session handler const sessionHandler = store.add(this._instantiationService.createInstance(AgentHostSessionHandler, { @@ -231,6 +246,8 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr connectionAuthority: 'local', isNewSession: sessionResource => listController.isNewSession(sessionResource), resolveAuthentication: (resources) => this._resolveAuthenticationInteractively(resources), + resolveMcpAuthentication: (mcpServer, resource, scopes) => this._resolveMcpAuthenticationInteractively(mcpServer, resource, scopes), + resolveMcpAuthenticationSilently: (mcpServer, resource, scopes) => this._resolveMcpAuthenticationSilently(mcpServer, resource, scopes), customizations, })); store.add(this._chatSessionsService.registerChatSessionContentProvider(sessionType, sessionHandler)); @@ -306,4 +323,52 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr } return false; } + + /** + * Interactively authenticates a single MCP server's protected resource + * and forwards the token scoped to that server. Used by the chat-input + * MCP-auth indicator so per-server tokens are pushed independently. + */ + private async _resolveMcpAuthenticationInteractively(mcpServer: URI, resource: ProtectedResourceMetadata, scopes: readonly string[] | undefined): Promise { + try { + return await resolveMcpServerAuthenticationInteractively(mcpServer, resource, scopes, { + authTokenCache: this._authTokenCache, + mcpAuthMemory: this._mcpAuthRegistry, + mcpAuthHostKey: 'local', + authenticationService: this._authenticationService, + logPrefix: '[AgentHost]', + logService: this._logService, + authenticate: request => this._loggedConnection!.authenticate(request), + }); + } catch (err) { + this._logService.error('[AgentHost] Interactive MCP authentication failed', err); + this._loggedConnection!.logError('resolveMcpAuthenticationInteractively', err); + } + return false; + } + + /** + * Silently re-authenticates a single MCP server using an existing + * OAuth session (no prompt). Triggered when an MCP server + * transitions to `AuthRequired` so previously-authenticated hosts + * recover without surfacing UI. Returns `false` when no matching + * session is available; the indicator click path then drives the + * interactive flow. + */ + private async _resolveMcpAuthenticationSilently(mcpServer: URI, resource: ProtectedResourceMetadata, scopes: readonly string[] | undefined): Promise { + try { + return await resolveMcpServerAuthenticationSilently(mcpServer, resource, scopes, { + authTokenCache: this._authTokenCache, + mcpAuthMemory: this._mcpAuthRegistry, + mcpAuthHostKey: 'local', + authenticationService: this._authenticationService, + logPrefix: '[AgentHost]', + logService: this._logService, + authenticate: request => this._loggedConnection!.authenticate(request), + }); + } catch (err) { + this._logService.warn('[AgentHost] Silent MCP authentication failed', err); + } + return false; + } } diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLocalCustomizations.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLocalCustomizations.ts index 39461a4ce69df..5041288fbafd1 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLocalCustomizations.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLocalCustomizations.ts @@ -217,9 +217,6 @@ export async function resolveCustomizationRefs( ): Promise { const enumerated = await enumerateLocalCustomizationsForHarness(promptsService, syncProvider, sessionType, CancellationToken.None); const enabled = enumerated.filter(e => !e.disabled); - if (enabled.length === 0) { - return []; - } const plugins = agentPluginService.plugins.get(); const pluginRefs = new Map(); @@ -243,6 +240,25 @@ export async function resolveCustomizationRefs( } } + // Plugins whose only contributions are MCP servers or hooks are not + // surfaced by `enumerateLocalCustomizationsForHarness` (which walks + // prompt files only). Pick them up here so the agent host still + // receives those plugins in the `customizations` set. + for (const plugin of plugins) { + const key = plugin.uri.toString(); + if (pluginRefs.has(key)) { + continue; + } + if (syncProvider.isDisabled(plugin.uri)) { + continue; + } + const hasNonPromptContent = plugin.mcpServerDefinitions.get().length > 0 + || plugin.hooks.get().length > 0; + if (hasNonPromptContent) { + pluginRefs.set(key, { uri: key as ProtocolURI, displayName: plugin.label }); + } + } + const refs: CustomizationRef[] = [...pluginRefs.values()]; if (looseFiles.length > 0) { const result = await bundler.bundle(looseFiles); diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthIndicatorContribution.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthIndicatorContribution.ts new file mode 100644 index 0000000000000..11236ac97d8cf --- /dev/null +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthIndicatorContribution.ts @@ -0,0 +1,82 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { Disposable, DisposableMap, DisposableStore } from '../../../../../../base/common/lifecycle.js'; +import { autorun } from '../../../../../../base/common/observable.js'; +import { McpServerStatusKind } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import { IWorkbenchContribution } from '../../../../../common/contributions.js'; +import { ChatContextKeys } from '../../../common/actions/chatContextKeys.js'; +import { IChatWidget, IChatWidgetService } from '../../chat.js'; +import { IAgentHostMcpAuthRegistry } from './agentHostMcpAuthRegistry.js'; + +/** + * Per-chat-widget binding of {@link ChatContextKeys.mcpAuthRequiredCount} + * to the count of MCP servers in `AuthRequired` state on the active AHP + * session. + * + * The action contributed to the chat-input toolbar (see + * `chatMcpAuthAction.ts`) gates its visibility on this key, so it + * appears only when at least one server needs authentication. + */ +export class AgentHostMcpAuthIndicatorContribution extends Disposable implements IWorkbenchContribution { + + static readonly ID = 'workbench.contrib.agentHostMcpAuthIndicator'; + + private readonly _widgetBindings = this._register(new DisposableMap()); + + constructor( + @IChatWidgetService private readonly _chatWidgetService: IChatWidgetService, + @IAgentHostMcpAuthRegistry private readonly _registry: IAgentHostMcpAuthRegistry, + ) { + super(); + + for (const widget of this._chatWidgetService.getAllWidgets()) { + this._bindWidget(widget); + } + this._register(this._chatWidgetService.onDidAddWidget(widget => this._bindWidget(widget))); + } + + private _bindWidget(widget: IChatWidget): void { + if (this._widgetBindings.has(widget)) { + return; + } + + const store = new DisposableStore(); + const countKey = ChatContextKeys.mcpAuthRequiredCount.bindTo(widget.scopedContextKeyService); + + // Single autorun covering session swap, registry register/unregister + // race (the session handler may register its entry AFTER the chat + // widget is created), and the entry's `mcpServers` observable. + const sync = () => { + autorunStore.clear(); + autorunStore.add(autorun(reader => { + const sessionResource = widget.viewModel?.sessionResource; + if (!sessionResource) { + countKey.set(0); + return; + } + const entry = this._registry.getEntry(sessionResource, reader); + if (!entry) { + countKey.set(0); + return; + } + const summaries = entry.mcpServers.read(reader); + let count = 0; + for (const s of summaries) { + if (s.status.kind === McpServerStatusKind.AuthRequired) { + count++; + } + } + countKey.set(count); + })); + }; + const autorunStore = store.add(new DisposableStore()); + + store.add(widget.onDidChangeViewModel(() => sync())); + sync(); + + this._widgetBindings.set(widget, store); + } +} diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthRegistry.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthRegistry.ts new file mode 100644 index 0000000000000..b7eef20500adc --- /dev/null +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostMcpAuthRegistry.ts @@ -0,0 +1,252 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { Disposable, IDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js'; +import { ResourceMap } from '../../../../../../base/common/map.js'; +import { IObservable, IReader, observableValue, transaction } from '../../../../../../base/common/observable.js'; +import { URI } from '../../../../../../base/common/uri.js'; +import { type McpServerSummary } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import { ILogService } from '../../../../../../platform/log/common/log.js'; +import { InstantiationType, registerSingleton } from '../../../../../../platform/instantiation/common/extensions.js'; +import { createDecorator } from '../../../../../../platform/instantiation/common/instantiation.js'; +import { IStorageService, StorageScope, StorageTarget } from '../../../../../../platform/storage/common/storage.js'; + +const STORAGE_KEY = 'agentHost.mcpAuth.consentedScopes'; + +interface IPersistedEntry { + /** Sorted-then-stored scope list; matches the entry written via {@link IAgentHostMcpAuthRegistry.remember}. */ + readonly scopes: readonly string[]; + /** Wall-clock time of the last successful authentication (ms). */ + readonly updatedAt: number; +} + +/** + * `host` → `resource URI` → entry. Hosts are stable identifiers + * (`'local'` for the in-process agent host, an address string for + * remote hosts) so that authorizing the same OAuth resource on a + * different host prompts the user again instead of silently reusing + * a token approved against another host. + */ +type IPersistedShape = Record>; + +/** + * Per-AHP-session view onto MCP server state needed by the chat input UI. + */ +export interface IAgentHostMcpAuthSessionEntry { + /** + * Observable list of MCP servers attached to this session, mirroring + * `SessionState.mcpServers`. Empty when no servers are registered. + */ + readonly mcpServers: IObservable; + + /** + * Drive interactive authentication for one server, or all servers in + * an `AuthRequired` state when `server` is omitted. Returns `true` + * when at least one authentication call succeeded. Implementations + * forward `McpServerStatusAuthRequired.requiredScopes` (when set) + * to the underlying authentication flow. + */ + authenticate(server?: McpServerSummary): Promise; +} + +export const IAgentHostMcpAuthRegistry = createDecorator('agentHostMcpAuthRegistry'); + +/** + * Workbench-scoped service that bridges AHP MCP-auth state to the chat + * UI. Combines two responsibilities: + * + * 1. **Per-session registry** — maps chat session resources to live + * `{ mcpServers, authenticate }` entries that the chat input UI + * looks up to render the indicator and drive interactive auth. + * Entries come and go with the session. + * 2. **Persistent scope memory** — records (via `IStorageService`) + * which OAuth scopes the user previously consented to for each + * protected resource. Survives reload because the OAuth sessions + * themselves are persisted by the auth provider; what we need to + * remember is *which* scopes the user approved so we can ask for + * them again silently. Without this, a server that re-emits + * `AuthRequired` after reload may end up prompting the user again + * because the new challenge advertised a narrower or different + * scope set than the one originally approved. + * + * Both pieces live here because they share the same domain (MCP auth + * state for AHP-backed sessions) and the same workbench lifetime, and + * because they're consumed in tandem — the auth helpers read/write + * the memory, the chat input UI reads the registry, and the indicator + * never appears for a server whose scopes are remembered and silently + * re-authenticated. + * + * Memory entries are keyed by `(host, resource)` — the agent-host + * identifier (`'local'` or a remote address) plus the canonical OAuth + * resource URI (`ProtectedResourceMetadata.resource`). The host + * dimension means authorizing the same OAuth resource on a different + * agent host prompts again instead of silently reusing a token + * approved against another host. Both axes are durable across + * reloads; chat session URIs and MCP server URIs are not. + */ +export interface IAgentHostMcpAuthRegistry { + readonly _serviceBrand: undefined; + + // ---- Per-session registry ------------------------------------------------ + + /** + * Registers an entry for `sessionResource`. Returns a disposable that + * removes the entry. If a different entry was previously registered + * for the same resource, this overwrites it; otherwise, the existing + * registration's disposal is left intact. + */ + registerSession(sessionResource: URI, entry: IAgentHostMcpAuthSessionEntry): IDisposable; + + /** + * Look up the entry for a chat session resource. When called within + * an autorun/derived, pass the `reader` so the computation re-fires + * when an entry is registered or unregistered for `sessionResource`. + */ + getEntry(sessionResource: URI, reader?: IReader): IAgentHostMcpAuthSessionEntry | undefined; + + // ---- Persistent scope memory -------------------------------------------- + + /** + * Records that the user successfully authenticated `resource` on + * `host` with `scopes`. Overwrites any prior entry for the + * (host, resource) pair; scopes are normalized to a sorted, deduped + * list so call-site differences (order, duplicates) don't produce + * divergent memo entries. + * + * `host` is a stable agent-host identifier (`'local'` for the + * in-process agent host, an address string for remote hosts). The + * (host, resource) keying ensures authorizing the same OAuth + * resource on a different agent host prompts the user again + * instead of silently reusing a token approved against another + * host. + */ + remember(host: string, resource: string, scopes: readonly string[]): void; + + /** + * Returns the scopes the user previously approved for `resource` + * on `host`, or `undefined` if no record exists. + */ + recall(host: string, resource: string): readonly string[] | undefined; + + /** + * Drops the record for `resource` on `host`. Call after a + * previously successful authentication is rejected (e.g. token + * revoked) so the silent path doesn't keep retrying with stale + * scopes. + */ + forget(host: string, resource: string): void; +} + +/** Exported for tests. Production code MUST use {@link IAgentHostMcpAuthRegistry}. */ +export class AgentHostMcpAuthRegistry extends Disposable implements IAgentHostMcpAuthRegistry { + declare readonly _serviceBrand: undefined; + + private readonly _entries = new ResourceMap(); + /** + * Version counter bumped on every register/unregister. Observers + * passing a `reader` to {@link getEntry} subscribe to this so the + * computation re-fires when an entry appears or disappears for the + * resource they care about. + */ + private readonly _entriesVersion = observableValue(this, 0); + + /** + * In-memory mirror of the persisted scope memory. Lazily loaded + * from storage on first read; mutated by `remember`/`forget`; + * flushed back to storage on `IStorageService.onWillSaveState`. + * + * Buffering avoids stringify+store round-trips on every + * authentication and lets the storage service batch the write + * with the rest of the workbench state. + */ + private _memoryCache: IPersistedShape | undefined; + /** Set when {@link _memoryCache} has unwritten mutations. */ + private _memoryDirty = false; + + constructor( + @IStorageService private readonly _storageService: IStorageService, + @ILogService private readonly _logService: ILogService, + ) { + super(); + + this._register(this._storageService.onWillSaveState(() => this._flushMemory())); + } + + registerSession(sessionResource: URI, entry: IAgentHostMcpAuthSessionEntry): IDisposable { + transaction(tx => { + this._entries.set(sessionResource, entry); + this._entriesVersion.set(this._entriesVersion.get() + 1, tx); + }); + return toDisposable(() => { + if (this._entries.get(sessionResource) === entry) { + transaction(tx => { + this._entries.delete(sessionResource); + this._entriesVersion.set(this._entriesVersion.get() + 1, tx); + }); + } + }); + } + + getEntry(sessionResource: URI, reader?: IReader): IAgentHostMcpAuthSessionEntry | undefined { + this._entriesVersion.read(reader); + return this._entries.get(sessionResource); + } + + remember(host: string, resource: string, scopes: readonly string[]): void { + const normalized = normalizeScopes(scopes); + const all = this._readMemory(); + const hostEntries = all[host] ?? (all[host] = {}); + hostEntries[resource] = { scopes: normalized, updatedAt: Date.now() }; + this._memoryDirty = true; + } + + recall(host: string, resource: string): readonly string[] | undefined { + return this._readMemory()[host]?.[resource]?.scopes; + } + + forget(host: string, resource: string): void { + const all = this._readMemory(); + const hostEntries = all[host]; + if (hostEntries && Object.prototype.hasOwnProperty.call(hostEntries, resource)) { + delete hostEntries[resource]; + if (Object.keys(hostEntries).length === 0) { + delete all[host]; + } + this._memoryDirty = true; + } + } + + private _readMemory(): IPersistedShape { + if (this._memoryCache) { + return this._memoryCache; + } + const raw = this._storageService.get(STORAGE_KEY, StorageScope.APPLICATION_SHARED); + if (!raw) { + return this._memoryCache = {}; + } + try { + const parsed = JSON.parse(raw); + this._memoryCache = (parsed && typeof parsed === 'object') ? parsed as IPersistedShape : {}; + } catch (err) { + this._logService.warn('[AgentHostMcpAuthRegistry] Failed to parse stored auth memory; resetting.', err); + this._memoryCache = {}; + } + return this._memoryCache; + } + + private _flushMemory(): void { + if (!this._memoryDirty || !this._memoryCache) { + return; + } + this._memoryDirty = false; + this._storageService.store(STORAGE_KEY, JSON.stringify(this._memoryCache), StorageScope.APPLICATION_SHARED, StorageTarget.MACHINE); + } +} + +function normalizeScopes(scopes: readonly string[]): string[] { + return [...new Set(scopes)].sort(); +} + +registerSingleton(IAgentHostMcpAuthRegistry, AgentHostMcpAuthRegistry, InstantiationType.Delayed); diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts index 1fc88af263f25..a79d3a1f13d8c 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -21,7 +21,7 @@ import { AgentProvider, AgentSession, type IAgentConnection } from '../../../../ import { IAgentSubscription, observableFromSubscription } from '../../../../../../platform/agentHost/common/state/agentSubscription.js'; import { SessionTruncatedAction } from '../../../../../../platform/agentHost/common/state/protocol/actions.js'; import { CompletionItemKind as AhpCompletionItemKind, type CompletionItem as AhpCompletionItem } from '../../../../../../platform/agentHost/common/state/protocol/commands.js'; -import { ConfirmationOptionKind, CustomizationRef, TerminalClaimKind, ToolResultContentType, type ConfirmationOption, type ProtectedResourceMetadata, type ToolDefinition } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import { ConfirmationOptionKind, CustomizationRef, McpServerStatusKind, TerminalClaimKind, ToolResultContentType, type ConfirmationOption, type McpServerSummary, type ProtectedResourceMetadata, type ToolDefinition } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; import { ActionType, SessionTurnStartedAction, type ClientSessionAction, type SessionAction, type SessionInputCompletedAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; import { AHP_AUTH_REQUIRED, ProtocolError } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { buildSubagentSessionUri, getToolFileEdits, getToolSubagentContent, MessageAttachmentKind, PendingMessageKind, ResponsePartKind, SessionInputAnswerState, SessionInputAnswerValueKind, SessionInputQuestionKind, SessionInputResponseKind, StateComponents, ToolCallCancellationReason, ToolCallConfirmationReason, ToolCallStatus, TurnState, type ICompletedToolCall, type MarkdownResponsePart, type MessageAttachment, type ModelSelection, type ReasoningResponsePart, type RootState, type SessionInputAnswer, type SessionInputRequest, type SessionState, type ToolCallResponsePart, type ToolCallState, type Turn } from '../../../../../../platform/agentHost/common/state/sessionState.js'; @@ -50,6 +50,8 @@ import { ILanguageModelsService } from '../../../common/languageModels.js'; import { ILanguageModelToolsService, IToolData, IToolInvocation, IToolResult, ToolInvocationPresentation } from '../../../common/tools/languageModelToolsService.js'; import { getAgentHostIcon } from '../agentSessions.js'; import { AgentHostEditingSession } from './agentHostEditingSession.js'; +import { IAgentHostActiveClientRegistry } from './agentHostActiveClientRegistry.js'; +import { IAgentHostMcpAuthRegistry, type IAgentHostMcpAuthSessionEntry } from './agentHostMcpAuthRegistry.js'; import { IAgentHostSessionWorkingDirectoryResolver } from './agentHostSessionWorkingDirectoryResolver.js'; import { activeTurnToProgress, completedToolCallToEditParts, completedToolCallToSerialized, finalizeToolInvocation, getTerminalContentUri, isSubagentTool, makeAhpTerminalToolSessionId, parseAhpTerminalToolSessionId, rawMarkdownToString, stringOrMarkdownToString, toolCallStateToInvocation, turnsToHistory, updateRunningToolSpecificData, userMessageToVariableData, type IToolCallFileEdit, type TurnModelLookup } from './stateToProgressAdapter.js'; @@ -210,6 +212,13 @@ class AgentHostChatSession extends Disposable implements IChatSession { readonly progressObs = observableValue('agentHostProgress', []); readonly isCompleteObs = observableValue('agentHostComplete', true); + /** + * Live MCP-server summaries for this session, mirroring + * `SessionState.mcpServers`. Updated by the owning session handler as + * the protocol state changes. + */ + readonly mcpServersObs = observableValue('agentHostMcpServers', []); + private readonly _onWillDispose = this._register(new Emitter()); readonly onWillDispose = this._onWillDispose.event; @@ -330,6 +339,28 @@ export interface IAgentHostSessionHandlerConfig { */ readonly resolveAuthentication?: (protectedResources: ProtectedResourceMetadata[]) => Promise; + /** + * Optional callback invoked when the user requests authentication for a + * specific MCP server (e.g. by clicking the MCP-auth indicator on the + * chat input). The token MUST be forwarded to the agent host scoped to + * `mcpServer` (i.e. {@link AuthenticateParams.server}). + * + * `scopes` carries the authoritative scopes for this challenge — for + * AHP this is `McpServerStatusAuthRequired.requiredScopes` when set, + * falling back to `resource.scopes_supported`. + */ + readonly resolveMcpAuthentication?: (mcpServer: URI, resource: ProtectedResourceMetadata, scopes: readonly string[] | undefined) => Promise; + + /** + * Optional callback used for the auto-recovery path: when an MCP + * server transitions to `AuthRequired`, the handler tries this once + * before surfacing the click-to-authenticate indicator. The + * implementation MUST NOT prompt the user — it should look up an + * existing OAuth session matching `scopes` and forward its token if + * available, otherwise return `false`. + */ + readonly resolveMcpAuthenticationSilently?: (mcpServer: URI, resource: ProtectedResourceMetadata, scopes: readonly string[] | undefined) => Promise; + /** * Observable set of agent-level customizations to include in the active * client set. When the value changes, active sessions are updated. @@ -356,6 +387,21 @@ function offsetToPosition(text: string, offset: number): IPosition { } return { lineNumber, column }; } + +/** + * Compact, deterministic fingerprint of an MCP server's + * `AuthRequired` challenge: the canonical resource URI plus the sorted + * scope list. Used to dedupe automatic silent re-auth attempts so a + * rejected silent token doesn't loop, while a subsequent `AuthRequired` + * with new scopes still re-arms the silent path. + */ +function mcpAuthChallengeFingerprint(summary: McpServerSummary): string { + if (summary.status.kind !== McpServerStatusKind.AuthRequired) { + return ''; + } + const scopes = [...(summary.status.requiredScopes ?? summary.status.resource.scopes_supported ?? [])].sort(); + return `${summary.status.resource.resource}|${scopes.join(' ')}`; +} export class AgentHostSessionHandler extends Disposable implements IChatSessionContentProvider { private readonly _activeSessions = new ResourceMap(); @@ -391,6 +437,8 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC @IConfigurationService private readonly _configurationService: IConfigurationService, @IChatWidgetService private readonly _chatWidgetService: IChatWidgetService, @ILanguageModelsService private readonly _languageModelsService: ILanguageModelsService, + @IAgentHostMcpAuthRegistry private readonly _mcpAuthRegistry: IAgentHostMcpAuthRegistry, + @IAgentHostActiveClientRegistry private readonly _activeClientRegistry: IAgentHostActiveClientRegistry, ) { super(); this._config = config; @@ -477,6 +525,19 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC })); } + // Publish our active-client snapshot under this handler's session + // type so the sessions provider can read it from + // `NewSession.eagerCreate` and pass it as the `activeClient` + // parameter on `connection.createSession`. That establishes the + // active client at the moment the backend session is created, + // avoiding a separate post-hoc dispatch for provisional + // (new-chat-input) sessions. + this._register(this._activeClientRegistry.register(config.sessionType, () => ({ + clientId: this._config.connection.clientId, + tools: this._clientToolsObs.get().map(toolDataToDefinition), + customizations: this._config.customizations?.get() ?? [], + }))); + this._registerAgent(); } @@ -650,11 +711,9 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC return true; }, ); - this._activeSessions.set(sessionResource, session); + this._registerActiveSession(sessionResource, resolvedSession, session); if (!isNewSession) { - this._ensurePendingMessageSubscription(sessionResource, resolvedSession); - // If there are historical turns with file edits, eagerly create // the editing session once the ChatModel is available so that // edit pills render with diff info on session restore. @@ -671,10 +730,6 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC if (activeTurnId && initialProgress !== undefined) { this._reconnectToActiveTurn(resolvedSession, activeTurnId, session, initialProgress); } - - // For existing sessions, start watching for server-initiated turns - // immediately. For new sessions, this is deferred to _createAndSubscribe. - this._watchForServerInitiatedTurns(resolvedSession, sessionResource); } return session; @@ -733,16 +788,20 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC // sessions provider involved, agent host not connected at // folder-pick time, or this session was created via a legacy/ // test path). Fall back to the original create-then-subscribe - // flow. + // flow. `_createAndSubscribe` will call + // {@link _ensureSessionSubscription} after the backend session + // exists, which wires MCP-server state into the chat session + // via the auto-wiring path described on + // {@link _ensureSessionSubscription}. await this._createAndSubscribe(request.sessionResource, this._createModelSelection(request.userSelectedModelId, request.modelConfiguration), undefined, request.agentHostSessionConfig); } else { // Eager-created session: take a refcounted subscription so the // handler observes state changes for the duration of the chat - // session, then wire up the per-turn machinery that - // `_createAndSubscribe` would normally set up. + // session. The `_ensureSessionSubscription` call also auto-wires + // the per-session bindings (MCP-server state, pending-message + // sync, server-initiated turn detection) via + // {@link _wireActiveSession} on first creation. this._ensureSessionSubscription(sessionKey); - this._ensurePendingMessageSubscription(request.sessionResource, resolvedSession); - this._watchForServerInitiatedTurns(resolvedSession, request.sessionResource); // Push the user-selected session config (e.g. isolation = worktree) // to the agent so its provisional record materializes with the @@ -2349,7 +2408,11 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC this._logService.trace(`[AgentHost] Created session: ${session.toString()}`); - // Subscribe to the new session's state + // Subscribe to the new session's state. The first call here also + // auto-wires the per-session bindings (MCP-server state, pending- + // message sync, server-initiated turn detection) into the chat + // session via {@link _wireActiveSession} \u2014 see + // {@link _ensureSessionSubscription}. const newSub = this._ensureSessionSubscription(session.toString()); if (!this._getSessionState(session.toString())) { // Wait for the subscription to hydrate @@ -2358,12 +2421,6 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC }); } - // Start syncing the chat model's pending requests to the protocol - this._ensurePendingMessageSubscription(sessionResource, session); - - // Start watching for server-initiated turns on this session - this._watchForServerInitiatedTurns(session, sessionResource); - return session; } @@ -2660,16 +2717,66 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC /** * Get or create a session subscription. The first call for a given URI * triggers a server subscribe; subsequent calls increment the refcount. + * + * On first-create, if a chat session for this backend URI is already + * registered in {@link _activeSessions}, the per-session bindings are + * auto-wired into it via {@link _wireActiveSession}. The complementary + * half of this pact lives in {@link _registerActiveSession}: when an + * active chat session is registered after the handler has already + * subscribed, that helper performs the wiring instead. Between the two, + * each session's bindings are wired exactly once, regardless of which + * event (subscribe vs. activate) happens first. + * + * Subagent child sessions and lazy seed-only subscriptions never appear + * in {@link _activeSessions}, so the auto-wire is a safe no-op for them. */ private _ensureSessionSubscription(sessionUri: string): IAgentSubscription { - let ref = this._sessionSubscriptions.get(sessionUri); - if (!ref) { - ref = this._config.connection.getSubscription(StateComponents.Session, URI.parse(sessionUri)); - this._sessionSubscriptions.set(sessionUri, ref); + const existing = this._sessionSubscriptions.get(sessionUri); + if (existing) { + return existing.object; + } + const ref = this._config.connection.getSubscription(StateComponents.Session, URI.parse(sessionUri)); + this._sessionSubscriptions.set(sessionUri, ref); + for (const [chatResource, chatSession] of this._activeSessions) { + if (this._resolveSessionUri(chatResource).toString() === sessionUri) { + this._wireActiveSession(chatResource, URI.parse(sessionUri), chatSession); + break; + } } return ref.object; } + /** + * Register a freshly-constructed chat session as active and, when the + * handler is already subscribed to its backend URI, wire its per-session + * bindings immediately via {@link _wireActiveSession}. When no + * subscription exists yet, wiring is deferred until the first + * {@link _ensureSessionSubscription} call for the backend URI (see the + * auto-wire path there). The dual trigger guarantees each session's + * bindings are wired exactly once, no matter which event — subscribe or + * activate — happens first. + */ + private _registerActiveSession(sessionResource: URI, backendSession: URI, session: AgentHostChatSession): void { + this._activeSessions.set(sessionResource, session); + if (this._sessionSubscriptions.has(backendSession.toString())) { + this._wireActiveSession(sessionResource, backendSession, session); + } + } + + /** + * Per-session bindings that all require both an active subscription on + * the backend URI *and* the chat session being live: MCP-server state + * routing, chat-model pending-request sync, and server-initiated turn + * detection. Called once per chat session from whichever of + * {@link _ensureSessionSubscription} or {@link _registerActiveSession} + * observes the second of the two preconditions arriving. + */ + private _wireActiveSession(sessionResource: URI, backendSession: URI, session: AgentHostChatSession): void { + this._wireMcpServerState(sessionResource, backendSession, session); + this._ensurePendingMessageSubscription(sessionResource, backendSession); + this._watchForServerInitiatedTurns(backendSession, sessionResource); + } + /** * Release a session subscription, decrementing refcount and unsubscribing * when it reaches zero. @@ -2694,6 +2801,113 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC return (value && !(value instanceof Error)) ? value : undefined; } + /** + * Wires the per-session MCP server state into `session.mcpServersObs` + * and registers the session with the workbench-scoped MCP auth + * registry so the chat input UI can look it up. Both bindings are + * torn down with the chat session. + * + * Only reachable through {@link _wireActiveSession}, which guarantees + * both the backend subscription and the active chat session are live + * before invocation — subscribing prior to `createSession` races the + * wire and fails with `AHP_SESSION_NOT_FOUND`, which silently buffers + * per-session envelopes without firing `onDidChange`. + */ + private _wireMcpServerState(sessionResource: URI, backendSession: URI, session: AgentHostChatSession): void { + const sub = this._ensureSessionSubscription(backendSession.toString()); + + // Per-server fingerprint of the most recent silent-auth attempt, + // keyed by `McpServerSummary.resource` (the server URI). We retry + // silent auth only when the fingerprint changes — i.e., a + // genuinely new `AuthRequired` episode (different resource or + // scopes) — so that a token rejection doesn't loop. Cleared when + // the server leaves `AuthRequired` so future challenges re-arm. + const lastSilentAttempt = new Map(); + + const update = () => { + const state = sub.value; + const servers = (state && !(state instanceof Error)) ? (state.mcpServers ?? []) : []; + session.mcpServersObs.set(servers, undefined); + + // React to AuthRequired transitions: try a silent auth pass + // for any server whose challenge fingerprint we haven't yet + // attempted in this episode. If it succeeds the host + // transitions the server back to Ready (no UI), otherwise + // the indicator surfaces normally. + const liveServerKeys = new Set(); + for (const summary of servers) { + liveServerKeys.add(summary.resource); + if (summary.status.kind !== McpServerStatusKind.AuthRequired) { + lastSilentAttempt.delete(summary.resource); + continue; + } + const fingerprint = mcpAuthChallengeFingerprint(summary); + if (lastSilentAttempt.get(summary.resource) === fingerprint) { + continue; + } + lastSilentAttempt.set(summary.resource, fingerprint); + this._tryMcpServerSilentAuth(summary); + } + // Drop memo entries for servers that disappeared. + for (const key of lastSilentAttempt.keys()) { + if (!liveServerKeys.has(key)) { + lastSilentAttempt.delete(key); + } + } + }; + update(); + session.registerDisposable(sub.onDidChange(update)); + + const entry: IAgentHostMcpAuthSessionEntry = { + mcpServers: session.mcpServersObs, + authenticate: server => this._authenticateMcpServers(session, server), + }; + session.registerDisposable(this._mcpAuthRegistry.registerSession(sessionResource, entry)); + } + + /** + * Drives interactive authentication for one or all MCP servers in an + * `AuthRequired` state on the given session. + */ + private async _authenticateMcpServers(session: AgentHostChatSession, target?: McpServerSummary): Promise { + const resolveMcp = this._config.resolveMcpAuthentication; + if (!resolveMcp) { + return false; + } + const summaries = session.mcpServersObs.get(); + const candidates = (target ? [target] : summaries.filter(s => s.status.kind === McpServerStatusKind.AuthRequired)); + let anySucceeded = false; + for (const summary of candidates) { + if (summary.status.kind !== McpServerStatusKind.AuthRequired) { + continue; + } + try { + const ok = await resolveMcp(URI.parse(summary.resource), summary.status.resource, summary.status.requiredScopes); + anySucceeded ||= ok; + } catch (err) { + this._logService.error(`[AgentHost] MCP authentication failed for ${summary.resource}`, err); + } + } + return anySucceeded; + } + + /** + * Fire-and-forget silent re-authentication for a single MCP server + * that just transitioned to (or refreshed) `AuthRequired`. Errors + * are swallowed; the indicator path remains the user-visible + * fallback. + */ + private _tryMcpServerSilentAuth(summary: McpServerSummary): void { + const resolveMcpSilently = this._config.resolveMcpAuthenticationSilently; + if (!resolveMcpSilently || summary.status.kind !== McpServerStatusKind.AuthRequired) { + return; + } + const status = summary.status; + Promise.resolve() + .then(() => resolveMcpSilently(URI.parse(summary.resource), status.resource, status.requiredScopes)) + .catch(err => this._logService.warn(`[AgentHost] Silent MCP authentication failed for ${summary.resource}`, err)); + } + /** * Read the current root state. */ diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/chatMcpAuthAction.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/chatMcpAuthAction.ts new file mode 100644 index 0000000000000..31be695d35368 --- /dev/null +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/chatMcpAuthAction.ts @@ -0,0 +1,196 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { IAction, Separator } from '../../../../../../base/common/actions.js'; +import { Codicon } from '../../../../../../base/common/codicons.js'; +import { localize, localize2 } from '../../../../../../nls.js'; +import { IAccessibilityService } from '../../../../../../platform/accessibility/common/accessibility.js'; +import { IMenuEntryActionViewItemOptions, MenuEntryActionViewItem } from '../../../../../../platform/actions/browser/menuEntryActionViewItem.js'; +import { Action2, MenuId, MenuItemAction, registerAction2 } from '../../../../../../platform/actions/common/actions.js'; +import { McpAuthRequiredReason, McpServerStatusKind, type McpServerSummary } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import { ContextKeyExpr, IContextKeyService } from '../../../../../../platform/contextkey/common/contextkey.js'; +import { IContextMenuService } from '../../../../../../platform/contextview/browser/contextView.js'; +import { IInstantiationService, ServicesAccessor } from '../../../../../../platform/instantiation/common/instantiation.js'; +import { IKeybindingService } from '../../../../../../platform/keybinding/common/keybinding.js'; +import { INotificationService } from '../../../../../../platform/notification/common/notification.js'; +import { IThemeService } from '../../../../../../platform/theme/common/themeService.js'; +import { ChatContextKeys } from '../../../common/actions/chatContextKeys.js'; +import { type IChatExecuteActionContext } from '../../actions/chatExecuteActions.js'; +import { IAgentHostMcpAuthRegistry, type IAgentHostMcpAuthSessionEntry } from './agentHostMcpAuthRegistry.js'; +import './media/chatMcpAuthAction.css'; + +const CHAT_CATEGORY = localize2('chat.category', 'Chat'); + +/** + * Toolbar action that surfaces an MCP icon next to the chat send button + * when the active session has at least one MCP server requiring + * authentication. Clicking opens a context menu listing each server, + * plus an "Allow All" entry that authenticates every server in + * `AuthRequired` state in one go. + */ +export class OpenMcpAuthAction extends Action2 { + + static readonly ID = 'workbench.action.chat.openMcpAuth'; + + constructor() { + super({ + id: OpenMcpAuthAction.ID, + title: localize2('chat.openMcpAuth.label', "Authenticate MCP Servers"), + tooltip: localize('chat.openMcpAuth.tooltip', "One or more MCP servers require authentication"), + f1: false, + category: CHAT_CATEGORY, + icon: Codicon.mcp, + menu: [{ + id: MenuId.ChatExecute, + when: ContextKeyExpr.notEquals(ChatContextKeys.mcpAuthRequiredCount.key, 0), + order: 2, + group: 'navigation', + }], + }); + } + + // `run` is a no-op fallback. The custom view item ({@link McpAuthActionViewItem}) + // handles clicks and shows a context menu anchored at the icon. This + // path is only reached if the action is invoked outside the toolbar + // (e.g. via the command palette, which is suppressed by `f1: false`). + override async run(_accessor: ServicesAccessor, ..._args: unknown[]): Promise { + // no-op + } +} + +/** + * Custom view item that renders the {@link OpenMcpAuthAction} icon and, + * on click, opens a context menu listing each MCP server requiring + * authentication for the current chat session. + */ +export class McpAuthActionViewItem extends MenuEntryActionViewItem { + + constructor( + action: MenuItemAction, + options: IMenuEntryActionViewItemOptions | undefined, + @IKeybindingService keybindingService: IKeybindingService, + @INotificationService notificationService: INotificationService, + @IContextKeyService contextKeyService: IContextKeyService, + @IThemeService themeService: IThemeService, + @IContextMenuService contextMenuService: IContextMenuService, + @IAccessibilityService accessibilityService: IAccessibilityService, + @IAgentHostMcpAuthRegistry private readonly _registry: IAgentHostMcpAuthRegistry, + ) { + super(action, options, keybindingService, notificationService, contextKeyService, themeService, contextMenuService, accessibilityService); + } + + override render(container: HTMLElement): void { + super.render(container); + container.classList.add('chat-mcp-auth-action'); + } + + override async onClick(event: MouseEvent): Promise { + event.preventDefault(); + event.stopPropagation(); + + const widget = (this._context as IChatExecuteActionContext | undefined)?.widget; + const sessionResource = widget?.viewModel?.sessionResource; + const entry = sessionResource ? this._registry.getEntry(sessionResource) : undefined; + if (!entry || !this.element) { + return; + } + showMcpAuthContextMenu(entry, this.element, this._contextMenuService); + } +} + +/** + * Opens the MCP authentication context menu anchored to `anchor` for + * the given session entry. Used by both the chat-input toolbar's + * {@link McpAuthActionViewItem} and the sessions window's + * `NewChatMcpAuthIndicator` so the two surfaces present identical + * affordances. + */ +export function showMcpAuthContextMenu( + entry: IAgentHostMcpAuthSessionEntry, + anchor: HTMLElement, + contextMenuService: IContextMenuService, +): void { + const summaries = entry.mcpServers.get(); + const pending = summaries.filter(s => s.status.kind === McpServerStatusKind.AuthRequired); + if (pending.length === 0) { + return; + } + + const actions: IAction[] = pending.map(summary => ({ + id: `chatMcpAuth.server.${summary.resource}`, + label: summary.label, + tooltip: describeAuthRequired(summary), + class: undefined, + enabled: true, + checked: undefined, + run: () => entry.authenticate(summary).catch(() => false), + })); + + if (pending.length > 1) { + actions.push(new Separator()); + actions.push({ + id: 'chatMcpAuth.allowAll', + label: localize('chat.openMcpAuth.allowAll', "Authenticate All"), + tooltip: '', + class: undefined, + enabled: true, + checked: undefined, + run: () => entry.authenticate().catch(() => false), + }); + } + + contextMenuService.showContextMenu({ + getAnchor: () => anchor, + getActions: () => actions, + }); +} + +function describeAuthRequired(summary: McpServerSummary): string { + if (summary.status.kind !== McpServerStatusKind.AuthRequired) { + return summary.label; + } + if (summary.status.description) { + return summary.status.description; + } + switch (summary.status.reason) { + case McpAuthRequiredReason.Expired: + return localize('chat.openMcpAuth.reason.expired', "Authentication expired"); + case McpAuthRequiredReason.InsufficientScope: + return localize('chat.openMcpAuth.reason.insufficientScope', "Additional permissions required"); + case McpAuthRequiredReason.Required: + default: + return localize('chat.openMcpAuth.reason.required', "Authentication required"); + } +} + +/** + * Returns true if the action describes the {@link OpenMcpAuthAction}, so + * the chat input toolbar can swap in {@link McpAuthActionViewItem} for + * its slot. + */ +export function isMcpAuthAction(action: IAction): action is MenuItemAction { + return action instanceof MenuItemAction && action.id === OpenMcpAuthAction.ID; +} + +/** + * Creates a {@link McpAuthActionViewItem} for `action` if it is the MCP + * auth toolbar action, otherwise returns undefined. Used as a branch in + * the chat input toolbar's `actionViewItemProvider`. + */ +export function createMcpAuthActionViewItem(action: IAction, options: IMenuEntryActionViewItemOptions | undefined, instantiationService: IInstantiationService): McpAuthActionViewItem | undefined { + if (!isMcpAuthAction(action)) { + return undefined; + } + return instantiationService.createInstance(McpAuthActionViewItem, action, options); +} + +export function registerMcpAuthActions(): void { + registerAction2(OpenMcpAuthAction); +} + +// The factory is exported so consumers can fetch it via the +// instantiation service as needed; the unused parameter silences the +// lint rule that complains about unused references. +export type { IAgentHostMcpAuthSessionEntry }; diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/media/chatMcpAuthAction.css b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/media/chatMcpAuthAction.css new file mode 100644 index 0000000000000..188fbae1fda8a --- /dev/null +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/media/chatMcpAuthAction.css @@ -0,0 +1,32 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +.monaco-action-bar .action-item.chat-mcp-auth-action > .action-label.codicon { + color: var(--vscode-notificationsWarningIcon-foreground); + animation: chat-mcp-auth-pulse 2s ease-in-out infinite; +} + +.monaco-action-bar .action-item.chat-mcp-auth-action:hover > .action-label.codicon, +.monaco-action-bar .action-item.chat-mcp-auth-action:focus-within > .action-label.codicon { + color: var(--vscode-notificationsWarningIcon-foreground); + animation: none; +} + +@keyframes chat-mcp-auth-pulse { + 0%, 100% { + opacity: 1; + transform: scale(1); + } + 50% { + opacity: 0.65; + transform: scale(1.08); + } +} + +@media (prefers-reduced-motion: reduce) { + .monaco-action-bar .action-item.chat-mcp-auth-action > .action-label.codicon { + animation: none; + } +} diff --git a/src/vs/workbench/contrib/chat/browser/chat.contribution.ts b/src/vs/workbench/contrib/chat/browser/chat.contribution.ts index 8c279b3fc2eae..f2a6a56788e2d 100644 --- a/src/vs/workbench/contrib/chat/browser/chat.contribution.ts +++ b/src/vs/workbench/contrib/chat/browser/chat.contribution.ts @@ -111,6 +111,10 @@ import { ChatDebugEditor } from './chatDebug/chatDebugEditor.js'; import { PromptsDebugContribution } from './promptsDebugContribution.js'; import { ChatDebugEditorInput, ChatDebugEditorInputSerializer } from './chatDebug/chatDebugEditorInput.js'; import './agentSessions/agentSessions.contribution.js'; +import './agentSessions/agentHost/agentHostActiveClientRegistry.js'; +import './agentSessions/agentHost/agentHostMcpAuthRegistry.js'; +import { AgentHostMcpAuthIndicatorContribution } from './agentSessions/agentHost/agentHostMcpAuthIndicatorContribution.js'; +import { registerMcpAuthActions } from './agentSessions/agentHost/chatMcpAuthAction.js'; import { ChatContextKeys } from '../common/actions/chatContextKeys.js'; @@ -2262,6 +2266,7 @@ registerWorkbenchContribution2(PromptLanguageFeaturesProvider.ID, PromptLanguage registerWorkbenchContribution2(ChatWindowNotifier.ID, ChatWindowNotifier, WorkbenchPhase.AfterRestored); registerWorkbenchContribution2(ChatRepoInfoContribution.ID, ChatRepoInfoContribution, WorkbenchPhase.Eventually); registerWorkbenchContribution2(AgentPluginRecommendations.ID, AgentPluginRecommendations, WorkbenchPhase.Eventually); +registerWorkbenchContribution2(AgentHostMcpAuthIndicatorContribution.ID, AgentHostMcpAuthIndicatorContribution, WorkbenchPhase.AfterRestored); registerChatActions(); registerChatAccessibilityActions(); @@ -2273,6 +2278,7 @@ registerChatFileTreeActions(); registerChatPromptNavigationActions(); registerChatTitleActions(); registerChatExecuteActions(); +registerMcpAuthActions(); registerChatQueueActions(); registerQuickChatActions(); registerChatExportActions(); diff --git a/src/vs/workbench/contrib/chat/browser/widget/input/chatInputPart.ts b/src/vs/workbench/contrib/chat/browser/widget/input/chatInputPart.ts index 1bb83a8fd5019..a1dc4a0a229a6 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/input/chatInputPart.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/input/chatInputPart.ts @@ -95,6 +95,7 @@ import { IChatAgentService } from '../../../common/participants/chatAgents.js'; import { ILanguageModelToolsService } from '../../../common/tools/languageModelToolsService.js'; import { ChatHistoryNavigator } from '../../../common/widget/chatWidgetHistoryService.js'; import { ChatSessionPrimaryPickerAction, ChatSubmitAction, IChatExecuteActionContext, OpenDelegationPickerAction, OpenModelPickerAction, OpenModePickerAction, OpenPermissionPickerAction, OpenSessionTargetPickerAction, OpenWorkspacePickerAction } from '../../actions/chatExecuteActions.js'; +import { createMcpAuthActionViewItem } from '../../agentSessions/agentHost/chatMcpAuthAction.js'; import { AgentSessionProviders, getAgentSessionProvider } from '../../agentSessions/agentSessions.js'; import { IAgentSessionsService } from '../../agentSessions/agentSessionsService.js'; import { ChatAttachmentModel } from '../../attachments/chatAttachmentModel.js'; @@ -2508,6 +2509,9 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge }, hoverDelegate, hiddenItemStrategy: HiddenItemStrategy.NoHide, + actionViewItemProvider: (action, options) => { + return createMcpAuthActionViewItem(action, options, this.instantiationService); + }, })); this.executeToolbar.getElement().classList.add('chat-execute-toolbar'); this.executeToolbar.context = { widget } satisfies IChatExecuteActionContext; diff --git a/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts b/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts index cc37dc7ede093..2a4f88ec90a30 100644 --- a/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts +++ b/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts @@ -20,6 +20,7 @@ export namespace ChatContextKeys { export const responseHasError = new RawContextKey('chatSessionResponseError', false, { type: 'boolean', description: localize('chatResponseErrored', "True when the chat response resulted in an error.") }); export const requestInProgress = new RawContextKey('chatSessionRequestInProgress', false, { type: 'boolean', description: localize('interactiveSessionRequestInProgress', "True when the current request is still in progress.") }); export const hasActiveRequest = new RawContextKey('chatSessionHasActiveRequest', false, { type: 'boolean', description: localize('chatSessionHasActiveRequest', "True when the current chat response has not completed, regardless of intermediate states like tool calls or elicitations.") }); + export const mcpAuthRequiredCount = new RawContextKey('chatSessionMcpAuthRequiredCount', 0, { type: 'number', description: localize('chatSessionMcpAuthRequiredCount', "Count of MCP servers attached to the active chat session that require authentication. Greater than 0 when at least one server is awaiting auth.") }); export const currentlyEditing = new RawContextKey('chatSessionCurrentlyEditing', false, { type: 'boolean', description: localize('interactiveSessionCurrentlyEditing', "True when the current request is being edited.") }); export const currentlyEditingInput = new RawContextKey('chatSessionCurrentlyEditingInput', false, { type: 'boolean', description: localize('interactiveSessionCurrentlyEditingInput', "True when the current request input at the bottom is being edited.") }); diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostAuth.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostAuth.test.ts index 29c0ecc5265b1..75920aedee9e6 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostAuth.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostAuth.test.ts @@ -10,7 +10,8 @@ import { type AgentInfo } from '../../../../../../platform/agentHost/common/stat import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; import { NullLogService } from '../../../../../../platform/log/common/log.js'; import { IAuthenticationService } from '../../../../../services/authentication/common/authentication.js'; -import { authenticateProtectedResources, resolveAuthenticationInteractively, resolveTokenForResource, AgentHostAuthTokenCache } from '../../../browser/agentSessions/agentHost/agentHostAuth.js'; +import { authenticateProtectedResources, resolveAuthenticationInteractively, resolveMcpServerAuthenticationInteractively, resolveMcpServerAuthenticationSilently, resolveTokenForResource, AgentHostAuthTokenCache, type IAgentHostAuthenticateRequest } from '../../../browser/agentSessions/agentHost/agentHostAuth.js'; +import { type IAgentHostMcpAuthRegistry } from '../../../browser/agentSessions/agentHost/agentHostMcpAuthRegistry.js'; function createMockAuthService(overrides: { getOrActivateProviderIdForServer?: (serverUri: URI, resourceUri: URI) => Promise; @@ -274,3 +275,274 @@ suite('resolveAuthenticationInteractively', () => { assert.deepStrictEqual(requests, [{ resource: protectedResource.resource, token: 'new-token' }]); }); }); + +/** + * In-memory stub of the persistent-scope-memory subset of + * {@link IAgentHostMcpAuthRegistry} used by the per-server MCP helpers. + * Keyed by `host|resource` so we can assert the helpers honor host + * scoping (a token approved against host A must NOT silently + * authenticate host B). + */ +class FakeMcpAuthMemory implements Pick { + readonly entries = new Map(); + readonly rememberCalls: { host: string; resource: string; scopes: readonly string[] }[] = []; + + private static key(host: string, resource: string): string { + return `${host}|${resource}`; + } + + recall(host: string, resource: string): readonly string[] | undefined { + return this.entries.get(FakeMcpAuthMemory.key(host, resource)); + } + + remember(host: string, resource: string, scopes: readonly string[]): void { + const sorted = [...new Set(scopes)].sort(); + this.entries.set(FakeMcpAuthMemory.key(host, resource), sorted); + this.rememberCalls.push({ host, resource, scopes: sorted }); + } + + forget(host: string, resource: string): void { + this.entries.delete(FakeMcpAuthMemory.key(host, resource)); + } +} + +suite('resolveMcpServerAuthenticationSilently', () => { + + const log = new NullLogService(); + const protectedResource: ProtectedResourceMetadata = { + resource: 'https://api.example.com', + authorization_servers: ['https://auth.example.com'], + scopes_supported: ['read'], + }; + const mcpServer = URI.parse('mcp:/session-1/server-1'); + + ensureNoDisposablesAreLeakedInTestSuite(); + + test('forwards token scoped to the MCP server URI on success', async () => { + const authService = createMockAuthService({ + getOrActivateProviderIdForServer: () => Promise.resolve('provider-1'), + getSessions: (_id, scopes) => Promise.resolve(scopes + ? [{ scopes: ['read'], accessToken: 'silent-token' }] + : []), + }); + const requests: IAgentHostAuthenticateRequest[] = []; + + const success = await resolveMcpServerAuthenticationSilently(mcpServer, protectedResource, undefined, { + authenticationService: authService, + logPrefix: '[Test]', + logService: log, + authenticate: async request => { requests.push(request); }, + }); + + assert.strictEqual(success, true); + assert.strictEqual(requests.length, 1); + assert.strictEqual(requests[0].token, 'silent-token'); + assert.strictEqual(requests[0].resource, protectedResource.resource); + assert.strictEqual(requests[0].server?.toString(), mcpServer.toString()); + }); + + test('returns false and does NOT prompt when no matching session exists', async () => { + let createSessionCalls = 0; + const authService = createMockAuthService({ + getOrActivateProviderIdForServer: () => Promise.resolve('provider-1'), + getSessions: () => Promise.resolve([]), + createSession: async () => { createSessionCalls++; return { accessToken: 'created' }; }, + }); + const requests: IAgentHostAuthenticateRequest[] = []; + + const success = await resolveMcpServerAuthenticationSilently(mcpServer, protectedResource, undefined, { + authenticationService: authService, + logPrefix: '[Test]', + logService: log, + authenticate: async request => { requests.push(request); }, + }); + + assert.strictEqual(success, false); + assert.strictEqual(requests.length, 0); + assert.strictEqual(createSessionCalls, 0); + }); + + test('prefers remembered scopes over current challenge scopes', async () => { + const requestedScopeSets: (readonly string[] | undefined)[] = []; + const authService = createMockAuthService({ + getOrActivateProviderIdForServer: () => Promise.resolve('provider-1'), + getSessions: (_id, scopes) => { + requestedScopeSets.push(scopes); + if (scopes && scopes.length === 2 && scopes.includes('read') && scopes.includes('write')) { + return Promise.resolve([{ scopes: ['read', 'write'], accessToken: 'wide-token' }]); + } + return Promise.resolve([]); + }, + }); + const memory = new FakeMcpAuthMemory(); + memory.entries.set('local|' + protectedResource.resource, ['read', 'write']); + + // The current challenge says only 'read' — but the user previously + // approved 'read'+'write' on this host, so the silent path should + // ask for those instead so the broader session is reused. + const success = await resolveMcpServerAuthenticationSilently(mcpServer, protectedResource, ['read'], { + mcpAuthMemory: memory, + mcpAuthHostKey: 'local', + authenticationService: authService, + logPrefix: '[Test]', + logService: log, + authenticate: async () => { }, + }); + + assert.strictEqual(success, true); + // First call to getSessions used the remembered, broader scopes. + assert.deepStrictEqual(requestedScopeSets[0], ['read', 'write']); + }); + + test('host-scoped: tokens approved against another host are NOT reused', async () => { + // Memory has a record for `host-a` but the silent attempt is for `host-b`. + // The silent helper must not consult `host-a`'s entry, so it falls + // back to the challenge's scopes — and since no matching session + // exists, the call returns false. + const memory = new FakeMcpAuthMemory(); + memory.entries.set('host-a|' + protectedResource.resource, ['read', 'write', 'admin']); + + const askedScopes: (readonly string[] | undefined)[] = []; + const authService = createMockAuthService({ + getOrActivateProviderIdForServer: () => Promise.resolve('provider-1'), + getSessions: (_id, scopes) => { + askedScopes.push(scopes); + return Promise.resolve([]); + }, + }); + + const success = await resolveMcpServerAuthenticationSilently(mcpServer, protectedResource, ['read'], { + mcpAuthMemory: memory, + mcpAuthHostKey: 'host-b', + authenticationService: authService, + logPrefix: '[Test]', + logService: log, + authenticate: async () => { }, + }); + + assert.strictEqual(success, false); + // Asked exactly once with the challenge scopes (no memory leakage). + assert.deepStrictEqual(askedScopes[0], ['read']); + }); + + test('records remembered scopes on success when host key is provided', async () => { + const authService = createMockAuthService({ + getOrActivateProviderIdForServer: () => Promise.resolve('provider-1'), + getSessions: (_id, scopes) => Promise.resolve(scopes + ? [{ scopes: ['read'], accessToken: 'silent-token' }] + : []), + }); + const memory = new FakeMcpAuthMemory(); + + await resolveMcpServerAuthenticationSilently(mcpServer, protectedResource, ['read'], { + mcpAuthMemory: memory, + mcpAuthHostKey: 'host-a', + authenticationService: authService, + logPrefix: '[Test]', + logService: log, + authenticate: async () => { }, + }); + + assert.deepStrictEqual(memory.rememberCalls, [{ host: 'host-a', resource: protectedResource.resource, scopes: ['read'] }]); + }); + + test('does NOT record when memory is provided but host key is not', async () => { + const authService = createMockAuthService({ + getOrActivateProviderIdForServer: () => Promise.resolve('provider-1'), + getSessions: (_id, scopes) => Promise.resolve(scopes + ? [{ scopes: ['read'], accessToken: 'silent-token' }] + : []), + }); + const memory = new FakeMcpAuthMemory(); + + await resolveMcpServerAuthenticationSilently(mcpServer, protectedResource, undefined, { + mcpAuthMemory: memory, + authenticationService: authService, + logPrefix: '[Test]', + logService: log, + authenticate: async () => { }, + }); + + assert.strictEqual(memory.rememberCalls.length, 0); + }); +}); + +suite('resolveMcpServerAuthenticationInteractively', () => { + + const log = new NullLogService(); + const protectedResource: ProtectedResourceMetadata = { + resource: 'https://api.example.com', + authorization_servers: ['https://auth.example.com'], + scopes_supported: ['read'], + }; + const mcpServer = URI.parse('mcp:/session-1/server-1'); + + ensureNoDisposablesAreLeakedInTestSuite(); + + test('uses an existing session before prompting', async () => { + let createSessionCalls = 0; + const authService = createMockAuthService({ + getOrActivateProviderIdForServer: () => Promise.resolve('provider-1'), + getSessions: (_id, scopes) => Promise.resolve(scopes + ? [{ scopes: ['read'], accessToken: 'silent-token' }] + : []), + createSession: async () => { createSessionCalls++; return { accessToken: 'new' }; }, + }); + const requests: IAgentHostAuthenticateRequest[] = []; + + const success = await resolveMcpServerAuthenticationInteractively(mcpServer, protectedResource, undefined, { + authenticationService: authService, + logPrefix: '[Test]', + logService: log, + authenticate: async request => { requests.push(request); }, + }); + + assert.strictEqual(success, true); + assert.strictEqual(createSessionCalls, 0); + assert.strictEqual(requests[0].token, 'silent-token'); + assert.strictEqual(requests[0].server?.toString(), mcpServer.toString()); + }); + + test('falls back to createSession with the requested scopes', async () => { + const createCalls: { scopes: string[] }[] = []; + const authService = createMockAuthService({ + getOrActivateProviderIdForServer: () => Promise.resolve('provider-1'), + getSessions: () => Promise.resolve([]), + createSession: async (_id, scopes) => { createCalls.push({ scopes: [...scopes] }); return { accessToken: 'new-token' }; }, + }); + const requests: IAgentHostAuthenticateRequest[] = []; + + const success = await resolveMcpServerAuthenticationInteractively(mcpServer, protectedResource, ['read', 'write'], { + authenticationService: authService, + logPrefix: '[Test]', + logService: log, + authenticate: async request => { requests.push(request); }, + }); + + assert.strictEqual(success, true); + assert.deepStrictEqual(createCalls, [{ scopes: ['read', 'write'] }]); + assert.strictEqual(requests[0].token, 'new-token'); + assert.strictEqual(requests[0].server?.toString(), mcpServer.toString()); + }); + + test('records the granted scopes in memory after createSession', async () => { + const authService = createMockAuthService({ + getOrActivateProviderIdForServer: () => Promise.resolve('provider-1'), + getSessions: () => Promise.resolve([]), + createSession: async () => ({ accessToken: 'new-token' }), + }); + const memory = new FakeMcpAuthMemory(); + + await resolveMcpServerAuthenticationInteractively(mcpServer, protectedResource, ['read', 'write'], { + mcpAuthMemory: memory, + mcpAuthHostKey: 'local', + authenticationService: authService, + logPrefix: '[Test]', + logService: log, + authenticate: async () => { }, + }); + + assert.deepStrictEqual(memory.rememberCalls, [{ host: 'local', resource: protectedResource.resource, scopes: ['read', 'write'] }]); + }); +}); + diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts index 091ecacb23841..27e798d6a1fe4 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts @@ -37,6 +37,8 @@ import { TestInstantiationService } from '../../../../../../platform/instantiati import { IOutputService } from '../../../../../services/output/common/output.js'; import { IWorkspaceContextService } from '../../../../../../platform/workspace/common/workspace.js'; import { AgentHostContribution, AgentHostSessionListController, AgentHostSessionHandler } from '../../../browser/agentSessions/agentHost/agentHostChatContribution.js'; +import { IAgentHostActiveClientRegistry } from '../../../browser/agentSessions/agentHost/agentHostActiveClientRegistry.js'; +import { IAgentHostMcpAuthRegistry } from '../../../browser/agentSessions/agentHost/agentHostMcpAuthRegistry.js'; import { AgentHostLanguageModelProvider } from '../../../browser/agentSessions/agentHost/agentHostLanguageModelProvider.js'; import { IFileService } from '../../../../../../platform/files/common/files.js'; import { TestFileService } from '../../../../../test/common/workbenchTestServices.js'; @@ -390,6 +392,17 @@ function createTestServices(disposables: DisposableStore, workingDirectoryResolv getTools: () => [], _serviceBrand: undefined, }); + instantiationService.stub(IAgentHostActiveClientRegistry, { + register: () => toDisposable(() => { }), + get: () => undefined, + }); + instantiationService.stub(IAgentHostMcpAuthRegistry, { + registerSession: () => toDisposable(() => { }), + getEntry: () => undefined, + remember: () => { }, + recall: () => undefined, + forget: () => { }, + }); instantiationService.stub(IOutputService, { getChannel: () => undefined }); instantiationService.stub(IWorkspaceContextService, { getWorkspace: () => ({ id: '', folders: [] }), getWorkspaceFolder: () => null }); instantiationService.stub(IChatEditingService, { diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts index d3540dd8ba520..06e33b08e592b 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostClientTools.test.ts @@ -30,6 +30,8 @@ import { IProductService } from '../../../../../../platform/product/common/produ import { TestInstantiationService } from '../../../../../../platform/instantiation/test/common/instantiationServiceMock.js'; import { IWorkspaceContextService } from '../../../../../../platform/workspace/common/workspace.js'; import { AgentHostSessionHandler, toolDataToDefinition, toolResultToProtocol } from '../../../browser/agentSessions/agentHost/agentHostSessionHandler.js'; +import { IAgentHostActiveClientRegistry } from '../../../browser/agentSessions/agentHost/agentHostActiveClientRegistry.js'; +import { IAgentHostMcpAuthRegistry } from '../../../browser/agentSessions/agentHost/agentHostMcpAuthRegistry.js'; import { IFileService } from '../../../../../../platform/files/common/files.js'; import { TestFileService } from '../../../../../test/common/workbenchTestServices.js'; import { ILabelService } from '../../../../../../platform/label/common/label.js'; @@ -455,6 +457,17 @@ suite('AgentHostClientTools', () => { isNewSession: () => false, }); instantiationService.stub(ILanguageModelToolsService, toolsService); + instantiationService.stub(IAgentHostActiveClientRegistry, { + register: () => toDisposable(() => { }), + get: () => undefined, + }); + instantiationService.stub(IAgentHostMcpAuthRegistry, { + registerSession: () => toDisposable(() => { }), + getEntry: () => undefined, + remember: () => { }, + forget: () => { }, + recall: () => undefined, + }); const handler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, { provider: 'copilot' as const, diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostMcpAuthRegistry.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostMcpAuthRegistry.test.ts new file mode 100644 index 0000000000000..42051be7a1b78 --- /dev/null +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostMcpAuthRegistry.test.ts @@ -0,0 +1,168 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { DisposableStore } from '../../../../../../base/common/lifecycle.js'; +import { observableValue } from '../../../../../../base/common/observable.js'; +import { URI } from '../../../../../../base/common/uri.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; +import { type McpServerSummary } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import { TestInstantiationService } from '../../../../../../platform/instantiation/test/common/instantiationServiceMock.js'; +import { ILogService, NullLogService } from '../../../../../../platform/log/common/log.js'; +import { IStorageService, WillSaveStateReason } from '../../../../../../platform/storage/common/storage.js'; +import { TestStorageService } from '../../../../../test/common/workbenchTestServices.js'; +import { IAgentHostMcpAuthRegistry, AgentHostMcpAuthRegistry, type IAgentHostMcpAuthSessionEntry } from '../../../browser/agentSessions/agentHost/agentHostMcpAuthRegistry.js'; + +function makeEntry(): IAgentHostMcpAuthSessionEntry { + return { + mcpServers: observableValue('test', []), + authenticate: async () => false, + }; +} + +function createRegistry(disposables: DisposableStore, storage?: TestStorageService): { registry: IAgentHostMcpAuthRegistry; storage: TestStorageService } { + const storageService = storage ?? disposables.add(new TestStorageService()); + const insta = disposables.add(new TestInstantiationService()); + insta.stub(IStorageService, storageService); + insta.stub(ILogService, new NullLogService()); + const registry = disposables.add(insta.createInstance(AgentHostMcpAuthRegistry)); + return { registry, storage: storageService }; +} + +suite('AgentHostMcpAuthRegistry — session entries', () => { + const disposables = new DisposableStore(); + teardown(() => disposables.clear()); + ensureNoDisposablesAreLeakedInTestSuite(); + + test('registerSession + getEntry round-trip', () => { + const { registry } = createRegistry(disposables); + const sessionResource = URI.parse('vscode-chat-session:/local/abc'); + const entry = makeEntry(); + + const registration = disposables.add(registry.registerSession(sessionResource, entry)); + assert.strictEqual(registry.getEntry(sessionResource), entry); + + registration.dispose(); + assert.strictEqual(registry.getEntry(sessionResource), undefined); + }); + + test('registerSession overwrites prior entry; disposing stale registration is a no-op', () => { + const { registry } = createRegistry(disposables); + const sessionResource = URI.parse('vscode-chat-session:/local/abc'); + const first = makeEntry(); + const second = makeEntry(); + + const firstReg = registry.registerSession(sessionResource, first); + const secondReg = disposables.add(registry.registerSession(sessionResource, second)); + + // Overwriting must replace the entry. + assert.strictEqual(registry.getEntry(sessionResource), second); + + // Disposing the FIRST registration must not clear the live SECOND entry. + firstReg.dispose(); + assert.strictEqual(registry.getEntry(sessionResource), second); + + secondReg.dispose(); + assert.strictEqual(registry.getEntry(sessionResource), undefined); + }); +}); + +suite('AgentHostMcpAuthRegistry — persistent scope memory', () => { + const disposables = new DisposableStore(); + teardown(() => disposables.clear()); + ensureNoDisposablesAreLeakedInTestSuite(); + + test('remember + recall is host-scoped', () => { + const { registry } = createRegistry(disposables); + registry.remember('host-a', 'https://api.example.com', ['read', 'write']); + + assert.deepStrictEqual(registry.recall('host-a', 'https://api.example.com'), ['read', 'write']); + assert.strictEqual(registry.recall('host-b', 'https://api.example.com'), undefined); + assert.strictEqual(registry.recall('host-a', 'https://other.example.com'), undefined); + }); + + test('remember normalizes scopes (sorted, deduped)', () => { + const { registry } = createRegistry(disposables); + registry.remember('host-a', 'https://api.example.com', ['write', 'read', 'read']); + assert.deepStrictEqual(registry.recall('host-a', 'https://api.example.com'), ['read', 'write']); + }); + + test('forget drops only the targeted (host, resource) entry', () => { + const { registry } = createRegistry(disposables); + registry.remember('host-a', 'https://api.example.com', ['read']); + registry.remember('host-a', 'https://other.example.com', ['read']); + registry.remember('host-b', 'https://api.example.com', ['read']); + + registry.forget('host-a', 'https://api.example.com'); + + assert.strictEqual(registry.recall('host-a', 'https://api.example.com'), undefined); + assert.deepStrictEqual(registry.recall('host-a', 'https://other.example.com'), ['read']); + assert.deepStrictEqual(registry.recall('host-b', 'https://api.example.com'), ['read']); + }); + + test('mutations are buffered in-memory and only flushed on onWillSaveState', () => { + const storage = new TestStorageService(); + disposables.add(storage); + const { registry } = createRegistry(disposables, storage); + + registry.remember('host-a', 'https://api.example.com', ['read']); + // Read-back through the registry sees the in-memory mutation + // even before flush, but storage itself is still empty. + assert.deepStrictEqual(registry.recall('host-a', 'https://api.example.com'), ['read']); + assert.strictEqual(storage.get('agentHost.mcpAuth.consentedScopes', /* APPLICATION_SHARED */ -2), undefined); + + // Flush: now storage holds the serialized state. + storage.testEmitWillSaveState(WillSaveStateReason.NONE); + const raw = storage.get('agentHost.mcpAuth.consentedScopes', -2); + assert.ok(raw, 'expected storage to be populated after flush'); + assert.deepStrictEqual(JSON.parse(raw!)['host-a']['https://api.example.com'].scopes, ['read']); + }); + + test('persisted state is reloaded by a fresh registry instance', () => { + const storage = new TestStorageService(); + disposables.add(storage); + + // Round 1: write through registry A. + const { registry: registryA } = createRegistry(disposables, storage); + registryA.remember('host-a', 'https://api.example.com', ['read']); + storage.testEmitWillSaveState(WillSaveStateReason.NONE); + + // Round 2: a brand-new registry sharing the same storage must + // see the previously-remembered scopes. + const disposables2 = new DisposableStore(); + try { + const { registry: registryB } = createRegistry(disposables2, storage); + assert.deepStrictEqual(registryB.recall('host-a', 'https://api.example.com'), ['read']); + } finally { + disposables2.dispose(); + } + }); + + test('flush is a no-op when no mutation occurred', () => { + const storage = new TestStorageService(); + disposables.add(storage); + const { registry } = createRegistry(disposables, storage); + + // Read-only: should not write storage. + registry.recall('host-a', 'https://api.example.com'); + storage.testEmitWillSaveState(WillSaveStateReason.NONE); + assert.strictEqual(storage.get('agentHost.mcpAuth.consentedScopes', -2), undefined); + }); + + test('forget clears empty host buckets', () => { + const storage = new TestStorageService(); + disposables.add(storage); + const { registry } = createRegistry(disposables, storage); + + registry.remember('host-a', 'https://api.example.com', ['read']); + registry.forget('host-a', 'https://api.example.com'); + storage.testEmitWillSaveState(WillSaveStateReason.NONE); + + const raw = storage.get('agentHost.mcpAuth.consentedScopes', -2); + assert.ok(raw); + const parsed = JSON.parse(raw!); + assert.strictEqual(parsed['host-a'], undefined, 'empty host bucket should be removed'); + }); +}); diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/resolveCustomizationRefs.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/resolveCustomizationRefs.test.ts index ccac5217dc7af..dc025a56427f3 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/resolveCustomizationRefs.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/resolveCustomizationRefs.test.ts @@ -170,4 +170,88 @@ suite('resolveCustomizationRefs - built-in skills', () => { // Use CancellationToken so the import isn't dead in the bundle. assert.ok(CancellationToken.None.isCancellationRequested === false); }); + + test('includes plugins whose only contribution is an MCP server', async () => { + const pluginUri = URI.file('/plugins/mcp-only'); + const mcpServerUri = URI.file('/plugins/mcp-only/mcp.json'); + const plugin: IAgentPlugin = { + uri: pluginUri, + label: 'MCP Only Plugin', + enablement: observableValue('enablement', 'enabled' as never), + remove: () => { /* no-op */ }, + hooks: observableValue('hooks', []), + commands: observableValue('commands', []), + skills: observableValue('skills', []), + agents: observableValue('agents', []), + instructions: observableValue('instructions', []), + mcpServerDefinitions: observableValue('mcpServers', [ + { name: 'test', configuration: {} as never, uri: mcpServerUri }, + ]), + }; + const promptsService = makePromptsService(new Map()); + const bundler = new FakeBundler(); + + const refs = await resolveCustomizationRefs( + promptsService, + new FakeSyncProvider(), + makeAgentPluginService([plugin]), + bundler as unknown as SyncedCustomizationBundler, + SessionType.CopilotCLI, + ); + + assert.strictEqual(bundler.received.length, 0); + assert.strictEqual(refs.length, 1); + assert.strictEqual(refs[0].displayName, 'MCP Only Plugin'); + assert.strictEqual(refs[0].uri, pluginUri.toString()); + }); + + test('omits MCP-only plugins that the user has disabled', async () => { + const pluginUri = URI.file('/plugins/mcp-only'); + const plugin: IAgentPlugin = { + uri: pluginUri, + label: 'MCP Only Plugin', + enablement: observableValue('enablement', 'enabled' as never), + remove: () => { /* no-op */ }, + hooks: observableValue('hooks', []), + commands: observableValue('commands', []), + skills: observableValue('skills', []), + agents: observableValue('agents', []), + instructions: observableValue('instructions', []), + mcpServerDefinitions: observableValue('mcpServers', [ + { name: 'test', configuration: {} as never, uri: URI.file('/plugins/mcp-only/mcp.json') }, + ]), + }; + const refs = await resolveCustomizationRefs( + makePromptsService(new Map()), + new FakeSyncProvider(new Set([pluginUri.toString()])), + makeAgentPluginService([plugin]), + new FakeBundler() as unknown as SyncedCustomizationBundler, + SessionType.CopilotCLI, + ); + assert.deepStrictEqual(refs, []); + }); + + test('skips plugins with no syncable content', async () => { + const pluginUri = URI.file('/plugins/empty'); + const plugin: IAgentPlugin = { + uri: pluginUri, + label: 'Empty Plugin', + enablement: observableValue('enablement', 'enabled' as never), + remove: () => { /* no-op */ }, + hooks: observableValue('hooks', []), + commands: observableValue('commands', []), + skills: observableValue('skills', []), + agents: observableValue('agents', []), + instructions: observableValue('instructions', []), + mcpServerDefinitions: observableValue('mcpServers', []), + }; + const refs = await resolveCustomizationRefs( + makePromptsService(new Map()), + new FakeSyncProvider(), + makeAgentPluginService([plugin]), + new FakeBundler() as unknown as SyncedCustomizationBundler, + SessionType.CopilotCLI, + ); + assert.deepStrictEqual(refs, []); + }); });