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>

test/discard-partially-built-session.test.ts

@@ -18,10 +18,11 @@
  * leaves that entry makes the next attempt wire nothing at all, and the user
  * gets a tab that never shows output.
  */
-import { mkdirSync, rmSync } from 'node:fs';
+import { mkdirSync } from 'node:fs';
 import { homedir } from 'node:os';
 import { join } from 'node:path';
-import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest';
+import { safeRmHomeTree } from './mocks/test-helpers.js';
 
 import { WebServer } from '../src/web/server.js';
 import { Session } from '../src/session.js';
@@ -55,9 +56,11 @@ function buildSession(): Session {
   });
 }
 
-beforeEach(() => {
+beforeAll(() => {
   mkdirSync(WORKSPACE, { recursive: true });
-  // Test mode: no port is opened and no CLI is launched.
+  // Test mode: no port is opened and no CLI is launched. One server for the file,
+  // stopped at the end: the constructor registers handlers on the module-level
+  // image, subagent, team and workflow watchers, and only stop() removes them.
   server = new WebServer(0, false, true);
   internals = server as unknown as ServerInternals;
   mux = new TmuxManager();
@@ -65,7 +68,11 @@ beforeEach(() => {
 
 afterEach(async () => {
   await internals.discardPartiallyBuiltSession(SESSION_ID).catch(() => {});
-  rmSync(WORKSPACE, { recursive: true, force: true });
+});
+
+afterAll(async () => {
+  await server.stop().catch(() => {});
+  safeRmHomeTree(WORKSPACE);
 });
 
 describe('discarding a session whose pane never started', () => {
@@ -85,25 +92,36 @@ describe('discarding a session whose pane never started', () => {
     await internals.setupSessionListeners(first);
     expect(internals.sessionListenerRefs.has(SESSION_ID)).toBe(true);
 
+    const firstRefs = internals.sessionListenerRefs.get(SESSION_ID);
     await internals.discardPartiallyBuiltSession(SESSION_ID);
     expect(internals.sessionListenerRefs.has(SESSION_ID)).toBe(false);
 
     // The retry reuses the id by design. `setupSessionListeners()` returns early
-    // while the refs are still there, so a session built now would run blind:
-    // no terminal output, no status updates, no exit broadcast.
+    // while the refs are still there, so a session built now would run blind: no
+    // terminal output, no status updates, no exit broadcast. Asserting a DIFFERENT
+    // refs object is what distinguishes wiring the retry from finding the corpse
+    // of the first attempt still in place.
     const retry = buildSession();
     await internals.registerSessionWithLayout(retry);
     await internals.setupSessionListeners(retry);
-    expect(internals.sessionListenerRefs.has(SESSION_ID)).toBe(true);
+    const retryRefs = internals.sessionListenerRefs.get(SESSION_ID);
+    expect(retryRefs).toBeDefined();
+    expect(retryRefs).not.toBe(firstRefs);
   });
 
   it('stops the run-summary tracker, whose interval would otherwise keep firing', async () => {
     const session = buildSession();
     await internals.registerSessionWithLayout(session);
     await internals.setupSessionListeners(session);
-    expect(internals.runSummaryTrackers.has(SESSION_ID)).toBe(true);
+    const tracker = internals.runSummaryTrackers.get(SESSION_ID) as { stop: () => void };
+    expect(tracker).toBeDefined();
+    // Dropping the map entry is not enough: the tracker arms a setInterval in its
+    // constructor, and only stop() clears it, so a discard that merely forgot the
+    // entry would leave the timer running for the life of the process.
+    const stopped = vi.spyOn(tracker, 'stop');
 
     await internals.discardPartiallyBuiltSession(SESSION_ID);
+    expect(stopped).toHaveBeenCalled();
     expect(internals.runSummaryTrackers.has(SESSION_ID)).toBe(false);
   });
 

test/reboot-restore.test.ts

@@ -198,15 +198,15 @@ describe('the plan the banner spends', () => {
   it('hands an entry to the first caller and nothing to the second', () => {
     const registry = new RebootRestoreRegistry();
     registry.set([entryFor('a'), entryFor('b')]);
-    expect(registry.take(all).map((e) => e.sessionId)).toEqual(['a', 'b']);
+    expect(registry.take(all, undefined, undefined).map((e) => e.sessionId)).toEqual(['a', 'b']);
     // The double-click: two panes on one conversation is what this prevents.
-    expect(registry.take(all)).toEqual([]);
+    expect(registry.take(all, undefined, undefined)).toEqual([]);
   });
 
   it('spends only the ids a caller asked for', () => {
     const registry = new RebootRestoreRegistry();
     registry.set([entryFor('a'), entryFor('b')]);
-    expect(registry.take(all, ['b']).map((e) => e.sessionId)).toEqual(['b']);
+    expect(registry.take(all, ['b'], undefined).map((e) => e.sessionId)).toEqual(['b']);
     expect(registry.list(all).map((e) => e.sessionId)).toEqual(['a']);
   });
 
@@ -215,7 +215,7 @@ describe('the plan the banner spends', () => {
     registry.set([entryFor('mine', 'alice'), entryFor('theirs', 'bob')]);
     const asAlice = (owner: string | undefined) => owner === 'alice';
     expect(registry.list(asAlice).map((e) => e.sessionId)).toEqual(['mine']);
-    expect(registry.take(asAlice).map((e) => e.sessionId)).toEqual(['mine']);
+    expect(registry.take(asAlice, undefined, 'alice').map((e) => e.sessionId)).toEqual(['mine']);
     // Bob's entry is still on offer for Bob.
     expect(registry.list(() => true).map((e) => e.sessionId)).toEqual(['theirs']);
   });
@@ -223,8 +223,8 @@ describe('the plan the banner spends', () => {
   it('puts back an entry that no pane was created for', () => {
     const registry = new RebootRestoreRegistry();
     registry.set([entryFor('a')]);
-    const taken = registry.take(all);
-    registry.restore(taken);
+    const taken = registry.take(all, undefined, undefined);
+    registry.releaseFlight(undefined, taken);
     expect(registry.list(all).map((e) => e.sessionId)).toEqual(['a']);
   });
 

test/routes/reboot-restore-rebuild-failure.test.ts

@@ -183,12 +183,23 @@ describe('a rebuild that succeeds', () => {
         callOrder.push(`reapply:${phase}`);
       }
     );
+    (ctx.setupSessionListeners as ReturnType<typeof vi.fn>).mockImplementation(async () => {
+      callOrder.push('setupSessionListeners');
+    });
     const app = await createHarness(ctx);
 
     await app.inject({ method: 'POST', url: '/api/reboot-restore/restore', payload: {} });
-    // The custom-model environment has to reach the process; the token totals
-    // must not land on a session whose pane never started.
-    expect(callOrder).toEqual(['reapply:before-spawn', 'startInteractive', 'reapply:after-spawn']);
+    // `setupSessionListeners()` READS the image-watcher flag that `before-spawn`
+    // restores, so the phase has to precede it or the session comes back
+    // reporting the watcher as on with nothing watching. The custom-model
+    // environment has to reach the process, and the token totals must not land
+    // on a session whose pane never started.
+    expect(callOrder).toEqual([
+      'reapply:before-spawn',
+      'setupSessionListeners',
+      'startInteractive',
+      'reapply:after-spawn',
+    ]);
     await app.close();
   });
 
@@ -282,6 +293,32 @@ describe('a failure before any entry is considered', () => {
 });
 
 describe('a dismiss that lands while a restore is running', () => {
+  it('wins when an admin is restoring the entries and their owner dismisses', async () => {
+    const theirs = offerEntry('theirs', 'bob');
+    rebootRestoreRegistry.set([theirs]);
+    // An admin may spend another user's entries, so the caller doing the restore
+    // and the owner of what is being restored are different people.
+    const taken = rebootRestoreRegistry.take(() => true, undefined, 'admin');
+    expect(taken.map((e) => e.sessionId)).toEqual(['theirs']);
+
+    // Bob dismisses his own banner. Nothing of his is in the plan any more, and
+    // the restore is running under a different name than his.
+    rebootRestoreRegistry.clear((owner) => owner === 'bob');
+    rebootRestoreRegistry.releaseFlight('admin', taken);
+
+    expect(rebootRestoreRegistry.list(() => true)).toEqual([]);
+  });
+
+  it('wins when an admin dismisses everything mid-restore', async () => {
+    rebootRestoreRegistry.set([offerEntry('theirs', 'bob')]);
+    const taken = rebootRestoreRegistry.take(() => true, undefined, 'admin');
+
+    rebootRestoreRegistry.clear(() => true);
+    rebootRestoreRegistry.releaseFlight('admin', taken);
+
+    expect(rebootRestoreRegistry.list(() => true)).toEqual([]);
+  });
+
   it('wins, rather than being undone when the route hands its entries back', async () => {
     rebootRestoreRegistry.set([offerEntry('a')]);
     const ctx = createMockRouteContext({ workspaceHooksEnabled: false });
@@ -304,15 +341,13 @@ describe('a dismiss that lands while a restore is running', () => {
   it('reaches an in-flight restore the dismisser can see, even once its entries are taken', async () => {
     const mine = offerEntry('mine', 'alice');
     rebootRestoreRegistry.set([mine]);
-    expect(rebootRestoreRegistry.beginSpending('alice')).toBe(true);
-    const generations = rebootRestoreRegistry.snapshotGenerations([mine]);
-    const taken = rebootRestoreRegistry.take((owner) => owner === 'alice');
+    const taken = rebootRestoreRegistry.take((owner) => owner === 'alice', undefined, 'alice');
+    expect(taken).toHaveLength(1);
 
-    // The plan is empty now, so a dismiss has nothing of Alice's to remove; the
-    // invalidation has to come from her claimed flight.
+    // The plan is empty now, so the dismiss has nothing of Alice's left in the
+    // plan; it has to reach the entry the restore is holding.
     rebootRestoreRegistry.clear((owner) => owner === 'alice');
-    rebootRestoreRegistry.restore(taken, generations);
-    rebootRestoreRegistry.endSpending('alice');
+    rebootRestoreRegistry.releaseFlight('alice', taken);
 
     expect(rebootRestoreRegistry.list(() => true)).toEqual([]);
   });
@@ -322,13 +357,8 @@ describe('a dismiss that lands while a restore is running', () => {
     const theirs = offerEntry('theirs', 'bob');
     rebootRestoreRegistry.set([mine, theirs]);
 
-    // Bob is mid-restore, holding his own entry. The claimed flight is what makes
-    // this the interesting case: a dismiss can no longer see Bob's entries in the
-    // plan, so the invalidation has to come from the in-flight set, filtered by
-    // what the dismissing user may access.
-    expect(rebootRestoreRegistry.beginSpending('bob')).toBe(true);
-    const bobsGenerations = rebootRestoreRegistry.snapshotGenerations([theirs]);
-    const bobsTaken = rebootRestoreRegistry.take((owner) => owner === 'bob');
+    // Bob is mid-restore, holding his own entry.
+    const bobsTaken = rebootRestoreRegistry.take((owner) => owner === 'bob', undefined, 'bob');
     expect(bobsTaken.map((e) => e.sessionId)).toEqual(['theirs']);
 
     // Alice dismisses her own banner meanwhile.
@@ -336,8 +366,7 @@ describe('a dismiss that lands while a restore is running', () => {
 
     // Bob's restore finishes and hands his entry back. Alice's dismiss covered
     // her entries, not his, so his offer survives.
-    rebootRestoreRegistry.restore(bobsTaken, bobsGenerations);
-    rebootRestoreRegistry.endSpending('bob');
+    rebootRestoreRegistry.releaseFlight('bob', bobsTaken);
     expect(rebootRestoreRegistry.list(() => true).map((e) => e.sessionId)).toEqual(['theirs']);
   });
 });