From d26022975d0ee7b4ef27ff334da04d5ad27421a7 Mon Sep 17 00:00:00 2001 From: Rob Lourens Date: Thu, 26 Mar 2026 16:22:13 -0700 Subject: [PATCH] Reconnect to in-progress remote agent host chat sessions (#304589) * Reconnect to in-progress remote agent host chat sessions When opening a remote agent host session that has an active (in-progress) turn, the chat UI now reconnects to it and streams ongoing progress instead of only showing completed turns as history. Key changes: - activeTurnToProgress() converts accumulated active turn state into IChatProgress[] for initial replay - provideChatSessionContent detects activeTurn on session state, includes it in history, and wires up live streaming via progressObs - _reconnectToActiveTurn(): streams incremental text/reasoning/tool call/permission updates, handles turn completion, dispatches turnCancelled on interrupt, resolves pending permissions interactively - Fixes live object identity (reuses ChatToolInvocation instances from initial progress), snapshot-to-listener race (immediate reconciliation), and proper cancellation dispatch (Written by Copilot) * Address Copilot review: fix empty initialProgress guard and handle completed tool calls between snapshots (Written by Copilot) * Fix test failures: add partId to delta action, add _meta.toolKind for terminal tool --------- Co-authored-by: Copilot Co-authored-by: Raymond Zhao <7199958+rzhao271@users.noreply.github.com> --- .../agentHost/agentHostSessionHandler.ts | 260 +++++++++++++++- .../agentHost/stateToProgressAdapter.ts | 43 ++- .../agentHostChatContribution.test.ts | 282 +++++++++++++++++- .../stateToProgressAdapter.test.ts | 136 ++++++++- 4 files changed, 706 insertions(+), 15 deletions(-) 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 f51a429171c..89bb620781b 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -4,10 +4,10 @@ *--------------------------------------------------------------------------------------------*/ import { Throttler } from '../../../../../../base/common/async.js'; -import { CancellationToken } from '../../../../../../base/common/cancellation.js'; +import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js'; import { Emitter } from '../../../../../../base/common/event.js'; import { MarkdownString } from '../../../../../../base/common/htmlContent.js'; -import { Disposable, DisposableMap, DisposableStore, MutableDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js'; +import { Disposable, DisposableMap, DisposableStore, MutableDisposable, type IDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js'; import { observableValue } from '../../../../../../base/common/observable.js'; import { URI } from '../../../../../../base/common/uri.js'; import { generateUuid } from '../../../../../../base/common/uuid.js'; @@ -18,7 +18,7 @@ import { ActionType, isSessionAction, type ISessionAction } from '../../../../.. import { SessionClientState } from '../../../../../../platform/agentHost/common/state/sessionClientState.js'; import { AHP_AUTH_REQUIRED, ProtocolError } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { getToolKind, getToolLanguage } from '../../../../../../platform/agentHost/common/state/sessionReducers.js'; -import { AttachmentType, PendingMessageKind, ResponsePartKind, ToolCallCancellationReason, ToolCallConfirmationReason, ToolCallStatus, TurnState, type IMessageAttachment } from '../../../../../../platform/agentHost/common/state/sessionState.js'; +import { AttachmentType, PendingMessageKind, ResponsePartKind, ToolCallCancellationReason, ToolCallConfirmationReason, ToolCallStatus, TurnState, type IMessageAttachment, type ISessionState } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { ExtensionIdentifier } from '../../../../../../platform/extensions/common/extensions.js'; import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js'; import { ILogService } from '../../../../../../platform/log/common/log.js'; @@ -30,10 +30,10 @@ import { ChatAgentLocation, ChatModeKind } from '../../../common/constants.js'; import { ChatToolInvocation } from '../../../common/model/chatProgressTypes/chatToolInvocation.js'; import { IChatAgentData, IChatAgentImplementation, IChatAgentRequest, IChatAgentResult, IChatAgentService } from '../../../common/participants/chatAgents.js'; import { getAgentHostIcon } from '../agentSessions.js'; -import { finalizeToolInvocation, toolCallStateToInvocation, turnsToHistory, type IToolCallFileEdit } from './stateToProgressAdapter.js'; +import { activeTurnToProgress, finalizeToolInvocation, toolCallStateToInvocation, turnsToHistory, type IToolCallFileEdit } from './stateToProgressAdapter.js'; // ============================================================================= -// AgentHostSessionHandler — renderer-side handler for a single agent host +// AgentHostSessionHandler - renderer-side handler for a single agent host // chat session type. Bridges the protocol state layer with the chat UI: // subscribes to session state, derives IChatProgress[] from immutable state // changes, and dispatches client actions (turnStarted, toolCallConfirmed, @@ -55,17 +55,24 @@ class AgentHostChatSession extends Disposable implements IChatSession { readonly onDidStartServerRequest = this._onDidStartServerRequest.event; readonly requestHandler: IChatSession['requestHandler']; - readonly interruptActiveResponseCallback: IChatSession['interruptActiveResponseCallback']; + interruptActiveResponseCallback: IChatSession['interruptActiveResponseCallback']; constructor( readonly sessionResource: URI, readonly history: readonly IChatSessionHistoryItem[], private readonly _sendRequest: (request: IChatAgentRequest, progress: (parts: IChatProgress[]) => void, token: CancellationToken) => Promise, + initialProgress: IChatProgress[] | undefined, onDispose: () => void, @ILogService private readonly _logService: ILogService, ) { super(); + const hasActiveTurn = initialProgress !== undefined; + if (hasActiveTurn) { + this.isCompleteObs.set(false, undefined); + this.progressObs.set(initialProgress, undefined); + } + this._register(toDisposable(() => this._onWillDispose.fire())); this._register(toDisposable(onDispose)); @@ -76,9 +83,34 @@ class AgentHostChatSession extends Disposable implements IChatSession { this.isCompleteObs.set(true, undefined); }; - this.interruptActiveResponseCallback = history.length > 0 ? undefined : async () => { + // Provide interrupt callback when reconnecting to an active turn or + // when this is a brand-new session (no history yet). + this.interruptActiveResponseCallback = (hasActiveTurn || history.length === 0) ? async () => { return true; - }; + } : undefined; + } + + /** + * Registers a disposable to be cleaned up when this session is disposed. + */ + registerDisposable(disposable: T): T { + return this._register(disposable); + } + + /** + * Appends new progress items to the observable. Used by the reconnection + * flow to stream ongoing state changes into the chat UI. + */ + appendProgress(items: IChatProgress[]): void { + const current = this.progressObs.get(); + this.progressObs.set([...current, ...items], undefined); + } + + /** + * Marks the active turn as complete. + */ + complete(): void { + this.isCompleteObs.set(true, undefined); } /** @@ -175,6 +207,8 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC let resolvedSession: URI | undefined; const isUntitled = resourceKey.startsWith('untitled-'); const history: IChatSessionHistoryItem[] = []; + let initialProgress: IChatProgress[] | undefined; + let activeTurnId: string | undefined; if (!isUntitled) { resolvedSession = this._resolveSessionUri(sessionResource); this._sessionToBackend.set(resourceKey, resolvedSession); @@ -185,6 +219,26 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC const sessionState = this._clientState.getSessionState(resolvedSession.toString()); if (sessionState) { history.push(...turnsToHistory(sessionState.turns, this._config.agentId)); + + // If there's an active turn, include its request in history + // with an empty response so the chat service creates a + // pending request, then provide accumulated progress via + // progressObs for live streaming. + if (sessionState.activeTurn) { + activeTurnId = sessionState.activeTurn.id; + history.push({ + type: 'request', + prompt: sessionState.activeTurn.userMessage.text, + participant: this._config.agentId, + }); + history.push({ + type: 'response', + parts: [], + participant: this._config.agentId, + }); + initialProgress = activeTurnToProgress(sessionState.activeTurn); + this._logService.info(`[AgentHost] Reconnecting to active turn ${activeTurnId} for session ${resolvedSession.toString()}`); + } } } } catch (err) { @@ -206,6 +260,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC this._ensurePendingMessageSubscription(resourceKey, sessionResource, backendSession); return this._handleTurn(backendSession, request, progress, token); }, + initialProgress, () => { this._activeSessions.delete(resourceKey); this._sessionToBackend.delete(resourceKey); @@ -220,9 +275,15 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC ); this._activeSessions.set(resourceKey, session); - // For existing (non-untitled) sessions, start watching for server-initiated turns - // immediately. For untitled sessions, this is deferred to _createAndSubscribe. if (resolvedSession) { + // If reconnecting to an active turn, wire up an ongoing state listener + // to stream new progress into the session's progressObs. + if (activeTurnId && initialProgress !== undefined) { + this._reconnectToActiveTurn(resolvedSession, activeTurnId, session, initialProgress); + } + + // For existing (non-untitled) sessions, start watching for server-initiated turns + // immediately. For untitled sessions, this is deferred to _createAndSubscribe. this._watchForServerInitiatedTurns(resolvedSession, sessionResource); } @@ -395,7 +456,11 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC private _watchForServerInitiatedTurns(backendSession: URI, sessionResource: URI): void { const resourceKey = sessionResource.path.substring(1); const sessionStr = backendSession.toString(); - let lastSeenTurnId: string | undefined; + + // Seed from the current state so we don't treat any pre-existing active + // turn (e.g. one being handled by _reconnectToActiveTurn) as new. + const currentState = this._clientState.getSessionState(sessionStr); + let lastSeenTurnId: string | undefined = currentState?.activeTurn?.id; let previousQueuedIds: Set | undefined; const disposables = new DisposableStore(); @@ -827,6 +892,179 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC }); } + // ---- Reconnection to active turn ---------------------------------------- + + /** + * Wires up an ongoing state listener that streams incremental progress + * from an already-running turn into the chat session's progressObs. + * This is the reconnection counterpart of {@link _handleTurn}, which + * handles newly-initiated turns. + */ + private _reconnectToActiveTurn( + backendSession: URI, + turnId: string, + chatSession: AgentHostChatSession, + initialProgress: IChatProgress[], + ): void { + const sessionKey = backendSession.toString(); + + // Extract live ChatToolInvocation objects from the initial progress + // array so we can update/finalize the same instances the chat UI holds. + const activeToolInvocations = new Map(); + for (const item of initialProgress) { + if (item instanceof ChatToolInvocation) { + activeToolInvocations.set(item.toolCallId, item); + } + } + + // Track last-emitted content lengths per response part to compute deltas. + // Seed from the current state so we only emit new content beyond what + // activeTurnToProgress already captured. + const lastEmittedLengths = new Map(); + const currentState = this._clientState.getSessionState(sessionKey); + if (currentState?.activeTurn) { + for (const rp of currentState.activeTurn.responseParts) { + if (rp.kind === ResponsePartKind.Markdown || rp.kind === ResponsePartKind.Reasoning) { + lastEmittedLengths.set(rp.id, rp.content.length); + } + } + } + + const reconnectDisposables = chatSession.registerDisposable(new DisposableStore()); + const throttler = new Throttler(); + reconnectDisposables.add(throttler); + + // Set up the interrupt callback so the user can actually cancel the + // remote turn. This dispatches session/turnCancelled to the server. + chatSession.interruptActiveResponseCallback = async () => { + this._logService.info(`[AgentHost] Reconnect cancellation requested for ${sessionKey}, dispatching turnCancelled`); + const cancelAction = { + type: ActionType.SessionTurnCancelled as const, + session: sessionKey, + turnId, + }; + const seq = this._clientState.applyOptimistic(cancelAction); + this._config.connection.dispatchAction(cancelAction, this._clientState.clientId, seq); + return true; + }; + + // Wire up awaitConfirmation for tool calls that were already pending + // confirmation at snapshot time so the user can approve/deny them. + const cts = new CancellationTokenSource(); + reconnectDisposables.add(toDisposable(() => cts.dispose(true))); + for (const [toolCallId, invocation] of activeToolInvocations) { + if (!IChatToolInvocation.isComplete(invocation)) { + this._awaitToolConfirmation(invocation, toolCallId, backendSession, turnId, cts.token); + } + } + + // Process state changes from the protocol layer. + const processStateChange = (sessionState: ISessionState) => { + const activeTurn = sessionState.activeTurn; + const isActive = activeTurn?.id === turnId; + const responseParts = isActive + ? activeTurn.responseParts + : sessionState.turns.find(t => t.id === turnId)?.responseParts; + + if (responseParts) { + for (const rp of responseParts) { + switch (rp.kind) { + case ResponsePartKind.Markdown: { + const lastLen = lastEmittedLengths.get(rp.id) ?? 0; + if (rp.content.length > lastLen) { + const delta = rp.content.substring(lastLen); + lastEmittedLengths.set(rp.id, rp.content.length); + chatSession.appendProgress([{ kind: 'markdownContent', content: new MarkdownString(delta, { supportHtml: true }) }]); + } + break; + } + case ResponsePartKind.Reasoning: { + const lastLen = lastEmittedLengths.get(rp.id) ?? 0; + if (rp.content.length > lastLen) { + const delta = rp.content.substring(lastLen); + lastEmittedLengths.set(rp.id, rp.content.length); + chatSession.appendProgress([{ kind: 'thinking', value: delta }]); + } + break; + } + case ResponsePartKind.ToolCall: { + const tc = rp.toolCall; + const toolCallId = tc.toolCallId; + let existing = activeToolInvocations.get(toolCallId); + + if (!existing) { + existing = toolCallStateToInvocation(tc); + activeToolInvocations.set(toolCallId, existing); + chatSession.appendProgress([existing]); + + if (tc.status === ToolCallStatus.PendingConfirmation) { + this._awaitToolConfirmation(existing, toolCallId, backendSession, turnId, cts.token); + } + } else if (tc.status === ToolCallStatus.PendingConfirmation) { + // Running -> PendingConfirmation (re-confirmation). + existing.didExecuteTool(undefined); + const confirmInvocation = toolCallStateToInvocation(tc); + activeToolInvocations.set(toolCallId, confirmInvocation); + chatSession.appendProgress([confirmInvocation]); + this._awaitToolConfirmation(confirmInvocation, toolCallId, backendSession, turnId, cts.token); + } else if (tc.status === ToolCallStatus.Running) { + existing.invocationMessage = typeof tc.invocationMessage === 'string' + ? tc.invocationMessage + : new MarkdownString(tc.invocationMessage.markdown); + if (getToolKind(tc) === 'terminal' && tc.toolInput) { + existing.toolSpecificData = { + kind: 'terminal', + commandLine: { original: tc.toolInput }, + language: getToolLanguage(tc) ?? 'shellscript', + }; + } + } + + // Finalize terminal-state tools + if (existing && (tc.status === ToolCallStatus.Completed || tc.status === ToolCallStatus.Cancelled) && !IChatToolInvocation.isComplete(existing)) { + activeToolInvocations.delete(toolCallId); + finalizeToolInvocation(existing, tc); + // Note: file edits from reconnection are not routed through + // the editing session pipeline as there is no active request + // context. The edits already happened on the remote. + } + break; + } + } + } + } + + // If the turn is no longer active, emit any error and finish. + if (!isActive) { + const lastTurn = sessionState.turns.find(t => t.id === turnId); + if (lastTurn?.state === TurnState.Error && lastTurn.error) { + chatSession.appendProgress([{ + kind: 'markdownContent', + content: new MarkdownString(`\n\nError: (${lastTurn.error.errorType}) ${lastTurn.error.message}`), + }]); + } + chatSession.complete(); + reconnectDisposables.dispose(); + } + }; + + // Attach the ongoing state listener + reconnectDisposables.add(this._clientState.onDidChangeSessionState(e => { + if (e.session !== sessionKey) { + return; + } + throttler.queue(async () => processStateChange(e.state)); + })); + + // Immediately reconcile against the current state to close any gap + // between snapshot time and listener registration. If the turn already + // completed in the interim, this will mark the session complete. + const latestState = this._clientState.getSessionState(sessionKey); + if (latestState) { + processStateChange(latestState); + } + } + // ---- File edit routing --------------------------------------------------- /** diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts index dceeac89e09..d1eacaa63ba 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts @@ -6,7 +6,7 @@ import { IMarkdownString, MarkdownString } from '../../../../../../base/common/htmlContent.js'; import { generateUuid } from '../../../../../../base/common/uuid.js'; import { URI } from '../../../../../../base/common/uri.js'; -import { ToolCallStatus, TurnState, ResponsePartKind, getToolFileEdits, getToolOutputText, type ICompletedToolCall, type IToolCallState, type ITurn } from '../../../../../../platform/agentHost/common/state/sessionState.js'; +import { ToolCallStatus, TurnState, ResponsePartKind, getToolFileEdits, getToolOutputText, type IActiveTurn, type ICompletedToolCall, type IToolCallState, type ITurn } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { getToolKind, getToolLanguage } from '../../../../../../platform/agentHost/common/state/sessionReducers.js'; import { type IChatProgress, type IChatTerminalToolInvocationData, type IChatToolInputInvocationData, type IChatToolInvocationSerialized, ToolConfirmKind } from '../../../common/chatService/chatService.js'; import { type IChatSessionHistoryItem } from '../../../common/chatSessionsService.js'; @@ -74,6 +74,47 @@ export function turnsToHistory(turns: readonly ITurn[], participantId: string): return history; } +/** + * Converts an active (in-progress) turn's accumulated state into progress + * items suitable for replaying into the chat UI when reconnecting to a + * session that is mid-turn. + * + * Returns serialized progress items for content already received (text, + * reasoning, completed tool calls) and live {@link ChatToolInvocation} + * objects for running tool calls and pending confirmations. + */ +export function activeTurnToProgress(activeTurn: IActiveTurn): IChatProgress[] { + const parts: IChatProgress[] = []; + + for (const rp of activeTurn.responseParts) { + switch (rp.kind) { + case ResponsePartKind.Markdown: + if (rp.content) { + parts.push({ kind: 'markdownContent', content: new MarkdownString(rp.content) }); + } + break; + case ResponsePartKind.Reasoning: + if (rp.content) { + parts.push({ kind: 'thinking', value: rp.content }); + } + break; + case ResponsePartKind.ToolCall: { + const tc = rp.toolCall; + if (tc.status === ToolCallStatus.Completed || tc.status === ToolCallStatus.Cancelled) { + parts.push(completedToolCallToSerialized(tc as ICompletedToolCall)); + } else if (tc.status === ToolCallStatus.Running || tc.status === ToolCallStatus.Streaming || tc.status === ToolCallStatus.PendingConfirmation) { + parts.push(toolCallStateToInvocation(tc)); + } + break; + } + case ResponsePartKind.ContentRef: + break; + } + } + + return parts; +} + /** * Converts a completed tool call from the protocol state into a serialized * tool invocation suitable for history replay. diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts index ed53e2ac21a..fa0d6674c79 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts @@ -16,7 +16,7 @@ import { IConfigurationService } from '../../../../../../platform/configuration/ import { IAgentCreateSessionConfig, IAgentHostService, IAgentSessionMetadata, AgentSession } from '../../../../../../platform/agentHost/common/agentService.js'; import type { IActionEnvelope, INotification, ISessionAction, IToolCallConfirmedAction, ITurnStartedAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; import type { IStateSnapshot } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; -import { SessionLifecycle, SessionStatus, TurnState, createSessionState, ROOT_STATE_URI, PolicyState, ResponsePartKind, type ISessionState, type ISessionSummary } from '../../../../../../platform/agentHost/common/state/sessionState.js'; +import { SessionLifecycle, SessionStatus, TurnState, ToolCallStatus, ToolCallConfirmationReason, createSessionState, createActiveTurn, ROOT_STATE_URI, PolicyState, ResponsePartKind, type ISessionState, type ISessionSummary } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { IDefaultAccountService } from '../../../../../../platform/defaultAccount/common/defaultAccount.js'; import { IAuthenticationService } from '../../../../../services/authentication/common/authentication.js'; import { IChatAgentData, IChatAgentImplementation, IChatAgentRequest, IChatAgentService } from '../../../common/participants/chatAgents.js'; @@ -1516,6 +1516,286 @@ suite('AgentHostChatContribution', () => { }); }); + // ---- Reconnection to active turn ---------------------------------------- + + suite('reconnection to active turn', () => { + + function makeSessionStateWithActiveTurn(sessionUri: string, overrides?: Partial<{ streamingText: string; reasoning: string }>): ISessionState { + const summary: ISessionSummary = { + resource: sessionUri, + provider: 'copilot', + title: 'Active Session', + status: SessionStatus.Idle, + createdAt: Date.now(), + modifiedAt: Date.now(), + }; + const activeTurnParts = []; + const reasoningText = overrides?.reasoning ?? ''; + if (reasoningText) { + activeTurnParts.push({ kind: ResponsePartKind.Reasoning as const, id: 'reasoning-1', content: reasoningText }); + } + activeTurnParts.push({ kind: ResponsePartKind.Markdown as const, id: 'md-active', content: overrides?.streamingText ?? 'Partial response so far' }); + return { + ...createSessionState(summary), + lifecycle: SessionLifecycle.Ready, + turns: [{ + id: 'turn-completed', + userMessage: { text: 'First message' }, + responseParts: [{ kind: ResponsePartKind.Markdown as const, id: 'md-1', content: 'First response' }], + usage: undefined, + state: TurnState.Complete, + }], + activeTurn: { + ...createActiveTurn('turn-active', { text: 'Second message' }), + responseParts: activeTurnParts, + }, + }; + } + + test('loads completed turns as history and active turn request/response', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const sessionUri = AgentSession.uri('copilot', 'reconnect-1'); + agentHostService.sessionStates.set(sessionUri.toString(), makeSessionStateWithActiveTurn(sessionUri.toString())); + + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/reconnect-1' }); + const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => session.dispose())); + + // Should have: completed turn (request + response) + active turn (request + empty response) = 4 + assert.strictEqual(session.history.length, 4); + assert.strictEqual(session.history[0].type, 'request'); + if (session.history[0].type === 'request') { + assert.strictEqual(session.history[0].prompt, 'First message'); + } + assert.strictEqual(session.history[2].type, 'request'); + if (session.history[2].type === 'request') { + assert.strictEqual(session.history[2].prompt, 'Second message'); + } + // Active turn response should be an empty placeholder + assert.strictEqual(session.history[3].type, 'response'); + if (session.history[3].type === 'response') { + assert.strictEqual(session.history[3].parts.length, 0); + } + }); + + test('sets isCompleteObs to false and populates progressObs for active turn', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const sessionUri = AgentSession.uri('copilot', 'reconnect-2'); + agentHostService.sessionStates.set(sessionUri.toString(), makeSessionStateWithActiveTurn(sessionUri.toString())); + + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/reconnect-2' }); + const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => session.dispose())); + + assert.strictEqual(session.isCompleteObs?.get(), false, 'Should not be complete when active turn exists'); + const progress = session.progressObs?.get() ?? []; + assert.ok(progress.length > 0, 'Should have initial progress from active turn'); + // Should contain the streaming text as markdown + const markdownPart = progress.find(p => p.kind === 'markdownContent') as IChatMarkdownContent | undefined; + assert.ok(markdownPart, 'Should have markdown content from streaming text'); + assert.strictEqual(markdownPart!.content.value, 'Partial response so far'); + }); + + test('provides interruptActiveResponseCallback when reconnecting', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const sessionUri = AgentSession.uri('copilot', 'reconnect-3'); + agentHostService.sessionStates.set(sessionUri.toString(), makeSessionStateWithActiveTurn(sessionUri.toString())); + + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/reconnect-3' }); + const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => session.dispose())); + + assert.ok(session.interruptActiveResponseCallback, 'Should provide interrupt callback'); + }); + + test('interrupt callback dispatches turnCancelled action', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const sessionUri = AgentSession.uri('copilot', 'reconnect-cancel'); + agentHostService.sessionStates.set(sessionUri.toString(), makeSessionStateWithActiveTurn(sessionUri.toString())); + + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/reconnect-cancel' }); + const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => session.dispose())); + + assert.ok(session.interruptActiveResponseCallback); + const result = await session.interruptActiveResponseCallback!(); + assert.strictEqual(result, true); + + // Should have dispatched a turnCancelled action + const cancelAction = agentHostService.dispatchedActions.find(d => d.action.type === 'session/turnCancelled'); + assert.ok(cancelAction, 'Should dispatch session/turnCancelled'); + }); + + test('streams new text deltas into progressObs after reconnection', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const sessionUri = AgentSession.uri('copilot', 'reconnect-stream'); + const sessionState = makeSessionStateWithActiveTurn(sessionUri.toString(), { streamingText: 'Before' }); + agentHostService.sessionStates.set(sessionUri.toString(), sessionState); + + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/reconnect-stream' }); + const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => session.dispose())); + + const initialLen = (session.progressObs?.get() ?? []).length; + + // Fire a delta action to simulate the server streaming more text + agentHostService.fireAction({ + action: { type: 'session/delta', session: sessionUri.toString(), turnId: 'turn-active', partId: 'md-active', content: ' and more' } as ISessionAction, + serverSeq: 1, + origin: undefined, + }); + + await timeout(10); + + const progress = session.progressObs?.get() ?? []; + assert.ok(progress.length > initialLen, 'Should have appended new progress items'); + // The last markdown part should be the delta + const lastMarkdown = [...progress].reverse().find(p => p.kind === 'markdownContent') as IChatMarkdownContent; + assert.ok(lastMarkdown, 'Should have a new markdown delta'); + assert.strictEqual(lastMarkdown.content.value, ' and more'); + }); + + test('marks session complete when turn finishes', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const sessionUri = AgentSession.uri('copilot', 'reconnect-complete'); + agentHostService.sessionStates.set(sessionUri.toString(), makeSessionStateWithActiveTurn(sessionUri.toString())); + + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/reconnect-complete' }); + const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => session.dispose())); + + assert.strictEqual(session.isCompleteObs?.get(), false); + + // Fire turnComplete to finish the active turn + agentHostService.fireAction({ + action: { type: 'session/turnComplete', session: sessionUri.toString(), turnId: 'turn-active' } as ISessionAction, + serverSeq: 1, + origin: undefined, + }); + + await timeout(10); + + assert.strictEqual(session.isCompleteObs?.get(), true, 'Should be complete after turnComplete'); + }); + + test('handles active turn with running tool call', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const sessionUri = AgentSession.uri('copilot', 'reconnect-tool'); + const sessionState = makeSessionStateWithActiveTurn(sessionUri.toString()); + sessionState.activeTurn!.responseParts.push({ + kind: ResponsePartKind.ToolCall, + toolCall: { + toolCallId: 'tc-running', + toolName: 'bash', + displayName: 'Bash', + invocationMessage: 'Running command', + status: ToolCallStatus.Running, + confirmed: ToolCallConfirmationReason.NotNeeded, + }, + }); + agentHostService.sessionStates.set(sessionUri.toString(), sessionState); + + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/reconnect-tool' }); + const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => session.dispose())); + + const progress = session.progressObs?.get() ?? []; + const toolInvocation = progress.find(p => p.kind === 'toolInvocation') as IChatToolInvocation | undefined; + assert.ok(toolInvocation, 'Should have a live tool invocation in progress'); + assert.strictEqual(toolInvocation!.toolCallId, 'tc-running'); + }); + + test('handles active turn with pending tool confirmation', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const sessionUri = AgentSession.uri('copilot', 'reconnect-perm'); + const sessionState = makeSessionStateWithActiveTurn(sessionUri.toString()); + sessionState.activeTurn!.responseParts.push({ + kind: ResponsePartKind.ToolCall, + toolCall: { + toolCallId: 'tc-pending', + toolName: 'bash', + displayName: 'Bash', + invocationMessage: 'Run command', + confirmationTitle: 'Clean up', + toolInput: 'rm -rf /tmp/test', + status: ToolCallStatus.PendingConfirmation, + }, + }); + agentHostService.sessionStates.set(sessionUri.toString(), sessionState); + + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/reconnect-perm' }); + const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => session.dispose())); + + const progress = session.progressObs?.get() ?? []; + const permInvocation = progress.find(p => p.kind === 'toolInvocation') as IChatToolInvocation | undefined; + assert.ok(permInvocation, 'Should have a live permission request in progress'); + + // Complete the turn so the awaitConfirmation promise and its internal + // DisposableStore are cleaned up before test teardown. + agentHostService.fireAction({ + action: { type: 'session/turnComplete', session: sessionUri.toString(), turnId: 'turn-active' } as ISessionAction, + serverSeq: 1, + origin: undefined, + }); + await timeout(10); + }); + + test('no active turn loads completed history only with isComplete true', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const sessionUri = AgentSession.uri('copilot', 'no-active-turn'); + agentHostService.sessionStates.set(sessionUri.toString(), { + ...createSessionState({ resource: sessionUri.toString(), provider: 'copilot', title: 'Done', status: SessionStatus.Idle, createdAt: Date.now(), modifiedAt: Date.now() }), + lifecycle: SessionLifecycle.Ready, + turns: [{ + id: 'turn-done', + userMessage: { text: 'Hello' }, + responseParts: [{ kind: ResponsePartKind.Markdown as const, id: 'md-1', content: 'Hi' }], + usage: undefined, + state: TurnState.Complete, + }], + }); + + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/no-active-turn' }); + const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => session.dispose())); + + assert.strictEqual(session.history.length, 2); + assert.strictEqual(session.isCompleteObs?.get(), true); + assert.deepStrictEqual(session.progressObs?.get(), []); + }); + + test('includes reasoning in initial progress', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const sessionUri = AgentSession.uri('copilot', 'reconnect-reasoning'); + agentHostService.sessionStates.set(sessionUri.toString(), makeSessionStateWithActiveTurn(sessionUri.toString(), { + streamingText: 'text', + reasoning: 'Let me think...', + })); + + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/reconnect-reasoning' }); + const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => session.dispose())); + + const progress = session.progressObs?.get() ?? []; + const thinking = progress.find(p => p.kind === 'thinking'); + assert.ok(thinking, 'Should have thinking progress from reasoning'); + const markdown = progress.find(p => p.kind === 'markdownContent') as IChatMarkdownContent; + assert.ok(markdown); + assert.strictEqual(markdown.content.value, 'text'); + }); + }); + // ---- Server-initiated turns ------------------------------------------- suite('server-initiated turns', () => { diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/stateToProgressAdapter.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/stateToProgressAdapter.test.ts index 800c3389593..64ee099bd7d 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/stateToProgressAdapter.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/stateToProgressAdapter.test.ts @@ -6,10 +6,10 @@ import assert from 'assert'; import { URI } from '../../../../../../base/common/uri.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; -import { ToolCallStatus, ToolCallConfirmationReason, ToolResultContentType, TurnState, ResponsePartKind, type ICompletedToolCall, type IToolCallRunningState, type ITurn, type IToolCallResponsePart, ToolCallCancellationReason } from '../../../../../../platform/agentHost/common/state/sessionState.js'; +import { ToolCallStatus, ToolCallConfirmationReason, ToolResultContentType, TurnState, ResponsePartKind, type IActiveTurn, type ICompletedToolCall, type IToolCallRunningState, type ITurn, type IToolCallResponsePart, ToolCallCancellationReason } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { IChatToolInvocationSerialized, type IChatMarkdownContent } from '../../../common/chatService/chatService.js'; import { ToolDataSource } from '../../../common/tools/languageModelToolsService.js'; -import { turnsToHistory, toolCallStateToInvocation, finalizeToolInvocation } from '../../../browser/agentSessions/agentHost/stateToProgressAdapter.js'; +import { turnsToHistory, activeTurnToProgress, toolCallStateToInvocation, finalizeToolInvocation } from '../../../browser/agentSessions/agentHost/stateToProgressAdapter.js'; // ---- Helper factories ------------------------------------------------------- @@ -372,4 +372,136 @@ suite('stateToProgressAdapter', () => { assert.strictEqual(fileEdits.length, 0); }); }); + + suite('activeTurnToProgress', () => { + + function createActiveTurnState(responseParts?: IActiveTurn['responseParts']): IActiveTurn { + return { + id: 'turn-active', + userMessage: { text: 'Do things' }, + responseParts: responseParts ?? [], + usage: undefined, + }; + } + + test('empty active turn produces empty progress', () => { + const result = activeTurnToProgress(createActiveTurnState()); + assert.deepStrictEqual(result, []); + }); + + test('produces markdown content for streamed text', () => { + const result = activeTurnToProgress(createActiveTurnState([ + { kind: ResponsePartKind.Markdown, id: 'md-1', content: 'Hello world' }, + ])); + assert.strictEqual(result.length, 1); + assert.strictEqual(result[0].kind, 'markdownContent'); + assert.strictEqual((result[0] as IChatMarkdownContent).content.value, 'Hello world'); + }); + + test('produces thinking progress for reasoning', () => { + const result = activeTurnToProgress(createActiveTurnState([ + { kind: ResponsePartKind.Reasoning, id: 'r-1', content: 'Let me think about this...' }, + ])); + assert.strictEqual(result.length, 1); + assert.strictEqual(result[0].kind, 'thinking'); + }); + + test('reasoning comes before streamed text when ordered that way', () => { + const result = activeTurnToProgress(createActiveTurnState([ + { kind: ResponsePartKind.Reasoning, id: 'r-1', content: 'Hmm...' }, + { kind: ResponsePartKind.Markdown, id: 'md-1', content: 'Result text' }, + ])); + assert.strictEqual(result.length, 2); + assert.strictEqual(result[0].kind, 'thinking'); + assert.strictEqual(result[1].kind, 'markdownContent'); + }); + + test('serializes completed tool calls', () => { + const result = activeTurnToProgress(createActiveTurnState([ + { + kind: ResponsePartKind.ToolCall, + toolCall: { + status: ToolCallStatus.Completed, + toolCallId: 'tc-done', + toolName: 'test_tool', + displayName: 'Test Tool', + invocationMessage: 'Ran test', + confirmed: ToolCallConfirmationReason.NotNeeded, + success: true, + pastTenseMessage: 'Ran test tool', + } as IToolCallResponsePart['toolCall'], + }, + ])); + assert.strictEqual(result.length, 1); + assert.strictEqual(result[0].kind, 'toolInvocationSerialized'); + }); + + test('creates live invocations for running tool calls', () => { + const result = activeTurnToProgress(createActiveTurnState([ + { + kind: ResponsePartKind.ToolCall, + toolCall: createToolCallState({ + toolCallId: 'tc-running', + status: ToolCallStatus.Running, + }), + }, + ])); + assert.strictEqual(result.length, 1); + // Live ChatToolInvocation - check it has the right toolCallId + const invocation = result[0] as { toolCallId?: string; kind?: string }; + assert.strictEqual(invocation.toolCallId, 'tc-running'); + }); + + test('creates confirmation invocations for pending tool confirmations', () => { + const result = activeTurnToProgress(createActiveTurnState([ + { + kind: ResponsePartKind.ToolCall, + toolCall: { + toolCallId: 'tc-pending', + toolName: 'bash', + displayName: 'Bash', + invocationMessage: 'Run command', + status: ToolCallStatus.PendingConfirmation, + confirmationTitle: 'Run command', + toolInput: 'echo hello', + _meta: { toolKind: 'terminal' }, + }, + }, + ])); + assert.strictEqual(result.length, 1); + // PendingConfirmation invocations have terminal toolSpecificData for shell tools + const invocation = result[0] as { toolSpecificData?: { kind: string } }; + assert.ok(invocation.toolSpecificData); + assert.strictEqual(invocation.toolSpecificData.kind, 'terminal'); + }); + + test('includes all parts in correct order', () => { + const result = activeTurnToProgress(createActiveTurnState([ + { kind: ResponsePartKind.Reasoning, id: 'r-1', content: 'Thinking...' }, + { kind: ResponsePartKind.Markdown, id: 'md-1', content: 'Output so far' }, + { + kind: ResponsePartKind.ToolCall, + toolCall: createToolCallState({ + toolCallId: 'tc-1', + status: ToolCallStatus.Running, + }), + }, + { + kind: ResponsePartKind.ToolCall, + toolCall: { + toolCallId: 'tc-2', + toolName: 'test_tool', + displayName: 'Test Tool', + invocationMessage: 'Confirm', + status: ToolCallStatus.PendingConfirmation, + confirmationTitle: 'Confirm', + }, + }, + ])); + // reasoning + text + tool call + pending confirmation = 4 items + assert.strictEqual(result.length, 4); + assert.strictEqual(result[0].kind, 'thinking'); + assert.strictEqual(result[1].kind, 'markdownContent'); + }); + }); });