From 858b9104c941a7fddaf722a3609535cfd0ab5892 Mon Sep 17 00:00:00 2001 From: Tarik Ermis Date: Sun, 9 Aug 2026 01:39:12 +0200 Subject: [PATCH 1/2] fix(agent-core): reconnect dropped streamable-HTTP MCP sessions on tool call When a streamable-HTTP MCP server's session drops mid-session (e.g. a server restart), the legacy engine never re-established it: every call to that server's tools failed for the rest of the session, and the failed server's tools were unregistered, so the model retried its work through other servers' tools instead of getting a loud, server-named error (#2742). Port the agent-core-v2 recovery pattern (#1991) to the legacy engine: - The wrapped MCP tool call now classifies the failure: server-answered errors are rethrown, ambiguous transport failures are probed with a ping and retried once in place when alive, and a dead transport triggers one shared reconnect through the connection manager before retrying on the fresh client. - McpConnectionManager gains reconnectAndJoin so parallel failing calls collapse into a single reconnect. - A server that flips to failed keeps its tools registered, so the next call can drive the reconnect and calls fail with the server's own error in the meantime. The default agent-core-v2 engine already recovers this way; this brings the legacy engine (KIMI_CODE_LEGACY_FLAG, SDK v1 client) to parity. --- .changeset/fix-mcp-http-session-reconnect.md | 5 + packages/agent-core/src/agent/tool/index.ts | 152 ++++++- packages/agent-core/src/mcp/client-http.ts | 5 + packages/agent-core/src/mcp/client-shared.ts | 53 ++- packages/agent-core/src/mcp/client-sse.ts | 5 + packages/agent-core/src/mcp/client-stdio.ts | 5 + .../agent-core/src/mcp/connection-manager.ts | 19 + packages/agent-core/src/mcp/types.ts | 5 + .../test/agent/llm-request-recorder.test.ts | 2 + .../test/agent/tool-select.e2e.test.ts | 2 + .../test/mcp/connection-manager.test.ts | 43 ++ .../test/mcp/tool-manager-mcp.test.ts | 409 ++++++++++++++++++ 12 files changed, 690 insertions(+), 15 deletions(-) create mode 100644 .changeset/fix-mcp-http-session-reconnect.md diff --git a/.changeset/fix-mcp-http-session-reconnect.md b/.changeset/fix-mcp-http-session-reconnect.md new file mode 100644 index 0000000000..26e6fb00a9 --- /dev/null +++ b/.changeset/fix-mcp-http-session-reconnect.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fix dropped streamable-HTTP MCP sessions never reconnecting on the legacy engine: tool calls now reconnect and retry transparently, and a failed server's tools stay registered and fail with the server's own error while it is down. diff --git a/packages/agent-core/src/agent/tool/index.ts b/packages/agent-core/src/agent/tool/index.ts index 6431b873a0..a0827aa0b0 100644 --- a/packages/agent-core/src/agent/tool/index.ts +++ b/packages/agent-core/src/agent/tool/index.ts @@ -7,16 +7,24 @@ import { collectLoadedDynamicToolNames, } from '../context/dynamic-tools'; import type { ContextMessage } from '../context/types'; -import { makeErrorPayload } from '../../errors'; -import type { ExecutableTool, ToolUpdate } from '../../loop'; +import { ErrorCodes, KimiError, makeErrorPayload } from '../../errors'; +import type { ExecutableTool, ExecutableToolContext, ToolUpdate } from '../../loop'; +import { errorMessage, isAbortError } from '../../loop/errors'; import { createMcpAuthTool } from '../../mcp/auth-tool'; import type { McpConnectionManager, McpServerEntry } from '../../mcp'; +import { + isMcpConnectionClosedError, + isMcpMalformedResultError, + isMcpTransportFailure, + probeMcpLiveness, +} from '../../mcp/client-shared'; import { mcpResultToExecutableOutput } from '../../mcp/output'; import { isMcpToolName, qualifyMcpToolName } from '../../mcp/tool-naming'; -import type { MCPClient, MCPToolDefinition } from '../../mcp/types'; +import type { MCPClient, MCPToolDefinition, MCPToolResult } from '../../mcp/types'; import { resolveSubagentTimeoutMs } from '../../session/subagent-host'; import { buildSubagentModelDescriptions } from '../../session/subagent-binding'; import { extendWorkspaceWithSkillRoots } from '../../skill'; +import { abortable } from '../../utils/abort'; import { fingerprint } from '../llm-request-logger'; import * as b from '../../tools/builtin'; import type { ToolStore, ToolStoreData, ToolStoreKey } from '../../tools/store'; @@ -327,11 +335,19 @@ export class ToolManager { // `args` has already been JSON-parsed and schema-validated by // the loop's preflight (`loop/tool-call.ts`), so the MCP // client gets a plain object directly. - const result = await client.callTool( - tool.name, - (args ?? {}) as Record, - context.signal, - ); + const mcpArgs = (args ?? {}) as Record; + let result: MCPToolResult; + try { + result = await client.callTool(tool.name, mcpArgs, context.signal); + } catch (error) { + result = await retryMcpCallAfterReconnect( + error, + client, + (activeClient) => activeClient.callTool(tool.name, mcpArgs, context.signal), + context, + this.mcpToolCallReconnect(serverName, client, context.signal), + ); + } return mcpResultToExecutableOutput(result, qualified, { originalsDir: this.agent.mediaOriginalsDir, telemetry: this.agent.telemetry, @@ -359,6 +375,22 @@ export class ToolManager { return true; } + /** + * Builds the `reconnect` callback for the wrapped-tool recovery path (see + * {@link retryMcpCallAfterReconnect}), or returns `undefined` when this + * manager has no connection manager behind it (tests wiring a bare fake + * client) — in that case calls keep their old fail-fast behavior. + */ + private mcpToolCallReconnect( + serverName: string, + staleClient: MCPClient, + signal: AbortSignal, + ): (() => Promise) | undefined { + const mcp = this.agent.mcp; + if (mcp === undefined) return undefined; + return () => abortable(joinHealedMcpClientOrReconnect(mcp, serverName, staleClient), signal); + } + private handleMcpServerStatusChange(mcp: McpConnectionManager, entry: McpServerEntry): void { if (entry.status === 'connected') { this.registerConnectedMcpServer(mcp, entry); @@ -369,12 +401,12 @@ export class ToolManager { return; } if (entry.status === 'failed') { - this.unregisterMcpServer(entry.name); - this.agent.emitEvent({ - type: 'tool.list.updated', - reason: 'mcp.failed', - serverName: entry.name, - }); + // Keep the tools registered: a dropped connection is recovered through + // the wrapped call's reconnect-and-retry path, and until then calls + // fail loudly with the server's error instead of vanishing from the + // tool list (which made the model retry through *other* servers' + // tools — see #2742). The tool list itself did not change, so no + // `tool.list.updated` event is emitted. return; } if (entry.status === 'disabled' || entry.status === 'pending') { @@ -1025,3 +1057,95 @@ export class ToolManager { .filter((tool) => !!tool); } } + +/** + * Recovery for a failed MCP tool call, mirroring agent-core-v2's + * `retryAfterReconnect` (`agent-core-v2/src/agent/mcp/tools/mcp.ts`): + * + * - The server answered (a JSON-RPC error, or a response that failed + * client-side schema validation) → the error is rethrown; reconnecting + * would not change the answer. + * - The failure is ambiguous (a raw fetch/socket error) → the client is + * probed with a ping: alive means a transient blip and the call is + * retried once in place; dead means the transport is gone. + * - The transport is provably dead (the SDK reported the connection closed, + * or the probe failed) → the server is reconnected once through + * `reconnect` and the call retried on the fresh client, so a dropped + * streamable-HTTP session (e.g. an MCP server restart) surfaces as a slow + * call instead of failing every call for the rest of the session (#2742). + * + * Retries are at-least-once: if the transport died after the server + * processed the call but before the response arrived, the retry may + * duplicate side effects. There is no protocol-level dedup across + * reconnects, so this trade-off is accepted deliberately. + */ +async function retryMcpCallAfterReconnect( + error: unknown, + client: MCPClient, + callTool: (activeClient: MCPClient) => Promise, + context: ExecutableToolContext, + reconnect: (() => Promise) | undefined, +): Promise { + const isUnrecoverable = (e: unknown): boolean => + context.signal.aborted || + isAbortError(e) || + !isMcpTransportFailure(e) || + isMcpMalformedResultError(e); + if (reconnect === undefined || isUnrecoverable(error)) { + throw error; + } + + let failure = error; + if (!isMcpConnectionClosedError(failure)) { + const alive = await probeMcpLiveness(client, context.signal); + context.signal.throwIfAborted(); + if (alive) { + try { + return await callTool(client); + } catch (retryError) { + if (isUnrecoverable(retryError)) { + throw retryError; + } + failure = retryError; + } + } + } + + context.onUpdate?.({ kind: 'status', text: 'MCP connection lost — reconnecting…' }); + let freshClient: MCPClient | undefined; + try { + freshClient = await reconnect(); + } catch (reconnectError) { + if (context.signal.aborted || isAbortError(reconnectError)) { + throw reconnectError; + } + throw new KimiError( + ErrorCodes.MCP_STARTUP_FAILED, + `${errorMessage(failure)} (reconnecting the MCP server also failed: ${errorMessage(reconnectError)})`, + { cause: reconnectError }, + ); + } + if (freshClient === undefined) { + throw failure; + } + return callTool(freshClient); +} + +/** + * Return the current client when the server already healed (a concurrent + * call finished the reconnect first), otherwise drive one shared reconnect + * through the manager — `reconnectAndJoin` dedupes parallel attempts — and + * return the client it produced. `undefined` means there is no fresh client + * to retry on, so the original failure should surface. + */ +async function joinHealedMcpClientOrReconnect( + mcp: McpConnectionManager, + serverName: string, + staleClient: MCPClient, +): Promise { + const healed = mcp.resolved(serverName)?.client; + if (healed !== undefined && healed !== staleClient) return healed; + await mcp.reconnectAndJoin(serverName); + const current = mcp.resolved(serverName)?.client; + return current !== undefined && current !== staleClient ? current : undefined; +} diff --git a/packages/agent-core/src/mcp/client-http.ts b/packages/agent-core/src/mcp/client-http.ts index f38b54a47d..6a13eecd32 100644 --- a/packages/agent-core/src/mcp/client-http.ts +++ b/packages/agent-core/src/mcp/client-http.ts @@ -7,6 +7,7 @@ import { buildRequestOptions, KIMI_MCP_CLIENT_NAME, KIMI_MCP_CLIENT_VERSION, + MCP_LIVENESS_PROBE_TIMEOUT_MS, toMcpToolDefinition, toMcpToolResult, type UnexpectedCloseListener, @@ -145,6 +146,10 @@ export class HttpMcpClient implements MCPClient { return toMcpToolResult(result); } + async ping(signal?: AbortSignal): Promise { + await this.client.ping(buildRequestOptions(MCP_LIVENESS_PROBE_TIMEOUT_MS, signal)); + } + private async closeStartedClient(): Promise { if (!this.started) return; this.started = false; diff --git a/packages/agent-core/src/mcp/client-shared.ts b/packages/agent-core/src/mcp/client-shared.ts index 18ee914e5d..9c7180e4c1 100644 --- a/packages/agent-core/src/mcp/client-shared.ts +++ b/packages/agent-core/src/mcp/client-shared.ts @@ -1,6 +1,7 @@ import { getCoreVersion } from '#/version'; +import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js'; -import type { MCPToolDefinition, MCPToolResult } from './types'; +import type { MCPClient, MCPToolDefinition, MCPToolResult } from './types'; export const KIMI_MCP_CLIENT_NAME = 'kimi-code'; // Resolved from agent-core's package.json so MCP servers see the real version @@ -31,6 +32,56 @@ export interface McpRequestOptions { readonly signal?: AbortSignal; } +/** + * True when the SDK reports the connection itself as gone (the transport was + * closed, so no in-flight request can ever complete). + */ +export function isMcpConnectionClosedError(error: unknown): boolean { + return ( + error instanceof Error && + (error as Error & { readonly code?: unknown }).code === ErrorCode.ConnectionClosed + ); +} + +/** + * True when a failed tool call might recover after a reconnect: either the + * connection is closed, or the error is a raw transport/fetch failure rather + * than a JSON-RPC answer from the server ({@link McpError}) — reconnecting + * would not change a server-side answer. + */ +export function isMcpTransportFailure(error: unknown): boolean { + if (!(error instanceof Error)) return false; + if (isMcpConnectionClosedError(error)) return true; + return !(error instanceof McpError); +} + +/** Bounded so a wedged server cannot stall the reconnect decision indefinitely. */ +export const MCP_LIVENESS_PROBE_TIMEOUT_MS = 5_000; + +/** Response failed client-side schema validation: the server answered, so it is alive. */ +export function isMcpMalformedResultError(error: unknown): boolean { + return error instanceof Error && error.name === 'ZodError'; +} + +/** + * Ping the server to decide whether its transport is still usable after an + * ambiguous failure. A malformed answer still proves liveness; a timeout or + * any transport-level failure means dead. + */ +export async function probeMcpLiveness(client: MCPClient, signal: AbortSignal): Promise { + try { + await client.ping(signal); + return true; + } catch (error) { + if (isMcpConnectionClosedError(error)) return false; + if (isMcpMalformedResultError(error)) return true; + if (error instanceof McpError) { + return (error as Error & { readonly code?: unknown }).code !== ErrorCode.RequestTimeout; + } + return false; + } +} + /** * Build the `RequestOptions` object accepted by MCP SDK requests, including * either a configured timeout, an in-flight abort signal, both, or neither. diff --git a/packages/agent-core/src/mcp/client-sse.ts b/packages/agent-core/src/mcp/client-sse.ts index 254c0786de..64c488ca23 100644 --- a/packages/agent-core/src/mcp/client-sse.ts +++ b/packages/agent-core/src/mcp/client-sse.ts @@ -7,6 +7,7 @@ import { buildRequestOptions, KIMI_MCP_CLIENT_NAME, KIMI_MCP_CLIENT_VERSION, + MCP_LIVENESS_PROBE_TIMEOUT_MS, toMcpToolDefinition, toMcpToolResult, type UnexpectedCloseListener, @@ -136,6 +137,10 @@ export class SseMcpClient implements MCPClient { return toMcpToolResult(result); } + async ping(signal?: AbortSignal): Promise { + await this.client.ping(buildRequestOptions(MCP_LIVENESS_PROBE_TIMEOUT_MS, signal)); + } + private async closeStartedClient(): Promise { if (!this.started) return; this.started = false; diff --git a/packages/agent-core/src/mcp/client-stdio.ts b/packages/agent-core/src/mcp/client-stdio.ts index 267056d428..c9bd0972ea 100644 --- a/packages/agent-core/src/mcp/client-stdio.ts +++ b/packages/agent-core/src/mcp/client-stdio.ts @@ -10,6 +10,7 @@ import { buildRequestOptions, KIMI_MCP_CLIENT_NAME, KIMI_MCP_CLIENT_VERSION, + MCP_LIVENESS_PROBE_TIMEOUT_MS, toMcpToolDefinition, toMcpToolResult, type UnexpectedCloseListener, @@ -162,6 +163,10 @@ export class StdioMcpClient implements MCPClient { return toMcpToolResult(result); } + async ping(signal?: AbortSignal): Promise { + await this.client.ping(buildRequestOptions(MCP_LIVENESS_PROBE_TIMEOUT_MS, signal)); + } + private async closeStartedClient(): Promise { if (!this.started) return; this.started = false; diff --git a/packages/agent-core/src/mcp/connection-manager.ts b/packages/agent-core/src/mcp/connection-manager.ts index 152d3a8131..720571c0cb 100644 --- a/packages/agent-core/src/mcp/connection-manager.ts +++ b/packages/agent-core/src/mcp/connection-manager.ts @@ -129,6 +129,7 @@ export interface McpConnectionManagerOptions { export class McpConnectionManager { private readonly entries = new Map(); private readonly listeners = new Set(); + private readonly inFlightReconnects = new Map>(); private initialLoad: Promise = Promise.resolve(); private initialLoadAttemptId = 0; private initialLoadStartedAt: number | undefined; @@ -314,6 +315,24 @@ export class McpConnectionManager { await this.connectOne(entry, attemptId); } + /** + * Reconnect that joins an already in-flight reconnect for the same server + * instead of starting a second one. Used by the tool-call recovery path, + * where several parallel calls to a dropped server can all decide to + * reconnect at once. + */ + reconnectAndJoin(name: string): Promise { + const existing = this.inFlightReconnects.get(name); + if (existing !== undefined) return existing; + const work = this.reconnect(name).finally(() => { + if (this.inFlightReconnects.get(name) === work) { + this.inFlightReconnects.delete(name); + } + }); + this.inFlightReconnects.set(name, work); + return work; + } + async shutdown(): Promise { const entries = Array.from(this.entries.values()); this.entries.clear(); diff --git a/packages/agent-core/src/mcp/types.ts b/packages/agent-core/src/mcp/types.ts index aedb555406..c811b36edd 100644 --- a/packages/agent-core/src/mcp/types.ts +++ b/packages/agent-core/src/mcp/types.ts @@ -88,6 +88,11 @@ export interface MCPClient { args: Record, signal?: AbortSignal, ): Promise; + /** + * Liveness probe used after a failed tool call to distinguish a transient + * blip from a dead transport before attempting a reconnect. + */ + ping(signal?: AbortSignal): Promise; } /** diff --git a/packages/agent-core/test/agent/llm-request-recorder.test.ts b/packages/agent-core/test/agent/llm-request-recorder.test.ts index 7779cc63f6..9443ad7f07 100644 --- a/packages/agent-core/test/agent/llm-request-recorder.test.ts +++ b/packages/agent-core/test/agent/llm-request-recorder.test.ts @@ -212,6 +212,7 @@ describe('mcp.tools_discovered records', () => { async callTool() { return { content: [{ type: 'text', text: 'ok' }], isError: false }; }, + async ping() {}, }; const entry: McpServerEntry = { name: input.serverName ?? 'grafana', @@ -375,6 +376,7 @@ describe('mcp.tools_discovered records', () => { async callTool() { return { content: [], isError: false }; }, + async ping() {}, }; ctx.agent.tools.registerMcpServer('graf.ana', occupant, [ { name: 'query_range', description: 'occupies the qualified name', parameters: {} }, diff --git a/packages/agent-core/test/agent/tool-select.e2e.test.ts b/packages/agent-core/test/agent/tool-select.e2e.test.ts index 2531149ef6..86cb2952ec 100644 --- a/packages/agent-core/test/agent/tool-select.e2e.test.ts +++ b/packages/agent-core/test/agent/tool-select.e2e.test.ts @@ -80,6 +80,7 @@ function grafanaClient(callLog: Array<[string, unknown]> = []): MCPClient { callLog.push([name, args]); return { content: [{ type: 'text', text: 'error_rate=0.02' }], isError: false }; }, + async ping() {}, }; } @@ -795,6 +796,7 @@ describe('disclosure mode — compaction', () => { async callTool() { return { content: [{ type: 'text', text: 'ok' }], isError: false }; }, + async ping() {}, }; ctx.agent.tools.registerMcpServer( 'grafana', diff --git a/packages/agent-core/test/mcp/connection-manager.test.ts b/packages/agent-core/test/mcp/connection-manager.test.ts index dabd666011..be21977079 100644 --- a/packages/agent-core/test/mcp/connection-manager.test.ts +++ b/packages/agent-core/test/mcp/connection-manager.test.ts @@ -297,6 +297,49 @@ describe('McpConnectionManager', () => { } }); + it('reconnectAndJoin joins an in-flight reconnect instead of starting a second one', async () => { + const cm = new McpConnectionManager(); + const seen: Array<{ name: string; status: McpServerEntry['status'] }> = []; + cm.onStatusChange((entry) => { + seen.push({ name: entry.name, status: entry.status }); + }); + const delayedMockServer = `setTimeout(() => import(${JSON.stringify(pathToFileURL(stdioFixture).href)}), 250)`; + + try { + await cm.connectAll({ + slow: { + transport: 'stdio', + command: process.execPath, + args: ['-e', delayedMockServer], + startupTimeoutMs: 5_000, + }, + }); + seen.length = 0; + + await Promise.all([cm.reconnectAndJoin('slow'), cm.reconnectAndJoin('slow')]); + + expect(cm.get('slow')?.status).toBe('connected'); + expect(seen.filter((event) => event.name === 'slow').map((event) => event.status)).toEqual([ + 'pending', + 'connected', + ]); + } finally { + await cm.shutdown(); + } + }, 20000); + + it('reconnectAndJoin rejects for unknown servers', async () => { + const cm = new McpConnectionManager(); + try { + await expect(cm.reconnectAndJoin('nope')).rejects.toBeInstanceOf(KimiError); + await expect(cm.reconnectAndJoin('nope')).rejects.toMatchObject({ + code: 'mcp.server_not_found', + }); + } finally { + await cm.shutdown(); + } + }); + it('shutdown clears entries and is idempotent', async () => { const cm = new McpConnectionManager(); await cm.connectAll({ alpha: stdioConfig() }); diff --git a/packages/agent-core/test/mcp/tool-manager-mcp.test.ts b/packages/agent-core/test/mcp/tool-manager-mcp.test.ts index f7f14bbe3f..808d62926e 100644 --- a/packages/agent-core/test/mcp/tool-manager-mcp.test.ts +++ b/packages/agent-core/test/mcp/tool-manager-mcp.test.ts @@ -1,8 +1,18 @@ +import { randomUUID } from 'node:crypto'; +import { createServer, type Server } from 'node:http'; +import type { AddressInfo } from 'node:net'; + import type { ContentPart, Tool } from '@moonshot-ai/kosong'; +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; +import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'; +import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js'; import { describe, expect, it } from 'vitest'; +import { z } from 'zod'; import type { Agent } from '../../src/agent'; import { ToolManager } from '../../src/agent/tool'; +import { KimiError } from '../../src/errors'; +import { McpConnectionManager, type McpServerEntry } from '../../src/mcp/connection-manager'; import type { MCPClient } from '../../src/mcp/types'; import { testAgent } from '../agent/harness/agent'; import { executeTool } from '../tools/fixtures/execute-tool'; @@ -53,6 +63,7 @@ function fakeClient(): MCPClient { } return { content: [{ type: 'text', text: 'ok' }], isError: false }; }, + async ping() {}, }; } @@ -133,6 +144,7 @@ describe('ToolManager MCP integration', () => { async callTool() { return { content: [{ type: 'text', text: 'ok' }], isError: false }; }, + async ping() {}, }; const result = tm.registerMcpServer('srv', colliding, await discoverTools(colliding)); @@ -160,6 +172,7 @@ describe('ToolManager MCP integration', () => { async callTool() { return { content: [{ type: 'text', text: 'first' }], isError: false }; }, + async ping() {}, }; const secondClient: MCPClient = { async listTools() { @@ -174,6 +187,7 @@ describe('ToolManager MCP integration', () => { async callTool() { return { content: [{ type: 'text', text: 'second' }], isError: false }; }, + async ping() {}, }; // Both servers collapse to the same sanitized form ("srv_a"), so the @@ -210,6 +224,7 @@ describe('ToolManager MCP integration', () => { async callTool() { return { content: [], isError: false }; }, + async ping() {}, }; tm.registerMcpServer('s', firstClient, await discoverTools(firstClient)); @@ -295,6 +310,7 @@ describe('ToolManager MCP integration', () => { isError: false, }; }, + async ping() {}, }; tm.registerMcpServer('s', client, await discoverTools(client)); const big = tm.loopTools.find((t) => t.name === 'mcp__s__big'); @@ -332,6 +348,7 @@ describe('ToolManager MCP integration', () => { isError: false, }; }, + async ping() {}, }; tm.registerMcpServer('s', client, await discoverTools(client)); const snap = tm.loopTools.find((t) => t.name === 'mcp__s__snap'); @@ -380,6 +397,7 @@ describe('ToolManager MCP integration', () => { isError: false, }; }, + async ping() {}, }; tm.registerMcpServer('s', client, await discoverTools(client)); const tool = tm.loopTools.find((t) => t.name === 'mcp__s__huge_img'); @@ -435,6 +453,7 @@ describe('ToolManager MCP integration', () => { isError: false, }; }, + async ping() {}, }; tm.registerMcpServer('s', client, await discoverTools(client)); const tool = tm.loopTools.find((t) => t.name === 'mcp__s__mixed'); @@ -483,6 +502,7 @@ describe('ToolManager MCP integration', () => { isError: false, }; }, + async ping() {}, }; tm.registerMcpServer('s', client, await discoverTools(client)); const tool = tm.loopTools.find((t) => t.name === 'mcp__s__mixed'); @@ -533,6 +553,7 @@ describe('ToolManager MCP integration', () => { receivedSignal = signal; return { content: [{ type: 'text', text: String(args['text']) }], isError: false }; }, + async ping() {}, }; tm.registerMcpServer('s', client, await discoverTools(client)); const echo = tm.loopTools.find((t) => t.name === 'mcp__s__echo'); @@ -661,3 +682,391 @@ describe('ToolManager MCP integration', () => { }); }); }); + +describe('ToolManager MCP reconnect-on-call', () => { + const toolCallContext = () => ({ + turnId: '1', + toolCallId: 'tc-1', + signal: new AbortController().signal, + }); + + function fakeAgentWithMcp(mcp: McpConnectionManager): Agent { + return { + records: { + observabilityReady: true, + logRecord() {}, + onOpened() {}, + }, + config: { + data: () => ({ provider: undefined }), + }, + goal: { + getGoal: () => ({ goal: null }), + }, + mcp, + emitEvent() {}, + } as unknown as Agent; + } + + interface FakeMcpManagerOptions { + readonly resolvedClient: () => MCPClient | undefined; + readonly onReconnectAndJoin?: () => void | Promise; + readonly onStatusListener?: (listener: (entry: McpServerEntry) => void) => void; + } + + function fakeMcpManager(options: FakeMcpManagerOptions): McpConnectionManager { + return { + list: () => [], + oauthService: undefined, + onStatusChange: (listener: (entry: McpServerEntry) => void) => { + options.onStatusListener?.(listener); + return () => {}; + }, + resolved: () => { + const client = options.resolvedClient(); + if (client === undefined) return undefined; + return { client, tools: [], rawTools: [], enabledNames: new Set() }; + }, + reconnectAndJoin: async () => { + await options.onReconnectAndJoin?.(); + }, + } as unknown as McpConnectionManager; + } + + function deadTransportClient(): MCPClient { + return { + async listTools() { + return []; + }, + async callTool() { + throw new Error('fetch failed'); + }, + async ping() { + throw new Error('fetch failed'); + }, + }; + } + + async function mcpEchoTool(tm: ToolManager, client: MCPClient) { + tm.setActiveTools(['mcp__*']); + tm.registerMcpServer('s', client, await discoverTools(fakeClient())); + const echo = tm.loopTools.find((t) => t.name === 'mcp__s__echo'); + expect(echo).toBeDefined(); + return echo!; + } + + it('reconnects a dead transport and retries the call on the fresh client', async () => { + const staleClient = deadTransportClient(); + const freshClient = fakeClient(); + let reconnects = 0; + const mcp = fakeMcpManager({ + resolvedClient: () => (reconnects === 0 ? staleClient : freshClient), + onReconnectAndJoin: () => { + reconnects += 1; + }, + }); + const tm = new ToolManager(fakeAgentWithMcp(mcp)); + const echo = await mcpEchoTool(tm, staleClient); + + const result = await executeTool(echo, { ...toolCallContext(), args: { text: 'hello world' } }); + expect(result.isError).toBe(false); + expect(result.output).toBe('hello world'); + expect(reconnects).toBe(1); + }); + + it('retries in place without reconnecting when the server is still alive', async () => { + let calls = 0; + let pings = 0; + const client: MCPClient = { + async listTools() { + return []; + }, + async callTool() { + calls += 1; + if (calls === 1) throw new Error('fetch failed'); + return { content: [{ type: 'text', text: 'recovered' }], isError: false }; + }, + async ping() { + pings += 1; + }, + }; + let reconnects = 0; + const mcp = fakeMcpManager({ + resolvedClient: () => undefined, + onReconnectAndJoin: () => { + reconnects += 1; + }, + }); + const tm = new ToolManager(fakeAgentWithMcp(mcp)); + const echo = await mcpEchoTool(tm, client); + + const result = await executeTool(echo, { ...toolCallContext(), args: { text: 'x' } }); + expect(result.isError).toBe(false); + expect(result.output).toBe('recovered'); + expect(pings).toBe(1); + expect(calls).toBe(2); + expect(reconnects).toBe(0); + }); + + it('rethrows server-answered MCP errors without probing or reconnecting', async () => { + let pings = 0; + const client: MCPClient = { + async listTools() { + return []; + }, + async callTool() { + throw new McpError(ErrorCode.InternalError, 'server said no'); + }, + async ping() { + pings += 1; + }, + }; + let reconnects = 0; + const mcp = fakeMcpManager({ + resolvedClient: () => undefined, + onReconnectAndJoin: () => { + reconnects += 1; + }, + }); + const tm = new ToolManager(fakeAgentWithMcp(mcp)); + const echo = await mcpEchoTool(tm, client); + + await expect(executeTool(echo, { ...toolCallContext(), args: {} })).rejects.toThrow( + 'server said no', + ); + expect(pings).toBe(0); + expect(reconnects).toBe(0); + }); + + it('surfaces both errors when the reconnect attempt itself fails', async () => { + const staleClient = deadTransportClient(); + const mcp = fakeMcpManager({ + resolvedClient: () => staleClient, + onReconnectAndJoin: () => { + throw new Error('connect ECONNREFUSED'); + }, + }); + const tm = new ToolManager(fakeAgentWithMcp(mcp)); + const echo = await mcpEchoTool(tm, staleClient); + + const rejection = await executeTool(echo, { ...toolCallContext(), args: {} }).then( + () => undefined, + (error: unknown) => error, + ); + expect(rejection).toBeInstanceOf(KimiError); + expect((rejection as KimiError).code).toBe('mcp.startup_failed'); + expect((rejection as Error).message).toContain('fetch failed'); + expect((rejection as Error).message).toContain('ECONNREFUSED'); + }); + + it('rethrows malformed results without probing or reconnecting', async () => { + let pings = 0; + const client: MCPClient = { + async listTools() { + return []; + }, + async callTool() { + const malformed = new Error('response failed schema validation'); + malformed.name = 'ZodError'; + throw malformed; + }, + async ping() { + pings += 1; + }, + }; + let reconnects = 0; + const mcp = fakeMcpManager({ + resolvedClient: () => undefined, + onReconnectAndJoin: () => { + reconnects += 1; + }, + }); + const tm = new ToolManager(fakeAgentWithMcp(mcp)); + const echo = await mcpEchoTool(tm, client); + + await expect(executeTool(echo, { ...toolCallContext(), args: {} })).rejects.toThrow( + 'response failed schema validation', + ); + expect(pings).toBe(0); + expect(reconnects).toBe(0); + }); + + it('surfaces the original error when the reconnect produces no fresh client', async () => { + const staleClient = deadTransportClient(); + let reconnects = 0; + const mcp = fakeMcpManager({ + resolvedClient: () => staleClient, + onReconnectAndJoin: () => { + reconnects += 1; + }, + }); + const tm = new ToolManager(fakeAgentWithMcp(mcp)); + const echo = await mcpEchoTool(tm, staleClient); + + await expect(executeTool(echo, { ...toolCallContext(), args: {} })).rejects.toThrow( + 'fetch failed', + ); + expect(reconnects).toBe(1); + }); + + it('does not reconnect when the call was aborted', async () => { + const controller = new AbortController(); + controller.abort(); + const staleClient = deadTransportClient(); + let reconnects = 0; + const mcp = fakeMcpManager({ + resolvedClient: () => staleClient, + onReconnectAndJoin: () => { + reconnects += 1; + }, + }); + const tm = new ToolManager(fakeAgentWithMcp(mcp)); + const echo = await mcpEchoTool(tm, staleClient); + + await expect( + executeTool(echo, { + turnId: '1', + toolCallId: 'tc-abort', + signal: controller.signal, + args: {}, + }), + ).rejects.toThrow('fetch failed'); + expect(reconnects).toBe(0); + }); + + it('skips the liveness probe when the connection is already closed', async () => { + let pings = 0; + const staleClient: MCPClient = { + async listTools() { + return []; + }, + async callTool() { + const closed = new Error('Connection closed') as Error & { code: number }; + closed.code = ErrorCode.ConnectionClosed; + throw closed; + }, + async ping() { + pings += 1; + }, + }; + const freshClient = fakeClient(); + let reconnects = 0; + const mcp = fakeMcpManager({ + resolvedClient: () => (reconnects === 0 ? staleClient : freshClient), + onReconnectAndJoin: () => { + reconnects += 1; + }, + }); + const tm = new ToolManager(fakeAgentWithMcp(mcp)); + const echo = await mcpEchoTool(tm, staleClient); + + const result = await executeTool(echo, { ...toolCallContext(), args: { text: 'again' } }); + expect(result.isError).toBe(false); + expect(result.output).toBe('again'); + expect(pings).toBe(0); + expect(reconnects).toBe(1); + }); + + it('keeps a failed server’s tools registered so the next call can drive the reconnect', async () => { + let statusListener: ((entry: McpServerEntry) => void) | undefined; + const mcp = fakeMcpManager({ + resolvedClient: () => undefined, + onStatusListener: (listener) => { + statusListener = listener; + }, + }); + const tm = new ToolManager(fakeAgentWithMcp(mcp)); + const client = fakeClient(); + tm.registerMcpServer('srv', client, await discoverTools(client)); + tm.setActiveTools(['mcp__*']); + expect(tm.loopTools.map((t) => t.name)).toContain('mcp__srv__echo'); + + statusListener?.({ + name: 'srv', + transport: 'http', + status: 'failed', + toolCount: 0, + error: 'connection dropped', + }); + + expect(tm.loopTools.map((t) => t.name)).toContain('mcp__srv__echo'); + }); + + it('recovers calls to a streamable-HTTP server that restarts mid-session (#2742)', async () => { + let server = await startHttpEchoServer(); + const url = `http://127.0.0.1:${server.port}/mcp`; + + const manager = new McpConnectionManager(); + const tm = new ToolManager(fakeAgentWithMcp(manager)); + tm.setActiveTools(['mcp__*']); + try { + await manager.connectAll({ srv: { transport: 'http', url } }); + await manager.waitForInitialLoad(); + expect(manager.get('srv')?.status).toBe('connected'); + + const echo = tm.loopTools.find((t) => t.name === 'mcp__srv__echo'); + expect(echo).toBeDefined(); + const context = toolCallContext(); + + const first = await executeTool(echo!, { ...context, args: { text: 'before' } }); + expect(first.isError).toBe(false); + expect(first.output).toBe('before'); + + // Restart the server mid-session: the replacement process has no memory + // of the old streamable-HTTP session. + await server.close(); + server = await startHttpEchoServer(server.port); + + // The drop heals transparently: the SAME tool on the SAME server + // answers — no error, and no other server's tool runs in its place. + const second = await executeTool(echo!, { ...context, args: { text: 'after' } }); + expect(second.isError).toBe(false); + expect(second.output).toBe('after'); + expect(manager.get('srv')?.status).toBe('connected'); + expect(tm.loopTools.map((t) => t.name)).toContain('mcp__srv__echo'); + + const third = await executeTool(echo!, { ...context, args: { text: 'third' } }); + expect(third.isError).toBe(false); + expect(third.output).toBe('third'); + } finally { + await manager.shutdown(); + await server.close(); + } + }, 20000); +}); + +async function startHttpEchoServer( + port = 0, +): Promise<{ port: number; close: () => Promise }> { + const mcpServer = new McpServer({ name: 'mock-http', version: '0.0.1' }); + mcpServer.registerTool( + 'echo', + { description: 'Echoes text', inputSchema: { text: z.string() } }, + ({ text }) => ({ content: [{ type: 'text', text }] }), + ); + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + }); + await mcpServer.connect(transport); + const httpServer: Server = createServer((req, res) => { + void transport.handleRequest(req, res); + }); + await new Promise((resolve) => httpServer.listen(port, '127.0.0.1', resolve)); + const boundPort = (httpServer.address() as AddressInfo).port; + return { + port: boundPort, + close: () => + new Promise((resolve, reject) => { + httpServer.close((err) => { + if (err) { + reject(err); + return; + } + resolve(); + }); + // Force-close the client's keep-alive / SSE sockets so `close()` + // returns even while a long-lived streamable-HTTP stream is open. + httpServer.closeAllConnections(); + }), + }; +} From 6db3f534d900dc2e66ed79c6f987df45672a67eb Mon Sep 17 00:00:00 2001 From: Tarik Ermis Date: Sun, 9 Aug 2026 07:10:49 +0200 Subject: [PATCH 2/2] fix(agent-core): keep MCP tools registered through reconnect pending state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses a P1 review finding on PR #2748: McpConnectionManager.reconnect() emits 'pending' before 'failed', and the v1 ToolManager unregistered a server's tools on 'pending' too — so one failed call-driven recovery stranded the session without tools and later calls could never drive another reconnect (the exact bug the PR fixes, #2742). Only 'disabled' now unregisters tools and emits mcp.disconnected; both re-registration paths (connected, needs-auth) already start by unregistering, so keeping tools through pending never leaves stale entries behind. Regression coverage: a real-HTTP-server test drives a call through a failed reconnect (server down), asserts the tools stay registered, no tool-list event fires, and the transport's own error surfaces byte-equivalently, then brings the server back and verifies the next call re-drives the reconnect and heals. --- .changeset/fix-mcp-http-session-reconnect.md | 2 +- packages/agent-core/src/agent/tool/index.ts | 10 +- .../test/mcp/connection-manager.test.ts | 12 ++- .../test/mcp/tool-manager-mcp.test.ts | 95 ++++++++++++++++++- 4 files changed, 109 insertions(+), 10 deletions(-) diff --git a/.changeset/fix-mcp-http-session-reconnect.md b/.changeset/fix-mcp-http-session-reconnect.md index 26e6fb00a9..10334b87d3 100644 --- a/.changeset/fix-mcp-http-session-reconnect.md +++ b/.changeset/fix-mcp-http-session-reconnect.md @@ -2,4 +2,4 @@ "@moonshot-ai/kimi-code": patch --- -Fix dropped streamable-HTTP MCP sessions never reconnecting on the legacy engine: tool calls now reconnect and retry transparently, and a failed server's tools stay registered and fail with the server's own error while it is down. +Fix dropped streamable-HTTP MCP sessions never reconnecting on the legacy engine: tool calls now reconnect and retry transparently, and a failed server's tools stay registered and fail with the server's own error while it is down. `tool.list.updated` with reason `mcp.disconnected` now fires only when a server is removed or disabled, not on transient reconnects. diff --git a/packages/agent-core/src/agent/tool/index.ts b/packages/agent-core/src/agent/tool/index.ts index a0827aa0b0..0af1fd36fc 100644 --- a/packages/agent-core/src/agent/tool/index.ts +++ b/packages/agent-core/src/agent/tool/index.ts @@ -409,7 +409,7 @@ export class ToolManager { // `tool.list.updated` event is emitted. return; } - if (entry.status === 'disabled' || entry.status === 'pending') { + if (entry.status === 'disabled') { const removed = this.unregisterMcpServer(entry.name); if (removed) { this.agent.emitEvent({ @@ -419,6 +419,14 @@ export class ToolManager { }); } } + // `pending` is deliberately NOT handled: it precedes every (re)connect + // attempt, so unregistering here would drop the tools mid-reconnect — + // and after a failed recovery they would stay gone for the rest of the + // session, leaving later calls unable to drive another reconnect (#2742). + // Keeping them is safe because both re-registration paths above start by + // unregistering the server's tools (`registerMcpServer` and + // `registerNeedsAuthMcpServer`), so a `connected` or `needs-auth` + // transition fully replaces the tool set rather than merging into it. } private registerNeedsAuthMcpServer(mcp: McpConnectionManager, entry: McpServerEntry): void { diff --git a/packages/agent-core/test/mcp/connection-manager.test.ts b/packages/agent-core/test/mcp/connection-manager.test.ts index be21977079..8d3a7e5888 100644 --- a/packages/agent-core/test/mcp/connection-manager.test.ts +++ b/packages/agent-core/test/mcp/connection-manager.test.ts @@ -1128,7 +1128,7 @@ describe('Session MCP startup', () => { } }, 7000); - it('emits tool.list.updated(mcp.disconnected) when reconnect drops the live tools', async () => { + it('keeps tools registered through a manual reconnect (no mcp.disconnected)', async () => { const tmp = await mkdtemp(join(tmpdir(), 'kimi-session-mcp-reconnect-')); const events: SessionRpcEvent[] = []; const session = new Session({ @@ -1165,16 +1165,18 @@ describe('Session MCP startup', () => { events.length = 0; await session.mcp.reconnect('good'); - // The reconnect cycle: pending (tools cleared) → connected (tools back). - // Both transitions must surface as tool.list.updated so SDK consumers - // watching that event don't see stale tools mid-cycle. + // The reconnect cycle keeps the existing tools registered through + // `pending` and only swaps them when `connected` re-registers: the + // tool list never gaps mid-cycle, so no mcp.disconnected is emitted + // and a failed attempt cannot strand the session without the tools + // (#2742). const disconnects = events.filter( (e) => e.type === 'tool.list.updated' && e.reason === 'mcp.disconnected', ); const connects = events.filter( (e) => e.type === 'tool.list.updated' && e.reason === 'mcp.connected', ); - expect(disconnects.length).toBeGreaterThanOrEqual(1); + expect(disconnects).toHaveLength(0); expect(connects.length).toBeGreaterThanOrEqual(1); } finally { await session.close(); diff --git a/packages/agent-core/test/mcp/tool-manager-mcp.test.ts b/packages/agent-core/test/mcp/tool-manager-mcp.test.ts index 808d62926e..f3b0b5c87d 100644 --- a/packages/agent-core/test/mcp/tool-manager-mcp.test.ts +++ b/packages/agent-core/test/mcp/tool-manager-mcp.test.ts @@ -1033,6 +1033,89 @@ describe('ToolManager MCP reconnect-on-call', () => { await server.close(); } }, 20000); + + it('keeps tools registered when a call-driven reconnect fails, so a later call re-drives it (#2742)', async () => { + let server = await startHttpEchoServer(); + const url = `http://127.0.0.1:${server.port}/mcp`; + + const manager = new McpConnectionManager(); + const agent = fakeAgentWithMcp(manager); + const emittedEvents: Array<{ type?: string; reason?: string }> = []; + agent.emitEvent = (event) => { + emittedEvents.push(event as { type?: string; reason?: string }); + }; + const tm = new ToolManager(agent); + tm.setActiveTools(['mcp__*']); + try { + await manager.connectAll({ srv: { transport: 'http', url } }); + await manager.waitForInitialLoad(); + expect(manager.get('srv')?.status).toBe('connected'); + + const echo = tm.loopTools.find((t) => t.name === 'mcp__srv__echo'); + expect(echo).toBeDefined(); + const context = toolCallContext(); + + const first = await executeTool(echo!, { ...context, args: { text: 'before' } }); + expect(first.isError).toBe(false); + expect(first.output).toBe('before'); + + // The server goes DOWN and stays down: the call drives a reconnect that + // fails, taking the entry through pending (where the tools used to be + // unregistered) into failed. + await server.close(); + const eventsBeforeDownCall = emittedEvents.length; + + const rejection = await executeTool(echo!, { ...context, args: { text: 'while down' } }).then( + () => undefined, + (error: unknown) => error, + ); + // The original transport error surfaces: a failed reconnect attempt is + // recorded as the entry's `failed` status, not thrown through. + expect(rejection).toBeInstanceOf(Error); + // The transport's own error text reaches the caller (failure + // classification stays byte-equivalent) — not a rewritten wrapper. + expect((rejection as Error).message).toMatch(/fetch failed|ECONNREFUSED|socket hang up/); + expect((rejection as Error).message).not.toContain( + 'reconnecting the MCP server also failed', + ); + expect(manager.get('srv')?.status).toBe('failed'); + // Regression: the failed recovery must not strand the session without + // the tools — they stay registered so the NEXT call can drive another + // reconnect attempt. + expect(tm.loopTools.map((t) => t.name)).toContain('mcp__srv__echo'); + // The other half of the removed `pending` branch: no tool-list update + // (in particular no `mcp.disconnected`) is emitted while the server is + // down — the tool list genuinely did not change. + expect( + emittedEvents.slice(eventsBeforeDownCall).filter((e) => e.type === 'tool.list.updated'), + ).toHaveLength(0); + + // The server comes back; the next call re-drives the reconnect and heals. + server = await startHttpEchoServer(server.port); + + const healed = await executeTool(echo!, { ...context, args: { text: 'back again' } }); + expect(healed.isError).toBe(false); + expect(healed.output).toBe('back again'); + expect(manager.get('srv')?.status).toBe('connected'); + expect(tm.loopTools.map((t) => t.name)).toContain('mcp__srv__echo'); + + // The re-listed handle (what the agent loop reads on the next prompt) + // works too — not just the stale one captured before the outage. + const relisted = tm.loopTools.find((t) => t.name === 'mcp__srv__echo'); + expect(relisted).toBeDefined(); + const relistedResult = await executeTool(relisted!, { + ...context, + args: { text: 'relisted' }, + }); + expect(relistedResult.isError).toBe(false); + expect(relistedResult.output).toBe('relisted'); + } finally { + await manager.shutdown(); + // Whichever server generation is current; earlier ones were already + // closed (idempotently) above. + await server.close(); + } + }, 20000); }); async function startHttpEchoServer( @@ -1053,10 +1136,14 @@ async function startHttpEchoServer( }); await new Promise((resolve) => httpServer.listen(port, '127.0.0.1', resolve)); const boundPort = (httpServer.address() as AddressInfo).port; + let closePromise: Promise | undefined; return { port: boundPort, - close: () => - new Promise((resolve, reject) => { + close: () => { + // Idempotent: cleanup paths must not mask an earlier test failure with + // ERR_SERVER_NOT_RUNNING. The promise is cached so a concurrent second + // call waits for the same shutdown instead of resolving early. + closePromise ??= new Promise((resolve, reject) => { httpServer.close((err) => { if (err) { reject(err); @@ -1067,6 +1154,8 @@ async function startHttpEchoServer( // Force-close the client's keep-alive / SSE sockets so `close()` // returns even while a long-lived streamable-HTTP stream is open. httpServer.closeAllConnections(); - }), + }); + return closePromise; + }, }; }