From 56fee9ac4fe6e3a71fdc181f92f8126ff7d90590 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Tue, 20 Oct 2026 13:53:39 -0700 Subject: [PATCH 1/3] agentHost: send each MCP OAuth token only to its target server MCP servers that share an OAuth protected resource and scopes received the same token. The host gave the first token to all pending authentication requests for that resource, so a server could use the credential of a different account. This change uses the new optional AHP `authenticate.target` (microsoft/agent-host-protocol#535) to identify the server for each token. - Backport `AuthenticateTarget` and `AuthenticateParams.target` into the generated protocol copy. - Workbench: send `target` (host session URI and MCP customization ID) for each MCP token. Keep provider-level tokens untargeted. Send untargeted tokens when an older host gives no target. - Agent host: reject a target that does not identify an MCP server in an existing session. Keep targeted tokens in their own slots. Do not give them to agent-level resources or `getAuthToken`. Keep the target when tokens are replayed. - Copilot: give a targeted token only to the pending request of the target server in the owning session. - Codex: keep targeted tokens per session and server, and reconnect only that session. - Mission Control relay: do not use a targeted MCP token as the relay identity. - Copilot: show a server as ready after OAuth when its new connection loads tools. The runtime does not send a `connected` status change after an OAuth reconnect, so the server stayed in "starting". Fixes https://github.com/microsoft/vscode/issues/342147 (Commit message generated by Copilot) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/vs/platform/agentHost/common/agent.ts | 4 +- .../common/agentHostExtensionProtocol.ts | 3 + .../common/state/protocol/common/commands.ts | 52 +++++++++++ .../agentHost/common/state/sessionProtocol.ts | 4 +- .../node/agentHostAuthenticationService.ts | 31 ++++--- .../node/agentHostAutomationService.ts | 3 +- .../platform/agentHost/node/agentService.ts | 14 ++- .../agentHost/node/codex/codexAgent.ts | 50 +++++++++-- .../agentHost/node/codex/codexMcpServers.ts | 29 ++++--- .../agentHost/node/copilot/copilotAgent.ts | 3 + .../node/copilot/copilotAgentSession.ts | 31 ++++++- .../missionControlAuthentication.ts | 2 +- .../node/agentHostAuthenticationScope.test.ts | 70 ++++++++++++++- .../node/agentHostAutomationService.test.ts | 11 +-- .../agentHostClientConnectionService.test.ts | 2 + .../agentHost/test/node/agentService.test.ts | 48 +++++++++- .../test/node/codex/codexAgent.test.ts | 77 +++++++++++++++- .../test/node/codex/codexMcpServers.test.ts | 15 ++++ .../agentHost/test/node/copilotAgent.test.ts | 53 ++++++++++- .../test/node/copilotAgentSession.test.ts | 84 +++++++++++++++++- .../node/missionControlAuthentication.test.ts | 17 ++++ .../common/agentHostSessionsProvider.ts | 3 + .../browser/baseAgentHostSessionsProvider.ts | 4 + .../browser/agentHostCustomizationService.ts | 5 ++ .../agentHostCustomizationService.test.ts | 3 + .../agentSessions/agentHost/agentHostAuth.ts | 24 +++-- .../agentHost/agentHostChatContribution.ts | 2 +- .../agentHostCustomizationService.ts | 10 ++- .../agentHost/agentHostSessionHandler.ts | 17 ++-- .../agentHost/stateToProgressAdapter.ts | 6 +- .../cloudSandboxConnectionCustomization.ts | 5 +- .../remoteAgentHostChatContribution.ts | 2 +- .../chat/common/chatService/chatService.ts | 2 + .../agentSessions/agentHostAuth.test.ts | 87 ++++++++++++++++--- .../agentHostClientTools.test.ts | 7 +- .../agentHostCustomizationService.test.ts | 50 +++++++++-- .../stateToProgressAdapter.test.ts | 8 +- ...loudSandboxConnectionCustomization.test.ts | 20 +++++ .../remoteAgentHostChatContribution.test.ts | 8 +- 39 files changed, 766 insertions(+), 100 deletions(-) diff --git a/src/vs/platform/agentHost/common/agent.ts b/src/vs/platform/agentHost/common/agent.ts index ce98fd4f33604b..06038a59015020 100644 --- a/src/vs/platform/agentHost/common/agent.ts +++ b/src/vs/platform/agentHost/common/agent.ts @@ -19,7 +19,7 @@ import type { CodexModelProvider, IAgentTurnTelemetryCorrelation, IAgentHostClie import type { AgentPermissionDecisionSource } from './meta/agentPermissionResponseMeta.js'; import type { IAgentChatUsage } from './meta/agentUsageMeta.js'; import type { IMcpServerConfiguration } from '../../mcp/common/mcpPlatformTypes.js'; -import type { ResolveSessionConfigResult, SessionConfigCompletionsResult } from './state/protocol/commands.js'; +import type { AuthenticateTarget, ResolveSessionConfigResult, SessionConfigCompletionsResult } from './state/protocol/commands.js'; import type { CanvasState } from './state/protocol/channels-canvas/state.js'; import { subagentChatUri, type ChatUriFormat } from './agentHostChatMinting.js'; import { ProtectedResourceMetadata, type BackgroundWork, type Changeset, type ChatInteractivity, type ChatOrigin, type ConfigSchema, type MessageAttachment, type ModelSelection, type AgentSelection, type SessionActiveClient, type ToolCallPendingConfirmationState, type ToolDefinition, ChangesSummary } from './state/protocol/state.js'; @@ -531,6 +531,8 @@ export interface AuthenticateParams { readonly token: string; /** The access token's remaining lifetime in seconds, when known. */ readonly expiresIn?: number; + /** The single consumer this token is for; when absent, tokens match by resource and scopes. */ + readonly target?: AuthenticateTarget; /** Optional client metadata. Hosts must remain usable when it is absent. */ readonly _meta?: Record; } diff --git a/src/vs/platform/agentHost/common/agentHostExtensionProtocol.ts b/src/vs/platform/agentHost/common/agentHostExtensionProtocol.ts index a713561f9f4743..012caa84cd8ab4 100644 --- a/src/vs/platform/agentHost/common/agentHostExtensionProtocol.ts +++ b/src/vs/platform/agentHost/common/agentHostExtensionProtocol.ts @@ -7,6 +7,7 @@ import { ValidatorBase, vBoolean, vEnum, vObj, vOptionalProp, vString, type Vali import type { IJSONSchema } from '../../../base/common/jsonSchema.js'; import type { IDevContainerAgentHostConnectResult } from './devContainerAgentHost.js'; import type { AgentHostDebugLogsArtifactKind, IAgentHostManagedSettingsDiagnostics, IAgentHostNetworkDiagnosticsInfo, IAgentHostNetworkFetchResult } from './agentService.js'; +import type { McpServerAuthenticateTarget } from './state/sessionProtocol.js'; import type { InitializeResult } from './state/protocol/common/commands.js'; import type { McpAuthRequirement } from './state/protocol/channels-session/state.js'; import { AgentHostArtifactRemovalCapabilityMetaKey } from './meta/agentHostArtifactRemovalMeta.js'; @@ -284,6 +285,8 @@ export interface IAgentHostWorkspaceTrustRequest { * `authenticate` command before responding. */ export interface IAgentHostMcpAuthenticationRequest { + /** Absent from hosts that predate targeted MCP authentication. */ + readonly target?: McpServerAuthenticateTarget; readonly serverName: string; readonly auth: McpAuthRequirement; } diff --git a/src/vs/platform/agentHost/common/state/protocol/common/commands.ts b/src/vs/platform/agentHost/common/state/protocol/common/commands.ts index 0df20bde69e813..57afd07d785242 100644 --- a/src/vs/platform/agentHost/common/state/protocol/common/commands.ts +++ b/src/vs/platform/agentHost/common/state/protocol/common/commands.ts @@ -1252,6 +1252,41 @@ export interface ResourceMkdirResult { // ─── authenticate ──────────────────────────────────────────────────────────── +/** + * Identifies the consumer of an authentication token. + * + * @category Commands + * @nonexhaustive + */ +export const enum AuthenticateTargetKind { + McpServer = 'mcpServer', +} + +/** + * An MCP server identified by host-published session and customization state. + * + * @category Commands + */ +export interface McpServerAuthenticateTarget { + kind: AuthenticateTargetKind.McpServer; + /** Session URI whose `SessionState.customizations` contains the MCP server. */ + session: URI; + /** + * The server's session-unique `McpServerCustomization.id`, echoed from the + * customization carrying `McpServerAuthRequiredState` or from + * `ToolCallAuthRequiredState.contributor.customizationId`. + */ + customizationId: string; +} + +/** + * Consumer of an authentication token, discriminated by `kind`. + * Extensible to support additional consumers, such as host-owned authentication. + * + * @category Commands + */ +export type AuthenticateTarget = McpServerAuthenticateTarget; + /** * Pushes a Bearer token for a protected resource. The `resource` field MUST * match a protected-resource identifier the client has discovered from the @@ -1324,6 +1359,23 @@ export interface AuthenticateParams extends BaseParams { * token. */ scopes?: string[]; + /** + * Optional consumer of the token. Clients SHOULD include this when answering + * a specific MCP server's `McpServerAuthRequiredState` or + * `ToolCallAuthRequiredState` challenge, and omit it for provider-level + * resources (`AgentInfo.protectedResources` or `AuthRequired` -32007 data). + * + * When present, the host MUST apply the token only to that MCP server in + * that session, never to sibling servers, other sessions, or agent-level + * resources, even if they share the resource and scopes. `resource` MUST + * still equal that server's advertised `McpAuthRequirement.resource.resource`. + * The host MUST reject an unknown session or customization, or a + * customization that is not an MCP server, with `InvalidParams` (-32602). + * + * Without a target, resource/scope matching is unchanged; the token may + * satisfy every matching pending challenge. Older hosts ignore this field. + */ + target?: AuthenticateTarget; } /** diff --git a/src/vs/platform/agentHost/common/state/sessionProtocol.ts b/src/vs/platform/agentHost/common/state/sessionProtocol.ts index 6c5c78f9fdeed4..2edade22bb83bb 100644 --- a/src/vs/platform/agentHost/common/state/sessionProtocol.ts +++ b/src/vs/platform/agentHost/common/state/sessionProtocol.ts @@ -44,6 +44,7 @@ export type { // Command params and results export type { + AuthenticateTarget, CreateSessionParams, DirectoryEntry, DispatchActionParams, @@ -54,6 +55,7 @@ export type { InitializeResult, ListSessionsParams, ListSessionsResult, + McpServerAuthenticateTarget, ReconnectParams, ReconnectReplayResult, ReconnectResult, @@ -84,7 +86,7 @@ export type { CreateResourceWatchResult, } from './protocol/channels-resource-watch/commands.js'; -export { ContentEncoding, ReconnectResultType, ResourceType, ResourceWriteMode } from './protocol/commands.js'; +export { AuthenticateTargetKind, ContentEncoding, ReconnectResultType, ResourceType, ResourceWriteMode } from './protocol/commands.js'; export { ResourceChangeType } from './protocol/channels-resource-watch/state.js'; export type { ResourceChange, ResourceWatchState } from './protocol/channels-resource-watch/state.js'; diff --git a/src/vs/platform/agentHost/node/agentHostAuthenticationService.ts b/src/vs/platform/agentHost/node/agentHostAuthenticationService.ts index 5d72da766f5a52..4d8ead95b6db7e 100644 --- a/src/vs/platform/agentHost/node/agentHostAuthenticationService.ts +++ b/src/vs/platform/agentHost/node/agentHostAuthenticationService.ts @@ -7,6 +7,8 @@ import { DeferredPromise } from '../../../base/common/async.js'; import { getExpirationTime, getRemainingTimeInSeconds, isExpired } from '../../../base/common/date.js'; import { Emitter, Event } from '../../../base/common/event.js'; import { Disposable, IDisposable, toDisposable } from '../../../base/common/lifecycle.js'; +import { getComparisonKey } from '../../../base/common/resources.js'; +import { URI } from '../../../base/common/uri.js'; import { createDecorator } from '../../instantiation/common/instantiation.js'; import type { ILogService } from '../../log/common/log.js'; import type { AuthenticateParams, AuthenticateResult, IAgent, IAgentHostAuthTokenRequest } from '../common/agent.js'; @@ -50,6 +52,7 @@ export interface IAgentHostAuthenticationController { } interface IStoredAuthToken { + readonly target?: AuthenticateParams['target']; readonly resource: string; readonly scopes: readonly string[]; readonly token: string; @@ -58,6 +61,7 @@ interface IStoredAuthToken { } interface IAuthenticationRequest { + readonly target?: AuthenticateParams['target']; readonly resource: string; readonly completed: DeferredPromise; } @@ -104,7 +108,7 @@ export class AgentHostAuthenticationService extends Disposable implements IAgent ...providerList.flatMap(provider => provider.getProtectedResources().map(resource => resource.resource)), ...(scope.principal ? [scope.principal.apiOrigin] : []), ]; - if (resources.includes(params.resource)) { + if (!params.target && resources.includes(params.resource)) { await scope.authenticate(params, resources); return { authenticated: true }; } @@ -123,9 +127,9 @@ export class AgentHostAuthenticationService extends Disposable implements IAgent return { authenticated }; } const scopes = normalizeAuthenticationScopes(params.scopes); - const key = this._key(params.resource, scopes); + const key = this._key(params.resource, scopes, params.target); const previousRequest = this._authenticationRequests.get(key); - const request: IAuthenticationRequest = { resource: params.resource, completed: new DeferredPromise() }; + const request: IAuthenticationRequest = { resource: params.resource, target: params.target, completed: new DeferredPromise() }; this._authenticationRequests.set(key, request); // Wake replayers only after the replacement request is visible. previousRequest?.completed.complete(); @@ -151,7 +155,7 @@ export class AgentHostAuthenticationService extends Disposable implements IAgent // server bind failure) MUST NOT prevent another provider from // accepting the same token. const matching = providerList.filter( - p => p.getProtectedResources().some(r => r.resource === params.resource), + p => !params.target && p.getProtectedResources().some(r => r.resource === params.resource), ); const settled = await Promise.allSettled( matching.map(p => p.authenticate(params.resource, params.token, params.expiresIn, params._meta)), @@ -201,10 +205,10 @@ export class AgentHostAuthenticationService extends Disposable implements IAgent this._revocations.set(key, (this._revocations.get(key) ?? 0) + 1); } else if (authenticated) { const account = readAuthenticationAccount(params); - this._tokens.set(key, { resource: params.resource, scopes, token: params.token, expiresAt, account }); + this._tokens.set(key, { resource: params.resource, scopes, token: params.token, expiresAt, account, ...(params.target ? { target: params.target } : {}) }); } const token = this._tokens.get(key)?.token; - if (previousToken !== token || authenticationAccountId(previous?.account) !== authenticationAccountId(this._tokens.get(key)?.account)) { + if (!params.target && (previousToken !== token || authenticationAccountId(previous?.account) !== authenticationAccountId(this._tokens.get(key)?.account))) { this._onDidChangeAuthToken.fire({ resource: params.resource, scopes, token }); } return { authenticated }; @@ -213,7 +217,7 @@ export class AgentHostAuthenticationService extends Disposable implements IAgent async replay(provider: IAgent): Promise { while (true) { const pending = [...this._authenticationRequests.values()].filter(request => - provider.handleAuthenticationToken || provider.getProtectedResources().some(resource => resource.resource === request.resource)); + provider.handleAuthenticationToken || (!request.target && provider.getProtectedResources().some(resource => resource.resource === request.resource))); if (pending.length === 0) { break; } @@ -232,9 +236,10 @@ export class AgentHostAuthenticationService extends Disposable implements IAgent const expiresIn = getRemainingTimeInSeconds(stored.expiresAt, now); const params: AuthenticateParams = { resource: stored.resource, scopes: stored.scopes, token: stored.token, expiresIn, + ...(stored.target ? { target: stored.target } : {}), ...(stored.account ? { _meta: authenticationAccountMeta(stored.account) } : {}), }; - if (protectedResources.has(stored.resource)) { + if (!stored.target && protectedResources.has(stored.resource)) { try { await provider.authenticate(stored.resource, stored.token, expiresIn); } catch (error) { @@ -407,7 +412,7 @@ export class AgentHostAuthenticationService extends Disposable implements IAgent } const request = { resource: resource.resource, scopes: resource.scopes_supported }; const stored = this._getAuthToken(request, false) - ?? [...this._tokens.values()].find(candidate => candidate.resource === resource.resource && !isExpired(candidate.expiresAt)); + ?? [...this._tokens.values()].find(candidate => !candidate.target && candidate.resource === resource.resource && !isExpired(candidate.expiresAt)); if (!stored) { if (resource.required !== false) { throw new ProtocolError(AhpErrorCodes.AuthRequired, 'Local session continuation requires authentication', { resources: [resource] }); @@ -453,7 +458,7 @@ export class AgentHostAuthenticationService extends Disposable implements IAgent const requested = new Set(scopes); let best: IStoredAuthToken | undefined; for (const candidate of this._tokens.values()) { - if (candidate.resource !== request.resource || candidate.scopes.length === 0 || (!includeExpired && isExpired(candidate.expiresAt))) { + if (candidate.target || candidate.resource !== request.resource || candidate.scopes.length === 0 || (!includeExpired && isExpired(candidate.expiresAt))) { continue; } if (!this._containsAll(candidate.scopes, requested)) { @@ -482,8 +487,10 @@ export class AgentHostAuthenticationService extends Disposable implements IAgent return true; } - private _key(resource: string, scopes: readonly string[]): string { - return `${resource}\x00${scopes.join('\x00')}`; + private _key(resource: string, scopes: readonly string[], target?: AuthenticateParams['target']): string { + return target + ? JSON.stringify([resource, scopes, target.kind, getComparisonKey(URI.parse(target.session)), target.customizationId]) + : `${resource}\x00${scopes.join('\x00')}`; } } diff --git a/src/vs/platform/agentHost/node/agentHostAutomationService.ts b/src/vs/platform/agentHost/node/agentHostAutomationService.ts index 77a0f2f13313a6..47b7fe1eae3977 100644 --- a/src/vs/platform/agentHost/node/agentHostAutomationService.ts +++ b/src/vs/platform/agentHost/node/agentHostAutomationService.ts @@ -38,6 +38,7 @@ import { IAgentHostProviderService } from './agentHostProviderService.js'; import { getModelTelemetryContext } from './agentHostTurnTelemetryContext.js'; import { IAgentHostClientConnectionService } from './agentHostClientConnectionService.js'; import { IAgentHostSystemWakeService } from './agentHostSystemWakeService.js'; +import { AuthenticateTargetKind } from '../common/state/sessionProtocol.js'; import { getMcpServerCustomizations } from './shared/mcpCustomizationController.js'; import { McpServerStatus, type McpAuthRequirement } from '../common/state/protocol/channels-session/state.js'; import type { IAgentHostMcpAuthenticationRequest } from '../common/agentHostExtensionProtocol.js'; @@ -994,7 +995,7 @@ export class AgentHostAutomationService extends Disposable implements IAgentHost } attempted.add(key); challenges.set(server.name, attempted); - void this._requestMcpAuthentication(session, { serverName: server.name, auth }); + void this._requestMcpAuthentication(session, { serverName: server.name, auth, target: { kind: AuthenticateTargetKind.McpServer, session, customizationId: server.id } }); } if (challenges.size > 0) { this._mcpAuthenticationChallenges.set(session, challenges); diff --git a/src/vs/platform/agentHost/node/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts index e501e3f8759baa..70b220a33f860e 100644 --- a/src/vs/platform/agentHost/node/agentService.ts +++ b/src/vs/platform/agentHost/node/agentService.ts @@ -47,7 +47,7 @@ import { AHP_CANVAS_SCHEME } from '../common/canvasUri.js'; import type { AutomationCapabilities } from '../common/state/protocol/common/commands.js'; import type { FetchAutomationRunsParams, FetchAutomationRunsResult, ListAutomationTriggerDefinitionsParams, ListAutomationTriggerDefinitionsResult, RunAutomationParams, RunAutomationResult } from '../common/state/protocol/channels-automation/commands.js'; import type { InvokeChangesetOperationParams, InvokeChangesetOperationResult } from '../common/state/protocol/channels-changeset/commands.js'; -import { AhpErrorCodes, AHP_SESSION_NOT_FOUND, ContentEncoding, JSON_RPC_INTERNAL_ERROR, JsonRpcErrorCodes, ProtocolError, ResourceChangeType, type CreateResourceWatchParams, type CreateResourceWatchResult, type ResourceCopyParams, type ResourceCopyResult, type ResourceDeleteParams, type ResourceDeleteResult, type ResourceListResult, type ResourceMkdirParams, type ResourceMkdirResult, type ResourceMoveParams, type ResourceMoveResult, type ResourceReadResult, type ResourceResolveParams, type ResourceResolveResult, type ResourceWatchState, type ResourceWriteParams, type ResourceWriteResult, type IStateSnapshot } from '../common/state/sessionProtocol.js'; +import { AhpErrorCodes, AHP_SESSION_NOT_FOUND, AuthenticateTargetKind, ContentEncoding, JSON_RPC_INTERNAL_ERROR, JsonRpcErrorCodes, ProtocolError, ResourceChangeType, type CreateResourceWatchParams, type CreateResourceWatchResult, type ResourceCopyParams, type ResourceCopyResult, type ResourceDeleteParams, type ResourceDeleteResult, type ResourceListResult, type ResourceMkdirParams, type ResourceMkdirResult, type ResourceMoveParams, type ResourceMoveResult, type ResourceReadResult, type ResourceResolveParams, type ResourceResolveResult, type ResourceWatchState, type ResourceWriteParams, type ResourceWriteResult, type IStateSnapshot } from '../common/state/sessionProtocol.js'; import { ChangesSummary, ChatInteractivity, ChatOriginKind, MessageAttachmentKind, PendingMessageKind, TerminalClaimKind, TerminalLifecycleStatus, type Annotation, type AnnotationEntry, type AnnotationOrigin, type AnnotationsState, type ChatOrigin, type ChatState, type Customization, type Message, type MessageAttachment, type MessageResourceAttachment, type TerminalState, type TextRange, type ToolResultTerminalContent } from '../common/state/protocol/state.js'; import type { ChatConfigChangedAction, ChatIsArchivedChangedAction, ChatIsReadChangedAction, ChatPendingMessageSetAction, ChatTurnStartedAction, SessionConfigChangedAction } from '../common/state/protocol/actions.js'; import { isAhpAutomationCatalogChannel, isAhpAutomationRunChannel, ISessionGitState, MessageKind, ResponsePartKind, SESSION_META_GITHUB_KEY, SESSION_META_GIT_KEY, SESSION_META_MULTI_ROOT_KEY, SESSION_META_SOURCE_CONTROL_KEY, AH_META_AUTO_ARCHIVED_AT_DB_KEY, AH_META_CREATED_BY_SESSION_DB_KEY, readSessionCreationReference, readSessionComparisonMetadata, readSessionSpawnDepth, withSessionSpawnDepth, withSessionCreationReference, parseSessionCreationReference, SessionLifecycle, SessionStatus, ToolCallStatus, ToolResultContentType, TurnState, AH_META_HAS_WORKSPACE_TRANSITIONS_DB_KEY, AH_META_WORKSPACE_CONVERSION_QUARANTINED_DB_KEY, AH_META_WORKSPACELESS_DB_KEY, AH_META_EHCLI_ADOPTED_DB_KEY, AH_META_IS_ARCHIVED_DB_KEY, AH_META_IS_DONE_DB_KEY, AH_META_IS_READ_DB_KEY, AH_META_DEFAULT_CHAT_IS_READ_DB_KEY, buildResourceWatchChannelUri, buildSubagentSessionUriPrefix, getErrorResponsePart, getSessionChatResource, isAhpChatChannel, isChatInSessionReadAggregate, isChatReadOnly, isSessionChatArchived, isSessionStatusArchived, isSessionStatusRead, isSubagentSession, needsSessionGitStateRefresh, parseResourceWatchChannelUri, parseSessionGitData, parseSessionMultiRootMetadata, parseSubagentSessionUri, readSessionExternal, readSessionGitHubState, readSessionGitState, readSessionMultiRootMetadata, readSessionSourceControlState, readSessionWorkspaceless, withSessionExternal, withSessionGitData, withSessionGitHubState, withSessionGitState, withSessionHasWorkspaceTransitions, withSessionMultiRootMetadata, withSessionSourceControlState, withSessionStatusFlag, withSessionWorkspaceless, withSessionEhcliAdopted, withSessionEhcliLastMigratedTurn, AH_META_EHCLI_LAST_TURN_DB_KEY, withSessionFolderPickerDecision, readSessionFolderPickerDecision, parseSessionFolderPickerDecision, SESSION_META_FOLDER_PICKER_KEY, getAllSessionRelatedPullRequestUrls, readSessionEhcliAdoptable, readSessionEhcliAdopted, readSessionGitHubData, parseSessionGitHubData, parseSessionGitHubState, readSessionGitHubStateInput, withMigratedSessionGitHubState, withReplacedFolderGitHubState, withMostRecentRelatedSessionPullRequest, SESSION_META_GITHUB_DATA_KEY, withWorkingDirectoryKey, withWorkingDirectoryScopeId, type ISessionGitHubState, type ISessionSourceControlState, type ChatSummary, type ISessionWithDefaultChat, type ProjectInfo, type SessionConfigState, type SessionSummary, type SessionSummaryMeta, type ToolResultSubagentContent, type Turn } from '../common/state/sessionState.js'; @@ -121,6 +121,7 @@ import { AgentHostSessionsV2CandidateResolution, AgentHostSessionsV2MigrationSer import { detachedWorktreeRecordUri, getPullRequestSessionBranchName, IAgentHostWorktreeIsolation, type IWorktreeHandOver, WORKTREE_META_REPOSITORY_ROOT, worktreeProjectFromRepositoryRoot } from './shared/worktreeIsolation.js'; import { IAgentHostProviderService } from './agentHostProviderService.js'; +import { findMcpServerName } from './shared/mcpCustomizationController.js'; import type { IAgentHostSessionLifecycleCandidate } from './agentHostSessionLifecycle.js'; import { IAgentHostCheckpointService } from '../common/agentHostCheckpointService.js'; import { IAgentHostReviewService } from '../common/agentHostReviewService.js'; @@ -1801,8 +1802,17 @@ export class AgentService extends Disposable implements IAgentService { // ---- auth --------------------------------------------------------------- async authenticate(params: AuthenticateParams, scope?: IAgentHostAuthenticationScope): Promise { + if (params.target) { + if (params.target.kind !== AuthenticateTargetKind.McpServer) { + throw new ProtocolError(JsonRpcErrorCodes.InvalidParams, 'Unsupported authentication target kind'); + } + const state = isAhpChatChannel(params.target.session) ? undefined : this._stateManager.getSessionState(params.target.session); + if (!state || findMcpServerName(state.customizations ?? [], params.target.customizationId) === undefined) { + throw new ProtocolError(JsonRpcErrorCodes.InvalidParams, 'Authentication target must identify an MCP server in an existing session'); + } + } const result = await this._providerService.authenticate(params, scope); - if (result.authenticated && !scope) { + if (result.authenticated && !scope && !params.target) { this._agentMergeController.refresh(); } return result; diff --git a/src/vs/platform/agentHost/node/codex/codexAgent.ts b/src/vs/platform/agentHost/node/codex/codexAgent.ts index 4476450b928da6..3f69027e17f3e8 100644 --- a/src/vs/platform/agentHost/node/codex/codexAgent.ts +++ b/src/vs/platform/agentHost/node/codex/codexAgent.ts @@ -12,6 +12,7 @@ import { DeferredPromise, disposableTimeout, Limiter, raceCancellationError, rac import { fetchResourceMetadata } from '../../../../base/common/oauth.js'; import { Emitter, Event } from '../../../../base/common/event.js'; import { Disposable, DisposableMap, DisposableStore, IDisposable, MutableDisposable, toDisposable } from '../../../../base/common/lifecycle.js'; +import { ResourceMap } from '../../../../base/common/map.js'; import { Schemas } from '../../../../base/common/network.js'; import { equals } from '../../../../base/common/objects.js'; import { type IObservable, observableValue } from '../../../../base/common/observable.js'; @@ -53,7 +54,7 @@ import { CodexChatDiscovery } from './codexChatDiscovery.js'; import { CodexChatHistory } from './codexChatHistory.js'; import type { InitializeResponse } from './protocol/generated/InitializeResponse.js'; import { applyMcpServerRuntimeStates, McpCustomizationController } from '../shared/mcpCustomizationController.js'; -import { buildCodexMcpReadResult, CodexMcpInventory, codexMcpListToInventory, codexMcpServersFromConfig, codexMcpToolsChanged, codexStartupErrorNeedsAuth, injectCodexMcpAuthTokens, inventoryToSdkServers, normalizeCodexMcpResourceUrl, toCodexMcpServerJson, translateCodexMcpStartupState, type ICodexMcpServerConfigJson } from './codexMcpServers.js'; +import { buildCodexMcpReadResult, CodexMcpInventory, codexMcpListToInventory, codexMcpServersFromConfig, codexMcpToolsChanged, codexStartupErrorNeedsAuth, injectCodexMcpAuthTokens, inventoryToSdkServers, normalizeCodexMcpResourceUrl, toCodexMcpServerJson, translateCodexMcpStartupState, type ICodexMcpAuthToken, type ICodexMcpServerConfigJson } from './codexMcpServers.js'; import { codexHooksToContainers, codexSelectedCapabilityRootCandidates, codexSkillsToContainers, discoverCodexWorkspaceAgents, discoverCodexWorkspaceInstructions, discoverCodexWorkspaceSkills, excludeCodexWorkspaceSkillDuplicates } from './codexCustomizations.js'; import { CodexClientCustomizationStore, codexAgentRoleToml, codexClientSkillInstructions, codexCustomizationConfig, codexMcpServersFromDefinitions, codexMcpServersFromPlugins, codexPluginMcpServerSources, parsedPluginChildren, type ICodexClientPlugin } from './codexClientCustomizations.js'; import { IAgentHostCustomizationEnablementService, targetForUnownedMcpServer } from '../agentHostCustomizationEnablementService.js'; @@ -1215,6 +1216,7 @@ export class CodexAgent extends Disposable implements IAgent { * shared app-server). */ private readonly _mcpAuthTokens = new Map(); + private readonly _mcpAuthTokensBySession = new ResourceMap>(); /** * Association from a normalized OAuth `resource` (what the workbench * authenticates) to the normalized MCP server URL(s) it unlocks. RFC 9728 @@ -1674,6 +1676,29 @@ export class CodexAgent extends Disposable implements IAgent { if (normalizedResource === undefined) { return false; } + if (params.target) { + const targetSession = URI.parse(params.target.session); + const session = this._sessionForMcpControl(targetSession); + if (!session || session.disposed || !isEqual(session.owner.get(), targetSession)) { + return false; + } + const serverName = this._resolveMcpServerName(session, params.target.customizationId); + const serverUrl = serverName === undefined ? undefined : this._httpMcpServerUrls(session).get(serverName); + if (serverName === undefined || serverUrl === undefined || (serverUrl !== normalizedResource + && !this._mcpAuthServerUrlsByResource.get(normalizedResource)?.has(serverUrl))) { + return false; + } + let tokens = this._mcpAuthTokensBySession.get(targetSession); + if (!tokens) { + tokens = new Map(); + this._mcpAuthTokensBySession.set(targetSession, tokens); + } + if (tokens.get(serverName)?.token !== params.token || tokens.get(serverName)?.url !== serverUrl) { + tokens.set(serverName, { url: serverUrl, token: params.token }); + await this._reconnectSessionsForMcpAuth(new Set([serverUrl]), targetSession); + } + return true; + } // The workbench authenticates the OAuth `resource`, which RFC 9728 // discovery may report as different from the configured server URL. // Resolve the server URL(s) this resource unlocks: the association @@ -1694,6 +1719,16 @@ export class CodexAgent extends Disposable implements IAgent { changed = true; } } + for (const session of this._sessions.values()) { + const tokens = this._mcpAuthTokensBySession.get(session.owner.get()); + if (tokens) { + for (const [name, url] of this._httpMcpServerUrls(session)) { + if (serverUrls.has(url)) { + changed = tokens.delete(name) || changed; + } + } + } + } if (!changed) { return true; } @@ -1721,9 +1756,9 @@ export class CodexAgent extends Disposable implements IAgent { * (`thread/resume` carries the same `config` field, loading history from the * rollout) on its next turn via {@link ICodexSession.needsResume}. */ - private async _reconnectSessionsForMcpAuth(normalizedUrls: ReadonlySet): Promise { + private async _reconnectSessionsForMcpAuth(normalizedUrls: ReadonlySet, targetSession?: URI): Promise { for (const session of this._sessions.values()) { - if (session.disposed || session.threadId === undefined) { + if (session.disposed || session.threadId === undefined || (targetSession && !isEqual(session.owner.get(), targetSession))) { continue; } if (![...this._httpMcpServerUrls(session).values()].some(url => normalizedUrls.has(url))) { @@ -2830,7 +2865,7 @@ export class CodexAgent extends Disposable implements IAgent { ? Object.fromEntries(Object.entries(configuredServers).filter(([name, server]) => !isAgentMergeRestrictedMcpServer(name, server))) : configuredServers; const builtInGitHub = this._builtInGitHubMcpServer(session, enabledConfiguredServers); - return injectCodexMcpAuthTokens({ ...builtInGitHub, ...enabledConfiguredServers }, this._mcpAuthTokens); + return injectCodexMcpAuthTokens({ ...builtInGitHub, ...enabledConfiguredServers }, this._mcpAuthTokens, this._mcpAuthTokensBySession.get(this._launchOwner(session))); } private _builtInGitHubMcpServer(session: ICodexSession, configuredServers: Record): Record { @@ -8757,7 +8792,11 @@ export class CodexAgent extends Disposable implements IAgent { // insufficient scopes). Drop it so the user is re-prompted // instead of getting stuck on a terminal error with no way to // re-authenticate. - if (this._mcpAuthTokens.delete(normalized)) { + const session = this._sessionForMcpThread(threadId); + const targetedTokens = session ? this._mcpAuthTokensBySession.get(this._launchOwner(session)) : undefined; + if (targetedTokens?.has(name)) { + targetedTokens.set(name, { url: normalized, token: '' }); + } else if (this._mcpAuthTokens.delete(normalized)) { this._logService.info(`[Codex] MCP server '${name}' rejected the stored token; clearing it to allow re-authentication`); } void this._surfaceMcpAuthRequired(client, threadId, name, url, error); @@ -8928,6 +8967,7 @@ export class CodexAgent extends Disposable implements IAgent { this._mcpPublisherSessionIdByConfiguration.clear(); this._publishedMcpTopLevelIdsByConfiguration.clear(); this._mcpAuthTokens.clear(); + this._mcpAuthTokensBySession.clear(); this._mcpAuthServerUrlsByResource.clear(); } diff --git a/src/vs/platform/agentHost/node/codex/codexMcpServers.ts b/src/vs/platform/agentHost/node/codex/codexMcpServers.ts index c3507d7b3cb546..66971596272f11 100644 --- a/src/vs/platform/agentHost/node/codex/codexMcpServers.ts +++ b/src/vs/platform/agentHost/node/codex/codexMcpServers.ts @@ -392,29 +392,34 @@ export function codexStartupErrorNeedsAuth(error: string | null | undefined): bo return /not logged in|mcp login|log in to|unauthori[sz]ed|requires? (?:authentication|authorization|login)|\b401\b/i.test(error); } +export interface ICodexMcpAuthToken { + readonly url: string; + readonly token: string; +} + /** - * Returns a copy of `servers` with `Authorization: Bearer ` injected - * into the `http_headers` of every http server whose (normalized) URL has a - * token in `tokensByNormalizedUrl`. stdio servers and servers without a token - * are passed through unchanged. Any existing authorization header is removed - * first -- case-insensitively, since HTTP header names are case-insensitive and - * a configured lowercase `authorization` would otherwise coexist with the - * injected value and leave a stale credential in the payload. + * Injects bearer credentials, with per-server tokens overriding shared URL tokens and empty overrides revoking only that server. + * Replaces existing authorization headers case-insensitively while leaving unrelated headers and stdio servers intact. */ export function injectCodexMcpAuthTokens( servers: Record, tokensByNormalizedUrl: ReadonlyMap, + tokensByServerName?: ReadonlyMap, ): Record { - if (tokensByNormalizedUrl.size === 0) { + if (tokensByNormalizedUrl.size === 0 && !tokensByServerName?.size) { return servers; } const out: Record = {}; for (const [name, server] of Object.entries(servers)) { const normalized = server.url !== undefined ? normalizeCodexMcpResourceUrl(server.url) : undefined; - const token = normalized !== undefined ? tokensByNormalizedUrl.get(normalized) : undefined; - out[name] = token !== undefined - ? { ...server, http_headers: { ...withoutAuthorizationHeaders(server.http_headers), Authorization: `Bearer ${token}` } } - : server; + const targeted = tokensByServerName?.get(name); + const targetedToken = normalized !== undefined && targeted?.url === normalized ? targeted.token : undefined; + const token = targetedToken ?? (normalized !== undefined ? tokensByNormalizedUrl.get(normalized) : undefined); + out[name] = targetedToken === '' + ? { ...server, http_headers: withoutAuthorizationHeaders(server.http_headers) } + : token !== undefined + ? { ...server, http_headers: { ...withoutAuthorizationHeaders(server.http_headers), Authorization: `Bearer ${token}` } } + : server; } return out; } diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts index a0328caddeecde..351d525df930a5 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts @@ -3108,6 +3108,9 @@ export class CopilotAgent extends Disposable implements IAgent { let handled = false; const sessions = new Set([...this._sessionsPendingRegistration.values(), ...this._allLiveSessions()]); for (const session of sessions) { + if (params.target && !isEqual(session.ownerSessionUri, URI.parse(params.target.session))) { + continue; + } const didHandle = await session.resolveMcpAuthentication(params); handled ||= didHandle; } diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts index 5a8a77346a3f7f..7226f0e9971732 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts @@ -1847,6 +1847,11 @@ export class CopilotAgentSession extends Disposable { /** Invalidates inventory snapshots when lifecycle events or authentication challenges arrive. */ private _mcpLifecycleVersion = 0; private _mcpInventoryRequestVersion = 0; + /** + * Servers reconnecting after an accepted OAuth token. The runtime reports their new connection's tools without a + * `connected` status change, so their next tools-loaded event is treated as connected. + */ + private readonly _mcpAwaitingOAuthReconnect = new Set(); /** * Configuration knowledge that MCP provenance is derived from. Only inventory and configuration reads update * these; lifecycle events never do, so a status change can never alter where a server appears to come from. @@ -4031,6 +4036,7 @@ export class CopilotAgentSession extends Disposable { } this._wrapperDisposables.clear(); this._wrapper = this._wrapperDisposables.add(wrapper); + this._mcpAwaitingOAuthReconnect.clear(); this._lastAppliedAgent = this._launchPlan.kind === 'create' ? this._launchPlan.resolvedAgentName ?? null : undefined; this._registeredByokConfig = wrapper.launchByokConfig; this._reportLocalModelSelection(this._launchPlan.kind === 'create' ? this._launchPlan.model?.id : this._launchPlan.fallback.model?.id); @@ -4618,12 +4624,19 @@ export class CopilotAgentSession extends Disposable { } /** - * Resolves only matching, currently pending SDK authentication callbacks. Connector-owned servers are never - * resolved with a forwarded MCP token: the runtime ignores it and keeps using the session credential. + * Resolves matching pending SDK callbacks, restricted to the targeted MCP server when provided. + * Connector-owned servers keep using the session credential, never a forwarded MCP token. */ async resolveMcpAuthentication(params: AuthenticateParams): Promise { + const serverName = params.target ? this._mcpCustomizations.serverNameForCustomizationId(params.target.customizationId) : undefined; + if (params.target && (serverName === undefined || !isEqual(this.ownerSessionUri, URI.parse(params.target.session)))) { + return false; + } let resolved = false; for (const [requestId, pending] of this._pendingMcpAuthRequests.entries()) { + if (params.target && pending.serverName !== serverName) { + continue; + } if (pending.resource.resource !== params.resource || !this._scopesSatisfy(params.scopes, pending.requiredScopes)) { continue; } @@ -4712,6 +4725,7 @@ export class CopilotAgentSession extends Disposable { description: request.wwwAuthenticateParams?.error, }; this._mcpLifecycleVersion++; + this._mcpAwaitingOAuthReconnect.delete(request.serverName); const toolCalls = this._activeMcpToolCalls(request.serverName); const copilotConnector = this._isCopilotConnectorServer(request.serverName, requiredScopes); const pending: IPendingMcpAuthRequest = { @@ -9720,6 +9734,9 @@ export class CopilotAgentSession extends Disposable { })); store.add(wrapper.onMcpServerStatusChanged(e => { this._mcpLifecycleVersion++; + if (e.data.status !== 'pending') { + this._mcpAwaitingOAuthReconnect.delete(e.data.serverName); + } this._observeMcpRestart(e.data.serverName, e.data.status); this._noteMcpSdkStatus(e.data.serverName, e.data.status); const requirement = this._lastMcpAuthRequirements.get(e.data.serverName); @@ -9865,7 +9882,12 @@ export class CopilotAgentSession extends Disposable { /** Promotes a connected server from Starting to Ready once its current connection has loaded tools. */ private _markMcpToolsLoaded(serverName: string): void { - if (this._store.isDisposed || this._mcpToolsLoaded.has(serverName)) { + if (this._store.isDisposed) { + return; + } + if (this._mcpAwaitingOAuthReconnect.delete(serverName)) { + this._noteMcpSdkStatus(serverName, 'connected'); + } else if (this._mcpToolsLoaded.has(serverName)) { return; } this._mcpToolsLoaded.add(serverName); @@ -9925,6 +9947,9 @@ export class CopilotAgentSession extends Disposable { } this._mcpLifecycleVersion++; this._disableMcpTokenCompletion(serverName, requestId); + // The reconnect is a new connection; tools from an earlier connection or a cached snapshot no longer count. + this._mcpToolsLoaded.delete(serverName); + this._mcpAwaitingOAuthReconnect.add(serverName); this._mcpCustomizations.applyOne({ name: serverName, state: { kind: McpServerStatus.Starting }, diff --git a/src/vs/platform/agentHost/node/missionControl/missionControlAuthentication.ts b/src/vs/platform/agentHost/node/missionControl/missionControlAuthentication.ts index 0850b159cb411f..fada0e388bea3f 100644 --- a/src/vs/platform/agentHost/node/missionControl/missionControlAuthentication.ts +++ b/src/vs/platform/agentHost/node/missionControl/missionControlAuthentication.ts @@ -343,7 +343,7 @@ export class MissionControlAuthentication extends Disposable { } const generation = this._generation; const resource = new URL(params.resource); - const identity = resource.origin === new URL(this._apiOrigin).origin; + const identity = !params.target && resource.origin === new URL(this._apiOrigin).origin; if (!identity && !this.authenticated) { const validation = this._identityValidation; if (!validation) { diff --git a/src/vs/platform/agentHost/test/node/agentHostAuthenticationScope.test.ts b/src/vs/platform/agentHost/test/node/agentHostAuthenticationScope.test.ts index e446685179c050..3c027c9c91683b 100644 --- a/src/vs/platform/agentHost/test/node/agentHostAuthenticationScope.test.ts +++ b/src/vs/platform/agentHost/test/node/agentHostAuthenticationScope.test.ts @@ -11,7 +11,7 @@ import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/c import { NullLogService } from '../../../log/common/log.js'; import type { AuthenticateParams, IAgent } from '../../common/agent.js'; import { captureCredentialSource, parseCredentialPrincipal, type IAgentHostCredentialSource } from '../../common/agentHostCredentials.js'; -import { AhpErrorCodes, JsonRpcErrorCodes, ProtocolError } from '../../common/state/sessionProtocol.js'; +import { AhpErrorCodes, AuthenticateTargetKind, JsonRpcErrorCodes, ProtocolError } from '../../common/state/sessionProtocol.js'; import { AgentHostAuthenticationScope } from '../../node/agentHostAuthenticationScope.js'; import { AgentHostAuthenticationService } from '../../node/agentHostAuthenticationService.js'; @@ -336,6 +336,74 @@ suite('AgentHostAuthenticationScope', () => { assert.deepStrictEqual(requests, []); }); + test('targeted MCP tokens stay out of agent credentials and retain independent concurrent replay slots', async () => { + const service = store.add(new AgentHostAuthenticationService(new NullLogService(), async () => { assert.fail('Targeted MCP tokens must not validate an agent identity'); })); + const target = (customizationId: string) => ({ kind: AuthenticateTargetKind.McpServer, session: 'ahp-session:/target', customizationId }); + const first: AuthenticateParams = { resource, scopes: ['read', 'write'], token: 'first-token', target: target('first') }; + const second: AuthenticateParams = { ...first, token: 'second-token', target: target('second') }; + const started = new DeferredPromise(); + const gate = new DeferredPromise(); + const completions: AuthenticateParams[] = []; + const providerTokens: string[] = []; + const changedTokens: (string | undefined)[] = []; + const provider = new class extends mock() { + override getProtectedResources() { return [{ resource }]; } + override async authenticate(_resource: string, token: string) { providerTokens.push(token); return true; } + override async handleAuthenticationToken(params: AuthenticateParams) { + completions.push(params); + if (params === first) { + started.complete(); + await gate.p; + } + return true; + } + }(); + store.add(service.onDidChangeAuthToken(event => changedTokens.push(event.token))); + const firstResult = service.authenticate(first, [provider]); + await started.p; + const secondResult = await service.authenticate(second, [provider]); + gate.complete(); + const result = await firstResult; + completions.length = 0; + await service.replay(provider); + const replayed = completions.map(({ token, target }) => ({ token, target })); + const lookups = { + exact: service.getAuthToken({ resource, scopes: ['read', 'write'] }), + subset: service.getAuthToken({ resource, scopes: ['read'] }), + unscoped: service.getAuthToken({ resource }), + account: service.getAuthAccount({ resource, scopes: ['read'] }), + }; + await assert.rejects(service.createLocalCredentialSource(principal, [{ resource }]), { code: AhpErrorCodes.AuthRequired }); + await service.authenticate({ ...first, token: '' }, [provider]); + completions.length = 0; + await service.replay(provider); + assert.deepStrictEqual({ result, secondResult, replayed, lookups, providerTokens, changedTokens, afterRevocation: completions.map(({ token, target }) => ({ token, target })) }, { + result: { authenticated: true }, secondResult: { authenticated: true }, + replayed: [{ token: 'second-token', target: second.target }, { token: 'first-token', target: first.target }], + lookups: { exact: undefined, subset: undefined, unscoped: undefined, account: undefined }, + providerTokens: [], changedTokens: [], afterRevocation: [{ token: 'second-token', target: second.target }], + }); + }); + + test('scoped targeted MCP authentication bypasses protected-resource capture', async () => { + const service = store.add(new AgentHostAuthenticationService(new NullLogService())); + const scope = createScope(); + const completions: AuthenticateParams[] = []; + const providerTokens: string[] = []; + const provider = new class extends mock() { + override getProtectedResources() { return [{ resource }]; } + override async authenticate(_resource: string, token: string) { providerTokens.push(token); return true; } + override async handleAuthenticationToken(params: AuthenticateParams) { completions.push(params); return true; } + }(); + const params: AuthenticateParams = { + resource, token: 'mcp-token', target: { kind: AuthenticateTargetKind.McpServer, session: 'ahp-session:/target', customizationId: 'mcp' }, + }; + const result = await service.authenticate(params, [provider], scope); + assert.deepStrictEqual({ result, completions, providerTokens, scoped: service.getAuthToken({ resource }, scope), legacy: service.getAuthToken({ resource }) }, { + result: { authenticated: true }, completions: [params], providerTokens: [], scoped: undefined, legacy: undefined, + }); + }); + test('MCP completion does not capture or replay an unrelated resource', async () => { const service = store.add(new AgentHostAuthenticationService(new NullLogService())); const completions: string[] = []; diff --git a/src/vs/platform/agentHost/test/node/agentHostAutomationService.test.ts b/src/vs/platform/agentHost/test/node/agentHostAutomationService.test.ts index b12a5293b75796..fd31ab1282fc46 100644 --- a/src/vs/platform/agentHost/test/node/agentHostAutomationService.test.ts +++ b/src/vs/platform/agentHost/test/node/agentHostAutomationService.test.ts @@ -25,6 +25,7 @@ import { createUnknownAgentHostClientTelemetryContext } from '../../common/agent import { AGENT_HOST_AUTOMATIONS_ENABLED_CONFIG_KEY } from '../../common/automationConfig.js'; import { SessionConfigKey } from '../../common/sessionConfigKeys.js'; import { ActionType } from '../../common/state/sessionActions.js'; +import { AuthenticateTargetKind } from '../../common/state/sessionProtocol.js'; import { AutomationMisfirePolicy, AutomationOperation, AutomationTriggerKind, type AutomationDefinition } from '../../common/state/protocol/channels-automation/state.js'; import { AutomationRunOriginKind, AutomationRunStatus, type AutomationRunState } from '../../common/state/protocol/channels-automation-run/state.js'; import type { RunAutomationParams } from '../../common/state/protocol/channels-automation/commands.js'; @@ -228,12 +229,12 @@ suite('AgentHostAutomationService', () => { publish(session, auth); service.dispose(); publish(session, changedAuth); + const requestsFor = (auth: McpAuthRequirement) => ['top', 'child'].map(serverName => ({ + serverName, auth, target: { kind: AuthenticateTargetKind.McpServer, session: session.toString(), customizationId: serverName }, + })); assert.deepStrictEqual({ inflightRequests, failedRequests, requests: authenticationRequests }, { - inflightRequests: [{ serverName: 'top', auth }, { serverName: 'child', auth }], - failedRequests: [{ serverName: 'top', auth }, { serverName: 'child', auth }], - requests: [auth, changedAuth, ...changedChallenges, auth, auth].flatMap(auth => [ - { serverName: 'top', auth }, { serverName: 'child', auth }, - ]), + inflightRequests: requestsFor(auth), failedRequests: requestsFor(auth), + requests: [auth, changedAuth, ...changedChallenges, auth, auth].flatMap(requestsFor), }); }); diff --git a/src/vs/platform/agentHost/test/node/agentHostClientConnectionService.test.ts b/src/vs/platform/agentHost/test/node/agentHostClientConnectionService.test.ts index 4e35f4583b4156..175cf192ef53e9 100644 --- a/src/vs/platform/agentHost/test/node/agentHostClientConnectionService.test.ts +++ b/src/vs/platform/agentHost/test/node/agentHostClientConnectionService.test.ts @@ -6,6 +6,7 @@ import assert from 'assert'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import type { IAgentHostMcpAuthenticationRequest } from '../../common/agentHostExtensionProtocol.js'; +import { AuthenticateTargetKind } from '../../common/state/sessionProtocol.js'; import { McpAuthRequiredReason } from '../../common/state/protocol/channels-session/state.js'; import { AgentHostClientConnectionService } from '../../node/agentHostClientConnectionService.js'; @@ -40,6 +41,7 @@ suite('AgentHostClientConnectionService', () => { const service = store.add(new AgentHostClientConnectionService()); const request: IAgentHostMcpAuthenticationRequest = { serverName: 'example', + target: { kind: AuthenticateTargetKind.McpServer, session: 'ahp-session:/example', customizationId: 'example-mcp' }, auth: { reason: McpAuthRequiredReason.Required, resource: { resource: 'https://mcp.example.com', authorization_servers: [] } }, }; const calls: { source: number; request: IAgentHostMcpAuthenticationRequest }[] = []; diff --git a/src/vs/platform/agentHost/test/node/agentService.test.ts b/src/vs/platform/agentHost/test/node/agentService.test.ts index 7786d6f7b5e5b8..da5265527de81b 100644 --- a/src/vs/platform/agentHost/test/node/agentService.test.ts +++ b/src/vs/platform/agentHost/test/node/agentService.test.ts @@ -71,7 +71,7 @@ import { AgentMergeConfigKey, readAgentMergeSessionState } from '../../common/ag import { SessionDatabase, sessionDatabaseMigrations } from '../../node/sessionDatabase.js'; import { ActionType, ActionEnvelope, NotificationType, type ChatAction, type ChatConfigChangedAction, type INotification, type SessionAction, type SessionSummaryChanges } from '../../common/state/sessionActions.js'; import { AH_META_AUTO_ARCHIVED_AT_DB_KEY, AH_META_CREATED_BY_SESSION_DB_KEY, AH_META_DEFAULT_CHAT_IS_READ_DB_KEY, AH_META_IS_READ_DB_KEY, AH_META_EHCLI_ADOPTED_DB_KEY, readSessionEhcliAdopted, AH_META_IS_ARCHIVED_DB_KEY, AH_META_WORKSPACE_CONVERSION_QUARANTINED_DB_KEY, AH_META_WORKSPACELESS_DB_KEY, ChangesetStatus, CustomizationType, MessageAttachmentKind, MessageKind, SessionActiveClient, ResponsePartKind, ROOT_STATE_URI, SESSION_META_EHCLI_ADOPTABLE_KEY, SESSION_META_FOLDER_PICKER_KEY, SESSION_META_MULTI_ROOT_KEY, SessionLifecycle, SessionSourceControlOutcome, SessionStatus, ToolCallCancellationReason, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType, TurnState, buildChatUri, buildDefaultChatUri, buildSubagentChatUri, buildSubagentSessionUri, createErrorResponsePart, customizationId, isMessageRequestHiddenFromTranscript, isSessionStatusArchived, isSessionStatusRead, isSubagentSession, parseChatUri, parseSubagentSessionUri, readSessionCreationReference, readSessionEhcliAdoptable, readSessionExternal, readSessionGitHubData, readSessionGitHubState, readSessionGitState, readWorkingDirectoryKeys, readWorkingDirectoryScopeIds, SESSION_META_GITHUB_DATA_KEY, readSessionMultiRootMetadata, readSessionFolderPickerDecision, readSessionSourceControlState, readSessionWorkspaceless, withSessionEhcliAdoptable, withSessionExternal, withSessionGitHubState, withSessionGitState, withSessionMultiRootMetadata, withSessionWorkspaceless, withWorkingDirectoryKey, withWorkingDirectoryScopeId, ChatOriginKind, type ChangesetState, type ISessionFolderPickerDecision, type ISessionWithDefaultChat, type MarkdownResponsePart, type SessionState, type SessionSummary, type SessionSummaryMeta, type ChatState, type ToolCallCompletedState, type ToolCallResponsePart, type Turn } from '../../common/state/sessionState.js'; -import { BackgroundWorkKind, ChatInteractivity, PendingMessageKind, type BackgroundWork, type Message, type MessageAttachment, TerminalClaimKind, type TerminalClaim } from '../../common/state/protocol/state.js'; +import { BackgroundWorkKind, ChatInteractivity, McpServerStatus, PendingMessageKind, type BackgroundWork, type Message, type MessageAttachment, TerminalClaimKind, type TerminalClaim } from '../../common/state/protocol/state.js'; import { isHostSnapshotAttachment, toHostSnapshotAttachmentMeta } from '../../common/meta/agentSnapshotAttachmentMeta.js'; import { readAgentMessageDelegationMeta } from '../../common/meta/agentMessageDelegationMeta.js'; import { readRemoteSessionDepth, readRemoteSessionOrigin, REMOTE_SESSION_ORIGIN_METADATA_KEY, withRemoteSessionOrigin } from '../../common/meta/agentRemoteSessionMeta.js'; @@ -114,7 +114,7 @@ import { OLDER_BUILD_WORKTREE_WITHHELD_KEY } from '../../node/agentHostSessionSt import { parseSessionWorktreeFolders, WORKTREE_FOLDERS_KEY } from '../../node/shared/sessionWorktreeFolders.js'; import type { IAgentHostPullRequestResolver } from '../../node/shared/pullRequestResolver.js'; import { ADDITIONAL_WORKTREES_METADATA_KEY, readSessionAdditionalWorktrees, writeSessionAdditionalWorktrees } from '../../node/shared/sessionAdditionalWorktrees.js'; -import { AhpErrorCodes, AHP_SESSION_NOT_FOUND, ContentEncoding, JSON_RPC_INTERNAL_ERROR, ProtocolError, type IStateSnapshot } from '../../common/state/sessionProtocol.js'; +import { AhpErrorCodes, AHP_SESSION_NOT_FOUND, AuthenticateTargetKind, ContentEncoding, JSON_RPC_INTERNAL_ERROR, JsonRpcErrorCodes, ProtocolError, type IStateSnapshot } from '../../common/state/sessionProtocol.js'; import type { IAgentServerToolHost } from '../../common/agentServerTools.js'; import { ArtifactServerToolName, SessionServerToolName } from '../../common/serverToolNames.js'; import { buildMcpChannel } from '../../node/shared/mcpCustomizationController.js'; @@ -23610,6 +23610,50 @@ suite('AgentService (node dispatcher)', () => { suite('authenticate', () => { + test('validates MCP authentication targets against top-level and child customizations', async () => { + registerTestAgentProvider(service, copilotAgent); + const session = await service.createSession({ legacyChatUris: false, provider: 'copilot' }); + getStateManager(service).dispatchServerAction(session.toString(), { + type: ActionType.SessionCustomizationsChanged, + customizations: [ + { type: CustomizationType.McpServer, id: 'top', uri: 'mcp:/top', name: 'top', state: { kind: McpServerStatus.Starting } }, + { + type: CustomizationType.Plugin, id: 'plugin', uri: 'file:///plugin', name: 'plugin', children: [ + { type: CustomizationType.McpServer, id: 'child', uri: 'mcp:/child', name: 'child', state: { kind: McpServerStatus.Starting } }, + ] + }, + ], + }); + const targets = ['top', 'child'].map(customizationId => ({ kind: AuthenticateTargetKind.McpServer, session: session.toString(), customizationId })); + const received: typeof targets = []; + const provider: IAgent = copilotAgent; + provider.handleAuthenticationToken = async params => { + if (params.target) { + received.push(params.target); + } + return true; + }; + const results = []; + for (const target of targets) { + results.push(await service.authenticate({ resource: 'https://api.github.com', token: 'mcp-token', target })); + } + const scope = disposables.add(new AgentHostAuthenticationScope(async params => params, () => { })); + for (const target of [ + { ...targets[0], session: 'ahp-session:/missing' }, + { ...targets[0], session: buildDefaultChatUri(session.toString()) }, + { ...targets[0], customizationId: 'missing' }, + { ...targets[0], customizationId: 'plugin' }, + { ...targets[0], kind: 'future' as AuthenticateTargetKind }, + ]) { + for (const authenticationScope of [undefined, scope]) { + await assert.rejects(service.authenticate({ resource: 'https://api.github.com', token: 'mcp-token', target }, authenticationScope), { code: JsonRpcErrorCodes.InvalidParams }); + } + } + assert.deepStrictEqual({ results, received, agentAuthentication: copilotAgent.authenticateCalls }, { + results: [{ authenticated: true }, { authenticated: true }], received: targets, agentAuthentication: [], + }); + }); + test('forwards optional authentication metadata to matching providers', async () => { const metadata = { 'vscode.authentication.entitlement': { trackingId: 'analytics', sku: 'free' } }; let received: Record | undefined; diff --git a/src/vs/platform/agentHost/test/node/codex/codexAgent.test.ts b/src/vs/platform/agentHost/test/node/codex/codexAgent.test.ts index d3c68131d740fa..755eb72c8b5f67 100644 --- a/src/vs/platform/agentHost/test/node/codex/codexAgent.test.ts +++ b/src/vs/platform/agentHost/test/node/codex/codexAgent.test.ts @@ -5,15 +5,18 @@ import assert from 'assert'; import { DeferredPromise } from '../../../../../base/common/async.js'; +import { ResourceMap } from '../../../../../base/common/map.js'; +import { isEqual } from '../../../../../base/common/resources.js'; import { DisposableStore } from '../../../../../base/common/lifecycle.js'; import { URI } from '../../../../../base/common/uri.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js'; import { NullLogService } from '../../../../../platform/log/common/log.js'; -import { AgentChatMigrationDeferred, AgentSession, CODEX_AGENT_PROVIDER_ID, type AgentProvider, type IAgentChatContext, type IAgentDiscoveredChat, type IAgentChatOwnerRef, fixedChatOwnerRef } from '../../../common/agent.js'; +import { AgentChatMigrationDeferred, AgentSession, CODEX_AGENT_PROVIDER_ID, type AgentProvider, type AuthenticateParams, type IAgentChatContext, type IAgentDiscoveredChat, type IAgentChatOwnerRef, fixedChatOwnerRef } from '../../../common/agent.js'; import { AgentSystemNotificationKind, toAgentSystemNotificationMeta } from '../../../common/meta/agentSystemNotificationMeta.js'; import { ActionType, type ChatAction } from '../../../common/state/sessionActions.js'; import { CustomizationEnablementKind, CustomizationType, McpServerStatus, type McpServerCustomization } from '../../../common/state/protocol/channels-session/state.js'; import { buildDefaultChatUri, ResponsePartKind, MessageAttachmentKind, MessageKind, type MessageAttachment, type PendingMessage } from '../../../common/state/sessionState.js'; +import { AuthenticateTargetKind } from '../../../common/state/sessionProtocol.js'; import { AgentHostStateManager } from '../../../node/agentHostStateManager.js'; import { getCustomizationEnablementKey, type CustomizationEnablementResolution, type ICustomizationEnablementTarget } from '../../../node/agentHostCustomizationEnablementService.js'; import { CodexAgent } from '../../../node/codex/codexAgent.js'; @@ -78,6 +81,10 @@ interface ICodexGitHubMcpHarness { }): Record; } +interface ICodexMcpAuthenticationHarness extends ICodexGitHubMcpHarness { + handleAuthenticationToken(params: AuthenticateParams): Promise; +} + interface ICodexGitHubEndpointChangeHarness { _handleGitHubEndpointChange(): void; } @@ -385,12 +392,76 @@ suite('CodexAgent', () => { }); }); + test('targeted MCP tokens and revocation affect only one server in one owning session', async () => { + const resource = 'https://mcp.slack.com'; + const serverUrl = 'https://mcp.slack.com/mcp'; + const sessions = ['first', 'second'].map(sessionId => ({ + sessionId, owner: fixedChatOwnerRef(URI.parse(`ahp-session:/${sessionId}`)), + workingDirectory: URI.file('/work'), disposed: false, threadId: `thread-${sessionId}`, firstTurnSent: true, + })); + const reloaded: string[] = []; + const globalTokens = new Map([[serverUrl, 'legacy-token']]); + const harness = Object.assign(Object.create(CodexAgent.prototype), { + _sessions: new Map(sessions.map(session => [session.sessionId, session])), + _configurationService: { + getRootValue: () => ({ + 'slack-a': { type: 'http', url: serverUrl }, 'slack-b': { type: 'http', url: serverUrl }, + }) + }, + _sessionMcpDiscoveries: new Map(), _enabledClientPlugins: () => [], _isMcpServerEnabledForSdk: () => true, + _githubMcpServerEnabled: false, _mcpAuthTokens: globalTokens, _mcpAuthTokensBySession: new ResourceMap(), + _mcpAuthServerUrlsByResource: new Map([[resource + '/', new Set([serverUrl])]]), + _sessionForMcpControl: (target: URI) => sessions.find(session => isEqual(session.owner.get(), target)), + _resolveMcpServerName: (_session: object, id: string) => id === 'id-a' ? 'slack-a' : id === 'id-b' ? 'slack-b' : undefined, + _httpMcpServerUrls: () => new Map([['slack-a', serverUrl], ['slack-b', serverUrl]]), + _markSessionForReload: (session: { sessionId: string }) => { reloaded.push(session.sessionId); }, + _logService: new NullLogService(), + }) as ICodexMcpAuthenticationHarness; + const authenticate = (session: string, customizationId: string, token: string) => harness.handleAuthenticationToken({ + resource, token, target: { kind: AuthenticateTargetKind.McpServer, session, customizationId }, + }); + const headers = () => sessions.map(session => Object.fromEntries(Object.entries(harness._buildSessionMcpServers(session)) + .map(([name, server]) => [name, server.http_headers?.Authorization]))); + const first = await authenticate('ahp-session:/first', 'id-a', 'token-a'); + const afterFirst = headers(); + const sibling = await authenticate('ahp-session:/first', 'id-b', 'token-b'); + const second = await authenticate('ahp-session:/second', 'id-a', 'other-account'); + const afterBothSessions = headers(); + const revoked = await authenticate('ahp-session:/first', 'id-a', ''); + const unknown = await authenticate('ahp-session:/missing', 'id-a', 'unknown'); + const unknownCustomization = await authenticate('ahp-session:/first', 'missing', 'unknown'); + const afterRevocation = headers(); + const legacy = await harness.handleAuthenticationToken({ resource, token: 'legacy-rotated' }); + assert.deepStrictEqual({ first, sibling, second, revoked, unknown, unknownCustomization, legacy, afterFirst, afterBothSessions, afterRevocation, afterLegacy: headers(), reloaded, globalTokens: [...globalTokens] }, { + first: true, sibling: true, second: true, revoked: true, unknown: false, unknownCustomization: false, legacy: true, + afterFirst: [ + { 'slack-a': 'Bearer token-a', 'slack-b': 'Bearer legacy-token' }, + { 'slack-a': 'Bearer legacy-token', 'slack-b': 'Bearer legacy-token' }, + ], + afterBothSessions: [ + { 'slack-a': 'Bearer token-a', 'slack-b': 'Bearer token-b' }, + { 'slack-a': 'Bearer other-account', 'slack-b': 'Bearer legacy-token' }, + ], + afterRevocation: [ + { 'slack-a': undefined, 'slack-b': 'Bearer token-b' }, + { 'slack-a': 'Bearer other-account', 'slack-b': 'Bearer legacy-token' }, + ], + afterLegacy: [ + { 'slack-a': 'Bearer legacy-rotated', 'slack-b': 'Bearer legacy-rotated' }, + { 'slack-a': 'Bearer legacy-rotated', 'slack-b': 'Bearer legacy-rotated' }, + ], + reloaded: ['first', 'first', 'second', 'first', 'first', 'second'], globalTokens: [[serverUrl, 'legacy-rotated']], + }); + }); + test('GitHub MCP injection respects unowned server enablement', () => { const createHarness = (enabled: boolean, customizationEnabled: boolean, token: string | undefined): ICodexGitHubMcpHarness => Object.assign(Object.create(CodexAgent.prototype), { _configurationService: { getRootValue: () => undefined }, _sessionMcpDiscoveries: new Map(), _enabledClientPlugins: () => [], _mcpAuthTokens: new Map(), + _mcpAuthTokensBySession: new ResourceMap(), + _launchOwner: () => URI.parse('ahp-session:/test'), _githubMcpServerEnabled: enabled, _githubToken: token, _gitHubMcpServerConfiguration: token ? createGitHubMcpServerConfiguration('https://api.githubcopilot.com') : undefined, @@ -425,6 +496,8 @@ suite('CodexAgent', () => { _sessionMcpDiscoveries: new Map(), _enabledClientPlugins: () => [], _mcpAuthTokens: new Map(), + _mcpAuthTokensBySession: new ResourceMap(), + _launchOwner: () => URI.parse('ahp-session:/test'), _githubMcpServerEnabled: true, _githubToken: 'token', _gitHubMcpServerConfiguration: createGitHubMcpServerConfiguration('https://api.githubcopilot.com'), @@ -446,6 +519,8 @@ suite('CodexAgent', () => { _sessionMcpDiscoveries: new Map(), _enabledClientPlugins: () => [], _mcpAuthTokens: new Map(), + _mcpAuthTokensBySession: new ResourceMap(), + _launchOwner: () => URI.parse('ahp-session:/test'), _githubMcpServerEnabled: true, _githubToken: 'token', _gitHubMcpServerConfiguration: createGitHubMcpServerConfiguration('https://api.githubcopilot.com'), diff --git a/src/vs/platform/agentHost/test/node/codex/codexMcpServers.test.ts b/src/vs/platform/agentHost/test/node/codex/codexMcpServers.test.ts index fe5e6b3eb65165..6774e528c22558 100644 --- a/src/vs/platform/agentHost/test/node/codex/codexMcpServers.test.ts +++ b/src/vs/platform/agentHost/test/node/codex/codexMcpServers.test.ts @@ -242,6 +242,21 @@ suite('codexMcpServers', () => { }); }); + test('per-server credentials override shared URL tokens and empty overrides revoke only that server', () => { + const url = 'https://mcp.slack.com/'; + assert.deepStrictEqual(injectCodexMcpAuthTokens({ + first: { url }, second: { url, http_headers: { authorization: 'Bearer configured', 'X-Test': 'keep' } }, + legacy: { url }, moved: { url: 'https://different.example.com/' }, stdio: { command: 'run' }, + }, new Map([[url, 'ambient']]), new Map([ + ['first', { url, token: 'targeted' }], ['second', { url, token: '' }], + ['moved', { url, token: 'old-endpoint' }], ['stdio', { url, token: 'ignored' }], + ])), { + first: { url, http_headers: { Authorization: 'Bearer targeted' } }, + second: { url, http_headers: { 'X-Test': 'keep' } }, + legacy: { url, http_headers: { Authorization: 'Bearer ambient' } }, moved: { url: 'https://different.example.com/' }, stdio: { command: 'run' }, + }); + }); + test('injectCodexMcpAuthTokens returns the input unchanged when there are no tokens', () => { const servers = { s: { url: 'https://mcp.eng.ms' } }; assert.strictEqual(injectCodexMcpAuthTokens(servers, new Map()), servers); diff --git a/src/vs/platform/agentHost/test/node/copilotAgent.test.ts b/src/vs/platform/agentHost/test/node/copilotAgent.test.ts index 4e70e40ecdc105..b0fdf2bf7c7381 100644 --- a/src/vs/platform/agentHost/test/node/copilotAgent.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotAgent.test.ts @@ -104,7 +104,7 @@ import { ICopilotChatDiscoveryScan } from '../../node/copilot/copilotChatDiscove import { COPILOT_DEFERRED_SDK_EXECUTION_METADATA_KEY, deferCopilotSdkExecution } from '../../node/copilot/copilotSessionExecutionMarker.js'; import { CopilotGitHubSessionCredentials } from '../../node/copilot/copilotGitHubCredentials.js'; import type { IAgentHostScopedCredential } from '../../common/agentHostCredentials.js'; -import { AhpErrorCodes, ProtocolError } from '../../common/state/sessionProtocol.js'; +import { AhpErrorCodes, AuthenticateTargetKind, ProtocolError } from '../../common/state/sessionProtocol.js'; import { GITHUB_MCP_SERVER_NAME } from '../../node/shared/githubMcpServer.js'; import { AGENT_HOST_FILE_LINK_INSTRUCTIONS } from '../../node/shared/fileLinkInstructions.js'; import { AGENT_HOST_WORKSPACELESS_INSTRUCTIONS } from '../../node/shared/workspacelessInstructions.js'; @@ -3197,6 +3197,57 @@ suite('CopilotAgent', () => { } }); + test('targeted tokens reach only their owning session during initialization and after registration', async () => { + const client = new McpAuthChallengeCopilotClient([]); + const { agent } = createTestAgentContext(disposables, { + copilotClient: client, + sessionDataService: disposables.add(new TestSessionDataService()), + }); + const sessions = ['target-a', 'target-b'].map(id => URI.parse(`ahp-session:/${id}`)); + const ids = new Map(); + const published = new DeferredPromise(); + disposables.add(agent.onDidChatProgress(signal => { + if (signal.kind === 'action' && signal.action.type === ActionType.SessionCustomizationUpdated + && signal.action.customization.type === CustomizationType.McpServer && signal.action.customization.name === 'example') { + ids.set(signal.resource.toString(), signal.action.customization.id); + if (ids.size === sessions.length) { published.complete(); } + } + })); + const resolver = spy(CopilotAgentSession.prototype, 'resolveMcpAuthentication'); + disposables.add(toDisposable(() => resolver.restore())); + try { + await agent.authenticate('https://api.github.com', 'github-token'); + const creations = sessions.map(session => { + const chat = defaultChatUri(session); + return agent.chats.createChat(chat, exactChatContext(session, chat, session), { + workingDirectories: [URI.file('/workspace')], deferBacking: false, + }); + }); + await published.p; + const authenticate = (index: number, token: string) => agent.handleAuthenticationToken({ + ...authenticationParams(), token, + target: { kind: AuthenticateTargetKind.McpServer, session: sessions[index].toString(), customizationId: ids.get(sessions[index].toString())! }, + }); + const first = await authenticate(0, 'token-a'); + await creations[0]; + const afterFirst = [...client.authenticationResponses]; + const second = await authenticate(1, 'token-b'); + await Promise.all(creations); + const afterRegistration = await authenticate(0, 'token-a'); + assert.deepStrictEqual({ + first, second, afterRegistration, afterFirst, responses: client.authenticationResponses, + deliveries: resolver.getCalls().map(call => (call.thisValue as CopilotAgentSession).ownerSessionUri.toString()), + }, { + first: true, second: true, afterRegistration: false, + afterFirst: [{ kind: 'token', accessToken: 'token-a' }], + responses: [{ kind: 'token', accessToken: 'token-a' }, { kind: 'token', accessToken: 'token-b' }], + deliveries: [sessions[0].toString(), sessions[1].toString(), sessions[0].toString()], + }); + } finally { + await disposeAgent(agent); + } + }); + test('one token authenticates every concurrently initializing new session', async () => { const client = new McpAuthChallengeCopilotClient([]); const { agent, authenticationService } = createTestAgentContext(disposables, { diff --git a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts index 654ca45783e3b2..e62fe5317a7fc3 100644 --- a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts @@ -78,7 +78,7 @@ import { COPILOT_DEFERRED_SDK_EXECUTION_METADATA_KEY } from '../../node/copilot/ import { SessionDataService } from '../../node/sessionDataService.js'; import type { ICopilotMcpServerInfo } from '../../node/copilot/copilotAgent.js'; import { CopilotGitHubCredentials, CopilotGitHubSessionCredentials } from '../../node/copilot/copilotGitHubCredentials.js'; -import { AhpErrorCodes } from '../../common/state/sessionProtocol.js'; +import { AhpErrorCodes, AuthenticateTargetKind } from '../../common/state/sessionProtocol.js'; import { CopilotGitHubTelemetryForwarder } from '../../node/copilot/copilotGitHubTelemetryForwarder.js'; import { ShellManager } from '../../node/copilot/copilotShellTools.js'; import { buildMcpChannel } from '../../node/shared/mcpCustomizationController.js'; @@ -27655,6 +27655,41 @@ Use the attached image as context. }); }); + test('targeted MCP authentication isolates servers sharing a protected resource and scopes', async () => { + const { session, runtime, waitForSignal } = await createAgentSession(disposables); + const resource = 'https://mcp.slack.com'; + const servers = ['slack-a', 'slack-b']; + const results: Array<{ serverName: string; result: Awaited> }> = []; + const requests = servers.map(serverName => runtime.handleMcpAuthRequest({ + requestId: `auth-${serverName}`, serverName, serverUrl: resource, reason: 'initial', + staticClientConfig: { clientId: `client-${serverName}` }, + resourceMetadata: JSON.stringify({ resource, authorization_servers: ['https://auth.slack.com'], scopes_supported: ['read'] }), + wwwAuthenticateParams: { scope: 'read' }, + }, { sessionId: 'test-session-1' }).then(result => { results.push({ serverName, result }); return result; })); + for (const serverName of servers) { + await waitForSignal(signal => isAction(signal, ActionType.SessionCustomizationUpdated) + && signal.action.type === ActionType.SessionCustomizationUpdated && signal.action.customization.name === serverName); + } + const [first, second] = servers.map(name => session.topLevelMcpCustomizations().find(server => server.name === name)!); + const authenticate = (customizationId: string, token: string) => session.resolveMcpAuthentication({ + resource, scopes: ['read'], token, + target: { kind: AuthenticateTargetKind.McpServer, session: session.ownerSessionUri.toString(), customizationId }, + }); + const firstHandled = await authenticate(first.id, 'token-a'); + await requests[0]; + const afterFirst = [...results]; + const secondHandled = await authenticate(second.id, 'token-b'); + await Promise.all(requests); + assert.deepStrictEqual({ firstHandled, secondHandled, afterFirst, results }, { + firstHandled: true, secondHandled: true, + afterFirst: [{ serverName: 'slack-a', result: { kind: 'token', accessToken: 'token-a' } }], + results: [ + { serverName: 'slack-a', result: { kind: 'token', accessToken: 'token-a' } }, + { serverName: 'slack-b', result: { kind: 'token', accessToken: 'token-b' } }, + ], + }); + }); + test('MCP authentication only clears auth-required after all matching server challenges resolve', async () => { const { session, runtime, signals, mockSession } = await createAgentSession(disposables); const firstAuth = runtime.handleMcpAuthRequest({ @@ -27798,6 +27833,53 @@ Use the attached image as context. assert.strictEqual(session.topLevelMcpCustomizations()[0]?.state.kind, McpServerStatus.Ready); }); + suite('readiness after OAuth reconnect without a connected status change', () => { + const serverName = 'slack'; + const resource = 'https://mcp.slack.com'; + + async function deliverToken(session: CopilotAgentSession, runtime: TestCopilotSessionRuntime, mockSession: MockCopilotSession): Promise { + const auth = runtime.handleMcpAuthRequest({ requestId: 'auth-slack', serverName, serverUrl: resource, reason: 'initial' }, { sessionId: 'test-session-1' }); + await timeout(0); + await session.resolveMcpAuthentication({ resource, scopes: [], token: 'token' }); + await auth; + mockSession.fire('mcp.oauth_completed', { requestId: 'auth-slack', outcome: 'token' } as SessionEventPayload<'mcp.oauth_completed'>['data']); + await timeout(0); + } + + test('the reconnect\'s tools-loaded event promotes the server to ready, even after a cached snapshot', async () => { + const { session, runtime, mockSession } = await createAgentSession(disposables); + // The inventory never reports the reconnect, so readiness cannot depend on it. + mockSession.mcpListResult = { servers: [{ name: serverName, status: 'pending' }] }; + fireMcpToolsLoaded(mockSession, serverName); + await deliverToken(session, runtime, mockSession); + const afterOAuth = session.topLevelMcpCustomizations()[0]?.state.kind; + fireMcpToolsLoaded(mockSession, serverName); + + assert.deepStrictEqual({ afterOAuth, afterToolsLoaded: session.topLevelMcpCustomizations()[0]?.state.kind }, { + afterOAuth: McpServerStatus.Starting, + afterToolsLoaded: McpServerStatus.Ready, + }); + }); + + test('a cached snapshot without an OAuth reconnect does not promote the server', async () => { + const { session, mockSession } = await createAgentSession(disposables); + mockSession.fire('session.mcp_server_status_changed', { serverName, status: 'pending' } as SessionEventPayload<'session.mcp_server_status_changed'>['data']); + fireMcpToolsLoaded(mockSession, serverName); + + assert.strictEqual(session.topLevelMcpCustomizations()[0]?.state.kind, McpServerStatus.Starting); + }); + + test('a rejected token clears the inference before the next tools-loaded event', async () => { + const { session, runtime, mockSession } = await createAgentSession(disposables); + mockSession.mcpListResult = { servers: [{ name: serverName, status: 'pending' }] }; + await deliverToken(session, runtime, mockSession); + mockSession.fire('session.mcp_server_status_changed', { serverName, status: 'needs-auth' } as SessionEventPayload<'session.mcp_server_status_changed'>['data']); + fireMcpToolsLoaded(mockSession, serverName); + + assert.notStrictEqual(session.topLevelMcpCustomizations()[0]?.state.kind, McpServerStatus.Ready); + }); + }); + test('ignores an older token completion while a newer authentication callback is pending', async () => { const serverName = 'workiq'; const resource = 'https://mcp.example.com'; diff --git a/src/vs/platform/agentHost/test/node/missionControlAuthentication.test.ts b/src/vs/platform/agentHost/test/node/missionControlAuthentication.test.ts index 6a91e3c1692fec..579fa34ddd440c 100644 --- a/src/vs/platform/agentHost/test/node/missionControlAuthentication.test.ts +++ b/src/vs/platform/agentHost/test/node/missionControlAuthentication.test.ts @@ -8,6 +8,7 @@ import { createHash, randomBytes } from 'crypto'; import sodium from 'libsodium-wrappers'; import { stub, restore, useFakeTimers } from 'sinon'; import { DeferredPromise } from '../../../../base/common/async.js'; +import { AuthenticateTargetKind } from '../../common/state/sessionProtocol.js'; import { AhpErrorCodes, JsonRpcErrorCodes } from '../../common/state/protocol/common/errors.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import { MissionControlAuthentication, MissionControlSealing, resolveMissionControlOwner, sealMissionControlCredential } from '../../node/missionControl/missionControlAuthentication.js'; @@ -316,6 +317,22 @@ suite('Mission Control sealed authentication', () => { }); } + test('targeted MCP tokens sharing the identity origin never reauthenticate the relay owner', async () => { + const sealing = store.add(new MissionControlSealing()); + const identityTokens: string[] = []; + const auth = store.add(new MissionControlAuthentication(sealing, '123', resource, async (_input, init) => { + identityTokens.push(new Headers(init?.headers).get('Authorization')!); + return Response.json({ id: 123, type: 'User' }); + }, false)); + const target = { kind: AuthenticateTargetKind.McpServer, session: 'ahp-session:/target', customizationId: 'mcp' }; + await assert.rejects(auth.authenticate({ resource, token: '', target }), /Authenticate the relay identity/); + await auth.authenticate({ resource, token: seal(sealing, 'owner-token') }); + const accepted = await auth.authenticate({ resource, token: seal(sealing, 'mcp-token', undefined, 'mcp-auth-token'), target }); + assert.deepStrictEqual({ accepted, identityTokens, authorized: auth.authenticated }, { + accepted: { resource, token: 'mcp-token', target }, identityTokens: ['Bearer owner-token'], authorized: true, + }); + }); + test('preserves repository and MCP credentials and scopes', async () => { const sealing = store.add(new MissionControlSealing()); const auth = store.add(new MissionControlAuthentication(sealing, '123', resource, async () => Response.json({ id: 123, type: 'User' }), false)); diff --git a/src/vs/sessions/common/agentHostSessionsProvider.ts b/src/vs/sessions/common/agentHostSessionsProvider.ts index 78dd8fb825ee51..fdf1fe383f1742 100644 --- a/src/vs/sessions/common/agentHostSessionsProvider.ts +++ b/src/vs/sessions/common/agentHostSessionsProvider.ts @@ -437,6 +437,9 @@ export interface IAgentHostSessionsProvider extends ISessionsProvider { */ getBackendChatResource(chatResource: URI): URI | undefined; + /** Returns the host-advertised session URI owning this session's customizations. */ + getBackendSessionResource(sessionId: string): URI | undefined; + } export const LOCAL_AGENT_HOST_PROVIDER_ID = 'local-agent-host'; diff --git a/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts b/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts index a0d6fc420b6152..c0926b9cc337a5 100644 --- a/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts +++ b/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts @@ -4254,6 +4254,10 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement return this._newSessions.get(sessionId); } + getBackendSessionResource(sessionId: string): URI | undefined { + return this._getBackendSessionUri(sessionId); + } + protected _getBackendSessionUri(sessionId: string): URI | undefined { const draft = this._newSessions.get(sessionId); if (draft) { diff --git a/src/vs/sessions/services/agentHost/browser/agentHostCustomizationService.ts b/src/vs/sessions/services/agentHost/browser/agentHostCustomizationService.ts index 0e07a1ffd90af5..be275c09d63b3e 100644 --- a/src/vs/sessions/services/agentHost/browser/agentHostCustomizationService.ts +++ b/src/vs/sessions/services/agentHost/browser/agentHostCustomizationService.ts @@ -68,7 +68,12 @@ export class AgentHostCustomizationService extends AbstractAgentHostCustomizatio if (!provider) { return undefined; } + const backendSession = provider.getBackendSessionResource(session.sessionId); + if (!backendSession) { + return undefined; + } return { + backendSession, customizations: provider.getCustomizations(session.sessionId), resourceUris: { fromAgentHost: resource => provider.mapAgentHostResource(resource), diff --git a/src/vs/sessions/services/agentHost/test/browser/agentHostCustomizationService.test.ts b/src/vs/sessions/services/agentHost/test/browser/agentHostCustomizationService.test.ts index 5b376c258b0e04..0f72e6d502a60e 100644 --- a/src/vs/sessions/services/agentHost/test/browser/agentHostCustomizationService.test.ts +++ b/src/vs/sessions/services/agentHost/test/browser/agentHostCustomizationService.test.ts @@ -38,6 +38,7 @@ suite('AgentHostCustomizationService', () => { override readonly id = 'agenthost-test'; override readonly onDidChangeCustomAgents = Event.None; override readonly onDidChangeCustomizations = Event.None; + override getBackendSessionResource(): URI { return URI.parse('other-agent://tenant/conversations/opaque-session'); } override getCustomizations(): Customization[] { return createCodexSkillCustomizations(); } override getWorkingDirectory(): string { return 'file:///workspace'; } override getWorkingDirectories(): readonly string[] { return ['file:///workspace']; } @@ -133,6 +134,7 @@ suite('AgentHostCustomizationService', () => { override readonly id = 'agenthost-test'; override readonly onDidChangeCustomAgents = Event.None; override readonly onDidChangeCustomizations = Event.None; + override getBackendSessionResource(): URI { return URI.parse('other-agent://tenant/conversations/opaque-session'); } override getCustomizations(): Customization[] { return [plugin]; } @@ -186,6 +188,7 @@ suite('AgentHostCustomizationService', () => { override readonly id = authority === 'local' ? 'local-agent-host' : 'agenthost-remote-test'; override readonly onDidChangeCustomAgents = Event.None; override readonly onDidChangeCustomizations = Event.None; + override getBackendSessionResource(): URI { return URI.parse('other-agent://tenant/conversations/opaque-session'); } override getCustomizations(): Customization[] { return [server]; } override getWorkingDirectory(): string { return roots[0].toString(); } override getWorkingDirectories(): readonly string[] { return roots.map(root => root.toString()); } 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 83f972d8238304..c767389087ad7d 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.ts @@ -19,6 +19,7 @@ import type { IDefaultAccount } from '../../../../../../base/common/defaultAccou import { deriveGitHubEndpoints } from '../../../../../../platform/github/common/githubEndpoints.js'; import { gitHubCopilotResource } from '../../../../../../platform/agentHost/common/githubEndpoints.js'; import { type McpAuthRequirement, type McpOAuthClient, type ModelSelection, type ProtectedResourceMetadata } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import type { AuthenticateTarget, McpServerAuthenticateTarget } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import type { IAgentConnection } from '../../../../../../platform/agentHost/common/agentService.js'; import { ICommandService } from '../../../../../../platform/commands/common/commands.js'; import { IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js'; @@ -88,12 +89,12 @@ export class AgentHostAuthTokenCache { private _globalGeneration = 0; /** - * Forwards a token once per resource/scope pair. Same-token callers share + * Forwards a token once per resource/scope/target combination. Same-token callers share * and await an in-flight authentication. */ - async authenticate(resource: string, scopes: readonly string[] | undefined, token: string, authenticate: () => Promise, accountId?: string, entitlement?: IAgentEntitlementHint): Promise { + async authenticate(resource: string, scopes: readonly string[] | undefined, token: string, authenticate: () => Promise, accountId?: string, entitlement?: IAgentEntitlementHint, target?: AuthenticateTarget): Promise { const credentialKey = JSON.stringify([token, accountId, entitlement?.trackingId, entitlement?.sku]); - const key = this._key(resource, scopes); + const key = this._key(resource, scopes, target); const globalGeneration = this._globalGeneration; const keyGeneration = this._keyGenerations.get(key) ?? 0; const pending = this._pendingAuthentications.get(key); @@ -114,7 +115,7 @@ export class AgentHostAuthTokenCache { if (!this._isCurrentGeneration(key, globalGeneration, keyGeneration)) { throw new CancellationError(); } - return this.authenticate(resource, scopes, token, authenticate, accountId, entitlement); + return this.authenticate(resource, scopes, token, authenticate, accountId, entitlement, target); } if (this._completedTokens.get(key) === credentialKey) { @@ -191,8 +192,8 @@ export class AgentHostAuthTokenCache { return this._globalGeneration === globalGeneration && (this._keyGenerations.get(key) ?? 0) === keyGeneration; } - private _key(resource: string, scopes: readonly string[] | undefined): string { - return `${resource}\x00${scopes ? [...new Set(scopes)].sort().join('\x00') : ''}`; + private _key(resource: string, scopes: readonly string[] | undefined, target?: AuthenticateTarget): string { + return `${resource}\x00${JSON.stringify([[...new Set(scopes)].sort(), ...(target ? [target.kind, target.session, target.customizationId] : [])])}`; } } @@ -482,6 +483,7 @@ async function resolveAuthenticationSessionForResource( } export interface IAgentHostAuthenticateRequest { + readonly target?: AuthenticateTarget; readonly resource: string; readonly scopes?: readonly string[]; /** An empty token revokes the credential previously forwarded for this resource and scope set. */ @@ -509,6 +511,7 @@ interface IResolvedAgentHostAuthenticationOptions extends IAgentHostAuthenticati interface IAgentHostAuthenticationForwardingOptions extends Pick { readonly getDefaultAccount?: () => IDefaultAccount | null; + readonly target?: AuthenticateTarget; } interface IForwardedAuthenticationSession { @@ -529,6 +532,8 @@ function resolveAuthenticationOptions(accessor: ServicesAccessor, options: IAgen } export interface IAgentHostMcpAuthenticationOptionsBase { + /** Echoes the host-advertised session URI and MCP customization id that issued the challenge, when known. */ + readonly target: McpServerAuthenticateTarget | undefined; readonly allowInteraction: boolean; readonly authTokenCache?: AgentHostAuthTokenCache; readonly logPrefix: string; @@ -559,7 +564,7 @@ async function forwardAuthenticationToken( ): Promise { throwIfAuthenticationStale(options); const token = session?.accessToken ?? ''; - const rejectedSession = options.authTokenCache?.getRejectedSession(resource, scopes); + const rejectedSession = options.target ? undefined : options.authTokenCache?.getRejectedSession(resource, scopes); // The lookup may have started before this session was quarantined. if (rejectedSession && rejectedSession.id === session?.id && rejectedSession.accessToken === token) { throw new CancellationError(); @@ -583,6 +588,7 @@ async function forwardAuthenticationToken( } } const request: IAgentHostAuthenticateRequest = { + ...(options.target ? { target: options.target } : {}), resource, scopes, token, @@ -590,7 +596,7 @@ async function forwardAuthenticationToken( ...((account || entitlement) ? { _meta: { ...(account ? authenticationAccountMeta(account) : {}), ...(entitlement ? agentEntitlementHintMeta(entitlement) : {}) } } : {}), }; if (options.authTokenCache) { - return options.authTokenCache.authenticate(resource, scopes ?? [], token, () => options.authenticate(request), authenticationAccountId(account), entitlement); + return options.authTokenCache.authenticate(resource, scopes ?? [], token, () => options.authenticate(request), authenticationAccountId(account), entitlement, options.target); } await options.authenticate(request); return true; @@ -979,8 +985,10 @@ export function autoAuthenticateMcpServer( agentHost: { readonly scheme: string; readonly authority: string }, serverName: string, auth: Pick, + target: McpServerAuthenticateTarget | undefined, ): Promise { return resolveMcpServerAuthentication(accessor, auth.resource, { + target, allowInteraction: false, logPrefix: '[AgentHost]', mcpServerId: agentHostMcpServerId(agentHost.authority, serverName, auth.resource.resource), 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 ef4a216787d90f..4e38e5ebc5e091 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostChatContribution.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostChatContribution.ts @@ -220,7 +220,7 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr if (this._agentHostService.registerMcpAuthenticationHandler) { // Chat session URIs have no authority; retain that identity for remembered MCP grants. store.add(this._agentHostService.registerMcpAuthenticationHandler(request => - this._instantiationService.invokeFunction(autoAuthenticateMcpServer, this._agentHostService, { scheme: AGENT_HOST_SCHEME, authority: '' }, request.serverName, request.auth))); + this._instantiationService.invokeFunction(autoAuthenticateMcpServer, this._agentHostService, { scheme: AGENT_HOST_SCHEME, authority: '' }, request.serverName, request.auth, request.target))); } store.add(this._agentHostService.onDidNotification(notification => { if (notification.type !== NotificationType.AuthRequired) { diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostCustomizationService.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostCustomizationService.ts index 017d1fad579953..49a6e7d4305ed0 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostCustomizationService.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostCustomizationService.ts @@ -20,6 +20,7 @@ import { getEffectiveAgents } from '../../../../../../platform/agentHost/common/ import { readMcpServerConnectorName, readMcpServerControllingSetting, readMcpServerDisplayName, readMcpServerSource, readMcpServerSourcePlugin } from '../../../../../../platform/agentHost/common/meta/mcpCustomizationMeta.js'; import { getCustomizationDisabledReason, isCustomizationEnabled, withCustomizationEnablement } from '../../../../../../platform/agentHost/common/customizationEnablement.js'; import { type IAgentSubscription } from '../../../../../../platform/agentHost/common/state/agentSubscription.js'; +import { AuthenticateTargetKind } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { ActionType } from '../../../../../../platform/agentHost/common/state/protocol/actions.js'; import { CustomizationEnablementKind, CustomizationType, McpServerCustomization, McpServerStatus, type Customization, type CustomizationEnablement, type McpServerState, type PluginCustomization, type RootConfigState, type SessionState } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; import { AgentCustomization, DEFAULT_CHAT_ID, ROOT_STATE_URI, StateComponents, getSessionChatResource, readSessionFolderPickerDecision, type ISessionFolderPickerDecision } from '../../../../../../platform/agentHost/common/state/sessionState.js'; @@ -38,7 +39,7 @@ import { isUntitledChatSession } from '../../../common/model/chatUri.js'; import { IAgentHostUntitledProvisionalSessionService } from './agentHostUntitledProvisionalSessionService.js'; import { IAgentHostActiveClientService } from './agentHostActiveClientService.js'; import { IAgentHostMcpServer } from '../../../../../../sessions/common/agentHostSessionsProvider.js'; -import { resolveMcpServerAuthentication, agentHostMcpServerId, autoAuthenticateMcpServer, authenticateCopilotConnectorCredential, isCopilotConnectorMcpAuthentication } from './agentHostAuth.js'; +import { resolveMcpServerAuthentication, agentHostMcpServerId, autoAuthenticateMcpServer, authenticateCopilotConnectorCredential, isCopilotConnectorMcpAuthentication, type IAgentHostAuthenticateRequest } from './agentHostAuth.js'; import { IOutputService } from '../../../../../services/output/common/output.js'; import { IMcpService } from '../../../../../contrib/mcp/common/mcpTypes.js'; import { MCP } from '../../../../mcp/common/modelContextProtocol.js'; @@ -174,6 +175,7 @@ export class NullAgentHostCustomizationService implements IAgentHostCustomizatio } export interface IAgentHostCustomizationTarget { + readonly backendSession: URI; readonly customizations: readonly Customization[]; readonly resourceUris: IAgentHostResourceUriMapper; readonly folderPickerDecision?: ISessionFolderPickerDecision; @@ -184,7 +186,7 @@ export interface IAgentHostCustomizationTarget { readonly clientWorkingDirectories?: readonly URI[]; readonly rootConfig?: RootConfigState; isBundledMcpServer(pluginUri: string, serverName: string): boolean; - authenticate(request: { resource: string; scopes?: readonly string[]; token: string }): Promise; + authenticate(request: IAgentHostAuthenticateRequest): Promise; handleMcpRequest?(channel: string, method: string, params: Record | undefined): Promise; setCustomizationEnablement(rawId: string, enablement: readonly CustomizationEnablement[]): void; startMcpServer(rawId: string): Promise; @@ -354,7 +356,7 @@ export abstract class AbstractAgentHostCustomizationService extends Disposable i const currentEntry = entry; this._instantiationService.invokeFunction(autoAuthenticateMcpServer, { authenticate: request => target.authenticate(request) as ReturnType, - }, { scheme: sessionResource.scheme, authority: sessionResource.authority }, server.name, state).catch(err => { + }, { scheme: sessionResource.scheme, authority: sessionResource.authority }, server.name, state, { kind: AuthenticateTargetKind.McpServer, session: target.backendSession.toString(), customizationId: server.id }).catch(err => { this._logService.warn(`[AgentHost] Failed to silently authenticate MCP server '${server.name}'`, err); return false; }).then(result => { @@ -499,6 +501,7 @@ export abstract class AbstractAgentHostCustomizationService extends Disposable i try { await target.startMcpServer(server.id); const forwarded = await this._instantiationService.invokeFunction(resolveMcpServerAuthentication, server.state.resource, { + target: { kind: AuthenticateTargetKind.McpServer, session: target.backendSession.toString(), customizationId: server.id }, allowInteraction: true, logPrefix: '[AgentHost]', mcpServerId: agentHostMcpServerId(sessionResource.authority, server.name, server.state.resource.resource), @@ -954,6 +957,7 @@ export class WorkbenchAgentHostCustomizationService extends AbstractAgentHostCus const rootState = target.connection.rootState.value; const channel = target.backendSession.toString(); return { + backendSession: target.backendSession, customizations: sessionState?.customizations ?? [], resourceUris: target.connection.resourceUris, folderPickerDecision: readSessionFolderPickerDecision(sessionState?._meta), 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 53d55e7434ca90..59c164f2796595 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -70,7 +70,7 @@ import { CompletionItemKind as AhpCompletionItemKind, ContentEncoding, type Comp import { BackgroundWorkKind, ConfirmationOptionKind, CustomizationType, JsonPrimitive, McpServerAuthRequiredState, McpServerStatus, SessionInputRequestKind, SessionLifecycle, TerminalClaimKind, ToolCallContributorKind, ToolResultContentType, type ConfirmationOption, type ProtectedResourceMetadata, type SessionActiveClient, type SessionInputRequest, type SessionToolClientExecutionRequest } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; import { compareProtocolVersions } from '../../../../../../platform/agentHost/common/state/protocol/version/registry.js'; import { ActionType, ChatTurnStartedAction, isChatAction, type ClientChatAction, type ClientSessionAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; -import { AhpErrorCodes, AHP_AUTH_REQUIRED, AHP_NOT_FOUND, JsonRpcErrorCodes, ProtocolError } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; +import { AuthenticateTargetKind, AhpErrorCodes, AHP_AUTH_REQUIRED, AHP_NOT_FOUND, JsonRpcErrorCodes, ProtocolError } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { ChangesetStatus, ChatOriginKind, getSessionChatResource, getErrorResponsePart, getInlineToolInput, getTurnError, isChatReadOnly, MessageAttachmentKind, MessageKind, PendingMessageKind, ResponsePartKind, ChatInputAnswerState, ChatInputAnswerValueKind, ChatInputQuestionKind, ChatInputResponseKind, SessionStatus, StateComponents, ToolCallCancellationReason, ToolCallConfirmationReason, ToolCallStatus, TurnState, mergeSessionWithDefaultChat, readSessionWorkspaceless, type ChatState, type ISessionFileDiff, type ISessionWithDefaultChat, type ICompletedToolCall, type InputRequestResponsePart, type MarkdownResponsePart, type Message, type MessageAttachment, type MessageAnnotationsAttachment, type MessageChatAttachment, type MessageResourceAttachment, type MessageEmbeddedResourceAttachment, type ModelSelection, type PendingMessage, type ReasoningResponsePart, type RootState, type ChatInputAnswer, type ChatInputQuestion, type ChatInputRequest, type ChatSummary, type SessionState, type StringOrMarkdown, type ToolCallPendingConfirmationState, type ToolCallResponsePart, type ToolCallRunningState, type ToolCallState, type ToolInput, type Turn, type UsageInfo } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { ExtensionIdentifier } from '../../../../../../platform/extensions/common/extensions.js'; import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js'; @@ -338,7 +338,7 @@ interface IOutputTerminalAttachment { readonly disposable: MutableDisposable; } -function getMcpAuthenticationRequiredServers(sessionResource: URI, state: ISessionWithDefaultChat | undefined): IChatMcpAuthenticationRequiredServer[] { +function getMcpAuthenticationRequiredServers(sessionResource: URI, backendSession: URI, state: ISessionWithDefaultChat | undefined): IChatMcpAuthenticationRequiredServer[] { const servers = state?.customizations?.flatMap(c => c.type === CustomizationType.McpServer ? [c] : c.children?.filter(c => c.type === CustomizationType.McpServer) ?? []) ?? []; @@ -354,6 +354,7 @@ function getMcpAuthenticationRequiredServers(sessionResource: URI, state: ISessi const state = server.state as McpServerAuthRequiredState; return { id: sessionResource.authority + '/' + server.id, + target: { kind: AuthenticateTargetKind.McpServer, session: backendSession.toString(), customizationId: server.id }, name: server.name, resource: state.resource.resource, oauthClient: state.oauthClient, @@ -2925,7 +2926,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC const sessionSub = this._ensureSessionSubscription(backendSession.toString()); let previousServers: readonly IChatMcpAuthenticationRequiredServer[] | undefined; const reconcile = () => { - const servers = getMcpAuthenticationRequiredServers(sessionResource, this._getSessionState(backendSession.toString(), chatURI)); + const servers = getMcpAuthenticationRequiredServers(sessionResource, backendSession, this._getSessionState(backendSession.toString(), chatURI)); if (equals(previousServers, servers)) { return; } @@ -3912,7 +3913,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } })); const mcpAuthRequired$ = derivedOpts({ equalsFn: equals }, reader => { - return getMcpAuthenticationRequiredServers(opts.sessionResource, mergedState$.read(reader)); + return getMcpAuthenticationRequiredServers(opts.sessionResource, opts.backendSession, mergedState$.read(reader)); }); const mcpStarting$ = derivedOpts({ equalsFn: equals }, reader => { const state = mergedState$.read(reader); @@ -4392,7 +4393,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC }, oauthClient: server.oauthClient, requiredScopes: server.requiredScopes ? [...server.requiredScopes] : undefined, - }).catch(err => { + }, server.target).catch(err => { this._logService.error(`[AgentHost] Failed to auto-authenticate MCP server '${server.name}'`, err); return false; }); @@ -4725,14 +4726,14 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC }, }, opts.chatURI); }; - const showAuthentication = (toolCall: ToolCallState & { status: ToolCallStatus.AuthRequired }) => invocation.setAuthenticationRequired(toolCallAuthenticationServer(toolCall, opts.sessionResource.authority, this._toolAuthenticationServerName(toolCall, opts)), onCancelAuthentication); + const showAuthentication = (toolCall: ToolCallState & { status: ToolCallStatus.AuthRequired }) => invocation.setAuthenticationRequired(toolCallAuthenticationServer(toolCall, opts.sessionResource.authority, opts.backendSession, this._toolAuthenticationServerName(toolCall, opts)), onCancelAuthentication); if (!silentAuthenticationAttempted) { // Try silent authentication once per tool call so an expired token // that is renewed automatically does not flash the authentication UI. // A later challenge means the silent credential was rejected, so it // shows the UI directly instead of retrying silently. silentAuthenticationAttempted = true; - const server = toolCallAuthenticationServer(tc, opts.sessionResource.authority); + const server = toolCallAuthenticationServer(tc, opts.sessionResource.authority, opts.backendSession); const serverName = this._customizationService.getMcpServers(opts.sessionResource).find(candidate => candidate.id === server.id)?.name ?? server.name; this._autoAuthenticateMcpServer(opts.sessionResource, { ...server, name: serverName }).then(() => { const current = part$.read(undefined).toolCall; @@ -4802,7 +4803,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } private _toolAuthenticationServerName(tc: ToolCallState & { status: ToolCallStatus.AuthRequired }, opts: IObserveTurnOptions): string | undefined { - const id = toolCallAuthenticationServer(tc, opts.sessionResource.authority).id; + const id = toolCallAuthenticationServer(tc, opts.sessionResource.authority, opts.backendSession).id; const server = this._customizationService.getMcpServers(opts.sessionResource).find(candidate => candidate.id === id); return server ? getMcpServerDisplayLabel(server) : undefined; } diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts index 0aeba5faa8fd53..c70c8cca00f8d6 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts @@ -39,6 +39,7 @@ import { isViewUnreviewedCommentsTool, isAddCommentTool } from '../../../../../. import { AGENT_HOST_SESSION_LINK_SCHEME, buildOpenSessionLinkUri, isCreateChatTool, isCreateSessionTool, isSendMessageTool, parseOpenSessionLinkChatId, parseOpenSessionLinkUri } from '../../../../../../platform/agentHost/common/openSessionLink.js'; import { getAgentHostChatId } from '../../../../../../platform/agentHost/common/agentHostChatIdentity.js'; import { parsePartialToolInputForDisplay } from '../../../../../../platform/agentHost/common/partialToolInput.js'; +import { AuthenticateTargetKind } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { MessageAttachmentKind, type FileEdit, type MessageAttachment, type ResponsePart, type StringOrMarkdown, type TextRange } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; import { normalizeFileEdit } from '../../../../../../platform/agentHost/common/fileEditDiff.js'; import { AgentSession } from '../../../../../../platform/agentHost/common/agentService.js'; @@ -2787,7 +2788,7 @@ export function toolCallStateToInvocation(tc: ToolCallState, subAgentInvocationI invocation.presentation = ToolInvocationPresentation.Hidden; } if (tc.status === ToolCallStatus.AuthRequired) { - invocation.setAuthenticationRequired(toolCallAuthenticationServer(tc, mcpServerAuthority, mcpServerName)); + invocation.setAuthenticationRequired(toolCallAuthenticationServer(tc, mcpServerAuthority, sessionResource, mcpServerName)); } applyToolCallProgress(invocation, tc); @@ -2849,10 +2850,11 @@ export function toolCallConfirmationMessages(tc: ToolCallPendingConfirmationStat }; } -export function toolCallAuthenticationServer(tc: ToolCallState & { status: ToolCallStatus.AuthRequired }, sessionAuthority: string, name?: string): IChatMcpAuthenticationRequiredServer { +export function toolCallAuthenticationServer(tc: ToolCallState & { status: ToolCallStatus.AuthRequired }, sessionAuthority: string, backendSession: URI, name?: string): IChatMcpAuthenticationRequiredServer { const metadata = readToolCallMeta(tc); return { id: `${sessionAuthority}/${tc.contributor.customizationId}`, + target: { kind: AuthenticateTargetKind.McpServer, session: backendSession.toString(), customizationId: tc.contributor.customizationId }, name: name ?? tc.auth.resource.resource_name ?? metadata.mcpServerName ?? tc.displayName, resource: tc.auth.resource.resource, oauthClient: tc.auth.oauthClient, diff --git a/src/vs/workbench/contrib/chat/browser/remoteAgentHost/cloudSandboxConnectionCustomization.ts b/src/vs/workbench/contrib/chat/browser/remoteAgentHost/cloudSandboxConnectionCustomization.ts index 6ea567b81f95c8..8334d10dc8a7f5 100644 --- a/src/vs/workbench/contrib/chat/browser/remoteAgentHost/cloudSandboxConnectionCustomization.ts +++ b/src/vs/workbench/contrib/chat/browser/remoteAgentHost/cloudSandboxConnectionCustomization.ts @@ -74,13 +74,14 @@ export function createCloudSandboxConnectionCustomization( if (reason !== AuthRequiredReason.Expired && isCloudSandboxSealedToken(request.token)) { return request; } - if (userLocal && new URL(request.resource).origin !== new URL(GITHUB_COPILOT_PROTECTED_RESOURCE.resource).origin) { + if (userLocal && (request.target || new URL(request.resource).origin !== new URL(GITHUB_COPILOT_PROTECTED_RESOURCE.resource).origin)) { if (!sealUserLocal) { throw new Error('User-local Mission Control sealing is unavailable.'); } return sealUserLocal(request); } - return resolveAuthentication(request.resource, request.scopes, reason === AuthRequiredReason.Expired); + const authentication = await resolveAuthentication(request.resource, request.scopes, reason === AuthRequiredReason.Expired); + return request.target ? { ...authentication, target: request.target } : authentication; }, renewAuthentication: userLocal ? undefined : resource => resolveAuthentication(resource.resource, resource.scopes_supported, true), backendSessionScheme: (provider: string): string | undefined => diff --git a/src/vs/workbench/contrib/chat/browser/remoteAgentHost/remoteAgentHostChatContribution.ts b/src/vs/workbench/contrib/chat/browser/remoteAgentHost/remoteAgentHostChatContribution.ts index f3607942598655..8250045ae6df92 100644 --- a/src/vs/workbench/contrib/chat/browser/remoteAgentHost/remoteAgentHostChatContribution.ts +++ b/src/vs/workbench/contrib/chat/browser/remoteAgentHost/remoteAgentHostChatContribution.ts @@ -297,7 +297,7 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc store.add(connection.registerMcpAuthenticationHandler(request => this._instantiationService.invokeFunction(autoAuthenticateMcpServer, { authenticate: this._authenticateCallback(address, connection) }, - { scheme: AGENT_HOST_SCHEME, authority: '' }, request.serverName, request.auth))); + { scheme: AGENT_HOST_SCHEME, authority: '' }, request.serverName, request.auth, request.target))); } connState.prepareSession = this._connectionCustomizations.get(address)?.createSessionPreparation?.(connection, store); diff --git a/src/vs/workbench/contrib/chat/common/chatService/chatService.ts b/src/vs/workbench/contrib/chat/common/chatService/chatService.ts index f2b1c6e2ac1131..6fc0592757abf8 100644 --- a/src/vs/workbench/contrib/chat/common/chatService/chatService.ts +++ b/src/vs/workbench/contrib/chat/common/chatService/chatService.ts @@ -39,6 +39,7 @@ import { IParsedChatRequest } from '../requestParser/chatParserTypes.js'; import { IChatParserContext } from '../requestParser/chatRequestParser.js'; import { IPreparedToolInvocation, IToolConfirmationMessages, IToolResult, IToolResultInputOutputDetails, ToolDataSource } from '../tools/languageModelToolsService.js'; import type { ChatToolInvocationSummary } from '../tools/toolInvocationSummary.js'; +import type { McpServerAuthenticateTarget } from '../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { ConfirmationOptionKind, type McpOAuthClient, type MessageOrigin } from '../../../../../platform/agentHost/common/state/protocol/state.js'; import { AgentFusionPhaseStatus } from '../../../../../platform/agentHost/common/meta/agentToolCallMeta.js'; import type { IAgentRuntimeModelConfiguration } from '../../../../../platform/agentHost/common/meta/agentModelConfigurationMeta.js'; @@ -1503,6 +1504,7 @@ export interface IChatMcpAuthenticationRequired { } export interface IChatMcpAuthenticationRequiredServer { + readonly target?: McpServerAuthenticateTarget; readonly id: string; readonly name: string; readonly resource: string; 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 82ee4e97be5622..c0dbd8ed2b3c99 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 @@ -12,6 +12,7 @@ import { DisposableStore } from '../../../../../../base/common/lifecycle.js'; import { isObject } from '../../../../../../base/common/types.js'; import { URI } from '../../../../../../base/common/uri.js'; import { type ProtectedResourceMetadata } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; +import { AuthenticateTargetKind, type McpServerAuthenticateTarget } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import type { IAgentConnection } from '../../../../../../platform/agentHost/common/agentService.js'; import { AgentHostMcpOAuthMetadataService, IAgentHostMcpOAuthMetadataService } from '../../../../../../platform/agentHost/common/agentHostMcpOAuthMetadataService.js'; import { authenticationAccountMeta, readAuthenticationAccount } from '../../../../../../platform/agentHost/common/meta/agentAuthenticationAccount.js'; @@ -37,9 +38,15 @@ import { IAuthenticationMcpUsageService } from '../../../../../services/authenti import { IAuthenticationService, type AuthenticationSession, type IAuthenticationProvider, type IAuthenticationProviderSessionOptions } from '../../../../../services/authentication/common/authentication.js'; import { IDynamicAuthenticationProviderStorageService } from '../../../../../services/authentication/common/dynamicAuthenticationProviderStorage.js'; import { CHAT_SETUP_ACTION_ID } from '../../../browser/actions/chatActions.js'; -import { AgentHostAuthenticationRecovery, authenticateCopilotConnectorCredential, autoAuthenticateMcpServer, authenticateProtectedResources, resolveAuthenticationInteractively, resolveSessionForResource, AgentHostAuthTokenCache, agentHostMcpServerId, resolveMcpServerAuthentication, modelRequiresAgentAuthentication, revokeAuthenticationForRemovedSessions, type IAgentHostAuthenticationOptions } from '../../../browser/agentSessions/agentHost/agentHostAuth.js'; +import { AgentHostAuthenticationRecovery, authenticateCopilotConnectorCredential, autoAuthenticateMcpServer, authenticateProtectedResources, resolveAuthenticationInteractively, resolveSessionForResource, AgentHostAuthTokenCache, agentHostMcpServerId, resolveMcpServerAuthentication, modelRequiresAgentAuthentication, revokeAuthenticationForRemovedSessions, type IAgentHostAuthenticationOptions, type IAgentHostAuthenticateRequest } from '../../../browser/agentSessions/agentHost/agentHostAuth.js'; import { createAgentModelByokMeta } from '../../../../../../platform/agentHost/common/agentModelByokMeta.js'; +const mcpAuthenticationTarget: McpServerAuthenticateTarget = { + kind: AuthenticateTargetKind.McpServer, + session: 'other-agent://tenant/conversations/opaque-session?revision=1', + customizationId: 'urn:example:mcp:account-a', +}; + class TestCommandService extends mock() { readonly calls: { commandId: string; args: unknown[] }[] = []; result: unknown = { success: true, dialogSkipped: false }; @@ -240,6 +247,24 @@ suite('AgentHostAuthTokenCache', () => { assert.deepStrictEqual({ results, authenticateCalls }, { results: [true, false], authenticateCalls: 1 }); }); + test('tracks identical MCP tokens independently by host session and customization id', async () => { + const cache = new AgentHostAuthTokenCache(); + const targets = [ + mcpAuthenticationTarget, + { ...mcpAuthenticationTarget, customizationId: 'opaque/account-b/server' }, + { ...mcpAuthenticationTarget, session: 'other-agent://tenant/conversations/another-session' }, + undefined, + ]; + const forwarded: (McpServerAuthenticateTarget | undefined)[] = []; + const results: boolean[] = []; + for (const target of [...targets, ...targets]) { + results.push(await cache.authenticate('https://mcp.slack.com', ['read'], 'shared-token', async () => { + forwarded.push(target); + }, undefined, undefined, target)); + } + assert.deepStrictEqual({ results, forwarded }, { results: [true, true, true, true, false, false, false, false], forwarded: targets }); + }); + test('preserves entitlement hints when waiting for an earlier same-token authentication', async () => { const cache = new AgentHostAuthTokenCache(); const pending = new DeferredPromise(); @@ -1017,14 +1042,14 @@ suite('autoAuthenticateMcpServer', () => { }, oauthClient: { clientId: 'docs-client' }, requiredScopes: ['read'], - }); + }, mcpAuthenticationTarget); results.push({ authenticated, serverIds, sessionRequests, authenticateRequests }); } assert.deepStrictEqual(results, [true, false].map(allowed => ({ authenticated: allowed, serverIds: [agentHostMcpServerId('', 'Docs', 'https://docs.example/mcp')], sessionRequests: [{ scopes: ['read'], silent: true }], - authenticateRequests: allowed ? [{ resource: 'https://docs.example/mcp', scopes: ['read'], token: 'docs-token' }] : [], + authenticateRequests: allowed ? [{ target: mcpAuthenticationTarget, resource: 'https://docs.example/mcp', scopes: ['read'], token: 'docs-token' }] : [], }))); }); @@ -1054,10 +1079,11 @@ suite('autoAuthenticateMcpServer', () => { const silent = await instantiationService.invokeFunction(autoAuthenticateMcpServer, connection, { scheme: 'vscode-agent-host', authority: '' }, 'connector-kusto', { resource, requiredScopes: ['write:plugin_gateway_connections'], - }); + }, mcpAuthenticationTarget); const interactive = await instantiationService.invokeFunction(resolveMcpServerAuthentication, resource, { allowInteraction: true, logPrefix: '[AgentHost]', + target: mcpAuthenticationTarget, mcpServerId: agentHostMcpServerId('', 'connector-kusto', resource.resource), mcpServerName: 'Kusto', mcpServerUrl: resource.resource, @@ -1113,6 +1139,7 @@ suite('resolveMcpServerAuthentication', () => { }, { allowInteraction: true, logPrefix: '[Test]', + target: mcpAuthenticationTarget, mcpServerId: 'server', mcpServerName: 'Server', mcpServerUrl: resource, @@ -1153,6 +1180,7 @@ suite('resolveMcpServerAuthentication', () => { }, { allowInteraction: false, logPrefix: '[AgentHost]', + target: mcpAuthenticationTarget, mcpServerId: 'server-id', mcpServerName: 'Example', mcpServerUrl: 'https://mcp.example.com', @@ -1193,6 +1221,7 @@ suite('resolveMcpServerAuthentication', () => { }, { allowInteraction: false, logPrefix: '[AgentHost]', + target: mcpAuthenticationTarget, mcpServerId: 'slack', mcpServerName: 'Slack', mcpServerUrl: 'https://mcp.slack.com', @@ -1233,6 +1262,7 @@ suite('resolveMcpServerAuthentication', () => { }, { allowInteraction: false, logPrefix: '[AgentHost]', + target: mcpAuthenticationTarget, mcpServerId: 'github', mcpServerName: 'GitHub', mcpServerUrl: 'https://api.githubcopilot.com/mcp', @@ -1278,6 +1308,7 @@ suite('resolveMcpServerAuthentication', () => { }, { allowInteraction: false, logPrefix: '[AgentHost]', + target: mcpAuthenticationTarget, mcpServerId: 'server-id', mcpServerName: 'Example', mcpServerUrl: 'https://mcp.example.com', @@ -1348,6 +1379,7 @@ suite('resolveMcpServerAuthentication', () => { }, { allowInteraction: false, logPrefix: '[AgentHost]', + target: mcpAuthenticationTarget, mcpServerId: 'avalara', mcpServerName: 'Avalara', mcpServerUrl: 'https://mcp.avalara.com/elr', @@ -1408,6 +1440,7 @@ suite('resolveMcpServerAuthentication', () => { }, { allowInteraction: false, logPrefix: '[AgentHost]', + target: mcpAuthenticationTarget, mcpServerId: 'notion', mcpServerName: 'notion', mcpServerUrl: 'https://mcp.notion.com/mcp', @@ -1430,6 +1463,7 @@ suite('resolveMcpServerAuthentication', () => { providerCreations: [{ clientId: 'notion-client-id', clientSecret: 'notion-client-secret' }], sessionRequests: [{ silent: true }], authenticateRequests: [{ + target: mcpAuthenticationTarget, resource: 'https://mcp.notion.com/mcp', scopes: [], token: 'notion-token', @@ -1437,13 +1471,13 @@ suite('resolveMcpServerAuthentication', () => { }); }); - test('serializes authentication transactions for different configured clients', async () => { + test('serializes OAuth clients and forwards distinct host-advertised targets for a shared MCP resource', async () => { const dynamicProviderId = 'https://mcp.example.com/ https://mcp.example.com/resource'; const firstSessionStarted = new DeferredPromise(); const firstSessionGate = new DeferredPromise(); const providerCreations: string[] = []; const sessionRequests: string[] = []; - const authenticateRequests: string[] = []; + const authenticateRequests: IAgentHostAuthenticateRequest[] = []; let activeClient: string | undefined; let providerActive = false; const authService = createMockAuthService({ @@ -1497,7 +1531,8 @@ suite('resolveMcpServerAuthentication', () => { const options = (clientId: string) => ({ allowInteraction: true, logPrefix: '[AgentHost]', - mcpServerId: 'example', + target: { ...mcpAuthenticationTarget, customizationId: `opaque/${clientId}/server` }, + mcpServerId: clientId, mcpServerName: 'Example', mcpServerUrl: 'https://mcp.example.com/resource', oauthClient: { clientId }, @@ -1510,8 +1545,8 @@ suite('resolveMcpServerAuthentication', () => { discoveryUrl: `${authorizationServer}/.well-known/oauth-authorization-server`, errors: [], }), - authenticate: async (request: { token: string }) => { - authenticateRequests.push(request.token); + authenticate: async (request: IAgentHostAuthenticateRequest) => { + authenticateRequests.push(request); }, }); @@ -1539,7 +1574,12 @@ suite('resolveMcpServerAuthentication', () => { results: [true, true], providerCreations: ['first-client', 'second-client'], sessionRequests: ['first-client', 'second-client'], - authenticateRequests: ['first-client-token', 'second-client-token'], + authenticateRequests: ['first-client', 'second-client'].map(clientId => ({ + target: { ...mcpAuthenticationTarget, customizationId: `opaque/${clientId}/server` }, + resource: protectedResource.resource, + scopes: [], + token: `${clientId}-token`, + })), }); }); @@ -1585,6 +1625,7 @@ suite('resolveMcpServerAuthentication', () => { }, { allowInteraction: false, logPrefix: '[AgentHost]', + target: mcpAuthenticationTarget, mcpServerId: 'slack', mcpServerName: 'Slack', mcpServerUrl: 'https://mcp.slack.com', @@ -1607,6 +1648,7 @@ suite('resolveMcpServerAuthentication', () => { result: true, providerCreations: ['slack-client-id'], authenticateRequests: [{ + target: mcpAuthenticationTarget, resource: 'https://mcp.slack.com', scopes: ['search:read.public'], token: 'slack-token', @@ -1702,6 +1744,7 @@ suite('resolveMcpServerAuthentication', () => { }, { allowInteraction: true, logPrefix: '[AgentHost]', + target: mcpAuthenticationTarget, mcpServerId: 'slack', mcpServerName: 'Slack', mcpServerUrl: 'https://mcp.slack.com', @@ -1753,11 +1796,13 @@ suite('resolveMcpServerAuthentication', () => { ], authenticateRequests: [ { + target: mcpAuthenticationTarget, resource: 'https://mcp.slack.com', scopes: ['search:read.public'], token: 'public-token', }, { + target: mcpAuthenticationTarget, resource: 'https://mcp.slack.com', scopes: ['search:read.public'], token: 'confidential-token', @@ -1831,6 +1876,28 @@ suite('authenticateProtectedResources', () => { const disposables = ensureNoDisposablesAreLeakedInTestSuite(); + test('provider-level protected resource authentication does not include an MCP target', async () => { + const authentication = createMockAuthService({ + getOrActivateProviderIdForServer: async () => 'example-provider', + getSessions: async () => [{ id: 'session', scopes: ['read'], accessToken: 'provider-token', account: { id: 'account', label: 'Account' } }], + }); + const instantiationService = createAuthInstantiationService(disposables, authentication); + const agents = [new class extends mock() { + override readonly protectedResources = [protectedResource]; + }()]; + const requests: IAgentHostAuthenticateRequest[] = []; + await instantiationService.invokeFunction(authenticateProtectedResources, agents, { + logPrefix: '[Test]', + authenticate: async request => { requests.push(request); }, + }); + assert.deepStrictEqual(requests, [{ + resource: protectedResource.resource, + scopes: ['read'], + token: 'provider-token', + _meta: authenticationAccountMeta({ providerId: 'example-provider', accountId: 'account' }), + }]); + }); + for (const defaultSession of ['session', 'connector-session']) { test(`forwards same-account entitlement changes when default account uses ${defaultSession}`, async () => { const baseSession = { id: 'session', scopes: ['read'], accessToken: 'token', account: { id: 'account', label: 'Account' } }; 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 5143d82e73ec9d..d1ecdeb044844e 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 @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { AuthenticateTargetKind, type AuthenticateTarget } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { renderMarkdown } from '../../../../../../base/browser/markdownRenderer.js'; import { DeferredPromise, timeout } from '../../../../../../base/common/async.js'; import { encodeBase64, VSBuffer } from '../../../../../../base/common/buffer.js'; @@ -2735,9 +2736,9 @@ suite('AgentHostClientTools', () => { instantiationService.stub(IAuthenticationMcpAccessService, { isAccessAllowedForUrl: () => true }); instantiationService.stub(IAuthenticationMcpService, { getAccountPreference: () => undefined }); instantiationService.stub(IAuthenticationMcpUsageService, { addAccountUsage: () => { } }); - const authenticateRequests: { resource: string; token: string }[] = []; + const authenticateRequests: { resource: string; token: string; target?: AuthenticateTarget }[] = []; connection.authenticate = async params => { - authenticateRequests.push({ resource: params.resource, token: params.token }); + authenticateRequests.push({ resource: params.resource, token: params.token, target: params.target }); connection.applySessionAction(chatURI, { type: ActionType.ChatToolCallAuthResolved, turnId: 'turn-1', @@ -2814,7 +2815,7 @@ suite('AgentHostClientTools', () => { authenticateRequests, }, { showedAuthentication: false, - authenticateRequests: [{ resource: 'https://mcp.notion.com/mcp', token: 'notion-token' }], + authenticateRequests: [{ resource: 'https://mcp.notion.com/mcp', token: 'notion-token', target: { kind: AuthenticateTargetKind.McpServer, session: 'copilot:/session-1', customizationId: 'notion-mcp' } }], }); })); diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostCustomizationService.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostCustomizationService.test.ts index b79a836375d6c2..dfdd6e9ad41732 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostCustomizationService.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostCustomizationService.test.ts @@ -37,6 +37,8 @@ import { IAuthenticationMcpUsageService } from '../../../../../services/authenti import { IDynamicAuthenticationProviderStorageService } from '../../../../../services/authentication/common/dynamicAuthenticationProviderStorage.js'; import { IMcpService } from '../../../../mcp/common/mcpTypes.js'; import { MCP } from '../../../../mcp/common/modelContextProtocol.js'; +import { AuthenticateTargetKind } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; +import type { IAgentHostAuthenticateRequest } from '../../../browser/agentSessions/agentHost/agentHostAuth.js'; import { AbstractAgentHostCustomizationService, IAgentHostCustomizationTarget, WorkbenchAgentHostCustomizationService } from '../../../browser/agentSessions/agentHost/agentHostCustomizationService.js'; import { IAgentHostUntitledProvisionalSessionService } from '../../../browser/agentSessions/agentHost/agentHostUntitledProvisionalSessionService.js'; import { IChatService } from '../../../common/chatService/chatService.js'; @@ -46,13 +48,14 @@ import { IAgentHostActiveClientService } from '../../../browser/agentSessions/ag import { assertCodexSkillItems, createCodexSkillCustomizations } from './agentHostSkillDiscoveryTestUtils.js'; class FakeTarget implements IAgentHostCustomizationTarget { + readonly backendSession = URI.parse('other-agent://tenant/conversations/opaque-session?revision=1'); readonly enablementChanges: { readonly rawId: string; readonly enablement: readonly CustomizationEnablement[] }[] = []; readonly startCalls: string[] = []; - readonly authenticateCalls: { resource: string; scopes?: readonly string[]; token: string }[] = []; + readonly authenticateCalls: IAgentHostAuthenticateRequest[] = []; readonly operationLog: string[] = []; authenticateResult: unknown = { authenticated: true }; startError: Error | undefined; - onAuthenticate: ((request: { resource: string; scopes?: readonly string[]; token: string }) => void) | undefined; + onAuthenticate: ((request: IAgentHostAuthenticateRequest) => void) | undefined; handleMcpRequest?: IAgentHostCustomizationTarget['handleMcpRequest']; constructor( @@ -67,7 +70,7 @@ class FakeTarget implements IAgentHostCustomizationTarget { return this._isBundledMcpServer(pluginUri, serverName); } - authenticate(request: { resource: string; scopes?: readonly string[]; token: string }): Promise { + authenticate(request: IAgentHostAuthenticateRequest): Promise { this.operationLog.push('authenticate'); this.authenticateCalls.push(request); this.onAuthenticate?.(request); @@ -633,6 +636,39 @@ suite('AbstractAgentHostCustomizationService', () => { }); }); + test('targets the advertised backend session and distinct customization ids for servers sharing a resource', async () => { + const sut = createSut(); + const session = URI.parse('vscode-agent-session:///frontend-session'); + const servers: McpServerCustomization[] = ['opaque/account-a/server', 'urn:example:mcp:account-b'].map(id => ({ + ...mcpServer(id, 'Slack'), + state: { + kind: McpServerStatus.AuthRequired, + reason: McpAuthRequiredReason.Required, + resource: { resource: 'https://mcp.slack.com', authorization_servers: ['https://auth.example.com'] }, + requiredScopes: ['read'], + }, + })); + const target = new FakeTarget(servers); + target.onAuthenticate = request => queueMicrotask(() => sut.updateTarget(target, target.customizations.map(customization => + customization.type === CustomizationType.McpServer && customization.id === request.target?.customizationId + ? { ...customization, state: { kind: McpServerStatus.Ready } } + : customization))); + sut.setTarget(session, target); + const results: boolean[] = []; + for (const server of servers) { + results.push(await sut.authenticateMcpServer(session, server.id)); + } + assert.deepStrictEqual({ results, requests: target.authenticateCalls }, { + results: [true, true], + requests: servers.map(server => ({ + target: { kind: AuthenticateTargetKind.McpServer, session: target.backendSession.toString(), customizationId: server.id }, + resource: 'https://mcp.slack.com', + scopes: ['read'], + token: 'token', + })), + }); + }); + test('starts an unchanged auth-required server before forwarding root authentication results and reports success once it restarts', async () => { const authenticationTargets: Array<{ id: string; name: string }> = []; const sut = createSut(undefined, { authenticationTargets }); @@ -667,8 +703,8 @@ suite('AbstractAgentHostCustomizationService', () => { startCalls: ['server-1', 'server-1'], operationLog: ['start', 'authenticate', 'start', 'authenticate'], authenticateCalls: [ - { resource: 'https://mcp.example.com', scopes: [], token: 'token' }, - { resource: 'https://mcp.example.com', scopes: [], token: 'token' }, + { target: { kind: AuthenticateTargetKind.McpServer, session: target.backendSession.toString(), customizationId: 'server-1' }, resource: 'https://mcp.example.com', scopes: [], token: 'token' }, + { target: { kind: AuthenticateTargetKind.McpServer, session: target.backendSession.toString(), customizationId: 'server-1' }, resource: 'https://mcp.example.com', scopes: [], token: 'token' }, ], authenticationTargets: [ { id: 'agent-host-mcp:/Server%20One/https%3A%2F%2Fmcp.example.com', name: 'Connector One (Connector)' }, @@ -686,8 +722,8 @@ suite('AbstractAgentHostCustomizationService', () => { failedAuthentication: false, operationLog: ['start', 'authenticate', 'start', 'authenticate', 'start'], authenticateCalls: [ - { resource: 'https://mcp.example.com', scopes: [], token: 'token' }, - { resource: 'https://mcp.example.com', scopes: [], token: 'token' }, + { target: { kind: AuthenticateTargetKind.McpServer, session: target.backendSession.toString(), customizationId: 'server-1' }, resource: 'https://mcp.example.com', scopes: [], token: 'token' }, + { target: { kind: AuthenticateTargetKind.McpServer, session: target.backendSession.toString(), customizationId: 'server-1' }, resource: 'https://mcp.example.com', scopes: [], token: 'token' }, ], }); }); diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/stateToProgressAdapter.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/stateToProgressAdapter.test.ts index 598432eb9ac8f2..e11a51b0870ce2 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/stateToProgressAdapter.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/stateToProgressAdapter.test.ts @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { AuthenticateTargetKind } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { Codicon } from '../../../../../../base/common/codicons.js'; import { autorun } from '../../../../../../base/common/observable.js'; import { hasKey } from '../../../../../../base/common/types.js'; @@ -2743,7 +2744,7 @@ suite('stateToProgressAdapter', () => { const invocation = rawToolCallStateToInvocation({ ...createToolCallState(), status: ToolCallStatus.AuthRequired, - contributor: { kind: ToolCallContributorKind.MCP, customizationId: 'mcp-1' }, + contributor: { kind: ToolCallContributorKind.MCP, customizationId: 'opaque/server-id:with/slashes' }, auth: { reason: McpAuthRequiredReason.InsufficientScope, oauthClient: { @@ -2758,7 +2759,7 @@ suite('stateToProgressAdapter', () => { }, requiredScopes: ['repo'], }, - }, undefined, URI.parse('agent-host-copilot://backend/session'), 'remote', 'frontend', undefined, undefined, 'Example MCP (Connector)'); + }, undefined, URI.parse('other-agent://tenant/conversations/opaque-session?revision=1'), 'remote', 'frontend', undefined, undefined, 'Example MCP (Connector)'); const state = invocation.state.get(); assert.strictEqual(state.type, IChatToolInvocation.StateKind.WaitingForAuthentication); @@ -2773,7 +2774,8 @@ suite('stateToProgressAdapter', () => { parameters: undefined, confirmationMessages: undefined, server: { - id: 'frontend/mcp-1', + id: 'frontend/opaque/server-id:with/slashes', + target: { kind: AuthenticateTargetKind.McpServer, session: URI.parse('other-agent://tenant/conversations/opaque-session?revision=1').toString(), customizationId: 'opaque/server-id:with/slashes' }, name: 'Example MCP (Connector)', resource: 'https://mcp.example.com', oauthClient: { diff --git a/src/vs/workbench/contrib/chat/test/browser/remoteAgentHost/cloudSandboxConnectionCustomization.test.ts b/src/vs/workbench/contrib/chat/test/browser/remoteAgentHost/cloudSandboxConnectionCustomization.test.ts index f8d8111a1e3628..86d7d545317493 100644 --- a/src/vs/workbench/contrib/chat/test/browser/remoteAgentHost/cloudSandboxConnectionCustomization.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/remoteAgentHost/cloudSandboxConnectionCustomization.test.ts @@ -7,6 +7,7 @@ import assert from 'assert'; import { mock } from '../../../../../../base/test/common/mock.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; import { cloudSandboxAddress, ICloudSandboxAgentHostService } from '../../../../../../platform/agentHost/common/cloudSandboxAgentHost.js'; +import { AuthenticateTargetKind, type McpServerAuthenticateTarget } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { AuthRequiredReason } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; import { createCloudSandboxConnectionCustomization } from '../../../browser/remoteAgentHost/cloudSandboxConnectionCustomization.js'; @@ -72,6 +73,25 @@ suite('CloudSandboxConnectionCustomization authentication', () => { }); }); + test('preserves an MCP target when replacing a credential with a sealed token', async () => { + const fixture = createFixture(); + const target: McpServerAuthenticateTarget = { kind: AuthenticateTargetKind.McpServer, session: 'other-agent://tenant/conversations/opaque-session', customizationId: 'opaque/github/server' }; + const request = { target, resource: 'https://api.github.com/mcp', scopes: ['repo'], token: 'plaintext' }; + const result = await fixture.authenticate(request); + assert.deepStrictEqual({ result, requests: fixture.requests }, { + result: { ...request, token: 'copilot-sealed.v1.key.cached' }, + requests: ['cached:env-1'], + }); + }); + + test('user-local MCP targets use MCP sealing even at the provider resource origin', async () => { + const service = new class extends mock() { }(); + const customization = createCloudSandboxConnectionCustomization(cloudSandboxAddress('env-1'), service, true, async request => ({ ...request, token: 'copilot-sealed.v1.mcp.fresh' })); + const target: McpServerAuthenticateTarget = { kind: AuthenticateTargetKind.McpServer, session: 'other-agent://tenant/conversations/opaque-session', customizationId: 'opaque/github/server' }; + const request = { target, resource: 'https://api.github.com', token: 'test-mcp-token' }; + assert.deepStrictEqual(await customization?.authenticate?.(request), { ...request, token: 'copilot-sealed.v1.mcp.fresh' }); + }); + for (const token of ['plaintext', 'copilot-sealed.v1.key.challenged']) { test(`renews an expired challenge instead of forwarding ${token}`, async () => { const fixture = createFixture(); diff --git a/src/vs/workbench/contrib/chat/test/browser/remoteAgentHost/remoteAgentHostChatContribution.test.ts b/src/vs/workbench/contrib/chat/test/browser/remoteAgentHost/remoteAgentHostChatContribution.test.ts index 5e693b01235b5b..925be8800f169b 100644 --- a/src/vs/workbench/contrib/chat/test/browser/remoteAgentHost/remoteAgentHostChatContribution.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/remoteAgentHost/remoteAgentHostChatContribution.test.ts @@ -14,6 +14,7 @@ import { IAgentConnection, type AuthenticateParams } from '../../../../../../pla import { IAgentHostMcpAuthenticationRequest } from '../../../../../../platform/agentHost/common/agentHostExtensionProtocol.js'; import { ICloudSandboxAgentHostService } from '../../../../../../platform/agentHost/common/cloudSandboxAgentHost.js'; import { IRemoteAgentHostConnectionInfo, IRemoteAgentHostService, RemoteAgentHostConnectionStatus } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { AuthenticateTargetKind } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { McpAuthRequiredReason } from '../../../../../../platform/agentHost/common/state/protocol/channels-session/state.js'; import { type INotification } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; import { IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js'; @@ -47,6 +48,7 @@ import { ILanguageModelsService } from '../../../common/languageModels.js'; suite('Remote Agent Host silent MCP authentication', () => { const store = ensureNoDisposablesAreLeakedInTestSuite(); const request: IAgentHostMcpAuthenticationRequest = { + target: { kind: AuthenticateTargetKind.McpServer, session: 'other-agent://tenant/conversations/opaque-session', customizationId: 'urn:example:mcp:docs' }, serverName: 'Docs', auth: { reason: McpAuthRequiredReason.Required, @@ -169,8 +171,8 @@ suite('Remote Agent Host silent MCP authentication', () => { const authenticated = await context.authenticate(request); assert.deepStrictEqual({ authenticated, sealed, sent: context.sent, silent: context.sessionOptions }, { authenticated: true, - sealed: [{ resource: 'https://docs.example/mcp', scopes: ['read'], token: 'docs-token' }], - sent: [{ resource: 'https://docs.example/mcp', scopes: ['read'], token: 'copilot-sealed.v1.mcp.ciphertext' }], + sealed: [{ target: request.target, resource: 'https://docs.example/mcp', scopes: ['read'], token: 'docs-token' }], + sent: [{ target: request.target, resource: 'https://docs.example/mcp', scopes: ['read'], token: 'copilot-sealed.v1.mcp.ciphertext' }], silent: [true], }); }); @@ -179,7 +181,7 @@ suite('Remote Agent Host silent MCP authentication', () => { const context = fixture(); const authenticated = await context.authenticate(request); assert.deepStrictEqual({ authenticated, sent: context.sent }, { - authenticated: true, sent: [{ resource: 'https://docs.example/mcp', scopes: ['read'], token: 'docs-token' }], + authenticated: true, sent: [{ target: request.target, resource: 'https://docs.example/mcp', scopes: ['read'], token: 'docs-token' }], }); }); From d50410c17c4aacd09ede0dd645563d99d9c18dd6 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Tue, 20 Oct 2026 13:54:19 -0700 Subject: [PATCH 2/3] agentHost: keep MCP OAuth sign-ins separate for each static client The workbench used one dynamic authentication provider for each authorization server and resource, and ignored the OAuth client ID. MCP servers with different static clients for the same resource used one provider. A sign-in with a different client removed the provider and its stored tokens, so a sign-in to one server signed out the other. After a reload, silent authentication did not find a token for that server. - Agent host MCP authentication: give each static OAuth client its own provider ID (` client:`). Use the legacy provider when its stored registration already has the same client, so existing sign-ins continue to work. - Do not remove a provider registration when a request uses a different client. Replace a registration only for the same client, for example when its secret changes. - Add an optional provider ID to `createDynamicAuthenticationProvider`, the provider host delegate, and `$registerDynamicAuthProvider`. The extension host provider then keeps its registration and tokens under that ID. Providers without a static client keep their IDs. Fixes https://github.com/microsoft/vscode/issues/342147 Fixes https://github.com/microsoft/vscode/issues/293533 (Commit message generated by Copilot) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../api/browser/mainThreadAuthentication.ts | 7 +- .../workbench/api/common/extHost.protocol.ts | 2 +- .../api/common/extHostAuthentication.ts | 13 +- .../api/node/extHostAuthentication.ts | 6 +- .../extHostAuthentication.integrationTest.ts | 1 + .../browser/extHostAuthentication.test.ts | 10 ++ ...ainThreadAuthentication.integrationTest.ts | 31 ++++ .../agentSessions/agentHost/agentHostAuth.ts | 31 +++- .../agentSessions/agentHostAuth.test.ts | 169 +++++++++++++++--- .../browser/authenticationService.ts | 12 +- .../authentication/common/authentication.ts | 7 +- .../authenticationQueryServiceMocks.ts | 2 +- 12 files changed, 243 insertions(+), 48 deletions(-) diff --git a/src/vs/workbench/api/browser/mainThreadAuthentication.ts b/src/vs/workbench/api/browser/mainThreadAuthentication.ts index 07a79c929e51ae..6aaccb13dbb2bc 100644 --- a/src/vs/workbench/api/browser/mainThreadAuthentication.ts +++ b/src/vs/workbench/api/browser/mainThreadAuthentication.ts @@ -178,8 +178,8 @@ export class MainThreadAuthentication extends Disposable implements MainThreadAu this._register(authenticationService.registerAuthenticationProviderHostDelegate({ // Prefer Node.js extension hosts when they're available. No CORS issues etc. priority: extHostContext.extensionHostKind === ExtensionHostKind.LocalWebWorker ? 0 : 1, - create: async (authorizationServer, serverMetadata, resource, overrideClientId, overrideClientSecret) => { - const authProviderId = getDynamicAuthenticationProviderId(authorizationServer, resource); + create: async (authorizationServer, serverMetadata, resource, overrideClientId, overrideClientSecret, providerId) => { + const authProviderId = providerId ?? getDynamicAuthenticationProviderId(authorizationServer, resource); const clientDetails = await this.dynamicAuthProviderStorageService.getClientRegistration(authProviderId); let clientId = overrideClientId ?? clientDetails?.clientId; const clientSecret = overrideClientId @@ -199,7 +199,8 @@ export class MainThreadAuthentication extends Disposable implements MainThreadAu resource, clientId, clientSecret, - initialTokens + initialTokens, + authProviderId ); }, createXaa: async (issuer) => { diff --git a/src/vs/workbench/api/common/extHost.protocol.ts b/src/vs/workbench/api/common/extHost.protocol.ts index 722cc9745f0faf..e763abe77f7a8f 100644 --- a/src/vs/workbench/api/common/extHost.protocol.ts +++ b/src/vs/workbench/api/common/extHost.protocol.ts @@ -2581,7 +2581,7 @@ export interface ExtHostAuthenticationShape { $removeSession(id: string, sessionId: string): Promise; $onDidChangeAuthenticationSessions(id: string, label: string, extensionIdFilter?: string[]): Promise; $onDidUnregisterAuthenticationProvider(id: string): Promise; - $registerDynamicAuthProvider(authorizationServer: UriComponents, serverMetadata: IAuthorizationServerMetadata, resource?: IAuthorizationProtectedResourceMetadata, clientId?: string, clientSecret?: string, initialTokens?: (IAuthorizationTokenResponse & { created_at: number })[]): Promise; + $registerDynamicAuthProvider(authorizationServer: UriComponents, serverMetadata: IAuthorizationServerMetadata, resource?: IAuthorizationProtectedResourceMetadata, clientId?: string, clientSecret?: string, initialTokens?: (IAuthorizationTokenResponse & { created_at: number })[], providerId?: string): Promise; $registerXaaAuthProvider(issuer: UriComponents, serverMetadata: IAuthorizationServerMetadata, clientId?: string, clientSecret?: string, initialTokens?: (IAuthorizationTokenResponse & { created_at: number })[]): Promise; $onDidChangeDynamicAuthProviderTokens(authProviderId: string, clientId: string, tokens?: (IAuthorizationTokenResponse & { created_at: number })[]): Promise; } diff --git a/src/vs/workbench/api/common/extHostAuthentication.ts b/src/vs/workbench/api/common/extHostAuthentication.ts index 99330661b60f0a..f236b17a8ab8d0 100644 --- a/src/vs/workbench/api/common/extHostAuthentication.ts +++ b/src/vs/workbench/api/common/extHostAuthentication.ts @@ -279,7 +279,8 @@ export class ExtHostAuthentication implements ExtHostAuthenticationShape { resourceMetadata: IAuthorizationProtectedResourceMetadata | undefined, clientId: string | undefined, clientSecret: string | undefined, - initialTokens: IAuthorizationToken[] | undefined + initialTokens: IAuthorizationToken[] | undefined, + providerId?: string ): Promise { if (!clientId) { const authorizationServer = URI.revive(authorizationServerComponents); @@ -322,7 +323,8 @@ export class ExtHostAuthentication implements ExtHostAuthenticationShape { clientId, clientSecret, this._onDidDynamicAuthProviderTokensChange, - initialTokens || [] + initialTokens || [], + providerId ); // Use the sequencer to ensure dynamic provider registration is serialized @@ -511,13 +513,14 @@ export class DynamicAuthProvider implements vscode.AuthenticationProvider { protected _clientSecret: string | undefined, onDidDynamicAuthProviderTokensChange: Emitter<{ authProviderId: string; clientId: string; tokens: IAuthorizationToken[] }>, initialTokens: IAuthorizationToken[], + providerId?: string, private readonly _fetch: typeof fetch = fetch, ) { const stringifiedServer = authorizationServer.toString(true); - // Auth Provider Id is a combination of the authorization server and the resource, if provided. - this.id = _resourceMetadata?.resource + // The default provider ID combines the authorization server and the resource, if provided. + this.id = providerId ?? (_resourceMetadata?.resource ? stringifiedServer + ' ' + _resourceMetadata?.resource - : stringifiedServer; + : stringifiedServer); // Auth Provider label is just the resource name if provided, otherwise the authority of the authorization server. this.label = _resourceMetadata?.resource_name ?? this.authorizationServer.authority; // Metadata is untrusted, so the token endpoint must also be Microsoft to earn Microsoft-specific treatment. diff --git a/src/vs/workbench/api/node/extHostAuthentication.ts b/src/vs/workbench/api/node/extHostAuthentication.ts index a6e458bf095ed1..68d5c954178ee3 100644 --- a/src/vs/workbench/api/node/extHostAuthentication.ts +++ b/src/vs/workbench/api/node/extHostAuthentication.ts @@ -39,7 +39,8 @@ export class NodeDynamicAuthProvider extends DynamicAuthProvider { clientId: string, clientSecret: string | undefined, onDidDynamicAuthProviderTokensChange: Emitter<{ authProviderId: string; clientId: string; tokens: any[] }>, - initialTokens: any[] + initialTokens: any[], + providerId?: string ) { super( extHostWindow, @@ -54,7 +55,8 @@ export class NodeDynamicAuthProvider extends DynamicAuthProvider { clientId, clientSecret, onDidDynamicAuthProviderTokensChange, - initialTokens + initialTokens, + providerId ); // Prepend Node-specific flows to the existing flows diff --git a/src/vs/workbench/api/test/browser/extHostAuthentication.integrationTest.ts b/src/vs/workbench/api/test/browser/extHostAuthentication.integrationTest.ts index cf5897160bdff9..c8cfde10794ceb 100644 --- a/src/vs/workbench/api/test/browser/extHostAuthentication.integrationTest.ts +++ b/src/vs/workbench/api/test/browser/extHostAuthentication.integrationTest.ts @@ -987,6 +987,7 @@ suite('ExtHostAuthentication', () => { { issuer: authorizationServer.toString(true), response_types_supported: ['code'], token_endpoint: 'https://issuer.example/token' }, undefined, 'original-client', undefined, disposables.add(new Emitter()), [{ access_token: 'expired-token', token_type: 'Bearer', scope: 'read', refresh_token: 'refresh-token', expires_in: 1, created_at: 0 }], + undefined, async () => { refreshRequests++; return new Response(JSON.stringify({ error: 'invalid_client' }), { status: 400 }); diff --git a/src/vs/workbench/api/test/browser/extHostAuthentication.test.ts b/src/vs/workbench/api/test/browser/extHostAuthentication.test.ts index 200596bd9c5bb5..5b0e125ba3627f 100644 --- a/src/vs/workbench/api/test/browser/extHostAuthentication.test.ts +++ b/src/vs/workbench/api/test/browser/extHostAuthentication.test.ts @@ -116,6 +116,7 @@ suite('DynamicAuthProvider', () => { function createProvider(options: ITestServer & { readonly resource: string; readonly initialTokens?: IAuthorizationToken[]; + readonly providerId?: string; readonly fetcher?: typeof fetch; readonly onDidOpenUri?: (uri: string) => void; }): TestDynamicAuthProvider { @@ -162,10 +163,19 @@ suite('DynamicAuthProvider', () => { undefined, disposables.add(new Emitter()), options.initialTokens ?? [], + options.providerId, options.fetcher, )); } + test('honors an explicit provider ID without changing the default identity', () => { + const providerId = 'https://auth.example.com/ https://workiq.example.com/mcp client:client-id'; + const scoped = createProvider({ ...servers.nonMicrosoft, resource, providerId }); + const legacy = createProvider({ ...servers.nonMicrosoft, resource }); + + assert.deepStrictEqual([scoped.id, legacy.id], [providerId, 'https://auth.example.com/ https://workiq.example.com/mcp']); + }); + test('does not rotate the client while silently refreshing a token', async () => { let fetchCalls = 0; const fetcher: typeof fetch = async () => { diff --git a/src/vs/workbench/api/test/browser/mainThreadAuthentication.integrationTest.ts b/src/vs/workbench/api/test/browser/mainThreadAuthentication.integrationTest.ts index f77a2ef4b30d10..fb49eb361cbe7f 100644 --- a/src/vs/workbench/api/test/browser/mainThreadAuthentication.integrationTest.ts +++ b/src/vs/workbench/api/test/browser/mainThreadAuthentication.integrationTest.ts @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { URI } from '../../../../base/common/uri.js'; import { IDialogService } from '../../../../platform/dialogs/common/dialogs.js'; import { TestDialogService } from '../../../../platform/dialogs/test/common/testDialogService.js'; import { TestInstantiationService } from '../../../../platform/instantiation/test/common/instantiationServiceMock.js'; @@ -105,6 +106,36 @@ suite('MainThreadAuthentication', () => { assert.ok(true, 'Registration and unregistration completed successfully'); }); + test('uses an explicit dynamic provider ID for registration and persisted tokens', async () => { + const authorizationServer = URI.parse('https://auth.example.com'); + const resource = { resource: 'https://resource.example.com/mcp' }; + const providerId = 'https://auth.example.com/ https://resource.example.com/mcp client:configured-client'; + const tokens = [{ access_token: 'stored-token', token_type: 'Bearer', created_at: 0 }]; + const registrations: Parameters[] = []; + rpcProtocol.set(ExtHostContext.ExtHostAuthentication, upcastPartial({ + $registerDynamicAuthProvider: async (...args) => { + registrations.push(args); + const id = args[6] ?? 'unexpected-legacy-id'; + await mainThreadAuthentication.$registerAuthenticationProvider({ id, label: 'Test', supportsMultipleAccounts: false }); + return id; + }, + $onDidChangeDynamicAuthProviderTokens: async () => { }, + })); + const storage = instantiationService.get(IDynamicAuthenticationProviderStorageService); + await storage.storeClientRegistration(providerId, authorizationServer.toString(true), 'configured-client', 'configured-secret'); + await storage.setSessionsForDynamicAuthProvider(providerId, 'configured-client', tokens); + const metadata = { issuer: authorizationServer.toString(true), response_types_supported: ['code'] }; + + const provider = await instantiationService.get(IAuthenticationService).createDynamicAuthenticationProvider(authorizationServer, metadata, resource, undefined, undefined, providerId); + await mainThreadAuthentication.$unregisterAuthenticationProvider(providerId); + await rpcProtocol.sync(); + + assert.deepStrictEqual({ providerId: provider?.id, registrations }, { + providerId, + registrations: [[authorizationServer.toJSON(), metadata, resource, 'configured-client', 'configured-secret', tokens, providerId]], + }); + }); + test('event suppression during explicit unregistration', async () => { let unregisterEventFired = false; let eventProviderId: string | undefined; 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 c767389087ad7d..171ff30a3112df 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.ts @@ -1035,12 +1035,13 @@ export async function resolveMcpServerAuthentication( const authenticationOperations = getMcpAuthenticationOperations(authenticationService); for (const authorizationServer of protectedResource.authorization_servers ?? []) { const authorizationServerUri = URI.parse(authorizationServer); - const providerOperationId = getDynamicAuthenticationProviderId(authorizationServerUri, protectedResource); + const providerOperationId = await getMcpAuthenticationProviderId(authorizationServerUri, protectedResource, options.oauthClient, dynamicAuthenticationProviderStorageService); const authenticated = await authenticationOperations.queue(providerOperationId, async () => { const providerId = await getOrCreateProviderForMcpResource( authorizationServerUri, protectedResource, options.oauthClient, + providerOperationId, authenticationService, dynamicAuthenticationProviderStorageService, logService, @@ -1108,10 +1109,29 @@ function isGitHubMcpResource(resource: ProtectedResourceMetadata): boolean { return resource.resource_name === 'GitHub MCP Server'; } +/** Reuses matching legacy sign-ins; other static clients get independent registrations and tokens. */ +async function getMcpAuthenticationProviderId( + authorizationServer: URI, + protectedResource: ProtectedResourceMetadata, + oauthClient: McpOAuthClient | undefined, + dynamicAuthenticationProviderStorageService: IDynamicAuthenticationProviderStorageService, +): Promise { + const legacyProviderId = getDynamicAuthenticationProviderId(authorizationServer, protectedResource); + if (!oauthClient) { + return legacyProviderId; + } + const registeredClient = await dynamicAuthenticationProviderStorageService.getClientRegistration(legacyProviderId); + if (registeredClient?.clientId === oauthClient.clientId && registeredClient.clientSecret === oauthClient.clientSecret) { + return legacyProviderId; + } + return `${legacyProviderId} client:${encodeURIComponent(oauthClient.clientId)}`; +} + async function getOrCreateProviderForMcpResource( authorizationServer: URI, protectedResource: ProtectedResourceMetadata, oauthClient: McpOAuthClient | undefined, + dynamicProviderId: string, authenticationService: IAuthenticationService, dynamicAuthenticationProviderStorageService: IDynamicAuthenticationProviderStorageService, logService: ILogService, @@ -1120,7 +1140,6 @@ async function getOrCreateProviderForMcpResource( authorizationServerMetadataFetcher: typeof fetchAuthorizationServerMetadata, ): Promise { const resourceUri = URI.parse(protectedResource.resource); - const dynamicProviderId = getDynamicAuthenticationProviderId(authorizationServer, protectedResource); let clientId = oauthClient?.clientId; let clientSecret = oauthClient?.clientSecret; if (oauthClient) { @@ -1135,8 +1154,10 @@ async function getOrCreateProviderForMcpResource( if (!allowCreation) { return undefined; } - if (isProviderActive) { - authenticationService.unregisterAuthenticationProvider(dynamicProviderId); + if (registeredClient?.clientId === oauthClient.clientId) { + if (isProviderActive) { + authenticationService.unregisterAuthenticationProvider(dynamicProviderId); + } await dynamicAuthenticationProviderStorageService.removeDynamicProvider(dynamicProviderId); } } @@ -1155,7 +1176,7 @@ async function getOrCreateProviderForMcpResource( try { const { metadata } = await authorizationServerMetadataFetcher(authorizationServer.toString(true)); - const provider = await authenticationService.createDynamicAuthenticationProvider(authorizationServer, metadata, protectedResource, clientId, clientSecret); + const provider = await authenticationService.createDynamicAuthenticationProvider(authorizationServer, metadata, protectedResource, clientId, clientSecret, oauthClient ? dynamicProviderId : undefined); return provider?.id; } catch (err) { logService.warn(`${logPrefix} Failed to create MCP auth provider for ${authorizationServer.toString(true)}`, err); 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 c0dbd8ed2b3c99..99d4b650421f74 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 @@ -1471,8 +1471,8 @@ suite('resolveMcpServerAuthentication', () => { }); }); - test('serializes OAuth clients and forwards distinct host-advertised targets for a shared MCP resource', async () => { - const dynamicProviderId = 'https://mcp.example.com/ https://mcp.example.com/resource'; + test('serializes authentication for the same OAuth client with distinct host-advertised targets', async () => { + const dynamicProviderId = 'https://mcp.example.com/ https://mcp.example.com/resource client:shared-client'; const firstSessionStarted = new DeferredPromise(); const firstSessionGate = new DeferredPromise(); const providerCreations: string[] = []; @@ -1494,7 +1494,7 @@ suite('resolveMcpServerAuthentication', () => { getSessions: async () => { const clientId = activeClient ?? ''; sessionRequests.push(clientId); - if (clientId === 'first-client') { + if (sessionRequests.length === 1) { firstSessionStarted.complete(); await firstSessionGate.p; } @@ -1518,7 +1518,7 @@ suite('resolveMcpServerAuthentication', () => { addAccountUsage: () => { }, }); instantiationService.stub(IDynamicAuthenticationProviderStorageService, { - getClientRegistration: () => Promise.resolve(activeClient ? { clientId: activeClient } : undefined), + getClientRegistration: providerId => Promise.resolve(providerId === dynamicProviderId && activeClient ? { clientId: activeClient } : undefined), removeDynamicProvider: async () => { activeClient = undefined; }, @@ -1535,7 +1535,7 @@ suite('resolveMcpServerAuthentication', () => { mcpServerId: clientId, mcpServerName: 'Example', mcpServerUrl: 'https://mcp.example.com/resource', - oauthClient: { clientId }, + oauthClient: { clientId: 'shared-client' }, scopes: [], authorizationServerMetadataFetcher: async (authorizationServer: string) => ({ metadata: { @@ -1568,21 +1568,128 @@ suite('resolveMcpServerAuthentication', () => { authenticateRequests, }, { beforeResolution: { - providerCreations: ['first-client'], - sessionRequests: ['first-client'], + providerCreations: ['shared-client'], + sessionRequests: ['shared-client'], }, results: [true, true], - providerCreations: ['first-client', 'second-client'], - sessionRequests: ['first-client', 'second-client'], + providerCreations: ['shared-client'], + sessionRequests: ['shared-client', 'shared-client'], authenticateRequests: ['first-client', 'second-client'].map(clientId => ({ target: { ...mcpAuthenticationTarget, customizationId: `opaque/${clientId}/server` }, resource: protectedResource.resource, scopes: [], - token: `${clientId}-token`, + token: 'shared-client-token', })), }); }); + for (const legacyRegistration of ['none', 'matching', 'different-secret']) { + test(`isolates static OAuth clients sharing a resource with ${legacyRegistration} legacy registration`, async () => { + const useLegacyProvider = legacyRegistration === 'matching'; + const legacyProviderId = 'https://mcp.slack.com/ https://mcp.slack.com'; + const firstProviderId = useLegacyProvider ? legacyProviderId : `${legacyProviderId} client:github-slack-client`; + const secondProviderId = `${legacyProviderId} client:slack-client`; + const registrations = new Map(); + const activeProviders = new Set(); + if (legacyRegistration !== 'none') { + registrations.set(legacyProviderId, { clientId: 'github-slack-client', clientSecret: legacyRegistration === 'matching' ? 'configured-secret' : 'previous-secret' }); + activeProviders.add(legacyProviderId); + } + const providerCreations: string[] = []; + const sessionRequests: { providerId: string; silent: boolean }[] = []; + const removedProviders: string[] = []; + const authenticateRequests: IAgentHostAuthenticateRequest[] = []; + const authService = createMockAuthService({ + isDynamicAuthenticationProvider: providerId => activeProviders.has(providerId), + createDynamicAuthenticationProvider: async (_authorizationServer, _metadata, _resource, clientId, clientSecret, providerId) => { + assert.ok(providerId && clientId); + providerCreations.push(providerId); + registrations.set(providerId, { clientId, clientSecret }); + activeProviders.add(providerId); + return { id: providerId }; + }, + unregisterAuthenticationProvider: providerId => { + removedProviders.push(providerId); + activeProviders.delete(providerId); + }, + getSessions: async (providerId, _scopes, options) => { + sessionRequests.push({ providerId, silent: options.silent }); + const clientId = activeProviders.has(providerId) ? registrations.get(providerId)?.clientId : undefined; + return clientId ? [{ + id: `${clientId}-session`, + scopes: [], + accessToken: `${clientId}-token`, + account: { id: 'account-id', label: 'Slack Account' }, + }] : []; + }, + }); + const instantiationService = disposables.add(new TestInstantiationService()); + instantiationService.stub(IAuthenticationService, authService); + instantiationService.stub(IAuthenticationMcpAccessService, { isAccessAllowedForUrl: () => true }); + instantiationService.stub(IAuthenticationMcpService, { getAccountPreference: () => 'Slack Account' }); + instantiationService.stub(IAuthenticationMcpUsageService, { addAccountUsage: () => { } }); + instantiationService.stub(IDynamicAuthenticationProviderStorageService, { + getClientRegistration: providerId => Promise.resolve(registrations.get(providerId)), + removeDynamicProvider: async providerId => { + removedProviders.push(providerId); + registrations.delete(providerId); + }, + }); + instantiationService.stub(ILogService, new NullLogService()); + const protectedResource = { + resource: 'https://mcp.slack.com', + authorization_servers: ['https://mcp.slack.com'], + }; + const authenticate = (serverName: string, clientId: string, allowInteraction: boolean) => instantiationService.invokeFunction(resolveMcpServerAuthentication, protectedResource, { + allowInteraction, + logPrefix: '[AgentHost]', + target: { ...mcpAuthenticationTarget, customizationId: `opaque/${serverName}/server` }, + mcpServerId: serverName, + mcpServerName: serverName, + mcpServerUrl: protectedResource.resource, + oauthClient: { clientId, clientSecret: clientId === 'github-slack-client' ? 'configured-secret' : undefined }, + scopes: [], + authorizationServerMetadataFetcher: async authorizationServer => ({ + metadata: { issuer: authorizationServer, response_types_supported: ['code'] }, + discoveryUrl: `${authorizationServer}/.well-known/oauth-authorization-server`, + errors: [], + }), + authenticate: async request => { authenticateRequests.push(request); }, + }); + + const results = [ + await authenticate('github-slack', 'github-slack-client', true), + await authenticate('slack', 'slack-client', true), + await authenticate('github-slack', 'github-slack-client', false), + ]; + activeProviders.clear(); + results.push(await authenticate('github-slack', 'github-slack-client', false)); + results.push(await authenticate('slack', 'slack-client', false)); + + assert.deepStrictEqual({ results, providerCreations, sessionRequests, removedProviders, authenticateRequests, registeredProviders: [...registrations.keys()] }, { + results: [true, true, true, true, true], + providerCreations: useLegacyProvider + ? [secondProviderId, firstProviderId, secondProviderId] + : [firstProviderId, secondProviderId, firstProviderId, secondProviderId], + sessionRequests: [ + { providerId: firstProviderId, silent: false }, + { providerId: secondProviderId, silent: false }, + { providerId: firstProviderId, silent: true }, + { providerId: firstProviderId, silent: true }, + { providerId: secondProviderId, silent: true }, + ], + removedProviders: [], + authenticateRequests: ['github-slack', 'slack', 'github-slack', 'github-slack', 'slack'].map(serverName => ({ + target: { ...mcpAuthenticationTarget, customizationId: `opaque/${serverName}/server` }, + resource: protectedResource.resource, + scopes: [], + token: `${serverName === 'github-slack' ? 'github-slack-client' : 'slack-client'}-token`, + })), + registeredProviders: legacyRegistration === 'different-secret' ? [legacyProviderId, firstProviderId, secondProviderId] : [firstProviderId, secondProviderId], + }); + }); + } + test('restores a persisted configured provider without user interaction', async () => { const dynamicProviderId = 'https://mcp.slack.com/ https://mcp.slack.com'; const providerCreations: string[] = []; @@ -1656,14 +1763,14 @@ suite('resolveMcpServerAuthentication', () => { }); }); - test('uses configured public and confidential clients when creating a dynamic provider', async () => { + test('keeps configured public and confidential clients separate when a client secret changes', async () => { const dynamicProviderId = 'https://mcp.slack.com/ https://mcp.slack.com'; - const providerCreations: { authorizationServer: string; resource: string | undefined; clientId: string | undefined; clientSecret: string | undefined }[] = []; + const providerCreations: { authorizationServer: string; resource: string | undefined; clientId: string | undefined; clientSecret: string | undefined; providerId: string | undefined }[] = []; const sessionRequests: { clientId: string | undefined; clientSecret: string | undefined }[] = []; const sessionCreations: { clientId: string | undefined; clientSecret: string | undefined }[] = []; const authenticateRequests: { resource: string; scopes?: readonly string[]; token: string }[] = []; const removedProviders: string[] = []; - let registeredClient: { clientId?: string; clientSecret?: string } | undefined; + const registeredClients = new Map(); let getSessionsCall = 0; const provider: IAuthenticationProvider = { id: dynamicProviderId, @@ -1694,21 +1801,23 @@ suite('resolveMcpServerAuthentication', () => { scopes: ['search:read.public'], }); }, - createDynamicAuthenticationProvider: async (authorizationServer, _metadata, resource, clientId, clientSecret) => { + createDynamicAuthenticationProvider: async (authorizationServer, _metadata, resource, clientId, clientSecret, providerId) => { providerCreations.push({ authorizationServer: authorizationServer.toString(true), resource: resource?.resource, clientId, clientSecret, + providerId, }); - registeredClient = { clientId, clientSecret }; - return { id: dynamicProviderId }; + const id = providerId ?? dynamicProviderId; + registeredClients.set(id, { clientId, clientSecret }); + return { id }; }, getProvider: () => provider, - isDynamicAuthenticationProvider: providerId => providerId === dynamicProviderId && registeredClient !== undefined, + isDynamicAuthenticationProvider: providerId => registeredClients.has(providerId), unregisterAuthenticationProvider: providerId => { removedProviders.push(providerId); - registeredClient = undefined; + registeredClients.delete(providerId); }, }); const instantiationService = disposables.add(new TestInstantiationService()); @@ -1725,7 +1834,7 @@ suite('resolveMcpServerAuthentication', () => { addAccountUsage: () => { }, }); instantiationService.stub(IDynamicAuthenticationProviderStorageService, { - getClientRegistration: () => Promise.resolve(registeredClient), + getClientRegistration: providerId => Promise.resolve(registeredClients.get(providerId)), removeDynamicProvider: async providerId => { removedProviders.push(providerId); }, @@ -1736,6 +1845,7 @@ suite('resolveMcpServerAuthentication', () => { for (const oauthClient of [ { clientId: 'public-client-id' }, { clientId: 'confidential-client-id', clientSecret: 'confidential-client-secret' }, + { clientId: 'confidential-client-id', clientSecret: 'updated-secret' }, ]) { results.push(await instantiationService.invokeFunction(resolveMcpServerAuthentication, { resource: 'https://mcp.slack.com', @@ -1772,27 +1882,38 @@ suite('resolveMcpServerAuthentication', () => { authenticateRequests, removedProviders, }, { - results: [true, true], + results: [true, true, true], providerCreations: [ { authorizationServer: 'https://mcp.slack.com/', resource: 'https://mcp.slack.com', clientId: 'public-client-id', clientSecret: undefined, + providerId: `${dynamicProviderId} client:public-client-id`, }, { authorizationServer: 'https://mcp.slack.com/', resource: 'https://mcp.slack.com', clientId: 'confidential-client-id', clientSecret: 'confidential-client-secret', + providerId: `${dynamicProviderId} client:confidential-client-id`, + }, + { + authorizationServer: 'https://mcp.slack.com/', + resource: 'https://mcp.slack.com', + clientId: 'confidential-client-id', + clientSecret: 'updated-secret', + providerId: `${dynamicProviderId} client:confidential-client-id`, }, ], sessionRequests: [ { clientId: 'public-client-id', clientSecret: undefined }, { clientId: 'confidential-client-id', clientSecret: 'confidential-client-secret' }, + { clientId: 'confidential-client-id', clientSecret: 'updated-secret' }, ], sessionCreations: [ { clientId: 'confidential-client-id', clientSecret: 'confidential-client-secret' }, + { clientId: 'confidential-client-id', clientSecret: 'updated-secret' }, ], authenticateRequests: [ { @@ -1807,8 +1928,14 @@ suite('resolveMcpServerAuthentication', () => { scopes: ['search:read.public'], token: 'confidential-token', }, + { + target: mcpAuthenticationTarget, + resource: 'https://mcp.slack.com', + scopes: ['search:read.public'], + token: 'confidential-token', + }, ], - removedProviders: [dynamicProviderId, dynamicProviderId], + removedProviders: [`${dynamicProviderId} client:confidential-client-id`, `${dynamicProviderId} client:confidential-client-id`], }); }); }); diff --git a/src/vs/workbench/services/authentication/browser/authenticationService.ts b/src/vs/workbench/services/authentication/browser/authenticationService.ts index 55fd7e16592766..9b85592ece2ca9 100644 --- a/src/vs/workbench/services/authentication/browser/authenticationService.ts +++ b/src/vs/workbench/services/authentication/browser/authenticationService.ts @@ -364,20 +364,20 @@ export class AuthenticationService extends Disposable implements IAuthentication return undefined; } - async createDynamicAuthenticationProvider(authorizationServer: URI, serverMetadata: IAuthorizationServerMetadata, resource: IAuthorizationProtectedResourceMetadata | undefined, clientId?: string, clientSecret?: string): Promise { + async createDynamicAuthenticationProvider(authorizationServer: URI, serverMetadata: IAuthorizationServerMetadata, resource: IAuthorizationProtectedResourceMetadata | undefined, clientId?: string, clientSecret?: string, providerId?: string): Promise { const delegate = this._delegates[0]; if (!delegate) { this._logService.error('No authentication provider host delegate found'); return undefined; } - const providerId = await delegate.create(authorizationServer, serverMetadata, resource, clientId, clientSecret); - const provider = this._authenticationProviders.get(providerId); + const createdProviderId = await delegate.create(authorizationServer, serverMetadata, resource, clientId, clientSecret, providerId); + const provider = this._authenticationProviders.get(createdProviderId); if (provider) { - this._logService.debug(`Created dynamic authentication provider: ${providerId}`); - this._dynamicAuthenticationProviderIds.add(providerId); + this._logService.debug(`Created dynamic authentication provider: ${createdProviderId}`); + this._dynamicAuthenticationProviderIds.add(createdProviderId); return provider; } - this._logService.error(`Failed to create dynamic authentication provider: ${providerId}`); + this._logService.error(`Failed to create dynamic authentication provider: ${createdProviderId}`); return undefined; } diff --git a/src/vs/workbench/services/authentication/common/authentication.ts b/src/vs/workbench/services/authentication/common/authentication.ts index 9913769d1a2d33..5713132973d5f6 100644 --- a/src/vs/workbench/services/authentication/common/authentication.ts +++ b/src/vs/workbench/services/authentication/common/authentication.ts @@ -122,7 +122,7 @@ export interface AllowedExtension { export interface IAuthenticationProviderHostDelegate { /** Priority for this delegate, delegates are tested in descending priority order */ readonly priority: number; - create(authorizationServer: URI, serverMetadata: IAuthorizationServerMetadata, resource: IAuthorizationProtectedResourceMetadata | undefined, clientId?: string, clientSecret?: string): Promise; + create(authorizationServer: URI, serverMetadata: IAuthorizationServerMetadata, resource: IAuthorizationProtectedResourceMetadata | undefined, clientId?: string, clientSecret?: string, providerId?: string): Promise; /** * Creates an XAA (enterprise-managed, ID-JAG) authentication provider for the given SSO issuer. * The returned string is the provider id. @@ -257,10 +257,9 @@ export interface IAuthenticationService { registerAuthenticationProviderHostDelegate(delegate: IAuthenticationProviderHostDelegate): IDisposable; /** - * Creates a dynamic authentication provider for the given server metadata - * @param serverMetadata The metadata for the server that is being authenticated against + * Creates a dynamic authentication provider. An explicit provider ID scopes its stored registration and tokens. */ - createDynamicAuthenticationProvider(authorizationServer: URI, serverMetadata: IAuthorizationServerMetadata, resourceMetadata: IAuthorizationProtectedResourceMetadata | undefined, clientId?: string, clientSecret?: string): Promise; + createDynamicAuthenticationProvider(authorizationServer: URI, serverMetadata: IAuthorizationServerMetadata, resourceMetadata: IAuthorizationProtectedResourceMetadata | undefined, clientId?: string, clientSecret?: string, providerId?: string): Promise; /** * Gets or creates a built-in XAA (enterprise-managed, ID-JAG) authentication provider for the given diff --git a/src/vs/workbench/services/authentication/test/browser/authenticationQueryServiceMocks.ts b/src/vs/workbench/services/authentication/test/browser/authenticationQueryServiceMocks.ts index ffe125e498c487..311d398ff6e651 100644 --- a/src/vs/workbench/services/authentication/test/browser/authenticationQueryServiceMocks.ts +++ b/src/vs/workbench/services/authentication/test/browser/authenticationQueryServiceMocks.ts @@ -357,7 +357,7 @@ export class TestAuthenticationService extends BaseTestService implements IAuthe unregisterDeclaredAuthenticationProvider(): void { } unregisterAuthenticationProvider(): void { } registerAuthenticationProviderHostDelegate(): IDisposable { return { dispose: () => { } }; } - createDynamicAuthenticationProvider(): Promise { return Promise.resolve(undefined); } + createDynamicAuthenticationProvider(..._args: Parameters): Promise { return Promise.resolve(undefined); } createOrGetXaaProvider(): Promise { return Promise.resolve(undefined); } async requestNewSession(): Promise { return createSession(); } async getSession(): Promise { return createSession(); } From 12bff69c2c69be6c2923eff9925b99db8a05a3a1 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Tue, 20 Oct 2026 14:10:30 -0700 Subject: [PATCH 3/3] agentHost: keep targeted MCP credentials separate for reconnect replay The protocol client keeps each accepted `authenticate` request to replay it after a reconnect. The replay cache used only the resource and scopes as its key. Targeted MCP tokens for servers that share these values replaced each other and the provider credential, and a revocation for one server removed the cached credential of another server. - Add the authentication target (kind, session, and customization ID) to the cache key. Keep the key for untargeted credentials unchanged. - Add a test that keeps the provider credential and a targeted credential when a different target is revoked. Fixes https://github.com/microsoft/vscode/issues/342147 (Commit message generated by Copilot) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../browser/agentHostProtocolClient.ts | 4 +++- .../agentHostProtocolClient.test.ts | 23 ++++++++++++++++++- 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts b/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts index cabffa328e0017..fc9cc3d5bd461e 100644 --- a/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts +++ b/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts @@ -1999,7 +1999,9 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect } private _authenticationKey(params: AuthenticateParams): string { - return `${params.resource}\0${JSON.stringify(params.scopes ?? [])}`; + const key = `${params.resource}\0${JSON.stringify(params.scopes ?? [])}`; + // Targeted MCP credentials are independent of each other and of the provider-level credential. + return params.target ? `${key}\0${JSON.stringify([params.target.kind, params.target.session, params.target.customizationId])}` : key; } /** diff --git a/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts b/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts index 4b6e5aec95d66b..0e30b880b89c41 100644 --- a/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts +++ b/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts @@ -40,7 +40,7 @@ import { deriveLegacyDefaultChat } from '../../common/agentHostLegacyDefaultChat import { isActionKnownToVersion, PROTOCOL_VERSION } from '../../common/state/protocol/version/registry.js'; import { ActionType, type ActionEnvelope, type ChatTurnCompleteAction, type ChatTurnStartedAction, type SessionActiveClientSetAction, type SessionActiveClientRemovedAction, type SessionTitleChangedAction } from '../../common/state/sessionActions.js'; import { chatReducer } from '../../common/state/sessionReducers.js'; -import { ProtocolError, type AhpServerNotification, type JsonRpcNotification, type JsonRpcRequest, type JsonRpcResponse, type ProtocolMessage } from '../../common/state/sessionProtocol.js'; +import { AuthenticateTargetKind, ProtocolError, type AhpServerNotification, type JsonRpcNotification, type JsonRpcRequest, type JsonRpcResponse, type ProtocolMessage } from '../../common/state/sessionProtocol.js'; import { hasKey } from '../../../../base/common/types.js'; import { mainWindow } from '../../../../base/browser/window.js'; import { AUTOMATION_CATALOG_URI, buildChatUri, buildDefaultChatUri, createChatState, createSessionState, CustomizationType, MessageAttachmentKind, MessageKind, PendingMessageKind, readSessionExternal, readSessionWorkspaceless, ROOT_STATE_URI, SessionStatus, StateComponents, TurnState, customizationId, withSessionExternal, withSessionWorkspaceless, type ChatState, type Message, type SessionState } from '../../common/state/sessionState.js'; @@ -1188,6 +1188,27 @@ suite('AgentHostProtocolClient', () => { assert.deepStrictEqual([...client['_authentication'].values()], []); }); + test('retains targeted MCP authentication independently for reconnect replay', async () => { + const { client, transport } = createClient(); + const resource = 'https://mcp.slack.com'; + const target = (customizationId: string) => ({ kind: AuthenticateTargetKind.McpServer, session: 'other-agent://tenant/session', customizationId }); + const send = async (params: AuthenticateParams) => { + const pending = client.authenticate(params); + const request = transport.sentMessages[transport.sentMessages.length - 1] as JsonRpcRequest; + transport.fireMessage({ jsonrpc: '2.0', id: request.id, result: { authenticated: true } }); + await pending; + }; + await send({ resource, scopes: ['read'], token: 'provider-token' }); + await send({ resource, scopes: ['read'], token: 'first-token', target: target('first') }); + await send({ resource, scopes: ['read'], token: 'second-token', target: target('second') }); + await send({ resource, scopes: ['read'], token: '', target: target('first') }); + + assert.deepStrictEqual([...client['_authentication'].values()].map(entry => [entry.params.token, entry.params.target?.customizationId]), [ + ['provider-token', undefined], + ['second-token', 'second'], + ]); + }); + test('listSessions carries the workspace-less marker and compatible working directories', async () => { // Regression: the sessions provider resolves a session's kind (quick // chat vs. workspace) from `_meta.workspaceless`, and after a window