Skip to content

Commit e78dfbd

Browse files
connor4312connor4312
authored
agentHost: fix other cases of client-provided tools getting stuck (#312942)
* agentHost: fix other cases of client-provided tools getting stuck * comments and ci
1 parent 6b2fbaf commit e78dfbd

9 files changed

Lines changed: 757 additions & 87 deletions

File tree

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,18 @@ export class AgentSideEffects extends Disposable {
124124
const agents = this._options.agents.read(reader);
125125
this._publishAgentInfos(agents, reader);
126126
}));
127+
128+
// Server-dispatched SessionToolCallComplete actions (e.g. from
129+
// the disconnect timeout in ProtocolServerHandler) bypass
130+
// handleAction, so the agent's SDK deferred never resolves.
131+
// Listen for these envelopes and notify the agent directly.
132+
this._register(this._stateManager.onDidEmitEnvelope(envelope => {
133+
if (!envelope.origin && envelope.action.type === ActionType.SessionToolCallComplete) {
134+
const action = envelope.action;
135+
const agent = this._options.getAgent(action.session);
136+
agent?.onClientToolCallComplete(URI.parse(action.session), action.toolCallId, action.result);
137+
}
138+
}));
127139
}
128140

129141
/**
@@ -622,7 +634,7 @@ export class AgentSideEffects extends Disposable {
622634
if (autoApproval !== undefined) {
623635
this._toolCallAgents.delete(`${sessionKey}:${e.toolCallId}`);
624636
agent.respondToPermissionRequest(e.toolCallId, true);
625-
return;
637+
e = { ...e, confirmationTitle: undefined }; // don't trigger confirmation
626638
}
627639
this._stateManager.dispatchServerAction(
628640
this._permissionManager.createToolReadyAction(e, sessionKey, turnId)

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

Lines changed: 108 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,13 +31,15 @@ import {
3131
type ReconnectParams,
3232
type IStateSnapshot,
3333
} from '../common/state/sessionProtocol.js';
34-
import { ROOT_STATE_URI, SessionStatus } from '../common/state/sessionState.js';
34+
import { ResponsePartKind, ROOT_STATE_URI, SessionStatus, ToolCallConfirmationReason, ToolCallStatus, ToolResultContentType, type SessionState } from '../common/state/sessionState.js';
3535
import type { IProtocolServer, IProtocolTransport } from '../common/state/sessionTransport.js';
3636
import { AgentHostStateManager } from './agentHostStateManager.js';
3737

3838
/** Default capacity of the server-side action replay buffer. */
3939
const REPLAY_BUFFER_CAPACITY = 1000;
4040

41+
const CLIENT_TOOL_CALL_DISCONNECT_TIMEOUT = 30_000;
42+
4143
/** Build a JSON-RPC success response suitable for transport.send(). */
4244
function jsonRpcSuccess(id: number, result: unknown): JsonRpcResponse {
4345
return { jsonrpc: '2.0', id, result };
@@ -100,6 +102,7 @@ export class ProtocolServerHandler extends Disposable {
100102

101103
private readonly _clients = new Map<string, IConnectedClient>();
102104
private readonly _replayBuffer: ActionEnvelope[] = [];
105+
private readonly _clientToolCallDisconnectTimeouts = new Map<string, ReturnType<typeof setTimeout>>();
103106

104107
private readonly _onDidChangeConnectionCount = this._register(new Emitter<number>());
105108

@@ -206,6 +209,7 @@ export class ProtocolServerHandler extends Disposable {
206209
this._logService.info(`[ProtocolServer] Client disconnected: ${client.clientId}, subscriptions=${client.subscriptions.size}`);
207210
this._clients.delete(client.clientId);
208211
this._rejectPendingReverseRequests(client.clientId);
212+
this._handleClientDisconnected(client.clientId);
209213
this._onDidChangeConnectionCount.fire(this._clients.size);
210214
}
211215
disposables.dispose();
@@ -256,6 +260,7 @@ export class ProtocolServerHandler extends Disposable {
256260
if (snapshot) {
257261
snapshots.push(snapshot);
258262
client.subscriptions.add(uri.toString());
263+
this._clearClientToolCallDisconnectTimeout(params.clientId, uri.toString());
259264
}
260265
}
261266
}
@@ -295,6 +300,7 @@ export class ProtocolServerHandler extends Disposable {
295300
const actions: ActionEnvelope[] = [];
296301
for (const sub of params.subscriptions) {
297302
client.subscriptions.add(sub.toString());
303+
this._clearClientToolCallDisconnectTimeout(params.clientId, sub.toString());
298304
}
299305
for (const envelope of this._replayBuffer) {
300306
if (envelope.serverSeq > params.lastSeenServerSeq) {
@@ -311,12 +317,108 @@ export class ProtocolServerHandler extends Disposable {
311317
if (snapshot) {
312318
snapshots.push(snapshot);
313319
client.subscriptions.add(sub);
320+
this._clearClientToolCallDisconnectTimeout(params.clientId, sub);
314321
}
315322
}
316323
return { client, response: { type: 'snapshot', snapshots } };
317324
}
318325
}
319326

327+
private _handleClientDisconnected(clientId: string): void {
328+
for (const session of this._stateManager.getSessionUris()) {
329+
const state = this._stateManager.getSessionState(session);
330+
const ownsPendingToolCall = state ? this._hasPendingClientToolCall(state, clientId) : false;
331+
if (state?.activeClient?.clientId === clientId) {
332+
this._stateManager.dispatchServerAction({
333+
type: ActionType.SessionActiveClientChanged,
334+
session,
335+
activeClient: null,
336+
});
337+
}
338+
if (state?.activeClient?.clientId === clientId || ownsPendingToolCall) {
339+
this._startClientToolCallDisconnectTimeout(clientId, session);
340+
}
341+
}
342+
}
343+
344+
private _hasPendingClientToolCall(state: ReturnType<AgentHostStateManager['getSessionState']>, clientId: string): boolean {
345+
const activeTurn = state?.activeTurn;
346+
if (!activeTurn) {
347+
return false;
348+
}
349+
return activeTurn.responseParts.some(part => part.kind === ResponsePartKind.ToolCall
350+
&& part.toolCall.toolClientId === clientId
351+
&& (part.toolCall.status === ToolCallStatus.Streaming || part.toolCall.status === ToolCallStatus.Running || part.toolCall.status === ToolCallStatus.PendingConfirmation));
352+
}
353+
354+
private _hasReplacementActiveClientTool(state: SessionState, clientId: string, toolName: string): boolean {
355+
const activeClient = state.activeClient;
356+
return activeClient !== undefined
357+
&& activeClient.clientId !== clientId
358+
&& activeClient.tools.some(tool => tool.name === toolName);
359+
}
360+
361+
private _startClientToolCallDisconnectTimeout(clientId: string, session: string): void {
362+
this._clearClientToolCallDisconnectTimeout(clientId, session);
363+
const key = this._clientToolCallDisconnectTimeoutKey(clientId, session);
364+
this._clientToolCallDisconnectTimeouts.set(key, setTimeout(() => {
365+
this._clientToolCallDisconnectTimeouts.delete(key);
366+
this._completeDisconnectedClientToolCalls(clientId, session);
367+
}, CLIENT_TOOL_CALL_DISCONNECT_TIMEOUT));
368+
}
369+
370+
private _clearClientToolCallDisconnectTimeout(clientId: string, session: string): void {
371+
const key = this._clientToolCallDisconnectTimeoutKey(clientId, session);
372+
const timeout = this._clientToolCallDisconnectTimeouts.get(key);
373+
if (timeout) {
374+
clearTimeout(timeout);
375+
this._clientToolCallDisconnectTimeouts.delete(key);
376+
}
377+
}
378+
379+
private _clientToolCallDisconnectTimeoutKey(clientId: string, session: string): string {
380+
return `${clientId}\n${session}`;
381+
}
382+
383+
private _completeDisconnectedClientToolCalls(clientId: string, session: string): void {
384+
const state = this._stateManager.getSessionState(session);
385+
const activeTurn = state?.activeTurn;
386+
if (!activeTurn) {
387+
return;
388+
}
389+
for (const part of activeTurn.responseParts) {
390+
if (part.kind !== ResponsePartKind.ToolCall) {
391+
continue;
392+
}
393+
const toolCall = part.toolCall;
394+
if (toolCall.toolClientId === clientId && (toolCall.status === ToolCallStatus.Streaming || toolCall.status === ToolCallStatus.Running || toolCall.status === ToolCallStatus.PendingConfirmation)) {
395+
const mayRetryWithReplacementClient = this._hasReplacementActiveClientTool(state, clientId, toolCall.toolName);
396+
if (toolCall.status === ToolCallStatus.Streaming) {
397+
this._stateManager.dispatchServerAction({
398+
type: ActionType.SessionToolCallReady,
399+
session,
400+
turnId: activeTurn.id,
401+
toolCallId: toolCall.toolCallId,
402+
invocationMessage: toolCall.invocationMessage ?? toolCall.displayName,
403+
confirmed: ToolCallConfirmationReason.NotNeeded,
404+
});
405+
}
406+
this._stateManager.dispatchServerAction({
407+
type: ActionType.SessionToolCallComplete,
408+
session,
409+
turnId: activeTurn.id,
410+
toolCallId: toolCall.toolCallId,
411+
result: {
412+
success: false,
413+
pastTenseMessage: `${toolCall.displayName} failed`,
414+
...(mayRetryWithReplacementClient ? { content: [{ type: ToolResultContentType.Text, text: `The client that was running ${toolCall.displayName} disconnected, but another active client now provides ${toolCall.displayName}. You may try calling the tool again.` }] } : {}),
415+
error: { message: `Client ${clientId} disconnected before completing ${toolCall.displayName}` },
416+
},
417+
});
418+
}
419+
}
420+
}
421+
320422
// ---- Requests (expect a response) ---------------------------------------
321423

322424
/**
@@ -328,6 +430,7 @@ export class ProtocolServerHandler extends Disposable {
328430
try {
329431
const snapshot = await this._agentService.subscribe(URI.parse(params.resource));
330432
client.subscriptions.add(params.resource);
433+
this._clearClientToolCallDisconnectTimeout(client.clientId, params.resource);
331434
return { snapshot };
332435
} catch (err) {
333436
if (err instanceof ProtocolError) {
@@ -606,6 +709,10 @@ export class ProtocolServerHandler extends Disposable {
606709
pending.reject(new Error('ProtocolServerHandler disposed'));
607710
}
608711
this._pendingReverseRequests.clear();
712+
for (const timeout of this._clientToolCallDisconnectTimeouts.values()) {
713+
clearTimeout(timeout);
714+
}
715+
this._clientToolCallDisconnectTimeouts.clear();
609716
this._replayBuffer.length = 0;
610717
super.dispose();
611718
}

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -961,6 +961,82 @@ suite('AgentSideEffects', () => {
961961
});
962962
});
963963

964+
// ---- tool_ready progress dispatch -----------------------------------
965+
966+
suite('tool_ready dispatches progress actions to advance tool call state', () => {
967+
968+
test('tool_ready for a non-permission tool dispatches SessionToolCallReady and advances state from Streaming to Running', () => {
969+
setupSession();
970+
startTurn('turn-1');
971+
disposables.add(sideEffects.registerProgressListener(agent));
972+
973+
// tool_start puts the tool call into Streaming state
974+
agent.fireProgress({
975+
session: sessionUri,
976+
type: 'tool_start',
977+
toolCallId: 'tc-ready-1',
978+
toolName: 'runTask',
979+
displayName: 'Run Task',
980+
invocationMessage: 'Running task...',
981+
toolClientId: 'test-client',
982+
});
983+
984+
const stateAfterStart = stateManager.getSessionState(sessionUri.toString());
985+
const partAfterStart = stateAfterStart?.activeTurn?.responseParts[0];
986+
assert.strictEqual(partAfterStart?.kind, ResponsePartKind.ToolCall);
987+
assert.strictEqual(partAfterStart?.kind === ResponsePartKind.ToolCall ? partAfterStart.toolCall.status : undefined, ToolCallStatus.Streaming);
988+
989+
// tool_ready without confirmationTitle should dispatch the ready
990+
// action and advance the tool call to Running
991+
agent.fireProgress({
992+
session: sessionUri,
993+
type: 'tool_ready',
994+
toolCallId: 'tc-ready-1',
995+
invocationMessage: 'Run Task',
996+
toolInput: '{"task":"build"}',
997+
});
998+
999+
const stateAfterReady = stateManager.getSessionState(sessionUri.toString());
1000+
const partAfterReady = stateAfterReady?.activeTurn?.responseParts[0];
1001+
assert.strictEqual(partAfterReady?.kind, ResponsePartKind.ToolCall);
1002+
assert.strictEqual(partAfterReady?.kind === ResponsePartKind.ToolCall ? partAfterReady.toolCall.status : undefined, ToolCallStatus.Running,
1003+
'tool call should advance from Streaming to Running after tool_ready');
1004+
});
1005+
1006+
test('tool_ready for a permission-gated tool dispatches SessionToolCallReady and advances state to PendingConfirmation', () => {
1007+
setupSession();
1008+
startTurn('turn-1');
1009+
disposables.add(sideEffects.registerProgressListener(agent));
1010+
1011+
agent.fireProgress({
1012+
session: sessionUri,
1013+
type: 'tool_start',
1014+
toolCallId: 'tc-perm-1',
1015+
toolName: 'write',
1016+
displayName: 'Write File',
1017+
invocationMessage: 'Writing file...',
1018+
toolClientId: 'test-client',
1019+
});
1020+
1021+
// tool_ready with confirmationTitle should dispatch the ready
1022+
// action and advance the tool call to PendingConfirmation
1023+
agent.fireProgress({
1024+
session: sessionUri,
1025+
type: 'tool_ready',
1026+
toolCallId: 'tc-perm-1',
1027+
invocationMessage: 'Write .env',
1028+
confirmationTitle: 'Write .env',
1029+
toolInput: '{"path":".env"}',
1030+
});
1031+
1032+
const state = stateManager.getSessionState(sessionUri.toString());
1033+
const part = state?.activeTurn?.responseParts[0];
1034+
assert.strictEqual(part?.kind, ResponsePartKind.ToolCall);
1035+
assert.strictEqual(part?.kind === ResponsePartKind.ToolCall ? part.toolCall.status : undefined, ToolCallStatus.PendingConfirmation,
1036+
'tool call should advance to PendingConfirmation for permission-gated tool_ready');
1037+
});
1038+
});
1039+
9641040
// ---- Session-level auto-approve (config) ----------------------------
9651041

9661042
suite('session config auto-approve', () => {

0 commit comments

Comments
 (0)