diff --git a/.changeset/fix-mcp-reconnect-state.md b/.changeset/fix-mcp-reconnect-state.md new file mode 100644 index 0000000000..361297ad48 --- /dev/null +++ b/.changeset/fix-mcp-reconnect-state.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fix MCP servers reporting as connected while their previous client is still closing during reconnect. diff --git a/packages/agent-core-v2/src/mcpCore/connection-manager.ts b/packages/agent-core-v2/src/mcpCore/connection-manager.ts index 854d73e18e..d80b20f66d 100644 --- a/packages/agent-core-v2/src/mcpCore/connection-manager.ts +++ b/packages/agent-core-v2/src/mcpCore/connection-manager.ts @@ -284,14 +284,15 @@ export class McpConnectionManager implements McpConnectionView { throw new Error2(ErrorCodes.MCP_SERVER_DISABLED, `MCP server is disabled: ${name}`); } const attemptId = this.beginConnectAttempt(entry); - await this.closeClient(entry); - if (!this.isCurrent(entry, attemptId)) return; + const closing = this.closeClient(entry); entry.status = 'pending'; entry.tools = undefined; entry.enabledNames = undefined; entry.rawTools = undefined; entry.error = undefined; this.emit(entry); + await closing; + if (!this.isCurrent(entry, attemptId)) return; await this.connectOne(entry, attemptId); } diff --git a/packages/agent-core-v2/test/mcpCore/connection-manager.test.ts b/packages/agent-core-v2/test/mcpCore/connection-manager.test.ts index 7f1784391b..7d74f8ea3a 100644 --- a/packages/agent-core-v2/test/mcpCore/connection-manager.test.ts +++ b/packages/agent-core-v2/test/mcpCore/connection-manager.test.ts @@ -30,6 +30,7 @@ import { z } from 'zod'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { Error2 } from '#/errors'; +import { StdioMcpClient } from '#/mcpCore/client-stdio'; import { KIMI_MCP_CLIENT_NAME } from '#/mcpCore/client-shared'; import { McpConnectionManager, type McpServerEntry } from '#/mcpCore/connection-manager'; import { McpOAuthService } from '#/mcpCore/oauth/service'; @@ -240,6 +241,54 @@ describe('McpConnectionManager', () => { } }, 15000); + it('reconnect reports pending while the previous client is still closing', async () => { + const cm = new McpConnectionManager(); + let releaseClose!: () => void; + let reportCloseStarted!: () => void; + const closeGate = new Promise((resolve) => { + releaseClose = resolve; + }); + const closeStarted = new Promise((resolve) => { + reportCloseStarted = resolve; + }); + const originalClose = StdioMcpClient.prototype.close; + const closeSpy = vi + .spyOn(StdioMcpClient.prototype, 'close') + .mockImplementation(async function (this: StdioMcpClient) { + reportCloseStarted(); + await closeGate; + await originalClose.call(this); + }); + const seen: Array<[McpServerEntry['status'], number, string | undefined]> = []; + let reconnect: Promise | undefined; + + try { + await cm.connectAll({ alpha: stdioConfig() }); + cm.onStatusChange((entry) => { + seen.push([entry.status, entry.toolCount, entry.error]); + }); + + reconnect = cm.reconnect('alpha'); + await closeStarted; + + expect(cm.get('alpha')).toMatchObject({ + status: 'pending', + toolCount: 0, + error: undefined, + }); + expect(cm.resolved('alpha')).toBeUndefined(); + expect(seen).toEqual([['pending', 0, undefined]]); + } finally { + releaseClose(); + await reconnect?.catch(() => undefined); + try { + await cm.shutdown(); + } finally { + closeSpy.mockRestore(); + } + } + }, 15000); + it('reconnect cycles a failed server back through pending and into connected when fixed', async () => { const cm = new McpConnectionManager(); try { diff --git a/packages/agent-core/src/mcp/connection-manager.ts b/packages/agent-core/src/mcp/connection-manager.ts index 7e46c1e3e4..8850ad1a89 100644 --- a/packages/agent-core/src/mcp/connection-manager.ts +++ b/packages/agent-core/src/mcp/connection-manager.ts @@ -304,14 +304,15 @@ export class McpConnectionManager { throw new KimiError(ErrorCodes.MCP_SERVER_DISABLED, `MCP server is disabled: ${name}`); } const attemptId = this.beginConnectAttempt(entry); - await this.closeClient(entry); - if (!this.isCurrent(entry, attemptId)) return; + const closing = this.closeClient(entry); entry.status = 'pending'; entry.tools = undefined; entry.rawTools = undefined; entry.enabledNames = undefined; entry.error = undefined; this.emit(entry); + await closing; + if (!this.isCurrent(entry, attemptId)) return; await this.connectOne(entry, attemptId); } diff --git a/packages/agent-core/test/mcp/connection-manager.test.ts b/packages/agent-core/test/mcp/connection-manager.test.ts index 2c699a7348..e6dccd208f 100644 --- a/packages/agent-core/test/mcp/connection-manager.test.ts +++ b/packages/agent-core/test/mcp/connection-manager.test.ts @@ -32,6 +32,7 @@ import type { import { z } from 'zod'; import { KimiError } from '../../src/errors'; +import { StdioMcpClient } from '../../src/mcp/client-stdio'; import { ProviderManager } from '../../src/session/provider-manager'; import { MCP_STARTUP_TIMEOUT_ENV, @@ -206,6 +207,54 @@ describe('McpConnectionManager', () => { } }, 15000); + it('reconnect reports pending while the previous client is still closing', async () => { + const cm = new McpConnectionManager(); + let releaseClose!: () => void; + let reportCloseStarted!: () => void; + const closeGate = new Promise((resolve) => { + releaseClose = resolve; + }); + const closeStarted = new Promise((resolve) => { + reportCloseStarted = resolve; + }); + const originalClose = StdioMcpClient.prototype.close; + const closeSpy = vi + .spyOn(StdioMcpClient.prototype, 'close') + .mockImplementation(async function (this: StdioMcpClient) { + reportCloseStarted(); + await closeGate; + await originalClose.call(this); + }); + const seen: Array<[McpServerEntry['status'], number, string | undefined]> = []; + let reconnect: Promise | undefined; + + try { + await cm.connectAll({ alpha: stdioConfig() }); + cm.onStatusChange((entry) => { + seen.push([entry.status, entry.toolCount, entry.error]); + }); + + reconnect = cm.reconnect('alpha'); + await closeStarted; + + expect(cm.get('alpha')).toMatchObject({ + status: 'pending', + toolCount: 0, + error: undefined, + }); + expect(cm.resolved('alpha')).toBeUndefined(); + expect(seen).toEqual([['pending', 0, undefined]]); + } finally { + releaseClose(); + await reconnect?.catch(() => undefined); + try { + await cm.shutdown(); + } finally { + closeSpy.mockRestore(); + } + } + }, 15000); + it('reconnect cycles a failed server back through pending and into connected when fixed', async () => { const cm = new McpConnectionManager(); try {