Fix premature disconnect force-fail of a live client tool call

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

@@ -193,11 +193,14 @@ interface IClientRecord {
 	/** Live connection while connected; `undefined` after disconnect (record retained for the grace window). */
 	connection: IConnectedClient | undefined;
 	/**
-	 * Epoch ms the client was last seen connected (handshake or disconnect).
-	 * `undefined` when the client has never connected. Drives the
-	 * disconnect-timeout grace window: a pending client tool call fails
-	 * `CLIENT_TOOL_CALL_DISCONNECT_TIMEOUT` ms after this point, never instantly
-	 * and never never.
+	 * Epoch ms a live frame was last received from this client (handshake,
+	 * disconnect, or any subsequent inbound message). `undefined` when the
+	 * client has never been seen. Drives the disconnect-timeout grace window:
+	 * a pending client tool call fails `CLIENT_TOOL_CALL_DISCONNECT_TIMEOUT` ms
+	 * after the client last sent anything — so a client that is still actively
+	 * sending frames is treated as alive even if its {@link connection} pointer
+	 * was transiently cleared (e.g. a reused clientId whose earlier transport
+	 * closed while a newer one is still live).
 	 */
 	lastSeenAt: number | undefined;
 	/**
@@ -338,6 +341,14 @@ export class ProtocolServerHandler extends Disposable {
 		let client: IConnectedClient | undefined;
 
 		disposables.add(transport.onMessage(msg => {
+			// Any inbound frame from an established client is fresh proof of
+			// life. Keep the per-client last-seen timestamp current from real
+			// traffic so the disconnect-grace machinery can distinguish a
+			// genuinely gone client from one that is still active (see
+			// IClientRecord.lastSeenAt).
+			if (client) {
+				this._markClientSeen(client.clientId);
+			}
 			if (isJsonRpcRequest(msg)) {
 				this._logService.trace(`[ProtocolServer] request: method=${msg.method} id=${msg.id}`);
 
@@ -798,11 +809,55 @@ export class ProtocolServerHandler extends Disposable {
 		const elapsed = Date.now() - record.lastSeenAt;
 		const delay = Math.max(0, CLIENT_TOOL_CALL_DISCONNECT_TIMEOUT - elapsed);
 		record.disconnectTimeouts.set(session, disposableTimeout(() => {
-			record.disconnectTimeouts.deleteAndDispose(session);
-			this._completeDisconnectedClientToolCalls(clientId, session);
+			// Re-verify the client is genuinely gone before force-failing its
+			// pending tool calls. A client that is still sending frames is
+			// alive and may yet deliver the result, even if its `connection`
+			// pointer reads undefined (e.g. a reused clientId whose earlier
+			// transport closed while this one is still live).
+			if (this._isClientGraceExpired(record)) {
+				record.disconnectTimeouts.deleteAndDispose(session);
+				this._completeDisconnectedClientToolCalls(clientId, session);
+				return;
+			}
+			// The client is still alive. Keep the grace timer running only
+			// while it still owns a pending tool call; otherwise let it go so
+			// we don't leave a perpetual timer ticking for an active client.
+			const state = this._stateManager.getSessionState(session);
+			if (state && this._hasPendingClientToolCall(state, clientId)) {
+				this._startClientToolCallDisconnectTimeout(clientId, session);
+			} else {
+				record.disconnectTimeouts.deleteAndDispose(session);
+			}
 		}, delay));
 	}
 
+	/**
+	 * Record that a live frame was just received from `clientId`. Keeps
+	 * {@link IClientRecord.lastSeenAt} current from real traffic (not just the
+	 * handshake), so the disconnect-grace machinery can tell a genuinely gone
+	 * client from one that is still active. Only updates an existing record;
+	 * never creates one for an unknown client.
+	 */
+	private _markClientSeen(clientId: string): void {
+		const record = this._clients.get(clientId);
+		if (record) {
+			record.lastSeenAt = Date.now();
+		}
+	}
+
+	/**
+	 * True when no live frame has been received from `record`'s client within
+	 * the disconnect grace window — i.e. the client is genuinely gone, not
+	 * merely between transports. A never-seen record counts as expired. The
+	 * decision is intentionally based on traffic recency rather than the
+	 * `connection` pointer, which can transiently read undefined for a client
+	 * that is still alive on another transport.
+	 */
+	private _isClientGraceExpired(record: IClientRecord): boolean {
+		return record.lastSeenAt === undefined
+			|| (Date.now() - record.lastSeenAt) >= CLIENT_TOOL_CALL_DISCONNECT_TIMEOUT;
+	}
+
 	/**
 	 * Scan a session for pending client tool calls whose owning client is not
 	 * currently connected, and arm the disconnect timeout for each such owner.

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

@@ -1000,6 +1000,116 @@ suite('ProtocolServerHandler', () => {
 		});
 	});
 
+	test('owned tool call is not failed while the owning client stays active on another transport', () => {
+		return runWithFakedTimers({ useFakeTimers: true }, async () => {
+			stateManager.createSession(makeSessionSummary());
+			stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, });
+			stateManager.dispatchServerAction(sessionUri, {
+				type: ActionType.SessionActiveClientChanged,
+				activeClient: {
+					clientId: 'client-tools',
+					tools: [{ name: 'runTask', description: 'Runs a task' }]
+				},
+			});
+			stateManager.dispatchServerAction(sessionUri, {
+				type: ActionType.ChatTurnStarted,
+				turnId: 'turn-1',
+				message: { text: 'run it', origin: { kind: MessageKind.User } },
+			});
+			stateManager.dispatchServerAction(sessionUri, {
+				type: ActionType.ChatToolCallStart,
+				turnId: 'turn-1',
+				toolCallId: 'tool-1',
+				toolName: 'runTask',
+				displayName: 'Run Task',
+				contributor: { kind: ToolCallContributorKind.Client, clientId: 'client-tools' },
+			});
+
+			// The same logical clientId is served by two transports (e.g. a
+			// window that reconnected before its previous transport's close was
+			// observed). The server tracks one connection per clientId, so the
+			// SECOND handshake becomes the tracked connection and the first is
+			// left live-but-untracked.
+			const liveTransport = connectClient('client-tools', [sessionUri]);
+			const trackedTransport = connectClient('client-tools', [sessionUri]);
+
+			// Closing the tracked transport clears the record's `connection`
+			// pointer and arms the disconnect-grace timeout for the pending
+			// tool call — even though `liveTransport` is still connected.
+			trackedTransport.simulateClose();
+
+			let part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0];
+			assert.strictEqual(part?.kind === ResponsePartKind.ToolCall ? part.toolCall.status : undefined, ToolCallStatus.Streaming);
+
+			// The live transport keeps sending frames across several grace
+			// windows. Each frame is fresh proof of life, so the tool call must
+			// never be force-failed.
+			for (let i = 0; i < 12; i++) {
+				await new Promise(r => setTimeout(r, 10_000));
+				liveTransport.simulateMessage(request(100 + i, 'ping'));
+			}
+
+			part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0];
+			assert.strictEqual(part?.kind === ResponsePartKind.ToolCall ? part.toolCall.status : undefined, ToolCallStatus.Streaming);
+
+			liveTransport.simulateClose();
+		});
+	});
+
+	test('owned tool call is failed once the client stops sending frames on every transport', () => {
+		return runWithFakedTimers({ useFakeTimers: true }, async () => {
+			stateManager.createSession(makeSessionSummary());
+			stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, });
+			stateManager.dispatchServerAction(sessionUri, {
+				type: ActionType.SessionActiveClientChanged,
+				activeClient: {
+					clientId: 'client-tools',
+					tools: [{ name: 'runTask', description: 'Runs a task' }]
+				},
+			});
+			stateManager.dispatchServerAction(sessionUri, {
+				type: ActionType.ChatTurnStarted,
+				turnId: 'turn-1',
+				message: { text: 'run it', origin: { kind: MessageKind.User } },
+			});
+			stateManager.dispatchServerAction(sessionUri, {
+				type: ActionType.ChatToolCallStart,
+				turnId: 'turn-1',
+				toolCallId: 'tool-1',
+				toolName: 'runTask',
+				displayName: 'Run Task',
+				contributor: { kind: ToolCallContributorKind.Client, clientId: 'client-tools' },
+			});
+
+			const liveTransport = connectClient('client-tools', [sessionUri]);
+			const trackedTransport = connectClient('client-tools', [sessionUri]);
+			trackedTransport.simulateClose();
+
+			// The live transport sends a few frames, then goes silent: the
+			// grace machinery must still fail the call once no frame has
+			// arrived for the full window (proving the fix does not simply
+			// disable the disconnect path for live-but-untracked transports).
+			liveTransport.simulateMessage(request(100, 'ping'));
+			await new Promise(r => setTimeout(r, 10_000));
+			liveTransport.simulateMessage(request(101, 'ping'));
+
+			await new Promise(r => setTimeout(r, 30_001));
+
+			const part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0];
+			assert.deepStrictEqual(part?.kind === ResponsePartKind.ToolCall ? {
+				status: part.toolCall.status,
+				success: part.toolCall.status === ToolCallStatus.Completed ? part.toolCall.success : undefined,
+				error: part.toolCall.status === ToolCallStatus.Completed ? part.toolCall.error?.message : undefined,
+			} : undefined, {
+				status: ToolCallStatus.Completed,
+				success: false,
+				error: 'Client client-tools disconnected before completing Run Task',
+			});
+
+			liveTransport.simulateClose();
+		});
+	});
+
 	test('client reconnect without session subscription does not clear tool call disconnect timeout', () => {
 		return runWithFakedTimers({ useFakeTimers: true }, async () => {
 			stateManager.createSession(makeSessionSummary());