@zgeoff/atc 2.11.0 → 2.12.0

package/src/store/state-store.ts

@@ -1,19 +1,24 @@
 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';
@@ -24,13 +29,14 @@ import type { TrailEntry } from './trail-entry';
 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 +76,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 +104,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,10 +119,19 @@ 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')
+      .where('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`rowid`)
       .select([
+        'session_id',
         'agent_session_id',
         'name',
         'cwd',
@@ -113,10 +139,14 @@ export class StateStore {
         'last_attached',
         'agent',
         'exited',
-        'parent',
+        'parent_session_id',
         'prompt',
         'result',
         'transcript_path',
+        'model',
+        'effort',
+        'target',
+        'target_identity',
       ])
       .execute();
 
@@ -124,17 +154,24 @@ 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 }),
       });
     }
 
@@ -164,27 +201,65 @@ 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,
+          })
+          .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 +268,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 +278,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> {
@@ -282,8 +357,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 +571,145 @@ 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();
+  }
+
   async stop(): Promise<void> {
     await this.db.destroy();
 
@@ -519,21 +733,187 @@ 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 {}
   }
 }
 
+// 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.
+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 +922,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 +998,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),