From bed56092a5dfcd170a85abe0e1b67fbe41a25f74 Mon Sep 17 00:00:00 2001 From: Alexander Dines Date: Fri, 25 Sep 2026 17:49:54 +0000 Subject: [PATCH] fix(claude): retire cancelled control requests during replay --- sdk/README.md | 4 + sdk/src/claude/replay-contract.test.ts | 187 +++++++++++++++++++++++++ sdk/src/claude/transport.test.ts | 14 +- sdk/src/claude/transport.ts | 7 +- sdk/src/shared/axon-frame-transport.ts | 9 +- 5 files changed, 215 insertions(+), 6 deletions(-) create mode 100644 sdk/src/claude/replay-contract.test.ts diff --git a/sdk/README.md b/sdk/README.md index 9fdb98e..3b10871 100644 --- a/sdk/README.md +++ b/sdk/README.md @@ -816,6 +816,10 @@ conn.onTimelineEvent((event) => { All modules subscribe to the Axon SSE stream. By default, `replay` is **`true`**: the connection queries the channel head and replays all events up to that point **without** dispatching to session updates, permission handlers, (Claude) control handlers, or (Codex) approval handlers — timeline listeners still receive every event. Unresolved permission/control/approval work is flushed after replay completes. +Claude replay retires a control request when it encounters either a matching user `control_response` or an agent `control_cancel_request`. A newer user query alone does not retire an earlier pending request. Choose `afterSequence` to include any pending request that must be recovered; a request before that bound is never fetched. Request identity is the protocol `request_id`, independent of the connection handling it. + +Claude and Codex retain buffered replay requests across SSE reconnects. A stream ending before the captured head does not dispatch those requests: the remaining history may still contain their responses or cancellations. + Set **`replay: false`** to restore the previous behavior: every replayed event invokes handlers immediately (useful if you rely on side effects during history replay). Pass **`afterSequence`** to subscribe starting **after** a known sequence number (earlier events are never requested — no full replay from the beginning). It composes with `replay`: diff --git a/sdk/src/claude/replay-contract.test.ts b/sdk/src/claude/replay-contract.test.ts new file mode 100644 index 0000000..cc5d6f6 --- /dev/null +++ b/sdk/src/claude/replay-contract.test.ts @@ -0,0 +1,187 @@ +import type { SDKControlResponse } from "@anthropic-ai/claude-agent-sdk"; +import type { Axon, Devbox } from "@runloop/api-client/sdk"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { + createControllableStream, + type MockAxonEvent, + makeAgentEvent, + makeSystemEvent, + makeUserEvent, +} from "../__test-utils__/mock-axon.js"; +import { ClaudeAxonConnection } from "./connection.js"; + +const question = (id: string, sequence: number) => + makeAgentEvent( + "control_request", + { + type: "control_request", + request_id: id, + request: { + subtype: "can_use_tool", + tool_name: "AskUserQuestion", + tool_use_id: `tool-${id}`, + input: { questions: [{ question: "Continue?", header: "Next", options: [] }] }, + }, + }, + sequence, + ); + +const cancel = (id: string, sequence: number) => + makeAgentEvent( + "control_cancel_request", + { type: "control_cancel_request", request_id: id }, + sequence, + ); + +/** Only the Axon boundary is fake; connection, transport and read loop are real. */ +class FakeAxon { + readonly id = "axon-replay-contract"; + readonly subscriptions: Array = []; + readonly streams: ReturnType[] = []; + readonly published: unknown[] = []; + readonly client = { get: async () => ({ total_count: this.history.at(-1)?.sequence }) }; + + constructor( + readonly history: MockAxonEvent[], + private readonly disconnectAfter?: number, + ) {} + + async subscribeSse(options?: { after_sequence: number }) { + this.subscriptions.push(options?.after_sequence); + const stream = createControllableStream(true); + stream.stream.controller?.abort.mockImplementation(() => stream.end()); + this.streams.push(stream); + for (const event of this.history) { + if ((event.sequence ?? 0) <= (options?.after_sequence ?? 0)) continue; + stream.push(event); + if (this.streams.length === 1 && event.sequence === this.disconnectAfter) { + stream.end(); + break; + } + } + return stream.stream; + } + + async publish(event: unknown) { + this.published.push(event); + } +} + +const connections: ClaudeAxonConnection[] = []; +afterEach(async () => { + await Promise.all(connections.splice(0).map((connection) => connection.disconnect())); +}); + +async function connect(axon: FakeAxon, afterSequence?: number) { + const dispatched: string[] = []; + const observed: number[] = []; + const connection = new ClaudeAxonConnection( + axon as unknown as Axon, + { id: "devbox-replay-contract" } as Devbox, + { afterSequence, retry: { baseDelayMs: 1, maxDelayMs: 1 } }, + ); + connections.push(connection); + connection.onAxonEvent((event) => observed.push(event.sequence)); + connection.onControlRequest("can_use_tool", (request) => { + dispatched.push(request.request_id); + // Model a parked application question. No response is persisted yet. + return new Promise(() => {}); + }); + await connection.connect(); + return { connection, dispatched, observed }; +} + +describe("Claude question replay contract", () => { + it.each([ + undefined, + 4, + ])("does not re-park cancelled history (afterSequence=%s)", async (bound) => { + const history = [ + question("old", 1), + cancel("old", 2), + makeSystemEvent("agent.interrupted", { agentId: "agent" }, 3), + makeAgentEvent("result", { type: "result" }, 4), + makeUserEvent("query", { type: "user" }, 5), + question("current", 6), + ]; + const axon = new FakeAxon(history); + const result = await connect(axon, bound); + await vi.waitFor(() => expect(result.dispatched).toEqual(["current"])); + expect(result.observed).toEqual(bound === undefined ? [1, 2, 3, 4, 5, 6] : [5, 6]); + expect(axon.subscriptions).toEqual([bound]); + expect(axon.published).toEqual([]); + }); + + it.each([ + "allow", + "deny", + "error", + ])("does not re-park an answered request (%s)", async (behavior) => { + const axon = new FakeAxon([ + question("answered", 1), + makeUserEvent( + "control_response", + { + type: "control_response", + response: { + request_id: "answered", + subtype: behavior === "error" ? "error" : "success", + response: { behavior }, + }, + }, + 2, + ), + question("pending", 3), + ]); + const result = await connect(axon); + await vi.waitFor(() => expect(result.dispatched).toEqual(["pending"])); + }); + + it("preserves a pending identity across a newer query and a new owner connection", async () => { + const axon = new FakeAxon([ + question("pending", 1), + makeUserEvent("query", { type: "user" }, 2), + ]); + const first = await connect(axon); + await vi.waitFor(() => expect(first.dispatched).toEqual(["pending"])); + await first.connection.disconnect(); + const second = await connect(axon); + await vi.waitFor(() => expect(second.dispatched).toEqual(["pending"])); + expect(axon.published).toEqual([]); + }); + + it("only retires the matching id on an agent cancellation", async () => { + const axon = new FakeAxon([ + question("pending", 1), + cancel("unrelated", 2), + makeUserEvent( + "control_cancel_request", + { type: "control_cancel_request", request_id: "pending" }, + 3, + ), + makeAgentEvent("control_cancel_request", { type: "control_cancel_request" }, 4), + ]); + const result = await connect(axon); + await vi.waitFor(() => expect(result.dispatched).toEqual(["pending"])); + }); + + it.each([ + true, + false, + ])("finishes replay after a dropped stream before dispatch (cancelled=%s)", async (cancelled) => { + const axon = new FakeAxon( + [ + question("old", 1), + cancelled ? cancel("old", 2) : makeUserEvent("query", { type: "user" }, 2), + question("current", 3), + ], + 1, + ); + const result = await connect(axon); + await vi.waitFor(() => + expect(result.dispatched).toEqual(cancelled ? ["current"] : ["old", "current"]), + ); + expect(axon.subscriptions).toEqual([undefined, 1]); + expect(result.observed).toEqual([1, 2, 3]); + }); +}); diff --git a/sdk/src/claude/transport.test.ts b/sdk/src/claude/transport.test.ts index 890b898..e865a9d 100644 --- a/sdk/src/claude/transport.test.ts +++ b/sdk/src/claude/transport.test.ts @@ -677,7 +677,7 @@ describe("AxonTransport", () => { expect(messages).toHaveLength(0); }); - it("flushes unresolved requests when stream ends before replay target", async () => { + it("retains unresolved requests when stream ends before replay target", async () => { const ctrl2 = createControllableStream(true); const mock = createMockAxon(ctrl2); const t = new AxonTransport(mock.axon as never, { replayTargetSequence: 100 }); @@ -695,8 +695,16 @@ describe("AxonTransport", () => { for await (const msg of t.readMessages()) { messages.push(msg); } - expect(messages).toHaveLength(1); - expect(messages[0]).toMatchObject({ request_id: "req-1" }); + expect(messages).toHaveLength(0); + + const resumed = createControllableStream(true); + mock.axon.subscribeSse.mockResolvedValueOnce(resumed.stream); + await t.reconnect(); + expect(mock.axon.subscribeSse).toHaveBeenLastCalledWith({ after_sequence: 1 }); + resumed.push(makeAgentEvent("assistant", { type: "assistant" }, 100)); + resumed.end(); + for await (const msg of t.readMessages()) messages.push(msg); + expect(messages).toEqual([{ request_id: "req-1", type: "can_use_tool" }]); }); it("fires onAxonEvent during replay", async () => { diff --git a/sdk/src/claude/transport.ts b/sdk/src/claude/transport.ts index 28978cb..2b10a12 100644 --- a/sdk/src/claude/transport.ts +++ b/sdk/src/claude/transport.ts @@ -194,8 +194,13 @@ export class AxonTransport implements Transport { allowInvalidOutbound: true, isReplayRequest: (event) => isFromAgent(event) && isControlRequest(event), requestId: (frame) => getRequestId(frame), - isReplayAnswer: (event) => isFromUser(event) && isControlResponse(event), + // Cancellation is sent by the agent and carries the request id at the + // top level, unlike the nested id in a user control response. + isReplayAnswer: (event) => + (isFromUser(event) && isControlResponse(event)) || + (isFromAgent(event) && event.event_type === "control_cancel_request"), answerId: (frame) => { + if (frame.type === "control_cancel_request") return getRequestId(frame); const response = isNonNullObject(frame.response) ? frame.response : undefined; return getRequestId(response); }, diff --git a/sdk/src/shared/axon-frame-transport.ts b/sdk/src/shared/axon-frame-transport.ts index 223a697..ee1f27a 100644 --- a/sdk/src/shared/axon-frame-transport.ts +++ b/sdk/src/shared/axon-frame-transport.ts @@ -37,6 +37,9 @@ export class AxonFrameTransport { private connected = false; private closed = false; private lastSequence?: number; + // A stream may disconnect before reaching the replay head. Keep correlation + // state alongside the resume cursor so the next subscription can finish it. + private readonly replayBuffer = new Map(); private readonly log: LogFn; constructor( @@ -86,7 +89,7 @@ export class AxonFrameTransport { async *readMessages(): AsyncGenerator { if (!this.stream) throw new Error("Transport not connected. Call connect() first."); const target = this.options.replayTargetSequence; - const replayBuffer = new Map(); + const replayBuffer = this.replayBuffer; const flush = function* (): Generator { for (const frame of replayBuffer.values()) yield frame; replayBuffer.clear(); @@ -123,7 +126,8 @@ export class AxonFrameTransport { if (frame !== undefined) yield frame; else this.log("read", `invalid frame for ${event.event_type}`); } - yield* flush(); + // EOF is not proof that replay reached its head: a later page may contain + // an answer or cancellation. Only flush at/past the captured target. } abortStream(): void { @@ -135,6 +139,7 @@ export class AxonFrameTransport { if (this.closed) return; this.closed = true; this.connected = false; + this.replayBuffer.clear(); this.abortStream(); }