Skip to content

Commit a38a5ba

Browse files
tanzhenxintanzhenxin
authored
Merge pull request #1285 from tt-a1i/fix/mcp-subprocess-cleanup
fix(core): properly cleanup MCP server subprocesses on exit
2 parents 9c97468 + 640196e commit a38a5ba

5 files changed

Lines changed: 287 additions & 24 deletions

File tree

‎packages/cli/src/gemini.tsx‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,9 @@ export async function main() {
339339
process.cwd(),
340340
argv.extensions,
341341
);
342+
343+
// Register cleanup for MCP clients as early as possible
344+
// This ensures MCP server subprocesses are properly terminated on exit
342345
registerCleanup(() => config.shutdown());
343346

344347
// FIXME: list extensions after the config initialize

‎packages/core/src/config/config.ts‎

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -814,13 +814,6 @@ export class Config {
814814
return this.sessionId;
815815
}
816816

817-
/**
818-
* Releases resources owned by the config instance.
819-
*/
820-
async shutdown(): Promise<void> {
821-
this.skillManager?.stopWatching();
822-
}
823-
824817
/**
825818
* Starts a new session and resets session-scoped services.
826819
*/
@@ -1027,6 +1020,28 @@ export class Config {
10271020
return this.toolRegistry;
10281021
}
10291022

1023+
/**
1024+
* Shuts down the Config and releases all resources.
1025+
* This method is idempotent and safe to call multiple times.
1026+
* It handles the case where initialization was not completed.
1027+
*/
1028+
async shutdown(): Promise<void> {
1029+
if (!this.initialized) {
1030+
// Nothing to clean up if not initialized
1031+
return;
1032+
}
1033+
try {
1034+
this.skillManager?.stopWatching();
1035+
1036+
if (this.toolRegistry) {
1037+
await this.toolRegistry.stop();
1038+
}
1039+
} catch (error) {
1040+
// Log but don't throw - cleanup should be best-effort
1041+
console.error('Error during Config shutdown:', error);
1042+
}
1043+
}
1044+
10301045
getPromptRegistry(): PromptRegistry {
10311046
return this.promptRegistry;
10321047
}

‎packages/core/src/tools/mcp-client-manager.test.ts‎

Lines changed: 178 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,16 @@ import { McpClientManager } from './mcp-client-manager.js';
99
import { McpClient } from './mcp-client.js';
1010
import type { ToolRegistry } from './tool-registry.js';
1111
import type { Config } from '../config/config.js';
12+
import type { PromptRegistry } from '../prompts/prompt-registry.js';
13+
import type { WorkspaceContext } from '../utils/workspaceContext.js';
1214

1315
vi.mock('./mcp-client.js', async () => {
1416
const originalModule = await vi.importActual('./mcp-client.js');
1517
return {
1618
...originalModule,
1719
McpClient: vi.fn(),
18-
populateMcpServerCommand: vi.fn(() => ({
19-
'test-server': {},
20-
})),
20+
// Return the input servers unchanged (identity function)
21+
populateMcpServerCommand: vi.fn((servers) => servers),
2122
};
2223
});
2324

@@ -73,4 +74,178 @@ describe('McpClientManager', () => {
7374
expect(mockedMcpClient.connect).not.toHaveBeenCalled();
7475
expect(mockedMcpClient.discover).not.toHaveBeenCalled();
7576
});
77+
78+
it('should disconnect all clients when stop is called', async () => {
79+
// Track disconnect calls across all instances
80+
const disconnectCalls: string[] = [];
81+
vi.mocked(McpClient).mockImplementation(
82+
(name: string) =>
83+
({
84+
connect: vi.fn(),
85+
discover: vi.fn(),
86+
disconnect: vi.fn().mockImplementation(() => {
87+
disconnectCalls.push(name);
88+
return Promise.resolve();
89+
}),
90+
getStatus: vi.fn(),
91+
}) as unknown as McpClient,
92+
);
93+
const mockConfig = {
94+
isTrustedFolder: () => true,
95+
getMcpServers: () => ({ 'test-server': {}, 'another-server': {} }),
96+
getMcpServerCommand: () => undefined,
97+
getPromptRegistry: () => ({}) as PromptRegistry,
98+
getWorkspaceContext: () => ({}) as WorkspaceContext,
99+
getDebugMode: () => false,
100+
} as unknown as Config;
101+
const manager = new McpClientManager(mockConfig, {} as ToolRegistry);
102+
// First connect to create the clients
103+
await manager.discoverAllMcpTools({
104+
isTrustedFolder: () => true,
105+
} as unknown as Config);
106+
107+
// Clear the disconnect calls from initial stop() in discoverAllMcpTools
108+
disconnectCalls.length = 0;
109+
110+
// Then stop
111+
await manager.stop();
112+
expect(disconnectCalls).toHaveLength(2);
113+
expect(disconnectCalls).toContain('test-server');
114+
expect(disconnectCalls).toContain('another-server');
115+
});
116+
117+
it('should be idempotent - stop can be called multiple times safely', async () => {
118+
const mockedMcpClient = {
119+
connect: vi.fn(),
120+
discover: vi.fn(),
121+
disconnect: vi.fn().mockResolvedValue(undefined),
122+
getStatus: vi.fn(),
123+
};
124+
vi.mocked(McpClient).mockReturnValue(
125+
mockedMcpClient as unknown as McpClient,
126+
);
127+
const mockConfig = {
128+
isTrustedFolder: () => true,
129+
getMcpServers: () => ({ 'test-server': {} }),
130+
getMcpServerCommand: () => undefined,
131+
getPromptRegistry: () => ({}) as PromptRegistry,
132+
getWorkspaceContext: () => ({}) as WorkspaceContext,
133+
getDebugMode: () => false,
134+
} as unknown as Config;
135+
const manager = new McpClientManager(mockConfig, {} as ToolRegistry);
136+
await manager.discoverAllMcpTools({
137+
isTrustedFolder: () => true,
138+
} as unknown as Config);
139+
140+
// Call stop multiple times - should not throw
141+
await manager.stop();
142+
await manager.stop();
143+
await manager.stop();
144+
});
145+
146+
it('should discover tools for a single server and track the client for stop', async () => {
147+
const mockedMcpClient = {
148+
connect: vi.fn(),
149+
discover: vi.fn(),
150+
disconnect: vi.fn().mockResolvedValue(undefined),
151+
getStatus: vi.fn(),
152+
};
153+
vi.mocked(McpClient).mockReturnValue(
154+
mockedMcpClient as unknown as McpClient,
155+
);
156+
157+
const mockConfig = {
158+
isTrustedFolder: () => true,
159+
getMcpServers: () => ({ 'test-server': {} }),
160+
getMcpServerCommand: () => undefined,
161+
getPromptRegistry: () => ({}) as PromptRegistry,
162+
getWorkspaceContext: () => ({}) as WorkspaceContext,
163+
getDebugMode: () => false,
164+
} as unknown as Config;
165+
const manager = new McpClientManager(mockConfig, {} as ToolRegistry);
166+
167+
await manager.discoverMcpToolsForServer(
168+
'test-server',
169+
{} as unknown as Config,
170+
);
171+
172+
expect(mockedMcpClient.connect).toHaveBeenCalledOnce();
173+
expect(mockedMcpClient.discover).toHaveBeenCalledOnce();
174+
175+
await manager.stop();
176+
expect(mockedMcpClient.disconnect).toHaveBeenCalledOnce();
177+
});
178+
179+
it('should replace an existing client when re-discovering a server', async () => {
180+
const firstClient = {
181+
connect: vi.fn(),
182+
discover: vi.fn(),
183+
disconnect: vi.fn().mockResolvedValue(undefined),
184+
getStatus: vi.fn(),
185+
};
186+
const secondClient = {
187+
connect: vi.fn(),
188+
discover: vi.fn(),
189+
disconnect: vi.fn().mockResolvedValue(undefined),
190+
getStatus: vi.fn(),
191+
};
192+
193+
vi.mocked(McpClient)
194+
.mockReturnValueOnce(firstClient as unknown as McpClient)
195+
.mockReturnValueOnce(secondClient as unknown as McpClient);
196+
197+
const mockConfig = {
198+
isTrustedFolder: () => true,
199+
getMcpServers: () => ({ 'test-server': {} }),
200+
getMcpServerCommand: () => undefined,
201+
getPromptRegistry: () => ({}) as PromptRegistry,
202+
getWorkspaceContext: () => ({}) as WorkspaceContext,
203+
getDebugMode: () => false,
204+
} as unknown as Config;
205+
const manager = new McpClientManager(mockConfig, {} as ToolRegistry);
206+
207+
await manager.discoverMcpToolsForServer(
208+
'test-server',
209+
{} as unknown as Config,
210+
);
211+
await manager.discoverMcpToolsForServer(
212+
'test-server',
213+
{} as unknown as Config,
214+
);
215+
216+
expect(firstClient.disconnect).toHaveBeenCalledOnce();
217+
expect(secondClient.connect).toHaveBeenCalledOnce();
218+
expect(secondClient.discover).toHaveBeenCalledOnce();
219+
220+
await manager.stop();
221+
expect(secondClient.disconnect).toHaveBeenCalledOnce();
222+
});
223+
224+
it('should no-op when discovering an unknown server', async () => {
225+
const mockedMcpClient = {
226+
connect: vi.fn(),
227+
discover: vi.fn(),
228+
disconnect: vi.fn().mockResolvedValue(undefined),
229+
getStatus: vi.fn(),
230+
};
231+
vi.mocked(McpClient).mockReturnValue(
232+
mockedMcpClient as unknown as McpClient,
233+
);
234+
235+
const mockConfig = {
236+
isTrustedFolder: () => true,
237+
getMcpServers: () => ({}),
238+
getMcpServerCommand: () => undefined,
239+
getPromptRegistry: () => ({}) as PromptRegistry,
240+
getWorkspaceContext: () => ({}) as WorkspaceContext,
241+
getDebugMode: () => false,
242+
} as unknown as Config;
243+
const manager = new McpClientManager(mockConfig, {} as ToolRegistry);
244+
245+
await manager.discoverMcpToolsForServer('unknown-server', {
246+
isTrustedFolder: () => true,
247+
} as unknown as Config);
248+
249+
expect(vi.mocked(McpClient)).not.toHaveBeenCalled();
250+
});
76251
});

‎packages/core/src/tools/mcp-client-manager.ts‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,73 @@ export class McpClientManager {
100100
this.discoveryState = MCPDiscoveryState.COMPLETED;
101101
}
102102

103+
/**
104+
* Connects to a single MCP server and discovers its tools/prompts.
105+
* The connected client is tracked so it can be closed by {@link stop}.
106+
*
107+
* This is primarily used for on-demand re-discovery flows (e.g. after OAuth).
108+
*/
109+
async discoverMcpToolsForServer(
110+
serverName: string,
111+
cliConfig: Config,
112+
): Promise<void> {
113+
const servers = populateMcpServerCommand(
114+
this.cliConfig.getMcpServers() || {},
115+
this.cliConfig.getMcpServerCommand(),
116+
);
117+
const serverConfig = servers[serverName];
118+
if (!serverConfig) {
119+
return;
120+
}
121+
122+
// Ensure we don't leak an existing connection for this server.
123+
const existingClient = this.clients.get(serverName);
124+
if (existingClient) {
125+
try {
126+
await existingClient.disconnect();
127+
} catch (error) {
128+
console.error(
129+
`Error stopping client '${serverName}': ${getErrorMessage(error)}`,
130+
);
131+
} finally {
132+
this.clients.delete(serverName);
133+
this.eventEmitter?.emit('mcp-client-update', this.clients);
134+
}
135+
}
136+
137+
// For SDK MCP servers, pass the sendSdkMcpMessage callback.
138+
const sdkCallback = isSdkMcpServerConfig(serverConfig)
139+
? this.sendSdkMcpMessage
140+
: undefined;
141+
142+
const client = new McpClient(
143+
serverName,
144+
serverConfig,
145+
this.toolRegistry,
146+
this.cliConfig.getPromptRegistry(),
147+
this.cliConfig.getWorkspaceContext(),
148+
this.cliConfig.getDebugMode(),
149+
sdkCallback,
150+
);
151+
152+
this.clients.set(serverName, client);
153+
this.eventEmitter?.emit('mcp-client-update', this.clients);
154+
155+
try {
156+
await client.connect();
157+
await client.discover(cliConfig);
158+
} catch (error) {
159+
// Log the error but don't throw: callers expect best-effort discovery.
160+
console.error(
161+
`Error during discovery for server '${serverName}': ${getErrorMessage(
162+
error,
163+
)}`,
164+
);
165+
} finally {
166+
this.eventEmitter?.emit('mcp-client-update', this.clients);
167+
}
168+
}
169+
103170
/**
104171
* Stops all running local MCP servers and closes all client connections.
105172
* This is the cleanup method to be called on application exit.

‎packages/core/src/tools/tool-registry.ts‎

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ import { Kind, BaseDeclarativeTool, BaseToolInvocation } from './tools.js';
1515
import type { Config } from '../config/config.js';
1616
import { spawn } from 'node:child_process';
1717
import { StringDecoder } from 'node:string_decoder';
18-
import { connectAndDiscover } from './mcp-client.js';
1918
import type { SendSdkMcpMessage } from './mcp-client.js';
2019
import { McpClientManager } from './mcp-client-manager.js';
2120
import { DiscoveredMCPTool } from './mcp-tool.js';
@@ -279,19 +278,10 @@ export class ToolRegistry {
279278

280279
this.config.getPromptRegistry().removePromptsByServer(serverName);
281280

282-
const mcpServers = this.config.getMcpServers() ?? {};
283-
const serverConfig = mcpServers[serverName];
284-
if (serverConfig) {
285-
await connectAndDiscover(
286-
serverName,
287-
serverConfig,
288-
this,
289-
this.config.getPromptRegistry(),
290-
this.config.getDebugMode(),
291-
this.config.getWorkspaceContext(),
292-
this.config,
293-
);
294-
}
281+
await this.mcpClientManager.discoverMcpToolsForServer(
282+
serverName,
283+
this.config,
284+
);
295285
}
296286

297287
private async discoverAndRegisterToolsFromCommand(): Promise<void> {
@@ -479,4 +469,17 @@ export class ToolRegistry {
479469
getTool(name: string): AnyDeclarativeTool | undefined {
480470
return this.tools.get(name);
481471
}
472+
473+
/**
474+
* Stops all MCP clients and cleans up resources.
475+
* This method is idempotent and safe to call multiple times.
476+
*/
477+
async stop(): Promise<void> {
478+
try {
479+
await this.mcpClientManager.stop();
480+
} catch (error) {
481+
// Log but don't throw - cleanup should be best-effort
482+
console.error('Error stopping MCP clients:', error);
483+
}
484+
}
482485
}

0 commit comments

Comments
 (0)