fix: Recover prompt stream after a failed turn (#706)

When a prompt errors mid-turn (e.g. a 400/529 from the API), the SDK
leaves a trailing session_state_changed: idle in the query iterator.
The next prompt's first query.next() consumed that stale idle and
short-circuited to end_turn with zero usage, shifting every subsequent
response one slot late until the session was restarted.

In the catch block, call query.interrupt() and drain the iterator
until the trailing idle is consumed (bounded so a misbehaving SDK
can't hang the next prompt). In the finally block, on the error path,
resolve all queued pending prompts as cancelled — the previous code
only resolved one (with handoff semantics), leaking the rest forever.

Closes #654
Supersedes #659, #660, #691

src/acp-agent.ts

@@ -781,6 +781,7 @@ export class ClaudeAcpAgent implements Agent {
 
     session.promptRunning = true;
     let handedOff = false;
+    let errored = false;
     let stopReason: StopReason = "end_turn";
 
     try {
@@ -1265,6 +1266,34 @@ export class ClaudeAcpAgent implements Agent {
       }
       throw new Error("Session did not end in result");
     } catch (error) {
+      errored = true;
+      // A failed turn typically leaves a trailing `session_state_changed: idle`
+      // (and possibly more) in the query iterator. If we don't drain it here,
+      // the next prompt's first `query.next()` consumes that stale idle and
+      // short-circuits to end_turn with zero usage
+      // Bounded so a misbehaving SDK can't hang the next prompt indefinitely.
+      try {
+        await session.query.interrupt();
+        const MAX_DRAIN = 100;
+        for (let i = 0; i < MAX_DRAIN; i++) {
+          const { value: m, done } = await session.query.next();
+          if (done || !m) break;
+          if (m.type === "system" && m.subtype === "session_state_changed" && m.state === "idle") {
+            break;
+          }
+          if (i === MAX_DRAIN - 1) {
+            this.logger.error(
+              `Session ${params.sessionId}: drained ${MAX_DRAIN} messages after error without observing idle`,
+            );
+          }
+        }
+      } catch (drainErr) {
+        this.logger.error(
+          `Session ${params.sessionId}: failed to drain query after prompt error:`,
+          drainErr,
+        );
+      }
+
       if (error instanceof RequestError || !(error instanceof Error)) {
         throw error;
       }
@@ -1289,10 +1318,19 @@ export class ClaudeAcpAgent implements Agent {
     } finally {
       if (!handedOff) {
         session.promptRunning = false;
-        // This usually should not happen, but in case the loop finishes
-        // without claude sending all message replays, we resolve the
-        // next pending prompt call to ensure no prompts get stuck.
-        if (session.pendingMessages.size > 0) {
+        if (errored) {
+          // The query stream was just drained — handing pending prompts off
+          // onto it would let them race with the recovery. Cancel them so
+          // each waiting prompt() returns stopReason: "cancelled" and the
+          // client can decide whether to retry.
+          for (const pending of session.pendingMessages.values()) {
+            pending.resolve(true);
+          }
+          session.pendingMessages.clear();
+        } else if (session.pendingMessages.size > 0) {
+          // This usually should not happen, but in case the loop finishes
+          // without claude sending all message replays, we resolve the
+          // next pending prompt call to ensure no prompts get stuck.
           const next = [...session.pendingMessages.entries()].sort(
             (a, b) => a[1].order - b[1].order,
           )[0];

src/tests/acp-agent.test.ts

@@ -3697,3 +3697,178 @@ describe("memory_recall handling", () => {
     expect(toolCall.update._meta.claudeCode.toolResponse).toEqual({ mode: "synthesize" });
   });
 });
+
+describe("post-error recovery", () => {
+  function createMockAgent() {
+    const mockClient = {
+      sessionUpdate: async () => {},
+    } as unknown as AgentSideConnection;
+    return new ClaudeAcpAgent(mockClient, { log: () => {}, error: () => {} });
+  }
+
+  function createResultMessage(overrides: {
+    subtype: "success" | "error_during_execution";
+    stop_reason: string | null;
+    is_error: boolean;
+    result?: string;
+    errors?: string[];
+  }) {
+    return {
+      type: "result" as const,
+      subtype: overrides.subtype,
+      stop_reason: overrides.stop_reason,
+      is_error: overrides.is_error,
+      result: overrides.result ?? "",
+      errors: overrides.errors ?? [],
+      duration_ms: 0,
+      duration_api_ms: 0,
+      num_turns: 1,
+      total_cost_usd: 0,
+      usage: {
+        input_tokens: 10,
+        output_tokens: 5,
+        cache_read_input_tokens: 0,
+        cache_creation_input_tokens: 0,
+      },
+      modelUsage: {},
+      permission_denials: [],
+      uuid: randomUUID(),
+      session_id: "test-session",
+    };
+  }
+
+  // Two-turn generator: turn 1 yields the caller-supplied `firstTurn`
+  // messages (including a trailing idle that the drain must consume).
+  // Turn 2 yields a clean success + idle, used to verify the next prompt
+  // sees real messages rather than the stale idle.
+  function injectTwoTurnSession(agent: ClaudeAcpAgent, firstTurn: unknown[]) {
+    const input = new Pushable<any>();
+    const interrupt = vi.fn(async () => {});
+    const close = vi.fn();
+    async function* messageGenerator() {
+      const iter = input[Symbol.asyncIterator]();
+
+      const first = await iter.next();
+      if (!first.done && first.value) {
+        yield {
+          type: "user",
+          message: first.value.message,
+          parent_tool_use_id: null,
+          uuid: first.value.uuid,
+          session_id: "test-session",
+          isReplay: true,
+        };
+      }
+      yield* firstTurn;
+
+      const second = await iter.next();
+      if (!second.done && second.value) {
+        yield {
+          type: "user",
+          message: second.value.message,
+          parent_tool_use_id: null,
+          uuid: second.value.uuid,
+          session_id: "test-session",
+          isReplay: true,
+        };
+      }
+      yield createResultMessage({ subtype: "success", stop_reason: null, is_error: false });
+      yield { type: "system", subtype: "session_state_changed", state: "idle" };
+    }
+    const gen = Object.assign(messageGenerator(), { interrupt, close });
+    agent.sessions["test-session"] = {
+      query: gen as any,
+      input,
+      cancelled: false,
+      cwd: "/test",
+      sessionFingerprint: JSON.stringify({ cwd: "/test", mcpServers: [] }),
+      modes: { currentModeId: "default", availableModes: [] },
+      models: { currentModelId: "default", availableModels: [] },
+      modelInfos: [],
+      settingsManager: { dispose: vi.fn() } as any,
+      accumulatedUsage: {
+        inputTokens: 0,
+        outputTokens: 0,
+        cachedReadTokens: 0,
+        cachedWriteTokens: 0,
+      },
+      configOptions: [],
+      promptRunning: false,
+      pendingMessages: new Map(),
+      nextPendingOrder: 0,
+      abortController: new AbortController(),
+      emitRawSDKMessages: false,
+      contextWindowSize: 200000,
+      taskState: new Map(),
+    };
+    return { interrupt };
+  }
+
+  it("drains a failed turn's trailing idle so the next prompt is not short-circuited", async () => {
+    const agent = createMockAgent();
+    const { interrupt } = injectTwoTurnSession(agent, [
+      createResultMessage({
+        subtype: "success",
+        stop_reason: "end_turn",
+        is_error: true,
+        result: "boom",
+      }),
+      // Trailing idle from the failed turn. Without draining, the next
+      // prompt's first query.next() would consume this and short-circuit
+      // to end_turn with zero usage (issue #654).
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+
+    await expect(
+      agent.prompt({
+        sessionId: "test-session",
+        prompt: [{ type: "text", text: "first" }],
+      }),
+    ).rejects.toThrow();
+
+    expect(interrupt).toHaveBeenCalled();
+
+    const second = await agent.prompt({
+      sessionId: "test-session",
+      prompt: [{ type: "text", text: "second" }],
+    });
+    expect(second.stopReason).toBe("end_turn");
+    expect(second.usage?.inputTokens).toBe(10);
+    expect(second.usage?.outputTokens).toBe(5);
+  });
+
+  it("cancels all queued pending prompts when a turn errors", async () => {
+    const agent = createMockAgent();
+    injectTwoTurnSession(agent, [
+      createResultMessage({
+        subtype: "success",
+        stop_reason: "end_turn",
+        is_error: true,
+        result: "boom",
+      }),
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+
+    // Simulate two prompts already queued behind the running turn. Both
+    // resolvers should fire with `true` (cancelled) when the running
+    // prompt errors, and the map should be cleared.
+    const session = agent.sessions["test-session"];
+    let resolveA!: (cancelled: boolean) => void;
+    let resolveB!: (cancelled: boolean) => void;
+    const pendingA = new Promise<boolean>((r) => (resolveA = r));
+    const pendingB = new Promise<boolean>((r) => (resolveB = r));
+    session.pendingMessages.set("uuid-a", { resolve: resolveA, order: 0 });
+    session.pendingMessages.set("uuid-b", { resolve: resolveB, order: 1 });
+
+    await expect(
+      agent.prompt({
+        sessionId: "test-session",
+        prompt: [{ type: "text", text: "first" }],
+      }),
+    ).rejects.toThrow();
+
+    await expect(pendingA).resolves.toBe(true);
+    await expect(pendingB).resolves.toBe(true);
+    expect(session.pendingMessages.size).toBe(0);
+  });
+});