Files
oh-my-pi/packages/coding-agent/test/mcp-reconnect.test.ts
T

389 lines
13 KiB
TypeScript

import { describe, expect, it, vi } from "bun:test";
import type { MCPReconnect } from "@oh-my-pi/pi-coding-agent/mcp/tool-bridge";
import {
DeferredMCPTool,
deduplicateMCPToolsByName,
isRetriableConnectionError,
MCPTool,
} from "@oh-my-pi/pi-coding-agent/mcp/tool-bridge";
import type { MCPServerConnection, MCPToolCallResult, MCPTransport } from "@oh-my-pi/pi-coding-agent/mcp/types";
import { ToolAbortError } from "@oh-my-pi/pi-coding-agent/tools/tool-errors";
import { logger } from "@oh-my-pi/pi-utils";
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
/** Create a minimal mock transport where `request` is controlled by the caller. */
function mockTransport(requestFn: (...args: Parameters<MCPTransport["request"]>) => Promise<unknown>): MCPTransport {
return {
connected: true,
request: requestFn as MCPTransport["request"],
async notify() {},
async close() {},
};
}
const TOOL_DEF = { name: "do_stuff", inputSchema: { type: "object" as const } };
function toolCallResult(text: string, isError = false): MCPToolCallResult {
return { content: [{ type: "text", text }], isError };
}
function makeConnection(transport: MCPTransport, name = "test-server"): MCPServerConnection {
return {
name,
config: { type: "stdio" as const, command: "echo" },
transport,
serverInfo: { name: "test", version: "1.0" },
capabilities: { tools: {} },
};
}
// ---------------------------------------------------------------------------
// deduplicateMCPToolsByName
// ---------------------------------------------------------------------------
describe("deduplicateMCPToolsByName", () => {
it("keeps the same collision winner after a reconnect reorders the tool list", () => {
const warn = vi.spyOn(logger, "warn").mockImplementation(() => {});
try {
const dotted = { name: "mcp__foo_bar_lookup", mcpServerName: "foo.bar", mcpToolName: "lookup" };
const underscored = { name: "mcp__foo_bar_lookup", mcpServerName: "foo_bar", mcpToolName: "lookup" };
const before = deduplicateMCPToolsByName([dotted, underscored]);
expect(before).toEqual([dotted]);
// Simulate MCPManager#replaceServerTools on the current winner: its
// tools are removed and re-appended, reordering it behind the loser.
// The minted name must not silently switch owners.
const after = deduplicateMCPToolsByName([underscored, dotted]);
expect(after).toEqual([dotted]);
} finally {
warn.mockRestore();
}
});
});
// ---------------------------------------------------------------------------
// isRetriableConnectionError
// ---------------------------------------------------------------------------
describe("isRetriableConnectionError", () => {
const retriable = [
"ECONNREFUSED",
"ECONNRESET",
"EPIPE",
"ENETUNREACH",
"EHOSTUNREACH",
"fetch failed",
"Transport not connected",
"network error",
"HTTP 404: Not Found",
"HTTP 502: Bad Gateway",
"HTTP 503: Service Unavailable",
"Transport closed",
];
for (const msg of retriable) {
it(`matches: ${msg}`, () => {
expect(isRetriableConnectionError(new Error(msg))).toBe(true);
});
}
const nonRetriable = [
"MCP error -32603: Server still initializing",
"HTTP 401: Unauthorized",
"HTTP 403: Forbidden",
"HTTP 400: Bad Request",
"Request timeout after 30000ms",
"SSE response timeout after 30000ms",
"Tool not found: do_stuff",
];
for (const msg of nonRetriable) {
it(`does not match: ${msg}`, () => {
expect(isRetriableConnectionError(new Error(msg))).toBe(false);
});
}
it("returns false for non-Error values", () => {
expect(isRetriableConnectionError("ECONNREFUSED")).toBe(false);
expect(isRetriableConnectionError(null)).toBe(false);
expect(isRetriableConnectionError(undefined)).toBe(false);
expect(isRetriableConnectionError({ message: "ECONNREFUSED" })).toBe(false);
});
});
// ---------------------------------------------------------------------------
// MCPTool.execute retry behavior
// ---------------------------------------------------------------------------
describe("MCPTool.execute retry on connection error", () => {
const noop = () => {};
const noCtx = {} as Parameters<MCPTool["execute"]>[3];
it("retries once on retriable error when reconnect succeeds", async () => {
let callCount = 0;
const failTransport = mockTransport(async () => {
callCount++;
throw new Error("ECONNREFUSED");
});
const successTransport = mockTransport(async () => {
callCount++;
return toolCallResult("ok");
});
const oldConn = makeConnection(failTransport);
const newConn = makeConnection(successTransport, "test-server-new");
const reconnect: MCPReconnect = async () => newConn;
const tool = new MCPTool(oldConn, TOOL_DEF, reconnect);
const result = await tool.execute("call-1", {}, noop, noCtx);
expect(callCount).toBe(2); // 1 fail + 1 retry
expect(result.details?.isError).toBeFalsy();
expect(result.content[0]).toEqual({ type: "text", text: "ok" });
});
it("retries on transport closed and rebinding succeeds", async () => {
let oldCalls = 0;
let newCalls = 0;
let reconnects = 0;
const closedTransport = mockTransport(async () => {
oldCalls++;
throw new Error("Transport closed");
});
const reopenedTransport = mockTransport(async () => {
newCalls++;
return toolCallResult("ok");
});
const oldConn = makeConnection(closedTransport);
const newConn = makeConnection(reopenedTransport, "test-server-transport-closed");
const reconnect: MCPReconnect = async () => {
reconnects++;
return newConn;
};
const tool = new MCPTool(oldConn, TOOL_DEF, reconnect);
const result = await tool.execute("call-1", {}, noop, noCtx);
expect(reconnects).toBe(1);
expect(oldCalls).toBe(1);
expect(newCalls).toBe(1);
expect(result.details?.isError).toBeFalsy();
expect(result.content[0]).toEqual({ type: "text", text: "ok" });
});
it("reuses refreshed connection on later call", async () => {
let oldCalls = 0;
let newCalls = 0;
let reconnects = 0;
const oldTransport = mockTransport(async () => {
oldCalls++;
throw new Error("ECONNREFUSED");
});
const newTransport = mockTransport(async () => {
newCalls++;
return toolCallResult("ok");
});
const oldConn = makeConnection(oldTransport);
const newConn = makeConnection(newTransport, "test-server-rebound");
const reconnect: MCPReconnect = async () => {
reconnects++;
return newConn;
};
const tool = new MCPTool(oldConn, TOOL_DEF, reconnect);
const first = await tool.execute("call-1", {}, noop, noCtx);
const second = await tool.execute("call-2", {}, noop, noCtx);
expect(oldCalls).toBe(1);
expect(newCalls).toBe(2);
expect(reconnects).toBe(1);
expect(first.details?.isError).toBeFalsy();
expect(second.details?.isError).toBeFalsy();
expect(first.content[0]).toEqual({ type: "text", text: "ok" });
expect(second.content[0]).toEqual({ type: "text", text: "ok" });
});
it("returns error result when reconnect returns null", async () => {
const failTransport = mockTransport(async () => {
throw new Error("ECONNRESET");
});
const reconnect: MCPReconnect = async () => null;
const tool = new MCPTool(makeConnection(failTransport), TOOL_DEF, reconnect);
const result = await tool.execute("call-1", {}, noop, noCtx);
expect(result.details?.isError).toBe(true);
expect(result.content[0]).toEqual({ type: "text", text: "MCP error: ECONNRESET" });
});
it("does not retry on non-retriable error", async () => {
let reconnectCalled = false;
const failTransport = mockTransport(async () => {
throw new Error("MCP error -32603: Internal error");
});
const reconnect: MCPReconnect = async () => {
reconnectCalled = true;
return null;
};
const tool = new MCPTool(makeConnection(failTransport), TOOL_DEF, reconnect);
const result = await tool.execute("call-1", {}, noop, noCtx);
expect(reconnectCalled).toBe(false);
expect(result.details?.isError).toBe(true);
});
it("does not retry when no reconnect callback", async () => {
const failTransport = mockTransport(async () => {
throw new Error("ECONNREFUSED");
});
const tool = new MCPTool(makeConnection(failTransport), TOOL_DEF); // no reconnect
const result = await tool.execute("call-1", {}, noop, noCtx);
expect(result.details?.isError).toBe(true);
expect(result.content[0]).toEqual({ type: "text", text: "MCP error: ECONNREFUSED" });
});
it("returns error from retry when retry also fails", async () => {
const failTransport = mockTransport(async () => {
throw new Error("ECONNREFUSED");
});
const retryFailTransport = mockTransport(async () => {
throw new Error("HTTP 503: Service Unavailable");
});
const reconnect: MCPReconnect = async () => makeConnection(retryFailTransport);
const tool = new MCPTool(makeConnection(failTransport), TOOL_DEF, reconnect);
const result = await tool.execute("call-1", {}, noop, noCtx);
expect(result.details?.isError).toBe(true);
expect(result.content[0]).toEqual({ type: "text", text: "MCP error: HTTP 503: Service Unavailable" });
});
it("preserves provider info from new connection on successful retry", async () => {
const failTransport = mockTransport(async () => {
throw new Error("fetch failed");
});
const successTransport = mockTransport(async () => toolCallResult("ok"));
const oldConn = makeConnection(failTransport);
oldConn._source = { provider: "old-provider", providerName: "Old", path: "/old", level: "user" };
const newConn = makeConnection(successTransport);
newConn._source = { provider: "new-provider", providerName: "New", path: "/new", level: "user" };
const tool = new MCPTool(oldConn, TOOL_DEF, async () => newConn);
const result = await tool.execute("call-1", {}, noop, noCtx);
expect(result.details?.provider).toBe("new-provider");
expect(result.details?.providerName).toBe("New");
});
it("falls back to original provider when new connection has no source", async () => {
const failTransport = mockTransport(async () => {
throw new Error("fetch failed");
});
const successTransport = mockTransport(async () => toolCallResult("ok"));
const oldConn = makeConnection(failTransport);
oldConn._source = { provider: "orig", providerName: "Original", path: "/orig", level: "user" };
const newConn = makeConnection(successTransport);
// newConn has no _source
const tool = new MCPTool(oldConn, TOOL_DEF, async () => newConn);
const result = await tool.execute("call-1", {}, noop, noCtx);
expect(result.details?.provider).toBe("orig");
expect(result.details?.providerName).toBe("Original");
});
it("reconnects once when a tool result carries an OAuth challenge", async () => {
let oldCalls = 0;
let newCalls = 0;
let challenge: unknown;
const oldTransport = mockTransport(async () => {
oldCalls++;
return {
...toolCallResult("authorize me", true),
_meta: { "mcp/www_authenticate": ['Bearer resource_metadata="https://mcp.example/meta"'] },
};
});
const newTransport = mockTransport(async () => {
newCalls++;
return toolCallResult("authorized");
});
const newConn = makeConnection(newTransport, "test-server-authorized");
const reconnect: MCPReconnect = async options => {
challenge = options?.authChallenge;
return newConn;
};
const tool = new MCPTool(makeConnection(oldTransport), TOOL_DEF, reconnect);
const result = await tool.execute("call-1", {}, noop, noCtx);
expect(oldCalls).toBe(1);
expect(newCalls).toBe(1);
expect(challenge).toEqual({
wwwAuthenticate: ['Bearer resource_metadata="https://mcp.example/meta"'],
});
expect(result.details?.isError).toBeFalsy();
expect(result.content[0]).toEqual({ type: "text", text: "authorized" });
});
it("preserves the OAuth challenge metadata when no reconnect handler exists", async () => {
const result = await new MCPTool(
makeConnection(
mockTransport(async () => ({
...toolCallResult("authorize me", true),
_meta: { "mcp/www_authenticate": ["Bearer"] },
})),
),
TOOL_DEF,
).execute("call-1", {}, noop, noCtx);
expect(result.details?.isError).toBe(true);
expect(result.details?.mcpMeta).toEqual({ "mcp/www_authenticate": ["Bearer"] });
});
});
describe("reconnect abort propagation", () => {
const noop = () => {};
const noCtx = {} as Parameters<MCPTool["execute"]>[3];
const noDeferredCtx = {} as Parameters<DeferredMCPTool["execute"]>[3];
it("throws ToolAbortError when MCPTool reconnect is aborted", async () => {
const failTransport = mockTransport(async () => {
throw new Error("ECONNRESET");
});
const { promise } = Promise.withResolvers<MCPServerConnection | null>();
const reconnect: MCPReconnect = async () => promise;
const tool = new MCPTool(makeConnection(failTransport), TOOL_DEF, reconnect);
const controller = new AbortController();
const pending = tool.execute("call-1", {}, noop, noCtx, controller.signal);
controller.abort();
await expect(pending).rejects.toBeInstanceOf(ToolAbortError);
});
it("throws ToolAbortError when DeferredMCPTool reconnect is aborted", async () => {
const getConnection = async () => {
throw new Error("MCP server not connected");
};
const { promise } = Promise.withResolvers<MCPServerConnection | null>();
const reconnect: MCPReconnect = async () => promise;
const tool = new DeferredMCPTool("test-server", TOOL_DEF, getConnection, undefined, reconnect);
const controller = new AbortController();
const pending = tool.execute("call-1", {}, noop, noDeferredCtx, controller.signal);
controller.abort();
await expect(pending).rejects.toBeInstanceOf(ToolAbortError);
});
});