From 94c7bf8213beb59e3181fdb61992a032fc65e9a2 Mon Sep 17 00:00:00 2001 From: Rob Lourens Date: Sun, 29 Mar 2026 19:07:20 -0700 Subject: [PATCH] Unify agentHost server-side dispatch: remove IProtocolSideEffectHandler (#306158) * Unify agentHost server-side dispatch: remove IProtocolSideEffectHandler Eliminate the IProtocolSideEffectHandler interface and make ProtocolServerHandler talk to IAgentService directly. This removes the duplicate adapter layer between the WebSocket protocol server and the real service implementation. Changes: - ProtocolServerHandler now takes IAgentService + SessionStateManager + IProtocolServerConfig instead of IProtocolSideEffectHandler - Deleted ~40-line inline adapter in agentHostMain.ts - agentHostServerMain.ts now uses AgentService instead of manually wiring SessionStateManager + AgentSideEffects - Removed implements IProtocolSideEffectHandler from AgentSideEffects - Removed dead methods from AgentSideEffects that were only needed by the deleted interface (handleCreateSession, handleDisposeSession, handleListSessions, handleGetResourceMetadata, handleAuthenticate, getDefaultDirectory) - Type conversions (URI<->string, IAgentSessionMetadata<->ISessionSummary) now happen at the protocol boundary in ProtocolServerHandler - Fixed dispatchAction double-dispatch: WS path previously dispatched to stateManager AND called handleAction (which dispatched again) - Extension methods (getResourceMetadata, authenticate, etc.) now call IAgentService directly instead of untyped fallbacks (Written by Copilot) * comments Co-authored-by: Copilot * Simplify further Co-authored-by: Copilot --------- Co-authored-by: Copilot --- .../platform/agentHost/node/agentHostMain.ts | 64 +--- .../agentHost/node/agentHostServerMain.ts | 46 +-- .../platform/agentHost/node/agentService.ts | 241 +++++++++++- .../agentHost/node/agentSideEffects.ts | 353 +----------------- .../agentHost/node/protocolServerHandler.ts | 130 ++++--- src/vs/platform/agentHost/protocol.md | 2 +- .../agentHost/test/node/agentService.test.ts | 120 +++++- .../test/node/agentSideEffects.test.ts | 311 +-------------- .../test/node/protocolServerHandler.test.ts | 89 +++-- 9 files changed, 512 insertions(+), 844 deletions(-) diff --git a/src/vs/platform/agentHost/node/agentHostMain.ts b/src/vs/platform/agentHost/node/agentHostMain.ts index dfe5ad6da74..01d1b3079ea 100644 --- a/src/vs/platform/agentHost/node/agentHostMain.ts +++ b/src/vs/platform/agentHost/node/agentHostMain.ts @@ -11,11 +11,10 @@ import { Emitter } from '../../../base/common/event.js'; import { DisposableStore } from '../../../base/common/lifecycle.js'; import { URI } from '../../../base/common/uri.js'; import * as os from 'os'; -import { AgentHostIpcChannels, AgentSession } from '../common/agentService.js'; -import { SessionStatus } from '../common/state/sessionState.js'; +import { AgentHostIpcChannels } from '../common/agentService.js'; import { AgentService } from './agentService.js'; import { CopilotAgent } from './copilot/copilotAgent.js'; -import { ProtocolServerHandler, type IProtocolSideEffectHandler } from './protocolServerHandler.js'; +import { ProtocolServerHandler } from './protocolServerHandler.js'; import { WebSocketProtocolServer } from './webSocketTransport.js'; import { NativeEnvironmentService } from '../../environment/node/environmentService.js'; import { parseArgs, OPTIONS } from '../../environment/node/argv.js'; @@ -144,58 +143,13 @@ async function startWebSocketServer(agentService: AgentService, logService: ILog logService, )); - // Create a side-effect handler that delegates to AgentService - const sideEffects: IProtocolSideEffectHandler = { - handleAction(action) { - agentService.dispatchAction(action, 'ws-server', 0); - }, - async handleCreateSession(command) { - await agentService.createSession({ - provider: command.provider, - model: command.model, - workingDirectory: command.workingDirectory ? URI.parse(command.workingDirectory) : undefined, - session: URI.parse(command.session), - }); - }, - handleDisposeSession(session) { - agentService.disposeSession(URI.parse(session)); - }, - async handleListSessions() { - const sessions = await agentService.listSessions(); - return sessions.map(s => ({ - resource: s.session.toString(), - provider: AgentSession.provider(s.session) ?? 'copilot', - title: s.summary ?? 'Session', - status: SessionStatus.Idle, - createdAt: s.startTime, - modifiedAt: s.modifiedTime, - workingDirectory: s.workingDirectory?.toString(), - })); - }, - handleGetResourceMetadata() { - return agentService.getResourceMetadataSync(); - }, - async handleAuthenticate(params) { - return agentService.authenticate(params); - }, - handleBrowseDirectory(uri) { - return agentService.browseDirectory(URI.parse(uri)); - }, - handleWriteFile(params) { - return agentService.writeFile(params); - }, - async handleRestoreSession(session) { - return agentService.restoreSession(URI.parse(session)); - }, - handleFetchContent(uri) { - return agentService.fetchContent(URI.parse(uri)); - }, - getDefaultDirectory() { - return URI.file(os.homedir()).toString(); - }, - }; - - const protocolHandler = disposables.add(new ProtocolServerHandler(agentService.stateManager, wsServer, sideEffects, logService)); + const protocolHandler = disposables.add(new ProtocolServerHandler( + agentService, + agentService.stateManager, + wsServer, + { defaultDirectory: URI.file(os.homedir()).toString() }, + logService, + )); disposables.add(protocolHandler.onDidChangeConnectionCount(onConnectionCountChanged)); const listenTarget = socketPath ?? `${host}:${port}`; diff --git a/src/vs/platform/agentHost/node/agentHostServerMain.ts b/src/vs/platform/agentHost/node/agentHostServerMain.ts index 8622b1c244b..a0b3487b79c 100644 --- a/src/vs/platform/agentHost/node/agentHostServerMain.ts +++ b/src/vs/platform/agentHost/node/agentHostServerMain.ts @@ -14,8 +14,8 @@ import { fileURLToPath } from 'url'; globalThis._VSCODE_FILE_ROOT = fileURLToPath(new URL('../../../..', import.meta.url)); import * as fs from 'fs'; +import * as os from 'os'; import { DisposableStore } from '../../../base/common/lifecycle.js'; -import { observableValue } from '../../../base/common/observable.js'; import { URI } from '../../../base/common/uri.js'; import { generateUuid } from '../../../base/common/uuid.js'; import { localize } from '../../../nls.js'; @@ -30,9 +30,7 @@ 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 { AgentSideEffects } from './agentSideEffects.js'; -import { SessionStateManager } from './sessionStateManager.js'; +import { AgentService } from './agentService.js'; import { WebSocketProtocolServer } from './webSocketTransport.js'; import { ProtocolServerHandler } from './protocolServerHandler.js'; import { FileService } from '../../files/common/fileService.js'; @@ -140,15 +138,6 @@ async function main(): Promise { 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(); - - // Observable agents list for root state - const registeredAgents = observableValue('agents', []); - // File service const fileService = disposables.add(new FileService(logService)); disposables.add(fileService.registerProvider(Schemas.file, disposables.add(new DiskFileSystemProvider(logService)))); @@ -156,22 +145,9 @@ async function main(): Promise { // Session data service const sessionDataService = new SessionDataService(URI.file(environmentService.userDataPath), fileService, logService); - // Shared side-effect handler - const sideEffects = disposables.add(new AgentSideEffects(stateManager, { - getAgent(session) { - const provider = AgentSession.provider(session); - return provider ? agents.get(provider) : agents.values().next().value; - }, - agents: registeredAgents, - sessionDataService, - }, logService, fileService)); - - function registerAgent(agent: IAgent): void { - agents.set(agent.id, agent); - disposables.add(sideEffects.registerProgressListener(agent)); - registeredAgents.set([...agents.values()], undefined); - logService.info(`[AgentHostServer] Registered agent: ${agent.id}`); - } + // Create the agent service (owns SessionStateManager + AgentSideEffects internally) + const agentService = new AgentService(logService, fileService, sessionDataService); + disposables.add(agentService); // Register agents if (!options.quiet) { @@ -184,7 +160,7 @@ async function main(): Promise { diServices.set(ISessionDataService, sessionDataService); const instantiationService = new InstantiationService(diServices); const copilotAgent = disposables.add(instantiationService.createInstance(CopilotAgent)); - registerAgent(copilotAgent); + agentService.registerProvider(copilotAgent); log('CopilotAgent registered'); } @@ -192,7 +168,7 @@ async function main(): Promise { // Dynamic import to avoid bundling test code in production import('../test/node/mockAgent.js').then(({ ScriptedMockAgent }) => { const mockAgent = disposables.add(new ScriptedMockAgent()); - registerAgent(mockAgent); + agentService.registerProvider(mockAgent); }).catch(err => { logService.error('[AgentHostServer] Failed to load mock agent', err); }); @@ -207,7 +183,13 @@ async function main(): Promise { }, logService)); // Wire up protocol handler - disposables.add(new ProtocolServerHandler(stateManager, wsServer, sideEffects, logService)); + disposables.add(new ProtocolServerHandler( + agentService, + agentService.stateManager, + wsServer, + { defaultDirectory: URI.file(os.homedir()).toString() }, + logService, + )); // Report ready function reportReady(addr: string): void { diff --git a/src/vs/platform/agentHost/node/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts index 13cc441565b..4850837a841 100644 --- a/src/vs/platform/agentHost/node/agentService.ts +++ b/src/vs/platform/agentHost/node/agentService.ts @@ -3,18 +3,21 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ +import { decodeBase64, VSBuffer } from '../../../base/common/buffer.js'; import { Emitter } from '../../../base/common/event.js'; import { Disposable, DisposableStore } from '../../../base/common/lifecycle.js'; import { observableValue } from '../../../base/common/observable.js'; import { URI } from '../../../base/common/uri.js'; -import { IFileService } from '../../files/common/files.js'; +import { generateUuid } from '../../../base/common/uuid.js'; +import { FileSystemProviderErrorCode, IFileService, toFileSystemProviderErrorCode } from '../../files/common/files.js'; import { ILogService } from '../../log/common/log.js'; -import { AgentProvider, AgentSession, IAgent, IAgentCreateSessionConfig, IAgentDescriptor, IAgentService, IAgentSessionMetadata, IAuthenticateParams, IAuthenticateResult, IResourceMetadata } from '../common/agentService.js'; +import { AgentProvider, AgentSession, IAgent, IAgentCreateSessionConfig, IAgentDescriptor, IAgentMessageEvent, IAgentService, IAgentSessionMetadata, IAgentToolCompleteEvent, IAgentToolStartEvent, IAuthenticateParams, IAuthenticateResult, IResourceMetadata } from '../common/agentService.js'; import { ISessionDataService } from '../common/sessionDataService.js'; import { ActionType, IActionEnvelope, INotification, ISessionAction } from '../common/state/sessionActions.js'; -import type { IBrowseDirectoryResult, IFetchContentResult, IStateSnapshot, IWriteFileParams, IWriteFileResult } from '../common/state/sessionProtocol.js'; -import { SessionStatus, type ISessionSummary } from '../common/state/sessionState.js'; +import { AhpErrorCodes, AHP_SESSION_NOT_FOUND, ContentEncoding, JSON_RPC_INTERNAL_ERROR, ProtocolError, type IBrowseDirectoryResult, type IDirectoryEntry, type IFetchContentResult, type IStateSnapshot, type IWriteFileParams, type IWriteFileResult } from '../common/state/sessionProtocol.js'; +import { ResponsePartKind, SessionStatus, ToolCallConfirmationReason, ToolCallStatus, TurnState, type IResponsePart, type ISessionSummary, type IToolCallCompletedState, type ITurn } from '../common/state/sessionState.js'; import { AgentSideEffects } from './agentSideEffects.js'; +import { ISessionDbUriFields, parseSessionDbUri } from './copilot/fileEditTracker.js'; import { SessionStateManager } from './sessionStateManager.js'; /** @@ -66,7 +69,7 @@ export class AgentService extends Disposable implements IAgentService { getAgent: session => this._findProviderForSession(session), sessionDataService: this._sessionDataService, agents: this._agents, - }, this._logService, this._fileService)); + }, this._logService)); } // ---- provider registration ---------------------------------------------- @@ -206,19 +209,125 @@ export class AgentService extends Disposable implements IAgentService { } async browseDirectory(uri: URI): Promise { - return this._sideEffects.handleBrowseDirectory(uri.toString()); + let stat; + try { + stat = await this._fileService.resolve(uri); + } catch { + throw new ProtocolError(AhpErrorCodes.NotFound, `Directory not found: ${uri.toString()}`); + } + + if (!stat.isDirectory) { + throw new ProtocolError(AhpErrorCodes.NotFound, `Not a directory: ${uri.toString()}`); + } + + const entries: IDirectoryEntry[] = (stat.children ?? []).map(child => ({ + name: child.name, + type: child.isDirectory ? 'directory' : 'file', + })); + return { entries }; } async restoreSession(session: URI): Promise { - return this._sideEffects.handleRestoreSession(session.toString()); + const sessionStr = session.toString(); + + // Already in state manager - nothing to do. + if (this._stateManager.getSessionState(sessionStr)) { + return; + } + + const agent = this._findProviderForSession(session); + if (!agent) { + throw new ProtocolError(AHP_SESSION_NOT_FOUND, `No agent for session: ${sessionStr}`); + } + + // Verify the session actually exists on the backend to avoid + // creating phantom sessions for made-up URIs. + let allSessions; + try { + allSessions = await agent.listSessions(); + } catch (err) { + if (err instanceof ProtocolError) { + throw err; + } + const message = err instanceof Error ? err.message : String(err); + throw new ProtocolError(JSON_RPC_INTERNAL_ERROR, `Failed to list sessions for ${sessionStr}: ${message}`); + } + const meta = allSessions.find(s => s.session.toString() === sessionStr); + if (!meta) { + throw new ProtocolError(AHP_SESSION_NOT_FOUND, `Session not found on backend: ${sessionStr}`); + } + + let messages; + try { + messages = await agent.getSessionMessages(session); + } catch (err) { + if (err instanceof ProtocolError) { + throw err; + } + const message = err instanceof Error ? err.message : String(err); + throw new ProtocolError(JSON_RPC_INTERNAL_ERROR, `Failed to restore session ${sessionStr}: ${message}`); + } + const turns = this._buildTurnsFromMessages(messages); + + const summary: ISessionSummary = { + resource: sessionStr, + provider: agent.id, + title: meta.summary ?? 'Session', + status: SessionStatus.Idle, + createdAt: meta.startTime, + modifiedAt: meta.modifiedTime, + workingDirectory: meta.workingDirectory?.toString(), + }; + + this._stateManager.restoreSession(summary, turns); + this._logService.info(`[AgentService] Restored session ${sessionStr} with ${turns.length} turns`); } async fetchContent(uri: URI): Promise { - return this._sideEffects.handleFetchContent(uri.toString()); + // Handle session-db: URIs that reference file-edit content stored + // in a per-session SQLite database. + const dbFields = parseSessionDbUri(uri.toString()); + if (dbFields) { + return this._fetchSessionDbContent(dbFields); + } + + try { + const content = await this._fileService.readFile(uri); + return { + data: content.value.toString(), + encoding: ContentEncoding.Utf8, + contentType: 'text/plain', + }; + } catch (_e) { + throw new ProtocolError(AhpErrorCodes.NotFound, `Content not found: ${uri.toString()}`); + } } async writeFile(params: IWriteFileParams): Promise { - return this._sideEffects.handleWriteFile(params); + const fileUri = typeof params.uri === 'string' ? URI.parse(params.uri) : URI.revive(params.uri); + let content: VSBuffer; + if (params.encoding === ContentEncoding.Base64) { + content = decodeBase64(params.data); + } else { + content = VSBuffer.fromString(params.data); + } + try { + if (params.createOnly) { + await this._fileService.createFile(fileUri, content, { overwrite: false }); + } else { + await this._fileService.writeFile(fileUri, content); + } + return {}; + } catch (e) { + const code = toFileSystemProviderErrorCode(e as Error); + if (code === FileSystemProviderErrorCode.FileExists) { + throw new ProtocolError(AhpErrorCodes.AlreadyExists, `File already exists: ${fileUri.toString()}`); + } + if (code === FileSystemProviderErrorCode.NoPermissions) { + throw new ProtocolError(AhpErrorCodes.PermissionDenied, `Permission denied: ${fileUri.toString()}`); + } + throw new ProtocolError(AhpErrorCodes.NotFound, `Failed to write file: ${fileUri.toString()}`); + } } async shutdown(): Promise { @@ -233,6 +342,120 @@ export class AgentService extends Disposable implements IAgentService { // ---- helpers ------------------------------------------------------------ + /** + * Reconstructs completed `ITurn[]` from a sequence of agent session + * messages. Each user-message starts a new turn; the assistant message + * closes it. + */ + private _buildTurnsFromMessages( + messages: readonly (IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[], + ): ITurn[] { + const turns: ITurn[] = []; + let currentTurn: { + id: string; + userMessage: { text: string }; + responseParts: IResponsePart[]; + pendingTools: Map; + } | undefined; + + const finalizeTurn = (turn: NonNullable, state: TurnState): void => { + turns.push({ + id: turn.id, + userMessage: turn.userMessage, + responseParts: turn.responseParts, + usage: undefined, + state, + }); + }; + + const startTurn = (id: string, text: string): NonNullable => ({ + id, + userMessage: { text }, + responseParts: [], + pendingTools: new Map(), + }); + + for (const msg of messages) { + if (msg.type === 'message' && msg.role === 'user') { + if (currentTurn) { + finalizeTurn(currentTurn, TurnState.Cancelled); + } + currentTurn = startTurn(msg.messageId, msg.content); + } else if (msg.type === 'message' && msg.role === 'assistant') { + if (!currentTurn) { + currentTurn = startTurn(msg.messageId, ''); + } + + if (msg.content) { + currentTurn.responseParts.push({ + kind: ResponsePartKind.Markdown, + id: generateUuid(), + content: msg.content, + }); + } + + if (!msg.toolRequests || msg.toolRequests.length === 0) { + finalizeTurn(currentTurn, TurnState.Complete); + currentTurn = undefined; + } + } else if (msg.type === 'tool_start') { + currentTurn?.pendingTools.set(msg.toolCallId, msg); + } else if (msg.type === 'tool_complete') { + if (currentTurn) { + const start = currentTurn.pendingTools.get(msg.toolCallId); + currentTurn.pendingTools.delete(msg.toolCallId); + + const tc: IToolCallCompletedState = { + status: ToolCallStatus.Completed, + toolCallId: msg.toolCallId, + toolName: start?.toolName ?? 'unknown', + displayName: start?.displayName ?? 'Unknown Tool', + invocationMessage: start?.invocationMessage ?? '', + toolInput: start?.toolInput, + success: msg.result.success, + pastTenseMessage: msg.result.pastTenseMessage, + content: msg.result.content, + error: msg.result.error, + confirmed: ToolCallConfirmationReason.NotNeeded, + _meta: start ? { + toolKind: start.toolKind, + language: start.language, + } : undefined, + }; + currentTurn.responseParts.push({ + kind: ResponsePartKind.ToolCall, + toolCall: tc, + }); + } + } + } + + if (currentTurn) { + finalizeTurn(currentTurn, TurnState.Cancelled); + } + + return turns; + } + + private async _fetchSessionDbContent(fields: ISessionDbUriFields): Promise { + const sessionUri = URI.parse(fields.sessionUri); + const ref = this._sessionDataService.openDatabase(sessionUri); + try { + const content = await ref.object.readFileEditContent(fields.toolCallId, fields.filePath); + if (!content) { + throw new ProtocolError(AhpErrorCodes.NotFound, `File edit not found: toolCallId=${fields.toolCallId}, filePath=${fields.filePath}`); + } + const bytes = fields.part === 'before' ? content.beforeContent : content.afterContent; + return { + data: new TextDecoder().decode(bytes), + encoding: ContentEncoding.Utf8, + contentType: 'text/plain', + }; + } finally { + ref.dispose(); + } + } + private _findProviderForSession(session: URI | string): IAgent | undefined { const key = typeof session === 'string' ? session : session.toString(); const providerId = this._sessionToProvider.get(key); diff --git a/src/vs/platform/agentHost/node/agentSideEffects.ts b/src/vs/platform/agentHost/node/agentSideEffects.ts index c8c59ef9ab5..782898884c4 100644 --- a/src/vs/platform/agentHost/node/agentSideEffects.ts +++ b/src/vs/platform/agentHost/node/agentSideEffects.ts @@ -3,36 +3,21 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import * as os from 'os'; -import { decodeBase64, VSBuffer } from '../../../base/common/buffer.js'; import { match as globMatch } from '../../../base/common/glob.js'; import { Disposable, DisposableStore, IDisposable } from '../../../base/common/lifecycle.js'; import { autorun, IObservable } from '../../../base/common/observable.js'; import { URI } from '../../../base/common/uri.js'; import { generateUuid } from '../../../base/common/uuid.js'; -import { FileSystemProviderErrorCode, IFileService, toFileSystemProviderErrorCode } from '../../files/common/files.js'; import { ILogService } from '../../log/common/log.js'; -import { IAgent, IAgentAttachment, IAgentMessageEvent, IAgentToolCompleteEvent, IAgentToolStartEvent, IAuthenticateParams, IAuthenticateResult, IResourceMetadata } from '../common/agentService.js'; +import { IAgent, IAgentAttachment } from '../common/agentService.js'; import { ISessionDataService } from '../common/sessionDataService.js'; import { ActionType, ISessionAction } from '../common/state/sessionActions.js'; -import { AhpErrorCodes, AHP_PROVIDER_NOT_FOUND, AHP_SESSION_NOT_FOUND, ContentEncoding, IBrowseDirectoryResult, ICreateSessionParams, IDirectoryEntry, IFetchContentResult, IWriteFileParams, IWriteFileResult, JSON_RPC_INTERNAL_ERROR, ProtocolError } from '../common/state/sessionProtocol.js'; import { PendingMessageKind, - ResponsePartKind, - SessionStatus, - ToolCallConfirmationReason, - ToolCallStatus, - TurnState, - type IResponsePart, type ISessionModelInfo, - type ISessionSummary, - type IToolCallCompletedState, - type ITurn, type URI as ProtocolURI, } from '../common/state/sessionState.js'; import { AgentEventMapper } from './agentEventMapper.js'; -import { ISessionDbUriFields, parseSessionDbUri } from './copilot/fileEditTracker.js'; -import type { IProtocolSideEffectHandler } from './protocolServerHandler.js'; import { SessionStateManager } from './sessionStateManager.js'; /** @@ -50,14 +35,14 @@ export interface IAgentSideEffectsOptions { /** * Shared implementation of agent side-effect handling. * - * Routes client-dispatched actions to the correct agent backend, handles - * session create/dispose/list operations, tracks pending permission requests, + * Routes client-dispatched actions to the correct agent backend, + * restores sessions from previous lifetimes, handles filesystem + * operations (browse/fetch/write), tracks pending permission requests, * and wires up agent progress events to the state manager. * - * Used by both the Electron utility-process path ({@link AgentService}) and - * the standalone WebSocket server (`agentHostServerMain`). + * Session create/dispose/list and auth are handled by {@link AgentService}. */ -export class AgentSideEffects extends Disposable implements IProtocolSideEffectHandler { +export class AgentSideEffects extends Disposable { /** Maps tool call IDs to the agent that owns them, for routing confirmations. */ private readonly _toolCallAgents = new Map(); @@ -68,7 +53,6 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH private readonly _stateManager: SessionStateManager, private readonly _options: IAgentSideEffectsOptions, private readonly _logService: ILogService, - private readonly _fileService: IFileService, ) { super(); @@ -187,7 +171,7 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH return disposables; } - // ---- IProtocolSideEffectHandler ----------------------------------------- + // ---- Side-effect handlers -------------------------------------------------- handleAction(action: ISessionAction): void { switch (action.type) { @@ -350,329 +334,6 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH }); } - async handleCreateSession(command: ICreateSessionParams): Promise { - const provider = command.provider; - if (!provider) { - throw new ProtocolError(AHP_PROVIDER_NOT_FOUND, 'No provider specified for session creation'); - } - const agent = this._options.agents.get().find(a => a.id === provider); - if (!agent) { - throw new ProtocolError(AHP_PROVIDER_NOT_FOUND, `No agent registered for provider: ${provider}`); - } - // Use the client-provided session URI per the protocol spec - const session = command.session; - await agent.createSession({ - provider, - model: command.model, - workingDirectory: command.workingDirectory ? URI.parse(command.workingDirectory) : undefined, - session: URI.parse(session), - }); - const summary: ISessionSummary = { - resource: session, - provider, - title: 'Session', - status: SessionStatus.Idle, - createdAt: Date.now(), - modifiedAt: Date.now(), - workingDirectory: command.workingDirectory, - }; - this._stateManager.createSession(summary); - this._stateManager.dispatchServerAction({ type: ActionType.SessionReady, session }); - } - - handleDisposeSession(session: ProtocolURI): void { - const agent = this._options.getAgent(session); - agent?.disposeSession(URI.parse(session)).catch(() => { }); - this._stateManager.deleteSession(session); - } - - async handleListSessions(): Promise { - const allSessions: ISessionSummary[] = []; - for (const agent of this._options.agents.get()) { - const sessions = await agent.listSessions(); - const provider = agent.id; - for (const s of sessions) { - allSessions.push({ - resource: s.session.toString(), - provider, - title: s.summary ?? 'Session', - status: SessionStatus.Idle, - createdAt: s.startTime, - modifiedAt: s.modifiedTime, - }); - } - } - return allSessions; - } - - /** - * Restores a session from a previous server lifetime into the state - * manager. Fetches the session's message history from the agent backend, - * reconstructs `ITurn[]`, and creates the session in the state manager. - * - * @throws {ProtocolError} if the session URI doesn't match any agent or - * the agent cannot retrieve the session messages. - */ - async handleRestoreSession(session: ProtocolURI): Promise { - // Already in state manager - nothing to do. - if (this._stateManager.getSessionState(session)) { - return; - } - - const agent = this._options.getAgent(session); - if (!agent) { - throw new ProtocolError(AHP_SESSION_NOT_FOUND, `No agent for session: ${session}`); - } - - // Verify the session actually exists on the backend to avoid - // creating phantom sessions for made-up URIs. - let allSessions; - try { - allSessions = await agent.listSessions(); - } catch (err) { - if (err instanceof ProtocolError) { - throw err; - } - const message = err instanceof Error ? err.message : String(err); - throw new ProtocolError(JSON_RPC_INTERNAL_ERROR, `Failed to list sessions for ${session}: ${message}`); - } - const meta = allSessions.find(s => s.session.toString() === session); - if (!meta) { - throw new ProtocolError(AHP_SESSION_NOT_FOUND, `Session not found on backend: ${session}`); - } - - const sessionUri = URI.parse(session); - let messages; - try { - messages = await agent.getSessionMessages(sessionUri); - } catch (err) { - if (err instanceof ProtocolError) { - throw err; - } - const message = err instanceof Error ? err.message : String(err); - throw new ProtocolError(JSON_RPC_INTERNAL_ERROR, `Failed to restore session ${session}: ${message}`); - } - const turns = this._buildTurnsFromMessages(messages); - - const summary: ISessionSummary = { - resource: session, - provider: agent.id, - title: meta.summary ?? 'Session', - status: SessionStatus.Idle, - createdAt: meta.startTime, - modifiedAt: meta.modifiedTime, - workingDirectory: meta.workingDirectory?.toString(), - }; - - this._stateManager.restoreSession(summary, turns); - this._logService.info(`[AgentSideEffects] Restored session ${session} with ${turns.length} turns`); - } - - /** - * Reconstructs completed `ITurn[]` from a sequence of agent session - * messages (user messages, assistant messages, tool starts, tool - * completions). Each user-message starts a new turn; the assistant - * message closes it. - */ - private _buildTurnsFromMessages( - messages: readonly (IAgentMessageEvent | IAgentToolStartEvent | IAgentToolCompleteEvent)[], - ): ITurn[] { - const turns: ITurn[] = []; - let currentTurn: { - id: string; - userMessage: { text: string }; - responseParts: IResponsePart[]; - pendingTools: Map; - } | undefined; - - const finalizeTurn = (turn: NonNullable, state: TurnState): void => { - turns.push({ - id: turn.id, - userMessage: turn.userMessage, - responseParts: turn.responseParts, - usage: undefined, - state, - }); - }; - - const startTurn = (id: string, text: string): NonNullable => ({ - id, - userMessage: { text }, - responseParts: [], - pendingTools: new Map(), - }); - - for (const msg of messages) { - if (msg.type === 'message' && msg.role === 'user') { - // Flush any in-progress turn (e.g. interrupted/cancelled - // turn that never got a closing assistant message). - if (currentTurn) { - finalizeTurn(currentTurn, TurnState.Cancelled); - } - currentTurn = startTurn(msg.messageId, msg.content); - } else if (msg.type === 'message' && msg.role === 'assistant') { - if (!currentTurn) { - currentTurn = startTurn(msg.messageId, ''); - } - - if (msg.content) { - currentTurn.responseParts.push({ - kind: ResponsePartKind.Markdown, - id: generateUuid(), - content: msg.content, - }); - } - - if (!msg.toolRequests || msg.toolRequests.length === 0) { - finalizeTurn(currentTurn, TurnState.Complete); - currentTurn = undefined; - } - } else if (msg.type === 'tool_start') { - currentTurn?.pendingTools.set(msg.toolCallId, msg); - } else if (msg.type === 'tool_complete') { - if (currentTurn) { - const start = currentTurn.pendingTools.get(msg.toolCallId); - currentTurn.pendingTools.delete(msg.toolCallId); - - const tc: IToolCallCompletedState = { - status: ToolCallStatus.Completed, - toolCallId: msg.toolCallId, - toolName: start?.toolName ?? 'unknown', - displayName: start?.displayName ?? 'Unknown Tool', - invocationMessage: start?.invocationMessage ?? '', - toolInput: start?.toolInput, - success: msg.result.success, - pastTenseMessage: msg.result.pastTenseMessage, - content: msg.result.content, - error: msg.result.error, - confirmed: ToolCallConfirmationReason.NotNeeded, - _meta: start ? { - toolKind: start.toolKind, - language: start.language, - } : undefined, - }; - currentTurn.responseParts.push({ - kind: ResponsePartKind.ToolCall, - toolCall: tc, - }); - } - } - } - - if (currentTurn) { - finalizeTurn(currentTurn, TurnState.Cancelled); - } - - return turns; - } - - handleGetResourceMetadata(): IResourceMetadata { - const resources = this._options.agents.get().flatMap(a => a.getProtectedResources()); - return { resources }; - } - - async handleAuthenticate(params: IAuthenticateParams): Promise { - for (const agent of this._options.agents.get()) { - const resources = agent.getProtectedResources(); - if (resources.some(r => r.resource === params.resource)) { - const accepted = await agent.authenticate(params.resource, params.token); - if (accepted) { - return { authenticated: true }; - } - } - } - return { authenticated: false }; - } - - async handleBrowseDirectory(uri: ProtocolURI): Promise { - let stat; - try { - stat = await this._fileService.resolve(URI.parse(uri)); - } catch { - throw new ProtocolError(AhpErrorCodes.NotFound, `Directory not found: ${uri.toString()}`); - } - - if (!stat.isDirectory) { - throw new ProtocolError(AhpErrorCodes.NotFound, `Not a directory: ${uri.toString()}`); - } - - const entries: IDirectoryEntry[] = (stat.children ?? []).map(child => ({ - name: child.name, - type: child.isDirectory ? 'directory' : 'file', - })); - return { entries }; - } - - getDefaultDirectory(): ProtocolURI { - return URI.file(os.homedir()).toString(); - } - - async handleFetchContent(uri: ProtocolURI): Promise { - // Handle session-db: URIs that reference file-edit content stored - // in a per-session SQLite database. - const dbFields = parseSessionDbUri(uri); - if (dbFields) { - return this._fetchSessionDbContent(dbFields); - } - - try { - const content = await this._fileService.readFile(URI.parse(uri)); - return { - data: content.value.toString(), - encoding: ContentEncoding.Utf8, - contentType: 'text/plain', - }; - } catch (_e) { - throw new ProtocolError(AhpErrorCodes.NotFound, `Content not found: ${uri}`); - } - } - - async handleWriteFile(params: IWriteFileParams): Promise { - const fileUri = typeof params.uri === 'string' ? URI.parse(params.uri) : URI.revive(params.uri); - let content: VSBuffer; - if (params.encoding === ContentEncoding.Base64) { - content = decodeBase64(params.data); - } else { - content = VSBuffer.fromString(params.data); - } - try { - if (params.createOnly) { - await this._fileService.createFile(fileUri, content, { overwrite: false }); - } else { - await this._fileService.writeFile(fileUri, content); - } - return {}; - } catch (e) { - const code = toFileSystemProviderErrorCode(e as Error); - if (code === FileSystemProviderErrorCode.FileExists) { - throw new ProtocolError(AhpErrorCodes.AlreadyExists, `File already exists: ${fileUri.toString()}`); - } - if (code === FileSystemProviderErrorCode.NoPermissions) { - throw new ProtocolError(AhpErrorCodes.PermissionDenied, `Permission denied: ${fileUri.toString()}`); - } - throw new ProtocolError(AhpErrorCodes.NotFound, `Failed to write file: ${fileUri.toString()}`); - } - } - - private async _fetchSessionDbContent(fields: ISessionDbUriFields): Promise { - const sessionUri = URI.parse(fields.sessionUri); - const ref = this._options.sessionDataService.openDatabase(sessionUri); - try { - const content = await ref.object.readFileEditContent(fields.toolCallId, fields.filePath); - if (!content) { - throw new ProtocolError(AhpErrorCodes.NotFound, `File edit not found: toolCallId=${fields.toolCallId}, filePath=${fields.filePath}`); - } - const bytes = fields.part === 'before' ? content.beforeContent : content.afterContent; - return { - data: new TextDecoder().decode(bytes), - encoding: ContentEncoding.Utf8, - contentType: 'text/plain', - }; - } finally { - ref.dispose(); - } - } - override dispose(): void { this._toolCallAgents.clear(); super.dispose(); diff --git a/src/vs/platform/agentHost/node/protocolServerHandler.ts b/src/vs/platform/agentHost/node/protocolServerHandler.ts index 232eb8c372c..640f5263103 100644 --- a/src/vs/platform/agentHost/node/protocolServerHandler.ts +++ b/src/vs/platform/agentHost/node/protocolServerHandler.ts @@ -5,12 +5,14 @@ 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 type { IAgentDescriptor, IAuthenticateParams, IAuthenticateResult, IResourceMetadata } from '../common/agentService.js'; +import { AgentSession, type IAgentService, type IAuthenticateParams } from '../common/agentService.js'; import type { ICommandMap } from '../common/state/protocol/messages.js'; import { IActionEnvelope, INotification, isSessionAction, type ISessionAction } from '../common/state/sessionActions.js'; import { MIN_PROTOCOL_VERSION, PROTOCOL_VERSION } from '../common/state/sessionCapabilities.js'; import { + AHP_PROVIDER_NOT_FOUND, AHP_SESSION_NOT_FOUND, AHP_UNSUPPORTED_PROTOCOL_VERSION, isJsonRpcNotification, @@ -18,17 +20,12 @@ import { JSON_RPC_INTERNAL_ERROR, ProtocolError, type IAhpServerNotification, - type IBrowseDirectoryResult, - type ICreateSessionParams, - type IFetchContentResult, type IInitializeParams, type IJsonRpcResponse, type IReconnectParams, type IStateSnapshot, - type IWriteFileParams, - type IWriteFileResult, } from '../common/state/sessionProtocol.js'; -import { ROOT_STATE_URI, type ISessionSummary, type URI } from '../common/state/sessionState.js'; +import { ROOT_STATE_URI, SessionStatus } from '../common/state/sessionState.js'; import type { IProtocolServer, IProtocolTransport } from '../common/state/sessionTransport.js'; import { SessionStateManager } from './sessionStateManager.js'; @@ -80,9 +77,17 @@ interface IConnectedClient { readonly disposables: DisposableStore; } +/** + * Configuration for protocol-level concerns outside of IAgentService. + */ +export interface IProtocolServerConfig { + /** Default directory returned to clients during the initialize handshake. */ + readonly defaultDirectory?: string; +} + /** * Server-side handler that manages protocol connections, routes JSON-RPC - * messages to the state manager, and broadcasts actions/notifications + * messages to the agent service, and broadcasts actions/notifications * to subscribed clients. */ export class ProtocolServerHandler extends Disposable { @@ -96,9 +101,10 @@ export class ProtocolServerHandler extends Disposable { readonly onDidChangeConnectionCount = this._onDidChangeConnectionCount.event; constructor( + private readonly _agentService: IAgentService, private readonly _stateManager: SessionStateManager, private readonly _server: IProtocolServer, - private readonly _sideEffectHandler: IProtocolSideEffectHandler, + private readonly _config: IProtocolServerConfig, @ILogService private readonly _logService: ILogService, ) { super(); @@ -168,10 +174,8 @@ export class ProtocolServerHandler extends Disposable { case 'dispatchAction': if (client) { this._logService.trace(`[ProtocolServer] dispatchAction: ${JSON.stringify(msg.params.action.type)}`); - const origin = { clientId: client.clientId, clientSeq: msg.params.clientSeq }; const action = msg.params.action as ISessionAction; - this._stateManager.dispatchClientAction(action, origin); - this._sideEffectHandler.handleAction(action); + this._agentService.dispatchAction(action, client.clientId, msg.params.clientSeq); } break; } @@ -234,7 +238,7 @@ export class ProtocolServerHandler extends Disposable { protocolVersion: PROTOCOL_VERSION, serverSeq: this._stateManager.serverSeq, snapshots, - defaultDirectory: this._sideEffectHandler.getDefaultDirectory?.(), + defaultDirectory: this._config.defaultDirectory, }, }; } @@ -293,33 +297,56 @@ export class ProtocolServerHandler extends Disposable { */ private readonly _requestHandlers: RequestHandlerMap = { subscribe: async (client, params) => { - let snapshot = this._stateManager.getSnapshot(params.resource); - if (!snapshot) { - // Session may exist on the agent backend but not in the - // current state manager (e.g. from a previous server - // lifetime). Try to restore it. - await this._sideEffectHandler.handleRestoreSession(params.resource); - snapshot = this._stateManager.getSnapshot(params.resource); - } - if (!snapshot) { + try { + const snapshot = await this._agentService.subscribe(URI.parse(params.resource)); + client.subscriptions.add(params.resource); + return { snapshot }; + } catch (err) { + if (err instanceof ProtocolError) { + throw err; + } throw new ProtocolError(AHP_SESSION_NOT_FOUND, `Resource not found: ${params.resource}`); } - client.subscriptions.add(params.resource); - return { snapshot }; }, createSession: async (_client, params) => { - await this._sideEffectHandler.handleCreateSession(params); + let createdSession: URI; + try { + createdSession = await this._agentService.createSession({ + provider: params.provider, + model: params.model, + workingDirectory: params.workingDirectory ? URI.parse(params.workingDirectory) : undefined, + session: URI.parse(params.session), + }); + } catch (err) { + if (err instanceof ProtocolError) { + throw err; + } + throw new ProtocolError(AHP_PROVIDER_NOT_FOUND, err instanceof Error ? err.message : String(err)); + } + // Verify the provider honored the client-chosen session URI per the protocol contract + if (createdSession.toString() !== URI.parse(params.session).toString()) { + this._logService.warn(`[ProtocolServer] createSession: provider returned URI ${createdSession.toString()} but client requested ${params.session}`); + } return null; }, disposeSession: async (_client, params) => { - this._sideEffectHandler.handleDisposeSession(params.session); + await this._agentService.disposeSession(URI.parse(params.session)); return null; }, writeFile: async (_client, params) => { - return this._sideEffectHandler.handleWriteFile(params); + return this._agentService.writeFile(params); }, listSessions: async () => { - const items = await this._sideEffectHandler.handleListSessions(); + const sessions = await this._agentService.listSessions(); + const items = sessions.map(s => ({ + resource: s.session.toString(), + provider: AgentSession.provider(s.session) ?? 'copilot', + title: s.summary ?? 'Session', + status: SessionStatus.Idle, + createdAt: s.startTime, + modifiedAt: s.modifiedTime, + workingDirectory: s.workingDirectory?.toString(), + })); return { items }; }, fetchTurns: async (_client, params) => { @@ -345,10 +372,10 @@ export class ProtocolServerHandler extends Disposable { }; }, browseDirectory: async (_client, params) => { - return this._sideEffectHandler.handleBrowseDirectory(params.uri); + return this._agentService.browseDirectory(URI.parse(params.uri)); }, fetchContent: async (_client, params) => { - return this._sideEffectHandler.handleFetchContent(params.uri); + return this._agentService.fetchContent(URI.parse(params.uri)); }, }; @@ -388,15 +415,20 @@ export class ProtocolServerHandler extends Disposable { private _handleExtensionRequest(method: string, params: unknown): Promise | undefined { switch (method) { case 'getResourceMetadata': - return Promise.resolve(this._sideEffectHandler.handleGetResourceMetadata()); - case 'authenticate': - return this._sideEffectHandler.handleAuthenticate(params as IAuthenticateParams); + return this._agentService.getResourceMetadata(); + case 'authenticate': { + const authParams = params as IAuthenticateParams; + if (!authParams || typeof authParams.resource !== 'string' || typeof authParams.token !== 'string') { + return Promise.reject(new ProtocolError(-32602, 'Invalid authenticate params')); + } + return this._agentService.authenticate(authParams); + } case 'refreshModels': - return this._sideEffectHandler.handleRefreshModels?.() ?? Promise.resolve(null); + return this._agentService.refreshModels(); case 'listAgents': - return Promise.resolve(this._sideEffectHandler.handleListAgents?.() ?? []); + return this._agentService.listAgents(); case 'shutdown': - return this._sideEffectHandler.handleShutdown?.() ?? Promise.resolve(null); + return this._agentService.shutdown(); default: return undefined; } @@ -441,29 +473,3 @@ export class ProtocolServerHandler extends Disposable { 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: ISessionAction): void; - handleCreateSession(command: ICreateSessionParams): Promise; - handleDisposeSession(session: URI): void; - handleListSessions(): Promise; - /** Restore a session from a previous server lifetime into the state manager. */ - handleRestoreSession(session: URI): Promise; - handleGetResourceMetadata(): IResourceMetadata; - handleAuthenticate(params: IAuthenticateParams): Promise; - handleBrowseDirectory(uri: URI): Promise; - handleFetchContent(uri: URI): Promise; - handleWriteFile(params: IWriteFileParams): Promise; - /** Returns the server's default browsing directory, if available. */ - getDefaultDirectory?(): URI; - /** Refresh models from all providers (VS Code extension method). */ - handleRefreshModels?(): Promise; - /** List agent descriptors (VS Code extension method). */ - handleListAgents?(): IAgentDescriptor[]; - /** Shut down all providers (VS Code extension method). */ - handleShutdown?(): Promise; -} diff --git a/src/vs/platform/agentHost/protocol.md b/src/vs/platform/agentHost/protocol.md index f089aba1bd9..56aa36d402e 100644 --- a/src/vs/platform/agentHost/protocol.md +++ b/src/vs/platform/agentHost/protocol.md @@ -40,7 +40,7 @@ Use this checklist when adding a new action, command, state field, or notificati 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`. +4. **Add the handler** in `protocolServerHandler.ts` request handler map. If the command requires I/O or agent interaction, it delegates to `IAgentService`. Implement the backing method in `AgentService` (or `AgentSideEffects` for operations that involve agent backends). 5. **Update `protocol.md`** — add the command to the Commands table. 6. **Verify the E2E test passes.** diff --git a/src/vs/platform/agentHost/test/node/agentService.test.ts b/src/vs/platform/agentHost/test/node/agentService.test.ts index bbe7695e5e9..81c9038f86b 100644 --- a/src/vs/platform/agentHost/test/node/agentService.test.ts +++ b/src/vs/platform/agentHost/test/node/agentService.test.ts @@ -4,14 +4,18 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { VSBuffer } from '../../../../base/common/buffer.js'; import { DisposableStore, toDisposable } from '../../../../base/common/lifecycle.js'; +import { Schemas } from '../../../../base/common/network.js'; import { URI } from '../../../../base/common/uri.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import { NullLogService } from '../../../log/common/log.js'; import { FileService } from '../../../files/common/fileService.js'; +import { InMemoryFileSystemProvider } from '../../../files/common/inMemoryFilesystemProvider.js'; import { AgentSession } from '../../common/agentService.js'; import { ISessionDataService } from '../../common/sessionDataService.js'; import { ActionType, IActionEnvelope } from '../../common/state/sessionActions.js'; +import { ResponsePartKind, SessionLifecycle, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType, TurnState, type IMarkdownResponsePart, type IToolCallCompletedState, type IToolCallResponsePart } from '../../common/state/sessionState.js'; import { AgentService } from '../../node/agentService.js'; import { MockAgent } from './mockAgent.js'; @@ -20,8 +24,9 @@ suite('AgentService (node dispatcher)', () => { const disposables = new DisposableStore(); let service: AgentService; let copilotAgent: MockAgent; + let fileService: FileService; - setup(() => { + setup(async () => { const nullSessionDataService: ISessionDataService = { _serviceBrand: undefined, getSessionDataDir: () => URI.parse('inmemory:/session-data'), @@ -30,7 +35,14 @@ suite('AgentService (node dispatcher)', () => { deleteSessionData: async () => { }, cleanupOrphanedData: async () => { }, }; - service = disposables.add(new AgentService(new NullLogService(), disposables.add(new FileService(new NullLogService())), nullSessionDataService)); + fileService = disposables.add(new FileService(new NullLogService())); + disposables.add(fileService.registerProvider(Schemas.inMemory, disposables.add(new InMemoryFileSystemProvider()))); + + // Seed a directory for browseDirectory tests + await fileService.createFolder(URI.from({ scheme: Schemas.inMemory, path: '/testDir' })); + await fileService.writeFile(URI.from({ scheme: Schemas.inMemory, path: '/testDir/file.txt' }), VSBuffer.fromString('hello')); + + service = disposables.add(new AgentService(new NullLogService(), fileService, nullSessionDataService)); copilotAgent = new MockAgent('copilot'); disposables.add(toDisposable(() => copilotAgent.dispose())); }); @@ -226,4 +238,108 @@ suite('AgentService (node dispatcher)', () => { assert.ok(copilotShutdown); }); }); + + // ---- restoreSession ------------------------------------------------- + + suite('restoreSession', () => { + + test('restores a session with message history', async () => { + service.registerProvider(copilotAgent); + const session = await copilotAgent.createSession(); + const sessions = await copilotAgent.listSessions(); + const sessionResource = sessions[0].session; + + copilotAgent.sessionMessages = [ + { type: 'message', session, role: 'user', messageId: 'msg-1', content: 'Hello', toolRequests: [] }, + { type: 'message', session, role: 'assistant', messageId: 'msg-2', content: 'Hi there!', toolRequests: [] }, + ]; + + await service.restoreSession(sessionResource); + + const state = service.stateManager.getSessionState(sessionResource.toString()); + assert.ok(state, 'session should be in state manager'); + assert.strictEqual(state!.lifecycle, SessionLifecycle.Ready); + assert.strictEqual(state!.turns.length, 1); + assert.strictEqual(state!.turns[0].userMessage.text, 'Hello'); + const mdPart = state!.turns[0].responseParts.find((p): p is IMarkdownResponsePart => p.kind === ResponsePartKind.Markdown); + assert.ok(mdPart); + assert.strictEqual(mdPart.content, 'Hi there!'); + assert.strictEqual(state!.turns[0].state, TurnState.Complete); + }); + + test('restores a session with tool calls', async () => { + service.registerProvider(copilotAgent); + const session = await copilotAgent.createSession(); + const sessions = await copilotAgent.listSessions(); + const sessionResource = sessions[0].session; + + copilotAgent.sessionMessages = [ + { type: 'message', session, role: 'user', messageId: 'msg-1', content: 'Run a command', toolRequests: [] }, + { type: 'message', session, role: 'assistant', messageId: 'msg-2', content: 'I will run a command.', toolRequests: [{ toolCallId: 'tc-1', name: 'shell' }] }, + { type: 'tool_start', session, toolCallId: 'tc-1', toolName: 'shell', displayName: 'Shell', invocationMessage: 'Running command...' }, + { type: 'tool_complete', session, toolCallId: 'tc-1', result: { success: true, pastTenseMessage: 'Ran command', content: [{ type: ToolResultContentType.Text, text: 'output' }] } }, + { type: 'message', session, role: 'assistant', messageId: 'msg-3', content: 'Done!', toolRequests: [] }, + ]; + + await service.restoreSession(sessionResource); + + const state = service.stateManager.getSessionState(sessionResource.toString()); + assert.ok(state); + const turn = state!.turns[0]; + const toolCallParts = turn.responseParts.filter((p): p is IToolCallResponsePart => p.kind === ResponsePartKind.ToolCall); + assert.strictEqual(toolCallParts.length, 1); + const tc = toolCallParts[0].toolCall as IToolCallCompletedState; + assert.strictEqual(tc.status, ToolCallStatus.Completed); + assert.strictEqual(tc.toolCallId, 'tc-1'); + assert.strictEqual(tc.confirmed, ToolCallConfirmationReason.NotNeeded); + }); + + test('flushes interrupted turns', async () => { + service.registerProvider(copilotAgent); + const session = await copilotAgent.createSession(); + const sessions = await copilotAgent.listSessions(); + const sessionResource = sessions[0].session; + + copilotAgent.sessionMessages = [ + { type: 'message', session, role: 'user', messageId: 'msg-1', content: 'Interrupted', toolRequests: [] }, + { type: 'message', session, role: 'user', messageId: 'msg-2', content: 'Retried', toolRequests: [] }, + { type: 'message', session, role: 'assistant', messageId: 'msg-3', content: 'Answer', toolRequests: [] }, + ]; + + await service.restoreSession(sessionResource); + + const state = service.stateManager.getSessionState(sessionResource.toString()); + assert.ok(state); + assert.strictEqual(state!.turns.length, 2); + assert.strictEqual(state!.turns[0].state, TurnState.Cancelled); + assert.strictEqual(state!.turns[1].state, TurnState.Complete); + }); + + test('throws when session is not found on backend', async () => { + service.registerProvider(copilotAgent); + await assert.rejects( + () => service.restoreSession(AgentSession.uri('copilot', 'nonexistent')), + /Session not found on backend/, + ); + }); + }); + + // ---- browseDirectory ------------------------------------------------ + + suite('browseDirectory', () => { + + test('throws when the directory does not exist', async () => { + await assert.rejects( + () => service.browseDirectory(URI.from({ scheme: Schemas.inMemory, path: '/nonexistent' })), + /Directory not found/, + ); + }); + + test('throws when the target is not a directory', async () => { + await assert.rejects( + () => service.browseDirectory(URI.from({ scheme: Schemas.inMemory, path: '/testDir/file.txt' })), + /Not a directory/, + ); + }); + }); }); diff --git a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts index b6ec18adc3c..e6b448f95e8 100644 --- a/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts +++ b/src/vs/platform/agentHost/test/node/agentSideEffects.test.ts @@ -16,7 +16,7 @@ import { NullLogService } from '../../../log/common/log.js'; import { AgentSession, IAgent } from '../../common/agentService.js'; import { ISessionDataService } from '../../common/sessionDataService.js'; import { ActionType, IActionEnvelope, ISessionAction } from '../../common/state/sessionActions.js'; -import { PendingMessageKind, ResponsePartKind, SessionLifecycle, SessionStatus, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType, TurnState, type IMarkdownResponsePart, type IToolCallCompletedState, type IToolCallResponsePart } from '../../common/state/sessionState.js'; +import { PendingMessageKind, ResponsePartKind, SessionStatus, ToolCallStatus, type IToolCallResponsePart } from '../../common/state/sessionState.js'; import { AgentSideEffects } from '../../node/agentSideEffects.js'; import { SessionStateManager } from '../../node/sessionStateManager.js'; import { MockAgent } from './mockAgent.js'; @@ -78,7 +78,7 @@ suite('AgentSideEffects', () => { deleteSessionData: async () => { }, cleanupOrphanedData: async () => { }, } satisfies ISessionDataService, - }, new NullLogService(), fileService)); + }, new NullLogService())); }); teardown(() => { @@ -113,7 +113,7 @@ suite('AgentSideEffects', () => { getAgent: () => undefined, agents: emptyAgents, sessionDataService: {} as ISessionDataService, - }, new NullLogService(), fileService)); + }, new NullLogService())); const envelopes: IActionEnvelope[] = []; disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e))); @@ -201,272 +201,6 @@ suite('AgentSideEffects', () => { }); }); - // ---- handleCreateSession -------------------------------------------- - - suite('handleCreateSession', () => { - - test('creates a session and dispatches session/ready', async () => { - const envelopes: IActionEnvelope[] = []; - disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e))); - - await sideEffects.handleCreateSession({ session: sessionUri.toString(), provider: 'mock' }); - - const ready = envelopes.find(e => e.action.type === ActionType.SessionReady); - assert.ok(ready, 'should dispatch session/ready'); - }); - - test('throws when no provider is specified', async () => { - await assert.rejects( - () => sideEffects.handleCreateSession({ session: sessionUri.toString() }), - /No provider specified/, - ); - }); - - test('throws when no agent matches provider', async () => { - const emptyAgents = observableValue('agents', []); - const noAgentSideEffects = disposables.add(new AgentSideEffects(stateManager, { - getAgent: () => undefined, - agents: emptyAgents, - sessionDataService: {} as ISessionDataService, - }, new NullLogService(), fileService)); - - await assert.rejects( - () => noAgentSideEffects.handleCreateSession({ session: sessionUri.toString(), provider: 'nonexistent' }), - /No agent registered/, - ); - }); - }); - - // ---- handleDisposeSession ------------------------------------------- - - suite('handleDisposeSession', () => { - - test('disposes the session on the agent and removes state', async () => { - setupSession(); - - sideEffects.handleDisposeSession(sessionUri.toString()); - - await new Promise(r => setTimeout(r, 10)); - - assert.strictEqual(agent.disposeSessionCalls.length, 1); - assert.strictEqual(stateManager.getSessionState(sessionUri.toString()), undefined); - }); - }); - - // ---- handleListSessions --------------------------------------------- - - suite('handleListSessions', () => { - - test('aggregates sessions from all agents', async () => { - await agent.createSession(); - const sessions = await sideEffects.handleListSessions(); - assert.strictEqual(sessions.length, 1); - assert.strictEqual(sessions[0].provider, 'mock'); - assert.strictEqual(sessions[0].title, 'Session'); - }); - }); - - // ---- handleRestoreSession ----------------------------------------------- - - suite('handleRestoreSession', () => { - - test('restores a session with message history into the state manager', async () => { - // Create a session on the agent backend (not in the state manager) - const session = await agent.createSession(); - const sessions = await agent.listSessions(); - const sessionResource = sessions[0].session.toString(); - - // Set up the agent's stored messages - agent.sessionMessages = [ - { type: 'message', session, role: 'user', messageId: 'msg-1', content: 'Hello', toolRequests: [] }, - { type: 'message', session, role: 'assistant', messageId: 'msg-2', content: 'Hi there!', toolRequests: [] }, - ]; - - // Before restore, state manager shouldn't have it - assert.strictEqual(stateManager.getSessionState(sessionResource), undefined); - - await sideEffects.handleRestoreSession(sessionResource); - - // After restore, state manager should have it - const state = stateManager.getSessionState(sessionResource); - assert.ok(state, 'session should be in state manager'); - assert.strictEqual(state!.lifecycle, SessionLifecycle.Ready); - assert.strictEqual(state!.turns.length, 1); - assert.strictEqual(state!.turns[0].userMessage.text, 'Hello'); - const mdPart = state!.turns[0].responseParts.find((p): p is IMarkdownResponsePart => p.kind === ResponsePartKind.Markdown); - assert.ok(mdPart, 'should have a markdown response part'); - assert.strictEqual(mdPart.content, 'Hi there!'); - assert.strictEqual(state!.turns[0].state, TurnState.Complete); - }); - - test('restores a session with tool calls', async () => { - const session = await agent.createSession(); - const sessions = await agent.listSessions(); - const sessionResource = sessions[0].session.toString(); - - agent.sessionMessages = [ - { type: 'message', session, role: 'user', messageId: 'msg-1', content: 'Run a command', toolRequests: [] }, - { type: 'message', session, role: 'assistant', messageId: 'msg-2', content: 'I will run a command.', toolRequests: [{ toolCallId: 'tc-1', name: 'shell' }] }, - { type: 'tool_start', session, toolCallId: 'tc-1', toolName: 'shell', displayName: 'Shell', invocationMessage: 'Running command...' }, - { type: 'tool_complete', session, toolCallId: 'tc-1', result: { success: true, pastTenseMessage: 'Ran command', content: [{ type: ToolResultContentType.Text, text: 'output' }] } }, - { type: 'message', session, role: 'assistant', messageId: 'msg-3', content: 'Done!', toolRequests: [] }, - ]; - - await sideEffects.handleRestoreSession(sessionResource); - - const state = stateManager.getSessionState(sessionResource); - assert.ok(state); - assert.strictEqual(state!.turns.length, 1); - - const turn = state!.turns[0]; - const toolCallParts = turn.responseParts.filter((p): p is IToolCallResponsePart => p.kind === ResponsePartKind.ToolCall); - assert.strictEqual(toolCallParts.length, 1); - const tc = toolCallParts[0].toolCall as IToolCallCompletedState; - assert.strictEqual(tc.status, ToolCallStatus.Completed); - assert.strictEqual(tc.toolCallId, 'tc-1'); - assert.strictEqual(tc.toolName, 'shell'); - assert.strictEqual(tc.displayName, 'Shell'); - assert.strictEqual(tc.success, true); - assert.strictEqual(tc.confirmed, ToolCallConfirmationReason.NotNeeded); - }); - - test('restores a session with multiple turns', async () => { - const session = await agent.createSession(); - const sessions = await agent.listSessions(); - const sessionResource = sessions[0].session.toString(); - - agent.sessionMessages = [ - { type: 'message', session, role: 'user', messageId: 'msg-1', content: 'First question', toolRequests: [] }, - { type: 'message', session, role: 'assistant', messageId: 'msg-2', content: 'First answer', toolRequests: [] }, - { type: 'message', session, role: 'user', messageId: 'msg-3', content: 'Second question', toolRequests: [] }, - { type: 'message', session, role: 'assistant', messageId: 'msg-4', content: 'Second answer', toolRequests: [] }, - ]; - - await sideEffects.handleRestoreSession(sessionResource); - - const state = stateManager.getSessionState(sessionResource); - assert.ok(state); - assert.strictEqual(state!.turns.length, 2); - assert.strictEqual(state!.turns[0].userMessage.text, 'First question'); - const mdPart0 = state!.turns[0].responseParts.find((p): p is IMarkdownResponsePart => p.kind === ResponsePartKind.Markdown); - assert.strictEqual(mdPart0?.content, 'First answer'); - assert.strictEqual(state!.turns[1].userMessage.text, 'Second question'); - const mdPart1 = state!.turns[1].responseParts.find((p): p is IMarkdownResponsePart => p.kind === ResponsePartKind.Markdown); - assert.strictEqual(mdPart1?.content, 'Second answer'); - }); - - test('flushes interrupted turns when user message arrives without closing assistant message', async () => { - const session = await agent.createSession(); - const sessions = await agent.listSessions(); - const sessionResource = sessions[0].session.toString(); - - agent.sessionMessages = [ - { type: 'message', session, role: 'user', messageId: 'msg-1', content: 'Interrupted question', toolRequests: [] }, - // No assistant message - the turn was interrupted - { type: 'message', session, role: 'user', messageId: 'msg-2', content: 'Retried question', toolRequests: [] }, - { type: 'message', session, role: 'assistant', messageId: 'msg-3', content: 'Answer', toolRequests: [] }, - ]; - - await sideEffects.handleRestoreSession(sessionResource); - - const state = stateManager.getSessionState(sessionResource); - assert.ok(state); - assert.strictEqual(state!.turns.length, 2); - assert.strictEqual(state!.turns[0].userMessage.text, 'Interrupted question'); - const mdPart0 = state!.turns[0].responseParts.find((p): p is IMarkdownResponsePart => p.kind === ResponsePartKind.Markdown); - assert.ok(!mdPart0 || mdPart0.content === '', 'interrupted turn should have empty response'); - assert.strictEqual(state!.turns[0].state, TurnState.Cancelled); - assert.strictEqual(state!.turns[1].userMessage.text, 'Retried question'); - const mdPart1 = state!.turns[1].responseParts.find((p): p is IMarkdownResponsePart => p.kind === ResponsePartKind.Markdown); - assert.strictEqual(mdPart1?.content, 'Answer'); - assert.strictEqual(state!.turns[1].state, TurnState.Complete); - }); - - test('is a no-op for a session already in the state manager', async () => { - setupSession(); - // Should not throw or create a duplicate - await sideEffects.handleRestoreSession(sessionUri.toString()); - assert.ok(stateManager.getSessionState(sessionUri.toString())); - }); - - test('throws when no agent found for session', async () => { - const noAgentSideEffects = disposables.add(new AgentSideEffects(stateManager, { - getAgent: () => undefined, - agents: observableValue('agents', []), - sessionDataService: {} as ISessionDataService, - }, new NullLogService(), fileService)); - - await assert.rejects( - () => noAgentSideEffects.handleRestoreSession('unknown://session-1'), - /No agent for session/, - ); - }); - - test('response parts include markdown segments', async () => { - const session = await agent.createSession(); - const sessions = await agent.listSessions(); - const sessionResource = sessions[0].session.toString(); - - agent.sessionMessages = [ - { type: 'message', session, role: 'user', messageId: 'msg-1', content: 'hello', toolRequests: [] }, - { type: 'message', session, role: 'assistant', messageId: 'msg-2', content: 'response text', toolRequests: [] }, - ]; - - await sideEffects.handleRestoreSession(sessionResource); - - const state = stateManager.getSessionState(sessionResource); - assert.ok(state); - assert.strictEqual(state!.turns[0].responseParts.length, 1); - assert.strictEqual(state!.turns[0].responseParts[0].kind, ResponsePartKind.Markdown); - assert.strictEqual(state!.turns[0].responseParts[0].content, 'response text'); - }); - - test('throws when session is not found on backend', async () => { - // Agent exists but session is not in listSessions - await assert.rejects( - () => sideEffects.handleRestoreSession(AgentSession.uri('mock', 'nonexistent').toString()), - /Session not found on backend/, - ); - }); - - test('preserves workingDirectory from agent metadata', async () => { - agent.sessionMetadataOverrides = { workingDirectory: URI.file('/home/user/project') }; - const session = await agent.createSession(); - const sessions = await agent.listSessions(); - const sessionResource = sessions[0].session.toString(); - - agent.sessionMessages = [ - { type: 'message', session, role: 'user', messageId: 'msg-1', content: 'hi', toolRequests: [] }, - { type: 'message', session, role: 'assistant', messageId: 'msg-2', content: 'hello', toolRequests: [] }, - ]; - - await sideEffects.handleRestoreSession(sessionResource); - - const state = stateManager.getSessionState(sessionResource); - assert.ok(state); - assert.strictEqual(state!.summary.workingDirectory, URI.file('/home/user/project').toString()); - }); - }); - - // ---- handleBrowseDirectory ------------------------------------------ - - suite('handleBrowseDirectory', () => { - - test('throws when the directory does not exist', async () => { - await assert.rejects( - () => sideEffects.handleBrowseDirectory(URI.from({ scheme: Schemas.inMemory, path: '/nonexistent' }).toString()), - /Directory not found/, - ); - }); - - test('throws when the target is not a directory', async () => { - await assert.rejects( - () => sideEffects.handleBrowseDirectory(URI.from({ scheme: Schemas.inMemory, path: '/testDir/file.txt' }).toString()), - /Not a directory/, - ); - }); - }); - // ---- agents observable -------------------------------------------------- suite('agents observable', () => { @@ -485,45 +219,6 @@ suite('AgentSideEffects', () => { }); }); - // ---- handleGetResourceMetadata / handleAuthenticate ----------------- - - suite('auth', () => { - - test('handleGetResourceMetadata aggregates resources from agents', () => { - agentList.set([agent], undefined); - - const metadata = sideEffects.handleGetResourceMetadata(); - assert.strictEqual(metadata.resources.length, 0, 'mock agent has no protected resources'); - }); - - test('handleGetResourceMetadata returns resources when agent declares them', () => { - const copilotAgent = new MockAgent('copilot'); - disposables.add(toDisposable(() => copilotAgent.dispose())); - agentList.set([copilotAgent], undefined); - - const metadata = sideEffects.handleGetResourceMetadata(); - assert.strictEqual(metadata.resources.length, 1); - assert.strictEqual(metadata.resources[0].resource, 'https://api.github.com'); - }); - - test('handleAuthenticate returns authenticated for matching resource', async () => { - const copilotAgent = new MockAgent('copilot'); - disposables.add(toDisposable(() => copilotAgent.dispose())); - agentList.set([copilotAgent], undefined); - - const result = await sideEffects.handleAuthenticate({ resource: 'https://api.github.com', token: 'test-token' }); - assert.deepStrictEqual(result, { authenticated: true }); - assert.deepStrictEqual(copilotAgent.authenticateCalls, [{ resource: 'https://api.github.com', token: 'test-token' }]); - }); - - test('handleAuthenticate returns not authenticated for non-matching resource', async () => { - agentList.set([agent], undefined); - - const result = await sideEffects.handleAuthenticate({ resource: 'https://unknown.example.com', token: 'test-token' }); - assert.deepStrictEqual(result, { authenticated: false }); - }); - }); - // ---- Pending message sync ----------------------------------------------- suite('pending message sync', () => { diff --git a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts index f4e2fae680e..c12af555f04 100644 --- a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts +++ b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts @@ -9,13 +9,14 @@ 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 { IAgentCreateSessionConfig, IAgentDescriptor, IAgentService, IAgentSessionMetadata, IAuthenticateParams, IAuthenticateResult, IResourceMetadata } from '../../common/agentService.js'; import { IFetchContentResult } from '../../common/state/protocol/commands.js'; import { ActionType, type ISessionAction } from '../../common/state/sessionActions.js'; import { PROTOCOL_VERSION } from '../../common/state/sessionCapabilities.js'; -import { isJsonRpcNotification, isJsonRpcResponse, JSON_RPC_INTERNAL_ERROR, ProtocolError, type IAhpNotification, type ICreateSessionParams, type IInitializeResult, type IProtocolMessage, type IReconnectResult, type IStateSnapshot } from '../../common/state/sessionProtocol.js'; +import { isJsonRpcNotification, isJsonRpcResponse, JSON_RPC_INTERNAL_ERROR, ProtocolError, type IAhpNotification, type IBrowseDirectoryResult, type IInitializeResult, type IProtocolMessage, type IReconnectResult, type IStateSnapshot, type IWriteFileParams, type IWriteFileResult } from '../../common/state/sessionProtocol.js'; import { SessionStatus, type ISessionSummary } from '../../common/state/sessionState.js'; import type { IProtocolServer, IProtocolTransport } from '../../common/state/sessionTransport.js'; -import { ProtocolServerHandler, type IProtocolSideEffectHandler } from '../../node/protocolServerHandler.js'; +import { ProtocolServerHandler } from '../../node/protocolServerHandler.js'; import { SessionStateManager } from '../../node/sessionStateManager.js'; // ---- Mock helpers ----------------------------------------------------------- @@ -64,24 +65,49 @@ class MockProtocolServer implements IProtocolServer { } } -class MockSideEffectHandler implements IProtocolSideEffectHandler { +class MockAgentService implements IAgentService { + declare readonly _serviceBrand: undefined; readonly handledActions: ISessionAction[] = []; readonly browsedUris: URI[] = []; readonly browseErrors = new Map(); - handleAction(action: ISessionAction): void { - this.handledActions.push(action); + private readonly _onDidAction = new Emitter(); + readonly onDidAction = this._onDidAction.event; + private readonly _onDidNotification = new Emitter(); + readonly onDidNotification = this._onDidNotification.event; + + private _stateManager!: SessionStateManager; + + /** Connect to the state manager so dispatchAction works correctly. */ + setStateManager(sm: SessionStateManager): void { + this._stateManager = sm; } - async handleCreateSession(_command: ICreateSessionParams): Promise { /* session created via state manager */ } - handleDisposeSession(_session: string): void { } - async handleListSessions(): Promise { return []; } - async handleRestoreSession(_session: string): Promise { } - handleGetResourceMetadata() { return { resources: [] }; } - handleWriteFile() { return Promise.resolve({}); } - async handleAuthenticate(_params: { resource: string; token: string }) { return { authenticated: true }; } - async handleBrowseDirectory(uri: string): Promise<{ entries: { name: string; type: 'file' | 'directory' }[] }> { - this.browsedUris.push(URI.parse(uri)); - const error = this.browseErrors.get(uri); + + dispatchAction(action: ISessionAction, clientId: string, clientSeq: number): void { + this.handledActions.push(action); + const origin = { clientId, clientSeq }; + this._stateManager.dispatchClientAction(action, origin); + } + async createSession(_config?: IAgentCreateSessionConfig): Promise { return URI.parse('copilot:///new-session'); } + async disposeSession(_session: URI): Promise { } + async listSessions(): Promise { return []; } + async subscribe(resource: URI): Promise { + const snapshot = this._stateManager.getSnapshot(resource.toString()); + if (!snapshot) { + throw new Error(`Cannot subscribe to unknown resource: ${resource.toString()}`); + } + return snapshot; + } + unsubscribe(_resource: URI): void { } + async shutdown(): Promise { } + async getResourceMetadata(): Promise { return { resources: [] }; } + async authenticate(_params: IAuthenticateParams): Promise { return { authenticated: true }; } + async refreshModels(): Promise { } + async listAgents(): Promise { return []; } + async writeFile(_params: IWriteFileParams): Promise { return {}; } + async browseDirectory(uri: URI): Promise { + this.browsedUris.push(uri); + const error = this.browseErrors.get(uri.toString()); if (error) { throw error; } @@ -92,12 +118,14 @@ class MockSideEffectHandler implements IProtocolSideEffectHandler { ], }; } - getDefaultDirectory(): string { - return URI.file('/home/testuser').toString(); - } - async handleFetchContent(_uri: string): Promise { + async fetchContent(_uri: URI): Promise { throw new Error('Not implemented'); } + + dispose(): void { + this._onDidAction.dispose(); + this._onDidNotification.dispose(); + } } // ---- Helpers ---------------------------------------------------------------- @@ -129,7 +157,7 @@ suite('ProtocolServerHandler', () => { let disposables: DisposableStore; let stateManager: SessionStateManager; let server: MockProtocolServer; - let sideEffects: MockSideEffectHandler; + let agentService: MockAgentService; let handler: ProtocolServerHandler; const sessionUri = URI.from({ scheme: 'copilot', path: '/test-session' }).toString(); @@ -160,11 +188,14 @@ suite('ProtocolServerHandler', () => { disposables = new DisposableStore(); stateManager = disposables.add(new SessionStateManager(new NullLogService())); server = disposables.add(new MockProtocolServer()); - sideEffects = new MockSideEffectHandler(); + agentService = new MockAgentService(); + agentService.setStateManager(stateManager); + disposables.add(agentService); disposables.add(handler = new ProtocolServerHandler( + agentService, stateManager, server, - sideEffects, + { defaultDirectory: URI.file('/home/testuser').toString() }, new NullLogService(), )); }); @@ -362,8 +393,8 @@ suite('ProtocolServerHandler', () => { transport.simulateMessage(request(2, 'browseDirectory', { uri: dirUri })); const resp = await responsePromise; - assert.strictEqual(sideEffects.browsedUris.length, 1); - assert.strictEqual(sideEffects.browsedUris[0].path, '/home/user/project'); + assert.strictEqual(agentService.browsedUris.length, 1); + assert.strictEqual(agentService.browsedUris[0].path, '/home/user/project'); assert.ok(resp); const result = (resp as unknown as { result: { entries: { name: string; uri: unknown; type: string }[] } }).result; @@ -379,7 +410,7 @@ suite('ProtocolServerHandler', () => { transport.sent.length = 0; const dirUri = URI.file('/missing').toString(); - sideEffects.browseErrors.set(dirUri, new ProtocolError(JSON_RPC_INTERNAL_ERROR, `Directory not found: ${dirUri}`)); + agentService.browseErrors.set(URI.file('/missing').toString(), new ProtocolError(JSON_RPC_INTERNAL_ERROR, `Directory not found: ${dirUri}`)); const responsePromise = waitForResponse(transport, 2); transport.simulateMessage(request(2, 'browseDirectory', { uri: dirUri })); const resp = await responsePromise as { error?: { code: number; message: string } }; @@ -416,9 +447,9 @@ suite('ProtocolServerHandler', () => { }); test('extension request preserves ProtocolError code and data', async () => { - // Override handleAuthenticate to throw a ProtocolError with data - const origHandler = sideEffects.handleAuthenticate; - sideEffects.handleAuthenticate = async () => { throw new ProtocolError(-32007, 'Auth required', { hint: 'sign in' }); }; + // Override authenticate to throw a ProtocolError with data + const origHandler = agentService.authenticate; + agentService.authenticate = async () => { throw new ProtocolError(-32007, 'Auth required', { hint: 'sign in' }); }; const transport = connectClient('client-auth-error'); transport.sent.length = 0; @@ -432,7 +463,7 @@ suite('ProtocolServerHandler', () => { assert.strictEqual(resp.error!.message, 'Auth required'); assert.deepStrictEqual(resp.error!.data, { hint: 'sign in' }); - sideEffects.handleAuthenticate = origHandler; + agentService.authenticate = origHandler; }); // ---- Connection count event -----------------------------------------