@zgeoff/atc 2.11.0 → 2.13.0

package/src/store/state-store.ts

@@ -1,36 +1,47 @@
 import { Database } from 'bun:sqlite';
+import { randomUUID } from 'node:crypto';
 import { existsSync, readFileSync } from 'node:fs';
 import { Kysely, SqliteAdapter, SqliteIntrospector, SqliteQueryCompiler, sql } from 'kysely';
-import type { Expression, ExpressionBuilder, SqlBool } from 'kysely';
+import type { Expression, ExpressionBuilder, SqlBool, Transaction } from 'kysely';
 import { toAgentID } from '../agents/agent-adapter';
 import type { AdapterEvent, AgentID } from '../agents/agent-adapter';
 import type { HookEvent } from '../daemon/hooks';
+import { DaemonError } from '../protocol/daemon-error';
 import type { AgentSessionID } from '../shared/agent-session-id';
+import type { DaemonID } from '../shared/daemon-id';
 import type { MessageID } from '../shared/message-id';
 import type { SessionID } from '../shared/session-id';
 import { toAgentSessionID } from '../shared/to-agent-session-id';
+import { toDaemonID } from '../shared/to-daemon-id';
 import { toMessageID } from '../shared/to-message-id';
 import { toSessionID } from '../shared/to-session-id';
 import { BunSqliteDriver } from './bun-sqlite-driver';
 import { parseFleetEntry } from './fleet-entry';
-import type { FleetEntry, FleetEntryUpdate } from './fleet-entry';
+import type { FleetEntry, FleetEntryUpdate, LegacyFleetEntry } from './fleet-entry';
+import type { EffectTarget, IdempotencyClaim, IdempotencyRecord } from './idempotency-record';
 import type { MessageOwner } from './message-owner';
 import type { MessageRecord } from './message-record';
 import { runMigrations } from './run-migrations';
 import type { StateStoreSchema } from './run-migrations';
 import type { TrailEntry } from './trail-entry';
+import type {
+  MaterializationUpdate,
+  SessionWorkspace,
+  WorkspaceMaterialization,
+} from './workspace-materialization';
 
 // Spelled as the partial index's predicate so SQLite can match them.
 const TRAIL_FILTER = sql<boolean>`kind IS NOT NULL AND kind != 'heartbeat'`;
 
 /**
- * One session's slice of the event trail: rows under its atc id, plus rows
- * under its agent session id, since a restore re-mints the atc id while the
- * agent session id carries on.
+ * Some sessions' slice of the event trail: rows under their atc ids, plus
+ * rows under their agent session ids, since rows written before atc session
+ * ids stayed stable across restores carry an earlier atc id. A scope of no
+ * sessions matches no row.
  */
 export interface EventScope {
-  readonly atcID: SessionID;
-  readonly agentSessionID: AgentSessionID | undefined;
+  readonly atcIDs: readonly SessionID[];
+  readonly agentSessionIDs: readonly AgentSessionID[];
 }
 
 export interface StoredEvent {
@@ -70,9 +81,14 @@ export class StateStore {
 
   private readonly db: Kysely<StateStoreSchema>;
 
-  private constructor(sqlite: Database, db: Kysely<StateStoreSchema>) {
+  // The daemon this store's fleet rows belong to, minted once by the
+  // migration that created session ownership.
+  readonly daemonID: DaemonID;
+
+  private constructor(sqlite: Database, db: Kysely<StateStoreSchema>, daemonID: DaemonID) {
     this.sqlite = sqlite;
     this.db = db;
+    this.daemonID = daemonID;
   }
 
   // Migrations run statements that cannot happen inside a constructor, so
@@ -93,7 +109,13 @@ export class StateStore {
 
     await runMigrations(db);
 
-    const store = new StateStore(sqlite, db);
+    const daemonIDRow = await db
+      .selectFrom('prefs')
+      .select('value')
+      .where('key', '=', 'daemon_id')
+      .executeTakeFirstOrThrow();
+
+    const store = new StateStore(sqlite, db, toDaemonID(daemonIDRow.value));
 
     if (legacyFleetPath !== undefined) {
       await store.adoptLegacyFleet(legacyFleetPath);
@@ -102,21 +124,48 @@ export class StateStore {
     return store;
   }
 
+  // Only the rows this daemon owns: another daemon's sessions are its own
+  // to restore.
   async loadFleet(): Promise<FleetEntry[]> {
     const rows = await this.db
       .selectFrom('fleet')
+
+      // A session's ready workspace comes back with its row.
+      .leftJoin('workspace_materialization as workspace', (join) =>
+        join
+          .onRef('workspace.session_id', '=', 'fleet.session_id')
+          .on('workspace.phase', '=', 'ready'),
+      )
+      .where('fleet.session_id', 'in', (eb) =>
+        eb.selectFrom('session_owner').select('session_id').where('daemon_id', '=', this.daemonID),
+      )
+
+      // Stored order, which a restore keeps for sessions with no recency.
+      .orderBy(sql`fleet.rowid`)
       .select([
-        'agent_session_id',
-        'name',
-        'cwd',
-        'pinned',
-        'last_attached',
-        'agent',
-        'exited',
-        'parent',
-        'prompt',
-        'result',
-        'transcript_path',
+        'fleet.session_id',
+        'fleet.agent_session_id',
+        'fleet.name',
+        'fleet.cwd',
+        'fleet.pinned',
+        'fleet.last_attached',
+        'fleet.agent',
+        'fleet.exited',
+        'fleet.parent_session_id',
+        'fleet.prompt',
+        'fleet.result',
+        'fleet.transcript_path',
+        'fleet.model',
+        'fleet.effort',
+        'fleet.target',
+        'fleet.target_identity',
+        'workspace.repo_url',
+        'workspace.sha',
+        'workspace.ref',
+        'workspace.materialized_at',
+        'workspace.withheld_env',
+        'fleet.desired',
+        'fleet.host_key',
       ])
       .execute();
 
@@ -124,17 +173,28 @@ export class StateStore {
 
     for (const row of rows) {
       entries.push({
-        agentSessionID: toAgentSessionID(row.agent_session_id),
+        sessionID: toSessionID(row.session_id),
         name: row.name,
         cwd: row.cwd,
+        ...(row.agent_session_id === null
+          ? {}
+          : { agentSessionID: toAgentSessionID(row.agent_session_id) }),
         agent: toAgentID(row.agent),
         ...(row.pinned === 0 ? {} : { pinned: true }),
         ...(row.last_attached === null ? {} : { lastAttachedAt: row.last_attached }),
         ...(row.exited === 0 ? {} : { exited: true }),
-        ...(row.parent === null ? {} : { parent: toAgentSessionID(row.parent) }),
+        ...(row.parent_session_id === null ? {} : { parent: toSessionID(row.parent_session_id) }),
         ...(row.prompt === null ? {} : { prompt: row.prompt }),
         ...(row.result === null ? {} : { result: row.result }),
         ...(row.transcript_path === null ? {} : { transcriptPath: row.transcript_path }),
+        ...(row.model === null ? {} : { model: row.model }),
+        ...(row.effort === null ? {} : { effort: row.effort }),
+        ...(row.target === null ? {} : { target: row.target }),
+        ...(row.target_identity === null ? {} : { targetIdentity: row.target_identity }),
+        ...buildWorkspaceField(row),
+        ...buildWithheldEnvField(row.withheld_env),
+        ...(row.desired === 'sleep' || row.desired === 'stop' ? { desired: row.desired } : {}),
+        ...(row.host_key === null ? {} : { hostKey: toSessionID(row.host_key) }),
       });
     }
 
@@ -164,27 +224,67 @@ export class StateStore {
     return recency;
   }
 
+  // Replaces this daemon's rows wholesale and leaves every other daemon's
+  // rows as they stand. A session another daemon owns, or one whose stored
+  // ownership epoch has moved past this daemon's, rejects the whole write
+  // with stale_epoch.
   async writeFleet(entries: readonly FleetEntry[]): Promise<void> {
-    await this.db.transaction().execute(async (trx) => {
-      await trx.deleteFrom('fleet').execute();
+    const kept = buildFleetWithoutReplacedRows(entries);
 
-      for (const entry of entries) {
+    await this.db.transaction().execute(async (trx) => {
+      await requireCurrentOwnership(
+        trx,
+        this.daemonID,
+        entries.map((entry) => entry.sessionID),
+      );
+
+      await trx
+        .deleteFrom('fleet')
+        .where('session_id', 'in', (eb) =>
+          eb
+            .selectFrom('session_owner')
+            .select('session_id')
+            .where('daemon_id', '=', this.daemonID),
+        )
+        .execute();
+
+      await trx.deleteFrom('session_owner').where('daemon_id', '=', this.daemonID).execute();
+
+      const updatedAt = Date.now();
+
+      for (const entry of kept) {
         await trx
           .insertInto('fleet')
           .values({
-            agent_session_id: entry.agentSessionID,
+            session_id: entry.sessionID,
+            agent_session_id: entry.agentSessionID ?? null,
             name: entry.name,
             cwd: entry.cwd,
             pinned: entry.pinned === true ? 1 : 0,
             last_attached: entry.lastAttachedAt ?? null,
             agent: entry.agent,
             exited: entry.exited === true ? 1 : 0,
-            parent: entry.parent ?? null,
+            parent_session_id: entry.parent ?? null,
             prompt: entry.prompt ?? null,
             result: entry.result ?? null,
             transcript_path: entry.transcriptPath ?? null,
+            model: entry.model ?? null,
+            effort: entry.effort ?? null,
+            target: entry.target ?? null,
+            target_identity: entry.targetIdentity ?? null,
+            desired: entry.desired ?? null,
+            host_key: entry.hostKey ?? null,
+          })
+          .execute();
+
+        await trx
+          .insertInto('session_owner')
+          .values({
+            session_id: entry.sessionID,
+            daemon_id: this.daemonID,
+            owner_epoch: OWNER_EPOCH,
+            updated_at: updatedAt,
           })
-          .orReplace()
           .execute();
       }
     });
@@ -193,7 +293,7 @@ export class StateStore {
   // Rewrites fields on one existing row and leaves every other row as it
   // stands. A session without a row yet is a no-op: the next fleet write
   // inserts it with these fields.
-  async updateFleetEntry(agentSessionID: AgentSessionID, fields: FleetEntryUpdate): Promise<void> {
+  async updateFleetEntry(sessionID: SessionID, fields: FleetEntryUpdate): Promise<void> {
     const values = {
       ...(fields.result === undefined ? {} : { result: fields.result }),
       ...(fields.transcriptPath === undefined ? {} : { transcript_path: fields.transcriptPath }),
@@ -203,11 +303,11 @@ export class StateStore {
       return;
     }
 
-    await this.db
-      .updateTable('fleet')
-      .set(values)
-      .where('agent_session_id', '=', agentSessionID)
-      .execute();
+    await this.db.transaction().execute(async (trx) => {
+      await requireCurrentOwnership(trx, this.daemonID, [sessionID]);
+
+      await trx.updateTable('fleet').set(values).where('session_id', '=', sessionID).execute();
+    });
   }
 
   async recordEvent(e: HookEvent, ev: Readonly<AdapterEvent> | null = null): Promise<void> {
@@ -231,8 +331,10 @@ export class StateStore {
       .execute();
   }
 
-  async recordTrailEntry(entry: TrailEntry): Promise<void> {
-    await this.db
+  // Returns whether the entry was written: a report whose id the trail
+  // already holds is left out.
+  async recordTrailEntry(entry: TrailEntry): Promise<boolean> {
+    const result = await this.db
       .insertInto('events')
       .values({
         ts: new Date(entry.at).toISOString(),
@@ -244,8 +346,12 @@ export class StateStore {
         session_id: entry.agentSessionID,
         kind: entry.kind,
         detail: entry.detail,
+        report_id: entry.kind === 'report' ? (entry.reportID ?? null) : null,
       })
-      .execute();
+      .onConflict((oc) => oc.column('report_id').doNothing())
+      .executeTakeFirst();
+
+    return result.numInsertedOrUpdatedRows !== 0n;
   }
 
   async collectEventsAfter(
@@ -282,8 +388,8 @@ export class StateStore {
     return buildStoredEvents(rows).toReversed();
   }
 
-  // Matches the session by either id, since a restore re-mints the atc id
-  // while the agent session id carries on.
+  // Matches the session by either id, since rows written before atc session
+  // ids stayed stable across restores carry an earlier atc id.
   async loadLastActivityAt(
     atcID: SessionID,
     agentSessionID: AgentSessionID | undefined,
@@ -496,6 +602,215 @@ export class StateStore {
       .execute();
   }
 
+  // Records a key in progress unless one is already held under the same
+  // principal and operation. Returns null when this call claimed the key,
+  // else the record the key already holds.
+  async claimIdempotencyKey(claim: IdempotencyClaim): Promise<IdempotencyRecord | null> {
+    const inserted = await this.db
+      .insertInto('idempotency')
+      .values({
+        principal: claim.principal,
+        operation: claim.operation,
+        key: claim.key,
+        payload_hash: claim.payloadHash,
+        state: 'in_progress',
+        effect_ref: claim.effectRef,
+        result: null,
+        created_at: claim.at,
+        updated_at: claim.at,
+        effect_target: null,
+        effect_target_identity: null,
+      })
+      .onConflict((oc) => oc.columns(['principal', 'operation', 'key']).doNothing())
+      .returning('key')
+      .executeTakeFirst();
+
+    if (inserted !== undefined) {
+      return null;
+    }
+
+    const row = await this.db
+      .selectFrom('idempotency')
+      .selectAll()
+      .where('principal', '=', claim.principal)
+      .where('operation', '=', claim.operation)
+      .where('key', '=', claim.key)
+      .executeTakeFirstOrThrow();
+
+    return toIdempotencyRecord(row);
+  }
+
+  async updateIdempotencyCompleted(
+    record: Pick<IdempotencyRecord, 'principal' | 'operation' | 'key'>,
+    result: string,
+    at: number,
+    effectTarget: EffectTarget | null = null,
+  ): Promise<void> {
+    await this.db
+      .updateTable('idempotency')
+      .set({
+        state: 'completed',
+        result,
+        updated_at: at,
+        effect_target: effectTarget?.target ?? null,
+        effect_target_identity: effectTarget?.targetIdentity ?? null,
+      })
+      .where('principal', '=', record.principal)
+      .where('operation', '=', record.operation)
+      .where('key', '=', record.key)
+      .execute();
+  }
+
+  // Keeps a claim whose effect may still stand as outcome_unknown, so a
+  // retry answers from it and never runs the effect again.
+  async updateIdempotencyOutcomeUnknown(
+    record: Pick<IdempotencyRecord, 'principal' | 'operation' | 'key'>,
+    at: number,
+  ): Promise<void> {
+    await this.db
+      .updateTable('idempotency')
+      .set({ state: 'outcome_unknown', updated_at: at })
+      .where('principal', '=', record.principal)
+      .where('operation', '=', record.operation)
+      .where('key', '=', record.key)
+      .execute();
+  }
+
+  // Drops a claim whose effect never started, so a retry runs it fresh.
+  async removeIdempotencyKey(
+    record: Pick<IdempotencyRecord, 'principal' | 'operation' | 'key'>,
+  ): Promise<void> {
+    await this.db
+      .deleteFrom('idempotency')
+      .where('principal', '=', record.principal)
+      .where('operation', '=', record.operation)
+      .where('key', '=', record.key)
+      .execute();
+  }
+
+  // Runs once as a daemon starts, before it serves a request: every key
+  // still in progress belonged to a daemon that stopped mid-effect, so its
+  // outcome is unknown. A key whose effect left its row behind completes;
+  // one whose effect left no trace stays unknown.
+  async reconcileIdempotencyKeys(at: number): Promise<void> {
+    await this.db.transaction().execute(async (trx) => {
+      await trx
+        .updateTable('idempotency')
+        .set({ state: 'outcome_unknown', updated_at: at })
+        .where('state', '=', 'in_progress')
+        .execute();
+
+      // A spawn that completes here records the target its fleet row holds,
+      // as a spawn that completes as it runs does.
+      await trx
+        .updateTable('idempotency')
+        .set((eb) => ({
+          state: 'completed',
+          updated_at: at,
+          effect_target: eb
+            .selectFrom('fleet')
+            .select('fleet.target')
+            .whereRef('fleet.session_id', '=', 'idempotency.effect_ref'),
+          effect_target_identity: eb
+            .selectFrom('fleet')
+            .select('fleet.target_identity')
+            .whereRef('fleet.session_id', '=', 'idempotency.effect_ref'),
+        }))
+        .where('state', '=', 'outcome_unknown')
+        .where('operation', '=', 'session.spawn')
+        .where('effect_ref', 'in', (eb) => eb.selectFrom('fleet').select('session_id'))
+        .execute();
+
+      await trx
+        .updateTable('idempotency')
+        .set({ state: 'completed', updated_at: at })
+        .where('state', '=', 'outcome_unknown')
+        .where('operation', '=', 'session.message')
+        .where('effect_ref', 'in', (eb) => eb.selectFrom('messages').select('id'))
+        .execute();
+    });
+  }
+
+  // Only completed keys expire: a key whose outcome is unknown stays, so a
+  // retry of it keeps answering outcome_unknown instead of running anew.
+  async removeExpiredIdempotencyKeys(before: number): Promise<void> {
+    await this.db
+      .deleteFrom('idempotency')
+      .where('state', '=', 'completed')
+      .where('updated_at', '<', before)
+      .execute();
+  }
+
+  // Records a materialization as it starts, in the resolving phase.
+  async createMaterialization(
+    row: Pick<
+      WorkspaceMaterialization,
+      'sessionID' | 'target' | 'dir' | 'sourceKind' | 'withheldEnv'
+    >,
+    at: number,
+  ): Promise<void> {
+    await this.db
+      .insertInto('workspace_materialization')
+      .values({
+        session_id: row.sessionID,
+        target: row.target,
+        dir: row.dir,
+        source_kind: row.sourceKind,
+        phase: 'resolving',
+        repo_url: null,
+        sha: null,
+        ref: null,
+        error_code: null,
+        started_at: at,
+        updated_at: at,
+        materialized_at: null,
+        withheld_env: JSON.stringify(row.withheldEnv),
+      })
+      .execute();
+  }
+
+  async updateMaterialization(
+    sessionID: SessionID,
+    fields: MaterializationUpdate,
+    at: number,
+  ): Promise<void> {
+    await this.db
+      .updateTable('workspace_materialization')
+      .set({
+        phase: fields.phase,
+        ...(fields.repoURL === undefined ? {} : { repo_url: fields.repoURL }),
+        ...(fields.sha === undefined ? {} : { sha: fields.sha }),
+        ...(fields.ref === undefined ? {} : { ref: fields.ref }),
+        ...(fields.errorCode === undefined ? {} : { error_code: fields.errorCode }),
+        ...(fields.materializedAt === undefined ? {} : { materialized_at: fields.materializedAt }),
+        updated_at: at,
+      })
+      .where('session_id', '=', sessionID)
+      .execute();
+  }
+
+  async findMaterialization(sessionID: SessionID): Promise<WorkspaceMaterialization | null> {
+    const row = await this.db
+      .selectFrom('workspace_materialization')
+      .selectAll()
+      .where('session_id', '=', sessionID)
+      .executeTakeFirst();
+
+    return row === undefined ? null : toWorkspaceMaterialization(row);
+  }
+
+  // Runs once as a daemon starts, before it serves a request: a
+  // materialization still short of ready or failed belonged to a daemon that
+  // stopped partway through it, and its workspace was never verified, so it
+  // fails as interrupted.
+  async reconcileMaterializations(at: number): Promise<void> {
+    await this.db
+      .updateTable('workspace_materialization')
+      .set({ phase: 'failed', error_code: 'workspace_interrupted', updated_at: at })
+      .where('phase', 'not in', ['ready', 'failed'])
+      .execute();
+  }
+
   async stop(): Promise<void> {
     await this.db.destroy();
 
@@ -519,21 +834,227 @@ export class StateStore {
         return;
       }
 
-      const entries: FleetEntry[] = [];
+      const legacy: LegacyFleetEntry[] = [];
 
       for (const entry of parsed) {
         const parsedEntry = parseFleetEntry(entry);
 
         if (parsedEntry !== undefined) {
-          entries.push(parsedEntry);
+          legacy.push(parsedEntry);
         }
       }
 
-      await this.writeFleet(entries);
+      await this.writeFleet(buildFleetFromLegacy(legacy));
     } catch {}
   }
 }
 
+// A fleet row's ready workspace as its entry field, or nothing for a row
+// that ran without one.
+function buildWorkspaceField(
+  row: Readonly<{
+    repo_url: string | null;
+    sha: string | null;
+    ref: string | null;
+    materialized_at: number | null;
+  }>,
+): { readonly workspace?: SessionWorkspace } {
+  if (row.repo_url === null || row.sha === null || row.materialized_at === null) {
+    return {};
+  }
+
+  return {
+    workspace: {
+      repoURL: row.repo_url,
+      sha: row.sha,
+      ...(row.ref === null ? {} : { ref: row.ref }),
+      materializedAt: row.materialized_at,
+    },
+  };
+}
+
+// Two sessions share an agent session id when one resumes the other's agent
+// session. The fleet keeps one row per agent session id, the entry written
+// last, and relinks the survivors as a one-level hierarchy of rows it holds:
+//
+// 1. A link to a dropped row follows to the row that replaced it. A row that
+//    replaced its own parent follows that parent's link instead, and a link
+//    to a row the fleet does not hold makes the row top-level.
+// 2. Where the links form a cycle, the row written first in the cycle
+//    becomes top-level.
+// 3. A row whose parent is itself a sub-session moves up to the top-level
+//    row above it.
+//
+// Every walk takes at most one step per entry, so it ends on any input.
+// A ready workspace's withheld variable names as their entry field, or
+// nothing for a row that withholds none.
+function buildWithheldEnvField(stored: string | null): { readonly withheldEnv?: string[] } {
+  const names = parseWithheldEnv(stored);
+
+  return names.length === 0 ? {} : { withheldEnv: names };
+}
+
+function parseWithheldEnv(stored: string | null): string[] {
+  const parsed: unknown = stored === null ? [] : JSON.parse(stored);
+
+  return Array.isArray(parsed)
+    ? parsed.filter((name): name is string => typeof name === 'string')
+    : [];
+}
+
+function buildFleetWithoutReplacedRows(entries: readonly FleetEntry[]): FleetEntry[] {
+  const survivors = new Map<AgentSessionID, SessionID>();
+
+  for (const entry of entries) {
+    if (entry.agentSessionID !== undefined) {
+      survivors.set(entry.agentSessionID, entry.sessionID);
+    }
+  }
+
+  const replaced = new Map<SessionID, SessionID>();
+
+  for (const entry of entries) {
+    const survivor =
+      entry.agentSessionID === undefined ? undefined : survivors.get(entry.agentSessionID);
+
+    if (survivor !== undefined && survivor !== entry.sessionID) {
+      replaced.set(entry.sessionID, survivor);
+    }
+  }
+
+  const kept = entries.filter((entry) => !replaced.has(entry.sessionID));
+
+  const order = new Map(kept.map((entry, index) => [entry.sessionID, index]));
+  const links = new Map(entries.map((entry) => [entry.sessionID, entry.parent]));
+
+  const bound = entries.length;
+
+  // The surviving row an entry's parent link resolves to, or undefined when
+  // it resolves to no row the fleet keeps.
+  const findSurvivingParent = (entry: FleetEntry): SessionID | undefined => {
+    let link = entry.parent;
+
+    for (let step = 0; step < bound && link !== undefined; step++) {
+      const target = replaced.get(link) ?? link;
+
+      if (target !== entry.sessionID) {
+        return order.has(target) ? target : undefined;
+      }
+
+      link = links.get(link);
+    }
+
+    return undefined;
+  };
+
+  const linked = new Map(kept.map((entry) => [entry.sessionID, findSurvivingParent(entry)]));
+
+  // The row written first in the cycle a walk up from this row enters, or
+  // undefined when the walk reaches a top-level row.
+  const findCycleStart = (sessionID: SessionID): SessionID | undefined => {
+    const seen: SessionID[] = [];
+    let current: SessionID | undefined = sessionID;
+
+    for (let step = 0; step <= bound && current !== undefined; step++) {
+      const at = seen.indexOf(current);
+
+      if (at !== -1) {
+        return seen.slice(at).toSorted((a, b) => (order.get(a) ?? 0) - (order.get(b) ?? 0))[0];
+      }
+
+      seen.push(current);
+
+      current = linked.get(current);
+    }
+
+    return undefined;
+  };
+
+  for (const entry of kept) {
+    const cycleStart = findCycleStart(entry.sessionID);
+
+    if (cycleStart !== undefined) {
+      linked.set(cycleStart, undefined);
+    }
+  }
+
+  // The top-level row a walk up from this row reaches.
+  const findTopLevel = (sessionID: SessionID): SessionID => {
+    let current = sessionID;
+
+    for (let step = 0; step < bound; step++) {
+      const parent = linked.get(current);
+
+      if (parent === undefined) {
+        return current;
+      }
+
+      current = parent;
+    }
+
+    return current;
+  };
+
+  const relinked: FleetEntry[] = [];
+
+  for (const entry of kept) {
+    const { parent: _parent, ...rest } = entry;
+    const top = findTopLevel(entry.sessionID);
+    const row = top === entry.sessionID ? rest : { ...rest, parent: top };
+
+    relinked.push(row);
+  }
+
+  return relinked;
+}
+
+// The ownership epoch this daemon writes under. Nothing moves a session to
+// a later epoch yet, so every row this daemon owns holds this one.
+const OWNER_EPOCH = 1;
+
+// Rejects a write that touches a session another daemon owns, or one whose
+// stored epoch has moved past this daemon's: the write is stale.
+async function requireCurrentOwnership(
+  trx: Transaction<StateStoreSchema>, // oxlint-disable-line prefer-readonly-parameter-types -- a kysely transaction bound to a live connection; not meaningfully freezable
+  daemonID: DaemonID,
+  sessionIDs: readonly SessionID[],
+): Promise<void> {
+  if (sessionIDs.length === 0) {
+    return;
+  }
+
+  const stale = await trx
+    .selectFrom('session_owner')
+    .select(['session_id', 'daemon_id', 'owner_epoch'])
+    .where('session_id', 'in', sessionIDs)
+    .where((eb) => eb.or([eb('daemon_id', '!=', daemonID), eb('owner_epoch', '!=', OWNER_EPOCH)]))
+    .executeTakeFirst();
+
+  if (stale !== undefined) {
+    throw new DaemonError(
+      'stale_epoch',
+      `session '${stale.session_id}' is owned by daemon '${stale.daemon_id}' at epoch ${stale.owner_epoch}; this daemon writes as '${daemonID}' at epoch ${OWNER_EPOCH}`,
+    );
+  }
+}
+
+// Mints each legacy entry an atc session id, then moves each sub-session
+// link from the parent's agent session id to the parent's minted id.
+function buildFleetFromLegacy(legacy: readonly LegacyFleetEntry[]): FleetEntry[] {
+  const minted = new Map(legacy.map((entry) => [entry.agentSessionID, toSessionID(randomUUID())]));
+
+  return legacy.map((entry) => {
+    const { parent, ...rest } = entry;
+    const parentID = parent === undefined ? undefined : minted.get(parent);
+
+    return {
+      ...rest,
+      sessionID: minted.get(entry.agentSessionID) ?? toSessionID(randomUUID()),
+      ...(parentID === undefined ? {} : { parent: parentID }),
+    };
+  });
+}
+
 function buildScopeMatch(
   eb: ExpressionBuilder<StateStoreSchema, 'events'>, // oxlint-disable-line prefer-readonly-parameter-types -- a kysely expression builder bound to a live query; not meaningfully freezable
   scope: EventScope | null,
@@ -542,9 +1063,13 @@ function buildScopeMatch(
     return eb.lit(true);
   }
 
+  if (scope.atcIDs.length === 0 && scope.agentSessionIDs.length === 0) {
+    return eb.lit(false);
+  }
+
   return eb.or([
-    eb('atc_id', '=', scope.atcID),
-    ...(scope.agentSessionID === undefined ? [] : [eb('session_id', '=', scope.agentSessionID)]),
+    ...(scope.atcIDs.length === 0 ? [] : [eb('atc_id', 'in', scope.atcIDs)]),
+    ...(scope.agentSessionIDs.length === 0 ? [] : [eb('session_id', 'in', scope.agentSessionIDs)]),
   ]);
 }
 
@@ -614,6 +1139,24 @@ function buildOwnerFilter(
     : eb.or([byAtcID, eb('agent_session_id', '=', owner.agentSessionID)]);
 }
 
+function toIdempotencyRecord(row: Readonly<StateStoreSchema['idempotency']>): IdempotencyRecord {
+  return {
+    principal: row.principal,
+    operation: row.operation,
+    key: row.key,
+    payloadHash: row.payload_hash,
+    state: row.state,
+    effectRef: row.effect_ref,
+    result: row.result,
+    effectTarget:
+      row.effect_target === null || row.effect_target_identity === null
+        ? null
+        : { target: row.effect_target, targetIdentity: row.effect_target_identity },
+    createdAt: row.created_at,
+    updatedAt: row.updated_at,
+  };
+}
+
 function toMessageRecord(row: Readonly<StateStoreSchema['messages']>): MessageRecord {
   return {
     id: toMessageID(row.id),
@@ -631,3 +1174,23 @@ function toMessageRecord(row: Readonly<StateStoreSchema['messages']>): MessageRe
     ...(row.turn_id === null ? {} : { turn: row.turn_id }),
   };
 }
+
+function toWorkspaceMaterialization(
+  row: Readonly<StateStoreSchema['workspace_materialization']>,
+): WorkspaceMaterialization {
+  return {
+    sessionID: toSessionID(row.session_id),
+    target: row.target,
+    dir: row.dir,
+    sourceKind: row.source_kind,
+    phase: row.phase,
+    repoURL: row.repo_url,
+    sha: row.sha,
+    ref: row.ref,
+    errorCode: row.error_code,
+    startedAt: row.started_at,
+    updatedAt: row.updated_at,
+    materializedAt: row.materialized_at,
+    withheldEnv: parseWithheldEnv(row.withheld_env),
+  };
+}