@zgeoff/atc 3.1.1 → 3.3.0

package/src/parse-systemd-unit.ts

@@ -0,0 +1,25 @@
+/**
+ * The name of the user service a process runs in, read from its
+ * `/proc/<pid>/cgroup` text: the first `<name>.service` below
+ * `user@<uid>.service`. A process in a login session, a scope, or a system
+ * service reads as null, since a user manager cannot restart it.
+ */
+export function parseSystemdUnit(cgroupText: string, uid: number): string | null {
+  for (const line of cgroupText.split('\n')) {
+    const path = line.slice(line.indexOf(':', line.indexOf(':') + 1) + 1);
+    const segments = path.split('/').filter((segment) => segment !== '');
+    const manager = segments.indexOf(`user@${uid}.service`);
+
+    if (manager === -1) {
+      continue;
+    }
+
+    const unit = segments.slice(manager + 1).find((segment) => segment.endsWith('.service'));
+
+    if (unit !== undefined) {
+      return unit;
+    }
+  }
+
+  return null;
+}

package/src/protocol/daemon-features.ts

@@ -59,6 +59,11 @@ export const DAEMON_FEATURES = [
   // can destroy its host answers `confirmation_required`.
   'session.forget',
 
+  // `session.forget` takes `refusePinned` and `refuseLive`, checked in the
+  // same step that forgets, and refuses with `session_pinned` or
+  // `session_live`.
+  'session.forget.preconditions',
+
   // `session.submit` exists.
   'session.submit',
 

package/src/protocol/protocol.ts

@@ -71,6 +71,8 @@ const ERROR_CODES = [
   'host_leased',
   'confirmation_required',
   'confirm_token_invalid',
+  'session_pinned',
+  'session_live',
   'already_answered',
   'too_slow',
   'stale_epoch',

package/src/protocol/request-param-schemas.ts

@@ -213,6 +213,12 @@ export const REQUEST_PARAM_SCHEMAS = {
       .string({ error: 'session.forget confirmToken must be a string' })
       .min(1, 'session.forget confirmToken must not be empty')
       .optional(),
+
+    // Refuse a pinned session, or a sub-session of a pinned one.
+    refusePinned: z.boolean({ error: 'session.forget refusePinned must be a boolean' }).optional(),
+
+    // Refuse a live session.
+    refuseLive: z.boolean({ error: 'session.forget refuseLive must be a boolean' }).optional(),
   }),
 
   // Owner-only: withdraw the grants of the runtime auth binding on the

package/src/read-daemon-process.ts

@@ -0,0 +1,50 @@
+import { readFileSync, readlinkSync, statSync } from 'node:fs';
+import { buildOwnDaemonProcess } from './build-own-daemon-process';
+import { collectRestartEnv } from './collect-restart-env';
+import type { DaemonProcess } from './daemon-process';
+import { parseDaemonFlags } from './parse-daemon-flags';
+
+/**
+ * Reads what a running daemon was started with from `/proc/<pid>`: its
+ * environment, its working directory, and the listener flags on its command
+ * line. Where `/proc` has no entry for it, as on macOS, the environment is
+ * this process's own and the flags are empty.
+ */
+export function readDaemonProcess(pid: number): DaemonProcess {
+  let environ: string;
+  let cmdline: string;
+
+  try {
+    environ = readFileSync(`/proc/${pid}/environ`, 'utf8');
+    cmdline = readFileSync(`/proc/${pid}/cmdline`, 'utf8');
+  } catch {
+    return buildOwnDaemonProcess();
+  }
+
+  const env: Record<string, string> = {};
+
+  for (const entry of environ.split('\0')) {
+    const eq = entry.indexOf('=');
+
+    if (eq > 0) {
+      env[entry.slice(0, eq)] = entry.slice(eq + 1);
+    }
+  }
+
+  return {
+    env: collectRestartEnv(env),
+    cwd: readCwd(pid),
+    flags: parseDaemonFlags(cmdline),
+    fromProc: true,
+  };
+}
+
+function readCwd(pid: number): string | null {
+  try {
+    const cwd = readlinkSync(`/proc/${pid}/cwd`);
+
+    return statSync(cwd).isDirectory() ? cwd : null;
+  } catch {
+    return null;
+  }
+}

package/src/read-fleet-snapshot.ts

@@ -0,0 +1,37 @@
+import { DaemonClient } from './client/daemon-client';
+import { readStoredRows } from './read-stored-rows';
+import type { StoredRow } from './read-stored-rows';
+import { getBuild } from './shared/get-build';
+
+const SNAPSHOT_DEADLINE_MS = 3000;
+
+/**
+ * Reads the stored fleet rows from the daemon at this socket, or null when
+ * it does not answer in time. The restart takes this before it stops the
+ * daemon, so the rows keep their original exit status.
+ */
+export async function readFleetSnapshot(socketPath: string): Promise<StoredRow[] | null> {
+  let client: DaemonClient;
+
+  try {
+    client = await DaemonClient.open(socketPath);
+  } catch {
+    return null;
+  }
+
+  const expiry = setTimeout(() => {
+    client.stop();
+  }, SNAPSHOT_DEADLINE_MS);
+
+  try {
+    await client.sendHello(getBuild());
+
+    return await readStoredRows(client);
+  } catch {
+    return null;
+  } finally {
+    clearTimeout(expiry);
+
+    client.stop();
+  }
+}

package/src/read-restart-result.ts

@@ -0,0 +1,17 @@
+import { readFileSync } from 'node:fs';
+import { join } from 'node:path';
+import { parseRestartResult } from './parse-restart-result';
+import type { RestartResult } from './parse-restart-result';
+import { restartsDir } from './shared/config';
+
+/**
+ * The result of the last finished restart, or null when none was written or
+ * the file holds anything else.
+ */
+export function readRestartResult(): RestartResult | null {
+  try {
+    return parseRestartResult(readFileSync(join(restartsDir, 'last.json'), 'utf8').trim());
+  } catch {
+    return null;
+  }
+}

package/src/read-stored-fleet-file.ts

@@ -0,0 +1,74 @@
+import { Database } from 'bun:sqlite';
+import type { StoredRow } from './read-stored-rows';
+
+/**
+ * Reads the stored fleet rows straight from the state database, read-only,
+ * for a daemon whose protocol this build cannot speak. Each column is read
+ * under the name the file's schema gives it, and one the schema lacks reads
+ * as its default, so an older daemon's database reads without a migration. On a schema that records
+ * session ownership, only the rows this state directory's daemon owns are
+ * read, as the daemon's own fleet list reads them. Returns null when the
+ * file is missing or cannot be read.
+ */
+export function readStoredFleetFile(dbPath: string): StoredRow[] | null {
+  let db: Database;
+
+  try {
+    db = new Database(dbPath, { readonly: true });
+  } catch {
+    return null;
+  }
+
+  try {
+    const columns = new Set(
+      db
+        .query<{ name: string }, []>('PRAGMA table_info(fleet)')
+        .all()
+        .map((column) => column.name),
+    );
+
+    // Before the store keyed rows by atc session id, a row held only the
+    // agent's session id, under one of two names; the new daemon assigns
+    // the atc id as it migrates, so such a row is matched by the agent's id.
+    const agentColumn = ['agent_session_id', 'claude_id'].find((name) => columns.has(name));
+
+    if (!columns.has('session_id') && agentColumn === undefined) {
+      return null;
+    }
+
+    const id = columns.has('session_id') ? 'fleet.session_id' : `fleet.${agentColumn}`;
+    const agentID = agentColumn === undefined ? 'NULL' : `fleet.${agentColumn}`;
+    const exited = columns.has('exited') ? 'fleet.exited' : '0';
+
+    const hasOwners =
+      db
+        .query<{ name: string }, []>(
+          "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'session_owner'",
+        )
+        .get() !== null;
+
+    const owned = hasOwners
+      ? " WHERE fleet.session_id IN (SELECT session_id FROM session_owner WHERE daemon_id = (SELECT value FROM prefs WHERE key = 'daemon_id'))"
+      : '';
+
+    const rows = db
+      .query<
+        { id: string; agent_session_id: string | null; name: string | null; exited: number | null },
+        []
+      >(
+        `SELECT ${id} AS id, ${agentID} AS agent_session_id, fleet.name, ${exited} AS exited FROM fleet${owned} ORDER BY fleet.rowid`,
+      )
+      .all();
+
+    return rows.map((row) => ({
+      id: row.id,
+      name: row.name ?? row.id,
+      exited: row.exited === 1,
+      agentSessionID: row.agent_session_id,
+    }));
+  } catch {
+    return null;
+  } finally {
+    db.close();
+  }
+}

package/src/read-stored-rows.ts

@@ -0,0 +1,36 @@
+import type { DaemonClient } from './client/daemon-client';
+import { isRecord } from './shared/report';
+
+export interface StoredRow {
+  readonly id: string;
+  readonly name: string;
+  readonly exited: boolean;
+
+  // The agent's own session id, which matches the row when its atc id is
+  // assigned only as the new daemon migrates the store; null when unknown.
+  readonly agentSessionID: string | null;
+}
+
+/**
+ * Reads the fleet rows a daemon has stored, exited ones included.
+ */
+export async function readStoredRows(
+  client: Pick<DaemonClient, 'sendRequest'>,
+): Promise<StoredRow[]> {
+  const listed = await client.sendRequest('fleet.list');
+
+  const fleet = listed['fleet'];
+
+  if (!Array.isArray(fleet)) {
+    return [];
+  }
+
+  return fleet
+    .filter((entry) => isRecord(entry))
+    .map((entry) => ({
+      id: String(entry['sessionID']),
+      name: String(entry['name']),
+      exited: entry['exited'] === true,
+      agentSessionID: typeof entry['agentSessionID'] === 'string' ? entry['agentSessionID'] : null,
+    }));
+}

package/src/remove-restart-result.ts

@@ -0,0 +1,11 @@
+import { rmSync } from 'node:fs';
+import { join } from 'node:path';
+import { restartsDir } from './shared/config';
+
+/**
+ * Removes the result of the last finished restart, so a reader finds a
+ * result only from a restart that ran after this call.
+ */
+export function removeRestartResult(): void {
+  rmSync(join(restartsDir, 'last.json'), { force: true });
+}

package/src/restart-daemon.ts

@@ -0,0 +1,149 @@
+import { resolve } from 'node:path';
+import { buildOwnDaemonProcess } from './build-own-daemon-process';
+import { buildRestartFailure } from './build-restart-failure';
+import { collectRestartPlan } from './collect-restart-plan';
+import { formatRestartPreflight } from './format-restart-preflight';
+import type { RestartResult } from './parse-restart-result';
+import { readDaemonProcess } from './read-daemon-process';
+import { readFleetSnapshot } from './read-fleet-snapshot';
+import { readStoredFleetFile } from './read-stored-fleet-file';
+import { runSystemctl } from './run-systemctl';
+import { dbFile } from './shared/config';
+import { getBuild } from './shared/get-build';
+import { loadListenerTokens } from './shared/load-listener-tokens';
+import { startReplacementDaemon } from './start-replacement-daemon';
+import { stopDaemonProcess } from './stop-daemon-process';
+import { verifyRestoredFleet } from './verify-restored-fleet';
+import { waitForReplacement } from './wait-for-replacement';
+
+export interface RestartOptions {
+  readonly runID: string;
+  readonly callerSession: string | null;
+  readonly listen: string | null;
+  readonly tokenFile: string | null;
+  readonly timeoutSeconds: number | null;
+}
+
+const START_WAIT_MS = 30_000;
+
+/**
+ * Stops the running daemon, starts exactly one in its place, restores the
+ * stored fleet on it, and reports what came back. The preflight prints
+ * before anything stops. A failure on the way is a result, never a throw.
+ */
+export async function restartDaemon(options: RestartOptions): Promise<RestartResult> {
+  const plan = await collectRestartPlan(options.callerSession);
+
+  for (const line of formatRestartPreflight(plan)) {
+    console.log(line);
+  }
+
+  const interrupted = (plan.sessions ?? [])
+    .filter((session) => session.state === 'running')
+    .map((session) => ({ name: session.name, id: session.id }));
+
+  // Taken while the old daemon still runs, so a row that the restore marks
+  // exited is not mistaken for one that was archived before. A daemon this
+  // build cannot speak to is read from its database instead.
+  const snapshot =
+    (plan.answer?.kind === 'ok' && plan.socketPath !== null
+      ? await readFleetSnapshot(plan.socketPath)
+      : null) ?? readStoredFleetFile(dbFile);
+
+  if (snapshot === null) {
+    console.log(
+      'cannot read the stored fleet before the stop: the report checks the rows the new daemon lists, and a row it marks exited counts as restored',
+    );
+  }
+
+  const buildFailure = (error: string): RestartResult =>
+    buildRestartFailure(options.runID, error, interrupted);
+
+  let expectedBuild: string | null = null;
+
+  if (plan.replacement.kind === 'unit') {
+    if (options.listen !== null || options.tokenFile !== null) {
+      console.log('--listen and --token-file are ignored: the unit decides how the daemon starts');
+    }
+
+    const restarted = await runSystemctl(['restart', plan.replacement.unit]);
+
+    if (restarted.code !== 0) {
+      return buildFailure(`systemctl restart ${plan.replacement.unit} failed: ${restarted.stderr}`);
+    }
+  } else {
+    expectedBuild = getBuild();
+
+    const old = plan.pid === null ? buildOwnDaemonProcess() : readDaemonProcess(plan.pid);
+
+    if (plan.pid !== null && !old.fromProc) {
+      console.log(
+        'cannot read /proc for the daemon: the replacement starts with this restart’s environment and without the old listener flags',
+      );
+    }
+
+    const listen = options.listen ?? old.flags.listen;
+    const tokenFile = options.tokenFile ?? old.flags.tokenFile;
+
+    if ((listen === null) !== (tokenFile === null)) {
+      return buildFailure('--listen and --token-file go together; the daemon was left running');
+    }
+
+    if (listen !== null && tokenFile !== null) {
+      const loaded = loadListenerTokens(resolve(old.cwd ?? process.cwd(), tokenFile));
+
+      if (!loaded.ok) {
+        return buildFailure(`${loaded.reason}; the daemon was left running`);
+      }
+    }
+
+    if (plan.pid !== null) {
+      const stopped = await stopDaemonProcess(plan.pid);
+
+      const stoppedLine =
+        stopped === 'killed'
+          ? `the daemon (pid ${plan.pid}) ignored SIGTERM for 10 s, so it was killed with SIGKILL`
+          : `stopped the daemon (pid ${plan.pid})`;
+
+      console.log(stoppedLine);
+    }
+
+    startReplacementDaemon(old, { listen: options.listen, tokenFile: options.tokenFile });
+  }
+
+  const replacement = await waitForReplacement(plan.pid, expectedBuild, START_WAIT_MS);
+
+  if (!replacement.ok) {
+    return buildFailure(replacement.reason);
+  }
+
+  try {
+    console.log(`daemon pid ${replacement.pid} is up (${replacement.build}); restoring the fleet`);
+
+    const verdict = await verifyRestoredFleet(replacement.client, options.timeoutSeconds, snapshot);
+
+    return {
+      runID: options.runID,
+      code: verdict.failed.length === 0 ? 0 : 1,
+      pid: replacement.pid,
+      build: replacement.build,
+      listenPort: replacement.listenPort,
+      restored: verdict.total - verdict.failed.length,
+      total: verdict.total,
+      failed: [...verdict.failed],
+      interrupted,
+      error: null,
+    };
+  } catch (error) {
+    const reason = error instanceof Error ? error.message : String(error);
+
+    return {
+      ...buildFailure(`the fleet could not be verified: ${reason}`),
+      pid: replacement.pid,
+      build: replacement.build,
+      listenPort: replacement.listenPort,
+    };
+  } finally {
+    replacement.client.stop();
+  }
+}

package/src/restart-plan.ts

@@ -0,0 +1,33 @@
+// What a restart is about to stop and what it starts in its place.
+export interface RestartPlan {
+  // The pid of the daemon to stop, or null when no live pid is recorded.
+  readonly pid: number | null;
+
+  // The socket that answered the handshake, or null when none did.
+  readonly socketPath: string | null;
+
+  // How the daemon answered this build's handshake.
+  readonly answer: DaemonAnswer | null;
+
+  // The session list the daemon returned, or null when it could not be read.
+  readonly sessions: readonly PlanSession[] | null;
+
+  // The session the restart was run from, when it runs inside one.
+  readonly callerSession: string | null;
+
+  readonly replacement: ReplacementPlan;
+}
+
+export type DaemonAnswer =
+  | { readonly kind: 'ok'; readonly build: string; readonly protocol: number }
+  | { readonly kind: 'refused'; readonly message: string };
+
+export interface PlanSession {
+  readonly id: string;
+  readonly name: string;
+  readonly state: string;
+}
+
+export type ReplacementPlan =
+  | { readonly kind: 'unit'; readonly unit: string; readonly execStart: string | null }
+  | { readonly kind: 'plain'; readonly build: string; readonly command: string };

package/src/run-daemon-restart-worker.ts

@@ -0,0 +1,77 @@
+import { mkdirSync } from 'node:fs';
+import { buildRestartFailure } from './build-restart-failure';
+import { formatRestartReport } from './format-restart-report';
+import type { RestartResult } from './parse-restart-result';
+import { readRestartResult } from './read-restart-result';
+import { removeRestartResult } from './remove-restart-result';
+import { restartDaemon } from './restart-daemon';
+import type { RestartOptions } from './restart-daemon';
+import { claimDaemonLock } from './shared/claim-daemon-lock';
+import { restartLockFile, restartsDir } from './shared/config';
+import { writeRestartResult } from './write-restart-result';
+
+// How long a restart that finds another in flight waits for it to finish.
+const JOIN_WAIT_MS = 15 * 60_000;
+
+/**
+ * Runs one restart as the handoff worker: it takes the restart lock and
+ * restarts the daemon, or, when another restart holds the lock, waits for
+ * that restart to finish and reports its result without touching the
+ * daemon. Prints each step to stdout, which the caller follows in the run
+ * log, and ends the log with the result record. Resolves the exit code.
+ */
+export async function runDaemonRestartWorker(options: RestartOptions): Promise<number> {
+  console.log(`worker pid ${process.pid}`);
+
+  mkdirSync(restartsDir, { recursive: true });
+
+  const lock = await claimDaemonLock(restartLockFile, 0);
+
+  let result: RestartResult;
+
+  if (lock === null) {
+    result = await waitForRestartInFlight(options.runID);
+  } else {
+    try {
+      // A joiner reads last.json once it gets the lock, so a result left by
+      // an earlier restart must not outlive this run's start.
+      removeRestartResult();
+
+      result = await restartDaemon(options).catch((error: unknown) => {
+        const reason = error instanceof Error ? error.message : String(error);
+
+        return buildRestartFailure(options.runID, reason);
+      });
+
+      writeRestartResult(result);
+    } finally {
+      lock.dispose();
+    }
+  }
+
+  for (const line of formatRestartReport(result)) {
+    console.log(line);
+  }
+
+  console.log(JSON.stringify(result));
+
+  return result.code;
+}
+
+async function waitForRestartInFlight(runID: string): Promise<RestartResult> {
+  console.log('another restart is in flight; this run joins it and reports its result');
+
+  const held = await claimDaemonLock(restartLockFile, JOIN_WAIT_MS);
+
+  if (held === null) {
+    return buildRestartFailure(runID, 'timed out waiting for the restart in flight to finish');
+  }
+
+  const finished = readRestartResult();
+
+  held.dispose();
+
+  return finished === null
+    ? buildRestartFailure(runID, 'the restart in flight left no result')
+    : { ...finished, runID };
+}