Microsoft/vscode
1
0
Fork 0
mirror of https://github.com/microsoft/vscode.git synced 2026-04-17 15:24:40 +01:00
Code Issues Packages Projects Releases Wiki Activity

wip

Browse Source
This commit is contained in:
Connor Peet
2026-03-31 09:14:16 -07:00
parent 1dadacc629
commit b7e5e30020
12 changed files with 815 additions and 1738 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/vs/platform/agentHost/common/agentService.ts
View File

@@ -358,12 +358,6 @@ export interface IAgent {
*/
authenticate(resource: string, token: string): Promise<boolean>;
/**
* Returns customizations (Open Plugins) this agent knows about.
* These are published in {@link IAgentInfo.customizations} in root state.
*/
getCustomizations?(): ICustomizationRef[];
/**
* Receives client-provided customization refs and syncs them (e.g. copies
* plugin files to local storage). Returns per-customization status with
@@ -371,13 +365,13 @@ export interface IAgent {
*
* The agent MAY defer a client restart until all active sessions are idle.
*/
setClientCustomizations?(clientId: string, customizations: ICustomizationRef[], progress?: (results: ISyncedCustomization[]) => void): Promise<ISyncedCustomization[]>;
setClientCustomizations(clientId: string, customizations: ICustomizationRef[], progress?: (results: ISyncedCustomization[]) => void): Promise<ISyncedCustomization[]>;
/**
* Notifies the agent that a customization has been toggled on or off.
* The agent MAY restart its client before the next message is sent.
*/
setCustomizationEnabled?(uri: string, enabled: boolean): void;
setCustomizationEnabled(uri: string, enabled: boolean): void;
/** Gracefully shut down all sessions. */
shutdown(): Promise<void>;

14
src/vs/platform/agentHost/node/agentHostMain.ts
View File

@@ -34,6 +34,10 @@ import { InstantiationService } from '../../instantiation/common/instantiationSe
import { ServiceCollection } from '../../instantiation/common/serviceCollection.js';
import { SessionDataService } from './sessionDataService.js';
import { ISessionDataService } from '../common/sessionDataService.js';
import { AgentHostClientFileSystemProvider } from '../common/agentHostClientFileSystemProvider.js';
import { AGENT_CLIENT_SCHEME } from '../common/agentClientUri.js';
import { IAgentPluginManager } from '../common/agentPluginManager.js';
import { AgentPluginManager } from './agentPluginManager.js';
// Entry point for the agent host utility process.
// Sets up IPC, logging, and registers agent providers (Copilot).
@@ -73,10 +77,12 @@ function startAgentHost(): void {
let agentService: AgentService;
try {
agentService = new AgentService(logService, fileService, sessionDataService);
const pluginManager = new AgentPluginManager(URI.file(environmentService.userDataPath), fileService, logService);
const diServices = new ServiceCollection();
diServices.set(ILogService, logService);
diServices.set(IFileService, fileService);
diServices.set(ISessionDataService, sessionDataService);
diServices.set(IAgentPluginManager, pluginManager);
const instantiationService = new InstantiationService(diServices);
agentService.registerProvider(instantiationService.createInstance(CopilotAgent));
} catch (err) {
@@ -97,7 +103,7 @@ function startAgentHost(): void {
server.registerChannel(AgentHostIpcChannels.ConnectionTracker, connectionTrackerChannel);
// Start WebSocket server for external clients if configured
startWebSocketServer(agentService, logService, disposables, count => connectionCountEmitter.fire(count)).catch(err => {
startWebSocketServer(agentService, fileService, logService, disposables, count => connectionCountEmitter.fire(count)).catch(err => {
logService.error('Failed to start WebSocket server', err);
});
@@ -114,7 +120,7 @@ function startAgentHost(): void {
* This reuses the same {@link AgentService} and {@link SessionStateManager}
* that the IPC channel uses, so both IPC and WebSocket clients share state.
*/
async function startWebSocketServer(agentService: AgentService, logService: ILogService, disposables: DisposableStore, onConnectionCountChanged: (count: number) => void): Promise<void> {
async function startWebSocketServer(agentService: AgentService, fileService: IFileService, logService: ILogService, disposables: DisposableStore, onConnectionCountChanged: (count: number) => void): Promise<void> {
const port = process.env['VSCODE_AGENT_HOST_PORT'];
const socketPath = process.env['VSCODE_AGENT_HOST_SOCKET_PATH'];
@@ -143,11 +149,15 @@ async function startWebSocketServer(agentService: AgentService, logService: ILog
logService,
));
const clientFileSystemProvider = disposables.add(new AgentHostClientFileSystemProvider());
disposables.add(fileService.registerProvider(AGENT_CLIENT_SCHEME, clientFileSystemProvider));
const protocolHandler = disposables.add(new ProtocolServerHandler(
agentService,
agentService.stateManager,
wsServer,
{ defaultDirectory: URI.file(os.homedir()).toString() },
clientFileSystemProvider,
logService,
));
disposables.add(protocolHandler.onDidChangeConnectionCount(onConnectionCountChanged));

7
src/vs/platform/agentHost/node/agentHostServerMain.ts
View File

@@ -39,6 +39,8 @@ import { DiskFileSystemProvider } from '../../files/node/diskFileSystemProvider.
import { Schemas } from '../../../base/common/network.js';
import { ISessionDataService } from '../common/sessionDataService.js';
import { SessionDataService } from './sessionDataService.js';
import { AgentHostClientFileSystemProvider } from '../common/agentHostClientFileSystemProvider.js';
import { AGENT_CLIENT_SCHEME } from '../common/agentClientUri.js';
/** Log to stderr so messages appear in the terminal alongside the process. */
function log(msg: string): void {
@@ -182,12 +184,17 @@ async function main(): Promise<void> {
: undefined,
}, logService));
const clientFileSystemProvider = disposables.add(new AgentHostClientFileSystemProvider());
disposables.add(fileService.registerProvider(AGENT_CLIENT_SCHEME, clientFileSystemProvider));
// Wire up protocol handler
disposables.add(new ProtocolServerHandler(
agentService,
agentService.stateManager,
wsServer,
{ defaultDirectory: URI.file(os.homedir()).toString() },
clientFileSystemProvider,
logService,
));

2
src/vs/platform/agentHost/node/agentService.ts
View File

@@ -205,7 +205,7 @@ export class AgentService extends Disposable implements IAgentService {
const state = this._stateManager.dispatchClientAction(action, origin);
this._logService.trace(`[AgentService] resulting state:`, state);
this._sideEffects.handleAction(action, clientId);
this._sideEffects.handleAction(action);
}
async browseDirectory(uri: URI): Promise<IBrowseDirectoryResult> {

424
src/vs/platform/agentHost/node/agentSideEffects.ts
View File

@@ -3,34 +3,23 @@
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/
import * as os from 'os';
import { match as globMatch } from '../../../base/common/glob.js';
import { Disposable, DisposableStore, IDisposable } from '../../../base/common/lifecycle.js';
import { VSBuffer } from '../../../base/common/buffer.js';
import { autorun, IObservable } from '../../../base/common/observable.js';
import { URI } from '../../../base/common/uri.js';
import { generateUuid } from '../../../base/common/uuid.js';
import { IFileService } 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, JSON_RPC_INTERNAL_ERROR, ProtocolError } from '../common/state/sessionProtocol.js';
import {
CustomizationStatus,
PendingMessageKind,
ResponsePartKind,
SessionStatus,
ToolCallConfirmationReason,
ToolCallStatus,
TurnState,
type IResponsePart,
type ISessionCustomization,
type ISessionModelInfo,
type ISessionSummary,
type IToolCallCompletedState,
type ITurn,
type URI as ProtocolURI,
} from '../common/state/sessionState.js';
import { AgentEventMapper } from './agentEventMapper.js';
import type { IProtocolSideEffectHandler } from './protocolServerHandler.js';
import { SessionStateManager } from './sessionStateManager.js';
/**
@@ -48,14 +37,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<string, string>();
@@ -66,7 +55,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();
@@ -94,12 +82,42 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH
} catch {
models = [];
}
const customizations = a.getCustomizations?.();
return { provider: d.provider, displayName: d.displayName, description: d.description, models, customizations };
return { provider: d.provider, displayName: d.displayName, description: d.description, models };
}));
this._stateManager.dispatchServerAction({ type: ActionType.RootAgentsChanged, agents: infos });
}
// ---- Edit auto-approve --------------------------------------------------
/**
* Default edit auto-approve patterns applied by the agent host.
* Matches the VS Code `chat.tools.edits.autoApprove` setting defaults.
*/
private static readonly _DEFAULT_EDIT_AUTO_APPROVE_PATTERNS: Readonly<Record<string, boolean>> = {
'**/*': true,
'**/.vscode/*.json': false,
'**/.git/**': false,
'**/{package.json,server.xml,build.rs,web.config,.gitattributes,.env}': false,
'**/*.{code-workspace,csproj,fsproj,vbproj,vcxproj,proj,targets,props}': false,
'**/*.lock': false,
'**/*-lock.{yaml,json}': false,
};
/**
* Returns whether a write to `filePath` should be auto-approved based on
* the built-in default patterns.
*/
private _shouldAutoApproveEdit(filePath: string): boolean {
const patterns = AgentSideEffects._DEFAULT_EDIT_AUTO_APPROVE_PATTERNS;
let approved = true;
for (const [pattern, isApproved] of Object.entries(patterns)) {
if (isApproved !== approved && globMatch(pattern, filePath)) {
approved = isApproved;
}
}
return approved;
}
// ---- Agent registration -------------------------------------------------
/**
@@ -125,6 +143,16 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH
const sessionKey = e.session.toString();
const turnId = this._stateManager.getActiveTurnId(sessionKey);
if (turnId) {
// Check if this is a write permission request that can be auto-approved
// based on the built-in default patterns.
if (e.type === 'tool_ready' && e.permissionKind === 'write' && e.permissionPath) {
if (this._shouldAutoApproveEdit(e.permissionPath)) {
this._logService.trace(`[AgentSideEffects] Auto-approving write to ${e.permissionPath}`);
agent.respondToPermissionRequest(e.toolCallId, true);
return;
}
}
const actions = agentMapper.mapProgressEventToActions(e, sessionKey, turnId);
if (actions) {
if (Array.isArray(actions)) {
@@ -145,7 +173,7 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH
return disposables;
}
// ---- IProtocolSideEffectHandler -----------------------------------------
// ---- Side-effect handlers --------------------------------------------------
handleAction(action: ISessionAction): void {
switch (action.type) {
@@ -206,22 +234,52 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH
});
break;
}
case ActionType.SessionPendingMessageSet:
case ActionType.SessionPendingMessageRemoved:
case ActionType.SessionQueuedMessagesReordered: {
this._syncPendingMessages(action.session);
break;
}
case ActionType.SessionActiveClientChanged: {
const customizations = action.activeClient?.customizations;
if (customizations && customizations.length > 0) {
const agent = this._options.getAgent(action.session);
if (agent?.setClientCustomizations) {
agent.setClientCustomizations(action.activeClient!.clientId, customizations, results => {
this._stateManager.dispatchServerAction({
type: ActionType.SessionCustomizationsChanged,
session: action.session,
customizations: results.map(r => r.customization),
});
}).catch(err => {
this._logService.error('[AgentSideEffects] setClientCustomizations failed', err);
});
}
const agent = this._options.getAgent(action.session);
const refs = action.activeClient?.customizations;
if (!agent?.setClientCustomizations || !refs?.length) {
break;
}
// Publish initial "loading" status for all customizations
const loading: ISessionCustomization[] = refs.map(r => ({
customization: r,
enabled: true,
status: CustomizationStatus.Loading,
}));
this._stateManager.dispatchServerAction({
type: ActionType.SessionCustomizationsChanged,
session: action.session,
customizations: loading,
});
agent.setClientCustomizations(
action.activeClient!.clientId,
refs,
(synced) => {
// Incremental progress: publish updated statuses
const statuses: ISessionCustomization[] = synced.map(s => s.customization);
this._stateManager.dispatchServerAction({
type: ActionType.SessionCustomizationsChanged,
session: action.session,
customizations: statuses,
});
},
).then(synced => {
// Final status
const statuses: ISessionCustomization[] = synced.map(s => s.customization);
this._stateManager.dispatchServerAction({
type: ActionType.SessionCustomizationsChanged,
session: action.session,
customizations: statuses,
});
}).catch(err => {
this._logService.error('[AgentSideEffects] setClientCustomizations failed', err);
});
break;
}
case ActionType.SessionCustomizationToggled: {
@@ -229,12 +287,6 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH
agent?.setCustomizationEnabled?.(action.uri, action.enabled);
break;
}
case ActionType.SessionPendingMessageSet:
case ActionType.SessionPendingMessageRemoved:
case ActionType.SessionQueuedMessagesReordered: {
this._syncPendingMessages(action.session);
break;
}
}
}
@@ -331,290 +383,6 @@ export class AgentSideEffects extends Disposable implements IProtocolSideEffectH
});
}
async handleCreateSession(command: ICreateSessionParams): Promise<void> {
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,
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.removeSession(session);
this._options.sessionDataService.deleteSessionData(URI.parse(session));
}
async handleListSessions(): Promise<ISessionSummary[]> {
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<void> {
// 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,
};
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<string, IAgentToolStartEvent>;
} | undefined;
let turnCounter = 0;
const finalizeTurn = (turn: NonNullable<typeof currentTurn>, state: TurnState): void => {
turns.push({
id: turn.id,
userMessage: turn.userMessage,
responseParts: turn.responseParts,
usage: undefined,
state,
});
};
const startTurn = (text: string): NonNullable<typeof currentTurn> => ({
id: `restored-${turnCounter++}`,
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.content);
} else if (msg.type === 'message' && msg.role === 'assistant') {
if (!currentTurn) {
currentTurn = startTurn('');
}
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<IAuthenticateResult> {
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<IBrowseDirectoryResult> {
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<IFetchContentResult> {
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(uri: ProtocolURI, data: string, encoding: ContentEncoding): Promise<void> {
try {
const content = encoding === ContentEncoding.Base64
? VSBuffer.wrap(Buffer.from(data, 'base64'))
: VSBuffer.fromString(data);
await this._fileService.writeFile(URI.parse(uri), content);
} catch (_e) {
throw new ProtocolError(AhpErrorCodes.NotFound, `Failed to write file: ${uri}`);
}
}
override dispose(): void {
this._toolCallAgents.clear();
super.dispose();

872
src/vs/platform/agentHost/node/copilot/copilotAgent.ts
View File

File diff suppressed because it is too large Load Diff

227
src/vs/platform/agentHost/node/protocolServerHandler.ts
View File

@@ -4,35 +4,32 @@
*--------------------------------------------------------------------------------------------*/
import { Emitter } from '../../../base/common/event.js';
import { isJsonRpcResponse } from '../../../base/common/jsonRpcProtocol.js';
import { Disposable, DisposableStore } from '../../../base/common/lifecycle.js';
import { ILogService } from '../../log/common/log.js';
import { hasKey } from '../../../base/common/types.js';
import { URI } from '../../../base/common/uri.js';
import { ILogService } from '../../log/common/log.js';
import { AHPFileSystemProvider } from '../common/agentHostFileSystemProvider.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,
ContentEncoding,
IJsonRpcRequest,
isJsonRpcNotification,
isJsonRpcRequest,
isJsonRpcResponse,
JSON_RPC_INTERNAL_ERROR,
ProtocolError,
type IAhpServerNotification,
type IBrowseDirectoryResult,
type ICreateSessionParams,
type IFetchContentResult,
type IInitializeParams,
type IJsonRpcRequest as IJsonRpcRequestType,
type IJsonRpcResponse,
type IProtocolMessage,
type IReconnectParams,
type IStateSnapshot,
} 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';
@@ -84,9 +81,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 {
@@ -100,9 +105,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,
private readonly _clientFileSystemProvider: AHPFileSystemProvider,
@ILogService private readonly _logService: ILogService,
) {
@@ -173,17 +179,21 @@ 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;
}
}
// Handle reverse RPC responses from the client
if (this._handleReverseResponse(msg)) {
return;
} else if (isJsonRpcResponse(msg)) {
const pending = this._pendingReverseRequests.get(msg.id);
if (pending) {
this._pendingReverseRequests.delete(msg.id);
if (hasKey(msg, { error: true })) {
pending.reject(new Error(msg.error?.message ?? 'Reverse RPC error'));
} else {
pending.resolve(msg.result);
}
}
}
}));
@@ -230,6 +240,7 @@ export class ProtocolServerHandler extends Disposable {
fetchContent: (uri) => this._sendReverseRequest(params.clientId, 'fetchContent', { uri: uri.toString() }),
}));
const snapshots: IStateSnapshot[] = [];
if (params.initialSubscriptions) {
for (const uri of params.initialSubscriptions) {
@@ -247,7 +258,7 @@ export class ProtocolServerHandler extends Disposable {
protocolVersion: PROTOCOL_VERSION,
serverSeq: this._stateManager.serverSeq,
snapshots,
defaultDirectory: this._sideEffectHandler.getDefaultDirectory?.(),
defaultDirectory: this._config.defaultDirectory,
},
};
}
@@ -269,12 +280,6 @@ export class ProtocolServerHandler extends Disposable {
this._clients.set(params.clientId, client);
this._onDidChangeConnectionCount.fire(this._clients.size);
// Register the client's filesystem connection for reverse RPC access
disposables.add(this._clientFileSystemProvider.registerAuthority(params.clientId, {
browseDirectory: (uri) => this._sendReverseRequest(params.clientId, 'browseDirectory', { uri: uri.toString() }),
fetchContent: (uri) => this._sendReverseRequest(params.clientId, 'fetchContent', { uri: uri.toString() }),
}));
const oldestBuffered = this._replayBuffer.length > 0 ? this._replayBuffer[0].serverSeq : this._stateManager.serverSeq;
const canReplay = params.lastSeenServerSeq >= oldestBuffered;
@@ -304,48 +309,6 @@ export class ProtocolServerHandler extends Disposable {
}
}
// ---- Reverse RPC (server → client requests) ----------------------------
private _reverseRequestId = 0;
private readonly _pendingReverseRequests = new Map<number, { resolve: (value: unknown) => void; reject: (reason: unknown) => void }>();
/**
* Sends a JSON-RPC request to a connected client and waits for the response.
* Used for reverse-RPC operations like reading client-side files.
*/
private _sendReverseRequest<T>(clientId: string, method: string, params: unknown): Promise<T> {
const client = this._clients.get(clientId);
if (!client) {
return Promise.reject(new Error(`Client ${clientId} is not connected`));
}
const id = ++this._reverseRequestId;
return new Promise<T>((resolve, reject) => {
this._pendingReverseRequests.set(id, { resolve: resolve as (value: unknown) => void, reject });
const request: IJsonRpcRequestType = { jsonrpc: '2.0', id, method, params };
client.transport.send(request);
});
}
/**
* Called when a JSON-RPC response arrives from a client (reverse RPC result).
*/
private _handleReverseResponse(msg: IProtocolMessage): boolean {
if (!isJsonRpcResponse(msg)) {
return false;
}
const pending = this._pendingReverseRequests.get(msg.id);
if (!pending) {
return false;
}
this._pendingReverseRequests.delete(msg.id);
if (hasKey(msg, { error: true })) {
pending.reject(new Error(msg.error?.message ?? 'Reverse RPC error'));
} else {
pending.resolve(msg.result);
}
return true;
}
// ---- Requests (expect a response) ---------------------------------------
/**
@@ -354,30 +317,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._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) => {
@@ -403,17 +392,36 @@ 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));
},
writeFile: async (_client, params) => {
await this._sideEffectHandler.handleWriteFile(params.uri, params.data, params.encoding);
return {};
}
};
// ---- Reverse RPC (server → client requests) ----------------------------
private _reverseRequestId = 0;
private readonly _pendingReverseRequests = new Map<number, { resolve: (value: unknown) => void; reject: (reason: unknown) => void }>();
/**
* Sends a JSON-RPC request to a connected client and waits for the response.
* Used for reverse-RPC operations like reading client-side files.
*/
private _sendReverseRequest<T>(clientId: string, method: string, params: unknown): Promise<T> {
const client = this._clients.get(clientId);
if (!client) {
return Promise.reject(new Error(`Client ${clientId} is not connected`));
}
const id = ++this._reverseRequestId;
return new Promise<T>((resolve, reject) => {
this._pendingReverseRequests.set(id, { resolve: resolve as (value: unknown) => void, reject });
const request: IJsonRpcRequest = { jsonrpc: '2.0', id, method, params };
client.transport.send(request);
});
}
private _handleRequest(client: IConnectedClient, method: string, params: unknown, id: number): void {
const handler = this._requestHandlers.hasOwnProperty(method) ? this._requestHandlers[method as RequestMethod] : undefined;
if (handler) {
@@ -450,15 +458,20 @@ export class ProtocolServerHandler extends Disposable {
private _handleExtensionRequest(method: string, params: unknown): Promise<unknown> | 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;
}
@@ -503,29 +516,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<void>;
handleDisposeSession(session: URI): void;
handleListSessions(): Promise<ISessionSummary[]>;
/** Restore a session from a previous server lifetime into the state manager. */
handleRestoreSession(session: URI): Promise<void>;
handleGetResourceMetadata(): IResourceMetadata;
handleAuthenticate(params: IAuthenticateParams): Promise<IAuthenticateResult>;
handleBrowseDirectory(uri: URI): Promise<IBrowseDirectoryResult>;
handleFetchContent(uri: URI): Promise<IFetchContentResult>;
handleWriteFile(uri: URI, data: string, encoding: ContentEncoding): Promise<void>;
/** Returns the server's default browsing directory, if available. */
getDefaultDirectory?(): URI;
/** Refresh models from all providers (VS Code extension method). */
handleRefreshModels?(): Promise<void>;
/** List agent descriptors (VS Code extension method). */
handleListAgents?(): IAgentDescriptor[];
/** Shut down all providers (VS Code extension method). */
handleShutdown?(): Promise<void>;
}

550
src/vs/platform/agentHost/test/node/agentSideEffects.test.ts
View File

@@ -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 ICustomizationRef, type IMarkdownResponsePart, type ISessionCustomization, type IToolCallCompletedState, type IToolCallResponsePart } from '../../common/state/sessionState.js';
import { PendingMessageKind, SessionStatus } from '../../common/state/sessionState.js';
import { AgentSideEffects } from '../../node/agentSideEffects.js';
import { SessionStateManager } from '../../node/sessionStateManager.js';
import { MockAgent } from './mockAgent.js';
@@ -74,10 +74,11 @@ suite('AgentSideEffects', () => {
_serviceBrand: undefined,
getSessionDataDir: () => URI.from({ scheme: Schemas.inMemory, path: '/session-data' }),
getSessionDataDirById: () => URI.from({ scheme: Schemas.inMemory, path: '/session-data' }),
openDatabase: () => { throw new Error('not implemented'); },
deleteSessionData: async () => { },
cleanupOrphanedData: async () => { },
} satisfies ISessionDataService,
}, new NullLogService(), fileService));
}, new NullLogService()));
});
teardown(() => {
@@ -112,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)));
@@ -200,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<readonly IAgent[]>('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<readonly IAgent[]>('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: '/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, '/home/user/project');
});
});
// ---- 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', () => {
@@ -484,186 +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 });
});
});
// ---- handleAction: session/activeClientChanged ----------------------
suite('handleAction — session/activeClientChanged', () => {
test('calls setClientCustomizations when active client has customizations', async () => {
setupSession();
const customizations: ICustomizationRef[] = [{
uri: 'https://plugins.example.com/my-plugin',
displayName: 'My Plugin',
nonce: 'abc',
}];
sideEffects.handleAction({
type: ActionType.SessionActiveClientChanged,
session: sessionUri.toString(),
activeClient: {
clientId: 'test-client',
displayName: 'Test Client',
tools: [],
customizations,
},
});
await new Promise(r => setTimeout(r, 10));
assert.deepStrictEqual(agent.setClientCustomizationsCalls, [customizations]);
});
test('dispatches customizationsChanged via progress callback', async () => {
setupSession();
const customizations: ICustomizationRef[] = [{
uri: 'https://plugins.example.com/plugin-a',
displayName: 'Plugin A',
}];
const envelopes: IActionEnvelope[] = [];
disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e)));
sideEffects.handleAction({
type: ActionType.SessionActiveClientChanged,
session: sessionUri.toString(),
activeClient: {
clientId: 'test-client',
tools: [],
customizations,
},
});
await new Promise(r => setTimeout(r, 50));
const customizationActions = envelopes.filter(e => e.action.type === ActionType.SessionCustomizationsChanged);
assert.ok(customizationActions.length >= 1, 'should dispatch at least one customizationsChanged');
const lastAction = customizationActions[customizationActions.length - 1];
assert.strictEqual(
(lastAction.action as { customizations: ISessionCustomization[] }).customizations.length,
1,
);
});
test('does not dispatch loading and does not call setClientCustomizations when no customizations provided', async () => {
setupSession();
const envelopes: IActionEnvelope[] = [];
disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e)));
sideEffects.handleAction({
type: ActionType.SessionActiveClientChanged,
session: sessionUri.toString(),
activeClient: {
clientId: 'test-client',
tools: [],
},
});
await new Promise(r => setTimeout(r, 10));
assert.strictEqual(agent.setClientCustomizationsCalls.length, 0);
const customizationActions = envelopes.filter(e => e.action.type === ActionType.SessionCustomizationsChanged);
assert.strictEqual(customizationActions.length, 0, 'should not dispatch customizationsChanged');
});
test('does not call setClientCustomizations when active client is null', async () => {
setupSession();
sideEffects.handleAction({
type: ActionType.SessionActiveClientChanged,
session: sessionUri.toString(),
activeClient: null,
});
await new Promise(r => setTimeout(r, 10));
assert.strictEqual(agent.setClientCustomizationsCalls.length, 0);
});
});
// ---- handleAction: session/customizationToggled ---------------------
suite('handleAction — session/customizationToggled', () => {
test('calls setCustomizationEnabled on the agent', () => {
setupSession();
sideEffects.handleAction({
type: ActionType.SessionCustomizationToggled,
session: sessionUri.toString(),
uri: 'https://plugins.example.com/toggle-me',
enabled: false,
});
assert.deepStrictEqual(agent.setCustomizationEnabledCalls, [{
uri: 'https://plugins.example.com/toggle-me',
enabled: false,
}]);
});
});
// ---- _publishAgentInfos (customizations in root state) --------------
suite('publishAgentInfos includes customizations', () => {
test('includes agent customizations in root/agentsChanged', async () => {
agent.customizations = [{
uri: 'https://plugins.example.com/root-plugin',
displayName: 'Root Plugin',
}];
const envelopes: IActionEnvelope[] = [];
disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e)));
agentList.set([agent], undefined);
// Wait for the async _publishAgentInfos to fire
await new Promise(r => setTimeout(r, 50));
const agentsChanged = envelopes.find(e => e.action.type === ActionType.RootAgentsChanged);
assert.ok(agentsChanged, 'should dispatch root/agentsChanged');
const agents = (agentsChanged!.action as { agents: { customizations?: unknown[] }[] }).agents;
assert.ok(agents[0].customizations);
assert.strictEqual(agents[0].customizations!.length, 1);
});
});
// ---- Pending message sync -----------------------------------------------
suite('pending message sync', () => {
@@ -863,4 +418,101 @@ suite('AgentSideEffects', () => {
assert.strictEqual(state?.steeringMessage, undefined);
});
});
// ---- handleAction: session/activeClientChanged ----------------------
suite('handleAction — session/activeClientChanged', () => {
test('calls setClientCustomizations and dispatches customizationsChanged', async () => {
setupSession();
const envelopes: IActionEnvelope[] = [];
disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e)));
const action: ISessionAction = {
type: ActionType.SessionActiveClientChanged,
session: sessionUri.toString(),
activeClient: {
clientId: 'test-client',
tools: [],
customizations: [
{ uri: 'file:///plugin-a', displayName: 'Plugin A' },
{ uri: 'file:///plugin-b', displayName: 'Plugin B' },
],
},
};
sideEffects.handleAction(action);
// Wait for async setClientCustomizations
await new Promise(r => setTimeout(r, 50));
assert.deepStrictEqual(agent.setClientCustomizationsCalls, [{
clientId: 'test-client',
customizations: [
{ uri: 'file:///plugin-a', displayName: 'Plugin A' },
{ uri: 'file:///plugin-b', displayName: 'Plugin B' },
],
}]);
const customizationActions = envelopes
.filter(e => e.action.type === ActionType.SessionCustomizationsChanged);
assert.ok(customizationActions.length >= 1, 'should dispatch at least one customizationsChanged');
});
test('skips when activeClient has no customizations', () => {
setupSession();
const envelopes: IActionEnvelope[] = [];
disposables.add(stateManager.onDidEmitEnvelope(e => envelopes.push(e)));
const action: ISessionAction = {
type: ActionType.SessionActiveClientChanged,
session: sessionUri.toString(),
activeClient: {
clientId: 'test-client',
tools: [],
},
};
sideEffects.handleAction(action);
assert.strictEqual(agent.setClientCustomizationsCalls.length, 0);
const customizationActions = envelopes
.filter(e => e.action.type === ActionType.SessionCustomizationsChanged);
assert.strictEqual(customizationActions.length, 0);
});
test('skips when activeClient is null', () => {
setupSession();
const action: ISessionAction = {
type: ActionType.SessionActiveClientChanged,
session: sessionUri.toString(),
activeClient: null,
};
sideEffects.handleAction(action);
assert.strictEqual(agent.setClientCustomizationsCalls.length, 0);
});
});
// ---- handleAction: session/customizationToggled ---------------------
suite('handleAction — session/customizationToggled', () => {
test('calls setCustomizationEnabled on the agent', () => {
setupSession();
const action: ISessionAction = {
type: ActionType.SessionCustomizationToggled,
session: sessionUri.toString(),
uri: 'file:///plugin-a',
enabled: false,
};
sideEffects.handleAction(action);
assert.deepStrictEqual(agent.setCustomizationEnabledCalls, [
{ uri: 'file:///plugin-a', enabled: false },
]);
});
});
});

2
src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts
View File

@@ -18,6 +18,7 @@ import { SessionStatus, type ISessionSummary } from '../../common/state/sessionS
import type { IProtocolServer, IProtocolTransport } from '../../common/state/sessionTransport.js';
import { ProtocolServerHandler } from '../../node/protocolServerHandler.js';
import { SessionStateManager } from '../../node/sessionStateManager.js';
import { AgentHostFileSystemProvider, AHPFileSystemProvider } from '../../common/agentHostFileSystemProvider.js';
// ---- Mock helpers -----------------------------------------------------------
@@ -196,6 +197,7 @@ suite('ProtocolServerHandler', () => {
stateManager,
server,
{ defaultDirectory: URI.file('/home/testuser').toString() },
disposables.add(new AgentHostFileSystemProvider()),
new NullLogService(),
));
});

118
src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHost.contribution.ts
View File

@@ -8,16 +8,19 @@ import { observableValue } from '../../../../base/common/observable.js';
import { isEqualOrParent } from '../../../../base/common/resources.js';
import { URI } from '../../../../base/common/uri.js';
import * as nls from '../../../../nls.js';
import { AGENT_HOST_SCHEME, agentHostAuthority, fromAgentHostUri } from '../../../../platform/agentHost/common/agentHostUri.js';
import { AGENT_HOST_LABEL_FORMATTER, agentHostAuthority } from '../../../../platform/agentHost/common/agentHostUri.js';
import { type AgentProvider, type IAgentConnection } from '../../../../platform/agentHost/common/agentService.js';
import { IRemoteAgentHostConnectionInfo, IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId } from '../../../../platform/agentHost/common/remoteAgentHostService.js';
import { IRemoteAgentHostConnectionInfo, IRemoteAgentHostEntry, IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId } from '../../../../platform/agentHost/common/remoteAgentHostService.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 ICustomizationRef, type IRootState } from '../../../../platform/agentHost/common/state/sessionState.js';
import { type URI as ProtocolURI } from '../../../../platform/agentHost/common/state/protocol/state.js';
import { IStorageService } from '../../../../platform/storage/common/storage.js';
import { Extensions as ConfigurationExtensions, IConfigurationRegistry } from '../../../../platform/configuration/common/configurationRegistry.js';
import { IConfigurationService } from '../../../../platform/configuration/common/configuration.js';
import { IDefaultAccountService } from '../../../../platform/defaultAccount/common/defaultAccount.js';
import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js';
import { ILabelService } from '../../../../platform/label/common/label.js';
import { ILogService } from '../../../../platform/log/common/log.js';
import { Registry } from '../../../../platform/registry/common/platform.js';
import { IWorkbenchContribution, registerWorkbenchContribution2, WorkbenchPhase } from '../../../../workbench/common/contributions.js';
@@ -28,15 +31,15 @@ import { LoggingAgentConnection } from '../../../../workbench/contrib/chat/brows
import { IChatSessionsService } from '../../../../workbench/contrib/chat/common/chatSessionsService.js';
import { ILanguageModelsService } from '../../../../workbench/contrib/chat/common/languageModels.js';
import { IAuthenticationService } from '../../../../workbench/services/authentication/common/authentication.js';
import { IAgentHostFileSystemService } from '../../../../workbench/services/agentHost/common/agentHostFileSystemService.js';
import { ICustomizationHarnessService } from '../../../../workbench/contrib/chat/common/customizationHarnessService.js';
import { IStorageService } from '../../../../platform/storage/common/storage.js';
import { IAgentPluginService } from '../../../../workbench/contrib/chat/common/plugins/agentPluginService.js';
import { PromptsType } from '../../../../workbench/contrib/chat/common/promptSyntax/promptTypes.js';
import { AgentCustomizationSyncProvider } from '../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentCustomizationSyncProvider.js';
import { ISessionsManagementService } from '../../../contrib/sessions/browser/sessionsManagementService.js';
import { ISessionsProvidersService } from '../../sessions/browser/sessionsProvidersService.js';
import { RemoteAgentHostSessionsProvider } from './remoteAgentHostSessionsProvider.js';
import { createRemoteAgentHarnessDescriptor, RemoteAgentCustomizationItemProvider, RemoteAgentSyncProvider } from './remoteAgentHostCustomizationHarness.js';
import { IAgentHostFileSystemService } from '../../../../workbench/services/agentHost/common/agentHostFileSystemService.js';
import { createRemoteAgentHarnessDescriptor, RemoteAgentCustomizationItemProvider } from './remoteAgentHostCustomizationHarness.js';
import { SyncedCustomizationBundler } from './syncedCustomizationBundler.js';
/** Per-connection state bundle, disposed when a connection is removed. */
@@ -75,6 +78,10 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
/** Per-connection state: client state + per-agent registrations. */
private readonly _connections = this._register(new DisposableMap<string, ConnectionState>());
/** Per-address sessions providers, registered for all configured entries. */
private readonly _providerStores = this._register(new DisposableMap<string, DisposableStore>());
private readonly _providerInstances = new Map<string, RemoteAgentHostSessionsProvider>();
constructor(
@IRemoteAgentHostService private readonly _remoteAgentHostService: IRemoteAgentHostService,
@IChatSessionsService private readonly _chatSessionsService: IChatSessionsService,
@@ -85,6 +92,8 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
@IDefaultAccountService private readonly _defaultAccountService: IDefaultAccountService,
@ISessionsManagementService private readonly _sessionsManagementService: ISessionsManagementService,
@ISessionsProvidersService private readonly _sessionsProvidersService: ISessionsProvidersService,
@ILabelService private readonly _labelService: ILabelService,
@IConfigurationService private readonly _configurationService: IConfigurationService,
@IAgentHostFileSystemService private readonly _agentHostFileSystemService: IAgentHostFileSystemService,
@ICustomizationHarnessService private readonly _customizationHarnessService: ICustomizationHarnessService,
@IStorageService private readonly _storageService: IStorageService,
@@ -92,17 +101,79 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
) {
super();
// Display agent-host URIs with the original file path
this._register(this._labelService.registerFormatter(AGENT_HOST_LABEL_FORMATTER));
// Reconcile providers when configured entries change
this._register(this._configurationService.onDidChangeConfiguration(e => {
if (e.affectsConfiguration(RemoteAgentHostsSettingId) || e.affectsConfiguration(RemoteAgentHostsEnabledSettingId)) {
this._reconcile();
}
}));
// Reconcile when connections change (added/removed/reconnected)
this._register(this._remoteAgentHostService.onDidChangeConnections(() => {
this._reconcileConnections();
this._reconcile();
}));
// Push auth token whenever the default account or sessions change
this._register(this._defaultAccountService.onDidChangeDefaultAccount(() => this._authenticateAllConnections()));
this._register(this._authenticationService.onDidChangeSessions(() => this._authenticateAllConnections()));
// Initial setup for already-connected remotes
// Initial setup for configured entries and connected remotes
this._reconcile();
}
private _reconcile(): void {
this._reconcileProviders();
this._reconcileConnections();
// Ensure every live connection is wired to its provider.
// This covers the case where a provider was recreated (e.g. name
// change) while a connection for that address already existed.
for (const [address, connState] of this._connections) {
const provider = this._providerInstances.get(address);
if (provider) {
const connectionInfo = this._remoteAgentHostService.connections.find(c => c.address === address);
provider.setConnection(connState.loggedConnection, connectionInfo?.defaultDirectory);
}
}
}
private _reconcileProviders(): void {
const enabled = this._configurationService.getValue<boolean>(RemoteAgentHostsEnabledSettingId);
const entries = enabled ? this._remoteAgentHostService.configuredEntries : [];
const desiredAddresses = new Set(entries.map(e => e.address));
// Remove providers no longer configured
for (const [address] of this._providerStores) {
if (!desiredAddresses.has(address)) {
this._providerStores.deleteAndDispose(address);
}
}
// Add or recreate providers for configured entries
for (const entry of entries) {
const existing = this._providerInstances.get(entry.address);
if (existing && existing.label !== (entry.name || entry.address)) {
// Name changed — recreate since ISessionsProvider.label is readonly
this._providerStores.deleteAndDispose(entry.address);
}
if (!this._providerStores.has(entry.address)) {
this._createProvider(entry);
}
}
}
private _createProvider(entry: IRemoteAgentHostEntry): void {
const store = new DisposableStore();
const provider = this._instantiationService.createInstance(
RemoteAgentHostSessionsProvider, { address: entry.address, name: entry.name });
store.add(provider);
store.add(this._sessionsProvidersService.registerProvider(provider));
this._providerInstances.set(entry.address, provider);
store.add(toDisposable(() => this._providerInstances.delete(entry.address)));
this._providerStores.set(entry.address, store);
}
private _reconcileConnections(): void {
@@ -112,6 +183,7 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
for (const [address] of this._connections) {
if (!currentAddresses.has(address)) {
this._logService.info(`[RemoteAgentHost] Removing contribution for ${address}`);
this._providerInstances.get(address)?.clearConnection();
this._connections.deleteAndDispose(address);
}
}
@@ -123,6 +195,7 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
// If the name changed, tear down and re-register with new name
if (existing.name !== connectionInfo.name) {
this._logService.info(`[RemoteAgentHost] Name changed for ${connectionInfo.address}: ${existing.name} -> ${connectionInfo.name}`);
this._providerInstances.get(connectionInfo.address)?.clearConnection();
this._connections.deleteAndDispose(connectionInfo.address);
this._setupConnection(connectionInfo);
}
@@ -182,12 +255,8 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
// Authenticate with this new connection
this._authenticateWithConnection(loggedConnection);
// Register a single sessions provider for the entire connection.
// It handles all agents discovered on this connection.
const sessionsProvider = this._instantiationService.createInstance(
RemoteAgentHostSessionsProvider, { connectionInfo, connection: loggedConnection });
store.add(sessionsProvider);
store.add(this._sessionsProvidersService.registerProvider(sessionsProvider));
// Wire connection to existing sessions provider
this._providerInstances.get(address)?.setConnection(loggedConnection, connectionInfo.defaultDirectory);
}
private _handleRootStateChange(address: string, loggedConnection: LoggingAgentConnection, rootState: IRootState): void {
@@ -241,11 +310,11 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
const displayName = configuredName || `${agent.displayName} (${address})`;
// Per-agent working directory cache, scoped to the agent store lifetime
const sessionWorkingDirs = new Map<string, string>();
const sessionWorkingDirs = new Map<string, URI>();
agentStore.add(toDisposable(() => sessionWorkingDirs.clear()));
// Capture the working directory from the active session for new sessions
const resolveWorkingDirectory = (resourceKey: string): string | undefined => {
const resolveWorkingDirectory = (resourceKey: string): URI | undefined => {
const cached = sessionWorkingDirs.get(resourceKey);
if (cached) {
return cached;
@@ -253,12 +322,8 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
const activeSession = this._sessionsManagementService.activeSession.get();
const repoUri = activeSession?.workspace.get()?.repositories[0]?.uri;
if (repoUri) {
// The repository URI may be wrapped as a vscode-agent-host:// URI.
// Unwrap to get the original filesystem path.
const originalUri = repoUri.scheme === AGENT_HOST_SCHEME ? fromAgentHostUri(repoUri) : repoUri;
const dir = originalUri.path;
sessionWorkingDirs.set(resourceKey, dir);
return dir;
sessionWorkingDirs.set(resourceKey, repoUri);
return repoUri;
}
return undefined;
};
@@ -276,16 +341,14 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
// Customization harness for this remote agent
const itemProvider = agentStore.add(new RemoteAgentCustomizationItemProvider(agent, connState.clientState));
const syncProvider = agentStore.add(new RemoteAgentSyncProvider(sessionType, this._storageService));
const syncProvider = agentStore.add(new AgentCustomizationSyncProvider(sessionType, this._storageService));
const harnessDescriptor = createRemoteAgentHarnessDescriptor(sessionType, displayName, itemProvider, syncProvider);
agentStore.add(this._customizationHarnessService.registerExternalHarness(harnessDescriptor));
// Bundler for packaging individual files into a virtual Open Plugin
const bundler = agentStore.add(this._instantiationService.createInstance(SyncedCustomizationBundler, sessionType));
// Agent-level customizations observable: resolves sync provider
// selections into ICustomizationRef[] and updates whenever the
// sync provider selection changes.
// Agent-level customizations observable
const customizations = observableValue<ICustomizationRef[]>('agentCustomizations', []);
const updateCustomizations = async () => {
const refs = await this._resolveCustomizations(syncProvider, bundler);
@@ -330,12 +393,11 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
*
* Entries are classified as either:
* - **Plugin**: A selected URI matches an installed plugin's root URI.
* The plugin directory is synced directly as an {@link ICustomizationRef}.
* - **Individual file**: All other selected files are bundled into a
* synthetic Open Plugin via {@link SyncedCustomizationBundler}.
*/
private async _resolveCustomizations(
syncProvider: RemoteAgentSyncProvider,
syncProvider: AgentCustomizationSyncProvider,
bundler: SyncedCustomizationBundler,
): Promise<ICustomizationRef[]> {
const entries = syncProvider.getSelectedEntries();
@@ -348,7 +410,6 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
const individualFiles: { uri: URI; type: PromptsType }[] = [];
for (const entry of entries) {
// Check if this URI matches an installed plugin's root directory
const plugin = plugins.find(p => isEqualOrParent(entry.uri, p.uri));
if (plugin) {
refs.push({
@@ -360,7 +421,6 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
}
}
// Bundle individual files into a synthetic plugin
if (individualFiles.length > 0) {
const result = await bundler.bundle(individualFiles);
if (result) {

252
src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts
View File

@@ -5,32 +5,36 @@
import { Throttler } from '../../../../../../base/common/async.js';
import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js';
import { Emitter } from '../../../../../../base/common/event.js';
import { Emitter, Event } from '../../../../../../base/common/event.js';
import { MarkdownString } from '../../../../../../base/common/htmlContent.js';
import { Disposable, DisposableMap, DisposableStore, MutableDisposable, type IDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js';
import { IObservable, observableValue } from '../../../../../../base/common/observable.js';
import { Disposable, DisposableResourceMap, DisposableStore, MutableDisposable, toDisposable, type IDisposable } from '../../../../../../base/common/lifecycle.js';
import { ResourceMap } from '../../../../../../base/common/map.js';
import { autorun, IObservable, observableValue } from '../../../../../../base/common/observable.js';
import { isEqual } from '../../../../../../base/common/resources.js';
import { URI } from '../../../../../../base/common/uri.js';
import { generateUuid } from '../../../../../../base/common/uuid.js';
import { localize } from '../../../../../../nls.js';
import { toAgentHostUri } from '../../../../../../platform/agentHost/common/agentHostUri.js';
import { AgentProvider, AgentSession, IAgentAttachment, type IAgentConnection } from '../../../../../../platform/agentHost/common/agentService.js';
import { ICustomizationRef } from '../../../../../../platform/agentHost/common/state/protocol/state.js';
import { ActionType, isSessionAction, type ISessionAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js';
import { SessionClientState } from '../../../../../../platform/agentHost/common/state/sessionClientState.js';
import { AHP_AUTH_REQUIRED, ProtocolError } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js';
import { getToolKind, getToolLanguage } from '../../../../../../platform/agentHost/common/state/sessionReducers.js';
import { AttachmentType, PendingMessageKind, ResponsePartKind, ToolCallCancellationReason, ToolCallConfirmationReason, ToolCallStatus, TurnState, type ICustomizationRef, type IMessageAttachment, type ISessionState } from '../../../../../../platform/agentHost/common/state/sessionState.js';
import { AttachmentType, getToolFileEdits, PendingMessageKind, ResponsePartKind, ToolCallCancellationReason, ToolCallConfirmationReason, ToolCallStatus, TurnState, type IMessageAttachment, type ISessionState, type IToolCallState, type ITurn } from '../../../../../../platform/agentHost/common/state/sessionState.js';
import { ExtensionIdentifier } from '../../../../../../platform/extensions/common/extensions.js';
import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js';
import { ILogService } from '../../../../../../platform/log/common/log.js';
import { IProductService } from '../../../../../../platform/product/common/productService.js';
import { IWorkspaceContextService } from '../../../../../../platform/workspace/common/workspace.js';
import { IChatProgress, IChatService, IChatToolInvocation, ToolConfirmKind, ChatRequestQueueKind } from '../../../common/chatService/chatService.js';
import { ChatRequestQueueKind, IChatProgress, IChatService, IChatToolInvocation, ToolConfirmKind } from '../../../common/chatService/chatService.js';
import { IChatSession, IChatSessionContentProvider, IChatSessionHistoryItem } from '../../../common/chatSessionsService.js';
import { ChatAgentLocation, ChatModeKind } from '../../../common/constants.js';
import { IChatEditingService } from '../../../common/editing/chatEditingService.js';
import { ChatToolInvocation } from '../../../common/model/chatProgressTypes/chatToolInvocation.js';
import { IChatAgentData, IChatAgentImplementation, IChatAgentRequest, IChatAgentResult, IChatAgentService } from '../../../common/participants/chatAgents.js';
import { getAgentHostIcon } from '../agentSessions.js';
import { activeTurnToProgress, finalizeToolInvocation, toolCallStateToInvocation, turnsToHistory, type IToolCallFileEdit } from './stateToProgressAdapter.js';
import { AgentHostEditingSession } from './agentHostEditingSession.js';
import { activeTurnToProgress, finalizeToolInvocation, toolCallStateToInvocation, turnsToHistory } from './stateToProgressAdapter.js';
// =============================================================================
// AgentHostSessionHandler - renderer-side handler for a single agent host
@@ -148,13 +152,14 @@ export interface IAgentHostSessionHandlerConfig {
* Optional callback to resolve a working directory for a new session.
* If not provided, falls back to the first workspace folder.
*/
readonly resolveWorkingDirectory?: (resourceKey: string) => string | undefined;
readonly resolveWorkingDirectory?: (resourceKey: string) => URI | undefined;
/**
* Optional callback invoked when the server rejects an operation because
* authentication is required. Should trigger interactive authentication
* and return true if the user authenticated successfully.
*/
readonly resolveAuthentication?: () => Promise<boolean>;
/**
* Observable set of agent-level customizations to include in the active
* client set. When the value changes, active sessions are updated.
@@ -164,13 +169,17 @@ export interface IAgentHostSessionHandlerConfig {
export class AgentHostSessionHandler extends Disposable implements IChatSessionContentProvider {
private readonly _activeSessions = new Map<string, AgentHostChatSession>();
private readonly _activeSessions = new ResourceMap<AgentHostChatSession>();
/** Maps UI resource keys to resolved backend session URIs. */
private readonly _sessionToBackend = new Map<string, URI>();
private readonly _sessionToBackend = new ResourceMap<URI>();
/** Per-session subscription to chat model pending request changes. */
private readonly _pendingMessageSubscriptions = this._register(new DisposableMap<string>());
private readonly _pendingMessageSubscriptions = this._register(new DisposableResourceMap());
/** Per-session subscription watching for server-initiated turns. */
private readonly _serverTurnWatchers = this._register(new DisposableMap<string>());
private readonly _serverTurnWatchers = this._register(new DisposableResourceMap());
/** Per-session writeFile listeners for agent host editing sessions. */
private readonly _editingSessionListeners = this._register(new DisposableResourceMap());
/** Historical turns with file edits, pending hydration into the editing session. */
private readonly _pendingHistoryTurns = new ResourceMap<readonly ITurn[]>();
/** Turn IDs dispatched by this client, used to distinguish server-originated turns. */
private readonly _clientDispatchedTurnIds = new Set<string>();
private readonly _config: IAgentHostSessionHandlerConfig;
@@ -182,6 +191,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
config: IAgentHostSessionHandlerConfig,
@IChatAgentService private readonly _chatAgentService: IChatAgentService,
@IChatService private readonly _chatService: IChatService,
@IChatEditingService private readonly _chatEditingService: IChatEditingService,
@ILogService private readonly _logService: ILogService,
@IProductService private readonly _productService: IProductService,
@IWorkspaceContextService private readonly _workspaceContextService: IWorkspaceContextService,
@@ -193,6 +203,20 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
// Create shared client state manager for this handler instance
this._clientState = this._register(new SessionClientState(config.connection.clientId, this._logService));
// Register an editing session provider for this handler's session type
this._register(this._chatEditingService.registerEditingSessionProvider(
config.sessionType,
{
createEditingSession: (chatSessionResource: URI) => {
return this._instantiationService.createInstance(
AgentHostEditingSession,
chatSessionResource,
config.connectionAuthority,
);
},
},
));
// Forward action envelopes from IPC to client state
this._register(config.connection.onDidAction(envelope => {
if (isSessionAction(envelope.action)) {
@@ -200,23 +224,38 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
}
}));
// When the customizations observable changes, re-dispatch
// activeClientChanged for sessions where this client is already
// the active client. This avoids overwriting another client's
// active status on sessions we're only observing.
if (config.customizations) {
this._register(autorun(reader => {
const refs = config.customizations!.read(reader);
for (const [, backendSession] of this._sessionToBackend) {
const state = this._clientState.getSessionState(backendSession.toString());
if (state?.activeClient?.clientId === this._clientState.clientId) {
this._dispatchActiveClient(backendSession, refs);
}
}
}));
}
this._registerAgent();
}
async provideChatSessionContent(sessionResource: URI, _token: CancellationToken): Promise<IChatSession> {
const resourceKey = sessionResource.path.substring(1);
// 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 isUntitled = sessionResource.path.substring(1).startsWith('untitled-');
const history: IChatSessionHistoryItem[] = [];
let initialProgress: IChatProgress[] | undefined;
let activeTurnId: string | undefined;
if (!isUntitled) {
resolvedSession = this._resolveSessionUri(sessionResource);
this._sessionToBackend.set(resourceKey, resolvedSession);
this._sessionToBackend.set(sessionResource, resolvedSession);
try {
const snapshot = await this._config.connection.subscribe(resolvedSession);
if (snapshot?.state) {
@@ -225,6 +264,16 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
if (sessionState) {
history.push(...turnsToHistory(sessionState.turns, this._config.agentId));
// Store turns with file edits so the editing session
// can be hydrated when it's created lazily.
const hasTurnsWithEdits = sessionState.turns.some(t =>
t.responseParts.some(rp => rp.kind === ResponsePartKind.ToolCall
&& rp.toolCall.status === ToolCallStatus.Completed
&& getToolFileEdits(rp.toolCall).length > 0));
if (hasTurnsWithEdits) {
this._pendingHistoryTurns.set(sessionResource, sessionState.turns);
}
// If there's an active turn, include its request in history
// with an empty response so the chat service creates a
// pending request, then provide accumulated progress via
@@ -249,6 +298,10 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
} catch (err) {
this._logService.warn(`[AgentHost] Failed to subscribe to existing session: ${resolvedSession.toString()}`, err);
}
// Claim the active client role with current customizations
const customizations = this._config.customizations?.get() ?? [];
this._dispatchActiveClient(resolvedSession, customizations);
}
const session = this._instantiationService.createInstance(
AgentHostChatSession,
@@ -258,19 +311,21 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
const backendSession = resolvedSession ?? await this._createAndSubscribe(sessionResource, request.userSelectedModelId);
if (!resolvedSession) {
resolvedSession = backendSession;
this._sessionToBackend.set(resourceKey, backendSession);
this._sessionToBackend.set(sessionResource, backendSession);
}
// For existing sessions, set up pending message sync on the first turn
// (after the ChatModel becomes available in the ChatService).
this._ensurePendingMessageSubscription(resourceKey, sessionResource, backendSession);
this._ensurePendingMessageSubscription(sessionResource, backendSession);
return this._handleTurn(backendSession, request, progress, token);
},
initialProgress,
() => {
this._activeSessions.delete(resourceKey);
this._sessionToBackend.delete(resourceKey);
this._pendingMessageSubscriptions.deleteAndDispose(resourceKey);
this._serverTurnWatchers.deleteAndDispose(resourceKey);
this._activeSessions.delete(sessionResource);
this._sessionToBackend.delete(sessionResource);
this._pendingMessageSubscriptions.deleteAndDispose(sessionResource);
this._serverTurnWatchers.deleteAndDispose(sessionResource);
this._editingSessionListeners.deleteAndDispose(sessionResource);
this._pendingHistoryTurns.delete(sessionResource);
if (resolvedSession) {
this._clientState.unsubscribe(resolvedSession.toString());
this._config.connection.unsubscribe(resolvedSession);
@@ -278,9 +333,20 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
}
},
);
this._activeSessions.set(resourceKey, session);
this._activeSessions.set(sessionResource, session);
if (resolvedSession) {
// If there are historical turns with file edits, eagerly create
// the editing session once the ChatModel is available so that
// edit pills render with diff info on session restore.
if (this._pendingHistoryTurns.has(sessionResource)) {
session.registerDisposable(Event.once(this._chatService.onDidCreateModel)(model => {
if (isEqual(model.sessionResource, sessionResource)) {
this._ensureEditingSession(sessionResource);
}
}));
}
// If reconnecting to an active turn, wire up an ongoing state listener
// to stream new progress into the session's progressObs.
if (activeTurnId && initialProgress !== undefined) {
@@ -334,16 +400,15 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
this._logService.info(`[AgentHost] _invokeAgent called for resource: ${request.sessionResource.toString()}`);
// Resolve or create backend session
const resourceKey = request.sessionResource.path.substring(1);
let resolvedSession = this._sessionToBackend.get(resourceKey);
let resolvedSession = this._sessionToBackend.get(request.sessionResource);
if (!resolvedSession) {
resolvedSession = await this._createAndSubscribe(request.sessionResource, request.userSelectedModelId);
this._sessionToBackend.set(resourceKey, resolvedSession);
this._sessionToBackend.set(request.sessionResource, resolvedSession);
}
await this._handleTurn(resolvedSession, request, progress, cancellationToken);
const activeSession = this._activeSessions.get(resourceKey);
const activeSession = this._activeSessions.get(request.sessionResource);
if (activeSession) {
activeSession.isCompleteObs.set(true, undefined);
}
@@ -448,14 +513,18 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
this._config.connection.dispatchAction(action, this._clientState.clientId, seq);
}
private _setActiveClient(session: URI): void {
/**
* Dispatches `session/activeClientChanged` to claim the active client
* role for this session and publish the current customizations.
*/
private _dispatchActiveClient(backendSession: URI, customizations: ICustomizationRef[]): void {
this._dispatchAction({
type: ActionType.SessionActiveClientChanged,
session: session.toString(),
session: backendSession.toString(),
activeClient: {
clientId: this._clientState.clientId,
tools: [],
customizations: this._config.customizations?.get() || [],
customizations,
},
});
}
@@ -471,7 +540,6 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
* if applicable, and pipes turn progress through `progressObs`.
*/
private _watchForServerInitiatedTurns(backendSession: URI, sessionResource: URI): void {
const resourceKey = sessionResource.path.substring(1);
const sessionStr = backendSession.toString();
// Seed from the current state so we don't treat any pre-existing active
@@ -479,6 +547,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
const currentState = this._clientState.getSessionState(sessionStr);
let lastSeenTurnId: string | undefined = currentState?.activeTurn?.id;
let previousQueuedIds: Set<string> | undefined;
let previousSteeringId: string | undefined = currentState?.steeringMessage?.id;
const disposables = new DisposableStore();
@@ -493,6 +562,13 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
// Track queued message IDs so we can detect which one was consumed
const currentQueuedIds = new Set((e.state.queuedMessages ?? []).map(m => m.id));
const currentSteeringId = e.state.steeringMessage?.id;
// Detect steering message removal or replacement regardless of turn changes
if (previousSteeringId && previousSteeringId !== currentSteeringId) {
this._chatService.removePendingRequest(sessionResource, previousSteeringId);
}
previousSteeringId = currentSteeringId;
const activeTurn = e.state.activeTurn;
if (!activeTurn || activeTurn.id === lastSeenTurnId) {
@@ -507,7 +583,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
return;
}
const chatSession = this._activeSessions.get(resourceKey);
const chatSession = this._activeSessions.get(sessionResource);
if (!chatSession) {
previousQueuedIds = currentQueuedIds;
return;
@@ -535,7 +611,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
this._trackServerTurnProgress(backendSession, activeTurn.id, chatSession, sessionResource, turnStore);
}));
this._serverTurnWatchers.set(resourceKey, disposables);
this._serverTurnWatchers.set(sessionResource, disposables);
}
/**
@@ -704,8 +780,6 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
}
}
this._setActiveClient(session);
// Dispatch session/turnStarted — the server will call sendMessage on
// the provider as a side effect.
const turnAction = {
@@ -830,7 +904,10 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
if (existing && (tc.status === ToolCallStatus.Completed || tc.status === ToolCallStatus.Cancelled) && !IChatToolInvocation.isComplete(existing)) {
const fileEdits = finalizeToolInvocation(existing, tc);
if (fileEdits.length > 0) {
await this._applyFileEdits(request.sessionResource, request, fileEdits, progress);
const editParts = this._hydrateFileEdits(request.sessionResource, request.requestId, tc);
if (editParts.length > 0) {
progress(editParts);
}
}
}
break;
@@ -1087,47 +1164,68 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
// ---- File edit routing ---------------------------------------------------
/**
* Routes file edits from completed tool calls through the editing session's
* external edits pipeline. Calls start/stop in sequence since the edit has
* already happened on the remote by the time we receive the tool completion.
* Ensures the chat model has an editing session and returns it if it's an
* {@link AgentHostEditingSession}. The editing session is created via the
* provider registered in the constructor if one doesn't exist yet.
*/
private async _applyFileEdits(
sessionResource: URI,
request: IChatAgentRequest,
fileEdits: IToolCallFileEdit[],
progress: (parts: IChatProgress[]) => void,
): Promise<void> {
const chatSession = this._chatService.getSession(sessionResource);
const editingSession = chatSession?.editingSession;
const response = chatSession?.getRequests().find(req => req.id === request.requestId)?.response;
if (!editingSession || !response) {
return;
private _ensureEditingSession(sessionResource: URI): AgentHostEditingSession | undefined {
const chatModel = this._chatService.getSession(sessionResource);
if (!chatModel) {
return undefined;
}
const authority = this._config.connectionAuthority;
const wrapUri = (uri: URI) => toAgentHostUri(uri, authority);
for (const edit of fileEdits) {
const operationId = this._nextOperationId++;
const resource = wrapUri(edit.resource);
const beforeUri = wrapUri(edit.beforeContentUri);
const afterUri = wrapUri(edit.afterContentUri);
const startProgress = await editingSession.startExternalEdits(
response, operationId, [resource], edit.undoStopId,
[beforeUri],
);
progress(startProgress);
const stopProgress = await editingSession.stopExternalEdits(
response, operationId,
[afterUri],
);
progress(stopProgress);
// Start the editing session if not already started — this will use
// our registered provider to create an AgentHostEditingSession.
if (!chatModel.editingSession) {
chatModel.startEditingSession();
}
const editingSession = chatModel.editingSession;
if (!(editingSession instanceof AgentHostEditingSession)) {
return undefined;
}
// Wire up the writeFile listener if not already done
if (!this._editingSessionListeners.has(sessionResource)) {
this._editingSessionListeners.set(sessionResource, editingSession.onDidRequestFileWrite(params => {
this._config.connection.writeFile(params).catch(err => {
this._logService.warn('[AgentHost] writeFile failed for undo/redo', err);
});
}));
// Hydrate from historical turns if this is the first time
// the editing session is accessed for this chat session.
const pendingTurns = this._pendingHistoryTurns.get(sessionResource);
if (pendingTurns) {
this._pendingHistoryTurns.delete(sessionResource);
for (const turn of pendingTurns) {
for (const rp of turn.responseParts) {
if (rp.kind === ResponsePartKind.ToolCall) {
editingSession.addToolCallEdits(turn.id, rp.toolCall);
}
}
}
}
}
return editingSession;
}
private _nextOperationId = 0;
/**
* Hydrates the editing session with file edits from a completed tool call
* and returns progress parts for the file edit pills.
*/
private _hydrateFileEdits(
sessionResource: URI,
requestId: string,
tc: IToolCallState,
): IChatProgress[] {
const editingSession = this._ensureEditingSession(sessionResource);
if (editingSession) {
return editingSession.addToolCallEdits(requestId, tc);
}
return [];
}
// ---- Session resolution -------------------------------------------------
@@ -1142,7 +1240,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
const rawModelId = this._extractRawModelId(modelId);
const resourceKey = sessionResource.path.substring(1);
const workingDirectory = this._config.resolveWorkingDirectory?.(resourceKey)
?? this._workspaceContextService.getWorkspace().folders[0]?.uri.fsPath;
?? this._workspaceContextService.getWorkspace().folders[0]?.uri;
this._logService.trace(`[AgentHost] Creating new session, model=${rawModelId ?? '(default)'}, provider=${this._config.provider}`);
@@ -1182,8 +1280,12 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
this._logService.error(`[AgentHost] Failed to subscribe to new session: ${session.toString()}`, err);
}
// Claim the active client role with current customizations
const customizations = this._config.customizations?.get() ?? [];
this._dispatchActiveClient(session, customizations);
// Start syncing the chat model's pending requests to the protocol
this._ensurePendingMessageSubscription(resourceKey, sessionResource, session);
this._ensurePendingMessageSubscription(sessionResource, session);
// Start watching for server-initiated turns on this session
this._watchForServerInitiatedTurns(session, sessionResource);
@@ -1195,13 +1297,13 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
* Ensures that the chat model's pending request changes are synced to the
* protocol for a given session. No-ops if already subscribed.
*/
private _ensurePendingMessageSubscription(resourceKey: string, sessionResource: URI, backendSession: URI): void {
if (this._pendingMessageSubscriptions.has(resourceKey)) {
private _ensurePendingMessageSubscription(sessionResource: URI, backendSession: URI): void {
if (this._pendingMessageSubscriptions.has(sessionResource)) {
return;
}
const chatModel = this._chatService?.getSession(sessionResource);
if (chatModel) {
this._pendingMessageSubscriptions.set(resourceKey, chatModel.onDidChangePendingRequests(() => {
this._pendingMessageSubscriptions.set(sessionResource, chatModel.onDidChangePendingRequests(() => {
this._syncPendingMessages(sessionResource, backendSession);
}));
}

75
src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts
View File

@@ -8,6 +8,7 @@ import { CancellationToken, CancellationTokenSource } from '../../../../../../ba
import { Emitter, Event } from '../../../../../../base/common/event.js';
import { DisposableStore, toDisposable } from '../../../../../../base/common/lifecycle.js';
import { URI } from '../../../../../../base/common/uri.js';
import { observableValue } from '../../../../../../base/common/observable.js';
import { mock, upcastPartial } from '../../../../../../base/test/common/mock.js';
import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js';
import { timeout } from '../../../../../../base/common/async.js';
@@ -16,6 +17,7 @@ import { IConfigurationService } from '../../../../../../platform/configuration/
import { IAgentCreateSessionConfig, IAgentHostService, IAgentSessionMetadata, AgentSession } from '../../../../../../platform/agentHost/common/agentService.js';
import type { IActionEnvelope, INotification, ISessionAction, IToolCallConfirmedAction, ITurnStartedAction } from '../../../../../../platform/agentHost/common/state/sessionActions.js';
import type { IStateSnapshot } from '../../../../../../platform/agentHost/common/state/sessionProtocol.js';
import type { ICustomizationRef } from '../../../../../../platform/agentHost/common/state/protocol/state.js';
import { SessionLifecycle, SessionStatus, TurnState, ToolCallStatus, ToolCallConfirmationReason, createSessionState, createActiveTurn, ROOT_STATE_URI, PolicyState, ResponsePartKind, type ISessionState, type ISessionSummary } from '../../../../../../platform/agentHost/common/state/sessionState.js';
import { IDefaultAccountService } from '../../../../../../platform/defaultAccount/common/defaultAccount.js';
import { IAuthenticationService } from '../../../../../services/authentication/common/authentication.js';
@@ -1932,4 +1934,77 @@ suite('AgentHostChatContribution', () => {
assert.strictEqual(serverRequestEvents.length, 0, 'Client-dispatched turns should not trigger onDidStartServerRequest');
});
});
// ---- Customizations dispatch ------------------------------------------
suite('customizations', () => {
test('dispatches activeClientChanged when a new session is created', async () => {
const { instantiationService, agentHostService } = createTestServices(disposables);
const customizations = observableValue<ICustomizationRef[]>('customizations', [
{ uri: 'file:///plugin-a', displayName: 'Plugin A' },
]);
const sessionHandler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, {
provider: 'copilot' as const,
agentId: 'agent-host-copilot',
sessionType: 'agent-host-copilot',
fullName: 'Agent Host - Copilot',
description: 'test',
connection: agentHostService,
connectionAuthority: 'local',
customizations,
}));
const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables);
fire({ type: 'session/turnComplete', session, turnId } as ISessionAction);
await turnPromise;
const activeClientAction = agentHostService.dispatchedActions.find(
d => d.action.type === 'session/activeClientChanged'
);
assert.ok(activeClientAction, 'should dispatch activeClientChanged');
const ac = activeClientAction!.action as { activeClient: { customizations?: ICustomizationRef[] } };
assert.strictEqual(ac.activeClient.customizations?.length, 1);
assert.strictEqual(ac.activeClient.customizations?.[0].uri, 'file:///plugin-a');
});
test('re-dispatches activeClientChanged when customizations observable changes', async () => {
const { instantiationService, agentHostService } = createTestServices(disposables);
const customizations = observableValue<ICustomizationRef[]>('customizations', []);
const sessionHandler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, {
provider: 'copilot' as const,
agentId: 'agent-host-copilot',
sessionType: 'agent-host-copilot',
fullName: 'Agent Host - Copilot',
description: 'test',
connection: agentHostService,
connectionAuthority: 'local',
customizations,
}));
// Create a session first
const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, disposables);
fire({ type: 'session/turnComplete', session, turnId } as ISessionAction);
await turnPromise;
agentHostService.dispatchedActions.length = 0;
// Update customizations
customizations.set([
{ uri: 'file:///plugin-b', displayName: 'Plugin B' },
], undefined);
const activeClientAction = agentHostService.dispatchedActions.find(
d => d.action.type === 'session/activeClientChanged'
);
assert.ok(activeClientAction, 'should re-dispatch activeClientChanged on change');
const ac = activeClientAction!.action as { activeClient: { customizations?: ICustomizationRef[] } };
assert.strictEqual(ac.activeClient.customizations?.length, 1);
assert.strictEqual(ac.activeClient.customizations?.[0].uri, 'file:///plugin-b');
});
});
});