From 03aa21518f036b414996fa6943ecb1cba2d5962f Mon Sep 17 00:00:00 2001 From: Elberth Date: Wed, 7 Oct 2026 23:33:55 +0000 Subject: [PATCH 1/4] fix(vscode-ide-companion): make IdeServer.stop() resolve while MCP sessions are open `stop()` awaited `http.Server.close()`, which only stops accepting new connections and waits for existing ones to drain. The CLI's StreamableHTTP client keeps a standalone `GET /mcp` SSE stream open for the life of the session, so the callback never fired, `stop()` never resolved, and the cleanup after the await (env collection, port file) never ran. `stop()` now: - detaches `server`/`transports` synchronously so repeated or concurrent calls are no-ops instead of racing on the same listener; - closes every MCP transport first (`Promise.allSettled`), which ends the SSE responses and fires `onclose` (clears keep-alive, evicts session); - calls `server.closeAllConnections()` alongside `server.close()` so idle keep-alive sockets cannot hold the callback; - runs the env-collection/port-file cleanup in `finally`. The keep-alive `missedPings >= 3` branch now closes the transport (as its log line already claimed) rather than only clearing the interval. Consecutive-miss semantics are unchanged. Fixes #28785 --- .../src/ide-server.test.ts | 230 ++++++++++++++++++ .../vscode-ide-companion/src/ide-server.ts | 61 +++-- 2 files changed, 269 insertions(+), 22 deletions(-) diff --git a/packages/vscode-ide-companion/src/ide-server.test.ts b/packages/vscode-ide-companion/src/ide-server.test.ts index 630640574ad..b614284fcd0 100644 --- a/packages/vscode-ide-companion/src/ide-server.test.ts +++ b/packages/vscode-ide-companion/src/ide-server.test.ts @@ -12,6 +12,10 @@ import * as path from 'node:path'; import * as http from 'node:http'; import { IDEServer } from './ide-server.js'; import type { DiffManager } from './diff-manager.js'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; +import type { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'; +import { IdeContextNotificationSchema } from '@google/gemini-cli-core/src/ide/types.js'; const { vscodeMock: baseVscodeMock } = await vi.hoisted( () => import('./utils/vscode-mock.js'), @@ -76,6 +80,11 @@ vi.mock('vscode', () => vscodeMock); vi.mock('./open-files-manager', () => { const OpenFilesManager = vi.fn(); OpenFilesManager.prototype.onDidChange = vi.fn(() => ({ dispose: vi.fn() })); + // A schema-valid IdeContext so the GET /mcp handler can emit the initial + // `ide/contextUpdate` notification instead of throwing on parse. + OpenFilesManager.prototype.state = { + workspaceState: { openFiles: [], isTrusted: true }, + }; return { OpenFilesManager }; }); @@ -316,6 +325,227 @@ describe('IDEServer', () => { expect(fs.unlink).toHaveBeenCalledWith(portFile); }); + describe('shutdown with active MCP sessions', () => { + type ServerTransports = Record; + const getTransports = () => + (ideServer as unknown as { transports: ServerTransports }).transports; + + /** + * Connects exactly the way IdeClient does. The SDK client opens a + * long-lived standalone `GET /mcp` SSE stream right after `initialize` + * (fire-and-forget), so we wait for the server's initial + * `ide/contextUpdate` push — which can only arrive over that stream — to + * know the socket is open end-to-end. + */ + const connectMcpClient = async (port: string) => { + const client = new Client({ name: 'test-client', version: '0.0.0' }); + const transport = new StreamableHTTPClientTransport( + new URL(`http://127.0.0.1:${port}/mcp`), + { + requestInit: { headers: { Authorization: 'Bearer test-auth-token' } }, + }, + ); + const sseStreamLive = new Promise((resolve) => { + client.setNotificationHandler(IdeContextNotificationSchema, () => + resolve(), + ); + }); + await client.connect(transport); + await Promise.race([ + sseStreamLive, + new Promise((_, reject) => + setTimeout( + () => reject(new Error('GET /mcp SSE stream never became live')), + 2_000, + ), + ), + ]); + return client; + }; + + const withDeadline = (promise: Promise, ms: number) => + Promise.race([ + promise.then(() => 'resolved' as const), + new Promise<'timed-out'>((resolve) => + setTimeout(() => resolve('timed-out'), ms), + ), + ]); + + let port: string; + let clients: Client[]; + + beforeEach(async () => { + clients = []; + await ideServer.start(mockContext); + port = getPortFromMock(mockContext.environmentVariableCollection.replace); + }); + + afterEach(async () => { + // Drain client sockets so a hung stop() cannot wedge the runner. + await Promise.allSettled(clients.map((c) => c.close())); + }); + + // Regression test for https://github.com/google-gemini/gemini-cli/issues/28785 + it('should resolve stop() while a session holds an open SSE stream', async () => { + clients.push(await connectMcpClient(port)); + expect(Object.keys(getTransports())).toHaveLength(1); + + const outcome = await withDeadline(ideServer.stop(), 2_000); + + expect(outcome, 'stop() hung while a client held GET /mcp open').toBe( + 'resolved', + ); + expect(mockLog).toHaveBeenCalledWith('IDE server shut down'); + }); + + it('should close sessions, clear keep-alive, and clean up on stop()', async () => { + const clearIntervalSpy = vi.spyOn(globalThis, 'clearInterval'); + clients.push(await connectMcpClient(port)); + const [sessionId] = Object.keys(getTransports()); + + await ideServer.stop(); + + expect(Object.keys(getTransports())).toHaveLength(0); + expect(mockLog).toHaveBeenCalledWith(`Session closed: ${sessionId}`); + expect(clearIntervalSpy).toHaveBeenCalled(); + expect( + mockContext.environmentVariableCollection.clear, + ).toHaveBeenCalled(); + expect(fs.unlink).toHaveBeenCalled(); + }); + + it('should close every session when multiple clients are connected', async () => { + // randomUUID is mocked to a constant, so give each session a unique id. + const { randomUUID } = await import('node:crypto'); + vi.mocked(randomUUID) + .mockReturnValueOnce('session-a' as ReturnType) + .mockReturnValueOnce('session-b' as ReturnType); + + clients.push(await connectMcpClient(port)); + clients.push(await connectMcpClient(port)); + expect(Object.keys(getTransports()).sort()).toEqual([ + 'session-a', + 'session-b', + ]); + + const outcome = await withDeadline(ideServer.stop(), 2_000); + + expect(outcome).toBe('resolved'); + expect(Object.keys(getTransports())).toHaveLength(0); + expect(mockLog).toHaveBeenCalledWith('Session closed: session-a'); + expect(mockLog).toHaveBeenCalledWith('Session closed: session-b'); + }); + + it('should make stop() idempotent', async () => { + clients.push(await connectMcpClient(port)); + + await ideServer.stop(); + await expect(ideServer.stop()).resolves.toBeUndefined(); + + expect( + vi + .mocked(mockLog) + .mock.calls.filter(([m]) => m === 'IDE server shut down'), + ).toHaveLength(1); + }); + + it('should tolerate concurrent stop() calls', async () => { + clients.push(await connectMcpClient(port)); + + await expect( + Promise.all([ideServer.stop(), ideServer.stop()]), + ).resolves.toEqual([undefined, undefined]); + + expect( + vi + .mocked(mockLog) + .mock.calls.filter(([m]) => m === 'IDE server shut down'), + ).toHaveLength(1); + expect(mockLog).not.toHaveBeenCalledWith( + expect.stringContaining('Error shutting down IDE server'), + ); + }); + + it('should still clean up when the HTTP server fails to close', async () => { + const server = (ideServer as unknown as { server: http.Server }).server; + const closeSpy = vi.spyOn(server, 'close').mockImplementation((cb) => { + cb?.(new Error('close failed')); + return server; + }); + + try { + await expect(ideServer.stop()).rejects.toThrow('close failed'); + + expect( + mockContext.environmentVariableCollection.clear, + ).toHaveBeenCalled(); + expect(fs.unlink).toHaveBeenCalled(); + // A follow-up stop() must not try to close the dead server again. + await expect(ideServer.stop()).resolves.toBeUndefined(); + } finally { + closeSpy.mockRestore(); + await new Promise((resolve) => server.close(() => resolve())); + } + }); + + describe('keep-alive', () => { + const KEEP_ALIVE_MS = 60_000; + + beforeEach(() => { + // Only fake the interval APIs so real sockets/fetch keep working. + vi.useFakeTimers({ toFake: ['setInterval', 'clearInterval'] }); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('should close the session after 3 consecutive missed pings', async () => { + clients.push(await connectMcpClient(port)); + const [sessionId] = Object.keys(getTransports()); + const transport = getTransports()[sessionId]; + const closeSpy = vi.spyOn(transport, 'close'); + vi.spyOn(transport, 'send').mockRejectedValue(new Error('EPIPE')); + + await vi.advanceTimersByTimeAsync(KEEP_ALIVE_MS * 2); + expect(closeSpy).not.toHaveBeenCalled(); + expect(getTransports()[sessionId]).toBe(transport); + + await vi.advanceTimersByTimeAsync(KEEP_ALIVE_MS); + expect(closeSpy).toHaveBeenCalledTimes(1); + expect(getTransports()[sessionId]).toBeUndefined(); + expect(mockLog).toHaveBeenCalledWith( + expect.stringContaining(`Session ${sessionId} missed 3 pings`), + ); + expect(mockLog).toHaveBeenCalledWith(`Session closed: ${sessionId}`); + + // Interval is gone: no further pings are attempted. + const sendCalls = vi.mocked(transport.send).mock.calls.length; + await vi.advanceTimersByTimeAsync(KEEP_ALIVE_MS * 2); + expect(transport.send).toHaveBeenCalledTimes(sendCalls); + }); + + it('should reset the missed-ping count on a successful ping', async () => { + clients.push(await connectMcpClient(port)); + const [sessionId] = Object.keys(getTransports()); + const transport = getTransports()[sessionId]; + const closeSpy = vi.spyOn(transport, 'close'); + vi.spyOn(transport, 'send') + .mockRejectedValueOnce(new Error('EPIPE')) + .mockRejectedValueOnce(new Error('EPIPE')) + .mockResolvedValueOnce(undefined) + .mockRejectedValueOnce(new Error('EPIPE')) + .mockRejectedValueOnce(new Error('EPIPE')); + + await vi.advanceTimersByTimeAsync(KEEP_ALIVE_MS * 5); + + expect(transport.send).toHaveBeenCalledTimes(5); + expect(closeSpy).not.toHaveBeenCalled(); + expect(getTransports()[sessionId]).toBe(transport); + }); + }); + }); + it.skipIf(process.platform !== 'win32')( 'should handle windows paths', async () => { diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index 39ef770079d..500ce28fd2e 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -241,7 +241,8 @@ export class IDEServer { this.log( `Session ${sessionId} missed ${missedPings} pings. Closing connection and cleaning up interval.`, ); - clearInterval(keepAlive); + // `onclose` clears the interval and evicts the session. + void transport.close(); } }); }, 60000); // 60 sec @@ -404,28 +405,44 @@ export class IDEServer { } async stop(): Promise { - if (this.server) { - await new Promise((resolve, reject) => { - this.server!.close((err?: Error) => { - if (err) { - this.log(`Error shutting down IDE server: ${err.message}`); - return reject(err); - } - this.log(`IDE server shut down`); - resolve(); + // Detach synchronously so concurrent or repeated stop() calls are no-ops + // rather than closing the same listener twice. + const server = this.server; + this.server = undefined; + const transports = Object.values(this.transports); + this.transports = {}; + + // Close every MCP session first. This ends their long-lived SSE + // responses and fires `onclose`, which clears the keep-alive interval. + await Promise.allSettled(transports.map((transport) => transport.close())); + + try { + if (server) { + await new Promise((resolve, reject) => { + server.close((err?: Error) => { + if (err) { + this.log(`Error shutting down IDE server: ${err.message}`); + return reject(err); + } + this.log(`IDE server shut down`); + resolve(); + }); + // `close()` only stops accepting new connections and waits for + // existing ones to drain. Drop any remaining sockets so the callback + // above fires promptly instead of hanging on an idle keep-alive. + server.closeAllConnections(); }); - }); - this.server = undefined; - } - - if (this.context) { - this.context.environmentVariableCollection.clear(); - } - if (this.portFile) { - try { - await fs.unlink(this.portFile); - } catch { - // Ignore errors if the file doesn't exist. + } + } finally { + if (this.context) { + this.context.environmentVariableCollection.clear(); + } + if (this.portFile) { + try { + await fs.unlink(this.portFile); + } catch { + // Ignore errors if the file doesn't exist. + } } } } From 6bb07a3a60595a4bd974b4c5b3ba9d1f5f1d739a Mon Sep 17 00:00:00 2001 From: Elberth Date: Thu, 8 Oct 2026 00:11:11 +0000 Subject: [PATCH 2/4] fix(vscode-ide-companion): address review on IdeServer shutdown - Catch and log a rejected `transport.close()` in the missed-pings branch instead of `void`-ing it, and clear the interval directly so a transport that cannot close is not retried every 60s. - Share the in-flight shutdown promise across concurrent `stop()` calls so a second caller resolves only after the listener, sockets and port file are actually gone, rather than instantly. --- .../src/ide-server.test.ts | 55 +++++++++++++++++-- .../vscode-ide-companion/src/ide-server.ts | 27 +++++++-- 2 files changed, 74 insertions(+), 8 deletions(-) diff --git a/packages/vscode-ide-companion/src/ide-server.test.ts b/packages/vscode-ide-companion/src/ide-server.test.ts index b614284fcd0..1cc94aa544b 100644 --- a/packages/vscode-ide-companion/src/ide-server.test.ts +++ b/packages/vscode-ide-companion/src/ide-server.test.ts @@ -449,13 +449,37 @@ describe('IDEServer', () => { ).toHaveLength(1); }); - it('should tolerate concurrent stop() calls', async () => { + it('should make concurrent stop() calls await the same in-flight shutdown', async () => { clients.push(await connectMcpClient(port)); + const server = (ideServer as unknown as { server: http.Server }).server; + const realClose = server.close.bind(server); + let releaseClose!: () => void; + const closeGate = new Promise((resolve) => { + releaseClose = resolve; + }); + vi.spyOn(server, 'close').mockImplementation((cb) => { + void closeGate.then(() => realClose(cb)); + return server; + }); - await expect( - Promise.all([ideServer.stop(), ideServer.stop()]), - ).resolves.toEqual([undefined, undefined]); + const first = ideServer.stop(); + const second = ideServer.stop(); + // The second caller must not resolve while the listener is still + // closing and cleanup (env collection, port file) has not run yet. + expect(await withDeadline(second, 50)).toBe('timed-out'); + expect( + mockContext.environmentVariableCollection.clear, + ).not.toHaveBeenCalled(); + + releaseClose(); + await expect(Promise.all([first, second])).resolves.toEqual([ + undefined, + undefined, + ]); + expect( + mockContext.environmentVariableCollection.clear, + ).toHaveBeenCalledTimes(1); expect( vi .mocked(mockLog) @@ -543,6 +567,29 @@ describe('IDEServer', () => { expect(closeSpy).not.toHaveBeenCalled(); expect(getTransports()[sessionId]).toBe(transport); }); + + it('should log and stop pinging if closing the evicted session fails', async () => { + clients.push(await connectMcpClient(port)); + const [sessionId] = Object.keys(getTransports()); + const transport = getTransports()[sessionId]; + vi.spyOn(transport, 'send').mockRejectedValue(new Error('EPIPE')); + const closeSpy = vi + .spyOn(transport, 'close') + .mockRejectedValue(new Error('close failed')); + + await vi.advanceTimersByTimeAsync(KEEP_ALIVE_MS * 3); + + expect(closeSpy).toHaveBeenCalledTimes(1); + expect(mockLog).toHaveBeenCalledWith( + `Failed to close transport for session ${sessionId}: close failed`, + ); + + // The interval must not keep retrying a transport that cannot close. + const sendCalls = vi.mocked(transport.send).mock.calls.length; + await vi.advanceTimersByTimeAsync(KEEP_ALIVE_MS * 2); + expect(transport.send).toHaveBeenCalledTimes(sendCalls); + expect(closeSpy).toHaveBeenCalledTimes(1); + }); }); }); diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index 500ce28fd2e..adfafb53a3d 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -128,6 +128,7 @@ export class IDEServer { private transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; private openFilesManager: OpenFilesManager | undefined; + private stopping: Promise | undefined; diffManager: DiffManager; constructor(log: (message: string) => void, diffManager: DiffManager) { @@ -241,8 +242,15 @@ export class IDEServer { this.log( `Session ${sessionId} missed ${missedPings} pings. Closing connection and cleaning up interval.`, ); - // `onclose` clears the interval and evicts the session. - void transport.close(); + clearInterval(keepAlive); + // `onclose` evicts the session from `this.transports`. + transport.close().catch((error: unknown) => { + const message = + error instanceof Error ? error.message : String(error); + this.log( + `Failed to close transport for session ${sessionId}: ${message}`, + ); + }); } }); }, 60000); // 60 sec @@ -405,8 +413,19 @@ export class IDEServer { } async stop(): Promise { - // Detach synchronously so concurrent or repeated stop() calls are no-ops - // rather than closing the same listener twice. + // Share the in-flight shutdown so a concurrent caller resolves only once + // the listener, sockets and port file are actually gone — not instantly. + if (!this.stopping) { + this.stopping = this.shutdown().finally(() => { + this.stopping = undefined; + }); + } + return this.stopping; + } + + private async shutdown(): Promise { + // Detach synchronously so a stop() issued after this one completes is a + // no-op rather than closing the same listener twice. const server = this.server; this.server = undefined; const transports = Object.values(this.transports); From 66578fb3cd63925877e943703aa33ad4d5aaa66a Mon Sep 17 00:00:00 2001 From: Elberth Date: Thu, 8 Oct 2026 00:16:43 +0000 Subject: [PATCH 3/4] refactor(vscode-ide-companion): track IdeServer shutdown with an explicit status Per review: model the shutdown lifecycle with `status: 'idle' | 'stopping' | 'stopped'` instead of the presence of a promise. `stopping` joins the in-flight shutdown, `stopped` is a no-op, and `start()` resets to `idle` so a restarted server remains stoppable. Adds a restart test. --- .../src/ide-server.test.ts | 23 +++++++++++++++ .../vscode-ide-companion/src/ide-server.ts | 28 +++++++++++++------ 2 files changed, 43 insertions(+), 8 deletions(-) diff --git a/packages/vscode-ide-companion/src/ide-server.test.ts b/packages/vscode-ide-companion/src/ide-server.test.ts index 1cc94aa544b..9572748c23a 100644 --- a/packages/vscode-ide-companion/src/ide-server.test.ts +++ b/packages/vscode-ide-companion/src/ide-server.test.ts @@ -449,6 +449,29 @@ describe('IDEServer', () => { ).toHaveLength(1); }); + it('should be stoppable again after a restart', async () => { + clients.push(await connectMcpClient(port)); + await ideServer.stop(); + + await ideServer.start(mockContext); + const newPort = vi + .mocked(mockContext.environmentVariableCollection.replace) + .mock.calls.filter(([k]) => k === 'GEMINI_CLI_IDE_SERVER_PORT') + .at(-1)?.[1] as string; + clients.push(await connectMcpClient(newPort)); + expect(Object.keys(getTransports())).toHaveLength(1); + + const outcome = await withDeadline(ideServer.stop(), 2_000); + + expect(outcome).toBe('resolved'); + expect(Object.keys(getTransports())).toHaveLength(0); + expect( + vi + .mocked(mockLog) + .mock.calls.filter(([m]) => m === 'IDE server shut down'), + ).toHaveLength(2); + }); + it('should make concurrent stop() calls await the same in-flight shutdown', async () => { clients.push(await connectMcpClient(port)); const server = (ideServer as unknown as { server: http.Server }).server; diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index adfafb53a3d..f036919aee6 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -128,7 +128,8 @@ export class IDEServer { private transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; private openFilesManager: OpenFilesManager | undefined; - private stopping: Promise | undefined; + private status: 'idle' | 'stopping' | 'stopped' = 'idle'; + private stopPromise: Promise | undefined; diffManager: DiffManager; constructor(log: (message: string) => void, diffManager: DiffManager) { @@ -138,6 +139,8 @@ export class IDEServer { start(context: vscode.ExtensionContext): Promise { return new Promise((resolve) => { + // A restart after stop() must be stoppable again. + this.status = 'idle'; this.context = context; this.authToken = randomUUID(); const sessionsWithInitialNotification = new Set(); @@ -413,14 +416,23 @@ export class IDEServer { } async stop(): Promise { - // Share the in-flight shutdown so a concurrent caller resolves only once - // the listener, sockets and port file are actually gone — not instantly. - if (!this.stopping) { - this.stopping = this.shutdown().finally(() => { - this.stopping = undefined; - }); + switch (this.status) { + case 'stopping': + // Join the in-flight shutdown so a concurrent caller resolves only + // once the listener, sockets and port file are actually gone. + return this.stopPromise; + case 'stopped': + return; + default: + this.status = 'stopping'; + this.stopPromise = this.shutdown().finally(() => { + // `start()` may have reset the status meanwhile; don't clobber it. + if (this.status === 'stopping') { + this.status = 'stopped'; + } + }); + return this.stopPromise; } - return this.stopping; } private async shutdown(): Promise { From c8b36a7254386d6724affd06da37cd4a58a14df7 Mon Sep 17 00:00:00 2001 From: Elberth Date: Thu, 8 Oct 2026 03:13:07 +0000 Subject: [PATCH 4/4] fix(vscode-ide-companion): isolate overlapping IdeServer stop/start cycles A shutdown that finishes after a newer start() must not touch the new server's state. Each shutdown now captures the lifecycle generation (bumped by start()) and only marks the server stopped / clears the env collection if it is still the current one; the port file is captured synchronously so only the old file is unlinked. --- .../src/ide-server.test.ts | 120 ++++++++++++++---- .../vscode-ide-companion/src/ide-server.ts | 31 +++-- 2 files changed, 116 insertions(+), 35 deletions(-) diff --git a/packages/vscode-ide-companion/src/ide-server.test.ts b/packages/vscode-ide-companion/src/ide-server.test.ts index 9572748c23a..cac6a8a1ee7 100644 --- a/packages/vscode-ide-companion/src/ide-server.test.ts +++ b/packages/vscode-ide-companion/src/ide-server.test.ts @@ -371,6 +371,42 @@ describe('IDEServer', () => { ), ]); + const getServer = () => + (ideServer as unknown as { server: http.Server }).server; + + /** Holds `server.close()` until the returned function is called. */ + const gateServerClose = (server: http.Server) => { + const realClose = server.close.bind(server); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + vi.spyOn(server, 'close').mockImplementation((cb) => { + void gate.then(() => realClose(cb)); + return server; + }); + return release; + }; + + const latestPort = () => + vi + .mocked(mockContext.environmentVariableCollection.replace) + .mock.calls.filter(([k]) => k === 'GEMINI_CLI_IDE_SERVER_PORT') + .at(-1)?.[1] as string; + + const portFileFor = (p: string) => + path.join( + '/tmp', + 'gemini', + 'ide', + `gemini-ide-server-${process.ppid}-${p}.json`, + ); + + const shutdownLogCount = () => + vi + .mocked(mockLog) + .mock.calls.filter(([m]) => m === 'IDE server shut down').length; + let port: string; let clients: Client[]; @@ -454,36 +490,19 @@ describe('IDEServer', () => { await ideServer.stop(); await ideServer.start(mockContext); - const newPort = vi - .mocked(mockContext.environmentVariableCollection.replace) - .mock.calls.filter(([k]) => k === 'GEMINI_CLI_IDE_SERVER_PORT') - .at(-1)?.[1] as string; - clients.push(await connectMcpClient(newPort)); + clients.push(await connectMcpClient(latestPort())); expect(Object.keys(getTransports())).toHaveLength(1); const outcome = await withDeadline(ideServer.stop(), 2_000); expect(outcome).toBe('resolved'); expect(Object.keys(getTransports())).toHaveLength(0); - expect( - vi - .mocked(mockLog) - .mock.calls.filter(([m]) => m === 'IDE server shut down'), - ).toHaveLength(2); + expect(shutdownLogCount()).toBe(2); }); it('should make concurrent stop() calls await the same in-flight shutdown', async () => { clients.push(await connectMcpClient(port)); - const server = (ideServer as unknown as { server: http.Server }).server; - const realClose = server.close.bind(server); - let releaseClose!: () => void; - const closeGate = new Promise((resolve) => { - releaseClose = resolve; - }); - vi.spyOn(server, 'close').mockImplementation((cb) => { - void closeGate.then(() => realClose(cb)); - return server; - }); + const releaseClose = gateServerClose(getServer()); const first = ideServer.stop(); const second = ideServer.stop(); @@ -503,16 +522,67 @@ describe('IDEServer', () => { expect( mockContext.environmentVariableCollection.clear, ).toHaveBeenCalledTimes(1); - expect( - vi - .mocked(mockLog) - .mock.calls.filter(([m]) => m === 'IDE server shut down'), - ).toHaveLength(1); + expect(shutdownLogCount()).toBe(1); expect(mockLog).not.toHaveBeenCalledWith( expect.stringContaining('Error shutting down IDE server'), ); }); + it('should not let a late-finishing stop() clobber a restarted server', async () => { + vi.mocked(fs.unlink).mockClear(); + clients.push(await connectMcpClient(port)); + const releaseOldClose = gateServerClose(getServer()); + const oldStop = ideServer.stop(); + + // Restart while the old shutdown is still draining. + await ideServer.start(mockContext); + const newPort = latestPort(); + expect(newPort).not.toBe(port); + clients.push(await connectMcpClient(newPort)); + + releaseOldClose(); + await oldStop; + + // The old shutdown removed only its own port file and left the new + // server's env vars alone. + expect(fs.unlink).toHaveBeenCalledWith(portFileFor(port)); + expect(fs.unlink).not.toHaveBeenCalledWith(portFileFor(newPort)); + expect( + mockContext.environmentVariableCollection.clear, + ).not.toHaveBeenCalled(); + + // ...and the new server is still fully stoppable. + expect(await withDeadline(ideServer.stop(), 2_000)).toBe('resolved'); + expect(shutdownLogCount()).toBe(2); + expect(fs.unlink).toHaveBeenCalledWith(portFileFor(newPort)); + expect( + mockContext.environmentVariableCollection.clear, + ).toHaveBeenCalledTimes(1); + }); + + it('should keep a second shutdown joinable when the first finishes later', async () => { + clients.push(await connectMcpClient(port)); + const releaseOldClose = gateServerClose(getServer()); + const oldStop = ideServer.stop(); + + await ideServer.start(mockContext); + clients.push(await connectMcpClient(latestPort())); + const releaseNewClose = gateServerClose(getServer()); + const newStop = ideServer.stop(); + + // First shutdown completes while the second is still in flight. + releaseOldClose(); + await oldStop; + + // The server must still report "stopping": a further stop() joins the + // second shutdown instead of returning instantly as "stopped". + expect(await withDeadline(ideServer.stop(), 50)).toBe('timed-out'); + + releaseNewClose(); + await expect(newStop).resolves.toBeUndefined(); + expect(shutdownLogCount()).toBe(2); + }); + it('should still clean up when the HTTP server fails to close', async () => { const server = (ideServer as unknown as { server: http.Server }).server; const closeSpy = vi.spyOn(server, 'close').mockImplementation((cb) => { diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index f036919aee6..d7bf0fbc3a0 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -130,6 +130,9 @@ export class IDEServer { private openFilesManager: OpenFilesManager | undefined; private status: 'idle' | 'stopping' | 'stopped' = 'idle'; private stopPromise: Promise | undefined; + /** Bumped by every start(); lets a late-finishing shutdown recognise that a + * newer server now owns the status, env vars and port file. */ + private generation = 0; diffManager: DiffManager; constructor(log: (message: string) => void, diffManager: DiffManager) { @@ -139,8 +142,10 @@ export class IDEServer { start(context: vscode.ExtensionContext): Promise { return new Promise((resolve) => { - // A restart after stop() must be stoppable again. + // A restart after stop() must be stoppable again, and any shutdown still + // in flight must not touch this server's state when it finishes. this.status = 'idle'; + this.generation++; this.context = context; this.authToken = randomUUID(); const sessionsWithInitialNotification = new Set(); @@ -423,25 +428,30 @@ export class IDEServer { return this.stopPromise; case 'stopped': return; - default: + default: { this.status = 'stopping'; - this.stopPromise = this.shutdown().finally(() => { - // `start()` may have reset the status meanwhile; don't clobber it. - if (this.status === 'stopping') { + const generation = this.generation; + this.stopPromise = this.shutdown(generation).finally(() => { + // A start() issued meanwhile owns the state now; leave it alone. + if (this.generation === generation) { this.status = 'stopped'; } }); return this.stopPromise; + } } } - private async shutdown(): Promise { + private async shutdown(generation: number): Promise { // Detach synchronously so a stop() issued after this one completes is a - // no-op rather than closing the same listener twice. + // no-op rather than closing the same listener twice, and so a start() + // issued meanwhile cannot have its port file removed by this shutdown. const server = this.server; this.server = undefined; const transports = Object.values(this.transports); this.transports = {}; + const portFile = this.portFile; + this.portFile = undefined; // Close every MCP session first. This ends their long-lived SSE // responses and fires `onclose`, which clears the keep-alive interval. @@ -465,12 +475,13 @@ export class IDEServer { }); } } finally { - if (this.context) { + // Only clear the env vars if no newer start() has written its own. + if (this.generation === generation && this.context) { this.context.environmentVariableCollection.clear(); } - if (this.portFile) { + if (portFile) { try { - await fs.unlink(this.portFile); + await fs.unlink(portFile); } catch { // Ignore errors if the file doesn't exist. }