diff --git a/src/vs/platform/agentHost/common/agent.ts b/src/vs/platform/agentHost/common/agent.ts index e50afee74b8e9..832b68fc44c9f 100644 --- a/src/vs/platform/agentHost/common/agent.ts +++ b/src/vs/platform/agentHost/common/agent.ts @@ -1171,6 +1171,12 @@ export interface IAgentChatAdoptionResult { readonly reason?: AgentChatAdoptionReason; } +/** Identifies the client that submitted a pending message. */ +export interface IAgentPendingMessageSender { + readonly clientId: string | undefined; + readonly clientContext: IAgentHostClientTelemetryContext; +} + /** * Implemented by each agent backend (e.g. Copilot SDK). * The {@link IAgentService} dispatches to the appropriate agent based on @@ -1215,7 +1221,7 @@ export interface IAgent { materializeChat(chat: URI, context: URI | IAgentChatContext, providerData: string | undefined): Promise; /** Optional steering hook for providers that can accept messages during an active turn. */ - setPendingMessages?(chat: URI, steeringMessage: PendingMessage | undefined, queuedMessages: readonly PendingMessage[]): void; + setPendingMessages?(chat: URI, steeringMessage: PendingMessage | undefined, queuedMessages: readonly PendingMessage[], steeringSender?: IAgentPendingMessageSender): void; /** Optional history mutation for providers with a native truncation operation. */ truncateChat?(chat: URI, turnId: string | undefined, context?: URI | IAgentChatContext): Promise; diff --git a/src/vs/platform/agentHost/node/agentSideEffects.ts b/src/vs/platform/agentHost/node/agentSideEffects.ts index 8365e08cb28a9..ce56b984549d5 100644 --- a/src/vs/platform/agentHost/node/agentSideEffects.ts +++ b/src/vs/platform/agentHost/node/agentSideEffects.ts @@ -363,6 +363,9 @@ export class AgentSideEffects extends Disposable { this._notifyClientToolCallComplete(sessionChannel, envelope.channel, action.toolCallId, action.result, 'server-envelope'); } } + if (!envelope.origin && envelope.action.type === ActionType.SessionActiveClientRemoved) { + this._removeActiveClient(envelope.channel, envelope.action.clientId); + } // A chat joining the catalog changes the session's authoritative // membership, so every already-contributing client is re-fanned-out // over the new set. Handled here (not `handleAction`) because every @@ -419,6 +422,13 @@ export class AgentSideEffects extends Disposable { } } + private _removeActiveClient(session: ProtocolURI, clientId: string): void { + const agent = this._options.getAgent(session); + for (const chat of getSessionChatsForFanOut(this._stateManager, session) ?? []) { + agent?.removeActiveClient(chat, this._chatContext(session, chat.toString()), clientId); + } + } + /** * Publishes agent descriptors using the last known model lists. */ @@ -1572,10 +1582,7 @@ export class AgentSideEffects extends Disposable { break; } case ActionType.SessionActiveClientRemoved: { - const agent = this._options.getAgent(channel); - for (const chat of getSessionChatsForFanOut(this._stateManager, channel) ?? []) { - agent?.removeActiveClient(chat, this._chatContext(channel, chat.toString()), action.clientId); - } + this._removeActiveClient(channel, action.clientId); break; } case ActionType.RootConfigChanged: { diff --git a/src/vs/platform/agentHost/node/chatContributions/queueDrain/queueDrainContribution.ts b/src/vs/platform/agentHost/node/chatContributions/queueDrain/queueDrainContribution.ts index 769cc4bec90ee..4af8f820d6de3 100644 --- a/src/vs/platform/agentHost/node/chatContributions/queueDrain/queueDrainContribution.ts +++ b/src/vs/platform/agentHost/node/chatContributions/queueDrain/queueDrainContribution.ts @@ -9,6 +9,7 @@ import { URI } from '../../../../../base/common/uri.js'; import { generateUuid } from '../../../../../base/common/uuid.js'; import { IInstantiationService } from '../../../../instantiation/common/instantiation.js'; import { ILogService } from '../../../../log/common/log.js'; +import type { IAgentPendingMessageSender } from '../../../common/agent.js'; import { AgentHostClientType } from '../../../common/agentHostClientInfo.js'; import { createUnknownAgentHostClientTelemetryContext } from '../../../common/agentHostTelemetry.js'; import { IAgentHostChatContributions, createChatMementoKey, type IAgentHostChatContribution, type IAgentHostChatContributionContext, type IAgentHostChatContributionHost, type IAppliedClientAction, type IQueuedMessageSender, type ITurnEnd } from '../../../common/agentHostChatContributionsService.js'; @@ -20,6 +21,7 @@ import { startTurn } from '../../agentHostTurnStarter.js'; import { ISessionWorkspaceConversionService } from '../sessionWorkspaceConversion/sessionWorkspaceConversionService.js'; const QueuedSender = createChatMementoKey('queueDrain.sender', () => undefined); +const SteeringSender = createChatMementoKey<{ readonly messageId: string; readonly sender: IAgentPendingMessageSender } | undefined>('queueDrain.steeringSender', () => undefined); /** Owns queued-message sender state and decides when a queued turn can be admitted. */ export class QueueDrainContribution extends Disposable implements IAgentHostChatContribution { @@ -53,12 +55,20 @@ export class QueueDrainContribution extends Disposable implements IAgentHostChat const action = observed.action; switch (action.type) { case ActionType.ChatPendingMessageSet: { - const queuedMessageExists = this._stateManager.getChatState(observed.channel)?.queuedMessages?.some(message => message.id === action.id) === true; - if (action.kind === PendingMessageKind.Queued && queuedMessageExists) { + const state = this._stateManager.getChatState(observed.channel); + if (action.kind === PendingMessageKind.Queued && state?.queuedMessages?.some(message => message.id === action.id) === true) { this._context.memento(QueuedSender, observed.channel, action.id).set({ clientId: observed.clientId, clientContext: observed.clientContext, }, undefined); + } else if (action.kind === PendingMessageKind.Steering && state?.steeringMessage?.id === action.id) { + this._context.memento(SteeringSender, observed.channel).set({ + messageId: action.id, + sender: { + clientId: observed.clientId, + clientContext: observed.clientContext, + }, + }, undefined); } this._syncPendingMessages(observed.channel); break; @@ -66,6 +76,11 @@ export class QueueDrainContribution extends Disposable implements IAgentHostChat case ActionType.ChatPendingMessageRemoved: { if (action.kind === PendingMessageKind.Queued) { this._context.deleteMemento(QueuedSender, observed.channel, action.id); + } else { + const steeringSender = this._context.memento(SteeringSender, observed.channel); + if (steeringSender.get()?.messageId === action.id) { + steeringSender.set(undefined, undefined); + } } this._syncPendingMessages(observed.channel); break; @@ -86,7 +101,13 @@ export class QueueDrainContribution extends Disposable implements IAgentHostChat return; } const session = parseRequiredSessionUriFromChatUri(channel); - this._providerService.getProviderForSession(session)?.setPendingMessages?.(URI.parse(channel), state.steeringMessage, []); + const steeringSender = this._context.memento(SteeringSender, channel).get(); + this._providerService.getProviderForSession(session)?.setPendingMessages?.( + URI.parse(channel), + state.steeringMessage, + [], + steeringSender && steeringSender.messageId === state.steeringMessage?.id ? steeringSender.sender : undefined, + ); this._tryConsumeNextQueuedMessage(channel); } diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts index cf1879240cfff..704503a4ed67d 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts @@ -44,7 +44,7 @@ import { CopilotCliConfigKey, CopilotCliVSCodeAssignmentContextKey, copilotCliCo import { AgentHostAutoApprovePolicyRestrictedConfigKey, AgentHostByokModelsEnabledConfigKey, AgentHostMcpServersConfigKey, AgentHostGitHubMcpServerEnabledConfigKey, AgentHostCopilotMultiRootEnabledConfigKey, AgentHostSessionSyncEnabledConfigKey, AgentHostSystemProxyEnabledConfigKey, AgentHostMigrateLegacyCopilotCliEnabledConfigKey, AgentHostProxyConfigKey, agentHostProxyConfigSchema, AutoApproveLevel, SessionMode, migrateLegacyAutopilotConfig, platformRootSchema, platformSessionSchema, type AgentHostMcpServers } from '../../common/agentHostSchema.js'; import { IAgentPluginManager, ISyncedCustomization } from '../../common/agentPluginManager.js'; import { decodeProviderData, encodeProviderData, type IPersistedChat } from '../agentChatBackings.js'; -import { AgentChatOperationContext, AgentSession, AgentSignal, AuthenticateParams, IActiveClient, IAgent, IAgentChatAdoptionResult, type IAgentAdoptedWorktree, IAgentChatConfigCompletionsParams, IAgentChatContext, IAgentChatDataChange, IAgentChatMetadata, IAgentChats, IAgentLegacyChat, IAgentCreateChatOptions, IAgentCreateChatResult, IAgentDescriptor, IAgentDiscoveredChat, IAgentHostManagedSettingsSnapshot, IAgentHostNetworkEndpoint, IAgentKnownSessionsFilter, IAgentMaterializeChatEvent, IAgentModelInfo, IAgentResolveChatConfigParams, IAgentSessionProjectInfo, IAgentSpawnChatEvent, IMcpNotification, SubagentChatSignal, resolveAgentChatContext, resolveAgentHostCustomizations, resolveAgentHostInstructions, resolveSubagentChatParent, type IAgentTurnDiagnosticSnapshot, type IAgentTurnTokenUsage } from '../../common/agent.js'; +import { AgentChatOperationContext, AgentSession, AgentSignal, AuthenticateParams, IActiveClient, IAgent, IAgentChatAdoptionResult, type IAgentAdoptedWorktree, IAgentChatConfigCompletionsParams, IAgentChatContext, IAgentChatDataChange, IAgentChatMetadata, IAgentChats, IAgentLegacyChat, IAgentCreateChatOptions, IAgentCreateChatResult, IAgentDescriptor, IAgentDiscoveredChat, IAgentHostManagedSettingsSnapshot, IAgentHostNetworkEndpoint, IAgentKnownSessionsFilter, IAgentMaterializeChatEvent, IAgentModelInfo, type IAgentPendingMessageSender, IAgentResolveChatConfigParams, IAgentSessionProjectInfo, IAgentSpawnChatEvent, IMcpNotification, SubagentChatSignal, resolveAgentChatContext, resolveAgentHostCustomizations, resolveAgentHostInstructions, resolveSubagentChatParent, type IAgentTurnDiagnosticSnapshot, type IAgentTurnTokenUsage } from '../../common/agent.js'; import { getReasoningEffortDescription, getReasoningEffortLabel, resolveDefaultReasoningEffort } from '../../common/reasoningEffort.js'; import { autoModeTiers, defaultAutoModeTier, getAutoModeTierDescription, getAutoModeTierLabel } from '../../common/autoModeTiers.js'; import { isAutoModel } from './modelIdentifiers.js'; @@ -4333,7 +4333,7 @@ export class CopilotAgent extends Disposable implements IAgent { }; } - setPendingMessages(chat: URI, steeringMessage: PendingMessage | undefined, _queuedMessages: readonly PendingMessage[]): void { + setPendingMessages(chat: URI, steeringMessage: PendingMessage | undefined, _queuedMessages: readonly PendingMessage[], steeringSender?: IAgentPendingMessageSender): void { const backing = this._chatBackings.get(chat.toString()); const target = backing ? this._findSessionBySdkId(backing.sdkSessionId) : undefined; if (!target) { @@ -4343,7 +4343,7 @@ export class CopilotAgent extends Disposable implements IAgent { // Steering: send with mode 'immediate' so the SDK injects it mid-turn if (steeringMessage) { - target.sendSteering(steeringMessage); + target.sendSteering(steeringMessage, steeringSender); } // Queued messages are consumed by the server (AgentSideEffects) diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts index 03f3625620f9e..ade35b628d3b7 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts @@ -43,7 +43,7 @@ import { getSessionSandboxOverrides } from '../sessionSandbox.js'; import { AgentHostSandboxConfigKey, sandboxConfigSchema } from '../../common/sandboxConfigSchema.js'; import { AgentHostAutoApprovePolicyRestrictedConfigKey, AgentHostGlobalAutoApproveEnabledConfigKey, AgentHostAutoReplyAnswer, AgentHostAutoReplyEnabledConfigKey, AgentHostDisableRepoInfoTelemetryConfigKey, platformRootSchema, platformSessionSchema } from '../../common/agentHostSchema.js'; import { createUnknownAgentHostClientTelemetryContext, type IAgentHostClientTelemetryContext } from '../../common/agentHostTelemetry.js'; -import { AgentSession, AgentSignal, AgentWorkingDirectoryChangedError, AuthenticateParams, IMcpNotification, type AgentSubagentTaskModelSource, type AgentTurnProviderCallState, type IAgentToolPendingConfirmationSignal, type IAgentTurnDiagnosticSnapshot, type IAgentTurnTokenUsage } from '../../common/agent.js'; +import { AgentSession, AgentSignal, AgentWorkingDirectoryChangedError, AuthenticateParams, IMcpNotification, type AgentSubagentTaskModelSource, type AgentTurnProviderCallState, type IAgentPendingMessageSender, type IAgentToolPendingConfirmationSignal, type IAgentTurnDiagnosticSnapshot, type IAgentTurnTokenUsage } from '../../common/agent.js'; import { isReasoningEffortLevel } from '../../common/reasoningEffort.js'; import { ObservedTokenUsage } from './observedTokenUsage.js'; import { META_DIFF_BASE_BRANCH } from '../../common/agentHostGitService.js'; @@ -796,6 +796,11 @@ class CopilotTurn extends Disposable { } } +interface IPendingSteering { + readonly pendingMessage: PendingMessage; + readonly sender: IAgentPendingMessageSender | undefined; +} + /** * Encapsulates a single Copilot SDK session and all its associated bookkeeping. * @@ -1099,7 +1104,7 @@ export class CopilotAgentSession extends Disposable { * `steering_consumed` signals so the chat UI's pending state still * clears in cleanup paths where we never observe the echo. */ - private readonly _pendingSteeringFlips = new Map(); + private readonly _pendingSteeringFlips = new Map(); /** Snapshot captured at session creation for refresh detection. */ private readonly _appliedSnapshot: IActiveClientSnapshot; @@ -1425,15 +1430,15 @@ export class CopilotAgentSession extends Disposable { * handler) can associate the SDK event id with the steering turn for * history.truncate / sessions.fork mapping. */ - private _beginSteeringTurn(steering: PendingMessage): string { + private _beginSteeringTurn(steering: IPendingSteering): string { this._completeActiveTurn(); const newTurnId = generateUuid(); this._emitAction({ type: ActionType.ChatTurnStarted, turnId: newTurnId, startedAt: new Date().toISOString(), - message: steering.message, - queuedMessageId: steering.id, + message: steering.pendingMessage.message, + queuedMessageId: steering.pendingMessage.id, }); // Mirror `resetTurnState` so per-turn counters/mappings (usage total, // streaming part ids) don't bleed from the preempted turn into the new @@ -1442,10 +1447,10 @@ export class CopilotAgentSession extends Disposable { // response: mark it `running` immediately rather than leaving it // `pending`, otherwise an abort during the steering turn would treat it // as a not-yet-started queued turn and leave it open. - this.resetTurnState(newTurnId); + this.resetTurnState(newTurnId, steering.sender?.clientId, steering.sender?.clientContext.clientType, steering.sender?.clientContext); const turn = this._currentTurn.value; if (turn) { - turn.messageCharLen = steering.message.text.length; + turn.messageCharLen = steering.pendingMessage.message.text.length; turn.markRunning(); } if (this._activeRootSdkTurnId) { @@ -1484,20 +1489,20 @@ export class CopilotAgentSession extends Disposable { * no buffered entry matches; the caller treats the `user.message` as * an ordinary echo and skips the turn flip. */ - private _takeMatchingPendingSteering(content: string): PendingMessage | undefined { + private _takeMatchingPendingSteering(content: string): IPendingSteering | undefined { if (this._pendingSteeringFlips.size === 0) { return undefined; } - let substringMatch: [string, PendingMessage] | undefined; - for (const [id, msg] of this._pendingSteeringFlips) { - if (msg.message.text === content) { + let substringMatch: [string, IPendingSteering] | undefined; + for (const [id, pending] of this._pendingSteeringFlips) { + if (pending.pendingMessage.message.text === content) { this._pendingSteeringFlips.delete(id); - return msg; + return pending; } - if (msg.message.text.length > 0 - && content.includes(msg.message.text) - && (!substringMatch || msg.message.text.length > substringMatch[1].message.text.length)) { - substringMatch = [id, msg]; + if (pending.pendingMessage.message.text.length > 0 + && content.includes(pending.pendingMessage.message.text) + && (!substringMatch || pending.pendingMessage.message.text.length > substringMatch[1].pendingMessage.message.text.length)) { + substringMatch = [id, pending]; } } if (substringMatch) { @@ -3434,7 +3439,7 @@ export class CopilotAgentSession extends Disposable { return this._configurationService.getRootValue(platformRootSchema, AgentHostAutoReplyEnabledConfigKey) === true; } - async sendSteering(steeringMessage: PendingMessage): Promise { + async sendSteering(steeringMessage: PendingMessage, sender?: IAgentPendingMessageSender): Promise { if (this._steeringMessagesInFlight.has(steeringMessage.id) || this._pendingSteeringFlips.has(steeringMessage.id)) { return; } @@ -3442,7 +3447,7 @@ export class CopilotAgentSession extends Disposable { this._logService.info(`[Copilot:${this.sessionId}] Sending steering message: "${steeringMessage.message.text.substring(0, 100)}"`); try { await this._reconcileMcpServerEnablement(); - this._pendingSteeringFlips.set(steeringMessage.id, steeringMessage); + this._pendingSteeringFlips.set(steeringMessage.id, { pendingMessage: steeringMessage, sender }); const sdkAttachments = await this._toSdkAttachments(steeringMessage.message.attachments); // Steering is injected into the active turn and never fires the SDK's `user-prompt-submitted` // hook, so the read-only snapshot signal can't ride `additionalContext` here. Fold it into the diff --git a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts index a70a389b9b8c6..383c1e5bfec4b 100644 --- a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts +++ b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts @@ -3445,13 +3445,22 @@ suite('AgentSideEffects', () => { message: { text: 'focus on tests', origin: { kind: MessageKind.User } }, }; stateManager.dispatchClientAction(defaultChatUri, action, { clientId: 'test', clientSeq: 1 }); - sideEffects.handleAction(defaultChatUri, action); + sideEffects.handleAction(defaultChatUri, action, 'client-editor', AgentHostClientType.EditorWindow); assert.strictEqual(agent.setPendingMessagesCalls.length, 1); - assert.deepStrictEqual(agent.setPendingMessagesCalls[0].steeringMessage, { id: 'steer-1', message: { text: 'focus on tests', origin: { kind: MessageKind.User } } }); - assert.deepStrictEqual(agent.setPendingMessagesCalls[0].queuedMessages, []); - // Steering is always addressed by a concrete chat channel URI. - assert.strictEqual(agent.setPendingMessagesCalls[0].chat.toString(), defaultChatUri); + assert.deepStrictEqual({ + chat: agent.setPendingMessagesCalls[0].chat.toString(), + steeringMessage: agent.setPendingMessagesCalls[0].steeringMessage, + queuedMessages: agent.setPendingMessagesCalls[0].queuedMessages, + senderClientId: agent.setPendingMessagesCalls[0].steeringSender?.clientId, + senderClientType: agent.setPendingMessagesCalls[0].steeringSender?.clientContext.clientType, + }, { + chat: defaultChatUri, + steeringMessage: { id: 'steer-1', message: { text: 'focus on tests', origin: { kind: MessageKind.User } } }, + queuedMessages: [], + senderClientId: 'client-editor', + senderClientType: AgentHostClientType.EditorWindow, + }); }); test('syncs a peer chat steering message addressed by the peer chat URI', () => { @@ -4120,16 +4129,40 @@ suite('AgentSideEffects', () => { }); }); - test('removes the active client when it is removed', () => { + test('removes the active client from the provider after server disconnect cleanup', () => { setupSession(); const peerChatUri = URI.parse(buildChatUri(sessionUri, 'peer-removal')); stateManager.addChat(sessionUri.toString(), peerChatUri.toString()); + const activeClientSet: SessionAction = { + type: ActionType.SessionActiveClientSet, + activeClient: { clientId: 'test-client', tools: [] }, + }; + stateManager.dispatchClientAction(sessionUri.toString(), activeClientSet, { clientId: 'test-client', clientSeq: 1 }); + sideEffects.handleAction(sessionUri.toString(), activeClientSet); - const action: SessionAction = { + stateManager.dispatchServerAction(sessionUri.toString(), { type: ActionType.SessionActiveClientRemoved, clientId: 'test-client', - }; - sideEffects.handleAction(sessionUri.toString(), action); + }); + + assert.deepStrictEqual(agent.removeActiveClientCalls.map(call => ({ + chat: call.chat.toString(), + clientId: call.clientId, + })), [ + { chat: defaultChatUri, clientId: 'test-client' }, + { chat: peerChatUri.toString(), clientId: 'test-client' }, + ]); + }); + + test('removes the active client from the provider after a client-dispatched removal', () => { + setupSession(); + const peerChatUri = URI.parse(buildChatUri(sessionUri, 'peer-removal')); + stateManager.addChat(sessionUri.toString(), peerChatUri.toString()); + + sideEffects.handleAction(sessionUri.toString(), { + type: ActionType.SessionActiveClientRemoved, + clientId: 'test-client', + }); assert.deepStrictEqual(agent.removeActiveClientCalls.map(call => ({ chat: call.chat.toString(), diff --git a/src/vs/platform/agentHost/test/node/chatContributions.test.ts b/src/vs/platform/agentHost/test/node/chatContributions.test.ts index 11cd72788de37..49c5102af4a8c 100644 --- a/src/vs/platform/agentHost/test/node/chatContributions.test.ts +++ b/src/vs/platform/agentHost/test/node/chatContributions.test.ts @@ -914,7 +914,11 @@ function createQueueDrainContributions(disposables: ReturnType pendingMessages.push(steeringMessage); + const pendingMessageSenders: (string | undefined)[] = []; + mockAgent.setPendingMessages = (_chat, steeringMessage, _queuedMessages, steeringSender) => { + pendingMessages.push(steeringMessage); + pendingMessageSenders.push(steeringSender?.clientId); + }; services.set(IAgentHostProviderService, createTestAgentHostProviderService(() => agent)); services.set(IAgentHostLocalTurns, new AgentHostLocalTurns(sessionDataService, logService)); const instantiationService = disposables.add(new InstantiationService(services, /*strict*/ true)); @@ -938,7 +942,7 @@ function createQueueDrainContributions(disposables: ReturnType & { readonly id: string })); disposables.add(service.registerContribution(SessionWorkspaceConversionContribution as unknown as IConstructorSignature & { readonly id: string })); disposables.add(service.registerContribution(QueueDrainContribution as unknown as IConstructorSignature & { readonly id: string })); - return { service, stateManager, session, chat, pendingMessages, admitted, titleController, telemetryService, clearAgent: () => agent = undefined, setConversionPending: (pending: boolean) => conversionPending = pending }; + return { service, stateManager, session, chat, pendingMessages, pendingMessageSenders, admitted, titleController, telemetryService, clearAgent: () => agent = undefined, setConversionPending: (pending: boolean) => conversionPending = pending }; } function appliedClientAction(channel: string, session: string, action: IAppliedClientAction['action'], clientId = 'client'): IAppliedClientAction { @@ -1193,6 +1197,52 @@ suite('AgentHostChatContributions', () => { }); }); + test('queue drain keeps steering sender ownership synchronized with replacement and removal', () => { + const queue = createQueueDrainContributions(disposables); + const first: IAppliedClientAction['action'] = { + type: ActionType.ChatPendingMessageSet, + kind: PendingMessageKind.Steering, + id: 'first', + message: { text: 'first', origin: { kind: MessageKind.User } }, + }; + const second: IAppliedClientAction['action'] = { + type: ActionType.ChatPendingMessageSet, + kind: PendingMessageKind.Steering, + id: 'second', + message: { text: 'second', origin: { kind: MessageKind.User } }, + }; + const removeFirst: IAppliedClientAction['action'] = { + type: ActionType.ChatPendingMessageRemoved, + kind: PendingMessageKind.Steering, + id: 'first', + }; + const removeSecond: IAppliedClientAction['action'] = { + type: ActionType.ChatPendingMessageRemoved, + kind: PendingMessageKind.Steering, + id: 'second', + }; + + queue.stateManager.dispatchServerAction(queue.chat, first); + queue.service.didApplyClientAction(appliedClientAction(queue.chat, queue.session, first, 'first-client')); + queue.stateManager.dispatchServerAction(queue.chat, second); + queue.service.didApplyClientAction(appliedClientAction(queue.chat, queue.session, second, 'second-client')); + queue.stateManager.dispatchServerAction(queue.chat, removeFirst); + queue.service.didApplyClientAction(appliedClientAction(queue.chat, queue.session, removeFirst, 'first-client')); + queue.service.didApplyClientAction(appliedClientAction(queue.chat, queue.session, { type: ActionType.ChatQueuedMessagesReordered, order: [] }, 'other-client')); + queue.stateManager.dispatchServerAction(queue.chat, removeSecond); + queue.service.didApplyClientAction(appliedClientAction(queue.chat, queue.session, removeSecond, 'second-client')); + queue.stateManager.dispatchServerAction(queue.chat, second); + queue.service.didApplyClientAction(appliedClientAction(queue.chat, queue.session, { type: ActionType.ChatQueuedMessagesReordered, order: [] }, 'other-client')); + + assert.deepStrictEqual({ + messages: queue.pendingMessages.map(message => message?.id), + senders: queue.pendingMessageSenders, + }, { + messages: ['first', 'second', 'second', 'second', undefined, 'second'], + senders: ['first-client', 'second-client', 'second-client', 'second-client', undefined, undefined], + }); + }); + test('queue drain defers stale queued actions until a resumable turn completes', () => { const queue = createQueueDrainContributions(disposables); queue.stateManager.dispatchServerAction(queue.chat, { diff --git a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts index c6b60e4fa37be..e5dc095a44dbc 100644 --- a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts @@ -30,7 +30,7 @@ import { NullTelemetryServiceShape } from '../../../telemetry/common/telemetryUt import { getTelemetryChatSessionId } from '../../common/agentTelemetryCorrelation.js'; import { AgentSession, type AgentSignal, type IAgentActionSignal, type IAgentToolPendingConfirmationSignal } from '../../common/agent.js'; import { AgentHostClientType } from '../../common/agentHostClientInfo.js'; -import { AgentHostClientConnectionKind, AgentHostLaunchKind, AgentHostTransportKind } from '../../common/agentHostTelemetry.js'; +import { AgentHostClientConnectionKind, AgentHostLaunchKind, AgentHostTransportKind, createUnknownAgentHostClientTelemetryContext } from '../../common/agentHostTelemetry.js'; import type { ChatInputRequestWithPlanReview } from '../../common/agentHostPlanReview.js'; import { AgentFeedbackAttachmentDisplayKind } from '../../common/meta/agentFeedbackAttachments.js'; import { ChatInputRequestPurpose, readChatInputRequestPurpose } from '../../common/meta/agentChatInputRequestMeta.js'; @@ -12634,6 +12634,37 @@ Use the attached image as context. ]); }); + test('client tool start after steering prefers the steering sender', async () => { + const activeClientToolSet = new ActiveClientToolSet(); + activeClientToolSet.set('original-client', snapshot.tools); + activeClientToolSet.set('steering-client', snapshot.tools); + const { session, mockSession, signals } = await createAgentSession(disposables, { clientSnapshot: snapshot, activeClientToolSet }); + session.resetTurnState('turn-original', 'original-client'); + + await session.sendSteering( + { id: 'steer-1', message: { text: 'focus on tests', origin: { kind: MessageKind.User } } }, + { + clientId: 'steering-client', + clientContext: createUnknownAgentHostClientTelemetryContext(AgentHostClientType.EditorWindow), + }, + ); + mockSession.fire('user.message', { + content: 'focus on tests', + interactionId: 'interaction-steer', + } as SessionEventPayload<'user.message'>['data']); + mockSession.fire('tool.execution_start', { + toolCallId: 'tc-steering', + toolName: 'my_tool', + arguments: {}, + } as SessionEventPayload<'tool.execution_start'>['data']); + + const start = signals.find((signal): signal is IAgentActionSignal => isAction(signal, ActionType.ChatToolCallStart)); + assert.deepStrictEqual(start && (start.action as ChatToolCallStartAction).contributor, { + kind: ToolCallContributorKind.Client, + clientId: 'steering-client', + }); + }); + test('completion arriving before the SDK handler registers still resolves', async () => { const { session, runtime } = await createAgentSession(disposables, { clientSnapshot: snapshot }); const tools = runtime.createClientSdkTools(); diff --git a/src/vs/platform/agentHost/test/node/e2e/captures/copilotcli-runtime-tools-removing-a-client-transfers-duplicate-tool-ownership-to-the-surviving-client.yaml b/src/vs/platform/agentHost/test/node/e2e/captures/copilotcli-runtime-tools-removing-a-client-transfers-duplicate-tool-ownership-to-the-surviving-client.yaml new file mode 100644 index 0000000000000..61a0440e885be --- /dev/null +++ b/src/vs/platform/agentHost/test/node/e2e/captures/copilotcli-runtime-tools-removing-a-client-transfers-duplicate-tool-ownership-to-the-surviving-client.yaml @@ -0,0 +1,35 @@ +version: 1 +dialect: responses +exchanges: + - request: + model: gpt-5.6-sol + system: ${system} + messages: + - role: user + content: Call route_probe exactly once, then reply with only its exact result. + response: + content: + - type: tool_use + id: toolcall_0 + name: route_probe + input: {} + stopReason: tool_use + - request: + model: gpt-5.6-sol + system: ${system} + messages: + - role: user + content: Call route_probe exactly once, then reply with only its exact result. + - role: assistant + content: + - type: tool_use + name: route_probe + input: {} + - role: user + content: + - type: tool_result + tool_use_id: toolcall_0 + content: SURVIVING_CLIENT_RESULT + response: + content: SURVIVING_CLIENT_RESULT + stopReason: end_turn diff --git a/src/vs/platform/agentHost/test/node/e2e/suites/copilotRuntimeToolsSuite.ts b/src/vs/platform/agentHost/test/node/e2e/suites/copilotRuntimeToolsSuite.ts index c6050b831cbde..275b942db52e1 100644 --- a/src/vs/platform/agentHost/test/node/e2e/suites/copilotRuntimeToolsSuite.ts +++ b/src/vs/platform/agentHost/test/node/e2e/suites/copilotRuntimeToolsSuite.ts @@ -11,9 +11,10 @@ import { URI } from '../../../../../../base/common/uri.js'; import { CopilotCliConfigKey } from '../../../../common/copilotCliConfig.js'; import type { ResourceReadResult, SubscribeResult } from '../../../../common/state/protocol/commands.js'; import { ContentEncoding } from '../../../../common/state/protocol/common/commands.js'; +import { PROTOCOL_VERSION } from '../../../../common/state/protocol/version/registry.js'; import { ActionType, type ChatErrorAction, type ChatToolCallCompleteAction, type ChatToolCallReadyAction, type ChatToolCallStartAction } from '../../../../common/state/sessionActions.js'; import { buildDefaultChatUri, getErrorResponsePart, getInlineToolInput, ROOT_STATE_URI, ToolCallConfirmationReason, ToolCallContributorKind, ToolResultContentType, TurnState, type ToolDefinition } from '../../../../common/state/sessionState.js'; -import { fetchSessionWithChat, getActionEnvelope, isActionNotification } from '../../serverIntegrationTestHelpers.js'; +import { fetchSessionWithChat, getActionEnvelope, isActionNotification, type TestProtocolClient } from '../../serverIntegrationTestHelpers.js'; import { createRealSession, dispatchTurn, driveTurnToCompletion, driveTurnWithModelToCompletion, textFromContent } from '../harness/agentHostE2ETestHarness.js'; import { anthropicMessageToSse } from '../harness/capiWireCodec.js'; import type { IAgentHostE2ETestContext } from './e2eTestContext.js'; @@ -33,6 +34,16 @@ export function defineCopilotRuntimeToolsTests(context: IAgentHostE2ETestContext return { sessionUri, workspace }; } + async function initializeAdditionalClient(clientId: string): Promise { + const client = await context.connectClient(); + await client.call('initialize', { + channel: ROOT_STATE_URI, + protocolVersions: [PROTOCOL_VERSION], + clientId, + }); + return client; + } + test('runtime tools: compacted shell output preserves the complete original', async function () { this.timeout(180_000); const { sessionUri, workspace } = await createSession('shell-compaction'); @@ -214,6 +225,105 @@ export function defineCopilotRuntimeToolsTests(context: IAgentHostE2ETestContext } }); + test('runtime tools: removing a client transfers duplicate tool ownership to the surviving client', async function () { + this.timeout(180_000); + const removedClientId = 'runtime-tool-owner-removed'; + const survivingClientId = 'runtime-tool-owner-surviving'; + let removedClient: TestProtocolClient | undefined; + let survivingClient: TestProtocolClient | undefined; + try { + removedClient = await initializeAdditionalClient(removedClientId); + survivingClient = await initializeAdditionalClient(survivingClientId); + const { sessionUri } = await createSession('runtime-tool-owner-cleanup'); + const chatUri = buildDefaultChatUri(sessionUri); + const tool: ToolDefinition = { + name: 'route_probe', + description: 'Returns the client tool owner marker.', + inputSchema: { type: 'object', properties: {} }, + }; + for (const [client, clientId] of [[removedClient, removedClientId], [survivingClient, survivingClientId]] as const) { + await client.call('subscribe', { channel: sessionUri }); + await client.call('subscribe', { channel: chatUri }); + client.dispatch({ + channel: sessionUri, + clientSeq: 1, + action: { + type: ActionType.SessionActiveClientSet, + activeClient: { clientId, tools: [tool] }, + }, + }); + await context.client.waitForNotification(n => { + if (!isActionNotification(n, ActionType.SessionActiveClientSet)) { + return false; + } + const action = getActionEnvelope(n).action as { readonly activeClient: { readonly clientId: string } }; + return action.activeClient.clientId === clientId; + }, 30_000); + } + + context.client.clearReceived(); + removedClient.notify('unsubscribe', { channel: sessionUri }); + await removedClient.call('ping', { channel: ROOT_STATE_URI }); + await context.client.waitForNotification(n => { + if (!isActionNotification(n, ActionType.SessionActiveClientRemoved)) { + return false; + } + const action = getActionEnvelope(n).action as { readonly clientId: string }; + return action.clientId === removedClientId; + }, 30_000); + + const [result, contributor] = await Promise.all([ + driveTurnWithModelToCompletion( + context.client, + sessionUri, + 'turn-runtime-tool-owner-cleanup', + 'Call route_probe exactly once, then reply with only its exact result.', + 'gpt-5.6-sol', + 1, + ), + (async () => { + const start = await context.client.waitForNotification(n => + isActionNotification(n, ActionType.ChatToolCallStart) + && (getActionEnvelope(n).action as ChatToolCallStartAction).toolName === tool.name, + 90_000, + ); + const startAction = getActionEnvelope(start).action as ChatToolCallStartAction; + await context.client.waitForNotification(n => + isActionNotification(n, ActionType.ChatToolCallReady) + && (getActionEnvelope(n).action as ChatToolCallReadyAction).toolCallId === startAction.toolCallId, + 90_000, + ); + survivingClient.dispatch({ + channel: chatUri, + clientSeq: 2, + action: { + type: ActionType.ChatToolCallComplete, + turnId: startAction.turnId, + toolCallId: startAction.toolCallId, + result: { + success: true, + pastTenseMessage: 'Returned the owner marker', + content: [{ type: ToolResultContentType.Text, text: 'SURVIVING_CLIENT_RESULT' }], + }, + }, + }); + return startAction.contributor; + })(), + ]); + + assert.deepStrictEqual({ + contributor, + response: result.responseText.trim(), + }, { + contributor: { kind: ToolCallContributorKind.Client, clientId: survivingClientId }, + response: 'SURVIVING_CLIENT_RESULT', + }); + } finally { + removedClient?.close(); + survivingClient?.close(); + } + }); + test('runtime tools: image client tool results preserve event delivery through turn completion', async function () { this.timeout(180_000); const clientId = 'runtime-image-tool'; diff --git a/src/vs/platform/agentHost/test/node/e2e/suites/mcpPluginSuite.ts b/src/vs/platform/agentHost/test/node/e2e/suites/mcpPluginSuite.ts index e600bf08455c0..0713fb0471aef 100644 --- a/src/vs/platform/agentHost/test/node/e2e/suites/mcpPluginSuite.ts +++ b/src/vs/platform/agentHost/test/node/e2e/suites/mcpPluginSuite.ts @@ -16,7 +16,7 @@ import { CustomizationEnablementKind, McpServerStatus } from '../../../../common import { ActionType, type ChatToolCallCompleteAction } from '../../../../common/state/sessionActions.js'; import { buildDefaultChatUri, ChatInputAnswerState, ChatInputAnswerValueKind, customizationId, CustomizationType, ResponsePartKind, ROOT_STATE_URI, type ChatInputAnswer, type ChatInputRequest, type ClientPluginCustomization, type McpServerCustomization, type PluginCustomization, type SessionState } from '../../../../common/state/sessionState.js'; import { createRealSession, driveTurnToCompletion, driveTurnWithAnswersToCompletion, driveTurnWithCancelledInputToCompletion, resolveGitHubToken, textFromContent } from '../harness/agentHostE2ETestHarness.js'; -import { fetchSessionWithChat, getActionEnvelope, isActionNotification } from '../../serverIntegrationTestHelpers.js'; +import { fetchSessionWithChat, getActionEnvelope, isActionNotification, type TestProtocolClient } from '../../serverIntegrationTestHelpers.js'; import { providerHostOnlyTest, type IAgentHostE2ETestContext } from './e2eTestContext.js'; const nodeRequire = createRequire(import.meta.url); @@ -34,6 +34,8 @@ interface IPluginSessionOptions { readonly hookExitCode?: number; readonly hookStdout?: string; readonly pluginName?: string; + readonly publisher?: TestProtocolClient; + readonly clientId?: string; } export function defineMcpPluginTests(context: IAgentHostE2ETestContext): void { @@ -172,8 +174,13 @@ export function defineMcpPluginTests(context: IAgentHostE2ETestContext): void { }, })); const pluginUri = URI.file(plugin).toString(); - const clientId = `mcp-plugin-${prefix}-${config.provider}`; - const sessionUri = await createRealSession(context.client, config, clientId, createdSessions, URI.file(workspace)); + const clientId = options.clientId ?? `mcp-plugin-${prefix}-${config.provider}`; + const sessionCreatorClientId = options.publisher ? `mcp-plugin-observer-${prefix}-${config.provider}` : clientId; + const sessionUri = await createRealSession(context.client, config, sessionCreatorClientId, createdSessions, URI.file(workspace)); + const publisher = options.publisher ?? context.client; + if (publisher !== context.client) { + await publisher.call('subscribe', { channel: sessionUri }); + } const customization: ClientPluginCustomization = { type: CustomizationType.Plugin, id: customizationId(pluginUri), @@ -182,7 +189,7 @@ export function defineMcpPluginTests(context: IAgentHostE2ETestContext): void { nonce: '1', enablement: [{ kind: CustomizationEnablementKind.Global, enabled: true }], }; - context.client.dispatch({ + publisher.dispatch({ channel: sessionUri, clientSeq: 1, action: { @@ -320,6 +327,44 @@ export function defineMcpPluginTests(context: IAgentHostE2ETestContext): void { }, 100, 100); }); + providerHostOnlyTest(context, 'unsubscribing an active client removes its provider customization', async function () { + const clientId = `mcp-plugin-unsubscribe-${config.provider}`; + const publisher = await context.connectClient(); + try { + await publisher.call('initialize', { + channel: ROOT_STATE_URI, + protocolVersions: [PROTOCOL_VERSION], + clientId, + }); + const { sessionUri, pluginUri } = await createPluginSession('unsubscribe', { publisher, clientId }); + const plugin = await pluginState(sessionUri, pluginUri); + context.client.clearReceived(); + + publisher.notify('unsubscribe', { channel: sessionUri }); + await publisher.call('ping', { channel: ROOT_STATE_URI }); + await context.client.waitForNotification(n => { + if (!isActionNotification(n, ActionType.SessionActiveClientRemoved)) { + return false; + } + const envelope = getActionEnvelope(n); + return envelope.channel === sessionUri + && (envelope.action as { readonly clientId: string }).clientId === clientId; + }, 30_000); + await retry(async () => { + const result = await context.client.call('subscribe', { channel: sessionUri }); + const state = result.snapshot!.state as SessionState; + if (state.activeClients.some(client => client.clientId === clientId)) { + throw new Error('Active client has not been removed'); + } + if (state.customizations?.some(customization => customization.id === plugin.id)) { + throw new Error('Plugin customization has not been removed'); + } + }, 100, 100); + } finally { + publisher.close(); + } + }); + const modelBackedEnabled = config.provider === 'copilotcli'; if (modelBackedEnabled) { // The SDK-owned runtime does not invoke hook callbacks on Windows. diff --git a/src/vs/platform/agentHost/test/node/mockAgent.ts b/src/vs/platform/agentHost/test/node/mockAgent.ts index e7f189b5c5255..3ace6e51d8caf 100644 --- a/src/vs/platform/agentHost/test/node/mockAgent.ts +++ b/src/vs/platform/agentHost/test/node/mockAgent.ts @@ -11,7 +11,7 @@ import { join } from '../../../../base/common/path.js'; import { URI } from '../../../../base/common/uri.js'; import { AgentHostClientType } from '../../common/agentHostClientInfo.js'; import { type ISyncedCustomization } from '../../common/agentPluginManager.js'; -import { AgentSession, type AgentChatMigrationResult, type AgentProvider, type AgentSignal, type IActiveClient, type IAgent, type IAgentActionSignal, type IAgentCapabilities, type IAgentChatConfigCompletionsParams, type IAgentChatContext, type IAgentChatMetadata, type IAgentChats, type IAgentCreateChatOptions, type IAgentCreateChatResult, type IAgentCreateSessionConfig, type IAgentDescriptor, type IAgentDiscoveredChat, type IAgentModelInfo, type IAgentResolveChatConfigParams, type IAgentSessionMetadata, type IAgentToolPendingConfirmationSignal, resolveAgentChatContext } from '../../common/agent.js'; +import { AgentSession, type AgentChatMigrationResult, type AgentProvider, type AgentSignal, type IActiveClient, type IAgent, type IAgentActionSignal, type IAgentCapabilities, type IAgentChatConfigCompletionsParams, type IAgentChatContext, type IAgentChatMetadata, type IAgentChats, type IAgentCreateChatOptions, type IAgentCreateChatResult, type IAgentCreateSessionConfig, type IAgentDescriptor, type IAgentDiscoveredChat, type IAgentModelInfo, type IAgentPendingMessageSender, type IAgentResolveChatConfigParams, type IAgentSessionMetadata, type IAgentToolPendingConfirmationSignal, resolveAgentChatContext } from '../../common/agent.js'; import { buildSubagentTurnsFromHistory, buildTurnsFromHistory, type IHistoryRecord } from './historyRecordFixtures.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'; @@ -78,7 +78,7 @@ export class MockAgent implements IAgent { readonly sendMessageCalls: IMockSendMessageCall[] = []; - readonly setPendingMessagesCalls: { chat: URI; steeringMessage: PendingMessage | undefined; queuedMessages: readonly PendingMessage[] }[] = []; + readonly setPendingMessagesCalls: { chat: URI; steeringMessage: PendingMessage | undefined; queuedMessages: readonly PendingMessage[]; steeringSender: IAgentPendingMessageSender | undefined }[] = []; readonly disposeSessionCalls: URI[] = []; readonly releaseSessionCalls: URI[] = []; readonly abortSessionCalls: URI[] = []; @@ -250,8 +250,8 @@ export class MockAgent implements IAgent { } } - setPendingMessages(chat: URI, steeringMessage: PendingMessage | undefined, queuedMessages: readonly PendingMessage[]): void { - this.setPendingMessagesCalls.push({ chat, steeringMessage, queuedMessages }); + setPendingMessages(chat: URI, steeringMessage: PendingMessage | undefined, queuedMessages: readonly PendingMessage[], steeringSender?: IAgentPendingMessageSender): void { + this.setPendingMessagesCalls.push({ chat, steeringMessage, queuedMessages, steeringSender }); } async getSessionMessages(session: URI): Promise { diff --git a/src/vs/platform/agentHost/test/node/providerIntegration/copilotMockLlm.integrationTest.ts b/src/vs/platform/agentHost/test/node/providerIntegration/copilotMockLlm.integrationTest.ts index 62ecbc201ada9..4d93bafae24cb 100644 --- a/src/vs/platform/agentHost/test/node/providerIntegration/copilotMockLlm.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/providerIntegration/copilotMockLlm.integrationTest.ts @@ -15,8 +15,9 @@ import { timeout } from '../../../../../base/common/async.js'; import { join } from '../../../../../base/common/path.js'; import { isWindows } from '../../../../../base/common/platform.js'; import { URI } from '../../../../../base/common/uri.js'; -import { ActionType, type ChatToolCallCompleteAction, type ChatToolCallReadyAction } from '../../../common/state/sessionActions.js'; -import { buildDefaultChatUri, ResponsePartKind, SessionStatus, type ISessionWithDefaultChat } from '../../../common/state/sessionState.js'; +import { ActionType, type ChatResponsePartAction, type ChatToolCallCompleteAction, type ChatToolCallReadyAction, type ChatToolCallStartAction, type ChatTurnStartedAction } from '../../../common/state/sessionActions.js'; +import { PROTOCOL_VERSION } from '../../../common/state/protocol/version/registry.js'; +import { buildDefaultChatUri, MessageKind, PendingMessageKind, ResponsePartKind, ROOT_STATE_URI, SessionStatus, ToolCallContributorKind, ToolResultContentType, type ISessionWithDefaultChat, type ToolDefinition } from '../../../common/state/sessionState.js'; import { ToolCallConfirmationReason } from '../../../common/state/protocol/channels-chat/state.js'; import { AgentHostSessionReleaseRetryMsEnvVar, AgentHostSessionResidencyLimitEnvVar } from '../../../common/agentService.js'; import { createProviderSession, dispatchTurn, type IAgentHostProviderTestConfig } from '../providerIntegrationTestHelpers.js'; @@ -30,6 +31,8 @@ const COPILOT_CONFIG: IAgentHostProviderTestConfig = { const DETACHED_SHELL_SCENARIO_ID = 'detached-shell-idle-release'; const DETACHED_SHELL_DELAY_MS = 6000; +const STEERING_OWNER_SCENARIO_ID = 'steering-client-tool-owner'; +const STEERING_RESPONSE_DELAY_MS = 10_000; function quoteShellArgument(value: string): string { return isWindows ? `'${value.replace(/'/g, '\'\'')}'` : `'${value.replace(/'/g, `'\\''`)}'`; @@ -50,6 +53,29 @@ suite('Agent Host Provider Integration — Copilot with Mock LLM', function () { mockLlm: true, homeDir: suiteHome, userDataDir: join(suiteHome, 'user-data'), + mockScenarios: [{ + id: STEERING_OWNER_SCENARIO_ID, + definition: { + type: 'multi-turn', + turns: [ + { + kind: 'content', + chunks: [ + { content: 'Initial response started.', delayMs: 0 }, + { content: ' Initial response finished.', delayMs: STEERING_RESPONSE_DELAY_MS }, + ], + }, + { + kind: 'tool-calls', + toolCalls: [{ + toolNamePattern: /^route_probe$/, + arguments: {}, + }], + }, + { kind: 'content', chunks: [{ content: 'STEERING_CLIENT_RESULT', delayMs: 0 }] }, + ], + }, + }], }); }); @@ -105,6 +131,153 @@ suite('Agent Host Provider Integration — Copilot with Mock LLM', function () { assert.ok(markdownText.trim().length > 0, `expected non-empty assistant markdown; got: ${JSON.stringify(markdownText)}`); assert.match(markdownText, new RegExp(`\\b${probeToken}\\b`, 'i'), `expected probe token in assistant markdown; got: ${JSON.stringify(markdownText)}`); }); + + test('routes a client tool after steering to the client that sent the steering message', async function () { + this.timeout(180_000); + const originalClientId = 'steering-original-client'; + const steeringClientId = 'steering-sender-client'; + const workspaceDir = await mkdtemp(`${tmpdir()}/test-mock-steering-owner`); + tempDirs.push(workspaceDir); + const sessionUri = await createProviderSession(client, COPILOT_CONFIG, originalClientId, createdSessions, URI.file(workspaceDir)); + const chatUri = buildDefaultChatUri(sessionUri); + const steeringClient = new TestProtocolClient(server.port); + const tools: ToolDefinition[] = [{ + name: 'route_probe', + description: 'Returns the steering owner marker.', + inputSchema: { type: 'object', properties: {} }, + }]; + await steeringClient.connect(); + try { + await steeringClient.call('initialize', { + channel: ROOT_STATE_URI, + protocolVersions: [PROTOCOL_VERSION], + clientId: steeringClientId, + }); + await steeringClient.call('subscribe', { channel: sessionUri }); + await steeringClient.call('subscribe', { channel: chatUri }); + + for (const [owner, clientId] of [[client, originalClientId], [steeringClient, steeringClientId]] as const) { + owner.dispatch({ + channel: sessionUri, + clientSeq: 1, + action: { + type: ActionType.SessionActiveClientSet, + activeClient: { clientId, tools }, + }, + }); + await client.waitForNotification(n => { + if (!isActionNotification(n, ActionType.SessionActiveClientSet)) { + return false; + } + const action = getActionEnvelope(n).action as { readonly activeClient: { readonly clientId: string } }; + return action.activeClient.clientId === clientId; + }, 30_000); + } + + dispatchTurn(client, sessionUri, 'turn-before-steering', `[scenario:${STEERING_OWNER_SCENARIO_ID}] Start the ownership test.`, 2); + await client.waitForNotification(n => + isActionNotification(n, ActionType.ChatResponsePart) + && (getActionEnvelope(n).action as ChatResponsePartAction).turnId === 'turn-before-steering' + && (getActionEnvelope(n).action as ChatResponsePartAction).part.kind === ResponsePartKind.Markdown, + 90_000, + ); + assert.strictEqual(client.receivedNotifications(n => + isActionNotification(n, ActionType.ChatTurnComplete) + && (getActionEnvelope(n).action as { readonly turnId: string }).turnId === 'turn-before-steering', + ).length, 0, 'steering must be submitted while the original turn is active'); + + steeringClient.dispatch({ + channel: chatUri, + clientSeq: 2, + action: { + type: ActionType.ChatPendingMessageSet, + kind: PendingMessageKind.Steering, + id: 'steering-owner-message', + message: { + text: 'Now call route_probe exactly once and reply with only its exact result.', + origin: { kind: MessageKind.User }, + }, + }, + }); + await client.waitForNotification(n => + isActionNotification(n, ActionType.ChatPendingMessageSet) + && (getActionEnvelope(n).action as { readonly id: string }).id === 'steering-owner-message', + 30_000, + ); + const steeringTurnNotification = await client.waitForNotification(n => + isActionNotification(n, ActionType.ChatTurnStarted) + && (getActionEnvelope(n).action as ChatTurnStartedAction).queuedMessageId === 'steering-owner-message', + 90_000, + ); + const steeringTurn = getActionEnvelope(steeringTurnNotification).action as ChatTurnStartedAction; + + const routeStartNotification = await client.waitForNotification(n => + isActionNotification(n, ActionType.ChatToolCallStart) + && (getActionEnvelope(n).action as ChatToolCallStartAction).toolName === 'route_probe', + 90_000, + ); + const routeStart = getActionEnvelope(routeStartNotification).action as ChatToolCallStartAction; + const routeReadyNotification = await client.waitForNotification(n => + isActionNotification(n, ActionType.ChatToolCallReady) + && (getActionEnvelope(n).action as ChatToolCallReadyAction).toolCallId === routeStart.toolCallId, + 90_000, + ); + const routeReady = getActionEnvelope(routeReadyNotification).action as ChatToolCallReadyAction; + let steeringClientSeq = 3; + if (!routeReady.confirmed) { + steeringClient.dispatch({ + channel: chatUri, + clientSeq: steeringClientSeq++, + action: { + type: ActionType.ChatToolCallConfirmed, + turnId: routeStart.turnId, + toolCallId: routeStart.toolCallId, + approved: true, + confirmed: ToolCallConfirmationReason.UserAction, + }, + }); + } + steeringClient.dispatch({ + channel: chatUri, + clientSeq: steeringClientSeq, + action: { + type: ActionType.ChatToolCallComplete, + turnId: routeStart.turnId, + toolCallId: routeStart.toolCallId, + result: { + success: true, + pastTenseMessage: 'Returned the steering owner marker', + content: [{ type: ToolResultContentType.Text, text: 'STEERING_CLIENT_RESULT' }], + }, + }, + }); + await client.waitForNotification(n => + isActionNotification(n, ActionType.ChatTurnComplete) + && (getActionEnvelope(n).action as { readonly turnId: string }).turnId === routeStart.turnId, + 90_000, + ); + + const state = await fetchSessionWithChat(client, sessionUri); + const completedSteeringTurn = state.turns.find(turn => turn.id === routeStart.turnId); + const response = completedSteeringTurn?.responseParts + .filter(part => part.kind === ResponsePartKind.Markdown) + .map(part => part.content) + .join('') ?? ''; + assert.deepStrictEqual({ + steeringContributor: routeStart.contributor, + steeringTurnId: steeringTurn.turnId, + routeTurnId: routeStart.turnId, + response: response.trim(), + }, { + steeringContributor: { kind: ToolCallContributorKind.Client, clientId: steeringClientId }, + steeringTurnId: routeStart.turnId, + routeTurnId: routeStart.turnId, + response: 'STEERING_CLIENT_RESULT', + }); + } finally { + steeringClient.close(); + } + }); }); /**