fix(agent): tighten idle-drain lifecycle from PR review

Address Greptile review comments on the idle drainer:

- interrupt(): await settleIdleDrainStop() after query.interrupt() so a drainer
  that is mid-forward (awaiting client.sessionUpdate) fully exits before
  interrupt() returns, closing the window where it could emit an event for an
  already-cancelled session.
- startIdleDrain scheduling: capture the session at schedule time and bail in
  the microtask if this.session was swapped (e.g. by refreshSession) before it
  runs, so the drainer can't tag new-query events with the old session id.
- Document the single-drainer invariant on the startIdleDrain guard: because a
  drainer holding a handoff keeps this.idleDrain non-null until
  settleIdleDrainStop clears it, no second drainer can be assigned in that
  window (the third review comment was self-correcting, not a bug).

Adds a test asserting cancel() while idle-draining awaits the drainer's exit.

Generated-By: PostHog Code
Task-Id: 39548658-f313-445b-bd1c-d07eaf9281cf

Author: jonathanlab

packages/agent/src/adapters/claude/claude-agent.idle-drain.test.ts

@@ -232,4 +232,21 @@ describe("ClaudeAcpAgent — between-turns idle drain (autonomous /loop turns)",
     // follow-up prompt drained the backlog).
     expect(texts.filter((t) => t === "autonomous update")).toHaveLength(1);
   });
+
+  it("cancel while idle-draining stops the drainer and returns cleanly", async () => {
+    const { agent } = makeAgent();
+    const sessionId = "s-idle-3";
+    const { query, input } = installFakeSession(agent, sessionId);
+
+    await runTurn(agent, query, input, sessionId, "loop", "answer", "r1");
+    // Drainer starts and blocks on query.next().
+    await tick();
+    const drainSlot = () =>
+      (agent as unknown as { idleDrain: unknown }).idleDrain;
+    expect(drainSlot()).not.toBeNull();
+
+    // Cancelling must await the drainer's exit (not return with it still live).
+    await agent.cancel({ sessionId });
+    expect(drainSlot()).toBeNull();
+  });
 });

packages/agent/src/adapters/claude/claude-agent.ts

@@ -1223,15 +1223,25 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
         !this.session.cancelled
       ) {
         const sessionId = params.sessionId;
-        queueMicrotask(() => this.startIdleDrain(sessionId));
+        const session = this.session;
+        queueMicrotask(() => {
+          // Bail if the session was swapped (e.g. by refreshSession) between
+          // scheduling and running: startIdleDrain reads this.session at
+          // execution time, so without this it could drain the new query while
+          // tagging events with the old session id.
+          if (this.session === session) {
+            this.startIdleDrain(sessionId);
+          }
+        });
       }
     }
   }
 
   // Called by BaseAcpAgent#cancel() to interrupt the session
   protected async interrupt(): Promise<void> {
     // A cancel/close may land while the between-turns idle drainer owns the
-    // query; stop it so `query.interrupt()` below isn't racing a live reader.
+    // query; flag it now so it stops at its next boundary, and fence its exit
+    // after `query.interrupt()` below (see end of method).
     this.requestIdleDrainStop();
     this.session.cancelled = true;
     for (const [, pending] of this.session.pendingMessages) {
@@ -1255,6 +1265,13 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
     }
 
     await this.session.query.interrupt();
+
+    // Fence the drainer's exit. `query.interrupt()` unblocks a drainer suspended
+    // in `query.next()`, but one mid-forward only exits after its in-flight
+    // `client.sessionUpdate()` settles. Awaiting here keeps interrupt() from
+    // returning while the drainer can still emit an event for a cancelled
+    // session. No-op when no drainer is running.
+    await this.settleIdleDrainStop();
   }
 
   /**
@@ -1355,6 +1372,11 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
    * (reverting to prior behavior), never breaks the session.
    */
   private startIdleDrain(sessionId: string): void {
+    // At most one drainer at a time. Combined with "drainers are only scheduled
+    // after a completed turn", this is what keeps the slot unambiguous: while a
+    // drainer holds a handoff for an in-flight prompt(), `this.idleDrain` stays
+    // non-null (its finally skips clearing), so no second drainer can be
+    // assigned before settleIdleDrainStop() clears it.
     if (this.idleDrain) return;
     if (this.session.promptRunning) return;
     if (this.session.abortController.signal.aborted) return;