Skip to content

Commit 886c556

Browse files
connor4312connor4312
committed
agentHost: adopt eager activeClient announcement
1 parent 8ba4c72 commit 886c556

10 files changed

Lines changed: 222 additions & 16 deletions

File tree

‎src/vs/platform/agentHost/browser/remoteAgentHostProtocolClient.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,7 @@ export class RemoteAgentHostProtocolClient extends Disposable implements IAgentC
188188
model: config?.model,
189189
workingDirectory: config?.workingDirectory ? fromAgentHostUri(config.workingDirectory).toString() : undefined,
190190
config: config?.config,
191+
activeClient: config?.activeClient,
191192
});
192193
return session;
193194
}

‎src/vs/platform/agentHost/common/agentService.ts‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import { createDecorator } from '../../instantiation/common/instantiation.js';
1212
import type { ISyncedCustomization } from './agentPluginManager.js';
1313
import type { IAgentSubscription } from './state/agentSubscription.js';
1414
import type { ICreateTerminalParams, IResolveSessionConfigResult, ISessionConfigCompletionsResult } from './state/protocol/commands.js';
15-
import { IProtectedResourceMetadata, type IConfigSchema, type IFileEdit, type IModelSelection, type IToolDefinition } from './state/protocol/state.js';
15+
import { IProtectedResourceMetadata, type IConfigSchema, type IFileEdit, type IModelSelection, type ISessionActiveClient, type IToolDefinition } from './state/protocol/state.js';
1616
import type { IActionEnvelope, INotification, ISessionAction, ITerminalAction } from './state/sessionActions.js';
1717
import type { IResourceCopyParams, IResourceCopyResult, IResourceDeleteParams, IResourceDeleteResult, IResourceListResult, IResourceMoveParams, IResourceMoveResult, IResourceReadResult, IResourceWriteParams, IResourceWriteResult, IStateSnapshot } from './state/sessionProtocol.js';
1818
import { AttachmentType, ComponentToState, SessionInputResponseKind, SessionStatus, StateComponents, type ICustomizationRef, type IPendingMessage, type IRootState, type ISessionInputAnswer, type ISessionInputRequest, type IToolCallResult, type IToolResultContent, type PolicyState, type StringOrMarkdown } from './state/sessionState.js';
@@ -125,6 +125,14 @@ export interface IAgentCreateSessionConfig {
125125
readonly session?: URI;
126126
readonly workingDirectory?: URI;
127127
readonly config?: Record<string, string>;
128+
/**
129+
* Eagerly claim the active client role for the new session. When provided,
130+
* the server initializes the session with this client as the active
131+
* client, equivalent to dispatching a `session/activeClientChanged`
132+
* action immediately after creation. The `clientId` MUST match the
133+
* connection's own `clientId`.
134+
*/
135+
readonly activeClient?: ISessionActiveClient;
128136
/** Fork from an existing session at a specific turn. */
129137
readonly fork?: {
130138
readonly session: URI;

‎src/vs/platform/agentHost/common/state/protocol/.ahp-version‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
ab467b2
1+
0947b17

‎src/vs/platform/agentHost/common/state/protocol/commands.ts‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
// allow-any-unicode-comment-file
77
// DO NOT EDIT -- auto-generated by scripts/sync-agent-host-protocol.ts
88

9-
import type { URI, ISnapshot, ISessionConfigSchema, ISessionSummary, IModelSelection, ITurn, ITerminalClaim } from './state.js';
9+
import type { URI, ISnapshot, ISessionConfigSchema, ISessionSummary, IModelSelection, ITurn, ITerminalClaim, ISessionActiveClient } from './state.js';
1010
import type { IActionEnvelope, IStateAction } from './actions.js';
1111

1212
export type { IConfigPropertySchema, IConfigSchema, ISessionConfigPropertySchema, ISessionConfigSchema } from './state.js';
@@ -198,6 +198,15 @@ export interface ICreateSessionParams {
198198
* Keys and values correspond to the schema returned by the server.
199199
*/
200200
config?: Record<string, string>;
201+
/**
202+
* Eagerly claim the active client role for the new session.
203+
*
204+
* When provided, the server initializes the session with this client as the
205+
* active client, equivalent to dispatching a `session/activeClientChanged`
206+
* action immediately after creation. The `clientId` MUST match the
207+
* `clientId` the creating client supplied in `initialize`.
208+
*/
209+
activeClient?: ISessionActiveClient;
201210
}
202211

203212
// ─── disposeSession ──────────────────────────────────────────────────────────

‎src/vs/platform/agentHost/node/copilot/copilotAgent.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -449,6 +449,13 @@ export class CopilotAgent extends Disposable implements IAgent {
449449

450450
const sessionId = config?.session ? AgentSession.id(config.session) : generateUuid();
451451
const sessionUri = AgentSession.uri(this.id, sessionId);
452+
if (config?.activeClient) {
453+
const ac = this._getOrCreateActiveClient(sessionUri);
454+
ac.updateTools(config.activeClient.clientId, config.activeClient.tools);
455+
if (config.activeClient.customizations?.length) {
456+
await this._plugins.sync(config.activeClient.clientId, config.activeClient.customizations);
457+
}
458+
}
452459
const activeClient = this._activeClients.get(sessionUri);
453460
const snapshot = activeClient ? await activeClient.snapshot() : undefined;
454461
const workingDirectory = await this._resolveSessionWorkingDirectory(config, sessionId);

‎src/vs/platform/agentHost/node/protocolServerHandler.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import { ILogService } from '../../log/common/log.js';
1212
import { AHPFileSystemProvider } from '../common/agentHostFileSystemProvider.js';
1313
import { AgentSession, type IAgentService } from '../common/agentService.js';
1414
import type { ICommandMap } from '../common/state/protocol/messages.js';
15+
import { ActionType } from '../common/state/protocol/actions.js';
1516
import { IActionEnvelope, INotification, isSessionAction, isTerminalAction, type ISessionAction } from '../common/state/sessionActions.js';
1617
import { MIN_PROTOCOL_VERSION, PROTOCOL_VERSION } from '../common/state/sessionCapabilities.js';
1718
import {
@@ -368,6 +369,19 @@ export class ProtocolServerHandler extends Disposable {
368369
if (createdSession.toString() !== URI.parse(params.session).toString()) {
369370
this._logService.warn(`[ProtocolServer] createSession: provider returned URI ${createdSession.toString()} but client requested ${params.session}`);
370371
}
372+
// If the client eagerly claimed the active client role, dispatch
373+
// `session/activeClientChanged` now so the claim is atomic with
374+
// creation.
375+
if (params.activeClient) {
376+
if (params.activeClient.clientId !== _client.clientId) {
377+
throw new ProtocolError(JSON_RPC_INTERNAL_ERROR, `createSession.activeClient.clientId must match the connection's clientId`);
378+
}
379+
this._agentService.dispatchAction({
380+
type: ActionType.SessionActiveClientChanged,
381+
session: createdSession.toString(),
382+
activeClient: params.activeClient,
383+
}, _client.clientId, 0);
384+
}
371385
return null;
372386
},
373387
disposeSession: async (_client, params) => {

‎src/vs/platform/agentHost/test/node/copilotAgent.test.ts‎

Lines changed: 71 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -127,14 +127,14 @@ class TestableCopilotAgent extends CopilotAgent {
127127
}
128128
}
129129

130-
function createTestAgent(disposables: Pick<DisposableStore, 'add'>, options?: { sessionDataService?: ISessionDataService; copilotClient?: ICopilotClient }): CopilotAgent {
130+
function createTestAgent(disposables: Pick<DisposableStore, 'add'>, options?: { sessionDataService?: ISessionDataService; copilotClient?: ICopilotClient; pluginManager?: IAgentPluginManager }): CopilotAgent {
131131
const services = new ServiceCollection();
132132
const logService = new NullLogService();
133133
const fileService = disposables.add(new FileService(logService));
134134
services.set(ILogService, logService);
135135
services.set(IFileService, fileService);
136136
services.set(ISessionDataService, options?.sessionDataService ?? createNullSessionDataService());
137-
services.set(IAgentPluginManager, new TestAgentPluginManager());
137+
services.set(IAgentPluginManager, options?.pluginManager ?? new TestAgentPluginManager());
138138
services.set(IAgentHostGitService, new TestAgentHostGitService());
139139
services.set(IAgentHostTerminalManager, new TestAgentHostTerminalManager());
140140
const instantiationService: IInstantiationService = disposables.add(new InstantiationService(services));
@@ -272,4 +272,73 @@ suite('CopilotAgent', () => {
272272
await disposeAgent(agent);
273273
}
274274
});
275+
276+
suite('createSession activeClient eager-claim', () => {
277+
278+
class SpyingPluginManager extends TestAgentPluginManager {
279+
public readonly calls: { clientId: string; customizations: ICustomizationRef[] }[] = [];
280+
281+
override async syncCustomizations(clientId: string, customizations: ICustomizationRef[], _progress?: (status: ISessionCustomization[]) => void): Promise<ISyncedCustomization[]> {
282+
this.calls.push({ clientId, customizations: [...customizations] });
283+
return [];
284+
}
285+
}
286+
287+
test('createSession seeds activeClient tools and syncs customizations', async () => {
288+
const sessionDataService = disposables.add(new TestSessionDataService());
289+
const client = new TestCopilotClient([]);
290+
const pluginManager = new SpyingPluginManager();
291+
// Fail fast inside the SDK factory so we don't need to wire up a
292+
// real raw session. The seeding of activeClient and the plugin
293+
// sync both happen before `client.createSession` is invoked.
294+
client.createSession = async () => { throw new Error('sentinel'); };
295+
296+
const agent = createTestAgent(disposables, { sessionDataService, copilotClient: client, pluginManager });
297+
try {
298+
await agent.authenticate('https://api.github.com', 'token');
299+
300+
const customizations: ICustomizationRef[] = [{ uri: 'file:///plugin-a', displayName: 'Plugin A' }];
301+
await assert.rejects(
302+
agent.createSession({
303+
session: AgentSession.uri('copilot', 'test-session'),
304+
workingDirectory: URI.file('/workspace'),
305+
activeClient: {
306+
clientId: 'client-1',
307+
tools: [{ name: 't1', description: 'd', inputSchema: { type: 'object' } }],
308+
customizations,
309+
},
310+
}),
311+
(err: Error) => /sentinel/.test(err.message),
312+
);
313+
314+
assert.deepStrictEqual(pluginManager.calls, [{ clientId: 'client-1', customizations }]);
315+
} finally {
316+
await disposeAgent(agent);
317+
}
318+
});
319+
320+
test('createSession without activeClient does not sync customizations', async () => {
321+
const sessionDataService = disposables.add(new TestSessionDataService());
322+
const client = new TestCopilotClient([]);
323+
const pluginManager = new SpyingPluginManager();
324+
client.createSession = async () => { throw new Error('sentinel'); };
325+
326+
const agent = createTestAgent(disposables, { sessionDataService, copilotClient: client, pluginManager });
327+
try {
328+
await agent.authenticate('https://api.github.com', 'token');
329+
330+
await assert.rejects(
331+
agent.createSession({
332+
session: AgentSession.uri('copilot', 'test-session-2'),
333+
workingDirectory: URI.file('/workspace'),
334+
}),
335+
(err: Error) => /sentinel/.test(err.message),
336+
);
337+
338+
assert.deepStrictEqual(pluginManager.calls, []);
339+
} finally {
340+
await disposeAgent(agent);
341+
}
342+
});
343+
});
275344
});

‎src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -577,4 +577,77 @@ suite('ProtocolServerHandler', () => {
577577
transport2.simulateClose();
578578
assert.deepStrictEqual(counts, [1, 1, 0]);
579579
});
580+
581+
// ---- createSession activeClient -------------------------------------
582+
583+
suite('createSession activeClient', () => {
584+
585+
test('forwards activeClient as SessionActiveClientChanged', async () => {
586+
const newSession = URI.parse('copilot:///eager-session').toString();
587+
// Pre-create the session so the client can subscribe at handshake
588+
// time. MockAgentService.createSession below is idempotent against
589+
// state-manager duplicates.
590+
stateManager.createSession(makeSessionSummary(newSession));
591+
592+
const transport = connectClient('client-1', [newSession]);
593+
transport.sent.length = 0;
594+
595+
const responsePromise = waitForResponse(transport, 2);
596+
transport.simulateMessage(request(2, 'createSession', {
597+
session: newSession,
598+
provider: 'copilot',
599+
activeClient: {
600+
clientId: 'client-1',
601+
tools: [{ name: 't1', description: 'd', inputSchema: { type: 'object' } }],
602+
customizations: [{ uri: 'file:///plugin-a', displayName: 'A' }],
603+
},
604+
}));
605+
await responsePromise;
606+
607+
const activeClientMsgs = findNotifications(transport.sent, 'action').filter(m => {
608+
const envelope = m.params as unknown as { action: { type: string } };
609+
return envelope.action.type === ActionType.SessionActiveClientChanged;
610+
});
611+
assert.strictEqual(activeClientMsgs.length, 1, 'should emit exactly one SessionActiveClientChanged');
612+
const envelope = activeClientMsgs[0].params as unknown as { action: { session: string; activeClient: { clientId: string; tools: { name: string }[]; customizations?: { uri: string }[] } } };
613+
assert.deepStrictEqual({
614+
session: envelope.action.session,
615+
clientId: envelope.action.activeClient.clientId,
616+
toolName: envelope.action.activeClient.tools[0].name,
617+
customizationUri: envelope.action.activeClient.customizations?.[0].uri,
618+
}, {
619+
session: newSession,
620+
clientId: 'client-1',
621+
toolName: 't1',
622+
customizationUri: 'file:///plugin-a',
623+
});
624+
});
625+
626+
test('rejects createSession when activeClient.clientId mismatches', async () => {
627+
const newSession = URI.parse('copilot:///mismatch-session').toString();
628+
stateManager.createSession(makeSessionSummary(newSession));
629+
630+
const transport = connectClient('client-1', [newSession]);
631+
transport.sent.length = 0;
632+
633+
const responsePromise = waitForResponse(transport, 2);
634+
transport.simulateMessage(request(2, 'createSession', {
635+
session: newSession,
636+
provider: 'copilot',
637+
activeClient: {
638+
clientId: 'other-client',
639+
tools: [],
640+
},
641+
}));
642+
const resp = await responsePromise as { result?: unknown; error?: { code: number; message: string } };
643+
644+
assert.ok(resp.error, 'response should be an error');
645+
assert.strictEqual(resp.result, undefined);
646+
const activeClientMsgs = findNotifications(transport.sent, 'action').filter(m => {
647+
const envelope = m.params as unknown as { action: { type: string } };
648+
return envelope.action.type === ActionType.SessionActiveClientChanged;
649+
});
650+
assert.strictEqual(activeClientMsgs.length, 0, 'no activeClient notification should have been emitted');
651+
});
652+
});
580653
});

‎src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2250,6 +2250,12 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
22502250
}
22512251
}
22522252

2253+
const activeClient = {
2254+
clientId: this._config.connection.clientId,
2255+
tools: this._clientToolsObs.get().map(toolDataToDefinition),
2256+
customizations: this._config.customizations?.get() ?? [],
2257+
};
2258+
22532259
let session: URI;
22542260
try {
22552261
session = await this._config.connection.createSession({
@@ -2258,6 +2264,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
22582264
workingDirectory,
22592265
fork,
22602266
config,
2267+
activeClient,
22612268
});
22622269
} catch (err) {
22632270
// If authentication is required (e.g. token expired), try interactive auth and retry once
@@ -2271,6 +2278,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
22712278
workingDirectory,
22722279
fork,
22732280
config,
2281+
activeClient,
22742282
});
22752283
} else {
22762284
throw new Error(localize('agentHost.authRequired', "Authentication is required to start a session. Please sign in and try again."));
@@ -2291,10 +2299,6 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
22912299
});
22922300
}
22932301

2294-
// Claim the active client role with current customizations
2295-
const customizations = this._config.customizations?.get() ?? [];
2296-
this._dispatchActiveClient(session, customizations);
2297-
22982302
// Start syncing the chat model's pending requests to the protocol
22992303
this._ensurePendingMessageSubscription(sessionResource, session);
23002304

‎src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts‎

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,26 @@ class MockAgentHostService extends mock<IAgentHostService>() {
8282
const id = `sdk-session-${this._nextId++}`;
8383
const session = AgentSession.uri('copilot', id);
8484
this._sessions.set(id, { session, startTime: Date.now(), modifiedTime: Date.now() });
85+
// Simulate the server's eager active-client claim: if the caller
86+
// provided activeClient, seed the session state so subscribers see it.
87+
if (config?.activeClient) {
88+
const summary: ISessionSummary = {
89+
resource: session.toString(),
90+
provider: 'copilot',
91+
title: 'Test',
92+
status: SessionStatus.Idle,
93+
createdAt: Date.now(),
94+
modifiedAt: Date.now(),
95+
};
96+
const initial: ISessionState = { ...createSessionState(summary), lifecycle: SessionLifecycle.Ready };
97+
const action: ISessionAction = {
98+
type: 'session/activeClientChanged' as const,
99+
session: session.toString(),
100+
activeClient: config.activeClient,
101+
} as ISessionAction;
102+
const withClient = sessionReducer(initial, action as Parameters<typeof sessionReducer>[1], () => { });
103+
this.sessionStates.set(session.toString(), withClient);
104+
}
85105
return session;
86106
}
87107

@@ -2491,13 +2511,14 @@ suite('AgentHostChatContribution', () => {
24912511
fire({ type: 'session/turnComplete', session, turnId } as ISessionAction);
24922512
await turnPromise;
24932513

2494-
const activeClientAction = agentHostService.dispatchedActions.find(
2495-
d => d.action.type === 'session/activeClientChanged'
2496-
);
2497-
assert.ok(activeClientAction, 'should dispatch activeClientChanged');
2498-
const ac = activeClientAction!.action as { activeClient: { customizations?: ICustomizationRef[] } };
2499-
assert.strictEqual(ac.activeClient.customizations?.length, 1);
2500-
assert.strictEqual(ac.activeClient.customizations?.[0].uri, 'file:///plugin-a');
2514+
// The active-client claim is now threaded through createSession
2515+
// rather than dispatched separately, so assert on createSessionCalls.
2516+
const createCall = agentHostService.createSessionCalls.at(-1);
2517+
assert.ok(createCall?.activeClient, 'createSession should carry activeClient');
2518+
assert.strictEqual(createCall!.activeClient!.clientId, agentHostService.clientId);
2519+
assert.ok(Array.isArray(createCall!.activeClient!.tools), 'activeClient.tools should be a defined array');
2520+
assert.strictEqual(createCall!.activeClient!.customizations?.length, 1);
2521+
assert.strictEqual(createCall!.activeClient!.customizations?.[0].uri, 'file:///plugin-a');
25012522
});
25022523

25032524
test('re-dispatches activeClientChanged when customizations observable changes', async () => {

0 commit comments

Comments
 (0)