diff --git a/src/vs/platform/agentHost/node/protocolServerHandler.ts b/src/vs/platform/agentHost/node/protocolServerHandler.ts index 0a2334d416b121..abacd8d88dda6e 100644 --- a/src/vs/platform/agentHost/node/protocolServerHandler.ts +++ b/src/vs/platform/agentHost/node/protocolServerHandler.ts @@ -875,9 +875,8 @@ export class ProtocolServerHandler extends Disposable { * any armed timers, so this is a no-op when the client is active. The delay * is the remaining grace measured from when the client disconnected — so a * client that disconnected a while before the call was issued gets the - * residual window rather than a fresh one, and a stamp from a long-dead - * client fails promptly. A never-connected client has its grace clock pinned - * to the first arm, so re-arms triggered by later orphaned tool calls in the + * residual window rather than a fresh one, and a stamp from a long-disconnected + * client fails promptly. Re-arms triggered by later orphaned tool calls in the * same session shrink the remaining window instead of resetting it. */ private _startClientToolCallDisconnectTimeout(clientId: string, session: string, chatChannel: string): void { @@ -895,20 +894,21 @@ export class ProtocolServerHandler extends Disposable { } /** - * Scan a session for pending client tool calls whose owning client is not - * currently connected, and arm the disconnect timeout for each such owner. + * Scan a chat for pending client tool calls owned by a disconnected client + * of this protocol server, and arm the disconnect timeout for each owner. * Called when a `ChatToolCallStart` / `ChatToolCallReady` envelope is * observed — covering calls issued for an already-gone client, which the * live disconnect path never sees. Ownerless client tool calls (no client * connected at stamp time) are failed immediately by the provider, so they - * never reach a pending state here. + * never reach a pending state here. Unknown client ids are ignored because + * they may belong to another transport such as local IPC. */ private _checkOrphanedClientToolCalls(session: string, chatChannel: string): void { const state = this._stateManager.getSessionState(chatChannel); const orphanOwners = new Set(); for (const { clientId } of this._pendingClientToolCalls(state)) { const ownerRecord = this._clients.get(clientId); - if (ownerRecord?.state !== 'active') { + if (ownerRecord?.state === 'grace') { orphanOwners.add(clientId); } } diff --git a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts index 6bab185da34173..eb7f1ad08cd787 100644 --- a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts +++ b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts @@ -1360,27 +1360,26 @@ suite('ProtocolServerHandler', () => { }); }); - test('client tool call stamped for a never-connected client fails after the grace period', () => { + test('client tool call stamped for a disconnected protocol client fails after the grace period', () => { return runWithFakedTimers({ useFakeTimers: true }, async () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); const chatUri = buildDefaultChatUri(sessionUri); + const transport = connectClient('disconnected-client', [sessionUri]); + transport.simulateClose(); stateManager.dispatchServerAction(chatUri, { type: ActionType.ChatTurnStarted, turnId: 'turn-1', startedAt: '2025-01-01T00:00:00.000Z', message: { text: 'run it', origin: { kind: MessageKind.User } }, }); - // Tool call stamped for a clientId that never connected (e.g. a - // stale stamp from a long-dead window). No disconnect event ever - // fires for it; the issuance-time orphan check must arm the timeout. stateManager.dispatchServerAction(chatUri, { type: ActionType.ChatToolCallStart, turnId: 'turn-1', toolCallId: 'tool-1', toolName: 'runTask', displayName: 'Run Task', - contributor: { kind: ToolCallContributorKind.Client, clientId: 'ghost-client' }, + contributor: { kind: ToolCallContributorKind.Client, clientId: 'disconnected-client' }, }); let part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0]; @@ -1396,15 +1395,50 @@ suite('ProtocolServerHandler', () => { } : undefined, { status: ToolCallStatus.Completed, success: false, - error: 'Client ghost-client disconnected before completing Run Task', + error: 'Client disconnected-client disconnected before completing Run Task', }); }); }); + test('client tool call owned by an active local IPC client is not treated as orphaned', () => { + return runWithFakedTimers({ useFakeTimers: true }, async () => { + stateManager.createSession(makeSessionSummary()); + stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); + stateManager.dispatchServerAction(sessionUri, { + type: ActionType.SessionActiveClientSet, + activeClient: { + clientId: 'local-client', + tools: [{ name: 'runTask', description: 'Runs a task' }] + }, + }); + stateManager.dispatchServerAction(defaultChatUri, { + type: ActionType.ChatTurnStarted, + turnId: 'turn-1', + startedAt: '2025-01-01T00:00:00.000Z', + message: { text: 'run it', origin: { kind: MessageKind.User } }, + }); + stateManager.dispatchServerAction(defaultChatUri, { + type: ActionType.ChatToolCallStart, + turnId: 'turn-1', + toolCallId: 'tool-1', + toolName: 'runTask', + displayName: 'Run Task', + contributor: { kind: ToolCallContributorKind.Client, clientId: 'local-client' }, + }); + + await new Promise(r => setTimeout(r, 30_001)); + + const part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0]; + assert.strictEqual(part?.kind === ResponsePartKind.ToolCall ? part.toolCall.status : undefined, ToolCallStatus.Streaming); + }); + }); + test('orphaned client tool call timeout is cleared when the owning client connects within the window', () => { return runWithFakedTimers({ useFakeTimers: true }, async () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); + const transport = connectClient('late-client', [sessionUri]); + transport.simulateClose(); stateManager.dispatchServerAction(defaultChatUri, { type: ActionType.ChatTurnStarted, turnId: 'turn-1', @@ -1420,8 +1454,7 @@ suite('ProtocolServerHandler', () => { contributor: { kind: ToolCallContributorKind.Client, clientId: 'late-client' }, }); - // The owning client connects (and subscribes) within the grace - // window — the subscribe path clears the armed timeout. + // The owning client reconnects within the grace window. connectClient('late-client', [sessionUri]); await new Promise(r => setTimeout(r, 30_001)); @@ -1435,27 +1468,25 @@ suite('ProtocolServerHandler', () => { return runWithFakedTimers({ useFakeTimers: true }, async () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); + const transport = connectClient('disconnected-client', [sessionUri]); + transport.simulateClose(); stateManager.dispatchServerAction(defaultChatUri, { type: ActionType.ChatTurnStarted, turnId: 'turn-1', startedAt: '2025-01-01T00:00:00.000Z', message: { text: 'run it', origin: { kind: MessageKind.User } }, }); - // First orphaned tool call (owner never connected) arms the grace timer. + // First orphaned tool call arms the grace timer. stateManager.dispatchServerAction(defaultChatUri, { type: ActionType.ChatToolCallStart, turnId: 'turn-1', toolCallId: 'tool-1', toolName: 'runTask', displayName: 'Run Task', - contributor: { kind: ToolCallContributorKind.Client, clientId: 'ghost-client' }, + contributor: { kind: ToolCallContributorKind.Client, clientId: 'disconnected-client' }, }); - // A second call for the same never-connected owner arrives partway - // through the window and re-arms the (shared) timer. The grace clock - // is pinned to the first arm, so the re-arm must NOT reset the - // deadline — otherwise the first call could be kept alive - // indefinitely. + // Re-arming for a later call must retain the original deadline. await new Promise(r => setTimeout(r, 20_000)); stateManager.dispatchServerAction(defaultChatUri, { type: ActionType.ChatToolCallStart, @@ -1463,7 +1494,7 @@ suite('ProtocolServerHandler', () => { toolCallId: 'tool-2', toolName: 'runTask', displayName: 'Run Task', - contributor: { kind: ToolCallContributorKind.Client, clientId: 'ghost-client' }, + contributor: { kind: ToolCallContributorKind.Client, clientId: 'disconnected-client' }, }); // 31s after the FIRST call: both must have failed.