From 233fcc85b1a7712f7f4a583b39e46816ba227cde Mon Sep 17 00:00:00 2001 From: Karthik Nadig Date: Sun, 18 Oct 2026 00:05:52 -0700 Subject: [PATCH 1/2] Stop delivering steering messages the user removed from the queue Removing a pending steering message during an Agent Host turn was refused by the host, and nothing withdrew the message from the Copilot SDK, so it was still delivered as its own turn. Let providers that can withdraw steering accept the removal, and have Copilot withdraw it from the SDK steering queue (or skip sending it if it is still being prepared). A message the model already received is kept and still shown as its own turn. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/vs/platform/agentHost/AGENTS.md | 2 +- src/vs/platform/agentHost/common/agent.ts | 7 ++ .../platform/agentHost/node/agentService.ts | 3 +- .../queueDrain/queueDrainContribution.ts | 13 +++ .../agentHost/node/copilot/copilotAgent.ts | 4 + .../node/copilot/copilotAgentSession.ts | 54 ++++++++++-- .../agentHost/test/node/agentService.test.ts | 35 ++++++++ .../test/node/chatContributions.test.ts | 35 +++++++- .../test/node/copilotAgentSession.test.ts | 86 +++++++++++++++++++ .../platform/agentHost/test/node/mockAgent.ts | 3 + 10 files changed, 233 insertions(+), 9 deletions(-) diff --git a/src/vs/platform/agentHost/AGENTS.md b/src/vs/platform/agentHost/AGENTS.md index df9c5c5e70f008..ad3aff2f26a275 100644 --- a/src/vs/platform/agentHost/AGENTS.md +++ b/src/vs/platform/agentHost/AGENTS.md @@ -557,7 +557,7 @@ Customization discovery (agents, skills, instructions and workspace MCP configur Steering received during turn preparation or configuration refresh is buffered for the addressed chat. A dispatched turn confirmed running by root SDK events accepts and flushes steering immediately, even if its admission RPC has not returned; successful admission also flushes it. Preparation or send failure retains buffered steering for a successful retry, rather than dispatching against an unready runtime. Pending updates replace or remove buffered steering; abort clears it and emits `steering_consumed` to clear the pending UI. The live SDK-session wrapper takes ownership synchronously before MCP reconciliation, so abort and disposal can clear steering throughout delivery. Chat disposal and host shutdown discard buffered delivery, and pending-message updates are ignored once host shutdown starts. -AHP stores steering and queued messages together in `queuedMessages`, distinguished by `PendingMessage.kind`. Steering entries use `IAgent.sendSteeringMessage` individually, bypassing Copilot's replaceable legacy-slot buffer. The host retries deferred handoff on provider progress and correlates consumption by message ID. Cancellation clears every pending steering entry, including prompts handed off but not yet consumed by the SDK, without draining queued follow-ups. +AHP stores steering and queued messages together in `queuedMessages`, distinguished by `PendingMessage.kind`. Steering entries use `IAgent.sendSteeringMessage` individually, bypassing Copilot's replaceable legacy-slot buffer. The host retries deferred handoff on provider progress and correlates consumption by message ID. Cancellation clears every pending steering entry, including prompts handed off but not yet consumed by the SDK, without draining queued follow-ups. A client may remove handed-off steering during a turn only when the provider implements `IAgent.withdrawSteeringMessage`; Copilot then withdraws it from the SDK with `rpc.queue.withdrawMessage`, or keeps it as its own turn if the model already received it. `CopilotSessionLauncher` sets `mcpOAuthTokenStorage: 'in-memory'` for created, resumed, and ephemeral SDK sessions. VS Code owns durable MCP credentials through `onMcpAuthRequest`; the runtime must not consult its persistent MCP OAuth keychain store. diff --git a/src/vs/platform/agentHost/common/agent.ts b/src/vs/platform/agentHost/common/agent.ts index 5dacb4f93473ae..6b233e565a6418 100644 --- a/src/vs/platform/agentHost/common/agent.ts +++ b/src/vs/platform/agentHost/common/agent.ts @@ -1716,6 +1716,13 @@ export interface IAgent { /** Hands off one independent submission without replacing the legacy slot; false defers until ready. */ sendSteeringMessage?(chat: URI, message: PendingMessage, sender?: IAgentPendingMessageSender): boolean | void; + /** + * Withdraws a steering message, previously handed off through {@link sendSteeringMessage}, that a + * client removed before the model received it. A message the model already received stays in its + * turn. Without this hook, the host rejects removing steering while a turn is active. + */ + withdrawSteeringMessage?(chat: URI, messageId: string): 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/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts index 982fbff37aefc7..6e6219e149040b 100644 --- a/src/vs/platform/agentHost/node/agentService.ts +++ b/src/vs/platform/agentHost/node/agentService.ts @@ -11098,9 +11098,10 @@ export class AgentService extends Disposable implements IAgentService { this._stateManager.rejectClientAction(channel, action, origin, 'The session provider does not support steering messages.'); return; } + // Steering handed to the provider cannot be edited, and can be removed only if the provider can withdraw it. if (state?.activeTurn && existing?.kind === PendingMessageKind.Steering && (action.type === ActionType.ChatPendingMessageRemoved - ? action.kind === PendingMessageKind.Steering + ? action.kind === PendingMessageKind.Steering && !this._providerService.getProviderForSession(sessionChannel)?.withdrawSteeringMessage : action.kind !== existing.kind || !equals(existing.message, action.message))) { this._stateManager.rejectClientAction(channel, action, origin, 'Cannot change a steering message while it is being processed.'); return; diff --git a/src/vs/platform/agentHost/node/chatContributions/queueDrain/queueDrainContribution.ts b/src/vs/platform/agentHost/node/chatContributions/queueDrain/queueDrainContribution.ts index 184af8576c814d..70474065cf1590 100644 --- a/src/vs/platform/agentHost/node/chatContributions/queueDrain/queueDrainContribution.ts +++ b/src/vs/platform/agentHost/node/chatContributions/queueDrain/queueDrainContribution.ts @@ -80,6 +80,9 @@ export class QueueDrainContribution extends Disposable implements IAgentHostChat ...(observed.credentials ? { credentials: observed.credentials } : {}), }, undefined); } + if (action.type === ActionType.ChatPendingMessageRemoved && action.kind === PendingMessageKind.Steering) { + this._withdrawSubmittedSteering(observed.channel, action.id); + } if (this._handlePendingMessageAction(observed.channel, action)) { this._tryConsumeNextQueuedMessage(observed.channel); } @@ -158,6 +161,16 @@ export class QueueDrainContribution extends Disposable implements IAgentHostChat && getPendingSteeringMessages(state).some(message => message.id === action.id); } + /** A client removed steering that was already handed to the provider, which would otherwise still deliver it. */ + private _withdrawSubmittedSteering(channel: ProtocolURI, id: string): void { + if (!this._context.memento(SubmittedSteering, channel).get().has(id)) { + return; + } + const session = this._chatOwnership.getChatOwner(channel); + const provider = session !== undefined ? this._providerService.getProviderForSession(session) : undefined; + provider?.withdrawSteeringMessage?.(URI.parse(channel), id); + } + private _syncPendingMessages(channel: ProtocolURI): void { const state = this._stateManager.getSessionState(channel); if (!state) { diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts index 9633782b62ae20..e9059bce901eca 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts @@ -6894,6 +6894,10 @@ export class CopilotAgent extends Disposable implements IAgent { return true; } + withdrawSteeringMessage(chat: URI, messageId: string): void { + void this._findChatByUri(chat)?.withdrawSteering(messageId); + } + private async _getChatMessages(chat: URI, sessionOrContext: URI | IAgentChatContext): Promise { if (this._isShuttingDown) { return []; diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts index f51eae8dcb9d03..8bf239ab2629ac 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts @@ -1022,6 +1022,8 @@ class CopilotTurn extends Disposable { interface IPendingSteering { readonly pendingMessage: PendingMessage; readonly sender: IAgentPendingMessageSender | undefined; + /** Set once the message is handed to the SDK; identifies it for withdrawal. */ + sdkSend?: { readonly prompt: string; readonly messageId: Promise }; } /** Where a transient chat activity comes from; see `CopilotAgentSession._publishActivity` for the precedence. */ @@ -5590,8 +5592,9 @@ export class CopilotAgentSession extends Disposable { } const steeringTurn = this._currentTurn.value; const abortToken = this._abortToken; + const pending: IPendingSteering = { pendingMessage: steeringMessage, sender }; this._steeringMessagesInFlight.add(steeringMessage.id); - this._pendingSteeringFlips.set(steeringMessage.id, { pendingMessage: steeringMessage, sender }); + this._pendingSteeringFlips.set(steeringMessage.id, pending); this._logService.info(`[Copilot:${this.sessionId}] Sending steering message: "${steeringMessage.message.text.substring(0, 100)}"`); try { await this._reconcileMcpServerEnablement(); @@ -5606,11 +5609,19 @@ export class CopilotAgentSession extends Disposable { : steeringMessage.message.text; const traceContext = this._otelService.getSessionTraceContext(this.sessionId, this.resourceUri.toString()); const execution = await this._executeSdkOperation( - () => this._otelService.withTraceContext(traceContext, () => this._wrapper.session.send({ - prompt: steeringPrompt, - attachments: sdkAttachments?.length ? sdkAttachments : undefined, - mode: 'immediate', - })), + async () => { + if (this._pendingSteeringFlips.get(steeringMessage.id) !== pending) { + this._logService.info(`[Copilot:${this.sessionId}] Steering message was withdrawn or cleared before it was sent: id=${steeringMessage.id}`); + return undefined; + } + const messageId = this._otelService.withTraceContext(traceContext, () => this._wrapper.session.send({ + prompt: steeringPrompt, + attachments: sdkAttachments?.length ? sdkAttachments : undefined, + mode: 'immediate', + })); + pending.sdkSend = { prompt: steeringPrompt, messageId }; + return messageId; + }, steeringTurn, abortToken, ); if (execution.kind === 'skipped' && this._pendingSteeringFlips.delete(steeringMessage.id)) { @@ -5626,6 +5637,37 @@ export class CopilotAgentSession extends Disposable { } } + /** + * Withdraws a steering message the user removed. A message still being prepared is never sent; one + * the SDK holds is withdrawn from its steering queue. If the SDK already delivered it to the model, + * the entry is kept so its `user.message` echo still shows it as its own turn. + */ + async withdrawSteering(id: string): Promise { + const pending = this._pendingSteeringFlips.get(id); + if (!pending) { + return; + } + if (!pending.sdkSend) { + this._pendingSteeringFlips.delete(id); + this._logService.info(`[Copilot:${this.sessionId}] Withdrew steering message before sending it: id=${id}`); + return; + } + try { + const messageId = await pending.sdkSend.messageId; + const result = await this._wrapper.session.rpc.queue.withdrawMessage({ messageId, expectedPrompt: pending.sdkSend.prompt }); + if (!result.removed) { + this._logService.info(`[Copilot:${this.sessionId}] Steering message was already delivered and could not be withdrawn: id=${id}`); + return; + } + if (this._pendingSteeringFlips.get(id) === pending) { + this._pendingSteeringFlips.delete(id); + } + this._logService.info(`[Copilot:${this.sessionId}] Withdrew steering message: id=${id}`); + } catch (err) { + this._logService.warn(`[Copilot:${this.sessionId}] Failed to withdraw steering message: id=${id}: ${getErrorMessage(err)}`); + } + } + async getMessages(): Promise { const result = await this._getMappedEvents(); return result.turns; diff --git a/src/vs/platform/agentHost/test/node/agentService.test.ts b/src/vs/platform/agentHost/test/node/agentService.test.ts index 0f97f8d6b25d3c..57bf135c3980e5 100644 --- a/src/vs/platform/agentHost/test/node/agentService.test.ts +++ b/src/vs/platform/agentHost/test/node/agentService.test.ts @@ -8322,6 +8322,41 @@ suite('AgentService (node dispatcher)', () => { }); }); + test('a user can remove a steering message mid-turn when the provider can withdraw it', async () => { + const { svc, agent, session } = await setup(); + const withdrawn: string[] = []; + agent.withdrawSteeringMessage = (_chat, id) => { withdrawn.push(id); }; + await dispatchTurnAndWait(svc, agent, session, []); + const chat = buildDefaultChatUri(session); + const accepted = Event.toPromise(Event.filter(svc.onDidAction, envelope => + envelope.action.type === ActionType.ChatPendingMessageSet && envelope.action.id === 'steer-a')); + await svc.dispatchAction(chat, { + type: ActionType.ChatPendingMessageSet, + kind: PendingMessageKind.Steering, id: 'steer-a', message: { text: 'Use tabs, not spaces', origin: { kind: MessageKind.User } }, + }, 'test-client', 2); + await accepted; + const rejections: string[] = []; + disposables.add(svc.onDidAction(envelope => { + if (envelope.rejectionReason) { + rejections.push(envelope.rejectionReason); + } + })); + const removed = Event.toPromise(Event.filter(svc.onDidAction, envelope => envelope.origin?.clientSeq === 3)); + svc.dispatchAction(chat, { type: ActionType.ChatPendingMessageRemoved, kind: PendingMessageKind.Steering, id: 'steer-a' }, 'test-client', 3); + await removed; + assert.deepStrictEqual({ + handedToAgent: agent.setPendingMessagesCalls.map(call => call.steeringMessage?.id), + stillQueued: getStateManager(svc).getChatState(chat)?.queuedMessages?.map(message => message.message.text), + rejections, + withdrawn, + }, { + handedToAgent: ['steer-a'], + stillQueued: undefined, + rejections: [], + withdrawn: ['steer-a'], + }); + }); + test('snapshots embedded text attachments as text files without retaining the payload in state', async () => { const { svc, agent, session, attachmentsRoot } = await setup(); const metadata = { kind: 'paste' }; diff --git a/src/vs/platform/agentHost/test/node/chatContributions.test.ts b/src/vs/platform/agentHost/test/node/chatContributions.test.ts index ff9af59957ca33..077a776d3ae7bd 100644 --- a/src/vs/platform/agentHost/test/node/chatContributions.test.ts +++ b/src/vs/platform/agentHost/test/node/chatContributions.test.ts @@ -1109,7 +1109,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, pendingMessageSenders, admitted, admittedCredentials, titleController, telemetryService, clearAgent: () => agent = undefined, setConversionPending: (pending: boolean) => conversionPending = pending, setSteeringReady: (ready: boolean) => steeringReady = ready }; + return { service, stateManager, session, chat, agent: mockAgent, pendingMessages, pendingMessageSenders, admitted, admittedCredentials, titleController, telemetryService, clearAgent: () => agent = undefined, setConversionPending: (pending: boolean) => conversionPending = pending, setSteeringReady: (ready: boolean) => steeringReady = ready }; } function appliedClientAction(channel: string, session: string, action: IAppliedClientAction['action'], clientId = 'client'): IAppliedClientAction { @@ -1787,6 +1787,39 @@ suite('AgentHostChatContributions', () => { }, { delivered: ['a', 'b'], senders: ['client-a', 'client-b'], pending: ['b'] }); }); + test('queue drain withdraws steering the user removes mid-turn, but not steering the agent consumed', () => { + const queue = createQueueDrainContributions(disposables); + const withdrawn: string[] = []; + queue.agent.withdrawSteeringMessage = (_chat, id) => { withdrawn.push(id); }; + queue.stateManager.dispatchServerAction(queue.chat, { + type: ActionType.ChatTurnStarted, turnId: 'active-turn', + startedAt: '2026-10-17T00:00:00.000Z', + message: { text: 'Working', origin: { kind: MessageKind.User } }, + }); + for (const id of ['removed-by-user', 'consumed-by-agent']) { + const steer: IAppliedClientAction['action'] = { + type: ActionType.ChatPendingMessageSet, + kind: PendingMessageKind.Steering, id, message: { text: id, origin: { kind: MessageKind.User } }, + }; + queue.stateManager.dispatchServerAction(queue.chat, steer); + queue.service.didApplyClientAction(appliedClientAction(queue.chat, queue.session, steer)); + } + const userRemoval: IAppliedClientAction['action'] = { type: ActionType.ChatPendingMessageRemoved, kind: PendingMessageKind.Steering, id: 'removed-by-user' }; + queue.stateManager.dispatchServerAction(queue.chat, userRemoval); + queue.service.didApplyClientAction(appliedClientAction(queue.chat, queue.session, userRemoval)); + const agentCleanup: IDispatchedAction['action'] = { type: ActionType.ChatPendingMessageRemoved, kind: PendingMessageKind.Steering, id: 'consumed-by-agent' }; + queue.stateManager.dispatchServerAction(queue.chat, agentCleanup); + queue.service.didDispatchAction(dispatchedAction(queue.chat, queue.session, agentCleanup)); + + assert.deepStrictEqual({ + handedToAgent: queue.pendingMessages.map(message => message?.id), + withdrawn, + }, { + handedToAgent: ['removed-by-user', 'consumed-by-agent'], + withdrawn: ['removed-by-user'], + }); + }); + test('queue drain holds independent steering until a turn starts', () => { const queue = createQueueDrainContributions(disposables); const steer: IDispatchedAction['action'] = { diff --git a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts index 3326a7f4569f1d..e66409e80fec15 100644 --- a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts @@ -169,6 +169,8 @@ class MockCopilotSession { readonly eventLogReadRequests: Parameters[0][] = []; eventLogReadGate: Promise | undefined; readonly sendRequests: unknown[] = []; + readonly withdrawMessageRequests: Parameters[0][] = []; + withdrawMessageResult: Awaited> = { removed: true, interrupted: false }; readonly sendMessagesRequests: unknown[] = []; sendMessagesError: Error | undefined; sendMessagesGate: Promise | undefined; @@ -482,6 +484,12 @@ class MockCopilotSession { } await this.sendMessagesGate; }, + queue: { + withdrawMessage: async (params: Parameters[0]) => { + this.withdrawMessageRequests.push(params); + return this.withdrawMessageResult; + }, + }, debug: { collectLogs: async (params: Parameters[0]) => { this.collectLogsCalls.push(params); @@ -13820,6 +13828,84 @@ Use the attached image as context. assert.strictEqual(mockSession.sendRequests.length, 1); }); + /** Like the runtime, folds every steering message the SDK still holds into the running turn and echoes it. */ + function deliverHeldSteering(mockSession: MockCopilotSession): void { + const withdrawn = new Set(mockSession.withdrawMessageRequests.map(request => request.messageId)); + mockSession.sendRequests.forEach((request, index) => { + const { prompt, mode } = request as { prompt: string; mode?: string }; + if (mode === 'immediate' && !withdrawn.has(`message-${index + 1}`)) { + mockSession.fire('user.message', { content: prompt, interactionId: `interaction-${index + 1}` } as SessionEventPayload<'user.message'>['data']); + } + }); + } + + function startedTurns(signals: readonly AgentSignal[]) { + return getActions(signals).flatMap(action => action.type === ActionType.ChatTurnStarted ? [{ text: action.message.text, queuedMessageId: action.queuedMessageId }] : []); + } + + test('a steering message the user removes before the model reads it is withdrawn and never becomes a turn', async () => { + const { session, mockSession, signals } = await createAgentSession(disposables); + session.resetTurnState('turn-original'); + + await session.sendSteering({ id: 'steer-1', message: { text: 'focus on tests', origin: { kind: MessageKind.User } } }); + await session.withdrawSteering('steer-1'); + deliverHeldSteering(mockSession); + + assert.deepStrictEqual({ + withdrawRequests: mockSession.withdrawMessageRequests, + startedTurns: startedTurns(signals), + }, { + withdrawRequests: [{ messageId: 'message-1', expectedPrompt: 'focus on tests' }], + startedTurns: [], + }); + }); + + test('a steering message the user removes while it is still being prepared is never sent', async () => { + const { session, mockSession, signals } = await createAgentSession(disposables); + await timeout(0); + session.resetTurnState('turn-original'); + const gate = new DeferredPromise(); + mockSession.mcpListGates.push(gate.p); + const previousListCalls = mockSession.mcpListCalls; + const steering = session.sendSteering({ id: 'steer-1', message: { text: 'focus on tests', origin: { kind: MessageKind.User } } }); + try { + await timeout(0); + assert.strictEqual(mockSession.mcpListCalls, previousListCalls + 1); + await session.withdrawSteering('steer-1'); + gate.complete(); + await steering; + deliverHeldSteering(mockSession); + + assert.deepStrictEqual({ + sendRequests: mockSession.sendRequests, + withdrawRequests: mockSession.withdrawMessageRequests, + startedTurns: startedTurns(signals), + }, { sendRequests: [], withdrawRequests: [], startedTurns: [] }); + } finally { + gate.complete(); + await steering; + } + }); + + test('a removed steering message the model already received still shows as its own turn', async () => { + const { session, mockSession, signals } = await createAgentSession(disposables); + session.resetTurnState('turn-original'); + mockSession.withdrawMessageResult = { removed: false, interrupted: false }; + + await session.sendSteering({ id: 'steer-1', message: { text: 'focus on tests', origin: { kind: MessageKind.User } } }); + await session.withdrawSteering('steer-1'); + // The SDK had already folded the message into the running turn, so it still echoes it. + mockSession.fire('user.message', { content: 'focus on tests', interactionId: 'interaction-steer' } as SessionEventPayload<'user.message'>['data']); + + assert.deepStrictEqual({ + withdrawRequests: mockSession.withdrawMessageRequests, + startedTurns: startedTurns(signals), + }, { + withdrawRequests: [{ messageId: 'message-1', expectedPrompt: 'focus on tests' }], + startedTurns: [{ text: 'focus on tests', queuedMessageId: 'steer-1' }], + }); + }); + test('fires steering_consumed on abort when the steering never reached its turn', async () => { const { session, signals } = await createAgentSession(disposables); diff --git a/src/vs/platform/agentHost/test/node/mockAgent.ts b/src/vs/platform/agentHost/test/node/mockAgent.ts index 8381c8d07032c8..2d90f77bffad9f 100644 --- a/src/vs/platform/agentHost/test/node/mockAgent.ts +++ b/src/vs/platform/agentHost/test/node/mockAgent.ts @@ -268,6 +268,9 @@ export class MockAgent implements IAgent { return this.setPendingMessages(chat, message, [], sender); } + /** Assign to model a provider that can withdraw steering it was already handed. */ + withdrawSteeringMessage?: (chat: URI, messageId: string) => void; + async getSessionMessages(session: URI): Promise { const subagentInfo = parseSubagentSessionUri(session); if (subagentInfo) { From 7b8604efbdea35799c93720a31a2a9ec3a33f0d0 Mon Sep 17 00:00:00 2001 From: Karthik Nadig Date: Sun, 18 Oct 2026 11:50:47 -0700 Subject: [PATCH 2/2] Match steering echoes by the SDK message id before falling back to text When two steering messages had the same text and one was being withdrawn, the echo for the other one could be matched to the withdrawn message. That showed the removed message as a turn and left the delivered one pending. The SDK echo carries the id that send returned, so prefer an exact id match and only fall back to matching by text for messages without a known id. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../node/copilot/copilotAgentSession.ts | 36 ++++++++++++++----- .../test/node/copilotAgentSession.test.ts | 35 +++++++++++++++++- 2 files changed, 61 insertions(+), 10 deletions(-) diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts index 8bf239ab2629ac..9746cef1a64b5c 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts @@ -1024,6 +1024,8 @@ interface IPendingSteering { readonly sender: IAgentPendingMessageSender | undefined; /** Set once the message is handed to the SDK; identifies it for withdrawal. */ sdkSend?: { readonly prompt: string; readonly messageId: Promise }; + /** The SDK's id for the message once `send` returns; its `user.message` echo carries the same id. */ + sdkMessageId?: string; } /** Where a transient chat activity comes from; see `CopilotAgentSession._publishActivity` for the precedence. */ @@ -2311,20 +2313,34 @@ export class CopilotAgentSession extends Disposable { } /** - * Pops the buffered steering message whose text is contained in the SDK - * `user.message` content we just observed. Matching by content (rather - * than just popping FIFO) keeps us robust against the SDK reordering - * or coalescing entries — concurrent steering messages with different - * texts are still matched to the correct one. Returns `undefined` if - * no buffered entry matches; the caller treats the `user.message` as - * an ordinary echo and skips the turn flip. + * Pops the buffered steering message delivered by the SDK `user.message` + * we just observed. The echo carries the id `send` returned, so an exact + * id match wins; that keeps identical texts apart, e.g. one being + * withdrawn while the other is delivered. Otherwise matches by content, + * skipping entries the SDK already identified as a different message. + * Matching by content (rather than just popping FIFO) keeps us robust + * against the SDK reordering or coalescing entries — concurrent steering + * messages with different texts are still matched to the correct one. + * Returns `undefined` if no buffered entry matches; the caller treats the + * `user.message` as an ordinary echo and skips the turn flip. */ - private _takeMatchingPendingSteering(content: string): IPendingSteering | undefined { + private _takeMatchingPendingSteering(content: string, messageId: string | undefined): IPendingSteering | undefined { if (this._pendingSteeringFlips.size === 0) { return undefined; } + if (messageId) { + for (const [id, pending] of this._pendingSteeringFlips) { + if (pending.sdkMessageId === messageId) { + this._pendingSteeringFlips.delete(id); + return pending; + } + } + } let substringMatch: [string, IPendingSteering] | undefined; for (const [id, pending] of this._pendingSteeringFlips) { + if (messageId && pending.sdkMessageId) { + continue; + } if (pending.pendingMessage.message.text === content) { this._pendingSteeringFlips.delete(id); return pending; @@ -5626,6 +5642,8 @@ export class CopilotAgentSession extends Disposable { ); if (execution.kind === 'skipped' && this._pendingSteeringFlips.delete(steeringMessage.id)) { this._fireSteeringConsumed(steeringMessage.id); + } else if (execution.kind === 'executed' && execution.value !== undefined) { + pending.sdkMessageId = execution.value; } } catch (err) { this._logService.error(`[Copilot:${this.sessionId}] Steering message failed`, err); @@ -7892,7 +7910,7 @@ export class CopilotAgentSession extends Disposable { this._currentTurn.value?.markRunning(); // The message is admitted only after required plugins are prepared. this._setManagedPluginActivity(undefined); - const steering = this._takeMatchingPendingSteering(e.data.content); + const steering = this._takeMatchingPendingSteering(e.data.content, e.data.messageId); if (steering) { this._beginSteeringTurn(steering); if (e.data.interactionId) { diff --git a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts index e66409e80fec15..662695b3a503d7 100644 --- a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts @@ -171,6 +171,7 @@ class MockCopilotSession { readonly sendRequests: unknown[] = []; readonly withdrawMessageRequests: Parameters[0][] = []; withdrawMessageResult: Awaited> = { removed: true, interrupted: false }; + withdrawMessageGate: Promise | undefined; readonly sendMessagesRequests: unknown[] = []; sendMessagesError: Error | undefined; sendMessagesGate: Promise | undefined; @@ -487,6 +488,7 @@ class MockCopilotSession { queue: { withdrawMessage: async (params: Parameters[0]) => { this.withdrawMessageRequests.push(params); + await this.withdrawMessageGate; return this.withdrawMessageResult; }, }, @@ -13834,7 +13836,7 @@ Use the attached image as context. mockSession.sendRequests.forEach((request, index) => { const { prompt, mode } = request as { prompt: string; mode?: string }; if (mode === 'immediate' && !withdrawn.has(`message-${index + 1}`)) { - mockSession.fire('user.message', { content: prompt, interactionId: `interaction-${index + 1}` } as SessionEventPayload<'user.message'>['data']); + mockSession.fire('user.message', { content: prompt, messageId: `message-${index + 1}`, interactionId: `interaction-${index + 1}` } as SessionEventPayload<'user.message'>['data']); } }); } @@ -13906,6 +13908,37 @@ Use the attached image as context. }); }); + test('an identical steering message delivered while another is being withdrawn becomes its own turn', async () => { + const { session, mockSession, signals } = await createAgentSession(disposables); + session.resetTurnState('turn-original'); + const withdrawGate = new DeferredPromise(); + mockSession.withdrawMessageGate = withdrawGate.p; + + await session.sendSteering({ id: 'steer-1', message: { text: 'focus on tests', origin: { kind: MessageKind.User } } }); + await session.sendSteering({ id: 'steer-2', message: { text: 'focus on tests', origin: { kind: MessageKind.User } } }); + const withdrawal = session.withdrawSteering('steer-1'); + try { + await timeout(0); + assert.strictEqual(mockSession.withdrawMessageRequests.length, 1); + // The SDK withdraws the first message but delivers the second before the withdrawal returns. + mockSession.fire('user.message', { content: 'focus on tests', messageId: 'message-2', interactionId: 'interaction-2' } as SessionEventPayload<'user.message'>['data']); + } finally { + withdrawGate.complete(); + await withdrawal; + } + await session.abort(); + + assert.deepStrictEqual({ + withdrawRequests: mockSession.withdrawMessageRequests, + startedTurns: startedTurns(signals), + leftOver: signals.filter(signal => signal.kind === 'steering_consumed'), + }, { + withdrawRequests: [{ messageId: 'message-1', expectedPrompt: 'focus on tests' }], + startedTurns: [{ text: 'focus on tests', queuedMessageId: 'steer-2' }], + leftOver: [], + }); + }); + test('fires steering_consumed on abort when the steering never reached its turn', async () => { const { session, signals } = await createAgentSession(disposables);