From e9299757f29949386f32f8b7e4367fa72f8d0720 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Sat, 7 Mar 2026 18:01:54 -0800 Subject: [PATCH] first working protocol version align more closely with protocol json rpc and some gaps --- .agents/skills/launch/SKILL.md | 17 + eslint.config.js | 1 + package-lock.json | 3 +- package.json | 2 + scripts/code-agent-host.js | 79 ++ scripts/code-agent-host.sh | 31 + src/vs/monaco.d.ts | 6 +- .../platform/agentHost/common/agentService.ts | 69 +- .../agentHost/common/state/sessionActions.ts | 27 +- .../common/state/sessionClientState.ts | 2 +- .../agentHost/common/state/sessionProtocol.ts | 275 ++-- .../agentHost/common/state/sessionReducers.ts | 51 +- .../agentHost/common/state/sessionState.ts | 44 +- .../common/state/sessionTransport.ts | 42 + .../agentHost/common/state/versions/v1.ts | 49 +- .../common/state/versions/versionRegistry.ts | 17 +- .../electron-browser/agentHostService.ts | 51 +- .../agentHost/node/agentEventMapper.ts | 160 +++ .../agentHost/node/agentHostServerMain.ts | 286 ++++ .../platform/agentHost/node/agentService.ts | 192 ++- .../agentHost/node/protocolServerHandler.ts | 334 +++++ .../agentHost/node/sessionStateManager.ts | 218 +++ .../agentHost/node/webSocketTransport.ts | 135 ++ src/vs/platform/agentHost/protocol.md | 137 +- .../test/node/agentEventMapper.test.ts | 221 +++ .../agentHost/test/node/agentService.test.ts | 61 +- .../node/createAndSendMessageAsLocalAgent.sh | 241 ++++ .../platform/agentHost/test/node/mockAgent.ts | 162 +++ .../test/node/protocolServerHandler.test.ts | 308 ++++ .../node/protocolWebSocket.integrationTest.ts | 663 +++++++++ .../test/node/sessionStateManager.test.ts | 163 +++ .../agentHost/agentHostChatContribution.ts | 103 +- .../agentHostLanguageModelProvider.ts | 69 +- .../agentHost/agentHostSessionHandler.ts | 581 ++++---- .../agentHostSessionListController.ts | 40 +- .../agentHost/stateToProgressAdapter.ts | 199 +++ .../agentHostChatContribution.test.ts | 1249 ++++++++--------- .../stateToProgressAdapter.test.ts | 298 ++++ 38 files changed, 5217 insertions(+), 1369 deletions(-) create mode 100644 scripts/code-agent-host.js create mode 100755 scripts/code-agent-host.sh create mode 100644 src/vs/platform/agentHost/common/state/sessionTransport.ts create mode 100644 src/vs/platform/agentHost/node/agentEventMapper.ts create mode 100644 src/vs/platform/agentHost/node/agentHostServerMain.ts create mode 100644 src/vs/platform/agentHost/node/protocolServerHandler.ts create mode 100644 src/vs/platform/agentHost/node/sessionStateManager.ts create mode 100644 src/vs/platform/agentHost/node/webSocketTransport.ts create mode 100644 src/vs/platform/agentHost/test/node/agentEventMapper.test.ts create mode 100755 src/vs/platform/agentHost/test/node/createAndSendMessageAsLocalAgent.sh create mode 100644 src/vs/platform/agentHost/test/node/mockAgent.ts create mode 100644 src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts create mode 100644 src/vs/platform/agentHost/test/node/protocolWebSocket.integrationTest.ts create mode 100644 src/vs/platform/agentHost/test/node/sessionStateManager.test.ts create mode 100644 src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts create mode 100644 src/vs/workbench/contrib/chat/test/browser/agentSessions/stateToProgressAdapter.test.ts diff --git a/.agents/skills/launch/SKILL.md b/.agents/skills/launch/SKILL.md index f967f6d0c1d..8000c393be0 100644 --- a/.agents/skills/launch/SKILL.md +++ b/.agents/skills/launch/SKILL.md @@ -348,3 +348,20 @@ Verify it's gone: # Confirm no process is listening on the debug port lsof -i :9224 # should return nothing ``` + +## Quick Test: Send a Chat Message as Local Agent + +There's a helper script that automates the full flow — launch Code OSS, switch to Local Agent mode, send a message, and print the response: + +```bash +# From the repo root: +./src/vs/platform/agent/test/node/createAndSendMessageAsLocalAgent.sh "Hello, what can you do?" + +# Options: +# --port CDP port (default: 9224) +# --timeout Response wait in seconds (default: 30) +# --no-kill Keep Code OSS running after +# --skip-launch Connect to already-running instance +``` + +This uses the JS mouse-event focus + `press`-per-key approach internally, handles session target switching, and cleans up on exit. diff --git a/eslint.config.js b/eslint.config.js index 50355b8289c..9411af56dfe 100644 --- a/eslint.config.js +++ b/eslint.config.js @@ -1499,6 +1499,7 @@ export default tseslint.config( 'vscode-regexpp', 'vscode-textmate', 'worker_threads', + 'ws', '@xterm/addon-clipboard', '@xterm/addon-image', '@xterm/addon-ligatures', diff --git a/package-lock.json b/package-lock.json index 736038a5cb4..2ee99b70b44 100644 --- a/package-lock.json +++ b/package-lock.json @@ -59,6 +59,7 @@ "vscode-oniguruma": "1.7.0", "vscode-regexpp": "^3.1.0", "vscode-textmate": "^9.3.2", + "ws": "^8.19.0", "yauzl": "^3.0.0", "yazl": "^2.4.3" }, @@ -81,6 +82,7 @@ "@types/wicg-file-system-access": "^2023.10.7", "@types/windows-foreground-love": "^0.3.0", "@types/winreg": "^1.2.30", + "@types/ws": "^8.18.1", "@types/yauzl": "^2.10.0", "@types/yazl": "^2.4.2", "@typescript-eslint/utils": "^8.45.0", @@ -21155,7 +21157,6 @@ "version": "8.19.0", "resolved": "https://registry.npmjs.org/ws/-/ws-8.19.0.tgz", "integrity": "sha512-blAT2mjOEIi0ZzruJfIhb3nps74PRWTCz1IjglWEEpQl5XS/UNama6u2/rjFkDDouqr4L67ry+1aGIALViWjDg==", - "dev": true, "license": "MIT", "engines": { "node": ">=10.0.0" diff --git a/package.json b/package.json index 054692f77e4..5940aad910c 100644 --- a/package.json +++ b/package.json @@ -129,6 +129,7 @@ "vscode-oniguruma": "1.7.0", "vscode-regexpp": "^3.1.0", "vscode-textmate": "^9.3.2", + "ws": "^8.19.0", "yauzl": "^3.0.0", "yazl": "^2.4.3" }, @@ -151,6 +152,7 @@ "@types/wicg-file-system-access": "^2023.10.7", "@types/windows-foreground-love": "^0.3.0", "@types/winreg": "^1.2.30", + "@types/ws": "^8.18.1", "@types/yauzl": "^2.10.0", "@types/yazl": "^2.4.2", "@typescript-eslint/utils": "^8.45.0", diff --git a/scripts/code-agent-host.js b/scripts/code-agent-host.js new file mode 100644 index 00000000000..f91eda13784 --- /dev/null +++ b/scripts/code-agent-host.js @@ -0,0 +1,79 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +// @ts-check + +const cp = require('child_process'); +const path = require('path'); +const minimist = require('minimist'); + +async function main() { + const args = minimist(process.argv.slice(2), { + boolean: ['help', 'no-launch'], + string: ['port'], + }); + + if (args.help) { + console.log( + 'Usage: ./scripts/code-agent-host.sh [options]\n' + + '\n' + + 'Options:\n' + + ' --port Port to listen on (default: 8081, or VSCODE_AGENT_HOST_PORT env)\n' + + ' --no-launch Start server without additional actions\n' + + ' --help Show this help message', + ); + return; + } + + const port = args.port || process.env['VSCODE_AGENT_HOST_PORT'] || '8081'; + const addr = await startServer(['--port', String(port)]); + console.log(`Agent Host server listening on ${addr}`); +} + +function startServer(programArgs) { + return new Promise((resolve, reject) => { + const env = { ...process.env }; + const entryPoint = path.join( + __dirname, + '..', + 'out', + 'vs', + 'platform', + 'agent', + 'node', + 'agentHostServerMain.js', + ); + + console.log( + `Starting agent host server: ${entryPoint} ${programArgs.join(' ')}`, + ); + const proc = cp.spawn(process.execPath, [entryPoint, ...programArgs], { + env, + stdio: [process.stdin, null, process.stderr], + }); + proc.stdout.on('data', (data) => { + const text = data.toString(); + process.stdout.write(text); + const m = text.match(/READY:(\d+)/); + if (m) { + resolve(`ws://127.0.0.1:${m[1]}`); + } + }); + + proc.on('exit', (code) => process.exit(code)); + + process.on('exit', () => proc.kill()); + process.on('SIGINT', () => { + proc.kill(); + process.exit(128 + 2); + }); + process.on('SIGTERM', () => { + proc.kill(); + process.exit(128 + 15); + }); + }); +} + +main(); diff --git a/scripts/code-agent-host.sh b/scripts/code-agent-host.sh new file mode 100755 index 00000000000..663f938ea1c --- /dev/null +++ b/scripts/code-agent-host.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash + +if [[ "$OSTYPE" == "darwin"* ]]; then + realpath() { [[ $1 = /* ]] && echo "$1" || echo "$PWD/${1#./}"; } + ROOT=$(dirname $(dirname $(realpath "$0"))) +else + ROOT=$(dirname $(dirname $(readlink -f $0))) +fi + +function code() { + pushd $ROOT + + # Get electron, compile, built-in extensions + if [[ -z "${VSCODE_SKIP_PRELAUNCH}" ]]; then + node build/lib/preLaunch.ts + fi + + NODE=$(node build/lib/node.ts) + if [ ! -e $NODE ];then + # Load remote node + npm run gulp node + fi + + popd + + NODE_ENV=development \ + VSCODE_DEV=1 \ + exec "$NODE" "$ROOT/scripts/code-agent-host.js" "$@" +} + +code "$@" diff --git a/src/vs/monaco.d.ts b/src/vs/monaco.d.ts index 5b0047e74a3..7d60866f371 100644 --- a/src/vs/monaco.d.ts +++ b/src/vs/monaco.d.ts @@ -5263,7 +5263,7 @@ declare namespace monaco.editor { export const EditorOptions: { acceptSuggestionOnCommitCharacter: IEditorOption; acceptSuggestionOnEnter: IEditorOption; - accessibilitySupport: IEditorOption; + accessibilitySupport: IEditorOption; accessibilityPageSize: IEditorOption; allowOverflow: IEditorOption; allowVariableLineHeights: IEditorOption; @@ -5326,7 +5326,7 @@ declare namespace monaco.editor { foldingMaximumRegions: IEditorOption; unfoldOnClickAfterEndOfLine: IEditorOption; fontFamily: IEditorOption; - fontInfo: IEditorOption; + fontInfo: IEditorOption; fontLigatures2: IEditorOption; fontSize: IEditorOption; fontWeight: IEditorOption; @@ -5366,7 +5366,7 @@ declare namespace monaco.editor { pasteAs: IEditorOption>>; parameterHints: IEditorOption>>; peekWidgetDefaultFocus: IEditorOption; - placeholder: IEditorOption; + placeholder: IEditorOption; definitionLinkOpensInPeek: IEditorOption; quickSuggestions: IEditorOption; quickSuggestionsDelay: IEditorOption; diff --git a/src/vs/platform/agentHost/common/agentService.ts b/src/vs/platform/agentHost/common/agentService.ts index 0099644da4a..3b30d5f8611 100644 --- a/src/vs/platform/agentHost/common/agentService.ts +++ b/src/vs/platform/agentHost/common/agentService.ts @@ -6,6 +6,8 @@ import { Event } from '../../../base/common/event.js'; import { URI } from '../../../base/common/uri.js'; import { createDecorator } from '../../instantiation/common/instantiation.js'; +import type { IActionEnvelope, INotification, ISessionAction } from './state/sessionActions.js'; +import type { IStateSnapshot } from './state/sessionProtocol.js'; // IPC contract between the renderer and the agent host utility process. // Defines all serializable event types, the IAgent provider interface, @@ -30,7 +32,7 @@ export interface IAgentSessionMetadata { readonly summary?: string; } -export type AgentProvider = 'copilot'; +export type AgentProvider = 'copilot' | 'local' | 'mock'; /** Metadata describing an agent backend, discovered over IPC. */ export interface IAgentDescriptor { @@ -247,7 +249,7 @@ export namespace AgentSession { */ export function provider(session: URI): AgentProvider | undefined { const scheme = session.scheme; - if (scheme === 'copilot') { + if (scheme === 'copilot' || scheme === 'local' || scheme === 'mock') { return scheme; } return undefined; @@ -283,6 +285,9 @@ export interface IAgent { /** Abort the current turn, stopping any in-flight processing. */ abortSession(session: URI): Promise; + /** Change the model for an existing session. */ + changeModel?(session: URI, model: string): Promise; + /** Respond to a pending permission request from the SDK. */ respondToPermissionRequest(requestId: string, approved: boolean): void; @@ -312,21 +317,25 @@ export const IAgentService = createDecorator('agentService'); /** * Service contract for communicating with the agent host process. Methods here * are proxied across MessagePort via `ProxyChannel`. + * + * State is synchronized via the subscribe/unsubscribe/dispatchAction protocol. + * Clients observe root state (agents, models) and session state via subscriptions, + * and mutate state by dispatching actions (e.g. session/turnStarted, session/turnCancelled). */ export interface IAgentService { readonly _serviceBrand: undefined; - /** Fires when the agent host streams progress for a session. */ - readonly onDidSessionProgress: Event; - /** Discover available agent backends from the agent host. */ listAgents(): Promise; /** Set the GitHub auth token used by the Copilot SDK. */ setAuthToken(token: string): Promise; - /** List available models from the agent. */ - listModels(): Promise; + /** + * Refresh the model list from all providers, publishing updated + * agents (with models) to root state via `root/agentsChanged`. + */ + refreshModels(): Promise; /** List all available sessions from the Copilot CLI. */ listSessions(): Promise; @@ -334,23 +343,43 @@ export interface IAgentService { /** Create a new session. Returns the session URI. */ createSession(config?: IAgentCreateSessionConfig): Promise; - /** Send a user message into an existing session. */ - sendMessage(session: URI, prompt: string, attachments?: IAgentAttachment[]): Promise; - - /** Retrieve all session events/messages for reconstruction, including tool invocations. */ - getSessionMessages(session: URI): Promise<(IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[]>; - /** Dispose a session in the agent host, freeing SDK resources. */ disposeSession(session: URI): Promise; - /** Abort the current turn in a session. */ - abortSession(session: URI): Promise; - - /** Respond to a pending permission request. */ - respondToPermissionRequest(requestId: string, approved: boolean): void; - /** Gracefully shut down all sessions and the underlying client. */ shutdown(): Promise; + + // ---- Protocol methods (sessions process protocol) ---------------------- + + /** + * Subscribe to state at the given URI. Returns a snapshot of the current + * state and the serverSeq at snapshot time. Subsequent actions for this + * resource arrive via {@link onDidAction}. + */ + subscribe(resource: URI): Promise; + + /** Unsubscribe from state updates for the given URI. */ + unsubscribe(resource: URI): void; + + /** + * Fires when the server applies an action to subscribable state. + * Clients use this alongside {@link subscribe} to keep their local + * state in sync. + */ + readonly onDidAction: Event; + + /** + * Fires when the server broadcasts an ephemeral notification + * (e.g. sessionAdded, sessionRemoved). + */ + readonly onDidNotification: Event; + + /** + * Dispatch a client-originated action to the server. The server applies + * it to state, triggers side effects, and echoes it back via + * {@link onDidAction} with the client's origin for reconciliation. + */ + dispatchAction(action: ISessionAction, clientId: string, clientSeq: number): void; } export const IAgentHostService = createDecorator('agentHostService'); @@ -361,6 +390,8 @@ export const IAgentHostService = createDecorator('agentHostSe */ export interface IAgentHostService extends IAgentService { + /** Unique identifier for this client window, used as the origin in action envelopes. */ + readonly clientId: string; readonly onAgentHostExit: Event; readonly onAgentHostStart: Event; diff --git a/src/vs/platform/agentHost/common/state/sessionActions.ts b/src/vs/platform/agentHost/common/state/sessionActions.ts index b331f86995d..8362d44564b 100644 --- a/src/vs/platform/agentHost/common/state/sessionActions.ts +++ b/src/vs/platform/agentHost/common/state/sessionActions.ts @@ -16,11 +16,9 @@ import { URI } from '../../../../base/common/uri.js'; import type { IAgentInfo, - ICompletedToolCall, IErrorInfo, IPermissionRequest, IResponsePart, - ISessionModelInfo, ISessionSummary, IToolCallState, IUsageInfo, @@ -59,18 +57,12 @@ export interface IActionOrigin { // ---- Root actions (server-only, mutate RootState) --------------------------- -export interface IModelsChangedAction { - readonly type: 'root/modelsChanged'; - readonly models: readonly ISessionModelInfo[]; -} - export interface IAgentsChangedAction { readonly type: 'root/agentsChanged'; readonly agents: readonly IAgentInfo[]; } export type IRootAction = - | IModelsChangedAction | IAgentsChangedAction; // ---- Session actions (mutate SessionState, scoped to a session URI) --------- @@ -126,7 +118,15 @@ export interface IToolCompleteAction extends ISessionActionBase { readonly type: 'session/toolComplete'; readonly turnId: string; readonly toolCallId: string; - readonly result: Omit; + readonly result: IToolCompleteResult; +} + +/** The data delivered with a tool completion event. */ +export interface IToolCompleteResult { + readonly success: boolean; + readonly pastTenseMessage: string; + readonly toolOutput?: string; + readonly error?: { readonly message: string; readonly code?: string }; } // -- Permissions -- @@ -189,6 +189,12 @@ export interface IReasoningAction extends ISessionActionBase { readonly content: string; } +/** Server-only. Dispatched when the session's model is changed. */ +export interface IModelChangedAction extends ISessionActionBase { + readonly type: 'session/modelChanged'; + readonly model: string; +} + export type ISessionAction = | ISessionReadyAction | ISessionCreationFailedAction @@ -204,7 +210,8 @@ export type ISessionAction = | ISessionErrorAction | ITitleChangedAction | IUsageAction - | IReasoningAction; + | IReasoningAction + | IModelChangedAction; // ---- Combined state action type --------------------------------------------- diff --git a/src/vs/platform/agentHost/common/state/sessionClientState.ts b/src/vs/platform/agentHost/common/state/sessionClientState.ts index 4a6bbfc3a86..3d26433161d 100644 --- a/src/vs/platform/agentHost/common/state/sessionClientState.ts +++ b/src/vs/platform/agentHost/common/state/sessionClientState.ts @@ -11,7 +11,7 @@ // or sends concurrent actions from other sources. // // This operates on two kinds of subscribable state: -// - Root state (agents, models) — server-only mutations, no write-ahead. +// - Root state (agents + their models) — server-only mutations, no write-ahead. // - Session state — mixed: some actions client-sendable (write-ahead), // others server-only. diff --git a/src/vs/platform/agentHost/common/state/sessionProtocol.ts b/src/vs/platform/agentHost/common/state/sessionProtocol.ts index 3e00e7d03a1..dde5a517479 100644 --- a/src/vs/platform/agentHost/common/state/sessionProtocol.ts +++ b/src/vs/platform/agentHost/common/state/sessionProtocol.ts @@ -3,199 +3,178 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -// Protocol messages for the sessions process client-server communication. -// See protocol.md -> Client-server protocol for the full design. +// Protocol messages using JSON-RPC 2.0 framing for the sessions process. +// See protocol.md for the full design. // -// These types define the wire format for handshake, URI-based subscription, -// commands, notifications, and reconnection. They are transport-agnostic — -// the actual transport (MessagePort, WebSocket, stdio) is plugged in separately. +// Client → Server messages are either: +// - Notifications (fire-and-forget): initialize, reconnect, unsubscribe, dispatchAction +// - Requests (expect a correlated response): subscribe, createSession, disposeSession, +// listSessions, fetchTurns, fetchContent +// +// Server → Client messages are either: +// - Notifications (pushed to clients): serverHello, reconnectResponse, action, notification +// - Responses (correlated to a client request by id) +import { hasKey } from '../../../../base/common/types.js'; import { URI } from '../../../../base/common/uri.js'; import type { IActionEnvelope, INotification, ISessionAction, IStateAction } from './sessionActions.js'; import type { IRootState, ISessionState, ISessionSummary } from './sessionState.js'; -// ---- Client → Server messages ----------------------------------------------- +// ---- JSON-RPC 2.0 base types ----------------------------------------------- -export interface IClientHello { - readonly type: 'clientHello'; - readonly protocolVersion: number; - readonly clientId: string; - /** Subscribe to these URIs as part of the handshake (saves a round-trip). */ - readonly initialSubscriptions?: readonly URI[]; +/** A JSON-RPC notification: has `method` but no `id`. */ +export interface IProtocolNotification { + readonly jsonrpc: '2.0'; + readonly method: string; + readonly params?: unknown; } -export interface IClientReconnect { - readonly type: 'clientReconnect'; - readonly clientId: string; - readonly lastSeenServerSeq: number; - /** URIs the client was subscribed to before disconnection. */ - readonly subscriptions: readonly URI[]; +/** A JSON-RPC request: has both `method` and `id`. */ +export interface IProtocolRequest { + readonly jsonrpc: '2.0'; + readonly id: number; + readonly method: string; + readonly params?: unknown; } -export interface ISubscribe { - readonly type: 'subscribe'; - /** URI to subscribe to (e.g. `agenthost:root` or `copilot:/`). */ - readonly resource: URI; +/** A JSON-RPC success response. */ +export interface IJsonRpcSuccessResponse { + readonly jsonrpc: '2.0'; + readonly id: number; + readonly result: unknown; } -export interface IUnsubscribe { - readonly type: 'unsubscribe'; - readonly resource: URI; +/** A JSON-RPC error response. */ +export interface IJsonRpcErrorResponse { + readonly jsonrpc: '2.0'; + readonly id: number; + readonly error: { + readonly code: number; + readonly message: string; + readonly data?: unknown; + }; } -/** - * A client-dispatched action. The server applies it to state and - * reacts with side effects (e.g., starting agent processing). - * Used for write-ahead actions like turnStarted, turnCancelled, - * permissionResolved. - */ -export interface IClientAction { - readonly type: 'action'; - readonly clientSeq: number; - readonly action: ISessionAction; +export type IJsonRpcResponse = IJsonRpcSuccessResponse | IJsonRpcErrorResponse; + +/** Any message that flows over the protocol transport. */ +export type IProtocolMessage = IProtocolNotification | IProtocolRequest | IJsonRpcResponse; + +// ---- Type guards ----------------------------------------------------------- + +export function isJsonRpcRequest(msg: IProtocolMessage): msg is IProtocolRequest { + return hasKey(msg, { id: true, method: true }); } -/** - * A command from the client requesting an imperative operation - * that doesn't map directly to a single state action. - */ -export interface IClientCommand { - readonly type: 'command'; - readonly command: ISessionCommand; +export function isJsonRpcNotification(msg: IProtocolMessage): msg is IProtocolNotification { + return hasKey(msg, { method: true }) && !hasKey(msg, { id: true }); } -export type IClientMessage = - | IClientHello - | IClientReconnect - | ISubscribe - | IUnsubscribe - | IClientAction - | IClientCommand; - -// ---- Commands (embedded in IClientCommand) ---------------------------------- - -export interface ICreateSessionCommand { - readonly type: 'createSession'; - /** URI the client has chosen for this session (client picks the ID). */ - readonly session: URI; - readonly provider?: string; - readonly model?: string; - readonly workingDirectory?: string; +export function isJsonRpcResponse(msg: IProtocolMessage): msg is IJsonRpcResponse { + return hasKey(msg, { id: true }) && !hasKey(msg, { method: true }); } -export interface IDisposeSessionCommand { - readonly type: 'disposeSession'; - readonly session: URI; -} +// ---- JSON-RPC error codes --------------------------------------------------- -export interface IFetchContentCommand { - readonly type: 'fetchContent'; - readonly uri: URI; -} +export const JSON_RPC_INTERNAL_ERROR = -32603; -export interface IFetchTurnsCommand { - readonly type: 'fetchTurns'; - readonly session: URI; - readonly startTurn: number; - readonly count: number; -} +// ---- Shared data types ------------------------------------------------------ -export interface IListSessionsCommand { - readonly type: 'listSessions'; -} - -export type ISessionCommand = - | ICreateSessionCommand - | IDisposeSessionCommand - | IFetchContentCommand - | IFetchTurnsCommand - | IListSessionsCommand; - -// ---- Server → Client messages ----------------------------------------------- - -export interface IServerHello { - readonly type: 'serverHello'; - readonly protocolVersion: number; - readonly serverSeq: number; - /** Snapshots for each URI in the client's `initialSubscriptions`. */ - readonly snapshots: readonly IStateSnapshot[]; -} - -/** - * Response to a subscribe request. Contains the state snapshot and - * the server sequence at snapshot time. The client processes subsequent - * actions with serverSeq > fromSeq. - */ +/** State snapshot returned by subscribe and included in handshake/reconnect. */ export interface IStateSnapshot { - readonly type: 'stateSnapshot'; readonly resource: URI; readonly state: IRootState | ISessionState; readonly fromSeq: number; } -/** - * A state-changing action broadcast to subscribed clients. - */ -export interface IActionMessage { - readonly type: 'action'; - readonly envelope: IActionEnvelope; +// ---- Client → Server: Notification params ----------------------------------- + +export interface IInitializeParams { + readonly protocolVersion: number; + readonly clientId: string; + readonly initialSubscriptions?: readonly URI[]; } -/** - * An ephemeral notification broadcast to all connected clients. - * Not stored in state, not replayed on reconnect. - */ -export interface INotificationMessage { - readonly type: 'notification'; - readonly notification: INotification; +export interface IReconnectParams { + readonly clientId: string; + readonly lastSeenServerSeq: number; + readonly subscriptions: readonly URI[]; } -/** - * Response to a fetchContent command. - */ -export interface IContentResponse { - readonly type: 'contentResponse'; - readonly uri: URI; - readonly data: string; // base64-encoded for binary safety over JSON - readonly mimeType?: string; +export interface IUnsubscribeParams { + readonly resource: URI; } -/** - * Response to a fetchTurns command. - */ -export interface ITurnsResponse { - readonly type: 'turnsResponse'; +export interface IDispatchActionParams { + readonly clientSeq: number; + readonly action: ISessionAction; +} + +// ---- Client → Server: Request params and results ---------------------------- + +export interface ISubscribeParams { + readonly resource: URI; +} +// Result: IStateSnapshot + +export interface ICreateSessionParams { + readonly session: URI; + readonly provider?: string; + readonly model?: string; + readonly workingDirectory?: string; +} +// Result: void (null) + +export interface IDisposeSessionParams { + readonly session: URI; +} +// Result: void (null) + +// listSessions: no params +export interface IListSessionsResult { + readonly sessions: readonly ISessionSummary[]; +} + +export interface IFetchTurnsParams { + readonly session: URI; + readonly startTurn: number; + readonly count: number; +} + +export interface IFetchTurnsResult { readonly session: URI; readonly startTurn: number; readonly turns: ISessionState['turns']; readonly totalTurns: number; } -/** - * Response to a listSessions command. - */ -export interface IListSessionsResponse { - readonly type: 'listSessionsResponse'; - readonly sessions: readonly ISessionSummary[]; +export interface IFetchContentParams { + readonly uri: URI; } -/** - * Sent on reconnection. Contains fresh snapshots for all previously - * subscribed URIs. Notifications are NOT replayed — the client should - * re-fetch the session list. - */ -export interface IReconnectResponse { - readonly type: 'reconnectResponse'; +export interface IFetchContentResult { + readonly uri: URI; + readonly data: string; // base64-encoded for binary safety + readonly mimeType?: string; +} + +// ---- Server → Client: Notification params ----------------------------------- + +export interface IServerHelloParams { + readonly protocolVersion: number; readonly serverSeq: number; readonly snapshots: readonly IStateSnapshot[]; } -export type IServerMessage = - | IServerHello - | IStateSnapshot - | IActionMessage - | INotificationMessage - | IContentResponse - | ITurnsResponse - | IListSessionsResponse - | IReconnectResponse; +export interface IReconnectResponseParams { + readonly serverSeq: number; + readonly snapshots: readonly IStateSnapshot[]; +} + +export interface IActionBroadcastParams { + readonly envelope: IActionEnvelope; +} + +export interface INotificationBroadcastParams { + readonly notification: INotification; +} diff --git a/src/vs/platform/agentHost/common/state/sessionReducers.ts b/src/vs/platform/agentHost/common/state/sessionReducers.ts index 4f016560a5d..df2103809ff 100644 --- a/src/vs/platform/agentHost/common/state/sessionReducers.ts +++ b/src/vs/platform/agentHost/common/state/sessionReducers.ts @@ -17,6 +17,7 @@ import type { IRootAction, ISessionAction } from './sessionActions.js'; import { type ICompletedToolCall, + type IErrorInfo, type IRootState, type ISessionState, type IToolCallState, @@ -36,9 +37,6 @@ import { */ export function rootReducer(state: IRootState, action: IRootAction): IRootState { switch (action.type) { - case 'root/modelsChanged': { - return { ...state, models: action.models }; - } case 'root/agentsChanged': { return { ...state, agents: action.agents }; } @@ -118,6 +116,9 @@ export function sessionReducer(state: ISessionState, action: ISessionAction): IS toolCalls.set(action.toolCallId, { ...toolCall, status: action.result.success ? ToolCallStatus.Completed : ToolCallStatus.Failed, + pastTenseMessage: action.result.pastTenseMessage, + toolOutput: action.result.toolOutput, + error: action.result.error, }); return { ...state, @@ -161,7 +162,9 @@ export function sessionReducer(state: ISessionState, action: ISessionAction): IS const mutable = new Map(toolCalls); mutable.set(resolved.toolCallId, { ...toolCall, - status: action.approved ? ToolCallStatus.Running : ToolCallStatus.Failed, + status: action.approved ? ToolCallStatus.Running : ToolCallStatus.Cancelled, + confirmed: action.approved ? 'user-action' : 'denied', + cancellationReason: action.approved ? undefined : 'denied', }); toolCalls = mutable; } @@ -178,18 +181,31 @@ export function sessionReducer(state: ISessionState, action: ISessionAction): IS return finalizeTurn(state, action.turnId, TurnState.Cancelled); } case 'session/error': { - return finalizeTurn(state, action.turnId, TurnState.Error); + return finalizeTurn(state, action.turnId, TurnState.Error, action.error); } case 'session/titleChanged': { return { ...state, - summary: { ...state.summary, title: action.title }, + summary: { ...state.summary, title: action.title, modifiedAt: Date.now() }, + }; + } + case 'session/modelChanged': { + return { + ...state, + summary: { ...state.summary, model: action.model, modifiedAt: Date.now() }, }; } case 'session/usage': { - // Usage is informational; stored on the active turn for now, - // then captured on the finalized Turn. - return state; + if (!state.activeTurn || state.activeTurn.id !== action.turnId) { + return state; + } + return { + ...state, + activeTurn: { + ...state.activeTurn, + usage: action.usage, + }, + }; } case 'session/reasoning': { if (!state.activeTurn || state.activeTurn.id !== action.turnId) { @@ -211,7 +227,7 @@ export function sessionReducer(state: ISessionState, action: ISessionAction): IS /** * Moves the active turn into the completed turns array and clears `activeTurn`. */ -function finalizeTurn(state: ISessionState, turnId: string, turnState: TurnState): ISessionState { +function finalizeTurn(state: ISessionState, turnId: string, turnState: TurnState, error?: IErrorInfo): ISessionState { if (!state.activeTurn || state.activeTurn.id !== turnId) { return state; } @@ -223,25 +239,32 @@ function finalizeTurn(state: ISessionState, turnId: string, turnState: TurnState toolCallId: tc.toolCallId, toolName: tc.toolName, displayName: tc.displayName, + invocationMessage: tc.invocationMessage, success: tc.status === ToolCallStatus.Completed, - pastTenseMessage: tc.invocationMessage, - toolOutput: tc.toolInput, + pastTenseMessage: tc.pastTenseMessage ?? tc.invocationMessage, + toolInput: tc.toolInput, + toolKind: tc.toolKind, + language: tc.language, + toolOutput: tc.toolOutput, + error: tc.error, }); } const finalizedTurn: ITurn = { id: active.id, userMessage: active.userMessage, + responseText: active.streamingText, responseParts: active.responseParts, toolCalls: completedToolCalls, - usage: undefined, + usage: active.usage, state: turnState, + error, }; return { ...state, turns: [...state.turns, finalizedTurn], activeTurn: undefined, - summary: { ...state.summary, status: SessionStatus.Idle }, + summary: { ...state.summary, status: SessionStatus.Idle, modifiedAt: Date.now() }, }; } diff --git a/src/vs/platform/agentHost/common/state/sessionState.ts b/src/vs/platform/agentHost/common/state/sessionState.ts index 212ea73a159..7ef6d763454 100644 --- a/src/vs/platform/agentHost/common/state/sessionState.ts +++ b/src/vs/platform/agentHost/common/state/sessionState.ts @@ -37,6 +37,7 @@ export interface ISessionSummary { readonly status: SessionStatus; readonly createdAt: number; readonly modifiedAt: number; + readonly model?: string; } // ---- Model info ------------------------------------------------------------- @@ -45,6 +46,9 @@ export interface ISessionModelInfo { readonly id: string; readonly provider: AgentProvider; readonly name: string; + readonly maxContextWindow?: number; + readonly supportsVision?: boolean; + readonly policyState?: 'enabled' | 'disabled' | 'unconfigured'; } // ---- Root state (subscribable at ROOT_STATE_URI) ---------------------------- @@ -56,13 +60,13 @@ export interface ISessionModelInfo { */ export interface IRootState { readonly agents: readonly IAgentInfo[]; - readonly models: readonly ISessionModelInfo[]; } export interface IAgentInfo { readonly provider: AgentProvider; readonly displayName: string; readonly description: string; + readonly models: readonly ISessionModelInfo[]; } // ---- Session lifecycle ------------------------------------------------------ @@ -109,10 +113,14 @@ export interface IMessageAttachment { export interface ITurn { readonly id: string; readonly userMessage: IUserMessage; + /** The final assistant response text (captured from streamingText on turn completion). */ + readonly responseText: string; readonly responseParts: readonly IResponsePart[]; readonly toolCalls: readonly ICompletedToolCall[]; readonly usage: IUsageInfo | undefined; readonly state: TurnState; + /** Error info if the turn ended with {@link TurnState.Error}. */ + readonly error?: IErrorInfo; } export const enum TurnState { @@ -132,6 +140,7 @@ export interface IActiveTurn { readonly toolCalls: ReadonlyMap; readonly pendingPermissions: ReadonlyMap; readonly reasoning: string; + readonly usage: IUsageInfo | undefined; } // ---- Response parts --------------------------------------------------------- @@ -162,12 +171,22 @@ export type IResponsePart = IMarkdownResponsePart | IContentRef; // ---- Tool calls ------------------------------------------------------------- export const enum ToolCallStatus { + /** Tool is actively executing. */ Running = 'running', + /** Waiting for user to approve before execution. */ PendingPermission = 'pending-permission', + /** Tool finished successfully. */ Completed = 'completed', + /** Tool failed with an error. */ Failed = 'failed', + /** Tool was denied or skipped by the user. */ + Cancelled = 'cancelled', } +/** + * Represents the full lifecycle state of a tool invocation within an active turn. + * Modeled after {@link IChatToolInvocation.State} to enable direct mapping to the chat UI. + */ export interface IToolCallState { readonly toolCallId: string; readonly toolName: string; @@ -176,16 +195,34 @@ export interface IToolCallState { readonly toolInput?: string; readonly toolKind?: 'terminal'; readonly language?: string; + readonly toolArguments?: string; readonly status: ToolCallStatus; + /** Parsed tool parameters (from toolArguments). */ + readonly parameters?: unknown; + /** How the tool was confirmed before execution (set after PendingPermission → Running). */ + readonly confirmed?: 'not-needed' | 'user-action' | 'setting' | 'denied' | 'skipped'; + /** Set when status transitions to Completed or Failed. */ + readonly pastTenseMessage?: string; + /** Set when status transitions to Completed or Failed. */ + readonly toolOutput?: string; + /** Set when status transitions to Failed. */ + readonly error?: { readonly message: string; readonly code?: string }; + /** Why the tool was cancelled (set when status is Cancelled). */ + readonly cancellationReason?: 'denied' | 'skipped'; } export interface ICompletedToolCall { readonly toolCallId: string; readonly toolName: string; readonly displayName: string; + readonly invocationMessage: string; readonly success: boolean; readonly pastTenseMessage: string; + readonly toolInput?: string; + readonly toolKind?: 'terminal'; + readonly language?: string; readonly toolOutput?: string; + readonly error?: { readonly message: string; readonly code?: string }; } // ---- Permission requests ---------------------------------------------------- @@ -197,6 +234,9 @@ export interface IPermissionRequest { readonly path?: string; readonly fullCommandText?: string; readonly intention?: string; + readonly serverName?: string; + readonly toolName?: string; + readonly rawRequest?: string; } // ---- Usage info ------------------------------------------------------------- @@ -221,7 +261,6 @@ export interface IErrorInfo { export function createRootState(): IRootState { return { agents: [], - models: [], }; } @@ -243,5 +282,6 @@ export function createActiveTurn(id: string, userMessage: IUserMessage): IActive toolCalls: new Map(), pendingPermissions: new Map(), reasoning: '', + usage: undefined, }; } diff --git a/src/vs/platform/agentHost/common/state/sessionTransport.ts b/src/vs/platform/agentHost/common/state/sessionTransport.ts new file mode 100644 index 00000000000..a876f59de88 --- /dev/null +++ b/src/vs/platform/agentHost/common/state/sessionTransport.ts @@ -0,0 +1,42 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +// Transport abstraction for the sessions process protocol. +// See protocol.md -> Client-server protocol for the full design. +// +// The transport is pluggable — the same protocol runs over MessagePort +// (ProxyChannel), WebSocket, or stdio. This module defines the contract; +// concrete implementations live in platform-specific folders. + +import { Event } from '../../../../base/common/event.js'; +import { IDisposable } from '../../../../base/common/lifecycle.js'; +import type { IProtocolMessage } from './sessionProtocol.js'; + +/** + * A bidirectional transport for protocol messages. Implementations handle + * serialization, framing, and connection management. + */ +export interface IProtocolTransport extends IDisposable { + /** Fires when a message is received from the remote end. */ + readonly onMessage: Event; + + /** Fires when the transport connection closes. */ + readonly onClose: Event; + + /** Send a message to the remote end. */ + send(message: IProtocolMessage): void; +} + +/** + * Server-side transport that accepts multiple client connections. + * Each connected client gets its own {@link IProtocolTransport}. + */ +export interface IProtocolServer extends IDisposable { + /** Fires when a new client connects. */ + readonly onConnection: Event; + + /** The port or address the server is listening on. */ + readonly address: string | undefined; +} diff --git a/src/vs/platform/agentHost/common/state/versions/v1.ts b/src/vs/platform/agentHost/common/state/versions/v1.ts index 8f618cb0a69..78a66215c64 100644 --- a/src/vs/platform/agentHost/common/state/versions/v1.ts +++ b/src/vs/platform/agentHost/common/state/versions/v1.ts @@ -17,19 +17,22 @@ import type { AgentProvider } from '../../agentService.js'; export interface IV1_RootState { readonly agents: readonly IV1_AgentInfo[]; - readonly models: readonly IV1_SessionModelInfo[]; } export interface IV1_AgentInfo { readonly provider: AgentProvider; readonly displayName: string; readonly description: string; + readonly models: readonly IV1_SessionModelInfo[]; } export interface IV1_SessionModelInfo { readonly id: string; readonly provider: AgentProvider; readonly name: string; + readonly maxContextWindow?: number; + readonly supportsVision?: boolean; + readonly policyState?: 'enabled' | 'disabled' | 'unconfigured'; } export interface IV1_SessionSummary { @@ -39,6 +42,7 @@ export interface IV1_SessionSummary { readonly status: 'idle' | 'in-progress' | 'error'; readonly createdAt: number; readonly modifiedAt: number; + readonly model?: string; } export interface IV1_SessionState { @@ -63,10 +67,12 @@ export interface IV1_MessageAttachment { export interface IV1_Turn { readonly id: string; readonly userMessage: IV1_UserMessage; + readonly responseText: string; readonly responseParts: readonly IV1_ResponsePart[]; readonly toolCalls: readonly IV1_CompletedToolCall[]; readonly usage: IV1_UsageInfo | undefined; readonly state: 'complete' | 'cancelled' | 'error'; + readonly error?: IV1_ErrorInfo; } export interface IV1_ActiveTurn { @@ -77,6 +83,7 @@ export interface IV1_ActiveTurn { readonly toolCalls: ReadonlyMap; readonly pendingPermissions: ReadonlyMap; readonly reasoning: string; + readonly usage: IV1_UsageInfo | undefined; } export interface IV1_MarkdownResponsePart { @@ -101,16 +108,28 @@ export interface IV1_ToolCallState { readonly toolInput?: string; readonly toolKind?: 'terminal'; readonly language?: string; - readonly status: 'running' | 'pending-permission' | 'completed' | 'failed'; + readonly toolArguments?: string; + readonly status: 'running' | 'pending-permission' | 'completed' | 'failed' | 'cancelled'; + readonly parameters?: unknown; + readonly confirmed?: 'not-needed' | 'user-action' | 'setting' | 'denied' | 'skipped'; + readonly pastTenseMessage?: string; + readonly toolOutput?: string; + readonly error?: { readonly message: string; readonly code?: string }; + readonly cancellationReason?: 'denied' | 'skipped'; } export interface IV1_CompletedToolCall { readonly toolCallId: string; readonly toolName: string; readonly displayName: string; + readonly invocationMessage: string; readonly success: boolean; readonly pastTenseMessage: string; + readonly toolInput?: string; + readonly toolKind?: 'terminal'; + readonly language?: string; readonly toolOutput?: string; + readonly error?: { readonly message: string; readonly code?: string }; } export interface IV1_PermissionRequest { @@ -120,6 +139,9 @@ export interface IV1_PermissionRequest { readonly path?: string; readonly fullCommandText?: string; readonly intention?: string; + readonly serverName?: string; + readonly toolName?: string; + readonly rawRequest?: string; } export interface IV1_UsageInfo { @@ -141,11 +163,6 @@ interface IV1_SessionActionBase { readonly session: URI; } -export interface IV1_ModelsChangedAction { - readonly type: 'root/modelsChanged'; - readonly models: readonly IV1_SessionModelInfo[]; -} - export interface IV1_AgentsChangedAction { readonly type: 'root/agentsChanged'; readonly agents: readonly IV1_AgentInfo[]; @@ -188,7 +205,14 @@ export interface IV1_ToolCompleteAction extends IV1_SessionActionBase { readonly type: 'session/toolComplete'; readonly turnId: string; readonly toolCallId: string; - readonly result: Omit; + readonly result: IV1_ToolCompleteResult; +} + +export interface IV1_ToolCompleteResult { + readonly success: boolean; + readonly pastTenseMessage: string; + readonly toolOutput?: string; + readonly error?: { readonly message: string; readonly code?: string }; } export interface IV1_PermissionRequestAction extends IV1_SessionActionBase { @@ -237,8 +261,12 @@ export interface IV1_ReasoningAction extends IV1_SessionActionBase { readonly content: string; } +export interface IV1_ModelChangedAction extends IV1_SessionActionBase { + readonly type: 'session/modelChanged'; + readonly model: string; +} + export type IV1_RootAction = - | IV1_ModelsChangedAction | IV1_AgentsChangedAction; export type IV1_SessionAction = @@ -256,7 +284,8 @@ export type IV1_SessionAction = | IV1_SessionErrorAction | IV1_TitleChangedAction | IV1_UsageAction - | IV1_ReasoningAction; + | IV1_ReasoningAction + | IV1_ModelChangedAction; export type IV1_StateAction = IV1_RootAction | IV1_SessionAction; diff --git a/src/vs/platform/agentHost/common/state/versions/versionRegistry.ts b/src/vs/platform/agentHost/common/state/versions/versionRegistry.ts index c1ba77edc2b..c650d977a9a 100644 --- a/src/vs/platform/agentHost/common/state/versions/versionRegistry.ts +++ b/src/vs/platform/agentHost/common/state/versions/versionRegistry.ts @@ -9,7 +9,7 @@ import type { IAgentsChangedAction, IDeltaAction, - IModelsChangedAction, + IModelChangedAction, INotification, IPermissionRequestAction, IPermissionResolvedAction, @@ -58,7 +58,7 @@ import type { IV1_ErrorInfo, IV1_MarkdownResponsePart, IV1_MessageAttachment, - IV1_ModelsChangedAction, + IV1_ModelChangedAction, IV1_PermissionRequest, IV1_PermissionRequestAction, IV1_PermissionResolvedAction, @@ -134,7 +134,6 @@ type _v1_ErrorInfo = AssertCompatible; // -- v1 action compatibility -- -type _v1_ModelsChanged = AssertCompatible; type _v1_AgentsChanged = AssertCompatible; type _v1_SessionReady = AssertCompatible; type _v1_CreationFailed = AssertCompatible; @@ -151,6 +150,7 @@ type _v1_SessionError = AssertCompatible; type _v1_Usage = AssertCompatible; type _v1_Reasoning = AssertCompatible; +type _v1_ModelChanged = AssertCompatible; // Suppress unused-variable warnings for compile-time-only checks. void (0 as unknown as @@ -159,11 +159,11 @@ void (0 as unknown as _v1_ActiveTurn & _v1_MarkdownResponsePart & _v1_ContentRef & _v1_ToolCallState & _v1_CompletedToolCall & _v1_PermissionRequest & _v1_UsageInfo & _v1_ErrorInfo & - _v1_ModelsChanged & _v1_AgentsChanged & _v1_SessionReady & _v1_CreationFailed & + _v1_AgentsChanged & _v1_SessionReady & _v1_CreationFailed & _v1_TurnStarted & _v1_Delta & _v1_ResponsePart & _v1_ToolStart & _v1_ToolComplete & _v1_PermissionRequestAction & _v1_PermissionResolved & _v1_TurnComplete & _v1_TurnCancelled & _v1_SessionError & _v1_TitleChanged & - _v1_Usage & _v1_Reasoning + _v1_Usage & _v1_Reasoning & _v1_ModelChanged ); // ---- Runtime action → version map ------------------------------------------- @@ -178,7 +178,6 @@ void (0 as unknown as /** Maps every action type string to the protocol version that introduced it. */ export const ACTION_INTRODUCED_IN: { readonly [K in IStateAction['type']]: number } = { // Root actions (v1) - 'root/modelsChanged': 1, 'root/agentsChanged': 1, // Session lifecycle (v1) 'session/ready': 1, @@ -201,6 +200,7 @@ export const ACTION_INTRODUCED_IN: { readonly [K in IStateAction['type']]: numbe 'session/titleChanged': 1, 'session/usage': 1, 'session/reasoning': 1, + 'session/modelChanged': 1, }; /** Maps every notification type string to the protocol version that introduced it. */ @@ -233,13 +233,14 @@ export function isNotificationKnownToVersion(notification: INotification, client // When you add a new protocol version, define its additions and extend the map. /** Action types introduced in v1. */ -type IRootAction_v1 = IV1_ModelsChangedAction | IV1_AgentsChangedAction; +type IRootAction_v1 = IV1_AgentsChangedAction; type ISessionAction_v1 = IV1_SessionReadyAction | IV1_SessionCreationFailedAction | IV1_TurnStartedAction | IV1_DeltaAction | IV1_ResponsePartAction | IV1_ToolStartAction | IV1_ToolCompleteAction | IV1_PermissionRequestAction | IV1_PermissionResolvedAction | IV1_TurnCompleteAction | IV1_TurnCancelledAction | IV1_SessionErrorAction - | IV1_TitleChangedAction | IV1_UsageAction | IV1_ReasoningAction; + | IV1_TitleChangedAction | IV1_UsageAction | IV1_ReasoningAction + | IV1_ModelChangedAction; /** * Maps protocol versions to their cumulative action type unions. diff --git a/src/vs/platform/agentHost/electron-browser/agentHostService.ts b/src/vs/platform/agentHost/electron-browser/agentHostService.ts index e2728cd1b2e..ae16987a056 100644 --- a/src/vs/platform/agentHost/electron-browser/agentHostService.ts +++ b/src/vs/platform/agentHost/electron-browser/agentHostService.ts @@ -6,14 +6,18 @@ import { DeferredPromise } from '../../../base/common/async.js'; import { Emitter } from '../../../base/common/event.js'; import { Disposable, DisposableStore } from '../../../base/common/lifecycle.js'; -import { URI } from '../../../base/common/uri.js'; +import { generateUuid } from '../../../base/common/uuid.js'; import { getDelayedChannel, ProxyChannel } from '../../../base/parts/ipc/common/ipc.js'; import { Client as MessagePortClient } from '../../../base/parts/ipc/common/ipc.mp.js'; import { acquirePort } from '../../../base/parts/ipc/electron-browser/ipc.mp.js'; import { InstantiationType, registerSingleton } from '../../instantiation/common/extensions.js'; import { IConfigurationService } from '../../configuration/common/configuration.js'; import { ILogService } from '../../log/common/log.js'; -import { AgentHostEnabledSettingId, AgentHostIpcChannels, IAgentAttachment, IAgentCreateSessionConfig, IAgentDescriptor, IAgentHostService, IAgentMessageEvent, IAgentModelInfo, IAgentProgressEvent, IAgentService, IAgentSessionMetadata, IAgentToolCompleteEvent, IAgentToolStartEvent } from '../common/agentService.js'; +import { AgentHostEnabledSettingId, AgentHostIpcChannels, IAgentCreateSessionConfig, IAgentDescriptor, IAgentHostService, IAgentService, IAgentSessionMetadata } from '../common/agentService.js'; +import type { IActionEnvelope, INotification, ISessionAction } from '../common/state/sessionActions.js'; +import type { IStateSnapshot } from '../common/state/sessionProtocol.js'; +import { revive } from '../../../base/common/marshalling.js'; +import { URI } from '../../../base/common/uri.js'; /** * Renderer-side implementation of {@link IAgentHostService} that connects @@ -24,6 +28,9 @@ import { AgentHostEnabledSettingId, AgentHostIpcChannels, IAgentAttachment, IAge class AgentHostServiceClient extends Disposable implements IAgentHostService { declare readonly _serviceBrand: undefined; + /** Unique identifier for this window, used in action envelope origin tracking. */ + readonly clientId = generateUuid(); + private readonly _clientEventually = new DeferredPromise(); private readonly _proxy: IAgentService; @@ -32,8 +39,11 @@ class AgentHostServiceClient extends Disposable implements IAgentHostService { private readonly _onAgentHostStart = this._register(new Emitter()); readonly onAgentHostStart = this._onAgentHostStart.event; - private readonly _onDidSessionProgress = this._register(new Emitter()); - readonly onDidSessionProgress = this._onDidSessionProgress.event; + private readonly _onDidAction = this._register(new Emitter()); + readonly onDidAction = this._onDidAction.event; + + private readonly _onDidNotification = this._register(new Emitter()); + readonly onDidNotification = this._onDidNotification.event; constructor( @ILogService private readonly _logService: ILogService, @@ -61,9 +71,11 @@ class AgentHostServiceClient extends Disposable implements IAgentHostService { const client = store.add(new MessagePortClient(port, `agentHost:window`)); this._clientEventually.complete(client); - store.add(this._proxy.onDidSessionProgress(e => { - // Events from ProxyChannel don't auto-revive nested URIs -- revive the session URI - this._onDidSessionProgress.fire({ ...e, session: URI.revive(e.session) }); + store.add(this._proxy.onDidAction(e => { + this._onDidAction.fire(revive(e)); + })); + store.add(this._proxy.onDidNotification(e => { + this._onDidNotification.fire(revive(e)); })); this._logService.info('[AgentHost:renderer] Direct MessagePort connection established'); this._onAgentHostStart.fire(); @@ -77,8 +89,8 @@ class AgentHostServiceClient extends Disposable implements IAgentHostService { listAgents(): Promise { return this._proxy.listAgents(); } - listModels(): Promise { - return this._proxy.listModels(); + refreshModels(): Promise { + return this._proxy.refreshModels(); } listSessions(): Promise { return this._proxy.listSessions(); @@ -86,24 +98,21 @@ class AgentHostServiceClient extends Disposable implements IAgentHostService { createSession(config?: IAgentCreateSessionConfig): Promise { return this._proxy.createSession(config); } - sendMessage(session: URI, prompt: string, attachments?: IAgentAttachment[]): Promise { - return this._proxy.sendMessage(session, prompt, attachments); - } - getSessionMessages(session: URI): Promise<(IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[]> { - return this._proxy.getSessionMessages(session); - } disposeSession(session: URI): Promise { return this._proxy.disposeSession(session); } - abortSession(session: URI): Promise { - return this._proxy.abortSession(session); - } - respondToPermissionRequest(requestId: string, approved: boolean): void { - this._proxy.respondToPermissionRequest(requestId, approved); - } shutdown(): Promise { return this._proxy.shutdown(); } + subscribe(resource: URI): Promise { + return this._proxy.subscribe(resource); + } + unsubscribe(resource: URI): void { + this._proxy.unsubscribe(resource); + } + dispatchAction(action: ISessionAction, clientId: string, clientSeq: number): void { + this._proxy.dispatchAction(action, clientId, clientSeq); + } async restartAgentHost(): Promise { // Restart is handled by the main process side } diff --git a/src/vs/platform/agentHost/node/agentEventMapper.ts b/src/vs/platform/agentHost/node/agentEventMapper.ts new file mode 100644 index 00000000000..c7639b180e6 --- /dev/null +++ b/src/vs/platform/agentHost/node/agentEventMapper.ts @@ -0,0 +1,160 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import type { + IAgentProgressEvent, + IAgentToolStartEvent, + IAgentToolCompleteEvent, + IAgentPermissionRequestEvent, + IAgentErrorEvent, + IAgentReasoningEvent, + IAgentUsageEvent, + IAgentDeltaEvent, + IAgentTitleChangedEvent, +} from '../common/agentService.js'; +import type { + ISessionAction, + IDeltaAction, + IToolStartAction, + IToolCompleteAction, + ITurnCompleteAction, + ISessionErrorAction, + IUsageAction, + ITitleChangedAction, + IPermissionRequestAction, + IReasoningAction, +} from '../common/state/sessionActions.js'; +import { ToolCallStatus } from '../common/state/sessionState.js'; +import { URI } from '../../../base/common/uri.js'; + +/** + * Maps a flat {@link IAgentProgressEvent} from the agent host into + * a protocol {@link ISessionAction} suitable for dispatch to the reducer. + * Returns `undefined` for events that have no corresponding action. + */ +export function mapProgressEventToAction(event: IAgentProgressEvent, session: URI, turnId: string): ISessionAction | undefined { + switch (event.type) { + case 'delta': + return { + type: 'session/delta', + session, + turnId, + content: (event as IAgentDeltaEvent).content, + } satisfies IDeltaAction; + + case 'tool_start': { + const e = event as IAgentToolStartEvent; + return { + type: 'session/toolStart', + session, + turnId, + toolCall: { + toolCallId: e.toolCallId, + toolName: e.toolName, + displayName: e.displayName, + invocationMessage: e.invocationMessage, + toolInput: e.toolInput, + toolKind: e.toolKind, + language: e.language, + toolArguments: e.toolArguments, + status: ToolCallStatus.Running, + }, + } satisfies IToolStartAction; + } + + case 'tool_complete': { + const e = event as IAgentToolCompleteEvent; + return { + type: 'session/toolComplete', + session, + turnId, + toolCallId: e.toolCallId, + result: { + success: e.success, + pastTenseMessage: e.pastTenseMessage, + toolOutput: e.toolOutput, + error: e.error, + }, + } satisfies IToolCompleteAction; + } + + case 'idle': + return { + type: 'session/turnComplete', + session, + turnId, + } satisfies ITurnCompleteAction; + + case 'error': { + const e = event as IAgentErrorEvent; + return { + type: 'session/error', + session, + turnId, + error: { + errorType: e.errorType, + message: e.message, + stack: e.stack, + }, + } satisfies ISessionErrorAction; + } + + case 'usage': { + const e = event as IAgentUsageEvent; + return { + type: 'session/usage', + session, + turnId, + usage: { + inputTokens: e.inputTokens, + outputTokens: e.outputTokens, + model: e.model, + cacheReadTokens: e.cacheReadTokens, + }, + } satisfies IUsageAction; + } + + case 'title_changed': + return { + type: 'session/titleChanged', + session, + title: (event as IAgentTitleChangedEvent).title, + } satisfies ITitleChangedAction; + + case 'permission_request': { + const e = event as IAgentPermissionRequestEvent; + return { + type: 'session/permissionRequest', + session, + turnId, + request: { + requestId: e.requestId, + permissionKind: e.permissionKind, + toolCallId: e.toolCallId, + path: e.path, + fullCommandText: e.fullCommandText, + intention: e.intention, + serverName: e.serverName, + toolName: e.toolName, + rawRequest: e.rawRequest, + }, + } satisfies IPermissionRequestAction; + } + + case 'reasoning': + return { + type: 'session/reasoning', + session, + turnId, + content: (event as IAgentReasoningEvent).content, + } satisfies IReasoningAction; + + case 'message': + return undefined; + + default: + return undefined; + } +} diff --git a/src/vs/platform/agentHost/node/agentHostServerMain.ts b/src/vs/platform/agentHost/node/agentHostServerMain.ts new file mode 100644 index 00000000000..a265d4a9dc4 --- /dev/null +++ b/src/vs/platform/agentHost/node/agentHostServerMain.ts @@ -0,0 +1,286 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +// Standalone agent host server with WebSocket protocol transport. +// Start with: node out/vs/platform/agentHost/node/agentHostServerMain.js [--port ] [--enable-mock-agent] + +import { DisposableStore } from '../../../base/common/lifecycle.js'; +import { URI } from '../../../base/common/uri.js'; +import { localize } from '../../../nls.js'; +import { NativeEnvironmentService } from '../../environment/node/environmentService.js'; +import { INativeEnvironmentService } from '../../environment/common/environment.js'; +import { parseArgs, OPTIONS } from '../../environment/node/argv.js'; +import { getLogLevel, ILogService, NullLogService } from '../../log/common/log.js'; +import { LogService } from '../../log/common/logService.js'; +import { LoggerService } from '../../log/node/loggerService.js'; +import product from '../../product/common/product.js'; +import { IProductService } from '../../product/common/productService.js'; +import { InstantiationService } from '../../instantiation/common/instantiationService.js'; +import { ServiceCollection } from '../../instantiation/common/serviceCollection.js'; +import { CopilotAgent } from './copilot/copilotAgent.js'; +import { AgentSession, type AgentProvider, type IAgent } from '../common/agentService.js'; +import { SessionStateManager } from './sessionStateManager.js'; +import { WebSocketProtocolServer } from './webSocketTransport.js'; +import { ProtocolServerHandler, type IProtocolSideEffectHandler } from './protocolServerHandler.js'; +import { mapProgressEventToAction } from './agentEventMapper.js'; +import { + ISessionModelInfo, + SessionStatus, type ISessionSummary +} from '../common/state/sessionState.js'; +import type { ISessionAction } from '../common/state/sessionActions.js'; +import type { ICreateSessionParams } from '../common/state/sessionProtocol.js'; + +// ---- Options ---------------------------------------------------------------- + +interface IServerOptions { + readonly port: number; + readonly enableMockAgent: boolean; + readonly quiet: boolean; +} + +function parseServerOptions(): IServerOptions { + const argv = process.argv.slice(2); + const envPort = parseInt(process.env['VSCODE_AGENT_HOST_PORT'] ?? '8081', 10); + const portIdx = argv.indexOf('--port'); + const port = portIdx >= 0 ? parseInt(argv[portIdx + 1], 10) : envPort; + const enableMockAgent = argv.includes('--enable-mock-agent'); + const quiet = argv.includes('--quiet'); + return { port, enableMockAgent, quiet }; +} + +// ---- Main ------------------------------------------------------------------- + +function main(): void { + const options = parseServerOptions(); + const disposables = new DisposableStore(); + + // Services — production logging unless --quiet + let logService: ILogService; + let loggerService: LoggerService | undefined; + + if (options.quiet) { + logService = new NullLogService(); + } else { + const services = new ServiceCollection(); + const productService: IProductService = { _serviceBrand: undefined, ...product }; + services.set(IProductService, productService); + const args = parseArgs(process.argv.slice(2), OPTIONS); + const environmentService = new NativeEnvironmentService(args, productService); + services.set(INativeEnvironmentService, environmentService); + loggerService = new LoggerService(getLogLevel(environmentService), environmentService.logsHome); + const logger = loggerService.createLogger('agenthost-server', { name: localize('agentHostServer', "Agent Host Server") }); + logService = disposables.add(new LogService(logger)); + services.set(ILogService, logService); + } + + logService.info('[AgentHostServer] Starting standalone agent host server'); + + // Create state manager + const stateManager = disposables.add(new SessionStateManager(logService)); + + // Agent registry — maps provider id to agent instance + const agents = new Map(); + + function registerAgent(agent: IAgent): void { + agents.set(agent.id, agent); + disposables.add(agent.onDidSessionProgress(e => { + const turnId = stateManager.getActiveTurnId(e.session); + if (turnId) { + const action = mapProgressEventToAction(e, e.session, turnId); + if (action) { + stateManager.dispatchServerAction(action); + } + } + })); + // Publish agent to root state (models fetched async) + publishAgentsToRootState(); + logService.info(`[AgentHostServer] Registered agent: ${agent.id}`); + } + + async function publishAgentsToRootState(): Promise { + const agentInfos = await Promise.all([...agents.values()].map(async a => { + const d = a.getDescriptor(); + let models: ISessionModelInfo[]; + try { + const rawModels = await a.listModels(); + models = rawModels.map(m => ({ + id: m.id, provider: m.provider, name: m.name, + maxContextWindow: m.maxContextWindow, supportsVision: m.supportsVision, + policyState: m.policyState, + })); + } catch { + models = []; + } + return { provider: d.provider, displayName: d.displayName, description: d.description, models }; + })); + stateManager.dispatchServerAction({ type: 'root/agentsChanged', agents: agentInfos }); + } + + function getAgent(session: URI): IAgent | undefined { + const provider = AgentSession.provider(session); + return provider ? agents.get(provider) : agents.values().next().value; + } + + // Register agents + if (!options.quiet) { + // Production agents (require DI) + const services = new ServiceCollection(); + const productService: IProductService = { _serviceBrand: undefined, ...product }; + services.set(IProductService, productService); + const args = parseArgs(process.argv.slice(2), OPTIONS); + const environmentService = new NativeEnvironmentService(args, productService); + services.set(INativeEnvironmentService, environmentService); + services.set(ILogService, logService); + const instantiationService = new InstantiationService(services); + const copilotAgent = disposables.add(instantiationService.createInstance(CopilotAgent)); + registerAgent(copilotAgent); + } + + if (options.enableMockAgent) { + // Dynamic import to avoid bundling test code in production + import('../test/node/mockAgent.js').then(({ ScriptedMockAgent }) => { + const mockAgent = disposables.add(new ScriptedMockAgent()); + registerAgent(mockAgent); + }).catch(err => { + logService.error('[AgentHostServer] Failed to load mock agent', err); + }); + } + + // WebSocket server + const wsServer = disposables.add(new WebSocketProtocolServer(options.port, logService)); + + // Side-effect handler — routes to the correct agent based on session URI + const sideEffects: IProtocolSideEffectHandler = { + handleAction(action: ISessionAction): void { + switch (action.type) { + case 'session/turnStarted': { + const agent = getAgent(action.session); + if (!agent) { + stateManager.dispatchServerAction({ + type: 'session/error', + session: action.session, + turnId: action.turnId, + error: { errorType: 'noAgent', message: 'No agent found for session' }, + }); + return; + } + const attachments = action.userMessage.attachments?.map(a => ({ + type: a.type, + path: a.path, + displayName: a.displayName, + })); + agent.sendMessage(action.session, action.userMessage.text, attachments).catch(err => { + logService.error('[AgentHostServer] sendMessage failed', err); + stateManager.dispatchServerAction({ + type: 'session/error', + session: action.session, + turnId: action.turnId, + error: { errorType: 'sendFailed', message: String(err) }, + }); + }); + break; + } + case 'session/permissionResolved': { + const agent = getAgent(action.session); + agent?.respondToPermissionRequest(action.requestId, action.approved); + break; + } + case 'session/turnCancelled': { + const agent = getAgent(action.session); + agent?.abortSession(action.session).catch(() => { }); + break; + } + case 'session/modelChanged': { + const agent = getAgent(action.session); + agent?.changeModel?.(action.session, action.model).catch(err => { + logService.error('[AgentHostServer] changeModel failed', err); + }); + break; + } + } + }, + async handleCreateSession(command: ICreateSessionParams): Promise { + const provider = (command.provider ?? agents.keys().next().value) as AgentProvider; + const agent = agents.get(provider); + if (!agent) { + throw new Error(`No agent registered for provider: ${provider}`); + } + const session = await agent.createSession({ + provider, + model: command.model, + workingDirectory: command.workingDirectory, + }); + const summary: ISessionSummary = { + resource: session, + provider, + title: 'Session', + status: SessionStatus.Idle, + createdAt: Date.now(), + modifiedAt: Date.now(), + }; + stateManager.createSession(summary); + stateManager.dispatchServerAction({ type: 'session/ready', session }); + }, + handleDisposeSession(session: URI): void { + const agent = getAgent(session); + agent?.disposeSession(session).catch(() => { }); + stateManager.removeSession(session); + }, + async handleListSessions(): Promise { + const allSessions: ISessionSummary[] = []; + for (const agent of agents.values()) { + const sessions = await agent.listSessions(); + const provider = agent.id; + for (const s of sessions) { + allSessions.push({ + resource: s.session, + provider, + title: s.summary ?? 'Session', + status: SessionStatus.Idle, + createdAt: s.startTime, + modifiedAt: s.modifiedTime, + }); + } + } + return allSessions; + }, + }; + + // Wire up protocol handler + disposables.add(new ProtocolServerHandler(stateManager, wsServer, sideEffects, logService)); + + // Report ready + const address = wsServer.address; + if (address) { + const listeningPort = address.split(':').pop(); + process.stdout.write(`READY:${listeningPort}\n`); + logService.info(`[AgentHostServer] WebSocket server listening on ws://${address}`); + } else { + const interval = setInterval(() => { + const addr = wsServer.address; + if (addr) { + clearInterval(interval); + const listeningPort = addr.split(':').pop(); + process.stdout.write(`READY:${listeningPort}\n`); + logService.info(`[AgentHostServer] WebSocket server listening on ws://${addr}`); + } + }, 10); + } + + // Keep alive until stdin closes or signal + process.stdin.resume(); + process.stdin.on('end', shutdown); + process.on('SIGTERM', shutdown); + process.on('SIGINT', shutdown); + + function shutdown(): void { + logService.info('[AgentHostServer] Shutting down...'); + disposables.dispose(); + loggerService?.dispose(); + process.exit(0); + } +} + +main(); diff --git a/src/vs/platform/agentHost/node/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts index e5fe65716aa..02bc38f51ce 100644 --- a/src/vs/platform/agentHost/node/agentService.ts +++ b/src/vs/platform/agentHost/node/agentService.ts @@ -7,7 +7,15 @@ import { Emitter } from '../../../base/common/event.js'; import { Disposable, DisposableStore } from '../../../base/common/lifecycle.js'; import { URI } from '../../../base/common/uri.js'; import { ILogService } from '../../log/common/log.js'; -import { AgentProvider, IAgentAttachment, IAgentCreateSessionConfig, IAgentModelInfo, IAgentProgressEvent, IAgentMessageEvent, IAgent, IAgentService, IAgentSessionMetadata, IAgentToolStartEvent, IAgentToolCompleteEvent, AgentSession, IAgentDescriptor } from '../common/agentService.js'; +import { AgentProvider, IAgentAttachment, IAgentCreateSessionConfig, IAgent, IAgentService, IAgentSessionMetadata, AgentSession, IAgentDescriptor } from '../common/agentService.js'; +import type { IActionEnvelope, INotification, ISessionAction } from '../common/state/sessionActions.js'; +import type { IStateSnapshot } from '../common/state/sessionProtocol.js'; +import { + ISessionModelInfo, + SessionStatus, type ISessionSummary +} from '../common/state/sessionState.js'; +import { mapProgressEventToAction } from './agentEventMapper.js'; +import { SessionStateManager } from './sessionStateManager.js'; /** * The agent service implementation that runs inside the agent-host utility @@ -17,8 +25,16 @@ import { AgentProvider, IAgentAttachment, IAgentCreateSessionConfig, IAgentModel export class AgentService extends Disposable implements IAgentService { declare readonly _serviceBrand: undefined; - private readonly _onDidSessionProgress = this._register(new Emitter()); - readonly onDidSessionProgress = this._onDidSessionProgress.event; + /** Protocol: fires when state is mutated by an action. */ + private readonly _onDidAction = this._register(new Emitter()); + readonly onDidAction = this._onDidAction.event; + + /** Protocol: fires for ephemeral notifications (sessionAdded/Removed). */ + private readonly _onDidNotification = this._register(new Emitter()); + readonly onDidNotification = this._onDidNotification.event; + + /** Authoritative state manager for the sessions process protocol. */ + private readonly _stateManager: SessionStateManager; /** Registered providers keyed by their {@link AgentProvider} id. */ private readonly _providers = new Map(); @@ -36,6 +52,9 @@ export class AgentService extends Disposable implements IAgentService { ) { super(); this._logService.info('AgentService initialized'); + this._stateManager = this._register(new SessionStateManager(_logService)); + this._register(this._stateManager.onDidEmitEnvelope(e => this._onDidAction.fire(e))); + this._register(this._stateManager.onDidEmitNotification(e => this._onDidNotification.fire(e))); } // ---- provider registration ---------------------------------------------- @@ -48,16 +67,27 @@ export class AgentService extends Disposable implements IAgentService { this._providers.set(provider.id, provider); this._providerSubscriptions.add( provider.onDidSessionProgress(e => { - // Track permission requests so respondToPermissionRequest can route + // Track permission requests so dispatchAction can route if (e.type === 'permission_request') { this._pendingPermissions.set(e.requestId, provider.id); } - this._onDidSessionProgress.fire(e); + + // Map to protocol action and dispatch through state manager + const turnId = this._stateManager.getActiveTurnId(e.session); + if (turnId) { + const action = mapProgressEventToAction(e, e.session, turnId); + if (action) { + this._stateManager.dispatchServerAction(action); + } + } }) ); if (!this._defaultProvider) { this._defaultProvider = provider.id; } + + // Update root state with current agents list + this._publishAgentsToRootState(); } // ---- auth --------------------------------------------------------------- @@ -87,14 +117,13 @@ export class AgentService extends Disposable implements IAgentService { return flat; } - async listModels(): Promise { - this._logService.trace('[AgentService] listModels called'); - const results = await Promise.all( - [...this._providers.values()].map(p => p.listModels()) - ); - const flat = results.flat(); - this._logService.trace(`[AgentService] listModels returned ${flat.length} models`); - return flat; + /** + * Refreshes the model list from all providers and publishes the updated + * agents (with their models) to root state via `root/agentsChanged`. + */ + async refreshModels(): Promise { + this._logService.trace('[AgentService] refreshModels called'); + await this._publishAgentsToRootState(); } async createSession(config?: IAgentCreateSessionConfig): Promise { @@ -107,28 +136,22 @@ export class AgentService extends Disposable implements IAgentService { const session = await provider.createSession(config); this._sessionToProvider.set(session.toString(), provider.id); this._logService.trace(`[AgentService] createSession returned: ${session.toString()}`); + + // Create state in the state manager + const summary: ISessionSummary = { + resource: session, + provider: provider.id, + title: 'New Session', + status: SessionStatus.Idle, + createdAt: Date.now(), + modifiedAt: Date.now(), + }; + this._stateManager.createSession(summary); + this._stateManager.dispatchServerAction({ type: 'session/ready', session }); + return session; } - async sendMessage(session: URI, prompt: string, attachments?: IAgentAttachment[]): Promise { - this._logService.trace(`[AgentService] sendMessage: session=${session.toString()}, prompt=${prompt.length} chars, attachments=${attachments?.length ?? 0}`); - const provider = this._getProviderForSession(session); - await provider.sendMessage(session, prompt, attachments); - this._logService.trace(`[AgentService] sendMessage returned for ${session.toString()}`); - } - - async getSessionMessages(session: URI): Promise<(IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[]> { - this._logService.trace(`[AgentService] getSessionMessages: ${session.toString()}`); - const provider = this._findProviderForSession(session); - if (!provider) { - this._logService.trace(`[AgentService] getSessionMessages: no provider found, returning empty`); - return []; - } - const messages = await provider.getSessionMessages(session); - this._logService.trace(`[AgentService] getSessionMessages returned ${messages.length} events`); - return messages; - } - async disposeSession(session: URI): Promise { this._logService.trace(`[AgentService] disposeSession: ${session.toString()}`); const provider = this._findProviderForSession(session); @@ -136,26 +159,76 @@ export class AgentService extends Disposable implements IAgentService { await provider.disposeSession(session); this._sessionToProvider.delete(session.toString()); } + this._stateManager.removeSession(session); } - async abortSession(session: URI): Promise { - this._logService.trace(`[AgentService] abortSession: ${session.toString()}`); - const provider = this._findProviderForSession(session); - if (provider) { - await provider.abortSession(session); + // ---- Protocol methods --------------------------------------------------- + + async subscribe(resource: URI): Promise { + this._logService.trace(`[AgentService] subscribe: ${resource.toString()}`); + const snapshot = this._stateManager.getSnapshot(resource); + if (!snapshot) { + throw new Error(`Cannot subscribe to unknown resource: ${resource.toString()}`); } + return snapshot; } - respondToPermissionRequest(requestId: string, approved: boolean): void { - this._logService.trace(`[AgentService] respondToPermissionRequest: ${requestId} approved=${approved}`); - const providerId = this._pendingPermissions.get(requestId); - if (!providerId) { - this._logService.warn(`[AgentService] No pending permission request for: ${requestId}`); - return; + unsubscribe(resource: URI): void { + this._logService.trace(`[AgentService] unsubscribe: ${resource.toString()}`); + // Server-side tracking of per-client subscriptions will be added + // in Phase 4 (multi-client). For now this is a no-op. + } + + dispatchAction(action: ISessionAction, clientId: string, clientSeq: number): void { + this._logService.trace(`[AgentService] dispatchAction: type=${action.type}, clientId=${clientId}, clientSeq=${clientSeq}`, action); + + const origin = { clientId, clientSeq }; + const state = this._stateManager.dispatchClientAction(action, origin); + this._logService.trace(`[AgentService] resulting state:`, state); + + // Trigger side effects based on the action type + switch (action.type) { + case 'session/turnStarted': { + const provider = this._findProviderForSession(action.session); + if (provider) { + const attachments = action.userMessage.attachments?.map(a => ({ + type: a.type, + path: a.path, + displayName: a.displayName, + }) satisfies IAgentAttachment); + provider.sendMessage(action.session, action.userMessage.text, attachments).catch(err => { + this._logService.error(`[AgentService] sendMessage failed for session/turnStarted`, err); + this._stateManager.dispatchServerAction({ + type: 'session/error', + session: action.session, + turnId: action.turnId, + error: { errorType: 'sendFailed', message: String(err) }, + }); + }); + } + break; + } + case 'session/permissionResolved': { + const providerId = this._pendingPermissions.get(action.requestId); + if (providerId) { + this._pendingPermissions.delete(action.requestId); + const permProvider = this._providers.get(providerId); + permProvider?.respondToPermissionRequest(action.requestId, action.approved); + } else { + this._logService.warn(`[AgentService] No pending permission request for: ${action.requestId}`); + } + break; + } + case 'session/turnCancelled': { + const provider = this._findProviderForSession(action.session); + if (provider) { + provider.abortSession(action.session).catch(err => { + this._logService.error(`[AgentService] abortSession failed for session/turnCancelled`, err); + }); + } + break; + } } - this._pendingPermissions.delete(requestId); - const provider = this._providers.get(providerId); - provider?.respondToPermissionRequest(requestId, approved); } async shutdown(): Promise { @@ -170,12 +243,27 @@ export class AgentService extends Disposable implements IAgentService { // ---- helpers ------------------------------------------------------------ - private _getProviderForSession(session: URI): IAgent { - const provider = this._findProviderForSession(session); - if (!provider) { - throw new Error(`No provider found for session: ${session.toString()}`); - } - return provider; + /** + * Fetches models from all providers and dispatches `root/agentsChanged` + * with the merged agent + model data. + */ + private async _publishAgentsToRootState(): Promise { + const agents = await Promise.all([...this._providers.values()].map(async p => { + const d = p.getDescriptor(); + let models: ISessionModelInfo[]; + try { + const rawModels = await p.listModels(); + models = rawModels.map(m => ({ + id: m.id, provider: m.provider, name: m.name, + maxContextWindow: m.maxContextWindow, supportsVision: m.supportsVision, + policyState: m.policyState, + })); + } catch { + models = []; + } + return { provider: d.provider, displayName: d.displayName, description: d.description, models }; + })); + this._stateManager.dispatchServerAction({ type: 'root/agentsChanged', agents }); } private _findProviderForSession(session: URI): IAgent | undefined { diff --git a/src/vs/platform/agentHost/node/protocolServerHandler.ts b/src/vs/platform/agentHost/node/protocolServerHandler.ts new file mode 100644 index 00000000000..cda85ebc0d7 --- /dev/null +++ b/src/vs/platform/agentHost/node/protocolServerHandler.ts @@ -0,0 +1,334 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { Disposable, DisposableStore } from '../../../base/common/lifecycle.js'; +import { URI } from '../../../base/common/uri.js'; +import { ILogService } from '../../log/common/log.js'; +import { IActionEnvelope, INotification, isSessionAction } from '../common/state/sessionActions.js'; +import { isActionKnownToVersion, PROTOCOL_VERSION } from '../common/state/sessionCapabilities.js'; +import { + isJsonRpcRequest, + isJsonRpcNotification, + JSON_RPC_INTERNAL_ERROR, + type ICreateSessionParams, + type IDispatchActionParams, + type IDisposeSessionParams, + type IFetchTurnsParams, + type IInitializeParams, + type IProtocolMessage, + type IReconnectParams, + type IStateSnapshot, + type ISubscribeParams, + type IUnsubscribeParams, +} from '../common/state/sessionProtocol.js'; +import { ROOT_STATE_URI } from '../common/state/sessionState.js'; +import type { IProtocolServer, IProtocolTransport } from '../common/state/sessionTransport.js'; +import { SessionStateManager } from './sessionStateManager.js'; + +/** Default capacity of the server-side action replay buffer. */ +const REPLAY_BUFFER_CAPACITY = 1000; + +/** + * Represents a connected protocol client with its subscription state. + */ +interface IConnectedClient { + readonly clientId: string; + readonly protocolVersion: number; + readonly transport: IProtocolTransport; + readonly subscriptions: Set; + readonly disposables: DisposableStore; +} + +/** + * Server-side handler that manages protocol connections, routes JSON-RPC + * messages to the state manager, and broadcasts actions/notifications + * to subscribed clients. + */ +export class ProtocolServerHandler extends Disposable { + + private readonly _clients = new Map(); + private readonly _replayBuffer: IActionEnvelope[] = []; + + constructor( + private readonly _stateManager: SessionStateManager, + private readonly _server: IProtocolServer, + private readonly _sideEffectHandler: IProtocolSideEffectHandler, + @ILogService private readonly _logService: ILogService, + ) { + super(); + + this._register(this._server.onConnection(transport => { + this._handleNewConnection(transport); + })); + + this._register(this._stateManager.onDidEmitEnvelope(envelope => { + this._replayBuffer.push(envelope); + if (this._replayBuffer.length > REPLAY_BUFFER_CAPACITY) { + this._replayBuffer.shift(); + } + this._broadcastAction(envelope); + })); + + this._register(this._stateManager.onDidEmitNotification(notification => { + this._broadcastNotification(notification); + })); + } + + // ---- Connection handling ------------------------------------------------- + + private _handleNewConnection(transport: IProtocolTransport): void { + const disposables = new DisposableStore(); + let client: IConnectedClient | undefined; + + disposables.add(transport.onMessage(msg => { + if (isJsonRpcRequest(msg)) { + // Request — expects a correlated response + if (!client) { + return; + } + this._handleRequest(client, msg.method, msg.params, msg.id); + } else if (isJsonRpcNotification(msg)) { + // Notification — fire-and-forget + switch (msg.method) { + case 'initialize': + client = this._handleInitialize(msg.params as IInitializeParams, transport, disposables); + break; + case 'reconnect': + client = this._handleReconnect(msg.params as IReconnectParams, transport, disposables); + break; + case 'unsubscribe': + if (client) { + client.subscriptions.delete((msg.params as IUnsubscribeParams).resource.toString()); + } + break; + case 'dispatchAction': + if (client) { + const params = msg.params as IDispatchActionParams; + const origin = { clientId: client.clientId, clientSeq: params.clientSeq }; + this._stateManager.dispatchClientAction(params.action, origin); + this._sideEffectHandler.handleAction(params.action); + } + break; + } + } + // Responses from the client (if any) are ignored on the server side. + })); + + disposables.add(transport.onClose(() => { + if (client) { + this._logService.info(`[ProtocolServer] Client disconnected: ${client.clientId}`); + this._clients.delete(client.clientId); + } + disposables.dispose(); + })); + + disposables.add(transport); + } + + // ---- Notifications (fire-and-forget) ------------------------------------ + + private _handleInitialize( + params: IInitializeParams, + transport: IProtocolTransport, + disposables: DisposableStore, + ): IConnectedClient { + this._logService.info(`[ProtocolServer] Initialize: clientId=${params.clientId}, version=${params.protocolVersion}`); + + const client: IConnectedClient = { + clientId: params.clientId, + protocolVersion: params.protocolVersion, + transport, + subscriptions: new Set(), + disposables, + }; + this._clients.set(params.clientId, client); + + const snapshots: IStateSnapshot[] = []; + if (params.initialSubscriptions) { + for (const uri of params.initialSubscriptions) { + const snapshot = this._stateManager.getSnapshot(uri); + if (snapshot) { + snapshots.push(snapshot); + client.subscriptions.add(uri.toString()); + } + } + } + + this._sendNotification(transport, 'serverHello', { + protocolVersion: PROTOCOL_VERSION, + serverSeq: this._stateManager.serverSeq, + snapshots, + }); + + return client; + } + + private _handleReconnect( + params: IReconnectParams, + transport: IProtocolTransport, + disposables: DisposableStore, + ): IConnectedClient { + this._logService.info(`[ProtocolServer] Reconnect: clientId=${params.clientId}, lastSeenSeq=${params.lastSeenServerSeq}`); + + const client: IConnectedClient = { + clientId: params.clientId, + protocolVersion: PROTOCOL_VERSION, + transport, + subscriptions: new Set(), + disposables, + }; + this._clients.set(params.clientId, client); + + const oldestBuffered = this._replayBuffer.length > 0 ? this._replayBuffer[0].serverSeq : this._stateManager.serverSeq; + const canReplay = params.lastSeenServerSeq >= oldestBuffered; + + if (canReplay) { + for (const sub of params.subscriptions) { + client.subscriptions.add(sub.toString()); + } + for (const envelope of this._replayBuffer) { + if (envelope.serverSeq > params.lastSeenServerSeq) { + if (this._isRelevantToClient(client, envelope)) { + this._sendNotification(transport, 'action', { envelope }); + } + } + } + } else { + const snapshots: IStateSnapshot[] = []; + for (const sub of params.subscriptions) { + const snapshot = this._stateManager.getSnapshot(sub); + if (snapshot) { + snapshots.push(snapshot); + client.subscriptions.add(sub.toString()); + } + } + this._sendNotification(transport, 'reconnectResponse', { + serverSeq: this._stateManager.serverSeq, + snapshots, + }); + } + + return client; + } + + // ---- Requests (expect a response) --------------------------------------- + + private _handleRequest(client: IConnectedClient, method: string, params: unknown, id: number): void { + this._handleRequestAsync(client, method, params).then(result => { + client.transport.send({ jsonrpc: '2.0', id, result: result ?? null }); + }).catch(err => { + this._logService.error(`[ProtocolServer] Request '${method}' failed`, err); + client.transport.send({ + jsonrpc: '2.0', + id, + error: { code: JSON_RPC_INTERNAL_ERROR, message: String(err?.message ?? err) }, + }); + }); + } + + private async _handleRequestAsync(client: IConnectedClient, method: string, params: unknown): Promise { + switch (method) { + case 'subscribe': { + const p = params as ISubscribeParams; + const snapshot = this._stateManager.getSnapshot(p.resource); + if (snapshot) { + client.subscriptions.add(p.resource.toString()); + } + return snapshot ?? null; + } + case 'createSession': { + await this._sideEffectHandler.handleCreateSession(params as ICreateSessionParams); + return null; + } + case 'disposeSession': { + this._sideEffectHandler.handleDisposeSession((params as IDisposeSessionParams).session); + return null; + } + case 'listSessions': { + const sessions = await this._sideEffectHandler.handleListSessions(); + return { sessions }; + } + case 'fetchTurns': { + const p = params as IFetchTurnsParams; + const state = this._stateManager.getSessionState(p.session); + if (state) { + const turns = state.turns; + const start = Math.max(0, p.startTurn); + const end = Math.min(turns.length, start + p.count); + return { + session: p.session, + startTurn: start, + turns: turns.slice(start, end), + totalTurns: turns.length, + }; + } + return { + session: p.session, + startTurn: p.startTurn, + turns: [], + totalTurns: 0, + }; + } + default: + throw new Error(`Unknown method: ${method}`); + } + } + + // ---- Broadcasting ------------------------------------------------------- + + private _sendNotification(transport: IProtocolTransport, method: string, params: unknown): void { + transport.send({ jsonrpc: '2.0', method, params }); + } + + private _broadcastAction(envelope: IActionEnvelope): void { + const msg: IProtocolMessage = { jsonrpc: '2.0', method: 'action', params: { envelope } }; + for (const client of this._clients.values()) { + if (this._isRelevantToClient(client, envelope)) { + client.transport.send(msg); + } + } + } + + private _broadcastNotification(notification: INotification): void { + const msg: IProtocolMessage = { jsonrpc: '2.0', method: 'notification', params: { notification } }; + for (const client of this._clients.values()) { + client.transport.send(msg); + } + } + + private _isRelevantToClient(client: IConnectedClient, envelope: IActionEnvelope): boolean { + const action = envelope.action; + if (!isActionKnownToVersion(action, client.protocolVersion)) { + return false; + } + if (action.type.startsWith('root/')) { + return client.subscriptions.has(ROOT_STATE_URI.toString()); + } + if (isSessionAction(action)) { + return client.subscriptions.has(action.session.toString()); + } + return false; + } + + override dispose(): void { + for (const client of this._clients.values()) { + client.disposables.dispose(); + } + this._clients.clear(); + this._replayBuffer.length = 0; + super.dispose(); + } +} + +/** + * Interface for side effects that the protocol server delegates to. + * These are operations that involve I/O, agent backends, etc. + */ +export interface IProtocolSideEffectHandler { + handleAction(action: import('../common/state/sessionActions.js').ISessionAction): void; + handleCreateSession(command: import('../common/state/sessionProtocol.js').ICreateSessionParams): Promise; + handleDisposeSession(session: URI): void; + handleListSessions(): Promise; +} diff --git a/src/vs/platform/agentHost/node/sessionStateManager.ts b/src/vs/platform/agentHost/node/sessionStateManager.ts new file mode 100644 index 00000000000..121daa48d8d --- /dev/null +++ b/src/vs/platform/agentHost/node/sessionStateManager.ts @@ -0,0 +1,218 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { Emitter, Event } from '../../../base/common/event.js'; +import { Disposable } from '../../../base/common/lifecycle.js'; +import { URI } from '../../../base/common/uri.js'; +import { ILogService } from '../../log/common/log.js'; +import { IActionEnvelope, IActionOrigin, INotification, ISessionAction, IRootAction, IStateAction, isRootAction, isSessionAction } from '../common/state/sessionActions.js'; +import { IStateSnapshot } from '../common/state/sessionProtocol.js'; +import { rootReducer, sessionReducer } from '../common/state/sessionReducers.js'; +import { createRootState, createSessionState, IRootState, ISessionState, ISessionSummary, ROOT_STATE_URI } from '../common/state/sessionState.js'; + +/** + * Server-side state manager for the sessions process protocol. + * + * Maintains the authoritative state tree (root + per-session), applies actions + * through pure reducers, assigns monotonic sequence numbers, and emits + * {@link IActionEnvelope}s for subscribed clients. + */ +export class SessionStateManager extends Disposable { + + private _serverSeq = 0; + + private _rootState: IRootState; + private readonly _sessionStates = new Map(); + + /** Tracks which session URI each active turn belongs to, keyed by turnId. */ + private readonly _activeTurnToSession = new Map(); + + private readonly _onDidEmitEnvelope = this._register(new Emitter()); + readonly onDidEmitEnvelope: Event = this._onDidEmitEnvelope.event; + + private readonly _onDidEmitNotification = this._register(new Emitter()); + readonly onDidEmitNotification: Event = this._onDidEmitNotification.event; + + constructor( + @ILogService private readonly _logService: ILogService, + ) { + super(); + this._rootState = createRootState(); + } + + // ---- State accessors ---------------------------------------------------- + + get rootState(): IRootState { + return this._rootState; + } + + getSessionState(session: URI): ISessionState | undefined { + return this._sessionStates.get(session.toString()); + } + + get serverSeq(): number { + return this._serverSeq; + } + + // ---- Snapshots ---------------------------------------------------------- + + /** + * Returns a state snapshot for a given resource URI. + * The `fromSeq` in the snapshot is the current serverSeq at snapshot time; + * the client should process subsequent envelopes with serverSeq > fromSeq. + */ + getSnapshot(resource: URI): IStateSnapshot | undefined { + const key = resource.toString(); + + if (key === ROOT_STATE_URI.toString()) { + return { + resource, + state: this._rootState, + fromSeq: this._serverSeq, + }; + } + + const sessionState = this._sessionStates.get(key); + if (!sessionState) { + return undefined; + } + + return { + resource, + state: sessionState, + fromSeq: this._serverSeq, + }; + } + + // ---- Session lifecycle -------------------------------------------------- + + /** + * Creates a new session in state with `lifecycle: 'creating'`. + * Returns the initial session state. + */ + createSession(summary: ISessionSummary): ISessionState { + const key = summary.resource.toString(); + if (this._sessionStates.has(key)) { + this._logService.warn(`[SessionStateManager] Session already exists: ${key}`); + return this._sessionStates.get(key)!; + } + + const state = createSessionState(summary); + this._sessionStates.set(key, state); + + this._logService.trace(`[SessionStateManager] Created session: ${key}`); + + this._onDidEmitNotification.fire({ + type: 'notify/sessionAdded', + summary, + }); + + return state; + } + + /** + * Removes a session from state and emits a sessionRemoved notification. + */ + removeSession(session: URI): void { + const key = session.toString(); + const state = this._sessionStates.get(key); + if (!state) { + return; + } + + // Clean up active turn tracking + if (state.activeTurn) { + this._activeTurnToSession.delete(state.activeTurn.id); + } + + this._sessionStates.delete(key); + this._logService.trace(`[SessionStateManager] Removed session: ${key}`); + + this._onDidEmitNotification.fire({ + type: 'notify/sessionRemoved', + session, + }); + } + + // ---- Turn tracking ------------------------------------------------------ + + /** + * Registers a mapping from turnId to session URI so that incoming + * provider events (which carry only session URI) can be associated + * with the correct active turn. + */ + getActiveTurnId(session: URI): string | undefined { + const state = this._sessionStates.get(session.toString()); + return state?.activeTurn?.id; + } + + // ---- Action dispatch ---------------------------------------------------- + + /** + * Dispatch a server-originated action (from the agent backend). + * The action is applied to state via the reducer and emitted as an + * envelope with no origin (server-produced). + */ + dispatchServerAction(action: IStateAction): void { + this._applyAndEmit(action, undefined); + } + + /** + * Dispatch a client-originated action (write-ahead from a renderer). + * The action is applied to state and emitted with the client's origin + * so the originating client can reconcile. + */ + dispatchClientAction(action: ISessionAction, origin: IActionOrigin): unknown { + return this._applyAndEmit(action, origin); + } + + // ---- Internal ----------------------------------------------------------- + + private _applyAndEmit(action: IStateAction, origin: IActionOrigin | undefined): unknown { + let resultingState: unknown = undefined; + // Apply to state + if (isRootAction(action)) { + this._rootState = rootReducer(this._rootState, action as IRootAction); + resultingState = this._rootState; + } + + if (isSessionAction(action)) { + const sessionAction = action as ISessionAction; + const key = sessionAction.session.toString(); + const state = this._sessionStates.get(key); + if (state) { + const newState = sessionReducer(state, sessionAction); + this._sessionStates.set(key, newState); + + // Track active turn for turn lifecycle + if (sessionAction.type === 'session/turnStarted') { + this._activeTurnToSession.set(sessionAction.turnId, key); + } else if ( + sessionAction.type === 'session/turnComplete' || + sessionAction.type === 'session/turnCancelled' || + sessionAction.type === 'session/error' + ) { + this._activeTurnToSession.delete(sessionAction.turnId); + } + + resultingState = newState; + } else { + this._logService.warn(`[SessionStateManager] Action for unknown session: ${key}, type=${action.type}`); + } + } + + // Emit envelope + const envelope: IActionEnvelope = { + action, + serverSeq: ++this._serverSeq, + origin, + }; + + this._logService.trace(`[SessionStateManager] Emitting envelope: seq=${envelope.serverSeq}, type=${action.type}${origin ? `, origin=${origin.clientId}:${origin.clientSeq}` : ''}`); + this._onDidEmitEnvelope.fire(envelope); + + return resultingState; + } +} diff --git a/src/vs/platform/agentHost/node/webSocketTransport.ts b/src/vs/platform/agentHost/node/webSocketTransport.ts new file mode 100644 index 00000000000..a56c2b8060c --- /dev/null +++ b/src/vs/platform/agentHost/node/webSocketTransport.ts @@ -0,0 +1,135 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +// WebSocket transport for the sessions process protocol. +// Uses JSON serialization with URI revival for cross-process communication. + +import { WebSocketServer, WebSocket } from 'ws'; +import { Emitter } from '../../../base/common/event.js'; +import { Disposable } from '../../../base/common/lifecycle.js'; +import { URI } from '../../../base/common/uri.js'; +import { ILogService } from '../../log/common/log.js'; +import type { IProtocolMessage } from '../common/state/sessionProtocol.js'; +import type { IProtocolServer, IProtocolTransport } from '../common/state/sessionTransport.js'; + +// ---- JSON serialization helpers --------------------------------------------- + +function uriReplacer(_key: string, value: unknown): unknown { + if (value instanceof URI) { + return value.toJSON(); + } + if (value instanceof Map) { + return { $type: 'Map', entries: [...value.entries()] }; + } + return value; +} + +function uriReviver(_key: string, value: unknown): unknown { + if (value && typeof value === 'object') { + const obj = value as Record; + if (obj.$mid === 1) { + return URI.revive(value as URI); + } + if (obj.$type === 'Map' && Array.isArray(obj.entries)) { + return new Map(obj.entries as [unknown, unknown][]); + } + } + return value; +} + +// ---- Per-connection transport ----------------------------------------------- + +/** + * Wraps a single WebSocket connection as an {@link IProtocolTransport}. + * Messages are serialized as JSON with URI revival. + */ +export class WebSocketProtocolTransport extends Disposable implements IProtocolTransport { + + private readonly _onMessage = this._register(new Emitter()); + readonly onMessage = this._onMessage.event; + + private readonly _onClose = this._register(new Emitter()); + readonly onClose = this._onClose.event; + + constructor(private readonly _ws: WebSocket) { + super(); + + this._ws.on('message', (data: Buffer | string) => { + try { + const text = typeof data === 'string' ? data : data.toString('utf-8'); + const message = JSON.parse(text, uriReviver) as IProtocolMessage; + this._onMessage.fire(message); + } catch { + // Malformed message — drop. No logger available at transport level. + } + }); + + this._ws.on('close', () => { + this._onClose.fire(); + }); + + this._ws.on('error', () => { + // Error always precedes close — closing is handled in the close handler. + this._onClose.fire(); + }); + } + + send(message: IProtocolMessage): void { + if (this._ws.readyState === WebSocket.OPEN) { + this._ws.send(JSON.stringify(message, uriReplacer)); + } + } + + override dispose(): void { + this._ws.close(); + super.dispose(); + } +} + +// ---- Server ----------------------------------------------------------------- + +/** + * WebSocket server that accepts client connections and wraps each one + * as an {@link IProtocolTransport}. + */ +export class WebSocketProtocolServer extends Disposable implements IProtocolServer { + + private readonly _wss: WebSocketServer; + + private readonly _onConnection = this._register(new Emitter()); + readonly onConnection = this._onConnection.event; + + get address(): string | undefined { + const addr = this._wss.address(); + if (!addr || typeof addr === 'string') { + return addr ?? undefined; + } + return `${addr.address}:${addr.port}`; + } + + constructor( + private readonly _port: number, + @ILogService private readonly _logService: ILogService, + ) { + super(); + this._wss = new WebSocketServer({ port: this._port, host: '127.0.0.1' }); + this._logService.info(`[WebSocketProtocol] Server listening on 127.0.0.1:${this._port}`); + + this._wss.on('connection', (ws) => { + this._logService.trace('[WebSocketProtocol] New client connection'); + const transport = new WebSocketProtocolTransport(ws); + this._onConnection.fire(transport); + }); + + this._wss.on('error', (err) => { + this._logService.error('[WebSocketProtocol] Server error', err); + }); + } + + override dispose(): void { + this._wss.close(); + super.dispose(); + } +} diff --git a/src/vs/platform/agentHost/protocol.md b/src/vs/platform/agentHost/protocol.md index dd93735c166..ed259d58fd1 100644 --- a/src/vs/platform/agentHost/protocol.md +++ b/src/vs/platform/agentHost/protocol.md @@ -2,6 +2,8 @@ > **Keep this document in sync with the code.** Changes to the state model, action types, protocol messages, or versioning strategy must be reflected here. Implementation lives in `common/state/`. +> **Pre-production.** This protocol is under active development and is not shipped yet. Breaking changes to wire types, actions, and state shapes are fine — do not worry about backward compatibility until the protocol is in production. The versioning machinery exists for future use. + For process architecture and IPC details, see [architecture.md](architecture.md). For design decisions, see [design.md](design.md). For the task backlog, see [backlog.md](backlog.md). ## Goal @@ -13,11 +15,62 @@ The sessions process is a portable, standalone server that multiple clients can 3. **Write-ahead with reconciliation** — clients optimistically apply their own actions locally, then reconcile when the server echoes them back alongside any concurrent actions from other clients or the server itself. 4. **Forward-compatible versioning** — newer clients can connect to older servers. A single protocol version number maps to a capabilities object; clients check capabilities before using features. +## Protocol development checklist + +Use this checklist when adding a new action, command, state field, or notification to the protocol. + +### Adding a new action type + +1. **Write an E2E test first** in `protocolWebSocket.integrationTest.ts` that exercises the action end-to-end through the WebSocket server. The test should fail until the implementation is complete. +2. **Add mock agent support** if the test needs a new prompt/behavior in `mockAgent.ts`. +3. **Define the action interface** in `sessionActions.ts`. Extend `ISessionActionBase` (for session-scoped) or define a standalone root action. Add it to the `ISessionAction` or `IRootAction` union. +4. **Add a reducer case** in `sessionReducers.ts`. The switch must remain exhaustive — the compiler will error if a case is missing. +5. **Add a v1 wire type** in `versions/v1.ts`. Mirror the action interface shape. Add it to the `IV1_SessionAction` or `IV1_RootAction` union. +6. **Register in `versionRegistry.ts`**: + - Import the new `IV1_*` type. + - Add an `AssertCompatible` check. + - Add the type to the `ISessionAction_v1` union. + - Add the type string to the suppress-warnings `void` expression. + - Add an entry to `ACTION_INTRODUCED_IN` (compiler enforces this). +7. **Update `protocol.md`** (this file) — add the action to the Actions table. +8. **Verify the E2E test passes.** + +### Adding a new command + +1. **Write an E2E test first** in `protocolWebSocket.integrationTest.ts`. The test should fail until the implementation is complete. +2. **Define the request params and result interfaces** in `sessionProtocol.ts`. +3. **Handle it in `protocolServerHandler.ts`** `_handleRequestAsync()`. The method returns the result; the caller wraps it in a JSON-RPC response or error automatically. +4. **Add the side-effect** in `IProtocolSideEffectHandler` if the command requires I/O or agent interaction. Implement it in `agentHostServerMain.ts`. +5. **Update `protocol.md`** — add the command to the Commands table. +6. **Verify the E2E test passes.** + +### Adding a new state field + +1. **Add the field** to the relevant interface in `sessionState.ts` (e.g. `ISessionSummary`, `IActiveTurn`, `ITurn`). +2. **Update the factory** (`createSessionState()`, `createActiveTurn()`) to initialize the field. +3. **Add to the v1 wire type** in `versions/v1.ts`. Optional fields are safe; required fields break the bidirectional `AssertCompatible` check (intentionally — add as optional or bump the protocol version). +4. **Update reducers** in `sessionReducers.ts` if the field needs to be mutated by actions. +5. **Update `finalizeTurn()`** if the field lives on `IActiveTurn` and should transfer to `ITurn` on completion. + +### Adding a new notification + +1. **Write an E2E test first** in `protocolWebSocket.integrationTest.ts`. +2. **Define the notification interface** in `sessionActions.ts`. Add it to the `INotification` union. +3. **Add to `NOTIFICATION_INTRODUCED_IN`** in `versionRegistry.ts`. +4. **Emit it** from `SessionStateManager` or the relevant server-side code. +5. **Verify the E2E test passes.** + +### Adding mock agent support (for testing) + +1. **Add a prompt case** in `mockAgent.ts` `sendMessage()` to trigger the behavior. +2. **Fire the corresponding `IAgentProgressEvent`** via `_fireSequence()` or manually through `_onDidSessionProgress`. + + ## URI-based subscriptions All state is identified by URIs. Clients subscribe to a URI to receive its current state snapshot and subsequent action updates. This is the single universal mechanism for state synchronization: -- **Root state** (`agenthost:root`) — always-present global state (agents, models). Clients subscribe to this on connect. +- **Root state** (`agenthost:root`) — always-present global state (agents and their models). Clients subscribe to this on connect. - **Session state** (`copilot:/`, etc.) — per-session state loaded on demand. Clients subscribe when opening a session. The `subscribe(uri)` / `unsubscribe(uri)` mechanism works identically for all resource types. @@ -31,6 +84,16 @@ Subscribable at `agenthost:root`. Contains global, lightweight data that all cli ``` RootState { agents: AgentInfo[] +} +``` + +Each `AgentInfo` includes the models available for that agent: + +``` +AgentInfo { + provider: string + displayName: string + description: string models: ModelInfo[] } ``` @@ -69,6 +132,7 @@ ActiveTurn { toolCalls: Map pendingPermissions: Map reasoning: string + usage: UsageInfo | undefined } ``` @@ -114,8 +178,7 @@ These mutate the root state. **All root actions are server-only** — clients ob | Type | Payload | When | |---|---|---| -| `root/modelsChanged` | `ModelInfo[]` | Available models changed | -| `root/agentsChanged` | `AgentInfo[]` | Available agent backends changed | +| `root/agentsChanged` | `AgentInfo[]` | Available agent backends or their models changed | ### Session actions @@ -140,6 +203,7 @@ When a client dispatches an action, the server applies it to the state and also | `session/titleChanged` | `title` | No | Session title updated | | `session/usage` | `turnId, UsageInfo` | No | Token usage report | | `session/reasoning` | `turnId, content` | No | Reasoning/thinking text | +| `session/modelChanged` | `model` | Yes | Model changed for this session | ### Notifications @@ -190,45 +254,88 @@ Clients interact with the server in two ways: ## Client-server protocol +The protocol uses **JSON-RPC 2.0** framing over the transport (WebSocket, MessagePort, etc.). + +### Message categories + +- **Client → Server notifications** (fire-and-forget): `initialize`, `reconnect`, `unsubscribe`, `dispatchAction` +- **Client → Server requests** (expect a correlated response): `subscribe`, `createSession`, `disposeSession`, `listSessions`, `fetchTurns`, `fetchContent` +- **Server → Client notifications** (pushed): `serverHello`, `reconnectResponse`, `action`, `notification` +- **Server → Client responses** (correlated to requests by `id`): success result or JSON-RPC error + ### Connection handshake ``` -1. Client → Server: ClientHello { protocolVersion, clientId, initialSubscriptions?: URI[] } -2. Server → Client: ServerHello { protocolVersion, serverSeq, snapshots[] } +1. Client → Server: { "jsonrpc": "2.0", "method": "initialize", "params": { protocolVersion, clientId, initialSubscriptions? } } +2. Server → Client: { "jsonrpc": "2.0", "method": "serverHello", "params": { protocolVersion, serverSeq, snapshots[] } } ``` `initialSubscriptions` allows the client to subscribe to root state (and any previously-open sessions on reconnect) in the same round-trip as the handshake. The server responds with snapshots for each. ### URI subscription -After handshake, clients can subscribe/unsubscribe at any time: +`subscribe` is a JSON-RPC **request** — the client receives the snapshot as the response result: ``` -Client → Server: Subscribe { resource: URI } -Server → Client: StateSnapshot { resource: URI, state, fromSeq } +Client → Server: { "jsonrpc": "2.0", "id": 1, "method": "subscribe", "params": { "resource": "copilot:/session-1" } } +Server → Client: { "jsonrpc": "2.0", "id": 1, "result": { "resource": ..., "state": ..., "fromSeq": 5 } } ``` After subscribing, the client receives all actions scoped to that URI with `serverSeq > fromSeq`. Multiple concurrent subscriptions are supported. +`unsubscribe` is a notification (no response needed): + ``` -Client → Server: Unsubscribe { resource: URI } +Client → Server: { "jsonrpc": "2.0", "method": "unsubscribe", "params": { "resource": "copilot:/session-1" } } ``` ### Action delivery -The server broadcasts `ActionEnvelope`s to subscribed clients: +The server broadcasts action envelopes as JSON-RPC notifications: + +``` +Server → Client: { "jsonrpc": "2.0", "method": "action", "params": { "envelope": { action, serverSeq, origin } } } +``` + - Root actions go to all clients subscribed to root state. - Session actions go to all clients subscribed to that session's URI. -Notifications go to all connected clients (no subscription required). +Protocol notifications (sessionAdded/sessionRemoved) are broadcast similarly: + +``` +Server → Client: { "jsonrpc": "2.0", "method": "notification", "params": { "notification": { type, ... } } } +``` + +### Commands as JSON-RPC requests + +Commands are JSON-RPC requests. The server returns a result or a JSON-RPC error: + +``` +Client → Server: { "jsonrpc": "2.0", "id": 2, "method": "createSession", "params": { session, provider?, model? } } +Server → Client: { "jsonrpc": "2.0", "id": 2, "result": null } +``` + +On failure: + +``` +Server → Client: { "jsonrpc": "2.0", "id": 2, "error": { "code": -32603, "message": "No agent for provider" } } +``` + +### Client-dispatched actions + +Actions are sent as notifications (fire-and-forget, write-ahead): + +``` +Client → Server: { "jsonrpc": "2.0", "method": "dispatchAction", "params": { clientSeq, action } } +``` ### Reconnection ``` -Client → Server: ClientReconnect { clientId, lastSeenServerSeq, subscriptions: URI[] } +Client → Server: { "jsonrpc": "2.0", "method": "reconnect", "params": { clientId, lastSeenServerSeq, subscriptions } } ``` -Server replays actions since `lastSeenServerSeq` from a bounded replay buffer. If the gap exceeds the buffer, sends fresh snapshots. Notifications are **not** replayed — the client should re-fetch the session list. +Server replays actions since `lastSeenServerSeq` from a bounded replay buffer. If the gap exceeds the buffer, sends fresh snapshots via a `reconnectResponse` notification. Notifications are **not** replayed — the client should re-fetch the session list. ## Write-ahead reconciliation @@ -303,7 +410,7 @@ The registry also maintains an exhaustive runtime map: ```typescript export const ACTION_INTRODUCED_IN: { readonly [K in IStateAction['type']]: number } = { - 'root/modelsChanged': 1, + 'root/agentsChanged': 1, 'session/turnStarted': 1, // ...every action type must have an entry }; @@ -383,7 +490,7 @@ src/vs/platform/agent/common/state/ ├── sessionState.ts # Immutable state types (RootState, SessionState, Turn, etc.) ├── sessionActions.ts # Action + notification discriminated unions, ActionEnvelope ├── sessionReducers.ts # Pure reducer functions (rootReducer, sessionReducer) -├── sessionProtocol.ts # Protocol messages (handshake, subscribe, reconnect, RPC) +├── sessionProtocol.ts # JSON-RPC message types, request params/results, type guards ├── sessionCapabilities.ts # Re-exports version constants + ProtocolCapabilities ├── sessionClientState.ts # Client-side state manager (confirmed + pending + reconciliation) └── versions/ diff --git a/src/vs/platform/agentHost/test/node/agentEventMapper.test.ts b/src/vs/platform/agentHost/test/node/agentEventMapper.test.ts new file mode 100644 index 00000000000..a152f4a454a --- /dev/null +++ b/src/vs/platform/agentHost/test/node/agentEventMapper.test.ts @@ -0,0 +1,221 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { URI } from '../../../../base/common/uri.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import type { + IAgentDeltaEvent, + IAgentErrorEvent, + IAgentIdleEvent, + IAgentMessageEvent, + IAgentPermissionRequestEvent, + IAgentReasoningEvent, + IAgentTitleChangedEvent, + IAgentToolCompleteEvent, + IAgentToolStartEvent, + IAgentUsageEvent, +} from '../../common/agentService.js'; +import type { + IDeltaAction, + IPermissionRequestAction, + IReasoningAction, + ISessionErrorAction, + ITitleChangedAction, + IToolCompleteAction, + IToolStartAction, + ITurnCompleteAction, + IUsageAction, +} from '../../common/state/sessionActions.js'; +import { ToolCallStatus } from '../../common/state/sessionState.js'; +import { mapProgressEventToAction } from '../../node/agentEventMapper.js'; + +suite('AgentEventMapper', () => { + + const session = URI.from({ scheme: 'copilot', path: '/test-session' }); + const turnId = 'turn-1'; + + ensureNoDisposablesAreLeakedInTestSuite(); + + test('delta event maps to session/delta action', () => { + const event: IAgentDeltaEvent = { + session, + type: 'delta', + messageId: 'msg-1', + content: 'hello world', + }; + + const action = mapProgressEventToAction(event, session, turnId); + assert.ok(action); + assert.strictEqual(action.type, 'session/delta'); + const delta = action as IDeltaAction; + assert.strictEqual(delta.content, 'hello world'); + assert.strictEqual(delta.session.toString(), session.toString()); + assert.strictEqual(delta.turnId, turnId); + }); + + test('tool_start event maps to session/toolStart action', () => { + const event: IAgentToolStartEvent = { + session, + type: 'tool_start', + toolCallId: 'tc-1', + toolName: 'readFile', + displayName: 'Read File', + invocationMessage: 'Reading file...', + toolInput: '/src/foo.ts', + toolKind: 'terminal', + language: 'shellscript', + }; + + const action = mapProgressEventToAction(event, session, turnId); + assert.ok(action); + assert.strictEqual(action.type, 'session/toolStart'); + const toolCall = (action as IToolStartAction).toolCall; + assert.strictEqual(toolCall.toolCallId, 'tc-1'); + assert.strictEqual(toolCall.toolName, 'readFile'); + assert.strictEqual(toolCall.displayName, 'Read File'); + assert.strictEqual(toolCall.invocationMessage, 'Reading file...'); + assert.strictEqual(toolCall.toolInput, '/src/foo.ts'); + assert.strictEqual(toolCall.toolKind, 'terminal'); + assert.strictEqual(toolCall.language, 'shellscript'); + assert.strictEqual(toolCall.status, ToolCallStatus.Running); + }); + + test('tool_complete event maps to session/toolComplete action', () => { + const event: IAgentToolCompleteEvent = { + session, + type: 'tool_complete', + toolCallId: 'tc-1', + success: true, + pastTenseMessage: 'Read file successfully', + toolOutput: 'file contents here', + }; + + const action = mapProgressEventToAction(event, session, turnId); + assert.ok(action); + assert.strictEqual(action.type, 'session/toolComplete'); + const complete = action as IToolCompleteAction; + assert.strictEqual(complete.toolCallId, 'tc-1'); + assert.strictEqual(complete.result.success, true); + assert.strictEqual(complete.result.pastTenseMessage, 'Read file successfully'); + assert.strictEqual(complete.result.toolOutput, 'file contents here'); + }); + + test('idle event maps to session/turnComplete action', () => { + const event: IAgentIdleEvent = { + session, + type: 'idle', + }; + + const action = mapProgressEventToAction(event, session, turnId); + assert.ok(action); + assert.strictEqual(action.type, 'session/turnComplete'); + const turnComplete = action as ITurnCompleteAction; + assert.strictEqual(turnComplete.session.toString(), session.toString()); + assert.strictEqual(turnComplete.turnId, turnId); + }); + + test('error event maps to session/error action', () => { + const event: IAgentErrorEvent = { + session, + type: 'error', + errorType: 'runtime', + message: 'Something went wrong', + stack: 'Error: Something went wrong\n at foo.ts:1', + }; + + const action = mapProgressEventToAction(event, session, turnId); + assert.ok(action); + assert.strictEqual(action.type, 'session/error'); + const errorAction = action as ISessionErrorAction; + assert.strictEqual(errorAction.error.errorType, 'runtime'); + assert.strictEqual(errorAction.error.message, 'Something went wrong'); + assert.strictEqual(errorAction.error.stack, 'Error: Something went wrong\n at foo.ts:1'); + }); + + test('usage event maps to session/usage action', () => { + const event: IAgentUsageEvent = { + session, + type: 'usage', + inputTokens: 100, + outputTokens: 50, + model: 'gpt-4', + cacheReadTokens: 25, + }; + + const action = mapProgressEventToAction(event, session, turnId); + assert.ok(action); + assert.strictEqual(action.type, 'session/usage'); + const usageAction = action as IUsageAction; + assert.strictEqual(usageAction.usage.inputTokens, 100); + assert.strictEqual(usageAction.usage.outputTokens, 50); + assert.strictEqual(usageAction.usage.model, 'gpt-4'); + assert.strictEqual(usageAction.usage.cacheReadTokens, 25); + }); + + test('title_changed event maps to session/titleChanged action', () => { + const event: IAgentTitleChangedEvent = { + session, + type: 'title_changed', + title: 'New Title', + }; + + const action = mapProgressEventToAction(event, session, turnId); + assert.ok(action); + assert.strictEqual(action.type, 'session/titleChanged'); + assert.strictEqual((action as ITitleChangedAction).title, 'New Title'); + }); + + test('permission_request event maps to session/permissionRequest action', () => { + const event: IAgentPermissionRequestEvent = { + session, + type: 'permission_request', + requestId: 'perm-1', + permissionKind: 'shell', + toolCallId: 'tc-2', + fullCommandText: 'rm -rf /', + intention: 'Delete all files', + rawRequest: '{}', + }; + + const action = mapProgressEventToAction(event, session, turnId); + assert.ok(action); + assert.strictEqual(action.type, 'session/permissionRequest'); + const req = (action as IPermissionRequestAction).request; + assert.strictEqual(req.requestId, 'perm-1'); + assert.strictEqual(req.permissionKind, 'shell'); + assert.strictEqual(req.toolCallId, 'tc-2'); + assert.strictEqual(req.fullCommandText, 'rm -rf /'); + assert.strictEqual(req.intention, 'Delete all files'); + }); + + test('reasoning event maps to session/reasoning action', () => { + const event: IAgentReasoningEvent = { + session, + type: 'reasoning', + content: 'Let me think about this...', + }; + + const action = mapProgressEventToAction(event, session, turnId); + assert.ok(action); + assert.strictEqual(action.type, 'session/reasoning'); + const reasoning = action as IReasoningAction; + assert.strictEqual(reasoning.content, 'Let me think about this...'); + assert.strictEqual(reasoning.turnId, turnId); + }); + + test('message event returns undefined', () => { + const event: IAgentMessageEvent = { + session, + type: 'message', + role: 'assistant', + messageId: 'msg-1', + content: 'Some full message', + }; + + const action = mapProgressEventToAction(event, session, turnId); + assert.strictEqual(action, undefined); + }); +}); diff --git a/src/vs/platform/agentHost/test/node/agentService.test.ts b/src/vs/platform/agentHost/test/node/agentService.test.ts index c68aa30c121..f394467fae9 100644 --- a/src/vs/platform/agentHost/test/node/agentService.test.ts +++ b/src/vs/platform/agentHost/test/node/agentService.test.ts @@ -10,6 +10,7 @@ import { URI } from '../../../../base/common/uri.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import { NullLogService } from '../../../log/common/log.js'; import { AgentSession, IAgent, IAgentCreateSessionConfig, IAgentDescriptor, IAgentMessageEvent, IAgentModelInfo, IAgentProgressEvent, IAgentSessionMetadata, IAgentToolCompleteEvent, IAgentToolStartEvent, AgentProvider } from '../../common/agentService.js'; +import { IActionEnvelope } from '../../common/state/sessionActions.js'; import { AgentService } from '../../node/agentService.js'; class MockAgent implements IAgent { @@ -107,16 +108,21 @@ suite('AgentService (node dispatcher)', () => { assert.throws(() => service.registerProvider(duplicate), /already registered/); }); - test('forwards progress events from registered providers', async () => { + test('maps progress events to protocol actions via onDidAction', async () => { service.registerProvider(copilotAgent); const session = await service.createSession({ provider: 'copilot' }); - const events: IAgentProgressEvent[] = []; - disposables.add(service.onDidSessionProgress(e => events.push(e))); + // Start a turn so there's an active turn to map events to + service.dispatchAction( + { type: 'session/turnStarted', session, turnId: 'turn-1', userMessage: { text: 'hello' } }, + 'test-client', 1, + ); + + const envelopes: IActionEnvelope[] = []; + disposables.add(service.onDidAction(e => envelopes.push(e))); copilotAgent.fireProgress({ session, type: 'delta', messageId: 'msg-1', content: 'hello' }); - assert.strictEqual(events.length, 1); - assert.strictEqual(events[0].type, 'delta'); + assert.ok(envelopes.some(e => e.action.type === 'session/delta')); }); }); @@ -161,39 +167,6 @@ suite('AgentService (node dispatcher)', () => { }); }); - // ---- sendMessage ---------------------------------------------------- - - suite('sendMessage', () => { - - test('dispatches to the correct provider based on session tracking', async () => { - service.registerProvider(copilotAgent); - - const session = await service.createSession({ provider: 'copilot' }); - await service.sendMessage(session, 'hello'); - - assert.strictEqual(copilotAgent.sendMessageCalls.length, 1); - assert.strictEqual(copilotAgent.sendMessageCalls[0].prompt, 'hello'); - }); - - test('infers provider from URI scheme for untracked sessions', async () => { - service.registerProvider(copilotAgent); - const session = AgentSession.uri('copilot', 'external-session'); - - await service.sendMessage(session, 'hello from untracked'); - - assert.strictEqual(copilotAgent.sendMessageCalls.length, 1); - }); - - test('falls back to default provider for unrecognized URI scheme', async () => { - service.registerProvider(copilotAgent); - const unknownSession = URI.from({ scheme: 'unknown', path: '/sess-1' }); - - // Should not throw -- falls back to the default provider - await service.sendMessage(unknownSession, 'hello'); - assert.strictEqual(copilotAgent.sendMessageCalls.length, 1); - }); - }); - // ---- disposeSession ------------------------------------------------- suite('disposeSession', () => { @@ -243,12 +216,16 @@ suite('AgentService (node dispatcher)', () => { assert.strictEqual(sessions.length, 1); }); - test('listModels aggregates models from all providers', async () => { + test('refreshModels publishes models in root state via agentsChanged', async () => { service.registerProvider(copilotAgent); - const models = await service.listModels(); - assert.strictEqual(models.length, 1); - assert.ok(models.some(m => m.provider === 'copilot')); + const envelopes: IActionEnvelope[] = []; + disposables.add(service.onDidAction(e => envelopes.push(e))); + + await service.refreshModels(); + + const agentsChanged = envelopes.find(e => e.action.type === 'root/agentsChanged'); + assert.ok(agentsChanged); }); }); diff --git a/src/vs/platform/agentHost/test/node/createAndSendMessageAsLocalAgent.sh b/src/vs/platform/agentHost/test/node/createAndSendMessageAsLocalAgent.sh new file mode 100755 index 00000000000..8a55277df90 --- /dev/null +++ b/src/vs/platform/agentHost/test/node/createAndSendMessageAsLocalAgent.sh @@ -0,0 +1,241 @@ +#!/usr/bin/env bash +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +# Launches Code OSS, switches to Local Agent mode, sends a chat message, +# waits for the response, and prints it to stdout. +# +# Usage: +# ./createAndSendMessageAsLocalAgent.sh "Hello, what can you do?" +# ./createAndSendMessageAsLocalAgent.sh --port 9225 "Explain this code" +# +# Options: +# --port CDP debugging port (default: 9224) +# --timeout Seconds to wait for response (default: 30) +# --no-kill Don't kill Code OSS after the test +# --skip-launch Assume Code OSS is already running on the given port +# +# Requires: agent-browser (npm install -g agent-browser, or use npx) + +set -e + +ROOT="$(cd "$(dirname "$0")/../../../../../.." && pwd)" +CDP_PORT=9224 +RESPONSE_TIMEOUT=30 +KILL_AFTER=true +SKIP_LAUNCH=false +MESSAGE="" + +# Parse arguments +while [[ $# -gt 0 ]]; do + case "$1" in + --port) + CDP_PORT="$2" + shift 2 + ;; + --timeout) + RESPONSE_TIMEOUT="$2" + shift 2 + ;; + --no-kill) + KILL_AFTER=false + shift + ;; + --skip-launch) + SKIP_LAUNCH=true + shift + ;; + -*) + echo "Unknown option: $1" >&2 + exit 1 + ;; + *) + MESSAGE="$1" + shift + ;; + esac +done + +if [ -z "$MESSAGE" ]; then + echo "Usage: $0 [--port ] [--timeout ] [--no-kill] [--skip-launch] " >&2 + exit 1 +fi + +AB="npx agent-browser" + +cleanup() { + if [ "$KILL_AFTER" = true ] && [ "$SKIP_LAUNCH" = false ]; then + $AB close 2>/dev/null || true + local PID + PID=$(lsof -t -i :"$CDP_PORT" 2>/dev/null || true) + if [ -n "$PID" ]; then + kill "$PID" 2>/dev/null || true + fi + fi +} +trap cleanup EXIT + +# ---- Step 1: Launch Code OSS ------------------------------------------------ + +if [ "$SKIP_LAUNCH" = false ]; then + # Check if already running + if lsof -i :"$CDP_PORT" >/dev/null 2>&1; then + echo "ERROR: Port $CDP_PORT already in use. Use --skip-launch or --port " >&2 + exit 1 + fi + + echo "Launching Code OSS on CDP port $CDP_PORT..." >&2 + cd "$ROOT" + VSCODE_SKIP_PRELAUNCH=1 ./scripts/code.sh --remote-debugging-port="$CDP_PORT" &>/dev/null & + + # Wait for it to start + echo "Waiting for Code OSS to start..." >&2 + for i in $(seq 1 20); do + if $AB connect "$CDP_PORT" 2>/dev/null; then + break + fi + sleep 2 + if [ "$i" -eq 20 ]; then + echo "ERROR: Code OSS did not start within 40 seconds" >&2 + exit 1 + fi + done +else + echo "Connecting to existing Code OSS on port $CDP_PORT..." >&2 + $AB connect "$CDP_PORT" 2>/dev/null || { + echo "ERROR: Cannot connect to Code OSS on port $CDP_PORT" >&2 + exit 1 + } +fi + +echo "Connected to Code OSS" >&2 + +# ---- Step 2: Switch to Local Agent mode ------------------------------------- + +# Check current session target +CURRENT_TARGET=$($AB snapshot -i 2>&1 | grep "Set Session Target" | head -1) + +if ! echo "$CURRENT_TARGET" | grep -q "Local Agent"; then + echo "Switching to Local Agent mode..." >&2 + + # Find and click the session target button + TARGET_REF=$($AB snapshot -i 2>&1 | grep "Set Session Target" | head -1 | grep -o 'ref=e[0-9]*' | head -1 | sed 's/ref=//') + if [ -z "$TARGET_REF" ]; then + echo "ERROR: Cannot find session target button" >&2 + exit 1 + fi + $AB click "@$TARGET_REF" 2>/dev/null + sleep 0.5 + + # Navigate to Local Agent via arrow keys + # Menu items: Local (checked), Copilot CLI, Cloud, Local Agent, ... + $AB press ArrowDown 2>/dev/null # Copilot CLI + $AB press ArrowDown 2>/dev/null # Cloud + $AB press ArrowDown 2>/dev/null # Local Agent + $AB press Enter 2>/dev/null + sleep 0.5 + + # Verify + VERIFY=$($AB snapshot -i 2>&1 | grep "Set Session Target" | head -1) + if echo "$VERIFY" | grep -q "Local Agent"; then + echo "Switched to Local Agent mode" >&2 + else + echo "WARNING: Could not confirm Local Agent mode. Current: $VERIFY" >&2 + fi +else + echo "Already in Local Agent mode" >&2 +fi + +# ---- Step 3: Focus chat input and type message ------------------------------ + +echo "Sending message: $MESSAGE" >&2 + +# Focus chat input via JavaScript mouse events (universal approach) +$AB eval ' +(() => { + const sidebar = document.querySelector(".part.auxiliarybar"); + if (!sidebar) return "no sidebar"; + const inputPart = sidebar.querySelector(".interactive-input-part"); + if (!inputPart) return "no input part"; + const editor = inputPart.querySelector(".monaco-editor"); + if (!editor) return "no editor"; + const rect = editor.getBoundingClientRect(); + const x = rect.x + rect.width / 2; + const y = rect.y + rect.height / 2; + editor.dispatchEvent(new MouseEvent("mousedown", { bubbles: true, clientX: x, clientY: y })); + editor.dispatchEvent(new MouseEvent("mouseup", { bubbles: true, clientX: x, clientY: y })); + editor.dispatchEvent(new MouseEvent("click", { bubbles: true, clientX: x, clientY: y })); + return "focused"; +})()' >/dev/null 2>&1 + +sleep 0.3 + +# Clear any existing text +$AB press Meta+a 2>/dev/null +$AB press Backspace 2>/dev/null + +# Type message character by character +for (( i=0; i<${#MESSAGE}; i++ )); do + CHAR="${MESSAGE:$i:1}" + case "$CHAR" in + " ") $AB press Space 2>/dev/null ;; + "?") $AB press Shift+/ 2>/dev/null ;; + "!") $AB press Shift+1 2>/dev/null ;; + ",") $AB press , 2>/dev/null ;; + ".") $AB press . 2>/dev/null ;; + "'") $AB press "'" 2>/dev/null ;; + '"') $AB press 'Shift+'"'" 2>/dev/null ;; + *) $AB press "$CHAR" 2>/dev/null ;; + esac +done + +# Verify text entered +ENTERED=$($AB eval ' +(() => { + const sidebar = document.querySelector(".part.auxiliarybar"); + const viewLines = sidebar?.querySelectorAll(".interactive-input-editor .view-line"); + return Array.from(viewLines || []).map(vl => vl.textContent).join(""); +})()' 2>&1 | tr -d '"') + +echo "Entered text: $ENTERED" >&2 + +# Send the message +$AB press Enter 2>/dev/null + +# ---- Step 4: Wait for response ---------------------------------------------- + +echo "Waiting for response (timeout: ${RESPONSE_TIMEOUT}s)..." >&2 + +RESPONSE="" +for i in $(seq 1 "$RESPONSE_TIMEOUT"); do + sleep 1 + RESPONSE=$($AB eval ' +(() => { + const sidebar = document.querySelector(".part.auxiliarybar"); + if (!sidebar) return ""; + const items = sidebar.querySelectorAll(".interactive-item-container"); + if (items.length < 2) return ""; + // Last item is the response + const lastItem = items[items.length - 1]; + const text = lastItem.textContent || ""; + // Check if it looks like a complete response (has content beyond the header) + if (text.length > 20) return text; + return ""; +})()' 2>&1 | sed 's/^"//;s/"$//') + + if [ -n "$RESPONSE" ]; then + break + fi +done + +if [ -z "$RESPONSE" ]; then + echo "ERROR: No response received within ${RESPONSE_TIMEOUT}s" >&2 + exit 1 +fi + +# ---- Step 5: Output response ------------------------------------------------ + +echo "---" >&2 +echo "$RESPONSE" diff --git a/src/vs/platform/agentHost/test/node/mockAgent.ts b/src/vs/platform/agentHost/test/node/mockAgent.ts new file mode 100644 index 00000000000..851faaff491 --- /dev/null +++ b/src/vs/platform/agentHost/test/node/mockAgent.ts @@ -0,0 +1,162 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { Emitter } from '../../../../base/common/event.js'; +import { URI } from '../../../../base/common/uri.js'; +import { AgentSession, type AgentProvider, type IAgent, type IAgentAttachment, type IAgentCreateSessionConfig, type IAgentDescriptor, type IAgentMessageEvent, type IAgentModelInfo, type IAgentProgressEvent, type IAgentSessionMetadata, type IAgentToolCompleteEvent, type IAgentToolStartEvent } from '../../common/agentService.js'; + +export class ScriptedMockAgent implements IAgent { + readonly id: AgentProvider = 'mock'; + + private readonly _onDidSessionProgress = new Emitter(); + readonly onDidSessionProgress = this._onDidSessionProgress.event; + + private readonly _sessions = new Map(); + private _nextId = 1; + + // Track pending permission requests + private readonly _pendingPermissions = new Map void>(); + // Track pending abort callbacks for slow responses + private readonly _pendingAborts = new Map void>(); + + getDescriptor(): IAgentDescriptor { + return { provider: 'mock', displayName: 'Mock Agent', description: 'Scripted test agent', requiresAuth: false }; + } + + async listModels(): Promise { + return [{ provider: 'mock', id: 'mock-model', name: 'Mock Model', maxContextWindow: 128000, supportsVision: false, supportsReasoningEffort: false }]; + } + + async listSessions(): Promise { + return [...this._sessions.values()].map(s => ({ session: s, startTime: Date.now(), modifiedTime: Date.now() })); + } + + async createSession(_config?: IAgentCreateSessionConfig): Promise { + const rawId = `mock-session-${this._nextId++}`; + const session = AgentSession.uri('mock', rawId); + this._sessions.set(rawId, session); + return session; + } + + async sendMessage(session: URI, prompt: string, _attachments?: IAgentAttachment[]): Promise { + switch (prompt) { + case 'hello': + this._fireSequence(session, [ + { type: 'delta', session, messageId: 'msg-1', content: 'Hello, world!' }, + { type: 'idle', session }, + ]); + break; + + case 'use-tool': + this._fireSequence(session, [ + { type: 'tool_start', session, toolCallId: 'tc-1', toolName: 'echo_tool', displayName: 'Echo Tool', invocationMessage: 'Running echo tool...' }, + { type: 'tool_complete', session, toolCallId: 'tc-1', success: true, pastTenseMessage: 'Ran echo tool', toolOutput: 'echoed' }, + { type: 'delta', session, messageId: 'msg-1', content: 'Tool done.' }, + { type: 'idle', session }, + ]); + break; + + case 'error': + this._fireSequence(session, [ + { type: 'error', session, errorType: 'test_error', message: 'Something went wrong' }, + ]); + break; + + case 'permission': { + // Fire permission_request, then wait for respondToPermissionRequest + const permEvent: IAgentProgressEvent = { + type: 'permission_request', + session, + requestId: 'perm-1', + permissionKind: 'shell', + fullCommandText: 'echo test', + intention: 'Run a test command', + rawRequest: JSON.stringify({ permissionKind: 'shell', fullCommandText: 'echo test', intention: 'Run a test command' }), + }; + setTimeout(() => this._onDidSessionProgress.fire(permEvent), 10); + this._pendingPermissions.set('perm-1', (approved) => { + if (approved) { + this._fireSequence(session, [ + { type: 'delta', session, messageId: 'msg-1', content: 'Allowed.' }, + { type: 'idle', session }, + ]); + } + }); + break; + } + + case 'with-usage': + this._fireSequence(session, [ + { type: 'delta', session, messageId: 'msg-1', content: 'Usage response.' }, + { type: 'usage', session, inputTokens: 100, outputTokens: 50, model: 'mock-model' }, + { type: 'idle', session }, + ]); + break; + + case 'slow': { + // Slow response for cancel testing — fires delta after a long delay + const timer = setTimeout(() => { + this._fireSequence(session, [ + { type: 'delta', session, messageId: 'msg-1', content: 'Slow response.' }, + { type: 'idle', session }, + ]); + }, 5000); + this._pendingAborts.set(session.toString(), () => clearTimeout(timer)); + break; + } + + default: + this._fireSequence(session, [ + { type: 'delta', session, messageId: 'msg-1', content: 'Unknown prompt: ' + prompt }, + { type: 'idle', session }, + ]); + break; + } + } + + async getSessionMessages(_session: URI): Promise<(IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[]> { + return []; + } + + async disposeSession(session: URI): Promise { + this._sessions.delete(AgentSession.id(session)); + } + + async abortSession(session: URI): Promise { + const callback = this._pendingAborts.get(session.toString()); + if (callback) { + this._pendingAborts.delete(session.toString()); + callback(); + } + } + + async changeModel(_session: URI, _model: string): Promise { + // Mock agent doesn't track model state + } + + respondToPermissionRequest(requestId: string, approved: boolean): void { + const callback = this._pendingPermissions.get(requestId); + if (callback) { + this._pendingPermissions.delete(requestId); + callback(approved); + } + } + + async setAuthToken(_token: string): Promise { } + + async shutdown(): Promise { } + + dispose(): void { + this._onDidSessionProgress.dispose(); + } + + private _fireSequence(session: URI, events: IAgentProgressEvent[]): void { + let delay = 0; + for (const event of events) { + delay += 10; + setTimeout(() => this._onDidSessionProgress.fire(event), delay); + } + } +} diff --git a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts new file mode 100644 index 00000000000..fbde7d8980c --- /dev/null +++ b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts @@ -0,0 +1,308 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { Emitter } from '../../../../base/common/event.js'; +import { DisposableStore } from '../../../../base/common/lifecycle.js'; +import { URI } from '../../../../base/common/uri.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import { NullLogService } from '../../../log/common/log.js'; +import type { ISessionAction } from '../../common/state/sessionActions.js'; +import { isJsonRpcNotification, isJsonRpcResponse, type ICreateSessionParams, type IProtocolMessage, type IProtocolNotification, type IServerHelloParams, type IStateSnapshot } from '../../common/state/sessionProtocol.js'; +import { SessionStatus, type ISessionSummary } from '../../common/state/sessionState.js'; +import { PROTOCOL_VERSION } from '../../common/state/sessionCapabilities.js'; +import type { IProtocolServer, IProtocolTransport } from '../../common/state/sessionTransport.js'; +import { ProtocolServerHandler, type IProtocolSideEffectHandler } from '../../node/protocolServerHandler.js'; +import { SessionStateManager } from '../../node/sessionStateManager.js'; + +// ---- Mock helpers ----------------------------------------------------------- + +class MockProtocolTransport implements IProtocolTransport { + private readonly _onMessage = new Emitter(); + readonly onMessage = this._onMessage.event; + private readonly _onClose = new Emitter(); + readonly onClose = this._onClose.event; + + readonly sent: IProtocolMessage[] = []; + + send(message: IProtocolMessage): void { + this.sent.push(message); + } + + simulateMessage(msg: IProtocolMessage): void { + this._onMessage.fire(msg); + } + + simulateClose(): void { + this._onClose.fire(); + } + + dispose(): void { + this._onMessage.dispose(); + this._onClose.dispose(); + } +} + +class MockProtocolServer implements IProtocolServer { + private readonly _onConnection = new Emitter(); + readonly onConnection = this._onConnection.event; + readonly address = 'mock://test'; + + simulateConnection(transport: IProtocolTransport): void { + this._onConnection.fire(transport); + } + + dispose(): void { + this._onConnection.dispose(); + } +} + +class MockSideEffectHandler implements IProtocolSideEffectHandler { + readonly handledActions: ISessionAction[] = []; + handleAction(action: ISessionAction): void { + this.handledActions.push(action); + } + async handleCreateSession(_command: ICreateSessionParams): Promise { } + handleDisposeSession(_session: URI): void { } + async handleListSessions(): Promise { return []; } +} + +// ---- Helpers ---------------------------------------------------------------- + +function notification(method: string, params?: unknown): IProtocolMessage { + return { jsonrpc: '2.0', method, params } as IProtocolMessage; +} + +function request(id: number, method: string, params?: unknown): IProtocolMessage { + return { jsonrpc: '2.0', id, method, params } as IProtocolMessage; +} + +function findNotification(sent: IProtocolMessage[], method: string): IProtocolNotification | undefined { + return sent.find(isJsonRpcNotification) as IProtocolNotification | undefined; +} + +function findNotifications(sent: IProtocolMessage[], method: string): IProtocolNotification[] { + return sent.filter(isJsonRpcNotification) as IProtocolNotification[]; +} + +function findResponse(sent: IProtocolMessage[], id: number): IProtocolMessage | undefined { + return sent.find(isJsonRpcResponse) as IProtocolMessage | undefined; +} + +// ---- Tests ------------------------------------------------------------------ + +suite('ProtocolServerHandler', () => { + + let disposables: DisposableStore; + let stateManager: SessionStateManager; + let server: MockProtocolServer; + let sideEffects: MockSideEffectHandler; + + const sessionUri = URI.from({ scheme: 'copilot', path: '/test-session' }); + + function makeSessionSummary(resource?: URI): ISessionSummary { + return { + resource: resource ?? sessionUri, + provider: 'copilot', + title: 'Test', + status: SessionStatus.Idle, + createdAt: Date.now(), + modifiedAt: Date.now(), + }; + } + + function connectClient(clientId: string, initialSubscriptions?: readonly URI[]): MockProtocolTransport { + const transport = new MockProtocolTransport(); + server.simulateConnection(transport); + transport.simulateMessage(notification('initialize', { + protocolVersion: PROTOCOL_VERSION, + clientId, + initialSubscriptions, + })); + return transport; + } + + setup(() => { + disposables = new DisposableStore(); + stateManager = disposables.add(new SessionStateManager(new NullLogService())); + server = disposables.add(new MockProtocolServer()); + sideEffects = new MockSideEffectHandler(); + disposables.add(new ProtocolServerHandler( + stateManager, + server, + sideEffects, + new NullLogService(), + )); + }); + + teardown(() => { + disposables.dispose(); + }); + + ensureNoDisposablesAreLeakedInTestSuite(); + + test('handshake sends serverHello notification', () => { + const transport = connectClient('client-1'); + + const hello = findNotification(transport.sent, 'serverHello'); + assert.ok(hello, 'should have sent serverHello'); + const params = hello.params as IServerHelloParams; + assert.strictEqual(params.protocolVersion, PROTOCOL_VERSION); + assert.strictEqual(params.serverSeq, stateManager.serverSeq); + }); + + test('handshake with initialSubscriptions returns snapshots', () => { + stateManager.createSession(makeSessionSummary()); + + const transport = connectClient('client-1', [sessionUri]); + + const hello = findNotification(transport.sent, 'serverHello'); + assert.ok(hello); + const params = hello.params as IServerHelloParams; + assert.strictEqual(params.snapshots.length, 1); + assert.strictEqual(params.snapshots[0].resource.toString(), sessionUri.toString()); + }); + + test('subscribe request returns snapshot', async () => { + stateManager.createSession(makeSessionSummary()); + + const transport = connectClient('client-1'); + transport.sent.length = 0; + + transport.simulateMessage(request(1, 'subscribe', { resource: sessionUri })); + + // Wait for async response + await new Promise(resolve => setTimeout(resolve, 10)); + + const resp = findResponse(transport.sent, 1); + assert.ok(resp, 'should have sent response'); + const snapshot = (resp as { result: IStateSnapshot }).result; + assert.strictEqual(snapshot.resource.toString(), sessionUri.toString()); + }); + + test('client action is dispatched and echoed', () => { + stateManager.createSession(makeSessionSummary()); + stateManager.dispatchServerAction({ type: 'session/ready', session: sessionUri }); + + const transport = connectClient('client-1', [sessionUri]); + transport.sent.length = 0; + + transport.simulateMessage(notification('dispatchAction', { + clientSeq: 1, + action: { + type: 'session/turnStarted', + session: sessionUri, + turnId: 'turn-1', + userMessage: { text: 'hello' }, + }, + })); + + const actionMsgs = findNotifications(transport.sent, 'action'); + const turnStarted = actionMsgs.find(m => { + const params = m.params as { envelope: { action: { type: string } } }; + return params.envelope.action.type === 'session/turnStarted'; + }); + assert.ok(turnStarted, 'should have echoed turnStarted'); + const envelope = (turnStarted!.params as { envelope: { origin: { clientId: string; clientSeq: number } } }).envelope; + assert.strictEqual(envelope.origin.clientId, 'client-1'); + assert.strictEqual(envelope.origin.clientSeq, 1); + }); + + test('actions are scoped to subscribed sessions', () => { + stateManager.createSession(makeSessionSummary()); + stateManager.dispatchServerAction({ type: 'session/ready', session: sessionUri }); + + const transportA = connectClient('client-a', [sessionUri]); + const transportB = connectClient('client-b'); + + transportA.sent.length = 0; + transportB.sent.length = 0; + + stateManager.dispatchServerAction({ + type: 'session/titleChanged', + session: sessionUri, + title: 'New Title', + }); + + assert.strictEqual(findNotifications(transportA.sent, 'action').length, 1); + assert.strictEqual(findNotifications(transportB.sent, 'action').length, 0); + }); + + test('notifications are broadcast to all clients', () => { + const transportA = connectClient('client-a'); + const transportB = connectClient('client-b'); + + transportA.sent.length = 0; + transportB.sent.length = 0; + + stateManager.createSession(makeSessionSummary()); + + assert.strictEqual(findNotifications(transportA.sent, 'notification').length, 1); + assert.strictEqual(findNotifications(transportB.sent, 'notification').length, 1); + }); + + test('reconnect replays missed actions', () => { + stateManager.createSession(makeSessionSummary()); + stateManager.dispatchServerAction({ type: 'session/ready', session: sessionUri }); + + const transport1 = connectClient('client-r', [sessionUri]); + const hello = findNotification(transport1.sent, 'serverHello'); + const helloSeq = (hello!.params as IServerHelloParams).serverSeq; + transport1.simulateClose(); + + stateManager.dispatchServerAction({ type: 'session/titleChanged', session: sessionUri, title: 'Title A' }); + stateManager.dispatchServerAction({ type: 'session/titleChanged', session: sessionUri, title: 'Title B' }); + + const transport2 = new MockProtocolTransport(); + server.simulateConnection(transport2); + transport2.simulateMessage(notification('reconnect', { + clientId: 'client-r', + lastSeenServerSeq: helloSeq, + subscriptions: [sessionUri], + })); + + const replayed = findNotifications(transport2.sent, 'action'); + assert.strictEqual(replayed.length, 2); + }); + + test('reconnect sends fresh snapshots when gap too large', () => { + stateManager.createSession(makeSessionSummary()); + stateManager.dispatchServerAction({ type: 'session/ready', session: sessionUri }); + + const transport1 = connectClient('client-g', [sessionUri]); + transport1.simulateClose(); + + for (let i = 0; i < 1100; i++) { + stateManager.dispatchServerAction({ type: 'session/titleChanged', session: sessionUri, title: `Title ${i}` }); + } + + const transport2 = new MockProtocolTransport(); + server.simulateConnection(transport2); + transport2.simulateMessage(notification('reconnect', { + clientId: 'client-g', + lastSeenServerSeq: 0, + subscriptions: [sessionUri], + })); + + const reconnectResp = findNotification(transport2.sent, 'reconnectResponse'); + assert.ok(reconnectResp, 'should receive a reconnectResponse'); + const params = reconnectResp!.params as { snapshots: IStateSnapshot[] }; + assert.ok(params.snapshots.length > 0, 'should contain snapshots'); + }); + + test('client disconnect cleans up', () => { + stateManager.createSession(makeSessionSummary()); + stateManager.dispatchServerAction({ type: 'session/ready', session: sessionUri }); + + const transport = connectClient('client-d', [sessionUri]); + transport.sent.length = 0; + + transport.simulateClose(); + + stateManager.dispatchServerAction({ type: 'session/titleChanged', session: sessionUri, title: 'After Disconnect' }); + + assert.strictEqual(transport.sent.length, 0); + }); +}); diff --git a/src/vs/platform/agentHost/test/node/protocolWebSocket.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocolWebSocket.integrationTest.ts new file mode 100644 index 00000000000..b8f49c502cb --- /dev/null +++ b/src/vs/platform/agentHost/test/node/protocolWebSocket.integrationTest.ts @@ -0,0 +1,663 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { ChildProcess, fork } from 'child_process'; +import { fileURLToPath } from 'url'; +import { WebSocket } from 'ws'; +import { URI } from '../../../../base/common/uri.js'; +import { PROTOCOL_VERSION } from '../../common/state/sessionCapabilities.js'; +import { + isJsonRpcNotification, + isJsonRpcResponse, + type IActionBroadcastParams, + type IFetchTurnsResult, + type IJsonRpcErrorResponse, + type IJsonRpcSuccessResponse, + type IListSessionsResult, + type INotificationBroadcastParams, + type IProtocolMessage, + type IProtocolNotification, + type IServerHelloParams, + type IStateSnapshot, +} from '../../common/state/sessionProtocol.js'; +import type { IDeltaAction, ISessionAddedNotification, ISessionRemovedNotification, IUsageAction } from '../../common/state/sessionActions.js'; +import type { ISessionState } from '../../common/state/sessionState.js'; + +// ---- JSON serialization helpers (mirror webSocketTransport.ts) -------------- + +function uriReplacer(_key: string, value: unknown): unknown { + if (value instanceof URI) { + return value.toJSON(); + } + if (value instanceof Map) { + return { $type: 'Map', entries: [...value.entries()] }; + } + return value; +} + +function uriReviver(_key: string, value: unknown): unknown { + if (value && typeof value === 'object') { + const obj = value as Record; + if (obj.$mid === 1) { + return URI.revive(value as URI); + } + if (obj.$type === 'Map' && Array.isArray(obj.entries)) { + return new Map(obj.entries as [unknown, unknown][]); + } + } + return value; +} + +// ---- JSON-RPC test client --------------------------------------------------- + +interface IPendingCall { + resolve: (result: unknown) => void; + reject: (err: Error) => void; +} + +class TestProtocolClient { + private readonly _ws: WebSocket; + private _nextId = 1; + private readonly _pendingCalls = new Map(); + private readonly _notifications: IProtocolNotification[] = []; + private readonly _notifWaiters: { predicate: (n: IProtocolNotification) => boolean; resolve: (n: IProtocolNotification) => void; reject: (err: Error) => void }[] = []; + + constructor(port: number) { + this._ws = new WebSocket(`ws://127.0.0.1:${port}`); + } + + async connect(): Promise { + return new Promise((resolve, reject) => { + this._ws.on('open', () => { + this._ws.on('message', (data: Buffer | string) => { + const text = typeof data === 'string' ? data : data.toString('utf-8'); + const msg = JSON.parse(text, uriReviver); + this._handleMessage(msg); + }); + resolve(); + }); + this._ws.on('error', reject); + }); + } + + private _handleMessage(msg: IProtocolMessage): void { + if (isJsonRpcResponse(msg)) { + // JSON-RPC response — resolve pending call + const pending = this._pendingCalls.get(msg.id); + if (pending) { + this._pendingCalls.delete(msg.id); + const errResp = msg as IJsonRpcErrorResponse; + if (errResp.error) { + pending.reject(new Error(errResp.error.message)); + } else { + pending.resolve((msg as IJsonRpcSuccessResponse).result); + } + } + } else if (isJsonRpcNotification(msg)) { + // JSON-RPC notification from server + const notif = msg; + // Check waiters first + for (let i = this._notifWaiters.length - 1; i >= 0; i--) { + if (this._notifWaiters[i].predicate(notif)) { + const waiter = this._notifWaiters.splice(i, 1)[0]; + waiter.resolve(notif); + } + } + this._notifications.push(notif); + } + } + + /** Send a JSON-RPC notification (fire-and-forget). */ + notify(method: string, params?: unknown): void { + const msg: IProtocolMessage = { jsonrpc: '2.0', method, params }; + this._ws.send(JSON.stringify(msg, uriReplacer)); + } + + /** Send a JSON-RPC request and await the response. */ + call(method: string, params?: unknown, timeoutMs = 5000): Promise { + const id = this._nextId++; + const msg: IProtocolMessage = { jsonrpc: '2.0', id, method, params }; + this._ws.send(JSON.stringify(msg, uriReplacer)); + + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + this._pendingCalls.delete(id); + reject(new Error(`Timeout waiting for response to ${method} (id=${id}, ${timeoutMs}ms)`)); + }, timeoutMs); + + this._pendingCalls.set(id, { + resolve: result => { clearTimeout(timer); resolve(result as T); }, + reject: err => { clearTimeout(timer); reject(err); }, + }); + }); + } + + /** Wait for a server notification matching a predicate. */ + waitForNotification(predicate: (n: IProtocolNotification) => boolean, timeoutMs = 5000): Promise { + const existing = this._notifications.find(predicate); + if (existing) { + return Promise.resolve(existing); + } + + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + const idx = this._notifWaiters.findIndex(w => w.resolve === resolve); + if (idx >= 0) { + this._notifWaiters.splice(idx, 1); + } + reject(new Error(`Timeout waiting for notification (${timeoutMs}ms)`)); + }, timeoutMs); + + this._notifWaiters.push({ + predicate, + resolve: n => { clearTimeout(timer); resolve(n); }, + reject, + }); + }); + } + + /** Return all received notifications matching a predicate. */ + receivedNotifications(predicate?: (n: IProtocolNotification) => boolean): IProtocolNotification[] { + return predicate ? this._notifications.filter(predicate) : [...this._notifications]; + } + + close(): void { + for (const w of this._notifWaiters) { + w.reject(new Error('Client closed')); + } + this._notifWaiters.length = 0; + for (const [, p] of this._pendingCalls) { + p.reject(new Error('Client closed')); + } + this._pendingCalls.clear(); + this._ws.close(); + } + + clearReceived(): void { + this._notifications.length = 0; + } +} + +// ---- Server process lifecycle ----------------------------------------------- + +async function startServer(): Promise<{ process: ChildProcess; port: number }> { + return new Promise((resolve, reject) => { + const serverPath = fileURLToPath(new URL('../../node/agentHostServerMain.js', import.meta.url)); + const child = fork(serverPath, ['--enable-mock-agent', '--quiet', '--port', '0'], { + stdio: ['pipe', 'pipe', 'pipe', 'ipc'], + }); + + const timeout = setTimeout(() => { + child.kill(); + reject(new Error('Server startup timed out')); + }, 10_000); + + child.stdout!.on('data', (data: Buffer) => { + const text = data.toString(); + const match = text.match(/READY:(\d+)/); + if (match) { + clearTimeout(timeout); + resolve({ process: child, port: parseInt(match[1], 10) }); + } + }); + + child.stderr!.on('data', (data: Buffer) => { + console.error('[TestServer]', data.toString()); + }); + + child.on('error', err => { + clearTimeout(timeout); + reject(err); + }); + + child.on('exit', code => { + clearTimeout(timeout); + reject(new Error(`Server exited prematurely with code ${code}`)); + }); + }); +} + +// ---- Helpers ---------------------------------------------------------------- + +let sessionCounter = 0; + +function nextSessionUri(): URI { + return URI.from({ scheme: 'mock', path: `/test-session-${++sessionCounter}` }); +} + +function isActionNotification(n: IProtocolNotification, actionType: string): boolean { + if (n.method !== 'action') { + return false; + } + const params = n.params as IActionBroadcastParams; + return params.envelope.action.type === actionType; +} + +function getActionParams(n: IProtocolNotification): IActionBroadcastParams { + return n.params as IActionBroadcastParams; +} + +/** Perform handshake, create a session, subscribe, and return its URI. */ +async function createAndSubscribeSession(c: TestProtocolClient, clientId: string): Promise { + c.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId }); + await c.waitForNotification(n => n.method === 'serverHello'); + + await c.call('createSession', { session: nextSessionUri(), provider: 'mock' }); + + const notif = await c.waitForNotification(n => + n.method === 'notification' && (n.params as INotificationBroadcastParams).notification.type === 'notify/sessionAdded' + ); + const realSessionUri = ((notif.params as INotificationBroadcastParams).notification as ISessionAddedNotification).summary.resource; + + await c.call('subscribe', { resource: realSessionUri }); + c.clearReceived(); + + return realSessionUri; +} + +function dispatchTurnStarted(c: TestProtocolClient, session: URI, turnId: string, text: string, clientSeq: number): void { + c.notify('dispatchAction', { + clientSeq, + action: { + type: 'session/turnStarted', + session, + turnId, + userMessage: { text }, + }, + }); +} + +// ---- Test suite ------------------------------------------------------------- + +suite('Protocol WebSocket E2E', function () { + + let server: { process: ChildProcess; port: number }; + let client: TestProtocolClient; + + suiteSetup(async function () { + this.timeout(15_000); + server = await startServer(); + }); + + suiteTeardown(function () { + server.process.kill(); + }); + + setup(async function () { + this.timeout(10_000); + client = new TestProtocolClient(server.port); + await client.connect(); + }); + + teardown(function () { + client.close(); + }); + + // 1. Handshake + test('handshake returns serverHello with protocol version', async function () { + this.timeout(5_000); + + client.notify('initialize', { + protocolVersion: PROTOCOL_VERSION, + clientId: 'test-handshake', + initialSubscriptions: [URI.from({ scheme: 'agenthost', path: '/root' })], + }); + + const hello = await client.waitForNotification(n => n.method === 'serverHello'); + const params = hello.params as IServerHelloParams; + assert.strictEqual(params.protocolVersion, PROTOCOL_VERSION); + assert.ok(params.serverSeq >= 0); + assert.ok(params.snapshots.length >= 1, 'should have root state snapshot'); + }); + + // 2. Create session + test('create session triggers sessionAdded notification', async function () { + this.timeout(10_000); + + client.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-create-session' }); + await client.waitForNotification(n => n.method === 'serverHello'); + + await client.call('createSession', { session: nextSessionUri(), provider: 'mock' }); + + const notif = await client.waitForNotification(n => + n.method === 'notification' && (n.params as INotificationBroadcastParams).notification.type === 'notify/sessionAdded' + ); + const notification = (notif.params as INotificationBroadcastParams).notification as ISessionAddedNotification; + assert.strictEqual(notification.summary.resource.scheme, 'mock'); + assert.strictEqual(notification.summary.provider, 'mock'); + }); + + // 3. Send message and receive response + test('send message and receive delta + turnComplete', async function () { + this.timeout(10_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-send-message'); + dispatchTurnStarted(client, sessionUri, 'turn-1', 'hello', 1); + + const delta = await client.waitForNotification(n => isActionNotification(n, 'session/delta')); + const deltaAction = getActionParams(delta).envelope.action; + assert.strictEqual(deltaAction.type, 'session/delta'); + if (deltaAction.type === 'session/delta') { + assert.strictEqual(deltaAction.content, 'Hello, world!'); + } + + await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + }); + + // 4. Tool invocation lifecycle + test('tool invocation: toolStart → toolComplete → delta → turnComplete', async function () { + this.timeout(10_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-tool-invocation'); + dispatchTurnStarted(client, sessionUri, 'turn-tool', 'use-tool', 1); + + await client.waitForNotification(n => isActionNotification(n, 'session/toolStart')); + const toolComplete = await client.waitForNotification(n => isActionNotification(n, 'session/toolComplete')); + const tcAction = getActionParams(toolComplete).envelope.action; + if (tcAction.type === 'session/toolComplete') { + assert.strictEqual(tcAction.result.success, true); + } + await client.waitForNotification(n => isActionNotification(n, 'session/delta')); + await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + }); + + // 5. Error handling + test('error prompt triggers session/error', async function () { + this.timeout(10_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-error'); + dispatchTurnStarted(client, sessionUri, 'turn-err', 'error', 1); + + const errorNotif = await client.waitForNotification(n => isActionNotification(n, 'session/error')); + const errorAction = getActionParams(errorNotif).envelope.action; + if (errorAction.type === 'session/error') { + assert.strictEqual(errorAction.error.message, 'Something went wrong'); + } + }); + + // 6. Permission flow + test('permission request → resolve → response', async function () { + this.timeout(10_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-permission'); + dispatchTurnStarted(client, sessionUri, 'turn-perm', 'permission', 1); + + await client.waitForNotification(n => isActionNotification(n, 'session/permissionRequest')); + + client.notify('dispatchAction', { + clientSeq: 2, + action: { + type: 'session/permissionResolved', + session: sessionUri, + turnId: 'turn-perm', + requestId: 'perm-1', + approved: true, + }, + }); + + const delta = await client.waitForNotification(n => isActionNotification(n, 'session/delta')); + const content = (getActionParams(delta).envelope.action as IDeltaAction).content; + assert.strictEqual(content, 'Allowed.'); + + await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + }); + + // 7. Session list + test('listSessions returns sessions', async function () { + this.timeout(10_000); + + client.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-list-sessions' }); + await client.waitForNotification(n => n.method === 'serverHello'); + + await client.call('createSession', { session: nextSessionUri(), provider: 'mock' }); + await client.waitForNotification(n => + n.method === 'notification' && (n.params as INotificationBroadcastParams).notification.type === 'notify/sessionAdded' + ); + + const result = await client.call('listSessions'); + assert.ok(Array.isArray(result.sessions)); + assert.ok(result.sessions.length >= 1, 'should have at least one session'); + }); + + // 8. Reconnect + test('reconnect replays missed actions', async function () { + this.timeout(15_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-reconnect'); + dispatchTurnStarted(client, sessionUri, 'turn-recon', 'hello', 1); + await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + + const allActions = client.receivedNotifications(n => n.method === 'action'); + assert.ok(allActions.length > 0); + const missedFromSeq = getActionParams(allActions[0]).envelope.serverSeq - 1; + + client.close(); + + const client2 = new TestProtocolClient(server.port); + await client2.connect(); + client2.notify('reconnect', { + clientId: 'test-reconnect', + lastSeenServerSeq: missedFromSeq, + subscriptions: [sessionUri], + }); + + await new Promise(resolve => setTimeout(resolve, 500)); + + const replayed = client2.receivedNotifications(); + assert.ok(replayed.length > 0, 'should receive replayed actions or reconnect response'); + const hasActions = replayed.some(n => n.method === 'action'); + const hasReconnect = replayed.some(n => n.method === 'reconnectResponse'); + assert.ok(hasActions || hasReconnect); + + client2.close(); + }); + + // ---- Gap tests: functionality bugs ---------------------------------------- + + test('usage info is captured on completed turn', async function () { + this.timeout(10_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-usage'); + dispatchTurnStarted(client, sessionUri, 'turn-usage', 'with-usage', 1); + + const usageNotif = await client.waitForNotification(n => isActionNotification(n, 'session/usage')); + const usageAction = getActionParams(usageNotif).envelope.action as IUsageAction; + assert.strictEqual(usageAction.usage.inputTokens, 100); + assert.strictEqual(usageAction.usage.outputTokens, 50); + + await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + + const snapshot = await client.call('subscribe', { resource: sessionUri }); + const state = snapshot.state as ISessionState; + assert.ok(state.turns.length >= 1); + const turn = state.turns[state.turns.length - 1]; + assert.ok(turn.usage); + assert.strictEqual(turn.usage!.inputTokens, 100); + assert.strictEqual(turn.usage!.outputTokens, 50); + }); + + test('modifiedAt updates on turn completion', async function () { + this.timeout(10_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-modifiedAt'); + + const initialSnapshot = await client.call('subscribe', { resource: sessionUri }); + const initialModifiedAt = (initialSnapshot.state as ISessionState).summary.modifiedAt; + + await new Promise(resolve => setTimeout(resolve, 50)); + + dispatchTurnStarted(client, sessionUri, 'turn-mod', 'hello', 1); + await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + + const updatedSnapshot = await client.call('subscribe', { resource: sessionUri }); + const updatedModifiedAt = (updatedSnapshot.state as ISessionState).summary.modifiedAt; + assert.ok(updatedModifiedAt >= initialModifiedAt); + }); + + test('createSession with invalid provider does not crash server', async function () { + this.timeout(10_000); + + client.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-invalid-create' }); + await client.waitForNotification(n => n.method === 'serverHello'); + + // This should return a JSON-RPC error + let gotError = false; + try { + await client.call('createSession', { session: nextSessionUri(), provider: 'nonexistent' }); + } catch { + gotError = true; + } + assert.ok(gotError, 'should have received an error for invalid provider'); + + // Server should still be functional + await client.call('createSession', { session: nextSessionUri(), provider: 'mock' }); + const notif = await client.waitForNotification(n => + n.method === 'notification' && (n.params as INotificationBroadcastParams).notification.type === 'notify/sessionAdded' + ); + assert.ok(notif); + }); + + test('fetchTurns returns completed turn history', async function () { + this.timeout(15_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-fetchTurns'); + + dispatchTurnStarted(client, sessionUri, 'turn-ft-1', 'hello', 1); + await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + + dispatchTurnStarted(client, sessionUri, 'turn-ft-2', 'hello', 2); + await new Promise(resolve => setTimeout(resolve, 200)); + await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + + const result = await client.call('fetchTurns', { session: sessionUri, startTurn: 0, count: 10 }); + assert.ok(result.turns.length >= 2); + assert.ok(result.totalTurns >= 2); + }); + + // ---- Gap tests: coverage --------------------------------------------------- + + test('dispose session sends sessionRemoved notification', async function () { + this.timeout(10_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-dispose'); + await client.call('disposeSession', { session: sessionUri }); + + const notif = await client.waitForNotification(n => + n.method === 'notification' && (n.params as INotificationBroadcastParams).notification.type === 'notify/sessionRemoved' + ); + const removed = (notif.params as INotificationBroadcastParams).notification as ISessionRemovedNotification; + assert.strictEqual(removed.session.toString(), sessionUri.toString()); + }); + + test('cancel turn stops in-progress processing', async function () { + this.timeout(10_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-cancel'); + dispatchTurnStarted(client, sessionUri, 'turn-cancel', 'slow', 1); + + client.notify('dispatchAction', { + clientSeq: 2, + action: { type: 'session/turnCancelled', session: sessionUri, turnId: 'turn-cancel' }, + }); + + await client.waitForNotification(n => isActionNotification(n, 'session/turnCancelled')); + + const snapshot = await client.call('subscribe', { resource: sessionUri }); + const state = snapshot.state as ISessionState; + assert.ok(state.turns.length >= 1); + assert.strictEqual(state.turns[state.turns.length - 1].state, 'cancelled'); + }); + + test('multiple sequential turns accumulate in history', async function () { + this.timeout(15_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-multi-turns'); + + dispatchTurnStarted(client, sessionUri, 'turn-m1', 'hello', 1); + await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + + dispatchTurnStarted(client, sessionUri, 'turn-m2', 'hello', 2); + await new Promise(resolve => setTimeout(resolve, 200)); + await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + + const snapshot = await client.call('subscribe', { resource: sessionUri }); + const state = snapshot.state as ISessionState; + assert.ok(state.turns.length >= 2, `expected >= 2 turns but got ${state.turns.length}`); + assert.strictEqual(state.turns[0].id, 'turn-m1'); + assert.strictEqual(state.turns[1].id, 'turn-m2'); + }); + + test('two clients on same session both see actions', async function () { + this.timeout(10_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-multi-client-1'); + + const client2 = new TestProtocolClient(server.port); + await client2.connect(); + client2.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-multi-client-2' }); + await client2.waitForNotification(n => n.method === 'serverHello'); + await client2.call('subscribe', { resource: sessionUri }); + client2.clearReceived(); + + dispatchTurnStarted(client, sessionUri, 'turn-mc', 'hello', 1); + + const d1 = await client.waitForNotification(n => isActionNotification(n, 'session/delta')); + const d2 = await client2.waitForNotification(n => isActionNotification(n, 'session/delta')); + assert.ok(d1); + assert.ok(d2); + + await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + await client2.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + + client2.close(); + }); + + test('unsubscribe stops receiving session actions', async function () { + this.timeout(10_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-unsubscribe'); + client.notify('unsubscribe', { resource: sessionUri }); + await new Promise(resolve => setTimeout(resolve, 100)); + client.clearReceived(); + + const client2 = new TestProtocolClient(server.port); + await client2.connect(); + client2.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-unsub-helper' }); + await client2.waitForNotification(n => n.method === 'serverHello'); + await client2.call('subscribe', { resource: sessionUri }); + + dispatchTurnStarted(client2, sessionUri, 'turn-unsub', 'hello', 1); + await client2.waitForNotification(n => isActionNotification(n, 'session/turnComplete')); + + await new Promise(resolve => setTimeout(resolve, 300)); + const sessionActions = client.receivedNotifications(n => isActionNotification(n, 'session/')); + assert.strictEqual(sessionActions.length, 0, 'unsubscribed client should not receive session actions'); + + client2.close(); + }); + + test('change model within session updates state', async function () { + this.timeout(10_000); + + const sessionUri = await createAndSubscribeSession(client, 'test-change-model'); + + client.notify('dispatchAction', { + clientSeq: 1, + action: { type: 'session/modelChanged', session: sessionUri, model: 'new-mock-model' }, + }); + + const modelChanged = await client.waitForNotification(n => isActionNotification(n, 'session/modelChanged')); + const action = getActionParams(modelChanged).envelope.action; + assert.strictEqual(action.type, 'session/modelChanged'); + if (action.type === 'session/modelChanged') { + assert.strictEqual((action as { model: string }).model, 'new-mock-model'); + } + + const snapshot = await client.call('subscribe', { resource: sessionUri }); + const state = snapshot.state as ISessionState; + assert.strictEqual(state.summary.model, 'new-mock-model'); + }); +}); diff --git a/src/vs/platform/agentHost/test/node/sessionStateManager.test.ts b/src/vs/platform/agentHost/test/node/sessionStateManager.test.ts new file mode 100644 index 00000000000..8bec27d0996 --- /dev/null +++ b/src/vs/platform/agentHost/test/node/sessionStateManager.test.ts @@ -0,0 +1,163 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { DisposableStore } from '../../../../base/common/lifecycle.js'; +import { URI } from '../../../../base/common/uri.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import { NullLogService } from '../../../log/common/log.js'; +import type { IActionEnvelope, INotification } from '../../common/state/sessionActions.js'; +import { ISessionSummary, ROOT_STATE_URI, SessionLifecycle, SessionStatus, type ISessionState } from '../../common/state/sessionState.js'; +import { SessionStateManager } from '../../node/sessionStateManager.js'; + +suite('SessionStateManager', () => { + + let disposables: DisposableStore; + let manager: SessionStateManager; + const sessionUri = URI.from({ scheme: 'copilot', path: '/test-session' }); + + function makeSessionSummary(resource?: URI): ISessionSummary { + return { + resource: resource ?? sessionUri, + provider: 'copilot', + title: 'Test', + status: SessionStatus.Idle, + createdAt: Date.now(), + modifiedAt: Date.now(), + }; + } + + setup(() => { + disposables = new DisposableStore(); + manager = disposables.add(new SessionStateManager(new NullLogService())); + }); + + teardown(() => { + disposables.dispose(); + }); + + ensureNoDisposablesAreLeakedInTestSuite(); + + test('createSession creates initial state with lifecycle Creating', () => { + const state = manager.createSession(makeSessionSummary()); + assert.strictEqual(state.lifecycle, SessionLifecycle.Creating); + assert.strictEqual(state.turns.length, 0); + assert.strictEqual(state.activeTurn, undefined); + assert.strictEqual(state.summary.resource.toString(), sessionUri.toString()); + }); + + test('getSnapshot returns undefined for unknown session', () => { + const unknown = URI.from({ scheme: 'copilot', path: '/unknown' }); + const snapshot = manager.getSnapshot(unknown); + assert.strictEqual(snapshot, undefined); + }); + + test('getSnapshot returns root snapshot', () => { + const snapshot = manager.getSnapshot(ROOT_STATE_URI); + assert.ok(snapshot); + assert.strictEqual(snapshot.resource.toString(), ROOT_STATE_URI.toString()); + assert.deepStrictEqual(snapshot.state, { agents: [] }); + }); + + test('getSnapshot returns session snapshot after creation', () => { + manager.createSession(makeSessionSummary()); + const snapshot = manager.getSnapshot(sessionUri); + assert.ok(snapshot); + assert.strictEqual(snapshot.resource.toString(), sessionUri.toString()); + assert.strictEqual((snapshot.state as ISessionState).lifecycle, SessionLifecycle.Creating); + }); + + test('dispatchServerAction applies action and emits envelope', () => { + manager.createSession(makeSessionSummary()); + + const envelopes: IActionEnvelope[] = []; + disposables.add(manager.onDidEmitEnvelope(e => envelopes.push(e))); + + manager.dispatchServerAction({ + type: 'session/ready', + session: sessionUri, + }); + + const state = manager.getSessionState(sessionUri); + assert.ok(state); + assert.strictEqual(state.lifecycle, SessionLifecycle.Ready); + + assert.strictEqual(envelopes.length, 1); + assert.strictEqual(envelopes[0].action.type, 'session/ready'); + assert.strictEqual(envelopes[0].serverSeq, 1); + assert.strictEqual(envelopes[0].origin, undefined); + }); + + test('serverSeq increments monotonically', () => { + manager.createSession(makeSessionSummary()); + + const envelopes: IActionEnvelope[] = []; + disposables.add(manager.onDidEmitEnvelope(e => envelopes.push(e))); + + manager.dispatchServerAction({ type: 'session/ready', session: sessionUri }); + manager.dispatchServerAction({ type: 'session/titleChanged', session: sessionUri, title: 'Updated' }); + + assert.strictEqual(envelopes.length, 2); + assert.strictEqual(envelopes[0].serverSeq, 1); + assert.strictEqual(envelopes[1].serverSeq, 2); + assert.ok(envelopes[1].serverSeq > envelopes[0].serverSeq); + }); + + test('dispatchClientAction includes origin in envelope', () => { + manager.createSession(makeSessionSummary()); + + const envelopes: IActionEnvelope[] = []; + disposables.add(manager.onDidEmitEnvelope(e => envelopes.push(e))); + + const origin = { clientId: 'renderer-1', clientSeq: 42 }; + manager.dispatchClientAction( + { type: 'session/ready', session: sessionUri }, + origin, + ); + + assert.strictEqual(envelopes.length, 1); + assert.deepStrictEqual(envelopes[0].origin, origin); + }); + + test('removeSession clears state and emits notification', () => { + manager.createSession(makeSessionSummary()); + + const notifications: INotification[] = []; + disposables.add(manager.onDidEmitNotification(n => notifications.push(n))); + + manager.removeSession(sessionUri); + + assert.strictEqual(manager.getSessionState(sessionUri), undefined); + assert.strictEqual(manager.getSnapshot(sessionUri), undefined); + assert.strictEqual(notifications.length, 1); + assert.strictEqual(notifications[0].type, 'notify/sessionRemoved'); + }); + + test('createSession emits sessionAdded notification', () => { + const notifications: INotification[] = []; + disposables.add(manager.onDidEmitNotification(n => notifications.push(n))); + + manager.createSession(makeSessionSummary()); + + assert.strictEqual(notifications.length, 1); + assert.strictEqual(notifications[0].type, 'notify/sessionAdded'); + }); + + test('getActiveTurnId returns active turn id after turnStarted', () => { + manager.createSession(makeSessionSummary()); + manager.dispatchServerAction({ type: 'session/ready', session: sessionUri }); + + assert.strictEqual(manager.getActiveTurnId(sessionUri), undefined); + + manager.dispatchServerAction({ + type: 'session/turnStarted', + session: sessionUri, + turnId: 'turn-1', + userMessage: { text: 'hello' }, + }); + + assert.strictEqual(manager.getActiveTurnId(sessionUri), 'turn-1'); + }); +}); diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostChatContribution.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostChatContribution.ts index 4a619d7efa3..cebdcabb4e5 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostChatContribution.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostChatContribution.ts @@ -6,7 +6,10 @@ import { Disposable, DisposableStore, toDisposable } from '../../../../../../base/common/lifecycle.js'; import { URI } from '../../../../../../base/common/uri.js'; import { IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js'; -import { IAgentHostService, AgentHostEnabledSettingId, IAgentDescriptor } from '../../../../../../platform/agentHost/common/agentService.js'; +import { IAgentHostService, AgentHostEnabledSettingId, type AgentProvider } from '../../../../../../platform/agentHost/common/agentService.js'; +import { isSessionAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; +import { SessionClientState } from '../../../../../../platform/agentHost/common/state/sessionClientState.js'; +import { ROOT_STATE_URI, type IAgentInfo, type IRootState } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { IDefaultAccountService } from '../../../../../../platform/defaultAccount/common/defaultAccount.js'; import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js'; import { ILogService, LogLevel } from '../../../../../../platform/log/common/log.js'; @@ -38,6 +41,10 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr private _outputChannel: IOutputChannel | undefined; private _isChannelRegistered = false; + private _clientState: SessionClientState | undefined; + private readonly _agentRegistrations = new Map(); + /** Model providers keyed by agent provider, for pushing model updates. */ + private readonly _modelProviders = new Map(); constructor( @IAgentHostService private readonly _agentHostService: IAgentHostService, @@ -57,7 +64,30 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr } this._setupIpcLogging(); - this._discoverAndRegisterAgents(); + + // Shared client state for protocol reconciliation + this._clientState = this._register(new SessionClientState(this._agentHostService.clientId)); + + // Forward action envelopes from the host to client state + this._register(this._agentHostService.onDidAction(envelope => { + // Only root actions are relevant here; session actions are + // handled by individual session handlers. + if (!isSessionAction(envelope.action)) { + this._clientState!.receiveEnvelope(envelope); + } + })); + + // Forward notifications to client state + this._register(this._agentHostService.onDidNotification(n => { + this._clientState!.receiveNotification(n); + })); + + // React to root state changes (agent discovery / removal) + this._register(this._clientState.onDidChangeRootState(rootState => { + this._handleRootStateChange(rootState); + })); + + this._initializeAndSubscribe(); } // ---- IPC output channel (trace-level only) ------------------------------ @@ -66,9 +96,12 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr this._updateOutputChannel(); this._register(this._logService.onDidChangeLogLevel(() => this._updateOutputChannel())); - // Subscribe to all progress events for IPC logging - this._register(this._agentHostService.onDidSessionProgress(e => { - this._traceIpc('event', 'onDidSessionProgress', e); + // Subscribe to action / notification streams for IPC logging + this._register(this._agentHostService.onDidAction(e => { + this._traceIpc('event', 'onDidAction', e); + })); + this._register(this._agentHostService.onDidNotification(e => { + this._traceIpc('event', 'onDidNotification', e); })); } @@ -121,22 +154,47 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr this._outputChannel.append(`[${timestamp}] [trace] ${arrow} ${method}${payload ? `\n${payload}` : ''}\n`); } - private async _discoverAndRegisterAgents(): Promise { + private async _initializeAndSubscribe(): Promise { try { - const agents = await this._agentHostService.listAgents(); + const snapshot = await this._agentHostService.subscribe(ROOT_STATE_URI); if (this._store.isDisposed) { return; } - for (const agent of agents) { - this._registerAgent(agent); - } + // Feed snapshot into client state — fires onDidChangeRootState + this._clientState!.handleSnapshot(ROOT_STATE_URI, snapshot.state, snapshot.fromSeq); } catch (err) { - this._logService.error(err, '[AgentHost] Failed to discover agents'); + this._logService.error('[AgentHost] Failed to subscribe to root state', err); } } - private _registerAgent(agent: IAgentDescriptor): void { - const store = this._register(new DisposableStore()); + private _handleRootStateChange(rootState: IRootState): void { + const incoming = new Set(rootState.agents.map(a => a.provider)); + + // Remove agents that are no longer present + for (const [provider, store] of this._agentRegistrations) { + if (!incoming.has(provider)) { + store.dispose(); + this._agentRegistrations.delete(provider); + this._modelProviders.delete(provider); + } + } + + // Register new agents and push model updates to existing ones + for (const agent of rootState.agents) { + if (!this._agentRegistrations.has(agent.provider)) { + this._registerAgent(agent); + } else { + // Push updated models to existing model provider + const modelProvider = this._modelProviders.get(agent.provider); + modelProvider?.updateModels(agent.models); + } + } + } + + private _registerAgent(agent: IAgentInfo): void { + const store = new DisposableStore(); + this._agentRegistrations.set(agent.provider, store); + this._register(store); const sessionType = `agent-host-${agent.provider}`; const agentId = sessionType; const vendor = sessionType; @@ -169,17 +227,18 @@ export class AgentHostContribution extends Disposable implements IWorkbenchContr const vendorDescriptor = { vendor, displayName: agent.displayName, configuration: undefined, managementCommand: undefined, when: undefined }; this._languageModelsService.deltaLanguageModelChatProviderDescriptors([vendorDescriptor], []); store.add(toDisposable(() => this._languageModelsService.deltaLanguageModelChatProviderDescriptors([], [vendorDescriptor]))); - const modelProvider = store.add(this._instantiationService.createInstance(AgentHostLanguageModelProvider, sessionType, vendor, agent.provider)); + const modelProvider = store.add(new AgentHostLanguageModelProvider(sessionType, vendor)); + modelProvider.updateModels(agent.models); + this._modelProviders.set(agent.provider, modelProvider); + store.add(toDisposable(() => this._modelProviders.delete(agent.provider))); store.add(this._languageModelsService.registerLanguageModelProvider(vendor, modelProvider)); - // Auth (only for agents that need it) - if (agent.requiresAuth) { - this._pushAuthToken().then(() => modelProvider.refresh()); - store.add(this._defaultAccountService.onDidChangeDefaultAccount(() => - this._pushAuthToken().then(() => modelProvider.refresh()))); - store.add(this._authenticationService.onDidChangeSessions(() => - this._pushAuthToken().then(() => modelProvider.refresh()))); - } + // Push auth token and refresh models from server + this._pushAuthToken().then(() => this._agentHostService.refreshModels()).catch(() => { /* best-effort */ }); + store.add(this._defaultAccountService.onDidChangeDefaultAccount(() => + this._pushAuthToken().then(() => this._agentHostService.refreshModels()).catch(() => { /* best-effort */ }))); + store.add(this._authenticationService.onDidChangeSessions(() => + this._pushAuthToken().then(() => this._agentHostService.refreshModels()).catch(() => { /* best-effort */ }))); } private async _pushAuthToken(): Promise { diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLanguageModelProvider.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLanguageModelProvider.ts index 88f44d2917c..546dd68513b 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLanguageModelProvider.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostLanguageModelProvider.ts @@ -7,63 +7,60 @@ import { CancellationToken } from '../../../../../../base/common/cancellation.js import { Emitter } from '../../../../../../base/common/event.js'; import { Disposable } from '../../../../../../base/common/lifecycle.js'; import { ExtensionIdentifier } from '../../../../../../platform/extensions/common/extensions.js'; -import { ILogService } from '../../../../../../platform/log/common/log.js'; -import { IAgentHostService } from '../../../../../../platform/agentHost/common/agentService.js'; +import { ISessionModelInfo } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { ILanguageModelChatProvider, ILanguageModelChatMetadataAndIdentifier } from '../../../common/languageModels.js'; /** * Exposes models available from the agent host process as selectable - * language models in the chat model picker. + * language models in the chat model picker. Models are provided from + * root state (via {@link IAgentInfo.models}) rather than via RPC. */ export class AgentHostLanguageModelProvider extends Disposable implements ILanguageModelChatProvider { private readonly _onDidChange = this._register(new Emitter()); readonly onDidChange = this._onDidChange.event; + private _models: readonly ISessionModelInfo[] = []; + constructor( private readonly _sessionType: string, private readonly _vendor: string, - private readonly _provider: string, - @IAgentHostService private readonly _agentHostService: IAgentHostService, - @ILogService private readonly _logService: ILogService, ) { super(); } - refresh(): void { + /** + * Called by {@link AgentHostContribution} when models change in root state. + */ + updateModels(models: readonly ISessionModelInfo[]): void { + this._models = models; this._onDidChange.fire(); } async provideLanguageModelChatInfo(_options: unknown, _token: CancellationToken): Promise { - try { - const models = await this._agentHostService.listModels(); - return models - .filter(m => m.provider === this._provider && m.policyState !== 'disabled') - .map(m => ({ - identifier: `${this._vendor}:${m.id}`, - metadata: { - extension: new ExtensionIdentifier('vscode.agent-host'), - name: m.name, - id: m.id, - vendor: this._vendor, - version: '1.0', - family: m.id, - maxInputTokens: m.maxContextWindow, - maxOutputTokens: 0, - isDefaultForLocation: {}, - isUserSelectable: true, - modelPickerCategory: undefined, - targetChatSessionType: this._sessionType, - capabilities: { - vision: m.supportsVision, - toolCalling: true, - agentMode: true, - }, + return this._models + .filter(m => m.policyState !== 'disabled') + .map(m => ({ + identifier: `${this._vendor}:${m.id}`, + metadata: { + extension: new ExtensionIdentifier('vscode.agent-host'), + name: m.name, + id: m.id, + vendor: this._vendor, + version: '1.0', + family: m.id, + maxInputTokens: m.maxContextWindow ?? 0, + maxOutputTokens: 0, + isDefaultForLocation: {}, + isUserSelectable: true, + modelPickerCategory: undefined, + targetChatSessionType: this._sessionType, + capabilities: { + vision: m.supportsVision ?? false, + toolCalling: true, + agentMode: true, }, - })); - } catch (err) { - this._logService.trace('[AgentHost] Models not available yet, will retry on next refresh'); - return []; - } + }, + })); } async sendChatRequest(): Promise { 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 a3a1a28960d..bb5d4f5b485 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -6,114 +6,35 @@ import { CancellationToken } from '../../../../../../base/common/cancellation.js'; import { Emitter } from '../../../../../../base/common/event.js'; import { MarkdownString } from '../../../../../../base/common/htmlContent.js'; -import { Disposable, toDisposable } from '../../../../../../base/common/lifecycle.js'; +import { Disposable, DisposableStore, toDisposable } from '../../../../../../base/common/lifecycle.js'; import { observableValue } from '../../../../../../base/common/observable.js'; +import { generateUuid } from '../../../../../../base/common/uuid.js'; import { URI } from '../../../../../../base/common/uri.js'; import { ExtensionIdentifier } from '../../../../../../platform/extensions/common/extensions.js'; import { ILogService } from '../../../../../../platform/log/common/log.js'; import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js'; import { IProductService } from '../../../../../../platform/product/common/productService.js'; import { IWorkspaceContextService } from '../../../../../../platform/workspace/common/workspace.js'; -import { IAgentHostService, IAgentAttachment, IAgentMessageEvent, IAgentToolCompleteEvent, IAgentToolStartEvent, AgentProvider, AgentSession, IAgentProgressEvent } from '../../../../../../platform/agentHost/common/agentService.js'; +import { IAgentHostService, IAgentAttachment, AgentProvider, AgentSession } from '../../../../../../platform/agentHost/common/agentService.js'; +import { isSessionAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; +import { SessionClientState } from '../../../../../../platform/agentHost/common/state/sessionClientState.js'; +import { ToolCallStatus, TurnState, type IMessageAttachment } from '../../../../../../platform/agentHost/common/state/sessionState.js'; import { ChatAgentLocation, ChatModeKind } from '../../../common/constants.js'; import { IChatAgentData, IChatAgentImplementation, IChatAgentRequest, IChatAgentResult, IChatAgentService } from '../../../common/participants/chatAgents.js'; -import { IChatProgress, IChatTerminalToolInvocationData, IChatToolInputInvocationData, IChatToolInvocation, IChatToolInvocationSerialized, ToolConfirmKind } from '../../../common/chatService/chatService.js'; +import { IChatProgress, IChatToolInvocation, ToolConfirmKind } from '../../../common/chatService/chatService.js'; import { ChatToolInvocation } from '../../../common/model/chatProgressTypes/chatToolInvocation.js'; -import { IPreparedToolInvocation, IToolConfirmationMessages, IToolData, ToolDataSource, ToolInvocationPresentation } from '../../../common/tools/languageModelToolsService.js'; import { IChatSession, IChatSessionContentProvider, IChatSessionHistoryItem } from '../../../common/chatSessionsService.js'; import { getAgentHostIcon } from '../agentSessions.js'; +import { turnsToHistory, toolCallStateToInvocation, permissionToConfirmation, finalizeToolInvocation } from './stateToProgressAdapter.js'; // ============================================================================= -// AgentHostSessionHandler - renderer-side handler for a single agent host -// chat session. Bridges the agent host IPC service with the chat UI: -// creates sessions, streams responses, manages tool invocations, and -// reconstructs history for session restore. +// 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, permissionResolved, +// turnCancelled) back to the server. // ============================================================================= -/** - * Converts a flat array of IPC events (messages + tool events) into - * request/response history items for the chat model. - */ -function buildHistory( - events: readonly (IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[], - history: IChatSessionHistoryItem[], - participantId: string, -): void { - let currentResponseParts: IChatProgress[] | undefined; - - for (const e of events) { - if (e.type === 'message') { - if (e.role === 'user') { - if (currentResponseParts) { - history.push({ type: 'response', parts: currentResponseParts, participant: participantId }); - currentResponseParts = undefined; - } - history.push({ type: 'request', prompt: e.content, participant: participantId }); - } else { - if (!currentResponseParts) { - currentResponseParts = []; - } - if (e.content) { - currentResponseParts.push({ kind: 'markdownContent', content: new MarkdownString(e.content) }); - } - } - } else if (e.type === 'tool_start') { - if (!currentResponseParts) { - currentResponseParts = []; - } - const toolSpecificData = (e.toolKind === 'terminal' && e.toolInput) - ? { kind: 'terminal' as const, commandLine: { original: e.toolInput }, language: e.language ?? 'shellscript' } - : undefined; - currentResponseParts.push({ - kind: 'toolInvocationSerialized', - toolCallId: e.toolCallId, - toolId: e.toolName, - source: ToolDataSource.Internal, - invocationMessage: new MarkdownString(e.invocationMessage), - originMessage: undefined, - pastTenseMessage: undefined, - isConfirmed: { type: ToolConfirmKind.ConfirmationNotNeeded }, - isComplete: false, - presentation: undefined, - toolSpecificData, - } satisfies IChatToolInvocationSerialized); - } else if (e.type === 'tool_complete') { - if (currentResponseParts) { - const idx = currentResponseParts.findIndex( - p => p.kind === 'toolInvocationSerialized' && p.toolCallId === e.toolCallId - ); - if (idx >= 0) { - const existing = currentResponseParts[idx] as IChatToolInvocationSerialized; - const isTerminal = existing.toolSpecificData?.kind === 'terminal'; - currentResponseParts[idx] = { - ...existing, - isComplete: true, - pastTenseMessage: isTerminal ? undefined : new MarkdownString(e.pastTenseMessage), - toolSpecificData: isTerminal - ? { - ...existing.toolSpecificData as IChatTerminalToolInvocationData, - terminalCommandOutput: e.toolOutput !== undefined ? { text: e.toolOutput } : undefined, - terminalCommandState: { exitCode: e.success ? 0 : 1 }, - } - : existing.toolSpecificData, - }; - } - } - } - } - - // Mark incomplete tool invocations as complete (orphaned tool_start without tool_complete) - if (currentResponseParts) { - for (let i = 0; i < currentResponseParts.length; i++) { - const part = currentResponseParts[i]; - if (part.kind === 'toolInvocationSerialized' && !part.isComplete) { - currentResponseParts[i] = { ...part, isComplete: true }; - } - } - history.push({ type: 'response', parts: currentResponseParts, participant: participantId }); - } -} - // ============================================================================= // Chat session // ============================================================================= @@ -131,7 +52,7 @@ class AgentHostChatSession extends Disposable implements IChatSession { constructor( readonly sessionResource: URI, readonly history: readonly IChatSessionHistoryItem[], - private readonly _sendRequest: (message: string, progress: (parts: IChatProgress[]) => void, token: CancellationToken) => Promise, + private readonly _sendRequest: (request: IChatAgentRequest, progress: (parts: IChatProgress[]) => void, token: CancellationToken) => Promise, onDispose: () => void, @ILogService private readonly _logService: ILogService, ) { @@ -143,7 +64,7 @@ class AgentHostChatSession extends Disposable implements IChatSession { this.requestHandler = async (request, progress, _history, cancellationToken) => { this._logService.info('[AgentHost] requestHandler called'); this.isCompleteObs.set(false, undefined); - await this._sendRequest(request.message, progress, cancellationToken); + await this._sendRequest(request, progress, cancellationToken); this.isCompleteObs.set(true, undefined); }; @@ -167,10 +88,14 @@ export interface IAgentHostSessionHandlerConfig { export class AgentHostSessionHandler extends Disposable implements IChatSessionContentProvider { - private readonly _resourceToSession = new Map(); private readonly _activeSessions = new Map(); + /** Maps UI resource keys to resolved backend session URIs. */ + private readonly _sessionToBackend = new Map(); private readonly _config: IAgentHostSessionHandlerConfig; + /** Client state manager shared across all sessions for this handler. */ + private readonly _clientState: SessionClientState; + constructor( config: IAgentHostSessionHandlerConfig, @IAgentHostService private readonly _agentHostService: IAgentHostService, @@ -182,35 +107,69 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC ) { super(); this._config = config; + + // Create shared client state manager for this handler instance + this._clientState = this._register(new SessionClientState(this._agentHostService.clientId)); + + // Forward action envelopes from IPC to client state + this._register(this._agentHostService.onDidAction(envelope => { + if (isSessionAction(envelope.action)) { + this._clientState.receiveEnvelope(envelope); + } + })); + this._registerAgent(); } async provideChatSessionContent(sessionResource: URI, _token: CancellationToken): Promise { const resourceKey = sessionResource.path.substring(1); - const resolvedSession = await this._resolveSession(sessionResource); - + // For untitled (new) sessions, defer backend session creation until the + // first request arrives so the user-selected model is available. + // For existing sessions we resolve immediately to load history. + let resolvedSession: URI | undefined; + const isUntitled = resourceKey.startsWith('untitled-'); const history: IChatSessionHistoryItem[] = []; - if (!resourceKey.startsWith('untitled-')) { - const events = await this._agentHostService.getSessionMessages(resolvedSession); - buildHistory(events, history, this._config.agentId); + if (!isUntitled) { + resolvedSession = this._resolveSessionUri(sessionResource); + this._sessionToBackend.set(resourceKey, resolvedSession); + try { + const snapshot = await this._agentHostService.subscribe(resolvedSession); + this._clientState.handleSnapshot(resolvedSession, snapshot.state, snapshot.fromSeq); + const sessionState = this._clientState.getSessionState(resolvedSession); + if (sessionState) { + history.push(...turnsToHistory(sessionState.turns, this._config.agentId)); + } + } catch (err) { + this._logService.warn(`[AgentHost] Failed to subscribe to existing session: ${resolvedSession.toString()}`, err); + } } const session = this._instantiationService.createInstance( AgentHostChatSession, sessionResource, history, - (message: string, progress: (parts: IChatProgress[]) => void, token: CancellationToken) => - this._sendAndStreamResponse(resolvedSession, message, [], progress, token), + async (request: IChatAgentRequest, progress: (parts: IChatProgress[]) => void, token: CancellationToken) => { + const backendSession = resolvedSession ?? await this._createAndSubscribe(sessionResource, request.userSelectedModelId); + resolvedSession = backendSession; + this._sessionToBackend.set(resourceKey, backendSession); + return this._handleTurn(backendSession, request, progress, token); + }, () => { this._activeSessions.delete(resourceKey); - this._resourceToSession.delete(sessionResource.toString()); - this._agentHostService.disposeSession(resolvedSession); + this._sessionToBackend.delete(resourceKey); + if (resolvedSession) { + this._clientState.unsubscribe(resolvedSession); + this._agentHostService.unsubscribe(resolvedSession); + this._agentHostService.disposeSession(resolvedSession); + } }, ); this._activeSessions.set(resourceKey, session); return session; } + // ---- Agent registration ------------------------------------------------- + private _registerAgent(): void { const agentData: IChatAgentData = { id: this._config.agentId, @@ -246,13 +205,17 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC cancellationToken: CancellationToken, ): Promise { this._logService.info(`[AgentHost] _invokeAgent called for resource: ${request.sessionResource.toString()}`); - const session = await this._resolveSession(request.sessionResource, request.userSelectedModelId); - this._logService.info(`[AgentHost] resolved session: ${session.toString()}`); - - const attachments = this._convertVariablesToAttachments(request); - await this._sendAndStreamResponse(session, request.message, attachments, progress, cancellationToken); + // Resolve or create backend session const resourceKey = request.sessionResource.path.substring(1); + let resolvedSession = this._sessionToBackend.get(resourceKey); + if (!resolvedSession) { + resolvedSession = await this._createAndSubscribe(request.sessionResource, request.userSelectedModelId); + this._sessionToBackend.set(resourceKey, resolvedSession); + } + + await this._handleTurn(resolvedSession, request, progress, cancellationToken); + const activeSession = this._activeSessions.get(resourceKey); if (activeSession) { activeSession.isCompleteObs.set(true, undefined); @@ -261,10 +224,11 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC return {}; } - private async _sendAndStreamResponse( + // ---- Turn handling (state-driven) --------------------------------------- + + private async _handleTurn( session: URI, - message: string, - attachments: IAgentAttachment[], + request: IChatAgentRequest, progress: (parts: IChatProgress[]) => void, cancellationToken: CancellationToken, ): Promise { @@ -272,7 +236,54 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC return; } + const turnId = generateUuid(); + const attachments = this._convertVariablesToAttachments(request); + const messageAttachments: IMessageAttachment[] = attachments.map(a => ({ + type: a.type, + path: a.path, + displayName: a.displayName, + })); + + // If the user selected a different model since the session was created + // (or since the last turn), dispatch a model change action first so the + // agent backend picks up the new model before processing the turn. + const rawModelId = this._extractRawModelId(request.userSelectedModelId); + if (rawModelId) { + const currentModel = this._clientState.getSessionState(session)?.summary.model; + if (currentModel !== rawModelId) { + const modelAction = { + type: 'session/modelChanged' as const, + session, + model: rawModelId, + }; + const modelSeq = this._clientState.applyOptimistic(modelAction); + this._agentHostService.dispatchAction(modelAction, this._clientState.clientId, modelSeq); + } + } + + // Dispatch session/turnStarted — the server will call sendMessage on + // the provider as a side effect. + const turnAction = { + type: 'session/turnStarted' as const, + session, + turnId, + userMessage: { + text: request.message, + attachments: messageAttachments.length > 0 ? messageAttachments : undefined, + }, + }; + const clientSeq = this._clientState.applyOptimistic(turnAction); + this._agentHostService.dispatchAction(turnAction, this._clientState.clientId, clientSeq); + + // Track live ChatToolInvocation/permission objects for this turn const activeToolInvocations = new Map(); + const activePermissions = new Map(); + + // Track last-emitted lengths to compute deltas from immutable state + let lastStreamedTextLen = 0; + let lastReasoningLen = 0; + + const turnDisposables = new DisposableStore(); let resolveDone: () => void; const done = new Promise(resolve => { resolveDone = resolve; }); @@ -283,237 +294,157 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC return; } finished = true; - this._finalizeOutstandingTools(activeToolInvocations); - listener.dispose(); + // Finalize any outstanding tool invocations + for (const [, invocation] of activeToolInvocations) { + invocation.didExecuteTool(undefined); + } + activeToolInvocations.clear(); + turnDisposables.dispose(); resolveDone(); }; - const sessionStr = session.toString(); - const listener = this._agentHostService.onDidSessionProgress(e => { - if (e.session.toString() !== sessionStr || cancellationToken.isCancellationRequested) { + // Listen to state changes and translate to IChatProgress[] + turnDisposables.add(this._clientState.onDidChangeSessionState(e => { + if (e.session.toString() !== session.toString() || cancellationToken.isCancellationRequested) { return; } - switch (e.type) { - case 'delta': - this._logService.trace(`[AgentHost:${sessionStr}] delta: ${e.content.length} chars`); - progress([{ kind: 'markdownContent', content: new MarkdownString(e.content) }]); - break; + const activeTurn = e.state.activeTurn; - case 'tool_start': { - this._logService.trace(`[AgentHost:${sessionStr}] tool_start: ${e.toolName} (${e.toolCallId}), kind=${e.toolKind ?? 'generic'}`); - const invocation = this._createToolInvocation(e); - activeToolInvocations.set(e.toolCallId, invocation); - progress([invocation]); - break; + if (!activeTurn || activeTurn.id !== turnId) { + // Turn completed (activeTurn cleared by reducer). + // Check if the finalized turn ended with an error and emit it. + const lastTurn = e.state.turns[e.state.turns.length - 1]; + if (lastTurn?.id === turnId && lastTurn.state === TurnState.Error && lastTurn.error) { + progress([{ kind: 'markdownContent', content: new MarkdownString(`\n\nError: (${lastTurn.error.errorType}) ${lastTurn.error.message}`) }]); } - - case 'tool_complete': { - this._logService.trace(`[AgentHost:${sessionStr}] tool_complete: ${e.toolCallId}, success=${e.success}`); - const invocation = activeToolInvocations.get(e.toolCallId); - if (invocation) { - activeToolInvocations.delete(e.toolCallId); - this._finalizeToolInvocation(invocation, e); - } else { - this._logService.trace(`[AgentHost:${sessionStr}] tool_complete for unknown toolCallId: ${e.toolCallId}`); - } - break; - } - - case 'idle': - this._logService.trace(`[AgentHost:${sessionStr}] idle, finishing`); + if (!finished) { finish(); - break; - - case 'error': - this._logService.error(`[AgentHost:${sessionStr}] error: (${e.errorType}) ${e.message}`); - progress([{ kind: 'markdownContent', content: new MarkdownString(`\n\nError: (${e.errorType}) ${e.message}`) }]); - finish(); - break; - - case 'usage': - this._logService.trace(`[AgentHost:${sessionStr}] usage: model=${e.model}, in=${e.inputTokens ?? '?'}, out=${e.outputTokens ?? '?'}`); - break; - - case 'title_changed': - this._logService.trace(`[AgentHost:${sessionStr}] title changed: ${e.title}`); - break; - - case 'reasoning': - this._logService.trace(`[AgentHost:${sessionStr}] reasoning: ${e.content.length} chars`); - progress([{ kind: 'thinking', value: e.content }]); - break; - - case 'permission_request': { - this._logService.info(`[AgentHost:${sessionStr}] permission_request: kind=${e.permissionKind}, requestId=${e.requestId}, path=${e.path ?? '(none)'}`); - - const confirmInvocation = this._createPermissionConfirmation(e); - progress([confirmInvocation]); - - IChatToolInvocation.awaitConfirmation(confirmInvocation, cancellationToken).then(reason => { - const approved = reason.type !== ToolConfirmKind.Denied && reason.type !== ToolConfirmKind.Skipped; - this._logService.info(`[AgentHost:${sessionStr}] permission response: requestId=${e.requestId}, approved=${approved} (kind=${reason.type})`); - this._agentHostService.respondToPermissionRequest(e.requestId, approved); - if (approved) { - confirmInvocation.didExecuteTool(undefined); - } else { - confirmInvocation.didExecuteTool({ content: [], toolResultError: 'User denied' }); - } - }); - break; } - - default: - this._logService.trace(`[AgentHost:${sessionStr}] unhandled event type: ${(e as IAgentProgressEvent).type}`); - break; + return; } - }); - const cancelListener = cancellationToken.onCancellationRequested(() => { - this._logService.info(`[AgentHost] Cancellation requested for ${sessionStr}, aborting...`); - this._agentHostService.abortSession(session).catch(err => { - this._logService.error(`[AgentHost] abortSession failed`, err); - }); - finish(); - cancelListener.dispose(); - }); + // Stream text deltas + if (activeTurn.streamingText.length > lastStreamedTextLen) { + const delta = activeTurn.streamingText.substring(lastStreamedTextLen); + lastStreamedTextLen = activeTurn.streamingText.length; + progress([{ kind: 'markdownContent', content: new MarkdownString(delta) }]); + } - try { - this._logService.info(`[AgentHost] Sending message to session ${session.toString()} (${attachments.length} attachments)`); - await this._agentHostService.sendMessage(session, message, attachments.length > 0 ? attachments : undefined); - this._logService.info(`[AgentHost] sendMessage returned for session ${session.toString()}`); - } catch (err) { - this._logService.error(`[AgentHost] [${session.toString()}] sendMessage failed`, err); + // Stream reasoning deltas + if (activeTurn.reasoning.length > lastReasoningLen) { + const delta = activeTurn.reasoning.substring(lastReasoningLen); + lastReasoningLen = activeTurn.reasoning.length; + progress([{ kind: 'thinking', value: delta }]); + } + + // Handle tool calls — create/finalize ChatToolInvocations + for (const [toolCallId, tc] of activeTurn.toolCalls) { + const existing = activeToolInvocations.get(toolCallId); + if (!existing) { + if (tc.status === ToolCallStatus.Running || tc.status === ToolCallStatus.PendingPermission) { + const invocation = toolCallStateToInvocation(tc); + activeToolInvocations.set(toolCallId, invocation); + progress([invocation]); + } + } else if (tc.status === ToolCallStatus.Completed || tc.status === ToolCallStatus.Failed) { + activeToolInvocations.delete(toolCallId); + finalizeToolInvocation(existing, tc); + } + } + + // Handle permission requests + for (const [requestId, perm] of activeTurn.pendingPermissions) { + if (activePermissions.has(requestId)) { + continue; + } + const confirmInvocation = permissionToConfirmation(perm); + activePermissions.set(requestId, confirmInvocation); + progress([confirmInvocation]); + + IChatToolInvocation.awaitConfirmation(confirmInvocation, cancellationToken).then(reason => { + const approved = reason.type !== ToolConfirmKind.Denied && reason.type !== ToolConfirmKind.Skipped; + this._logService.info(`[AgentHost] Permission response: requestId=${requestId}, approved=${approved}`); + const resolveAction = { + type: 'session/permissionResolved' as const, + session, + turnId, + requestId, + approved, + }; + const seq = this._clientState.applyOptimistic(resolveAction); + this._agentHostService.dispatchAction(resolveAction, this._clientState.clientId, seq); + if (approved) { + confirmInvocation.didExecuteTool(undefined); + } else { + confirmInvocation.didExecuteTool({ content: [], toolResultError: 'User denied' }); + } + }).catch(err => { + this._logService.warn(`[AgentHost] Permission confirmation failed for requestId=${requestId}`, err); + }); + } + })); + + turnDisposables.add(cancellationToken.onCancellationRequested(() => { + this._logService.info(`[AgentHost] Cancellation requested for ${session.toString()}, dispatching turnCancelled`); + const cancelAction = { + type: 'session/turnCancelled' as const, + session, + turnId, + }; + const seq = this._clientState.applyOptimistic(cancelAction); + this._agentHostService.dispatchAction(cancelAction, this._clientState.clientId, seq); finish(); - } + })); await done; - cancelListener.dispose(); } - private async _resolveSession(sessionResource: URI, model?: string): Promise { - if (sessionResource.scheme === this._config.sessionType && !sessionResource.path.startsWith('/untitled-')) { - // Convert UI resource scheme (e.g. agent-host) to provider URI scheme (e.g. copilot) - const rawId = sessionResource.path.substring(1); - const session = AgentSession.uri(this._config.provider, rawId); - this._logService.trace(`[AgentHost] Resolved existing session: ${sessionResource.toString()} -> ${session.toString()}`); - return session; - } + // ---- Session resolution ------------------------------------------------- - const key = sessionResource.toString(); - const existing = this._resourceToSession.get(key); - if (existing) { - this._logService.trace(`[AgentHost] Reusing mapped session: ${key} -> ${existing.toString()}`); - return existing; - } + /** Maps a UI session resource to a backend provider URI. */ + private _resolveSessionUri(sessionResource: URI): URI { + const rawId = sessionResource.path.substring(1); + return AgentSession.uri(this._config.provider, rawId); + } - this._logService.trace(`[AgentHost] Creating new session for resource ${key}, model=${model ?? '(default)'}, provider=${this._config.provider}`); + /** Creates a new backend session and subscribes to its state. */ + private async _createAndSubscribe(sessionResource: URI, modelId?: string): Promise { + const rawModelId = this._extractRawModelId(modelId); const workspaceFolder = this._workspaceContextService.getWorkspace().folders[0]; + + this._logService.trace(`[AgentHost] Creating new session, model=${rawModelId ?? '(default)'}, provider=${this._config.provider}`); const session = await this._agentHostService.createSession({ - model, + model: rawModelId, provider: this._config.provider, workingDirectory: workspaceFolder?.uri.fsPath, }); - this._logService.trace(`[AgentHost] Created new session: ${session.toString()}`); - this._resourceToSession.set(key, session); + this._logService.trace(`[AgentHost] Created session: ${session.toString()}`); + + // Subscribe to the new session's state + try { + const snapshot = await this._agentHostService.subscribe(session); + this._clientState.handleSnapshot(session, snapshot.state, snapshot.fromSeq); + } catch (err) { + this._logService.error(`[AgentHost] Failed to subscribe to new session: ${session.toString()}`, err); + } + return session; } - private _createToolInvocation(event: IAgentToolStartEvent): ChatToolInvocation { - const toolData: IToolData = { - id: event.toolName, - source: ToolDataSource.Internal, - displayName: event.displayName, - modelDescription: event.toolName, - }; - let parameters: unknown; - if (event.toolArguments) { - try { - parameters = JSON.parse(event.toolArguments); - } catch { - // malformed JSON - } + /** + * Extracts the raw model id from a language-model service identifier. + * E.g. "agent-host-copilot:claude-sonnet-4-20250514" → "claude-sonnet-4-20250514". + */ + private _extractRawModelId(languageModelIdentifier: string | undefined): string | undefined { + if (!languageModelIdentifier) { + return undefined; } - - const invocation = new ChatToolInvocation(undefined, toolData, event.toolCallId, undefined, parameters); - invocation.invocationMessage = new MarkdownString(event.invocationMessage); - - if (event.toolKind === 'terminal' && event.toolInput) { - invocation.toolSpecificData = { - kind: 'terminal', - commandLine: { original: event.toolInput }, - language: event.language ?? 'shellscript', - } satisfies IChatTerminalToolInvocationData; + const prefix = this._config.sessionType + ':'; + if (languageModelIdentifier.startsWith(prefix)) { + return languageModelIdentifier.substring(prefix.length); } - - return invocation; - } - - private _createPermissionConfirmation(event: import('../../../../../../platform/agentHost/common/agentService.js').IAgentPermissionRequestEvent): ChatToolInvocation { - let title: string; - let toolSpecificData: IChatTerminalToolInvocationData | IChatToolInputInvocationData | undefined; - - switch (event.permissionKind) { - case 'shell': { - title = event.intention ?? 'Run command'; - toolSpecificData = event.fullCommandText ? { - kind: 'terminal', - commandLine: { original: event.fullCommandText }, - language: 'shellscript', - } : undefined; - break; - } - case 'write': { - title = event.path ? `Edit ${event.path}` : 'Edit file'; - let rawInput: unknown; - try { rawInput = JSON.parse(event.rawRequest); } catch { rawInput = { path: event.path }; } - toolSpecificData = { kind: 'input', rawInput }; - break; - } - case 'mcp': { - const toolTitle = event.toolName ?? 'MCP Tool'; - title = event.serverName ? `${event.serverName}: ${toolTitle}` : toolTitle; - let rawInput: unknown; - try { rawInput = JSON.parse(event.rawRequest); } catch { rawInput = { serverName: event.serverName, toolName: event.toolName }; } - toolSpecificData = { kind: 'input', rawInput }; - break; - } - case 'read': { - title = event.intention ?? 'Read file'; - let rawInput: unknown; - try { rawInput = JSON.parse(event.rawRequest); } catch { rawInput = { path: event.path, intention: event.intention }; } - toolSpecificData = { kind: 'input', rawInput }; - break; - } - default: { - title = 'Permission request'; - let rawInput: unknown; - try { rawInput = JSON.parse(event.rawRequest); } catch { rawInput = {}; } - toolSpecificData = { kind: 'input', rawInput }; - break; - } - } - - const confirmationMessages: IToolConfirmationMessages = { - title: new MarkdownString(title), - message: new MarkdownString(''), - }; - - const toolData: IToolData = { - id: `permission_${event.permissionKind}`, - source: ToolDataSource.Internal, - displayName: title, - modelDescription: '', - }; - - const preparedInvocation: IPreparedToolInvocation = { - invocationMessage: new MarkdownString(title), - confirmationMessages, - presentation: ToolInvocationPresentation.HiddenAfterComplete, - toolSpecificData, - }; - - return new ChatToolInvocation(preparedInvocation, toolData, event.requestId, undefined, undefined); + return languageModelIdentifier; } private _convertVariablesToAttachments(request: IChatAgentRequest): IAgentAttachment[] { @@ -542,34 +473,14 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC return attachments; } - private _finalizeToolInvocation(invocation: ChatToolInvocation, event: IAgentToolCompleteEvent): void { - if (invocation.toolSpecificData?.kind === 'terminal') { - const terminalData = invocation.toolSpecificData as IChatTerminalToolInvocationData; - invocation.toolSpecificData = { - ...terminalData, - terminalCommandOutput: event.toolOutput !== undefined ? { text: event.toolOutput } : undefined, - terminalCommandState: { exitCode: event.success ? 0 : 1 }, - }; - } else { - invocation.pastTenseMessage = new MarkdownString(event.pastTenseMessage); - } - - invocation.didExecuteTool(!event.success ? { content: [], toolResultError: event.error?.message } : undefined); - } - - private _finalizeOutstandingTools(activeToolInvocations: Map): void { - for (const [id, invocation] of activeToolInvocations) { - invocation.didExecuteTool(undefined); - activeToolInvocations.delete(id); - } - } + // ---- Lifecycle ---------------------------------------------------------- override dispose(): void { for (const [, session] of this._activeSessions) { session.dispose(); } this._activeSessions.clear(); - this._resourceToSession.clear(); + this._sessionToBackend.clear(); super.dispose(); } } diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionListController.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionListController.ts index db54a126cfe..4f9338886d2 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionListController.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionListController.ts @@ -3,18 +3,20 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { CancellationToken } from '../../../../../../base/common/cancellation.js'; +import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js'; import { Emitter } from '../../../../../../base/common/event.js'; import { Disposable } from '../../../../../../base/common/lifecycle.js'; import { URI } from '../../../../../../base/common/uri.js'; import { IProductService } from '../../../../../../platform/product/common/productService.js'; import { IAgentHostService, AgentSession } from '../../../../../../platform/agentHost/common/agentService.js'; +import { isSessionAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; import { ChatSessionStatus, IChatSessionItem, IChatSessionItemController } from '../../../common/chatSessionsService.js'; import { getAgentHostIcon } from '../agentSessions.js'; /** * Provides session list items for the chat sessions sidebar by querying - * active sessions from the agent host process. + * active sessions from the agent host process. Listens to protocol + * notifications for incremental updates. */ export class AgentHostSessionListController extends Disposable implements IChatSessionItemController { @@ -30,6 +32,40 @@ export class AgentHostSessionListController extends Disposable implements IChatS @IProductService private readonly _productService: IProductService, ) { super(); + + // React to protocol notifications for session list changes + this._register(this._agentHostService.onDidNotification(n => { + if (n.type === 'notify/sessionAdded' && n.summary.provider === this._provider) { + const rawId = AgentSession.id(n.summary.resource); + this._items.push({ + resource: URI.from({ scheme: this._sessionType, path: `/${rawId}` }), + label: n.summary.title ?? `Session ${rawId.substring(0, 8)}`, + iconPath: getAgentHostIcon(this._productService), + status: ChatSessionStatus.Completed, + timing: { + created: n.summary.createdAt, + lastRequestStarted: n.summary.modifiedAt, + lastRequestEnded: n.summary.modifiedAt, + }, + }); + this._onDidChangeChatSessionItems.fire(); + } else if (n.type === 'notify/sessionRemoved') { + const removedId = AgentSession.id(n.session); + const idx = this._items.findIndex(item => item.resource.path === `/${removedId}`); + if (idx >= 0) { + this._items.splice(idx, 1); + this._onDidChangeChatSessionItems.fire(); + } + } + })); + + // Refresh on turnComplete actions for metadata updates (title, timing) + this._register(this._agentHostService.onDidAction(e => { + if (e.action.type === 'session/turnComplete' && isSessionAction(e.action) && AgentSession.provider(e.action.session) === this._provider) { + const cts = new CancellationTokenSource(); + this.refresh(cts.token).finally(() => cts.dispose()); + } + })); } get items(): readonly IChatSessionItem[] { diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts new file mode 100644 index 00000000000..fcdb959660b --- /dev/null +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/stateToProgressAdapter.ts @@ -0,0 +1,199 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { MarkdownString } from '../../../../../../base/common/htmlContent.js'; +import { ToolCallStatus, TurnState, type ICompletedToolCall, type IPermissionRequest, type IToolCallState, type ITurn } from '../../../../../../platform/agentHost/common/state/sessionState.js'; +import { type IChatProgress, type IChatTerminalToolInvocationData, type IChatToolInputInvocationData, type IChatToolInvocationSerialized, ToolConfirmKind } from '../../../common/chatService/chatService.js'; +import { type IChatSessionHistoryItem } from '../../../common/chatSessionsService.js'; +import { ChatToolInvocation } from '../../../common/model/chatProgressTypes/chatToolInvocation.js'; +import { type IPreparedToolInvocation, type IToolConfirmationMessages, type IToolData, ToolDataSource, ToolInvocationPresentation } from '../../../common/tools/languageModelToolsService.js'; + +/** + * Converts completed turns from the protocol state into session history items. + */ +export function turnsToHistory(turns: readonly ITurn[], participantId: string): IChatSessionHistoryItem[] { + const history: IChatSessionHistoryItem[] = []; + for (const turn of turns) { + // Request + history.push({ type: 'request', prompt: turn.userMessage.text, participant: participantId }); + + // Response parts + const parts: IChatProgress[] = []; + + // Assistant response text + if (turn.responseText) { + parts.push({ kind: 'markdownContent', content: new MarkdownString(turn.responseText) }); + } + + // Completed tool calls + for (const tc of turn.toolCalls) { + parts.push(completedToolCallToSerialized(tc)); + } + + // Error message for failed turns + if (turn.state === TurnState.Error && turn.error) { + parts.push({ kind: 'markdownContent', content: new MarkdownString(`\n\nError: (${turn.error.errorType}) ${turn.error.message}`) }); + } + + history.push({ type: 'response', parts, participant: participantId }); + } + return history; +} + +/** + * Converts a completed tool call from the protocol state into a serialized + * tool invocation suitable for history replay. + */ +function completedToolCallToSerialized(tc: ICompletedToolCall): IChatToolInvocationSerialized { + const isTerminal = tc.toolKind === 'terminal'; + + let toolSpecificData: IChatTerminalToolInvocationData | undefined; + if (isTerminal && tc.toolInput) { + toolSpecificData = { + kind: 'terminal', + commandLine: { original: tc.toolInput }, + language: tc.language ?? 'shellscript', + terminalCommandOutput: tc.toolOutput !== undefined ? { text: tc.toolOutput } : undefined, + terminalCommandState: { exitCode: tc.success ? 0 : 1 }, + }; + } + + return { + kind: 'toolInvocationSerialized', + toolCallId: tc.toolCallId, + toolId: tc.toolName, + source: ToolDataSource.Internal, + invocationMessage: new MarkdownString(tc.invocationMessage), + originMessage: undefined, + pastTenseMessage: isTerminal ? undefined : new MarkdownString(tc.pastTenseMessage), + isConfirmed: { type: ToolConfirmKind.ConfirmationNotNeeded }, + isComplete: true, + presentation: undefined, + toolSpecificData, + }; +} + +/** + * Creates a live {@link ChatToolInvocation} from the protocol's tool-call + * state. Used during active turns to represent running tool calls in the UI. + */ +export function toolCallStateToInvocation(tc: IToolCallState): ChatToolInvocation { + const toolData: IToolData = { + id: tc.toolName, + source: ToolDataSource.Internal, + displayName: tc.displayName, + modelDescription: tc.toolName, + }; + + let parameters: unknown; + if (tc.toolArguments) { + try { parameters = JSON.parse(tc.toolArguments); } catch { /* malformed JSON */ } + } + + const invocation = new ChatToolInvocation(undefined, toolData, tc.toolCallId, undefined, parameters); + invocation.invocationMessage = new MarkdownString(tc.invocationMessage); + + if (tc.toolKind === 'terminal' && tc.toolInput) { + invocation.toolSpecificData = { + kind: 'terminal', + commandLine: { original: tc.toolInput }, + language: tc.language ?? 'shellscript', + } satisfies IChatTerminalToolInvocationData; + } + + return invocation; +} + +/** + * Creates a {@link ChatToolInvocation} with confirmation messages from a + * protocol permission request. The resulting invocation starts in the + * waiting-for-confirmation state. + */ +export function permissionToConfirmation(perm: IPermissionRequest): ChatToolInvocation { + let title: string; + let toolSpecificData: IChatTerminalToolInvocationData | IChatToolInputInvocationData | undefined; + + switch (perm.permissionKind) { + case 'shell': { + title = perm.intention ?? 'Run command'; + toolSpecificData = perm.fullCommandText ? { + kind: 'terminal', + commandLine: { original: perm.fullCommandText }, + language: 'shellscript', + } : undefined; + break; + } + case 'write': { + title = perm.path ? `Edit ${perm.path}` : 'Edit file'; + let rawInput: unknown; + try { rawInput = perm.rawRequest ? JSON.parse(perm.rawRequest) : { path: perm.path }; } catch { rawInput = { path: perm.path }; } + toolSpecificData = { kind: 'input', rawInput }; + break; + } + case 'mcp': { + const toolTitle = perm.toolName ?? 'MCP Tool'; + title = perm.serverName ? `${perm.serverName}: ${toolTitle}` : toolTitle; + let rawInput: unknown; + try { rawInput = perm.rawRequest ? JSON.parse(perm.rawRequest) : { serverName: perm.serverName, toolName: perm.toolName }; } catch { rawInput = { serverName: perm.serverName, toolName: perm.toolName }; } + toolSpecificData = { kind: 'input', rawInput }; + break; + } + case 'read': { + title = perm.intention ?? 'Read file'; + let rawInput: unknown; + try { rawInput = perm.rawRequest ? JSON.parse(perm.rawRequest) : { path: perm.path, intention: perm.intention }; } catch { rawInput = { path: perm.path, intention: perm.intention }; } + toolSpecificData = { kind: 'input', rawInput }; + break; + } + default: { + title = 'Permission request'; + let rawInput: unknown; + try { rawInput = perm.rawRequest ? JSON.parse(perm.rawRequest) : {}; } catch { rawInput = {}; } + toolSpecificData = { kind: 'input', rawInput }; + break; + } + } + + const confirmationMessages: IToolConfirmationMessages = { + title: new MarkdownString(title), + message: new MarkdownString(''), + }; + + const toolData: IToolData = { + id: `permission_${perm.permissionKind}`, + source: ToolDataSource.Internal, + displayName: title, + modelDescription: '', + }; + + const preparedInvocation: IPreparedToolInvocation = { + invocationMessage: new MarkdownString(title), + confirmationMessages, + presentation: ToolInvocationPresentation.HiddenAfterComplete, + toolSpecificData, + }; + + return new ChatToolInvocation(preparedInvocation, toolData, perm.requestId, undefined, undefined); +} + +/** + * Updates a live {@link ChatToolInvocation} with completion data from the + * protocol's tool-call state, transitioning it to the completed state. + */ +export function finalizeToolInvocation(invocation: ChatToolInvocation, tc: IToolCallState): void { + if (invocation.toolSpecificData?.kind === 'terminal') { + const terminalData = invocation.toolSpecificData as IChatTerminalToolInvocationData; + invocation.toolSpecificData = { + ...terminalData, + terminalCommandOutput: tc.toolOutput !== undefined ? { text: tc.toolOutput } : undefined, + terminalCommandState: { exitCode: tc.status === ToolCallStatus.Completed ? 0 : 1 }, + }; + } else if (tc.pastTenseMessage) { + invocation.pastTenseMessage = new MarkdownString(tc.pastTenseMessage); + } + + const isFailure = tc.status === ToolCallStatus.Failed; + invocation.didExecuteTool(isFailure ? { content: [], toolResultError: tc.error?.message } : undefined); +} 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 6c2ac5cddec..b3889c6e9ee 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 @@ -13,7 +13,10 @@ import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/ import { timeout } from '../../../../../../base/common/async.js'; import { ILogService, NullLogService } from '../../../../../../platform/log/common/log.js'; import { IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js'; -import { IAgentAttachment, IAgentHostService, IAgentMessageEvent, IAgentModelInfo, IAgentProgressEvent, IAgentSessionMetadata, IAgentToolCompleteEvent, IAgentToolStartEvent, AgentSession } from '../../../../../../platform/agentHost/common/agentService.js'; +import { IAgentCreateSessionConfig, IAgentHostService, IAgentSessionMetadata, AgentSession } from '../../../../../../platform/agentHost/common/agentService.js'; +import type { IActionEnvelope, INotification, IPermissionResolvedAction, ISessionAction, ITurnStartedAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; +import type { IStateSnapshot } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js'; +import { SessionLifecycle, SessionStatus, ToolCallStatus, TurnState, createSessionState, 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'; @@ -34,16 +37,16 @@ import { AgentHostLanguageModelProvider } from '../../../browser/agentSessions/a class MockAgentHostService extends mock() { declare readonly _serviceBrand: undefined; - private readonly _onDidSessionProgress = new Emitter(); - override readonly onDidSessionProgress = this._onDidSessionProgress.event; + private readonly _onDidAction = new Emitter(); + override readonly onDidAction = this._onDidAction.event; + private readonly _onDidNotification = new Emitter(); + override readonly onDidNotification = this._onDidNotification.event; override readonly onAgentHostExit = Event.None; override readonly onAgentHostStart = Event.None; private _nextId = 1; private readonly _sessions = new Map(); - private readonly _sessionMessages = new Map(); - public sendMessageCalls: { session: URI; prompt: string; attachments?: IAgentAttachment[] }[] = []; - public models: IAgentModelInfo[] = []; + public createSessionCalls: IAgentCreateSessionConfig[] = []; public agents = [{ provider: 'copilot' as const, displayName: 'Agent Host - Copilot', description: 'test', requiresAuth: true }]; override async setAuthToken(_token: string): Promise { } @@ -56,9 +59,7 @@ class MockAgentHostService extends mock() { return this.agents; } - override async listModels(): Promise { - return this.models; - } + override async refreshModels(): Promise { } override async createSession(): Promise { const id = `sdk-session-${this._nextId++}`; @@ -67,29 +68,51 @@ class MockAgentHostService extends mock() { return session; } - override async sendMessage(session: URI, prompt: string, attachments?: IAgentAttachment[]): Promise { - this.sendMessageCalls.push({ session, prompt, attachments }); - } - - override async getSessionMessages(session: URI): Promise<(IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[]> { - return this._sessionMessages.get(AgentSession.id(session)) ?? []; - } - override async disposeSession(_session: URI): Promise { } - public abortSessionCalls: URI[] = []; - override async abortSession(session: URI): Promise { this.abortSessionCalls.push(session); } - public permissionResponses: { requestId: string; approved: boolean }[] = []; - override respondToPermissionRequest(requestId: string, approved: boolean): void { this.permissionResponses.push({ requestId, approved }); } override async shutdown(): Promise { } override async restartAgentHost(): Promise { } - // Test helpers - fireProgress(event: IAgentProgressEvent): void { - this._onDidSessionProgress.fire(event); + // Protocol methods + public override readonly clientId = 'test-window-1'; + public dispatchedActions: { action: ISessionAction; clientId: string; clientSeq: number }[] = []; + public sessionStates = new Map(); + override async subscribe(resource: URI): Promise { + const existingState = this.sessionStates.get(resource.toString()); + if (existingState) { + return { resource, state: existingState, fromSeq: 0 }; + } + // Root state subscription + if (resource.scheme === 'agenthost') { + return { + resource, + state: { + agents: this.agents.map(a => ({ provider: a.provider, displayName: a.displayName, description: a.description, models: [] })), + }, + fromSeq: 0, + }; + } + const summary: ISessionSummary = { + resource, + provider: 'copilot', + title: 'Test', + status: SessionStatus.Idle, + createdAt: Date.now(), + modifiedAt: Date.now(), + }; + return { + resource, + state: { ...createSessionState(summary), lifecycle: SessionLifecycle.Ready }, + fromSeq: 0, + }; + } + override unsubscribe(_resource: URI): void { } + override dispatchAction(action: ISessionAction, clientId: string, clientSeq: number): void { + this.dispatchedActions.push({ action, clientId, clientSeq }); } - setSessionMessages(sessionId: string, messages: (IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[]): void { - this._sessionMessages.set(sessionId, messages); + // Test helpers + fireAction(envelope: IActionEnvelope): void { + this._onDidAction.fire(envelope); } addSession(meta: IAgentSessionMetadata): void { @@ -97,7 +120,8 @@ class MockAgentHostService extends mock() { } dispose(): void { - this._onDidSessionProgress.dispose(); + this._onDidAction.dispose(); + this._onDidNotification.dispose(); } } @@ -162,7 +186,7 @@ function createContribution(disposables: DisposableStore) { return { contribution, listController, sessionHandler, agentHostService, chatAgentService }; } -function makeRequest(overrides: Partial<{ message: string; sessionResource: URI; variables: IChatAgentRequest['variables'] }> = {}): IChatAgentRequest { +function makeRequest(overrides: Partial<{ message: string; sessionResource: URI; variables: IChatAgentRequest['variables']; userSelectedModelId: string }> = {}): IChatAgentRequest { return upcastPartial({ sessionResource: overrides.sessionResource ?? URI.from({ scheme: 'untitled', path: '/chat-1' }), requestId: 'req-1', @@ -170,6 +194,7 @@ function makeRequest(overrides: Partial<{ message: string; sessionResource: URI; message: overrides.message ?? 'Hello', variables: overrides.variables ?? { variables: [] }, location: ChatAgentLocation.Chat, + userSelectedModelId: overrides.userSelectedModelId, }); } @@ -181,6 +206,66 @@ function textOf(value: string | IMarkdownString | undefined): string | undefined return typeof value === 'string' ? value : value.value; } +/** + * Start a turn through the state-driven flow. Creates a chat session, + * starts the requestHandler (non-blocking), and waits for the first action + * to be dispatched. Returns helpers to fire server action envelopes. + */ +async function startTurn( + sessionHandler: AgentHostSessionHandler, + agentHostService: MockAgentHostService, + ds: DisposableStore, + overrides?: Partial<{ + message: string; + sessionResource: URI; + variables: IChatAgentRequest['variables']; + userSelectedModelId: string; + cancellationToken: CancellationToken; + }>, +) { + const sessionResource = overrides?.sessionResource ?? URI.from({ scheme: 'agent-host-copilot', path: '/untitled-turntest' }); + const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + ds.add(toDisposable(() => chatSession.dispose())); + + const collected: IChatProgress[][] = []; + const seq = { v: 1 }; + + const turnPromise = chatSession.requestHandler!( + makeRequest({ + message: overrides?.message ?? 'Hello', + sessionResource, + variables: overrides?.variables, + userSelectedModelId: overrides?.userSelectedModelId, + }), + (parts) => collected.push(parts), + [], + overrides?.cancellationToken ?? CancellationToken.None, + ); + + await timeout(10); + + const lastDispatch = agentHostService.dispatchedActions[agentHostService.dispatchedActions.length - 1]; + const session = (lastDispatch?.action as ITurnStartedAction)?.session; + const turnId = (lastDispatch?.action as ITurnStartedAction)?.turnId; + + const fire = (action: ISessionAction) => { + agentHostService.fireAction({ action, serverSeq: seq.v++, origin: undefined }); + }; + + // Echo the turnStarted action to clear the pending write-ahead entry. + // Without this, the optimistic state replay would re-add activeTurn after + // the server's turnComplete clears it, preventing the turn from finishing. + if (lastDispatch) { + agentHostService.fireAction({ + action: lastDispatch.action, + serverSeq: seq.v++, + origin: { clientId: agentHostService.clientId, clientSeq: lastDispatch.clientSeq }, + }); + } + + return { turnPromise, collected, chatSession, session, turnId, fire }; +} + suite('AgentHostChatContribution', () => { const disposables = new DisposableStore(); @@ -246,88 +331,120 @@ suite('AgentHostChatContribution', () => { suite('session ID resolution', () => { test('creates new SDK session for untitled resource', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { message: 'Hello' }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; - const origSend = agentHostService.sendMessage.bind(agentHostService); - agentHostService.sendMessage = async (session: URI, prompt: string) => { - await origSend(session, prompt); - agentHostService.fireProgress({ session, type: 'idle' }); - }; - - await agent.impl.invoke( - makeRequest({ message: 'Hello' }), - () => { }, [], CancellationToken.None, - ); - - assert.strictEqual(agentHostService.sendMessageCalls.length, 1); - assert.strictEqual(agentHostService.sendMessageCalls[0].prompt, 'Hello'); - assert.ok(AgentSession.id(agentHostService.sendMessageCalls[0].session).startsWith('sdk-session-')); + assert.strictEqual(agentHostService.dispatchedActions.length, 1); + assert.strictEqual(agentHostService.dispatchedActions[0].action.type, 'session/turnStarted'); + assert.strictEqual((agentHostService.dispatchedActions[0].action as ITurnStartedAction).userMessage.text, 'Hello'); + assert.ok(AgentSession.id(session).startsWith('sdk-session-')); }); test('reuses SDK session for same resource on second message', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const resource = URI.from({ scheme: 'untitled', path: '/chat-reuse' }); + const resource = URI.from({ scheme: 'agent-host-copilot', path: '/untitled-reuse' }); + const chatSession = await sessionHandler.provideChatSessionContent(resource, CancellationToken.None); + disposables.add(toDisposable(() => chatSession.dispose())); - agentHostService.sendMessage = async (session: URI, prompt: string) => { - agentHostService.sendMessageCalls.push({ session, prompt }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; - - await agent.impl.invoke( + // First turn + const turn1Promise = chatSession.requestHandler!( makeRequest({ message: 'First', sessionResource: resource }), () => { }, [], CancellationToken.None, ); + await timeout(10); + const dispatch1 = agentHostService.dispatchedActions[0]; + const action1 = dispatch1.action as ITurnStartedAction; + // Echo the turnStarted to clear pending write-ahead + agentHostService.fireAction({ action: dispatch1.action, serverSeq: 1, origin: { clientId: agentHostService.clientId, clientSeq: dispatch1.clientSeq } }); + agentHostService.fireAction({ action: { type: 'session/turnComplete', session: action1.session, turnId: action1.turnId } as ISessionAction, serverSeq: 2, origin: undefined }); + await turn1Promise; - await agent.impl.invoke( + // Second turn + const turn2Promise = chatSession.requestHandler!( makeRequest({ message: 'Second', sessionResource: resource }), () => { }, [], CancellationToken.None, ); + await timeout(10); + const dispatch2 = agentHostService.dispatchedActions[1]; + const action2 = dispatch2.action as ITurnStartedAction; + agentHostService.fireAction({ action: dispatch2.action, serverSeq: 3, origin: { clientId: agentHostService.clientId, clientSeq: dispatch2.clientSeq } }); + agentHostService.fireAction({ action: { type: 'session/turnComplete', session: action2.session, turnId: action2.turnId } as ISessionAction, serverSeq: 4, origin: undefined }); + await turn2Promise; - assert.strictEqual(agentHostService.sendMessageCalls.length, 2); - assert.strictEqual(agentHostService.sendMessageCalls[0].session.toString(), agentHostService.sendMessageCalls[1].session.toString()); + assert.strictEqual(agentHostService.dispatchedActions.length, 2); + assert.strictEqual( + (agentHostService.dispatchedActions[0].action as ITurnStartedAction).session.toString(), + (agentHostService.dispatchedActions[1].action as ITurnStartedAction).session.toString(), + ); }); test('uses sessionId from agent-host scheme resource', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const resource = URI.from({ scheme: 'agent-host-copilot', path: '/existing-session-42' }); + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + message: 'Hi', + sessionResource: URI.from({ scheme: 'agent-host-copilot', path: '/existing-session-42' }), + }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; - agentHostService.sendMessage = async (session: URI, prompt: string) => { - agentHostService.sendMessageCalls.push({ session, prompt }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; - - await agent.impl.invoke( - makeRequest({ message: 'Hi', sessionResource: resource }), - () => { }, [], CancellationToken.None, - ); - - assert.strictEqual(AgentSession.id(agentHostService.sendMessageCalls[0].session), 'existing-session-42'); + assert.strictEqual(AgentSession.id(session), 'existing-session-42'); }); test('agent-host scheme with untitled path creates new session via mapping', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const resource = URI.from({ scheme: 'agent-host-copilot', path: '/untitled-abc123' }); - - agentHostService.sendMessage = async (session: URI, prompt: string) => { - agentHostService.sendMessageCalls.push({ session, prompt }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; - - await agent.impl.invoke( - makeRequest({ message: 'Hi', sessionResource: resource }), - () => { }, [], CancellationToken.None, - ); + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + message: 'Hi', + sessionResource: URI.from({ scheme: 'agent-host-copilot', path: '/untitled-abc123' }), + }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; // Should create a new SDK session, not use "untitled-abc123" literally - assert.ok(AgentSession.id(agentHostService.sendMessageCalls[0].session).startsWith('sdk-session-')); + assert.ok(AgentSession.id(session).startsWith('sdk-session-')); + }); + test('passes raw model id extracted from language model identifier', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + message: 'Hi', + userSelectedModelId: 'agent-host-copilot:claude-sonnet-4-20250514', + }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; + + assert.strictEqual(agentHostService.createSessionCalls.length, 1); + assert.strictEqual(agentHostService.createSessionCalls[0].model, 'claude-sonnet-4-20250514'); + }); + + test('passes model id as-is when no vendor prefix', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + message: 'Hi', + userSelectedModelId: 'gpt-4o', + }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; + + assert.strictEqual(agentHostService.createSessionCalls.length, 1); + assert.strictEqual(agentHostService.createSessionCalls[0].model, 'gpt-4o'); + }); + + test('does not create backend session eagerly for untitled sessions', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); + + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/untitled-deferred' }); + const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => session.dispose())); + + // No backend session should have been created yet + assert.strictEqual(agentHostService.createSessionCalls.length, 0); }); }); @@ -336,23 +453,15 @@ suite('AgentHostChatContribution', () => { suite('progress routing', () => { test('delta events become markdownContent progress', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ session, type: 'delta', messageId: 'msg-1', content: 'hello ' }); - agentHostService.fireProgress({ session, type: 'delta', messageId: 'msg-1', content: 'world' }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + fire({ type: 'session/delta', session, turnId, content: 'hello ' } as ISessionAction); + fire({ type: 'session/delta', session, turnId, content: 'world' } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); + await turnPromise; assert.strictEqual(collected.length, 2); assert.strictEqual(collected[0][0].kind, 'markdownContent'); @@ -362,45 +471,38 @@ suite('AgentHostChatContribution', () => { }); test('tool_start events become toolInvocation progress', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ session, type: 'tool_start', toolCallId: 'tc-1', toolName: 'read_file', displayName: 'Read File', invocationMessage: 'Reading file' }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + fire({ + type: 'session/toolStart', session, turnId, + toolCall: { toolCallId: 'tc-1', toolName: 'read_file', displayName: 'Read File', invocationMessage: 'Reading file', status: ToolCallStatus.Running }, + } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); + await turnPromise; assert.strictEqual(collected.length, 1); assert.strictEqual(collected[0][0].kind, 'toolInvocation'); }); test('tool_complete event transitions toolInvocation to completed', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ session, type: 'tool_start', toolCallId: 'tc-2', toolName: 'bash', displayName: 'Bash', invocationMessage: 'Running Bash command' }); - agentHostService.fireProgress({ session, type: 'tool_complete', toolCallId: 'tc-2', success: true, pastTenseMessage: 'Ran Bash command' }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + fire({ + type: 'session/toolStart', session, turnId, + toolCall: { toolCallId: 'tc-2', toolName: 'bash', displayName: 'Bash', invocationMessage: 'Running Bash command', status: ToolCallStatus.Running }, + } as ISessionAction); + fire({ + type: 'session/toolComplete', session, turnId, toolCallId: 'tc-2', + result: { success: true, pastTenseMessage: 'Ran Bash command' }, + } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); + await turnPromise; assert.strictEqual(collected.length, 1); const invocation = collected[0][0] as IChatToolInvocation; @@ -410,23 +512,21 @@ suite('AgentHostChatContribution', () => { }); test('tool_complete with failure sets error state', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ session, type: 'tool_start', toolCallId: 'tc-3', toolName: 'bash', displayName: 'Bash', invocationMessage: 'Running Bash command' }); - agentHostService.fireProgress({ session, type: 'tool_complete', toolCallId: 'tc-3', success: false, pastTenseMessage: '"Bash" failed', toolOutput: 'command not found', error: { message: 'command not found' } }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + fire({ + type: 'session/toolStart', session, turnId, + toolCall: { toolCallId: 'tc-3', toolName: 'bash', displayName: 'Bash', invocationMessage: 'Running Bash command', status: ToolCallStatus.Running }, + } as ISessionAction); + fire({ + type: 'session/toolComplete', session, turnId, toolCallId: 'tc-3', + result: { success: false, pastTenseMessage: '"Bash" failed', toolOutput: 'command not found', error: { message: 'command not found' } }, + } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); + await turnPromise; assert.strictEqual(collected.length, 1); const invocation = collected[0][0] as IChatToolInvocation; @@ -435,45 +535,35 @@ suite('AgentHostChatContribution', () => { }); test('malformed toolArguments does not throw', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ session, type: 'tool_start', toolCallId: 'tc-bad', toolName: 'bash', displayName: 'Bash', invocationMessage: 'Running Bash command', toolArguments: '{not valid json' }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + fire({ + type: 'session/toolStart', session, turnId, + toolCall: { toolCallId: 'tc-bad', toolName: 'bash', displayName: 'Bash', invocationMessage: 'Running Bash command', status: ToolCallStatus.Running, toolArguments: '{not valid json' }, + } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); + await turnPromise; assert.strictEqual(collected.length, 1); assert.strictEqual(collected[0][0].kind, 'toolInvocation'); }); test('outstanding tool invocations are completed on idle', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - // tool_start without tool_complete - agentHostService.fireProgress({ session, type: 'tool_start', toolCallId: 'tc-orphan', toolName: 'bash', displayName: 'Bash', invocationMessage: 'Running Bash command' }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + // tool_start without tool_complete + fire({ + type: 'session/toolStart', session, turnId, + toolCall: { toolCallId: 'tc-orphan', toolName: 'bash', displayName: 'Bash', invocationMessage: 'Running Bash command', status: ToolCallStatus.Running }, + } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); + await turnPromise; assert.strictEqual(collected.length, 1); const invocation = collected[0][0] as IChatToolInvocation; @@ -482,23 +572,20 @@ suite('AgentHostChatContribution', () => { }); test('events from other sessions are ignored', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ session: AgentSession.uri('copilot', 'other-session'), type: 'delta', messageId: 'msg-x', content: 'wrong' }); - agentHostService.fireProgress({ session, type: 'delta', messageId: 'msg-y', content: 'right' }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + // Delta from a different session — will be ignored (session not subscribed) + agentHostService.fireAction({ + action: { type: 'session/delta', session: AgentSession.uri('copilot', 'other-session'), turnId, content: 'wrong' } as ISessionAction, + serverSeq: 100, + origin: undefined, + }); + fire({ type: 'session/delta', session, turnId, content: 'right' } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); + await turnPromise; assert.strictEqual(collected.length, 1); assert.strictEqual((collected[0][0] as IChatMarkdownContent).content.value, 'right'); @@ -510,44 +597,38 @@ suite('AgentHostChatContribution', () => { suite('cancellation', () => { test('cancellation resolves the agent invoke', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; const cts = new CancellationTokenSource(); disposables.add(cts); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - cts.cancel(); - }; + const { turnPromise } = await startTurn(sessionHandler, agentHostService, disposables, { + cancellationToken: cts.token, + }); - const result = await agent.impl.invoke( - makeRequest(), - () => { }, [], cts.token, - ); + cts.cancel(); + await turnPromise; - assert.ok(result); + assert.ok(agentHostService.dispatchedActions.some(a => a.action.type === 'session/turnCancelled')); }); test('cancellation force-completes outstanding tool invocations', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; const cts = new CancellationTokenSource(); disposables.add(cts); - const collected: IChatProgress[][] = []; - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ session, type: 'tool_start', toolCallId: 'tc-cancel', toolName: 'bash', displayName: 'Bash', invocationMessage: 'Running Bash command' }); - cts.cancel(); - }; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + cancellationToken: cts.token, + }); - await agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], cts.token, - ); + fire({ + type: 'session/toolStart', session, turnId, + toolCall: { toolCallId: 'tc-cancel', toolName: 'bash', displayName: 'Bash', invocationMessage: 'Running Bash command', status: ToolCallStatus.Running }, + } as ISessionAction); + + cts.cancel(); + await turnPromise; assert.strictEqual(collected.length, 1); const invocation = collected[0][0] as IChatToolInvocation; @@ -556,23 +637,20 @@ suite('AgentHostChatContribution', () => { }); test('cancellation calls abortSession on the agent host service', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; const cts = new CancellationTokenSource(); disposables.add(cts); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - cts.cancel(); - }; + const { turnPromise } = await startTurn(sessionHandler, agentHostService, disposables, { + cancellationToken: cts.token, + }); - await agent.impl.invoke( - makeRequest(), - () => { }, [], cts.token, - ); + cts.cancel(); + await turnPromise; - assert.strictEqual(agentHostService.abortSessionCalls.length, 1); + // Cancellation now dispatches session/turnCancelled action + assert.ok(agentHostService.dispatchedActions.some(a => a.action.type === 'session/turnCancelled')); }); }); @@ -581,26 +659,27 @@ suite('AgentHostChatContribution', () => { suite('error events', () => { test('error event renders error message and finishes the request', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ session, type: 'error', errorType: 'test_error', message: 'Something went wrong' }); - }; + agentHostService.fireAction({ + action: { + type: 'session/error', + session, + turnId, + error: { errorType: 'test_error', message: 'Something went wrong' }, + } as ISessionAction, + serverSeq: 99, + origin: undefined, + }); - await agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); + await turnPromise; // Should have received the error message and the request should have finished - assert.strictEqual(collected.length, 1); - assert.strictEqual(collected[0][0].kind, 'markdownContent'); - assert.ok((collected[0][0] as IChatMarkdownContent).content.value.includes('Something went wrong')); + assert.ok(collected.length >= 1); + const errorPart = collected.flat().find(p => p.kind === 'markdownContent' && (p as IChatMarkdownContent).content.value.includes('Something went wrong')); + assert.ok(errorPart, 'Should have found a markdownContent part containing the error message'); }); }); @@ -609,32 +688,16 @@ suite('AgentHostChatContribution', () => { suite('permission requests', () => { test('permission_request event shows confirmation and responds when confirmed', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - // Simulate a permission request - agentHostService.fireProgress({ - session, - type: 'permission_request', - requestId: 'perm-1', - permissionKind: 'shell', - fullCommandText: 'echo hello', - rawRequest: '{}', - }); - }; + // Simulate a permission request + fire({ + type: 'session/permissionRequest', session, turnId, + request: { requestId: 'perm-1', permissionKind: 'shell', fullCommandText: 'echo hello', rawRequest: '{}' }, + } as ISessionAction); - // Start the invoke but don't await yet -- we need to confirm the permission - const invokePromise = agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); - - // Wait for the permission confirmation to appear await timeout(10); // The permission request should have produced a ChatToolInvocation in WaitingForConfirmation state @@ -645,42 +708,26 @@ suite('AgentHostChatContribution', () => { // Confirm the permission IChatToolInvocation.confirmWith(permInvocation, { type: ToolConfirmKind.UserAction }); - // Now fire idle to complete the request await timeout(10); - const session = agentHostService.sendMessageCalls[0].session; - agentHostService.fireProgress({ session, type: 'idle' }); - await invokePromise; + // The handler should have dispatched session/permissionResolved + assert.ok(agentHostService.dispatchedActions.some( + a => a.action.type === 'session/permissionResolved' && (a.action as IPermissionResolvedAction).requestId === 'perm-1' && (a.action as IPermissionResolvedAction).approved === true + )); - // The handler should have approved the permission request - assert.strictEqual(agentHostService.permissionResponses.length, 1); - assert.strictEqual(agentHostService.permissionResponses[0].requestId, 'perm-1'); - assert.strictEqual(agentHostService.permissionResponses[0].approved, true); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; }); test('permission_request denied when user skips', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ - session, - type: 'permission_request', - requestId: 'perm-2', - permissionKind: 'write', - path: '/tmp/test.txt', - rawRequest: '{}', - }); - }; - - const invokePromise = agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); + fire({ + type: 'session/permissionRequest', session, turnId, + request: { requestId: 'perm-2', permissionKind: 'write', path: '/tmp/test.txt', rawRequest: '{}' }, + } as ISessionAction); await timeout(10); @@ -689,40 +736,24 @@ suite('AgentHostChatContribution', () => { IChatToolInvocation.confirmWith(permInvocation, { type: ToolConfirmKind.Denied }); await timeout(10); - const session = agentHostService.sendMessageCalls[0].session; - agentHostService.fireProgress({ session, type: 'idle' }); - await invokePromise; + assert.ok(agentHostService.dispatchedActions.some( + a => a.action.type === 'session/permissionResolved' && (a.action as IPermissionResolvedAction).requestId === 'perm-2' && (a.action as IPermissionResolvedAction).approved === false + )); - assert.strictEqual(agentHostService.permissionResponses.length, 1); - assert.strictEqual(agentHostService.permissionResponses[0].requestId, 'perm-2'); - assert.strictEqual(agentHostService.permissionResponses[0].approved, false); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; }); test('shell permission shows terminal-style confirmation data', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ - session, - type: 'permission_request', - requestId: 'perm-shell', - permissionKind: 'shell', - fullCommandText: 'echo hello', - intention: 'Print greeting', - rawRequest: '{}', - }); - }; - - const invokePromise = agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); + fire({ + type: 'session/permissionRequest', session, turnId, + request: { requestId: 'perm-shell', permissionKind: 'shell', fullCommandText: 'echo hello', intention: 'Print greeting', rawRequest: '{}' }, + } as ISessionAction); await timeout(10); const permInvocation = collected[0][0] as IChatToolInvocation; @@ -732,34 +763,19 @@ suite('AgentHostChatContribution', () => { IChatToolInvocation.confirmWith(permInvocation, { type: ToolConfirmKind.UserAction }); await timeout(10); - agentHostService.fireProgress({ session: agentHostService.sendMessageCalls[0].session, type: 'idle' }); - await invokePromise; + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; }); test('read permission shows input-style confirmation data', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ - session, - type: 'permission_request', - requestId: 'perm-read', - permissionKind: 'read', - path: '/workspace/file.ts', - intention: 'Read file contents', - rawRequest: '{"kind":"read","path":"/workspace/file.ts"}', - }); - }; - - const invokePromise = agent.impl.invoke( - makeRequest(), - (parts) => collected.push(parts), - [], CancellationToken.None, - ); + fire({ + type: 'session/permissionRequest', session, turnId, + request: { requestId: 'perm-read', permissionKind: 'read', path: '/workspace/file.ts', intention: 'Read file contents', rawRequest: '{"kind":"read","path":"/workspace/file.ts"}' }, + } as ISessionAction); await timeout(10); const permInvocation = collected[0][0] as IChatToolInvocation; @@ -767,8 +783,8 @@ suite('AgentHostChatContribution', () => { IChatToolInvocation.confirmWith(permInvocation, { type: ToolConfirmKind.UserAction }); await timeout(10); - agentHostService.fireProgress({ session: agentHostService.sendMessageCalls[0].session, type: 'idle' }); - await invokePromise; + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; }); }); @@ -779,10 +795,20 @@ suite('AgentHostChatContribution', () => { test('loads user and assistant messages into history', async () => { const { sessionHandler, agentHostService } = createContribution(disposables); - agentHostService.setSessionMessages('sess-1', [ - { session: AgentSession.uri('copilot', 'sess-1'), type: 'message', messageId: 'msg-u1', content: 'What is 2+2?', role: 'user' }, - { session: AgentSession.uri('copilot', 'sess-1'), type: 'message', messageId: 'msg-a1', content: '4', role: 'assistant' }, - ]); + const sessionUri = AgentSession.uri('copilot', 'sess-1'); + agentHostService.sessionStates.set(sessionUri.toString(), { + ...createSessionState({ resource: sessionUri, provider: 'copilot', title: 'Test', status: SessionStatus.Idle, createdAt: Date.now(), modifiedAt: Date.now() }), + lifecycle: SessionLifecycle.Ready, + turns: [{ + id: 'turn-1', + userMessage: { text: 'What is 2+2?' }, + responseText: '4', + responseParts: [], + toolCalls: [], + usage: undefined, + state: TurnState.Complete, + }], + }); const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/sess-1' }); const session = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); @@ -799,6 +825,7 @@ suite('AgentHostChatContribution', () => { const response = session.history[1]; assert.strictEqual(response.type, 'response'); if (response.type === 'response') { + assert.strictEqual(response.parts.length, 1); assert.strictEqual((response.parts[0] as IChatMarkdownContent).content.value, '4'); } }); @@ -819,24 +846,26 @@ suite('AgentHostChatContribution', () => { suite('tool invocation rendering', () => { test('bash tool renders as terminal command block with output', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ - session, type: 'tool_start', toolCallId: 'tc-shell', toolName: 'bash', - displayName: 'Bash', invocationMessage: 'Running `echo hello`', - toolInput: 'echo hello', toolKind: 'terminal', + fire({ + type: 'session/toolStart', session, turnId, + toolCall: { + toolCallId: 'tc-shell', toolName: 'bash', displayName: 'Bash', + invocationMessage: 'Running `echo hello`', toolInput: 'echo hello', + toolKind: 'terminal', status: ToolCallStatus.Running, toolArguments: JSON.stringify({ command: 'echo hello' }), - }); - agentHostService.fireProgress({ session, type: 'tool_complete', toolCallId: 'tc-shell', success: true, pastTenseMessage: 'Ran `echo hello`', toolOutput: 'hello\n', result: { content: 'hello\n' } }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + }, + } as ISessionAction); + fire({ + type: 'session/toolComplete', session, turnId, toolCallId: 'tc-shell', + result: { success: true, pastTenseMessage: 'Ran `echo hello`', toolOutput: 'hello\n' }, + } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke(makeRequest(), (parts) => collected.push(parts), [], CancellationToken.None); + await turnPromise; const invocation = collected[0][0] as IChatToolInvocation; const termData = invocation.toolSpecificData as IChatTerminalToolInvocationData; @@ -862,28 +891,26 @@ suite('AgentHostChatContribution', () => { }); test('bash tool failure sets exit code 1 and error output', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ - session, type: 'tool_start', toolCallId: 'tc-fail', toolName: 'bash', - displayName: 'Bash', invocationMessage: 'Running `bad_cmd`', - toolInput: 'bad_cmd', toolKind: 'terminal', + fire({ + type: 'session/toolStart', session, turnId, + toolCall: { + toolCallId: 'tc-fail', toolName: 'bash', displayName: 'Bash', + invocationMessage: 'Running `bad_cmd`', toolInput: 'bad_cmd', + toolKind: 'terminal', status: ToolCallStatus.Running, toolArguments: JSON.stringify({ command: 'bad_cmd' }), - }); - agentHostService.fireProgress({ - session, type: 'tool_complete', toolCallId: 'tc-fail', success: false, - pastTenseMessage: '"Bash" failed', toolOutput: 'command not found: bad_cmd', - error: { message: 'command not found: bad_cmd' }, - }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + }, + } as ISessionAction); + fire({ + type: 'session/toolComplete', session, turnId, toolCallId: 'tc-fail', + result: { success: false, pastTenseMessage: '"Bash" failed', toolOutput: 'command not found: bad_cmd', error: { message: 'command not found: bad_cmd' } }, + } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke(makeRequest(), (parts) => collected.push(parts), [], CancellationToken.None); + await turnPromise; const invocation = collected[0][0] as IChatToolInvocation; const termData = invocation.toolSpecificData as IChatTerminalToolInvocationData; @@ -899,23 +926,25 @@ suite('AgentHostChatContribution', () => { }); test('generic tool has invocation message and no toolSpecificData', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ - session, type: 'tool_start', toolCallId: 'tc-gen', toolName: 'custom_tool', - displayName: 'custom_tool', invocationMessage: 'Using "custom_tool"', + fire({ + type: 'session/toolStart', session, turnId, + toolCall: { + toolCallId: 'tc-gen', toolName: 'custom_tool', displayName: 'custom_tool', + invocationMessage: 'Using "custom_tool"', status: ToolCallStatus.Running, toolArguments: JSON.stringify({ input: 'data' }), - }); - agentHostService.fireProgress({ session, type: 'tool_complete', toolCallId: 'tc-gen', success: true, pastTenseMessage: 'Used "custom_tool"' }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + }, + } as ISessionAction); + fire({ + type: 'session/toolComplete', session, turnId, toolCallId: 'tc-gen', + result: { success: true, pastTenseMessage: 'Used "custom_tool"' }, + } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke(makeRequest(), (parts) => collected.push(parts), [], CancellationToken.None); + await turnPromise; const invocation = collected[0][0] as IChatToolInvocation; assert.deepStrictEqual({ @@ -930,23 +959,25 @@ suite('AgentHostChatContribution', () => { }); test('bash tool without arguments has no terminal data', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ - session, type: 'tool_start', toolCallId: 'tc-noargs', toolName: 'bash', - displayName: 'Bash', invocationMessage: 'Running Bash command', - toolKind: 'terminal', - }); - agentHostService.fireProgress({ session, type: 'tool_complete', toolCallId: 'tc-noargs', success: true, pastTenseMessage: 'Ran Bash command' }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + fire({ + type: 'session/toolStart', session, turnId, + toolCall: { + toolCallId: 'tc-noargs', toolName: 'bash', displayName: 'Bash', + invocationMessage: 'Running Bash command', toolKind: 'terminal', + status: ToolCallStatus.Running, + }, + } as ISessionAction); + fire({ + type: 'session/toolComplete', session, turnId, toolCallId: 'tc-noargs', + result: { success: true, pastTenseMessage: 'Ran Bash command' }, + } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke(makeRequest(), (parts) => collected.push(parts), [], CancellationToken.None); + await turnPromise; const invocation = collected[0][0] as IChatToolInvocation; assert.deepStrictEqual({ @@ -961,23 +992,25 @@ suite('AgentHostChatContribution', () => { }); test('view tool shows file path in messages', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; - const collected: IChatProgress[][] = []; + const { turnPromise, collected, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async (session: URI) => { - agentHostService.sendMessageCalls.push({ session, prompt: '' }); - agentHostService.fireProgress({ - session, type: 'tool_start', toolCallId: 'tc-view', toolName: 'view', - displayName: 'View File', invocationMessage: 'Reading /tmp/test.txt', + fire({ + type: 'session/toolStart', session, turnId, + toolCall: { + toolCallId: 'tc-view', toolName: 'view', displayName: 'View File', + invocationMessage: 'Reading /tmp/test.txt', status: ToolCallStatus.Running, toolArguments: JSON.stringify({ file_path: '/tmp/test.txt' }), - }); - agentHostService.fireProgress({ session, type: 'tool_complete', toolCallId: 'tc-view', success: true, pastTenseMessage: 'Read /tmp/test.txt' }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; + }, + } as ISessionAction); + fire({ + type: 'session/toolComplete', session, turnId, toolCallId: 'tc-view', + result: { success: true, pastTenseMessage: 'Read /tmp/test.txt' }, + } as ISessionAction); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); - await agent.impl.invoke(makeRequest(), (parts) => collected.push(parts), [], CancellationToken.None); + await turnPromise; const invocation = collected[0][0] as IChatToolInvocation; assert.deepStrictEqual({ @@ -996,14 +1029,25 @@ suite('AgentHostChatContribution', () => { test('tool_start and tool_complete appear as toolInvocationSerialized in history', async () => { const { sessionHandler, agentHostService } = createContribution(disposables); - const session = AgentSession.uri('copilot', 'tool-hist'); + const sessionUri = AgentSession.uri('copilot', 'tool-hist'); - agentHostService.setSessionMessages('tool-hist', [ - { session, type: 'message', messageId: 'u1', content: 'run ls', role: 'user' }, - { session, type: 'tool_start', toolCallId: 'tc-1', toolName: 'bash', displayName: 'Bash', invocationMessage: 'Running `ls`', toolInput: 'ls', toolKind: 'terminal' }, - { session, type: 'tool_complete', toolCallId: 'tc-1', success: true, pastTenseMessage: 'Ran `ls`', toolOutput: 'file1\nfile2' }, - { session, type: 'message', messageId: 'a1', content: 'Here are the files.', role: 'assistant' }, - ]); + agentHostService.sessionStates.set(sessionUri.toString(), { + ...createSessionState({ resource: sessionUri, provider: 'copilot', title: 'Test', status: SessionStatus.Idle, createdAt: Date.now(), modifiedAt: Date.now() }), + lifecycle: SessionLifecycle.Ready, + turns: [{ + id: 'turn-1', + userMessage: { text: 'run ls' }, + state: TurnState.Complete, + responseParts: [], + usage: undefined, + toolCalls: [{ + toolCallId: 'tc-1', toolName: 'bash', displayName: 'Bash', + invocationMessage: 'Running `ls`', toolInput: 'ls', toolKind: 'terminal' as const, + success: true, pastTenseMessage: 'Ran `ls`', toolOutput: 'file1\nfile2', + }], + responseText: '', + }], + } as ISessionState); const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/tool-hist' }); const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); @@ -1015,8 +1059,7 @@ suite('AgentHostChatContribution', () => { const response = chatSession.history[1]; assert.strictEqual(response.type, 'response'); if (response.type === 'response') { - // tool invocation + markdown content - assert.strictEqual(response.parts.length, 2); + assert.strictEqual(response.parts.length, 1); const toolPart = response.parts[0] as IChatToolInvocationSerialized; assert.strictEqual(toolPart.kind, 'toolInvocationSerialized'); assert.strictEqual(toolPart.toolCallId, 'tc-1'); @@ -1031,14 +1074,21 @@ suite('AgentHostChatContribution', () => { test('orphaned tool_start is marked complete in history', async () => { const { sessionHandler, agentHostService } = createContribution(disposables); - const session = AgentSession.uri('copilot', 'orphan-tool'); + const sessionUri = AgentSession.uri('copilot', 'orphan-tool'); - agentHostService.setSessionMessages('orphan-tool', [ - { session, type: 'message', messageId: 'u1', content: 'do something', role: 'user' }, - { session, type: 'tool_start', toolCallId: 'tc-orphan', toolName: 'read_file', displayName: 'Read File', invocationMessage: 'Reading file' }, - // No tool_complete for tc-orphan - { session, type: 'message', messageId: 'a1', content: 'Done', role: 'assistant' }, - ]); + agentHostService.sessionStates.set(sessionUri.toString(), { + ...createSessionState({ resource: sessionUri, provider: 'copilot', title: 'Test', status: SessionStatus.Idle, createdAt: Date.now(), modifiedAt: Date.now() }), + lifecycle: SessionLifecycle.Ready, + turns: [{ + id: 'turn-1', + userMessage: { text: 'do something' }, + state: TurnState.Complete, + responseParts: [], + responseText: '', + usage: undefined, + toolCalls: [{ toolCallId: 'tc-orphan', toolName: 'read_file', displayName: 'Read File', invocationMessage: 'Reading file', success: false, pastTenseMessage: 'Reading file' }], + }], + } as ISessionState); const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/orphan-tool' }); const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); @@ -1055,14 +1105,21 @@ suite('AgentHostChatContribution', () => { test('non-terminal tool_complete sets pastTenseMessage in history', async () => { const { sessionHandler, agentHostService } = createContribution(disposables); - const session = AgentSession.uri('copilot', 'generic-tool'); + const sessionUri = AgentSession.uri('copilot', 'generic-tool'); - agentHostService.setSessionMessages('generic-tool', [ - { session, type: 'message', messageId: 'u1', content: 'search', role: 'user' }, - { session, type: 'tool_start', toolCallId: 'tc-g', toolName: 'grep', displayName: 'Grep', invocationMessage: 'Searching...' }, - { session, type: 'tool_complete', toolCallId: 'tc-g', success: true, pastTenseMessage: 'Searched for pattern' }, - { session, type: 'message', messageId: 'a1', content: 'Found it.', role: 'assistant' }, - ]); + agentHostService.sessionStates.set(sessionUri.toString(), { + ...createSessionState({ resource: sessionUri, provider: 'copilot', title: 'Test', status: SessionStatus.Idle, createdAt: Date.now(), modifiedAt: Date.now() }), + lifecycle: SessionLifecycle.Ready, + turns: [{ + id: 'turn-1', + userMessage: { text: 'search' }, + state: TurnState.Complete, + responseParts: [], + usage: undefined, + responseText: '', + toolCalls: [{ toolCallId: 'tc-g', toolName: 'grep', displayName: 'Grep', invocationMessage: 'Searching...', success: true, pastTenseMessage: 'Searched for pattern' }], + }], + } as ISessionState); const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/generic-tool' }); const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); @@ -1076,10 +1133,15 @@ suite('AgentHostChatContribution', () => { } }); - test('empty events produce empty history', async () => { + test('empty session produces empty history', async () => { const { sessionHandler, agentHostService } = createContribution(disposables); - agentHostService.setSessionMessages('empty-sess', []); + const sessionUri = AgentSession.uri('copilot', 'empty-sess'); + agentHostService.sessionStates.set(sessionUri.toString(), { + ...createSessionState({ resource: sessionUri, provider: 'copilot', title: 'Test', status: SessionStatus.Idle, createdAt: Date.now(), modifiedAt: Date.now() }), + lifecycle: SessionLifecycle.Ready, + turns: [], + } as ISessionState); const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/empty-sess' }); const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); @@ -1089,25 +1151,28 @@ suite('AgentHostChatContribution', () => { }); }); - // ---- sendMessage error handling -------------------------------------- + // ---- Server error handling ---------------------------------------------- - suite('sendMessage error handling', () => { + suite('server error handling', () => { - test('sendMessage failure resolves the agent invoke without throwing', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + test('server-side error resolves the agent invoke without throwing', async () => { + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; + const { turnPromise, session, turnId } = await startTurn(sessionHandler, agentHostService, disposables); - agentHostService.sendMessage = async () => { - throw new Error('connection lost'); - }; + // Simulate a server-side error (e.g. sendMessage failure on the server) + agentHostService.fireAction({ + action: { + type: 'session/error', + session, + turnId, + error: { errorType: 'connection_error', message: 'connection lost' }, + } as ISessionAction, + serverSeq: 99, + origin: undefined, + }); - const result = await agent.impl.invoke( - makeRequest({ message: 'Hello' }), - () => { }, [], CancellationToken.None, - ); - - assert.ok(result); + await turnPromise; }); }); @@ -1136,13 +1201,11 @@ suite('AgentHostChatContribution', () => { suite('language model provider', () => { test('maps models with correct metadata', async () => { - const { instantiationService, agentHostService } = createTestServices(disposables); + const provider = disposables.add(new AgentHostLanguageModelProvider('agent-host-copilot', 'agent-host-copilot')); + provider.updateModels([ + { provider: 'copilot', id: 'gpt-4o', name: 'GPT-4o', maxContextWindow: 128000, supportsVision: true }, + ]); - agentHostService.models = [ - { provider: 'copilot', id: 'gpt-4o', name: 'GPT-4o', maxContextWindow: 128000, supportsVision: true, supportsReasoningEffort: false }, - ]; - - const provider = disposables.add(instantiationService.createInstance(AgentHostLanguageModelProvider, 'agent-host-copilot', 'agent-host-copilot', 'copilot')); const models = await provider.provideLanguageModelChatInfo({}, CancellationToken.None); assert.strictEqual(models.length, 1); @@ -1153,51 +1216,29 @@ suite('AgentHostChatContribution', () => { assert.strictEqual(models[0].metadata.targetChatSessionType, 'agent-host-copilot'); }); - test('filters out models from other providers', async () => { - const { instantiationService, agentHostService } = createTestServices(disposables); - - agentHostService.models = [ - { provider: 'copilot', id: 'gpt-4o', name: 'GPT-4o', maxContextWindow: 128000, supportsVision: false, supportsReasoningEffort: false }, - { provider: 'copilot', id: 'other-model', name: 'Other Model', maxContextWindow: 200000, supportsVision: false, supportsReasoningEffort: false }, - ]; - - // Create a provider that filters to a different vendor, simulating cross-provider filtering - const provider = disposables.add(instantiationService.createInstance(AgentHostLanguageModelProvider, 'agent-host-copilot', 'agent-host-copilot', 'not-copilot')); - const models = await provider.provideLanguageModelChatInfo({}, CancellationToken.None); - - assert.strictEqual(models.length, 0); - }); - test('filters out disabled models', async () => { - const { instantiationService, agentHostService } = createTestServices(disposables); + const provider = disposables.add(new AgentHostLanguageModelProvider('agent-host-copilot', 'agent-host-copilot')); + provider.updateModels([ + { provider: 'copilot', id: 'gpt-4o', name: 'GPT-4o', maxContextWindow: 128000, supportsVision: false, policyState: 'enabled' }, + { provider: 'copilot', id: 'gpt-3.5', name: 'GPT-3.5', maxContextWindow: 16000, supportsVision: false, policyState: 'disabled' }, + ]); - agentHostService.models = [ - { provider: 'copilot', id: 'gpt-4o', name: 'GPT-4o', maxContextWindow: 128000, supportsVision: false, supportsReasoningEffort: false, policyState: 'enabled' }, - { provider: 'copilot', id: 'gpt-3.5', name: 'GPT-3.5', maxContextWindow: 16000, supportsVision: false, supportsReasoningEffort: false, policyState: 'disabled' }, - ]; - - const provider = disposables.add(instantiationService.createInstance(AgentHostLanguageModelProvider, 'agent-host-copilot', 'agent-host-copilot', 'copilot')); const models = await provider.provideLanguageModelChatInfo({}, CancellationToken.None); assert.strictEqual(models.length, 1); assert.strictEqual(models[0].metadata.name, 'GPT-4o'); }); - test('returns empty array on error', async () => { - const { instantiationService, agentHostService } = createTestServices(disposables); + test('returns empty when no models set', async () => { + const provider = disposables.add(new AgentHostLanguageModelProvider('agent-host-copilot', 'agent-host-copilot')); - agentHostService.listModels = async () => { throw new Error('not connected'); }; - - const provider = disposables.add(instantiationService.createInstance(AgentHostLanguageModelProvider, 'agent-host-copilot', 'agent-host-copilot', 'copilot')); const models = await provider.provideLanguageModelChatInfo({}, CancellationToken.None); assert.strictEqual(models.length, 0); }); test('sendChatRequest throws', async () => { - const { instantiationService } = createTestServices(disposables); - - const provider = disposables.add(instantiationService.createInstance(AgentHostLanguageModelProvider, 'agent-host-copilot', 'agent-host-copilot', 'copilot')); + const provider = disposables.add(new AgentHostLanguageModelProvider('agent-host-copilot', 'agent-host-copilot')); await assert.rejects(() => provider.sendChatRequest(), /do not support direct chat requests/); }); @@ -1208,192 +1249,144 @@ suite('AgentHostChatContribution', () => { suite('attachment context', () => { test('file variable with file:// URI becomes file attachment', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + message: 'check this file', + variables: { + variables: [ + upcastPartial({ kind: 'file', id: 'v-file', name: 'test.ts', value: URI.file('/workspace/test.ts') }), + ], + }, + }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; - agentHostService.sendMessage = async (session: URI, prompt: string, attachments?: IAgentAttachment[]) => { - agentHostService.sendMessageCalls.push({ session, prompt, attachments }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; - - await agent.impl.invoke( - makeRequest({ - message: 'check this file', - variables: { - variables: [ - upcastPartial({ kind: 'file', id: 'v-file', name: 'test.ts', value: URI.file('/workspace/test.ts') }), - ], - }, - }), - () => { }, [], CancellationToken.None, - ); - - assert.strictEqual(agentHostService.sendMessageCalls.length, 1); - const call = agentHostService.sendMessageCalls[0]; - assert.deepStrictEqual(call.attachments, [ + assert.strictEqual(agentHostService.dispatchedActions.length, 1); + const turnAction = agentHostService.dispatchedActions[0].action as ITurnStartedAction; + assert.deepStrictEqual(turnAction.userMessage.attachments, [ { type: 'file', path: URI.file('/workspace/test.ts').fsPath, displayName: 'test.ts' }, ]); }); test('directory variable with file:// URI becomes directory attachment', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + message: 'check this dir', + variables: { + variables: [ + upcastPartial({ kind: 'directory', id: 'v-dir', name: 'src', value: URI.file('/workspace/src') }), + ], + }, + }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; - agentHostService.sendMessage = async (session: URI, prompt: string, attachments?: IAgentAttachment[]) => { - agentHostService.sendMessageCalls.push({ session, prompt, attachments }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; - - await agent.impl.invoke( - makeRequest({ - message: 'check this dir', - variables: { - variables: [ - upcastPartial({ kind: 'directory', id: 'v-dir', name: 'src', value: URI.file('/workspace/src') }), - ], - }, - }), - () => { }, [], CancellationToken.None, - ); - - assert.strictEqual(agentHostService.sendMessageCalls.length, 1); - assert.deepStrictEqual(agentHostService.sendMessageCalls[0].attachments, [ + assert.strictEqual(agentHostService.dispatchedActions.length, 1); + const turnAction = agentHostService.dispatchedActions[0].action as ITurnStartedAction; + assert.deepStrictEqual(turnAction.userMessage.attachments, [ { type: 'directory', path: URI.file('/workspace/src').fsPath, displayName: 'src' }, ]); }); test('implicit selection variable becomes selection attachment', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + message: 'explain this', + variables: { + variables: [ + upcastPartial({ kind: 'implicit', id: 'v-implicit', name: 'selection', isFile: true as const, isSelection: true, uri: URI.file('/workspace/foo.ts'), enabled: true, value: undefined }), + ], + }, + }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; - agentHostService.sendMessage = async (session: URI, prompt: string, attachments?: IAgentAttachment[]) => { - agentHostService.sendMessageCalls.push({ session, prompt, attachments }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; - - await agent.impl.invoke( - makeRequest({ - message: 'explain this', - variables: { - variables: [ - upcastPartial({ kind: 'implicit', id: 'v-implicit', name: 'selection', isFile: true as const, isSelection: true, uri: URI.file('/workspace/foo.ts'), enabled: true, value: undefined }), - ], - }, - }), - () => { }, [], CancellationToken.None, - ); - - assert.strictEqual(agentHostService.sendMessageCalls.length, 1); - assert.deepStrictEqual(agentHostService.sendMessageCalls[0].attachments, [ + assert.strictEqual(agentHostService.dispatchedActions.length, 1); + const turnAction = agentHostService.dispatchedActions[0].action as ITurnStartedAction; + assert.deepStrictEqual(turnAction.userMessage.attachments, [ { type: 'selection', path: URI.file('/workspace/foo.ts').fsPath, displayName: 'selection' }, ]); }); test('non-file URIs are skipped', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + message: 'check this', + variables: { + variables: [ + upcastPartial({ kind: 'file', id: 'v-file', name: 'untitled', value: URI.from({ scheme: 'untitled', path: '/foo' }) }), + ], + }, + }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; - agentHostService.sendMessage = async (session: URI, prompt: string, attachments?: IAgentAttachment[]) => { - agentHostService.sendMessageCalls.push({ session, prompt, attachments }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; - - await agent.impl.invoke( - makeRequest({ - message: 'check this', - variables: { - variables: [ - upcastPartial({ kind: 'file', id: 'v-file', name: 'untitled', value: URI.from({ scheme: 'untitled', path: '/foo' }) }), - ], - }, - }), - () => { }, [], CancellationToken.None, - ); - - assert.strictEqual(agentHostService.sendMessageCalls.length, 1); + assert.strictEqual(agentHostService.dispatchedActions.length, 1); + const turnAction = agentHostService.dispatchedActions[0].action as ITurnStartedAction; // No attachments because it's not a file:// URI - assert.strictEqual(agentHostService.sendMessageCalls[0].attachments, undefined); + assert.strictEqual(turnAction.userMessage.attachments, undefined); }); test('tool variables are skipped', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + message: 'use tools', + variables: { + variables: [ + upcastPartial({ kind: 'tool', id: 'v-tool', name: 'myTool', value: { id: 'tool-1' } }), + ], + }, + }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; - agentHostService.sendMessage = async (session: URI, prompt: string, attachments?: IAgentAttachment[]) => { - agentHostService.sendMessageCalls.push({ session, prompt, attachments }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; - - await agent.impl.invoke( - makeRequest({ - message: 'use tools', - variables: { - variables: [ - upcastPartial({ kind: 'tool', id: 'v-tool', name: 'myTool', value: { id: 'tool-1' } }), - ], - }, - }), - () => { }, [], CancellationToken.None, - ); - - assert.strictEqual(agentHostService.sendMessageCalls.length, 1); - assert.strictEqual(agentHostService.sendMessageCalls[0].attachments, undefined); + assert.strictEqual(agentHostService.dispatchedActions.length, 1); + const turnAction = agentHostService.dispatchedActions[0].action as ITurnStartedAction; + assert.strictEqual(turnAction.userMessage.attachments, undefined); }); test('mixed variables extracts only supported types', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + message: 'mixed', + variables: { + variables: [ + upcastPartial({ kind: 'file', id: 'v-file', name: 'a.ts', value: URI.file('/workspace/a.ts') }), + upcastPartial({ kind: 'tool', id: 'v-tool', name: 'myTool', value: { id: 'tool-1' } }), + upcastPartial({ kind: 'directory', id: 'v-dir', name: 'lib', value: URI.file('/workspace/lib') }), + upcastPartial({ kind: 'file', id: 'v-file', name: 'remote.ts', value: URI.from({ scheme: 'vscode-remote', path: '/remote/file.ts' }) }), + ], + }, + }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; - agentHostService.sendMessage = async (session: URI, prompt: string, attachments?: IAgentAttachment[]) => { - agentHostService.sendMessageCalls.push({ session, prompt, attachments }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; - - await agent.impl.invoke( - makeRequest({ - message: 'mixed', - variables: { - variables: [ - upcastPartial({ kind: 'file', id: 'v-file', name: 'a.ts', value: URI.file('/workspace/a.ts') }), - upcastPartial({ kind: 'tool', id: 'v-tool', name: 'myTool', value: { id: 'tool-1' } }), - upcastPartial({ kind: 'directory', id: 'v-dir', name: 'lib', value: URI.file('/workspace/lib') }), - upcastPartial({ kind: 'file', id: 'v-file', name: 'remote.ts', value: URI.from({ scheme: 'vscode-remote', path: '/remote/file.ts' }) }), - ], - }, - }), - () => { }, [], CancellationToken.None, - ); - - assert.strictEqual(agentHostService.sendMessageCalls.length, 1); - assert.deepStrictEqual(agentHostService.sendMessageCalls[0].attachments, [ + assert.strictEqual(agentHostService.dispatchedActions.length, 1); + const turnAction = agentHostService.dispatchedActions[0].action as ITurnStartedAction; + assert.deepStrictEqual(turnAction.userMessage.attachments, [ { type: 'file', path: URI.file('/workspace/a.ts').fsPath, displayName: 'a.ts' }, { type: 'directory', path: URI.file('/workspace/lib').fsPath, displayName: 'lib' }, ]); }); test('no variables results in no attachments argument', async () => { - const { chatAgentService, agentHostService } = createContribution(disposables); + const { sessionHandler, agentHostService } = createContribution(disposables); - const agent = chatAgentService.registeredAgents.get('agent-host-copilot')!; + const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables, { + message: 'Hello', + }); + fire({ type: 'session/turnComplete', session, turnId } as ISessionAction); + await turnPromise; - agentHostService.sendMessage = async (session: URI, prompt: string, attachments?: IAgentAttachment[]) => { - agentHostService.sendMessageCalls.push({ session, prompt, attachments }); - agentHostService.fireProgress({ session, type: 'idle' }); - }; - - await agent.impl.invoke( - makeRequest({ message: 'Hello' }), - () => { }, [], CancellationToken.None, - ); - - assert.strictEqual(agentHostService.sendMessageCalls.length, 1); - assert.strictEqual(agentHostService.sendMessageCalls[0].attachments, undefined); + assert.strictEqual(agentHostService.dispatchedActions.length, 1); + const turnAction = agentHostService.dispatchedActions[0].action as ITurnStartedAction; + assert.strictEqual(turnAction.userMessage.attachments, undefined); }); }); 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 new file mode 100644 index 00000000000..c7e54d2622a --- /dev/null +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/stateToProgressAdapter.test.ts @@ -0,0 +1,298 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; +import { ToolCallStatus, TurnState, type ICompletedToolCall, type IPermissionRequest, type IToolCallState, type ITurn } 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, permissionToConfirmation, finalizeToolInvocation } from '../../../browser/agentSessions/agentHost/stateToProgressAdapter.js'; + +// ---- Helper factories ------------------------------------------------------- + +function createToolCallState(overrides?: Partial): IToolCallState { + return { + toolCallId: 'tc-1', + toolName: 'test_tool', + displayName: 'Test Tool', + invocationMessage: 'Running test tool...', + status: ToolCallStatus.Running, + ...overrides, + }; +} + +function createCompletedToolCall(overrides?: Partial): ICompletedToolCall { + return { + toolCallId: 'tc-1', + toolName: 'test_tool', + displayName: 'Test Tool', + invocationMessage: 'Running test tool...', + success: true, + pastTenseMessage: 'Ran test tool', + ...overrides, + }; +} + +function createTurn(overrides?: Partial): ITurn { + return { + id: 'turn-1', + userMessage: { text: 'Hello' }, + responseText: '', + responseParts: [], + toolCalls: [], + usage: undefined, + state: TurnState.Complete, + ...overrides, + }; +} + +function createPermission(overrides?: Partial): IPermissionRequest { + return { + requestId: 'perm-1', + permissionKind: 'shell', + ...overrides, + }; +} + +// ---- Tests ------------------------------------------------------------------ + +suite('stateToProgressAdapter', () => { + + ensureNoDisposablesAreLeakedInTestSuite(); + + suite('turnsToHistory', () => { + + test('empty turns produces empty history', () => { + const result = turnsToHistory([], 'p'); + assert.deepStrictEqual(result, []); + }); + + test('single turn produces request + response pair', () => { + const turn = createTurn({ + userMessage: { text: 'Do something' }, + toolCalls: [createCompletedToolCall()], + }); + + const history = turnsToHistory([turn], 'participant-1'); + assert.strictEqual(history.length, 2); + + // Request + assert.strictEqual(history[0].type, 'request'); + assert.strictEqual(history[0].prompt, 'Do something'); + assert.strictEqual(history[0].participant, 'participant-1'); + + // Response + assert.strictEqual(history[1].type, 'response'); + assert.strictEqual(history[1].participant, 'participant-1'); + assert.strictEqual(history[1].parts.length, 1); + + const serialized = history[1].parts[0] as IChatToolInvocationSerialized; + assert.strictEqual(serialized.kind, 'toolInvocationSerialized'); + assert.strictEqual(serialized.toolCallId, 'tc-1'); + assert.strictEqual(serialized.toolId, 'test_tool'); + assert.strictEqual(serialized.isComplete, true); + }); + + test('terminal tool call in history has correct terminal data', () => { + const turn = createTurn({ + toolCalls: [createCompletedToolCall({ + toolKind: 'terminal', + toolInput: 'echo hello', + language: 'shellscript', + toolOutput: 'hello', + success: true, + })], + }); + + const history = turnsToHistory([turn], 'p'); + const response = history[1]; + assert.strictEqual(response.type, 'response'); + if (response.type !== 'response') { return; } + const serialized = response.parts[0] as IChatToolInvocationSerialized; + + assert.ok(serialized.toolSpecificData); + assert.strictEqual(serialized.toolSpecificData.kind, 'terminal'); + const termData = serialized.toolSpecificData as { kind: 'terminal'; commandLine: { original: string }; terminalCommandOutput: { text: string }; terminalCommandState: { exitCode: number } }; + assert.strictEqual(termData.commandLine.original, 'echo hello'); + assert.strictEqual(termData.terminalCommandOutput.text, 'hello'); + assert.strictEqual(termData.terminalCommandState.exitCode, 0); + }); + + test('turn with responseText produces markdown content in history', () => { + const turn = createTurn({ + responseText: 'Hello world', + }); + + const history = turnsToHistory([turn], 'p'); + assert.strictEqual(history.length, 2); + + const response = history[1]; + assert.strictEqual(response.type, 'response'); + if (response.type !== 'response') { return; } + assert.strictEqual(response.parts.length, 1); + assert.strictEqual(response.parts[0].kind, 'markdownContent'); + assert.strictEqual((response.parts[0] as IChatMarkdownContent).content.value, 'Hello world'); + }); + + test('error turn produces error message in history', () => { + const turn = createTurn({ + state: TurnState.Error, + error: { errorType: 'test', message: 'boom' }, + }); + + const history = turnsToHistory([turn], 'p'); + const response = history[1]; + assert.strictEqual(response.type, 'response'); + if (response.type !== 'response') { return; } + const errorPart = response.parts.find(p => p.kind === 'markdownContent' && (p as IChatMarkdownContent).content.value.includes('boom')); + assert.ok(errorPart, 'Should have a markdownContent part containing the error message'); + }); + + test('failed tool in history has exitCode 1', () => { + const turn = createTurn({ + toolCalls: [createCompletedToolCall({ + toolKind: 'terminal', + toolInput: 'bad-command', + toolOutput: 'error', + success: false, + })], + }); + + const history = turnsToHistory([turn], 'p'); + const response = history[1]; + assert.strictEqual(response.type, 'response'); + if (response.type !== 'response') { return; } + const serialized = response.parts[0] as IChatToolInvocationSerialized; + + assert.ok(serialized.toolSpecificData); + assert.strictEqual(serialized.toolSpecificData.kind, 'terminal'); + const termData = serialized.toolSpecificData as { kind: 'terminal'; terminalCommandState: { exitCode: number } }; + assert.strictEqual(termData.terminalCommandState.exitCode, 1); + }); + }); + + suite('toolCallStateToInvocation', () => { + + test('creates ChatToolInvocation for running tool', () => { + const tc = createToolCallState({ + toolCallId: 'tc-42', + toolName: 'my_tool', + displayName: 'My Tool', + invocationMessage: 'Doing stuff', + status: ToolCallStatus.Running, + }); + + const invocation = toolCallStateToInvocation(tc); + assert.strictEqual(invocation.toolCallId, 'tc-42'); + assert.strictEqual(invocation.toolId, 'my_tool'); + assert.strictEqual(invocation.source, ToolDataSource.Internal); + }); + + test('sets terminal toolSpecificData', () => { + const tc = createToolCallState({ + toolKind: 'terminal', + toolInput: 'ls -la', + }); + + const invocation = toolCallStateToInvocation(tc); + assert.ok(invocation.toolSpecificData); + assert.strictEqual(invocation.toolSpecificData.kind, 'terminal'); + const termData = invocation.toolSpecificData as { kind: 'terminal'; commandLine: { original: string } }; + assert.strictEqual(termData.commandLine.original, 'ls -la'); + }); + + test('parses toolArguments as parameters', () => { + const tc = createToolCallState({ + toolArguments: '{"path":"test.ts"}', + }); + + const invocation = toolCallStateToInvocation(tc); + assert.deepStrictEqual(invocation.parameters, { path: 'test.ts' }); + }); + }); + + suite('permissionToConfirmation', () => { + + test('shell permission has terminal data', () => { + const perm = createPermission({ + permissionKind: 'shell', + fullCommandText: 'rm -rf /', + intention: 'Delete everything', + }); + + const invocation = permissionToConfirmation(perm); + assert.ok(invocation.toolSpecificData); + assert.strictEqual(invocation.toolSpecificData.kind, 'terminal'); + const termData = invocation.toolSpecificData as { kind: 'terminal'; commandLine: { original: string } }; + assert.strictEqual(termData.commandLine.original, 'rm -rf /'); + }); + + test('mcp permission uses server + tool name as title', () => { + const perm = createPermission({ + permissionKind: 'mcp', + serverName: 'My Server', + toolName: 'my_tool', + }); + + const invocation = permissionToConfirmation(perm); + const message = typeof invocation.invocationMessage === 'string' ? invocation.invocationMessage : invocation.invocationMessage.value; + assert.ok(message.includes('My Server: my_tool')); + }); + + test('write permission has input data', () => { + const perm = createPermission({ + permissionKind: 'write', + path: '/test.ts', + rawRequest: '{"path":"/test.ts","content":"hello"}', + }); + + const invocation = permissionToConfirmation(perm); + assert.ok(invocation.toolSpecificData); + assert.strictEqual(invocation.toolSpecificData.kind, 'input'); + }); + }); + + suite('finalizeToolInvocation', () => { + + test('finalizes terminal tool with output and exit code', () => { + const tc = createToolCallState({ + toolKind: 'terminal', + toolInput: 'echo hi', + status: ToolCallStatus.Running, + }); + const invocation = toolCallStateToInvocation(tc); + + const completedTc = createToolCallState({ + toolKind: 'terminal', + toolInput: 'echo hi', + status: ToolCallStatus.Completed, + toolOutput: 'output text', + }); + + finalizeToolInvocation(invocation, completedTc); + + assert.ok(invocation.toolSpecificData); + assert.strictEqual(invocation.toolSpecificData.kind, 'terminal'); + const termData = invocation.toolSpecificData as { kind: 'terminal'; terminalCommandOutput: { text: string }; terminalCommandState: { exitCode: number } }; + assert.strictEqual(termData.terminalCommandOutput.text, 'output text'); + assert.strictEqual(termData.terminalCommandState.exitCode, 0); + }); + + test('finalizes failed tool with error message', () => { + const tc = createToolCallState({ + status: ToolCallStatus.Running, + }); + const invocation = toolCallStateToInvocation(tc); + + const failedTc = createToolCallState({ + status: ToolCallStatus.Failed, + error: { message: 'timeout' }, + }); + + // Should not throw + finalizeToolInvocation(invocation, failedTc); + }); + }); +});