sorokin.oleg.edu/learnharness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 6 Wiki Activity

fix: scope held Claude processes and honor MCP cancellation

Browse Source
This commit is contained in:
amosblomqvist
2026-10-28 17:19:51 +02:00
parent 6ca431644d
commit 64a36bebd2
7 changed files with 192 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent/claude-code/live-processes.ts
+4 -4
View File
@@ -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;
src/agent/claude-code/stream.ts
+23 -11
View File
@@ -392,6 +392,7 @@ class ClaudeProcess implements LiveProcess {
private sourceCalls: ToolCall[] = [];
private sourceComplete = false;
private admitted: PendingToolCall[] = [];
private readonly cancelledIds = new Set<string>();
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);
src/agent/claude-code/tool-bridge.ts
+22 -7
View File
@@ -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<PendingToolCall, (result: CallToolResult) => void>();
private waiters: ToolCallWaiter[] = [];
add(call: PendingToolCall): Promise<CallToolResult> {
const result = new Promise<CallToolResult>((resolve) => this.results.set(call, resolve));
add(call: PendingToolCall, signal: AbortSignal, onCancel?: (call: PendingToolCall) => void): Promise<CallToolResult> {
const result = new Promise<CallToolResult>((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<string>();
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;
});