Skip to content

Commit 1507749

Browse files
tt-a1itt-a1i
committed
fix(core): properly cleanup MCP server subprocesses on exit
Fixes #1213 MCP server subprocesses were not being terminated when qwen-code exits, causing orphaned processes to accumulate. Root cause: `mcpClientManager.stop()` was never called during the exit cleanup flow. Changes: - Add `ToolRegistry.stop()` to stop all MCP clients - Add `Config.shutdown()` as a unified cleanup entry point - Register cleanup in gemini.tsx immediately after config creation - Add tests for stop behavior and idempotency The cleanup now covers all execution modes (interactive, non-interactive, stream-json) since it's registered at the CLI entry point.
1 parent 105ad74 commit 1507749

5 files changed

Lines changed: 288 additions & 17 deletions

File tree

‎packages/cli/src/gemini.tsx‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -343,6 +343,10 @@ export async function main() {
343343
argv,
344344
);
345345

346+
// Register cleanup for MCP clients as early as possible
347+
// This ensures MCP server subprocesses are properly terminated on exit
348+
registerCleanup(() => config.shutdown());
349+
346350
if (config.getListExtensions()) {
347351
console.log('Installed extensions:');
348352
for (const extension of extensions) {

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -845,6 +845,26 @@ export class Config {
845845
return this.toolRegistry;
846846
}
847847

848+
/**
849+
* Shuts down the Config and releases all resources.
850+
* This method is idempotent and safe to call multiple times.
851+
* It handles the case where initialization was not completed.
852+
*/
853+
async shutdown(): Promise<void> {
854+
if (!this.initialized) {
855+
// Nothing to clean up if not initialized
856+
return;
857+
}
858+
try {
859+
if (this.toolRegistry) {
860+
await this.toolRegistry.stop();
861+
}
862+
} catch (error) {
863+
// Log but don't throw - cleanup should be best-effort
864+
console.error('Error during Config shutdown:', error);
865+
}
866+
}
867+
848868
getPromptRegistry(): PromptRegistry {
849869
return this.promptRegistry;
850870
}

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

Lines changed: 180 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,8 @@ vi.mock('./mcp-client.js', async () => {
1717
return {
1818
...originalModule,
1919
McpClient: vi.fn(),
20-
populateMcpServerCommand: vi.fn(() => ({
21-
'test-server': {},
22-
})),
20+
// Return the input servers unchanged (identity function)
21+
populateMcpServerCommand: vi.fn((servers) => servers),
2322
};
2423
});
2524

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

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

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,73 @@ export class McpClientManager {
113113
this.discoveryState = MCPDiscoveryState.COMPLETED;
114114
}
115115

116+
/**
117+
* Connects to a single MCP server and discovers its tools/prompts.
118+
* The connected client is tracked so it can be closed by {@link stop}.
119+
*
120+
* This is primarily used for on-demand re-discovery flows (e.g. after OAuth).
121+
*/
122+
async discoverMcpToolsForServer(
123+
serverName: string,
124+
cliConfig: Config,
125+
): Promise<void> {
126+
const servers = populateMcpServerCommand(
127+
this.mcpServers,
128+
this.mcpServerCommand,
129+
);
130+
const serverConfig = servers[serverName];
131+
if (!serverConfig) {
132+
return;
133+
}
134+
135+
// Ensure we don't leak an existing connection for this server.
136+
const existingClient = this.clients.get(serverName);
137+
if (existingClient) {
138+
try {
139+
await existingClient.disconnect();
140+
} catch (error) {
141+
console.error(
142+
`Error stopping client '${serverName}': ${getErrorMessage(error)}`,
143+
);
144+
} finally {
145+
this.clients.delete(serverName);
146+
this.eventEmitter?.emit('mcp-client-update', this.clients);
147+
}
148+
}
149+
150+
// For SDK MCP servers, pass the sendSdkMcpMessage callback.
151+
const sdkCallback = isSdkMcpServerConfig(serverConfig)
152+
? this.sendSdkMcpMessage
153+
: undefined;
154+
155+
const client = new McpClient(
156+
serverName,
157+
serverConfig,
158+
this.toolRegistry,
159+
this.promptRegistry,
160+
this.workspaceContext,
161+
this.debugMode,
162+
sdkCallback,
163+
);
164+
165+
this.clients.set(serverName, client);
166+
this.eventEmitter?.emit('mcp-client-update', this.clients);
167+
168+
try {
169+
await client.connect();
170+
await client.discover(cliConfig);
171+
} catch (error) {
172+
// Log the error but don't throw: callers expect best-effort discovery.
173+
console.error(
174+
`Error during discovery for server '${serverName}': ${getErrorMessage(
175+
error,
176+
)}`,
177+
);
178+
} finally {
179+
this.eventEmitter?.emit('mcp-client-update', this.clients);
180+
}
181+
}
182+
116183
/**
117184
* Stops all running local MCP servers and closes all client connections.
118185
* 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';
@@ -283,19 +282,10 @@ export class ToolRegistry {
283282

284283
this.config.getPromptRegistry().removePromptsByServer(serverName);
285284

286-
const mcpServers = this.config.getMcpServers() ?? {};
287-
const serverConfig = mcpServers[serverName];
288-
if (serverConfig) {
289-
await connectAndDiscover(
290-
serverName,
291-
serverConfig,
292-
this,
293-
this.config.getPromptRegistry(),
294-
this.config.getDebugMode(),
295-
this.config.getWorkspaceContext(),
296-
this.config,
297-
);
298-
}
285+
await this.mcpClientManager.discoverMcpToolsForServer(
286+
serverName,
287+
this.config,
288+
);
299289
}
300290

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

0 commit comments

Comments
 (0)