From 220f6ee4aaf1be9ccf107291140fdeca9840fe5c Mon Sep 17 00:00:00 2001 From: Rob Lourens Date: Fri, 19 Jun 2026 15:54:34 -0700 Subject: [PATCH 1/2] Fix premature disconnect force-fail of a live client tool call A client tool call could be force-failed as "disconnected" while the owning client was in fact still connected and actively sending frames, destroying the tool call (and its pending confirmation) before the user ever saw the confirmation prompt. The disconnect-grace machinery decided a client was gone from two signals: `IClientRecord.connection === undefined` (arm condition) and `lastSeenAt` (grace-window delay). Two gaps made it misfire: - `lastSeenAt` was only updated at handshake/disconnect, never on ordinary inbound frames, so for a long-lived chatty client it was stale by minutes. The arm delay `max(0, TIMEOUT - elapsed)` then collapsed to 0 and the timeout fired immediately. - `connection` is a single pointer per clientId. When a clientId is reused across overlapping transports, closing the most-recent transport clears `connection` even though another transport for the same client is still live, so `connection === undefined` does not imply the client is gone. Fix: track liveness from real traffic by bumping `lastSeenAt` on every inbound frame, and re-verify liveness when the timeout fires (recency-based, not connection-based) before force-failing. If the client is still alive, re-arm only while it still owns a pending tool call so an active client never leaves a perpetual timer running. Adds regression tests: a live client on a second transport survives despite `connection === undefined`, and a silent owner still fails after the grace window. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../agentHost/node/protocolServerHandler.ts | 69 +++++++++-- .../test/node/protocolServerHandler.test.ts | 110 ++++++++++++++++++ 2 files changed, 172 insertions(+), 7 deletions(-) diff --git a/src/vs/platform/agentHost/node/protocolServerHandler.ts b/src/vs/platform/agentHost/node/protocolServerHandler.ts index 0ad18cbc11b..4f60efadd4c 100644 --- a/src/vs/platform/agentHost/node/protocolServerHandler.ts +++ b/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. diff --git a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts index dcae0709357..65a37427485 100644 --- a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts +++ b/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()); From 9333d12443db958197a037c69f7e3516242b9c04 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Mon, 22 Jun 2026 13:04:42 -0700 Subject: [PATCH 2/2] agent host: handle overlapping client transports Track multiple live transports per clientId so closing a newer overlapping transport falls back to an older live one instead of treating the logical client as disconnected. This prevents pending client tool calls from being force-failed while the owning client is still connected. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../agentHost/node/protocolServerHandler.ts | 216 ++++++++---------- .../test/node/protocolServerHandler.test.ts | 132 ++++++++--- 2 files changed, 193 insertions(+), 155 deletions(-) diff --git a/src/vs/platform/agentHost/node/protocolServerHandler.ts b/src/vs/platform/agentHost/node/protocolServerHandler.ts index 4f60efadd4c..b0284f108f3 100644 --- a/src/vs/platform/agentHost/node/protocolServerHandler.ts +++ b/src/vs/platform/agentHost/node/protocolServerHandler.ts @@ -183,24 +183,25 @@ interface IConnectedClient { /** * Per-client server-side record, keyed by clientId in - * {@link ProtocolServerHandler._clients}. Unlike {@link IConnectedClient}, the - * record OUTLIVES the connection: when a client disconnects, `connection` is - * cleared to `undefined` but the record is retained (until pruned) so the - * tool-call disconnect-grace machinery can still compute the remaining window - * and hold any armed timeouts. + * {@link ProtocolServerHandler._clients}. Unlike {@link IConnectedClient}, + * the record OUTLIVES individual transports: multiple overlapping transports + * for the same logical client are held oldest-first, with the active transport + * at the end. When the last transport disconnects, the record is retained + * (until pruned) so the tool-call disconnect-grace machinery can compute the + * remaining window and hold any armed timeouts. */ interface IClientRecord { - /** Live connection while connected; `undefined` after disconnect (record retained for the grace window). */ - connection: IConnectedClient | undefined; /** - * 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). + * All live transports for this client, oldest first. The active connection + * is the last entry (most recent wins). Older entries are kept so that if a + * reconnecting client registers `A`, then `B`, then `B` closes first, we can + * fall back to `A` instead of treating the client as disconnected. + */ + readonly connections: IConnectedClient[]; + /** + * Epoch ms when the client last had no live transports. `undefined` while at + * least one connection is active, or when the client has never connected. + * Drives the disconnect-timeout grace window for disconnected records only. */ lastSeenAt: number | undefined; /** @@ -272,8 +273,8 @@ export class ProtocolServerHandler extends Disposable { /** * Per-client records keyed by clientId. Holds both connected clients - * (`connection` set) and recently-disconnected ones retained for the - * tool-call disconnect-grace window (`connection === undefined`). See + * (`connections` non-empty) and recently-disconnected ones retained for the + * tool-call disconnect-grace window (`connections.length === 0`). See * {@link IClientRecord}. */ private readonly _clients = new Map(); @@ -341,14 +342,6 @@ 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}`); @@ -423,7 +416,7 @@ export class ProtocolServerHandler extends Disposable { } } else if (isJsonRpcResponse(msg)) { const pending = this._pendingReverseRequests.get(msg.id); - if (pending) { + if (pending && pending.client === client) { this._pendingReverseRequests.delete(msg.id); if (hasKey(msg, { error: true })) { pending.reject(new ProtocolError( @@ -440,26 +433,20 @@ export class ProtocolServerHandler extends Disposable { disposables.add(transport.onClose(() => { const record = client ? this._clients.get(client.clientId) : undefined; - if (client && record && record.connection === client) { - this._logService.info(`[ProtocolServer] Client disconnected: ${client.clientId}, subscriptions=${client.subscriptions.size}`); - // Treat disconnect as an implicit unsubscribe of every channel the - // client held, so the server-side refcount can drop to zero and any - // idle restored session state can be evicted. OTLP subscriptions - // have no server-side state to release, so the per-client map is - // simply discarded. - for (const sub of client.subscriptions.values()) { - if (sub.kind === ChannelKind.State) { - this._agentService.unsubscribe(URI.parse(sub.uri), client.clientId); - } else if (sub.kind === ChannelKind.ResourceWatch) { - this._agentService.onResourceWatchUnsubscribed(sub.uri); + if (client && record) { + const connectionIndex = record.connections.indexOf(client); + if (connectionIndex !== -1) { + const subscriptionCount = client.subscriptions.size; + record.connections.splice(connectionIndex, 1); + this._releaseClientSubscriptions(client, record); + this._rejectPendingReverseRequestsForConnection(client); + if (record.connections.length === 0) { + this._logService.info(`[ProtocolServer] Client disconnected: ${client.clientId}, subscriptions=${subscriptionCount}`); + record.lastSeenAt = Date.now(); + this._handleClientDisconnected(client.clientId); + this._onDidChangeConnectionCount.fire(this._connectedClientCount); } } - client.subscriptions.clear(); - record.connection = undefined; - record.lastSeenAt = Date.now(); - this._rejectPendingReverseRequests(client.clientId); - this._handleClientDisconnected(client.clientId); - this._onDidChangeConnectionCount.fire(this._connectedClientCount); } disposables.dispose(); })); @@ -505,8 +492,8 @@ export class ProtocolServerHandler extends Disposable { disposables, }; const record = this._ensureClientRecord(params.clientId); - record.connection = client; - record.lastSeenAt = Date.now(); + record.connections.push(client); + record.lastSeenAt = undefined; this._pruneClientRecords(); this._onDidChangeConnectionCount.fire(this._connectedClientCount); @@ -622,8 +609,8 @@ export class ProtocolServerHandler extends Disposable { disposables, }; const record = this._ensureClientRecord(params.clientId); - record.connection = client; - record.lastSeenAt = Date.now(); + record.connections.push(client); + record.lastSeenAt = undefined; this._pruneClientRecords(); this._onDidChangeConnectionCount.fire(this._connectedClientCount); @@ -791,12 +778,12 @@ export class ProtocolServerHandler extends Disposable { /** * Arm (or re-arm) the per-(clientId, session) timeout that fails pending - * client tool calls owned by `clientId` if it does not (re)connect. The - * delay is the remaining grace measured from when the client was last - * seen — 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 client never seen at all has its - * grace clock pinned to the first arm, so re-arms triggered by later + * client tool calls owned by `clientId` if it does not reconnect and + * resubscribe. 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 client never seen at all + * has its grace clock pinned to the first arm, so re-arms triggered by later * orphaned tool calls in the same session shrink the remaining window * instead of resetting it. */ @@ -809,55 +796,11 @@ 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(() => { - // 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); - } + record.disconnectTimeouts.deleteAndDispose(session); + this._completeDisconnectedClientToolCalls(clientId, 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. @@ -872,7 +815,7 @@ export class ProtocolServerHandler extends Disposable { const orphanOwners = new Set(); for (const { clientId } of this._pendingClientToolCalls(state)) { const ownerRecord = this._clients.get(clientId); - if (!ownerRecord || ownerRecord.connection === undefined) { + if (!ownerRecord || ownerRecord.connections.length === 0) { orphanOwners.add(clientId); } } @@ -883,22 +826,54 @@ export class ProtocolServerHandler extends Disposable { /** * Get the existing per-client record or create an empty one. A freshly - * created record has no connection and `lastSeenAt === undefined`. + * created record has no connections and `lastSeenAt === undefined`. */ private _ensureClientRecord(clientId: string): IClientRecord { let record = this._clients.get(clientId); if (!record) { - record = { connection: undefined, lastSeenAt: undefined, disconnectTimeouts: new DisposableMap() }; + record = { connections: [], lastSeenAt: undefined, disconnectTimeouts: new DisposableMap() }; this._clients.set(clientId, record); } return record; } + private _getActiveClient(clientId: string): IConnectedClient | undefined { + const connections = this._clients.get(clientId)?.connections; + return connections?.[connections.length - 1]; + } + + private _getActiveClientFromRecord(record: IClientRecord): IConnectedClient | undefined { + return record.connections[record.connections.length - 1]; + } + + private _releaseClientSubscriptions(client: IConnectedClient, record: IClientRecord): void { + for (const sub of client.subscriptions.values()) { + if (sub.kind === ChannelKind.State) { + if (this._hasSubscriptionInOtherConnection(record, client, sub.uri)) { + continue; + } + this._agentService.unsubscribe(URI.parse(sub.uri), client.clientId); + } else if (sub.kind === ChannelKind.ResourceWatch) { + this._agentService.onResourceWatchUnsubscribed(sub.uri); + } + } + client.subscriptions.clear(); + } + + private _hasSubscriptionInOtherConnection(record: IClientRecord, client: IConnectedClient, uri: string): boolean { + for (const other of record.connections) { + if (other !== client && other.subscriptions.has(uri)) { + return true; + } + } + return false; + } + /** Number of records that currently hold a live connection. */ private get _connectedClientCount(): number { let count = 0; for (const record of this._clients.values()) { - if (record.connection) { + if (record.connections.length > 0) { count++; } } @@ -916,7 +891,7 @@ export class ProtocolServerHandler extends Disposable { private _pruneClientRecords(): void { const cutoff = Date.now() - CLIENT_TOOL_CALL_DISCONNECT_TIMEOUT * 10; for (const [clientId, record] of this._clients) { - if (record.connection === undefined + if (record.connections.length === 0 && record.disconnectTimeouts.size === 0 && (record.lastSeenAt === undefined || record.lastSeenAt < cutoff)) { this._clients.delete(clientId); @@ -1220,7 +1195,7 @@ export class ProtocolServerHandler extends Disposable { // ---- Reverse RPC (server → client requests) ---------------------------- private _reverseRequestId = 0; - private readonly _pendingReverseRequests = new Map void; reject: (reason: unknown) => void }>(); + private readonly _pendingReverseRequests = new Map void; reject: (reason: unknown) => void }>(); /** * Sends a JSON-RPC request to a connected client and waits for the response. @@ -1228,26 +1203,27 @@ export class ProtocolServerHandler extends Disposable { * Rejects if the client disconnects or the server is disposed. */ private _sendReverseRequest(clientId: string, method: string, params: unknown): Promise { - const client = this._clients.get(clientId)?.connection; + const client = this._getActiveClient(clientId); if (!client) { return Promise.reject(new Error(`Client ${clientId} is not connected`)); } const id = ++this._reverseRequestId; return new Promise((resolve, reject) => { - this._pendingReverseRequests.set(id, { clientId, resolve: resolve as (value: unknown) => void, reject }); + this._pendingReverseRequests.set(id, { client, resolve: resolve as (value: unknown) => void, reject }); const request: JsonRpcRequest = { jsonrpc: '2.0', id, method, params }; client.transport.send(request); }); } /** - * Rejects and clears all pending reverse-RPC requests for a given client. + * Rejects and clears all pending reverse-RPC requests sent over a given + * connection. */ - private _rejectPendingReverseRequests(clientId: string): void { + private _rejectPendingReverseRequestsForConnection(client: IConnectedClient): void { for (const [id, pending] of this._pendingReverseRequests) { - if (pending.clientId === clientId) { + if (pending.client === client) { this._pendingReverseRequests.delete(id); - pending.reject(new Error(`Client ${clientId} disconnected`)); + pending.reject(new Error(`Client ${client.clientId} disconnected`)); } } } @@ -1321,7 +1297,7 @@ export class ProtocolServerHandler extends Disposable { this._logService.trace(`[ProtocolServer] Broadcasting action: ${envelope.action.type}`); const msg: AhpServerNotification<'action'> = { jsonrpc: '2.0', method: 'action', params: envelope }; for (const record of this._clients.values()) { - const client = record.connection; + const client = this._getActiveClientFromRecord(record); if (client && this._isRelevantToClient(client, envelope)) { client.transport.send(msg); } @@ -1336,7 +1312,7 @@ export class ProtocolServerHandler extends Disposable { // eslint-disable-next-line local/code-no-dangerous-type-assertions const msg = { jsonrpc: '2.0', method: type, params } as AhpServerNotification; for (const record of this._clients.values()) { - record.connection?.transport.send(msg); + this._getActiveClientFromRecord(record)?.transport.send(msg); } } @@ -1357,7 +1333,7 @@ export class ProtocolServerHandler extends Disposable { // eslint-disable-next-line local/code-no-dangerous-type-assertions const msg = { jsonrpc: '2.0' as const, method: notification.method, params } as unknown as AhpServerNotification; for (const record of this._clients.values()) { - record.connection?.transport.send(msg); + this._getActiveClientFromRecord(record)?.transport.send(msg); } } @@ -1380,6 +1356,10 @@ export class ProtocolServerHandler extends Disposable { } client.subscriptions.delete(classified.uri); if (sub.kind === ChannelKind.State) { + const record = this._clients.get(client.clientId); + if (record && this._hasSubscriptionInOtherConnection(record, client, sub.uri)) { + return; + } this._agentService.unsubscribe(URI.parse(sub.uri), client.clientId); } else if (sub.kind === ChannelKind.ResourceWatch) { this._agentService.onResourceWatchUnsubscribed(sub.uri); @@ -1396,7 +1376,7 @@ export class ProtocolServerHandler extends Disposable { private _broadcastOtlpLog(record: IOtlpLogRecord): void { const payload = toResourceLogsPayload(record); for (const clientRecord of this._clients.values()) { - const client = clientRecord.connection; + const client = this._getActiveClientFromRecord(clientRecord); if (!client) { continue; } @@ -1437,7 +1417,9 @@ export class ProtocolServerHandler extends Disposable { override dispose(): void { for (const record of this._clients.values()) { - record.connection?.disposables.dispose(); + for (const connection of [...record.connections]) { + connection.disposables.dispose(); + } record.disconnectTimeouts.dispose(); } this._clients.clear(); diff --git a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts index 65a37427485..7a9df3e27ea 100644 --- a/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts +++ b/src/vs/platform/agentHost/test/node/protocolServerHandler.test.ts @@ -881,6 +881,54 @@ suite('ProtocolServerHandler', () => { assert.deepStrictEqual(result, [['after-reconnect.txt', FileType.File]]); }); + test('overlapping reconnect keeps earlier reverse-RPC requests alive until that transport closes', async () => { + const transport1 = connectClient('client-fs-overlap'); + const reverseRequestPromise = Event.toPromise(Event.filter(transport1.onDidSend, msg => isJsonRpcRequest(msg) && msg.method === 'resourceList')); + const readPromise = fileSystemProvider.readdir(agentHostUri('client-fs-overlap', '/workspace')); + const reverseRequest = await reverseRequestPromise; + assert.ok(isJsonRpcRequest(reverseRequest)); + + const transport2 = new MockProtocolTransport(); + server.simulateConnection(transport2); + const reconnectRespPromise = waitForResponse(transport2, 1); + transport2.simulateMessage(request(1, 'reconnect', { + clientId: 'client-fs-overlap', + lastSeenServerSeq: 0, + subscriptions: [], + })); + await reconnectRespPromise; + + transport1.simulateMessage({ + jsonrpc: '2.0', + id: reverseRequest.id, + result: { entries: [{ name: 'from-original-transport.txt', type: 'file' as const }] }, + }); + + const result = await readPromise; + assert.deepStrictEqual(result, [['from-original-transport.txt', FileType.File]]); + }); + + test('closing an older overlapping transport rejects its pending reverse-RPC requests', async () => { + const transport1 = connectClient('client-fs-overlap-close'); + const reverseRequestPromise = Event.toPromise(Event.filter(transport1.onDidSend, msg => isJsonRpcRequest(msg) && msg.method === 'resourceList')); + const readPromise = fileSystemProvider.readdir(agentHostUri('client-fs-overlap-close', '/workspace')); + await reverseRequestPromise; + + const transport2 = new MockProtocolTransport(); + server.simulateConnection(transport2); + const reconnectRespPromise = waitForResponse(transport2, 1); + transport2.simulateMessage(request(1, 'reconnect', { + clientId: 'client-fs-overlap-close', + lastSeenServerSeq: 0, + subscriptions: [], + })); + await reconnectRespPromise; + + transport1.simulateClose(); + + await assert.rejects(readPromise, /Client client-fs-overlap-close disconnected/); + }); + test('client disconnect cleans up', () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); @@ -1000,7 +1048,7 @@ suite('ProtocolServerHandler', () => { }); }); - test('owned tool call is not failed while the owning client stays active on another transport', () => { + test('owned tool call is not failed when closing the latest overlapping transport falls back to an older one', () => { return runWithFakedTimers({ useFakeTimers: true }, async () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); @@ -1025,38 +1073,24 @@ suite('ProtocolServerHandler', () => { 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]); + const fallbackTransport = connectClient('client-tools', [sessionUri]); + const latestTransport = 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(); + latestTransport.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')); - } + await new Promise(r => setTimeout(r, 30_001)); part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0]; assert.strictEqual(part?.kind === ResponsePartKind.ToolCall ? part.toolCall.status : undefined, ToolCallStatus.Streaming); - liveTransport.simulateClose(); + fallbackTransport.simulateClose(); }); }); - test('owned tool call is failed once the client stops sending frames on every transport', () => { + test('owned tool call is failed after the last overlapping transport closes', () => { return runWithFakedTimers({ useFakeTimers: true }, async () => { stateManager.createSession(makeSessionSummary()); stateManager.dispatchServerAction(sessionUri, { type: ActionType.SessionReady, }); @@ -1081,21 +1115,18 @@ suite('ProtocolServerHandler', () => { 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')); + const fallbackTransport = connectClient('client-tools', [sessionUri]); + const latestTransport = connectClient('client-tools', [sessionUri]); + latestTransport.simulateClose(); await new Promise(r => setTimeout(r, 30_001)); + let part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0]; + assert.strictEqual(part?.kind === ResponsePartKind.ToolCall ? part.toolCall.status : undefined, ToolCallStatus.Streaming); - const part = stateManager.getSessionState(sessionUri)?.activeTurn?.responseParts[0]; + fallbackTransport.simulateClose(); + await new Promise(r => setTimeout(r, 30_001)); + + 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, @@ -1105,8 +1136,6 @@ suite('ProtocolServerHandler', () => { success: false, error: 'Client client-tools disconnected before completing Run Task', }); - - liveTransport.simulateClose(); }); }); @@ -1492,7 +1521,7 @@ suite('ProtocolServerHandler', () => { const transport1 = connectClient('client-rc'); assert.deepStrictEqual(counts, [1]); - // Reconnect with same clientId (new transport) + // Reconnect with same clientId (new active transport) const transport2 = new MockProtocolTransport(); server.simulateConnection(transport2); transport2.simulateMessage(request(1, 'reconnect', { @@ -1500,10 +1529,11 @@ suite('ProtocolServerHandler', () => { lastSeenServerSeq: 0, subscriptions: [], })); - // Count is unchanged because same clientId was overwritten + // Count is unchanged because the logical clientId is already connected. assert.deepStrictEqual(counts, [1, 1]); - // Old transport closes - should NOT decrement since it's stale + // Old transport closes - should NOT decrement because the newer + // transport is still connected. transport1.simulateClose(); assert.deepStrictEqual(counts, [1, 1]); @@ -1800,5 +1830,31 @@ suite('ProtocolServerHandler', () => { transport.simulateClose(); assert.deepStrictEqual(agentService.watchUnsubscribeCalls, [watchChannel]); }); + + test('overlapping transports release each resource-watch subscription', async () => { + const watchChannel = 'ahp-resource-watch:/mock-watch-overlap'; + agentService.liveWatchDescriptors.set(watchChannel, { root: 'file:///root', recursive: false }); + + const transport1 = connectClient('client-watch-overlap'); + const subPromise1 = waitForResponse(transport1, 200); + transport1.simulateMessage(request(200, 'subscribe', { channel: watchChannel })); + await subPromise1; + + const transport2 = connectClient('client-watch-overlap'); + const subPromise2 = waitForResponse(transport2, 201); + transport2.simulateMessage(request(201, 'subscribe', { channel: watchChannel })); + await subPromise2; + + transport2.simulateClose(); + transport1.simulateClose(); + + assert.deepStrictEqual({ + subscribes: agentService.watchSubscribeCalls, + unsubscribes: agentService.watchUnsubscribeCalls, + }, { + subscribes: [watchChannel, watchChannel], + unsubscribes: [watchChannel, watchChannel], + }); + }); }); });