diff --git a/CHANGELOG.md b/CHANGELOG.md index b2f96818a..76e11c036 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- **Scale-out, phase A: the foundation for API pods that share one database and runs that outlive them** (`docs/plans/scale-out.md`). Every API process has an id (`OPTIO_INSTANCE_ID`, the pod's name from the chart, plus a suffix new at every boot; `GET /api/health` says which instance answered, to a signed-in request only). One migration adds every table and column the plan's phases build on: the attachment columns on runs (`exec_state`, `exec_pid`, `consumed_bytes`, `attached_by`, `attach_lease_until`, on tasks, Job runs, PR-review runs and agent turns), `leases`, `ws_upgrade_tokens`, `inbound_webhook_deliveries`, `glance_state`, `installed_skill_files`, `ticket_sync_claims`, and one active PR review per PR URL. `services/lease-service.ts` is the one way an instance claims something for a while. The **run protocol** (`packages/container-runtime/src/run-protocol.ts`) is on every container runtime — `startRun`, `attachRun`, `deliverStdin`, `killRun`: a start script launches the agent under a supervisor in its own session in the run's home, with stdin fed from a file through a FIFO until the `__OPTIO_STDIN_EOF__` line, its output appended to `output.ndjson` and its exit code to `exit`; the supervisor writes its own pid file (` `) once its TERM trap is in place, and attach, kill and the pooled start guard trust a pid only when the live process's start time matches it; stdin bytes travel over the exec's stdin (`head -c `), never in a script or its environment; an attach is `tail --pid` from a byte offset; the exec that started a run can go away and any instance can attach later. `tini` is the pods' pid 1 (`images/base.Dockerfile`, every init script), so finished runs are reaped — **the agent images must be rebuilt for the protocol**. The repo-cleanup worker removes the run homes of finished runs from every pod (`run-home-sweep-service.ts`), the backstop for the attached worker that removes them once a run is terminal and consumed (phase B). `buildTaskStartScript` / `buildPooledStartScript` build the start scripts beside the exec scripts the workers still use (they switch in phase B). The fake runtime keeps its pods on disk (`OPTIO_FAKE_RUNTIME_DIR`) and plays a protocol run as a detached process (`fake-agent.mjs`, every event with a `seq`), and `test-utils/e2e/api-cluster.ts` boots N real API servers on one database for the e2e tier: two boot together, one is SIGKILLed, a third joins. + - **Android: every trigger type in the automations and agent trigger sheets.** A Local automation's Add trigger sheet and a persistent agent's New trigger sheet take all fourteen trigger types — GitLab, Jira, Pylon, PagerDuty, Sentry, Alertmanager and Datadog events included, with their event kinds, identity and filters — through the same rows as the New work form (one trigger editor in `core:ui`, `triggers/`). A Pylon / Alertmanager / Datadog trigger's own URL and shared secret are shown once, with copy buttons, right after it is created, in the sheets and in the New work form. - **Every agent runtime gets the work's MCP servers.** Connections' tools and MCP servers reached only Claude Code (`.mcp.json`) and Codex (`config.toml`); a Job on Gemini, a persistent agent on Copilot, or a Task on OpenCode or Cursor had the connections' credentials and notes but no tools. Each runtime now gets the same servers in the file it reads, established from the versions the agent image installs: Gemini CLI in the user `settings.json` of a `GEMINI_CLI_HOME` of the run's own (the adapter's settings merged in, each server `trust: true`), OpenCode in an `OPENCODE_CONFIG` file merged before the project's own, GitHub Copilot CLI through `--additional-mcp-config @` (`type: "local"`, every tool enabled), and Cursor in the project's `.cursor/mcp.json` with `--approve-mcps`. See `docs/connections.md` → "Which runtime reads which file". - The per-run files live in one **run home**, `/home/agent/optio/runs/` (`OPTIO_RUN_HOME`), named after the task, Job run, or agent turn — so a retry lands in the same place — and removed when the run's script exits and again when the task's worktree is cleaned up. Codex's per-run `CODEX_HOME` moves there too; it used to be a random directory that was never removed from the repo pod's home volume. diff --git a/apps/api/e2e/api-cluster.e2e.test.ts b/apps/api/e2e/api-cluster.e2e.test.ts new file mode 100644 index 000000000..32bc8daff --- /dev/null +++ b/apps/api/e2e/api-cluster.e2e.test.ts @@ -0,0 +1,117 @@ +/** + * E2E: two real API servers on one database (the scale-out harness, + * src/test-utils/e2e/api-cluster.ts). Both boot (migrations under the + * advisory lock, workers registered twice), work made through one is read + * through the other, a run completes with both taking BullMQ jobs, and one + * server's SIGKILL leaves the other serving and starting new runs. + * + * The runs here still go through the classic exec path (the workers move to + * the run protocol in phase B); the fake runtime keeps its pods on disk + * under the cluster's directory, which both servers read. + */ +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { existsSync, readdirSync } from "node:fs"; +import { join } from "node:path"; +import { startApiCluster, type ApiCluster } from "../src/test-utils/e2e/api-cluster.js"; +import { waitFor } from "../src/test-utils/e2e/api-server.js"; + +let cluster: ApiCluster; + +beforeAll(async () => { + cluster = await startApiCluster({ size: 2 }); +}, 240_000); + +afterAll(async () => { + await cluster?.stop(); +}); + +interface RunBody { + run: { id: string; state: string; costUsd: string | null; errorMessage: string | null }; +} + +async function createJob(i: number, name: string): Promise { + const { status, body } = await cluster.client(i)<{ workflow: { id: string } }>("/api/jobs", { + method: "POST", + body: JSON.stringify({ + name, + promptTemplate: "Say hello from the cluster", + agentRuntime: "claude-code", + }), + }); + expect(status).toBe(201); + return body.workflow.id; +} + +async function startRun(i: number, workflowId: string): Promise { + const { status, body } = await cluster.client(i)(`/api/jobs/${workflowId}/runs`, { + method: "POST", + body: JSON.stringify({}), + }); + expect(status).toBe(201); + return body.run.id; +} + +async function waitForRun(i: number, runId: string, states: string[]): Promise { + return waitFor( + async () => { + const { body } = await cluster.client(i)(`/api/workflow-runs/${runId}`); + return states.includes(body.run.state) ? body.run : null; + }, + { timeoutMs: 90_000, label: `run ${runId} → ${states.join("|")}` }, + ); +} + +describe("api cluster e2e", () => { + it("both servers are healthy; the public probe names no instance", async () => { + const [a, b] = await Promise.all([ + cluster.client(0)<{ healthy: boolean; instanceId?: string }>("/api/health"), + cluster.client(1)<{ healthy: boolean; instanceId?: string }>("/api/health"), + ]); + expect(a.body.healthy).toBe(true); + expect(b.body.healthy).toBe(true); + // Auth is disabled here, so no request is signed in: nothing identifying. + expect(a.body).not.toHaveProperty("instanceId"); + expect(b.body).not.toHaveProperty("instanceId"); + expect(cluster.servers.map((s) => s.name)).toEqual(["api-0", "api-1"]); + }); + + it("a Job created through A is read through B, and its run completes", async () => { + const workflowId = await createJob(0, "cluster job"); + const { status, body } = await cluster.client(1)<{ workflow: { id: string; name: string } }>( + `/api/jobs/${workflowId}`, + ); + expect(status).toBe(200); + expect(body.workflow.name).toBe("cluster job"); + + const runId = await startRun(0, workflowId); + const run = await waitForRun(1, runId, ["completed", "failed"]); + expect(run.state).toBe("completed"); + expect(run.costUsd).toBe("0.0123"); + // The fake kept its pod on disk, where both servers look. + expect(existsSync(join(cluster.fakeRuntimeDir, "containers"))).toBe(true); + expect(readdirSync(join(cluster.fakeRuntimeDir, "containers")).length).toBeGreaterThan(0); + }); + + it("after A is SIGKILLed, B keeps serving and runs new work", async () => { + await cluster.kill(0, "SIGKILL"); + expect(cluster.live().map((s) => s.name)).toEqual(["api-1"]); + await expect(cluster.client(0)("/api/health")).rejects.toThrow(/killed/); + + const workflowId = await createJob(1, "after the kill"); + const runId = await startRun(1, workflowId); + const run = await waitForRun(1, runId, ["completed", "failed"]); + expect(run.state).toBe("completed"); + + const { body } = await cluster.any()<{ healthy: boolean }>("/api/health"); + expect(body.healthy).toBe(true); + }); + + it("a third server joins the same database and work", async () => { + const c = await cluster.add(); + expect(c.name).toBe("api-2"); + const { body } = await cluster.client(2)<{ workflows: Array<{ name: string }> }>("/api/jobs"); + expect(body.workflows.map((w) => w.name)).toEqual( + expect.arrayContaining(["cluster job", "after the kill"]), + ); + }, 150_000); +}); diff --git a/apps/api/src/db/migrate-safe.test.ts b/apps/api/src/db/migrate-safe.test.ts new file mode 100644 index 000000000..1b6519b69 --- /dev/null +++ b/apps/api/src/db/migrate-safe.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, it } from "vitest"; +import { viewSourceColumns, withoutColumns } from "./migrate-safe.js"; +import { RUN_VIEWS_SQL } from "./run-views.js"; + +const VIEW = `CREATE VIEW "workflow_runs" AS + SELECT "id", "kind", "work_id" AS "workflow_id", "trigger_id", + "recovery_required", "exec_state", "consumed_bytes", + "created_at", "updated_at" + FROM "tasks" WHERE "kind" = 'standalone' + WITH CASCADED CHECK OPTION`; + +describe("run views between historical migrations", () => { + it("knows the source columns a view selects, aliases included as their source", () => { + expect(viewSourceColumns(VIEW)).toEqual([ + "id", + "kind", + "work_id", + "trigger_id", + "recovery_required", + "exec_state", + "consumed_bytes", + "created_at", + "updated_at", + ]); + expect( + viewSourceColumns(`CREATE VIEW "repo_tasks" AS SELECT * FROM "tasks" WHERE "kind" = 'repo'`), + ).toEqual([]); + expect( + viewSourceColumns(`ALTER VIEW "workflow_runs" ALTER COLUMN "kind" SET DEFAULT 'standalone'`), + ).toEqual([]); + }); + + it("leaves out the columns a later migration adds, keeping every other item intact", () => { + const out = withoutColumns(VIEW, ["exec_state", "consumed_bytes", "recovery_required"]); + expect(out).toContain( + `SELECT "id", "kind", "work_id" AS "workflow_id", "trigger_id", "created_at", "updated_at" FROM "tasks" WHERE "kind" = 'standalone'`, + ); + expect(out).toContain("WITH CASCADED CHECK OPTION"); + expect(withoutColumns(VIEW, [])).toBe(VIEW); + expect(withoutColumns(`SELECT * FROM "tasks"`, ["exec_state"])).toBe(`SELECT * FROM "tasks"`); + }); + + it("the real views' later columns need no list: each is derived from the SQL", () => { + const columns = RUN_VIEWS_SQL.flatMap(viewSourceColumns); + for (const c of [ + "exec_state", + "exec_pid", + "consumed_bytes", + "attached_by", + "attach_lease_until", + "recovery_required", + ]) { + expect(columns).toContain(c); + } + const stripped = RUN_VIEWS_SQL.map((s) => withoutColumns(s, ["attach_lease_until"])); + expect(stripped.join("\n")).not.toContain("attach_lease_until"); + expect(stripped.join("\n")).toContain('"attached_by"'); + }); +}); diff --git a/apps/api/src/db/migrate-safe.ts b/apps/api/src/db/migrate-safe.ts index 857955dbe..b233a37a2 100644 --- a/apps/api/src/db/migrate-safe.ts +++ b/apps/api/src/db/migrate-safe.ts @@ -33,6 +33,67 @@ interface MigrationEntry { const ADVISORY_LOCK_ID = 8_675_309; // arbitrary, unique to optio migrations +/** The items of a view statement's select list, trimmed. */ +function selectList(stmt: string): { list: string[]; from: string } | null { + const m = stmt.match(/SELECT([\s\S]*?)(\sFROM\s)/i); + if (!m) return null; + return { + list: m[1] + .split(",") + .map((s) => s.trim()) + .filter(Boolean), + from: m[2], + }; +} + +/** The source column of a select item (`"col"` or `"col" AS "alias"`), or null (`*`, expressions). */ +function sourceColumn(item: string): string | null { + const m = item.match(/^"([^"]+)"(?:\s+AS\s+"[^"]+")?$/i); + return m ? m[1] : null; +} + +/** The `tasks` columns a view statement selects (aliases are not columns). */ +export function viewSourceColumns(stmt: string): string[] { + const parsed = selectList(stmt); + if (!parsed) return []; + const names = new Set(); + for (const item of parsed.list) { + const col = sourceColumn(item); + if (col) names.add(col); + } + return [...names]; +} + +/** + * The view statement without the given source columns. A fresh install + * remakes the views between every historical migration, before a later + * migration has added a column the views select: the column is left out + * until it exists, and nothing has to list which columns those are. + */ +export function withoutColumns(stmt: string, missing: string[]): string { + if (missing.length === 0) return stmt; + const parsed = selectList(stmt); + if (!parsed) return stmt; + const kept = parsed.list.filter((item) => { + const col = sourceColumn(item); + return !col || !missing.includes(col); + }); + return stmt.replace( + /SELECT([\s\S]*?)(\sFROM\s)/i, + () => `SELECT ${kept.join(", ")}${parsed.from}`, + ); +} + +/** The view columns the `tasks` table does not have yet. */ +async function missingRunViewColumns(tx: { execute: Database["execute"] }): Promise { + const rows = await tx.execute<{ column_name: string }>(sql` + SELECT column_name FROM information_schema.columns + WHERE table_schema = current_schema() AND table_name = 'tasks'`); + const present = new Set(rows.map((r) => r.column_name)); + const wanted = new Set(RUN_VIEWS_SQL.flatMap(viewSourceColumns)); + return [...wanted].filter((c) => !present.has(c)); +} + /** Drop the run views if they are views (before the runs table, `workflow_runs` was a table). */ const DROP_RUN_VIEWS = sql.raw(`DO $$ BEGIN IF EXISTS (SELECT 1 FROM pg_views WHERE schemaname = current_schema() AND viewname = 'workflow_runs') THEN @@ -100,13 +161,10 @@ export async function migrateSafe(db: Database, migrationsFolder: string): Promi await tx.execute(DROP_RUN_VIEWS); // A fresh install passes through historical migrations before new // run columns exist. Keep the intermediate views compatible until - // the migration adding the column has run. - const [{ present }] = await tx.execute<{ present: boolean }>(sql` - SELECT EXISTS (SELECT 1 FROM information_schema.columns - WHERE table_schema = current_schema() AND table_name = 'tasks' - AND column_name = 'recovery_required') AS present`); + // the migration adding each column has run. + const missing = await missingRunViewColumns(tx); for (const stmt of RUN_VIEWS_SQL) { - await tx.execute(sql.raw(present ? stmt : stmt.replace('"recovery_required", ', ""))); + await tx.execute(sql.raw(withoutColumns(stmt, missing))); } } await tx.execute( diff --git a/apps/api/src/db/migrations/1792010000_scale_out.sql b/apps/api/src/db/migrations/1792010000_scale_out.sql new file mode 100644 index 000000000..5d2c2bc5c --- /dev/null +++ b/apps/api/src/db/migrations/1792010000_scale_out.sql @@ -0,0 +1,135 @@ +-- Scale-out (docs/plans/scale-out.md): every API pod shares one database, and +-- a run outlives the pod that started it. This migration carries every schema +-- change of the plan, so the phases that follow (re-attachable runs, +-- coordination, sessions / glance / skills) touch no schema. + +-- Runs: which instance is streaming the run, until when, how much of the +-- pod-side output file it has turned into log rows, and the supervisor the +-- agent runs under in its pod (run-protocol.ts). On `tasks` (every repo task +-- and Job run; the run views follow on their own), PR-review runs, and +-- persistent-agent turns. +ALTER TABLE "tasks" ADD COLUMN IF NOT EXISTS "exec_state" text; +--> statement-breakpoint +ALTER TABLE "tasks" ADD COLUMN IF NOT EXISTS "exec_pid" integer; +--> statement-breakpoint +ALTER TABLE "tasks" ADD COLUMN IF NOT EXISTS "consumed_bytes" bigint DEFAULT 0 NOT NULL; +--> statement-breakpoint +ALTER TABLE "tasks" ADD COLUMN IF NOT EXISTS "attached_by" text; +--> statement-breakpoint +ALTER TABLE "tasks" ADD COLUMN IF NOT EXISTS "attach_lease_until" timestamp with time zone; +--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "tasks_attach_lease_idx" ON "tasks" ("attach_lease_until") + WHERE "exec_state" = 'started'; +--> statement-breakpoint +ALTER TABLE "pr_review_runs" ADD COLUMN IF NOT EXISTS "exec_state" text; +--> statement-breakpoint +ALTER TABLE "pr_review_runs" ADD COLUMN IF NOT EXISTS "exec_pid" integer; +--> statement-breakpoint +ALTER TABLE "pr_review_runs" ADD COLUMN IF NOT EXISTS "consumed_bytes" bigint DEFAULT 0 NOT NULL; +--> statement-breakpoint +ALTER TABLE "pr_review_runs" ADD COLUMN IF NOT EXISTS "attached_by" text; +--> statement-breakpoint +ALTER TABLE "pr_review_runs" ADD COLUMN IF NOT EXISTS "attach_lease_until" timestamp with time zone; +--> statement-breakpoint +ALTER TABLE "persistent_agent_turns" ADD COLUMN IF NOT EXISTS "exec_state" text; +--> statement-breakpoint +ALTER TABLE "persistent_agent_turns" ADD COLUMN IF NOT EXISTS "exec_pid" integer; +--> statement-breakpoint +ALTER TABLE "persistent_agent_turns" ADD COLUMN IF NOT EXISTS "consumed_bytes" bigint DEFAULT 0 NOT NULL; +--> statement-breakpoint +ALTER TABLE "persistent_agent_turns" ADD COLUMN IF NOT EXISTS "attached_by" text; +--> statement-breakpoint +ALTER TABLE "persistent_agent_turns" ADD COLUMN IF NOT EXISTS "attach_lease_until" timestamp with time zone; +--> statement-breakpoint + +-- Leases: the one way an instance claims something for a while (a poller's +-- turn, a chat lock, a Local host). services/lease-service.ts. +CREATE TABLE IF NOT EXISTS "leases" ( + "key" text PRIMARY KEY, + "holder" text NOT NULL, + "expires_at" timestamp with time zone NOT NULL +); +--> statement-breakpoint + +-- Single-use WebSocket upgrade tokens, minted by one instance and consumed +-- by whichever one the browser's upgrade lands on. +CREATE TABLE IF NOT EXISTS "ws_upgrade_tokens" ( + "token_hash" text PRIMARY KEY, + "user_id" uuid NOT NULL, + "workspace_id" uuid, + "expires_at" timestamp with time zone NOT NULL +); +--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "ws_upgrade_tokens_expires_idx" ON "ws_upgrade_tokens" ("expires_at"); +--> statement-breakpoint + +-- Inbound webhook deliveries already handled, by source and the provider's +-- delivery id: a claim is an insert that returns a row. Swept after a day. +-- (`webhook_deliveries` is the outbound webhooks' delivery log.) +CREATE TABLE IF NOT EXISTS "inbound_webhook_deliveries" ( + "source" text NOT NULL, + "delivery_id" text NOT NULL, + "received_at" timestamp with time zone DEFAULT now() NOT NULL, + PRIMARY KEY ("source", "delivery_id") +); +--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "inbound_webhook_deliveries_received_idx" ON "inbound_webhook_deliveries" ("received_at"); +--> statement-breakpoint + +-- The glance / push bookkeeping that lived in one process's memory: one row +-- per key (`attention:`, `needs-you:`, `running:`, +-- `host-offline:`), and timers as rows with a due time +-- (`watch-end:`, `snooze:`), swept when due. +CREATE TABLE IF NOT EXISTS "glance_state" ( + "key" text PRIMARY KEY, + "value" jsonb NOT NULL, + "due_at" timestamp with time zone +); +--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "glance_state_due_idx" ON "glance_state" ("due_at") WHERE "due_at" IS NOT NULL; +--> statement-breakpoint + +-- Marketplace skills' files, written by the sync worker and read at spawn, +-- so a skill is there on every instance (the cache PVC was ReadWriteOnce). +CREATE TABLE IF NOT EXISTS "installed_skill_files" ( + "skill_id" uuid NOT NULL REFERENCES "installed_skills"("id") ON DELETE CASCADE, + "path" text NOT NULL, + "content" bytea NOT NULL, + "executable" boolean DEFAULT false NOT NULL, + "resolved_sha" text, + PRIMARY KEY ("skill_id", "path") +); +--> statement-breakpoint + +-- The ticket sync's claim: one task per (source, ticket, repo) created by the +-- sync, however many sweeps overlap. Its own table rather than an index on +-- tasks: a person may start a second task from the same issue on purpose. +CREATE TABLE IF NOT EXISTS "ticket_sync_claims" ( + "ticket_source" text NOT NULL, + "ticket_external_id" text NOT NULL, + "repo_url" text NOT NULL, + "task_id" uuid, + "created_at" timestamp with time zone DEFAULT now() NOT NULL, + PRIMARY KEY ("ticket_source", "ticket_external_id", "repo_url") +); +--> statement-breakpoint + +-- External PR reviews: one active review per PR URL, so two overlapping +-- sweeps can't both launch one. Older active duplicates (left by overlapping +-- sweeps before this) are cancelled first, newest kept, along with their +-- runs that have not finished. +WITH "dups" AS ( + SELECT "id" FROM ( + SELECT "id", row_number() OVER (PARTITION BY "pr_url" ORDER BY "updated_at" DESC, "created_at" DESC) AS rn + FROM "pr_reviews" WHERE "state" IN ('queued', 'waiting_ci', 'reviewing', 'ready') + ) d WHERE rn > 1 +), "cancelled_reviews" AS ( + UPDATE "pr_reviews" SET "state" = 'cancelled', "updated_at" = now() + WHERE "id" IN (SELECT "id" FROM "dups") RETURNING "id" +) +UPDATE "pr_review_runs" SET "state" = 'cancelled', "updated_at" = now() +WHERE "pr_review_id" IN (SELECT "id" FROM "cancelled_reviews") + AND "state" IN ('queued', 'provisioning', 'running'); +--> statement-breakpoint +CREATE UNIQUE INDEX IF NOT EXISTS "pr_reviews_active_pr_url_key" ON "pr_reviews" ("pr_url") + WHERE "state" IN ('queued', 'waiting_ci', 'reviewing', 'ready'); diff --git a/apps/api/src/db/migrations/meta/_journal.json b/apps/api/src/db/migrations/meta/_journal.json index 81f30f8d9..e2f4ffdc1 100644 --- a/apps/api/src/db/migrations/meta/_journal.json +++ b/apps/api/src/db/migrations/meta/_journal.json @@ -834,6 +834,13 @@ "when": 1792000000000, "tag": "1792000000_local_terminal_pins", "breakpoints": true + }, + { + "idx": 119, + "version": "7", + "when": 1792010000000, + "tag": "1792010000_scale_out", + "breakpoints": true } ] } diff --git a/apps/api/src/db/run-views.ts b/apps/api/src/db/run-views.ts index e09243284..16473c959 100644 --- a/apps/api/src/db/run-views.ts +++ b/apps/api/src/db/run-views.ts @@ -30,6 +30,7 @@ const VIEWS = [ "retry_count", "recovery_required", "started_at", "completed_at" AS "finished_at", "control_intent", "reconcile_backoff_until", "reconcile_attempts", "workspace_id", "owner_user_id", "prompt", "agent_type", "run_target", "max_retries", "last_activity_at", + "exec_state", "exec_pid", "consumed_bytes", "attached_by", "attach_lease_until", "created_at", "updated_at" FROM "tasks" WHERE "kind" = 'standalone' WITH CASCADED CHECK OPTION`, diff --git a/apps/api/src/db/scale-out-schema.int.test.ts b/apps/api/src/db/scale-out-schema.int.test.ts new file mode 100644 index 000000000..ffb1c1087 --- /dev/null +++ b/apps/api/src/db/scale-out-schema.int.test.ts @@ -0,0 +1,167 @@ +/** + * The scale-out migration (1792010000_scale_out) against the real, migrated + * database: every column and table the plan's phases build on is there, the + * run views carry the new run columns, and the constraints behave. + */ +import { describe, expect, it } from "vitest"; +import { sql } from "drizzle-orm"; +import { db } from "./client.js"; +import { + insertTask, + insertWorkflow, + insertWorkflowRun, +} from "../test-utils/integration/fixtures.js"; + +async function columns(table: string): Promise> { + const rows = await db.execute<{ column_name: string }>(sql` + SELECT column_name FROM information_schema.columns + WHERE table_schema = current_schema() AND table_name = ${table}`); + return new Set(rows.map((r) => r.column_name)); +} + +const ATTACH_COLUMNS = [ + "exec_state", + "exec_pid", + "consumed_bytes", + "attached_by", + "attach_lease_until", +]; + +describe("scale-out schema", () => { + it("every run table and view has the attachment columns", async () => { + for (const table of [ + "tasks", + "repo_tasks", + "workflow_runs", + "pr_review_runs", + "persistent_agent_turns", + ]) { + const cols = await columns(table); + for (const c of ATTACH_COLUMNS) expect(cols.has(c), `${table}.${c}`).toBe(true); + } + }); + + it("the new tables exist with their keys", async () => { + const rows = await db.execute<{ table_name: string }>(sql` + SELECT table_name FROM information_schema.tables WHERE table_schema = current_schema()`); + const names = new Set(rows.map((r) => r.table_name)); + for (const t of [ + "leases", + "ws_upgrade_tokens", + "inbound_webhook_deliveries", + "glance_state", + "installed_skill_files", + "ticket_sync_claims", + ]) { + expect(names.has(t), t).toBe(true); + } + // The outbound webhook log is untouched. + expect((await columns("webhook_deliveries")).has("webhook_id")).toBe(true); + const idx = await db.execute<{ indexname: string }>(sql` + SELECT indexname FROM pg_indexes WHERE schemaname = current_schema()`); + const indexes = new Set(idx.map((r) => r.indexname)); + for (const i of [ + "pr_reviews_active_pr_url_key", + "tasks_attach_lease_idx", + "ws_upgrade_tokens_expires_idx", + "inbound_webhook_deliveries_received_idx", + "glance_state_due_idx", + ]) { + expect(indexes.has(i), i).toBe(true); + } + }); + + it("consumed_bytes defaults to 0 and the columns round-trip through the views", async () => { + const task = await insertTask(); + expect(task.consumedBytes).toBe(0); + expect(task.execState).toBeNull(); + await db.execute(sql` + UPDATE "repo_tasks" SET "exec_state" = 'started', "exec_pid" = 42, "consumed_bytes" = 1234, + "attached_by" = 'api-0:abc', "attach_lease_until" = now() + interval '1 minute' + WHERE "id" = ${task.id}`); + const [row] = await db.execute<{ + exec_state: string; + consumed_bytes: string; + attached_by: string; + }>(sql` + SELECT "exec_state", "consumed_bytes", "attached_by" FROM "tasks" WHERE "id" = ${task.id}`); + expect(row).toMatchObject({ + exec_state: "started", + consumed_bytes: "1234", + attached_by: "api-0:abc", + }); + + const wf = await insertWorkflow(); + const run = await insertWorkflowRun(wf.id); + await db.execute(sql` + UPDATE "workflow_runs" SET "exec_state" = 'started', "consumed_bytes" = 7 WHERE "id" = ${run.id}`); + const [viewRow] = await db.execute<{ exec_state: string; consumed_bytes: string }>(sql` + SELECT "exec_state", "consumed_bytes" FROM "workflow_runs" WHERE "id" = ${run.id}`); + expect(viewRow).toMatchObject({ exec_state: "started", consumed_bytes: "7" }); + }); + + it("only one active PR review per PR URL; finished ones may repeat", async () => { + const insert = (state: string) => + db.execute(sql` + INSERT INTO "pr_reviews" ("pr_url", "pr_number", "repo_owner", "repo_name", "repo_url", "head_sha", "state") + VALUES ('https://github.com/o/r/pull/7', 7, 'o', 'r', 'https://github.com/o/r', 'abc', ${state}::pr_review_state)`); + await insert("submitted"); + await insert("cancelled"); + await insert("reviewing"); + const violates = (err: unknown) => { + const cause = (err as { cause?: { message?: string } }).cause; + return /pr_reviews_active_pr_url_key/.test(cause?.message ?? (err as Error).message); + }; + await expect(insert("queued")).rejects.toSatisfy(violates); + await expect(insert("ready")).rejects.toSatisfy(violates); + }); + + it("a ticket sync claim is taken once; the second insert returns no row", async () => { + const claim = () => + db.execute<{ repo_url: string }>(sql` + INSERT INTO "ticket_sync_claims" ("ticket_source", "ticket_external_id", "repo_url") + VALUES ('github', '12', 'https://github.com/o/r') + ON CONFLICT DO NOTHING RETURNING "repo_url"`); + expect((await claim()).length).toBe(1); + expect((await claim()).length).toBe(0); + }); + + it("inbound webhook deliveries dedupe by (source, delivery id)", async () => { + const claim = () => + db.execute<{ source: string }>(sql` + INSERT INTO "inbound_webhook_deliveries" ("source", "delivery_id") VALUES ('slack', 'Ev123') + ON CONFLICT DO NOTHING RETURNING "source"`); + expect((await claim()).length).toBe(1); + expect((await claim()).length).toBe(0); + const other = await db.execute<{ source: string }>(sql` + INSERT INTO "inbound_webhook_deliveries" ("source", "delivery_id") VALUES ('linear', 'Ev123') + ON CONFLICT DO NOTHING RETURNING "source"`); + expect(other.length).toBe(1); + }); + + it("glance state upserts by key and keeps a due time", async () => { + await db.execute(sql` + INSERT INTO "glance_state" ("key", "value", "due_at") VALUES ('watch-end:u1', '{"providers":["apns"]}', now() + interval '2 minutes') + ON CONFLICT ("key") DO UPDATE SET "value" = EXCLUDED."value", "due_at" = EXCLUDED."due_at"`); + await db.execute(sql` + INSERT INTO "glance_state" ("key", "value", "due_at") VALUES ('watch-end:u1', '{"providers":["fcm"]}', NULL) + ON CONFLICT ("key") DO UPDATE SET "value" = EXCLUDED."value", "due_at" = EXCLUDED."due_at"`); + const [row] = await db.execute<{ value: { providers: string[] }; due_at: string | null }>(sql` + SELECT "value", "due_at" FROM "glance_state" WHERE "key" = 'watch-end:u1'`); + expect(row.value.providers).toEqual(["fcm"]); + expect(row.due_at).toBeNull(); + }); + + it("installed skill files cascade with their skill", async () => { + const [skill] = await db.execute<{ id: string }>(sql` + INSERT INTO "installed_skills" ("name", "source_url") VALUES ('it-skill', 'https://github.com/o/skills.git') RETURNING "id"`); + await db.execute(sql` + INSERT INTO "installed_skill_files" ("skill_id", "path", "content", "executable") + VALUES (${skill.id}, 'SKILL.md', ${Buffer.from("# hi")}, false)`); + await db.execute(sql`DELETE FROM "installed_skills" WHERE "id" = ${skill.id}`); + const left = await db.execute( + sql`SELECT 1 FROM "installed_skill_files" WHERE "skill_id" = ${skill.id}`, + ); + expect(left.length).toBe(0); + }); +}); diff --git a/apps/api/src/db/schema.ts b/apps/api/src/db/schema.ts index bc8116ec9..9506c8dd2 100644 --- a/apps/api/src/db/schema.ts +++ b/apps/api/src/db/schema.ts @@ -201,6 +201,16 @@ function runColumns() { // Control plane: durable reconcile backoff for transient world-read failures. reconcileBackoffUntil: timestamp("reconcile_backoff_until", { withTimezone: true }), reconcileAttempts: integer("reconcile_attempts").notNull().default(0), + // Scale-out (docs/plans/scale-out.md): the run protocol and attachment. + // exec_state: null (not started) | "started" (the supervisor runs; exec_pid + // set) | "exited". consumed_bytes: how much of the pod-side output file + // is log rows. attached_by / attach_lease_until: the instance streaming + // the run, and until when. + execState: text("exec_state").$type<"started" | "exited">(), + execPid: integer("exec_pid"), + consumedBytes: bigint("consumed_bytes", { mode: "number" }).notNull().default(0), + attachedBy: text("attached_by"), + attachLeaseUntil: timestamp("attach_lease_until", { withTimezone: true }), createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), updatedAt: timestamp("updated_at", { withTimezone: true }).notNull().defaultNow(), startedAt: timestamp("started_at", { withTimezone: true }), @@ -966,6 +976,16 @@ export const workflowRuns = pgTable("workflow_runs", { maxRetries: integer("max_retries"), // Stall detection: the attempt's last agent event. lastActivityAt: timestamp("last_activity_at", { withTimezone: true }), + // Scale-out (docs/plans/scale-out.md): the run protocol and attachment. + // exec_state: null (not started) | "started" (the supervisor runs; exec_pid + // set) | "exited". consumed_bytes: how much of the pod-side output file + // is log rows. attached_by / attach_lease_until: the instance streaming + // the run, and until when. + execState: text("exec_state").$type<"started" | "exited">(), + execPid: integer("exec_pid"), + consumedBytes: bigint("consumed_bytes", { mode: "number" }).notNull().default(0), + attachedBy: text("attached_by"), + attachLeaseUntil: timestamp("attach_lease_until", { withTimezone: true }), createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), updatedAt: timestamp("updated_at", { withTimezone: true }).notNull().defaultNow(), }); @@ -1304,6 +1324,94 @@ export const installedSkills = pgTable( ], ); +/** + * A marketplace skill's files, written by the sync worker at each resolved + * sha and read at spawn (docs/plans/scale-out.md: the cache PVC was + * ReadWriteOnce, so a skill was only there on the one instance that synced). + */ +export const installedSkillFiles = pgTable( + "installed_skill_files", + { + skillId: uuid("skill_id") + .notNull() + .references(() => installedSkills.id, { onDelete: "cascade" }), + path: text("path").notNull(), + content: customType<{ data: Buffer; driverData: Buffer }>({ + dataType: () => "bytea", + })("content").notNull(), + executable: boolean("executable").notNull().default(false), + resolvedSha: text("resolved_sha"), + }, + (table) => [primaryKey({ columns: [table.skillId, table.path] })], +); + +// ── Scale-out: shared coordination state (docs/plans/scale-out.md §3) ──────── + +/** What one instance holds for a while (services/lease-service.ts). */ +export const leases = pgTable("leases", { + key: text("key").primaryKey(), + holder: text("holder").notNull(), + expiresAt: timestamp("expires_at", { withTimezone: true }).notNull(), +}); + +/** Single-use WebSocket upgrade tokens, usable on any instance. */ +export const wsUpgradeTokens = pgTable( + "ws_upgrade_tokens", + { + tokenHash: text("token_hash").primaryKey(), + userId: uuid("user_id").notNull(), + workspaceId: uuid("workspace_id"), + expiresAt: timestamp("expires_at", { withTimezone: true }).notNull(), + }, + (table) => [index("ws_upgrade_tokens_expires_idx").on(table.expiresAt)], +); + +/** + * Inbound webhook deliveries already handled (by source and the provider's + * delivery id): an insert that returns a row is the claim. Not to be + * confused with `webhookDeliveries`, the outbound webhooks' delivery log. + */ +export const inboundWebhookDeliveries = pgTable( + "inbound_webhook_deliveries", + { + source: text("source").notNull(), + deliveryId: text("delivery_id").notNull(), + receivedAt: timestamp("received_at", { withTimezone: true }).notNull().defaultNow(), + }, + (table) => [ + primaryKey({ columns: [table.source, table.deliveryId] }), + index("inbound_webhook_deliveries_received_idx").on(table.receivedAt), + ], +); + +/** Glance / push bookkeeping by key; a timer is a row with `due_at`. */ +export const glanceState = pgTable( + "glance_state", + { + key: text("key").primaryKey(), + value: jsonb("value").$type>().notNull(), + dueAt: timestamp("due_at", { withTimezone: true }), + }, + (table) => [ + index("glance_state_due_idx") + .on(table.dueAt) + .where(sql`${table.dueAt} IS NOT NULL`), + ], +); + +/** The ticket sync's claim on (source, ticket, repo): one task per claim, however many sweeps overlap. */ +export const ticketSyncClaims = pgTable( + "ticket_sync_claims", + { + ticketSource: text("ticket_source").notNull(), + ticketExternalId: text("ticket_external_id").notNull(), + repoUrl: text("repo_url").notNull(), + taskId: uuid("task_id"), + createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), + }, + (table) => [primaryKey({ columns: [table.ticketSource, table.ticketExternalId, table.repoUrl] })], +); + // ── API Keys (CLI personal access tokens + user-created keys) ───────────────── export const apiKeys = pgTable( @@ -1431,6 +1539,10 @@ export const prReviews = pgTable( index("pr_reviews_workspace_idx").on(table.workspaceId), index("pr_reviews_state_idx").on(table.state), index("pr_reviews_pr_url_idx").on(table.prUrl), + // One active review per PR URL: overlapping sweeps can't both launch one. + uniqueIndex("pr_reviews_active_pr_url_key") + .on(table.prUrl) + .where(sql`${table.state} IN ('queued', 'waiting_ci', 'reviewing', 'ready')`), index("pr_reviews_repo_url_idx").on(table.repoUrl), index("pr_reviews_updated_idx").on(table.updatedAt.desc()), ], @@ -1469,6 +1581,16 @@ export const prReviewRuns = pgTable( modelUsed: text("model_used"), metadata: jsonb("metadata").$type>(), lastActivityAt: timestamp("last_activity_at", { withTimezone: true }), + // Scale-out (docs/plans/scale-out.md): the run protocol and attachment. + // exec_state: null (not started) | "started" (the supervisor runs; exec_pid + // set) | "exited". consumed_bytes: how much of the pod-side output file + // is log rows. attached_by / attach_lease_until: the instance streaming + // the run, and until when. + execState: text("exec_state").$type<"started" | "exited">(), + execPid: integer("exec_pid"), + consumedBytes: bigint("consumed_bytes", { mode: "number" }).notNull().default(0), + attachedBy: text("attached_by"), + attachLeaseUntil: timestamp("attach_lease_until", { withTimezone: true }), startedAt: timestamp("started_at", { withTimezone: true }), completedAt: timestamp("completed_at", { withTimezone: true }), createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), @@ -1845,6 +1967,16 @@ export const persistentAgentTurns = pgTable( outputTokens: integer("output_tokens"), sessionId: text("session_id"), summary: text("summary"), + // Scale-out (docs/plans/scale-out.md): the run protocol and attachment. + // exec_state: null (not started) | "started" (the supervisor runs; exec_pid + // set) | "exited". consumed_bytes: how much of the pod-side output file + // is log rows. attached_by / attach_lease_until: the instance streaming + // the run, and until when. + execState: text("exec_state").$type<"started" | "exited">(), + execPid: integer("exec_pid"), + consumedBytes: bigint("consumed_bytes", { mode: "number" }).notNull().default(0), + attachedBy: text("attached_by"), + attachLeaseUntil: timestamp("attach_lease_until", { withTimezone: true }), startedAt: timestamp("started_at", { withTimezone: true }), finishedAt: timestamp("finished_at", { withTimezone: true }), createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), diff --git a/apps/api/src/index.ts b/apps/api/src/index.ts index 53c0de4b9..d57813165 100644 --- a/apps/api/src/index.ts +++ b/apps/api/src/index.ts @@ -2,6 +2,7 @@ import "dotenv/config"; import { parseIntEnv } from "@optio/shared"; import { initTelemetry, shutdownTelemetry, registerMetricCallbacks } from "./telemetry.js"; import { logger } from "./logger.js"; +import { INSTANCE_ID } from "./services/instance.js"; // Prevent Redis connection errors from crashing the process process.on("uncaughtException", (err) => { @@ -159,6 +160,7 @@ async function main() { const migrationsPath = join(dirname(fileURLToPath(import.meta.url)), "db", "migrations"); const applied = await migrateSafe(db, migrationsPath); logger.info({ applied }, "Database migrations applied"); + logger.info({ instanceId: INSTANCE_ID }, "API instance"); // Heal contradictory (scope='global', workspace_id IS NOT NULL) secret rows // left over from issue #509. Idempotent — a no-op once the data is clean. diff --git a/apps/api/src/routes/health.test.ts b/apps/api/src/routes/health.test.ts index f342edc18..ceef6d40d 100644 --- a/apps/api/src/routes/health.test.ts +++ b/apps/api/src/routes/health.test.ts @@ -22,8 +22,8 @@ import { healthRoutes, _resetHealthCache } from "./health.js"; // ─── Helpers ─── -async function buildTestApp(): Promise { - return buildRouteTestApp(healthRoutes, { user: null }); +async function buildTestApp(user: unknown = null): Promise { + return buildRouteTestApp(healthRoutes, { user: user as never }); } describe("GET /api/health", () => { @@ -96,3 +96,29 @@ describe("GET /api/health", () => { expect(res.json().maxConcurrent).toBeTypeOf("number"); }); }); + +describe("GET /api/health instance id", () => { + beforeEach(() => { + vi.clearAllMocks(); + _resetHealthCache(); + mockDbExecute.mockResolvedValue([{ "?column?": 1 }]); + mockCheckRuntimeHealth.mockResolvedValue(true); + }); + + it("names the instance only for a signed-in request", async () => { + const anonymous = await buildTestApp(null); + const res = await anonymous.inject({ method: "GET", url: "/api/health" }); + expect(res.statusCode).toBe(200); + expect(res.json()).not.toHaveProperty("instanceId"); + + const signedIn = await buildTestApp({ + id: "u1", + provider: "github", + email: "u@x", + displayName: "U", + }); + const res2 = await signedIn.inject({ method: "GET", url: "/api/health" }); + expect(res2.statusCode).toBe(200); + expect(res2.json().instanceId).toMatch(/^[^:]+:[0-9a-f]{6}$/); + }); +}); diff --git a/apps/api/src/routes/health.ts b/apps/api/src/routes/health.ts index feddaee6a..99f37a9e0 100644 --- a/apps/api/src/routes/health.ts +++ b/apps/api/src/routes/health.ts @@ -6,6 +6,7 @@ import { checkRuntimeHealth } from "../services/container-service.js"; import { sql } from "drizzle-orm"; import { parseIntEnv } from "@optio/shared"; import { isOtelEnabled } from "../telemetry.js"; +import { INSTANCE_ID } from "../services/instance.js"; // Cache runtime health to avoid slow k8s API calls on every health check. // The UI polls this frequently — a 30s TTL is sufficient. @@ -25,6 +26,12 @@ const HealthResponseSchema = z checks: z.record(z.boolean()), maxConcurrent: z.number().int(), otelEnabled: z.boolean(), + /** + * Which API process answered (services/instance.ts): `:`. Only for a signed-in request; the public probe gets nothing + * that identifies a pod. + */ + instanceId: z.string().optional(), }) .describe("API health probe result"); @@ -51,7 +58,7 @@ export async function healthRoutes(rawApp: FastifyInstance) { }, }, }, - async (_req, reply) => { + async (req, reply) => { const checks: Record = {}; try { @@ -81,9 +88,13 @@ export async function healthRoutes(rawApp: FastifyInstance) { // state but should not cause liveness/readiness probes to fail. const healthy = checks.database; const maxConcurrent = parseIntEnv("OPTIO_MAX_CONCURRENT", 5); - reply - .status(healthy ? 200 : 503) - .send({ healthy, checks, maxConcurrent, otelEnabled: isOtelEnabled() }); + reply.status(healthy ? 200 : 503).send({ + healthy, + checks, + maxConcurrent, + otelEnabled: isOtelEnabled(), + ...(req.user ? { instanceId: INSTANCE_ID } : {}), + }); }, ); } diff --git a/apps/api/src/services/instance.ts b/apps/api/src/services/instance.ts new file mode 100644 index 000000000..618d4047b --- /dev/null +++ b/apps/api/src/services/instance.ts @@ -0,0 +1,18 @@ +/** + * Which API process this is (docs/plans/scale-out.md). Every API pod shares + * one Postgres and one Redis; what tells them apart is this id: the holder + * of a lease, the instance a run is attached to, the name in a log line. + * + * `OPTIO_INSTANCE_ID` names the pod (the chart sets it from the pod's name); + * the process always appends a suffix of its own, so a restarted pod is a + * new instance and nothing it held is mistaken for the old one's. Without + * the env var the host name stands in for the pod name. + */ +import { randomBytes } from "node:crypto"; +import { hostname } from "node:os"; + +const base = process.env.OPTIO_INSTANCE_ID?.trim() || hostname(); +const suffix = randomBytes(3).toString("hex"); + +/** This process's id: `:<6 hex, new at every boot>`. */ +export const INSTANCE_ID: string = `${base}:${suffix}`; diff --git a/apps/api/src/services/lease-service.int.test.ts b/apps/api/src/services/lease-service.int.test.ts new file mode 100644 index 000000000..a066d65b5 --- /dev/null +++ b/apps/api/src/services/lease-service.int.test.ts @@ -0,0 +1,159 @@ +/** + * Leases against the real database: two instances contending for one key, + * expiry, renewal, release, seizure, and `withLease` giving up cleanly when + * the lease is taken from under it. + */ +import { afterEach, describe, expect, it } from "vitest"; +import { sql } from "drizzle-orm"; +import { db } from "../db/client.js"; +import { + LeaseLostError, + acquireLease, + leaseHolder, + releaseLease, + renewLease, + seizeLease, + sweepLeases, + withLease, +} from "./lease-service.js"; + +const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); +const A = "api-a:111111"; +const B = "api-b:222222"; + +afterEach(async () => { + await db.execute(sql`DELETE FROM "leases"`); +}); + +describe("leases", () => { + it("exactly one of two contending instances gets a free key", async () => { + const key = "poller:test"; + const [a, b] = await Promise.all([acquireLease(key, 5_000, A), acquireLease(key, 5_000, B)]); + expect([a, b].filter(Boolean)).toHaveLength(1); + const holder = await leaseHolder(key); + expect(holder).toBe(a ? A : B); + }); + + it("the holder can re-acquire (extend); another instance cannot until it expires", async () => { + const key = "poller:extend"; + expect(await acquireLease(key, 300, A)).toBe(true); + expect(await acquireLease(key, 300, B)).toBe(false); + expect(await acquireLease(key, 300, A)).toBe(true); + await sleep(400); + expect(await acquireLease(key, 5_000, B)).toBe(true); + expect(await leaseHolder(key)).toBe(B); + }); + + it("renew extends only a live lease of the holder; release ends it", async () => { + const key = "poller:renew"; + expect(await acquireLease(key, 400, A)).toBe(true); + expect(await renewLease(key, 5_000, B)).toBe(false); + expect(await renewLease(key, 5_000, A)).toBe(true); + await releaseLease(key, B); // not the holder: no effect + expect(await leaseHolder(key)).toBe(A); + await releaseLease(key, A); + expect(await leaseHolder(key)).toBeNull(); + expect(await renewLease(key, 5_000, A)).toBe(false); + expect(await acquireLease(key, 5_000, B)).toBe(true); + }); + + it("an expired lease cannot be renewed back to life", async () => { + const key = "poller:expired"; + expect(await acquireLease(key, 200, A)).toBe(true); + await sleep(300); + expect(await renewLease(key, 5_000, A)).toBe(false); + }); + + it("seize takes a live lease from its holder", async () => { + const key = "local-host:h1"; + expect(await acquireLease(key, 5_000, A)).toBe(true); + await seizeLease(key, 5_000, B); + expect(await leaseHolder(key)).toBe(B); + expect(await renewLease(key, 5_000, A)).toBe(false); + }); + + it("rejects empty and oversized keys", async () => { + await expect(acquireLease("", 1000, A)).rejects.toThrow(/Invalid lease key/); + await expect(acquireLease("x".repeat(201), 1000, A)).rejects.toThrow(/Invalid lease key/); + }); + + it("sweeps leases that expired more than a day ago", async () => { + await db.execute(sql` + INSERT INTO "leases" ("key", "holder", "expires_at") + VALUES ('old', 'x', now() - interval '2 days'), ('recent', 'x', now() - interval '1 hour')`); + expect(await sweepLeases()).toBe(1); + const rows = await db.execute<{ key: string }>(sql`SELECT "key" FROM "leases" ORDER BY "key"`); + expect(rows.map((r) => r.key)).toEqual(["recent"]); + }); +}); + +describe("withLease", () => { + it("runs fn under the lease, renews while it runs, and releases afterwards", async () => { + const key = "poller:with"; + const result = await withLease( + key, + async () => { + await sleep(700); + expect(await leaseHolder(key)).toBe(A); + return "done"; + }, + { ttlMs: 500, renewEveryMs: 150, holder: A }, + ); + expect(result).toBe("done"); + expect(await leaseHolder(key)).toBeNull(); + }); + + it("returns null without running fn when another instance holds the lease", async () => { + const key = "poller:busy"; + expect(await acquireLease(key, 5_000, B)).toBe(true); + let ran = false; + const result = await withLease( + key, + async () => { + ran = true; + return 1; + }, + { holder: A }, + ); + expect(result).toBeNull(); + expect(ran).toBe(false); + }); + + it("aborts fn and rejects with LeaseLostError when the lease is taken away", async () => { + const key = "poller:lost"; + let aborted = false; + const run = withLease( + key, + async (signal) => { + await new Promise((resolve) => { + signal.addEventListener("abort", () => { + aborted = true; + resolve(); + }); + }); + return "should not count"; + }, + { ttlMs: 1_000, renewEveryMs: 100, holder: A }, + ); + await sleep(150); + await seizeLease(key, 5_000, B); + await expect(run).rejects.toBeInstanceOf(LeaseLostError); + expect(aborted).toBe(true); + // The loser never touches the lease B now holds. + expect(await leaseHolder(key)).toBe(B); + }); + + it("fn's own error propagates and the lease is released", async () => { + const key = "poller:throws"; + await expect( + withLease( + key, + async () => { + throw new Error("boom"); + }, + { holder: A }, + ), + ).rejects.toThrow("boom"); + expect(await leaseHolder(key)).toBeNull(); + }); +}); diff --git a/apps/api/src/services/lease-service.ts b/apps/api/src/services/lease-service.ts new file mode 100644 index 000000000..1b1bf4409 --- /dev/null +++ b/apps/api/src/services/lease-service.ts @@ -0,0 +1,161 @@ +/** + * Leases (docs/plans/scale-out.md §3): the one way an API instance claims + * something for a while when several instances share the database — a + * poller's turn, a per-user chat lock, a Local host's owner. + * + * A lease is a row in `leases` (`key`, `holder`, `expires_at`). Acquiring is + * one statement, an insert that becomes an update only when the row has + * expired or already belongs to the same holder, so two instances can race + * for a key and exactly one gets it. Renewing and releasing are updates + * guarded by the holder. Nothing here waits: a caller that loses a lease + * does something else, or tries again later. + */ +import { sql } from "drizzle-orm"; +import { db } from "../db/client.js"; +import { INSTANCE_ID } from "./instance.js"; + +const MAX_KEY = 200; + +function assertKey(key: string): string { + if (!key || key.length > MAX_KEY) throw new Error(`Invalid lease key: ${JSON.stringify(key)}`); + return key; +} + +/** + * Takes the lease for `holder` (this instance by default) for `ttlMs` when + * nobody holds it, it has expired, or `holder` already holds it (which + * extends it). True iff the lease is now held by `holder`. + */ +export async function acquireLease( + key: string, + ttlMs: number, + holder: string = INSTANCE_ID, +): Promise { + const rows = await db.execute<{ key: string }>(sql` + INSERT INTO "leases" ("key", "holder", "expires_at") + VALUES (${assertKey(key)}, ${holder}, now() + make_interval(secs => ${ttlMs / 1000})) + ON CONFLICT ("key") DO UPDATE + SET "holder" = EXCLUDED."holder", "expires_at" = EXCLUDED."expires_at" + WHERE "leases"."expires_at" < now() OR "leases"."holder" = EXCLUDED."holder" + RETURNING "key"`); + return rows.length > 0; +} + +/** Extends a lease `holder` holds. False when it was lost (expired and taken, or released). */ +export async function renewLease( + key: string, + ttlMs: number, + holder: string = INSTANCE_ID, +): Promise { + const rows = await db.execute<{ key: string }>(sql` + UPDATE "leases" + SET "expires_at" = now() + make_interval(secs => ${ttlMs / 1000}) + WHERE "key" = ${assertKey(key)} AND "holder" = ${holder} AND "expires_at" >= now() + RETURNING "key"`); + return rows.length > 0; +} + +/** Gives a lease up right away, if `holder` holds it. */ +export async function releaseLease(key: string, holder: string = INSTANCE_ID): Promise { + await db.execute(sql` + UPDATE "leases" SET "expires_at" = now() + WHERE "key" = ${assertKey(key)} AND "holder" = ${holder}`); +} + +/** Who holds a lease right now, or null when nobody does. */ +export async function leaseHolder(key: string): Promise { + const rows = await db.execute<{ holder: string }>(sql` + SELECT "holder" FROM "leases" WHERE "key" = ${assertKey(key)} AND "expires_at" > now()`); + return rows[0]?.holder ?? null; +} + +/** + * Takes a lease by force — for "newest daemon wins" (the Local relay): the + * previous holder finds out at its next renew. + */ +export async function seizeLease( + key: string, + ttlMs: number, + holder: string = INSTANCE_ID, +): Promise { + await db.execute(sql` + INSERT INTO "leases" ("key", "holder", "expires_at") + VALUES (${assertKey(key)}, ${holder}, now() + make_interval(secs => ${ttlMs / 1000})) + ON CONFLICT ("key") DO UPDATE + SET "holder" = EXCLUDED."holder", "expires_at" = EXCLUDED."expires_at"`); +} + +/** Drops leases that expired more than a day ago (housekeeping). */ +export async function sweepLeases(): Promise { + const rows = await db.execute<{ key: string }>(sql` + DELETE FROM "leases" WHERE "expires_at" < now() - interval '1 day' RETURNING "key"`); + return rows.length; +} + +export class LeaseLostError extends Error { + constructor(public readonly key: string) { + super(`lease ${key} was lost`); + this.name = "LeaseLostError"; + } +} + +export interface WithLeaseOptions { + /** How long each acquire / renew holds the lease (default 60 s). */ + ttlMs?: number; + /** How often to renew while `fn` runs (default a third of the TTL). */ + renewEveryMs?: number; + holder?: string; +} + +/** + * Runs `fn` while holding the lease, renewing it on an interval. Resolves + * `null` without running `fn` when the lease is held elsewhere. If a renew + * fails while `fn` runs, `fn`'s `signal` aborts and, once `fn` settles, the + * call rejects with `LeaseLostError`; `fn` should stop what it is doing when + * the signal fires. The lease is released when `fn` finishes. + */ +export async function withLease( + key: string, + fn: (signal: AbortSignal) => Promise, + opts: WithLeaseOptions = {}, +): Promise { + const ttlMs = opts.ttlMs ?? 60_000; + const renewEveryMs = opts.renewEveryMs ?? Math.max(250, Math.floor(ttlMs / 3)); + const holder = opts.holder ?? INSTANCE_ID; + if (!(await acquireLease(key, ttlMs, holder))) return null; + + const controller = new AbortController(); + let lost = false; + let renewing: Promise = Promise.resolve(); + const timer = setInterval(() => { + renewing = renewing.then(async () => { + if (lost) return; + let ok = false; + try { + ok = await renewLease(key, ttlMs, holder); + } catch { + // A failed renew (the database away for a moment) is not a lost + // lease: the next tick tries again, and the TTL covers the gap. + return; + } + if (!ok) { + lost = true; + controller.abort(new LeaseLostError(key)); + } + }); + }, renewEveryMs); + timer.unref?.(); + + try { + const result = await fn(controller.signal); + if (lost) throw new LeaseLostError(key); + return result; + } catch (err) { + if (lost && !(err instanceof LeaseLostError)) throw new LeaseLostError(key); + throw err; + } finally { + clearInterval(timer); + await renewing.catch(() => {}); + if (!lost) await releaseLease(key, holder).catch(() => {}); + } +} diff --git a/apps/api/src/services/persistent-agent-pool-service.ts b/apps/api/src/services/persistent-agent-pool-service.ts index c71e269cc..0fba9d1e4 100644 --- a/apps/api/src/services/persistent-agent-pool-service.ts +++ b/apps/api/src/services/persistent-agent-pool-service.ts @@ -23,7 +23,7 @@ import { PersistentAgentPodLifecycle, parseIntEnv, type RepoImageConfig } from " import { logger } from "../logger.js"; import { resolveImage } from "./repo-pool-service.js"; import * as podPool from "./agent-pod-pool.js"; -import { buildPooledExecScript } from "../utils/pod-env.js"; +import { buildPooledExecScript, KEEP_POD_ALIVE } from "../utils/pod-env.js"; const POD_PROVISION_TIMEOUT_MS = parseIntEnv("OPTIO_PERSISTENT_AGENT_POD_PROVISION_MS", 120_000); @@ -203,7 +203,7 @@ function buildInitScript(): string { "mkdir -p /workspace/turns", "touch /workspace/.ready", "echo '[optio] Persistent agent pod ready'", - "exec sleep infinity", + KEEP_POD_ALIVE, ].join("\n"); } diff --git a/apps/api/src/services/pr-review-service.int.test.ts b/apps/api/src/services/pr-review-service.int.test.ts new file mode 100644 index 000000000..b8ae60961 --- /dev/null +++ b/apps/api/src/services/pr-review-service.int.test.ts @@ -0,0 +1,63 @@ +/** + * The review row a launch builds on, against the real database: an active + * review is picked over a historical cancelled duplicate of the same PR URL + * whatever their order; with no active one, the newest. + */ +import { describe, expect, it } from "vitest"; +import { sql } from "drizzle-orm"; +import { db } from "../db/client.js"; +import { prReviews } from "../db/schema.js"; +import { findReviewByUrl } from "./pr-review-service.js"; + +async function insertReview(prUrl: string, state: string, updatedAt: Date) { + const [row] = await db + .insert(prReviews) + .values({ + prUrl, + prNumber: 7, + repoOwner: "o", + repoName: "r", + repoUrl: "https://github.com/o/r", + headSha: "abc", + state: state as never, + updatedAt, + }) + .returning(); + return row; +} + +describe("findReviewByUrl", () => { + it("prefers the active review over a newer cancelled duplicate", async () => { + const url = "https://github.com/o/r/pull/7"; + const older = new Date(Date.now() - 3_600_000); + const newer = new Date(); + const active = await insertReview(url, "reviewing", older); + await insertReview(url, "cancelled", newer); + expect((await findReviewByUrl(url))?.id).toBe(active.id); + }); + + it("prefers the active review over an older cancelled duplicate", async () => { + const url = "https://github.com/o/r/pull/8"; + await insertReview(url, "cancelled", new Date(Date.now() - 3_600_000)); + const active = await insertReview(url, "ready", new Date()); + expect((await findReviewByUrl(url))?.id).toBe(active.id); + }); + + it("with no active review, picks the most recently updated", async () => { + const url = "https://github.com/o/r/pull/9"; + await insertReview(url, "cancelled", new Date(Date.now() - 3_600_000)); + const newest = await insertReview(url, "submitted", new Date()); + expect((await findReviewByUrl(url))?.id).toBe(newest.id); + expect(await findReviewByUrl("https://github.com/o/r/pull/404")).toBeUndefined(); + }); + + it("the unique index keeps a second active review out", async () => { + const url = "https://github.com/o/r/pull/10"; + await insertReview(url, "queued", new Date()); + await expect(insertReview(url, "reviewing", new Date())).rejects.toSatisfy((err: unknown) => { + const cause = (err as { cause?: { message?: string } }).cause; + return /pr_reviews_active_pr_url_key/.test(cause?.message ?? (err as Error).message); + }); + await db.execute(sql`DELETE FROM "pr_reviews" WHERE "pr_url" = ${url}`); + }); +}); diff --git a/apps/api/src/services/pr-review-service.ts b/apps/api/src/services/pr-review-service.ts index b55517da8..3c11ed134 100644 --- a/apps/api/src/services/pr-review-service.ts +++ b/apps/api/src/services/pr-review-service.ts @@ -312,6 +312,38 @@ interface LaunchPrReviewInput { * pr_reviews row already exists for the PR URL, promote it to reviewing * and spawn a rereview run rather than duplicating. */ +/** The states a review is "in flight" in: at most one of these per PR URL (`pr_reviews_active_pr_url_key`). */ +export const ACTIVE_PR_REVIEW_STATES: readonly PrReviewState[] = [ + PrReviewState.QUEUED, + PrReviewState.WAITING_CI, + PrReviewState.REVIEWING, + PrReviewState.READY, +]; + +/** + * The review row a launch builds on for a PR URL: the active one when there + * is one, else the most recently updated. A historical cancelled duplicate + * (overlapping sweeps before the unique index) is never picked over a live + * review. + */ +export async function findReviewByUrl(prUrl: string): Promise { + const [row] = await db + .select() + .from(prReviews) + .where(eq(prReviews.prUrl, prUrl)) + .orderBy( + desc( + sql`${prReviews.state} IN (${sql.join( + ACTIVE_PR_REVIEW_STATES.map((s) => sql`${s}`), + sql`, `, + )})`, + ), + desc(prReviews.updatedAt), + ) + .limit(1); + return row; +} + export async function launchPrReview(input: LaunchPrReviewInput) { const parsed = parsePrUrl(input.prUrl); if (!parsed) throw new Error("Invalid PR URL"); @@ -332,7 +364,7 @@ export async function launchPrReview(input: LaunchPrReviewInput) { const origin = input.origin ?? "manual"; // Upsert pr_reviews row. - const [existing] = await db.select().from(prReviews).where(eq(prReviews.prUrl, input.prUrl)); + const existing = await findReviewByUrl(input.prUrl); let review: PrReview; if (existing) { // Reset to queued/reviewing with fresh head_sha. Keep origin sticky — diff --git a/apps/api/src/services/repo-pool-service.test.ts b/apps/api/src/services/repo-pool-service.test.ts index 2ec459061..59795a660 100644 --- a/apps/api/src/services/repo-pool-service.test.ts +++ b/apps/api/src/services/repo-pool-service.test.ts @@ -153,6 +153,8 @@ import { execTaskInRepoPod, type RepoPod, } from "./repo-pool-service.js"; +import { buildTaskExecScript, buildTaskStartScript } from "./repo-pool-service.js"; +import { superviseAgent } from "../utils/pod-env.js"; const REPO_URL = "https://github.com/org/repo"; @@ -1138,3 +1140,69 @@ describe("execTaskInRepoPod", () => { } }); }); + +describe("task scripts: exec (inline agent) and start (the run protocol)", () => { + const taskId = "0f2e1c1a-1111-4222-8333-444444444444"; + const env = { OPTIO_REPO_URL: "https://github.com/o/r", OPTIO_REPO_BRANCH: "main" }; + const agent = ["export X='1'", "claude -p --output-format stream-json"]; + + it("both do the same setup: lock, worktree, credential helper, setup files, git exclude", () => { + const exec = buildTaskExecScript(taskId, env, agent, { runToken: "tok" }); + const start = buildTaskStartScript(taskId, env, agent, { runToken: "tok" }); + for (const script of [exec, start]) { + expect(script).toContain(`exec 7>'/workspace/.task-${taskId}.lock'`); + expect(script).toContain( + `git worktree add /workspace/tasks/${taskId} -b optio/task-${taskId} "origin/main"`, + ); + expect(script).toContain(`cd /workspace/tasks/${taskId}`); + expect(script).toContain( + "git config --local credential.helper '/usr/local/bin/optio-git-credential'", + ); + expect(script).toContain(`echo "tok" > /workspace/tasks/${taskId}/.optio-run-token`); + expect(script).toContain(`export OPTIO_TASK_ID="${taskId}"`); + expect(script).toContain(`grep -qxF '.optio/' "$EXCLUDE_FILE"`); + } + const setupEnd = `grep -qxF '.optio-cache/' "$EXCLUDE_FILE" 2>/dev/null || echo '.optio-cache/' >> "$EXCLUDE_FILE"`; + // The start script prunes the stale `-wt` worktree up front (the exec + // script's EXIT trap did it at the end); otherwise the setup is the same. + const prune = `git worktree remove --force /workspace/tasks/${taskId}-wt 2>/dev/null || true\ngit worktree prune 2>/dev/null || true\n`; + expect(start).toContain(prune); + expect(exec.slice(0, exec.indexOf(setupEnd))).toBe( + start.slice(0, start.indexOf(setupEnd)).replace(prune, ""), + ); + }); + + it("the exec script runs the agent inline with the EXIT trap and the EPIPE watchdog, as before", () => { + const script = buildTaskExecScript(taskId, env, agent); + expect(script).toContain("_optio_main_pid=$$"); + expect(script).toContain( + `git worktree remove --force /workspace/tasks/${taskId}-wt 2>/dev/null || true; git worktree prune 2>/dev/null || true;`, + ); + expect(script).toContain("kill -TERM $_optio_main_pid"); + expect(script).toContain( + "set +e\nexport X='1'\nclaude -p --output-format stream-json\nAGENT_EXIT=$?", + ); + expect(script).toContain("[ $AGENT_EXIT -eq 0 ] && touch /home/agent/.optio-env-ready"); + expect(script.trim().endsWith("exit $AGENT_EXIT")).toBe(true); + }); + + it("the start script hands the agent to the supervisor in the task's run home and marks the env warm on success", () => { + const script = buildTaskStartScript(taskId, env, agent); + expect(script).toContain( + superviseAgent(`/home/agent/optio/runs/${taskId}`, agent, { markEnvReady: true }).join("\n"), + ); + expect(script).not.toContain("_optio_main_pid"); + expect(script).not.toContain(" EXIT"); + expect(script).not.toContain("AGENT_EXIT=$?"); + expect(script).not.toContain("kill -TERM $_optio_main_pid"); + expect(script.trim().endsWith('echo "__OPTIO_RUN_STARTED__:$_optio_sup"')).toBe(true); + }); + + it("a same-pod retry resets the worktree in both", () => { + for (const build of [buildTaskExecScript, buildTaskStartScript]) { + const script = build(taskId, env, agent, { resetWorktree: true }); + expect(script).toContain("[optio] Resetting existing worktree for retry..."); + expect(script).not.toContain(`rm -rf /workspace/tasks/${taskId}`); + } + }); +}); diff --git a/apps/api/src/services/repo-pool-service.ts b/apps/api/src/services/repo-pool-service.ts index 93638ce17..fd27ddab8 100644 --- a/apps/api/src/services/repo-pool-service.ts +++ b/apps/api/src/services/repo-pool-service.ts @@ -32,8 +32,13 @@ import { } from "./envoy-sidecar.js"; import { parseIntEnv } from "@optio/shared"; import { withSpan } from "../telemetry/spans.js"; -import { buildEnvExports, RUN_WORK_SETUP_COMMANDS, WRITE_SETUP_FILES } from "../utils/pod-env.js"; -import { REMOVE_RUN_HOME, removeRunHome } from "../utils/harness-config.js"; +import { + buildEnvExports, + RUN_WORK_SETUP_COMMANDS, + superviseAgent, + WRITE_SETUP_FILES, +} from "../utils/pod-env.js"; +import { REMOVE_RUN_HOME, removeRunHome, runHome } from "../utils/harness-config.js"; import { resolveSecretsForSetup, workspaceRestrictsPodSecrets } from "./secret-service.js"; /** The run env a repo pod starts with: git sign-in and the repo's setup. */ @@ -755,6 +760,229 @@ async function createRepoPodViaStatefulSet( } } +export interface TaskScriptOptions { + /** On a same-pod retry: reset the existing worktree instead of recreating it. */ + resetWorktree?: boolean; + /** The run token written to the worktree (a fresh one when omitted). */ + runToken?: string; + /** + * Remove the task's internal `-wt` worktree and prune stale worktree + * records after the worktree setup (what the exec script's EXIT trap does + * at the end; the start script has no trap, so it does it up front). + */ + pruneStaleWorktrees?: boolean; +} + +/** + * The setup every task run does in its repo pod before its agent starts: + * the task lock, the repo lock and worktree, the credential helper and CLI + * wrappers, the setup files, and git's exclude list. Shared by the exec + * script (the agent inline, in the exec's lifetime) and the start script + * (the agent under the run supervisor, outliving the exec). + */ +function taskSetupLines( + taskId: string, + env: Record, + opts: TaskScriptOptions, +): string[] { + // Build the exec command. Env values (including the task prompt) are + // embedded as inert single-quoted exports — see buildEnvExports. + const envExports = buildEnvExports({ + ...env, + OPTIO_TASK_ID: taskId, + REPO_INIT_TIMEOUT_SECS: String(Math.ceil(REPO_INIT_TIMEOUT_MS / 1000)), + }); + const runToken = opts.runToken ?? randomUUID(); + + // Build worktree setup commands based on whether we're resetting or creating fresh + const worktreeSetup = opts.resetWorktree + ? [ + `if [ -d "/workspace/tasks/${taskId}" ]; then`, + ` echo "[optio] Resetting existing worktree for retry..."`, + ` cd /workspace/tasks/${taskId}`, + ` git checkout -- . 2>/dev/null || true`, + ` git clean -fd 2>/dev/null || true`, + ` cd /workspace/repo`, + ` echo "[optio] Worktree reset complete"`, + `else`, + ` echo "[optio] No existing worktree found, creating fresh..."`, + ` git branch -D optio/task-${taskId} 2>/dev/null || true`, + ` if ! git worktree add /workspace/tasks/${taskId} -b optio/task-${taskId} "origin/${env.OPTIO_REPO_BRANCH ?? "main"}" 2>/dev/null; then`, + ` echo "[optio] Cleaning up stale worktree references..."`, + ` git worktree remove --force /workspace/tasks/${taskId}-wt 2>/dev/null || true`, + ` for wt_path in $(git worktree list --porcelain | grep -B1 "branch refs/heads/optio/task-${taskId}$" | grep "^worktree " | cut -d" " -f2-); do`, + ` git worktree remove --force "$wt_path" 2>/dev/null || true`, + ` done`, + ` git worktree prune`, + ` git branch -D optio/task-${taskId} 2>/dev/null || true`, + ` git worktree add /workspace/tasks/${taskId} -b optio/task-${taskId} "origin/${env.OPTIO_REPO_BRANCH ?? "main"}"`, + ` fi`, + `fi`, + ] + : [ + `git worktree remove --force /workspace/tasks/${taskId} 2>/dev/null || true`, + `rm -rf /workspace/tasks/${taskId}`, + removeRunHome(taskId), + `if [ "\${OPTIO_RESTART_FROM_BRANCH:-}" = "true" ] && git rev-parse --verify origin/optio/task-${taskId} >/dev/null 2>&1; then`, + ` echo "[optio] Force-restart: checking out existing PR branch"`, + ` for wt_path in $(git worktree list --porcelain | grep -B1 "branch refs/heads/optio/task-${taskId}$" | grep "^worktree " | cut -d" " -f2-); do`, + ` git worktree remove --force "$wt_path" 2>/dev/null || true`, + ` done`, + ` git worktree prune`, + ` git branch -D optio/task-${taskId} 2>/dev/null || true`, + ` git worktree add /workspace/tasks/${taskId} -b optio/task-${taskId} origin/optio/task-${taskId}`, + `else`, + ` git branch -D optio/task-${taskId} 2>/dev/null || true`, + ` if ! git worktree add /workspace/tasks/${taskId} -b optio/task-${taskId} "origin/${env.OPTIO_REPO_BRANCH ?? "main"}" 2>/dev/null; then`, + ` echo "[optio] Cleaning up stale worktree references..."`, + ` git worktree remove --force /workspace/tasks/${taskId}-wt 2>/dev/null || true`, + ` for wt_path in $(git worktree list --porcelain | grep -B1 "branch refs/heads/optio/task-${taskId}$" | grep "^worktree " | cut -d" " -f2-); do`, + ` git worktree remove --force "$wt_path" 2>/dev/null || true`, + ` done`, + ` git worktree prune`, + ` git branch -D optio/task-${taskId} 2>/dev/null || true`, + ` git worktree add /workspace/tasks/${taskId} -b optio/task-${taskId} "origin/${env.OPTIO_REPO_BRANCH ?? "main"}"`, + ` fi`, + `fi`, + ]; + + const setup = [ + "set -e", + ...envExports, + `exec 7>${shellQuote(`/workspace/.task-${taskId}.lock`)}`, + "flock -n 7 || { echo 'Previous task process is still running; inspect it before retrying.' >&2; exit 75; }", + `echo "[optio] Waiting for repo to be ready..."`, + `for i in $(seq 1 \${REPO_INIT_TIMEOUT_SECS}); do [ -f /workspace/.ready ] && break; sleep 1; done`, + `[ -f /workspace/.ready ] || { echo "[optio] ERROR: repo not ready after \${REPO_INIT_TIMEOUT_SECS}s (increase OPTIO_REPO_INIT_TIMEOUT_MS to extend)"; exit 1; }`, + `echo "[optio] Repo ready"`, + // Use task-scoped credential URL for git operations (user-scoped token). + // Override the pod-level URL which returns an installation token. + `if [ -n "\${OPTIO_GIT_TASK_CREDENTIAL_URL:-}" ]; then`, + ` export OPTIO_GIT_CREDENTIAL_URL="\${OPTIO_GIT_TASK_CREDENTIAL_URL}"`, + `fi`, + // Set up gh CLI wrapper with PATH prepend (no root required) + `if [ -f /usr/local/bin/optio-gh-wrapper ]; then`, + ` mkdir -p /home/agent/.local/bin`, + ` cp /usr/local/bin/optio-gh-wrapper /home/agent/.local/bin/gh 2>/dev/null || true`, + ` chmod +x /home/agent/.local/bin/gh 2>/dev/null || true`, + ` export PATH="/home/agent/.local/bin:$PATH"`, + `fi`, + // Set up glab CLI wrapper and authenticate for GitLab repos + `if [ -n "\${GITLAB_TOKEN:-}" ]; then`, + ` if [ -f /usr/local/bin/optio-glab-wrapper ]; then`, + ` mkdir -p /home/agent/.local/bin`, + ` cp /usr/local/bin/optio-glab-wrapper /home/agent/.local/bin/glab 2>/dev/null || true`, + ` chmod +x /home/agent/.local/bin/glab 2>/dev/null || true`, + ` export PATH="/home/agent/.local/bin:$PATH"`, + ` fi`, + ` GITLAB_HOST="gitlab.com"`, + ` case "\${OPTIO_REPO_URL:-}" in *://*) GITLAB_HOST=$(echo "\${OPTIO_REPO_URL}" | sed -E 's|.*://([^/]+).*|\\1|');; esac`, + ` glab auth login --hostname "\${GITLAB_HOST}" --token "\${GITLAB_TOKEN}" 2>/dev/null || true`, + `fi`, + `ENV_FRESH="true"`, + `[ -f /home/agent/.optio-env-ready ] && ENV_FRESH="false"`, + `export ENV_FRESH`, + `if [ "$ENV_FRESH" = "true" ]; then echo "[optio] Fresh environment — tools may need to be installed"; else echo "[optio] Warm environment — tools from previous tasks should be available"; fi`, + `echo "[optio] Acquiring repo lock..."`, + `exec 9>/workspace/.repo-lock`, + `flock 9`, + `echo "[optio] Repo lock acquired"`, + `cd /workspace/repo`, + `git fetch origin`, + `git checkout "${env.OPTIO_REPO_BRANCH ?? "main"}" 2>/dev/null || true`, + `git reset --hard "origin/${env.OPTIO_REPO_BRANCH ?? "main"}"`, + ...worktreeSetup, + ...(opts.pruneStaleWorktrees + ? [ + `git worktree remove --force /workspace/tasks/${taskId}-wt 2>/dev/null || true`, + `git worktree prune 2>/dev/null || true`, + ] + : []), + `if [ -f /workspace/repo/.gitmodules ]; then git -C /workspace/tasks/${taskId} submodule update --init --recursive 2>&1 || true; fi`, + `flock -u 9`, + `exec 9>&-`, + `cd /workspace/tasks/${taskId}`, + // Configure git at worktree scope so concurrent tasks don't interfere + `if [ -n "\${OPTIO_GIT_CREDENTIAL_URL:-}" ] && [ -f /usr/local/bin/optio-git-credential ]; then`, + ` git config --local credential.helper '/usr/local/bin/optio-git-credential'`, + ` echo "[optio] Worktree credential helper configured"`, + `fi`, + `git config --local user.name "\${GITHUB_APP_BOT_NAME:-Optio Agent}"`, + `git config --local user.email "\${GITHUB_APP_BOT_EMAIL:-optio-agent@noreply.github.com}"`, + `echo "${runToken}" > /workspace/tasks/${taskId}/.optio-run-token`, + `export OPTIO_TASK_ID="${taskId}"`, + ...WRITE_SETUP_FILES, + // Exclude Optio runtime files from git tracking using the local exclude file + // (never committed, unlike .gitignore modifications) + `EXCLUDE_FILE="$(git rev-parse --git-dir)/info/exclude"`, + `mkdir -p "$(dirname "$EXCLUDE_FILE")"`, + `grep -qxF '.optio/' "$EXCLUDE_FILE" 2>/dev/null || echo '.optio/' >> "$EXCLUDE_FILE"`, + `grep -qxF '.optio-run-token' "$EXCLUDE_FILE" 2>/dev/null || echo '.optio-run-token' >> "$EXCLUDE_FILE"`, + `grep -qxF '.optio-cache/' "$EXCLUDE_FILE" 2>/dev/null || echo '.optio-cache/' >> "$EXCLUDE_FILE"`, + ]; + return setup; +} + +/** + * The classic exec script: runs the agent inline, kills it when the exec + * stream breaks (the EPIPE watchdog), and removes the run home on exit. + * Kept until every worker is on the run protocol (docs/plans/scale-out.md). + */ +export function buildTaskExecScript( + taskId: string, + env: Record, + agentCommand: string[], + opts: TaskScriptOptions = {}, +): string { + const setup = taskSetupLines(taskId, env, opts); + const script = [ + ...setup, + // EXIT trap: kill child processes (agent + watchdog) then clean up internal worktrees. + // This ensures orphaned agent processes are killed when the exec stream is severed + // (e.g. API pod restart closes the SPDY connection but kubelet doesn't send SIGHUP). + `_optio_main_pid=$$`, + // The run's own home (its runtime's MCP config, Codex / Gemini session + // files) goes with the run; a retry writes it afresh (harness-config.ts). + `trap 'kill $(jobs -p) 2>/dev/null; wait 2>/dev/null; cd /workspace/repo 2>/dev/null; git worktree remove --force /workspace/tasks/${taskId}-wt 2>/dev/null || true; git worktree prune 2>/dev/null || true; ${REMOVE_RUN_HOME}' EXIT`, + // Background heartbeat: detect broken stdout pipe (EPIPE) from severed exec stream. + // Writes an empty line every 30s (skipped by the NDJSON parser). If stdout is broken + // (API pod died), sends SIGTERM to the main script which triggers the EXIT trap. + `(trap '' PIPE; while sleep 30; do printf '\\n' 2>/dev/null || { kill -TERM $_optio_main_pid 2>/dev/null; exit; }; done) &`, + ...RUN_WORK_SETUP_COMMANDS, + `set +e`, + ...agentCommand, + `AGENT_EXIT=$?`, + `[ $AGENT_EXIT -eq 0 ] && touch /home/agent/.optio-env-ready`, + `exit $AGENT_EXIT`, + ]; + return script.join("\n"); +} + +/** + * The start script of the run protocol (run-protocol.ts): the same setup, + * the stale `-wt` worktree and worktree records pruned (the exec script's + * trap did that at exit), then the agent launched under the supervisor in + * the task's run home, so the exec returns as soon as the agent runs and any + * API instance can attach later. No trap, no watchdog. The run home stays + * until the attached worker removes it — when the run is terminal and + * `consumed_bytes` equals the output file's size (phase B) — with the + * run-home sweep (`run-home-sweep-service.ts`) as the backstop. + */ +export function buildTaskStartScript( + taskId: string, + env: Record, + agentCommand: string[], + opts: TaskScriptOptions = {}, +): string { + const setup = taskSetupLines(taskId, env, { ...opts, pruneStaleWorktrees: true }); + const home = runHome(taskId); + return [ + ...setup, + ...RUN_WORK_SETUP_COMMANDS, + ...superviseAgent(home.podDir, agentCommand, { markEnvReady: true }), + ].join("\n"); +} + /** * Execute a task in a repo pod using a git worktree. * Returns an ExecSession for streaming output. @@ -789,152 +1017,9 @@ export async function execTaskInRepoPod( .set({ worktreeState: "active", lastPodId: pod.id, updatedAt: new Date() }) .where(eq(tasks.id, taskId)); - // Build the exec command. Env values (including the task prompt) are - // embedded as inert single-quoted exports — see buildEnvExports. - const envExports = buildEnvExports({ - ...env, - OPTIO_TASK_ID: taskId, - REPO_INIT_TIMEOUT_SECS: String(Math.ceil(REPO_INIT_TIMEOUT_MS / 1000)), + const script = buildTaskExecScript(taskId, env, agentCommand, { + resetWorktree: opts?.resetWorktree, }); - const runToken = randomUUID(); - - // Build worktree setup commands based on whether we're resetting or creating fresh - const worktreeSetup = opts?.resetWorktree - ? [ - `if [ -d "/workspace/tasks/${taskId}" ]; then`, - ` echo "[optio] Resetting existing worktree for retry..."`, - ` cd /workspace/tasks/${taskId}`, - ` git checkout -- . 2>/dev/null || true`, - ` git clean -fd 2>/dev/null || true`, - ` cd /workspace/repo`, - ` echo "[optio] Worktree reset complete"`, - `else`, - ` echo "[optio] No existing worktree found, creating fresh..."`, - ` git branch -D optio/task-${taskId} 2>/dev/null || true`, - ` if ! git worktree add /workspace/tasks/${taskId} -b optio/task-${taskId} "origin/${env.OPTIO_REPO_BRANCH ?? "main"}" 2>/dev/null; then`, - ` echo "[optio] Cleaning up stale worktree references..."`, - ` git worktree remove --force /workspace/tasks/${taskId}-wt 2>/dev/null || true`, - ` for wt_path in $(git worktree list --porcelain | grep -B1 "branch refs/heads/optio/task-${taskId}$" | grep "^worktree " | cut -d" " -f2-); do`, - ` git worktree remove --force "$wt_path" 2>/dev/null || true`, - ` done`, - ` git worktree prune`, - ` git branch -D optio/task-${taskId} 2>/dev/null || true`, - ` git worktree add /workspace/tasks/${taskId} -b optio/task-${taskId} "origin/${env.OPTIO_REPO_BRANCH ?? "main"}"`, - ` fi`, - `fi`, - ] - : [ - `git worktree remove --force /workspace/tasks/${taskId} 2>/dev/null || true`, - `rm -rf /workspace/tasks/${taskId}`, - removeRunHome(taskId), - `if [ "\${OPTIO_RESTART_FROM_BRANCH:-}" = "true" ] && git rev-parse --verify origin/optio/task-${taskId} >/dev/null 2>&1; then`, - ` echo "[optio] Force-restart: checking out existing PR branch"`, - ` for wt_path in $(git worktree list --porcelain | grep -B1 "branch refs/heads/optio/task-${taskId}$" | grep "^worktree " | cut -d" " -f2-); do`, - ` git worktree remove --force "$wt_path" 2>/dev/null || true`, - ` done`, - ` git worktree prune`, - ` git branch -D optio/task-${taskId} 2>/dev/null || true`, - ` git worktree add /workspace/tasks/${taskId} -b optio/task-${taskId} origin/optio/task-${taskId}`, - `else`, - ` git branch -D optio/task-${taskId} 2>/dev/null || true`, - ` if ! git worktree add /workspace/tasks/${taskId} -b optio/task-${taskId} "origin/${env.OPTIO_REPO_BRANCH ?? "main"}" 2>/dev/null; then`, - ` echo "[optio] Cleaning up stale worktree references..."`, - ` git worktree remove --force /workspace/tasks/${taskId}-wt 2>/dev/null || true`, - ` for wt_path in $(git worktree list --porcelain | grep -B1 "branch refs/heads/optio/task-${taskId}$" | grep "^worktree " | cut -d" " -f2-); do`, - ` git worktree remove --force "$wt_path" 2>/dev/null || true`, - ` done`, - ` git worktree prune`, - ` git branch -D optio/task-${taskId} 2>/dev/null || true`, - ` git worktree add /workspace/tasks/${taskId} -b optio/task-${taskId} "origin/${env.OPTIO_REPO_BRANCH ?? "main"}"`, - ` fi`, - `fi`, - ]; - - const script = [ - "set -e", - ...envExports, - `exec 7>${shellQuote(`/workspace/.task-${taskId}.lock`)}`, - "flock -n 7 || { echo 'Previous task process is still running; inspect it before retrying.' >&2; exit 75; }", - `echo "[optio] Waiting for repo to be ready..."`, - `for i in $(seq 1 \${REPO_INIT_TIMEOUT_SECS}); do [ -f /workspace/.ready ] && break; sleep 1; done`, - `[ -f /workspace/.ready ] || { echo "[optio] ERROR: repo not ready after \${REPO_INIT_TIMEOUT_SECS}s (increase OPTIO_REPO_INIT_TIMEOUT_MS to extend)"; exit 1; }`, - `echo "[optio] Repo ready"`, - // Use task-scoped credential URL for git operations (user-scoped token). - // Override the pod-level URL which returns an installation token. - `if [ -n "\${OPTIO_GIT_TASK_CREDENTIAL_URL:-}" ]; then`, - ` export OPTIO_GIT_CREDENTIAL_URL="\${OPTIO_GIT_TASK_CREDENTIAL_URL}"`, - `fi`, - // Set up gh CLI wrapper with PATH prepend (no root required) - `if [ -f /usr/local/bin/optio-gh-wrapper ]; then`, - ` mkdir -p /home/agent/.local/bin`, - ` cp /usr/local/bin/optio-gh-wrapper /home/agent/.local/bin/gh 2>/dev/null || true`, - ` chmod +x /home/agent/.local/bin/gh 2>/dev/null || true`, - ` export PATH="/home/agent/.local/bin:$PATH"`, - `fi`, - // Set up glab CLI wrapper and authenticate for GitLab repos - `if [ -n "\${GITLAB_TOKEN:-}" ]; then`, - ` if [ -f /usr/local/bin/optio-glab-wrapper ]; then`, - ` mkdir -p /home/agent/.local/bin`, - ` cp /usr/local/bin/optio-glab-wrapper /home/agent/.local/bin/glab 2>/dev/null || true`, - ` chmod +x /home/agent/.local/bin/glab 2>/dev/null || true`, - ` export PATH="/home/agent/.local/bin:$PATH"`, - ` fi`, - ` GITLAB_HOST="gitlab.com"`, - ` case "\${OPTIO_REPO_URL:-}" in *://*) GITLAB_HOST=$(echo "\${OPTIO_REPO_URL}" | sed -E 's|.*://([^/]+).*|\\1|');; esac`, - ` glab auth login --hostname "\${GITLAB_HOST}" --token "\${GITLAB_TOKEN}" 2>/dev/null || true`, - `fi`, - `ENV_FRESH="true"`, - `[ -f /home/agent/.optio-env-ready ] && ENV_FRESH="false"`, - `export ENV_FRESH`, - `if [ "$ENV_FRESH" = "true" ]; then echo "[optio] Fresh environment — tools may need to be installed"; else echo "[optio] Warm environment — tools from previous tasks should be available"; fi`, - `echo "[optio] Acquiring repo lock..."`, - `exec 9>/workspace/.repo-lock`, - `flock 9`, - `echo "[optio] Repo lock acquired"`, - `cd /workspace/repo`, - `git fetch origin`, - `git checkout "${env.OPTIO_REPO_BRANCH ?? "main"}" 2>/dev/null || true`, - `git reset --hard "origin/${env.OPTIO_REPO_BRANCH ?? "main"}"`, - ...worktreeSetup, - `if [ -f /workspace/repo/.gitmodules ]; then git -C /workspace/tasks/${taskId} submodule update --init --recursive 2>&1 || true; fi`, - `flock -u 9`, - `exec 9>&-`, - `cd /workspace/tasks/${taskId}`, - // Configure git at worktree scope so concurrent tasks don't interfere - `if [ -n "\${OPTIO_GIT_CREDENTIAL_URL:-}" ] && [ -f /usr/local/bin/optio-git-credential ]; then`, - ` git config --local credential.helper '/usr/local/bin/optio-git-credential'`, - ` echo "[optio] Worktree credential helper configured"`, - `fi`, - `git config --local user.name "\${GITHUB_APP_BOT_NAME:-Optio Agent}"`, - `git config --local user.email "\${GITHUB_APP_BOT_EMAIL:-optio-agent@noreply.github.com}"`, - `echo "${runToken}" > /workspace/tasks/${taskId}/.optio-run-token`, - `export OPTIO_TASK_ID="${taskId}"`, - ...WRITE_SETUP_FILES, - // Exclude Optio runtime files from git tracking using the local exclude file - // (never committed, unlike .gitignore modifications) - `EXCLUDE_FILE="$(git rev-parse --git-dir)/info/exclude"`, - `mkdir -p "$(dirname "$EXCLUDE_FILE")"`, - `grep -qxF '.optio/' "$EXCLUDE_FILE" 2>/dev/null || echo '.optio/' >> "$EXCLUDE_FILE"`, - `grep -qxF '.optio-run-token' "$EXCLUDE_FILE" 2>/dev/null || echo '.optio-run-token' >> "$EXCLUDE_FILE"`, - `grep -qxF '.optio-cache/' "$EXCLUDE_FILE" 2>/dev/null || echo '.optio-cache/' >> "$EXCLUDE_FILE"`, - // EXIT trap: kill child processes (agent + watchdog) then clean up internal worktrees. - // This ensures orphaned agent processes are killed when the exec stream is severed - // (e.g. API pod restart closes the SPDY connection but kubelet doesn't send SIGHUP). - `_optio_main_pid=$$`, - // The run's own home (its runtime's MCP config, Codex / Gemini session - // files) goes with the run; a retry writes it afresh (harness-config.ts). - `trap 'kill $(jobs -p) 2>/dev/null; wait 2>/dev/null; cd /workspace/repo 2>/dev/null; git worktree remove --force /workspace/tasks/${taskId}-wt 2>/dev/null || true; git worktree prune 2>/dev/null || true; ${REMOVE_RUN_HOME}' EXIT`, - // Background heartbeat: detect broken stdout pipe (EPIPE) from severed exec stream. - // Writes an empty line every 30s (skipped by the NDJSON parser). If stdout is broken - // (API pod died), sends SIGTERM to the main script which triggers the EXIT trap. - `(trap '' PIPE; while sleep 30; do printf '\\n' 2>/dev/null || { kill -TERM $_optio_main_pid 2>/dev/null; exit; }; done) &`, - ...RUN_WORK_SETUP_COMMANDS, - `set +e`, - ...agentCommand, - `AGENT_EXIT=$?`, - `[ $AGENT_EXIT -eq 0 ] && touch /home/agent/.optio-env-ready`, - `exit $AGENT_EXIT`, - ].join("\n"); return rt.exec(handle, ["bash", "-c", script], { tty: false }); }, diff --git a/apps/api/src/services/run-home-sweep-service.test.ts b/apps/api/src/services/run-home-sweep-service.test.ts new file mode 100644 index 000000000..b982b6b70 --- /dev/null +++ b/apps/api/src/services/run-home-sweep-service.test.ts @@ -0,0 +1,38 @@ +import { describe, expect, it } from "vitest"; +import { parseRunHomeList, removable } from "./run-home-sweep-service.js"; + +describe("run-home sweep", () => { + it("parses the pod's listing, skipping names that are not run ids", () => { + expect( + parseRunHomeList( + [ + "0f2e1c1a-1111-4222-8333-444444444444 1234", + "../etc 0", + "turn-7 0", + "", + "weird name 9", + "abc", + ].join("\n"), + ), + ).toEqual([ + { runId: "0f2e1c1a-1111-4222-8333-444444444444", outputBytes: 1234 }, + { runId: "turn-7", outputBytes: 0 }, + ]); + }); + + it("removes a non-protocol run's home at once, a protocol run's once drained or an hour later", () => { + const now = Date.now(); + const recent = new Date(now - 60_000); + const old = new Date(now - 2 * 60 * 60_000); + const base = { id: "r", finishedAt: recent }; + expect(removable({ ...base, protocol: false, consumedBytes: 0 }, 500, now)).toBe(true); + expect(removable({ ...base, protocol: true, consumedBytes: 500 }, 500, now)).toBe(true); + expect(removable({ ...base, protocol: true, consumedBytes: 499 }, 500, now)).toBe(false); + expect( + removable({ ...base, protocol: true, consumedBytes: 0, finishedAt: old }, 500, now), + ).toBe(true); + expect( + removable({ ...base, protocol: true, consumedBytes: 0, finishedAt: null }, 500, now), + ).toBe(false); + }); +}); diff --git a/apps/api/src/services/run-home-sweep-service.ts b/apps/api/src/services/run-home-sweep-service.ts new file mode 100644 index 000000000..ad2280b90 --- /dev/null +++ b/apps/api/src/services/run-home-sweep-service.ts @@ -0,0 +1,187 @@ +/** + * The run-home sweep (docs/plans/scale-out.md §1, "Where things live"): the + * backstop that removes `/home/agent/optio/runs/` directories of + * finished runs from every ready pod. The attached worker removes a run's + * directory itself once the run is terminal and its output consumed; this + * sweep catches what a worker that died never got to, and the directories + * of runs that never used the protocol (the exec path, which has nothing to + * drain). A run home carries the run's setup files, some sensitive, so no + * finished run's directory outlives the sweep's cycle. + */ +import { and, inArray, isNotNull, sql } from "drizzle-orm"; +import { shellQuote } from "@optio/shared"; +import { db } from "../db/client.js"; +import { agentPods, persistentAgentTurns, prReviewRuns, workRuns } from "../db/schema.js"; +import { getRuntime } from "./container-service.js"; +import { podHandle } from "./agent-pod-pool.js"; +import { RUN_HOMES_DIR } from "../utils/harness-config.js"; +import { logger } from "../logger.js"; + +const SAFE_RUN_ID = /^[A-Za-z0-9][A-Za-z0-9_-]{0,127}$/; +/** A terminal run whose output was never drained is still removed after this long. */ +const UNDRAINED_GRACE_MS = 60 * 60_000; + +const TERMINAL_TASK_STATES = ["completed", "failed", "cancelled"] as const; +const TERMINAL_REVIEW_RUN_STATES = ["completed", "failed", "cancelled"] as const; + +/** Reads a whole stream as a string. */ +function collect(stream: NodeJS.ReadableStream): Promise { + return new Promise((resolve, reject) => { + let buf = ""; + stream.on("data", (c: Buffer | string) => (buf += c.toString())); + stream.on("end", () => resolve(buf)); + stream.on("error", reject); + }); +} + +/** `ls` of the pod's run homes with each output file's size: ` ` per line. */ +const LIST_RUN_HOMES = [ + `for d in ${RUN_HOMES_DIR}/*; do`, + ` [ -d "$d" ] || continue`, + ` printf '%s %s\\n' "$(basename "$d")" "$(stat -c %s "$d/output.ndjson" 2>/dev/null || echo 0)"`, + `done`, +].join("\n"); + +export interface RunHomeEntry { + runId: string; + outputBytes: number; +} + +export function parseRunHomeList(text: string): RunHomeEntry[] { + const out: RunHomeEntry[] = []; + for (const line of text.split("\n")) { + const parts = line.trim().split(/\s+/); + if (parts.length !== 2) continue; + const [id, bytes] = parts; + if (!SAFE_RUN_ID.test(id)) continue; + out.push({ runId: id, outputBytes: Number.parseInt(bytes, 10) || 0 }); + } + return out; +} + +interface FinishedRun { + id: string; + /** Whether the run went through the protocol (an attached worker drains it). */ + protocol: boolean; + consumedBytes: number; + finishedAt: Date | null; +} + +/** Which of the given ids are finished runs, across every run table. */ +export async function finishedRuns(ids: string[]): Promise> { + const found = new Map(); + if (ids.length === 0) return found; + const tasks = await db + .select({ + id: workRuns.id, + execState: workRuns.execState, + consumedBytes: workRuns.consumedBytes, + completedAt: workRuns.completedAt, + updatedAt: workRuns.updatedAt, + }) + .from(workRuns) + .where(and(inArray(workRuns.id, ids), inArray(workRuns.state, [...TERMINAL_TASK_STATES]))); + for (const t of tasks) { + found.set(t.id, { + id: t.id, + protocol: t.execState !== null, + consumedBytes: Number(t.consumedBytes ?? 0), + finishedAt: t.completedAt ?? t.updatedAt, + }); + } + const reviews = await db + .select({ + id: prReviewRuns.id, + execState: prReviewRuns.execState, + consumedBytes: prReviewRuns.consumedBytes, + completedAt: prReviewRuns.completedAt, + updatedAt: prReviewRuns.updatedAt, + }) + .from(prReviewRuns) + .where( + and( + inArray(prReviewRuns.id, ids), + inArray(prReviewRuns.state, [...TERMINAL_REVIEW_RUN_STATES]), + ), + ); + for (const r of reviews) { + found.set(r.id, { + id: r.id, + protocol: r.execState !== null, + consumedBytes: Number(r.consumedBytes ?? 0), + finishedAt: r.completedAt ?? r.updatedAt, + }); + } + const turns = await db + .select({ + id: persistentAgentTurns.id, + execState: persistentAgentTurns.execState, + consumedBytes: persistentAgentTurns.consumedBytes, + finishedAt: persistentAgentTurns.finishedAt, + }) + .from(persistentAgentTurns) + .where(and(inArray(persistentAgentTurns.id, ids), isNotNull(persistentAgentTurns.finishedAt))); + for (const t of turns) { + found.set(t.id, { + id: t.id, + protocol: t.execState !== null, + consumedBytes: Number(t.consumedBytes ?? 0), + finishedAt: t.finishedAt, + }); + } + return found; +} + +/** Whether a finished run's home may go: drained, never protocol, or finished long ago. */ +export function removable(run: FinishedRun, outputBytes: number, now = Date.now()): boolean { + if (!run.protocol) return true; + if (run.consumedBytes >= outputBytes) return true; + return run.finishedAt !== null && now - run.finishedAt.getTime() > UNDRAINED_GRACE_MS; +} + +/** + * Sweeps every ready pod once. Returns how many run homes were removed. + * Unknown directories (no run row at all) are left alone. + */ +export async function sweepFinishedRunHomes(): Promise { + const pods = await db + .select() + .from(agentPods) + .where(and(sql`${agentPods.state} = 'ready'`, isNotNull(agentPods.podName))); + const rt = getRuntime(); + let removed = 0; + for (const pod of pods) { + const handle = podHandle(pod); + let listing: string; + try { + const session = await rt.exec(handle, ["bash", "-c", LIST_RUN_HOMES], { tty: false }); + listing = await collect(session.stdout); + } catch (err) { + logger.debug({ err, pod: pod.podName }, "Run-home sweep: could not list the pod's run homes"); + continue; + } + const entries = parseRunHomeList(listing); + if (entries.length === 0) continue; + const finished = await finishedRuns(entries.map((e) => e.runId)); + const gone = entries.filter((e) => { + const run = finished.get(e.runId); + return run !== undefined && removable(run, e.outputBytes); + }); + if (gone.length === 0) continue; + const script = gone + .map((e) => `rm -rf ${shellQuote(`${RUN_HOMES_DIR}/${e.runId}`)}`) + .join("\n"); + try { + const session = await rt.exec(handle, ["bash", "-c", script], { tty: false }); + await collect(session.stdout); + removed += gone.length; + logger.info( + { pod: pod.podName, runIds: gone.map((e) => e.runId) }, + "Run-home sweep: removed finished runs' directories", + ); + } catch (err) { + logger.warn({ err, pod: pod.podName }, "Run-home sweep: could not remove run homes"); + } + } + return removed; +} diff --git a/apps/api/src/services/workflow-pool-service.ts b/apps/api/src/services/workflow-pool-service.ts index e59c986bc..7d7b396b6 100644 --- a/apps/api/src/services/workflow-pool-service.ts +++ b/apps/api/src/services/workflow-pool-service.ts @@ -16,7 +16,7 @@ import { logger } from "../logger.js"; import { resolveImage } from "./repo-pool-service.js"; import { getWorkloadManager, isStatefulSetEnabled } from "./k8s-workload-service.js"; import * as podPool from "./agent-pod-pool.js"; -import { buildPooledExecScript } from "../utils/pod-env.js"; +import { buildPooledExecScript, KEEP_POD_ALIVE } from "../utils/pod-env.js"; import { updatedAtMatches } from "../utils/pg-timestamp.js"; const IDLE_TIMEOUT_MS = parseIntEnv("OPTIO_WORKFLOW_POD_IDLE_MS", 600000); // 10 min default @@ -60,7 +60,7 @@ function buildInitScript(): string { "mkdir -p /workspace/runs", "touch /workspace/.ready", "echo '[optio] Workflow pod ready'", - "exec sleep infinity", + KEEP_POD_ALIVE, ].join("\n"); } diff --git a/apps/api/src/test-utils/e2e/api-cluster.ts b/apps/api/src/test-utils/e2e/api-cluster.ts new file mode 100644 index 000000000..a1c862a41 --- /dev/null +++ b/apps/api/src/test-utils/e2e/api-cluster.ts @@ -0,0 +1,148 @@ +/** + * Several REAL API servers on one database, one Redis and one fake-runtime + * directory — the scale-out e2e tier (docs/plans/scale-out.md, Testing). + * Each server is a `startApiServer()` process with its own port and + * `OPTIO_INSTANCE_ID`; the fake runtime keeps its pods and runs on disk + * (`OPTIO_FAKE_RUNTIME_DIR`), so a run one server started is there for the + * next one. Kill a server with SIGKILL and the rest keep serving; add one + * and it joins the same work. + * + * Usage: + * let cluster: ApiCluster; + * beforeAll(async () => { cluster = await startApiCluster({ size: 2 }); }, 200_000); + * afterAll(() => cluster.stop()); + * const { body } = await cluster.client(0)("/api/jobs", { method: "POST", ... }); + * await cluster.kill(0); // SIGKILL, waits for the process to go + * await cluster.any()("/api/health"); // round-robins over the live servers + */ +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { startApiServer, type ApiServerHandle, type StartApiServerOptions } from "./api-server.js"; + +export interface ApiClusterOptions { + /** How many servers to boot (≥ 1). */ + size: number; + /** The fake runtime's directory (default: a fresh temp dir, removed on stop). */ + fakeRuntimeDir?: string; + /** Extra env for every server (a server's own `OPTIO_INSTANCE_ID` wins). */ + env?: Record; + logLevel?: string; + readyTimeoutMs?: number; +} + +export type ClusterClient = ( + path: string, + init?: RequestInit, +) => Promise<{ status: number; body: T }>; + +export interface ClusterServer { + /** `api-`: the server's `OPTIO_INSTANCE_ID` base (the process appends its boot suffix). */ + name: string; + handle: ApiServerHandle; + /** False once `kill()` took it down. */ + alive: boolean; +} + +export interface ApiCluster { + readonly servers: ClusterServer[]; + readonly fakeRuntimeDir: string; + /** The live servers, in boot order. */ + live(): ClusterServer[]; + /** A client bound to server `i` (by boot index, killed or not — a killed one fails). */ + client(i: number): ClusterClient; + /** A client that round-robins over the live servers. */ + any(): ClusterClient; + /** Kills server `i` (SIGKILL by default) and waits for its process to be gone. */ + kill(i: number, signal?: NodeJS.Signals): Promise; + /** Boots another server on the same database and fake directory. */ + add(): Promise; + /** Stops every live server and removes a temp fake directory. */ + stop(): Promise; +} + +function makeClient(baseUrl: () => string): ClusterClient { + return async (path: string, init?: RequestInit) => { + const res = await fetch(`${baseUrl()}${path}`, { + headers: { "content-type": "application/json" }, + ...init, + }); + const text = await res.text(); + let body: T; + try { + body = JSON.parse(text) as T; + } catch { + body = text as unknown as T; + } + return { status: res.status, body }; + }; +} + +export async function startApiCluster(opts: ApiClusterOptions): Promise { + if (!Number.isInteger(opts.size) || opts.size < 1) throw new Error("size must be ≥ 1"); + const ownsDir = !opts.fakeRuntimeDir; + const fakeRuntimeDir = opts.fakeRuntimeDir ?? mkdtempSync(join(tmpdir(), "optio-e2e-fake-")); + const servers: ClusterServer[] = []; + let next = 0; + let rr = 0; + + const boot = async (): Promise => { + const name = `api-${next++}`; + const serverOpts: StartApiServerOptions = { + logLevel: opts.logLevel, + readyTimeoutMs: opts.readyTimeoutMs, + env: { + ...opts.env, + OPTIO_FAKE_RUNTIME_DIR: fakeRuntimeDir, + OPTIO_INSTANCE_ID: name, + }, + }; + const handle = await startApiServer(serverOpts); + const server: ClusterServer = { name, handle, alive: true }; + handle.proc.once("exit", () => { + server.alive = false; + }); + servers.push(server); + return server; + }; + + // Boot together: the migration runner's advisory lock serializes what must be. + await Promise.all(Array.from({ length: opts.size }, () => boot())); + servers.sort((a, b) => Number(a.name.slice(4)) - Number(b.name.slice(4))); + + const live = () => servers.filter((s) => s.alive); + + return { + servers, + fakeRuntimeDir, + live, + client(i) { + const server = servers[i]; + if (!server) throw new Error(`no server ${i}`); + return makeClient(() => { + if (!server.alive) throw new Error(`server ${server.name} was killed`); + return server.handle.baseUrl; + }); + }, + any() { + return makeClient(() => { + const alive = live(); + if (alive.length === 0) throw new Error("no live server"); + const server = alive[rr++ % alive.length]; + return server.handle.baseUrl; + }); + }, + async kill(i, signal = "SIGKILL") { + const server = servers[i]; + if (!server) throw new Error(`no server ${i}`); + if (!server.alive) return; + await server.handle.kill(signal); + server.alive = false; + }, + add: boot, + async stop() { + await Promise.all(live().map((s) => s.handle.stop())); + if (ownsDir) rmSync(fakeRuntimeDir, { recursive: true, force: true }); + }, + }; +} diff --git a/apps/api/src/test-utils/e2e/api-server.ts b/apps/api/src/test-utils/e2e/api-server.ts index 7b9043bf3..4e57b2688 100644 --- a/apps/api/src/test-utils/e2e/api-server.ts +++ b/apps/api/src/test-utils/e2e/api-server.ts @@ -37,7 +37,14 @@ export interface ApiServerHandle { proc: ChildProcess; /** Captured stdout+stderr (ring buffer) — for debugging failed scenarios. */ logs(): string; + /** Graceful stop: SIGTERM, then SIGKILL after 10 s. */ stop(): Promise; + /** + * Sends `signal` to the server's process group and waits for the process + * to be gone (a signal other than SIGKILL is followed by SIGKILL after + * 15 s). The scale-out e2e tier kills servers mid-run with it. + */ + kill(signal: NodeJS.Signals): Promise; } export interface StartApiServerOptions { @@ -175,23 +182,29 @@ async function startApiServerOnce(opts: StartApiServerOptions): Promise setTimeout(r, 250)); } + const kill = (signal: NodeJS.Signals, forceAfterMs: number) => + new Promise((resolve) => { + if (exited) return resolve(); + const killTimer = + signal === "SIGKILL" + ? null + : setTimeout(() => { + killGroup("SIGKILL"); + }, forceAfterMs); + proc.once("exit", () => { + if (killTimer) clearTimeout(killTimer); + resolve(); + }); + killGroup(signal); + }); + return { baseUrl, port, proc, logs: () => lines.join("\n"), - stop: () => - new Promise((resolve) => { - if (exited) return resolve(); - const killTimer = setTimeout(() => { - killGroup("SIGKILL"); - }, 10_000); - proc.once("exit", () => { - clearTimeout(killTimer); - resolve(); - }); - killGroup("SIGTERM"); - }), + stop: () => kill("SIGTERM", 10_000), + kill: (signal) => kill(signal, 15_000), }; } diff --git a/apps/api/src/utils/pod-env.test.ts b/apps/api/src/utils/pod-env.test.ts index 14f08e8c6..c81b81352 100644 --- a/apps/api/src/utils/pod-env.test.ts +++ b/apps/api/src/utils/pod-env.test.ts @@ -12,11 +12,19 @@ import { import { tmpdir } from "node:os"; import { join } from "node:path"; import { + attachRunScript, buildEnvExports, buildPooledExecScript, + buildPooledStartScript, CHECKOUT_REPO, + deliverStdinScript, + killRunScript, POOLED_CHECKOUT_DIR, + RUN_FILES, + RUN_STARTED_MARKER, RUN_WORK_SETUP_COMMANDS, + STDIN_EOF_SENTINEL, + superviseAgent, WRITE_SETUP_FILES, } from "./pod-env.js"; @@ -419,3 +427,51 @@ describe("WRITE_SETUP_FILES — a repo's own config", () => { } }); }); + +describe("buildPooledStartScript (the run protocol)", () => { + const runId = "0f2e1c1a-1111-4222-8333-444444444444"; + const base = { + env: { A: "1" }, + runId, + workDir: "/workspace/turns/t1", + agentCommand: ["run-agent"], + }; + + it("does the pooled setup, then hands the agent to the supervisor in the run's home", () => { + const script = buildPooledStartScript(base); + expect(script).toContain(`mkdir -p '/workspace/turns/t1'`); + expect(script).toContain(`cd '/workspace/turns/t1'`); + expect(script.indexOf(WRITE_SETUP_FILES[0])).toBeLessThan(script.indexOf("run-agent")); + expect(script.indexOf(RUN_WORK_SETUP_COMMANDS[0])).toBeLessThan(script.indexOf("run-agent")); + expect(script).toContain( + superviseAgent(`/home/agent/optio/runs/${runId}`, ["run-agent"]).join("\n"), + ); + expect(script.trim().endsWith(`echo "${RUN_STARTED_MARKER}:$_optio_sup"`)).toBe(true); + }); + + it("keeps the run lock, refuses a second start of a live run, and never traps or removes the home", () => { + const script = buildPooledStartScript(base); + expect(script).toMatch(/exec 7>\/workspace\/\.run-[0-9a-f]+\.lock\nflock -n 7/); + expect(script).toContain(`if optio_run_pid '/home/agent/optio/runs/${runId}' >/dev/null; then`); + expect(script).toContain("optio_run_pid() {"); + expect(script).toContain(`read -r _optio_pid _optio_start < "$1/${RUN_FILES.pid}"`); + expect(script).not.toContain("trap '"); + expect(script).not.toContain("__OPTIO_RUN_EXIT__"); + expect(script).not.toContain("rm -rf"); + // The exec script's inline agent + exit line are not here. + expect(script).not.toContain("AGENT_EXIT=$?"); + }); + + it("with a repo, works in the pod's one checkout of it", () => { + const script = buildPooledStartScript({ ...base, checkout: true }); + expect(script).toContain(CHECKOUT_REPO.join("\n")); + expect(script).toContain(`cd '${POOLED_CHECKOUT_DIR}'`); + }); + + it("re-exports the protocol's pieces beside the other script fragments", () => { + expect(STDIN_EOF_SENTINEL).toBe("__OPTIO_STDIN_EOF__"); + expect(attachRunScript("/home/agent/optio/runs/x", 0)).toContain("tail -c +1 "); + expect(deliverStdinScript("/home/agent/optio/runs/x", 3)).toContain("head -c 3 >>"); + expect(killRunScript("/home/agent/optio/runs/x", "TERM")).toContain("kill -TERM"); + }); +}); diff --git a/apps/api/src/utils/pod-env.ts b/apps/api/src/utils/pod-env.ts index a37131a74..dabb848e9 100644 --- a/apps/api/src/utils/pod-env.ts +++ b/apps/api/src/utils/pod-env.ts @@ -17,7 +17,47 @@ import { createHash } from "node:crypto"; */ import { shellQuote } from "@optio/shared"; -import { REMOVE_RUN_HOME } from "./harness-config.js"; +import { + RUN_ALIVE_HELPER, + RUN_DIR_ENV, + RUN_EXIT_MARKER, + RUN_FILES, + RUN_LOST_MARKER, + RUN_MARK_ENV_READY_ENV, + RUN_STARTED_MARKER, + RUN_STDIN_BYTES_ENV, + STDIN_DELIVERED_MARKER, + STDIN_EOF_SENTINEL, + SUPERVISOR_SCRIPT, + attachRunScript, + deliverStdinScript, + killRunScript, + startRunPrelude, + superviseAgent, +} from "@optio/container-runtime"; +import { REMOVE_RUN_HOME, runHome } from "./harness-config.js"; + +// The run protocol's script fragments (docs/plans/scale-out.md §1) live with +// the runtimes that run them (packages/container-runtime/src/run-protocol.ts) +// and are re-exported here, next to the other exec script pieces. +export { + RUN_ALIVE_HELPER, + RUN_DIR_ENV, + RUN_EXIT_MARKER, + RUN_FILES, + RUN_LOST_MARKER, + RUN_MARK_ENV_READY_ENV, + RUN_STARTED_MARKER, + RUN_STDIN_BYTES_ENV, + STDIN_DELIVERED_MARKER, + STDIN_EOF_SENTINEL, + SUPERVISOR_SCRIPT, + attachRunScript, + deliverStdinScript, + killRunScript, + startRunPrelude, + superviseAgent, +}; export const VALID_ENV_NAME = /^[A-Za-z_][A-Za-z0-9_]*$/; @@ -247,3 +287,61 @@ export function buildPooledExecScript(input: { `exit $AGENT_EXIT`, ].join("\n"); } + +/** + * The last line of every pod's init script: keep the pod alive under a pid + * 1 that reaps. A run's supervisor is reparented to pid 1 when the exec + * that started it ends (run-protocol.ts), and a bare `sleep` never calls + * wait(), so every finished run would stay a zombie. `tini` is in the agent + * image (images/base.Dockerfile); an older image falls back to `sleep`. + */ +export const KEEP_POD_ALIVE = + "if command -v tini >/dev/null 2>&1; then exec tini -- sleep infinity; else exec sleep infinity; fi"; + +/** + * The run protocol's start script for a pooled pod (a Job run or a + * persistent-agent turn): the same setup as `buildPooledExecScript`, then + * the agent launched under the supervisor in the run's own home + * (`runHome(runId)`), so the exec returns once the agent runs and any API + * instance can attach to it later (run-protocol.ts). The run lock stays: + * a second start of the same run while its supervisor lives exits 75. No + * trap removes the run home: the attached worker removes it when the run + * is terminal and `consumed_bytes` equals the output file's size (phase B), + * with the run-home sweep (`run-home-sweep-service.ts`) as the backstop. + */ +export function buildPooledStartScript(input: { + env: Record; + /** Optio's id for the run (a Job run or turn id): names its run home. */ + runId: string; + /** The run's directory; ignored with `checkout` (the run works in the checkout). */ + workDir: string; + agentCommand: string[]; + label?: string; + /** Work in a checkout of `$OPTIO_REPO_URL` (`CHECKOUT_REPO`). */ + checkout?: boolean; +}): string { + const dir = input.checkout ? POOLED_CHECKOUT_DIR : input.workDir; + const what = input.label ?? "pod"; + const lock = createHash("sha256").update(input.workDir).digest("hex").slice(0, 32); + const home = runHome(input.runId); + return [ + "set -e", + ...buildEnvExports(input.env), + ...(input.label ? [`echo "[optio] Waiting for ${what} to be ready..."`] : []), + `for i in $(seq 1 120); do [ -f /workspace/.ready ] && break; sleep 1; done`, + `[ -f /workspace/.ready ] || { echo "[optio] ERROR: ${what} not ready after 120s"; exit 1; }`, + ...(input.label ? [`echo "[optio] ${what[0].toUpperCase()}${what.slice(1)} ready"`] : []), + // The lock is held by the start exec only while it sets up: the + // supervisor's pid file is what says a run is live, and only when its + // process is the one that wrote it (RUN_ALIVE_HELPER). + `exec 7>/workspace/.run-${lock}.lock`, + "flock -n 7 || { echo 'Previous run is still active; inspect before retrying.' >&2; exit 75; }", + ...RUN_ALIVE_HELPER, + `if optio_run_pid ${shellQuote(home.podDir)} >/dev/null; then echo 'Previous run is still active; inspect before retrying.' >&2; exit 75; fi`, + ...(input.checkout ? CHECKOUT_REPO : [`mkdir -p ${shellQuote(dir)}`]), + `cd ${shellQuote(dir)}`, + ...WRITE_SETUP_FILES, + ...RUN_WORK_SETUP_COMMANDS, + ...superviseAgent(home.podDir, input.agentCommand), + ].join("\n"); +} diff --git a/apps/api/src/workers/repo-cleanup-worker.ts b/apps/api/src/workers/repo-cleanup-worker.ts index b15b54f9c..432f94309 100644 --- a/apps/api/src/workers/repo-cleanup-worker.ts +++ b/apps/api/src/workers/repo-cleanup-worker.ts @@ -34,6 +34,7 @@ import { instrumentWorkerProcessor } from "../telemetry/instrument-worker.js"; import { getBullMQOptions } from "../services/redis-config.js"; import { removeRunHome } from "../utils/harness-config.js"; +import { sweepFinishedRunHomes } from "../services/run-home-sweep-service.js"; const bullmqOpts = getBullMQOptions(); @@ -492,6 +493,15 @@ export function startRepoCleanupWorker() { logger.warn({ err }, "Failed to clean up idle persistent agent pods"); } + // Run homes of finished runs (the run protocol's directories in the + // pods): the attached worker removes them; this is the backstop. + try { + const homes = await sweepFinishedRunHomes(); + if (homes > 0) logger.info({ homes }, "Removed finished run homes"); + } catch (err) { + logger.warn({ err }, "Failed to sweep run homes"); + } + // Detect zombie workflow_runs (running but pod terminated/gone) try { const zombieRuns = await cleanupZombieWorkflowRuns(); diff --git a/docs/plans/scale-out.md b/docs/plans/scale-out.md new file mode 100644 index 000000000..84bb17d93 --- /dev/null +++ b/docs/plans/scale-out.md @@ -0,0 +1,322 @@ +# Scale-out: state in Postgres, runs that outlive an API pod + +Optio runs one API/web pod. The Helm chart enforces it, because the process +holds state that a second replica would not see, and because every agent run +is an open Kubernetes exec stream held by that one process: when the pod goes +away, so does the run. This plan makes the API process disposable. Any number +of API pods share one Postgres and one Redis; the pod a run is attached to can +be killed, and another pod picks the run up where it was, with every log line +accounted for once. + +**Status (2026-10-10):** phase A (foundation) built; B, C1, C2, D, E to come. +Phases land as separate PRs; each leaves the system working on one replica. + +## Vocabulary + +| Word | Meaning | +| ------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **Instance** | One API process. `OPTIO_INSTANCE_ID` = `:` (the pod name in Kubernetes, with a boot suffix so a restarted pod is a new instance). | +| **Run** | One agent execution Optio supervises: a Repo Task run, a Job run, a persistent-agent turn, a PR-review run. All are `tasks` rows except reviews (`pr_review_runs`) and turns (`persistent_agent_turns`). | +| **Run directory** | The run's directory in its pod, `/home/agent/optio/runs/` (the run home of #664), holding the agent's output, stdin, pid and exit files. | +| **Supervisor** | The pod-side shell that starts the agent detached from the exec session and records its exit. Nothing Optio starts in a pod is a child of an exec any more. | +| **Attached** | A run some instance is currently streaming: `attached_by` is that instance and `attach_lease_until` is in the future. Renewed every 20 s; expires 60 s after the last renew. | +| **Consumed offset** | `consumed_bytes` on the run: how much of the output file has been turned into log rows. Committed with those rows in one transaction, at line boundaries only. | +| **Lease** | A row in `leases` (`key`, `holder`, `expires_at`): the one way an instance claims something for a while — a poller's turn, a chat lock, a Local host. | + +## Starting point (verified in code, 2026-10-10) + +What ties the system to one process, from the inventory that preceded this plan: + +- **Runs are exec streams.** `execTaskInRepoPod`, `execRunInPod` and + `execTurnInPod` open one Kubernetes exec per run and the BullMQ processor + reads its stdout for the run's whole life (`task-worker.ts` ~681–830, + `workflow-worker.ts` ~305+, `persistent-agent-worker.ts` ~302+; PR reviews + reuse the task path). The whole output is also kept in memory (`allLogs`). + The pod-side script SIGTERMs the agent when stdout breaks (the EPIPE + watchdog in `repo-pool-service.ts`), because nobody would read it otherwise; + pooled runs have no watchdog and orphan instead. Boot runs unconditional + sweeps that fail or flag every in-flight run (`recoverInterruptedExecutions`, + `reconcileOrphanedTasks`, `recoverSessionTurns`), since nothing records which + process owns a run. +- **Claude's stdin matters.** Claude Code runs with `--input-format stream-json`: + the first user message and every mid-run message are written to the exec's + stdin, and stdin is closed on the agent's terminal event so it exits. +- **In-memory state that diverges across replicas:** WebSocket upgrade tokens + (`session-service.ts wsUpgradeTokens`; `/ws` has no ingress affinity), the + per-IP WS cap, the Optio-chat per-user lock, Slack/Linear delivery dedupe + (`event-ingress.ts recentDeliveryIds`; the other sources use Redis `SET NX`), + the claim mutexes around count-then-claim (`claimLockChain` in the task and + workflow workers), GitHub token refresh locks, config-sync serialization, + glance/push bookkeeping (`glance-service.ts`, which says it assumes one + replica), session-chat listeners and interrupts (`session-turn-service.ts`), + session-sharing sockets, the whole Optio Local relay (`local-relay.ts`). +- **Pollers that must not overlap:** ticket sync (dedupes by scanning 500 + tasks), the schedule checker (advances `next_fire_at` after firing, no CAS), + external PR review (`pr_reviews_pr_url_idx` is not unique), and boot's + `cleanRepeatJobs`, which deletes every repeat schedule and lets workers re-add + them. There is no leader election, advisory lock (outside migrations and + secret heals), `FOR UPDATE` or `SKIP LOCKED` anywhere. +- **Already cross-replica:** BullMQ job delivery with CAS claims, the reconciler + (version CAS in `reconcile-executor.ts`), every log/event stream (Redis + pub/sub per task, Job run, review, agent; `/ws/events`), mid-task messages + (`optio:task-messages:{id}`), pod slots (DB), rate limits, the setup token, + session-terminal sockets (tmux in the pod), chat-turn exclusivity (a unique + partial index). +- The agent image is Ubuntu 24.04 (GNU coreutils: `tail --pid`, `flock`, + `setsid`, `python3` are there). The run home from #664 is + `/home/agent/optio/runs/`, deterministic per run, removed by the exec's + EXIT trap and again with the worktree. +- BullMQ is 5.52 (`upsertJobScheduler` exists). The fake runtime + (`packages/container-runtime/src/fake.ts`) kept its containers in a + per-process `Map` and played agents in-process, so it could not survive the + API process it ran in; phase A gave it a directory and a process per run. +- Found while building phase A: `ticket_external_id` is written by five paths + besides the sync (the API, the Issues page, ticket triggers, local + terminals, ticket events), so a unique index on tasks would reject work a + person starts on purpose; the sync gets a claim table instead. The outbound + webhooks already own `webhook_deliveries`, so the inbound dedupe table is + `inbound_webhook_deliveries`. The protocol's script fragments have to live + with the runtimes (`packages/container-runtime/src/run-protocol.ts`, the + API imports the runtime package and not the other way round); `pod-env.ts` + re-exports them beside the other script pieces. A `RunExit` has a third + kind, `detached`, for an attach whose exec dropped without a marker (the + attaching instance died, or closed it): neither exited nor lost. + +## Design + +### 1. The run protocol: start, attach, deliver, kill + +Every run goes through four pod-side operations, each a short exec. The +`ContainerRuntime` interface gains them, so the Kubernetes runtime and the fake +runtime implement the same protocol and the workers never build a long-lived +exec again. + +``` +startRun(handle, { runId, script, initialStdin }) → { pid } exits as soon as the agent is running +attachRun(handle, { runId, fromByte }) → { output, exit } output from fromByte until the agent exits; exit resolves exited: | lost | detached +deliverStdin(handle, { runId, line }) → void appends one line to stdin.ndjson (any instance may call it) +killRun(handle, { runId, signal }) → void kills the supervisor's process group +``` + +**Start.** The existing setup script (repo lock, worktree, credential helper, +setup files, setup commands) stays as it is up to the agent command. Instead +of running the agent inline, it writes `initialStdin` (the first stream-json +user message, built by the worker as today) to `$RUN_DIR/stdin.ndjson`, then +launches the supervisor detached — `setsid`, stdio to files, `&`, `disown` — +and prints `__OPTIO_RUN_STARTED__:` and exits. The supervisor: + +``` +( tail -n +1 -f stdin.ndjson | while IFS= read -r l; do [ "$l" = "__OPTIO_STDIN_EOF__" ] && break; printf '%s\n' "$l"; done ) \ + | >> output.ndjson 2>> stderr.log +echo $? > exit +``` + +The `while read` loop closes the agent's stdin when it sees the EOF sentinel, +which is how the worker ends a Claude turn (it appends the sentinel where it +called `stdin.end()` today). `pid` holds the supervisor's pid, and `setsid` +makes it a process-group leader so `killRun` can take the whole tree. The +EXIT trap no longer kills children or removes the run directory; the +EPIPE watchdog is gone. `touch /home/agent/.optio-env-ready` moves into the +supervisor, after a zero exit. + +**Attach.** `tail -c +$((fromByte+1)) --pid= -f output.ndjson`, which +streams from the offset and ends when the supervisor exits, flushing what is +left; then the script prints `__OPTIO_RUN_EXIT__:` from the `exit` file +(waiting up to a few seconds for it). A pid that is gone with no `exit` file +(the pod itself restarted) prints `__OPTIO_RUN_LOST__`, which the worker +treats as a pod death. Attach always streams from `consumed_bytes`; the worker +**rebuilds its in-memory derivations from byte 0** on a re-attach (the PR +tool-call tracker, the session id, the terminal-event flag) by reading the +file once without persisting, so nothing derived is ever persisted beside the +offset. Output is MBs at most. + +**Deliver.** Mid-run user messages stop needing the attached instance: the +route that receives a message appends the stream-json line to +`stdin.ndjson` through `deliverStdin` and marks it delivered. The Redis +message channel stays only to notify viewers. + +**Kill.** `killRun` sends the signal to the process group in `pid`; this is +what cancel and the reconciler's stall action use, from any instance. + +**Where things live.** `RUN_DIR` is the run home. It is no longer removed by +the exec's EXIT trap. The **attached worker removes the run directory when +the run is terminal and `consumed_bytes` equals the output file's size** +(phase B, for tasks, Job runs, turns and reviews). The backstop, from phase +A on, is the repo-cleanup worker's sweep (`run-home-sweep-service.ts`): on +each health cycle it lists `/home/agent/optio/runs` in every ready pod and +removes the directories of runs that are terminal and consumed (or terminal +for over an hour), through a small exec. Sensitive setup files never outlive +a run by more than that cycle. The exec script's EXIT trap also pruned the +task's stale `-wt` worktree; the start script does that right after the +worktree setup instead. + +**Stdin bytes never travel in a script.** An exec's argv is at most 128 KiB, +so the start script primes `stdin.ndjson` with `head -c ` from the exec's +own stdin (saved on fd 3 by the runtime's prelude), and a delivered line goes +the same way (`head -c >> stdin.ndjson`, then `__OPTIO_STDIN_DELIVERED__` +on stdout). `pid` holds ` ` and is written by the supervisor +itself once its TERM trap is in place; attach, kill and the pooled start +guard trust a pid only when the live process's start time matches, so a pid +reused after a pod restart reads as not running. + +**The pod's pid 1 must reap.** The supervisor is reparented to pid 1 when the +start exec ends; a bare `sleep infinity` never calls `wait()`, so every +finished run would stay a zombie. `tini` is in the agent image and every pod +init script ends with `exec tini -- sleep infinity` (falling back to `sleep` +on an older image): **the agent images must be rebuilt for the protocol** (the +fake runtime is unaffected). + +### 2. Ownership: columns, lease, reconcile + +Run rows (`runColumns()` in `schema.ts`, so both views pick them up; and the +same columns on `pr_review_runs` and `persistent_agent_turns`): + +| Column | Meaning | +| -------------------- | ------------------------------------------------------------------------------- | +| `exec_state` | `null` (not started), `started` (supervisor launched; `exec_pid` set), `exited` | +| `exec_pid` | The supervisor's pid | +| `consumed_bytes` | Bytes of `output.ndjson` turned into log rows | +| `attached_by` | Instance streaming the run, or null | +| `attach_lease_until` | When that attachment expires | + +**Claiming an attachment** is one CAS: +`UPDATE … SET attached_by = $me, attach_lease_until = now() + 60s WHERE id = $id AND (attached_by IS NULL OR attach_lease_until < now())`. +The attached worker renews every 20 s. On SIGTERM it **detaches**: closes the +exec, sets `attach_lease_until = now()`, and lets the process exit, so another +instance picks the run up in seconds rather than after a lease expiry. + +**The reconciler** (one new action per run kind, `reattach`, plus the existing +`requeue`) decides from the snapshot: + +- `running`/`provisioning`, `exec_state = started`, attachment expired → **reattach**: enqueue the kind's BullMQ job with a fresh id (`attach::`). The handler claims the attachment (CAS above) and, if it loses, exits. +- `provisioning`, `exec_state = null`, attachment expired → the run never got going: **requeue** through the existing retry path (same-pod affinity, worktree reset). +- Stall (`lastActivityAt` old) only counts while attached; an unattached run is re-attached first. +- `__OPTIO_RUN_LOST__` and a pod that is gone → the existing pod-death transition. + +A repeat job every 30 s (`attach-sweep`, idempotent, under a lease) enqueues +reconcile for runs whose attachment expired, so pickup does not wait for the +5-minute resync. + +**Boot** no longer fails anything. `recoverInterruptedExecutions` and the +fail/flag half of `reconcileOrphanedTasks` go; what remains is re-enqueueing +`queued` work (deduped against BullMQ), count repair, and a resync kick. +`recoverSessionTurns` stays until chat turns move to the protocol (see Out of +scope). + +**BullMQ** job lifetime still equals attachment lifetime, but stall +re-delivery is no longer how recovery happens: the attach handler is +idempotent through the lease CAS, so a stale re-delivery is a no-op. +`lockDuration` / `stalledInterval` drop to 60 s / 30 s for the run queues. + +**Repeat jobs** use `upsertJobScheduler` with stable ids; every instance +upserts the same schedulers, and boot's `cleanRepeatJobs` is deleted. + +### 3. State moves: what goes to Postgres, what is a cache, what is transport + +Rule: **correctness state goes to Postgres**; a **bounded-staleness cache** +may stay per instance when every instance can rebuild it and a stale read +costs nothing but freshness; **transport** (fan-out of live frames) is Redis +pub/sub, as it already is for logs and events. + +| Today (in process) | Home | +| ----------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| WS upgrade tokens | `ws_upgrade_tokens (token_hash pk, user_id, workspace_id, expires_at)`; consumed with `DELETE … RETURNING`; swept by the attach-sweep job | +| Slack / Linear in-memory dedupe, Redis `SET NX` dedupe | `inbound_webhook_deliveries (source, delivery_id, received_at; pk (source, delivery_id))` (`webhook_deliveries` is the outbound log); `INSERT … ON CONFLICT DO NOTHING RETURNING` claims; rows older than 24 h swept | +| `claimLockChain` (tasks, workflows) | `pg_advisory_xact_lock(hashtext('claim:tasks'))` and per-repo / per-Job keys inside one transaction around count + CAS | +| Optio-chat per-user lock, config-sync `running`, GitHub token refresh locks | `leases` (`acquireLease` / `renewLease` / `releaseLease`: insert-or-update-where-expired, returning the row iff acquired) | +| Pollers: ticket sync, schedule checker, external PR review, repo cleanup, skill sync, token validation | Each sweep runs under a lease (`poller:`), **and** the work is made safe to overlap: a `ticket_sync_claims` row per (source, ticket, repo) the sync inserts before creating a task (not a unique index on tasks: a person may start a second task from the same issue on purpose), `next_fire_at` advanced by CAS before firing, a unique partial index on active reviews per PR URL | +| Glance / push bookkeeping (`lastAttention`, `lastNeedsYou`, `lastRunning`, end and snooze timers, `hostOfflineAlerted`) | `glance_state (key pk, value jsonb, due_at)`; timers become `due_at` rows swept by the attach-sweep job. The APNs/FCM 1-second coalescer stays per instance (worst case one duplicate push, which the collapse id folds) | +| Session-chat listeners, interrupt, end-session | Frames fan out on `optio:session:{id}` (the channel that exists with no subscriber); interrupt and end are `control_intent` on the turn, acted on by the attached instance; the turn's exec itself moves to the run protocol in a later phase | +| Session-sharing socket close on revoke | Publish a `session:share_revoked` event; every instance closes its sockets on receipt (the 10 s DB recheck stays) | +| Installed-skills cache (ReadWriteOnce PVC) | `installed_skill_files (skill_id, path, content bytea, executable)` written by the sync worker; spawn reads the DB; the PVC goes | +| Per-IP WS cap, health/version/setup/scan caches, gauges | Stay per instance (caches). Documented as per replica | + +### 4. Optio Local relay (phase D) + +A daemon is one socket on one instance; a viewer may be on another. The +owning instance holds `leases` key `local-host:`; "newest daemon wins" +takes the lease by force and the old holder closes its socket when its renew +fails. Terminal output is published on `optio:local:term:{id}:out` (binary +frames as buffers), which viewers subscribe to on any instance; everything +inbound (input, resize, view reports, attach, spawn, kill, dirs, limits, +credentials, backfill, capability queries) is published on +`optio:local:host:{id}:in` and replies on `optio:local:reply:{requestId}`; +the owner consumes them and keeps the grid arbiter. `isHostOnline` reads the +lease. Until this phase lands, the chart keeps `api.replicas` at 1. + +### 5. Helm and ops + +`OPTIO_INSTANCE_ID` from the downward API pod name plus a boot suffix. The +replica validation lifts in phase D. The deployment stays `Recreate`: +migrations still run at boot under the advisory lock and there is no +mixed-version window; since runs now survive the gap, Recreate costs an API +outage of seconds, not any work. Rolling updates (expand/contract migrations) +are out of scope. + +## Phases + +| Phase | PR | Contents | Depends on | +| ----- | ------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------- | +| A | foundation | This document. `OPTIO_INSTANCE_ID`. **Every schema change of the whole plan in one migration** (run columns on the three run tables, `leases`, `ws_upgrade_tokens`, `inbound_webhook_deliveries`, `glance_state`, `installed_skill_files`, `ticket_sync_claims`, the unique partial index on active PR reviews). `lease-service.ts`. The run protocol on `ContainerRuntime` with the Kubernetes implementation (script builders in `pod-env.ts` / `repo-pool-service.ts`) and the **multi-process fake runtime**: containers and runs persisted under `OPTIO_FAKE_RUNTIME_DIR`, fake agents as detached child processes writing `output.ndjson`, so a fake run survives the API process. `test-utils/e2e/api-cluster.ts` starts N API servers on one database and one fake-runtime dir. | — | +| B | re-attachable runs | Workers for Repo Tasks, Job runs, agent turns and PR reviews on the protocol: start, attach from `consumed_bytes`, log rows + offset in one transaction, derivations rebuilt from byte 0, `allLogs` removed, deliver via stdin file, kill via pid, lease renew and SIGTERM detach. Reconciler `reattach` / `requeue`, attach-sweep job, boot sweeps reduced, BullMQ schedulers upserted. The kill-the-API e2e suite. | A | +| C1 | coordination | Advisory-lock claims, poller leases and unique keys, trigger CAS, webhook dedupe table, WS tokens table, chat lock. | A | +| C2 | sessions, glance, skills | Session chat over pub/sub with control intents, share revocation event, glance state table, skills in DB and the PVC removed from the chart. | A | +| D | Local relay | §4, and the chart's replica validation lifted. | B, C1 | +| E | verification | The manual pass on a real cluster (below) and the docs: `docs/reconciliation.md`, `docs/production-eks.md`, `CLAUDE.md`. | B, C1, C2 (D for Local) | + +B, C1 and C2 touch different files and run in parallel. + +## Testing + +**Unit.** Lease CAS semantics; offset accounting at line boundaries; the +start / attach / deliver / kill script builders; the fake runtime's protocol. + +**Integration (real Postgres + Redis).** Two concurrent claimers under the +advisory lock never overshoot a limit of N; two instances acquiring one lease; +`webhook_deliveries` claims; trigger `next_fire_at` CAS with two sweeps; the +unique partial indexes; attach CAS with two claimers. + +**Pipeline e2e (`api-cluster.ts`, fake runtime).** Each case asserts the run's +final state, that its log rows are contiguous with no duplicates (a counter in +the fake agent's output), and that cost was added once. + +1. One server; a Repo Task running; `SIGKILL` the server; start a new one → the run completes. +2. Two servers; a run started through A; A killed → B re-attaches within 10 s; the run completes. +3. Three servers; ten concurrent runs across kinds; one server killed → all complete; the global and per-repo limits are never exceeded (polled from the DB throughout). +4. Cancel sent to the instance that is **not** attached kills the agent. +5. A mid-run message sent to a non-attached instance reaches the agent (the fake echoes stdin). +6. A WS upgrade token minted on A is accepted by B. +7. The same webhook delivery posted to A and B fires once. +8. A schedule trigger due now with two servers sweeping fires once. +9. `SIGTERM` to the attached server: another picks the run up in under 5 s. +10. The pod goes away mid-run (fake `destroy`) → the run fails with the pod-death reason, not a silent hang. +11. Job runs and persistent-agent turns through cases 1, 4 and 5. +12. Boot with runs in flight fails nothing. + +**Web e2e.** The existing suite, unchanged, must pass. + +**Live (manual, before merging E).** A second Helm release in its own +namespace (`optio-scale`, bundled Postgres and Redis, its own NodePorts), +never the `optio` namespace in daily use. `api.replicas: 3`. Real Repo Tasks +with a real agent (one short LLM run each, as `/live-e2e` does): + +- `kubectl delete pod` of the attached API pod mid-run → the run completes and opens its PR; the UI's log view shows no gap. +- Scale 1 → 3 → 1 while runs are in flight. +- Cancel and a mid-run message from the UI while attached elsewhere (the NodePort round-robins). +- A Job run and a persistent-agent turn through the same. +- Phase D: an isolated daemon (the probe recipe, never a live session) against the 3-replica release: spawn, input, resize, kill, dirs from a browser that lands on another pod. + +## Decisions (to confirm with Jon) + +- **State home.** Correctness state in Postgres, caches per instance, transport in Redis. Chosen for one transactional truth beside the run rows; Redis stays what it is today, a queue and a bus. +- **Attach replays derivations from byte 0** rather than persisting tracker state. Simpler and cannot drift; output is small. +- **Recreate stays.** Runs survive it now, so the only cost is seconds of API downtime per deploy; rolling updates would impose expand/contract migrations on every schema change. +- **Session chat turns** keep today's "interrupted on restart" behavior in this plan; moving them onto the run protocol is the obvious follow-up and the design covers it. +- **The relay is last** and gates the replica count; a single relay-role deployment was the alternative and was rejected because the chart would still have to explain two kinds of API pod. + +## Out of scope + +- Rolling updates and expand/contract migrations. +- Session-chat turns on the run protocol (follow-up). +- Local runs (Optio Local on a machine): the daemon already owns the process; the relay phase covers their viewers. +- Postgres as a bottleneck (log insert rate, snapshot queries); measure after this lands. diff --git a/helm/optio/templates/api-deployment.yaml b/helm/optio/templates/api-deployment.yaml index e9a3c97df..0ead893dd 100644 --- a/helm/optio/templates/api-deployment.yaml +++ b/helm/optio/templates/api-deployment.yaml @@ -150,6 +150,12 @@ spec: valueFrom: fieldRef: fieldPath: spec.serviceAccountName + # Which API pod this is (services/instance.ts): the holder of + # leases and the instance a run is attached to. + - name: OPTIO_INSTANCE_ID + valueFrom: + fieldRef: + fieldPath: metadata.name resources: {{- toYaml .Values.api.resources | nindent 12 }} # Database migrations run at boot, before the server listens. The diff --git a/images/base.Dockerfile b/images/base.Dockerfile index e542af70b..bdb6da87d 100644 --- a/images/base.Dockerfile +++ b/images/base.Dockerfile @@ -15,8 +15,10 @@ FROM ubuntu:24.04@sha256:186072bba1b2f436cbb91ef2567abca677337cfc786c86e107d25b7 ENV DEBIAN_FRONTEND=noninteractive # System essentials +# tini: the pods' pid 1 (`exec tini -- sleep infinity` ends every init +# script), so the run supervisors reparented to it are reaped (run protocol). RUN apt-get update && apt-get install -y \ - git curl wget jq unzip tmux \ + git curl wget jq unzip tmux tini \ ca-certificates gnupg \ openssh-client \ && rm -rf /var/lib/apt/lists/* diff --git a/packages/container-runtime/src/docker.ts b/packages/container-runtime/src/docker.ts index dc8e1305c..3c465c609 100644 --- a/packages/container-runtime/src/docker.ts +++ b/packages/container-runtime/src/docker.ts @@ -1,7 +1,18 @@ import Docker from "dockerode"; import { Readable, Writable } from "node:stream"; import type { ContainerSpec, ContainerHandle, ContainerStatus, ExecSession } from "@optio/shared"; -import type { ContainerRuntime, LogOptions, ExecOptions } from "./types.js"; +import type { + ContainerRuntime, + LogOptions, + ExecOptions, + RunAttachInput, + RunAttachment, + RunKillInput, + RunStartInput, + RunStartResult, + RunStdinInput, +} from "./types.js"; +import { execRunProtocol } from "./run-protocol.js"; export interface DockerRuntimeOptions { socketPath?: string; @@ -133,9 +144,27 @@ export class DockerContainerRuntime implements ContainerRuntime { read() {}, }); - duplex.on("data", (chunk: Buffer) => { - stdout.push(chunk); - }); + if (opts?.tty ?? true) { + duplex.on("data", (chunk: Buffer) => { + stdout.push(chunk); + }); + } else { + // Without a TTY Docker multiplexes both streams over the one socket + // (8-byte frame headers): split them, so stderr arrives as stderr. + const out = new Writable({ + write(chunk: Buffer, _enc, cb) { + stdout.push(chunk); + cb(); + }, + }); + const err = new Writable({ + write(chunk: Buffer, _enc, cb) { + stderr.push(chunk); + cb(); + }, + }); + this.docker.modem.demuxStream(duplex, out, err); + } duplex.on("end", () => { stdout.push(null); stderr.push(null); @@ -170,6 +199,24 @@ export class DockerContainerRuntime implements ContainerRuntime { }; } + // The run protocol (run-protocol.ts), as short execs into the container. + + startRun(handle: ContainerHandle, input: RunStartInput): Promise { + return execRunProtocol.startRun(this, handle, input); + } + + attachRun(handle: ContainerHandle, input: RunAttachInput): Promise { + return execRunProtocol.attachRun(this, handle, input); + } + + deliverStdin(handle: ContainerHandle, input: RunStdinInput): Promise { + return execRunProtocol.deliverStdin(this, handle, input); + } + + killRun(handle: ContainerHandle, input: RunKillInput): Promise { + return execRunProtocol.killRun(this, handle, input); + } + async destroy(handle: ContainerHandle): Promise { const container = this.docker.getContainer(handle.id); try { diff --git a/packages/container-runtime/src/fake-agent.mjs b/packages/container-runtime/src/fake-agent.mjs new file mode 100644 index 000000000..d065554bd --- /dev/null +++ b/packages/container-runtime/src/fake-agent.mjs @@ -0,0 +1,228 @@ +#!/usr/bin/env node +/** + * The fake runtime's agent, as a process of its own (fake.ts `startRun`): + * plays the tape fake-tape.mjs plans into the run directory's + * `output.ndjson`, reads `stdin.ndjson` the way the real supervisor feeds an + * agent, and writes `exit` when it ends. It is spawned detached, so it + * outlives the API process that started it — which is what the run protocol + * is for (docs/plans/scale-out.md §1). + * + * node fake-agent.mjs + * + * `tape.json` in the run directory carries what the parent knew at start: + * the exec script (directives can hide in its base64 blobs), the container's + * env, the first stdin prompt when one was given, the PR number, and the + * timeouts. Like the real supervisor, this process writes its own pid file + * (` `) only once its TERM handler is installed, so a + * signal can never arrive before it; the parent waits for that file. + * + * Every JSON event carries a monotonically increasing `seq`, so a consumer + * that re-attached can prove its log rows are contiguous and never + * duplicated. Like claude with `--input-format stream-json`, the agent waits + * for the EOF sentinel after its result; a run that never sends one ends + * with exit 97 and a line in stderr.log, loudly. `[[mock:echo-stdin]]` turns + * every later stdin line into an assistant text event. + */ +import { + appendFileSync, + closeSync, + existsSync, + fstatSync, + openSync, + readFileSync, + readSync, + renameSync, + writeFileSync, +} from "node:fs"; +import { join } from "node:path"; +import { extractPromptText, noPromptResult, planAgentRun, planCommandRun } from "./fake-tape.mjs"; + +const runDir = process.argv[2]; +if (!runDir) { + process.stderr.write("usage: fake-agent.mjs \n"); + process.exit(2); +} + +const FILES = { + output: "output.ndjson", + stdin: "stdin.ndjson", + stderr: "stderr.log", + exit: "exit", + pid: "pid", +}; +const STDIN_EOF_SENTINEL = "__OPTIO_STDIN_EOF__"; + +const tape = JSON.parse(readFileSync(join(runDir, "tape.json"), "utf8")); + +let seq = 0; +let finished = false; + +function emit(event) { + seq += 1; + appendFileSync(join(runDir, FILES.output), JSON.stringify({ ...event, seq }) + "\n"); +} +function emitRaw(line) { + appendFileSync(join(runDir, FILES.output), line + "\n"); +} +function note(msg) { + appendFileSync(join(runDir, FILES.stderr), `[fake-agent] ${msg}\n`); +} +function finish(code) { + if (finished) return; + finished = true; + writeFileSync(join(runDir, FILES.exit), `${code}\n`); + process.exit(code); +} + +// A SIGTERM ends the run the way the real supervisor's trap does (143); a +// SIGKILL leaves no exit file, and an attach reports the run lost. +process.on("SIGTERM", () => finish(143)); +process.on("SIGINT", () => finish(130)); +process.on("SIGHUP", () => finish(129)); + +/** This process's identity beside its pid: its start time on Linux, a token elsewhere. */ +export function processStartToken() { + try { + const stat = readFileSync("/proc/self/stat", "utf8"); + return stat.slice(stat.lastIndexOf(")") + 2).split(" ")[19]; + } catch { + return `t${process.hrtime.bigint().toString(36)}`; + } +} + +// Only now — handlers in place — does the parent get to know us. +writeFileSync(join(runDir, `${FILES.pid}.tmp`), `${process.pid} ${processStartToken()}\n`); +renameSync(join(runDir, `${FILES.pid}.tmp`), join(runDir, FILES.pid)); + +const sleep = (ms) => new Promise((r) => setTimeout(r, ms)); + +// ── stdin: new complete lines of stdin.ndjson since the last read ────────── + +let stdinOffset = 0; +let stdinPartial = ""; +function readStdinLines() { + const path = join(runDir, FILES.stdin); + if (!existsSync(path)) return []; + const fd = openSync(path, "r"); + try { + const size = fstatSync(fd).size; + if (size <= stdinOffset) return []; + const buf = Buffer.alloc(size - stdinOffset); + readSync(fd, buf, 0, buf.length, stdinOffset); + stdinOffset = size; + const parts = (stdinPartial + buf.toString("utf8")).split("\n"); + stdinPartial = parts.pop() ?? ""; + return parts.filter((l) => l.length > 0); + } finally { + closeSync(fd); + } +} + +function writeLine(line) { + if ("event" in line) emit(line.event); + else emitRaw(line.raw); +} + +async function playCommand() { + const plan = planCommandRun(tape.script ?? ""); + for (const line of plan.lines) writeLine(line); + finish(plan.exitCode); +} + +async function playAgent() { + const script = tape.script ?? ""; + const specEnv = tape.specEnv ?? {}; + let echoStdin = false; + let eofSeen = false; + let sessionId = tape.sessionId; + + // Every stdin line after the prompt: the sentinel, or (with echo-stdin) a + // message to acknowledge as an assistant event. + const consumeStdin = (lines) => { + for (const line of lines) { + if (line === STDIN_EOF_SENTINEL) { + eofSeen = true; + continue; + } + if (echoStdin) { + const text = extractPromptText(line) ?? line; + emit({ + type: "assistant", + session_id: sessionId, + message: { model: "fake-model", content: [{ type: "text", text: `stdin: ${text}` }] }, + }); + } + } + }; + + // The prompt: given at start, or the first user message on stdin. + let prompt = tape.initialPrompt ?? null; + if (prompt === null) { + const timeoutMs = tape.promptTimeoutMs ?? 10_000; + const deadline = Date.now() + timeoutMs; + while (prompt === null) { + for (const line of readStdinLines()) { + const text = extractPromptText(line); + if (text !== null && prompt === null) prompt = text; + else consumeStdin([line]); + } + if (prompt !== null) break; + if (Date.now() > deadline) { + emit(noPromptResult(sessionId, timeoutMs)); + finish(1); + return; + } + await sleep(50); + } + } + + const plan = planAgentRun({ script, prompt, specEnv, prNumber: tape.prNumber, sessionId }); + sessionId = plan.sessionId; + if (plan.kind === "silent") { + finish(0); + return; + } + echoStdin = plan.echoStdin; + + emit(plan.init); + + if (plan.hang) { + // Never finishes: only a kill ends it. + for (;;) { + consumeStdin(readStdinLines()); + await sleep(100); + } + } + + if (plan.sleepMs > 0) { + const until = Date.now() + plan.sleepMs; + while (Date.now() < until) { + consumeStdin(readStdinLines()); + await sleep(Math.min(50, until - Date.now())); + } + } + + for (const line of plan.lines) writeLine(line); + + // Like claude: the turn is over, but the process waits for stdin to close. + const eofTimeoutMs = tape.eofTimeoutMs ?? 60_000; + const eofDeadline = Date.now() + eofTimeoutMs; + while (!eofSeen) { + consumeStdin(readStdinLines()); + if (eofSeen) break; + if (Date.now() > eofDeadline) { + note( + `no ${STDIN_EOF_SENTINEL} arrived within ${eofTimeoutMs}ms of the result — the worker never closed the agent's stdin`, + ); + finish(97); + return; + } + await sleep(50); + } + finish(plan.exitCode); +} + +(tape.kind === "command" ? playCommand() : playAgent()).catch((err) => { + note(`crashed: ${err?.stack ?? err}`); + finish(70); +}); diff --git a/packages/container-runtime/src/fake-protocol.test.ts b/packages/container-runtime/src/fake-protocol.test.ts new file mode 100644 index 000000000..428c12841 --- /dev/null +++ b/packages/container-runtime/src/fake-protocol.test.ts @@ -0,0 +1,329 @@ +/** + * The fake runtime's run protocol: a run is a process of its own, written to + * files under the runtime's directory, so a different runtime instance (and + * a different API process — the e2e tier proves that one) attaches to it. + */ +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { existsSync, mkdtempSync, readFileSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import type { ContainerHandle, ContainerSpec } from "@optio/shared"; +import { FakeContainerRuntime } from "./fake.js"; +import { RUN_FILES, STDIN_EOF_SENTINEL } from "./run-protocol.js"; +import type { RunAttachment } from "./types.js"; + +const AGENT_SCRIPT = [ + "set -e", + "export OPTIO_REPO_URL='https://github.com/e2e-org/e2e-repo'", + "claude -p --input-format stream-json --output-format stream-json --verbose", +].join("\n"); + +const RUN_DIR = "/home/agent/optio/runs"; + +function userMessage(text: string): string { + return JSON.stringify({ + type: "user", + message: { role: "user", content: [{ type: "text", text }] }, + }); +} + +function spec(env: Record = {}): ContainerSpec { + return { + image: "optio-agent:latest", + command: ["sleep", "infinity"], + env, + workDir: "/workspace", + labels: {}, + }; +} + +interface Event { + type?: string; + subtype?: string; + seq?: number; + is_error?: boolean; + message?: { content?: Array<{ type?: string; text?: string }> }; +} + +async function drain(attachment: RunAttachment): Promise<{ text: string; events: Event[] }> { + let text = ""; + for await (const chunk of attachment.output) text += chunk.toString(); + const events: Event[] = []; + for (const line of text.split("\n")) { + if (!line) continue; + try { + events.push(JSON.parse(line) as Event); + } catch { + // raw line + } + } + return { text, events }; +} + +const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); + +let dir: string; +let a: FakeContainerRuntime; +let handle: ContainerHandle; +const started: string[] = []; + +beforeEach(async () => { + dir = mkdtempSync(join(tmpdir(), "optio-fake-protocol-")); + a = new FakeContainerRuntime({ dir }); + handle = await a.create(spec({ OPTIO_REPO_URL: "https://github.com/e2e-org/e2e-repo" })); +}); + +afterEach(async () => { + // Never leave a detached fake agent behind. + for (const runId of started.splice(0)) + await a.killRun(handle, { runId, runDir: RUN_DIR, signal: "KILL" }).catch(() => {}); + rmSync(dir, { recursive: true, force: true }); +}); + +async function start(prompt: string, script = AGENT_SCRIPT, rt = a): Promise { + const runId = randomUUID(); + started.push(runId); + const result = await rt.startRun(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + script, + initialStdin: `${userMessage(prompt)}\n`, + }); + expect(result.pid).toBeGreaterThan(0); + return runId; +} + +describe("FakeContainerRuntime run protocol", () => { + it("a run started by one instance is attached by another on the same directory, with contiguous seq", async () => { + const runId = await start("hello [[mock:cost:0.5]]"); + const b = new FakeContainerRuntime({ dir }); + expect((await b.status(handle)).state).toBe("running"); + + const attachment = await b.attachRun(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + fromByte: 0, + }); + // The agent waits for stdin to close after its result, like claude. + await sleep(300); + await b.deliverStdin(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + line: STDIN_EOF_SENTINEL, + }); + const { events } = await drain(attachment); + expect(await attachment.exit).toEqual({ kind: "exited", code: 0 }); + + expect(events.map((e) => e.type)).toEqual(["system", "assistant", "result"]); + expect(events.map((e) => e.seq)).toEqual([1, 2, 3]); + expect(events[2]).toMatchObject({ subtype: "success", total_cost_usd: 0.5 }); + expect(readFileSync(join(b.runDirectory(runId), RUN_FILES.exit), "utf8").trim()).toBe("0"); + // The agent wrote its own pid file, pid and identity token, once its TERM handler was up. + expect(readFileSync(join(b.runDirectory(runId), RUN_FILES.pid), "utf8").trim()).toMatch( + /^\d+ \S+$/, + ); + }); + + it("re-attaching from the consumed offset yields the rest exactly once", async () => { + const runId = await start("slow [[mock:sleep:600]]"); + const first = await a.attachRun(handle, { runId, runDir: `${RUN_DIR}/${runId}`, fromByte: 0 }); + let firstText = ""; + first.output.on("data", (c: Buffer) => (firstText += c.toString())); + await sleep(250); + first.close(); + expect(await first.exit).toEqual({ kind: "detached", reason: "closed" }); + // Only whole lines count as consumed. + const consumed = firstText.lastIndexOf("\n") + 1; + expect(consumed).toBeGreaterThan(0); + expect(JSON.parse(firstText.slice(0, consumed).trim())).toMatchObject({ + type: "system", + seq: 1, + }); + + const second = await a.attachRun(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + fromByte: consumed, + }); + await sleep(800); + await a.deliverStdin(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + line: STDIN_EOF_SENTINEL, + }); + const { events } = await drain(second); + expect(await second.exit).toEqual({ kind: "exited", code: 0 }); + expect(events.map((e) => e.seq)).toEqual([2, 3]); + }); + + it("delivers mid-run stdin to the agent, which echoes it with [[mock:echo-stdin]]", async () => { + const runId = await start("chatty [[mock:echo-stdin]] [[mock:sleep:400]]"); + const attachment = await a.attachRun(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + fromByte: 0, + }); + await sleep(100); + await a.deliverStdin(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + line: userMessage("are you there?"), + }); + await sleep(600); + await a.deliverStdin(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + line: STDIN_EOF_SENTINEL, + }); + const { events } = await drain(attachment); + const texts = events.flatMap((e) => e.message?.content?.map((c) => c.text) ?? []); + expect(texts).toContain("stdin: are you there?"); + expect(await attachment.exit).toEqual({ kind: "exited", code: 0 }); + }); + + it("a failing tape exits 1 after its error result", async () => { + const runId = await start("boom [[mock:fail]]"); + const attachment = await a.attachRun(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + fromByte: 0, + }); + await sleep(200); + await a.deliverStdin(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + line: STDIN_EOF_SENTINEL, + }); + const { events } = await drain(attachment); + expect(events.at(-1)).toMatchObject({ type: "result", is_error: true }); + expect(await attachment.exit).toEqual({ kind: "exited", code: 1 }); + }); + + it("[[mock:silent]] writes nothing and exits 0", async () => { + const runId = await start("quiet [[mock:silent]]"); + const attachment = await a.attachRun(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + fromByte: 0, + }); + const { text } = await drain(attachment); + expect(text).toBe(""); + expect(await attachment.exit).toEqual({ kind: "exited", code: 0 }); + }); + + it("killRun TERM ends a hanging run with exit 143; KILL leaves it lost", async () => { + const hung = await start("stuck [[mock:hang]]"); + const attachment = await a.attachRun(handle, { + runId: hung, + runDir: `${RUN_DIR}/${hung}`, + fromByte: 0, + }); + await sleep(200); + expect( + await a.killRun(handle, { runId: hung, runDir: `${RUN_DIR}/${hung}`, signal: "TERM" }), + ).toBe(true); + await drain(attachment); + expect(await attachment.exit).toEqual({ kind: "exited", code: 143 }); + expect( + await a.killRun(handle, { runId: hung, runDir: `${RUN_DIR}/${hung}`, signal: "TERM" }), + ).toBe(false); + + const killed = await start("stuck again [[mock:hang]]"); + const second = await a.attachRun(handle, { + runId: killed, + runDir: `${RUN_DIR}/${killed}`, + fromByte: 0, + }); + await sleep(200); + expect( + await a.killRun(handle, { runId: killed, runDir: `${RUN_DIR}/${killed}`, signal: "KILL" }), + ).toBe(true); + await drain(second); + expect(await second.exit).toEqual({ kind: "lost" }); + expect(existsSync(join(a.runDirectory(killed), RUN_FILES.exit))).toBe(false); + }, 15_000); + + it("a kill-style utility exec and destroy reach protocol runs too", async () => { + const runId = await start("stuck [[mock:hang]]"); + await sleep(150); + const session = await a.exec(handle, ["bash", "-c", "pkill -TERM -f claude"]); + session.close(); + await sleep(300); + expect(await a.killRun(handle, { runId, runDir: `${RUN_DIR}/${runId}`, signal: "TERM" })).toBe( + false, + ); + expect(readFileSync(join(a.runDirectory(runId), RUN_FILES.exit), "utf8").trim()).toBe("143"); + + const other = await start("stuck [[mock:hang]]"); + await sleep(150); + await a.destroy(handle); + await sleep(200); + expect( + await a.killRun(handle, { runId: other, runDir: `${RUN_DIR}/${other}`, signal: "TERM" }), + ).toBe(false); + expect((await a.status(handle)).state).toBe("unknown"); + }); + + it("a command run prints its output and exit status and exits", async () => { + const script = "export OPTIO_COMMAND='echo hi'\nbash -lc \"$OPTIO_COMMAND\""; + const runId = randomUUID(); + started.push(runId); + await a.startRun(handle, { runId, runDir: `${RUN_DIR}/${runId}`, script, initialStdin: "" }); + const attachment = await a.attachRun(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + fromByte: 0, + }); + const { text } = await drain(attachment); + expect(text).toContain("fake command output\n[optio:exit] 0"); + expect(await attachment.exit).toEqual({ kind: "exited", code: 0 }); + }); + + it("refuses agents it cannot play, bad run ids, and multi-line stdin", async () => { + await expect( + a.startRun(handle, { + runId: randomUUID(), + runDir: RUN_DIR, + script: "set -e\n codex exec", + initialStdin: "", + }), + ).rejects.toThrow(/only plays claude-code/); + await expect( + a.startRun(handle, { + runId: "../etc", + runDir: RUN_DIR, + script: AGENT_SCRIPT, + initialStdin: "", + }), + ).rejects.toThrow(/Invalid run id/); + const runId = await start("x"); + await expect( + a.deliverStdin(handle, { runId, runDir: `${RUN_DIR}/${runId}`, line: "a\nb" }), + ).rejects.toThrow(/newline/); + await a.deliverStdin(handle, { + runId, + runDir: `${RUN_DIR}/${runId}`, + line: STDIN_EOF_SENTINEL, + }); + }); + + it("attaching to a run that was never started reports it lost", async () => { + const attachment = await a.attachRun(handle, { + runId: randomUUID(), + runDir: RUN_DIR, + fromByte: 0, + }); + const { text } = await drain(attachment); + expect(text).toBe(""); + expect(await attachment.exit).toEqual({ kind: "lost" }); + }); + + it("a non-persistent runtime keeps containers in memory and runs in a private directory", async () => { + const mem = new FakeContainerRuntime(); + expect(mem.persistent).toBe(false); + const h = await mem.create(spec()); + expect((await new FakeContainerRuntime({ dir: mem.dir }).status(h)).state).toBe("unknown"); + rmSync(mem.dir, { recursive: true, force: true }); + }); +}); diff --git a/packages/container-runtime/src/fake-tape.d.mts b/packages/container-runtime/src/fake-tape.d.mts new file mode 100644 index 000000000..1e747f02b --- /dev/null +++ b/packages/container-runtime/src/fake-tape.d.mts @@ -0,0 +1,45 @@ +// Types for fake-tape.mjs (see there). Kept by hand: the module is plain +// JavaScript so fake-agent.mjs can run it with node alone. + +export const AGENT_EXEC_MARKER: string; +export const COMMAND_RUN_MARKER: string; +export const UNSUPPORTED_AGENT_MARKERS: string[]; +export const CURSOR_EXEC_MARKER: string; + +export function extractScriptExport(script: string, name: string): string | null; +export function extractPromptText(line: string): string | null; +export function directive(prompt: string, name: string): boolean; +export function directiveArg(prompt: string, name: string): string | null; +export function collectDirectiveHaystack(script: string, prompt: string): string; + +export type TapeLine = { event: Record } | { raw: string }; + +export interface CommandPlan { + lines: TapeLine[]; + exitCode: number; + /** Whether the exec script expects `__OPTIO_RUN_EXIT__:` on stdout. */ + exitMarker: boolean; +} +export function planCommandRun(script: string): CommandPlan; + +export type AgentPlan = + | { kind: "silent"; sessionId: string } + | { + kind: "play"; + sessionId: string; + init: Record; + hang: boolean; + sleepMs: number; + lines: TapeLine[]; + exitCode: number; + echoStdin: boolean; + }; +export function planAgentRun(input: { + script: string; + prompt: string; + specEnv?: Record; + prNumber?: number; + sessionId?: string; +}): AgentPlan; + +export function noPromptResult(sessionId: string, timeoutMs: number): Record; diff --git a/packages/container-runtime/src/fake-tape.mjs b/packages/container-runtime/src/fake-tape.mjs new file mode 100644 index 000000000..42939457b --- /dev/null +++ b/packages/container-runtime/src/fake-tape.mjs @@ -0,0 +1,256 @@ +/** + * What the fake runtime plays, in one place for both ways of playing it: + * fake.ts (the in-process `exec` session) and fake-agent.mjs (the detached + * process of the run protocol). Plain JavaScript so the agent process needs + * no loader; fake.ts imports it through fake-tape.d.mts. + * + * `planAgentRun()` turns an exec script, the prompt, and the container's env + * into a plan: the init event, whether to hang, how long to sleep, then the + * lines to write (events and raw lines) and the exit code. The directives + * are the ones fake.ts documents; echo-stdin and the EOF wait are protocol + * behaviors the agent process adds on top. + */ +import { randomUUID } from "node:crypto"; + +export const AGENT_EXEC_MARKER = "--output-format stream-json"; +/** A Job's shell command (apps/api/src/services/command-run.ts). */ +export const COMMAND_RUN_MARKER = 'bash -lc "$OPTIO_COMMAND"'; +/** Markers of non-claude agent CLIs the fake cannot play — fail loudly. */ +export const UNSUPPORTED_AGENT_MARKERS = [ + " codex ", + " copilot ", + " gemini ", + " opencode ", + " openclaw ", +]; +/** + * Cursor's stream-json events are claude-shaped (system:init / assistant / + * result), so the fake plays cursor-agent execs with the standard tape. The + * prompt arrives as a positional `"$OPTIO_PROMPT"` (no stdin priming), set by + * the exec script's `export OPTIO_PROMPT='...'` line — extract it from there. + */ +export const CURSOR_EXEC_MARKER = "cursor-agent "; + +/** The value an exec script exports for `name` (`export NAME='…'`), or null. */ +export function extractScriptExport(script, name) { + const m = script.match(new RegExp(`export ${name}='([^']*(?:'\\\\''[^']*)*)'`)); + return m ? m[1].replaceAll("'\\''", "'") : null; +} + +/** The text of a stream-json user message line, or null for any other line. */ +export function extractPromptText(line) { + try { + const msg = JSON.parse(line); + if (msg.type !== "user") return null; + const blocks = msg.message?.content ?? []; + return blocks + .filter((b) => b.type === "text" && typeof b.text === "string") + .map((b) => b.text) + .join("\n"); + } catch { + return null; + } +} + +export function directive(prompt, name) { + return prompt.includes(`[[mock:${name}]]`); +} + +export function directiveArg(prompt, name) { + const m = prompt.match(new RegExp(`\\[\\[mock:${name}:([^\\]]+)\\]\\]`)); + return m ? m[1] : null; +} + +/** + * Directives can hide inside the exec script rather than the stdin prompt: + * for REPO tasks the stdin prompt is the rendered coding template (the raw + * task prompt only exists in the task FILE, which travels base64-encoded in + * the script — env blob → OPTIO_SETUP_FILES → file content). Decode base64 + * runs up to two levels deep and search everything. + */ +export function collectDirectiveHaystack(script, prompt) { + const parts = [prompt]; + const decodeOnce = (s) => { + try { + const text = Buffer.from(s, "base64").toString("utf8"); + // Reject binary-looking decodes. + return /[\x00-\x08\x0e-\x1f]/.test(text.slice(0, 200)) ? null : text; + } catch { + return null; + } + }; + for (const run of script.match(/[A-Za-z0-9+/=]{80,}/g) ?? []) { + const level1 = decodeOnce(run); + if (!level1) continue; + parts.push(level1); + for (const nested of level1.match(/[A-Za-z0-9+/=]{80,}/g) ?? []) { + const level2 = decodeOnce(nested); + if (level2) parts.push(level2); + } + } + return parts.join("\n"); +} + +/** The lines of a command run (`bash -lc "$OPTIO_COMMAND"`): raw output, then the status line. */ +export function planCommandRun(script) { + const fail = directive(script, "fail"); + const exitCode = fail ? 1 : 0; + if (directive(script, "silent")) return { lines: [], exitCode, exitMarker: false }; + return { + lines: [ + { raw: "[optio] Running command..." }, + { raw: "fake command output" }, + { raw: `[optio:exit] ${exitCode}` }, + ], + exitCode, + exitMarker: script.includes("__OPTIO_RUN_EXIT__:"), + }; +} + +/** + * The plan for an agent run. + * + * { kind: "silent" } → nothing at all, exit 0 + * { kind: "play", init, hang, sleepMs, + * lines: Array<{ event } | { raw }>, + * exitCode, echoStdin } + * + * `prNumber` names the PR `[[mock:pr]]` opens; `sessionId` is new when not given. + */ +export function planAgentRun({ script, prompt, specEnv = {}, prNumber = 1, sessionId }) { + const haystack = collectDirectiveHaystack(script, prompt); + const sid = sessionId ?? randomUUID(); + if (directive(haystack, "silent")) return { kind: "silent", sessionId: sid }; + + const init = { type: "system", subtype: "init", session_id: sid, model: "fake-model", tools: [] }; + const hang = directive(haystack, "hang"); + const sleepMs = Number(directiveArg(haystack, "sleep") ?? 0); + const lines = []; + const event = (e) => lines.push({ event: e }); + const raw = (l) => lines.push({ raw: l }); + + event({ + type: "assistant", + session_id: sid, + message: { + model: "fake-model", + usage: { input_tokens: 100, output_tokens: 25 }, + content: [{ type: "text", text: `Mock agent handled: ${prompt.slice(0, 120)}` }], + }, + }); + + for (const m of haystack.matchAll(/\[\[mock:env:([A-Z0-9_]+)\]\]/g)) { + const value = extractScriptExport(script, m[1]) ?? specEnv[m[1]] ?? ""; + raw(`env ${m[1]}=${value}`); + } + + // [[mock:file:PATH]] → the setup file the exec script would write, as + // `file PATH=`; `` when the script carries no such file. + for (const m of haystack.matchAll(/\[\[mock:file:([^\]]+)\]\]/g)) { + const blob = extractScriptExport(script, "OPTIO_SETUP_FILES") ?? specEnv.OPTIO_SETUP_FILES; + let content = ""; + if (blob) { + try { + const files = JSON.parse(Buffer.from(blob, "base64").toString("utf8")); + const want = m[1]; + const f = files.find((x) => + want.startsWith("*") ? x.path.endsWith(want.slice(1)) : x.path === want, + ); + if (f) { + content = f.contentBase64 + ? Buffer.from(f.contentBase64, "base64").toString("utf8") + : (f.content ?? ""); + } + } catch { + content = ""; + } + } + // One line: the content as a JSON string literal (log lines are split on newlines). + raw(`file ${m[1]}=${JSON.stringify(content)}`); + } + + // [[mock:script:TEXT]] → whether the exec script carries TEXT — outside + // the prompt's own export line, which carries the directive itself. + const scriptSansPrompt = script.replace(/export OPTIO_PROMPT='[^']*(?:'\\''[^']*)*'/, ""); + for (const m of haystack.matchAll(/\[\[mock:script:([^\]]+)\]\]/g)) { + raw(`script ${m[1]}=${scriptSansPrompt.includes(m[1]) ? "present" : "absent"}`); + } + + const repoUrl = (specEnv.OPTIO_REPO_URL ?? "https://github.com/mock/repo").replace(/\.git$/, ""); + const toolCall = (id, command, output) => { + event({ + type: "assistant", + session_id: sid, + message: { + model: "fake-model", + content: [{ type: "tool_use", id, name: "Bash", input: { command } }], + }, + }); + event({ + type: "user", + session_id: sid, + message: { content: [{ type: "tool_result", tool_use_id: id, content: output }] }, + }); + }; + + if (directive(haystack, "pr-mention")) { + const mentioned = `${repoUrl}/pull/9999`; + event({ + type: "assistant", + session_id: sid, + message: { + model: "fake-model", + content: [{ type: "text", text: `This is similar to ${mentioned} (see #9999).` }], + }, + }); + toolCall(`toolu_view_${randomUUID().slice(0, 8)}`, "gh pr view 9999", `url: ${mentioned}`); + } + + if (directive(haystack, "pr")) { + const prUrl = `${repoUrl}/pull/${prNumber}`; + toolCall( + `toolu_create_${randomUUID().slice(0, 8)}`, + "git push -u origin HEAD && gh pr create --fill", + `${prUrl}\n`, + ); + raw(`Opened pull request: ${prUrl}`); + } + + const isError = directive(haystack, "fail"); + const cost = Number(directiveArg(haystack, "cost") ?? 0.0123); + event({ + type: "result", + subtype: isError ? "error_during_execution" : "success", + is_error: isError, + result: isError ? "Mock agent failure" : "Mock agent success", + total_cost_usd: cost, + num_turns: 1, + duration_ms: 5, + session_id: sid, + }); + + return { + kind: "play", + sessionId: sid, + init, + hang, + sleepMs, + lines, + exitCode: isError ? 1 : 0, + echoStdin: directive(haystack, "echo-stdin"), + }; +} + +/** The error result the fake emits when no prompt ever arrived. */ +export function noPromptResult(sessionId, timeoutMs) { + return { + type: "result", + subtype: "error_during_execution", + is_error: true, + result: `FakeContainerRuntime: no prompt arrived on stdin within ${timeoutMs}ms — prompt delivery is broken`, + total_cost_usd: 0, + num_turns: 0, + duration_ms: timeoutMs, + session_id: sessionId, + }; +} diff --git a/packages/container-runtime/src/fake.ts b/packages/container-runtime/src/fake.ts index f36a28f85..c646b0fa1 100644 --- a/packages/container-runtime/src/fake.ts +++ b/packages/container-runtime/src/fake.ts @@ -41,6 +41,20 @@ * [[mock:script:TEXT]] → also print `script TEXT=present|absent`, whether * the exec script contains TEXT (a launch flag, a * cleanup line) + * [[mock:echo-stdin]] → (run protocol only) every stdin line delivered + * after the prompt becomes an assistant text event + * `stdin: ` — proves a mid-run message + * reached the agent + * + * The run protocol (startRun / attachRun / deliverStdin / killRun, + * run-protocol.ts) is played by a process of its own: `startRun` writes the + * run's directory under the runtime's directory and spawns fake-agent.mjs + * detached, which plays the same tape into `output.ndjson` and writes + * `exit` when it ends. With `OPTIO_FAKE_RUNTIME_DIR` set (or `dir` given), + * containers are kept on disk there too, so a second API process — or the + * one that replaces a killed one — sees the same pods and attaches to the + * same runs; that is what the scale-out e2e tier kills servers to prove. + * Every JSON event the agent writes carries a `seq` field. * * Non-agent execs (worktree cleanup, orphan kills, health probes) return an * immediately-ending empty session; kill-style scripts (pkill/kill) also @@ -49,29 +63,59 @@ * are played from the inline prompt. If no prompt ever arrives, the run * fails loudly — silent success would mask broken prompt delivery. */ -import { randomBytes, randomUUID } from "node:crypto"; -import { PassThrough, Writable } from "node:stream"; +import { randomBytes, randomInt, randomUUID } from "node:crypto"; +import { spawn } from "node:child_process"; +import { + appendFileSync, + closeSync, + existsSync, + fstatSync, + mkdirSync, + mkdtempSync, + openSync, + readdirSync, + readFileSync, + readSync, + rmSync, + writeFileSync, +} from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { fileURLToPath } from "node:url"; +import { PassThrough, Readable, Writable } from "node:stream"; import type { ContainerSpec, ContainerHandle, ContainerStatus, ExecSession } from "@optio/shared"; -import type { ContainerRuntime, ExecOptions, LogOptions } from "./types.js"; - -const AGENT_EXEC_MARKER = "--output-format stream-json"; -/** A Job's shell command (apps/api/src/services/command-run.ts). */ -const COMMAND_RUN_MARKER = 'bash -lc "$OPTIO_COMMAND"'; -/** Markers of non-claude agent CLIs the fake cannot play — fail loudly. */ -const UNSUPPORTED_AGENT_MARKERS = [" codex ", " copilot ", " gemini ", " opencode ", " openclaw "]; -/** - * Cursor's stream-json events are claude-shaped (system:init / assistant / - * result), so the fake plays cursor-agent execs with the standard tape. The - * prompt arrives as a positional `"$OPTIO_PROMPT"` (no stdin priming), set by - * the exec script's `export OPTIO_PROMPT='...'` line — extract it from there. - */ -const CURSOR_EXEC_MARKER = "cursor-agent "; - -/** The value an exec script exports for `name` (`export NAME='…'`), or null. */ -function extractScriptExport(script: string, name: string): string | null { - const m = script.match(new RegExp(`export ${name}='([^']*(?:'\\\\''[^']*)*)'`)); - return m ? m[1].replaceAll("'\\''", "'") : null; -} +import type { + ContainerRuntime, + ExecOptions, + LogOptions, + RunAttachInput, + RunAttachment, + RunExit, + RunKillInput, + RunStartInput, + RunStartResult, + RunStdinInput, +} from "./types.js"; +import { RUN_FILES, RUN_SIGNALS } from "./run-protocol.js"; +import { + AGENT_EXEC_MARKER, + COMMAND_RUN_MARKER, + CURSOR_EXEC_MARKER, + UNSUPPORTED_AGENT_MARKERS, + directive, + extractPromptText, + extractScriptExport, + noPromptResult, + planAgentRun, + planCommandRun, + type TapeLine, +} from "./fake-tape.mjs"; + +const FAKE_AGENT_PATH = fileURLToPath(new URL("./fake-agent.mjs", import.meta.url)); +const SAFE_RUN_ID = /^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$/; +const ATTACH_POLL_MS = 50; +/** How long after the agent's pid is gone with no exit file an attach reports the run lost. */ +const LOST_AFTER_MS = Number(process.env.OPTIO_FAKE_LOST_AFTER_MS ?? 3_000); /** * How long exec waits for a prompt on stdin before failing the run — a @@ -85,77 +129,65 @@ interface FakeContainer { createdAt: Date; } -function extractPromptText(line: string): string | null { - try { - const msg = JSON.parse(line) as { - type?: string; - message?: { content?: Array<{ type?: string; text?: string }> }; - }; - if (msg.type !== "user") return null; - const blocks = msg.message?.content ?? []; - return blocks - .filter((b) => b.type === "text" && typeof b.text === "string") - .map((b) => b.text) - .join("\n"); - } catch { - return null; - } +export interface FakeRuntimeOptions { + /** + * Where containers and runs live on disk. Given (or `OPTIO_FAKE_RUNTIME_DIR` + * set), the runtime is persistent: another instance on the same directory + * sees the same containers and runs. Otherwise a private temp directory + * holds the runs and containers stay in memory, as the old tests expect. + */ + dir?: string; } -function directive(prompt: string, name: string): boolean { - return prompt.includes(`[[mock:${name}]]`); -} +export class FakeContainerRuntime implements ContainerRuntime { + private containers = new Map(); + /** Live agent-session closers per container — for kill/destroy emulation. */ + private sessionClosers = new Map void>>(); + private prCounter = 0; + /** Where runs (and, when persistent, containers) live. */ + readonly dir: string; + /** Whether containers are kept on disk, visible to other instances. */ + readonly persistent: boolean; + + constructor(opts: FakeRuntimeOptions = {}) { + const dir = opts.dir ?? process.env.OPTIO_FAKE_RUNTIME_DIR; + this.persistent = Boolean(dir); + this.dir = dir ?? mkdtempSync(join(tmpdir(), "optio-fake-runtime-")); + mkdirSync(join(this.dir, "containers"), { recursive: true }); + mkdirSync(join(this.dir, "runs"), { recursive: true }); + } -function directiveArg(prompt: string, name: string): string | null { - const m = prompt.match(new RegExp(`\\[\\[mock:${name}:([^\\]]+)\\]\\]`)); - return m ? m[1] : null; -} + private containerFile(id: string): string { + return join(this.dir, "containers", `${id}.json`); + } -/** - * Directives can hide inside the exec script rather than the stdin prompt: - * for REPO tasks the stdin prompt is the rendered coding template (the raw - * task prompt only exists in the task FILE, which travels base64-encoded in - * the script — env blob → OPTIO_SETUP_FILES → file content). Decode base64 - * runs up to two levels deep and search everything. - */ -function collectDirectiveHaystack(script: string, prompt: string): string { - const parts = [prompt]; - const decodeOnce = (s: string): string | null => { + private getContainer(id: string): FakeContainer | undefined { + const inMemory = this.containers.get(id); + if (inMemory || !this.persistent) return inMemory; try { - const text = Buffer.from(s, "base64").toString("utf8"); - // Reject binary-looking decodes. - return /[\x00-\x08\x0e-\x1f]/.test(text.slice(0, 200)) ? null : text; + const raw = JSON.parse(readFileSync(this.containerFile(id), "utf8")) as { + spec: ContainerSpec; + createdAt: string; + }; + return { spec: raw.spec, createdAt: new Date(raw.createdAt) }; } catch { - return null; - } - }; - for (const run of script.match(/[A-Za-z0-9+/=]{80,}/g) ?? []) { - const level1 = decodeOnce(run); - if (!level1) continue; - parts.push(level1); - for (const nested of level1.match(/[A-Za-z0-9+/=]{80,}/g) ?? []) { - const level2 = decodeOnce(nested); - if (level2) parts.push(level2); + return undefined; } } - return parts.join("\n"); -} - -export class FakeContainerRuntime implements ContainerRuntime { - private containers = new Map(); - /** Live agent-session closers per container — for kill/destroy emulation. */ - private sessionClosers = new Map void>>(); - private prCounter = 0; async create(spec: ContainerSpec): Promise { const id = `fake-${randomBytes(6).toString("hex")}`; const name = spec.name ?? id; - this.containers.set(id, { spec, createdAt: new Date() }); + const container = { spec, createdAt: new Date() }; + this.containers.set(id, container); + if (this.persistent) { + writeFileSync(this.containerFile(id), JSON.stringify(container, null, 2)); + } return { id, name }; } async status(handle: ContainerHandle): Promise { - const c = this.containers.get(handle.id); + const c = this.getContainer(handle.id); if (!c) return { state: "unknown", reason: "fake container not found" }; return { state: "running", startedAt: c.createdAt }; } @@ -185,6 +217,7 @@ export class FakeContainerRuntime implements ContainerRuntime { // hanging runs are actually reapable, mirroring the real pod. if (/\bpkill\b|\bkill\b/.test(script)) { for (const closeSession of this.sessionClosers.get(handle.id) ?? []) closeSession(); + this.killRunsOf(handle.id, "TERM"); } return this.utilitySession(); } @@ -195,6 +228,259 @@ export class FakeContainerRuntime implements ContainerRuntime { for (const closeSession of this.sessionClosers.get(handle.id) ?? []) closeSession(); this.sessionClosers.delete(handle.id); this.containers.delete(handle.id); + this.killRunsOf(handle.id, "KILL"); + if (this.persistent) rmSync(this.containerFile(handle.id), { force: true }); + } + + // ── The run protocol, played by fake-agent.mjs ────────────────────────── + + private runDir(runId: string): string { + if (!SAFE_RUN_ID.test(runId)) throw new Error(`Invalid run id: ${JSON.stringify(runId)}`); + return join(this.dir, "runs", runId); + } + + /** + * The run's agent pid, when the pid file's process is alive and — on + * Linux, where /proc says — the one that wrote it (its start time + * matches; the real supervisor's pid file works the same way). A reused + * pid after a restart reads as not running. + */ + private runPid(runDir: string): number | null { + let pid: number; + let token: string | undefined; + try { + [pid, token] = readFileSync(join(runDir, RUN_FILES.pid), "utf8") + .trim() + .split(/\s+/) + .map((x, i) => (i === 0 ? Number.parseInt(x, 10) : x)) as [number, string | undefined]; + } catch { + return null; + } + if (!Number.isFinite(pid) || !token) return null; + if (!this.pidAlive(pid)) return null; + try { + const stat = readFileSync(`/proc/${pid}/stat`, "utf8"); + const start = stat.slice(stat.lastIndexOf(")") + 2).split(" ")[19]; + if (start !== token) return null; + } catch { + // No /proc here (macOS): the pid alone has to do. + } + return pid; + } + + private pidAlive(pid: number): boolean { + try { + process.kill(pid, 0); + return true; + } catch { + return false; + } + } + + /** Signals every run of a container (a kill-style utility exec, destroy). */ + private killRunsOf(containerId: string, signal: "TERM" | "KILL"): void { + let runs: string[] = []; + try { + runs = readdirSync(join(this.dir, "runs")); + } catch { + return; + } + for (const runId of runs) { + const dir = join(this.dir, "runs", runId); + let owner: string; + try { + owner = readFileSync(join(dir, "container"), "utf8").trim(); + } catch { + continue; + } + if (owner !== containerId) continue; + const pid = this.runPid(dir); + if (pid !== null && this.pidAlive(pid)) this.signal(pid, signal); + } + } + + private signal(pid: number, signal: RunKillInput["signal"]): boolean { + for (const target of [-pid, pid]) { + try { + process.kill(target, `SIG${signal}`); + return true; + } catch { + // not a group leader, or already gone: try the next form + } + } + return false; + } + + async startRun(handle: ContainerHandle, input: RunStartInput): Promise { + const container = this.getContainer(handle.id); + if (!container) throw new Error(`fake container ${handle.id} not found`); + const script = input.script; + let tape: Record; + if (script.includes(COMMAND_RUN_MARKER)) { + tape = { kind: "command", script }; + } else if (!script.includes(AGENT_EXEC_MARKER) && !script.includes(CURSOR_EXEC_MARKER)) { + throw new Error( + "FakeContainerRuntime only plays claude-code stream-json agents; " + + "use agentType/agentRuntime claude-code in e2e tests", + ); + } else { + // The prompt: cursor passes it positionally; claude gets it as the + // first stream-json user message of the run's stdin. + let initialPrompt: string | null = null; + if (script.includes(CURSOR_EXEC_MARKER)) { + initialPrompt = + extractScriptExport(script, "OPTIO_PROMPT") || (container.spec.env?.OPTIO_PROMPT ?? ""); + } else { + for (const line of input.initialStdin.split("\n")) { + const text = extractPromptText(line); + if (text !== null) { + initialPrompt = text; + break; + } + } + } + tape = { + kind: "agent", + script, + specEnv: container.spec.env ?? {}, + initialPrompt, + sessionId: randomUUID(), + promptTimeoutMs: PROMPT_TIMEOUT_MS, + eofTimeoutMs: Number(process.env.OPTIO_FAKE_EOF_TIMEOUT_MS ?? 60_000), + // Distinct across processes sharing a directory: no counter to race on. + prNumber: this.persistent ? randomInt(1, 9000) : ++this.prCounter, + }; + } + + const dir = this.runDir(input.runId); + mkdirSync(dir, { recursive: true }); + for (const f of [RUN_FILES.exit, RUN_FILES.pid]) rmSync(join(dir, f), { force: true }); + writeFileSync(join(dir, "container"), handle.id); + writeFileSync(join(dir, "tape.json"), JSON.stringify(tape, null, 2)); + writeFileSync(join(dir, RUN_FILES.output), ""); + writeFileSync(join(dir, RUN_FILES.stderr), ""); + writeFileSync(join(dir, RUN_FILES.stdin), input.initialStdin); + + // Detached, in a session of its own, and unreferenced: it outlives us. + const child = spawn(process.execPath, [FAKE_AGENT_PATH, dir], { + detached: true, + stdio: "ignore", + env: { ...process.env }, + }); + child.unref(); + if (!child.pid) throw new Error("could not spawn the fake agent"); + // Like the real supervisor, the agent writes its own pid file once its + // TERM handler is installed; only then is the run "started". + const deadline = Date.now() + 10_000; + while (!existsSync(join(dir, RUN_FILES.pid))) { + if (Date.now() > deadline || !this.pidAlive(child.pid)) { + throw new Error("the fake agent did not start"); + } + await new Promise((r) => setTimeout(r, 50)); + } + const pid = this.runPid(dir) ?? child.pid; + return { pid, output: "[fake] agent started\n" }; + } + + async attachRun(_handle: ContainerHandle, input: RunAttachInput): Promise { + const dir = this.runDir(input.runId); + const outputPath = join(dir, RUN_FILES.output); + const exitPath = join(dir, RUN_FILES.exit); + const output = new Readable({ read() {} }); + let offset = Math.max(0, input.fromByte); + let closed = false; + let settle: (exit: RunExit) => void = () => {}; + const exit = new Promise((resolve) => { + settle = resolve; + }); + const end = (value: RunExit) => { + if (closed) return; + closed = true; + output.push(null); + settle(value); + }; + + const pid = this.runPid(dir); + if (!existsSync(outputPath) || (pid === null && !existsSync(exitPath))) { + end({ kind: "lost" }); + return { output, exit, close: () => {} }; + } + + let deadSince: number | null = null; + const pump = (): boolean => { + let fd: number; + try { + fd = openSync(outputPath, "r"); + } catch { + return false; + } + try { + const size = fstatSync(fd).size; + if (size > offset) { + const buf = Buffer.alloc(size - offset); + readSync(fd, buf, 0, buf.length, offset); + offset = size; + output.push(buf); + return true; + } + return false; + } finally { + closeSync(fd); + } + }; + const tick = () => { + if (closed) return; + pump(); + if (existsSync(exitPath)) { + // The agent writes exit after its last output line: drain, then end. + pump(); + const code = Number.parseInt(readFileSync(exitPath, "utf8"), 10); + end({ kind: "exited", code: Number.isFinite(code) ? code : 1 }); + return; + } + if (pid === null || !this.pidAlive(pid)) { + deadSince ??= Date.now(); + if (Date.now() - deadSince > LOST_AFTER_MS) { + pump(); + end({ kind: "lost" }); + return; + } + } + setTimeout(tick, ATTACH_POLL_MS).unref?.(); + }; + tick(); + + return { + output, + exit, + close: () => end({ kind: "detached", reason: "closed" }), + }; + } + + async deliverStdin(_handle: ContainerHandle, input: RunStdinInput): Promise { + if (input.line.includes("\n") || input.line.includes("\r")) { + throw new Error("a stdin line must not contain a newline"); + } + const dir = this.runDir(input.runId); + if (!existsSync(dir)) throw new Error(`fake run ${input.runId} not found`); + appendFileSync(join(dir, RUN_FILES.stdin), `${input.line}\n`); + } + + async killRun(_handle: ContainerHandle, input: RunKillInput): Promise { + if (!RUN_SIGNALS.includes(input.signal)) throw new Error(`Unsupported signal: ${input.signal}`); + const dir = this.runDir(input.runId); + // A run with an exit file is over, whatever its process is still doing + // in its last milliseconds (the real kill script keys on the pid alone, + // but a bash supervisor is gone the instant it writes its exit). + if (existsSync(join(dir, RUN_FILES.exit))) return false; + const pid = this.runPid(dir); + if (pid === null || !this.pidAlive(pid)) return false; + return this.signal(pid, input.signal); + } + + /** Test helper: the run directory the fake keeps for a run. */ + runDirectory(runId: string): string { + return this.runDir(runId); } async ping(): Promise { @@ -212,13 +498,9 @@ export class FakeContainerRuntime implements ContainerRuntime { const stdin = new Writable({ write: (_c, _e, cb) => cb() }); // The command travels in the script itself (`export OPTIO_COMMAND='…'`). setImmediate(() => { - if (!directive(script, "silent")) { - stdout.write("[optio] Running command...\n"); - stdout.write("fake command output\n"); - stdout.write(`[optio:exit] ${directive(script, "fail") ? 1 : 0}\n`); - if (script.includes("__OPTIO_RUN_EXIT__:")) - stdout.write(`__OPTIO_RUN_EXIT__:${directive(script, "fail") ? 1 : 0}\n`); - } + const plan = planCommandRun(script); + for (const line of plan.lines) stdout.write(`${"raw" in line ? line.raw : ""}\n`); + if (plan.exitMarker) stdout.write(`__OPTIO_RUN_EXIT__:${plan.exitCode}\n`); stdout.end(); stderr.end(); }); @@ -242,7 +524,7 @@ export class FakeContainerRuntime implements ContainerRuntime { } private agentSession(containerId: string, script: string): ExecSession { - const spec = this.containers.get(containerId)?.spec; + const spec = this.getContainer(containerId)?.spec; const stdout = new PassThrough(); const stderr = new PassThrough(); const sessionId = randomUUID(); @@ -267,146 +549,37 @@ export class FakeContainerRuntime implements ContainerRuntime { stderr.end(); }; + const writeLine = (line: TapeLine) => { + if ("event" in line) emit(line.event); + else emitRaw(line.raw); + }; const run = async (echoPrompt: string) => { // Directives may live in the stdin prompt (standalone jobs, persistent // agents) or buried in the exec script's task file (repo tasks). - const prompt = collectDirectiveHaystack(script, echoPrompt); - if (directive(prompt, "silent")) { + const plan = planAgentRun({ + script, + prompt: echoPrompt, + specEnv: spec?.env ?? {}, + prNumber: ++this.prCounter, + sessionId, + }); + if (plan.kind === "silent") { finish(); return; } - - emit({ - type: "system", - subtype: "init", - session_id: sessionId, - model: "fake-model", - tools: [], - }); - - if (directive(prompt, "hang")) { + emit(plan.init); + if (plan.hang) { return; // stream stays open until close() } - - const sleepMs = Number(directiveArg(prompt, "sleep") ?? 0); - if (sleepMs > 0) { + if (plan.sleepMs > 0) { await new Promise((r) => { - sleepTimer = setTimeout(r, sleepMs); + sleepTimer = setTimeout(r, plan.sleepMs); sleepTimer.unref?.(); }); if (closed) return; } - - emit({ - type: "assistant", - session_id: sessionId, - message: { - model: "fake-model", - usage: { input_tokens: 100, output_tokens: 25 }, - content: [{ type: "text", text: `Mock agent handled: ${echoPrompt.slice(0, 120)}` }], - }, - }); - - for (const m of prompt.matchAll(/\[\[mock:env:([A-Z0-9_]+)\]\]/g)) { - const value = extractScriptExport(script, m[1]) ?? spec?.env?.[m[1]] ?? ""; - emitRaw(`env ${m[1]}=${value}`); - } - - // [[mock:file:PATH]] → print a setup file the exec script would write - // (`OPTIO_SETUP_FILES`), as `file PATH=`; `` when the - // script carries no such file. - for (const m of prompt.matchAll(/\[\[mock:file:([^\]]+)\]\]/g)) { - const blob = - extractScriptExport(script, "OPTIO_SETUP_FILES") ?? spec?.env?.OPTIO_SETUP_FILES; - let content = ""; - if (blob) { - try { - const files = JSON.parse(Buffer.from(blob, "base64").toString("utf8")) as Array<{ - path: string; - content?: string; - contentBase64?: string; - }>; - const want = m[1]; - const f = files.find((x) => - want.startsWith("*") ? x.path.endsWith(want.slice(1)) : x.path === want, - ); - if (f) { - content = f.contentBase64 - ? Buffer.from(f.contentBase64, "base64").toString("utf8") - : (f.content ?? ""); - } - } catch { - content = ""; - } - } - // One line: the content as a JSON string literal (log lines are split on newlines). - emitRaw(`file ${m[1]}=${JSON.stringify(content)}`); - } - - // [[mock:script:TEXT]] → whether the exec script carries TEXT — outside - // the prompt's own export line, which carries the directive itself. - const scriptSansPrompt = script.replace(/export OPTIO_PROMPT='[^']*(?:'\\''[^']*)*'/, ""); - for (const m of prompt.matchAll(/\[\[mock:script:([^\]]+)\]\]/g)) { - emitRaw(`script ${m[1]}=${scriptSansPrompt.includes(m[1]) ? "present" : "absent"}`); - } - - const repoUrl = (spec?.env?.OPTIO_REPO_URL ?? "https://github.com/mock/repo").replace( - /\.git$/, - "", - ); - const toolCall = (id: string, command: string, output: string) => { - emit({ - type: "assistant", - session_id: sessionId, - message: { - model: "fake-model", - content: [{ type: "tool_use", id, name: "Bash", input: { command } }], - }, - }); - emit({ - type: "user", - session_id: sessionId, - message: { content: [{ type: "tool_result", tool_use_id: id, content: output }] }, - }); - }; - - if (directive(prompt, "pr-mention")) { - const mentioned = `${repoUrl}/pull/9999`; - emit({ - type: "assistant", - session_id: sessionId, - message: { - model: "fake-model", - content: [{ type: "text", text: `This is similar to ${mentioned} (see #9999).` }], - }, - }); - toolCall(`toolu_view_${randomUUID().slice(0, 8)}`, "gh pr view 9999", `url: ${mentioned}`); - } - - if (directive(prompt, "pr")) { - this.prCounter += 1; - const prUrl = `${repoUrl}/pull/${this.prCounter}`; - toolCall( - `toolu_create_${randomUUID().slice(0, 8)}`, - "git push -u origin HEAD && gh pr create --fill", - `${prUrl}\n`, - ); - emitRaw(`Opened pull request: ${prUrl}`); - } - - const isError = directive(prompt, "fail"); - const cost = Number(directiveArg(prompt, "cost") ?? 0.0123); - emit({ - type: "result", - subtype: isError ? "error_during_execution" : "success", - is_error: isError, - result: isError ? "Mock agent failure" : "Mock agent success", - total_cost_usd: cost, - num_turns: 1, - duration_ms: 5, - session_id: sessionId, - }); - finish(isError ? 1 : 0); + for (const line of plan.lines) writeLine(line); + finish(plan.exitCode); }; const handlePrompt = (prompt: string) => { @@ -426,16 +599,7 @@ export class FakeContainerRuntime implements ContainerRuntime { const promptTimer = setTimeout(() => { if (promptHandled) return; promptHandled = true; - emit({ - type: "result", - subtype: "error_during_execution", - is_error: true, - result: `FakeContainerRuntime: no prompt arrived on stdin within ${PROMPT_TIMEOUT_MS}ms — prompt delivery is broken`, - total_cost_usd: 0, - num_turns: 0, - duration_ms: PROMPT_TIMEOUT_MS, - session_id: sessionId, - }); + emit(noPromptResult(sessionId, PROMPT_TIMEOUT_MS)); finish(); }, PROMPT_TIMEOUT_MS); promptTimer.unref?.(); diff --git a/packages/container-runtime/src/index.ts b/packages/container-runtime/src/index.ts index beaf1da6c..04ee22feb 100644 --- a/packages/container-runtime/src/index.ts +++ b/packages/container-runtime/src/index.ts @@ -1,4 +1,40 @@ -export type { ContainerRuntime, LogOptions, ExecOptions } from "./types.js"; +export type { + ContainerRuntime, + LogOptions, + ExecOptions, + RunSignal, + RunStartInput, + RunStartResult, + RunAttachInput, + RunAttachment, + RunExit, + RunStdinInput, + RunKillInput, +} from "./types.js"; +export { + RUN_STARTED_MARKER, + RUN_EXIT_MARKER, + RUN_LOST_MARKER, + STDIN_DELIVERED_MARKER, + STDIN_EOF_SENTINEL, + RUN_FILES, + RUN_DIR_ENV, + RUN_MARK_ENV_READY_ENV, + RUN_STDIN_BYTES_ENV, + RUN_SIGNALS, + RUN_ALIVE_HELPER, + SUPERVISOR_SCRIPT, + startRunPrelude, + superviseAgent, + attachRunScript, + deliverStdinScript, + killRunScript, + normalizeFromByte, + parseRunExitLine, + parseRunStartedPid, + RunStartError, + RunDeliverError, +} from "./run-protocol.js"; export { DockerContainerRuntime } from "./docker.js"; export type { DockerRuntimeOptions } from "./docker.js"; export { diff --git a/packages/container-runtime/src/kubernetes.test.ts b/packages/container-runtime/src/kubernetes.test.ts index f062e3076..73a731449 100644 --- a/packages/container-runtime/src/kubernetes.test.ts +++ b/packages/container-runtime/src/kubernetes.test.ts @@ -887,3 +887,64 @@ describe("KubernetesContainerRuntime", () => { }); }); }); + +describe("KubernetesContainerRuntime run protocol", () => { + it("startRun, attachRun, deliverStdin and killRun are short execs of the protocol scripts", async () => { + const runtime = new KubernetesContainerRuntime("test-ns"); + const { PassThrough, Writable } = await import("node:stream"); + const calls: string[][] = []; + const play: Array<{ stdout: string; stderr: string }> = [ + { stdout: "[optio] ready\n__OPTIO_RUN_STARTED__:12\n", stderr: "" }, + { stdout: "line\n", stderr: "__OPTIO_RUN_EXIT__:0\n" }, + { stdout: "__OPTIO_STDIN_DELIVERED__\n", stderr: "" }, + { stdout: "killed\n", stderr: "" }, + ]; + vi.spyOn(runtime, "exec").mockImplementation(async (_h, command) => { + calls.push(command); + const tape = play.shift()!; + const stdout = new PassThrough(); + const stderr = new PassThrough(); + setImmediate(() => { + stdout.end(tape.stdout); + stderr.end(tape.stderr); + }); + return { + stdin: new Writable({ write: (_c, _e, cb) => cb() }), + stdout, + stderr, + resize: () => {}, + close: () => {}, + }; + }); + const handle = { id: "p", name: "p" }; + const runDir = "/home/agent/optio/runs/r1"; + expect( + await runtime.startRun(handle, { + runId: "r1", + runDir, + script: "echo hi", + initialStdin: "x\n", + }), + ).toEqual({ + pid: 12, + output: "[optio] ready\n", + }); + const attachment = await runtime.attachRun(handle, { runId: "r1", runDir, fromByte: 5 }); + let out = ""; + for await (const c of attachment.output) out += c.toString(); + expect(out).toBe("line\n"); + expect(await attachment.exit).toEqual({ kind: "exited", code: 0 }); + await runtime.deliverStdin(handle, { runId: "r1", runDir, line: "hello" }); + expect(await runtime.killRun(handle, { runId: "r1", runDir, signal: "TERM" })).toBe(true); + expect(calls.map((c) => c.slice(0, 2))).toEqual([ + ["bash", "-c"], + ["bash", "-c"], + ["bash", "-c"], + ["bash", "-c"], + ]); + expect(calls[0][2]).toBe("exec 3<&0 0 { + return execRunProtocol.startRun(this, handle, input); + } + + attachRun(handle: ContainerHandle, input: RunAttachInput): Promise { + return execRunProtocol.attachRun(this, handle, input); + } + + deliverStdin(handle: ContainerHandle, input: RunStdinInput): Promise { + return execRunProtocol.deliverStdin(this, handle, input); + } + + killRun(handle: ContainerHandle, input: RunKillInput): Promise { + return execRunProtocol.killRun(this, handle, input); + } + async destroy(handle: ContainerHandle): Promise { try { await this.coreApi.deleteNamespacedPod({ diff --git a/packages/container-runtime/src/run-protocol.linux.test.ts b/packages/container-runtime/src/run-protocol.linux.test.ts new file mode 100644 index 000000000..358b57279 --- /dev/null +++ b/packages/container-runtime/src/run-protocol.linux.test.ts @@ -0,0 +1,211 @@ +/** + * The protocol scripts, run for real: start a supervisor with bash, attach + * with `tail --pid`, deliver stdin over the exec's stdin, close it with the + * sentinel, read the exit, kill a hanging run, and refuse a stale pid file. + * Needs GNU coreutils, util-linux (`setsid`) and /proc, so it runs on Linux + * (CI) and skips on macOS, where the Docker check in docs/plans/scale-out.md + * covers the same ground. + */ +import { describe, expect, it } from "vitest"; +import { execFileSync, spawn, spawnSync } from "node:child_process"; +import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { + RUN_FILES, + STDIN_DELIVERED_MARKER, + STDIN_EOF_SENTINEL, + attachRunScript, + deliverStdinScript, + killRunScript, + startRunPrelude, + superviseAgent, +} from "./run-protocol.js"; + +function hasGnuTools(): boolean { + if (process.platform !== "linux") return false; + const tail = spawnSync("tail", ["--version"], { encoding: "utf8" }); + const setsid = spawnSync("setsid", ["--version"], { encoding: "utf8" }); + return tail.status === 0 && /GNU coreutils/.test(tail.stdout) && setsid.status === 0; +} + +const describeLinux = hasGnuTools() ? describe : describe.skip; + +const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); + +/** Runs a script with bash, feeding `input` on stdin, returning stdout. */ +function bash(script: string, input: Buffer | string = ""): string { + return execFileSync("bash", ["-c", script], { + encoding: "utf8", + input, + maxBuffer: 64 * 1024 * 1024, + stdio: ["pipe", "pipe", "pipe"], + }); +} + +/** The start script as the runtime runs it: prelude + script, stdin bytes fed. */ +function start(runDir: string, agent: string[], initialStdin: string, opts = {}): number { + const bytes = Buffer.from(initialStdin, "utf8"); + const script = [ + ...startRunPrelude(bytes.length), + "set -e", + ...superviseAgent(runDir, agent, opts), + ].join("\n"); + const out = bash(script, bytes); + const pid = Number(out.match(/__OPTIO_RUN_STARTED__:(\d+)/)?.[1]); + expect(pid).toBeGreaterThan(0); + return pid; +} + +function deliver(runDir: string, line: string): void { + const bytes = Buffer.from(`${line}\n`, "utf8"); + const out = bash(deliverStdinScript(runDir, bytes.length), bytes); + expect(out).toContain(STDIN_DELIVERED_MARKER); +} + +/** Runs the attach script, resolving with stdout, stderr once it ends. */ +function attach(script: string): Promise<{ stdout: string; stderr: string }> { + return new Promise((resolve, reject) => { + const child = spawn("bash", ["-c", script], { stdio: ["ignore", "pipe", "pipe"] }); + let stdout = ""; + let stderr = ""; + child.stdout.on("data", (c: Buffer) => (stdout += c.toString())); + child.stderr.on("data", (c: Buffer) => (stderr += c.toString())); + child.on("error", reject); + child.on("close", () => resolve({ stdout, stderr })); + }); +} + +const ECHO_AGENT = [ + "n=0", + `while IFS= read -r l; do n=$((n+1)); printf '{"len":%d,"n":%d}\\n' "\${#l}" "$n"; done`, + `printf '{"done":true,"n":%d}\\n' "$n"`, + "exit 7", +]; + +describeLinux("the run protocol scripts under bash", () => { + it("starts a detached agent, streams its output, feeds stdin until the sentinel, records its exit", async () => { + const dir = mkdtempSync(join(tmpdir(), "optio-proto-")); + const runDir = join(dir, "run"); + try { + const pid = start(runDir, ECHO_AGENT, '{"first":true}\n'); + const [filePid, startTime] = readFileSync(join(runDir, RUN_FILES.pid), "utf8") + .trim() + .split(" "); + expect(Number(filePid)).toBe(pid); + expect(startTime).toMatch(/^\d+$/); + + const attached = attach(attachRunScript(runDir, 0)); + await sleep(300); + deliver(runDir, '{"msg":"second"}'); + await sleep(300); + const sentAt = Date.now(); + deliver(runDir, STDIN_EOF_SENTINEL); + const { stdout, stderr } = await attached; + // The sentinel closes the agent's stdin at once. + expect(Date.now() - sentAt).toBeLessThan(3_000); + expect(stdout).toBe('{"len":14,"n":1}\n{"len":16,"n":2}\n{"done":true,"n":2}\n'); + expect(stderr.trim()).toBe("__OPTIO_RUN_EXIT__:7"); + expect(readFileSync(join(runDir, RUN_FILES.exit), "utf8").trim()).toBe("7"); + + // A second attach from an offset replays exactly the rest. + const offset = stdout.indexOf("\n") + 1; + const again = await attach(attachRunScript(runDir, offset)); + expect(again.stdout).toBe(stdout.slice(offset)); + expect(again.stderr.trim()).toBe("__OPTIO_RUN_EXIT__:7"); + + // Nothing of the run is left running. + await sleep(1500); + expect(spawnSync("kill", ["-0", String(pid)]).status).not.toBe(0); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }, 30_000); + + it("carries a 300 KB first message and a 200 KB delivered line byte-exact", async () => { + const dir = mkdtempSync(join(tmpdir(), "optio-proto-")); + const runDir = join(dir, "run"); + try { + const big = `{"big":"${"x".repeat(300_000)}"}`; + start(runDir, ECHO_AGENT, `${big}\n`); + const attached = attach(attachRunScript(runDir, 0)); + await sleep(300); + const line = `{"d":"${"é".repeat(100_000)}"}`; // 200 KB of UTF-8, 100 K characters + deliver(runDir, line); + await sleep(500); + deliver(runDir, STDIN_EOF_SENTINEL); + const { stdout, stderr } = await attached; + expect(stderr.trim()).toBe("__OPTIO_RUN_EXIT__:7"); + expect(stdout).toBe( + `{"len":${big.length},"n":1}\n{"len":${line.length},"n":2}\n{"done":true,"n":2}\n`, + ); + expect(readFileSync(join(runDir, RUN_FILES.stdin), "utf8")).toBe( + `${big}\n${line}\n${STDIN_EOF_SENTINEL}\n`, + ); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }, 60_000); + + it("kills a hanging run's whole process group and records 143; a SIGKILLed supervisor is lost", async () => { + const dir = mkdtempSync(join(tmpdir(), "optio-proto-")); + const runDir = join(dir, "run"); + try { + const pid = start(runDir, ["echo '{\"init\":1}'", "sleep 1000"], ""); + const attached = attach(attachRunScript(runDir, 0)); + await sleep(300); + expect(bash(killRunScript(runDir, "TERM")).trim()).toBe("killed"); + const { stdout, stderr } = await attached; + expect(stdout).toBe('{"init":1}\n'); + expect(stderr.trim()).toBe("__OPTIO_RUN_EXIT__:143"); + expect(bash(killRunScript(runDir, "TERM")).trim()).toBe("none"); + await sleep(200); + expect(spawnSync("pgrep", ["-f", "sleep 1000"]).status).not.toBe(0); + expect(spawnSync("kill", ["-0", String(pid)]).status).not.toBe(0); + + rmSync(runDir, { recursive: true, force: true }); + const pid2 = start(runDir, ["echo '{\"init\":1}'", "sleep 1000"], ""); + execFileSync("kill", ["-KILL", "--", `-${pid2}`]); + await sleep(300); + const lost = await attach(attachRunScript(runDir, 0)); + expect(lost.stderr.trim()).toBe("__OPTIO_RUN_LOST__"); + expect(existsSync(join(runDir, RUN_FILES.exit))).toBe(false); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }, 30_000); + + it("a pid file naming a live process with another start time is a dead run", async () => { + const dir = mkdtempSync(join(tmpdir(), "optio-proto-")); + const runDir = join(dir, "run"); + try { + start(runDir, ["echo '{\"init\":1}'", "sleep 1000"], ""); + const pid = Number(readFileSync(join(runDir, RUN_FILES.pid), "utf8").split(" ")[0]); + // Pretend the pid was reused after a pod restart: same pid, other start time. + writeFileSync(join(runDir, RUN_FILES.pid), `${pid} 1\n`); + expect(bash(killRunScript(runDir, "TERM")).trim()).toBe("none"); + const lost = await attach(attachRunScript(runDir, 0)); + expect(lost.stdout).toBe('{"init":1}\n'); + expect(lost.stderr.trim()).toBe("__OPTIO_RUN_LOST__"); + // The real process is still there: end it. + execFileSync("kill", ["-KILL", "--", `-${pid}`]); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }, 30_000); + + it("a zero exit is recorded as 0", async () => { + const dir = mkdtempSync(join(tmpdir(), "optio-proto-")); + const runDir = join(dir, "run"); + try { + start(runDir, ["cat >/dev/null", "exit 0"], `${STDIN_EOF_SENTINEL}\n`, { + markEnvReady: true, + }); + const { stderr } = await attach(attachRunScript(runDir, 0)); + expect(stderr.trim()).toBe("__OPTIO_RUN_EXIT__:0"); + expect(readFileSync(join(runDir, RUN_FILES.exit), "utf8").trim()).toBe("0"); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }, 30_000); +}); diff --git a/packages/container-runtime/src/run-protocol.test.ts b/packages/container-runtime/src/run-protocol.test.ts new file mode 100644 index 000000000..c3cf5a701 --- /dev/null +++ b/packages/container-runtime/src/run-protocol.test.ts @@ -0,0 +1,454 @@ +import { describe, expect, it } from "vitest"; +import { PassThrough, Writable } from "node:stream"; +import type { ContainerHandle, ExecSession } from "@optio/shared"; +import { + RUN_ALIVE_HELPER, + RUN_EXIT_MARKER, + RUN_FILES, + RUN_LOST_MARKER, + RUN_STARTED_MARKER, + RunDeliverError, + RunStartError, + STDIN_DELIVERED_MARKER, + STDIN_EOF_SENTINEL, + SUPERVISOR_SCRIPT, + attachRunScript, + deliverStdinScript, + execRunProtocol, + killRunScript, + normalizeFromByte, + parseRunExitLine, + parseRunStartedPid, + startRunPrelude, + superviseAgent, +} from "./run-protocol.js"; + +const RUN_DIR = "/home/agent/optio/runs/0f2e1c1a-1111-4222-8333-444444444444"; +const HANDLE: ContainerHandle = { id: "pod-1", name: "pod-1" }; +const STARTTIME = `sed 's/.*) //' "/proc/$_optio_pid/stat" 2>/dev/null | cut -d' ' -f20`; + +describe("the pid helper", () => { + const helper = RUN_ALIVE_HELPER.join("\n"); + + it("trusts a pid only when the live process's start time matches the file's", () => { + expect(helper).toContain(`read -r _optio_pid _optio_start < "$1/${RUN_FILES.pid}" || return 1`); + expect(helper).toContain('[ -n "$_optio_pid" ] && [ -n "$_optio_start" ] || return 1'); + expect(helper).toContain('kill -0 "$_optio_pid" 2>/dev/null || return 1'); + expect(helper).toContain(`[ "$(${STARTTIME})" = "$_optio_start" ] || return 1`); + expect(helper).toContain(`printf '%s' "$_optio_pid"`); + }); + + it("reads the start time past the (comm), which may hold spaces", () => { + // field 22 of /proc//stat is field 20 once `pid (comm) ` is gone + expect(STARTTIME).toContain("sed 's/.*) //'"); + expect(STARTTIME).toContain("cut -d' ' -f20"); + }); +}); + +describe("startRunPrelude", () => { + it("saves the exec's stdin on fd 3, points fd 0 at /dev/null, and says how many bytes to take", () => { + expect(startRunPrelude(300_000)).toEqual([ + "exec 3<&0 0 startRunPrelude(-1)).toThrow(); + expect(() => startRunPrelude(1.5)).toThrow(); + }); +}); + +describe("superviseAgent", () => { + const lines = superviseAgent(RUN_DIR, [ + "export FOO='bar'", + "claude -p --output-format stream-json", + ]); + const script = lines.join("\n"); + + it("exports the run directory and writes the agent command verbatim in a quoted heredoc", () => { + expect(lines[0]).toBe(`export OPTIO_RUN_DIR='${RUN_DIR}'`); + expect(script).toContain(`cat > "$OPTIO_RUN_DIR"/${RUN_FILES.agent} <<'__OPTIO_AGENT_SH__'`); + expect(script).toContain( + "export FOO='bar'\nclaude -p --output-format stream-json\n__OPTIO_AGENT_SH__", + ); + }); + + it("writes the supervisor as a file and launches it in its own session", () => { + expect(script).toContain( + `cat > "$OPTIO_RUN_DIR"/${RUN_FILES.supervisor} <<'__OPTIO_SUPERVISOR_SH__'`, + ); + expect(script).toContain(SUPERVISOR_SCRIPT.trimEnd()); + expect(script).toContain( + `setsid bash "$OPTIO_RUN_DIR"/${RUN_FILES.supervisor} > /dev/null 2>&1 < /dev/null &`, + ); + }); + + it("waits for the supervisor's own pid file before reporting the run started", () => { + expect(script).toContain( + `for i in $(seq 1 200); do [ -s "$OPTIO_RUN_DIR"/${RUN_FILES.pid} ] && break; sleep 0.05; done`, + ); + expect(script).toContain( + `[ -s "$OPTIO_RUN_DIR"/${RUN_FILES.pid} ] || { echo "[optio] ERROR: the run supervisor did not start" >&2; exit 1; }`, + ); + expect(script).toContain(`read -r _optio_sup _optio_start < "$OPTIO_RUN_DIR"/${RUN_FILES.pid}`); + expect(lines.at(-1)).toBe(`echo "${RUN_STARTED_MARKER}:$_optio_sup"`); + // The start script itself never writes the pid file. + expect(script).not.toContain(`> "$OPTIO_RUN_DIR"/${RUN_FILES.pid}`); + }); + + it("primes stdin from the exec's stdin (fd 3), never from the script or env", () => { + expect(script).toContain( + `if [ -n "\${OPTIO_RUN_STDIN_BYTES:-}" ]; then head -c "$OPTIO_RUN_STDIN_BYTES" <&3 > "$OPTIO_RUN_DIR"/${RUN_FILES.stdin}; exec 3<&-; else : > "$OPTIO_RUN_DIR"/${RUN_FILES.stdin}; fi`, + ); + expect(script).not.toContain("OPTIO_RUN_STDIN="); + }); + + it("starts a fresh attempt: truncated output, no stale exit, pid or FIFO", () => { + expect(script).toContain( + `rm -f "$OPTIO_RUN_DIR"/${RUN_FILES.exit} "$OPTIO_RUN_DIR"/${RUN_FILES.pid} "$OPTIO_RUN_DIR"/${RUN_FILES.stdinPipe}`, + ); + expect(script).toContain(`: > "$OPTIO_RUN_DIR"/${RUN_FILES.output}`); + expect(script).toContain(`mkfifo "$OPTIO_RUN_DIR"/${RUN_FILES.stdinPipe}`); + }); + + it("never kills children or removes the run directory on exit, and has no EPIPE watchdog", () => { + // The only trap is the supervisor's own TERM handler, inside its heredoc. + const outsideHeredocs = script.replace(/<<'(__OPTIO_\w+__)'\n[\s\S]*?\n\1/g, ""); + expect(outsideHeredocs).not.toContain("trap"); + expect(script).not.toContain(" EXIT"); + expect(script).not.toContain("rm -rf"); + expect(script).not.toContain("kill -TERM $_optio_main_pid"); + expect(script).not.toContain("printf '\\n'"); + }); + + it("marks the environment warm only when asked", () => { + expect(script).not.toContain("OPTIO_RUN_MARK_ENV_READY=1"); + expect(superviseAgent(RUN_DIR, ["true"], { markEnvReady: true }).join("\n")).toContain( + "export OPTIO_RUN_MARK_ENV_READY=1", + ); + }); + + it("rejects a run directory outside the safe shape and an agent command holding the heredoc delimiter", () => { + expect(() => superviseAgent("/home/agent/../etc", ["true"])).toThrow(/Invalid run directory/); + expect(() => superviseAgent("relative/dir", ["true"])).toThrow(/Invalid run directory/); + expect(() => superviseAgent(RUN_DIR, ["echo __OPTIO_AGENT_SH__"])).toThrow(/heredoc/); + }); +}); + +describe("the supervisor", () => { + it("installs its TERM trap and launches the agent before writing its pid file", () => { + const trapAt = SUPERVISOR_SCRIPT.indexOf("trap on_term TERM INT HUP"); + const agentAt = SUPERVISOR_SCRIPT.indexOf(`bash "$d/${RUN_FILES.agent}"`); + const pidAt = SUPERVISOR_SCRIPT.indexOf(`> "$d/${RUN_FILES.pid}.tmp"`); + expect(trapAt).toBeGreaterThan(0); + expect(agentAt).toBeGreaterThan(trapAt); + expect(pidAt).toBeGreaterThan(agentAt); + expect(SUPERVISOR_SCRIPT).toContain( + `printf '%s %s\\n' "$$" "$(sed 's/.*) //' "/proc/$$/stat" 2>/dev/null | cut -d' ' -f20)" > "$d/${RUN_FILES.pid}.tmp" && mv "$d/${RUN_FILES.pid}.tmp" "$d/${RUN_FILES.pid}"`, + ); + }); + + it("feeds stdin through a FIFO that only the loop holds, so the sentinel closes it at once", () => { + expect(SUPERVISOR_SCRIPT).toContain( + `bash "$d/${RUN_FILES.agent}" < "$d/${RUN_FILES.stdinPipe}" >> "$d/${RUN_FILES.output}" 2>> "$d/${RUN_FILES.stderr}" &`, + ); + expect(SUPERVISOR_SCRIPT).toContain( + `tail -n +1 --pid="$agent" -f "$d/${RUN_FILES.stdin}" | ( while IFS= read -r l; do if [ "$l" = "${STDIN_EOF_SENTINEL}" ]; then break; fi; printf '%s\\n' "$l"; done > "$d/${RUN_FILES.stdinPipe}" ) &`, + ); + expect(SUPERVISOR_SCRIPT).toContain(`wait "$agent"`); + expect(SUPERVISOR_SCRIPT).toContain(`echo "$code" > "$d/${RUN_FILES.exit}"`); + }); + + it("records 143 when terminated, after passing the signal to the agent", () => { + expect(SUPERVISOR_SCRIPT).toContain( + `on_term() { kill -TERM "$agent" 2>/dev/null; wait "$agent" 2>/dev/null; echo 143 > "$d/${RUN_FILES.exit}"; exit 143; }`, + ); + }); +}); + +describe("attachRunScript", () => { + it("tails the output file from the byte after the offset while the live supervisor runs", () => { + const s = attachRunScript(RUN_DIR, 1234); + expect(s).toContain(`d='${RUN_DIR}'`); + expect(s).toContain(RUN_ALIVE_HELPER.join("\n")); + expect(s).toContain(`if pid=$(optio_run_pid "$d"); then`); + expect(s).toContain(`tail -c +1235 --pid="$pid" -f "$d/${RUN_FILES.output}"`); + expect(s).toContain(`tail -c +1235 "$d/${RUN_FILES.output}"`); + }); + + it("reports the exit code, or a lost run, on stderr only", () => { + const s = attachRunScript(RUN_DIR, 0); + expect(s).toContain("tail -c +1 "); + expect(s).toContain(`echo "${RUN_EXIT_MARKER}:$(cat "$d/${RUN_FILES.exit}")" >&2`); + expect(s).toContain(`echo "${RUN_LOST_MARKER}" >&2`); + expect(s.match(new RegExp(RUN_LOST_MARKER, "g"))).toHaveLength(2); + }); + + it("normalizes the offset", () => { + expect(normalizeFromByte(-5)).toBe(0); + expect(normalizeFromByte(7.9)).toBe(7); + expect(normalizeFromByte(Number.NaN)).toBe(0); + expect(normalizeFromByte(Number.POSITIVE_INFINITY)).toBe(0); + expect(attachRunScript(RUN_DIR, Number.NaN)).toContain("tail -c +1 "); + }); +}); + +describe("deliverStdinScript", () => { + it("appends exactly the given bytes from the exec's stdin and prints the delivered marker", () => { + const s = deliverStdinScript(RUN_DIR, 42); + expect(s).toContain(`d='${RUN_DIR}'`); + expect(s).toContain(`[ -d "$d" ] || { echo "no run directory at $d" >&2; exit 1; }`); + expect(s).toContain( + `head -c 42 >> "$d/${RUN_FILES.stdin}" && echo "${STDIN_DELIVERED_MARKER}"`, + ); + expect(s).not.toContain("OPTIO_STDIN_LINE"); + }); + + it("refuses a byte count below one", () => { + expect(() => deliverStdinScript(RUN_DIR, 0)).toThrow(/byteLength/); + expect(() => deliverStdinScript(RUN_DIR, 2.5)).toThrow(/byteLength/); + }); +}); + +describe("killRunScript", () => { + it("signals the live supervisor's process group, falling back to the pid", () => { + const s = killRunScript(RUN_DIR, "TERM"); + expect(s).toContain(RUN_ALIVE_HELPER.join("\n")); + expect(s).toContain(`if pid=$(optio_run_pid "$d"); then`); + expect(s).toContain( + `kill -TERM -- -"$pid" 2>/dev/null || kill -TERM "$pid" 2>/dev/null || true`, + ); + expect(s).toContain("echo killed"); + expect(s).toContain("echo none"); + }); + + it("only knows the four signals", () => { + expect(() => killRunScript(RUN_DIR, "USR1" as never)).toThrow(/Unsupported signal/); + }); +}); + +describe("markers", () => { + it("parses exit and lost lines", () => { + expect(parseRunExitLine(`${RUN_EXIT_MARKER}:0`)).toEqual({ kind: "exited", code: 0 }); + expect(parseRunExitLine(` ${RUN_EXIT_MARKER}:143 \n`)).toEqual({ kind: "exited", code: 143 }); + expect(parseRunExitLine(`${RUN_EXIT_MARKER}:x`)).toEqual({ kind: "exited", code: 1 }); + expect(parseRunExitLine(RUN_LOST_MARKER)).toEqual({ kind: "lost" }); + expect(parseRunExitLine("tail: warning")).toBeNull(); + }); + + it("parses the started pid", () => { + expect(parseRunStartedPid(`[optio] Repo ready\n${RUN_STARTED_MARKER}:4242\n`)).toBe(4242); + expect(parseRunStartedPid("[optio] ERROR: repo not ready")).toBeNull(); + }); +}); + +type StubSession = ExecSession & { + stdoutW: PassThrough; + stderrW: PassThrough; + stdinData: Buffer[]; +}; + +/** A stub exec: records the command and what was written to stdin, plays scripted output. */ +function stubExec(play: (cmd: string[], session: StubSession) => void) { + const calls: string[][] = []; + const sessions: StubSession[] = []; + const rt = { + async exec(_h: ContainerHandle, command: string[]): Promise { + calls.push(command); + const stdout = new PassThrough(); + const stderr = new PassThrough(); + const stdinData: Buffer[] = []; + let stdinEnded = false; + const stdin = new Writable({ + write: (c: Buffer, _e, cb) => { + stdinData.push(Buffer.from(c)); + cb(); + }, + final: (cb) => { + stdinEnded = true; + cb(); + }, + }); + let closed = false; + const session: StubSession = { + stdin, + stdout, + stderr, + stdoutW: stdout, + stderrW: stderr, + stdinData, + resize: () => {}, + close: () => { + if (closed) return; + closed = true; + stdout.end(); + stderr.end(); + }, + }; + Object.defineProperty(session, "stdinEnded", { get: () => stdinEnded }); + sessions.push(session); + setImmediate(() => play(command, session)); + return session; + }, + }; + return { rt, calls, sessions }; +} + +describe("execRunProtocol", () => { + it("startRun prepends the prelude, feeds the initial stdin bytes over the exec's stdin, and parses the pid", async () => { + const { rt, calls, sessions } = stubExec((_cmd, s) => { + // The script reads stdin before printing: play after a tick. + setTimeout(() => { + s.stdoutW.end(`[optio] Repo ready\n${RUN_STARTED_MARKER}:777\n`); + s.stderrW.end(); + }, 10); + }); + const initialStdin = `{"type":"user","text":"héllo"}\n`; // 2 bytes for é + const result = await execRunProtocol.startRun(rt, HANDLE, { + script: "set -e\necho hi", + initialStdin, + }); + expect(result).toEqual({ pid: 777, output: "[optio] Repo ready\n" }); + const bytes = Buffer.byteLength(initialStdin, "utf8"); + expect(bytes).toBe(initialStdin.length + 1); + expect(calls[0].slice(0, 2)).toEqual(["bash", "-c"]); + expect(calls[0][2]).toBe(`${startRunPrelude(bytes).join("\n")}\nset -e\necho hi`); + expect(calls[0][2]).not.toContain("héllo"); + expect(Buffer.concat(sessions[0].stdinData).toString("utf8")).toBe(initialStdin); + expect((sessions[0] as unknown as { stdinEnded: boolean }).stdinEnded).toBe(true); + }); + + it("startRun fails with the script's output when no pid was printed", async () => { + const { rt } = stubExec((_cmd, s) => { + s.stdoutW.end("[optio] Waiting for repo to be ready...\n[optio] ERROR: repo not ready\n"); + s.stderrW.end("bash: line 3: git: not found\n"); + }); + await expect( + execRunProtocol.startRun(rt, HANDLE, { script: "exit 1", initialStdin: "" }), + ).rejects.toSatisfy((err: unknown) => { + expect(err).toBeInstanceOf(RunStartError); + expect((err as Error).message).toContain("ended without launching the agent"); + expect((err as Error).message).toContain("repo not ready"); + expect((err as RunStartError).output).toContain("git: not found"); + return true; + }); + }); + + it("attachRun streams stdout as the output and resolves the exit from the stderr marker", async () => { + const { rt, calls } = stubExec((_cmd, s) => { + s.stdoutW.write('{"seq":1}\n'); + s.stdoutW.write('{"seq":2}\n'); + s.stderrW.write("tail: a warning\n"); + s.stderrW.write(`${RUN_EXIT_MARKER}:3\n`); + s.stdoutW.end(); + s.stderrW.end(); + }); + const attachment = await execRunProtocol.attachRun(rt, HANDLE, { + runDir: RUN_DIR, + fromByte: 10, + }); + let out = ""; + for await (const chunk of attachment.output) out += chunk.toString(); + expect(out).toBe('{"seq":1}\n{"seq":2}\n'); + expect(await attachment.exit).toEqual({ kind: "exited", code: 3 }); + expect(calls[0][2]).toBe(attachRunScript(RUN_DIR, 10)); + }); + + it("attachRun's exit settles only once the output has fully ended, however early the marker came", async () => { + const { rt } = stubExec((_cmd, s) => { + // The marker first; stdout keeps flowing for a while. + s.stderrW.end(`${RUN_EXIT_MARKER}:0\n`); + s.stdoutW.write("early\n"); + setTimeout(() => { + s.stdoutW.write("late\n"); + s.stdoutW.end(); + }, 60); + }); + const attachment = await execRunProtocol.attachRun(rt, HANDLE, { + runDir: RUN_DIR, + fromByte: 0, + }); + let seen = ""; + attachment.output.on("data", (c: Buffer) => (seen += c.toString())); + const exit = await attachment.exit; + expect(exit).toEqual({ kind: "exited", code: 0 }); + expect(seen).toBe("early\nlate\n"); + }); + + it("attachRun resolves lost from the marker and detached when the exec drops without one", async () => { + const lost = stubExec((_cmd, s) => { + s.stderrW.end(`${RUN_LOST_MARKER}\n`); + s.stdoutW.end(); + }); + expect( + await ( + await execRunProtocol.attachRun(lost.rt, HANDLE, { runDir: RUN_DIR, fromByte: 0 }) + ).exit, + ).toEqual({ kind: "lost" }); + + const dropped = stubExec((_cmd, s) => { + s.stdoutW.write("partial"); + s.stdoutW.end(); + s.stderrW.end(); + }); + expect( + await ( + await execRunProtocol.attachRun(dropped.rt, HANDLE, { runDir: RUN_DIR, fromByte: 0 }) + ).exit, + ).toEqual({ kind: "detached", reason: "stream-ended" }); + + const closedByUs = stubExec(() => {}); + const attachment = await execRunProtocol.attachRun(closedByUs.rt, HANDLE, { + runDir: RUN_DIR, + fromByte: 0, + }); + attachment.close(); + expect(await attachment.exit).toEqual({ kind: "detached", reason: "closed" }); + }); + + it("deliverStdin feeds the line over stdin and judges by the delivered marker on stdout", async () => { + const ok = stubExec((_cmd, s) => { + setTimeout(() => { + s.stdoutW.end(`${STDIN_DELIVERED_MARKER}\n`); + s.stderrW.end("some stderr noise\n"); + }, 10); + }); + const line = `{"type":"user","text":"ünïcode"}`; + await execRunProtocol.deliverStdin(ok.rt, HANDLE, { runDir: RUN_DIR, line }); + const bytes = Buffer.byteLength(`${line}\n`, "utf8"); + expect(ok.calls[0][2]).toBe(deliverStdinScript(RUN_DIR, bytes)); + expect(ok.calls[0][2]).not.toContain("ünïcode"); + expect(Buffer.concat(ok.sessions[0].stdinData).toString("utf8")).toBe(`${line}\n`); + + const bad = stubExec((_cmd, s) => { + s.stdoutW.end(); + s.stderrW.end("no run directory at /home/agent/optio/runs/x\n"); + }); + await expect( + execRunProtocol.deliverStdin(bad.rt, HANDLE, { runDir: RUN_DIR, line: "hello" }), + ).rejects.toSatisfy((err: unknown) => { + expect(err).toBeInstanceOf(RunDeliverError); + expect((err as Error).message).toContain("no run directory"); + return true; + }); + await expect( + execRunProtocol.deliverStdin(bad.rt, HANDLE, { runDir: RUN_DIR, line: "a\nb" }), + ).rejects.toThrow(/newline/); + }); + + it("killRun reports whether anything was signalled", async () => { + const killed = stubExec((_cmd, s) => { + s.stdoutW.end("killed\n"); + s.stderrW.end(); + }); + expect( + await execRunProtocol.killRun(killed.rt, HANDLE, { runDir: RUN_DIR, signal: "TERM" }), + ).toBe(true); + expect(killed.calls[0][2]).toBe(killRunScript(RUN_DIR, "TERM")); + const none = stubExec((_cmd, s) => { + s.stdoutW.end("none\n"); + s.stderrW.end(); + }); + expect( + await execRunProtocol.killRun(none.rt, HANDLE, { runDir: RUN_DIR, signal: "KILL" }), + ).toBe(false); + }); +}); diff --git a/packages/container-runtime/src/run-protocol.ts b/packages/container-runtime/src/run-protocol.ts new file mode 100644 index 000000000..825ee0f70 --- /dev/null +++ b/packages/container-runtime/src/run-protocol.ts @@ -0,0 +1,451 @@ +/** + * The run protocol (docs/plans/scale-out.md §1): how Optio starts an agent in + * a pod so that no API process has to stay attached to it, and how any API + * process attaches later. + * + * Every run has a directory in its pod (`runDir`, the run home) with: + * + * agent.sh the agent command, as the start script wrote it + * supervisor.sh the supervisor: runs agent.sh detached from the exec + * stdin.ndjson every stdin line the run was given; `__OPTIO_STDIN_EOF__` + * closes the agent's stdin + * stdin.pipe the FIFO the agent reads stdin from + * output.ndjson the agent's stdout, appended as it runs + * stderr.log the agent's stderr + * pid ` ` of the supervisor (a process-group + * leader, `setsid`), written by the supervisor itself once + * its TERM trap is in place; starttime (field 22 of + * /proc//stat) tells a live supervisor from a reused + * pid after a pod restart + * exit the agent's exit code, written when it ends + * + * The four operations are short execs: + * + * start the caller's setup script, then `superviseAgent()`: prime stdin + * from the exec's own stdin, launch the supervisor, wait for its + * pid file, print `__OPTIO_RUN_STARTED__:` + * attach `tail -c + --pid= -f output.ndjson`, then + * `__OPTIO_RUN_EXIT__:` (or `__OPTIO_RUN_LOST__`) on stderr + * deliver `head -c >> stdin.ndjson` from the exec's stdin, then + * `__OPTIO_STDIN_DELIVERED__` on stdout + * kill signal the supervisor's process group + * + * Stdin bytes travel over the exec's stdin, never in the script or its + * environment (an argv is at most 128 KiB). The pod's pid 1 must reap: the + * supervisor is reparented to it when the start exec ends (`tini` in the + * agent image, images/base.Dockerfile). + * + * The Kubernetes and Docker runtimes run these scripts through `exec`; the + * fake runtime plays them with local processes and the same files. + */ +import { shellQuote } from "@optio/shared"; +import type { ContainerHandle, ExecSession } from "@optio/shared"; +import { Readable } from "node:stream"; +import type { + ContainerRuntime, + RunAttachment, + RunExit, + RunSignal, + RunStartResult, +} from "./types.js"; + +export const RUN_STARTED_MARKER = "__OPTIO_RUN_STARTED__"; +export const RUN_EXIT_MARKER = "__OPTIO_RUN_EXIT__"; +export const RUN_LOST_MARKER = "__OPTIO_RUN_LOST__"; +export const STDIN_DELIVERED_MARKER = "__OPTIO_STDIN_DELIVERED__"; +/** The stdin line that closes the agent's stdin (how a Claude turn is ended). */ +export const STDIN_EOF_SENTINEL = "__OPTIO_STDIN_EOF__"; + +/** The files of a run directory. */ +export const RUN_FILES = { + agent: "agent.sh", + supervisor: "supervisor.sh", + stdin: "stdin.ndjson", + stdinPipe: "stdin.pipe", + output: "output.ndjson", + stderr: "stderr.log", + pid: "pid", + exit: "exit", +} as const; + +/** The env var the supervisor reads the run directory from. */ +export const RUN_DIR_ENV = "OPTIO_RUN_DIR"; +/** Set to `1` for the supervisor to mark the pod's environment warm on success. */ +export const RUN_MARK_ENV_READY_ENV = "OPTIO_RUN_MARK_ENV_READY"; +/** How many bytes of the exec's stdin the start script primes `stdin.ndjson` with. */ +export const RUN_STDIN_BYTES_ENV = "OPTIO_RUN_STDIN_BYTES"; +/** The exec's stdin, saved by the start prelude for the start script to read. */ +const STDIN_FD = 3; +/** How long the start script waits for the supervisor's pid file (50 ms polls). */ +const START_WAIT_POLLS = 200; + +export const RUN_SIGNALS: readonly RunSignal[] = ["TERM", "KILL", "INT", "HUP"]; + +const SAFE_RUN_DIR = /^\/[A-Za-z0-9_./-]+$/; + +function assertRunDir(runDir: string): string { + if (!SAFE_RUN_DIR.test(runDir) || runDir.includes("..")) { + throw new Error(`Invalid run directory: ${JSON.stringify(runDir)}`); + } + return runDir; +} + +/** The byte offset an attach starts from, as a whole non-negative number. */ +export function normalizeFromByte(fromByte: number): number { + return Number.isFinite(fromByte) ? Math.max(0, Math.floor(fromByte)) : 0; +} + +/** A process's start time: field 22 of /proc//stat, read past the `(comm)` (which may hold spaces). */ +const PROC_STARTTIME = (pid: string) => + `sed 's/.*) //' "/proc/${pid}/stat" 2>/dev/null | cut -d' ' -f20`; + +/** + * The shell helper every script that trusts a pid file shares: + * `optio_run_pid ` prints the supervisor's pid when the process is + * alive AND is the one that wrote the file (its start time matches); a dead, + * absent, or reused pid prints nothing and fails. One helper for attach, + * kill, and the pooled start guard. + */ +export const RUN_ALIVE_HELPER: readonly string[] = [ + `optio_run_pid() {`, + ` [ -s "$1/${RUN_FILES.pid}" ] || return 1`, + ` read -r _optio_pid _optio_start < "$1/${RUN_FILES.pid}" || return 1`, + ` [ -n "$_optio_pid" ] && [ -n "$_optio_start" ] || return 1`, + ` kill -0 "$_optio_pid" 2>/dev/null || return 1`, + ` [ "$(${PROC_STARTTIME("$_optio_pid")})" = "$_optio_start" ] || return 1`, + ` printf '%s' "$_optio_pid"`, + `}`, +]; + +/** + * The supervisor, as a file: launches the agent with its stdin fed from + * `stdin.ndjson` through a FIFO, writes its pid file once the TERM trap and + * the agent are in place (so a signal can never arrive before the trap), + * waits for the agent, and records its exit. Written by the start script + * with a quoted heredoc, so nothing in it expands at write time. + */ +export const SUPERVISOR_SCRIPT: string = [ + `#!/bin/bash`, + `d="\$${RUN_DIR_ENV}"`, + `on_term() { kill -TERM "$agent" 2>/dev/null; wait "$agent" 2>/dev/null; echo 143 > "$d/${RUN_FILES.exit}"; exit 143; }`, + `trap on_term TERM INT HUP`, + `bash "$d/${RUN_FILES.agent}" < "$d/${RUN_FILES.stdinPipe}" >> "$d/${RUN_FILES.output}" 2>> "$d/${RUN_FILES.stderr}" &`, + `agent=$!`, + // The feeder: every line of stdin.ndjson, as it arrives, until the EOF + // sentinel. Only the loop holds the FIFO, so `break` closes the agent's + // stdin at once, whatever the tail is still doing; `--pid` ends the tail + // when the agent is gone. + `tail -n +1 --pid="$agent" -f "$d/${RUN_FILES.stdin}" | ( while IFS= read -r l; do if [ "$l" = "${STDIN_EOF_SENTINEL}" ]; then break; fi; printf '%s\\n' "$l"; done > "$d/${RUN_FILES.stdinPipe}" ) &`, + `printf '%s %s\\n' "$$" "$(${PROC_STARTTIME("$$")})" > "$d/${RUN_FILES.pid}.tmp" && mv "$d/${RUN_FILES.pid}.tmp" "$d/${RUN_FILES.pid}"`, + `wait "$agent"`, + `code=$?`, + `echo "$code" > "$d/${RUN_FILES.exit}"`, + `if [ "$code" -eq 0 ] && [ "\${${RUN_MARK_ENV_READY_ENV}:-}" = "1" ]; then touch /home/agent/.optio-env-ready; fi`, + ``, +].join("\n"); + +/** + * The lines a runtime puts before a start script: the exec's stdin saved on + * fd 3 for `superviseAgent()` to prime the run's stdin from (fd 0 then + * points at /dev/null, so nothing in the setup can eat those bytes), and + * how many bytes to take. + */ +export function startRunPrelude(stdinBytes: number): string[] { + if (!Number.isInteger(stdinBytes) || stdinBytes < 0) throw new Error("stdinBytes must be ≥ 0"); + return [`exec ${STDIN_FD}<&0 0 ${D}/${RUN_FILES.agent} <<'__OPTIO_AGENT_SH__'`, + agentScript, + `__OPTIO_AGENT_SH__`, + `cat > ${D}/${RUN_FILES.supervisor} <<'__OPTIO_SUPERVISOR_SH__'`, + SUPERVISOR_SCRIPT.trimEnd(), + `__OPTIO_SUPERVISOR_SH__`, + // A fresh attempt: no exit or pid from a previous one, empty output. + `rm -f ${D}/${RUN_FILES.exit} ${D}/${RUN_FILES.pid} ${D}/${RUN_FILES.stdinPipe}`, + `: > ${D}/${RUN_FILES.output}`, + `: > ${D}/${RUN_FILES.stderr}`, + `if [ -n "\${${RUN_STDIN_BYTES_ENV}:-}" ]; then head -c "\$${RUN_STDIN_BYTES_ENV}" <&${STDIN_FD} > ${D}/${RUN_FILES.stdin}; exec ${STDIN_FD}<&-; else : > ${D}/${RUN_FILES.stdin}; fi`, + `mkfifo ${D}/${RUN_FILES.stdinPipe}`, + `setsid bash ${D}/${RUN_FILES.supervisor} > /dev/null 2>&1 < /dev/null &`, + `for i in $(seq 1 ${START_WAIT_POLLS}); do [ -s ${D}/${RUN_FILES.pid} ] && break; sleep 0.05; done`, + `[ -s ${D}/${RUN_FILES.pid} ] || { echo "[optio] ERROR: the run supervisor did not start" >&2; exit 1; }`, + `read -r _optio_sup _optio_start < ${D}/${RUN_FILES.pid}`, + `echo "${RUN_STARTED_MARKER}:$_optio_sup"`, + ]; +} + +/** + * The attach script: the output file from byte `fromByte` on, following it + * while the supervisor lives, then the exit marker on stderr (stdout stays + * pure file bytes, so the consumer's offset accounting is exact). A pid file + * whose process is gone, or belongs to another process now (a pod restart + * reused the pid), counts as not running. + */ +export function attachRunScript(runDir: string, fromByte: number): string { + const d = shellQuote(assertRunDir(runDir)); + const from = normalizeFromByte(fromByte); + return [ + `d=${d}`, + ...RUN_ALIVE_HELPER, + `if [ ! -f "$d/${RUN_FILES.output}" ]; then echo "${RUN_LOST_MARKER}" >&2; exit 0; fi`, + `if pid=$(optio_run_pid "$d"); then`, + ` tail -c +${from + 1} --pid="$pid" -f "$d/${RUN_FILES.output}"`, + `else`, + ` tail -c +${from + 1} "$d/${RUN_FILES.output}"`, + `fi`, + `for i in $(seq 1 20); do [ -f "$d/${RUN_FILES.exit}" ] && break; sleep 0.5; done`, + `if [ -f "$d/${RUN_FILES.exit}" ]; then echo "${RUN_EXIT_MARKER}:$(cat "$d/${RUN_FILES.exit}")" >&2; else echo "${RUN_LOST_MARKER}" >&2; fi`, + ].join("\n"); +} + +/** + * The deliver script: `byteLength` bytes of the exec's stdin appended to the + * run's stdin file, then the delivered marker on stdout. The bytes never + * touch the script or its environment. + */ +export function deliverStdinScript(runDir: string, byteLength: number): string { + const d = shellQuote(assertRunDir(runDir)); + if (!Number.isInteger(byteLength) || byteLength < 1) throw new Error("byteLength must be ≥ 1"); + return [ + `d=${d}`, + `[ -d "$d" ] || { echo "no run directory at $d" >&2; exit 1; }`, + `head -c ${byteLength} >> "$d/${RUN_FILES.stdin}" && echo "${STDIN_DELIVERED_MARKER}"`, + ].join("\n"); +} + +/** The kill script: the signal to the supervisor's process group, when it is the live supervisor. */ +export function killRunScript(runDir: string, signal: RunSignal): string { + const d = shellQuote(assertRunDir(runDir)); + if (!RUN_SIGNALS.includes(signal)) throw new Error(`Unsupported signal: ${signal}`); + return [ + `d=${d}`, + ...RUN_ALIVE_HELPER, + `if pid=$(optio_run_pid "$d"); then`, + ` kill -${signal} -- -"$pid" 2>/dev/null || kill -${signal} "$pid" 2>/dev/null || true`, + ` echo killed`, + `else`, + ` echo none`, + `fi`, + ].join("\n"); +} + +/** Reads a whole stream as a string. */ +function collect(stream: NodeJS.ReadableStream): Promise { + return new Promise((resolve, reject) => { + let buf = ""; + stream.on("data", (chunk: Buffer | string) => { + buf += chunk.toString(); + }); + stream.on("end", () => resolve(buf)); + stream.on("error", reject); + }); +} + +/** Writes `bytes` to the exec's stdin and closes it. */ +function feedStdin(session: ExecSession, bytes: Buffer): void { + try { + if (bytes.length > 0) session.stdin.write(bytes); + session.stdin.end(); + } catch { + // The exec may have ended already; its output says what happened. + } +} + +/** Parses `__OPTIO_RUN_EXIT__:` / `__OPTIO_RUN_LOST__` out of a line, or null. */ +export function parseRunExitLine(line: string): RunExit | null { + const trimmed = line.trim(); + if (trimmed === RUN_LOST_MARKER) return { kind: "lost" }; + if (trimmed.startsWith(`${RUN_EXIT_MARKER}:`)) { + const code = Number.parseInt(trimmed.slice(RUN_EXIT_MARKER.length + 1), 10); + return { kind: "exited", code: Number.isFinite(code) ? code : 1 }; + } + return null; +} + +/** Parses the started marker out of a start script's output, or null. */ +export function parseRunStartedPid(output: string): number | null { + const m = output.match(new RegExp(`${RUN_STARTED_MARKER}:(\\d+)`)); + return m ? Number.parseInt(m[1], 10) : null; +} + +export class RunStartError extends Error { + constructor( + message: string, + public readonly output: string, + ) { + super(message); + this.name = "RunStartError"; + } +} + +/** Thrown when a stdin line could not be delivered (the script failed). */ +export class RunDeliverError extends Error { + constructor( + message: string, + public readonly output: string, + ) { + super(message); + this.name = "RunDeliverError"; + } +} + +/** + * The exec-based implementation of the protocol, shared by the runtimes + * that run scripts in a container (Kubernetes, Docker). + */ +export const execRunProtocol = { + async startRun( + rt: Pick, + handle: ContainerHandle, + input: { script: string; initialStdin: string }, + ): Promise { + const stdin = Buffer.from(input.initialStdin, "utf8"); + const script = [...startRunPrelude(stdin.length), input.script].join("\n"); + const session = await rt.exec(handle, ["bash", "-c", script], { tty: false }); + const collected = Promise.all([collect(session.stdout), collect(session.stderr)]); + feedStdin(session, stdin); + const [stdout, stderr] = await collected; + const pid = parseRunStartedPid(stdout); + if (pid === null) { + const output = [stdout, stderr].filter(Boolean).join("\n"); + throw new RunStartError( + `the run's start script ended without launching the agent:\n${output.trim().split("\n").slice(-20).join("\n")}`, + output, + ); + } + return { pid, output: stdout.replace(new RegExp(`${RUN_STARTED_MARKER}:\\d+\\n?`), "") }; + }, + + async attachRun( + rt: Pick, + handle: ContainerHandle, + input: { runDir: string; fromByte: number }, + ): Promise { + const session: ExecSession = await rt.exec( + handle, + ["bash", "-c", attachRunScript(input.runDir, input.fromByte)], + { tty: false }, + ); + let detached = false; + // The marker, from stderr; or "detached" when stderr ends without one. + const marker = new Promise((resolve) => { + let buf = ""; + let settled = false; + const settle = (value: RunExit) => { + if (settled) return; + settled = true; + resolve(value); + }; + session.stderr.on("data", (chunk: Buffer | string) => { + buf += chunk.toString(); + let idx: number; + while ((idx = buf.indexOf("\n")) >= 0) { + const line = buf.slice(0, idx); + buf = buf.slice(idx + 1); + const parsed = parseRunExitLine(line); + if (parsed) settle(parsed); + } + }); + session.stderr.on("end", () => { + const parsed = parseRunExitLine(buf); + if (parsed) settle(parsed); + // A severed exec (the attach was closed, or the connection dropped) + // ends the streams without a marker: that is "detached", not lost. + settle({ kind: "detached", reason: detached ? "closed" : "stream-ended" }); + }); + session.stderr.on("error", () => settle({ kind: "detached", reason: "stream-ended" })); + }); + // The output must be fully delivered before the exit is reported, so a + // consumer that awaits `exit` has seen every byte. A consumer that + // closes the attachment abandons the output: that settles at once. + let outputDone: () => void = () => {}; + const outputEnded = new Promise((resolve) => { + outputDone = resolve; + // "end" needs a reader; "finish" is the writer's end (the exec closed + // its stream), which is what "every byte is in" means here. + session.stdout.on("end", () => resolve()); + (session.stdout as Readable & { on(ev: "finish", fn: () => void): unknown }).on( + "finish", + () => resolve(), + ); + session.stdout.on("close", () => resolve()); + session.stdout.on("error", () => resolve()); + }); + const exit = Promise.all([marker, outputEnded]).then(([value]) => value); + return { + output: session.stdout as Readable, + exit, + close() { + detached = true; + session.close(); + outputDone(); + }, + }; + }, + + async deliverStdin( + rt: Pick, + handle: ContainerHandle, + input: { runDir: string; line: string }, + ): Promise { + if (input.line.includes("\n") || input.line.includes("\r")) { + throw new Error("a stdin line must not contain a newline"); + } + const bytes = Buffer.from(`${input.line}\n`, "utf8"); + const session = await rt.exec( + handle, + ["bash", "-c", deliverStdinScript(input.runDir, bytes.length)], + { tty: false }, + ); + const collected = Promise.all([collect(session.stdout), collect(session.stderr)]); + feedStdin(session, bytes); + const [stdout, stderr] = await collected; + if (!stdout.includes(STDIN_DELIVERED_MARKER)) { + const output = [stdout, stderr].filter(Boolean).join("\n"); + throw new RunDeliverError( + `could not append to the run's stdin: ${output.trim().split("\n").slice(-5).join(" | ") || "no output"}`, + output, + ); + } + }, + + async killRun( + rt: Pick, + handle: ContainerHandle, + input: { runDir: string; signal: RunSignal }, + ): Promise { + const session = await rt.exec( + handle, + ["bash", "-c", killRunScript(input.runDir, input.signal)], + { tty: false }, + ); + const stdout = await collect(session.stdout); + return stdout.includes("killed"); + }, +}; diff --git a/packages/container-runtime/src/types.ts b/packages/container-runtime/src/types.ts index 2d3a49df5..50b0fefd0 100644 --- a/packages/container-runtime/src/types.ts +++ b/packages/container-runtime/src/types.ts @@ -1,4 +1,5 @@ import type { ContainerSpec, ContainerHandle, ContainerStatus, ExecSession } from "@optio/shared"; +import type { Readable } from "node:stream"; export interface LogOptions { follow?: boolean; @@ -12,6 +13,75 @@ export interface ExecOptions { rows?: number; } +/** A signal `killRun` may send to a run's process group. */ +export type RunSignal = "TERM" | "KILL" | "INT" | "HUP"; + +export interface RunStartInput { + /** Optio's id for the run (a task, Job run, turn or review run id). */ + runId: string; + /** The run's directory in the container (its run home). */ + runDir: string; + /** + * The start script: the caller's setup, ending with `superviseAgent()` + * (run-protocol.ts), which launches the agent detached and prints + * `__OPTIO_RUN_STARTED__:`. + */ + script: string; + /** + * The run's first stdin lines (for Claude Code, the first stream-json user + * message), written to the run's stdin file before the agent starts. + */ + initialStdin: string; +} + +export interface RunStartResult { + /** The supervisor's pid: a process-group leader in the container. */ + pid: number; + /** What the start script printed before the agent launched (setup lines). */ + output: string; +} + +export interface RunAttachInput { + runId: string; + runDir: string; + /** Stream the output file from this byte offset (0 = from the start). */ + fromByte: number; +} + +/** How an attachment ended. */ +export type RunExit = + /** The agent exited; `exit` holds its code. */ + | { kind: "exited"; code: number } + /** The run's files or supervisor are gone (the pod restarted). */ + | { kind: "lost" } + /** The attachment was closed, or its exec dropped, while the agent may still run. */ + | { kind: "detached"; reason: "closed" | "stream-ended" }; + +export interface RunAttachment { + /** The output file's bytes from `fromByte` on, exactly, as they arrive. */ + output: Readable; + /** + * How the attachment ended. Settles only after `output` has ended, so a + * consumer that awaits it has seen every byte the attachment delivered. + */ + exit: Promise; + /** Detach: sever the exec. The agent keeps running. */ + close(): void; +} + +export interface RunStdinInput { + runId: string; + runDir: string; + /** One line, without its newline. */ + line: string; +} + +export interface RunKillInput { + runId: string; + runDir: string; + signal: RunSignal; +} + export interface ContainerRuntime { create(spec: ContainerSpec): Promise; status(handle: ContainerHandle): Promise; @@ -19,4 +89,16 @@ export interface ContainerRuntime { exec(handle: ContainerHandle, command: string[], opts?: ExecOptions): Promise; destroy(handle: ContainerHandle): Promise; ping(): Promise; + + // The run protocol (run-protocol.ts): an agent run that no API process has + // to stay attached to. + + /** Starts a run: runs the start script, which launches the agent detached. */ + startRun(handle: ContainerHandle, input: RunStartInput): Promise; + /** Attaches to a run's output from a byte offset until the agent exits. */ + attachRun(handle: ContainerHandle, input: RunAttachInput): Promise; + /** Appends a line to the run's stdin. */ + deliverStdin(handle: ContainerHandle, input: RunStdinInput): Promise; + /** Signals the run's process group; true when something was there to signal. */ + killRun(handle: ContainerHandle, input: RunKillInput): Promise; } diff --git a/scripts/repo-init.sh b/scripts/repo-init.sh index 8cfd582e5..a14dd4ec0 100755 --- a/scripts/repo-init.sh +++ b/scripts/repo-init.sh @@ -163,5 +163,11 @@ fi touch /workspace/.ready echo "[optio] Repo pod ready — waiting for tasks" -# Keep the pod alive -exec sleep infinity +# Keep the pod alive under a pid 1 that reaps: a run's supervisor is +# reparented here when the exec that started it ends (run protocol), and a +# bare `sleep` would leave every finished run a zombie. +if command -v tini >/dev/null 2>&1; then + exec tini -- sleep infinity +else + exec sleep infinity +fi