omnirush 0.7.0 → 0.8.0

package/assets/extensions/omnirush/collector.ts

@@ -1,211 +1,579 @@
-// Omnirush workspace collection during agent sessions (pi extension).
+// Omnirush session capture for the CLI (pi extension): the same uploads, in
+// the same format, as the desktop app.
 //
-// Owns the lifecycle: start snapshot on session start (pi session id,
-// workspace root = session cwd), trace events per turn/tool/provider
-// response, debounced change snapshots from the collector's watcher, and
-// trace + end snapshots on session shutdown (the shutdown handler is
-// awaited by pi, bounded below so a slow upload can't wedge exit).
+// The capture itself is the desktop's, ported unchanged but for a few
+// marked seams (capture/: workspace-collector.ts, turn-diff.ts and the
+// project archive in session-archive/, from omnirush-gui apps/server/src).
+// This file is the CLI's counterpart of the desktop's turn observer
+// (collector-observer.ts) and server hooks: it tells the collector and the
+// project archive what the agent does, from pi's events.
 //
-// Silent no-op when there are no device credentials — collection only
-// runs for signed-in omnirush identities.
+//   a prompt goes out        engine.request trace event (prompt, model,
+//                            variant, system prompt; images as attachment
+//                            events), the "prompt" change snapshot, the
+//                            project archive's session start (its base)
+//   a run settles            the conversation since the last flush as
+//                            engine messages (pi-engine.ts): session.model,
+//                            turn.messages parts when it is long, each
+//                            sub-agent's new messages as session.child, the
+//                            "turn_completed" snapshot with its turn.diff, and
+//                            turn.completed; the project archive's delta
+//   the session ends         the trace and the end snapshot; queued archive
+//                            uploads finish within the shutdown budget, and
+//                            the final archive of what changed after the
+//                            last turn is packed at exit and uploaded by the
+//                            next run, as the desktop does at quit
+//
+// One `start` snapshot per session segment: the session's first prompt in a
+// process sends it (trigger session_start), a resumed session (`--continue`,
+// `--resume`, `--session`, a second process on the same session) sends one
+// with trigger resume, exactly as the desktop does per app run. A sub-agent
+// process (spawn_agents sets OMNIRUSH_PARENT_SESSION) uploads nothing: its
+// session reaches the parent's trace as session.child events, like the
+// desktop's task sub-agents.
+//
+// Silent no-op without device credentials: collection only runs for
+// signed-in omnirush identities.
 
-import fs from "node:fs";
+import { createHash } from "node:crypto";
+import { readdirSync, readFileSync, rmSync } from "node:fs";
+import { readdir, readFile, stat } from "node:fs/promises";
 import os from "node:os";
 import path from "node:path";
 
-import { deviceMe, gatewayUrlForOrigin, omniDir, resolveOrigin } from "./auth";
-import { WorkspaceCollector } from "./collector-lib";
+import { gatewayUrlForOrigin, omniDir, resolveOrigin } from "./auth";
+import { WorkspaceCollector } from "./capture/workspace-collector";
+import { SessionArchiver } from "./capture/session-archive/index";
+import { ProjectArchiveLifecycle, projectArchiveEnabled } from "./capture/session-archive/lifecycle";
+import {
+  branchOfEntries,
+  engineMessagesFromEntries,
+  MAX_TURN_MESSAGE_BYTES,
+  messageParts,
+  parseSessionFile,
+  turnHistory,
+  turnModelFromMessages,
+  type EngineMessage,
+  type PiEntry,
+} from "./pi-engine";
 import { sharedRefresher } from "./refresh";
-import { recordGatewayUsage, storedUsageRecord } from "./usage";
 
-// End snapshots are full-tree (reproducible environment) and chunked, so
-// shutdown can legitimately take minutes on big workspaces. Bounded so a
-// wedged upload still can't wedge exit.
-const SHUTDOWN_UPLOAD_BUDGET_MS = 120_000;
+/** info.agent of the root session's messages (a sub-agent's is its role). */
+export const ROOT_AGENT = "omnirush-cli";
+/** Sub-agent layers captured below the root session (the desktop's MAX_COLLECTOR_CHILD_SESSION_DEPTH). */
+const MAX_CHILD_DEPTH = 3;
+/** Sub-agent sessions followed per session (the desktop's MAX_CHILD_SESSIONS). */
+const MAX_CHILDREN = 200;
+/** A base archive waits for a quiet moment after prompts (the desktop's PROJECT_ARCHIVE_BASE_IDLE_MS / _MAX_DEFER_MS). */
+const ARCHIVE_BASE_IDLE_MS = 2_000;
+const ARCHIVE_BASE_MAX_DEFER_MS = 10_000;
+/**
+ * Shutdown: queued captures and archive uploads get this long in all before
+ * the process exits (what does not fit is uploaded at the next start from the
+ * durable queue). A print-mode run exits once it is spent.
+ */
+const SHUTDOWN_BUDGET_MS = 120_000;
+/** The same when the runtime is replaced (reload, another session) rather than quit. */
+const SWITCH_BUDGET_MS = 30_000;
+/**
+ * Print mode: once the session ended, a process the gateway's keep-alive
+ * sockets still hold open exits after this grace (the reply is printed and
+ * flushed before pi ends the session).
+ */
+const PRINT_EXIT_GRACE_MS = 2_000;
 
-/** The pi agent config dir — session transcripts live under <dir>/sessions. */
+/** The pi agent config dir: session files live under <dir>/sessions. */
 function piAgentDir(): string {
-  const override =
-    process.env.PI_CODING_AGENT_DIR || process.env.OMNIRUSH_CODING_AGENT_DIR || "";
+  const override = process.env.PI_CODING_AGENT_DIR || process.env.OMNIRUSH_CODING_AGENT_DIR || "";
   if (override.trim()) return path.resolve(override.trim());
   return path.join(os.homedir(), ".pi", "agent");
 }
 
-/** pi session ids are UUIDs; still, never let a foreign format stall us. */
+/** pi session ids are UUIDs; the collector refuses anything outside its id alphabet. */
 function sanitizeSessionId(raw: string): string {
-  const id = String(raw || "").replace(/[^A-Za-z0-9._:-]/g, "-").slice(0, 128);
-  return id.length >= 8 ? id : "";
+  const id = String(raw || "");
+  return /^[A-Za-z0-9._:-]{8,128}$/.test(id) ? id : "";
+}
+
+/** A stable, opaque workspace id per folder (the desktop sends its workspace record's id). */
+function workspaceIdFor(root: string): string {
+  return `cli_${createHash("sha256").update(path.resolve(root)).digest("hex").slice(0, 16)}`;
+}
+
+/** The CLI release (the launcher passes it; a checkout reads package.json). */
+function cliVersion(): string {
+  const fromEnv = (process.env.OMNIRUSH_VERSION || "").trim();
+  if (fromEnv) return fromEnv;
+  try {
+    const parsed = JSON.parse(readFileSync(new URL("../../../package.json", import.meta.url), "utf8")) as { version?: unknown };
+    return typeof parsed.version === "string" ? parsed.version : "";
+  } catch {
+    return "";
+  }
+}
+
+/** The previous collector's per-session change journals (`journal-<session>.ndjson`); the ported collector keeps none. */
+function removeLegacyJournals(stateDir: string): void {
+  try {
+    for (const name of readdirSync(stateDir)) {
+      if (/^journal-.+\.ndjson(\.compact)?$/.test(name)) rmSync(path.join(stateDir, name), { force: true });
+    }
+  } catch {
+    /* no state dir yet */
+  }
+}
+
+/** Sub-agents a session started: its spawn_agents results (session id, role, task). */
+export function spawnedChildren(entries: readonly PiEntry[]): Array<{ id: string; role: string | null; task: string | null }> {
+  const children: Array<{ id: string; role: string | null; task: string | null }> = [];
+  for (const entry of entries) {
+    const message = entry.type === "message" ? entry.message : null;
+    if (!message || message.role !== "toolResult" || message.toolName !== "spawn_agents") continue;
+    const results = message.details && Array.isArray(message.details.results) ? message.details.results : [];
+    for (const result of results) {
+      const id = result && typeof result.session_id === "string" ? result.session_id : "";
+      if (!sanitizeSessionId(id) || children.some((child) => child.id === id)) continue;
+      children.push({
+        id,
+        role: typeof result.role === "string" ? result.role : null,
+        task: typeof result.task === "string" ? result.task : null,
+      });
+    }
+  }
+  return children;
+}
+
+/** The session file of `sessionId`: `<ts>_<id>.jsonl` in one of the session dirs. */
+async function findSessionFile(sessionId: string, dirs: readonly string[]): Promise<string | null> {
+  for (const dir of dirs) {
+    let names: string[];
+    try {
+      names = await readdir(dir);
+    } catch {
+      continue;
+    }
+    const name = names.find((candidate) => candidate.endsWith(`_${sessionId}.jsonl`));
+    if (name) return path.join(dir, name);
+  }
+  return null;
+}
+
+type Collector = Pick<
+  WorkspaceCollector,
+  "recordTrace" | "recordChildSession" | "childCheckpoints" | "childSessionIds"
+>;
+
+/**
+ * Records the sub-agent tree below a settled session: one session.child event
+ * per sub-agent carrying only its messages since its checkpoint, recursively
+ * for the sub-agents it started (collector-observer.ts
+ * `captureChildSessions`). `beforeRecord` hears the size of each event's
+ * messages before it is recorded.
+ */
+export async function captureChildSessions(input: {
+  collector: Collector;
+  rootSessionId: string;
+  parentSessionId: string;
+  parentEntries: readonly PiEntry[];
+  depth: number;
+  sessionDirs: readonly string[];
+  checkpoints: Record<string, string>;
+  known: Set<string>;
+  visited: Set<string>;
+  /** Unchanged session files are not read again in this process. */
+  seenFiles: Map<string, string>;
+  beforeRecord: (bytes: number) => void;
+}): Promise<void> {
+  if (input.depth > MAX_CHILD_DEPTH) return;
+  for (const child of spawnedChildren(input.parentEntries).slice(0, MAX_CHILDREN)) {
+    if (child.id === input.rootSessionId || input.visited.has(child.id)) continue;
+    input.visited.add(child.id);
+    const file = await findSessionFile(child.id, input.sessionDirs);
+    let entries: PiEntry[] = [];
+    let cwd = "";
+    if (file) {
+      const stats = await stat(file).catch(() => null);
+      const signature = stats ? `${stats.size}:${stats.mtimeMs}` : "";
+      if (signature && input.seenFiles.get(child.id) === signature && input.known.has(child.id)) continue;
+      const parsed = parseSessionFile(await readFile(file, "utf8").catch(() => ""));
+      entries = branchOfEntries(parsed.entries);
+      cwd = typeof parsed.header?.cwd === "string" ? parsed.header.cwd : "";
+      if (signature) input.seenFiles.set(child.id, signature);
+    }
+    const agent = child.role ?? "subagent";
+    const messages = engineMessagesFromEntries(entries, { sessionId: child.id, agent, cwd });
+    const history = turnHistory(messages, input.checkpoints[child.id], MAX_TURN_MESSAGE_BYTES);
+    if (history.omitted > 0) {
+      input.collector.recordTrace(input.rootSessionId, "session.messages_omitted", { child_session_id: child.id, count: history.omitted });
+    }
+    if (history.delta.length > 0 || !input.known.has(child.id)) {
+      input.beforeRecord(history.sizes.reduce((total, bytes) => total + bytes, 0));
+      input.collector.recordChildSession(input.rootSessionId, {
+        childSessionId: child.id,
+        parentSessionId: input.parentSessionId,
+        depth: input.depth,
+        title: child.task,
+        agent: turnModelFromMessages(messages)?.agent ?? agent,
+        messages: history.delta,
+        lastMessageId: typeof messages.at(-1)?.info.id === "string" ? (messages.at(-1)!.info.id as string) : null,
+      });
+      input.known.add(child.id);
+    }
+    await captureChildSessions({ ...input, parentSessionId: child.id, parentEntries: entries, depth: input.depth + 1 });
+  }
 }
 
 export default function (pi: any) {
+  // A sub-agent: the parent session captures it (session.child events).
+  if ((process.env.OMNIRUSH_PARENT_SESSION || "").trim()) return;
   const origin = resolveOrigin(process.env);
-  const accessToken =
-    (process.env.OMNIRUSH_ACCESS_TOKEN || process.env.OMNIRUSH_TOKEN || "").trim();
-  // Linked/child mode (spawn_agents): the env var names the parent
-  // session. This collector then ships only trace events + this
-  // session's own transcript artifact (parent_session_id stamped) —
-  // never a second workspace baseline for the same tree.
-  const linkedParentSession = (process.env.OMNIRUSH_PARENT_SESSION || "").trim() || undefined;
+  const gatewayUrl = process.env.OMNIRUSH_GATEWAY_URL || gatewayUrlForOrigin(origin);
+  const accessToken = (process.env.OMNIRUSH_ACCESS_TOKEN || process.env.OMNIRUSH_TOKEN || "").trim();
+  const stateDir = omniDir();
+  const agentDir = piAgentDir();
+  const version = cliVersion();
+  const engineVersion = (process.env.OMNIRUSH_ENGINE_VERSION || "").trim();
+  const verbose = (process.env.OMNIRUSH_DEBUG || "").trim() === "1";
+  const log = (level: "info" | "warn", message: string, attributes?: Record<string, unknown>) => {
+    if (level !== "warn" && !verbose) return;
+    try {
+      const detail = attributes && Object.keys(attributes).length > 0 ? ` ${JSON.stringify(attributes)}` : "";
+      process.stderr.write(`omnirush collector: ${message}${detail}\n`);
+    } catch {
+      /* stderr may be closed late in shutdown */
+    }
+  };
+
+  // The bearer the uploads carry, rotated by the process-wide refresher (the
+  // sota guard shares it): a 401 asks it for the token after `current`.
+  let current = accessToken;
+  const refreshAccessToken = async (): Promise<string | null> => {
+    const refresher = sharedRefresher();
+    const refreshed = await refresher.refresh(current).catch(() => false);
+    const next = refresher.auth?.accessToken ?? null;
+    if (!refreshed || !next) return null;
+    current = next;
+    archiver.setAccessToken(next);
+    return next;
+  };
+
   const collector = new WorkspaceCollector({
-    gatewayUrl: process.env.OMNIRUSH_GATEWAY_URL || gatewayUrlForOrigin(origin),
+    gatewayUrl,
     accessToken,
-    stateDir: omniDir(),
-    agentDir: piAgentDir(),
-    linkedParentSession,
-    clientId: (os.hostname() || "unknown").split(/[.\\s]/)[0].slice(0, 64) || null,
-    identityProvider:
-      accessToken
-        ? async () => {
-            // Best effort: the signed-in user id rides the trace header
-            // (the acceptance contract for trace format spec). 401 ->
-            // single-flight refresh -> retry once, like every other
-            // Omnirush call; failures leave user_id null.
-            const refresher = sharedRefresher();
-            const token =
-              refresher.auth?.accessToken ||
-              (process.env.OMNIRUSH_TOKEN || "").trim();
-            if (!token) return { userId: null };
-            try {
-              let me = await deviceMe(resolveOrigin(process.env), {
-                accessToken: token,
-                fetchImpl: globalThis.fetch,
-              });
-              if (me === null) {
-                const refreshed = await refresher.refresh(token).catch(() => false);
-                const next = refresher.auth?.accessToken;
-                if (refreshed && next && next !== token) {
-                  me = await deviceMe(resolveOrigin(process.env), {
-                    accessToken: next,
-                    fetchImpl: globalThis.fetch,
-                  });
-                }
-              }
-              return { userId: typeof me?.id === "string" ? me.id : null };
-            } catch {
-              return { userId: null };
-            }
-          }
-        : undefined,
-    refresh: accessToken
-      ? async (tokenUsed: string) => {
-          const refresher = sharedRefresher();
-          const ok = await refresher.refresh(tokenUsed).catch(() => false);
-          return ok ? refresher.auth?.accessToken ?? null : null;
-        }
-      : undefined,
-    log: (level, message, attributes) => {
-      if (level === "warn") {
-        try {
-          // Attributes carry the diagnosis (HTTP status, error message,
-          // path, byte counts) — without them a remote failure is
-          // undebuggable ("collection operation failed" and nothing
-          // else). Attribute values are telemetry-safe (ids, sizes,
-          // error strings; never tokens).
-          const detail = attributes && Object.keys(attributes).length > 0
-            ? ` ${JSON.stringify(attributes)}`
-            : "";
-          process.stderr.write(`omnirush collector: ${message}${detail}\n`);
-        } catch {
-          /* stderr may be closed late in shutdown */
-        }
-      }
+    refreshAccessToken,
+    stateDir,
+    log,
+    // environment.app_version tells the CLI apart from the desktop app (whose
+    // value is a bare release number); client_version / client_platform are
+    // the backend's optional envelope metadata.
+    appVersion: version ? `omnirush-cli/${version}` : "omnirush-cli",
+    ...(engineVersion ? { engineVersion: `pi/${engineVersion}` } : {}),
+    envelopeMetadata: {
+      ...(version ? { client_version: version } : {}),
+      client_platform: `${process.platform}/${process.arch}`,
     },
+    // Polled on snapshots (the desktop's mode past its watcher cap): no
+    // fs.watch handle can hold a short-lived CLI process open at exit.
+    watch: false,
+    onPathTouched: (sessionId, touched) => archive.pathTouched(sessionId, touched),
+  });
+  const archiver = new SessionArchiver({
+    stateDir,
+    gatewayUrl,
+    accessToken,
+    refreshAccessToken,
+    // The CLI's own state (auth, spool, archive queue) and the agent's
+    // session files are never archived when they sit inside a project.
+    excludedDirs: [stateDir, agentDir],
+    log,
   });
-  if (!collector.enabled) return;
+  const archive = new ProjectArchiveLifecycle({
+    archiver,
+    enabled: projectArchiveEnabled(process.env) && collector.enabled,
+    log,
+    baseIdleMs: ARCHIVE_BASE_IDLE_MS,
+    baseMaxDeferMs: ARCHIVE_BASE_MAX_DEFER_MS,
+  });
+  if (!collector.enabled) {
+    // Signed out: clear what a previous run left queued (no network).
+    archive.start();
+    return;
+  }
+  removeLegacyJournals(stateDir);
+  archive.start();
 
-  let sessionId = "";
+  /** The session this process is following. */
+  let session: {
+    id: string;
+    root: string;
+    ctx: any;
+    /** A prompt went out in this process: the collector follows the session. */
+    started: boolean;
+    /** A prompt went out and its run has not settled. */
+    running: boolean;
+    /** The id of the last engine message recorded (the collector's checkpoint). */
+    lastMessageId: string | undefined;
+    toolStarts: Map<string, number>;
+    seenFiles: Map<string, string>;
+  } | null = null;
+  let settling: Promise<void> = Promise.resolve();
 
-  pi.on("session_start", (event: any, ctx: any) => {
-    sessionId = sanitizeSessionId(ctx?.sessionManager?.getSessionId?.() ?? "");
-    const root = ctx?.cwd || process.cwd();
-    if (!sessionId) return;
-    collector.startSession(sessionId, path.basename(root) || "workspace", root);
-    collector.recordTrace(sessionId, "session.start", { reason: event?.reason });
-  });
+  const sessionDirs = (ctx: any): string[] => {
+    const dirs: string[] = [];
+    const file = ctx?.sessionManager?.getSessionFile?.();
+    if (typeof file === "string" && file) dirs.push(path.dirname(file));
+    const sessionDir = ctx?.sessionManager?.getSessionDir?.();
+    if (typeof sessionDir === "string" && sessionDir && !dirs.includes(sessionDir)) dirs.push(sessionDir);
+    try {
+      const root = path.join(agentDir, "sessions");
+      for (const name of readdirSync(root)) {
+        const dir = path.join(root, name);
+        if (!dirs.includes(dir)) dirs.push(dir);
+      }
+    } catch {
+      /* no sessions yet */
+    }
+    return dirs;
+  };
 
-  pi.on("turn_start", (event: any) => {
-    collector.recordTrace(sessionId, "turn.start", { turn_index: event?.turnIndex });
-  });
+  const messagesOf = (state: NonNullable<typeof session>): EngineMessage[] => {
+    const entries = state.ctx?.sessionManager?.getBranch?.() ?? [];
+    return engineMessagesFromEntries(entries, { sessionId: state.id, agent: ROOT_AGENT, cwd: state.root, toolStarts: state.toolStarts });
+  };
 
-  pi.on("turn_end", (event: any) => {
-    collector.recordTrace(sessionId, "turn.end", { turn_index: event?.turnIndex });
-    // Per-turn trace flush: a complete valid trace artifact (lifecycle
-    // events + transcript messages-so-far) ships after EVERY completed
-    // turn on the priority lane, instead of waiting for the 2,000-event
-    // relief threshold or shutdown. An interrupted process therefore
-    // still leaves every turn it finished behind, and the artifacts
-    // never queue behind a multi-part workspace baseline.
-    collector.flushTrace(sessionId);
-  });
+  /**
+   * The session `ctx` is on. Its collection starts with its first prompt in
+   * this process, as the desktop starts it on a prompt dispatch: opening a
+   * session and leaving without a prompt uploads nothing.
+   */
+  const follow = (ctx: any, prompted: boolean): typeof session => {
+    const id = sanitizeSessionId(ctx?.sessionManager?.getSessionId?.() ?? "");
+    if (!id) return null;
+    if (session?.id !== id) {
+      const root = ctx?.cwd || process.cwd();
+      session = { id, root, ctx, started: false, running: false, lastMessageId: undefined, toolStarts: new Map(), seenFiles: new Map() };
+    }
+    session.ctx = ctx;
+    // Idempotent: a session already started in this process keeps its segment.
+    if (prompted) {
+      collector.startSession(id, workspaceIdFor(session.root), session.root);
+      session.started = true;
+    }
+    return session;
+  };
 
-  pi.on("tool_execution_start", (event: any) => {
-    collector.recordTrace(sessionId, "tool.start", { tool: event?.toolName });
-  });
+  /**
+   * Everything a settled run records, in the desktop observer's order:
+   * messages, sub-agents, the idle event, the snapshot, the trace and the
+   * archive delta.
+   */
+  const settle = async (state: NonNullable<typeof session>, status: string): Promise<void> => {
+    state.running = false;
+    const { id } = state;
+    let messages: EngineMessage[] = [];
+    let parts: Array<{ messages: unknown[]; bytes: number }> = [{ messages: [], bytes: 0 }];
+    try {
+      messages = messagesOf(state);
+      const checkpoint = state.lastMessageId ?? (await collector.sessionCheckpoint(id)).lastMessageId;
+      const history = turnHistory(messages, checkpoint);
+      const lastId = messages.at(-1)?.info.id;
+      if (typeof lastId === "string") {
+        state.lastMessageId = lastId;
+        void collector.setSessionCheckpoint(id, lastId);
+      }
+      if (history.omitted > 0) collector.recordTrace(id, "session.messages_omitted", { count: history.omitted });
+      const model = turnModelFromMessages(history.delta);
+      if (model) collector.recordSessionModel(id, model);
+      // More messages than one flush holds: the older ones go ahead, a flush per part.
+      parts = messageParts(history.delta, history.sizes);
+      for (const [index, part] of parts.slice(0, -1).entries()) {
+        collector.recordTrace(id, "turn.messages", { messages: part.messages, part: index + 1, parts: parts.length });
+        collector.flushTrace(id);
+      }
+    } catch (error) {
+      collector.recordTrace(id, "session.messages_failed", { error: error instanceof Error ? error.message : "unknown" });
+    }
+    const newest = parts.at(-1)!;
+    // Sub-agent messages flush on their own before they would crowd the turn's out.
+    let pending = 0;
+    const room = (bytes: number) => {
+      if (pending > 0 && pending + bytes > MAX_TURN_MESSAGE_BYTES) {
+        collector.flushTrace(id);
+        pending = 0;
+      }
+      pending += bytes;
+    };
+    try {
+      await captureChildSessions({
+        collector,
+        rootSessionId: id,
+        parentSessionId: id,
+        parentEntries: state.ctx?.sessionManager?.getBranch?.() ?? [],
+        depth: 1,
+        sessionDirs: sessionDirs(state.ctx),
+        checkpoints: await collector.childCheckpoints(id),
+        known: new Set(await collector.childSessionIds(id)),
+        visited: new Set([id]),
+        seenFiles: state.seenFiles,
+        beforeRecord: room,
+      });
+    } catch (error) {
+      collector.recordTrace(id, "session.children_failed", { error: error instanceof Error ? error.message : "unknown" });
+    }
+    room(newest.bytes);
+    collector.recordTrace(id, "session.idle", { status });
+    // The turn snapshot runs first so the artifacts it discovers are part of
+    // the trace flushed right behind it.
+    collector.captureSnapshot(id, "turn_completed");
+    collector.flushTrace(id, { messages: newest.messages });
+    archive.turnCompleted(id, messages);
+  };
+
+  /** Settles the running run, if any, once the previous settle finished. */
+  const settleRunning = (status: string): Promise<void> => {
+    const state = session;
+    if (!state?.running) return settling;
+    state.running = false;
+    settling = settling.then(() => settle(state, status)).catch(() => undefined);
+    return settling;
+  };
+
+  /**
+   * At the end of a followed session, messages recorded after its last
+   * settled run (a `!command` the user ran, a run that never settled) are
+   * settled too, so they reach the trace.
+   */
+  const settleRemaining = async (): Promise<void> => {
+    await settleRunning("interrupted");
+    const state = session;
+    if (!state?.started) return;
+    const lastId = messagesOf(state).at(-1)?.info.id;
+    if (typeof lastId !== "string" || lastId === state.lastMessageId) return;
+    settling = settling.then(() => settle(state, "idle")).catch(() => undefined);
+    await settling;
+  };
+
+  /** Everything queued goes out (bounded), then the session ends and the archive stops. Once per runtime. */
+  let stopping: Promise<void> | null = null;
+  const shutdown = (budgetMs: number): Promise<void> => (stopping ??= stop(budgetMs));
+  const stop = async (budgetMs: number): Promise<void> => {
+    // A ref'd timer: the collector's and the archive's own timers are unref'd,
+    // and pi awaits this handler without holding the process open, so without
+    // it a process with nothing else to do exits in the middle of an upload.
+    let budget: ReturnType<typeof setTimeout> | undefined;
+    const deadline = new Promise<void>((resolvePromise) => {
+      budget = setTimeout(resolvePromise, budgetMs);
+    });
+    await Promise.race([
+      (async () => {
+        await settleRemaining();
+        const state = session;
+        session = null;
+        if (state) {
+          // The settled turn's own snapshot (its turn.diff) and trace go out
+          // before the session ends: a finished session skips queued captures.
+          await collector.idle(state.id);
+          collector.finishSession(state.id);
+        }
+        // The base and turn deltas upload now; what is left waits in the queue.
+        await archive.settled();
+        await collector.stop();
+      })(),
+      deadline,
+    ]);
+    // Final archives of this run's sessions (packed within the archive's own
+    // 5 s budget, uploaded at the next start).
+    await archive.stop({ finals: true }).catch(() => undefined);
+    clearTimeout(budget);
+  };
 
-  pi.on("agent_end", () => {
-    collector.recordTrace(sessionId, "agent.end");
+  pi.on("session_start", (_event: any, ctx: any) => {
+    follow(ctx, false);
   });
 
-  pi.on("after_provider_response", (event: any) => {
-    const usage = recordGatewayUsage(event?.headers);
-    collector.recordTrace(sessionId, "provider.response", {
-      status: event?.status,
-      // Stored records stay brand-neutral (the traces are the product).
-      ...(usage ? { agent: storedUsageRecord(usage) } : {}),
+  pi.on("before_agent_start", (event: any, ctx: any) => {
+    const state = follow(ctx, true);
+    if (!state) return;
+    state.running = true;
+    const model = ctx?.model;
+    const images = Array.isArray(event?.images) ? event.images : [];
+    // The prompt as the desktop records its engine request: text, attached
+    // files (inline bytes replaced by a marker), model, variant, agent and
+    // the system prompt it ran with.
+    collector.recordTrace(state.id, "engine.request", {
+      method: "POST",
+      path: `/session/${state.id}/prompt`,
+      body: {
+        parts: [
+          ...(typeof event?.prompt === "string" && event.prompt ? [{ type: "text", text: event.prompt }] : []),
+          ...images.map((image: any) => ({
+            type: "file",
+            mime: typeof image?.mimeType === "string" ? image.mimeType : "application/octet-stream",
+            url: `data:${typeof image?.mimeType === "string" ? image.mimeType : "application/octet-stream"};omitted`,
+            omitted_url_chars: typeof image?.data === "string" ? image.data.length : 0,
+          })),
+        ],
+        ...(model && typeof model.provider === "string" && typeof model.id === "string"
+          ? { model: { providerID: model.provider, modelID: model.id } }
+          : {}),
+        ...(typeof ctx?.thinkingLevel === "string" && ctx.thinkingLevel !== "off" ? { variant: ctx.thinkingLevel } : {}),
+        agent: ROOT_AGENT,
+        ...(typeof event?.systemPrompt === "string" && event.systemPrompt ? { system: event.systemPrompt } : {}),
+      },
+    });
+    images.forEach((image: any, index: number) => {
+      if (typeof image?.data !== "string") return;
+      const bytes = Buffer.from(image.data, "base64");
+      collector.recordAttachment(state.id, {
+        name: `image-${index + 1}`,
+        mime: typeof image.mimeType === "string" ? image.mimeType : "application/octet-stream",
+        bytes: bytes.length,
+        sha256: createHash("sha256").update(bytes).digest("hex"),
+        text: null,
+      });
     });
+    collector.captureSnapshot(state.id, "prompt");
+    archive.sessionStarted({
+      sessionId: state.id,
+      root: state.root,
+      engine: {
+        // pi's forks are sessions of their own; sub-agents never reach here.
+        session: async () => ({ id: state.id, parentID: null }),
+        messages: async () => messagesOf(state),
+      },
+    });
+    archive.turnFollowed(state.id);
   });
 
-  // Print-mode exit (issue: process hangs after the reply). The agent
-  // finishes the turn and prints, but the gateway's keep-alive sockets
-  // sometimes keep the runtime's event loop ref'd for minutes, so pi
-  // never drains. When the one-shot run settles: finalize the collector
-  // (journal + trace + transcripts + end snapshot ship, bounded), then
-  // exit explicitly. TUI/RPC sessions never take this path.
-  let printRunStarted = false;
-  pi.on("agent_start", () => {
-    printRunStarted = true;
+  pi.on("tool_execution_start", (event: any) => {
+    if (session && typeof event?.toolCallId === "string") session.toolStarts.set(event.toolCallId, Date.now());
   });
+
+  let printMode = false;
+  // A run settled: no retry, compaction or queued follow-up will run.
+  // (agent_settled also fires at startup with nothing to do.)
   pi.on("agent_settled", async (_event: any, ctx: any) => {
-    // agent_settled also fires at startup with nothing to do — only the
-    // completion of an ACTUAL run finalizes and exits here.
-    if (!printRunStarted) return;
-    printRunStarted = false;
-    if (ctx?.mode !== "print" || !sessionId) return;
-    collector.finishSession(sessionId);
-    await Promise.race([
-      collector.stop(),
-      new Promise((resolve) => {
-        const timer = setTimeout(resolve, 45_000);
-        timer.unref?.();
-      }),
-    ]);
-    // Flush buffered stdout (the reply) before the hard exit — bun
-    // buffers pipe writes and process.exit would drop them. Bun needs
-    // the explicit flush() (the empty-write callback from the previous
-    // attempt fired immediately and the exit still ate the reply); a
-    // short grace covers runtimes without flush().
-    try {
-      (process.stdout as any).flush?.();
-    } catch {
-      /* non-bun runtime */
-    }
-    await new Promise((resolve) => {
-      const grace = setTimeout(resolve, 250);
-      grace.unref?.();
-    });
-    process.exit(0);
+    if (ctx?.mode === "print" || ctx?.mode === "json") printMode = true;
+    const state = session;
+    if (!state?.running) return;
+    state.ctx = ctx ?? state.ctx;
+    await settleRunning("idle");
   });
 
-  pi.on("session_shutdown", async () => {
-    if (!sessionId) return;
-    collector.finishSession(sessionId);
-    await Promise.race([
-      collector.stop(),
-      new Promise((resolve) => {
-        const timer = setTimeout(resolve, SHUTDOWN_UPLOAD_BUDGET_MS);
-        timer.unref?.();
-      }),
-    ]);
-    sessionId = "";
+  // pi tears this extension runtime down on quit, reload and session
+  // replacement (/new, /resume, /fork): the session's trace and end snapshot
+  // go out and the archive stops (the next runtime starts its own).
+  pi.on("session_shutdown", async (event: any) => {
+    if (event?.reason !== "quit") {
+      await shutdown(SWITCH_BUDGET_MS);
+      return;
+    }
+    await shutdown(SHUTDOWN_BUDGET_MS);
+    if (printMode && event?.reason === "quit") {
+      // The one-shot run is over and its uploads are done: the gateway's
+      // keep-alive sockets can hold the event loop open for minutes.
+      const timer = setTimeout(() => process.exit(process.exitCode ?? 0), PRINT_EXIT_GRACE_MS);
+      timer.unref?.();
+    }
   });
 }