From eca37987beb68ae82a992892813aa378a626ccdf Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Mon, 30 Mar 2026 15:38:09 -0700 Subject: [PATCH] agentHost: propagate steering for cli sessions (#306488) * agentHost: propagate steering for cli sessions Closes #305323 * comments --- .../platform/agentHost/common/agentService.ts | 9 ++++++- .../agentHost/node/agentSideEffects.ts | 23 +++++++++------- .../node/copilot/copilotAgentSession.ts | 19 +++++++++----- .../test/node/agentSideEffects.test.ts | 19 ++++++++++++-- .../test/node/copilotAgentSession.test.ts | 26 +++++++++++++++++++ .../agentHost/agentHostSessionHandler.ts | 8 ++++++ 6 files changed, 85 insertions(+), 19 deletions(-) diff --git a/src/vs/platform/agentHost/common/agentService.ts b/src/vs/platform/agentHost/common/agentService.ts index 421e5b39cd5..f044baa13e1 100644 --- a/src/vs/platform/agentHost/common/agentService.ts +++ b/src/vs/platform/agentHost/common/agentService.ts @@ -251,6 +251,12 @@ export interface IAgentReasoningEvent extends IAgentProgressEventBase { readonly content: string; } +/** A steering message was consumed (sent to the model). */ +export interface IAgentSteeringConsumedEvent extends IAgentProgressEventBase { + readonly type: 'steering_consumed'; + readonly id: string; +} + export type IAgentProgressEvent = | IAgentDeltaEvent | IAgentMessageEvent @@ -261,7 +267,8 @@ export type IAgentProgressEvent = | IAgentTitleChangedEvent | IAgentErrorEvent | IAgentUsageEvent - | IAgentReasoningEvent; + | IAgentReasoningEvent + | IAgentSteeringConsumedEvent; // ---- Session URI helpers ---------------------------------------------------- diff --git a/src/vs/platform/agentHost/node/agentSideEffects.ts b/src/vs/platform/agentHost/node/agentSideEffects.ts index 782898884c4..843e9ca32db 100644 --- a/src/vs/platform/agentHost/node/agentSideEffects.ts +++ b/src/vs/platform/agentHost/node/agentSideEffects.ts @@ -167,6 +167,16 @@ export class AgentSideEffects extends Disposable { if (e.type === 'idle') { this._tryConsumeNextQueuedMessage(sessionKey); } + + // Steering message was consumed by the agent — remove from protocol state + if (e.type === 'steering_consumed') { + this._stateManager.dispatchServerAction({ + type: ActionType.SessionPendingMessageRemoved, + session: sessionKey, + kind: PendingMessageKind.Steering, + id: e.id, + }); + } })); return disposables; } @@ -258,16 +268,9 @@ export class AgentSideEffects extends Disposable { [], ); - // Steering messages are consumed immediately by the agent; - // remove from protocol state so clients see the consumption. - if (state.steeringMessage) { - this._stateManager.dispatchServerAction({ - type: ActionType.SessionPendingMessageRemoved, - session, - kind: PendingMessageKind.Steering, - id: state.steeringMessage.id, - }); - } + // Steering message removal is now dispatched by the agent + // via the 'steering_consumed' progress event once the message + // has actually been sent to the model. // If the session is idle, try to consume the next queued message this._tryConsumeNextQueuedMessage(session); diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts index a4046ad2467..9b79f53d587 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgentSession.ts @@ -194,14 +194,21 @@ export class CopilotAgentSession extends Disposable { this._logService.info(`[Copilot:${this.sessionId}] session.send() returned`); } - sendSteering(steeringMessage: IPendingMessage): void { + async sendSteering(steeringMessage: IPendingMessage): Promise { this._logService.info(`[Copilot:${this.sessionId}] Sending steering message: "${steeringMessage.userMessage.text.substring(0, 100)}"`); - this._wrapper.session.send({ - prompt: steeringMessage.userMessage.text, - mode: 'immediate', - }).catch(err => { + try { + await this._wrapper.session.send({ + prompt: steeringMessage.userMessage.text, + mode: 'immediate', + }); + this._onDidSessionProgress.fire({ + session: this.sessionUri, + type: 'steering_consumed', + id: steeringMessage.id, + }); + } catch (err) { this._logService.error(`[Copilot:${this.sessionId}] Steering message failed`, err); - }); + } } async getMessages(): Promise<(IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[]> { diff --git a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts index e6b448f95e8..3995b6caef6 100644 --- a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts +++ b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts @@ -390,8 +390,9 @@ suite('AgentSideEffects', () => { assert.strictEqual(state?.queuedMessages?.[0].id, 'q-wait'); }); - test('dispatches SessionPendingMessageRemoved for steering messages', () => { + test('dispatches SessionPendingMessageRemoved for steering messages on steering_consumed', () => { setupSession(); + disposables.add(sideEffects.registerProgressListener(agent)); const envelopes: IActionEnvelope[] = []; disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e))); @@ -406,7 +407,21 @@ suite('AgentSideEffects', () => { stateManager.dispatchClientAction(action, { clientId: 'test', clientSeq: 1 }); sideEffects.handleAction(action); - const removal = envelopes.find(e => + // Removal is not dispatched synchronously; it waits for the agent + let removal = envelopes.find(e => + e.action.type === ActionType.SessionPendingMessageRemoved && + (e.action as { kind: PendingMessageKind }).kind === PendingMessageKind.Steering + ); + assert.strictEqual(removal, undefined, 'should not dispatch removal until steering_consumed'); + + // Simulate the agent consuming the steering message + agent.fireProgress({ + session: sessionUri, + type: 'steering_consumed', + id: 'steer-rm', + }); + + removal = envelopes.find(e => e.action.type === ActionType.SessionPendingMessageRemoved && (e.action as { kind: PendingMessageKind }).kind === PendingMessageKind.Steering ); diff --git a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts index 7ff57d4e8df..8664bb984fa 100644 --- a/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotAgentSession.test.ts @@ -207,6 +207,32 @@ suite('CopilotAgentSession', () => { }); }); + // ---- sendSteering ---- + + suite('sendSteering', () => { + + test('fires steering_consumed after send resolves', async () => { + const { session, progressEvents } = await createAgentSession(disposables); + + await session.sendSteering({ id: 'steer-1', userMessage: { text: 'focus on tests' } }); + + const consumed = progressEvents.find(e => e.type === 'steering_consumed'); + assert.ok(consumed, 'should fire steering_consumed event'); + assert.strictEqual((consumed as { id: string }).id, 'steer-1'); + }); + + test('does not fire steering_consumed when send fails', async () => { + const { session, mockSession, progressEvents } = await createAgentSession(disposables); + + mockSession.send = async () => { throw new Error('send failed'); }; + + await session.sendSteering({ id: 'steer-fail', userMessage: { text: 'will fail' } }); + + const consumed = progressEvents.find(e => e.type === 'steering_consumed'); + assert.strictEqual(consumed, undefined, 'should not fire steering_consumed on failure'); + }); + }); + // ---- event mapping ---- suite('event mapping', () => { 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 105c45831df..2df40d981a5 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -504,6 +504,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC const currentState = this._clientState.getSessionState(sessionStr); let lastSeenTurnId: string | undefined = currentState?.activeTurn?.id; let previousQueuedIds: Set | undefined; + let previousSteeringId: string | undefined = currentState?.steeringMessage?.id; const disposables = new DisposableStore(); @@ -518,6 +519,13 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC // Track queued message IDs so we can detect which one was consumed const currentQueuedIds = new Set((e.state.queuedMessages ?? []).map(m => m.id)); + const currentSteeringId = e.state.steeringMessage?.id; + + // Detect steering message removal or replacement regardless of turn changes + if (previousSteeringId && previousSteeringId !== currentSteeringId) { + this._chatService.removePendingRequest(sessionResource, previousSteeringId); + } + previousSteeringId = currentSteeringId; const activeTurn = e.state.activeTurn; if (!activeTurn || activeTurn.id === lastSeenTurnId) {