diff --git a/src/agent/claude-code/live-processes.ts b/src/agent/claude-code/live-processes.ts index 1a68672..ebb729b 100644 --- a/src/agent/claude-code/live-processes.ts +++ b/src/agent/claude-code/live-processes.ts @@ -1,4 +1,4 @@ -export type LiveProcess = { readonly owner: symbol; toolCallIds: string[]; kill(): void }; +export type LiveProcess = { readonly owner: symbol; readonly sessionId: string | undefined; toolCallIds: string[]; kill(): void }; // Tools usually call back within seconds, but a question tool holds its call until the learner // answers. Past this, the process is killed and the answer starts a fresh one with a full replay. @@ -27,11 +27,11 @@ function hasSameIds(ids: string[], otherIds: string[]): boolean { return ids.length === otherIds.length && new Set(otherIds).size === otherIds.length && ids.every((id) => otherIds.includes(id)); } -export function getLiveProcess(toolCallIds: string[], owner: symbol): LiveProcess | null { - if (toolCallIds.length === 0) return null; +export function getLiveProcess(toolCallIds: string[], owner: symbol, sessionId: string | undefined): LiveProcess | null { + if (!sessionId || toolCallIds.length === 0) return null; for (const process of liveProcesses.keys()) { - if (process.owner === owner && hasSameIds(process.toolCallIds, toolCallIds)) return process; + if (process.owner === owner && process.sessionId === sessionId && hasSameIds(process.toolCallIds, toolCallIds)) return process; } return null; diff --git a/src/agent/claude-code/stream.ts b/src/agent/claude-code/stream.ts index cd27167..8c9c89a 100644 --- a/src/agent/claude-code/stream.ts +++ b/src/agent/claude-code/stream.ts @@ -392,6 +392,7 @@ class ClaudeProcess implements LiveProcess { private sourceCalls: ToolCall[] = []; private sourceComplete = false; private admitted: PendingToolCall[] = []; + private readonly cancelledIds = new Set(); private pendingError: string | undefined; constructor( @@ -400,6 +401,7 @@ class ClaudeProcess implements LiveProcess { private readonly files: TurnFiles, private turn: ClaudeTurn, readonly owner: symbol, + readonly sessionId: string | undefined, private readonly configuration: unknown, private conversation: unknown[], ) {} @@ -422,9 +424,9 @@ class ClaudeProcess implements LiveProcess { this.setSignal(signal); } - canContinue(owner: symbol, configuration: unknown, messages: Message[], results: ToolResultMessage[]): boolean { + canContinue(owner: symbol, sessionId: string | undefined, configuration: unknown, messages: Message[], results: ToolResultMessage[]): boolean { const conversation = conversationForReuse(messages); - return !this.stopping && owner === this.owner + return !this.stopping && owner === this.owner && !!sessionId && sessionId === this.sessionId && isDeepStrictEqual(configuration, this.configuration) && isDeepStrictEqual(conversation.slice(0, this.conversation.length), this.conversation) && conversation.length === this.conversation.length + this.admitted.length @@ -462,7 +464,7 @@ class ClaudeProcess implements LiveProcess { // A parallel request can fail after the admitted slice was handed to // Pi. Retain that failure for its continuation instead of silently // replaying and hiding the compatibility error. - if (this.turn.finished && this.toolCallIds.length > 0) this.pendingError = message; + if (this.sessionId && this.turn.finished && this.toolCallIds.length > 0) this.pendingError = message; this.turn.fail(message, "error"); this.kill(); if (this.pendingError) addLiveProcess(this); @@ -477,6 +479,13 @@ class ClaudeProcess implements LiveProcess { return true; } + cancelCall(call: PendingToolCall): void { + this.cancelledIds.add(call.id); + // Already admitted work belongs to Pi and cannot be undone. A queued + // request, however, must never become executable after cancellation. + this.sourceCalls = this.sourceCalls.filter((intent) => intent.id !== call.id); + } + /** Only authenticated MCP requests, never stdout intent alone, admit side effects. */ flushAdmissions(): void { if (this.stopping || this.turn.finished || !this.sourceComplete) return; @@ -496,7 +505,8 @@ class ClaudeProcess implements LiveProcess { this.toolCallIds = pending.map((call) => call.id); this.turn.admit(calls); this.conversation.push(...structuredClone(conversationForReuse([this.turn.getMessage()]))); - addLiveProcess(this); + if (this.sessionId) addLiveProcess(this); + else this.kill(); // No trusted scope: return the batch, then use fresh replay. } private readonly abort = () => { @@ -534,11 +544,12 @@ class ClaudeProcess implements LiveProcess { this.turn.handle(event); if (this.turn.finished) { this.kill(); return; } if (event.type === "messageStop" && this.turn.waitingForAdmission) { - this.sourceCalls = this.turn.getIntents(); - if (this.sourceCalls.length === 0 || new Set(this.sourceCalls.map((call) => call.id)).size !== this.sourceCalls.length) { + const intents = this.turn.getIntents(); + if (intents.length === 0 || new Set(intents.map((call) => call.id)).size !== intents.length) { this.fail("claude CLI emitted an invalid tool batch"); return; } + this.sourceCalls = intents.filter((call) => !this.cancelledIds.has(call.id)); this.sourceComplete = true; this.flushAdmissions(); } @@ -559,17 +570,17 @@ class ClaudeProcess implements LiveProcess { } } -function continueLiveProcess(turn: ClaudeTurn, owner: symbol, configuration: unknown, messages: Message[], signal: AbortSignal | undefined): boolean { +function continueLiveProcess(turn: ClaudeTurn, owner: symbol, sessionId: string | undefined, configuration: unknown, messages: Message[], signal: AbortSignal | undefined): boolean { const assistantIndex = messages.findLastIndex((message) => message.role === "assistant"); const assistant = messages[assistantIndex]; if (assistant?.role !== "assistant") return false; const ids = assistant.content.flatMap((block) => block.type === "toolCall" ? [block.id] : []); - const process = getLiveProcess(ids, owner); + const process = getLiveProcess(ids, owner, sessionId); if (!(process instanceof ClaudeProcess)) return false; if (process.consumeFailure(owner, turn)) return true; const trailing = messages.slice(assistantIndex + 1); const results = trailing.filter((message) => message.role === "toolResult"); - if (trailing.length !== results.length || !process.canContinue(owner, configuration, messages, results)) { + if (trailing.length !== results.length || !process.canContinue(owner, sessionId, configuration, messages, results)) { process.kill(); return false; } @@ -604,6 +615,7 @@ async function startProcess( const bridge = await createToolBridge(input.tools, { onCall: () => cliProcess?.flushAdmissions(), onError: (error) => cliProcess?.fail(error.message), + onCancel: (call) => cliProcess?.cancelCall(call), }); let files: TurnFiles; try { @@ -639,7 +651,7 @@ async function startProcess( return; } - cliProcess = new ClaudeProcess(child, bridge, files, turn, owner, configuration, structuredClone(conversationForReuse(context.messages))); + cliProcess = new ClaudeProcess(child, bridge, files, turn, owner, options?.sessionId, configuration, structuredClone(conversationForReuse(context.messages))); cliProcess.start(toStreamJsonLine(input.userContent), options?.signal); } @@ -650,7 +662,7 @@ export function createClaudeCodeStreamFunction(runClaude: RunClaude): StreamFunc const turn = new ClaudeTurn(stream, model); const configuration = configurationForReuse(model, context, toEffort(options?.reasoning)); - if (continueLiveProcess(turn, owner, configuration, context.messages, options?.signal)) return stream; + if (continueLiveProcess(turn, owner, options?.sessionId, configuration, context.messages, options?.signal)) return stream; startProcess(runClaude, turn, model, context, options, owner, configuration).catch((error: unknown) => { console.error("claude CLI turn failed:", error); diff --git a/src/agent/claude-code/tool-bridge.ts b/src/agent/claude-code/tool-bridge.ts index db9f4da..71dc6de 100644 --- a/src/agent/claude-code/tool-bridge.ts +++ b/src/agent/claude-code/tool-bridge.ts @@ -22,7 +22,7 @@ const TOKEN_BYTES = 32; export type PendingToolCall = { id: string; name: string; arguments: JsonObject }; -type BridgeCallbacks = { onCall(): void; onError(error: Error): void }; +type BridgeCallbacks = { onCall(): void; onError(error: Error): void; onCancel?(call: PendingToolCall): void }; export const MCP_IDENTITY_ERROR = "claude CLI compatibility error: MCP calls must include a unique claudecode/toolUseId"; @@ -67,8 +67,22 @@ class ToolCallQueue { private readonly results = new Map void>(); private waiters: ToolCallWaiter[] = []; - add(call: PendingToolCall): Promise { - const result = new Promise((resolve) => this.results.set(call, resolve)); + add(call: PendingToolCall, signal: AbortSignal, onCancel?: (call: PendingToolCall) => void): Promise { + const result = new Promise((resolve) => { + const cancel = () => { + const settle = this.results.get(call); + if (!settle) return; + settle({ content: [{ type: "text", text: "MCP request cancelled" }], isError: true }); + onCancel?.(call); + }; + this.results.set(call, (result) => { + this.results.delete(call); + signal.removeEventListener("abort", cancel); + resolve(result); + }); + signal.addEventListener("abort", cancel, { once: true }); + if (signal.aborted) cancel(); + }); this.wakeWaiters(); @@ -84,7 +98,6 @@ class ToolCallQueue { resolve(call: PendingToolCall, result: ToolResultMessage): void { this.results.get(call)?.(toCallToolResult(result)); - this.results.delete(call); } list(): PendingToolCall[] { @@ -92,7 +105,9 @@ class ToolCallQueue { } clear(): void { - this.results.clear(); + for (const settle of this.results.values()) { + settle({ content: [{ type: "text", text: "MCP bridge closed" }], isError: true }); + } for (const waiter of this.waiters) waiter.reject(new Error("claude-code tool bridge closed")); this.waiters = []; @@ -116,7 +131,7 @@ function createMcpServer(tools: Tool[], queue: ToolCallQueue, callbacks?: Bridge server.setRequestHandler(ListToolsRequestSchema, () => ({ tools: tools.map(toMcpTool) })); const seenIds = new Set(); - server.setRequestHandler(CallToolRequestSchema, (request) => { + server.setRequestHandler(CallToolRequestSchema, (request, extra) => { const id = request.params._meta?.["claudecode/toolUseId"]; if (typeof id !== "string" || !id.trim() || seenIds.has(id)) { const error = new McpError(ErrorCode.InvalidParams, MCP_IDENTITY_ERROR); @@ -129,7 +144,7 @@ function createMcpServer(tools: Tool[], queue: ToolCallQueue, callbacks?: Bridge throw error; } seenIds.add(id); - const result = queue.add({ id, name: request.params.name, arguments: (request.params.arguments ?? {}) as JsonObject }); + const result = queue.add({ id, name: request.params.name, arguments: (request.params.arguments ?? {}) as JsonObject }, extra.signal, callbacks?.onCancel); callbacks?.onCall(); return result; }); diff --git a/tests/claude-code-harness.test.ts b/tests/claude-code-harness.test.ts index d300b35..aab05bb 100644 --- a/tests/claude-code-harness.test.ts +++ b/tests/claude-code-harness.test.ts @@ -78,9 +78,9 @@ function createFakeRun(t: TestContext, scenarios: string[][]): FakeRun { const child = spawn(process.execPath, [FAKE_CLAUDE, ...args], { cwd: options.cwd, env: { ...options.env, FAKE_CLAUDE_SCENARIO: scenarioFile, FAKE_CLAUDE_STDIN: stdinFile }, - stdio: "pipe", - }); - t.after(() => child.kill()); + stdio: ["pipe", "pipe", "pipe", "ipc"], + }) as ChildProcessWithoutNullStreams; + t.after(() => child.kill("SIGKILL")); children.push(child); stdinFiles.push(stdinFile); @@ -104,7 +104,7 @@ function createAgentFactory(fakeRun: FakeRun): Parameters createTeacherAgent(undefined, { engine, messages: options.messages }); + return (options) => createTeacherAgent(undefined, { engine, chatId: options.chatId, messages: options.messages }); } async function runChatTurn(chatId: string, input: object, agentFactory: Parameters[1]): Promise { @@ -196,6 +196,40 @@ test("a question pauses later calls from the same CLI response until the learner await waitForExit(fakeRun.children[0]!); }); +test("an individually cancelled MCP question queued behind a learner wait never executes", async (t) => { + const cancelledId = "toolu_cancelled_question"; + const question = "This cancelled question must never be posed"; + const scenario = toAskScenario(); + scenario.splice(4, 0, + streamEvent({ type: "content_block_start", index: 1, content_block: { type: "tool_use", id: cancelledId, name: "mcp__learnharness__ask_user_question", input: {} } }), + streamEvent({ type: "content_block_delta", index: 1, delta: { type: "input_json_delta", partial_json: JSON.stringify({ question }) } }), + streamEvent({ type: "content_block_stop", index: 1 }), + ); + scenario.splice(scenario.length - 2, 2, + JSON.stringify({ fake: "start-call", id: ASK_ID, name: "ask_user_question", arguments: ASK_INPUT }), + JSON.stringify({ fake: "wait-for-message" }), + JSON.stringify({ fake: "start-call", id: cancelledId, name: "ask_user_question", arguments: { question } }), + JSON.stringify({ fake: "cancel-call", id: cancelledId }), + JSON.stringify({ fake: "checkpoint" }), + JSON.stringify({ fake: "reply" }), + ); + const fakeRun = createFakeRun(t, [scenario]); + const factory = createAgentFactory(fakeRun); + const chatId = createClaudeCodeChat(t); + await runChatTurn(chatId, { type: "text", text: "teach me sockets" }, factory); + assert.equal(danglingQuestionOf(chatStore.getChat(chatId)?.messages ?? [])?.toolCallId, ASK_ID); + const child = fakeRun.children[0]!; + const cancelledOnServer = once(child, "message", { signal: AbortSignal.timeout(TURN_TIMEOUT_MS) }); + child.send("queue and cancel B"); + assert.deepEqual((await cancelledOnServer)[0], { fake: "checkpoint" }); + await runChatTurn(chatId, toAnswerInput(), factory); + assert.equal(fakeRun.children.length, 1); + assert.equal(danglingQuestionOf(chatStore.getChat(chatId)?.messages ?? []), null); + assert.equal(JSON.stringify(chatStore.getChat(chatId)?.messages).includes(cancelledId), false, "cancelled B must never reach native tool execution or persistence"); + assert.ok(getLastAssistantText(chatId).includes(ANSWER_TEXT)); + await waitForExit(child); +}); + test("a question answered after the idle timeout starts a fresh CLI process that replays the question and answer", async (t) => { liveProcessIdleTimeoutMs.current = 300; t.after(() => (liveProcessIdleTimeoutMs.current = LIVE_PROCESS_IDLE_TIMEOUT_MS)); diff --git a/tests/claude-code-stream.test.ts b/tests/claude-code-stream.test.ts index efdd322..0a61add 100644 --- a/tests/claude-code-stream.test.ts +++ b/tests/claude-code-stream.test.ts @@ -45,6 +45,7 @@ const HOLD = JSON.stringify({ fake: "hold" }); const REPLY = JSON.stringify({ fake: "reply" }); const CLEANUP_TIMEOUT_MS = 2000; +const SESSION_ID = "stream-test-chat"; const READ_TOOL: Tool = { name: "read", description: "Read a file", parameters: { type: "object", properties: {} } }; const HI: Message = { role: "user", content: "hi", timestamp: 0 }; @@ -156,7 +157,7 @@ async function runScenario( messages: [{ role: "user", content: "hi", timestamp: 0 }], tools: [READ_TOOL, { name: "get_weather", description: "Weather", parameters: { type: "object", properties: {} } }], }); - const events = await collectEvents(createClaudeCodeStreamFunction(runClaude)(model, context, options)); + const events = await collectEvents(createClaudeCodeStreamFunction(runClaude)(model, context, { sessionId: SESSION_ID, ...options })); return { events, calls }; } @@ -210,7 +211,7 @@ async function runTurn( ): Promise { const context = normalizeContext({ systemPrompt: "You are a tutor.", messages, tools: [READ_TOOL] }); - const done = getLastEvent(await collectEvents(streamFunction(getModel(), context, options))); + const done = getLastEvent(await collectEvents(streamFunction(getModel(), context, { sessionId: SESSION_ID, ...options }))); assert.ok(done.type === "done", `turn ended with ${done.type}`); @@ -852,7 +853,7 @@ test("a parallel identity failure after an admitted slice is reported on continu const message = await runTurn(stream, [HI]); getOnlyCall(calls).child.send("continue"); await waitForExit(getOnlyCall(calls).child); - const events = await collectEvents(stream(getModel(), normalizeContext({ systemPrompt: "You are a tutor.", messages: [HI, message, toToolResult(first, "A")], tools: [READ_TOOL] }))); + const events = await collectEvents(stream(getModel(), normalizeContext({ systemPrompt: "You are a tutor.", messages: [HI, message, toToolResult(first, "A")], tools: [READ_TOOL] }), { sessionId: SESSION_ID })); assert.match(getErrorMessage(events) ?? "", /compatibility error.*claudecode\/toolUseId/); assert.equal(calls.length, 1, "must not hide the compatibility failure by spawning again"); await assertCleanedUp(getOnlyCall(calls)); @@ -867,7 +868,7 @@ test("output after MCP admission fails continuation instead of silently replayin const message = await runTurn(stream, [HI]); getOnlyCall(calls).child.send("continue"); await waitForExit(getOnlyCall(calls).child); - const events = await collectEvents(stream(getModel(), normalizeContext({ systemPrompt: "You are a tutor.", messages: [HI, message, toToolResult(call, "A")], tools: [READ_TOOL] }))); + const events = await collectEvents(stream(getModel(), normalizeContext({ systemPrompt: "You are a tutor.", messages: [HI, message, toToolResult(call, "A")], tools: [READ_TOOL] }), { sessionId: SESSION_ID })); assert.match(getErrorMessage(events) ?? "", /moved on while.*waiting for results/); assert.equal(calls.length, 1); await assertCleanedUp(getOnlyCall(calls)); @@ -927,13 +928,55 @@ for (const change of ["model", "effort", "system prompt", "tool schema", "tool d }], }); const model = change === "model" ? CLAUDE_CODE_MODELS[1]! : getModel(); - const events = await collectEvents(stream(model, context, change === "effort" ? { reasoning: "high" } : undefined)); + const events = await collectEvents(stream(model, context, { sessionId: SESSION_ID, ...(change === "effort" ? { reasoning: "high" as const } : {}) })); assert.equal(getLastEvent(events).type, "done"); assert.equal(fresh.calls.length, 1, "must start a fresh CLI rather than silently use stale configuration"); await assertCleanedUp(getOnlyCall(waiting.calls)); }); } +for (const identicalHistory of [false, true]) test(`one registered stream isolates two sessions with colliding tool IDs and ${identicalHistory ? "identical" : "different"} histories`, async (t) => { + const call: ToolUse = { id: "toolu_shared_singleton", name: "read", arguments: { path: "same" } }; + const one = createFakeRun(t, [...toToolUseMessage([call]), callTool(call), REPLY]); + const two = createFakeRun(t, [...toToolUseMessage([call]), callTool(call), REPLY]); + let spawns = 0; + const stream = createClaudeCodeStreamFunction((args, options) => { + const run = [one, two][spawns++]; + assert.ok(run, "a valid continuation must not spawn another CLI"); + return run.runClaude(args, options); + }); + const historyOne = [HI]; + const historyTwo: Message[] = identicalHistory ? [HI] : [{ role: "user", content: "another conversation", timestamp: 0 }]; + const messageOne = await runTurn(stream, historyOne, { sessionId: "chat-one" }); + const messageTwo = await runTurn(stream, historyTwo, { sessionId: "chat-two" }); + // B returns first while A is earlier in the registry: IDs and even full + // history must not let B resolve A's authenticated MCP request. + const replyTwo = await runTurn(stream, [...historyTwo, messageTwo, toToolResult(call, "from two")], { sessionId: "chat-two" }); + const replyOne = await runTurn(stream, [...historyOne, messageOne, toToolResult(call, "from one")], { sessionId: "chat-one" }); + assert.equal(getText(replyTwo), "from two"); + assert.equal(getText(replyOne), "from one"); + assert.equal(spawns, 2); + await assertCleanedUp(getOnlyCall(one.calls)); + await assertCleanedUp(getOnlyCall(two.calls)); +}); + +test("a caller without a session ID uses fresh replay and retains no held process", async (t) => { + const call: ToolUse = { id: "toolu_no_scope", name: "read", arguments: { path: "a" } }; + const waiting = createFakeRun(t, [...toToolUseMessage([call]), callTool(call), REPLY]); + const stdinDir = mkdtempSync(join(tmpdir(), "claude-unscoped-replay-")); + t.after(() => rmSync(stdinDir, { recursive: true, force: true })); + const stdinFile = join(stdinDir, "stdin.jsonl"); + const fresh = createFakeRun(t, readFixtureLines("text.jsonl"), { FAKE_CLAUDE_STDIN: stdinFile }); + const stream = createClaudeCodeStreamFunction((args, options) => (waiting.calls.length === 0 ? waiting : fresh).runClaude(args, options)); + const message = await runTurn(stream, [HI], { sessionId: undefined }); + await assertCleanedUp(getOnlyCall(waiting.calls)); + const reply = await runTurn(stream, [HI, message, toToolResult(call, "the preserved answer")], { sessionId: undefined }); + assert.equal(reply.stopReason, "stop"); + assert.equal(fresh.calls.length, 1); + assert.match(readFileSync(stdinFile, "utf8"), /the preserved answer/); + await assertCleanedUp(getOnlyCall(fresh.calls)); +}); + for (const sharedIdentity of [false, true]) test(`two concurrent stream owners keep separate processes${sharedIdentity ? " even with identical call IDs and arguments" : ""}`, async (t) => { const readOne: ToolUse = { id: "toolu_concurrent_one", name: "read", arguments: { path: "one" } }; const readTwo: ToolUse = sharedIdentity ? { ...readOne } : { id: "toolu_concurrent_two", name: "read", arguments: { path: "two" } }; diff --git a/tests/claude-code-tool-bridge.test.ts b/tests/claude-code-tool-bridge.test.ts index 9b44dfb..ff66fd9 100644 --- a/tests/claude-code-tool-bridge.test.ts +++ b/tests/claude-code-tool-bridge.test.ts @@ -136,6 +136,33 @@ test("tools/call stays pending until resolved and returns its text and image con assert.deepEqual(bridge.getPendingToolCalls(), []); }); +test("individual MCP cancellation removes only that request and leaves the other request resolvable", async (t) => { + let acknowledgeCancel!: () => void; + const cancelled = new Promise((resolve) => { acknowledgeCancel = resolve; }); + const cancellations: string[] = []; + const bridge = await createToolBridge([READ_TOOL], { + onCall() {}, onError(error) { throw error; }, + onCancel(call) { cancellations.push(call.id); acknowledgeCancel(); }, + }); + t.after(() => bridge.close()); + const client = await connectClient(bridge); + t.after(() => client.close()); + const first = client.callTool({ name: "read", arguments: { path: "a" }, _meta: { "claudecode/toolUseId": "call-a" } }); + const controller = new AbortController(); + const second = client.callTool({ name: "read", arguments: { path: "b" }, _meta: { "claudecode/toolUseId": "call-b" } }, undefined, { signal: controller.signal }); + const rejected = assert.rejects(second, /cancelled B/); + await bridge.waitForToolCalls(2); + controller.abort("cancelled B"); + await cancelled; + await rejected; + assert.deepEqual(cancellations, ["call-b"]); + const pending = bridge.getPendingToolCalls(); + assert.deepEqual(pending.map((call) => call.id), ["call-a"]); + bridge.resolveToolCall(pending[0]!, toolResult([{ type: "text", text: "A" }])); + assert.deepEqual((await first).content, [{ type: "text", text: "A" }]); + assert.deepEqual(bridge.getPendingToolCalls(), []); +}); + test("waitForToolCalls resolves once that many calls are pending, in arrival order", async (t) => { const bridge = await createToolBridge([READ_TOOL, SCREENSHOT_TOOL]); t.after(() => bridge.close()); diff --git a/tests/fixtures/claude-code/fake-claude.mjs b/tests/fixtures/claude-code/fake-claude.mjs index d1fef86..7438444 100644 --- a/tests/fixtures/claude-code/fake-claude.mjs +++ b/tests/fixtures/claude-code/fake-claude.mjs @@ -6,17 +6,20 @@ * {"fake":"exit","code":1,"stderr":"..."} write stderr and exit with that code * {"fake":"call","id":"toolu_...","name":"read","arguments":{...}} call the MCP bridge and wait for its result * {"fake":"start-call",...} start the same request without waiting for its result + * {"fake":"cancel-call","id":"toolu_..."} cancel an individual pending MCP request + * {"fake":"checkpoint"} notify the parent after previous sends finish * {"fake":"wait-for-message"} pause until the test releases the child over IPC * {"fake":"ignore-term"} ignore SIGTERM to exercise kill escalation * {"fake":"reply"} stream one assistant message whose text is every result so far, joined by " | " * FAKE_CLAUDE_STDIN, when set, receives the stdin it read. */ import { readFileSync, writeFileSync } from "node:fs"; -import { once } from "node:events"; +import { on } from "node:events"; import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; +const parentMessages = on(process, "message"); const HOLD_INTERVAL_MS = 60_000; const BRIDGE_SERVER_NAME = "learnharness"; @@ -28,11 +31,16 @@ function streamEvent(event) { return JSON.stringify({ type: "stream_event", event }); } +let lastSend = Promise.resolve(); + async function connectBridge(mcpConfig) { const { url, headers } = JSON.parse(mcpConfig).mcpServers[BRIDGE_SERVER_NAME]; const client = new Client({ name: "fake-claude", version: "1.0.0" }); - await client.connect(new StreamableHTTPClientTransport(new URL(url), { requestInit: { headers } })); + const transport = new StreamableHTTPClientTransport(new URL(url), { requestInit: { headers } }); + const send = transport.send.bind(transport); + transport.send = (...args) => (lastSend = send(...args)); + await client.connect(transport); return client; } @@ -54,12 +62,25 @@ if (process.env.FAKE_CLAUDE_STDIN) writeFileSync(process.env.FAKE_CLAUDE_STDIN, let bridge = null; const results = []; const pendingResults = []; +const controllers = new Map(); +const cancelled = new Set(); const lines = readFileSync(process.env.FAKE_CLAUDE_SCENARIO, "utf8").split("\n").filter(Boolean); for (const line of lines) { const directive = JSON.parse(line); + if (directive.fake === "cancel-call") { + cancelled.add(directive.id); + controllers.get(directive.id)?.abort("cancelled by fake CLI"); + await lastSend; + continue; + } + if (directive.fake === "checkpoint") { + await lastSend; + process.send?.({ fake: "checkpoint" }); + continue; + } if (directive.fake === "wait-for-message") { - await once(process, "message"); + await parentMessages.next(); continue; } if (directive.fake === "hold") { @@ -78,14 +99,17 @@ for (const line of lines) { } if (directive.fake === "call" || directive.fake === "start-call") { bridge ??= await connectBridge(mcpConfig); + const controller = new AbortController(); + controllers.set(directive.id, controller); const request = bridge.callTool({ name: directive.name, arguments: directive.arguments, ...(directive.id !== undefined ? { _meta: { "claudecode/toolUseId": directive.id } } : {}), - }); + }, undefined, { signal: controller.signal }); if (directive.fake === "start-call") { - pendingResults.push(request); + pendingResults.push({ id: directive.id, request }); void request.catch(() => {}); + await lastSend; } else { const result = await request; results.push(result.content.map((block) => block.text).join("")); @@ -93,7 +117,7 @@ for (const line of lines) { continue; } if (directive.fake === "reply") { - for (const result of await Promise.all(pendingResults.splice(0))) results.push(result.content.map((block) => block.text).join("")); + for (const result of await Promise.all(pendingResults.splice(0).filter(({ id }) => !cancelled.has(id)).map(({ request }) => request))) results.push(result.content.map((block) => block.text).join("")); await writeReply(results.join(" | ")); continue; }