fix(sessions): let a dismiss reach the entries a restore is holding

Fourth review of the reboot-restore branch, and the third to find a defect
in the previous round's fix. This one is the same shape as its predecessor:
a counter keyed on one thing, compared against a set keyed on another.

The generation counter was indexed by the entry's owner, while the in-flight
set holds the caller doing the restoring. Those are the same person exactly
when a user restores their own sessions, which is every case the tests
covered. The route deliberately supports the other case: an admin may spend
another user's entries. So when an admin restored Bob's sessions and Bob
dismissed the banner, nothing matched, the entries came back, and a plan Bob
had explicitly dismissed was re-armed for another twenty-four hours.

Rather than reconcile the two key spaces, the counter is gone. `take()` now
parks the entries it hands out, remembering which caller is spending them,
and they stay parked until that restore ends. A dismiss filters the parked
entries by `canAccess(entry.owner)` — the same predicate it already applies
to the plan — so it reaches them wherever they are. `releaseFlight()` puts
back only what is still parked. Expiry and a fresh boot plan unpark
everything, for the same reason. There is one key space now, the entry's
owner, and the spender is only ever used to tell two concurrent flights
apart. That removes `generations`, `snapshotGenerations()`, `bump()`,
`bumpAll()` and the argument threaded through the route.

The discard grew the teardown it still lacked. A rebuild can fail after
startInteractive() resolved, and a restored workspace still carries
Codeman's hooks, so the CLI can post a hook event within milliseconds; the
transcript watcher that starts from it, the attachment registry, the wait
registry and the approvals inbox all outlive the listeners and would meet
the retry, which reuses the session id by design. Its steps also run in
reverse order now, so no live listener can reach a tracker that has already
stopped, and the mux kill has its own guard, because stop() kills the pane
in its last block after destroying four trackers.

Tests. The run-summary test named an interval and asserted a map entry, so
dropping stop() left it green; it now spies on stop(). Nothing pinned that
before-spawn must precede setupSessionListeners, which reads the flag that
phase restores, so swapping the two lines was silent; the ordering test now
includes the listener setup. The retry assertion was a tautology and now
asserts a different refs object. Both strengthened tests were verified by
reverting their fix. Two new tests cover the admin-restores-another-owner
cases this round was about. The server in the discard test is built once and
stopped, since its constructor registers handlers on module-level watchers,
and the workspace is removed through safeRmHomeTree.

Refs #411

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

src/web/reboot-restore-registry.ts

@@ -48,16 +48,18 @@ export class RebootRestoreRegistry {
   /** When the boot pass built the plan, in ms since the epoch. */
   private builtAt = 0;
   /**
-   * Per owner, bumped by anything that invalidates that owner's entries while a
-   * restore is already holding them. A Dismiss arriving mid-restore must win:
-   * without this the route's `finally` would put its unspent entries back and
-   * resurrect the offer the user just cleared, with a fresh 24-hour life.
+   * Entries handed to a restore that has not finished, by session id, each
+   * remembering which caller is spending it.
    *
-   * Keyed by owner rather than global, because `clear()` is ownership-scoped. A
-   * single counter would let one user's Dismiss discard another user's unspent
-   * entries, and the plan is in-memory, so those offers would be gone for good.
+   * A taken entry is still part of the offer until its restore resolves it, so
+   * it has to stay reachable by everything that can invalidate an offer. Holding
+   * the entries themselves — rather than a counter to compare against later —
+   * means `clear()` filters them by the SAME `canAccess(entry.owner)` predicate
+   * it already applies to the plan. A counter cannot do that, because the caller
+   * spending an entry need not be its owner: an admin may restore another user's
+   * sessions, and then the spender and the owner are different keys.
    */
-  private generations = new Map<string | undefined, number>();
+  private parked = new Map<string, { entry: RebootRestoreEntry; spender: string | undefined }>();
   /**
    * Owners with a restore in flight, between its take and its last pane.
    * Keyed by owner so one user's restore does not turn another user's click into
@@ -70,17 +72,8 @@ export class RebootRestoreRegistry {
   set(entries: readonly RebootRestoreEntry[]): void {
     this.entries = new Map(entries.map((entry) => [entry.sessionId, entry]));
     this.builtAt = entries.length > 0 ? Date.now() : 0;
-    this.bumpAll();
-  }
-
-  /**
-   * The generations of the owners of `entries`, for a caller that will hand some
-   * of them back later. Pass the result to {@link restore}.
-   */
-  snapshotGenerations(entries: readonly RebootRestoreEntry[]): Map<string | undefined, number> {
-    const snapshot = new Map<string | undefined, number>();
-    for (const entry of entries) snapshot.set(entry.owner, this.generations.get(entry.owner) ?? 0);
-    return snapshot;
+    // A fresh boot plan supersedes anything an in-flight restore still holds.
+    this.parked.clear();
   }
 
   /**
@@ -104,7 +97,11 @@ export class RebootRestoreRegistry {
    *
    * @param sessionIds The ids to spend, or undefined for every visible entry.
    */
-  take(canAccess: (owner: string | undefined) => boolean, sessionIds?: readonly string[]): RebootRestoreEntry[] {
+  take(
+    canAccess: (owner: string | undefined) => boolean,
+    sessionIds: readonly string[] | undefined,
+    spender: string | undefined
+  ): RebootRestoreEntry[] {
     this.dropIfExpired();
     const wanted = sessionIds ? new Set(sessionIds) : undefined;
     const taken: RebootRestoreEntry[] = [];
@@ -112,6 +109,9 @@ export class RebootRestoreRegistry {
       if (wanted && !wanted.has(entry.sessionId)) continue;
       if (!canAccess(entry.owner)) continue;
       this.entries.delete(entry.sessionId);
+      // Parked rather than forgotten: until this restore resolves the entry, a
+      // dismiss still has to be able to reach and cancel it.
+      this.parked.set(entry.sessionId, { entry, spender });
       taken.push(entry);
     }
     return taken;
@@ -126,18 +126,19 @@ export class RebootRestoreRegistry {
    * hand is NOT put back, because that one cannot stop being true, and an entry
    * the banner keeps re-offering forever is noise only Dismiss can clear.
    */
-  restore(entries: readonly RebootRestoreEntry[], generations?: ReadonlyMap<string | undefined, number>): void {
+  releaseFlight(spender: string | undefined, keep: readonly RebootRestoreEntry[]): void {
+    const wanted = new Set(keep.map((entry) => entry.sessionId));
     let added = 0;
-    for (const entry of entries) {
-      // A dismiss (or a fresh boot plan) for THIS entry's owner since the caller
-      // took it means it is no longer wanted back. Another owner's dismiss is
-      // none of this entry's business.
-      if (generations) {
-        const taken = generations.get(entry.owner);
-        if (taken !== undefined && taken !== (this.generations.get(entry.owner) ?? 0)) continue;
+    for (const [sessionId, held] of [...this.parked]) {
+      if (held.spender !== spender) continue;
+      this.parked.delete(sessionId);
+      // Still parked means nothing cancelled it while the restore ran. A dismiss,
+      // an expiry or a fresh boot plan removes it from `parked`, and then it does
+      // not come back however the restore ended.
+      if (wanted.has(sessionId)) {
+        this.entries.set(sessionId, held.entry);
+        added += 1;
       }
-      this.entries.set(entry.sessionId, entry);
-      added += 1;
     }
     if (added > 0 && this.builtAt === 0) this.builtAt = Date.now();
   }
@@ -146,16 +147,17 @@ export class RebootRestoreRegistry {
   clear(canAccess: (owner: string | undefined) => boolean): number {
     const removable = [...this.entries.values()].filter((entry) => canAccess(entry.owner));
     for (const entry of removable) this.entries.delete(entry.sessionId);
+    // Entries a restore is holding are dismissed by the same rule, so a dismiss
+    // that lands mid-restore wins. Judged on the ENTRY's owner, exactly as above,
+    // rather than on who happens to be restoring it.
+    let parkedRemoved = 0;
+    for (const [sessionId, held] of [...this.parked]) {
+      if (!canAccess(held.entry.owner)) continue;
+      this.parked.delete(sessionId);
+      parkedRemoved += 1;
+    }
     if (this.entries.size === 0) this.builtAt = 0;
-    // A restore in flight for these owners must not put their entries back. The
-    // in-flight owners are the ones that matter and the ones the plan can no
-    // longer name: `take()` has already removed their entries, so a dismiss that
-    // lands mid-restore sees nothing of theirs to remove. The bump is limited to
-    // owners this caller could see, so it cannot reach anyone else's restore.
-    const invalidated = new Set(removable.map((entry) => entry.owner));
-    for (const owner of this.spending) if (canAccess(owner)) invalidated.add(owner);
-    for (const owner of invalidated) this.bump(owner);
-    return removable.length;
+    return removable.length + parkedRemoved;
   }
 
   /**
@@ -176,27 +178,16 @@ export class RebootRestoreRegistry {
   /** Test hook: forget everything, including the single-flight claim. */
   reset(): void {
     this.entries.clear();
+    this.parked.clear();
     this.builtAt = 0;
     this.spending.clear();
-    this.generations.clear();
-  }
-
-  private bump(owner: string | undefined): void {
-    this.generations.set(owner, (this.generations.get(owner) ?? 0) + 1);
-  }
-
-  /** Invalidate every owner's in-flight returns, including owners not yet seen. */
-  private bumpAll(): void {
-    for (const owner of new Set([...this.entries.values()].map((entry) => entry.owner))) this.bump(owner);
-    for (const owner of [...this.generations.keys()]) this.bump(owner);
   }
 
   private dropIfExpired(): void {
     if (this.builtAt > 0 && Date.now() - this.builtAt > PLAN_TTL_MS) {
-      // Bump before clearing, while the owners are still known: a restore that
-      // took entries just before the expiry must not hand them back afterwards
-      // and give an expired plan another full day of life.
-      this.bumpAll();
+      // A restore that took entries just before the expiry must not hand them
+      // back afterwards and give an expired plan another full day of life.
+      this.parked.clear();
       this.entries.clear();
       this.builtAt = 0;
     }

src/web/routes/reboot-restore-routes.ts

@@ -92,8 +92,7 @@ export function registerRebootRestoreRoutes(app: FastifyInstance, ctx: RebootRes
     if (!rebootRestoreRegistry.beginSpending(owner)) {
       return reply.code(409).send(createErrorResponse(ApiErrorCode.CONFLICT, 'A reboot restore is already running'));
     }
-    const taken = rebootRestoreRegistry.take(canAccess, body.sessionIds);
-    const generations = rebootRestoreRegistry.snapshotGenerations(taken);
+    const taken = rebootRestoreRegistry.take(canAccess, body.sessionIds, owner);
     // Entries nothing built a pane for, returned to the plan on every exit path
     // including a throw. Without this a failure between here and the loop would
     // spend the offer and rebuild nothing, and the plan cannot be rebuilt.
@@ -252,9 +251,9 @@ export function registerRebootRestoreRoutes(app: FastifyInstance, ctx: RebootRes
     } finally {
       // Anything that never became a pane goes back on offer, including after a
       // throw, so a transient failure costs a retry rather than the whole plan.
-      // Passing the generations makes a Dismiss that landed mid-restore win, for
-      // the owners it actually covered.
-      rebootRestoreRegistry.restore([...unspent], generations);
+      // Ends the flight: entries still parked for it come back if they are in
+      // `unspent`, and a Dismiss that unparked them meanwhile wins.
+      rebootRestoreRegistry.releaseFlight(owner, [...unspent]);
       rebootRestoreRegistry.endSpending(owner);
     }
   });

src/web/server.ts

@@ -3010,26 +3010,42 @@ export class WebServer extends EventEmitter {
     if (!session) return;
     this.sessions.delete(sessionId);
 
-    // --- the inverse of setupSessionListeners(), in its order ---
-    const summaryTracker = this.runSummaryTrackers.get(sessionId);
-    if (summaryTracker) {
-      summaryTracker.stop();
-      this.runSummaryTrackers.delete(sessionId);
-    }
-    // An fs.watch on the workspace (or on @fix_plan.md) that nothing else closes.
-    session.ralphTracker.stopWatchingFixPlan();
-    // An FSWatcher on the workspace, likewise.
-    imageWatcher.unwatchSession(sessionId);
+    // --- the inverse of setupSessionListeners(), in reverse order ---
+    // Listeners first: while they are attached, one of them can still reach a
+    // tracker this is about to stop.
     const listeners = this.sessionListenerRefs.get(sessionId);
     if (listeners) {
       detachSessionListeners(session, listeners);
       this.sessionListenerRefs.delete(sessionId);
     }
+    // An FSWatcher on the workspace that nothing else closes.
+    imageWatcher.unwatchSession(sessionId);
+    // An fs.watch on the workspace (or on @fix_plan.md), likewise.
+    session.ralphTracker.stopWatchingFixPlan();
+    const summaryTracker = this.runSummaryTrackers.get(sessionId);
+    if (summaryTracker) {
+      summaryTracker.stop();
+      this.runSummaryTrackers.delete(sessionId);
+    }
+
+    // --- what anything else may have attached to this id in the meantime ---
+    // A rebuild can fail AFTER startInteractive() resolved, and a restored
+    // workspace still carries Codeman's hooks, so the CLI can post a hook event
+    // within milliseconds. Each of these outlives the listeners and would
+    // otherwise meet the retry, which reuses the same session id by design.
+    this.stopTranscriptWatcher(sessionId);
+    attachmentRegistry.clearSession(sessionId);
+    sessionWaits.notifySignal(sessionId, 'exit');
+    sessionWaits.cancelAll(sessionId);
+    approvalInbox.resolveForSession(sessionId, 'session_ended');
 
     // --- the inverse of the construction itself ---
     this.sse.cleanupSessionBatches(sessionId);
     this.persistDeb.cancelKey(sessionId);
     fileStreamManager.closeSessionStreams(sessionId);
+    // `lastRecordedTokens` is deliberately NOT deleted: the `after-spawn` phase
+    // seeds it as the daily-usage baseline for these restored totals, and the
+    // retry reuses the id, so dropping it would count them as new usage.
     // The per-session custom-model config dir carries the endpoint's API key, and
     // `before-spawn` may already have written it. Nothing else would ever remove
     // it: the stale sweep only touches state.json. A retry rewrites it.
@@ -3039,6 +3055,9 @@ export class WebServer extends EventEmitter {
       await session.stop(true);
     } catch (err) {
       console.warn(`[Server] stopping a partially built session failed: ${getErrorMessage(err)}`);
+      // `stop()` kills the mux session in its last block, after destroying its
+      // trackers, so a throw on the way there leaves the pane running.
+      await this.mux.killSession(sessionId).catch(() => {});
     }
     try {
       await this.tabLayouts.sessionsRemoved([{ id: sessionId, owner: session.owner }]);