@zgeoff/atc 2.15.1 → 2.15.2

package/src/daemon/drain-inbox.ts

@@ -0,0 +1,85 @@
+import { PROTOCOL_V } from '../protocol/protocol';
+import type { SessionID } from '../shared/session-id';
+import type { MessageOwner } from '../store/message-owner';
+import type { MessageRecord } from '../store/message-record';
+import type { TapClient } from './daemon-context';
+import type { TapGeneration, TapRegistry } from './tap-registry';
+
+/**
+ * What a drain reads from: the taps, the owner a linked tap reads messages
+ * under (the session's atc id and agent session id, or null for a session
+ * the daemon does not hold), and the store's pending messages for an owner.
+ */
+export interface InboxSource {
+  readonly taps: TapRegistry<TapClient>;
+  readonly findLinkedOwner: (sessionID: SessionID) => MessageOwner | null;
+  readonly collectPendingMessages: (owner: MessageOwner) => Promise<MessageRecord[]>;
+}
+
+/**
+ * Hands the session's tap the oldest pending message it has not taken. The
+ * backlog is read before anything is sent, so a tap's ok response is always
+ * queued ahead of its first message, and one message goes per call: the
+ * tap's ack calls it again, so the backlog never outgrows the connection's
+ * outbound queue. The drain is bound to the tap it found before the read: a
+ * tap another attach made while the read waited gets nothing from it.
+ */
+// oxlint-disable-next-line prefer-readonly-parameter-types -- the drain claims deliveries in the tap registry it is given
+export async function drainInbox(sessionID: SessionID, source: InboxSource): Promise<void> {
+  const linkedOwner = source.findLinkedOwner(sessionID);
+  const tap = source.taps.findTap(sessionID);
+
+  if (linkedOwner === null || tap === null) {
+    return;
+  }
+
+  // A tap a principal holds takes only the messages sent to this atc id:
+  // another session that shares the agent session id may be out of its
+  // reach.
+  const owner = tap.linked ? linkedOwner : { atcID: linkedOwner.atcID };
+
+  try {
+    const pending = await source.collectPendingMessages(owner);
+
+    for (const record of pending) {
+      if (sendInboxMessage(sessionID, record, tap, source.taps)) {
+        return;
+      }
+    }
+  } catch {}
+}
+
+// Hands one pending message to the tap the drain was made for, once, and
+// reports whether it did. An unlinked tap takes only a message sent to the
+// session's own atc id. The event goes to the tap connection alone: it never
+// reaches other clients, the events socket, or hooks.
+function sendInboxMessage(
+  sessionID: SessionID,
+  record: MessageRecord,
+  tap: TapGeneration,
+
+  // oxlint-disable-next-line prefer-readonly-parameter-types -- a delivery claims the message in the registry
+  taps: TapRegistry<TapClient>,
+): boolean {
+  if (!tap.linked && record.atcID !== sessionID) {
+    return false;
+  }
+
+  const client = taps.claimDelivery(sessionID, record.id, tap.generation);
+
+  if (client === null) {
+    return false;
+  }
+
+  client.sendEvent({
+    v: PROTOCOL_V,
+    ev: 'InboxMessage',
+    s: sessionID,
+    message: record.id,
+    from: record.from,
+    text: record.text,
+    sentAt: record.sentAt,
+  });
+
+  return true;
+}

package/src/daemon/is-tree-in-reach.ts

@@ -0,0 +1,32 @@
+import type { SessionID } from '../shared/session-id';
+import type { TargetAccess, TargetGrant } from './target-access';
+
+// A session as the tree check reads it: its id, its parent if it has one,
+// and the target it is bound to.
+interface TreeMember extends TargetGrant {
+  readonly id: SessionID;
+  readonly parent: SessionID | null;
+}
+
+/**
+ * Whether the access reaches every session in the tree the given session
+ * belongs to: its top-level session and each sub-session of that one. A
+ * session the list does not hold, or one whose parent it does not hold,
+ * is out of reach.
+ */
+export function isTreeInReach(
+  sessions: readonly TreeMember[],
+  id: SessionID,
+  access: TargetAccess,
+): boolean {
+  const session = sessions.find((x) => x.id === id);
+
+  if (session === undefined) {
+    return false;
+  }
+
+  const root = session.parent ?? session.id;
+  const tree = sessions.filter((x) => x.id === root || x.parent === root);
+
+  return tree.some((x) => x.id === root) && tree.every((x) => access.canUse(x));
+}

package/src/daemon/sessions.ts

@@ -479,8 +479,15 @@ export class SessionManager {
   // Adopts a headless session back into a terminal: a fresh PTY resumes the
   // same agent session id. On a remote host the host wakes first, and a
   // harness still running inside it is attached rather than started again;
-  // every other session left asleep on that host comes back with it.
-  async adoptTerminal(id: SessionID, cols: number, rows: number): Promise<Session | null> {
+  // every other session left asleep on that host comes back with it. The
+  // caller's check runs again after each await, and a failed check leaves
+  // the session as it was.
+  async adoptTerminal(
+    id: SessionID,
+    cols: number,
+    rows: number,
+    canProceed: () => boolean = () => true,
+  ): Promise<Session | null> {
     const s = this.sessions.find((x) => x.id === id);
 
     if (!s || s.pty !== null || s.agentSessionID === undefined || this.adopting.has(id)) {
@@ -489,7 +496,7 @@ export class SessionManager {
 
     const settled = await this.isRollbackSettled(id);
 
-    if (!settled) {
+    if (!settled || !canProceed()) {
       return null;
     }
 
@@ -516,9 +523,11 @@ export class SessionManager {
       this.adopting.delete(id);
     }
 
-    // A kill, a second adopt, or a failed spawn's rollback can land while
-    // the host wakes.
+    // A kill, a second adopt, a failed spawn's rollback, or a change that
+    // takes the session out of the caller's reach can land while the host
+    // wakes.
     if (
+      !canProceed() ||
       s.pty !== null ||
       !this.sessions.includes(s) ||
       this.rollingBack.has(id) ||
@@ -729,6 +738,7 @@ export class SessionManager {
     id: SessionID = mintSessionID(),
     target = 'local',
     materialized: MaterializedSpawn | null = null,
+    requireInReach: () => void = () => {},
   ): Promise<Session> {
     const adapter = this.findAdapter(agent);
 
@@ -751,6 +761,10 @@ export class SessionManager {
       ...overrides,
     });
 
+    // The caller's check runs again once the host is ready, before the
+    // harness starts.
+    requireInReach();
+
     const binding = this.mintBridgeBinding(id, target, execution.identity, hostKey);
 
     const pty = provider.spawnHarness({
@@ -1248,7 +1262,8 @@ export class SessionManager {
   // every harness on it kept inside, and fails whole when the host stays
   // awake. A dead session on a target that can destroy its host is not
   // forgotten by a kill: forgetting it destroys the host, which takes a
-  // confirmed forget.
+  // confirmed forget. The set is taken before the first await, so a
+  // sub-session spawned while the kill waits on a host is not part of it.
   async kill(id: SessionID): Promise<void> {
     const s = this.sessions.find((x) => x.id === id);
 
@@ -1256,11 +1271,15 @@ export class SessionManager {
       return;
     }
 
+    const children = this.collectChildren(id);
+
     if (s.pty) {
       await this.stopHarness(s);
 
-      for (const child of this.collectChildren(id)) {
-        await this.tryStopHarness(child);
+      for (const child of children) {
+        if (this.sessions.includes(child) && child.parent === id) {
+          await this.tryStopHarness(child);
+        }
       }
     } else {
       if (this.findProvider(s)?.capabilities.destroy === true) {
@@ -1271,8 +1290,7 @@ export class SessionManager {
         );
       }
 
-      this.updateForgottenChildren(id);
-      this.remove(s);
+      this.removeWithChildren(s, children);
     }
 
     await this.writeFleet();
@@ -1288,7 +1306,9 @@ export class SessionManager {
    * kept asleep in that host is refused until the host wakes. Its dead
    * sub-sessions on other hosts go with it, unless their own target can
    * destroy their host, and its live ones become top-level. A failed destroy
-   * throws before anything is forgotten.
+   * throws before anything is forgotten. The sub-sessions it may forget are
+   * taken before the first await; one spawned while the forget waits on a
+   * host becomes top-level.
    */
   async forget(id: SessionID): Promise<boolean> {
     const s = this.sessions.find((x) => x.id === id);
@@ -1297,6 +1317,7 @@ export class SessionManager {
       return false;
     }
 
+    const children = this.collectChildren(id);
     const provider = this.findProvider(s);
     const destroys = provider !== null && provider.capabilities.destroy && s.hostKey === s.id;
 
@@ -1330,8 +1351,7 @@ export class SessionManager {
       this.killTerminal(s);
     }
 
-    this.updateForgottenChildren(id);
-    this.remove(s);
+    this.removeWithChildren(s, children);
 
     await this.writeFleet();
 
@@ -1340,22 +1360,42 @@ export class SessionManager {
     return destroys;
   }
 
-  // A forgotten parent's dead sub-sessions go with it, except one whose own
-  // target can destroy its host: forgetting that one destroys the host, which
-  // takes its own confirmed forget. Every sub-session that stays becomes
-  // top-level.
-  private updateForgottenChildren(id: SessionID): void {
-    for (const child of this.collectChildren(id)) {
+  // Removes a forgotten session and its dead sub-sessions, except one whose
+  // own target can destroy its host: forgetting that one destroys the host,
+  // which takes its own confirmed forget. Only a sub-session among the given
+  // ones may go. Every sub-session that stays becomes top-level before the
+  // session's removal is announced, and is announced after it, so no
+  // announcement ever finds a sub-session whose parent is gone or a parent
+  // standing without the sub-sessions that leave with it.
+  private removeWithChildren(s: Session, forgettable: readonly Session[]): void {
+    const kept: Session[] = [];
+    const dropped: Session[] = [];
+
+    for (const child of this.collectChildren(s.id)) {
       const live = child.pty !== null || (child.kind === 'headless' && child.state !== 'exited');
 
-      if (live || this.findProvider(child)?.capabilities.destroy === true) {
+      if (
+        live ||
+        !forgettable.includes(child) ||
+        this.findProvider(child)?.capabilities.destroy === true
+      ) {
         child.parent = null;
 
-        this.onEvent('state', child);
+        kept.push(child);
       } else {
-        this.remove(child);
+        dropped.push(child);
       }
     }
+
+    this.remove(s);
+
+    for (const child of kept) {
+      this.onEvent('state', child);
+    }
+
+    for (const child of dropped) {
+      this.remove(child);
+    }
   }
 
   /**

package/src/daemon/tap-registry.ts

@@ -4,6 +4,23 @@ import type { SessionID } from '../shared/session-id';
 interface TapEntry<TClient> {
   readonly client: TClient;
   readonly claimed: Set<MessageID>;
+
+  // Whether the tap also takes messages held under the session's agent
+  // session id, as opposed to its atc id alone.
+  readonly linked: boolean;
+
+  // Which attach made this tap; every attach makes a new one.
+  readonly generation: number;
+}
+
+/**
+ * A session's tap as a drain finds it before it waits: the attach that made
+ * it and whether it is linked. A delivery made for it lands only while the
+ * same attach holds the tap.
+ */
+export interface TapGeneration {
+  readonly generation: number;
+  readonly linked: boolean;
 }
 
 /**
@@ -13,14 +30,24 @@ interface TapEntry<TClient> {
 export class TapRegistry<TClient> {
   private readonly bySession = new Map<SessionID, TapEntry<TClient>>();
 
+  private nextGeneration = 0;
+
   /**
    * Makes the client the session's tap and returns the client it displaced,
-   * or null when there was none or it is the same connection.
+   * or null when there was none or it is the same connection. A linked tap
+   * also takes the messages held under the session's agent session id.
    */
-  attach(sessionID: SessionID, client: TClient): TClient | null {
+  attach(sessionID: SessionID, client: TClient, linked = true): TClient | null {
     const previous = this.bySession.get(sessionID)?.client ?? null;
 
-    this.bySession.set(sessionID, { client, claimed: new Set() });
+    this.nextGeneration += 1;
+
+    this.bySession.set(sessionID, {
+      client,
+      claimed: new Set(),
+      linked,
+      generation: this.nextGeneration,
+    });
 
     return previous === client ? null : previous;
   }
@@ -45,6 +72,12 @@ export class TapRegistry<TClient> {
     return previous;
   }
 
+  findTap(sessionID: SessionID): TapGeneration | null {
+    const entry = this.bySession.get(sessionID);
+
+    return entry === undefined ? null : { generation: entry.generation, linked: entry.linked };
+  }
+
   hasTap(sessionID: SessionID): boolean {
     return this.bySession.has(sessionID);
   }
@@ -53,10 +86,15 @@ export class TapRegistry<TClient> {
     return this.bySession.get(sessionID)?.client === client;
   }
 
-  claimDelivery(sessionID: SessionID, messageID: MessageID): TClient | null {
+  /**
+   * Claims the message for the session's tap and returns its client, or
+   * null when the tap already took it or another attach replaced the tap
+   * the delivery was made for.
+   */
+  claimDelivery(sessionID: SessionID, messageID: MessageID, generation: number): TClient | null {
     const entry = this.bySession.get(sessionID);
 
-    if (entry === undefined || entry.claimed.has(messageID)) {
+    if (entry === undefined || entry.generation !== generation || entry.claimed.has(messageID)) {
       return null;
     }
 

package/src/store/run-migrations.ts

@@ -56,6 +56,8 @@ interface EventsTable {
 
 interface SpawnHistoryTable {
   cwd: string;
+  target: string;
+  target_identity: string;
   last_spawn: number;
 }
 
@@ -152,8 +154,9 @@ export interface StateStoreSchema {
 // agent_session_id under its Claude-era name plus name and cwd, and each
 // later step adds one column the daemon grew to depend on, until the
 // rebuild that keys the table by the atc session id. events later gains
-// columns of its own, as does messages, while spawn_history and prefs have
-// carried one shape since they were added.
+// columns of its own, as does messages, spawn_history is rebuilt to key each
+// directory by the target it was spawned on, and prefs has carried one shape
+// since it was added.
 const MIGRATIONS: Record<string, Migration> = {
   '001_create_initial_schema': {
     async up(db: Kysely<StateStoreSchema>) {
@@ -446,6 +449,29 @@ const MIGRATIONS: Record<string, Migration> = {
       await db.schema.alterTable('events').addColumn('report_text', 'text').execute();
     },
   },
+  '024_rebuild_spawn_history_keyed_by_target': {
+    async up(db: Kysely<StateStoreSchema>) {
+      // The ledger records this step only after it returns, so a crash in
+      // between leaves a rebuilt table behind a ledger that still lacks it.
+      const columns = await sql<ColumnInfoRow>`PRAGMA table_info(spawn_history)`.execute(db);
+
+      if (columns.rows.some((column) => column.name === 'target')) {
+        return;
+      }
+
+      await sql`BEGIN IMMEDIATE`.execute(db);
+
+      try {
+        await updateSpawnHistoryKeyToTarget(db);
+
+        await sql`COMMIT`.execute(db);
+      } catch (error) {
+        await sql`ROLLBACK`.execute(db);
+
+        throw error;
+      }
+    },
+  },
 };
 
 const PROVIDER: MigrationProvider = {
@@ -457,8 +483,9 @@ const PROVIDER: MigrationProvider = {
  * `Migrator`, one additive step at a time. A database from before the
  * ladder existed is recognized at whichever shape it stopped at by a
  * baselining pass that records the steps its columns already satisfy, so
- * only what is genuinely missing runs. One step rebuilds the fleet table
- * under a new key and carries every row across; no step drops a row.
+ * only what is genuinely missing runs. Two steps rebuild a table under a
+ * new key, the fleet and the spawn history, and carry every row across; no
+ * step drops a row.
  */
 export async function runMigrations(db: Kysely<StateStoreSchema>): Promise<void> {
   // A baselined step can land out of order relative to a step that turned
@@ -652,6 +679,30 @@ interface ColumnInfoRow {
   type: string;
 }
 
+// The identity of the local target with no options. Every spawn from before
+// spawns recorded their target ran there.
+const LEGACY_SPAWN_TARGET_IDENTITY = 'local-pty:44136fa355b3678a';
+
+async function updateSpawnHistoryKeyToTarget(db: Kysely<StateStoreSchema>): Promise<void> {
+  await sql`
+    CREATE TABLE spawn_history_rebuilt (
+      cwd TEXT NOT NULL,
+      target TEXT NOT NULL,
+      target_identity TEXT NOT NULL,
+      last_spawn INTEGER NOT NULL,
+      PRIMARY KEY (cwd, target, target_identity)
+    )
+  `.execute(db);
+
+  await sql`
+    INSERT INTO spawn_history_rebuilt (cwd, target, target_identity, last_spawn)
+    SELECT cwd, 'local', ${LEGACY_SPAWN_TARGET_IDENTITY}, last_spawn FROM spawn_history
+  `.execute(db);
+
+  await sql`DROP TABLE spawn_history`.execute(db);
+  await sql`ALTER TABLE spawn_history_rebuilt RENAME TO spawn_history`.execute(db);
+}
+
 async function collectFleetColumns(db: Kysely<StateStoreSchema>): Promise<ReadonlySet<string>> {
   const result = await sql<ColumnInfoRow>`PRAGMA table_info(fleet)`.execute(db);
 

package/src/store/state-store.ts

@@ -64,6 +64,27 @@ export interface StoredEvent {
   readonly label?: string;
 }
 
+// The target a spawn ran on: its name and the identity the spawned session
+// was bound to there.
+interface SpawnTarget {
+  readonly target: string;
+  readonly targetIdentity: string;
+}
+
+// A directory a spawn ran in and the target it ran on.
+export interface SpawnDir {
+  readonly cwd: string;
+  readonly grant: SpawnTarget;
+}
+
+/**
+ * Another message one turn answered, and the atc id it was sent to.
+ */
+export interface TurnSibling {
+  readonly id: MessageID;
+  readonly atcID: SessionID;
+}
+
 /**
  * One report as the trail holds it. A report recorded before the trail kept
  * whole texts has only its preview, so its text is that preview and
@@ -513,24 +534,38 @@ export class StateStore {
     return row?.ts === null || row?.ts === undefined ? null : Date.parse(row.ts);
   }
 
-  async recordSpawnDir(cwd: string): Promise<void> {
+  async recordSpawnDir(cwd: string, grant: SpawnTarget): Promise<void> {
     await this.db
       .insertInto('spawn_history')
-      .values({ cwd, last_spawn: Date.now() })
+      .values({
+        cwd,
+        target: grant.target,
+        target_identity: grant.targetIdentity,
+        last_spawn: Date.now(),
+      })
       .onConflict((oc) =>
-        oc.column('cwd').doUpdateSet((eb) => ({ last_spawn: eb.ref('excluded.last_spawn') })),
+        oc
+          .columns(['cwd', 'target', 'target_identity'])
+          .doUpdateSet((eb) => ({ last_spawn: eb.ref('excluded.last_spawn') })),
       )
       .execute();
   }
 
-  async collectSpawnDirs(): Promise<string[]> {
+  /**
+   * Each directory a spawn ran in, with the target it ran on, most recent
+   * first. A directory spawned on several targets has one entry for each.
+   */
+  async collectSpawnDirs(): Promise<SpawnDir[]> {
     const rows = await this.db
       .selectFrom('spawn_history')
-      .select('cwd')
+      .select(['cwd', 'target', 'target_identity'])
       .orderBy('last_spawn', 'desc')
       .execute();
 
-    return rows.map((row) => row.cwd);
+    return rows.map((row) => ({
+      cwd: row.cwd,
+      grant: { target: row.target, targetIdentity: row.target_identity },
+    }));
   }
 
   async loadLastUsedAgent(): Promise<AgentID> {
@@ -607,7 +642,7 @@ export class StateStore {
 
   // The other messages of the same session the given message's turn
   // answered, oldest first; none when the message has no turn.
-  async collectTurnSiblings(record: MessageRecord): Promise<MessageID[]> {
+  async collectTurnSiblings(record: MessageRecord): Promise<TurnSibling[]> {
     if (record.turn === undefined) {
       return [];
     }
@@ -619,7 +654,7 @@ export class StateStore {
 
     const rows = await this.db
       .selectFrom('messages')
-      .select('id')
+      .select(['id', 'atc_id'])
       .where('turn_id', '=', record.turn)
       .where('id', '!=', record.id)
       .where((eb) => buildOwnerFilter(eb, owner))
@@ -627,7 +662,7 @@ export class StateStore {
       .orderBy(sql`rowid`, 'asc')
       .execute();
 
-    return rows.map((row) => toMessageID(row.id));
+    return rows.map((row) => ({ id: toMessageID(row.id), atcID: toSessionID(row.atc_id) }));
   }
 
   async updateMessageDelivered(