From e719715c9f64c4905021ce48ee9bb11ac39450f5 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Tue, 23 Jun 2026 20:31:31 -0700 Subject: [PATCH] agentHost: support multiple active clients per session Adopt the agent host protocol update that moves from a single active client per session to multiple active clients, and make the agent implementations model and merge them correctly. Protocol + UI/server: - SessionState.activeClient? -> activeClients[]; session/activeClientChanged is split into session/activeClientSet (upsert by clientId) and session/activeClientRemoved; session/activeClientToolsChanged carries a clientId. - The UI adds itself via session/activeClientSet and never removes itself. - The server removes a client from activeClients (and cancels its in-flight client tool calls) on unsubscribe, and on reconnect when the client does not resubscribe to a session where it was still active. Disconnect keeps the client active during the grace window; the grace timeout removes it and fails its pending tool calls if it never returns. IAgent implementations: - Replace setClientTools/setClientCustomizations with a per-client handle API: getOrCreateActiveClient(session, {clientId, displayName}) -> IActiveClient (mutable readonly-array tools/customizations accessors) and removeActiveClient(session, clientId). - Add shared node infra ActiveClientToolSet (per-session, clientId-keyed tool registry with merge-by-name + ownerOf) and adopt it in Copilot, Claude and Codex so multiple active clients' tools/customizations are stored, merged (deduped, first-inserted client wins) for the SDK, and tool calls are stamped to the owning client. - Guard teardown races: Copilot invalidates an in-flight customization sync on removeClient; Claude removes tools synchronously and serializes customization removal through the session sequencer, and prunes cached handles on session disposal. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../browser/remoteAgentHostProtocolClient.ts | 4 +- .../platform/agentHost/common/agentService.ts | 56 +++- .../common/state/protocol/.ahp-version | 2 +- .../state/protocol/action-origin.generated.ts | 11 +- .../protocol/channels-session/actions.ts | 49 ++- .../protocol/channels-session/commands.ts | 2 +- .../protocol/channels-session/reducer.ts | 40 ++- .../state/protocol/channels-session/state.ts | 18 +- .../common/state/protocol/common/actions.ts | 8 +- .../common/state/protocol/version/registry.ts | 3 +- .../agentHost/common/state/sessionActions.ts | 3 +- .../agentHost/common/state/sessionState.ts | 1 + .../agentHost/node/activeClientState.ts | 95 ++++++ .../node/agentHostTelemetryReporter.ts | 16 +- .../platform/agentHost/node/agentService.ts | 4 +- .../agentHost/node/agentSideEffects.ts | 26 +- .../agentHost/node/claude/claudeAgent.ts | 101 +++++- .../node/claude/claudeAgentSession.ts | 30 +- .../node/claude/claudeMapSessionEvents.ts | 8 +- .../node/claude/claudeSdkMessageRouter.ts | 12 +- .../agentHost/node/claude/claudeSdkOptions.ts | 4 +- .../node/claude/claudeSdkPipeline.ts | 14 +- .../claudeSessionClientToolsModel.ts | 114 ++++--- .../claudeSessionClientCustomizationsModel.ts | 39 ++- .../agentHost/node/codex/codexAgent.ts | 93 ++++-- .../node/codex/codexMapAppServerEvents.ts | 18 +- .../agentHost/node/copilot/copilotAgent.ts | 301 +++++++++++++----- .../node/copilot/copilotAgentSession.ts | 37 ++- .../node/copilot/copilotSessionLauncher.ts | 17 +- .../agentHost/node/protocolServerHandler.ts | 87 ++++- .../test/common/agentSubscription.test.ts | 1 + .../remoteAgentHostProtocolClient.test.ts | 24 +- .../agentHost/test/node/agentService.test.ts | 10 +- .../test/node/agentSideEffects.test.ts | 49 +-- .../agentHost/test/node/claudeAgent.test.ts | 46 +-- .../test/node/claudeMapSessionEvents.test.ts | 2 +- .../claudeSessionClientToolsModel.test.ts | 85 ++--- .../codex/codexMapAppServerEvents.test.ts | 13 +- .../agentHost/test/node/copilotAgent.test.ts | 66 +++- .../test/node/copilotAgentSession.test.ts | 48 +-- ...deSessionClientCustomizationsModel.test.ts | 45 ++- .../platform/agentHost/test/node/mockAgent.ts | 49 ++- .../protocol/codexRealSdk.integrationTest.ts | 4 +- .../test/node/protocol/realSdkTestHelpers.ts | 2 +- .../test/node/protocolServerHandler.test.ts | 186 ++++++++++- .../agentHost/test/node/reducers.test.ts | 1 + .../browser/baseAgentHostSessionsProvider.ts | 25 +- .../localAgentHostSessionsProvider.test.ts | 17 +- .../remoteAgentHostSessionsProvider.test.ts | 1 + .../agentHost/agentHostLocalCustomizations.ts | 4 +- .../agentHost/agentHostSessionHandler.ts | 29 +- .../agentHostChatContribution.test.ts | 56 ++-- .../agentHostClientTools.test.ts | 2 +- 53 files changed, 1393 insertions(+), 585 deletions(-) diff --git a/src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts b/src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts index f89891f13e3ab5..854606a862256f 100644 --- a/src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts +++ b/src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts @@ -904,11 +904,11 @@ export class RemoteAgentHostProtocolClient extends Disposable implements IAgentC /** * Inspect an outgoing client-dispatched action and grant implicit reads * for any customization URIs it carries. Today this covers - * `SessionActiveClientChanged`, which is the only client-dispatched + * `SessionActiveClientSet`, which is the only client-dispatched * action that ships customization URIs to the host. */ private _grantImplicitReadsForOutgoingAction(action: SessionAction | ChatAction | TerminalAction | ClientAnnotationsAction | IRootConfigChangedAction): void { - if (action.type === ActionType.SessionActiveClientChanged && action.activeClient?.customizations) { + if (action.type === ActionType.SessionActiveClientSet && action.activeClient.customizations) { this._grantImplicitReadsForCustomizations(action.activeClient.customizations); } } diff --git a/src/vs/platform/agentHost/common/agentService.ts b/src/vs/platform/agentHost/common/agentService.ts index f1d771be12fd23..7d5604931c5617 100644 --- a/src/vs/platform/agentHost/common/agentService.ts +++ b/src/vs/platform/agentHost/common/agentService.ts @@ -12,7 +12,6 @@ import type { IObservable } from '../../../base/common/observable.js'; import { URI } from '../../../base/common/uri.js'; import type { IConfigurationService } from '../../configuration/common/configuration.js'; import { createDecorator } from '../../instantiation/common/instantiation.js'; -import type { ISyncedCustomization } from './agentPluginManager.js'; import type { IAgentServerToolHost } from './agentServerTools.js'; import type { IActiveSubscriptionInfo, IAgentSubscription } from './state/agentSubscription.js'; import type { IRemoteWatchHandle } from './agentHostFileSystemProvider.js'; @@ -636,7 +635,7 @@ export interface IAgentCreateSessionConfig { /** * Eagerly claim the active client role for the new session. When provided, * the server initializes the session with this client as the active - * client, equivalent to dispatching a `session/activeClientChanged` + * client, equivalent to dispatching a `session/activeClientSet` * action immediately after creation. The `clientId` MUST match the * connection's own `clientId`. */ @@ -841,6 +840,30 @@ export interface IMcpNotification { readonly params?: Record; } +/** + * A per-session handle for one active client's contributions (tools and + * plugin customizations) to an agent session, obtained via + * {@link IAgent.getOrCreateActiveClient}. + * + * `tools` and `customizations` are mutable accessor properties: assigning a + * new array replaces this client's contribution wholesale and triggers the + * agent's internal reaction (refreshing the merged tool set exposed to the + * model, or kicking off an asynchronous customization sync). The arrays are + * `readonly` so callers cannot mutate them in place and silently bypass the + * setter. The agent merges the contributions of all active clients on a + * session, deduplicating as needed. + */ +export interface IActiveClient { + /** Client identifier (matches `clientId` from `initialize`). */ + readonly clientId: string; + /** Human-readable client name (e.g. `"VS Code"`), if provided. */ + readonly displayName: string | undefined; + /** This client's tools. Assigning replaces the set (full replacement). */ + tools: readonly ToolDefinition[]; + /** This client's plugin customizations. Assigning replaces the set and starts an internal sync. */ + customizations: readonly ClientPluginCustomization[]; +} + /** * Implemented by each agent backend (e.g. Copilot SDK). * The {@link IAgentService} dispatches to the appropriate agent based on @@ -1016,27 +1039,26 @@ export interface IAgent { onArchivedChanged?(session: URI, isArchived: boolean): Promise; /** - * Receives client-provided customization refs for a session and syncs them - * (e.g. copies plugin files to local storage). The agent publishes - * customization state actions as the sync progresses. + * Get (or lazily create) the per-session handle for an active client, + * identified by `clientId`. Mutating the returned {@link IActiveClient}'s + * `tools` / `customizations` updates only that client's contribution; the + * agent merges the contributions of all active clients when exposing them + * to the model. A session MAY have several active clients at once. * - * The agent MAY defer a client restart until all active sessions are idle. + * @param session The session URI this client contributes to. + * @param client The client's `clientId` and optional human-readable name. */ - setClientCustomizations(session: URI, clientId: string, customizations: ClientPluginCustomization[]): Promise; + getOrCreateActiveClient(session: URI, client: { readonly clientId: string; readonly displayName?: string }): IActiveClient; /** - * Receives client-provided tool definitions to make available in a - * specific session. The agent registers these as custom tools so the - * LLM can call them; execution is routed back to the owning client. - * - * Always called on `activeClientChanged`, even with an empty array, - * to clear a previous client's tools. + * Remove an active client from a session, clearing its tool and + * customization contributions. No-op when no active client matches + * `clientId`. * - * @param session The session URI this tool set applies to. - * @param clientId The client that owns these tools. - * @param tools The tool definitions (full replacement). + * @param session The session the client is leaving. + * @param clientId The client to remove. */ - setClientTools(session: URI, clientId: string | undefined, tools: ToolDefinition[]): void; + removeActiveClient(session: URI, clientId: string): void; /** * Called when a client completes a client-provided tool call. diff --git a/src/vs/platform/agentHost/common/state/protocol/.ahp-version b/src/vs/platform/agentHost/common/state/protocol/.ahp-version index a486d1ec6eb203..ec859badb2194d 100644 --- a/src/vs/platform/agentHost/common/state/protocol/.ahp-version +++ b/src/vs/platform/agentHost/common/state/protocol/.ahp-version @@ -1 +1 @@ -77c6312 +0259a7e diff --git a/src/vs/platform/agentHost/common/state/protocol/action-origin.generated.ts b/src/vs/platform/agentHost/common/state/protocol/action-origin.generated.ts index ebf7a03c9c2663..acd1a33760d04a 100644 --- a/src/vs/platform/agentHost/common/state/protocol/action-origin.generated.ts +++ b/src/vs/platform/agentHost/common/state/protocol/action-origin.generated.ts @@ -9,7 +9,7 @@ // Generated from types/actions.ts — do not edit // Run `npm run generate` to regenerate. -import { ActionType, type StateAction, type RootAgentsChangedAction, type RootActiveSessionsChangedAction, type RootTerminalsChangedAction, type RootConfigChangedAction, type SessionReadyAction, type SessionCreationFailedAction, type SessionChatAddedAction, type SessionChatRemovedAction, type SessionChatUpdatedAction, type SessionDefaultChatChangedAction, type SessionTitleChangedAction, type SessionModelChangedAction, type SessionAgentChangedAction, type SessionServerToolsChangedAction, type SessionActiveClientChangedAction, type SessionActiveClientToolsChangedAction, type SessionCustomizationsChangedAction, type SessionCustomizationToggledAction, type SessionCustomizationUpdatedAction, type SessionCustomizationRemovedAction, type SessionMcpServerStateChangedAction, type SessionIsReadChangedAction, type SessionIsArchivedChangedAction, type SessionActivityChangedAction, type SessionChangesetsChangedAction, type SessionConfigChangedAction, type SessionMetaChangedAction, type ChatTurnStartedAction, type ChatDeltaAction, type ChatResponsePartAction, type ChatToolCallStartAction, type ChatToolCallDeltaAction, type ChatToolCallReadyAction, type ChatToolCallConfirmedAction, type ChatToolCallCompleteAction, type ChatToolCallResultConfirmedAction, type ChatToolCallContentChangedAction, type ChatTurnCompleteAction, type ChatTurnCancelledAction, type ChatErrorAction, type ChatUsageAction, type ChatReasoningAction, type ChatPendingMessageSetAction, type ChatPendingMessageRemovedAction, type ChatQueuedMessagesReorderedAction, type ChatInputRequestedAction, type ChatInputAnswerChangedAction, type ChatInputCompletedAction, type ChatTruncatedAction, type ChangesetStatusChangedAction, type ChangesetFileSetAction, type ChangesetFileRemovedAction, type ChangesetContentChangedAction, type ChangesetOperationsChangedAction, type ChangesetOperationStatusChangedAction, type ChangesetClearedAction, type AnnotationsSetAction, type AnnotationsUpdatedAction, type AnnotationsRemovedAction, type AnnotationsEntrySetAction, type AnnotationsEntryRemovedAction, type TerminalDataAction, type TerminalInputAction, type TerminalResizedAction, type TerminalClaimedAction, type TerminalTitleChangedAction, type TerminalCwdChangedAction, type TerminalExitedAction, type TerminalClearedAction, type TerminalCommandDetectionAvailableAction, type TerminalCommandExecutedAction, type TerminalCommandFinishedAction, type ResourceWatchChangedAction } from './actions.js'; +import { ActionType, type StateAction, type RootAgentsChangedAction, type RootActiveSessionsChangedAction, type RootTerminalsChangedAction, type RootConfigChangedAction, type SessionReadyAction, type SessionCreationFailedAction, type SessionChatAddedAction, type SessionChatRemovedAction, type SessionChatUpdatedAction, type SessionDefaultChatChangedAction, type SessionTitleChangedAction, type SessionModelChangedAction, type SessionAgentChangedAction, type SessionServerToolsChangedAction, type SessionActiveClientSetAction, type SessionActiveClientRemovedAction, type SessionActiveClientToolsChangedAction, type SessionCustomizationsChangedAction, type SessionCustomizationToggledAction, type SessionCustomizationUpdatedAction, type SessionCustomizationRemovedAction, type SessionMcpServerStateChangedAction, type SessionIsReadChangedAction, type SessionIsArchivedChangedAction, type SessionActivityChangedAction, type SessionChangesetsChangedAction, type SessionConfigChangedAction, type SessionMetaChangedAction, type ChatTurnStartedAction, type ChatDeltaAction, type ChatResponsePartAction, type ChatToolCallStartAction, type ChatToolCallDeltaAction, type ChatToolCallReadyAction, type ChatToolCallConfirmedAction, type ChatToolCallCompleteAction, type ChatToolCallResultConfirmedAction, type ChatToolCallContentChangedAction, type ChatTurnCompleteAction, type ChatTurnCancelledAction, type ChatErrorAction, type ChatUsageAction, type ChatReasoningAction, type ChatPendingMessageSetAction, type ChatPendingMessageRemovedAction, type ChatQueuedMessagesReorderedAction, type ChatInputRequestedAction, type ChatInputAnswerChangedAction, type ChatInputCompletedAction, type ChatTruncatedAction, type ChangesetStatusChangedAction, type ChangesetFileSetAction, type ChangesetFileRemovedAction, type ChangesetContentChangedAction, type ChangesetOperationsChangedAction, type ChangesetOperationStatusChangedAction, type ChangesetClearedAction, type AnnotationsSetAction, type AnnotationsUpdatedAction, type AnnotationsRemovedAction, type AnnotationsEntrySetAction, type AnnotationsEntryRemovedAction, type TerminalDataAction, type TerminalInputAction, type TerminalResizedAction, type TerminalClaimedAction, type TerminalTitleChangedAction, type TerminalCwdChangedAction, type TerminalExitedAction, type TerminalClearedAction, type TerminalCommandDetectionAvailableAction, type TerminalCommandExecutedAction, type TerminalCommandFinishedAction, type ResourceWatchChangedAction } from './actions.js'; // ─── Root vs Session vs Chat vs Terminal vs Changeset Action Unions ───────────────── @@ -46,7 +46,8 @@ export type SessionAction = | SessionModelChangedAction | SessionAgentChangedAction | SessionServerToolsChangedAction - | SessionActiveClientChangedAction + | SessionActiveClientSetAction + | SessionActiveClientRemovedAction | SessionActiveClientToolsChangedAction | SessionCustomizationsChangedAction | SessionCustomizationToggledAction @@ -66,7 +67,8 @@ export type ClientSessionAction = | SessionTitleChangedAction | SessionModelChangedAction | SessionAgentChangedAction - | SessionActiveClientChangedAction + | SessionActiveClientSetAction + | SessionActiveClientRemovedAction | SessionActiveClientToolsChangedAction | SessionCustomizationToggledAction | SessionIsReadChangedAction @@ -268,7 +270,8 @@ export const IS_CLIENT_DISPATCHABLE: { readonly [K in StateAction['type']]: bool [ActionType.SessionModelChanged]: true, [ActionType.SessionAgentChanged]: true, [ActionType.SessionServerToolsChanged]: false, - [ActionType.SessionActiveClientChanged]: true, + [ActionType.SessionActiveClientSet]: true, + [ActionType.SessionActiveClientRemoved]: true, [ActionType.SessionActiveClientToolsChanged]: true, [ActionType.SessionCustomizationsChanged]: false, [ActionType.SessionCustomizationToggled]: true, diff --git a/src/vs/platform/agentHost/common/state/protocol/channels-session/actions.ts b/src/vs/platform/agentHost/common/state/protocol/channels-session/actions.ts index 5fa60999e29620..88600a10caec3a 100644 --- a/src/vs/platform/agentHost/common/state/protocol/channels-session/actions.ts +++ b/src/vs/platform/agentHost/common/state/protocol/channels-session/actions.ts @@ -238,29 +238,50 @@ export interface SessionServerToolsChangedAction { } /** - * The active client for this session has changed. + * An active client for this session was added or updated. * - * A client dispatches this action with its own `SessionActiveClient` to claim - * the active role, or with `null` to release it. The server SHOULD reject if - * another client is already active. The server SHOULD automatically dispatch - * this action with `activeClient: null` when the active client disconnects. + * Upsert semantics keyed by {@link SessionActiveClient.clientId | `clientId`}: + * a client dispatches this action with its own `SessionActiveClient` to claim + * the active role or refresh its entry, replacing any existing entry that has + * the same `clientId`. Multiple clients may be active at once. Use + * {@link SessionActiveClientRemovedAction | `session/activeClientRemoved`} to + * release the role. The server SHOULD automatically dispatch that removal when + * an active client disconnects. * * @category Session Actions * @version 1 * @clientDispatchable */ -export interface SessionActiveClientChangedAction { - type: ActionType.SessionActiveClientChanged; - /** The new active client, or `null` to unset */ - activeClient: SessionActiveClient | null; +export interface SessionActiveClientSetAction { + type: ActionType.SessionActiveClientSet; + /** The active client to add or update, matched by `clientId`. */ + activeClient: SessionActiveClient; } /** - * The active client's tool list has changed. + * An active client was removed from this session. * - * Full-replacement semantics: the `tools` array replaces the active client's - * previous tools entirely. The server SHOULD reject if the dispatching client - * is not the current active client. + * Releases the active role for the client identified by `clientId`. No-op when + * no active client matches. The server SHOULD dispatch this automatically when + * an active client disconnects. + * + * @category Session Actions + * @version 1 + * @clientDispatchable + */ +export interface SessionActiveClientRemovedAction { + type: ActionType.SessionActiveClientRemoved; + /** The `clientId` of the active client to remove. */ + clientId: string; +} + +/** + * An active client's tool list has changed. + * + * Full-replacement semantics: the `tools` array replaces the named active + * client's previous tools entirely. The active client is identified by + * `clientId`; the action is a no-op when no active client matches. The server + * SHOULD reject if the dispatching client is not the named active client. * * @category Session Actions * @version 1 @@ -268,6 +289,8 @@ export interface SessionActiveClientChangedAction { */ export interface SessionActiveClientToolsChangedAction { type: ActionType.SessionActiveClientToolsChanged; + /** The `clientId` of the active client whose tools changed. */ + clientId: string; /** Updated client tools list (full replacement) */ tools: ToolDefinition[]; } diff --git a/src/vs/platform/agentHost/common/state/protocol/channels-session/commands.ts b/src/vs/platform/agentHost/common/state/protocol/channels-session/commands.ts index 291c23314c929e..268856b57cc2e9 100644 --- a/src/vs/platform/agentHost/common/state/protocol/channels-session/commands.ts +++ b/src/vs/platform/agentHost/common/state/protocol/channels-session/commands.ts @@ -88,7 +88,7 @@ export interface CreateSessionParams extends BaseParams { * Eagerly claim the active client role for the new session. * * When provided, the server initializes the session with this client as the - * active client, equivalent to dispatching a `session/activeClientChanged` + * active client, equivalent to dispatching a `session/activeClientSet` * action immediately after creation. The `clientId` MUST match the * `clientId` the creating client supplied in `initialize`. */ diff --git a/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts b/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts index 8b75efaeff3ea4..3fee0f7d0b45ac 100644 --- a/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts +++ b/src/vs/platform/agentHost/common/state/protocol/channels-session/reducer.ts @@ -153,20 +153,38 @@ export function sessionReducer(state: SessionState, action: SessionAction, log?: case ActionType.SessionServerToolsChanged: return { ...state, serverTools: action.tools }; - case ActionType.SessionActiveClientChanged: - return { - ...state, - activeClient: action.activeClient ?? undefined, - }; + case ActionType.SessionActiveClientSet: { + const list = state.activeClients; + const idx = list.findIndex(c => c.clientId === action.activeClient.clientId); + if (idx < 0) { + return { ...state, activeClients: [...list, action.activeClient] }; + } + const updated = list.slice(); + updated[idx] = action.activeClient; + return { ...state, activeClients: updated }; + } - case ActionType.SessionActiveClientToolsChanged: - if (!state.activeClient) { + case ActionType.SessionActiveClientRemoved: { + const list = state.activeClients; + const idx = list.findIndex(c => c.clientId === action.clientId); + if (idx < 0) { return state; } - return { - ...state, - activeClient: { ...state.activeClient, tools: action.tools }, - }; + const updated = list.slice(); + updated.splice(idx, 1); + return { ...state, activeClients: updated }; + } + + case ActionType.SessionActiveClientToolsChanged: { + const list = state.activeClients; + const idx = list.findIndex(c => c.clientId === action.clientId); + if (idx < 0) { + return state; + } + const updated = list.slice(); + updated[idx] = { ...updated[idx], tools: action.tools }; + return { ...state, activeClients: updated }; + } // ── Customizations ────────────────────────────────────────────────── diff --git a/src/vs/platform/agentHost/common/state/protocol/channels-session/state.ts b/src/vs/platform/agentHost/common/state/protocol/channels-session/state.ts index 25d0f865d63df5..d02d3f802685d4 100644 --- a/src/vs/platform/agentHost/common/state/protocol/channels-session/state.ts +++ b/src/vs/platform/agentHost/common/state/protocol/channels-session/state.ts @@ -63,8 +63,13 @@ export interface SessionState { creationError?: ErrorInfo; /** Tools provided by the server (agent host) for this session */ serverTools?: ToolDefinition[]; - /** The client currently providing tools and interactive capabilities to this session */ - activeClient?: SessionActiveClient; + /** + * The clients currently providing tools and interactive capabilities to this + * session. If multiple tools or customizations are provided by the same + * active client, an agent host MAY deduplicate them when exposed to a model, + * with a preference given to the client that started the turn. + */ + activeClients: SessionActiveClient[]; /** Catalog of chats in this session. */ chats: ChatSummary[]; /** @@ -91,7 +96,7 @@ export interface SessionState { * also appear as children of a container. * * Client-published plugins arrive via - * {@link SessionActiveClient.customizations | `activeClient.customizations`} + * {@link SessionActiveClient.customizations | `activeClients[].customizations`} * and the host propagates them into this list (typically with the * container's `clientId` set and `children` populated). Clients * publish in container shape only; bare MCP servers at the top level @@ -117,10 +122,11 @@ export interface SessionState { } /** - * The client currently providing tools and interactive capabilities to a session. + * A client currently providing tools and interactive capabilities to a session. * - * Only one client may be active per session at a time. The server SHOULD - * automatically unset the active client if that client disconnects. + * A session MAY have several active clients at once; entries in + * {@link SessionState.activeClients} are keyed by `clientId`. The server SHOULD + * automatically remove an active client when that client disconnects. * * @category Session State */ diff --git a/src/vs/platform/agentHost/common/state/protocol/common/actions.ts b/src/vs/platform/agentHost/common/state/protocol/common/actions.ts index 36932e3c5bc4f2..382e47ca51261a 100644 --- a/src/vs/platform/agentHost/common/state/protocol/common/actions.ts +++ b/src/vs/platform/agentHost/common/state/protocol/common/actions.ts @@ -10,7 +10,7 @@ import type { URI } from './state.js'; import type { RootAgentsChangedAction, RootActiveSessionsChangedAction, RootTerminalsChangedAction, RootConfigChangedAction } from '../channels-root/actions.js'; -import type { SessionReadyAction, SessionCreationFailedAction, SessionChatAddedAction, SessionChatRemovedAction, SessionChatUpdatedAction, SessionDefaultChatChangedAction, SessionTitleChangedAction, SessionModelChangedAction, SessionAgentChangedAction, SessionServerToolsChangedAction, SessionActiveClientChangedAction, SessionActiveClientToolsChangedAction, SessionCustomizationsChangedAction, SessionCustomizationToggledAction, SessionCustomizationUpdatedAction, SessionCustomizationRemovedAction, SessionMcpServerStateChangedAction, SessionIsReadChangedAction, SessionIsArchivedChangedAction, SessionActivityChangedAction, SessionChangesetsChangedAction, SessionConfigChangedAction, SessionMetaChangedAction } from '../channels-session/actions.js'; +import type { SessionReadyAction, SessionCreationFailedAction, SessionChatAddedAction, SessionChatRemovedAction, SessionChatUpdatedAction, SessionDefaultChatChangedAction, SessionTitleChangedAction, SessionModelChangedAction, SessionAgentChangedAction, SessionServerToolsChangedAction, SessionActiveClientSetAction, SessionActiveClientRemovedAction, SessionActiveClientToolsChangedAction, SessionCustomizationsChangedAction, SessionCustomizationToggledAction, SessionCustomizationUpdatedAction, SessionCustomizationRemovedAction, SessionMcpServerStateChangedAction, SessionIsReadChangedAction, SessionIsArchivedChangedAction, SessionActivityChangedAction, SessionChangesetsChangedAction, SessionConfigChangedAction, SessionMetaChangedAction } from '../channels-session/actions.js'; import type { ChatTurnStartedAction, ChatDeltaAction, ChatResponsePartAction, ChatToolCallStartAction, ChatToolCallDeltaAction, ChatToolCallReadyAction, ChatToolCallConfirmedAction, ChatToolCallCompleteAction, ChatToolCallResultConfirmedAction, ChatToolCallContentChangedAction, ChatTurnCompleteAction, ChatTurnCancelledAction, ChatErrorAction, ChatUsageAction, ChatReasoningAction, ChatPendingMessageSetAction, ChatPendingMessageRemovedAction, ChatQueuedMessagesReorderedAction, ChatInputRequestedAction, ChatInputAnswerChangedAction, ChatInputCompletedAction, ChatTruncatedAction } from '../channels-chat/actions.js'; @@ -57,7 +57,8 @@ export const enum ActionType { SessionModelChanged = 'session/modelChanged', SessionAgentChanged = 'session/agentChanged', SessionServerToolsChanged = 'session/serverToolsChanged', - SessionActiveClientChanged = 'session/activeClientChanged', + SessionActiveClientSet = 'session/activeClientSet', + SessionActiveClientRemoved = 'session/activeClientRemoved', SessionActiveClientToolsChanged = 'session/activeClientToolsChanged', ChatPendingMessageSet = 'chat/pendingMessageSet', ChatPendingMessageRemoved = 'chat/pendingMessageRemoved', @@ -153,7 +154,8 @@ export type StateAction = | SessionModelChangedAction | SessionAgentChangedAction | SessionServerToolsChangedAction - | SessionActiveClientChangedAction + | SessionActiveClientSetAction + | SessionActiveClientRemovedAction | SessionActiveClientToolsChangedAction | SessionCustomizationsChangedAction | SessionCustomizationToggledAction diff --git a/src/vs/platform/agentHost/common/state/protocol/version/registry.ts b/src/vs/platform/agentHost/common/state/protocol/version/registry.ts index 3cb11e65c52de4..d7bb394ea74c2f 100644 --- a/src/vs/platform/agentHost/common/state/protocol/version/registry.ts +++ b/src/vs/platform/agentHost/common/state/protocol/version/registry.ts @@ -90,7 +90,8 @@ export const ACTION_INTRODUCED_IN: { readonly [K in StateAction['type']]: string [ActionType.SessionModelChanged]: '0.1.0', [ActionType.SessionAgentChanged]: '0.2.0', [ActionType.SessionServerToolsChanged]: '0.1.0', - [ActionType.SessionActiveClientChanged]: '0.1.0', + [ActionType.SessionActiveClientSet]: '0.5.0', + [ActionType.SessionActiveClientRemoved]: '0.5.0', [ActionType.SessionActiveClientToolsChanged]: '0.1.0', [ActionType.SessionCustomizationsChanged]: '0.1.0', [ActionType.SessionCustomizationToggled]: '0.1.0', diff --git a/src/vs/platform/agentHost/common/state/sessionActions.ts b/src/vs/platform/agentHost/common/state/sessionActions.ts index b82768f8aeb757..a8b72183236896 100644 --- a/src/vs/platform/agentHost/common/state/sessionActions.ts +++ b/src/vs/platform/agentHost/common/state/sessionActions.ts @@ -45,7 +45,8 @@ export { type ChatUsageAction, type SessionAgentChangedAction, type SessionServerToolsChangedAction, - type SessionActiveClientChangedAction, + type SessionActiveClientSetAction, + type SessionActiveClientRemovedAction, type SessionActiveClientToolsChangedAction, type SessionCustomizationsChangedAction, type SessionCustomizationToggledAction, diff --git a/src/vs/platform/agentHost/common/state/sessionState.ts b/src/vs/platform/agentHost/common/state/sessionState.ts index 7f4387ff458b6c..29755aa48b5d35 100644 --- a/src/vs/platform/agentHost/common/state/sessionState.ts +++ b/src/vs/platform/agentHost/common/state/sessionState.ts @@ -475,6 +475,7 @@ export function createSessionState(summary: SessionSummary): SessionState { return { summary, lifecycle: SessionLifecycle.Creating, + activeClients: [], chats: [], defaultChat: undefined, }; diff --git a/src/vs/platform/agentHost/node/activeClientState.ts b/src/vs/platform/agentHost/node/activeClientState.ts index 444f018b08cd77..b94c5bf8bd089c 100644 --- a/src/vs/platform/agentHost/node/activeClientState.ts +++ b/src/vs/platform/agentHost/node/activeClientState.ts @@ -52,6 +52,101 @@ export function structuralToolsEqual( return true; } +/** + * A per-session registry of the tools contributed by each active client, + * keyed by `clientId` and kept in insertion order. Backs the multi-active-client + * tool model shared by the agent-host providers (Copilot, Claude, Codex): + * each provider stores one of these per session and exposes the + * {@link merged} view to its SDK while routing tool calls back to the + * {@link ownerOf | owning client}. + * + * Deduplication of {@link merged} is by tool `name`, first-inserted-client + * wins, so the merged order and the owner of any given tool name are + * deterministic regardless of how many clients contribute it. + */ +export class ActiveClientToolSet { + private readonly _byClient = new Map(); + + /** Number of clients currently contributing tools. */ + get size(): number { + return this._byClient.size; + } + + /** Whether `clientId` currently contributes tools. */ + has(clientId: string): boolean { + return this._byClient.has(clientId); + } + + /** The client ids currently contributing tools, in insertion order. */ + clientIds(): IterableIterator { + return this._byClient.keys(); + } + + /** This client's contributed tools, or an empty array when absent. */ + get(clientId: string): readonly ToolDefinition[] { + return this._byClient.get(clientId) ?? []; + } + + /** + * Replace `clientId`'s contributed tools (full replacement). A new + * `clientId` is appended after existing ones; re-setting an existing + * `clientId` keeps its insertion position so merged ordering and tool + * ownership stay stable across updates. + */ + set(clientId: string, tools: readonly ToolDefinition[]): void { + this._byClient.set(clientId, tools); + } + + /** Remove `clientId`'s contribution. Returns whether anything was removed. */ + delete(clientId: string): boolean { + return this._byClient.delete(clientId); + } + + /** + * The union of every client's tools, deduplicated by `name` with the + * first-inserted contributor winning. Order follows client insertion + * order, then per-client tool order. + */ + merged(): readonly ToolDefinition[] { + const seen = new Set(); + const result: ToolDefinition[] = []; + for (const tools of this._byClient.values()) { + for (const tool of tools) { + if (seen.has(tool.name)) { + continue; + } + seen.add(tool.name); + result.push(tool); + } + } + return result; + } + + /** + * The `clientId` that owns the merged tool named `toolName` (the + * first-inserted contributor of that name), or `undefined` when no active + * client provides it. Used to stamp client tool calls with their owning + * client at invocation time. + */ + ownerOf(toolName: string): string | undefined { + for (const [clientId, tools] of this._byClient) { + if (tools.some(tool => tool.name === toolName)) { + return clientId; + } + } + return undefined; + } + + /** + * Structural comparison of the current {@link merged} tools against a + * previously-applied snapshot (`name + description + inputSchema`, + * order-insensitive). Returns `true` when no SDK restart is required. + */ + structuralEquals(applied: readonly ToolDefinition[] | undefined): boolean { + return structuralToolsEqual(this.merged(), applied); + } +} + /** * Live, mutable holder for the active client's identity (`clientId`) and the * structural tool snapshot it contributes. Shared between the Copilot and diff --git a/src/vs/platform/agentHost/node/agentHostTelemetryReporter.ts b/src/vs/platform/agentHost/node/agentHostTelemetryReporter.ts index 6a1cadce7ecb99..66c259290ec259 100644 --- a/src/vs/platform/agentHost/node/agentHostTelemetryReporter.ts +++ b/src/vs/platform/agentHost/node/agentHostTelemetryReporter.ts @@ -28,9 +28,9 @@ export type IAgentHostUserMessageSentClassification = { source: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether the user message was sent directly or from the queued-message flow.' }; isSubagentSession: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'Whether the message was sent to a subagent session.' }; turnCount: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The number of completed turns in the session when the message was sent.' }; - activeClientId?: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The active client identifier for the session, if any.' }; - activeClientToolCount?: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The number of tools provided by the active client, if any.' }; - activeClientCustomizationCount?: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The number of customizations provided by the active client, if any.' }; + activeClientId?: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The identifier of the first active client for the session, if any.' }; + activeClientToolCount?: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The total number of tools provided by the active clients, if any.' }; + activeClientCustomizationCount?: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The total number of customizations provided by the active clients, if any.' }; attachmentCount: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; isMeasurement: true; comment: 'The number of attachments included with the user message.' }; owner: 'roblourens'; comment: 'Tracks user messages sent from the agent host process to an agent provider.'; @@ -76,17 +76,17 @@ export class AgentHostTelemetryReporter { userMessageSent(provider: string, session: string, sessionState: ISessionWithDefaultChat | undefined, source: AgentHostUserMessageSentSource, attachments: readonly MessageAttachment[] | undefined): void { const attachmentCount = attachments?.length ?? 0; - const activeClient = sessionState?.activeClient; + const activeClients = sessionState?.activeClients ?? []; this._telemetryService.publicLog2('agentHost.userMessageSent', { provider, agentSessionId: AgentSession.id(session), source, isSubagentSession: isSubagentSession(session), turnCount: sessionState?.turns.length ?? 0, - ...(activeClient ? { - activeClientId: activeClient.clientId, - activeClientToolCount: activeClient.tools.length, - activeClientCustomizationCount: activeClient.customizations?.length ?? 0, + ...(activeClients.length > 0 ? { + activeClientId: activeClients[0].clientId, + activeClientToolCount: activeClients.reduce((sum, client) => sum + client.tools.length, 0), + activeClientCustomizationCount: activeClients.reduce((sum, client) => sum + (client.customizations?.length ?? 0), 0), } : {}), attachmentCount, }); diff --git a/src/vs/platform/agentHost/node/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts index 57d7d82fbe210a..5075b90fa7dbcf 100644 --- a/src/vs/platform/agentHost/node/agentService.ts +++ b/src/vs/platform/agentHost/node/agentService.ts @@ -640,7 +640,7 @@ export class AgentService extends Disposable implements IAgentService { const state = this._stateManager.createSession(summary); state.config = sessionConfig; this._stateManager.seedDefaultChatTurns(summary.resource, sourceTurns); - state.activeClient = config.activeClient; + state.activeClients = config.activeClient ? [config.activeClient] : []; if (initialCustomizations && initialCustomizations.length > 0) { state.customizations = [...initialCustomizations]; } @@ -654,7 +654,7 @@ export class AgentService extends Disposable implements IAgentService { const summary = this._buildInitialSummary(provider, session, config, created, ''); const state = this._stateManager.createSession(summary, { emitNotification: !created.provisional }); state.config = sessionConfig; - state.activeClient = config?.activeClient; + state.activeClients = config?.activeClient ? [config.activeClient] : []; if (initialCustomizations && initialCustomizations.length > 0) { state.customizations = [...initialCustomizations]; } diff --git a/src/vs/platform/agentHost/node/agentSideEffects.ts b/src/vs/platform/agentHost/node/agentSideEffects.ts index 0cb3f7fa50f59d..9ee1c43f203583 100644 --- a/src/vs/platform/agentHost/node/agentSideEffects.ts +++ b/src/vs/platform/agentHost/node/agentSideEffects.ts @@ -913,19 +913,23 @@ export class AgentSideEffects extends Disposable { this._changesets.onSessionTruncated(channel); break; } - case ActionType.SessionActiveClientChanged: { + case ActionType.SessionActiveClientSet: { const agent = this._options.getAgent(channel); if (!agent) { break; } - // Always forward client tools, even if empty, to clear previous client's tools - const clientId = action.activeClient?.clientId; - agent.setClientTools(URI.parse(channel), clientId, action.activeClient?.tools ?? []); - - const refs = action.activeClient?.customizations ?? []; - agent.setClientCustomizations(URI.parse(channel), clientId ?? '', refs).catch(err => { - this._logService.error('[AgentSideEffects] setClientCustomizations failed', err); + const activeClient = action.activeClient; + const handle = agent.getOrCreateActiveClient(URI.parse(channel), { + clientId: activeClient.clientId, + displayName: activeClient.displayName, }); + handle.tools = activeClient.tools; + handle.customizations = activeClient.customizations ?? []; + break; + } + case ActionType.SessionActiveClientRemoved: { + const agent = this._options.getAgent(channel); + agent?.removeActiveClient(URI.parse(channel), action.clientId); break; } case ActionType.RootConfigChanged: { @@ -942,9 +946,9 @@ export class AgentSideEffects extends Disposable { const agent = this._options.getAgent(channel); if (agent) { const sessionState = this._stateManager.getSessionState(channel); - const toolClientId = sessionState?.activeClient?.clientId; - if (toolClientId) { - agent.setClientTools(URI.parse(channel), toolClientId, action.tools); + const isActiveClient = sessionState?.activeClients.some(c => c.clientId === action.clientId); + if (isActiveClient) { + agent.getOrCreateActiveClient(URI.parse(channel), { clientId: action.clientId }).tools = action.tools; } } break; diff --git a/src/vs/platform/agentHost/node/claude/claudeAgent.ts b/src/vs/platform/agentHost/node/claude/claudeAgent.ts index 5ca688efdb6149..326c7d36ee5373 100644 --- a/src/vs/platform/agentHost/node/claude/claudeAgent.ts +++ b/src/vs/platform/agentHost/node/claude/claudeAgent.ts @@ -23,7 +23,7 @@ import { createSchema, platformSessionSchema, schemaProperty } from '../../commo import { ClaudePermissionMode, ClaudeSessionConfigKey, narrowClaudePermissionMode } from '../../common/claudeSessionConfigKeys.js'; import { createClaudeThinkingLevelSchema, isClaudeEffortLevel } from '../../common/claudeModelConfig.js'; import { SessionConfigKey } from '../../common/sessionConfigKeys.js'; -import { AgentProvider, AgentSession, AgentSignal, GITHUB_COPILOT_PROTECTED_RESOURCE, GITHUB_REPO_PROTECTED_RESOURCE, IAgent, IAgentCreateSessionConfig, IAgentCreateSessionResult, IAgentDescriptor, IAgentMaterializeSessionEvent, IAgentModelInfo, IAgentResolveSessionConfigParams, IAgentSessionConfigCompletionsParams, IAgentSessionMetadata, IAgentSessionProjectInfo } from '../../common/agentService.js'; +import { AgentProvider, AgentSession, AgentSignal, GITHUB_COPILOT_PROTECTED_RESOURCE, GITHUB_REPO_PROTECTED_RESOURCE, IActiveClient, IAgent, IAgentCreateSessionConfig, IAgentCreateSessionResult, IAgentDescriptor, IAgentMaterializeSessionEvent, IAgentModelInfo, IAgentResolveSessionConfigParams, IAgentSessionConfigCompletionsParams, IAgentSessionMetadata, IAgentSessionProjectInfo } from '../../common/agentService.js'; import { ActionType } from '../../common/state/sessionActions.js'; import type { ResolveSessionConfigResult, SessionConfigCompletionsResult } from '../../common/state/protocol/commands.js'; import { AHP_AUTH_REQUIRED, ProtocolError } from '../../common/state/sessionProtocol.js'; @@ -123,6 +123,40 @@ function toAgentModelInfo(m: CCAModel, provider: AgentProvider): IAgentModelInfo // provisionalConfig). The legacy `IClaudeProvisionalSession` map shape // was retired in Phase 10.5 Step 3a. +/** + * Claude active-client handle. Tools read/write through the live session's + * {@link SessionClientToolsModel}; customization assignment kicks off the + * agent's async sync (via the provided closure). The handle caches the last + * assigned customization inputs so the getter reflects what the client most + * recently published. + */ +class ClaudeActiveClientHandle implements IActiveClient { + private _customizations: readonly ClientPluginCustomization[] = []; + + constructor( + readonly clientId: string, + readonly displayName: string | undefined, + private readonly _getTools: () => readonly ToolDefinition[], + private readonly _setTools: (tools: readonly ToolDefinition[]) => void, + private readonly _syncCustomizations: (customizations: readonly ClientPluginCustomization[]) => void, + ) { } + + get tools(): readonly ToolDefinition[] { + return this._getTools(); + } + set tools(tools: readonly ToolDefinition[]) { + this._setTools(tools); + } + + get customizations(): readonly ClientPluginCustomization[] { + return this._customizations; + } + set customizations(customizations: readonly ClientPluginCustomization[]) { + this._customizations = customizations; + this._syncCustomizations(customizations); + } +} + /** * Phase 4 skeleton {@link IAgent} provider for the Claude Agent SDK. * @@ -182,6 +216,9 @@ export class ClaudeAgent extends Disposable implements IAgent { */ private readonly _sessions = this._register(new DisposableMap()); + /** Stable active-client handles, keyed by `${sessionId}\0${clientId}`. */ + private readonly _activeClientHandles = new Map(); + /** * Phase 6: fired once per session when {@link _materializeProvisional} * promotes a provisional record into a real {@link ClaudeAgentSession}. @@ -560,6 +597,7 @@ export class ClaudeAgent extends Disposable implements IAgent { sess.abortController.abort(); } this._sessions.deleteAndDispose(sessionId); + this._pruneActiveClientHandles(sessionId); }); } @@ -790,6 +828,7 @@ export class ClaudeAgent extends Disposable implements IAgent { await Promise.all(sessionIds.map(sessionId => this._disposeSequencer.queue(sessionId, async () => { this._sessions.deleteAndDispose(sessionId); + this._pruneActiveClientHandles(sessionId); }) )); })(); @@ -938,14 +977,47 @@ export class ClaudeAgent extends Disposable implements IAgent { this._serverToolHost = host; } - setClientTools(session: URI, clientId: string | undefined, tools: ToolDefinition[]): void { + getOrCreateActiveClient(session: URI, client: { readonly clientId: string; readonly displayName?: string }): IActiveClient { const sessionId = AgentSession.id(session); - this._logService.info(`[Claude:${sessionId}] setClientTools clientId=${clientId} tools=[${tools.map(t => t.name).join(', ') || '(none)'}]`); - const sess = this._findAnySession(sessionId); - if (!sess) { - return; + const key = `${sessionId}\u0000${client.clientId}`; + let handle = this._activeClientHandles.get(key); + if (!handle) { + handle = new ClaudeActiveClientHandle( + client.clientId, + client.displayName, + () => this._findAnySession(sessionId)?.getClientTools(client.clientId) ?? [], + tools => { + this._logService.info(`[Claude:${sessionId}] active client ${client.clientId} tools=[${tools.map(t => t.name).join(', ') || '(none)'}]`); + this._findAnySession(sessionId)?.setClientTools(client.clientId, tools); + }, + customizations => { void this.syncClientCustomizations(session, client.clientId, [...customizations]); }, + ); + this._activeClientHandles.set(key, handle); + } + return handle; + } + + removeActiveClient(session: URI, clientId: string): void { + const sessionId = AgentSession.id(session); + this._activeClientHandles.delete(`${sessionId}\u0000${clientId}`); + // Tools are written synchronously, so remove them immediately. The + // customization sync runs inside the session sequencer, so serialize + // its removal there too — otherwise a late in-flight sync could + // resurrect the removed client's customizations after it has left. + this._findAnySession(sessionId)?.removeClientTools(clientId); + void this._sessionSequencer.queue(sessionId, async () => { + this._findAnySession(sessionId)?.removeClientCustomizations(clientId); + }).catch(() => { /* session torn down */ }); + } + + /** Drop cached active-client handles belonging to a session being torn down. */ + private _pruneActiveClientHandles(sessionId: string): void { + const prefix = `${sessionId}\u0000`; + for (const key of [...this._activeClientHandles.keys()]) { + if (key.startsWith(prefix)) { + this._activeClientHandles.delete(key); + } } - sess.setClientTools(tools, clientId); } onClientToolCallComplete(session: URI, toolCallId: string, result: ToolCallResult): void { @@ -964,26 +1036,25 @@ export class ClaudeAgent extends Disposable implements IAgent { entry?.session.completeClientToolCall(toolCallId, result); } - async setClientCustomizations(session: URI, clientId: string, customizations: ClientPluginCustomization[]): Promise { + async syncClientCustomizations(session: URI, clientId: string, customizations: ClientPluginCustomization[]): Promise { const sessionId = AgentSession.id(session); const sess = this._findAnySession(sessionId); if (!sess) { - this._logService.warn(`[Claude:${sessionId}] setClientCustomizations: session not found`); + this._logService.warn(`[Claude:${sessionId}] syncClientCustomizations: session not found`); return []; } // Run inside the session sequencer so that a fire-and-forget - // `setClientCustomizations` from `AgentSideEffects` cannot race - // ahead of a first `sendMessage`: if `sendMessage` is already - // queued, the sync runs first or queues behind it; either way - // the materialize call reads the most recently adopted plugin - // set, never an empty one mid-sync. + // customization sync cannot race ahead of a first `sendMessage`: if + // `sendMessage` is already queued, the sync runs first or queues + // behind it; either way the materialize call reads the most recently + // adopted plugin set, never an empty one mid-sync. return this._sessionSequencer.queue(sessionId, async () => { const synced = await this._pluginManager.syncCustomizations( clientId, customizations, status => this._fireCustomizationUpdated(session, { customization: status }), ); - sess.adoptClientCustomizations(synced); + sess.adoptClientCustomizations(clientId, synced); return synced; }); } diff --git a/src/vs/platform/agentHost/node/claude/claudeAgentSession.ts b/src/vs/platform/agentHost/node/claude/claudeAgentSession.ts index 25fe525b31635c..b82f08cdedb4ec 100644 --- a/src/vs/platform/agentHost/node/claude/claudeAgentSession.ts +++ b/src/vs/platform/agentHost/node/claude/claudeAgentSession.ts @@ -349,7 +349,7 @@ export class ClaudeAgentSession extends Disposable { this.abortController, dbRef, this.subagents, - this.toolDiff.model.state.get().clientId, + (toolName: string) => this.toolDiff.model.ownerOf(toolName), )); } catch (err) { dbRef.dispose(); @@ -741,12 +741,24 @@ export class ClaudeAgentSession extends Disposable { // #region Phase 10 — client tools - /** Replace the registered client tools snapshot. */ - setClientTools(tools: readonly ToolDefinition[], clientId?: string): void { - this.toolDiff.model.setTools(tools, clientId); - if (this._pipeline) { - this._pipeline.setClientId(this.toolDiff.model.state.get().clientId); - } + /** Replace a client's registered tools (full replacement). */ + setClientTools(clientId: string, tools: readonly ToolDefinition[]): void { + this.toolDiff.model.setTools(clientId, tools); + } + + /** This client's registered tools (empty when absent). */ + getClientTools(clientId: string): readonly ToolDefinition[] { + return this.toolDiff.model.getTools(clientId); + } + + /** Remove a client's tool contribution from this session. */ + removeClientTools(clientId: string): void { + this.toolDiff.model.removeClient(clientId); + } + + /** Remove a client's customization contribution from this session. */ + removeClientCustomizations(clientId: string): void { + this.clientCustomizationsDiff.model.removeClient(clientId); } /** @@ -803,8 +815,8 @@ export class ClaudeAgentSession extends Disposable { * the resulting snapshot down here. Flips the client-side dirty bit * so the next {@link send} pre-flight reloads SDK plugins. */ - adoptClientCustomizations(synced: readonly ISyncedCustomization[]): void { - this.clientCustomizationsDiff.model.setSyncedCustomizations(synced); + adoptClientCustomizations(clientId: string, synced: readonly ISyncedCustomization[]): void { + this.clientCustomizationsDiff.model.setSyncedCustomizations(clientId, synced); } /** Toggle a **client-pushed** customization on/off for this session. */ diff --git a/src/vs/platform/agentHost/node/claude/claudeMapSessionEvents.ts b/src/vs/platform/agentHost/node/claude/claudeMapSessionEvents.ts index c066e432758481..09b89d71dc19ec 100644 --- a/src/vs/platform/agentHost/node/claude/claudeMapSessionEvents.ts +++ b/src/vs/platform/agentHost/node/claude/claudeMapSessionEvents.ts @@ -219,7 +219,7 @@ export function mapSDKMessageToAgentSignals( state: ClaudeMapperState, logService: ILogService, registry: SubagentRegistry, - clientId?: string, + clientToolOwner?: (toolName: string) => string | undefined, ): AgentSignal[] { if (logService.getLevel() <= LogLevel.Trace) { try { @@ -232,7 +232,7 @@ export function mapSDKMessageToAgentSignals( switch (message.type) { case 'stream_event': return tagWithParent( - mapStreamEvent(message.event, session, turnId, state, logService, message.parent_tool_use_id, registry, clientId), + mapStreamEvent(message.event, session, turnId, state, logService, message.parent_tool_use_id, registry, clientToolOwner), session, message.parent_tool_use_id, registry, @@ -504,7 +504,7 @@ function mapStreamEvent( logService: ILogService, parentToolUseId: string | null, registry: SubagentRegistry, - clientId: string | undefined, + clientToolOwner: ((toolName: string) => string | undefined) | undefined, ): AgentSignal[] { switch (event.type) { case 'message_start': @@ -581,7 +581,7 @@ function mapStreamEvent( // produced by `buildClaudeToolMeta` because // `getClaudeToolKind('Task') === 'subagent'`. const meta = buildClaudeToolMeta(toolName); - const toolClientId = isClientTool ? clientId : undefined; + const toolClientId = isClientTool ? clientToolOwner?.(toolName) : undefined; return [{ kind: 'action', session, diff --git a/src/vs/platform/agentHost/node/claude/claudeSdkMessageRouter.ts b/src/vs/platform/agentHost/node/claude/claudeSdkMessageRouter.ts index 450c96a4835bb6..ff3b46bc9130b3 100644 --- a/src/vs/platform/agentHost/node/claude/claudeSdkMessageRouter.ts +++ b/src/vs/platform/agentHost/node/claude/claudeSdkMessageRouter.ts @@ -36,25 +36,25 @@ export class ClaudeSdkMessageRouter extends Disposable { private readonly _editObserver: ClaudeFileEditObserver; private readonly _mapperState = new ClaudeMapperState(); - private _clientId: string | undefined; + private _clientToolOwner: ((toolName: string) => string | undefined) | undefined; constructor( private readonly _sessionUri: URI, dbRef: IReference, private readonly _subagents: SubagentRegistry, - clientId: string | undefined = undefined, + clientToolOwner: ((toolName: string) => string | undefined) | undefined = undefined, @IInstantiationService instantiationService: IInstantiationService, @ILogService private readonly _logService: ILogService, ) { super(); - this._clientId = clientId; + this._clientToolOwner = clientToolOwner; this._editObserver = this._register( instantiationService.createInstance(ClaudeFileEditObserver, _sessionUri.toString(), dbRef), ); } - setClientId(clientId: string | undefined): void { - this._clientId = clientId; + setClientToolOwner(clientToolOwner: ((toolName: string) => string | undefined) | undefined): void { + this._clientToolOwner = clientToolOwner; } async handle(message: SDKMessage, turnId: string | undefined): Promise { @@ -74,7 +74,7 @@ export class ClaudeSdkMessageRouter extends Disposable { this._mapperState, this._logService, this._subagents, - this._clientId, + this._clientToolOwner, ); for (const signal of signals) { this._onDidProduceSignal.fire(signal); diff --git a/src/vs/platform/agentHost/node/claude/claudeSdkOptions.ts b/src/vs/platform/agentHost/node/claude/claudeSdkOptions.ts index 6c6d15be9e17f8..9242f2c8177beb 100644 --- a/src/vs/platform/agentHost/node/claude/claudeSdkOptions.ts +++ b/src/vs/platform/agentHost/node/claude/claudeSdkOptions.ts @@ -141,8 +141,8 @@ export async function buildClientMcpServers( registry: PendingRequestRegistry, sdkService: IClaudeAgentSdkService, ): Promise | undefined> { - const { tools } = toolDiff.consume(); - if (!tools || tools.length === 0) { + const tools = toolDiff.consume(); + if (tools.length === 0) { return undefined; } const server = await buildClientToolMcpServer(tools, id => registry.register(id), sdkService); diff --git a/src/vs/platform/agentHost/node/claude/claudeSdkPipeline.ts b/src/vs/platform/agentHost/node/claude/claudeSdkPipeline.ts index c617234be86a03..80cab8bff49bab 100644 --- a/src/vs/platform/agentHost/node/claude/claudeSdkPipeline.ts +++ b/src/vs/platform/agentHost/node/claude/claudeSdkPipeline.ts @@ -165,7 +165,7 @@ export class ClaudeSdkPipeline extends Disposable { abortController: AbortController, dbRef: IReference, subagents: SubagentRegistry, - clientId: string | undefined = undefined, + clientToolOwner: ((toolName: string) => string | undefined) | undefined = undefined, @IInstantiationService instantiationService: IInstantiationService, @ILogService private readonly _logService: ILogService, ) { @@ -184,7 +184,7 @@ export class ClaudeSdkPipeline extends Disposable { }), )); this._router = this._register(instantiationService.createInstance( - ClaudeSdkMessageRouter, sessionUri, dbRef, subagents, clientId, + ClaudeSdkMessageRouter, sessionUri, dbRef, subagents, clientToolOwner, )); this._register(this._router.onDidProduceSignal(s => this._onDidProduceSignal.fire(s))); // Dispose chain → abort → SDK cleanup. Reads the *current* @@ -211,13 +211,11 @@ export class ClaudeSdkPipeline extends Disposable { } /** - * Phase 10 — update the workbench `clientId` that the stream mapper - * stamps onto subsequent `ChatToolCallStart` events. Called by the - * session whenever {@link SessionClientToolsModel} receives a new - * clientId via `setClientTools`. + * Phase 10 — update the resolver the stream mapper uses to stamp the + * owning workbench `clientId` onto subsequent `ChatToolCallStart` events. */ - setClientId(clientId: string | undefined): void { - this._router.setClientId(clientId); + setClientToolOwner(clientToolOwner: ((toolName: string) => string | undefined) | undefined): void { + this._router.setClientToolOwner(clientToolOwner); } /** Attach the rematerializer hook for abort / crash recovery. Optional — tests that exercise only the dispose path skip this. */ diff --git a/src/vs/platform/agentHost/node/claude/clientTools/claudeSessionClientToolsModel.ts b/src/vs/platform/agentHost/node/claude/clientTools/claudeSessionClientToolsModel.ts index f309b1aba2f6af..de679d950e3229 100644 --- a/src/vs/platform/agentHost/node/claude/clientTools/claudeSessionClientToolsModel.ts +++ b/src/vs/platform/agentHost/node/claude/clientTools/claudeSessionClientToolsModel.ts @@ -6,63 +6,66 @@ import { Disposable } from '../../../../../base/common/lifecycle.js'; import { autorun, IObservable, ISettableObservable, observableValueOpts } from '../../../../../base/common/observable.js'; import type { ToolDefinition } from '../../../common/state/protocol/state.js'; -import { structuralToolsEqual } from '../../activeClientState.js'; +import { ActiveClientToolSet, structuralToolsEqual } from '../../activeClientState.js'; /** - * Combined snapshot of the workbench-registered client-tool definitions - * and the workbench `clientId` that owns them. The two travel as one - * value so consumers can read both with a single `.get()` and so that an - * update to either field is observed as a single change. - */ -export interface ISessionClientToolsState { - readonly tools: readonly ToolDefinition[] | undefined; - readonly clientId: string | undefined; -} - -const INITIAL_STATE: ISessionClientToolsState = { tools: undefined, clientId: undefined }; - -/** - * Pure state holder for the workbench-registered client-tool snapshot - * and the workbench `clientId` that owns it. Exposes the pair as a - * single {@link IObservable} so consumers can react to changes without - * polling. + * Pure state holder for the workbench-registered client-tool snapshots + * contributed by potentially several active clients, keyed by `clientId`. + * Exposes the merged tool set as a single {@link IObservable} (deduplicated + * by name, first-inserted client wins) so consumers can react to changes + * without polling, and {@link ownerOf} so tool calls can be routed back to + * the contributing client. * - * The `state` observable dedupes structurally-equivalent writes: tool - * snapshots compare on `name + description + inputSchema` - * (order-insensitive, `undefined` equivalent to `[]`); `clientId` compares strictly. - * A re-send of the same `(tools, clientId)` pair therefore does NOT fire - * downstream subscribers. + * The `merged` observable dedupes structurally-equivalent writes: tool + * snapshots compare on `name + description + inputSchema` (order-insensitive, + * `undefined` equivalent to `[]`). A re-send of a structurally identical set + * therefore does NOT fire downstream subscribers. * * Knows nothing about diffing or the SDK — pair with - * {@link SessionClientToolsDiff} to track "has the snapshot changed + * {@link SessionClientToolsDiff} to track "has the merged snapshot changed * since the last successful SDK build". */ export class SessionClientToolsModel { - private readonly _state: ISettableObservable = observableValueOpts( - { owner: this, equalsFn: stateEqual }, - INITIAL_STATE, + private readonly _toolSet = new ActiveClientToolSet(); + private readonly _merged: ISettableObservable = observableValueOpts( + { owner: this, equalsFn: (a, b) => structuralToolsEqual(a, b) }, + [], ); - readonly state: IObservable = this._state; + readonly merged: IObservable = this._merged; + + /** Replace `clientId`'s contributed tools (full replacement). */ + setTools(clientId: string, tools: readonly ToolDefinition[]): void { + this._toolSet.set(clientId, tools); + this._merged.set(this._toolSet.merged(), undefined); + } - setTools(tools: readonly ToolDefinition[] | undefined, clientId?: string): void { - const current = this._state.get(); - this._state.set({ - tools, - clientId: clientId ?? current.clientId, - }, undefined); + /** This client's contributed tools (empty when absent). */ + getTools(clientId: string): readonly ToolDefinition[] { + return this._toolSet.get(clientId); + } + + /** Remove a client's tool contribution. */ + removeClient(clientId: string): void { + if (this._toolSet.delete(clientId)) { + this._merged.set(this._toolSet.merged(), undefined); + } + } + + /** The `clientId` that owns the merged tool named `toolName`, or `undefined`. */ + ownerOf(toolName: string): string | undefined { + return this._toolSet.ownerOf(toolName); } } /** - * Tracks "has {@link SessionClientToolsModel.state} changed since the - * last successful {@link build}?". Subscribes to the model's observable - * and flips a private dirty bit on every change; {@link build} captures - * the current snapshot, hands it to the supplied builder, and clears - * the bit on success — preserving the C6 pin invariant from the - * previous `ClientToolDiff` implementation: a `setTools` call that - * races the in-flight builder re-flips the bit via the autorun, so the - * next sendMessage detects the stale set and triggers another rebind. + * Tracks "has {@link SessionClientToolsModel.merged} changed since the + * last successful {@link consume}?". Subscribes to the model's observable + * and flips a private dirty bit on every change; {@link consume} captures + * the current merged snapshot and clears the bit — preserving the C6 pin + * invariant from the previous `ClientToolDiff` implementation: a `setTools` + * call that races the in-flight builder re-flips the bit via the autorun, so + * the next sendMessage detects the stale set and triggers another rebind. * * On builder throw the bit is left set — the SDK is still running with * the previous snapshot, so the next sendMessage should retry. @@ -77,21 +80,18 @@ export class SessionClientToolsDiff extends Disposable { // dirty before any `setTools` has happened. private _ignoreNextFire = true; // Structural tool snapshot last marked applied (via {@link consume}). - // The dirty bit only flips when the live tools differ structurally from - // this — a `clientId`-only change (same tools, new window) must NOT - // trigger a yield-restart even though the observable still fires. - private _lastAppliedTools: readonly ToolDefinition[] | undefined = undefined; + private _lastAppliedTools: readonly ToolDefinition[] = []; constructor() { super(); this._register(autorun(reader => { - const state = this.model.state.read(reader); + const merged = this.model.merged.read(reader); if (this._ignoreNextFire) { this._ignoreNextFire = false; - this._lastAppliedTools = state.tools; + this._lastAppliedTools = merged; return; } - if (!structuralToolsEqual(state.tools, this._lastAppliedTools)) { + if (!structuralToolsEqual(merged, this._lastAppliedTools)) { this._dirty = true; } })); @@ -102,18 +102,18 @@ export class SessionClientToolsDiff extends Disposable { } /** - * Read the current state and mark it as the applied snapshot. A - * subsequent {@link SessionClientToolsModel.setTools} re-flips dirty + * Read the current merged tool set and mark it as the applied snapshot. + * A subsequent {@link SessionClientToolsModel.setTools} re-flips dirty * via the autorun, so callers do NOT need to compare snapshots * themselves to detect a race. If the caller's downstream work * (e.g. SDK rebuild) fails, call {@link markDirty} to surface the * stale state so the next sendMessage retries. */ - consume(): ISessionClientToolsState { - const state = this.model.state.get(); + consume(): readonly ToolDefinition[] { + const merged = this.model.merged.get(); this._dirty = false; - this._lastAppliedTools = state.tools; - return state; + this._lastAppliedTools = merged; + return merged; } /** @@ -125,7 +125,3 @@ export class SessionClientToolsDiff extends Disposable { this._dirty = true; } } - -function stateEqual(a: ISessionClientToolsState, b: ISessionClientToolsState): boolean { - return a.clientId === b.clientId && structuralToolsEqual(a.tools, b.tools); -} diff --git a/src/vs/platform/agentHost/node/claude/customizations/claudeSessionClientCustomizationsModel.ts b/src/vs/platform/agentHost/node/claude/customizations/claudeSessionClientCustomizationsModel.ts index ace08dfade4cba..d9b2da83bdd23c 100644 --- a/src/vs/platform/agentHost/node/claude/customizations/claudeSessionClientCustomizationsModel.ts +++ b/src/vs/platform/agentHost/node/claude/customizations/claudeSessionClientCustomizationsModel.ts @@ -46,6 +46,9 @@ const INITIAL_STATE: ISessionCustomizationsState = { synced: [], enablement: new */ export class SessionClientCustomizationsModel { + /** Per-client synced customizations, keyed by `clientId`, merged into `state.synced`. */ + private readonly _byClient = new Map(); + private readonly _state: ISettableObservable = observableValueOpts( { owner: this, equalsFn: stateEqual }, INITIAL_STATE, @@ -77,10 +80,40 @@ export class SessionClientCustomizationsModel { }, ); - /** Replace the client-pushed customization snapshot for this session. */ - setSyncedCustomizations(synced: readonly ISyncedCustomization[]): void { + /** + * The union of every client's synced customizations, deduplicated by + * customization `id` with the first-inserted client winning. Order + * follows client insertion order. + */ + private _mergedSynced(): readonly ISyncedCustomization[] { + const seen = new Set(); + const result: ISyncedCustomization[] = []; + for (const synced of this._byClient.values()) { + for (const item of synced) { + if (seen.has(item.customization.id)) { + continue; + } + seen.add(item.customization.id); + result.push(item); + } + } + return result; + } + + /** Replace a single client's pushed customization snapshot for this session. */ + setSyncedCustomizations(clientId: string, synced: readonly ISyncedCustomization[]): void { + this._byClient.set(clientId, synced); + const cur = this._state.get(); + this._state.set({ synced: this._mergedSynced(), enablement: cur.enablement }, undefined); + } + + /** Remove a client's pushed customizations from this session. */ + removeClient(clientId: string): void { + if (!this._byClient.delete(clientId)) { + return; + } const cur = this._state.get(); - this._state.set({ synced, enablement: cur.enablement }, undefined); + this._state.set({ synced: this._mergedSynced(), enablement: cur.enablement }, undefined); } /** Toggle a client-pushed customization on/off for this session. */ diff --git a/src/vs/platform/agentHost/node/codex/codexAgent.ts b/src/vs/platform/agentHost/node/codex/codexAgent.ts index 1caa7d44c151a3..2bcf1e88a9b0ee 100644 --- a/src/vs/platform/agentHost/node/codex/codexAgent.ts +++ b/src/vs/platform/agentHost/node/codex/codexAgent.ts @@ -18,15 +18,15 @@ import { ILogService } from '../../../log/common/log.js'; import { IProductService } from '../../../product/common/productService.js'; import { createSchema, platformSessionSchema, schemaProperty, type SessionMode } from '../../common/agentHostSchema.js'; import { getReasoningEffortDescription, getReasoningEffortLabel } from '../../common/reasoningEffort.js'; -import { AgentHostCodexAgentBinaryArgsEnvVar, AgentHostCodexAgentCodexHomeEnvVar, AgentHostCodexAgentSdkRootEnvVar, AgentSession, AgentSignal, GITHUB_COPILOT_PROTECTED_RESOURCE, GITHUB_REPO_PROTECTED_RESOURCE, IAgent, IAgentCreateSessionConfig, IAgentCreateSessionResult, IAgentDescriptor, IAgentMaterializeSessionEvent, IAgentModelInfo, IAgentResolveSessionConfigParams, IAgentSessionConfigCompletionsParams, IAgentSessionMetadata, IMcpNotification, type AgentProvider } from '../../common/agentService.js'; +import { AgentHostCodexAgentBinaryArgsEnvVar, AgentHostCodexAgentCodexHomeEnvVar, AgentHostCodexAgentSdkRootEnvVar, AgentSession, AgentSignal, GITHUB_COPILOT_PROTECTED_RESOURCE, GITHUB_REPO_PROTECTED_RESOURCE, IActiveClient, IAgent, IAgentCreateSessionConfig, IAgentCreateSessionResult, IAgentDescriptor, IAgentMaterializeSessionEvent, IAgentModelInfo, IAgentResolveSessionConfigParams, IAgentSessionConfigCompletionsParams, IAgentSessionMetadata, IMcpNotification, type AgentProvider } from '../../common/agentService.js'; import { SessionConfigKey } from '../../common/sessionConfigKeys.js'; import { AHP_AUTH_REQUIRED, ProtocolError } from '../../common/state/sessionProtocol.js'; import { ActionType, type SessionAction, type ChatAction } from '../../common/state/sessionActions.js'; import type { ConfigSchema, ModelSelection, ProtectedResourceMetadata, ToolDefinition } from '../../common/state/protocol/state.js'; import type { ResolveSessionConfigResult, SessionConfigCompletionsResult } from '../../common/state/protocol/commands.js'; import { type ClientPluginCustomization, type MessageAttachment, type PendingMessage, type ChatInputAnswer, ChatInputResponseKind, type PolicyState, type ToolCallResult, ToolResultContentType, type Turn } from '../../common/state/sessionState.js'; -import type { ISyncedCustomization } from '../../common/agentPluginManager.js'; import type { IAgentServerToolHost } from '../../common/agentServerTools.js'; +import { ActiveClientToolSet } from '../activeClientState.js'; import { McpCustomizationController } from '../shared/mcpCustomizationController.js'; import { buildCodexMcpReadResult, codexMcpListToInventory, codexMcpToolsChanged, inventoryToSdkServers, translateCodexMcpStartupState, type ICodexMcpServerEntry } from './codexMcpServers.js'; import { buildElicitationRequest, cancelledElicitationResponse, declinedElicitationResponse, elicitationResponseFromAnswers } from './codexElicitationMapper.js'; @@ -345,13 +345,12 @@ interface ICodexSession { */ readonly pendingSteeringFlips: Map; /** - * Client-provided tool definitions for this session (the active - * workbench client's tools), registered with codex as `dynamicTools` - * at `thread/start`. `undefined` until the first {@link setClientTools}. + * Client-provided tool definitions for this session, keyed by the + * contributing workbench client. The merged set is registered with codex + * as `dynamicTools` at `thread/start`. Empty until the first active client + * sets its tools. */ - clientTools: readonly ToolDefinition[] | undefined; - /** Workbench client id that owns {@link clientTools}. */ - toolsClientId: string | undefined; + readonly clientToolSet: ActiveClientToolSet; /** * Parked deferreds for in-flight client-tool calls (codex * `item/tool/call`), keyed by the host-side toolCallId. Resolved by @@ -488,6 +487,37 @@ function toolsSignature(tools: readonly ToolDefinition[] | undefined): string { .join('\u0001'); } +/** + * Codex active-client handle. Writes flow into the owning session's + * {@link ActiveClientToolSet}; the session is resolved lazily so writes that + * arrive before (or after) the session exists are gracefully dropped, matching + * the prior `setClientTools` early-return behavior. Codex has no client + * customization layer, so `customizations` is inert. + */ +class CodexActiveClientHandle implements IActiveClient { + constructor( + private readonly _getSession: () => ICodexSession | undefined, + readonly clientId: string, + readonly displayName: string | undefined, + private readonly _onToolsSet: (tools: readonly ToolDefinition[]) => void, + ) { } + + get tools(): readonly ToolDefinition[] { + return this._getSession()?.clientToolSet.get(this.clientId) ?? []; + } + set tools(tools: readonly ToolDefinition[]) { + this._getSession()?.clientToolSet.set(this.clientId, tools); + this._onToolsSet(tools); + } + + get customizations(): readonly ClientPluginCustomization[] { + return []; + } + set customizations(_customizations: readonly ClientPluginCustomization[]) { + // Codex does not support client-contributed customizations. + } +} + /** * Map a resolved approval decision to the {@link FileChangeApprovalDecision} * subset. The host's boolean response only yields `accept`/`decline`; the @@ -1027,7 +1057,7 @@ export class CodexAgent extends Disposable implements IAgent { */ private _buildDynamicTools(session: ICodexSession): DynamicToolSpec[] | undefined { const serverTools = this._serverToolHost?.definitions ?? []; - const clientTools = session.clientTools ?? []; + const clientTools = session.clientToolSet.merged(); // Server tools first; a server tool name shadows a colliding client tool // (the agent host owns those names) and matches the routing order below. const seen = new Set(); @@ -1076,7 +1106,7 @@ export class CodexAgent extends Disposable implements IAgent { if (toolCallId === undefined) { return { result: this._toolFailure(`No pending client tool call for ${params.tool} (callId ${params.callId})`) }; } - if (!session.clientTools || session.toolsClientId === undefined) { + if (session.clientToolSet.size === 0) { return { result: this._toolFailure(`No client available to run ${params.tool}`) }; } try { @@ -1584,17 +1614,17 @@ export class CodexAgent extends Disposable implements IAgent { }; } + const clientToolSet = new ActiveClientToolSet(); const session: ICodexSession = { sessionId, threadId: undefined, sessionUri, workingDirectory: config.workingDirectory, - mapState: createCodexSessionMapState(new Set(this._serverToolHost?.toolNames ?? [])), + mapState: createCodexSessionMapState(new Set(this._serverToolHost?.toolNames ?? []), clientToolSet), pendingCommandApprovals: new PendingRequestRegistry(), acceptedForSession: new Set(), pendingSteeringFlips: new Map(), - clientTools: undefined, - toolsClientId: undefined, + clientToolSet, pendingClientToolCalls: new PendingRequestRegistry(), pendingUserInputs: new PendingRequestRegistry(), materializedToolsSig: undefined, @@ -1684,7 +1714,7 @@ export class CodexAgent extends Disposable implements IAgent { return; } session.threadId = threadId; - session.materializedToolsSig = toolsSignature(session.clientTools); + session.materializedToolsSig = toolsSignature(session.clientToolSet.merged()); this._logService.info(`[Codex DEBUG] materialized session=${session.sessionUri.toString()} threadId=${session.threadId}`); this._sessionIdByThreadId.set(session.threadId, session.sessionId); // Advertise the agent host's server tools on this session so clients see @@ -1705,7 +1735,7 @@ export class CodexAgent extends Disposable implements IAgent { private async _restartThreadWithCurrentTools(session: ICodexSession): Promise { const conn = this._connection; const oldThreadId = session.threadId; - this._logService.info(`[Codex:${session.sessionId}] restarting thread ${oldThreadId} to apply client tools [${(session.clientTools ?? []).map(t => t.name).join(', ') || '(none)'}]`); + this._logService.info(`[Codex:${session.sessionId}] restarting thread ${oldThreadId} to apply client tools [${session.clientToolSet.merged().map(t => t.name).join(', ') || '(none)'}]`); if (conn.kind === 'ready' && oldThreadId !== undefined) { this._sessionIdByThreadId.delete(oldThreadId); try { @@ -1820,7 +1850,7 @@ export class CodexAgent extends Disposable implements IAgent { // was prewarmed (or otherwise started) before the current client tools // were known, restart it now — before any turn commits history, so // nothing is lost — so the tools land in `dynamicTools`. - if (!session.firstTurnSent && !session.needsResume && toolsSignature(session.clientTools) !== session.materializedToolsSig) { + if (!session.firstTurnSent && !session.needsResume && toolsSignature(session.clientToolSet.merged()) !== session.materializedToolsSig) { try { await this._restartThreadWithCurrentTools(session); this._persistMaterializedSession(session); @@ -2136,17 +2166,17 @@ export class CodexAgent extends Disposable implements IAgent { if (!this._sessions.has(sessionId)) { const workingDirectory = read.thread.cwd ? URI.file(read.thread.cwd) : undefined; const threadId = read.thread.id; + const clientToolSet = new ActiveClientToolSet(); this._sessions.set(sessionId, { sessionId, threadId, sessionUri: session, workingDirectory, - mapState: createCodexSessionMapState(new Set(this._serverToolHost?.toolNames ?? [])), + mapState: createCodexSessionMapState(new Set(this._serverToolHost?.toolNames ?? []), clientToolSet), pendingCommandApprovals: new PendingRequestRegistry(), acceptedForSession: new Set(), pendingSteeringFlips: new Map(), - clientTools: undefined, - toolsClientId: undefined, + clientToolSet, pendingClientToolCalls: new PendingRequestRegistry(), pendingUserInputs: new PendingRequestRegistry(), materializedToolsSig: undefined, @@ -2259,16 +2289,19 @@ export class CodexAgent extends Disposable implements IAgent { this._serverToolHost = host; } - setClientTools(session: URI, clientId: string | undefined, tools: ToolDefinition[]): void { + getOrCreateActiveClient(session: URI, client: { readonly clientId: string; readonly displayName?: string }): IActiveClient { const sessionId = AgentSession.id(session); - const sess = this._sessions.get(sessionId); - if (!sess) { - return; - } - sess.clientTools = tools.length > 0 ? tools : undefined; - sess.toolsClientId = clientId; - sess.mapState.toolsClientId = sess.clientTools ? clientId : undefined; - this._logService.info(`[Codex:${sessionId}] setClientTools clientId=${clientId ?? '(none)'} tools=[${tools.map(t => t.name).join(', ') || '(none)'}]`); + return new CodexActiveClientHandle( + () => this._sessions.get(sessionId), + client.clientId, + client.displayName, + tools => this._logService.info(`[Codex:${sessionId}] active client ${client.clientId} tools=[${tools.map(t => t.name).join(', ') || '(none)'}]`), + ); + } + + removeActiveClient(session: URI, clientId: string): void { + const sessionId = AgentSession.id(session); + this._sessions.get(sessionId)?.clientToolSet.delete(clientId); } onClientToolCallComplete(session: URI, toolCallId: string, result: ToolCallResult): void { @@ -2279,10 +2312,6 @@ export class CodexAgent extends Disposable implements IAgent { sess?.pendingClientToolCalls.respondOrBuffer(toolCallId, result); } - setClientCustomizations(_session: URI, _clientId: string, _customizations: ClientPluginCustomization[]): Promise { - return Promise.resolve([]); - } - setCustomizationEnabled(_uri: string, _enabled: boolean): void { // no-op; customizations not yet wired for codex. } diff --git a/src/vs/platform/agentHost/node/codex/codexMapAppServerEvents.ts b/src/vs/platform/agentHost/node/codex/codexMapAppServerEvents.ts index a307d76a1a94e4..b19caf405e459b 100644 --- a/src/vs/platform/agentHost/node/codex/codexMapAppServerEvents.ts +++ b/src/vs/platform/agentHost/node/codex/codexMapAppServerEvents.ts @@ -8,6 +8,7 @@ import { toToolCallMeta } from '../../common/meta/agentToolCallMeta.js'; import { ActionType, type SessionAction, type ChatAction } from '../../common/state/sessionActions.js'; import { MessageKind, ResponsePartKind, ToolCallConfirmationReason, ToolCallContributorKind, ToolResultContentType, TurnState } from '../../common/state/sessionState.js'; import { extractForwardedErrorInfo } from '../shared/forwardedChatError.js'; +import { ActiveClientToolSet } from '../activeClientState.js'; import type { AgentMessageDeltaNotification } from './protocol/generated/v2/AgentMessageDeltaNotification.js'; import type { CommandExecutionOutputDeltaNotification } from './protocol/generated/v2/CommandExecutionOutputDeltaNotification.js'; import type { FileChangeOutputDeltaNotification } from './protocol/generated/v2/FileChangeOutputDeltaNotification.js'; @@ -50,12 +51,12 @@ export interface ICodexSessionMapState { /** Current turn id (per `turn/started`). */ currentTurnId: string | undefined; /** - * Workbench client id that owns the session's client-provided - * (`dynamicTools`) tools. Set so `dynamicToolCall` tool-call starts - * carry a `Client` contributor and the workbench routes execution back - * to that client. `undefined` when no client tools are registered. + * Live registry of the session's client-provided (`dynamicTools`) tools, + * keyed by contributing workbench client. A `dynamicToolCall` tool-call + * start is stamped with the owning client (so the workbench routes + * execution back to it) resolved via {@link ActiveClientToolSet.ownerOf}. */ - toolsClientId: string | undefined; + clientToolSet: ActiveClientToolSet; /** * Names of the agent host's server tools (executed in-process). A * `dynamicToolCall` for one of these omits the `Client` contributor so the @@ -72,13 +73,13 @@ export interface ICodexToolCallEntry { output: string; } -export function createCodexSessionMapState(serverToolNames: ReadonlySet = new Set()): ICodexSessionMapState { +export function createCodexSessionMapState(serverToolNames: ReadonlySet = new Set(), clientToolSet: ActiveClientToolSet = new ActiveClientToolSet()): ICodexSessionMapState { return { itemToPartId: new Map(), itemToToolCall: new Map(), itemToReasoningPartId: new Map(), currentTurnId: undefined, - toolsClientId: undefined, + clientToolSet, serverToolNames, }; } @@ -454,6 +455,7 @@ export function mapItemStarted( // they carry no `Client` contributor; only client-provided tools route // execution back to the owning workbench client. const isServerTool = params.item.namespace === null && state.serverToolNames.has(params.item.tool); + const ownerClientId = isServerTool ? undefined : state.clientToolSet.ownerOf(params.item.tool); state.itemToToolCall.set(params.item.id, { toolCallId, turnId: params.turnId, @@ -467,7 +469,7 @@ export function mapItemStarted( toolCallId, toolName, displayName: params.item.tool, - ...(state.toolsClientId && !isServerTool ? { contributor: { kind: ToolCallContributorKind.Client, clientId: state.toolsClientId } } : {}), + ...(ownerClientId ? { contributor: { kind: ToolCallContributorKind.Client, clientId: ownerClientId } } : {}), }, { type: ActionType.ChatToolCallDelta, diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts index a186cc08ccf685..dac696b6e9e961 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts @@ -32,7 +32,7 @@ import { createAgentModelPricingMeta } from '../../common/agentModelPricing.js'; import { AgentHostConfigKey, agentHostCustomizationConfigSchema, toContainerCustomization } from '../../common/agentHostCustomizationConfig.js'; import { AgentHostMcpServersConfigKey, AgentHostSessionSyncEnabledConfigKey, AutoApproveLevel, ISchemaProperty, SessionMode, createSchema, migrateLegacyAutopilotConfig, platformRootSchema, platformSessionSchema, schemaProperty, type AgentHostMcpServers } from '../../common/agentHostSchema.js'; import { IAgentPluginManager, ISyncedCustomization } from '../../common/agentPluginManager.js'; -import { AgentSession, AgentSignal, GITHUB_COPILOT_PROTECTED_RESOURCE, GITHUB_REPO_PROTECTED_RESOURCE, IAgent, IAgentCreateChatOptions, IAgentCreateSessionConfig, IAgentCreateSessionResult, IAgentDescriptor, IAgentMaterializeSessionEvent, IAgentModelInfo, IAgentResolveSessionConfigParams, IAgentSessionConfigCompletionsParams, IAgentSessionMetadata, IAgentSessionProjectInfo, IMcpNotification } from '../../common/agentService.js'; +import { AgentSession, AgentSignal, GITHUB_COPILOT_PROTECTED_RESOURCE, GITHUB_REPO_PROTECTED_RESOURCE, IActiveClient, IAgent, IAgentCreateChatOptions, IAgentCreateSessionConfig, IAgentCreateSessionResult, IAgentDescriptor, IAgentMaterializeSessionEvent, IAgentModelInfo, IAgentResolveSessionConfigParams, IAgentSessionConfigCompletionsParams, IAgentSessionMetadata, IAgentSessionProjectInfo, IMcpNotification } from '../../common/agentService.js'; import { getEffectiveAgents } from '../../common/customAgents.js'; import { getReasoningEffortDescription, getReasoningEffortLabel } from '../../common/reasoningEffort.js'; import type { IAgentServerToolHost } from '../../common/agentServerTools.js'; @@ -43,7 +43,7 @@ import type { ResolveSessionConfigResult, SessionConfigCompletionsResult } from import { ProtectedResourceMetadata, type AgentSelection, type ChildCustomizationType, type ConfigPropertySchema, type ConfigSchema, type ModelSelection, type ToolDefinition } from '../../common/state/protocol/state.js'; import { ActionType, type SessionAction } from '../../common/state/sessionActions.js'; import { AgentCustomization, CustomizationLoadStatus, CustomizationType, ResponsePartKind, RuleCustomization, ChatInputResponseKind, SkillCustomization, customizationId, buildChatUri, isDefaultChatUri, parseChatUri, parseSubagentSessionUri, type ChildCustomization, type ClientPluginCustomization, type Customization, type DirectoryCustomization, type HookCustomization, type MessageAttachment, type PendingMessage, type PluginCustomization, type PolicyState, type ResponsePart, type ChatInputAnswer, type ToolCallResult, type Turn } from '../../common/state/sessionState.js'; -import { ActiveClientState } from '../activeClientState.js'; +import { ActiveClientToolSet } from '../activeClientState.js'; import { IAgentConfigurationService } from '../agentConfigurationService.js'; import { IAgentHostCompletions } from '../agentHostCompletions.js'; import { IAgentHostGitService, META_DIFF_BASE_BRANCH } from '../../common/agentHostGitService.js'; @@ -1123,13 +1123,15 @@ export class CopilotAgent extends Disposable implements IAgent { // of activeClient state isn't engaged until materialization. if (sessionConfig.activeClient) { const ac = this._getOrCreateActiveClient(sessionUri, workingDirectory); - ac.updateTools(sessionConfig.activeClient.clientId, sessionConfig.activeClient.tools); - if (sessionConfig.activeClient.customizations !== undefined) { + const seeded = sessionConfig.activeClient; + ac.toolSet.set(seeded.clientId, seeded.tools); + ac.getOrCreateHandle(seeded.clientId, seeded.displayName); + if (seeded.customizations !== undefined) { // Provisional eager-create: no session-state listener is // hooked up yet, so suppress action events. The session // reads the final view via its initial snapshot once it // materializes. - await ac.pluginController.sync(sessionConfig.activeClient.clientId, sessionConfig.activeClient.customizations, { quiet: true }); + await ac.pluginController.sync(seeded.clientId, seeded.customizations, { quiet: true }); } } @@ -1193,7 +1195,7 @@ export class CopilotAgent extends Disposable implements IAgent { const customizationDirectory = provisional.workingDirectory; // Always create an ActiveClient so the snapshot includes host + // session-discovered customizations, even when no client has called - // `setClientCustomizations` / `setClientTools` yet. + // active-client handle yet. const activeClient = this._getOrCreateActiveClient(sessionUri, customizationDirectory); const snapshot = await activeClient.snapshot(); const workingDirectory = await this._resolveSessionWorkingDirectory(materializedConfig, sessionId, prompt); @@ -1209,7 +1211,7 @@ export class CopilotAgent extends Disposable implements IAgent { workingDirectory, resolvedAgentName, snapshot, - activeClientState: activeClient.state, + activeClientToolSet: activeClient.toolSet, shellManager, githubToken: this._githubToken, model: provisional.model, @@ -1318,18 +1320,24 @@ export class CopilotAgent extends Disposable implements IAgent { return { items: branches.map(branch => ({ value: branch, label: branch })) }; } - async setClientCustomizations(session: URI, clientId: string, customizations: ClientPluginCustomization[]): Promise { - const directory = await this._getSessionCustomizationDirectory(session); - const activeClient = this._getOrCreateActiveClient(session, directory); - return activeClient.pluginController.sync(clientId, customizations); + getOrCreateActiveClient(session: URI, client: { readonly clientId: string; readonly displayName?: string }): IActiveClient { + const activeClient = this._getOrCreateActiveClient(session, undefined); + // Anchor the customization directory (best-effort, idempotent) so + // session-discovered customizations surface alongside this client's, + // mirroring the previous eager resolution in `setClientCustomizations`. + if (!activeClient.pluginController.directory) { + this._getSessionCustomizationDirectory(session).then( + directory => activeClient.pluginController.setDirectory(directory), + () => { /* best-effort anchoring */ }, + ); + } + return activeClient.getOrCreateHandle(client.clientId, client.displayName); } - setClientTools(session: URI, clientId: string | undefined, tools: ToolDefinition[]): void { + removeActiveClient(session: URI, clientId: string): void { const sessionId = AgentSession.id(session); - const activeClient = this._getOrCreateActiveClient(session, undefined); - const hasCachedEntry = this._sessions.has(sessionId); - this._logService.info(`[Copilot:${sessionId}] setClientTools: clientId=${clientId ?? '(none)'}, tools=[${tools.map(t => t.name).join(', ') || '(none)'}], hasCachedSdkSession=${hasCachedEntry}`); - activeClient.updateTools(clientId, tools); + this._logService.info(`[Copilot:${sessionId}] removeActiveClient: clientId=${clientId}`); + this._activeClients.get(session)?.removeClient(clientId); } onClientToolCallComplete(session: URI, toolCallId: string, result: ToolCallResult): void { @@ -1391,7 +1399,7 @@ export class CopilotAgent extends Disposable implements IAgent { const hadCachedEntry = !!entry; this._logService.info(`[Copilot:${sessionId}] sendMessage: cachedEntry=${hadCachedEntry}, hasActiveClient=${!!activeClient}, activeClientId=${activeClient ? '(set)' : '(none)'}`); if (entry && activeClient && await activeClient.requiresRestart(entry.appliedSnapshot)) { - this._logService.info(`[Copilot:${sessionId}] Session config changed (requiresRestart=true), refreshing session. clientId=${activeClient.state.clientId ?? '(none)'}`); + this._logService.info(`[Copilot:${sessionId}] Session config changed (requiresRestart=true), refreshing session. clients=[${[...activeClient.toolSet.clientIds()].join(', ') || '(none)'}]`); this._sessions.deleteAndDispose(sessionId); entry = undefined; } @@ -1684,11 +1692,10 @@ export class CopilotAgent extends Disposable implements IAgent { const chatSdkId = generateUuid(); // Peer chats share the owning session's ActiveClient so that // client tool / customization updates (which are keyed by the - // session URI via `setClientTools` / `setClientCustomizations`) - // reach the additional chat's SDK conversation. Keying it by the - // chat URI instead would snapshot empty/stale tools and never see - // subsequent updates, and would also leak (nothing disposes a - // chat-keyed ActiveClient). + // session URI via the active-client handles) reach the additional + // chat's SDK conversation. Keying it by the chat URI instead would + // snapshot empty/stale tools and never see subsequent updates, and + // would also leak (nothing disposes a chat-keyed ActiveClient). const activeClient = this._getOrCreateActiveClient(session, workingDirectory); const snapshot = await activeClient.snapshot(); const shellManager = this._instantiationService.createInstance(ShellManager, chat, workingDirectory); @@ -1699,7 +1706,7 @@ export class CopilotAgent extends Disposable implements IAgent { workingDirectory, resolvedAgentName: undefined, snapshot, - activeClientState: activeClient.state, + activeClientToolSet: activeClient.toolSet, shellManager, githubToken: this._githubToken, model: options?.model, @@ -1813,7 +1820,7 @@ export class CopilotAgent extends Disposable implements IAgent { workingDirectory, resolvedAgentName: undefined, snapshot, - activeClientState: activeClient.state, + activeClientToolSet: activeClient.toolSet, shellManager, githubToken: this._githubToken, fallback: { model: info.model }, @@ -2021,7 +2028,7 @@ export class CopilotAgent extends Disposable implements IAgent { workingDirectory: launchPlan.workingDirectory, customizationDirectory, clientSnapshot: launchPlan.snapshot, - activeClientState: launchPlan.activeClientState, + activeClientToolSet: launchPlan.activeClientToolSet, resolveMcpChildId: name => findMcpChildId(activeClient.pluginController.getCustomizations(), name), serverToolHost: this._serverToolHost, }, @@ -2108,7 +2115,7 @@ export class CopilotAgent extends Disposable implements IAgent { const customizationDirectory = storedMetadata.customizationDirectory ?? storedMetadata.workingDirectory; // Always create an ActiveClient so the snapshot includes host + // session-discovered customizations, even when no client has called - // `setClientCustomizations` / `setClientTools` yet. + // active-client handle yet. const activeClient = this._getOrCreateActiveClient(sessionUri, customizationDirectory); const snapshot = await activeClient.snapshot(); const sessionMetadata = await client.getSessionMetadata(sessionId).catch(err => { @@ -2129,7 +2136,7 @@ export class CopilotAgent extends Disposable implements IAgent { workingDirectory, resolvedAgentName, snapshot, - activeClientState: activeClient.state, + activeClientToolSet: activeClient.toolSet, shellManager, githubToken: this._githubToken, fallback: { @@ -2924,6 +2931,22 @@ class PluginController extends Disposable { } } +/** + * Per-client slice of {@link SessionPluginController} customization state. + * One entry exists per active client that has contributed customizations to + * the session. + */ +interface IClientCustomizationState { + /** Monotonic revision used to detect and ignore stale in-flight syncs for this client. */ + revision: number; + /** This client's resolved customizations (Loading/Loaded/Error per item). */ + customizations: readonly IResolvedCustomization[]; + /** This client's in-flight (or settled) sync promise. */ + sync: Promise; + /** The raw inputs last passed to {@link SessionPluginController.sync} for this client. */ + inputs: readonly ClientPluginCustomization[]; +} + /** * Per-session view over {@link PluginController}. * @@ -2944,11 +2967,14 @@ class SessionPluginController extends Disposable { readonly onDidPublish = this._onDidPublish.event; private readonly _enablement = new Map(); - private _clientCustomizations: readonly IResolvedCustomization[] = []; - private _clientSync: Promise = Promise.resolve([]); - private _clientRevision = 0; - /** Last clientId seen via {@link sync}; used by {@link retryFailedClientSyncIfNeeded}. */ - private _clientId: string | undefined; + /** + * Per-client customization state, keyed by `clientId`. Each active client + * contributing customizations to this session has one entry; the published + * customization list is the union across all entries (deduplicated by URI, + * first-inserted client wins). Insertion order is preserved so the merged + * order stays stable across updates. + */ + private readonly _clients = new Map(); private readonly _sessionDiscovered: MutableDisposable = this._register(new MutableDisposable()); @@ -2979,7 +3005,7 @@ class SessionPluginController extends Disposable { public getCustomizations(): readonly Customization[] { const result: Customization[] = [ ...this._parent.hostCustomizations().map(item => this._applyEnablement(item.customization)), - ...this._clientCustomizations.map(item => this._applyEnablement(item.customization)), + ...this._flattenClientCustomizations().map(item => this._applyEnablement(item.customization)), ]; const entry = this._discoveredEntry(); const discovered = entry?.currentCustomizations() ?? []; @@ -2989,9 +3015,29 @@ class SessionPluginController extends Disposable { return result; } + /** + * The union of every active client's resolved customizations, + * deduplicated by URI with the first-inserted client winning. Order + * follows client insertion order, then per-client order. + */ + private _flattenClientCustomizations(): readonly IResolvedCustomization[] { + const seen = new Set(); + const result: IResolvedCustomization[] = []; + for (const client of this._clients.values()) { + for (const item of client.customizations) { + if (seen.has(item.customization.uri)) { + continue; + } + seen.add(item.customization.uri); + result.push(item); + } + } + return result; + } + /** * Settled variant of {@link getCustomizations}: awaits the in-flight - * host sync, the in-flight client sync, and the discovered entry's + * host sync, every in-flight client sync, and the discovered entry's * initial scan + parse before snapshotting the list. Callers that * publish customizations into session state at session creation time * MUST use this — the synchronous variant can return an empty list @@ -3002,7 +3048,7 @@ class SessionPluginController extends Disposable { const entry = this._discoveredEntry(); await Promise.all([ this._parent.hostSync().catch(err => this._parent.logService.warn('[Copilot:SessionPluginController] Host customization update failed', err)), - this._clientSync.catch(err => this._parent.logService.warn('[Copilot:SessionPluginController] Client customization sync failed', err)), + ...[...this._clients.values()].map(client => client.sync.catch(err => this._parent.logService.warn('[Copilot:SessionPluginController] Client customization sync failed', err))), entry?.whenSettled(), ]); return this.getCustomizations(); @@ -3011,15 +3057,15 @@ class SessionPluginController extends Disposable { /** Returns the parsed plugins currently enabled for this session, awaiting any pending sync. */ public async getAppliedPlugins(): Promise { const entry = this._discoveredEntry(); - const [host, client] = await Promise.all([ + const [host] = await Promise.all([ this._parent.hostSync().catch(err => { this._parent.logService.warn('[Copilot:SessionPluginController] Host customization update failed', err); return this._parent.hostCustomizations(); }), - this._clientSync.catch(err => { + ...[...this._clients.values()].map(client => client.sync.catch(err => { this._parent.logService.warn('[Copilot:SessionPluginController] Client customization sync failed', err); - return this._clientCustomizations; - }), + return client.customizations; + })), entry?.whenSettled(), ]); @@ -3030,7 +3076,7 @@ class SessionPluginController extends Disposable { return [ ...host.filter(item => !!item.plugin && this._isEnabled(item.customization)) .map(item => ({ ...item.plugin!, pluginDir: item.pluginDir })), - ...client.filter(item => !!item.plugin && this._isEnabled(item.customization)) + ...this._flattenClientCustomizations().filter(item => !!item.plugin && this._isEnabled(item.customization)) .map(item => ({ ...item.plugin!, pluginDir: item.pluginDir })), ...sessionPlugins, ]; @@ -3048,7 +3094,10 @@ class SessionPluginController extends Disposable { } /** - * Sync the published client customizations for this session. + * Sync the published customizations for a single client of this session, + * keyed by `clientId`. Replaces only that client's slice; other clients' + * customizations are untouched. The published session-state list is the + * union across all clients. * * @param quiet when `true`, suppress {@link onDidPublish} events for * this sync. Used during eager-create paths where there is no @@ -3057,9 +3106,14 @@ class SessionPluginController extends Disposable { */ public sync(clientId: string, customizations: ClientPluginCustomization[], options?: { quiet?: boolean }) { const quiet = options?.quiet === true; - this._clientId = clientId; - const revision = ++this._clientRevision; - this._clientCustomizations = customizations.map(customization => ({ + let client = this._clients.get(clientId); + if (!client) { + client = { revision: 0, customizations: [], sync: Promise.resolve([]), inputs: [] }; + this._clients.set(clientId, client); + } + const revision = ++client.revision; + client.inputs = customizations; + client.customizations = customizations.map(customization => ({ customization: { ...customization, clientId, @@ -3074,7 +3128,7 @@ class SessionPluginController extends Disposable { }); } const published = new Map(); - for (const customization of this._clientCustomizations) { + for (const customization of client.customizations) { const enabled = this._applyEnablement(customization.customization); published.set(enabled.uri, enabled); } @@ -3092,13 +3146,13 @@ class SessionPluginController extends Disposable { } }; - const prev = this._clientSync; - const promise = this._clientSync = prev.catch(err => { + const prev = client.sync; + const promise = client.sync = prev.catch(err => { this._parent.logService.warn('[Copilot:SessionPluginController] Previous customization sync failed', err); }).then(async () => { const inputByUri = new Map(customizations.map(c => [c.uri, c])); const result = await this._parent.pluginManager.syncCustomizations(clientId, customizations, status => { - if (revision !== this._clientRevision) { + if (revision !== client.revision) { return; } publishUpdate({ @@ -3108,8 +3162,8 @@ class SessionPluginController extends Disposable { }); const resolved = await Promise.all(result.map(item => this._parent.resolveSyncedCustomization(item, clientId, inputByUri.get(item.customization.uri)))); - if (revision === this._clientRevision) { - this._clientCustomizations = resolved; + if (revision === client.revision) { + client.customizations = resolved; for (const item of resolved) { publishUpdate(item); } @@ -3124,31 +3178,57 @@ class SessionPluginController extends Disposable { } /** - * Re-issue the last client sync if any previously-synced customization - * is currently in an error state. Used to recover from transient - * sync failures (e.g. a `vscode-agent-host://` connection drop during - * reconnection) at message boundaries. Re-syncs **only** the errored - * items and always non-quiet so listeners observe recovery. + * Remove a client's customization contribution from this session, + * publishing the updated (union) customization list so the removed + * client's plugins disappear from session state. */ - public async retryFailedClientSyncIfNeeded(): Promise { - await this._clientSync.catch(() => { }); - if (!this._clientId) { - return; - } - const errored = this._clientCustomizations.filter(item => - item.customization.load?.kind === CustomizationLoadStatus.Error - && item.input !== undefined - ); - if (errored.length === 0) { + public removeClient(clientId: string): void { + const client = this._clients.get(clientId); + if (!client) { return; } - const inputs = errored.map(item => item.input!); - this._parent.logService.info(`[Copilot:SessionPluginController] Retrying ${inputs.length} previously-failed client customization(s)`); - await this.sync(this._clientId, inputs).catch(err => { - this._parent.logService.warn('[Copilot:SessionPluginController] Retried client customization sync failed', err); + // Invalidate any in-flight sync for this client by bumping its + // revision so the late continuation's `revision === client.revision` + // guards fail and it does not re-publish the removed client's + // customizations. + client.revision++; + this._clients.delete(clientId); + this._onDidPublish.fire({ + type: ActionType.SessionCustomizationsChanged, + customizations: [...this.getCustomizations()], }); } + /** The raw input customizations last synced for `clientId` (empty when absent). */ + public clientInputs(clientId: string): readonly ClientPluginCustomization[] { + return this._clients.get(clientId)?.inputs ?? []; + } + + /** + * Re-issue each client's last sync if any of its previously-synced + * customizations is currently in an error state. Used to recover from + * transient sync failures (e.g. a `vscode-agent-host://` connection drop + * during reconnection) at message boundaries. Re-syncs **only** the + * errored items and always non-quiet so listeners observe recovery. + */ + public async retryFailedClientSyncIfNeeded(): Promise { + await Promise.all([...this._clients.values()].map(client => client.sync.catch(() => { }))); + for (const [clientId, client] of [...this._clients]) { + const errored = client.customizations.filter(item => + item.customization.load?.kind === CustomizationLoadStatus.Error + && item.input !== undefined + ); + if (errored.length === 0) { + continue; + } + const inputs = errored.map(item => item.input!); + this._parent.logService.info(`[Copilot:SessionPluginController] Retrying ${inputs.length} previously-failed client customization(s) for ${clientId}`); + await this.sync(clientId, inputs).catch(err => { + this._parent.logService.warn('[Copilot:SessionPluginController] Retried client customization sync failed', err); + }); + } + } + private _discoveredEntry(): SessionDiscoveredEntry | undefined { if (!this._directory) { return undefined; @@ -3179,23 +3259,59 @@ class SessionPluginController extends Disposable { } /** - * Tracks per-session active client contributions (tools and plugins). - * Owns the session's {@link SessionPluginController}, which is the - * authoritative source for both the plugin snapshot (host + client + - * session-discovered) and per-session action events. Disposing this - * tears down the controller and any disk watchers it created. + * A per-(session, clientId) handle returned by + * {@link CopilotAgent.getOrCreateActiveClient}. Reads/writes flow straight + * through to the owning session's {@link ActiveClient} (the multi-client + * container), so assigning `tools` / `customizations` updates only this + * client's slice. + */ +class CopilotActiveClientHandle implements IActiveClient { + constructor( + private readonly _owner: ActiveClient, + readonly clientId: string, + readonly displayName: string | undefined, + ) { } + + get tools(): readonly ToolDefinition[] { + return this._owner.toolSet.get(this.clientId); + } + set tools(tools: readonly ToolDefinition[]) { + this._owner.toolSet.set(this.clientId, tools); + } + + get customizations(): readonly ClientPluginCustomization[] { + return this._owner.pluginController.clientInputs(this.clientId); + } + set customizations(customizations: readonly ClientPluginCustomization[]) { + // Fire-and-forget: progress and the settled result flow out via the + // controller's `onDidPublish` session actions, not the setter. + this._owner.pluginController.sync(this.clientId, [...customizations]).catch(() => { /* logged inside sync */ }); + } +} + +/** + * Tracks per-session active client contributions (tools and plugins) across + * potentially several active clients. Owns the session's + * {@link SessionPluginController}, which is the authoritative source for both + * the plugin snapshot (host + all clients + session-discovered) and + * per-session action events, and the {@link ActiveClientToolSet} that merges + * every client's tools. Disposing this tears down the controller and any disk + * watchers it created. */ class ActiveClient extends Disposable { /** - * Live holder of the owning `clientId` and contributed tools. Shared by - * reference with the session's {@link CopilotAgentSession} so a window - * reload (new `clientId`, identical tools) is reflected at tool-call - * stamp time without restarting the SDK session. + * Live, multi-client registry of contributed tools. Shared by reference + * with the session's {@link CopilotAgentSession} so a window reload (new + * `clientId`, identical tools) is reflected at tool-call stamp time without + * restarting the SDK session, and so tool calls are attributed to the + * contributing client. */ - readonly state = new ActiveClientState(); + readonly toolSet = new ActiveClientToolSet(); public readonly pluginController: SessionPluginController; + private readonly _handles = new Map(); + constructor( private readonly _sessionUri: URI, pluginController: SessionPluginController, @@ -3211,13 +3327,26 @@ class ActiveClient extends Disposable { })); } - updateTools(clientId: string | undefined, tools: readonly ToolDefinition[]): void { - this.state.update(clientId, tools); + /** Get (or lazily create) the stable handle for `clientId`. */ + getOrCreateHandle(clientId: string, displayName: string | undefined): CopilotActiveClientHandle { + let handle = this._handles.get(clientId); + if (!handle) { + handle = new CopilotActiveClientHandle(this, clientId, displayName); + this._handles.set(clientId, handle); + } + return handle; + } + + /** Drop a client's tool and customization contributions from this session. */ + removeClient(clientId: string): void { + this._handles.delete(clientId); + this.toolSet.delete(clientId); + this.pluginController.removeClient(clientId); } async snapshot(): Promise { return { - tools: this.state.tools, + tools: this.toolSet.merged(), plugins: await this.pluginController.getAppliedPlugins(), mcpServers: this._getMcpServers(), }; @@ -3232,9 +3361,9 @@ class ActiveClient extends Disposable { /** * Returns `true` when the SDK session must be disposed and resumed to * pick up a changed config. Compares ONLY plugins and the structural - * tool set (name + description + inputSchema). The `clientId` is - * deliberately excluded — a clientId-only change is reflected live via - * {@link state} and never requires a restart. + * (merged) tool set (name + description + inputSchema). The owning + * `clientId`s are deliberately excluded — a clientId-only change is + * reflected live via {@link toolSet} and never requires a restart. */ async requiresRestart(snap: IActiveClientSnapshot): Promise { const plugins = await this.pluginController.getAppliedPlugins(); @@ -3244,6 +3373,6 @@ class ActiveClient extends Disposable { if (!equals(snap.mcpServers, this._getMcpServers())) { return true; } - return !this.state.structuralEquals({ tools: snap.tools }); + return !this.toolSet.structuralEquals(snap.tools); } } diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts index 6db6f3d74a415d..e211df7b51c63c 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts @@ -41,7 +41,7 @@ import { IAgentConfigurationService } from '../agentConfigurationService.js'; import type { IExitPlanModeResponse } from './copilotAgent.js'; import { CopilotSessionWrapper } from './copilotSessionWrapper.js'; import { clientToolNamesFromSnapshot, type CopilotSessionLaunchPlan, type IActiveClientSnapshot, type ICopilotSessionLauncher, type ICopilotSessionRuntime } from './copilotSessionLauncher.js'; -import { ActiveClientState } from '../activeClientState.js'; +import { ActiveClientToolSet } from '../activeClientState.js'; import { PendingRequestRegistry } from '../../common/pendingRequestRegistry.js'; import { buildCopilotSystemNotification } from './copilotSystemNotification.js'; import { isPromptInvokedCopilotSlashCommand, isRuntimeCopilotSlashCommand, parseLeadingSlashCommand } from './copilotSlashCommandCompletionProvider.js'; @@ -313,13 +313,15 @@ export interface ICopilotAgentSessionOptions { */ readonly resolveMcpChildId: (serverName: string) => string | undefined; /** - * Live holder of the owning client's identity, shared by reference with - * the agent's per-session {@link ActiveClient}. Read at tool-call stamp - * time so a window reload (new `clientId`, identical tools) stamps with - * the current id. When omitted, a fresh state seeded from - * {@link clientSnapshot} is used (test / standalone path). + * Live registry of every active client's tool contributions, shared by + * reference with the agent's per-session {@link ActiveClient}. Read at + * tool-call stamp time so a window reload (new `clientId`, identical + * tools) stamps with the current owning id, and so each tool call is + * attributed to whichever client contributed it. When omitted, a fresh + * empty registry is used (test / standalone path) and client tool calls + * are left unstamped. */ - readonly activeClientState?: ActiveClientState; + readonly activeClientToolSet?: ActiveClientToolSet; /** * Server-side host for the agent host's server tools. When provided, the * session advertises the server tools (feedback "comments" today, more in @@ -499,7 +501,7 @@ export class CopilotAgentSession extends Disposable { * subsequent client tool calls with the current id rather than the one * frozen into {@link _appliedSnapshot}. */ - private readonly _activeClientState: ActiveClientState; + private readonly _activeClientToolSet: ActiveClientToolSet; /** Tool names that are client-provided, derived from snapshot. */ private readonly _clientToolNames: ReadonlySet; /** Deferred promises for pending client tool calls, keyed by toolCallId. */ @@ -563,15 +565,11 @@ export class CopilotAgentSession extends Disposable { this._appliedSnapshot = options.clientSnapshot ?? { tools: [], plugins: [], mcpServers: {} }; this._clientToolNames = clientToolNamesFromSnapshot(this._appliedSnapshot); - // Share the agent's live ActiveClientState when provided so clientId - // changes are observed at stamp time. Standalone / test construction - // seeds a private instance with the applied tools and no owning client. - if (options.activeClientState) { - this._activeClientState = options.activeClientState; - } else { - this._activeClientState = new ActiveClientState(); - this._activeClientState.update(undefined, this._appliedSnapshot.tools); - } + // Share the agent's live ActiveClientToolSet when provided so client + // contributions (and owner identity) are observed at stamp time. + // Standalone / test construction uses a fresh empty registry, which + // leaves client tool calls unstamped (no owning client). + this._activeClientToolSet = options.activeClientToolSet ?? new ActiveClientToolSet(); this._databaseRef = sessionDataService.openDatabase(options.sessionUri); this._register(toDisposable(() => this._databaseRef.dispose())); @@ -2212,8 +2210,9 @@ export class CopilotAgentSession extends Disposable { let contributor: { readonly kind: ToolCallContributorKind.Client; readonly clientId: string } | { readonly kind: ToolCallContributorKind.MCP; readonly customizationId: string } | undefined; const isClientTool = this._clientToolNames.has(e.data.toolName); - if (isClientTool && this._activeClientState.clientId) { - contributor = { kind: ToolCallContributorKind.Client, clientId: this._activeClientState.clientId }; + const ownerClientId = isClientTool ? this._activeClientToolSet.ownerOf(e.data.toolName) : undefined; + if (ownerClientId) { + contributor = { kind: ToolCallContributorKind.Client, clientId: ownerClientId }; } else if (e.data.mcpServerName) { const customizationId = this._mcpCustomizations.customizationIdForServer(e.data.mcpServerName); if (customizationId !== undefined) { diff --git a/src/vs/platform/agentHost/node/copilot/copilotSessionLauncher.ts b/src/vs/platform/agentHost/node/copilot/copilotSessionLauncher.ts index a28716b1e4e99e..51d53d2dc6ad36 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotSessionLauncher.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotSessionLauncher.ts @@ -15,7 +15,7 @@ import { AgentHostSandboxConfigKey, sandboxConfigSchema } from '../../common/san import { IAgentConfigurationService } from '../agentConfigurationService.js'; import { IAgentHostTerminalManager } from '../agentHostTerminalManager.js'; import type { ModelSelection, ToolDefinition } from '../../common/state/protocol/state.js'; -import type { ActiveClientState } from '../activeClientState.js'; +import type { ActiveClientToolSet } from '../activeClientState.js'; import { CopilotSessionWrapper } from './copilotSessionWrapper.js'; import { ShellManager, createShellTools, type IUnsandboxedCommandConfirmationRequest } from './copilotShellTools.js'; import { toSdkHooks, toSdkInstructionDirectories, toSdkMcpServers, toSdkMcpServersFromConfigMap, toSdkSessionCustomAgents, toSdkSkillDirectories } from './copilotPluginConverters.js'; @@ -54,9 +54,9 @@ type CopilotSessionLaunchConfig = ResumeSessionConfig & { * Immutable snapshot of the active client's structural contributions at * session creation time. Used to detect when the session needs to be * refreshed. Root MCP servers participate in restart detection because they - * are merged into the SDK session config. The owning `clientId` is + * are merged into the SDK session config. The owning `clientId`s are * deliberately NOT part of this snapshot: client identity is tracked live via - * {@link ActiveClientState} so a window + * {@link ActiveClientToolSet} so a window * reload (new `clientId`, identical tools/plugins) does not force a restart. */ export interface IActiveClientSnapshot { @@ -106,12 +106,13 @@ interface ICopilotSessionLaunchBase { readonly resolvedAgentName: string | undefined; readonly snapshot: IActiveClientSnapshot; /** - * Live, long-lived holder of the owning client's identity. Read at - * tool-call stamp time so a window reload (new `clientId`, identical - * tools) stamps subsequent client tool calls with the current id - * rather than the one frozen into {@link snapshot} at creation. + * Live, long-lived registry of every active client's tool contributions. + * Read at tool-call stamp time so a window reload (new `clientId`, + * identical tools) stamps subsequent client tool calls with the current + * owning id rather than the one frozen into {@link snapshot} at creation, + * and so a tool call is attributed to whichever client contributed it. */ - readonly activeClientState: ActiveClientState; + readonly activeClientToolSet: ActiveClientToolSet; readonly shellManager: ShellManager | undefined; readonly githubToken: string | undefined; } diff --git a/src/vs/platform/agentHost/node/protocolServerHandler.ts b/src/vs/platform/agentHost/node/protocolServerHandler.ts index 02e4b92a899320..968de598cd43fa 100644 --- a/src/vs/platform/agentHost/node/protocolServerHandler.ts +++ b/src/vs/platform/agentHost/node/protocolServerHandler.ts @@ -714,6 +714,8 @@ export class ProtocolServerHandler extends Disposable { } })); + this._reconcileActiveClientsAfterReconnect(client); + if (canReplay) { const actions: ActionEnvelope[] = []; for (const envelope of this._replayBuffer) { @@ -728,22 +730,82 @@ export class ProtocolServerHandler extends Disposable { return { type: 'snapshot', snapshots: snapshots.filter((s): s is IStateSnapshot => s !== undefined) }; } + /** + * Release a client from every session where it is still an active client + * but did not resubscribe during a reconnect. The set of resubscribed + * sessions is gathered from every live connection the client currently + * holds (not just the reconnecting one) so an overlapping connection that + * still subscribes to a session keeps the client active there. + */ + private _reconcileActiveClientsAfterReconnect(client: IConnectedClient): void { + const record = this._clients.get(client.clientId); + const resubscribed = new Set(); + for (const connection of record?.connections ?? [client]) { + for (const sub of connection.subscriptions.values()) { + if (sub.kind === ChannelKind.State) { + resubscribed.add(sub.uri); + } + } + } + for (const session of this._stateManager.getSessionUris()) { + if (resubscribed.has(session)) { + continue; + } + const state = this._stateManager.getSessionState(session); + if (state && this._isActiveClient(state, client.clientId)) { + this._releaseActiveClientForSession(session, client.clientId); + } + } + } + private _handleClientDisconnected(clientId: string): void { for (const session of this._stateManager.getSessionUris()) { const state = this._stateManager.getSessionState(session); + const isActive = state ? this._isActiveClient(state, clientId) : false; const ownsPendingToolCall = state ? this._hasPendingClientToolCall(state, clientId) : false; - if (state?.activeClient?.clientId === clientId) { - this._stateManager.dispatchServerAction(session, { - type: ActionType.SessionActiveClientChanged, - activeClient: null, - }); - } - if (state?.activeClient?.clientId === clientId || ownsPendingToolCall) { + // Keep the client marked active during the grace window so a quick + // reconnect that resubscribes can retain its slot. The disconnect + // timeout removes the active client (and fails its pending tool + // calls) if it never returns; an explicit unsubscribe or a + // reconnect without resubscription removes it sooner. + if (isActive || ownsPendingToolCall) { this._startClientToolCallDisconnectTimeout(clientId, session); } } } + /** Whether `clientId` is one of the session's active clients. */ + private _isActiveClient(state: SessionState, clientId: string): boolean { + return state.activeClients.some(c => c.clientId === clientId); + } + + /** + * Remove `clientId` from a session's active clients, if present. Dispatched + * as a server action so the removal is reflected in state and broadcast to + * the remaining subscribers. + */ + private _removeActiveClient(session: string, clientId: string): void { + const state = this._stateManager.getSessionState(session); + if (state && this._isActiveClient(state, clientId)) { + this._stateManager.dispatchServerAction(session, { + type: ActionType.SessionActiveClientRemoved, + clientId, + }); + } + } + + /** + * Release a client from a session: clear its pending disconnect timeout, + * fail any client tool calls it still owns, and remove it from the active + * clients. Used by the explicit-unsubscribe and reconnect-reconciliation + * paths to drop a client that has left a session. + */ + private _releaseActiveClientForSession(session: string, clientId: string): void { + this._clearClientToolCallDisconnectTimeout(clientId, session); + this._completeDisconnectedClientToolCalls(clientId, session); + this._removeActiveClient(session, clientId); + } + /** * Yields every still-pending client-contributed tool call in `state`'s * active turn, paired with its owning `clientId`. Single source of truth @@ -779,10 +841,9 @@ export class ProtocolServerHandler extends Disposable { } private _hasReplacementActiveClientTool(state: SessionState, clientId: string, toolName: string): boolean { - const activeClient = state.activeClient; - return activeClient !== undefined - && activeClient.clientId !== clientId - && activeClient.tools.some(tool => tool.name === toolName); + return state.activeClients.some(client => + client.clientId !== clientId + && client.tools.some(tool => tool.name === toolName)); } /** @@ -805,8 +866,7 @@ export class ProtocolServerHandler extends Disposable { const elapsed = Date.now() - record.lastSeenAt; const delay = Math.max(0, CLIENT_TOOL_CALL_DISCONNECT_TIMEOUT - elapsed); record.disconnectTimeouts.set(session, disposableTimeout(() => { - record.disconnectTimeouts.deleteAndDispose(session); - this._completeDisconnectedClientToolCalls(clientId, session); + this._releaseActiveClientForSession(session, clientId); }, delay)); } @@ -1372,6 +1432,7 @@ export class ProtocolServerHandler extends Disposable { return; } this._agentService.unsubscribe(URI.parse(sub.uri), client.clientId); + this._releaseActiveClientForSession(sub.uri, client.clientId); } else if (sub.kind === ChannelKind.ResourceWatch) { this._agentService.onResourceWatchUnsubscribed(sub.uri); } diff --git a/src/vs/platform/agentHost/test/common/agentSubscription.test.ts b/src/vs/platform/agentHost/test/common/agentSubscription.test.ts index 0274bd41abf8a6..62914cc1beae3e 100644 --- a/src/vs/platform/agentHost/test/common/agentSubscription.test.ts +++ b/src/vs/platform/agentHost/test/common/agentSubscription.test.ts @@ -35,6 +35,7 @@ function makeSessionState(sessionUri: string, overrides?: Partial) project: { uri: 'file:///test-project', displayName: 'Test Project' }, }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], ...overrides, }; diff --git a/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostProtocolClient.test.ts b/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostProtocolClient.test.ts index 7dc3039be810cd..8e79883bfc329d 100644 --- a/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostProtocolClient.test.ts +++ b/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostProtocolClient.test.ts @@ -18,7 +18,7 @@ import { AgentHostPermissionMode, AgentHostResourcePermissionError, IAgentHostRe import { ContentEncoding, ReconnectResultType } from '../../common/state/protocol/commands.js'; import { AhpErrorCodes } from '../../common/state/protocol/errors.js'; import { PROTOCOL_VERSION } from '../../common/state/protocol/version/registry.js'; -import { ActionType, type SessionActiveClientChangedAction, type SessionTitleChangedAction } from '../../common/state/sessionActions.js'; +import { ActionType, type SessionActiveClientSetAction, type SessionActiveClientRemovedAction, type SessionTitleChangedAction } from '../../common/state/sessionActions.js'; import { 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'; @@ -392,9 +392,9 @@ suite('RemoteAgentHostProtocolClient', () => { transport.fireMessage({ jsonrpc: '2.0', id: 1, result: { entries: [] } }); // Late notification — must not fan out as an action event. - const lateAction: SessionActiveClientChangedAction = { - type: ActionType.SessionActiveClientChanged, - activeClient: null, + const lateAction: SessionActiveClientRemovedAction = { + type: ActionType.SessionActiveClientRemoved, + clientId: 'c1', }; transport.fireMessage({ jsonrpc: '2.0', @@ -692,13 +692,13 @@ suite('RemoteAgentHostProtocolClient', () => { return { service, calls }; } - test('SessionActiveClientChanged dispatches implicit reads for each customization', () => { + test('SessionActiveClientSet dispatches implicit reads for each customization', () => { const { service, calls } = createCapturingPermissionService(); const { client } = createClient(undefined, service); const sessionUri = URI.parse('ahp-session:/test'); client.dispatch(sessionUri.toString(), { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'c1', tools: [], @@ -724,7 +724,7 @@ suite('RemoteAgentHostProtocolClient', () => { const sessionUri = URI.parse('ahp-session:/test'); client.dispatch(sessionUri.toString(), { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'c1', tools: [], @@ -746,8 +746,8 @@ suite('RemoteAgentHostProtocolClient', () => { const { client } = createClient(undefined, service); const sessionUri = URI.parse('ahp-session:/test'); - const action: SessionActiveClientChangedAction = { - type: ActionType.SessionActiveClientChanged, + const action: SessionActiveClientSetAction = { + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'c1', tools: [], @@ -763,14 +763,14 @@ suite('RemoteAgentHostProtocolClient', () => { assert.strictEqual(calls.length, 1); }); - test('null activeClient does not crash', () => { + test('active client removal does not crash', () => { const { service, calls } = createCapturingPermissionService(); const { client } = createClient(undefined, service); const sessionUri = URI.parse('ahp-session:/test'); client.dispatch(sessionUri.toString(), { - type: ActionType.SessionActiveClientChanged, - activeClient: null, + type: ActionType.SessionActiveClientRemoved, + clientId: 'c1', }); assert.strictEqual(calls.length, 0); diff --git a/src/vs/platform/agentHost/test/node/agentService.test.ts b/src/vs/platform/agentHost/test/node/agentService.test.ts index 387d4f4e3ce1d9..5db1c8a2dfd308 100644 --- a/src/vs/platform/agentHost/test/node/agentService.test.ts +++ b/src/vs/platform/agentHost/test/node/agentService.test.ts @@ -1583,11 +1583,11 @@ suite('AgentService (node dispatcher)', () => { const session = await service.createSession({ provider: 'copilot', activeClient }); assert.deepStrictEqual({ - activeClient: service.stateManager.getSessionState(session.toString())?.activeClient, - dispatchedActiveClientChanged: envelopes.some(e => e.action.type === ActionType.SessionActiveClientChanged), + activeClients: service.stateManager.getSessionState(session.toString())?.activeClients, + dispatchedActiveClientSet: envelopes.some(e => e.action.type === ActionType.SessionActiveClientSet), }, { - activeClient, - dispatchedActiveClientChanged: false, + activeClients: [activeClient], + dispatchedActiveClientSet: false, }); }); @@ -1596,7 +1596,7 @@ suite('AgentService (node dispatcher)', () => { const session = await service.createSession({ provider: 'copilot' }); - assert.strictEqual(service.stateManager.getSessionState(session.toString())?.activeClient, undefined); + assert.deepStrictEqual(service.stateManager.getSessionState(session.toString())?.activeClients, []); }); }); diff --git a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts index 01a10fbcedb89b..b0488ac0f74d1c 100644 --- a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts +++ b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts @@ -244,7 +244,7 @@ suite('AgentSideEffects', () => { test('logs telemetry when sending a direct user message', () => { setupSession(); const activeClientAction: SessionAction = { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'test-client', tools: [{ name: 'testTool', inputSchema: { type: 'object' } }], @@ -1135,9 +1135,9 @@ suite('AgentSideEffects', () => { }); }); - // ---- handleAction: session/activeClientChanged ---------------------- + // ---- handleAction: session/activeClientSet ---------------------- - suite('handleAction — session/activeClientChanged', () => { + suite('handleAction — session/activeClientSet', () => { setup(() => { disposables.add(sideEffects.registerProgressListener(agent)); @@ -1155,7 +1155,7 @@ suite('AgentSideEffects', () => { disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e))); const action: SessionAction = { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'test-client', tools: [], @@ -1187,7 +1187,7 @@ suite('AgentSideEffects', () => { const pluginAClient: ClientPluginCustomization = { type: CustomizationType.Plugin, id: customizationId('file:///plugin-a'), uri: 'file:///plugin-a', name: 'Plugin A', enabled: true }; let currentCustomizations: readonly Customization[] = []; agent.getSessionCustomizations = async () => currentCustomizations; - agent.setClientCustomizations = async (session, clientId, customizations) => { + agent.syncClientCustomizations = (session, clientId, customizations) => { agent.setClientCustomizationsCalls.push({ clientId, customizations }); const loading: PluginCustomization = { ...pluginAClient, load: { kind: CustomizationLoadStatus.Loading } }; currentCustomizations = [loading]; @@ -1199,17 +1199,19 @@ suite('AgentSideEffects', () => { customizations: [...currentCustomizations], }, }); - await new Promise(resolve => setTimeout(resolve, 0)); - const loaded: PluginCustomization = { ...pluginAClient, load: { kind: CustomizationLoadStatus.Loaded } }; - currentCustomizations = [loaded]; - agent.fireProgress({ - kind: 'action', - session, - action: { - type: ActionType.SessionCustomizationUpdated, - customization: loaded, - }, - }); + void (async () => { + await new Promise(resolve => setTimeout(resolve, 0)); + const loaded: PluginCustomization = { ...pluginAClient, load: { kind: CustomizationLoadStatus.Loaded } }; + currentCustomizations = [loaded]; + agent.fireProgress({ + kind: 'action', + session, + action: { + type: ActionType.SessionCustomizationUpdated, + customization: loaded, + }, + }); + })(); return currentCustomizations.map(customization => ({ customization: customization as PluginCustomization })); }; @@ -1217,7 +1219,7 @@ suite('AgentSideEffects', () => { disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e))); sideEffects.handleAction(sessionUri.toString(), { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'test-client', tools: [], @@ -1249,7 +1251,7 @@ suite('AgentSideEffects', () => { disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e))); const action: SessionAction = { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'test-client', tools: [] @@ -1270,18 +1272,17 @@ suite('AgentSideEffects', () => { }); }); - test('clears client customizations when activeClient is null', () => { + test('removes the active client when it is removed', () => { setupSession(); const action: SessionAction = { - type: ActionType.SessionActiveClientChanged, - activeClient: null, + type: ActionType.SessionActiveClientRemoved, + clientId: 'test-client', }; sideEffects.handleAction(sessionUri.toString(), action); - assert.deepStrictEqual(agent.setClientCustomizationsCalls, [{ - clientId: '', - customizations: [], + assert.deepStrictEqual(agent.removeActiveClientCalls, [{ + clientId: 'test-client', }]); }); }); diff --git a/src/vs/platform/agentHost/test/node/claudeAgent.test.ts b/src/vs/platform/agentHost/test/node/claudeAgent.test.ts index 63a299aaf01fa5..328b29c51366f7 100644 --- a/src/vs/platform/agentHost/test/node/claudeAgent.test.ts +++ b/src/vs/platform/agentHost/test/node/claudeAgent.test.ts @@ -3398,7 +3398,7 @@ suite('ClaudeAgent', () => { const sessionId = AgentSession.id(created.session); const tools: ToolDefinition[] = [{ name: 'echo', description: 'Echo back', inputSchema: { type: 'object', properties: { msg: { type: 'string' } }, required: ['msg'] } }]; - agent.setClientTools(created.session, 'client-1', tools); + agent.getOrCreateActiveClient(created.session, { clientId: 'client-1' }).tools = tools; sdk.nextQueryMessages = [makeSystemInitMessage(sessionId), makeResultSuccess(sessionId)]; await agent.sendMessage(created.session, 'go', undefined, 'turn-1'); @@ -3432,7 +3432,7 @@ suite('ClaudeAgent', () => { await agent.sendMessage(created.session, 'first', undefined, 'turn-1'); assert.strictEqual(sdk.startupCallCount, 1, 'first materialize'); - agent.setClientTools(created.session, 'client-1', [{ name: 'echo', inputSchema: { type: 'object' } }]); + agent.getOrCreateActiveClient(created.session, { clientId: 'client-1' }).tools = [{ name: 'echo', inputSchema: { type: 'object' } }]; sdk.queryAdvance = undefined; advance.complete(); await agent.sendMessage(created.session, 'second', undefined, 'turn-2'); @@ -3463,11 +3463,11 @@ suite('ClaudeAgent', () => { ]; const tools: ToolDefinition[] = [{ name: 'echo', description: 'e', inputSchema: { type: 'object' } }]; - agent.setClientTools(created.session, 'c1', tools); + agent.getOrCreateActiveClient(created.session, { clientId: 'c1' }).tools = tools; await agent.sendMessage(created.session, 'first', undefined, 'turn-1'); assert.strictEqual(sdk.startupCallCount, 1, 'first materialize'); - agent.setClientTools(created.session, 'c1', [{ name: 'echo', description: 'e', inputSchema: { type: 'object' } }]); + agent.getOrCreateActiveClient(created.session, { clientId: 'c1' }).tools = [{ name: 'echo', description: 'e', inputSchema: { type: 'object' } }]; advance.complete(); await agent.sendMessage(created.session, 'second', undefined, 'turn-2'); @@ -3477,7 +3477,7 @@ suite('ClaudeAgent', () => { test('setClientTools on an unknown session id is silently dropped', () => { const { agent } = createTestContext(disposables); assert.doesNotThrow(() => { - agent.setClientTools(URI.parse('claude:/never-existed'), 'c1', [{ name: 't', inputSchema: { type: 'object' } }]); + agent.getOrCreateActiveClient(URI.parse('claude:/never-existed'), { clientId: 'c1' }).tools = [{ name: 't', inputSchema: { type: 'object' } }]; }); }); @@ -3486,7 +3486,7 @@ suite('ClaudeAgent', () => { await agent.authenticate(GITHUB_COPILOT_PROTECTED_RESOURCE.resource, 'tok'); const created = await agent.createSession({ workingDirectory: URI.file('/work') }); const sessionId = AgentSession.id(created.session); - agent.setClientTools(created.session, 'c1', [{ name: 'echo', inputSchema: { type: 'object' } }]); + agent.getOrCreateActiveClient(created.session, { clientId: 'c1' }).tools = [{ name: 'echo', inputSchema: { type: 'object' } }]; sdk.nextQueryMessages = [makeSystemInitMessage(sessionId), makeResultSuccess(sessionId)]; await agent.sendMessage(created.session, 'go', undefined, 'turn-1'); @@ -3521,7 +3521,7 @@ suite('ClaudeAgent', () => { await agent.authenticate(GITHUB_COPILOT_PROTECTED_RESOURCE.resource, 'tok'); const created = await agent.createSession({ workingDirectory: URI.file('/work') }); const sessionId = AgentSession.id(created.session); - agent.setClientTools(created.session, 'c1', [{ name: 'echo', inputSchema: { type: 'object' } }]); + agent.getOrCreateActiveClient(created.session, { clientId: 'c1' }).tools = [{ name: 'echo', inputSchema: { type: 'object' } }]; sdk.nextQueryMessages = [makeSystemInitMessage(sessionId), makeResultSuccess(sessionId)]; await agent.sendMessage(created.session, 'go', undefined, 'turn-1'); await assert.doesNotReject(agent.disposeSession(created.session)); @@ -3532,11 +3532,11 @@ suite('ClaudeAgent', () => { await agent.authenticate(GITHUB_COPILOT_PROTECTED_RESOURCE.resource, 'tok'); const created = await agent.createSession({ workingDirectory: URI.file('/work') }); const sessionId = AgentSession.id(created.session); - agent.setClientTools(created.session, 'c1', [{ name: 'echo', inputSchema: { type: 'object' } }]); + agent.getOrCreateActiveClient(created.session, { clientId: 'c1' }).tools = [{ name: 'echo', inputSchema: { type: 'object' } }]; sdk.nextQueryMessages = [makeSystemInitMessage(sessionId), makeResultSuccess(sessionId)]; await agent.sendMessage(created.session, 'first', undefined, 'turn-1'); // Change tools to force a rebind path (must use yield-restart, NOT Query.setMcpServers). - agent.setClientTools(created.session, 'c1', [{ name: 'echo2', inputSchema: { type: 'object' } }]); + agent.getOrCreateActiveClient(created.session, { clientId: 'c1' }).tools = [{ name: 'echo2', inputSchema: { type: 'object' } }]; sdk.nextQueryMessages = [makeSystemInitMessage(sessionId), makeResultSuccess(sessionId)]; await agent.sendMessage(created.session, 'second', undefined, 'turn-2'); // If `Query.setMcpServers` had been called, `FakeQuery.setMcpServers` would have thrown. @@ -3550,7 +3550,7 @@ suite('ClaudeAgent', () => { const sessionId = AgentSession.id(created.session); // Initial snapshot before materialize starts. - agent.setClientTools(created.session, 'c1', [{ name: 'first', inputSchema: { type: 'object' } }]); + agent.getOrCreateActiveClient(created.session, { clientId: 'c1' }).tools = [{ name: 'first', inputSchema: { type: 'object' } }]; // Pause startup #1 so we can inject an update during the gap. const startupReached = new DeferredPromise(); @@ -3567,7 +3567,7 @@ suite('ClaudeAgent', () => { // Wait until the materializer has snapshotted ['first'] into the diff // and is paused inside `sdk.startup`. THEN inject the update. await startupReached.p; - agent.setClientTools(created.session, 'c1', [{ name: 'second', inputSchema: { type: 'object' } }]); + agent.getOrCreateActiveClient(created.session, { clientId: 'c1' }).tools = [{ name: 'second', inputSchema: { type: 'object' } }]; startupGate.complete(); await send; @@ -3616,7 +3616,7 @@ suite('ClaudeAgent', () => { // update. Pre-fix the call hit the silent-drop branch because no // provisional was registered for the resume. await startupReached.p; - agent.setClientTools(sessionUri, 'c1', [{ name: 'resumed', inputSchema: { type: 'object' } }]); + agent.getOrCreateActiveClient(sessionUri, { clientId: 'c1' }).tools = [{ name: 'resumed', inputSchema: { type: 'object' } }]; startupGate.complete(); await send; @@ -3648,7 +3648,7 @@ suite('ClaudeAgent', () => { assert.strictEqual(sdk.startupCallCount, 1); // Stage a rebind whose startup will reject. - agent.setClientTools(created.session, 'c1', [{ name: 'echo', inputSchema: { type: 'object' } }]); + agent.getOrCreateActiveClient(created.session, { clientId: 'c1' }).tools = [{ name: 'echo', inputSchema: { type: 'object' } }]; sdk.startupRejection = new Error('simulated rebind startup failure'); sdk.queryAdvance = undefined; advance.complete(); @@ -5000,7 +5000,7 @@ suite('ClaudeAgent — Phase 11 customizations', () => { } })); - const synced = await agent.setClientCustomizations(created.session, 'client-1', [ + const synced = await agent.syncClientCustomizations(created.session, 'client-1', [ makeClientCustomization('https://a', 'A'), makeClientCustomization('https://b', 'B'), ]); @@ -5019,8 +5019,8 @@ suite('ClaudeAgent — Phase 11 customizations', () => { const s2 = await agent.createSession({ session: AgentSession.uri('claude', 'b'), workingDirectory: URI.file('/work') }); pm.syncResult = [makeSyncedRef('https://shared', '/p/shared')]; - await agent.setClientCustomizations(s1.session, 'c', [makeClientCustomization('https://shared', 'S')]); - await agent.setClientCustomizations(s2.session, 'c', [makeClientCustomization('https://shared', 'S')]); + await agent.syncClientCustomizations(s1.session, 'c', [makeClientCustomization('https://shared', 'S')]); + await agent.syncClientCustomizations(s2.session, 'c', [makeClientCustomization('https://shared', 'S')]); // One fire per per-session diff change confirms fan-out. let changes = 0; @@ -5039,9 +5039,9 @@ suite('ClaudeAgent — Phase 11 customizations', () => { const s2 = await agent.createSession({ session: AgentSession.uri('claude', 'two'), workingDirectory: URI.file('/work') }); pm.syncResult = [makeSyncedRef('https://shared', '/p/shared'), makeSyncedRef('https://a', '/p/a')]; - await agent.setClientCustomizations(s1.session, 'c', []); + await agent.syncClientCustomizations(s1.session, 'c', []); pm.syncResult = [makeSyncedRef('https://shared', '/p/shared'), makeSyncedRef('https://b', '/p/b')]; - await agent.setClientCustomizations(s2.session, 'c', []); + await agent.syncClientCustomizations(s2.session, 'c', []); // `IAgent.getCustomizations()` is the provider-level catalogue // (host-configured), NOT an aggregator across sessions. Claude has @@ -5058,7 +5058,7 @@ suite('ClaudeAgent — Phase 11 customizations', () => { const created = await agent.createSession({ workingDirectory: URI.file('/work') }); assert.strictEqual(created.provisional, true); - await agent.setClientCustomizations(created.session, 'c', [makeClientCustomization('https://a', 'A')]); + await agent.syncClientCustomizations(created.session, 'c', [makeClientCustomization('https://a', 'A')]); const customizations = await agent.getSessionCustomizations!(created.session); // The client-pushed customization, plus the curated read-only built-ins @@ -5075,7 +5075,7 @@ suite('ClaudeAgent — Phase 11 customizations', () => { await agent.authenticate(GITHUB_COPILOT_PROTECTED_RESOURCE.resource, 'tok'); const created = await agent.createSession({ workingDirectory: URI.file('/work') }); - await agent.setClientCustomizations(created.session, 'c', [makeClientCustomization('https://a', 'A')]); + await agent.syncClientCustomizations(created.session, 'c', [makeClientCustomization('https://a', 'A')]); // Disable the client-pushed entry; the projection must reflect it. agent.setCustomizationEnabled(customizationId('https://a'), false); // Disabling a DISCOVERED entry's id must be a no-op — the enablement @@ -5110,7 +5110,7 @@ suite('ClaudeAgent — Phase 11 customizations', () => { // pre-flight rebinds so `Options.plugins` on the new Query // includes the new path. pm.syncResult = [makeSyncedRef('https://a', '/p/a')]; - await agent.setClientCustomizations(created.session, 'c', [makeClientCustomization('https://a', 'A')]); + await agent.syncClientCustomizations(created.session, 'c', [makeClientCustomization('https://a', 'A')]); const firstQuery = sdk.warmQueries[0].produced!; const p2 = agent.sendMessage(created.session, 'second', undefined, 'turn-2'); @@ -5140,7 +5140,7 @@ suite('ClaudeAgent — Phase 11 customizations', () => { makeSystemInitMessage(sessionId), makeResultSuccess(sessionId), ]; pm.syncResult = [makeSyncedRef('https://x', '/p/x')]; - await agent.setClientCustomizations(created.session, 'c', [makeClientCustomization('https://x', 'X')]); + await agent.syncClientCustomizations(created.session, 'c', [makeClientCustomization('https://x', 'X')]); await agent.sendMessage(created.session, 'first', undefined, 'turn-1'); const session = agent.getSessionForTesting(created.session)!; // First-turn materialize consumed the dirty bit from the sync @@ -5182,7 +5182,7 @@ suite('ClaudeAgent — Phase 11 customizations', () => { const sessionId = AgentSession.id(created.session); sdk.nextQueryMessages = [makeSystemInitMessage(sessionId), makeResultSuccess(sessionId)]; - await agent.setClientCustomizations(created.session, 'c', [makeClientCustomization('https://a', 'A')]); + await agent.syncClientCustomizations(created.session, 'c', [makeClientCustomization('https://a', 'A')]); await agent.sendMessage(created.session, 'first', undefined, 'turn-1'); const customizations = await agent.getSessionCustomizations!(created.session); diff --git a/src/vs/platform/agentHost/test/node/claudeMapSessionEvents.test.ts b/src/vs/platform/agentHost/test/node/claudeMapSessionEvents.test.ts index 131b7a240a613e..4ea7b392e7bd2f 100644 --- a/src/vs/platform/agentHost/test/node/claudeMapSessionEvents.test.ts +++ b/src/vs/platform/agentHost/test/node/claudeMapSessionEvents.test.ts @@ -253,7 +253,7 @@ suite('claudeMapSessionEvents — direct mapper tests', () => { state, log, r(), - CLIENT_ID, + () => CLIENT_ID, ); assert.deepStrictEqual(signals, [{ diff --git a/src/vs/platform/agentHost/test/node/clientTools/claudeSessionClientToolsModel.test.ts b/src/vs/platform/agentHost/test/node/clientTools/claudeSessionClientToolsModel.test.ts index 4b6a5d9f840ff0..94f8900c0e3928 100644 --- a/src/vs/platform/agentHost/test/node/clientTools/claudeSessionClientToolsModel.test.ts +++ b/src/vs/platform/agentHost/test/node/clientTools/claudeSessionClientToolsModel.test.ts @@ -19,53 +19,53 @@ suite('SessionClientToolsDiff', () => { const disposables = ensureNoDisposablesAreLeakedInTestSuite(); - test('fresh diff: state empty, no difference', () => { + test('fresh diff: merged empty, no difference', () => { const diff = disposables.add(new SessionClientToolsDiff()); - assert.deepStrictEqual(diff.model.state.get(), { tools: undefined, clientId: undefined }); + assert.deepStrictEqual(diff.model.merged.get(), []); assert.strictEqual(diff.hasDifference, false); }); - test('setTools(undefined → []) does NOT flip dirty (undefined ≡ [])', () => { + test('setTools(c1, []) does NOT flip dirty (undefined ≡ [])', () => { const diff = disposables.add(new SessionClientToolsDiff()); - diff.model.setTools([]); + diff.model.setTools('c1', []); assert.strictEqual(diff.hasDifference, false); }); test('setTools with a real snapshot flips dirty', () => { const diff = disposables.add(new SessionClientToolsDiff()); - diff.model.setTools([tool()]); + diff.model.setTools('c1', [tool()]); assert.strictEqual(diff.hasDifference, true); }); - test('consume() returns the current state and clears dirty', () => { + test('consume() returns the merged tools and clears dirty', () => { const diff = disposables.add(new SessionClientToolsDiff()); - diff.model.setTools([tool()], 'c1'); - const state = diff.consume(); - assert.deepStrictEqual(state, { tools: [tool()], clientId: 'c1' }); + diff.model.setTools('c1', [tool()]); + const merged = diff.consume(); + assert.deepStrictEqual(merged, [tool()]); assert.strictEqual(diff.hasDifference, false); }); test('setTools with a structurally-equal snapshot does NOT re-flip dirty after consume', () => { const diff = disposables.add(new SessionClientToolsDiff()); - diff.model.setTools([tool()]); + diff.model.setTools('c1', [tool()]); diff.consume(); - diff.model.setTools([{ name: 'echo', description: 'echoes', inputSchema: { type: 'object', properties: { msg: { type: 'string' } }, required: ['msg'] } }]); + diff.model.setTools('c1', [{ name: 'echo', description: 'echoes', inputSchema: { type: 'object', properties: { msg: { type: 'string' } }, required: ['msg'] } }]); assert.strictEqual(diff.hasDifference, false); }); test('C6: setTools racing async work after consume re-flips dirty via autorun', async () => { const diff = disposables.add(new SessionClientToolsDiff()); - diff.model.setTools([tool({ name: 'original' })]); - const state = diff.consume(); - assert.deepStrictEqual(state.tools, [tool({ name: 'original' })]); + diff.model.setTools('c1', [tool({ name: 'original' })]); + const merged = diff.consume(); + assert.deepStrictEqual(merged, [tool({ name: 'original' })]); await Promise.resolve(); - diff.model.setTools([tool({ name: 'racer' })]); + diff.model.setTools('c1', [tool({ name: 'racer' })]); assert.strictEqual(diff.hasDifference, true); }); test('markDirty re-flips after a failed downstream build', () => { const diff = disposables.add(new SessionClientToolsDiff()); - diff.model.setTools([tool()]); + diff.model.setTools('c1', [tool()]); diff.consume(); assert.strictEqual(diff.hasDifference, false); diff.markDirty(); @@ -74,53 +74,62 @@ suite('SessionClientToolsDiff', () => { test('hasDifference detects rename / description / inputSchema; ignores title', () => { const diff = disposables.add(new SessionClientToolsDiff()); - diff.model.setTools([tool({ name: 'a' })]); + diff.model.setTools('c1', [tool({ name: 'a' })]); diff.consume(); - diff.model.setTools([tool({ name: 'b' })]); + diff.model.setTools('c1', [tool({ name: 'b' })]); assert.strictEqual(diff.hasDifference, true); diff.consume(); - diff.model.setTools([tool({ name: 'b', description: 'new' })]); + diff.model.setTools('c1', [tool({ name: 'b', description: 'new' })]); assert.strictEqual(diff.hasDifference, true); diff.consume(); - diff.model.setTools([tool({ name: 'b', description: 'new', inputSchema: { type: 'object', properties: { msg: { type: 'number' } }, required: ['msg'] } })]); + diff.model.setTools('c1', [tool({ name: 'b', description: 'new', inputSchema: { type: 'object', properties: { msg: { type: 'number' } }, required: ['msg'] } })]); assert.strictEqual(diff.hasDifference, true); diff.consume(); - diff.model.setTools([tool({ name: 'b', description: 'new', inputSchema: { type: 'object', properties: { msg: { type: 'number' } }, required: ['msg'] }, title: 'X' })]); + diff.model.setTools('c1', [tool({ name: 'b', description: 'new', inputSchema: { type: 'object', properties: { msg: { type: 'number' } }, required: ['msg'] }, title: 'X' })]); assert.strictEqual(diff.hasDifference, false, 'title is outside the diff scope'); }); test('order-insensitive: reordering tools does not flip dirty', () => { const diff = disposables.add(new SessionClientToolsDiff()); - diff.model.setTools([tool({ name: 'a' }), tool({ name: 'b' })]); + diff.model.setTools('c1', [tool({ name: 'a' }), tool({ name: 'b' })]); diff.consume(); - diff.model.setTools([tool({ name: 'b' }), tool({ name: 'a' })]); + diff.model.setTools('c1', [tool({ name: 'b' }), tool({ name: 'a' })]); assert.strictEqual(diff.hasDifference, false); }); - test('setTools writes clientId only when supplied', () => { + test('ownerOf reflects which client contributed a tool; getTools returns a client slice', () => { const diff = disposables.add(new SessionClientToolsDiff()); - diff.model.setTools([tool()], 'c1'); - assert.strictEqual(diff.model.state.get().clientId, 'c1'); - diff.model.setTools([tool({ name: 'echo2' })]); - assert.strictEqual(diff.model.state.get().clientId, 'c1', 'omitted clientId leaves previous value'); - diff.model.setTools([tool()], 'c2'); - assert.strictEqual(diff.model.state.get().clientId, 'c2'); + diff.model.setTools('c1', [tool({ name: 'a' })]); + diff.model.setTools('c2', [tool({ name: 'b' })]); + assert.strictEqual(diff.model.ownerOf('a'), 'c1'); + assert.strictEqual(diff.model.ownerOf('b'), 'c2'); + assert.strictEqual(diff.model.ownerOf('missing'), undefined); + assert.deepStrictEqual(diff.model.getTools('c1'), [tool({ name: 'a' })]); + assert.deepStrictEqual(diff.model.getTools('c2'), [tool({ name: 'b' })]); }); - test('clientId change alone does NOT flip dirty (same tools, new window)', () => { + test('merged unions multiple clients and dedupes by name (first client wins)', () => { const diff = disposables.add(new SessionClientToolsDiff()); - diff.model.setTools([tool()], 'c1'); + diff.model.setTools('c1', [tool({ name: 'shared', description: 'from c1' }), tool({ name: 'a' })]); + diff.model.setTools('c2', [tool({ name: 'shared', description: 'from c2' }), tool({ name: 'b' })]); + assert.deepStrictEqual(diff.model.merged.get().map(t => t.name), ['shared', 'a', 'b']); + assert.strictEqual(diff.model.ownerOf('shared'), 'c1', 'first-inserted client wins the shared name'); + }); + + test('removeClient drops that client and re-flips dirty when the merged set changes', () => { + const diff = disposables.add(new SessionClientToolsDiff()); + diff.model.setTools('c1', [tool({ name: 'a' })]); + diff.model.setTools('c2', [tool({ name: 'b' })]); diff.consume(); assert.strictEqual(diff.hasDifference, false); - // A window reload re-pushes identical tools under a new clientId. The - // observable still updates clientId, but no SDK yield-restart is - // required — only structural tool changes flip the dirty bit. - diff.model.setTools([tool()], 'c2'); - assert.strictEqual(diff.hasDifference, false); - assert.strictEqual(diff.model.state.get().clientId, 'c2'); + + diff.model.removeClient('c2'); + assert.strictEqual(diff.hasDifference, true); + assert.deepStrictEqual(diff.model.merged.get().map(t => t.name), ['a']); + assert.strictEqual(diff.model.ownerOf('b'), undefined); }); }); diff --git a/src/vs/platform/agentHost/test/node/codex/codexMapAppServerEvents.test.ts b/src/vs/platform/agentHost/test/node/codex/codexMapAppServerEvents.test.ts index a17f1fba090b4e..3409f81ea884e8 100644 --- a/src/vs/platform/agentHost/test/node/codex/codexMapAppServerEvents.test.ts +++ b/src/vs/platform/agentHost/test/node/codex/codexMapAppServerEvents.test.ts @@ -8,6 +8,7 @@ import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/tes import { createCodexSessionMapState, extractUserInputText, mapAgentMessageDelta, mapCommandExecutionOutputDelta, mapFileChangePatchUpdated, mapItemCompleted, mapItemStarted, mapMcpToolCallProgress, mapReasoningSummaryPartAdded, mapReasoningSummaryTextDelta, mapReasoningTextDelta, mapTokenUsageUpdated, mapTurnCompleted, mapTurnStarted, resetCodexTurnMapState, turnStateFromStatus } from '../../../node/codex/codexMapAppServerEvents.js'; import { ActionType } from '../../../common/state/sessionActions.js'; import { MessageKind, ResponsePartKind, ToolCallConfirmationReason, ToolCallContributorKind, ToolResultContentType, TurnState } from '../../../common/state/sessionState.js'; +import { ActiveClientToolSet } from '../../../node/activeClientState.js'; suite('codexMapAppServerEvents', () => { @@ -393,9 +394,10 @@ suite('codexMapAppServerEvents', () => { }); }); - test('dynamicToolCall item carries a Client contributor when toolsClientId is set', () => { - const state = createCodexSessionMapState(); - state.toolsClientId = 'win-7'; + test('dynamicToolCall item carries a Client contributor when a client owns the tool', () => { + const toolSet = new ActiveClientToolSet(); + toolSet.set('win-7', [{ name: 'get_magic_word' }]); + const state = createCodexSessionMapState(new Set(), toolSet); const startActions = mapItemStarted(state, { item: { type: 'dynamicToolCall', id: 'dyn_2', namespace: null, tool: 'get_magic_word', arguments: {}, status: 'inProgress', contentItems: null, success: null, durationMs: null } as never, threadId: 'thr_1', turnId: 'turn_a', startedAtMs: 0, @@ -416,8 +418,9 @@ suite('codexMapAppServerEvents', () => { // A server tool is registered under its bare name and executes // in-process, so it must not carry a Client contributor even when a // workbench client owns the (other) client tools. - const state = createCodexSessionMapState(new Set(['addComment'])); - state.toolsClientId = 'win-7'; + const toolSet = new ActiveClientToolSet(); + toolSet.set('win-7', [{ name: 'get_magic_word' }]); + const state = createCodexSessionMapState(new Set(['addComment']), toolSet); const startActions = mapItemStarted(state, { item: { type: 'dynamicToolCall', id: 'dyn_3', namespace: null, tool: 'addComment', arguments: {}, status: 'inProgress', contentItems: null, success: null, durationMs: null } as never, threadId: 'thr_1', turnId: 'turn_a', startedAtMs: 0, diff --git a/src/vs/platform/agentHost/test/node/copilotAgent.test.ts b/src/vs/platform/agentHost/test/node/copilotAgent.test.ts index 108989120593c7..ccc9e079ee0eb5 100644 --- a/src/vs/platform/agentHost/test/node/copilotAgent.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotAgent.test.ts @@ -50,7 +50,7 @@ import type { CopilotSessionLaunchPlan, IActiveClientSnapshot } from '../../node import { ShellManager } from '../../node/copilot/copilotShellTools.js'; import { SessionDatabase } from '../../node/sessionDatabase.js'; import { createNullSessionDataService } from '../common/sessionTestHelpers.js'; -import { ActiveClientState } from '../../node/activeClientState.js'; +import { ActiveClientToolSet } from '../../node/activeClientState.js'; import { ICopilotApiService, type ICopilotApiServiceRequestOptions, type ICopilotUtilityChatCompletionRequest } from '../../node/shared/copilotApiService.js'; class TestAgentPluginManager implements IAgentPluginManager { @@ -477,7 +477,7 @@ function createAgentSessionThroughAgent(agent: CopilotAgent, instantiationServic }, resumeSession: async () => new MockCopilotSession() as unknown as CopilotSession, }, - activeClientState: new ActiveClientState(), + activeClientToolSet: new ActiveClientToolSet(), sessionId: 'test-session-1', workingDirectory: undefined, resolvedAgentName: undefined, @@ -1325,7 +1325,7 @@ suite('CopilotAgent', () => { await agent.authenticate('https://api.github.com', 'token'); const session = AgentSession.uri('copilotcli', 'sync-customizations-test'); - await agent.setClientCustomizations(session, 'client-1', [{ type: CustomizationType.Plugin, id: customizationId(pluginDir.toString()), uri: pluginDir.toString(), name: 'Plugin A', enabled: true }]); + agent.getOrCreateActiveClient(session, { clientId: 'client-1' }).customizations = [{ type: CustomizationType.Plugin, id: customizationId(pluginDir.toString()), uri: pluginDir.toString(), name: 'Plugin A', enabled: true }]; // Wait for the deferred resolution chain in PluginController.sync. await new Promise(r => setTimeout(r, 50)); @@ -2105,7 +2105,7 @@ suite('CopilotAgent', () => { suite('client tool refresh on reload (#319516)', () => { /** Minimal structural view of the agent's private per-session ActiveClient. */ type TestActiveClient = { - readonly state: { readonly clientId: string | undefined }; + readonly toolSet: { ownerOf(toolName: string): string | undefined }; snapshot(): Promise; requiresRestart(snap: IActiveClientSnapshot): Promise; }; @@ -2113,33 +2113,34 @@ suite('CopilotAgent', () => { function getActiveClient(agent: CopilotAgent, session: URI): TestActiveClient { const activeClients = (agent as unknown as { _activeClients: { get(s: URI): TestActiveClient | undefined } })._activeClients; const activeClient = activeClients.get(session); - assert.ok(activeClient, 'expected an ActiveClient to exist after setClientTools'); + assert.ok(activeClient, 'expected an ActiveClient to exist after registering client tools'); return activeClient; } const tools: ToolDefinition[] = [{ name: 'my_tool', description: 'A test tool', inputSchema: { type: 'object', properties: {} } }]; - test('clientId-only change (reload) does NOT require a restart and updates the live clientId', async () => { + test('clientId-only change (reload) does NOT require a restart and updates the live owner', async () => { const agent = createTestAgent(disposables); try { const session = AgentSession.uri('copilotcli', 'reload-session'); // Window A registers its tools; this is the snapshot the SDK // session would be created with. - agent.setClientTools(session, 'client-A', tools); + agent.getOrCreateActiveClient(session, { clientId: 'client-A' }).tools = tools; const activeClient = getActiveClient(agent, session); const appliedSnapshot = await activeClient.snapshot(); - assert.strictEqual(activeClient.state.clientId, 'client-A'); + assert.strictEqual(activeClient.toolSet.ownerOf('my_tool'), 'client-A'); // Window A reloads: window B reconnects with a new clientId but - // the identical tool list. - agent.setClientTools(session, 'client-B', [...tools]); + // the identical tool list. The reload removes A then adds B. + agent.removeActiveClient(session, 'client-A'); + agent.getOrCreateActiveClient(session, { clientId: 'client-B' }).tools = [...tools]; // Root-cause assertions: the cached SDK session must be reused - // (no restart) AND the live clientId must now be window B's, so + // (no restart) AND the live owner must now be window B's, so // the next client tool call is stamped with a live owner. assert.strictEqual(await activeClient.requiresRestart(appliedSnapshot), false); - assert.strictEqual(activeClient.state.clientId, 'client-B'); + assert.strictEqual(activeClient.toolSet.ownerOf('my_tool'), 'client-B'); } finally { await disposeAgent(agent); } @@ -2150,19 +2151,56 @@ suite('CopilotAgent', () => { try { const session = AgentSession.uri('copilotcli', 'tools-change-session'); - agent.setClientTools(session, 'client-A', tools); + agent.getOrCreateActiveClient(session, { clientId: 'client-A' }).tools = tools; const activeClient = getActiveClient(agent, session); const appliedSnapshot = await activeClient.snapshot(); // A genuinely different tool set (added tool) must restart so the // SDK session is rebuilt with the new tools. - agent.setClientTools(session, 'client-A', [...tools, { name: 'second_tool', description: 'another', inputSchema: { type: 'object', properties: {} } }]); + agent.getOrCreateActiveClient(session, { clientId: 'client-A' }).tools = [...tools, { name: 'second_tool', description: 'another', inputSchema: { type: 'object', properties: {} } }]; assert.strictEqual(await activeClient.requiresRestart(appliedSnapshot), true); } finally { await disposeAgent(agent); } }); + + test('multiple active clients merge their tools and removal isolates per client', async () => { + const agent = createTestAgent(disposables); + try { + const session = AgentSession.uri('copilotcli', 'multi-client-session'); + + // Two clients each contribute their own tool plus a shared one. + agent.getOrCreateActiveClient(session, { clientId: 'client-A' }).tools = [ + { name: 'shared', description: 'from A', inputSchema: { type: 'object', properties: {} } }, + { name: 'a_tool', description: 'A only', inputSchema: { type: 'object', properties: {} } }, + ]; + agent.getOrCreateActiveClient(session, { clientId: 'client-B' }).tools = [ + { name: 'shared', description: 'from B', inputSchema: { type: 'object', properties: {} } }, + { name: 'b_tool', description: 'B only', inputSchema: { type: 'object', properties: {} } }, + ]; + const activeClient = getActiveClient(agent, session); + + // The SDK snapshot merges both clients, deduping the shared name + // in favor of the first-inserted client (A), and ownership maps + // each tool to its contributing client. + const merged = await activeClient.snapshot(); + assert.deepStrictEqual(merged.tools.map(t => t.name), ['shared', 'a_tool', 'b_tool']); + assert.strictEqual(activeClient.toolSet.ownerOf('shared'), 'client-A'); + assert.strictEqual(activeClient.toolSet.ownerOf('a_tool'), 'client-A'); + assert.strictEqual(activeClient.toolSet.ownerOf('b_tool'), 'client-B'); + + // Removing client A keeps B's contribution and hands the shared + // tool to B (now the sole provider). + agent.removeActiveClient(session, 'client-A'); + const afterRemoval = await activeClient.snapshot(); + assert.deepStrictEqual(afterRemoval.tools.map(t => t.name), ['shared', 'b_tool']); + assert.strictEqual(activeClient.toolSet.ownerOf('shared'), 'client-B'); + assert.strictEqual(activeClient.toolSet.ownerOf('a_tool'), undefined); + } finally { + await disposeAgent(agent); + } + }); }); suite('_resumeSession dedup', () => { diff --git a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts index 7cff7183609133..17354c67b2682d 100644 --- a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts @@ -27,7 +27,7 @@ import { ISessionDataService } from '../../common/sessionDataService.js'; import { ActionType, type ChatDeltaAction, type ChatErrorAction, type ChatInputRequestedAction, type ChatResponsePartAction, type ChatToolCallCompleteAction, type ChatToolCallReadyAction, type ChatToolCallStartAction, type ChatTurnCompleteAction } from '../../common/state/sessionActions.js'; import { MessageAttachmentKind, MessageKind, ResponsePartKind, ChatInputAnswerState, ChatInputAnswerValueKind, ChatInputQuestionKind, ChatInputResponseKind, ToolCallContributorKind, ToolCallStatus, ToolResultContentType, type ToolDefinition, type ToolResultContent, type ToolResultFileEditContent } from '../../common/state/sessionState.js'; import { CopilotAgentSession } from '../../node/copilot/copilotAgentSession.js'; -import { ActiveClientState } from '../../node/activeClientState.js'; +import { ActiveClientToolSet } from '../../node/activeClientState.js'; import { type CopilotSessionLaunchPlan, type IActiveClientSnapshot, type ICopilotSessionLauncher, type ICopilotSessionRuntime } from '../../node/copilot/copilotSessionLauncher.js'; import { CopilotSessionWrapper } from '../../node/copilot/copilotSessionWrapper.js'; import { buildCopilotSystemNotification } from '../../node/copilot/copilotSystemNotification.js'; @@ -227,7 +227,7 @@ function getInputRequest(signal: AgentSignal): ChatInputRequestedAction['request async function createAgentSession(disposables: DisposableStore, options?: { clientSnapshot?: IActiveClientSnapshot; - activeClientState?: ActiveClientState; + activeClientToolSet?: ActiveClientToolSet; environmentServiceRegistration?: 'native' | 'none'; logService?: ILogService; telemetryService?: ITelemetryService; @@ -284,7 +284,7 @@ async function createAgentSession(disposables: DisposableStore, options?: { createSession: async () => mockSession as unknown as CopilotSession, resumeSession: async () => mockSession as unknown as CopilotSession, }, - activeClientState: new ActiveClientState(), + activeClientToolSet: new ActiveClientToolSet(), sessionId: 'test-session-1', workingDirectory: options?.workingDirectory, resolvedAgentName: undefined, @@ -357,7 +357,7 @@ async function createAgentSession(disposables: DisposableStore, options?: { launchPlan, shellManager: undefined, clientSnapshot: options?.clientSnapshot, - activeClientState: options?.activeClientState, + activeClientToolSet: options?.activeClientToolSet, resolveMcpChildId: () => undefined, workingDirectory: options?.workingDirectory, serverToolHost: options?.serverToolHost, @@ -1830,8 +1830,8 @@ suite('CopilotAgentSession', () => { test('client tool telemetry does not use clientId as toolExtensionId', async () => { const telemetryService = new RecordingTelemetryService(); const tools = [{ name: 'my_tool', description: 'A test tool', inputSchema: { type: 'object', properties: {} } }] as const; - const activeClientState = new ActiveClientState(); - activeClientState.update('test-client', tools); + const activeClientToolSet = new ActiveClientToolSet(); + activeClientToolSet.set('test-client', tools); const { mockSession } = await createAgentSession(disposables, { telemetryService, clientSnapshot: { @@ -1839,7 +1839,7 @@ suite('CopilotAgentSession', () => { plugins: [], mcpServers: {}, }, - activeClientState, + activeClientToolSet, }); mockSession.fire('tool.execution_start', { @@ -3064,11 +3064,11 @@ suite('CopilotAgentSession', () => { mcpServers: {}, }; - /** Builds a live ActiveClientState seeded with the given owning clientId and the snapshot's tools. */ - const activeClientStateWith = (clientId: string): ActiveClientState => { - const state = new ActiveClientState(); - state.update(clientId, snapshot.tools); - return state; + /** Builds a live ActiveClientToolSet seeded with the given owning clientId and the snapshot's tools. */ + const activeClientToolSetWith = (clientId: string): ActiveClientToolSet => { + const toolSet = new ActiveClientToolSet(); + toolSet.set(clientId, snapshot.tools); + return toolSet; }; test('client tool started with no connected client fails immediately', async () => { @@ -3103,7 +3103,7 @@ suite('CopilotAgentSession', () => { test('client tool handler waits for completion without emitting tool_ready', async () => { - const { session, runtime, mockSession, signals } = await createAgentSession(disposables, { clientSnapshot: snapshot, activeClientState: activeClientStateWith('test-client') }); + const { session, runtime, mockSession, signals } = await createAgentSession(disposables, { clientSnapshot: snapshot, activeClientToolSet: activeClientToolSetWith('test-client') }); // SDK emits tool.execution_start — tool_start fires immediately mockSession.fire('tool.execution_start', { @@ -3141,9 +3141,9 @@ suite('CopilotAgentSession', () => { }); test('client tool handler does not emit tool_ready (permission flow owns it)', async () => { - const activeClientState = new ActiveClientState(); - activeClientState.update('client-perm', snapshot.tools); - const { session, runtime, mockSession, signals, waitForSignal } = await createAgentSession(disposables, { clientSnapshot: snapshot, activeClientState }); + const activeClientToolSet = new ActiveClientToolSet(); + activeClientToolSet.set('client-perm', snapshot.tools); + const { session, runtime, mockSession, signals, waitForSignal } = await createAgentSession(disposables, { clientSnapshot: snapshot, activeClientToolSet }); // SDK emits tool.execution_start — tool_start fires immediately mockSession.fire('tool.execution_start', { @@ -3200,7 +3200,7 @@ suite('CopilotAgentSession', () => { // ChatToolCallReady to the subagent session and emits a // stray ready against the parent session (no preceding // ChatToolCallStart). - const { session, runtime, mockSession, signals, waitForSignal } = await createAgentSession(disposables, { clientSnapshot: snapshot, activeClientState: activeClientStateWith('test-client') }); + const { session, runtime, mockSession, signals, waitForSignal } = await createAgentSession(disposables, { clientSnapshot: snapshot, activeClientToolSet: activeClientToolSetWith('test-client') }); mockSession.fire('subagent.started', { toolCallId: 'tc-parent-subagent', @@ -3448,10 +3448,10 @@ suite('CopilotAgentSession', () => { } }); - test('client tool start stamps the LIVE clientId from the shared ActiveClientState', async () => { - const activeClientState = new ActiveClientState(); - activeClientState.update('client-A', snapshot.tools); - const { mockSession, signals } = await createAgentSession(disposables, { clientSnapshot: snapshot, activeClientState }); + test('client tool start stamps the owning clientId from the shared ActiveClientToolSet', async () => { + const activeClientToolSet = new ActiveClientToolSet(); + activeClientToolSet.set('client-A', snapshot.tools); + const { mockSession, signals } = await createAgentSession(disposables, { clientSnapshot: snapshot, activeClientToolSet }); mockSession.fire('tool.execution_start', { toolCallId: 'tc-live-1', @@ -3459,8 +3459,10 @@ suite('CopilotAgentSession', () => { arguments: {}, } as SessionEventPayload<'tool.execution_start'>['data']); - // A window reload re-pushes the same tools under a new clientId. - activeClientState.update('client-B', snapshot.tools); + // A window reload removes the old client and re-pushes the same + // tools under a new clientId. + activeClientToolSet.delete('client-A'); + activeClientToolSet.set('client-B', snapshot.tools); mockSession.fire('tool.execution_start', { toolCallId: 'tc-live-2', toolName: 'my_tool', diff --git a/src/vs/platform/agentHost/test/node/customizations/claudeSessionClientCustomizationsModel.test.ts b/src/vs/platform/agentHost/test/node/customizations/claudeSessionClientCustomizationsModel.test.ts index 1c7791b42d0a5f..0f3ce0d81db72d 100644 --- a/src/vs/platform/agentHost/test/node/customizations/claudeSessionClientCustomizationsModel.test.ts +++ b/src/vs/platform/agentHost/test/node/customizations/claudeSessionClientCustomizationsModel.test.ts @@ -39,14 +39,14 @@ suite('SessionClientCustomizationsDiff', () => { const diff = disposables.add(new SessionClientCustomizationsDiff()); let fires = 0; disposables.add(diff.onDidChange(() => fires++)); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a' })]); assert.strictEqual(diff.hasDifference, true); assert.strictEqual(fires, 1); }); test('enabledPluginPaths excludes entries without pluginDir', () => { const diff = disposables.add(new SessionClientCustomizationsDiff()); - diff.model.setSyncedCustomizations([ + diff.model.setSyncedCustomizations('c1', [ synced('https://a', { dir: '/p/a' }), synced('https://b'), ]); @@ -55,7 +55,7 @@ suite('SessionClientCustomizationsDiff', () => { test('setEnabled(false) removes from enabled paths and flips dirty exactly when value changes', () => { const diff = disposables.add(new SessionClientCustomizationsDiff()); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a' })]); diff.consume(); assert.strictEqual(diff.hasDifference, false); @@ -74,13 +74,13 @@ suite('SessionClientCustomizationsDiff', () => { test('default enablement is true (absent entry counts as enabled)', () => { const diff = disposables.add(new SessionClientCustomizationsDiff()); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a' })]); assert.strictEqual(diff.model.enabledPluginPaths.get().length, 1); }); test('setEnabled(true) is a no-op for default-enabled entries', () => { const diff = disposables.add(new SessionClientCustomizationsDiff()); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a' })]); diff.consume(); let fires = 0; disposables.add(diff.onDidChange(() => fires++)); @@ -91,7 +91,7 @@ suite('SessionClientCustomizationsDiff', () => { test('consume returns current paths and clears dirty', () => { const diff = disposables.add(new SessionClientCustomizationsDiff()); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a' })]); const paths = diff.consume(); assert.strictEqual(paths.length, 1); assert.strictEqual(diff.hasDifference, false); @@ -99,7 +99,7 @@ suite('SessionClientCustomizationsDiff', () => { test('markDirty re-flips after failed downstream reload', () => { const diff = disposables.add(new SessionClientCustomizationsDiff()); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a' })]); diff.consume(); assert.strictEqual(diff.hasDifference, false); diff.markDirty(); @@ -108,18 +108,18 @@ suite('SessionClientCustomizationsDiff', () => { test('structurally-equivalent re-send is deduped (no fire, no dirty)', () => { const diff = disposables.add(new SessionClientCustomizationsDiff()); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a' })]); diff.consume(); let fires = 0; disposables.add(diff.onDidChange(() => fires++)); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a' })]); assert.strictEqual(fires, 0); assert.strictEqual(diff.hasDifference, false); }); test('toggling enablement of customization without pluginDir still flips dirty (no-restart optimisation intentionally given up: rebind is cheap and correctness > efficiency)', () => { const diff = disposables.add(new SessionClientCustomizationsDiff()); - diff.model.setSyncedCustomizations([synced('https://a')]); + diff.model.setSyncedCustomizations('c1', [synced('https://a')]); diff.consume(); diff.model.setEnabled(customizationId('https://a'), false); assert.strictEqual(diff.hasDifference, true); @@ -127,20 +127,37 @@ suite('SessionClientCustomizationsDiff', () => { test('nonce change at same URI / pluginDir flips dirty', () => { const diff = disposables.add(new SessionClientCustomizationsDiff()); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a', nonce: 'v1' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a', nonce: 'v1' })]); diff.consume(); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a', nonce: 'v2' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a', nonce: 'v2' })]); assert.strictEqual(diff.hasDifference, true); }); test('name change at same URI flips dirty (state observable fires for workbench refetch)', () => { const diff = disposables.add(new SessionClientCustomizationsDiff()); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a', name: 'A' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a', name: 'A' })]); diff.consume(); let fires = 0; disposables.add(diff.onDidChange(() => fires++)); - diff.model.setSyncedCustomizations([synced('https://a', { dir: '/p/a', name: 'A renamed' })]); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a', name: 'A renamed' })]); assert.strictEqual(fires, 1); assert.strictEqual(diff.hasDifference, true); }); + + test('merges multiple clients and dedupes by id (first client wins); removeClient drops a client', () => { + const diff = disposables.add(new SessionClientCustomizationsDiff()); + diff.model.setSyncedCustomizations('c1', [synced('https://a', { dir: '/p/a' })]); + diff.model.setSyncedCustomizations('c2', [synced('https://b', { dir: '/p/b' })]); + assert.deepStrictEqual( + diff.model.enabledPluginPaths.get().map(u => u.fsPath), + [URI.file('/p/a').fsPath, URI.file('/p/b').fsPath], + ); + + diff.model.removeClient('c1'); + assert.deepStrictEqual( + diff.model.enabledPluginPaths.get().map(u => u.fsPath), + [URI.file('/p/b').fsPath], + ); + }); }); + diff --git a/src/vs/platform/agentHost/test/node/mockAgent.ts b/src/vs/platform/agentHost/test/node/mockAgent.ts index 5daeb348755abc..3baa03f3824e83 100644 --- a/src/vs/platform/agentHost/test/node/mockAgent.ts +++ b/src/vs/platform/agentHost/test/node/mockAgent.ts @@ -9,9 +9,9 @@ import { observableValue } from '../../../../base/common/observable.js'; import type { IAuthorizationProtectedResourceMetadata } from '../../../../base/common/oauth.js'; import { URI } from '../../../../base/common/uri.js'; import { type ISyncedCustomization } from '../../common/agentPluginManager.js'; -import { AgentSession, type AgentProvider, type AgentSignal, type IAgent, type IAgentActionSignal, type IAgentCreateSessionConfig, type IAgentCreateSessionResult, type IAgentDescriptor, type IAgentModelInfo, type IAgentResolveSessionConfigParams, type IAgentSessionConfigCompletionsParams, type IAgentSessionMetadata, type IAgentToolPendingConfirmationSignal } from '../../common/agentService.js'; +import { AgentSession, type AgentProvider, type AgentSignal, type IActiveClient, type IAgent, type IAgentActionSignal, type IAgentCreateSessionConfig, type IAgentCreateSessionResult, type IAgentDescriptor, type IAgentModelInfo, type IAgentResolveSessionConfigParams, type IAgentSessionConfigCompletionsParams, type IAgentSessionMetadata, type IAgentToolPendingConfirmationSignal } from '../../common/agentService.js'; import { buildSubagentTurnsFromHistory, buildTurnsFromHistory, type IHistoryRecord } from './historyRecordFixtures.js'; -import { ProtectedResourceMetadata, ToolCallContributorKind, type AgentSelection, type MessageAttachment, type ModelSelection } from '../../common/state/protocol/state.js'; +import { ProtectedResourceMetadata, ToolCallContributorKind, type AgentSelection, type MessageAttachment, type ModelSelection, type ToolDefinition } from '../../common/state/protocol/state.js'; import type { ResolveSessionConfigResult, SessionConfigCompletionsResult } from '../../common/state/protocol/commands.js'; import { ActionType } from '../../common/state/sessionActions.js'; import { ResponsePartKind, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType, CustomizationLoadStatus, parseSubagentSessionUri, type ClientPluginCustomization, type Customization, type PendingMessage, type StringOrMarkdown, type ToolCallResult, type Turn, type UsageInfo } from '../../common/state/sessionState.js'; @@ -57,6 +57,8 @@ export class MockAgent implements IAgent { readonly changeAgentCalls: { session: URI; agent: AgentSelection | undefined; chat?: URI }[] = []; readonly authenticateCalls: { resource: string; token: string }[] = []; readonly setClientCustomizationsCalls: { clientId: string; customizations: ClientPluginCustomization[] }[] = []; + readonly setClientToolsCalls: { clientId: string; tools: readonly ToolDefinition[] }[] = []; + readonly removeActiveClientCalls: { clientId: string }[] = []; readonly setCustomizationEnabledCalls: { id: string; enabled: boolean }[] = []; /** Configurable return value for getCustomizations. */ customizations: Customization[] = []; @@ -176,7 +178,7 @@ export class MockAgent implements IAgent { return this.customizations; } - async setClientCustomizations(session: URI, clientId: string, customizations: ClientPluginCustomization[]): Promise { + syncClientCustomizations(session: URI, clientId: string, customizations: ClientPluginCustomization[]): ISyncedCustomization[] { this.setClientCustomizationsCalls.push({ clientId, customizations }); const results: ISyncedCustomization[] = customizations.map(c => ({ customization: { @@ -199,7 +201,29 @@ export class MockAgent implements IAgent { this.setCustomizationEnabledCalls.push({ id, enabled }); } - setClientTools(): void { } + getOrCreateActiveClient(session: URI, client: { readonly clientId: string; readonly displayName?: string }): IActiveClient { + const self = this; + let tools: readonly ToolDefinition[] = []; + let customizations: readonly ClientPluginCustomization[] = []; + return { + clientId: client.clientId, + displayName: client.displayName, + get tools() { return tools; }, + set tools(value: readonly ToolDefinition[]) { + tools = value; + self.setClientToolsCalls.push({ clientId: client.clientId, tools: value }); + }, + get customizations() { return customizations; }, + set customizations(value: readonly ClientPluginCustomization[]) { + customizations = value; + self.syncClientCustomizations(session, client.clientId, [...value]); + }, + }; + } + + removeActiveClient(_session: URI, clientId: string): void { + this.removeActiveClientCalls.push({ clientId }); + } onClientToolCallComplete(): void { } @@ -713,16 +737,25 @@ export class ScriptedMockAgent implements IAgent { } } - async setClientCustomizations() { - return []; + getOrCreateActiveClient(_session: URI, client: { readonly clientId: string; readonly displayName?: string }): IActiveClient { + let tools: readonly ToolDefinition[] = []; + let customizations: readonly ClientPluginCustomization[] = []; + return { + clientId: client.clientId, + displayName: client.displayName, + get tools() { return tools; }, + set tools(value: readonly ToolDefinition[]) { tools = value; }, + get customizations() { return customizations; }, + set customizations(value: readonly ClientPluginCustomization[]) { customizations = value; }, + }; } + removeActiveClient(): void { } + setCustomizationEnabled() { } - setClientTools(): void { } - private didCompleteToolCalls = new Set(); onClientToolCallComplete(session: URI, toolCallId: string, result: ToolCallResult): void { diff --git a/src/vs/platform/agentHost/test/node/protocol/codexRealSdk.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/codexRealSdk.integrationTest.ts index 0515308523f8e2..c303e34221a76d 100644 --- a/src/vs/platform/agentHost/test/node/protocol/codexRealSdk.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/codexRealSdk.integrationTest.ts @@ -150,7 +150,7 @@ defineSharedRealSdkTests(CODEX_CONFIG); channel: session, clientSeq: 1, action: { - type: 'session/activeClientChanged', + type: 'session/activeClientSet', activeClient: { clientId: 'tool-client', tools: [{ @@ -228,7 +228,7 @@ defineSharedRealSdkTests(CODEX_CONFIG); channel: session, clientSeq: 1, action: { - type: 'session/activeClientChanged', + type: 'session/activeClientSet', activeClient: { clientId: 'tool-client-2', tools: [{ diff --git a/src/vs/platform/agentHost/test/node/protocol/realSdkTestHelpers.ts b/src/vs/platform/agentHost/test/node/protocol/realSdkTestHelpers.ts index 352d8a3cd674ea..8857234859205d 100644 --- a/src/vs/platform/agentHost/test/node/protocol/realSdkTestHelpers.ts +++ b/src/vs/platform/agentHost/test/node/protocol/realSdkTestHelpers.ts @@ -805,7 +805,7 @@ export function defineSharedRealSdkTests(config: IRealSdkProviderConfig): void { channel: sessionUri, clientSeq: 1, action: { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: `real-sdk-worktree-${config.provider}`, displayName: 'Test Client', diff --git a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts index c37134136cf9a5..f0d2f17a0eb214 100644 --- a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts +++ b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts @@ -958,12 +958,12 @@ suite('ProtocolServerHandler', () => { assert.strictEqual(transport.sent.length, 0); }); - test('client disconnect clears active client and fails owned tool calls after grace period', () => { + test('client disconnect retains active client during grace, then removes it and fails owned tool calls after grace period', () => { return runWithFakedTimers({ useFakeTimers: true }, async () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); stateManager.dispatchServerAction(sessionUri, { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'client-tools', tools: [{ name: 'runTask', description: 'Runs a task' }] @@ -994,13 +994,18 @@ suite('ProtocolServerHandler', () => { const transport = connectClient('client-tools', [sessionUri]); transport.simulateClose(); - assert.strictEqual(stateManager.getSessionState(sessionUri)?.activeClient, undefined); + // The active client is retained during the grace window so a quick + // reconnect can keep its slot. + assert.deepStrictEqual(stateManager.getSessionState(sessionUri)?.activeClients.map(c => c.clientId), ['client-tools']); let part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0]; assert.strictEqual(part?.kind, ResponsePartKind.ToolCall); assert.strictEqual(part?.kind === ResponsePartKind.ToolCall ? part.toolCall.status : undefined, ToolCallStatus.Running); await new Promise(r => setTimeout(r, 30_001)); + // After the grace window the active client is removed and its + // pending tool call is failed. + assert.deepStrictEqual(stateManager.getSessionState(sessionUri)?.activeClients, []); part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0]; assert.strictEqual(part?.kind, ResponsePartKind.ToolCall); assert.deepStrictEqual(part?.kind === ResponsePartKind.ToolCall ? { @@ -1020,7 +1025,7 @@ suite('ProtocolServerHandler', () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); stateManager.dispatchServerAction(sessionUri, { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'client-tools', tools: [{ name: 'runTask', description: 'Runs a task' }] @@ -1068,7 +1073,7 @@ suite('ProtocolServerHandler', () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); stateManager.dispatchServerAction(sessionUri, { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'client-tools', tools: [{ name: 'runTask', description: 'Runs a task' }] @@ -1110,7 +1115,7 @@ suite('ProtocolServerHandler', () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); stateManager.dispatchServerAction(sessionUri, { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'client-tools', tools: [{ name: 'runTask', description: 'Runs a task' }] @@ -1159,7 +1164,7 @@ suite('ProtocolServerHandler', () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); stateManager.dispatchServerAction(sessionUri, { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'client-tools', tools: [{ name: 'runTask', description: 'Runs a task' }] @@ -1217,7 +1222,7 @@ suite('ProtocolServerHandler', () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); stateManager.dispatchServerAction(sessionUri, { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'client-tools', tools: [{ name: 'runTask', description: 'Runs a task' }] @@ -1269,7 +1274,7 @@ suite('ProtocolServerHandler', () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); stateManager.dispatchServerAction(sessionUri, { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'client-tools', tools: [{ name: 'runTask', description: 'Runs a task' }] @@ -1300,7 +1305,7 @@ suite('ProtocolServerHandler', () => { const transport = connectClient('client-tools', [sessionUri]); transport.simulateClose(); stateManager.dispatchServerAction(sessionUri, { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { clientId: 'client-replacement', tools: [{ name: 'runTask', description: 'Runs a task' }] @@ -1436,6 +1441,167 @@ suite('ProtocolServerHandler', () => { }); }); + test('unsubscribe removes the active client and fails its owned tool calls', () => { + stateManager.createSession(makeSessionSummary()); + stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.SessionActiveClientSet, + activeClient: { + clientId: 'client-tools', + tools: [{ name: 'runTask', description: 'Runs a task' }] + }, + }); + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.ChatTurnStarted, + turnId: 'turn-1', + message: { text: 'run it', origin: { kind: MessageKind.User } }, + }); + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.ChatToolCallStart, + turnId: 'turn-1', + toolCallId: 'tool-1', + toolName: 'runTask', + displayName: 'Run Task', + contributor: { kind: ToolCallContributorKind.Client, clientId: 'client-tools' }, + }); + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.ChatToolCallReady, + turnId: 'turn-1', + toolCallId: 'tool-1', + invocationMessage: 'Run Task', + toolInput: '{}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }); + + const transport = connectClient('client-tools', [sessionUri]); + transport.simulateMessage(notification('unsubscribe', { channel: sessionUri })); + + assert.deepStrictEqual(stateManager.getSessionState(sessionUri)?.activeClients, []); + const part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0]; + assert.deepStrictEqual(part?.kind === ResponsePartKind.ToolCall ? { + status: part.toolCall.status, + success: part.toolCall.status === ToolCallStatus.Completed ? part.toolCall.success : undefined, + error: part.toolCall.status === ToolCallStatus.Completed ? part.toolCall.error?.message : undefined, + } : undefined, { + status: ToolCallStatus.Completed, + success: false, + error: 'Client client-tools disconnected before completing Run Task', + }); + + transport.simulateClose(); + }); + + test('reconnect without resubscription removes the active client and fails its owned tool calls', async () => { + stateManager.createSession(makeSessionSummary()); + stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); + + const transport1 = connectClient('client-tools', [sessionUri]); + const initSeq = (findResponse(transport1.sent, 1) as { result: InitializeResult }).result.serverSeq; + + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.SessionActiveClientSet, + activeClient: { + clientId: 'client-tools', + tools: [{ name: 'runTask', description: 'Runs a task' }] + }, + }); + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.ChatTurnStarted, + turnId: 'turn-1', + message: { text: 'run it', origin: { kind: MessageKind.User } }, + }); + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.ChatToolCallStart, + turnId: 'turn-1', + toolCallId: 'tool-1', + toolName: 'runTask', + displayName: 'Run Task', + contributor: { kind: ToolCallContributorKind.Client, clientId: 'client-tools' }, + }); + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.ChatToolCallReady, + turnId: 'turn-1', + toolCallId: 'tool-1', + invocationMessage: 'Run Task', + toolInput: '{}', + confirmed: ToolCallConfirmationReason.NotNeeded, + }); + + transport1.simulateClose(); + + // Reconnect, but do NOT resubscribe to the session. + const transport2 = new MockProtocolTransport(); + server.simulateConnection(transport2); + const reconnectRespPromise = waitForResponse(transport2, 1); + transport2.simulateMessage(request(1, 'reconnect', { + clientId: 'client-tools', + lastSeenServerSeq: initSeq, + subscriptions: [], + })); + await reconnectRespPromise; + + assert.deepStrictEqual(stateManager.getSessionState(sessionUri)?.activeClients, []); + const part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0]; + assert.deepStrictEqual(part?.kind === ResponsePartKind.ToolCall ? { + status: part.toolCall.status, + success: part.toolCall.status === ToolCallStatus.Completed ? part.toolCall.success : undefined, + error: part.toolCall.status === ToolCallStatus.Completed ? part.toolCall.error?.message : undefined, + } : undefined, { + status: ToolCallStatus.Completed, + success: false, + error: 'Client client-tools disconnected before completing Run Task', + }); + + transport2.simulateClose(); + }); + + test('reconnect with resubscription keeps the active client and its owned tool calls', async () => { + stateManager.createSession(makeSessionSummary()); + stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); + + const transport1 = connectClient('client-tools', [sessionUri]); + const initSeq = (findResponse(transport1.sent, 1) as { result: InitializeResult }).result.serverSeq; + + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.SessionActiveClientSet, + activeClient: { + clientId: 'client-tools', + tools: [{ name: 'runTask', description: 'Runs a task' }] + }, + }); + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.ChatTurnStarted, + turnId: 'turn-1', + message: { text: 'run it', origin: { kind: MessageKind.User } }, + }); + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.ChatToolCallStart, + turnId: 'turn-1', + toolCallId: 'tool-1', + toolName: 'runTask', + displayName: 'Run Task', + contributor: { kind: ToolCallContributorKind.Client, clientId: 'client-tools' }, + }); + + transport1.simulateClose(); + + const transport2 = new MockProtocolTransport(); + server.simulateConnection(transport2); + const reconnectRespPromise = waitForResponse(transport2, 1); + transport2.simulateMessage(request(1, 'reconnect', { + clientId: 'client-tools', + lastSeenServerSeq: initSeq, + subscriptions: [sessionUri], + })); + await reconnectRespPromise; + + assert.deepStrictEqual(stateManager.getSessionState(sessionUri)?.activeClients.map(c => c.clientId), ['client-tools']); + const part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0]; + assert.strictEqual(part?.kind === ResponsePartKind.ToolCall ? part.toolCall.status : undefined, ToolCallStatus.Streaming); + + transport2.simulateClose(); + }); + test('handshake includes defaultDirectory from side effects', () => { const transport = connectClient('client-home'); diff --git a/src/vs/platform/agentHost/test/node/reducers.test.ts b/src/vs/platform/agentHost/test/node/reducers.test.ts index 2c54ef35288c25..109183f86bf753 100644 --- a/src/vs/platform/agentHost/test/node/reducers.test.ts +++ b/src/vs/platform/agentHost/test/node/reducers.test.ts @@ -22,6 +22,7 @@ function makeSession(): SessionState { project: { uri: 'file:///test-project', displayName: 'Test Project' }, }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], }; } diff --git a/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts b/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts index cfcb6ffd6f904e..73318d40865b45 100644 --- a/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts +++ b/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts @@ -26,7 +26,7 @@ import { KNOWN_MODE_VALUES, SessionConfigKey } from '../../../../../platform/age import { migrateLegacyAutopilotConfig } from '../../../../../platform/agentHost/common/agentHostSchema.js'; import type { IAgentSubscription } from '../../../../../platform/agentHost/common/state/agentSubscription.js'; import { ResolveSessionConfigResult } from '../../../../../platform/agentHost/common/state/protocol/commands.js'; -import { AgentCustomization, AgentSelection, ChangesSummary, type ChangesetFile, Customization, CustomizationType, ModelSelection, SessionStatus as ProtocolSessionStatus, RootConfigState, RootState, SessionActiveClient, SessionState, SessionSummary, type Changeset } from '../../../../../platform/agentHost/common/state/protocol/state.js'; +import { AgentCustomization, AgentSelection, ChangesSummary, type ChangesetFile, type ClientPluginCustomization, Customization, CustomizationType, ModelSelection, SessionStatus as ProtocolSessionStatus, RootConfigState, RootState, SessionActiveClient, SessionState, SessionSummary, type Changeset } from '../../../../../platform/agentHost/common/state/protocol/state.js'; import { ActionType, isChatAction, isSessionAction, NotificationType } from '../../../../../platform/agentHost/common/state/sessionActions.js'; import { AgentInfo, buildChatUri, buildDefaultChatUri, isDefaultChatUri, parseChatUri, readSessionGitHubState, readSessionGitState, ROOT_STATE_URI, SessionMeta, SessionSummaryMeta, StateComponents, type ChatSummary, type ISessionGitState } from '../../../../../platform/agentHost/common/state/sessionState.js'; import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js'; @@ -836,15 +836,9 @@ function customizationsChanged(previous: SessionState, state: SessionState): boo if (previous.customizations !== state.customizations) { return true; } - const previousActiveCustomizations = previous.activeClient?.customizations; - const currentActiveCustomizations = state.activeClient?.customizations; - if (previousActiveCustomizations === currentActiveCustomizations) { - return false; - } - if (!previousActiveCustomizations || !currentActiveCustomizations) { - return true; - } - return arrayEquals(previousActiveCustomizations, currentActiveCustomizations, (a, b) => { + const previousActiveCustomizations = flattenActiveClientCustomizations(previous); + const currentActiveCustomizations = flattenActiveClientCustomizations(state); + return !arrayEquals(previousActiveCustomizations, currentActiveCustomizations, (a, b) => { if (a.nonce !== undefined && a.nonce === b.nonce) { return true; } @@ -852,6 +846,17 @@ function customizationsChanged(previous: SessionState, state: SessionState): boo }); } +/** Flattens the customizations contributed by every active client of a session. */ +function flattenActiveClientCustomizations(state: SessionState): ClientPluginCustomization[] { + const result: ClientPluginCustomization[] = []; + for (const client of state.activeClients) { + if (client.customizations) { + result.push(...client.customizations); + } + } + return result; +} + // ============================================================================ // NewSession — bundles the in-flight new-session state // ============================================================================ diff --git a/src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts b/src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts index 448dba16170a6d..102268a8dda7d0 100644 --- a/src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts +++ b/src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts @@ -1034,6 +1034,7 @@ suite('LocalAgentHostSessionsProvider', () => { modifiedAt: 0, }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], customizations: [{ type: CustomizationType.Plugin, @@ -1158,6 +1159,7 @@ suite('LocalAgentHostSessionsProvider', () => { modifiedAt: 0, }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], customizations: [{ type: CustomizationType.Plugin, @@ -1172,7 +1174,7 @@ suite('LocalAgentHostSessionsProvider', () => { assert.ok(fired > afterRoot, 'expected event to fire on session state customization change'); // A second state update with the SAME customizations reference must - // NOT fire — only churn in `customizations` / `activeClient.customizations` + // NOT fire — only churn in `customizations` / `activeClients[].customizations` // counts. const afterFirstCustomization = fired; agentHost.setSessionState('cust-events', 'copilotcli', { @@ -1185,6 +1187,7 @@ suite('LocalAgentHostSessionsProvider', () => { modifiedAt: 0, }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], // Same identity as before: customizations: (provider as unknown as { _lastSessionStates: Map })._lastSessionStates.get(session!.sessionId)?.customizations, @@ -1227,6 +1230,7 @@ suite('LocalAgentHostSessionsProvider', () => { modifiedAt: 0, }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], customizations, }; @@ -1271,6 +1275,7 @@ suite('LocalAgentHostSessionsProvider', () => { modifiedAt: 0, }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], customizations: [{ type: CustomizationType.Plugin, @@ -1770,6 +1775,7 @@ suite('LocalAgentHostSessionsProvider', () => { modifiedAt: 0, }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats, ...(opts?.defaultChat ? { defaultChat: opts.defaultChat } : {}), }; @@ -2244,6 +2250,7 @@ suite('LocalAgentHostSessionsProvider', () => { const fakeState: SessionState = { summary: { resource: AgentSession.uri('copilotcli', 'seed-1').toString(), provider: 'copilotcli', title: 'Seeded Session', status: ProtocolSessionStatus.Idle, createdAt: 0, modifiedAt: 0 }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], config, }; @@ -2275,6 +2282,7 @@ suite('LocalAgentHostSessionsProvider', () => { const fullState: SessionState = { summary: { resource: AgentSession.uri('copilotcli', 'seed-schema').toString(), provider: 'copilotcli', title: 'Schema Preserve Session', status: ProtocolSessionStatus.Idle, createdAt: 0, modifiedAt: 0 }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], config: { schema: { @@ -2392,6 +2400,7 @@ suite('LocalAgentHostSessionsProvider', () => { agentHost.setSessionState('pr-sticky', 'copilotcli', { summary: { resource: AgentSession.uri('copilotcli', 'pr-sticky').toString(), provider: 'copilotcli', title: 'PR Session', status: ProtocolSessionStatus.Idle, createdAt: 0, modifiedAt: 0 }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], _meta: { git: { hasGitHubRemote: true, githubOwner: 'owner', githubRepo: 'repo', branchName: 'feature' } }, }); @@ -2448,6 +2457,7 @@ suite('LocalAgentHostSessionsProvider', () => { const fakeState: SessionState = { summary: { resource: AgentSession.uri('copilotcli', 'rep-1').toString(), provider: 'copilotcli', title: 'Replace Session', status: ProtocolSessionStatus.Idle, createdAt: 0, modifiedAt: 0 }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], config, }; @@ -2500,6 +2510,7 @@ suite('LocalAgentHostSessionsProvider', () => { const fakeState: SessionState = { summary: { resource: AgentSession.uri('copilotcli', 'policy-write').toString(), provider: 'copilotcli', title: 'Policy Write Session', status: ProtocolSessionStatus.Idle, createdAt: 0, modifiedAt: 0 }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], config, }; @@ -2556,6 +2567,7 @@ suite('LocalAgentHostSessionsProvider', () => { const fakeState: SessionState = { summary: { resource: AgentSession.uri('copilotcli', 'schema-write').toString(), provider: 'copilotcli', title: 'Schema Write Session', status: ProtocolSessionStatus.Idle, createdAt: 0, modifiedAt: 0 }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], config, }; @@ -2623,6 +2635,7 @@ suite('LocalAgentHostSessionsProvider', () => { const fakeState: SessionState = { summary: { resource: AgentSession.uri('copilotcli', 'rep-2').toString(), provider: 'copilotcli', title: 'No-op Session', status: ProtocolSessionStatus.Idle, createdAt: 0, modifiedAt: 0 }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], config, }; @@ -2649,6 +2662,7 @@ suite('LocalAgentHostSessionsProvider', () => { const fakeState: SessionState = { summary: { resource: AgentSession.uri('copilotcli', 'cfg-merge').toString(), provider: 'copilotcli', title: 'Merge Session', status: ProtocolSessionStatus.Idle, createdAt: 0, modifiedAt: 0 }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], config: { schema: { @@ -2689,6 +2703,7 @@ suite('LocalAgentHostSessionsProvider', () => { const fakeState: SessionState = { summary: { resource: AgentSession.uri('copilotcli', 'cfg-replace').toString(), provider: 'copilotcli', title: 'Replace Session', status: ProtocolSessionStatus.Idle, createdAt: 0, modifiedAt: 0 }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], config: { schema: { diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts index 47acddd7d16a05..e6f97874fb78e8 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts @@ -1059,6 +1059,7 @@ suite('RemoteAgentHostSessionsProvider', () => { const fakeState: SessionState = { summary: { resource: AgentSession.uri('copilotcli', 'seed-1').toString(), provider: 'copilotcli', title: 'Seeded Session', status: ProtocolSessionStatus.Idle, createdAt: 0, modifiedAt: 0 }, lifecycle: SessionLifecycle.Ready, + activeClients: [], chats: [], config, }; diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLocalCustomizations.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLocalCustomizations.ts index ce1ab650492de0..817f98ee448679 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLocalCustomizations.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLocalCustomizations.ts @@ -59,7 +59,7 @@ export interface ILocalCustomizationFile { * * This is the single source of truth used by both the AI Customization view * (to render disable affordances) and the agent host wire (to compute the - * `customizations` set published via `activeClientChanged`). + * `customizations` set published via `activeClientSet`). * * Built-in skills bundled with the Agents app (only present when the * sessions-aware prompts service is in play) are also enumerated so that @@ -186,7 +186,7 @@ export function collectNonPluginMcpServers(mcpService: IMcpService): ISyncableMc } /** - * Resolves the customization refs to include in an `activeClientChanged` + * Resolves the customization refs to include in an `activeClientSet` * message. * * Every eligible local file is synced unless the user opted out. Files 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 287bb7b8cdb241..877c694cf9e8fb 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -485,12 +485,14 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC this._register(autorun(reader => { const defs = this._activeClientService.clientTools.read(reader); + const clientId = this._config.connection.clientId; for (const [sessionResource] of this._activeSessions) { const backendSession = this._resolveSessionUri(sessionResource); const state = this._getSessionState(backendSession.toString()); - if (state?.activeClient?.clientId === this._config.connection.clientId) { + if (state?.activeClients.some(c => c.clientId === clientId)) { this._dispatchAction(backendSession, { type: ActionType.SessionActiveClientToolsChanged, + clientId, tools: [...defs], }); } @@ -533,10 +535,12 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC const customizationsObs = this._activeClientService.getCustomizations(config.sessionType); this._register(autorun(reader => { const refs = customizationsObs.read(reader); + const clientId = this._config.connection.clientId; for (const [sessionResource] of this._activeSessions) { const backendSession = this._resolveSessionUri(sessionResource); const state = this._getSessionState(backendSession.toString()); - if (state?.activeClient?.clientId === this._config.connection.clientId && !equals(state.activeClient.customizations ?? [], refs)) { + const existing = state?.activeClients.find(c => c.clientId === clientId); + if (existing && !equals(existing.customizations ?? [], refs)) { this._dispatchActiveClient(backendSession, [...refs]); } } @@ -1153,24 +1157,25 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC private _ensureActiveClientForMessage(backendSession: URI): void { const state = this._getSessionState(backendSession.toString()); const activeClient = this._getCurrentActiveClient(); - if (equals(state?.activeClient, activeClient)) { + const existing = state?.activeClients.find(c => c.clientId === activeClient.clientId); + if (equals(existing, activeClient)) { return; } this._dispatchAction(backendSession, { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient, }); } /** - * Dispatches `session/activeClientChanged` to claim the active client - * role for this session and publish the current customizations and - * client-provided tools. + * Dispatches `session/activeClientSet` to add this connection as an + * active client for this session and publish the current customizations + * and client-provided tools. This client never removes itself. */ private _dispatchActiveClient(backendSession: URI, customizations: ClientPluginCustomization[]): void { const current = this._getCurrentActiveClient(); this._dispatchAction(backendSession, { - type: ActionType.SessionActiveClientChanged, + type: ActionType.SessionActiveClientSet, activeClient: { ...current, customizations }, }); } @@ -1328,10 +1333,10 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC return; } - // Claim the active-client slot for this connection before the turn - // goes out. We only claim on turn start (not on session open) so - // that opening a session doesn't dispossess another client that's - // in the middle of a turn. + // Add this connection as an active client for the session before the + // turn goes out. We only do this on turn start (not on session open) + // so that opening a session doesn't eagerly register this client while + // another client is in the middle of a turn. this._ensureActiveClientForMessage(session); // Model and agent selections are dispatched to the per-chat turn diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts index 4540c7fd87a669..8ad4b4ee23dec5 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts @@ -154,7 +154,7 @@ class MockAgentHostService extends mock() { const state: SessionState = { ...createSessionState(summary), lifecycle: SessionLifecycle.Ready, - activeClient: config.activeClient, + activeClients: [config.activeClient], }; this.sessionStates.set(session.toString(), state); } @@ -171,7 +171,7 @@ class MockAgentHostService extends mock() { public dispatchedActions: { channel: string; action: SessionAction | ChatAction | TerminalAction | ClientAnnotationsAction | IRootConfigChangedAction; clientId: string; clientSeq: number }[] = []; /** Returns dispatched actions filtered to turn-related types only - * (excludes lifecycle actions like activeClientChanged). */ + * (excludes lifecycle actions like activeClientSet). */ get turnActions() { return this.dispatchedActions.filter(d => d.action.type === 'chat/turnStarted'); } @@ -342,7 +342,11 @@ class MockAgentHostService extends mock() { // logic (e.g. customization re-dispatch) sees the correct activeClient. // Turn lifecycle actions (turnStarted, toolCallConfirmed, etc.) are applied // later via fireAction when the server echoes them back. - if (isSessionAction(action) && action.type === 'session/activeClientChanged') { + if (isSessionAction(action) && ( + action.type === 'session/activeClientSet' + || action.type === 'session/activeClientRemoved' + || action.type === 'session/activeClientToolsChanged' + )) { const entry = this._liveSubscriptions.get(channel.toString()); if (entry) { const noop = () => { }; @@ -715,7 +719,7 @@ async function startTurn( const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); ds.add(toDisposable(() => chatSession.dispose())); - // Clear any lifecycle actions (e.g. activeClientChanged from customization setup) + // Clear any lifecycle actions (e.g. activeClientSet from customization setup) // so tests only see turn-related dispatches. agentHostService.dispatchedActions.length = 0; @@ -742,7 +746,7 @@ async function startTurn( await timeout(10); - // Filter for turn-related dispatches only (skip activeClientChanged etc.) + // Filter for turn-related dispatches only (skip activeClientSet etc.) const turnDispatches = agentHostService.dispatchedActions.filter(d => d.action.type === 'chat/turnStarted'); const lastDispatch = turnDispatches[turnDispatches.length - 1] ?? agentHostService.dispatchedActions[agentHostService.dispatchedActions.length - 1]; const session = lastDispatch?.channel.toString(); @@ -5483,7 +5487,7 @@ suite('AgentHostChatContribution', () => { suite('customizations', () => { - test('dispatches activeClientChanged when a new session is created', async () => { + test('dispatches activeClientSet when a new session is created', async () => { const { instantiationService, agentHostService, chatAgentService, seedActiveClient } = createTestServices(disposables); const customizations = observableValue('customizations', [ @@ -5515,7 +5519,7 @@ suite('AgentHostChatContribution', () => { assert.strictEqual(createCall!.activeClient!.customizations?.[0].uri, 'file:///plugin-a'); }); - test('re-dispatches activeClientChanged when customizations observable changes', async () => { + test('re-dispatches activeClientSet when customizations observable changes', async () => { const { instantiationService, agentHostService, chatAgentService, seedActiveClient } = createTestServices(disposables); const customizations = observableValue('customizations', []); @@ -5544,15 +5548,15 @@ suite('AgentHostChatContribution', () => { ], undefined); const activeClientAction = agentHostService.dispatchedActions.find( - d => d.action.type === 'session/activeClientChanged' + d => d.action.type === 'session/activeClientSet' ); - assert.ok(activeClientAction, 'should re-dispatch activeClientChanged on change'); + assert.ok(activeClientAction, 'should re-dispatch activeClientSet on change'); const ac = activeClientAction!.action as { activeClient: { customizations?: ClientPluginCustomization[] } }; assert.strictEqual(ac.activeClient.customizations?.length, 1); assert.strictEqual(ac.activeClient.customizations?.[0].uri, 'file:///plugin-b'); }); - test('does not dispatch activeClientChanged when an existing session is restored and this client is already active', async () => { + test('does not dispatch activeClientSet when an existing session is restored and this client is already active', async () => { const { instantiationService, agentHostService } = createTestServices(disposables); const sessionResource = AgentSession.uri('copilot', 'existing-session'); const summary: SessionSummary = { @@ -5566,11 +5570,11 @@ suite('AgentHostChatContribution', () => { agentHostService.sessionStates.set(sessionResource.toString(), { ...createSessionState(summary), lifecycle: SessionLifecycle.Ready, - activeClient: { + activeClients: [{ clientId: agentHostService.clientId, tools: [], customizations: [], - }, + }], }); const sessionHandler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, { @@ -5587,12 +5591,12 @@ suite('AgentHostChatContribution', () => { disposables.add(toDisposable(() => chatSession.dispose())); assert.strictEqual( - agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientChanged').length, + agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientSet').length, 0, ); }); - test('dispatches activeClientChanged on first turn when restoring a session where another client is active', async () => { + test('dispatches activeClientSet on first turn when restoring a session where another client is active', async () => { const { instantiationService, agentHostService, chatAgentService } = createTestServices(disposables); const sessionResource = AgentSession.uri('copilot', 'existing-session'); const summary: SessionSummary = { @@ -5606,10 +5610,10 @@ suite('AgentHostChatContribution', () => { agentHostService.sessionStates.set(sessionResource.toString(), { ...createSessionState(summary), lifecycle: SessionLifecycle.Ready, - activeClient: { + activeClients: [{ clientId: 'other-client', tools: [], - }, + }], }); const sessionHandler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, { @@ -5629,7 +5633,7 @@ suite('AgentHostChatContribution', () => { const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); disposables.add(toDisposable(() => chatSession.dispose())); assert.strictEqual( - agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientChanged').length, + agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientSet').length, 0, 'no dispatch expected on session open', ); @@ -5639,12 +5643,12 @@ suite('AgentHostChatContribution', () => { fire({ type: 'chat/turnComplete', session, turnId } as ChatAction); await turnPromise; - const activeClientActions = agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientChanged'); + const activeClientActions = agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientSet'); assert.strictEqual(activeClientActions.length, 1); assert.strictEqual(activeClientActions[0].channel, sessionResource.toString()); }); - test('dispatches activeClientChanged on first turn when restoring a session where current client customizations are stale', async () => { + test('dispatches activeClientSet on first turn when restoring a session where current client customizations are stale', async () => { const { instantiationService, agentHostService, chatAgentService, seedActiveClient } = createTestServices(disposables); const customizations = observableValue('customizations', [ { type: CustomizationType.Plugin, id: 'file:///plugin-new', uri: 'file:///plugin-new', name: 'Plugin New', enabled: true }, @@ -5662,11 +5666,11 @@ suite('AgentHostChatContribution', () => { agentHostService.sessionStates.set(sessionResource.toString(), { ...createSessionState(summary), lifecycle: SessionLifecycle.Ready, - activeClient: { + activeClients: [{ clientId: agentHostService.clientId, tools: [], customizations: [{ type: CustomizationType.Plugin, id: 'file:///plugin-old', uri: 'file:///plugin-old', name: 'Plugin Old', enabled: true }], - }, + }], }); const sessionHandler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, { @@ -5683,7 +5687,7 @@ suite('AgentHostChatContribution', () => { const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); disposables.add(toDisposable(() => chatSession.dispose())); assert.strictEqual( - agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientChanged').length, + agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientSet').length, 0, 'no dispatch expected on session open', ); @@ -5693,11 +5697,11 @@ suite('AgentHostChatContribution', () => { fire({ type: 'chat/turnComplete', session, turnId } as ChatAction); await turnPromise; - const activeClientActions = agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientChanged'); + const activeClientActions = agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientSet'); assert.strictEqual(activeClientActions.length, 1); - const activeClientAction = activeClientActions[0].action; - assert.strictEqual(activeClientAction.type, 'session/activeClientChanged'); - assert.deepStrictEqual(activeClientAction.activeClient?.customizations, [ + const activeClientAction = activeClientActions[0].action as { type: string; activeClient: { customizations?: ClientPluginCustomization[] } }; + assert.strictEqual(activeClientAction.type, 'session/activeClientSet'); + assert.deepStrictEqual(activeClientAction.activeClient.customizations, [ { type: CustomizationType.Plugin, id: 'file:///plugin-new', uri: 'file:///plugin-new', name: 'Plugin New', enabled: true }, ]); }); 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 ddc87f7852a8ab..44ecf748b508fb 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 @@ -540,7 +540,7 @@ suite('AgentHostClientTools', () => { test('maps allowlisted tool data to protocol definitions', async () => { const { connection } = createHandlerWithMocks(disposables, [testRunTestsTool, testRunTaskTool, testUnlistedTool]); - // The handler dispatches activeClientChanged in the constructor when + // The handler dispatches activeClientSet in the constructor when // customizations observable fires, but here it fires during provideChatSessionContent. // Verify tools are built correctly by checking what would be dispatched. assert.ok(connection);