Skip to content

Commit ec12e30

Browse files
vincentkocvincentkoc
committed
fix(agents): bound MCP catalog fanout
1 parent 44f9e17 commit ec12e30

4 files changed

Lines changed: 272 additions & 222 deletions

File tree

‎scripts/repro-94162-timing.mjs‎

Lines changed: 0 additions & 61 deletions
This file was deleted.

‎src/agents/agent-bundle-mcp-runtime.test.ts‎

Lines changed: 97 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import http from "node:http";
44
import os from "node:os";
55
import path from "node:path";
66
import { afterEach, describe, expect, it, vi } from "vitest";
7+
import { cleanupTempDirs, makeTempDir } from "../../test/helpers/temp-dir.js";
78
import { createBundleMcpJsonSchemaValidator } from "./agent-bundle-mcp-runtime.js";
89
import { cleanupBundleMcpHarness } from "./agent-bundle-mcp-test-harness.js";
910
import {
@@ -26,6 +27,8 @@ vi.mock("./embedded-agent-mcp.js", () => ({
2627
}),
2728
}));
2829

30+
const tempDirs: string[] = [];
31+
2932
type RuntimeFactoryOptions = NonNullable<
3033
Parameters<typeof testing.createSessionMcpRuntimeManager>[0]
3134
>;
@@ -37,10 +40,12 @@ async function writeListToolsMcpServer(params: {
3740
filePath: string;
3841
logPath: string;
3942
delayMs?: number;
43+
initializeDelayMs?: number;
4044
hang?: boolean;
4145
inputSchema?: unknown;
4246
tools?: Array<{ name: string; description?: string; inputSchema?: unknown }>;
4347
capabilities?: Record<string, unknown>;
48+
notifyListChangedOnInitialized?: boolean;
4449
listToolsMethodNotFound?: boolean;
4550
callToolIsError?: boolean;
4651
callToolJsonRpcError?: boolean;
@@ -53,8 +58,10 @@ import fs from "node:fs/promises";
5358
5459
const logPath = ${JSON.stringify(params.logPath)};
5560
const delayMs = ${params.delayMs ?? 0};
61+
const initializeDelayMs = ${params.initializeDelayMs ?? 0};
5662
const hang = ${params.hang === true};
5763
const capabilities = ${JSON.stringify(params.capabilities ?? { tools: {} })};
64+
const notifyListChangedOnInitialized = ${params.notifyListChangedOnInitialized === true};
5865
const listToolsMethodNotFound = ${params.listToolsMethodNotFound === true};
5966
const tools = ${JSON.stringify(
6067
params.tools ?? [
@@ -84,18 +91,27 @@ function handle(message) {
8491
}
8592
log("recv " + String(message.method ?? "unknown"));
8693
if (message.method === "initialize") {
87-
send({
94+
const response = {
8895
jsonrpc: "2.0",
8996
id: message.id,
9097
result: {
9198
protocolVersion: message.params?.protocolVersion ?? "2025-03-26",
9299
capabilities,
93100
serverInfo: { name: "test-list-tools", version: "1.0.0" },
94101
},
95-
});
102+
};
103+
if (initializeDelayMs > 0) {
104+
setTimeout(() => send(response), initializeDelayMs);
105+
} else {
106+
send(response);
107+
}
96108
return;
97109
}
98110
if (message.method === "notifications/initialized") {
111+
if (notifyListChangedOnInitialized) {
112+
log("notify tools/list_changed");
113+
send({ jsonrpc: "2.0", method: "notifications/tools/list_changed" });
114+
}
99115
return;
100116
}
101117
if (message.method === "tools/list") {
@@ -281,6 +297,7 @@ function makeRuntime(
281297
}
282298

283299
afterEach(async () => {
300+
cleanupTempDirs(tempDirs);
284301
await cleanupBundleMcpHarness();
285302
});
286303

@@ -2036,16 +2053,11 @@ process.stdin.on("end", () => {
20362053
},
20372054
);
20382055

2039-
// FIX #94162: regression test — parallel MCP catalog loading must complete
2040-
// in approx max(server delays) not sum(server delays).
2041-
// Under the original sequential for-await loop, this test would take
2042-
// ~1200ms (200+400+600) and fail the <900ms assertion.
2043-
// Under the parallel Promise.allSettled fix, it completes in ~600ms.
20442056
it(
20452057
"parallelizes MCP server catalog loading across multiple slow servers",
20462058
{ timeout: LIST_TOOLS_TEST_DEADLINE_MS },
20472059
async () => {
2048-
const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "bundle-mcp-parallel-"));
2060+
const tempDir = makeTempDir(tempDirs, "bundle-mcp-parallel-");
20492061
const delays = [200, 400, 600];
20502062
const serverPaths = delays.map((delay, i) => {
20512063
const serverPath = path.join(tempDir, `slow-server-${i}.mjs`);
@@ -2084,7 +2096,7 @@ process.stdin.on("end", () => {
20842096
try {
20852097
const sumDelays = delays.reduce((a, b) => a + b, 0);
20862098
const maxDelay = Math.max(...delays);
2087-
const tolerance = 350; // ms of overhead for process spawn, JSON serialization, etc.
2099+
const parallelBudgetMs = maxDelay + 500;
20882100

20892101
const t0 = performance.now();
20902102
const catalog = await runtime.getCatalog();
@@ -2098,15 +2110,86 @@ process.stdin.on("end", () => {
20982110
"slow_tool",
20992111
]);
21002112

2101-
// Wall time must be LESS than sum of all delays (proves parallelism).
2102-
// Under the original sequential loop this would exceed sumDelays + tolerance.
2103-
expect(wallTime).toBeLessThan(sumDelays + tolerance);
2113+
// Sequential listing would have to wait roughly sumDelays before overhead;
2114+
// parallel listing should stay near the slowest server plus launch overhead.
2115+
expect(wallTime).toBeLessThan(parallelBudgetMs);
2116+
expect(parallelBudgetMs).toBeLessThan(sumDelays);
21042117

2105-
// Wall time must be at least the max delay (proves we wait for the slowest).
21062118
expect(wallTime).toBeGreaterThanOrEqual(maxDelay * 0.7);
21072119
} finally {
21082120
await runtime.dispose();
2109-
await fs.rm(tempDir, { recursive: true, force: true });
2121+
}
2122+
},
2123+
);
2124+
2125+
it(
2126+
"awaits in-progress MCP session connections after catalog invalidation",
2127+
{ timeout: LIST_TOOLS_TEST_DEADLINE_MS },
2128+
async () => {
2129+
const tempDir = makeTempDir(tempDirs, "bundle-mcp-inflight-connect-");
2130+
const invalidatingServer = {
2131+
serverName: "invalidatingServer",
2132+
serverPath: path.join(tempDir, "invalidating-server.mjs"),
2133+
logPath: path.join(tempDir, "invalidating-server.log"),
2134+
};
2135+
const slowConnectServer = {
2136+
serverName: "slowConnectServer",
2137+
serverPath: path.join(tempDir, "slow-connect-server.mjs"),
2138+
logPath: path.join(tempDir, "slow-connect-server.log"),
2139+
};
2140+
2141+
await writeListToolsMcpServer({
2142+
filePath: invalidatingServer.serverPath,
2143+
logPath: invalidatingServer.logPath,
2144+
capabilities: { tools: { listChanged: true } },
2145+
notifyListChangedOnInitialized: true,
2146+
});
2147+
await writeListToolsMcpServer({
2148+
filePath: slowConnectServer.serverPath,
2149+
logPath: slowConnectServer.logPath,
2150+
initializeDelayMs: 500,
2151+
});
2152+
2153+
testing.setBundleMcpCatalogListTimeoutMsForTest(4_000);
2154+
2155+
const runtime = await getOrCreateSessionMcpRuntime({
2156+
sessionId: "session-inflight-connect-test",
2157+
sessionKey: "agent:test:session-inflight-connect-test",
2158+
workspaceDir: "/workspace",
2159+
cfg: {
2160+
mcp: {
2161+
servers: Object.fromEntries(
2162+
[invalidatingServer, slowConnectServer].map(({ serverName, serverPath }) => [
2163+
serverName,
2164+
{
2165+
command: process.execPath,
2166+
args: [serverPath],
2167+
connectionTimeoutMs: 2_000,
2168+
},
2169+
]),
2170+
),
2171+
},
2172+
},
2173+
});
2174+
2175+
try {
2176+
const firstCatalog = runtime.getCatalog();
2177+
await waitForFileText(
2178+
invalidatingServer.logPath,
2179+
"notify tools/list_changed",
2180+
LIST_TOOLS_SERVER_LOG_TIMEOUT_MS,
2181+
);
2182+
2183+
const secondCatalog = await runtime.getCatalog();
2184+
await firstCatalog;
2185+
2186+
expect(Object.keys(secondCatalog.servers).toSorted()).toEqual([
2187+
invalidatingServer.serverName,
2188+
slowConnectServer.serverName,
2189+
]);
2190+
expect(secondCatalog.diagnostics ?? []).toEqual([]);
2191+
} finally {
2192+
await runtime.dispose();
21102193
}
21112194
},
21122195
);

0 commit comments

Comments
 (0)