@saadnvd1/agent-os 0.3.824 → 0.3.826

package/lib/fork/ledger.ts

@@ -0,0 +1,324 @@
+/**
+ * The children ledger: what a session has forked, kept in the database so
+ * it survives compaction, restarts and fresh conversations. A child settles
+ * once, by a conditional UPDATE; whichever caller's update changed a row
+ * acts on it.
+ */
+
+import { randomUUID } from "crypto";
+import { db } from "../db";
+import { FORK_LIMITS } from "./options";
+import type { ChildOutcome, ChildRow, GroupRow, ResultKind } from "./types";
+
+const OPEN = `state != 'settled'`;
+
+export const newGroupId = () => `g${randomUUID().slice(0, 6)}`;
+
+export interface NewChild {
+  childId: string;
+  label: string;
+  instruction: string;
+  worktree: boolean;
+}
+
+export function groupByRequest(requestId: string): GroupRow | undefined {
+  return db
+    .prepare(`SELECT * FROM fork_groups WHERE request_id = ?`)
+    .get(requestId) as GroupRow | undefined;
+}
+
+export const getGroup = (id: string): GroupRow | undefined =>
+  db.prepare(`SELECT * FROM fork_groups WHERE id = ?`).get(id) as
+    GroupRow | undefined;
+
+// The group and its children in one transaction, so a crash leaves both or
+// neither. A repeat of the same request id is the caller's check, made first
+// (the column is UNIQUE, so a race fails here rather than doubling).
+export function insertGroup(input: {
+  id: string;
+  parentId: string;
+  requestId: string;
+  startedBy: "agent" | "user";
+  context: "fork" | "brief";
+  delivery: "wake" | "next";
+  forkItemId?: string | null;
+  children: NewChild[];
+  now?: number;
+}): void {
+  const now = input.now ?? Date.now();
+  db.transaction(() => {
+    db.prepare(
+      `INSERT INTO fork_groups (id, parent_session_id, request_id, started_by, context, delivery, fork_item_id, created_at)
+       VALUES (?, ?, ?, ?, ?, ?, ?, ?)`
+    ).run(
+      input.id,
+      input.parentId,
+      input.requestId,
+      input.startedBy,
+      input.context,
+      input.delivery,
+      input.forkItemId ?? null,
+      now
+    );
+    const add = db.prepare(
+      `INSERT INTO session_children (child_session_id, parent_session_id, group_id, position, label, instruction, worktree, state, created_at)
+       VALUES (?, ?, ?, ?, ?, ?, ?, 'queued', ?)`
+    );
+    input.children.forEach((c, i) =>
+      add.run(
+        c.childId,
+        input.parentId,
+        input.id,
+        i + 1,
+        c.label,
+        c.instruction,
+        c.worktree ? 1 : 0,
+        now
+      )
+    );
+  })();
+}
+
+export const getChild = (id: string): ChildRow | undefined =>
+  db
+    .prepare(`SELECT * FROM session_children WHERE child_session_id = ?`)
+    .get(id) as ChildRow | undefined;
+
+export const isChild = (sessionId: string): boolean => !!getChild(sessionId);
+
+// A parent's children, oldest first: the order the numbers in `aos
+// children` follow, across groups.
+export function childrenOf(parentId: string): ChildRow[] {
+  return db
+    .prepare(
+      `SELECT * FROM session_children WHERE parent_session_id = ?
+       ORDER BY created_at, group_id, position`
+    )
+    .all(parentId) as ChildRow[];
+}
+
+export const numberOf = (parentId: string, childId: string): number =>
+  childrenOf(parentId).findIndex((c) => c.child_session_id === childId) + 1;
+
+export const openChildren = (parentId: string): ChildRow[] =>
+  db
+    .prepare(
+      `SELECT * FROM session_children WHERE parent_session_id = ? AND ${OPEN}
+       ORDER BY created_at, group_id, position`
+    )
+    .all(parentId) as ChildRow[];
+
+export const groupChildren = (groupId: string): ChildRow[] =>
+  db
+    .prepare(
+      `SELECT * FROM session_children WHERE group_id = ? ORDER BY position`
+    )
+    .all(groupId) as ChildRow[];
+
+// A child named the way an agent or a person would: its number in `aos
+// children`, its label, or its session id (or a prefix of it).
+export function findChild(parentId: string, ref: string): ChildRow {
+  const all = childrenOf(parentId);
+  const text = ref.trim();
+  if (/^\d+$/.test(text)) {
+    const c = all[Number(text) - 1];
+    if (c) return c;
+    throw new Error(`No child ${text}: you have ${all.length}`);
+  }
+  const byId = all.filter((c) => c.child_session_id.startsWith(text));
+  if (byId.length === 1) return byId[0];
+  const byLabel = all.filter(
+    (c) => c.label.toLowerCase() === text.toLowerCase()
+  );
+  if (byLabel.length === 1) return byLabel[0];
+  if (byLabel.length > 1)
+    throw new Error(
+      `${byLabel.length} children are named ${text}: use their numbers`
+    );
+  throw new Error(`No child "${text}" of this session`);
+}
+
+export function setState(
+  childId: string,
+  state: ChildRow["state"],
+  hold: string | null = null,
+  now = Date.now()
+): void {
+  if (state === "settled") throw new Error("Settle through settleChild");
+  db.prepare(
+    `UPDATE session_children SET state = ?, hold = ?,
+       started_at = CASE WHEN ? = 'running' AND started_at IS NULL THEN ? ELSE started_at END
+     WHERE child_session_id = ? AND state != 'settled'`
+  ).run(state, hold, state, now, childId);
+}
+
+export const markFellBack = (childId: string): void => {
+  db.prepare(
+    `UPDATE session_children SET fell_back = 1 WHERE child_session_id = ?`
+  ).run(childId);
+};
+
+export const setCut = (childId: string, cut: string | null): void => {
+  db.prepare(
+    `UPDATE session_children SET cut_message_id = ? WHERE child_session_id = ?`
+  ).run(cut, childId);
+};
+
+// First line, cut: what the parent's agent sees.
+export function summaryOf(text: string): string {
+  const line =
+    text
+      .trim()
+      .split("\n")
+      .find((l) => l.trim()) ?? "";
+  const one = line.replace(/\s+/g, " ").trim();
+  const max = FORK_LIMITS.summaryChars;
+  return one.length <= max ? one : `${one.slice(0, max - 1)}…`;
+}
+
+export const capResult = (text: string): string =>
+  text.length <= FORK_LIMITS.resultChars
+    ? text
+    : `${text.slice(0, FORK_LIMITS.resultChars)}\n… (cut)`;
+
+// True for the one caller that settled it.
+export function settleChild(
+  childId: string,
+  s: {
+    outcome: ChildOutcome;
+    summary: string;
+    result?: string | null;
+    kind: ResultKind;
+    tokens?: { in: number; out: number; cacheRead: number } | null;
+    now?: number;
+  }
+): boolean {
+  const now = s.now ?? Date.now();
+  const changed =
+    db
+      .prepare(
+        `UPDATE session_children SET state = 'settled', hold = NULL, idle_since = NULL,
+           outcome = ?, summary = ?, result = ?, result_kind = ?, settled_at = ?,
+           started_at = COALESCE(started_at, ?),
+           tokens_in = ?, tokens_out = ?, cache_read = ?
+         WHERE child_session_id = ? AND state != 'settled'`
+      )
+      .run(
+        s.outcome,
+        s.summary,
+        s.result ? capResult(s.result) : null,
+        s.kind,
+        now,
+        now,
+        s.tokens?.in ?? null,
+        s.tokens?.out ?? null,
+        s.tokens?.cacheRead ?? null,
+        childId
+      ).changes === 1;
+  if (changed)
+    db.prepare(
+      `UPDATE fork_groups SET last_settled_at = ? WHERE id =
+         (SELECT group_id FROM session_children WHERE child_session_id = ?)`
+    ).run(now, childId);
+  return changed;
+}
+
+// A return after settling: the new text stands, shown as updated, and is
+// sent again only if the first hadn't been (a delivered one stays delivered).
+export function updateResult(
+  childId: string,
+  summary: string,
+  result: string | null,
+  outcome: ChildOutcome
+): void {
+  db.prepare(
+    `UPDATE session_children SET summary = ?, result = ?, outcome = ?, result_kind = 'explicit'
+     WHERE child_session_id = ? AND state = 'settled' AND outcome != 'cancelled'`
+  ).run(summary, result ? capResult(result) : null, outcome, childId);
+}
+
+export const markIdle = (childId: string, since: number | null): void => {
+  db.prepare(
+    `UPDATE session_children SET idle_since = ? WHERE child_session_id = ?`
+  ).run(since, childId);
+};
+
+export const bumpHeavy = (childId: string): void => {
+  db.prepare(
+    `UPDATE session_children SET heavy_runs = heavy_runs + 1 WHERE child_session_id = ?`
+  ).run(childId);
+};
+
+export const markUsed = (
+  childId: string,
+  how: "opened" | "mentioned"
+): void => {
+  db.prepare(
+    `UPDATE session_children SET used = ?
+     WHERE child_session_id = ? AND (used IS NULL OR (used = 'mentioned' AND ? = 'opened'))`
+  ).run(how, childId, how);
+};
+
+// Settled and not yet handed to the parent's agent, oldest first.
+export const undelivered = (parentId: string): ChildRow[] =>
+  db
+    .prepare(
+      `SELECT * FROM session_children
+       WHERE parent_session_id = ? AND settled_at IS NOT NULL AND delivered_at IS NULL
+       ORDER BY settled_at, position`
+    )
+    .all(parentId) as ChildRow[];
+
+export const parentsWithUndelivered = (): string[] =>
+  (
+    db
+      .prepare(
+        `SELECT DISTINCT parent_session_id AS id FROM session_children
+         WHERE settled_at IS NOT NULL AND delivered_at IS NULL`
+      )
+      .all() as { id: string }[]
+  ).map((r) => r.id);
+
+// Marks these rows as one delivery; only rows nobody claimed change. The
+// rows that now carry the id are returned, so a crash between the claim and
+// the send just sends the same id again.
+export function claimDelivery(ids: string[], deliveryId: string): ChildRow[] {
+  const claim = db.prepare(
+    `UPDATE session_children SET delivery_id = ?
+     WHERE child_session_id = ? AND delivery_id IS NULL AND delivered_at IS NULL`
+  );
+  db.transaction(() => ids.forEach((id) => claim.run(deliveryId, id)))();
+  return db
+    .prepare(
+      `SELECT * FROM session_children WHERE delivery_id = ? ORDER BY settled_at, position`
+    )
+    .all(deliveryId) as ChildRow[];
+}
+
+export const markDelivered = (deliveryId: string, now = Date.now()): number =>
+  db
+    .prepare(
+      `UPDATE session_children SET delivered_at = ? WHERE delivery_id = ? AND delivered_at IS NULL`
+    )
+    .run(now, deliveryId).changes;
+
+// A result already in a wait's wake message, so it isn't sent again.
+export const markDeliveredIds = (ids: string[], now = Date.now()): void => {
+  const mark = db.prepare(
+    `UPDATE session_children SET delivered_at = ? WHERE child_session_id = ? AND delivered_at IS NULL`
+  );
+  db.transaction(() => ids.forEach((id) => mark.run(now, id)))();
+};
+
+export const unsettledWithDelivery = (parentId: string): string | null =>
+  (
+    db
+      .prepare(
+        `SELECT delivery_id AS d FROM session_children
+         WHERE parent_session_id = ? AND delivery_id IS NOT NULL AND delivered_at IS NULL LIMIT 1`
+      )
+      .get(parentId) as { d: string } | undefined
+  )?.d ?? null;
+
+export const allChildren = (): ChildRow[] =>
+  db.prepare(`SELECT * FROM session_children`).all() as ChildRow[];

package/lib/fork/options.ts

@@ -0,0 +1,46 @@
+import { experimentOn } from "../experiments";
+
+// The server's limits for forks. Clients read them from /api/client-config
+// and never carry their own (lib/config/client-config.ts).
+export const FORK_LIMITS = {
+  // Children one `aos fork` call may start.
+  maxPerFork: 8,
+  // Children of one parent not yet settled, all calls together.
+  maxOpenPerParent: 16,
+  // Children of one parent running at once; the rest wait their turn.
+  maxRunningPerParent: 4,
+  // Characters of a result's first line the parent's agent sees.
+  summaryChars: 300,
+  // Characters of a full result kept.
+  resultChars: 64_000,
+  // Characters of one instruction.
+  instructionChars: 4_000,
+  // Results that settle within this long of the first go in one delivery,
+  // and a delivery never waits longer than the cap after the first.
+  batchWindowMs: 20_000,
+  batchCapMs: 2 * 60_000,
+  // A child idle this long with no `aos return` settles with its final
+  // answer, so one about to return gets there first.
+  implicitGraceMs: 45_000,
+  // Fork children's builds and test runs at once on this machine
+  // (lib/fork/heavy-lease); 2 while the machine reads green.
+  heavySlots: 1,
+  heavySlotsCalm: 2,
+} as const;
+
+export const DELIVERY_MODES = ["wake", "next"] as const;
+export type DeliveryMode = (typeof DELIVERY_MODES)[number];
+
+export const forksOn = (): boolean => experimentOn("forks");
+
+export const FORKS_OFF = "Parallel forks are off (Settings → Experimental)";
+
+// What /api/client-config serves; no client carries these.
+export function forkConfig() {
+  return {
+    enabled: forksOn(),
+    maxPerFork: FORK_LIMITS.maxPerFork,
+    summaryChars: FORK_LIMITS.summaryChars,
+    deliveryModes: [...DELIVERY_MODES],
+  };
+}

package/lib/fork/point.ts

@@ -0,0 +1,41 @@
+/**
+ * Where a conversation is cut for a fork. A parent that runs `aos fork` is
+ * inside a turn whose own tool call has no result yet, so the cut is the
+ * last entry after which every tool call so far has its result: nothing
+ * dangling, and every earlier step of the turn kept.
+ */
+
+export interface TranscriptEntry {
+  type: string;
+  uuid: string;
+  message?: unknown;
+  parent_tool_use_id?: string | null;
+}
+
+interface Block {
+  type?: string;
+  id?: string;
+  tool_use_id?: string;
+}
+
+const blocks = (entry: TranscriptEntry): Block[] => {
+  const content = (entry.message as { content?: unknown } | undefined)?.content;
+  return Array.isArray(content) ? (content as Block[]) : [];
+};
+
+// The uuid of the entry to cut at, or null when nothing is safe to keep.
+export function cutPoint(entries: TranscriptEntry[]): string | null {
+  const open = new Set<string>();
+  let cut: string | null = null;
+  for (const entry of entries) {
+    if (entry.parent_tool_use_id) continue;
+    if (entry.type !== "user" && entry.type !== "assistant") continue;
+    for (const block of blocks(entry)) {
+      if (block.type === "tool_use" && block.id) open.add(block.id);
+      if (block.type === "tool_result" && block.tool_use_id)
+        open.delete(block.tool_use_id);
+    }
+    if (open.size === 0) cut = entry.uuid;
+  }
+  return cut;
+}

package/lib/fork/scratch.ts

@@ -0,0 +1,67 @@
+/**
+ * A fork child's own folders under the data root: a scratch folder for temp
+ * work (scripts, repros, notes), and, when it was asked for, a throwaway
+ * checkout. Both are removed when the child settles or is dropped, and a
+ * sweep removes what a crash left.
+ */
+
+import fs from "fs";
+import path from "path";
+import { agentosHome } from "../agentos-home";
+import { git } from "../git-exec";
+
+export const forksRoot = () => agentosHome("forks");
+export const forkHome = (childId: string) => agentosHome("forks", childId);
+export const scratchOf = (childId: string) =>
+  path.join(forkHome(childId), "scratch");
+export const worktreeOf = (childId: string) =>
+  path.join(forkHome(childId), "wt");
+
+// Made on demand, so a settled child that is talked to again has one.
+export async function ensureScratch(childId: string): Promise<string> {
+  const dir = scratchOf(childId);
+  await fs.promises.mkdir(dir, { recursive: true });
+  return dir;
+}
+
+// The folder's real repository, for taking a worktree out of its list.
+async function repoOf(wt: string): Promise<string | null> {
+  try {
+    const common = (
+      await git(wt, ["rev-parse", "--path-format=absolute", "--git-common-dir"])
+    ).trim();
+    return path.dirname(common);
+  } catch {
+    return null;
+  }
+}
+
+// Everything of the child's, the worktree first so git forgets it.
+export async function removeFork(childId: string): Promise<void> {
+  const home = forkHome(childId);
+  const wt = worktreeOf(childId);
+  if (fs.existsSync(wt)) {
+    const repo = await repoOf(wt);
+    if (repo)
+      await git(repo, ["worktree", "remove", "--force", "--", wt]).catch(
+        () => undefined
+      );
+  }
+  await fs.promises.rm(home, { recursive: true, force: true, maxRetries: 3 });
+}
+
+// Folders of children that are settled or gone: what a crash between a
+// settle and its cleanup left. `keep` says whose folders are still in use.
+export async function sweepForks(
+  keep: (childId: string) => boolean
+): Promise<number> {
+  const root = forksRoot();
+  const names = await fs.promises.readdir(root).catch(() => [] as string[]);
+  let removed = 0;
+  for (const name of names) {
+    if (keep(name)) continue;
+    await removeFork(name);
+    removed++;
+  }
+  return removed;
+}

package/lib/fork/settle.ts

@@ -0,0 +1,215 @@
+/**
+ * How a child ends: it returns (`aos return`), its turn ends without one and
+ * its last reply stands in after a grace, it is cancelled, or it can't be
+ * run. A child settles once; the settle that changed the row does the rest.
+ */
+
+import fs from "fs";
+import path from "path";
+import { db, queries, type Session } from "../db";
+import { git } from "../git-exec";
+import { statusOf } from "../orchestrator/facts";
+import { openWaits } from "../bus/waits/store";
+import { lastItems } from "../chat/store";
+import { stopChat } from "../chat/runner";
+import type { UsageTotals } from "../chat/events";
+import {
+  getChild,
+  markIdle,
+  setState,
+  settleChild,
+  summaryOf,
+  updateResult,
+  childrenOf,
+  findChild,
+  markDeliveredIds,
+} from "./ledger";
+import { FORK_LIMITS } from "./options";
+import { forksRoot, removeFork, worktreeOf } from "./scratch";
+import type { ChildOutcome, ChildRow, ResultKind } from "./types";
+
+const sessionOf = (id: string) =>
+  queries.getSession(db).get(id) as Session | undefined;
+
+function tokensOf(childId: string) {
+  const raw = (
+    db.prepare(`SELECT chat_usage FROM sessions WHERE id = ?`).get(childId) as
+      { chat_usage: string | null } | undefined
+  )?.chat_usage;
+  if (!raw) return null;
+  try {
+    const u = JSON.parse(raw) as UsageTotals;
+    return {
+      in: u.inputTokens,
+      out: u.outputTokens,
+      cacheRead: u.cacheReadTokens,
+    };
+  } catch {
+    return null;
+  }
+}
+
+export const patchPath = (childId: string) =>
+  path.join(forksRoot(), `${childId}.patch`);
+
+// What a child changed in its own checkout, kept as a patch the parent can
+// apply: its checkout goes with it.
+async function keepPatch(childId: string): Promise<string | null> {
+  const wt = worktreeOf(childId);
+  if (!fs.existsSync(wt)) return null;
+  try {
+    await git(wt, ["add", "-A"]);
+    const diff = await git(wt, ["diff", "--cached", "HEAD", "--binary"], {
+      maxBuffer: 64 * 1024 * 1024,
+    });
+    if (!diff.trim()) return null;
+    const file = patchPath(childId);
+    await fs.promises.mkdir(path.dirname(file), { recursive: true });
+    await fs.promises.writeFile(file, diff);
+    return file;
+  } catch {
+    return null;
+  }
+}
+
+export interface Settlement {
+  outcome: ChildOutcome;
+  summary: string;
+  result?: string | null;
+  kind: ResultKind;
+}
+
+// Settles, tidies up and starts what was waiting for the room. True when
+// this call was the one that settled it.
+export async function settleAndDeliver(
+  childId: string,
+  s: Settlement
+): Promise<boolean> {
+  const row = getChild(childId);
+  if (!row || row.state === "settled") return false;
+  let result = s.result ?? null;
+  if (row.worktree && s.outcome !== "cancelled") {
+    const patch = await keepPatch(childId);
+    if (patch)
+      result =
+        `${result ?? ""}\n\nIts changes are in ${patch} (apply with \`git apply ${patch}\`).`.trim();
+  }
+  const changed = settleChild(childId, {
+    ...s,
+    result,
+    tokens: tokensOf(childId),
+  });
+  if (!changed) return false;
+  // Cancelled by its parent or the user: they know, so nothing is sent.
+  if (s.outcome === "cancelled") markDeliveredIds([childId]);
+  await removeFork(childId).catch(() => undefined);
+  const { admitQueued } = await import("./start");
+  await admitQueued(row.parent_session_id).catch(() => undefined);
+  const { kickForks } = await import("./tick");
+  kickForks();
+  return true;
+}
+
+// `aos return`. After settling it replaces the result, shown as updated.
+export async function returnChild(
+  childId: string,
+  text: string,
+  failed: boolean
+): Promise<"settled" | "updated"> {
+  const row = getChild(childId);
+  if (!row) throw new Error("Only a fork child can return a result");
+  const body = text.trim();
+  if (!body) throw new Error('Give the result: aos return "<result>"');
+  const outcome: ChildOutcome = failed ? "failed" : "ok";
+  if (
+    await settleAndDeliver(childId, {
+      outcome,
+      summary: summaryOf(body),
+      result: body,
+      kind: "explicit",
+    })
+  )
+    return "settled";
+  updateResult(childId, summaryOf(body), body, outcome);
+  return "updated";
+}
+
+export async function cancelChild(childId: string): Promise<boolean> {
+  const row = getChild(childId);
+  if (!row || row.state === "settled") return false;
+  if (row.state !== "queued") stopChat(childId);
+  return settleAndDeliver(childId, {
+    outcome: "cancelled",
+    summary: "cancelled",
+    kind: "implicit",
+  });
+}
+
+export async function cancelChildren(
+  parentId: string,
+  ref?: string
+): Promise<ChildRow[]> {
+  const targets = ref
+    ? [findChild(parentId, ref)]
+    : childrenOf(parentId).filter((c) => c.state !== "settled");
+  const done: ChildRow[] = [];
+  for (const c of targets)
+    if (await cancelChild(c.child_session_id)) done.push(c);
+  return done;
+}
+
+const lastReply = (childId: string): string => {
+  const said = lastItems(childId, 40)
+    .reverse()
+    .find((i) => i.kind === "assistant" && i.text.trim());
+  return said?.kind === "assistant" ? said.text.trim() : "";
+};
+
+// One look at a child that hasn't settled: needing an answer, still working,
+// gone, or idle long enough that its last reply is its result.
+export async function lookAtChild(row: ChildRow, now: number): Promise<void> {
+  if (row.state === "queued" || row.state === "settled") return;
+  const session = sessionOf(row.child_session_id);
+  if (!session || session.archived_at) {
+    await settleAndDeliver(row.child_session_id, {
+      outcome: "failed",
+      summary: "its session was closed before it returned",
+      kind: "implicit",
+    });
+    return;
+  }
+  const { status, activity } = await statusOf(session);
+  if (status === "waiting") {
+    setState(
+      row.child_session_id,
+      "needs_input",
+      activity ?? "waiting for an answer"
+    );
+    markIdle(row.child_session_id, null);
+    return;
+  }
+  if (status === "running" || status === "dead") {
+    if (row.state !== "running" || row.hold)
+      setState(row.child_session_id, "running");
+    markIdle(row.child_session_id, null);
+    return;
+  }
+  if (row.state === "starting") return;
+  if (row.state !== "running" || row.hold)
+    setState(row.child_session_id, "running");
+  if (row.idle_since == null) {
+    markIdle(row.child_session_id, now);
+    return;
+  }
+  if (now - row.idle_since < FORK_LIMITS.implicitGraceMs) return;
+  // It asked its parent something and is waiting for the answer.
+  if (openWaits().some((w) => w.waiter_id === row.child_session_id)) return;
+  const reply = lastReply(row.child_session_id);
+  if (!reply) return;
+  await settleAndDeliver(row.child_session_id, {
+    outcome: "ok",
+    summary: summaryOf(reply),
+    result: reply,
+    kind: "implicit",
+  });
+}