diff --git a/.github/workflows/deploy-try.yml b/.github/workflows/deploy-try.yml index 4eddaa8d..22db4e1d 100644 --- a/.github/workflows/deploy-try.yml +++ b/.github/workflows/deploy-try.yml @@ -78,19 +78,22 @@ jobs: env: CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }} CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} - run: ../../node_modules/.bin/vp exec wrangler deploy --env prod --minify --dry-run + run: ../../node_modules/.bin/vp exec wrangler deploy --env prod --minify --dry-run --experimental-provision=false - name: Dry-run Web Worker working-directory: apps/web env: CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }} CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} - run: ../../node_modules/.bin/vp exec wrangler deploy --env prod --dry-run + run: ../../node_modules/.bin/vp exec wrangler deploy --env prod --dry-run --experimental-provision=false - name: Deploy env: CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }} CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} + CLOUDFLARE_ZONE_ID: ${{ secrets.CLOUDFLARE_ZONE_ID }} + MOSOO_PROTOCOL_V3_SMOKE_AGENT_ID: ${{ vars.MOSOO_PROTOCOL_V3_SMOKE_AGENT_ID }} + MOSOO_PROTOCOL_V3_SMOKE_TOKEN: ${{ secrets.MOSOO_PROTOCOL_V3_SMOKE_TOKEN }} VITE_MOSOO_DEPLOYMENT_MODE: cloud VITE_MOSOO_ENVIRONMENT: production VITE_POSTHOG_API_HOST: https://us.i.posthog.com diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index df8d0df0..4e6b2a83 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -52,7 +52,7 @@ Minimum generated-file rules: Required tools: -- `bun >= 1.4.0-canary.1` with text `bun.lock` support. +- `bun >= 1.4.0` with text `bun.lock` support. - `just >= 1.51` The repository already pins Vite Plus and Git hook tooling dependencies. Human-facing repository operations use `just`; the `justfile` delegates to `bun run`, which resolves the pinned local Vite Plus binary after `bun install`. A global Vite Plus install is optional for direct shell use: `curl -fsSL https://vite.plus | bash`. @@ -164,6 +164,12 @@ just test-file apps/api/tests/session-run-cancel.test.ts mosoo uses generated Drizzle migrations for local and production D1: - The schema source of truth is `pkgs/db/src/schema/**`. +- `just check` generates that schema in memory with `drizzle-kit/api` and + requires it to match the latest checked-in snapshot without writing temp + migration output. +- The production schema test also applies the complete migration chain to an + in-memory SQLite database and requires its managed catalog to match that same + snapshot. - The chain is append-only. Once a migration SQL file, Drizzle snapshot, or journal entry is merged or applied, do not modify, delete, rename, or regenerate it. @@ -191,14 +197,40 @@ D1 state and reapplies the chain. Never delete or regenerate established files under `pkgs/db/drizzle`. Production D1 is not reset during deploy. `just deploy-api` runs -`apps/api/bin/deploy-prod.ts`, whose first remote action is applying pending D1 -migrations. It then verifies that every table in the latest Drizzle snapshot -exists in production — the DEPLOY-D1-001 missing-table guard — ensures the -environment-artifact queue, builds the Driver, and deploys the API Worker. -The guard does not compare columns, indexes, constraints, or extra live tables. +`apps/api/bin/deploy-prod.ts`, which builds the Driver and dry-runs the API +Worker before its first remote mutation. It then acquires the durable D1 deploy +lease before it ensures production queues or applies ordinary pending +migrations. Every D1, Queue, API Worker, or smoke mutation renews and verifies +that exact owner before and after the mutation through an independent final D1 +read; the canonical lease table may have no triggers. The protocol v3 breaking migrations +use the one-shot paused-queue, D1 admission-gate, drain, legacy terminal +integrity preflight, Time Travel bookmark, deploy, live Driver handshake smoke, +durable queue-resume phase, and Queue API readback sequence documented in the +production runbook. +Migration `0013` is allowed to proceed only when its read-only inventory reports zero historical loss or conflict candidates, and the migration repeats that guard before its first rewrite. +Run `bun apps/api/bin/deploy-prod.ts --protocol-v3-lossy-migration-inventory` before scheduling a release with `0013` pending. +The legacy preflight requires a provably collision-free terminal identity and +fully committed Run, Session, assistant, cursor, and permission state; migration +`0014` normalizes only that proven identity and never guesses ambiguous history. +The migration itself repeats those fail-closed checks and requires a short-lived +authorization bound to the bookmark, exact candidate manifest, and current +deploy lease owner, so applying `0014` directly cannot bypass the cutover. +Run `bun apps/api/bin/deploy-prod.ts --protocol-v3-legacy-inventory` as the +read-only human gate before scheduling this one-shot production cutover. +After the breaking migration is recorded, deployment recovery is roll-forward +only with the exact same v3 release. +The bookmark is an emergency destructive backup, not a routine rollback: a D1 +Time Travel restore loses every later D1 write and does not restore Durable +Object storage. +After migration, the script compares every managed table's columns, primary key, +defaults, nullability, named indexes and partial predicates, foreign keys, CHECK +expressions, autoincrement state, and exact migration-owned trigger definitions +with the latest migration contract before it deploys the API Worker. +Unknown extra application tables fail closed. The eight retained Channels and +WeChat tables are explicitly unmanaged legacy data because #577 removed their +runtime subsystem without an approved destructive data migration. No other automated migration gate exists today: there is no trusted-Git-range -check, no clean-worktree check, and no dry-run inside the deploy script. Run -`just check` and the +or clean-worktree check. Run `just check` and the [Production Deploy Verification](./docs/production-deploy-verification.md) runbook before deploying. @@ -407,7 +439,7 @@ just deploy-web # Web only — runs no gate; does not touch D1 Only `just deploy` runs the repository gate before publishing; `just deploy-api` and `just deploy-web` publish directly. The API deploy applies pending remote D1 -migrations as its first remote action (see +migrations only after its local Driver build and API dry-run (see [Database And Migrations](#database-and-migrations)). API production config lives in `apps/api/wrangler.toml`; web production config lives in `apps/web/wrangler.toml`. Cloudflare routes send `cloud.mosoo.ai/api/*` to the API Worker and `cloud.mosoo.ai/*` to the console Web Worker. The legacy console host redirects Web traffic to `cloud.mosoo.ai` but keeps `try.mosoo.ai/api/*` as a direct compatibility route. The public landing page and blog on `mosoo.ai/*` are owned by `langgenius/mosoo-website`. diff --git a/apps/api/bin/d1-json.ts b/apps/api/bin/d1-json.ts new file mode 100644 index 00000000..fd76049b --- /dev/null +++ b/apps/api/bin/d1-json.ts @@ -0,0 +1,74 @@ +export type D1JsonRow = Readonly>; + +function parseJsonArrayAt(raw: string, start: number): unknown { + let depth = 0; + let escaped = false; + let inString = false; + for (let index = start; index < raw.length; index += 1) { + const character = raw[index]; + if (inString) { + if (escaped) escaped = false; + else if (character === "\\") escaped = true; + else if (character === '"') inString = false; + continue; + } + if (character === '"') inString = true; + else if (character === "[") depth += 1; + else if (character === "]") { + depth -= 1; + if (depth === 0) { + try { + return JSON.parse(raw.slice(start, index + 1)); + } catch { + return null; + } + } + } + } + return null; +} + +export function parseD1JsonResults(raw: string): D1JsonRow[][] { + for (let start = raw.indexOf("["); start !== -1; start = raw.indexOf("[", start + 1)) { + const value = parseJsonArrayAt(raw, start); + if (!Array.isArray(value) || value.length === 0) continue; + + const statements: D1JsonRow[][] = []; + let valid = true; + for (const result of value) { + if ( + typeof result !== "object" || + result === null || + !("results" in result) || + !("success" in result) || + result.success !== true || + !Array.isArray(result.results) + ) { + valid = false; + break; + } + const rows: D1JsonRow[] = []; + for (const row of result.results) { + if (typeof row !== "object" || row === null || Array.isArray(row)) { + valid = false; + break; + } + rows.push(row as D1JsonRow); + } + if (!valid) break; + statements.push(rows); + } + if (valid) return statements; + } + throw new Error("D1 output has no successful JSON statement results."); +} + +export function requireSingleD1Row(raw: string): D1JsonRow { + const statements = parseD1JsonResults(raw); + if (statements.length !== 1) { + throw new Error("D1 JSON output must contain exactly one statement result."); + } + const rows = statements[0]; + if (rows?.length !== 1) throw new Error("D1 JSON output must contain exactly one row."); + return rows[0]; +} diff --git a/apps/api/bin/deploy-prod.ts b/apps/api/bin/deploy-prod.ts index a9014198..6503696f 100755 --- a/apps/api/bin/deploy-prod.ts +++ b/apps/api/bin/deploy-prod.ts @@ -1,11 +1,143 @@ #!/usr/bin/env bun -import type { BunRuntime } from "../../../config/bun-script-types"; +import { DRIVER_PROTOCOL_VERSION } from "@mosoo/agent-driver/boot"; +import { RUNTIME_EVENT_SCHEMA_VERSION } from "@mosoo/agent-driver/events"; +import Cloudflare from "cloudflare"; + +import type { BunRuntime, BunSpawnSyncResult } from "../../../config/bun-script-types"; +import { APP_DEPLOYMENT_PROBE_SCRIPT_NAME } from "../src/modules/apps/application/app-deployment-gateway"; +import { parseD1JsonResults } from "./d1-json"; +import { + acquireProdDeployLeaseStatements, + assertProdDeployLeaseOwned, + assertProdDeployLeaseReleased, + releaseProdDeployLeaseStatements, + verifyProdDeployLeaseStatements, +} from "./prod-deploy-lease"; +import type { ProdSchemaCatalog } from "./prod-schema-guard"; import { - extractTableNames, - findMissingProdTables, - getLatestSnapshotFilename, - parseExpectedTableNames, + assertProdSchemaMatches, + createProdSchemaCatalogFromIntrospectionRows, + createProdSchemaIntrospectionStatements, + parseGeneratedProdSchemaCatalog, } from "./prod-schema-guard"; +import { + apiWorkerDeployArgs, + apiWorkerDryRunArgs, + assertProdWfpNamespaceExact, + assertProdWfpReadOnlyInfrastructure, + assertProdWfpWorkflowBoundary, + assertProdWfpWorkflowDeployGate, + probeProdWfpWebSocket, + PROD_API_WORKER_NAME, + PROD_APP_DEPLOYMENT_DOMAIN, + PROD_APP_DEPLOYMENT_WORKFLOW, + PROD_APP_DEPLOYMENT_WORKFLOW_BINDING, + PROD_APP_DEPLOYMENT_WORKFLOW_CLASS, + PROD_APP_DISPATCH_NAMESPACE, + PROD_APP_WILDCARD_DNS, + PROD_APP_WILDCARD_ROUTE, + PROD_WFP_WRITE_CANARY_SCRIPT, + shouldProbeProdWfpBeforeMutation, + verifyProdWfpGatewayProbe, + verifyProdWfpWriteCanary, +} from "./prod-wfp-preflight"; +import type { + ProdWfpInfrastructureConfig, + ProdWfpReadOnlyInventory, + ProdWfpWorkflowMode, + WfpNamespaceInventory, + WfpScriptInventory, + WfpWorkerBindingInventory, + WfpWorkflowInventory, + WfpWriteCanaryClient, +} from "./prod-wfp-preflight"; +import type { + ProtocolV3ContainerApplication, + ProtocolV3ContainerInstance, + ProtocolV3CutoverState, + ProtocolV3LegacyTerminalSourceInventory, + ProtocolV3LossyMigrationInventory, +} from "./protocol-v3-cutover"; +import { + ACCEPT_PROTOCOL_V3_QUEUE_RESUME_SQL, + assertCutoverMigrationJournalAudited, + assertProtocolV3SmokeAgent, + assertProtocolV3Release, + assertProtocolV3WorkerVersion, + assertProtocolV3LegacyTerminalIntegrity, + assertProtocolV3LegacyTerminalSourceInventory, + assertProtocolV3LossyMigrationInventory, + assertProtocolV3RuntimeAuthorityPreflight, + authorizeProtocolV3LegacyRewriteSql, + beginProtocolV3MigrationSql, + CLOSE_PROTOCOL_V3_SMOKE_WINDOW_SQL, + collectProtocolV3ContainerApplications, + collectProtocolV3ContainerInstances, + completeProtocolV3QueueResume, + ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL, + ENTER_PROTOCOL_V3_DRAIN_SQL, + ENTER_PROTOCOL_V3_QUEUES_RESUMING_SQL, + findPendingProdMigrations, + installProtocolV3PostMigrationCutoverSql, + installProtocolV3CutoverSql, + isProtocolV3ContainerRolloutConverged, + isProtocolV3CutoverDrained, + isProtocolV3RuntimeDrained, + isProtocolV3SmokeReady, + openProtocolV3SmokeWindowSql, + parseProtocolV3CommandFreeze, + parseProtocolV3ContainerManifestDigest, + parseProtocolV3CutoverDrain, + parseProtocolV3CutoverObjects, + parseProtocolV3CutoverProbe, + parseProtocolV3CutoverState, + parseProtocolV3LegacyTerminalIntegrity, + parseProtocolV3LegacyRewriteAuthorization, + parseProtocolV3LegacyTerminalSourceInventory, + parseProtocolV3LossyMigrationInventory, + parseProtocolV3RuntimeAuthorityPreflight, + parseProtocolV3SmokeStatus, + parseCleanGitTreeOid, + parseProtocolV3WorkerDeployment, + parseStoredProtocolV3SmokeRequestKey, + parseStoredProtocolV3SmokeSession, + parseStoredProtocolV3CutoverBookmark, + parseTimeTravelBookmark, + PROD_APPLIED_MIGRATIONS_SQL, + PROTOCOL_V3_CUTOVER_BOOKMARK_SQL, + PROTOCOL_V3_CUTOVER_DRAIN_SQL, + PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + PROTOCOL_V3_CUTOVER_OBJECTS_SQL, + PROTOCOL_V3_CUTOVER_PROBE_SQL, + PROTOCOL_V3_COMMAND_FREEZE_SQL, + PROTOCOL_V3_LEGACY_TERMINAL_INTEGRITY_SQL, + PROTOCOL_V3_LEGACY_TERMINAL_SOURCE_INVENTORY_SQL, + PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_SQL, + PROTOCOL_V3_LOSSY_MIGRATION_INVENTORY_SQL, + PROTOCOL_V3_MIGRATION, + PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + PROTOCOL_V3_POST_MIGRATION_CUTOVER_DRAIN_SQL, + PROTOCOL_V3_POST_MIGRATION_UNSAFE_SANDBOXES_SQL, + PROTOCOL_V3_CUTOVER_QUEUE_NAMES, + PROTOCOL_V3_RUNTIME_AUTHORITY_MIGRATION, + PROTOCOL_V3_SESSION_CLEANUP_MIGRATION, + PROTOCOL_V3_SESSION_EVENT_MIGRATION, + recoverProtocolV3CutoverFailure, + PROTOCOL_V3_SMOKE_REQUEST_KEY_SQL, + PROTOCOL_V3_SMOKE_SESSION_SQL, + PROTOCOL_V3_UNSAFE_SANDBOXES_SQL, + protocolV3SmokeAgentSql, + protocolV3SmokeStatusSql, + protocolV3RuntimeAuthorityPreflightSql, + protocolV3ReleaseTag, + protocolV3ContainerImageTag, + REMOVE_PROTOCOL_V3_CUTOVER_SQL, + storeProtocolV3CutoverBookmarkSql, + storeProtocolV3RolloutSql, + storeProtocolV3SmokeRequestKeySql, + storeProtocolV3SmokeSessionSql, + updateAndVerifyProtocolV3QueueDelivery, +} from "./protocol-v3-cutover"; declare const Bun: BunRuntime; @@ -14,6 +146,28 @@ const apiDir = `${scriptDir}/..`; const repoRoot = `${apiDir}/../..`; const D1_BINDING = "DB"; const ENV = "prod"; +const EXPECTED_DRIVER_PROTOCOL_VERSION = 3; +const EXPECTED_RUNTIME_EVENT_SCHEMA_VERSION = "2026-08-29"; +const PROD_CONTAINER_APPLICATION = "mosoo-api-prod-sandbox-prod"; +const PROD_PUBLIC_API_URL = "https://cloud.mosoo.ai/api/v1"; +const PROD_GRAPHQL_URL = "https://cloud.mosoo.ai/api/graphql"; +const PROD_HEALTH_URL = "https://cloud.mosoo.ai/api/health?deep=1"; +const CUTOVER_DRAIN_TIMEOUT_MS = 15 * 60 * 1_000; +const CUTOVER_ROLLOUT_TIMEOUT_MS = 10 * 60 * 1_000; +const CUTOVER_SMOKE_TIMEOUT_MS = 5 * 60 * 1_000; +const REMOTE_MUTATION_TIMEOUT_MS = 10 * 60 * 1_000; +const POLL_INTERVAL_MS = 5_000; +const PLATFORM_ID_PATTERN = /^[0-7][0-9A-HJKMNP-TV-Z]{25}$/iu; +const CLOUDFLARE_ACCOUNT_ID_PATTERN = /^[0-9a-f]{32}$/iu; +const CLOUDFLARE_ZONE_ID_PATTERN = /^[0-9a-f]{32}$/iu; +const PROD_ZONE_NAME = "mosoo.ai"; +const PROD_APP_DISPATCH_BINDING = "APP_DEPLOYMENT_DISPATCHER"; +const APP_DEPLOYMENT_SCRIPT_LEDGER_PROBE_SQL = `SELECT EXISTS ( + SELECT 1 + FROM sqlite_schema + WHERE type = 'table' AND name = 'app_deployment_script' +) AS ledger_present`; +const LEGACY_APP_RESOURCE_NAME_PATTERN = /^app-[0-7][0-9a-hjkmnp-tv-z]{25}$/u; function writeStdout(message: string): void { process.stdout.write(`${message}\n`); @@ -30,7 +184,7 @@ function run(args: string[], cwd = apiDir): void { stdout: "inherit", }); if (result.exitCode !== 0) { - process.exit(result.exitCode); + throw new Error(`wrangler ${args.join(" ")} exited with ${result.exitCode}.`); } } @@ -42,72 +196,292 @@ function runVp(args: string[], cwd = repoRoot): void { stdout: "inherit", }); if (result.exitCode !== 0) { - process.exit(result.exitCode); + throw new Error(`vp ${args.join(" ")} exited with ${result.exitCode}.`); } } -function applyD1Migrations(): void { - run(["d1", "migrations", "apply", D1_BINDING, "--remote", "--env", ENV]); +function captureGit(args: string[]): string { + const result = Bun.spawnSync(["git", ...args], { cwd: repoRoot }); + if (result.exitCode !== 0) { + throw new Error( + `git ${args.join(" ")} exited with ${result.exitCode}: ${result.stderr.toString("utf8")}`, + ); + } + return result.stdout.toString("utf8"); } -const MIGRATION_META_DIR = `${repoRoot}/pkgs/db/drizzle/meta`; +function readLocalMigrationNames(): string[] { + return captureGit(["ls-files", "--", "pkgs/db/drizzle/*.sql"]) + .trim() + .split("\n") + .filter(Boolean) + .map((path) => path.slice("pkgs/db/drizzle/".length)); +} -/** - * Refuse to deploy a Worker whose schema references tables missing from prod. - * Runs AFTER `applyD1Migrations`, so on the correct append-only path every - * table is already present (DEPLOY-D1-001). - */ -async function loadExpectedMigrationTables(): Promise { - const journal = await Bun.file(`${MIGRATION_META_DIR}/_journal.json`).text(); - const snapshotFilename = getLatestSnapshotFilename(journal); - const snapshot = await Bun.file(`${MIGRATION_META_DIR}/${snapshotFilename}`).text(); +function readCleanReleaseTreeOid(): string { + return parseCleanGitTreeOid( + captureGit(["rev-parse", "HEAD^{tree}"]), + captureGit(["status", "--porcelain=v1", "--untracked-files=all", "--ignore-submodules=none"]), + `${captureGit(["ls-files", "-v"])}${captureGit([ + "submodule", + "foreach", + "--quiet", + "--recursive", + "git ls-files -v", + ])}`, + `${captureGit([ + "ls-files", + "--others", + "--ignored", + "--exclude-standard", + "--", + ".env", + ".env.*", + ".dev.vars", + ".dev.vars.*", + "apps/api/.env", + "apps/api/.env.*", + "apps/api/.dev.vars", + "apps/api/.dev.vars.*", + "apps/web/.env", + "apps/web/.env.*", + ])}${captureGit([ + "submodule", + "foreach", + "--quiet", + "--recursive", + "git ls-files --others --ignored --exclude-standard -- .env .env.* .dev.vars .dev.vars.*", + ])}`, + Object.keys(process.env), + ); +} - return parseExpectedTableNames(snapshot); +function assertReleaseTreeUnchanged(releaseTreeOid: string): void { + if (readCleanReleaseTreeOid() !== releaseTreeOid) { + throw new Error("Production release tree changed during deployment."); + } } -async function assertProdSchemaMatchesMigrations(expectedTables: readonly string[]): Promise { - const result = Bun.spawnSync( - [ - wranglerBin, - "d1", - "execute", - D1_BINDING, - "--remote", - "--env", - ENV, - "--json", - "--command", - "SELECT name FROM sqlite_master WHERE type='table'", - ], - { cwd: apiDir }, - ); +function captureWrangler(args: string[], timeout?: number): string { + const result = Bun.spawnSync([wranglerBin, ...args], { + cwd: apiDir, + ...(timeout === undefined ? {} : { timeout }), + }); const stdout = result.stdout.toString("utf8"); const stderr = result.stderr.toString("utf8"); if (result.exitCode !== 0) { throw new Error( - `wrangler d1 execute (schema check) exited with ${result.exitCode}\nstderr: ${stderr}\nstdout: ${stdout}`, + `wrangler ${args.join(" ")} exited with ${result.exitCode}\nstderr: ${stderr}\nstdout: ${stdout}`, ); } - const missingTables = findMissingProdTables(expectedTables, extractTableNames(stdout)); + return stdout; +} - if (missingTables.length > 0) { +function readTaggedContainerImageDigest( + application: ProtocolV3ContainerApplication, + workerVersionId: string, +): string { + const docker = process.env["WRANGLER_DOCKER_BIN"]?.trim() || "docker"; + const tag = protocolV3ContainerImageTag(application.imageRepository, workerVersionId); + const result = Bun.spawnSync([docker, "manifest", "inspect", "-v", tag], { + cwd: apiDir, + timeout: 60_000, + }); + if (result.exitCode !== 0) { throw new Error( + `Container manifest inspection failed for ${tag}: ${result.stderr.toString("utf8")}`, + ); + } + return parseProtocolV3ContainerManifestDigest(result.stdout.toString("utf8")); +} + +function executeRawProdD1Json(sql: string, timeout?: number): string { + return captureWrangler( + ["d1", "execute", D1_BINDING, "--remote", "--env", ENV, "--json", "--command", sql], + timeout, + ); +} + +let prodDeployLeaseOwner: string | null = null; + +function executeProdDeployLease(statements: readonly string[]): string { + return executeRawProdD1Json(`${statements.join(";\n")};`, REMOTE_MUTATION_TIMEOUT_MS); +} + +function executeIdempotentProdDeployLeaseWithRetry(statements: readonly string[]): string { + try { + return executeProdDeployLease(statements); + } catch { + return executeProdDeployLease(statements); + } +} + +function executeOwnedProdDeployLease(statements: readonly string[], owner: string): void { + assertProdDeployLeaseOwned(executeIdempotentProdDeployLeaseWithRetry(statements), owner); +} + +export function acquireProdDeployLease( + owner: string, + execute: (statements: readonly string[]) => string, +): void { + assertProdDeployLeaseOwned(execute(acquireProdDeployLeaseStatements(owner)), owner); + prodDeployLeaseOwner = owner; +} + +function verifyProdDeployLease(): void { + if (prodDeployLeaseOwner === null) { + throw new Error("Production remote mutation requires an acquired deploy lease."); + } + executeOwnedProdDeployLease(verifyProdDeployLeaseStatements(), prodDeployLeaseOwner); +} + +function releaseProdDeployLease(): void { + if (prodDeployLeaseOwner === null) return; + const owner = prodDeployLeaseOwner; + try { + assertProdDeployLeaseReleased( + executeIdempotentProdDeployLeaseWithRetry(releaseProdDeployLeaseStatements(owner)), + ); + } finally { + prodDeployLeaseOwner = null; + } +} + +function executeProdMutation(args: string[], inheritOutput = false): BunSpawnSyncResult { + verifyProdDeployLease(); + let result: BunSpawnSyncResult; + try { + result = Bun.spawnSync([wranglerBin, ...args], { + cwd: apiDir, + ...(inheritOutput + ? { stderr: "inherit" as const, stdin: "inherit" as const, stdout: "inherit" as const } + : {}), + timeout: REMOTE_MUTATION_TIMEOUT_MS, + }); + } catch (mutationError) { + try { + verifyProdDeployLease(); + } catch (verificationError) { + throw new AggregateError( + [mutationError, verificationError], + "Production mutation and deploy lease verification both failed.", + { cause: verificationError }, + ); + } + throw mutationError; + } + let verificationError: unknown = null; + try { + verifyProdDeployLease(); + } catch (error) { + verificationError = error; + } + if (verificationError !== null) { + throw new AggregateError( + [ + new Error( + `wrangler ${args.join(" ")} completed with ${result.exitCode} before deploy lease verification failed.`, + ), + verificationError, + ], + "Production mutation ownership could not be verified.", + ); + } + return result; +} + +function runProdMutation(args: string[]): void { + const result = executeProdMutation(args, true); + if (result.exitCode !== 0) { + throw new Error(`wrangler ${args.join(" ")} exited with ${result.exitCode}.`); + } +} + +async function runOwnedProdAsyncMutation(mutation: () => Promise): Promise { + verifyProdDeployLease(); + let result: T; + try { + result = await mutation(); + } catch (mutationError) { + try { + verifyProdDeployLease(); + } catch (verificationError) { + throw new AggregateError( + [mutationError, verificationError], + "Production mutation and deploy lease verification both failed.", + { cause: verificationError }, + ); + } + throw mutationError; + } + try { + verifyProdDeployLease(); + } catch (verificationError) { + throw new AggregateError( [ - "✗ Prod D1 schema drift: the latest migration snapshot defines tables absent from prod.", - ` Missing: ${missingTables.join(", ")}`, - " Cause: a migration expected to create a table was skipped, rewritten,", - " or incomplete (DEPLOY-D1-001). Fix the history with a new reviewed", - " migration; never rewrite an applied migration, then re-run the deploy.", - ].join("\n"), + new Error("Cloudflare mutation completed before deploy lease verification failed."), + verificationError, + ], + "Production mutation ownership could not be verified.", + { cause: verificationError }, ); } + return result; +} + +function executeProdD1Json(sql: string): string { + return executeRawProdD1Json(sql); +} - writeStdout(` prod schema OK (${expectedTables.length} migration tables present)`); +function executeProdD1(sql: string): void { + runProdMutation([ + "d1", + "execute", + D1_BINDING, + "--remote", + "--env", + ENV, + "--yes", + "--command", + sql, + ]); } -const REQUIRED_PROD_QUEUES: readonly string[] = ["environment-artifact-build"]; +function sleep(milliseconds: number): Promise { + return new Promise((resolve) => setTimeout(resolve, milliseconds)); +} + +function applyD1Migrations(): void { + runProdMutation(["d1", "migrations", "apply", D1_BINDING, "--remote", "--env", ENV]); +} + +function loadExpectedProdSchema(): ProdSchemaCatalog { + const result = Bun.spawnSync(["bun", "pkgs/db/scripts/check-schema.ts", "--catalog"], { + cwd: repoRoot, + }); + const stdout = result.stdout.toString("utf8"); + const stderr = result.stderr.toString("utf8"); + + if (result.exitCode !== 0) { + throw new Error( + `Drizzle schema freshness check exited with ${result.exitCode}\nstderr: ${stderr}\nstdout: ${stdout}`, + ); + } + return parseGeneratedProdSchemaCatalog(stdout); +} + +function assertProdSchemaMatchesSnapshot(expected: ProdSchemaCatalog): void { + const tableNames = expected.tables.map(({ name }) => name); + const statements = createProdSchemaIntrospectionStatements(tableNames); + const live = createProdSchemaCatalogFromIntrospectionRows( + parseD1JsonResults(executeProdD1Json(statements.join(";\n"))), + tableNames, + ); + assertProdSchemaMatches(expected, live); + writeStdout(` prod schema OK (${expected.tables.length} tables fully match)`); +} function listProdQueues(): string[] { const result = Bun.spawnSync([wranglerBin, "queues", "list"], { cwd: apiDir }); @@ -140,9 +514,7 @@ function ensureQueueExists(queueName: string, existingQueues: readonly string[]) return; } - const result = Bun.spawnSync([wranglerBin, "queues", "create", queueName], { - cwd: apiDir, - }); + const result = executeProdMutation(["queues", "create", queueName]); const stdout = result.stdout.toString("utf8"); const stderr = result.stderr.toString("utf8"); @@ -164,32 +536,1556 @@ function ensureQueueExists(queueName: string, existingQueues: readonly string[]) function ensureRequiredProdQueues(): void { const listing = listProdQueues(); - for (const queueName of REQUIRED_PROD_QUEUES) { + for (const queueName of PROTOCOL_V3_CUTOVER_QUEUE_NAMES) { ensureQueueExists(queueName, listing); } } -const expectedMigrationTables = await loadExpectedMigrationTables().catch((error: unknown) => { - writeStdout(error instanceof Error ? error.message : String(error)); - process.exit(1); -}); +function assertProtocolV3Artifacts(): void { + const protocolVersion: number = DRIVER_PROTOCOL_VERSION; + const eventSchemaVersion: string = RUNTIME_EVENT_SCHEMA_VERSION; + + if ( + protocolVersion !== EXPECTED_DRIVER_PROTOCOL_VERSION || + eventSchemaVersion !== EXPECTED_RUNTIME_EVENT_SCHEMA_VERSION + ) { + throw new Error( + `Production requires Driver protocol ${EXPECTED_DRIVER_PROTOCOL_VERSION} and event schema ${EXPECTED_RUNTIME_EVENT_SCHEMA_VERSION}; got ${protocolVersion} and ${eventSchemaVersion}.`, + ); + } +} + +function runLocalPreflight(releaseTreeOid: string): void { + assertProtocolV3Artifacts(); + + writeStdout("▶ Building Driver before any production mutation"); + runVp(["run", "--filter", "agent-driver", "build"]); + + writeStdout("▶ Dry-running the API Worker before any production mutation"); + run(apiWorkerDryRunArgs()); + + assertReleaseTreeUnchanged(releaseTreeOid); +} + +interface ProtocolV3SmokeConfig { + readonly agentId: string; + readonly token: string; +} + +interface ProdQueueApiConfig { + readonly accountId: string; + readonly apiToken: string; + readonly zoneId: string; +} + +function readProdQueueApiConfig(): ProdQueueApiConfig { + const accountId = process.env["CLOUDFLARE_ACCOUNT_ID"]?.trim() ?? ""; + const apiToken = process.env["CLOUDFLARE_API_TOKEN"]?.trim() ?? ""; + const zoneId = process.env["CLOUDFLARE_ZONE_ID"]?.trim() ?? ""; + + if ( + !CLOUDFLARE_ACCOUNT_ID_PATTERN.test(accountId) || + !CLOUDFLARE_ZONE_ID_PATTERN.test(zoneId) || + !apiToken + ) { + throw new Error( + "Production deploy requires exact CLOUDFLARE_ACCOUNT_ID, CLOUDFLARE_ZONE_ID, and CLOUDFLARE_API_TOKEN values.", + ); + } + + return { accountId, apiToken, zoneId }; +} + +function prodWfpInfrastructureConfig( + config: ProdQueueApiConfig, + requireProbeOnlyNamespace: boolean, + workflowRepairAllowed = false, +): ProdWfpInfrastructureConfig { + return { + accountId: config.accountId, + apiWorkerName: PROD_API_WORKER_NAME, + dispatchBinding: PROD_APP_DISPATCH_BINDING, + dispatchNamespace: PROD_APP_DISPATCH_NAMESPACE, + requireProbeOnlyNamespace, + wildcardDns: PROD_APP_WILDCARD_DNS, + wildcardRoute: PROD_APP_WILDCARD_ROUTE, + workflowBinding: PROD_APP_DEPLOYMENT_WORKFLOW_BINDING, + workflowClass: PROD_APP_DEPLOYMENT_WORKFLOW_CLASS, + workflowName: PROD_APP_DEPLOYMENT_WORKFLOW, + workflowRepairAllowed, + zoneId: config.zoneId, + zoneName: PROD_ZONE_NAME, + }; +} + +function readAppDeploymentScriptLedgerPresent(): boolean { + const row = parseD1JsonResults(executeProdD1Json(APP_DEPLOYMENT_SCRIPT_LEDGER_PROBE_SQL))[0]?.[0]; + const value = row?.["ledger_present"]; + if (value !== 0 && value !== 1) { + throw new Error("Production D1 did not return an exact App deployment script ledger probe."); + } + return value === 1; +} + +async function collectAsync(items: AsyncIterable): Promise { + const values: T[] = []; + for await (const item of items) values.push(item); + return values; +} + +function normalizeProdWfpNamespace( + namespace: Awaited< + ReturnType + >, +): WfpNamespaceInventory { + return { + id: namespace.namespace_id ?? null, + name: namespace.namespace_name ?? null, + scriptCount: namespace.script_count ?? null, + trustedWorkers: namespace.trusted_workers ?? null, + }; +} + +async function readProdWfpNamespace( + client: Cloudflare, + config: ProdQueueApiConfig, +): Promise { + const namespace = await client.workersForPlatforms.dispatch.namespaces.get( + PROD_APP_DISPATCH_NAMESPACE, + { account_id: config.accountId }, + ); + return normalizeProdWfpNamespace(namespace); +} + +async function readProdWfpProbeScript( + client: Cloudflare, + config: ProdQueueApiConfig, +): Promise { + try { + const script = await client.workersForPlatforms.dispatch.namespaces.scripts.get( + APP_DEPLOYMENT_PROBE_SCRIPT_NAME, + { + account_id: config.accountId, + dispatch_namespace: PROD_APP_DISPATCH_NAMESPACE, + }, + ); + return { + id: script.script?.id ?? null, + namespace: script.dispatch_namespace ?? null, + }; + } catch (error) { + if (!isCloudflareStatus(error, 404)) throw error; + return { id: null, namespace: null }; + } +} + +async function readProdAppDeploymentWorkflow( + client: Cloudflare, + config: ProdQueueApiConfig, +): Promise { + try { + const workflow = await client.workflows.get(PROD_APP_DEPLOYMENT_WORKFLOW, { + account_id: config.accountId, + }); + return { + className: workflow.class_name, + id: workflow.id, + name: workflow.name, + scriptName: workflow.script_name, + }; + } catch (error) { + if (!isCloudflareStatus(error, 404)) throw error; + return { className: null, id: null, name: null, scriptName: null }; + } +} + +type ProdWorkerSettings = Awaited< + ReturnType +>; + +function normalizeProdWorkerBindings( + bindings: ProdWorkerSettings["bindings"], +): WfpWorkerBindingInventory[] { + return (bindings ?? []).map((binding) => ({ + className: binding.type === "workflow" ? (binding.class_name ?? null) : null, + name: binding.name ?? null, + namespace: binding.type === "dispatch_namespace" ? binding.namespace : null, + scriptName: binding.type === "workflow" ? (binding.script_name ?? PROD_API_WORKER_NAME) : null, + type: binding.type ?? null, + workflowName: binding.type === "workflow" ? binding.workflow_name : null, + })); +} + +async function readProdAppDeploymentWorkflowBoundary( + client: Cloudflare, + config: ProdQueueApiConfig, +): Promise<{ + bindings: WfpWorkerBindingInventory[]; + workflow: WfpWorkflowInventory; +}> { + const [workflow, settings] = await Promise.all([ + readProdAppDeploymentWorkflow(client, config), + client.workers.scripts.scriptAndVersionSettings.get(PROD_API_WORKER_NAME, { + account_id: config.accountId, + }), + ]); + return { bindings: normalizeProdWorkerBindings(settings.bindings), workflow }; +} + +function isLegacyAppName(value: string | undefined): value is string { + return value !== undefined && LEGACY_APP_RESOURCE_NAME_PATTERN.test(value); +} + +function isLegacyAppHostname(value: string | undefined): value is string { + if (value === undefined) return false; + const suffix = `.${PROD_APP_DEPLOYMENT_DOMAIN}`; + return value.endsWith(suffix) && isLegacyAppName(value.slice(0, -suffix.length)); +} + +async function readProdWfpReadOnlyInventory( + client: Cloudflare, + config: ProdQueueApiConfig, +): Promise { + const [ + namespace, + probeScript, + workflow, + zone, + dnsRecords, + certificatePacks, + routes, + workerSettings, + pagesProjects, + workerScripts, + workerDomains, + ] = await Promise.all([ + readProdWfpNamespace(client, config), + readProdWfpProbeScript(client, config), + readProdAppDeploymentWorkflow(client, config), + client.zones.get({ zone_id: config.zoneId }), + collectAsync(client.dns.records.list({ zone_id: config.zoneId })), + collectAsync(client.ssl.certificatePacks.list({ zone_id: config.zoneId })), + collectAsync(client.workers.routes.list({ zone_id: config.zoneId })), + client.workers.scripts.scriptAndVersionSettings.get(PROD_API_WORKER_NAME, { + account_id: config.accountId, + }), + collectAsync(client.pages.projects.list({ account_id: config.accountId })), + collectAsync(client.workers.scripts.list({ account_id: config.accountId })), + collectAsync(client.workers.domains.list({ account_id: config.accountId })), + ]); + + const legacyResources: string[] = []; + for (const project of pagesProjects) { + if (isLegacyAppName(project.name)) { + legacyResources.push( + `Pages project ${project.name}: wrangler pages project delete ${project.name}`, + ); + } + for await (const domain of client.pages.projects.domains.list(project.name, { + account_id: config.accountId, + })) { + if (isLegacyAppHostname(domain.name)) { + legacyResources.push( + `Pages domain ${project.name}/${domain.name}: DELETE /accounts/$CLOUDFLARE_ACCOUNT_ID/pages/projects/${project.name}/domains/${domain.name}`, + ); + } + } + } + for (const script of workerScripts) { + if (isLegacyAppName(script.id)) { + legacyResources.push( + `classic Worker script ${script.id}: wrangler delete --name ${script.id}`, + ); + } + } + for (const domain of workerDomains) { + if (isLegacyAppHostname(domain.hostname) || isLegacyAppName(domain.service)) { + legacyResources.push( + `classic Worker domain ${domain.id} ${domain.hostname}: DELETE /accounts/$CLOUDFLARE_ACCOUNT_ID/workers/domains/${domain.id}`, + ); + } + } + for (const route of routes) { + const hostname = route.pattern.endsWith("/*") ? route.pattern.slice(0, -2) : undefined; + if (isLegacyAppHostname(hostname) || isLegacyAppName(route.script)) { + legacyResources.push( + `classic Worker route ${route.id} ${route.pattern}: DELETE /zones/$CLOUDFLARE_ZONE_ID/workers/routes/${route.id}`, + ); + } + } + + return { + apiWorkerBindings: normalizeProdWorkerBindings(workerSettings.bindings), + certificates: certificatePacks.map((pack) => ({ + certificates: pack.certificates.map((certificate) => ({ + expiresOn: certificate.expires_on ?? null, + hosts: certificate.hosts, + status: certificate.status, + })), + hosts: pack.hosts, + status: pack.status, + })), + dnsRecords: dnsRecords.map((record) => ({ + name: record.name ?? null, + proxied: record.proxied ?? null, + type: record.type ?? null, + })), + legacyResources: legacyResources.toSorted(), + namespace, + probeScript, + routes: routes.map((route) => ({ + pattern: route.pattern ?? null, + script: route.script ?? null, + })), + workflow, + zone: { + accountId: zone.account.id ?? null, + id: zone.id ?? null, + name: zone.name ?? null, + status: zone.status ?? null, + }, + }; +} + +function isCloudflareStatus(error: unknown, status: number): boolean { + return ( + typeof error === "object" && + error !== null && + "status" in error && + Reflect.get(error, "status") === status + ); +} + +function createProdWfpWriteCanaryClient( + client: Cloudflare, + config: ProdQueueApiConfig, +): WfpWriteCanaryClient { + const params = { + account_id: config.accountId, + dispatch_namespace: PROD_APP_DISPATCH_NAMESPACE, + } as const; + + return { + async deleteScript(scriptName) { + await runOwnedProdAsyncMutation(async () => { + try { + await client.workersForPlatforms.dispatch.namespaces.scripts.delete(scriptName, params); + } catch (error) { + if (!isCloudflareStatus(error, 404)) throw error; + } + }); + }, + async readScript(scriptName) { + try { + const script = await client.workersForPlatforms.dispatch.namespaces.scripts.get( + scriptName, + params, + ); + return { + id: script.script?.id ?? null, + namespace: script.dispatch_namespace ?? null, + }; + } catch (error) { + if (isCloudflareStatus(error, 404)) return null; + throw error; + } + }, + async uploadScript(scriptName) { + const result = await runOwnedProdAsyncMutation(() => + client.workersForPlatforms.dispatch.namespaces.scripts.update(scriptName, { + ...params, + files: [ + new File( + [ + 'export default { fetch() { return new Response("mosoo release write canary"); } };', + ], + "canary.mjs", + { type: "application/javascript+module" }, + ), + ], + metadata: { + compatibility_date: "2026-06-01", + main_module: "canary.mjs", + }, + }), + ); + return { id: result.id ?? null }; + }, + }; +} + +function readProtocolV3SmokeConfig(): ProtocolV3SmokeConfig { + const agentId = process.env["MOSOO_PROTOCOL_V3_SMOKE_AGENT_ID"]?.trim() ?? ""; + const token = process.env["MOSOO_PROTOCOL_V3_SMOKE_TOKEN"]?.trim() ?? ""; + + if (!PLATFORM_ID_PATTERN.test(agentId) || !token) { + throw new Error( + "Production deploy requires a dedicated MOSOO_PROTOCOL_V3_SMOKE_AGENT_ID and MOSOO_PROTOCOL_V3_SMOKE_TOKEN.", + ); + } + + return { agentId, token }; +} + +function probeProtocolV3Cutover() { + return parseProtocolV3CutoverProbe(executeProdD1Json(PROTOCOL_V3_CUTOVER_PROBE_SQL)); +} + +function readPendingProdMigrations(localMigrationNames: readonly string[]): string[] { + return findPendingProdMigrations( + executeProdD1Json(PROD_APPLIED_MIGRATIONS_SQL), + localMigrationNames, + ); +} + +function printProtocolV3LossyMigrationInventory( + inventory: ProtocolV3LossyMigrationInventory, +): void { + writeStdout( + ` migration 0013 loss candidates: total=${inventory.totalCandidates} orphan_effects=${inventory.orphanEffects} attempt_completion_time_fabrications=${inventory.attemptCompletionTimeFabrications} command_payload_conflicts=${inventory.commandPayloadConflicts} MCP_arguments=${inventory.mcpArgumentOmissions} input_text=${inventory.inputTextOmissions} input_results=${inventory.inputStartResultOmissions} control_reason=${inventory.controlReasonOmissions} permission_payload_rewrites=${inventory.permissionPayloadRewrites} MCP_result_omissions=${inventory.mcpResultOmissions} MCP_result_conflicts=${inventory.mcpResultConflicts} provider_receipts=${inventory.providerReceiptLosses} MCP_terminal_conflicts=${inventory.mcpCommandTerminalConflicts} command_errors=${inventory.commandErrorOmissions} Session_Run_errors=${inventory.sessionRunErrorOmissions}`, + ); + for (const candidate of inventory.candidateIds) { + writeStdout(` ${candidate.category}: ${candidate.id}`); + } + if (inventory.totalCandidates > inventory.candidateIds.length) { + writeStdout( + ` ... ${inventory.totalCandidates - inventory.candidateIds.length} more candidates`, + ); + } +} + +function verifyProdLossyMigrationInventory(): void { + const inventory = parseProtocolV3LossyMigrationInventory( + executeProdD1Json(PROTOCOL_V3_LOSSY_MIGRATION_INVENTORY_SQL), + ); + printProtocolV3LossyMigrationInventory(inventory); + assertProtocolV3LossyMigrationInventory(inventory); +} + +function printProtocolV3LegacyTerminalSourceInventory( + inventory: ProtocolV3LegacyTerminalSourceInventory, +): void { + for (const [kind, counts] of [ + ["cancelled", inventory.cancelled], + ["completed", inventory.completed], + ["failed", inventory.failed], + ] as const) { + writeStdout( + ` run.${kind}: total=${counts.total} canonical=${counts.canonical} rewrite_candidates=${counts.noncanonical} canonical_target_collisions=${counts.canonicalTargetCollisions}`, + ); + } + writeStdout( + ` invalid terminal links=${inventory.invalidTerminalLinks} status/kind mismatches=${inventory.mismatchedTerminalEvents} multiple terminal Runs=${inventory.multipleTerminalRuns}`, + ); +} + +function verifyProdLegacyTerminalSourceInventory(): void { + const inventory = parseProtocolV3LegacyTerminalSourceInventory( + executeProdD1Json(PROTOCOL_V3_LEGACY_TERMINAL_SOURCE_INVENTORY_SQL), + ); + printProtocolV3LegacyTerminalSourceInventory(inventory); + assertProtocolV3LegacyTerminalSourceInventory(inventory); +} + +interface LegacyRewriteProof { + readonly candidateCount: number; + readonly candidateManifestJson: string; +} + +function preflightProdLegacyTerminalIntegrity(): LegacyRewriteProof { + verifyProdLegacyTerminalSourceInventory(); + + const integrity = parseProtocolV3LegacyTerminalIntegrity( + executeProdD1Json(PROTOCOL_V3_LEGACY_TERMINAL_INTEGRITY_SQL), + ); + assertProtocolV3LegacyTerminalIntegrity(integrity); + writeStdout( + ` legacy terminal integrity is unambiguous; migration 0014 will normalize ${integrity.noncanonicalTerminalSources} terminal source identities and ${integrity.repairableFailedRuns} failed Run errors, label ${integrity.legacyMaterializedMessages} messages as materialized, backfill ${integrity.legacyStreamRows} row-local stream identities, and preserve ${integrity.legacyTerminalEvents} terminal events as semantic_hash=NULL legacy history`, + ); + return { + candidateCount: integrity.noncanonicalTerminalSources, + candidateManifestJson: integrity.rewriteCandidateManifestJson, + }; +} + +function authorizeProdLegacyTerminalRewrite( + proof: LegacyRewriteProof, + releaseTreeOid: string, + bookmark: string, +): void { + if (prodDeployLeaseOwner === null) { + throw new Error("Protocol v3 legacy rewrite requires the production deploy lease."); + } + const owner = prodDeployLeaseOwner; + executeProdD1( + authorizeProtocolV3LegacyRewriteSql(owner, proof.candidateCount, proof.candidateManifestJson), + ); + const authorization = parseProtocolV3LegacyRewriteAuthorization( + executeProdD1Json(PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_SQL), + ); + if ( + authorization.deployOwner !== owner || + authorization.bookmark !== bookmark || + authorization.candidateCount !== proof.candidateCount || + authorization.candidateManifestJson !== proof.candidateManifestJson || + authorization.releaseTreeOid !== releaseTreeOid || + authorization.expiresAt <= Math.floor(Date.now() / 1000) + ) { + throw new Error("Protocol v3 legacy rewrite authorization did not persist exactly."); + } +} + +function readProtocolV3CutoverState(): ProtocolV3CutoverState { + return parseProtocolV3CutoverState(executeProdD1Json(PROTOCOL_V3_COMMAND_FREEZE_SQL)); +} + +function installProtocolV3CutoverGate( + releaseTreeOid: string, + postRuntimeAuthorityMigration: boolean, +): ProtocolV3CutoverState { + executeProdD1( + postRuntimeAuthorityMigration + ? installProtocolV3PostMigrationCutoverSql(releaseTreeOid) + : installProtocolV3CutoverSql(releaseTreeOid), + ); + assertProtocolV3CutoverGateExact(); + const state = readProtocolV3CutoverState(); + assertProtocolV3Release(state, releaseTreeOid); + return state; +} + +function assertProtocolV3CutoverGateExact(): void { + const objects = parseProtocolV3CutoverObjects(executeProdD1Json(PROTOCOL_V3_CUTOVER_OBJECTS_SQL)); + + const validCount = + objects.objectCount === PROTOCOL_V3_CUTOVER_OBJECT_COUNT || + objects.objectCount === PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT; + if (!validCount || objects.exactObjectCount !== objects.objectCount) { + throw new Error( + `Protocol v3 cutover gate is invalid (${objects.objectCount} observed, ${objects.exactObjectCount} exact objects).`, + ); + } +} + +function removeProtocolV3CutoverGate(): void { + executeProdD1(REMOVE_PROTOCOL_V3_CUTOVER_SQL); + const objects = parseProtocolV3CutoverObjects(executeProdD1Json(PROTOCOL_V3_CUTOVER_OBJECTS_SQL)); + if (objects.objectCount !== 0 || probeProtocolV3Cutover().gatePresent) { + throw new Error( + `Protocol v3 cutover gate cleanup is incomplete (${objects.objectCount} reserved objects remain).`, + ); + } +} + +function enableProtocolV3CommandFreeze(): void { + executeProdD1(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + if (!parseProtocolV3CommandFreeze(executeProdD1Json(PROTOCOL_V3_COMMAND_FREEZE_SQL))) { + throw new Error("Protocol v3 final Driver command freeze was not enabled."); + } +} + +function beginProtocolV3Migration(releaseTreeOid: string): ProtocolV3CutoverState { + executeProdD1(beginProtocolV3MigrationSql(releaseTreeOid)); + const state = readProtocolV3CutoverState(); + if (!state.migrationStarted) { + throw new Error("Protocol v3 migration intent did not persist exactly."); + } + return state; +} + +function enterProtocolV3QueuesResuming(): ProtocolV3CutoverState { + executeProdD1(ENTER_PROTOCOL_V3_QUEUES_RESUMING_SQL); + const state = readProtocolV3CutoverState(); + if (state.phase !== "queues_resuming" || !state.enabled || !state.commandFreeze) { + throw new Error("Protocol v3 queue-resume recovery phase was not persisted."); + } + return state; +} + +function acceptProtocolV3QueueResume(): void { + executeProdD1(ACCEPT_PROTOCOL_V3_QUEUE_RESUME_SQL); + const state = readProtocolV3CutoverState(); + if (state.phase !== "queues_resuming" || state.enabled || !state.commandFreeze) { + throw new Error("Protocol v3 queue-resume acceptance was not persisted."); + } +} + +function enterProtocolV3Drain(): void { + executeProdD1(ENTER_PROTOCOL_V3_DRAIN_SQL); + if (parseProtocolV3CommandFreeze(executeProdD1Json(PROTOCOL_V3_COMMAND_FREEZE_SQL))) { + throw new Error("Protocol v3 drain gate did not keep control commands available."); + } +} + +async function updateAndVerifyProdQueues( + config: ProdQueueApiConfig, + action: "pause" | "resume", +): Promise { + const client = new Cloudflare({ apiToken: config.apiToken }); + const queues: Array<{ id: string; name: string }> = []; + + for await (const queue of client.queues.list({ account_id: config.accountId })) { + if (queue.queue_name && queue.queue_id) { + queues.push({ id: queue.queue_id, name: queue.queue_name }); + } + } + + await updateAndVerifyProtocolV3QueueDelivery( + { + list: async () => queues, + mutate: (queueName, queueAction) => { + runProdMutation(["queues", `${queueAction}-delivery`, queueName]); + }, + read: async (queueId) => { + const queue = await client.queues.get(queueId, { account_id: config.accountId }); + return { + deliveryPaused: queue.settings?.delivery_paused, + name: queue.queue_name ?? "", + }; + }, + }, + action, + ); + for (const queueName of PROTOCOL_V3_CUTOVER_QUEUE_NAMES) { + writeStdout(` queue ${queueName} delivery is ${action}d`); + } +} + +function resumeAndVerifyProdQueues(config: ProdQueueApiConfig): Promise { + return updateAndVerifyProdQueues(config, "resume"); +} + +function pauseAndVerifyProdQueues(config: ProdQueueApiConfig): Promise { + return updateAndVerifyProdQueues(config, "pause"); +} + +async function waitForProtocolV3Drain( + postRuntimeAuthorityMigration: boolean, + requireCommands = true, +): Promise { + const deadline = Date.now() + CUTOVER_DRAIN_TIMEOUT_MS; + + while (true) { + const state = parseProtocolV3CutoverDrain( + executeProdD1Json( + postRuntimeAuthorityMigration + ? PROTOCOL_V3_POST_MIGRATION_CUTOVER_DRAIN_SQL + : PROTOCOL_V3_CUTOVER_DRAIN_SQL, + ), + ); + + writeStdout( + ` drain: runs=${state.activeRuns} appDeploymentRuns=${state.activeAppDeploymentRuns} drivers=${state.liveDrivers} effects=${state.unsettledEffects} driverCommands=${state.nonterminalCommands} apiCommands=${state.nonterminalApiCommands} sandboxes=${state.unsafeSandboxes} sandboxSessions=${state.unsafeSandboxSessions} backups=${state.unsafeSandboxBackups} backupStaging=${state.unsafeSandboxBackupStaging} environmentArtifactStaging=${state.unsafeEnvironmentArtifactBackupStaging} sessions=${state.unsafeSessions}`, + ); + if (requireCommands ? isProtocolV3CutoverDrained(state) : isProtocolV3RuntimeDrained(state)) { + return; + } + if (Date.now() >= deadline) { + const unsafeSandboxes = parseD1JsonResults( + executeProdD1Json( + postRuntimeAuthorityMigration + ? PROTOCOL_V3_POST_MIGRATION_UNSAFE_SANDBOXES_SQL + : PROTOCOL_V3_UNSAFE_SANDBOXES_SQL, + ), + ) + .flatMap((rows) => rows) + .map((row) => `${String(row.id)}(${String(row.status)})`) + .join(", "); + throw new Error( + `Protocol v3 ${requireCommands ? "final command" : "runtime"} drain did not finish before the 15-minute safety timeout.${unsafeSandboxes.length === 0 ? "" : ` Unsafe sandboxes: ${unsafeSandboxes}. Use the supported lifecycle hibernate/checkpoint path, then retry.`}`, + ); + } + + await sleep(POLL_INTERVAL_MS); + } +} + +function readStoredProtocolV3Bookmark(): string | null { + return parseStoredProtocolV3CutoverBookmark(executeProdD1Json(PROTOCOL_V3_CUTOVER_BOOKMARK_SQL)); +} + +function printProtocolV3EmergencyBookmark(bookmark: string): void { + writeStdout(` emergency destructive backup bookmark: ${bookmark}`); + writeStdout( + " After migration, automated recovery is roll-forward only: rerun this exact v3 release.", + ); + writeStdout( + " D1 Time Travel would discard every D1 write after this bookmark and does not restore Durable Object storage.", + ); + writeStdout( + " Manual v2 recovery requires a global maintenance write stop, a reconciled inventory of post-bookmark writes, and explicit human approval.", + ); +} + +function readExactProdWorkerVersion(releaseTreeOid: string): string { + const deployment = parseProtocolV3WorkerDeployment( + captureWrangler(["deployments", "status", "--env", ENV, "--json"]), + ); + assertProtocolV3WorkerVersion( + captureWrangler(["versions", "view", deployment.versionId, "--env", ENV, "--json"]), + deployment.versionId, + releaseTreeOid, + ); + return deployment.versionId; +} + +function ensureProtocolV3Bookmark(): string { + const stored = readStoredProtocolV3Bookmark(); + if (stored !== null) return stored; + + const bookmark = parseTimeTravelBookmark( + captureWrangler(["d1", "time-travel", "info", D1_BINDING, "--env", ENV, "--json"]), + ); + executeProdD1(storeProtocolV3CutoverBookmarkSql(bookmark)); + const persisted = readStoredProtocolV3Bookmark(); + + if (persisted !== bookmark) { + throw new Error( + "Failed to persist the pre-migration Time Travel bookmark in the cutover gate.", + ); + } + + return bookmark; +} + +async function readProdContainerApplication( + config: ProdQueueApiConfig, +): Promise { + const client = new Cloudflare({ apiToken: config.apiToken }); + const applications = await collectProtocolV3ContainerApplications((pageToken) => + client.get(`/accounts/${config.accountId}/containers/dash/applications`, { + query: { + per_page: 100, + ...(pageToken === null ? {} : { page_token: pageToken }), + }, + timeout: 10_000, + }), + ); + const matches = applications.filter( + (application) => application.name === PROD_CONTAINER_APPLICATION, + ); + if (matches.length !== 1) { + throw new Error( + `Expected exactly one Container application ${PROD_CONTAINER_APPLICATION}; found ${matches.length}.`, + ); + } + return matches[0]; +} + +function readProdContainerInstances(applicationId: string): Promise { + return collectProtocolV3ContainerInstances((pageToken) => + captureWrangler([ + "containers", + "instances", + applicationId, + "--json", + "--per-page", + "100", + ...(pageToken === null ? [] : ["--page-token", pageToken]), + "--env", + ENV, + ]), + ); +} + +async function waitForContainerRollout( + config: ProdQueueApiConfig, + expectedImageDigest: string, + previousVersion: number, + expectedVersion: number | null, +): Promise { + const deadline = Date.now() + CUTOVER_ROLLOUT_TIMEOUT_MS; + + while (true) { + const application = await readProdContainerApplication(config); + const instances = await readProdContainerInstances(application.id); + + if (expectedVersion !== null && application.version !== expectedVersion) { + throw new Error( + `Production Container application moved from release version ${expectedVersion} to ${application.version}.`, + ); + } + if ( + expectedVersion === null && + (application.version < previousVersion || application.version > previousVersion + 1) + ) { + throw new Error("Production Container application advanced outside this release rollout."); + } + if ( + application.imageDigest !== expectedImageDigest && + application.version !== previousVersion + ) { + throw new Error("Production Container application advanced to a foreign image."); + } + + if ( + application.imageDigest === expectedImageDigest && + isProtocolV3ContainerRolloutConverged(application, instances) + ) { + writeStdout( + ` container rollout converged at application version ${application.version} (${instances.length} known instances)`, + ); + return application; + } + if (Date.now() >= deadline) { + throw new Error( + "Protocol v3 Container rollout did not converge before the 10-minute timeout.", + ); + } + + writeStdout( + ` waiting for ${instances.length} Container instances to reach the target version`, + ); + await sleep(POLL_INTERVAL_MS); + } +} + +async function assertPublishedProtocolV3Release( + state: ProtocolV3CutoverState, + releaseTreeOid: string, + config: ProdQueueApiConfig, +): Promise { + assertProtocolV3Release(state, releaseTreeOid); + if ( + state.workerVersionId === null || + state.containerApplicationVersion === null || + state.containerImageDigest === null + ) { + throw new Error("Protocol v3 rollout metadata is not durably bound to the cutover marker."); + } + if (readExactProdWorkerVersion(releaseTreeOid) !== state.workerVersionId) { + throw new Error("Production Worker moved away from the bound protocol v3 version."); + } + const application = await readProdContainerApplication(config); + if ( + application.version !== state.containerApplicationVersion || + application.imageDigest !== state.containerImageDigest || + readTaggedContainerImageDigest(application, state.workerVersionId) !== + state.containerImageDigest + ) { + throw new Error( + "Production Container application moved away from the bound protocol v3 version.", + ); + } +} + +async function verifyProdHealth(): Promise { + let lastError: unknown = new Error("Production health check was not attempted."); + + for (let attempt = 1; attempt <= 12; attempt += 1) { + try { + const response = await fetch(PROD_HEALTH_URL, { signal: AbortSignal.timeout(10_000) }); + const body = (await response.json()) as { name?: unknown; ok?: unknown }; + + if (response.ok && body.name === "mosoo" && body.ok === true) { + writeStdout(" production Worker and D1 health check passed"); + return; + } + + lastError = new Error(`health returned HTTP ${response.status}`); + } catch (error) { + lastError = error; + } + + if (attempt < 12) await sleep(POLL_INTERVAL_MS); + } + + throw new Error("Production health check failed after deployment.", { cause: lastError }); +} + +function requireJsonRecord(value: unknown, label: string): Record { + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new Error(`${label} must be a JSON object.`); + } + return value as Record; +} + +async function fetchProdMutation(input: string, init: RequestInit): Promise { + return runOwnedProdAsyncMutation(() => + fetch(input, { + ...init, + signal: init.signal ?? AbortSignal.timeout(REMOTE_MUTATION_TIMEOUT_MS), + }), + ); +} + +async function fetchProtocolV3SmokeJson( + config: ProtocolV3SmokeConfig, + path: string, + init: RequestInit, +): Promise> { + const headers = new Headers(init.headers); + headers.set("Accept", "application/json"); + headers.set("Authorization", `Bearer ${config.token}`); + if (init.body !== undefined) headers.set("Content-Type", "application/json"); + + const response = await fetchProdMutation(`${PROD_PUBLIC_API_URL}${path}`, { + ...init, + headers, + signal: AbortSignal.timeout(30_000), + }); + const body = requireJsonRecord(await response.json(), "Protocol v3 smoke response"); + + if (!response.ok) { + throw new Error( + `Protocol v3 smoke ${init.method ?? "GET"} ${path} returned HTTP ${response.status}.`, + ); + } + + return body; +} + +async function readProtocolV3SmokeAccountId(config: ProtocolV3SmokeConfig): Promise { + const response = await fetch(PROD_GRAPHQL_URL, { + body: JSON.stringify({ query: "query ProtocolV3SmokeViewer { viewer { account { id } } }" }), + headers: { + Accept: "application/json", + Authorization: `Bearer ${config.token}`, + "Content-Type": "application/json", + }, + method: "POST", + signal: AbortSignal.timeout(30_000), + }); + const body = requireJsonRecord(await response.json(), "Protocol v3 smoke viewer response"); + + if (!response.ok || (Array.isArray(body.errors) && body.errors.length > 0)) { + throw new Error(`Protocol v3 smoke viewer query returned HTTP ${response.status}.`); + } + + const data = requireJsonRecord(body.data, "Protocol v3 smoke viewer data"); + const viewer = requireJsonRecord(data.viewer, "Protocol v3 smoke viewer"); + const account = requireJsonRecord(viewer.account, "Protocol v3 smoke account"); + const accountId = account.id; + + if (typeof accountId !== "string" || !PLATFORM_ID_PATTERN.test(accountId)) { + throw new Error("Protocol v3 smoke PAT did not resolve to a valid account ID."); + } + + return accountId; +} + +async function createProtocolV3SmokeThread( + config: ProtocolV3SmokeConfig, + idempotencyKey: string, +): Promise { + let lastError: unknown = new Error("Protocol v3 smoke create was not attempted."); + + for (let attempt = 1; attempt <= 3; attempt += 1) { + try { + const body = await fetchProtocolV3SmokeJson( + config, + `/agents/${encodeURIComponent(config.agentId)}/threads`, + { + body: JSON.stringify({ userId: idempotencyKey }), + headers: { "Idempotency-Key": idempotencyKey }, + method: "POST", + }, + ); + const thread = requireJsonRecord(body.thread, "Protocol v3 smoke Thread"); + const threadId = thread.id; + + if (body.run !== null) { + throw new Error("Protocol v3 empty smoke Thread unexpectedly created a Run."); + } + if (typeof threadId !== "string" || !PLATFORM_ID_PATTERN.test(threadId)) { + throw new Error("Protocol v3 smoke create response is missing a valid Thread ID."); + } + + return threadId; + } catch (error) { + lastError = error; + if (attempt < 3) await sleep(POLL_INTERVAL_MS); + } + } + + throw new Error("Protocol v3 smoke Thread creation failed after idempotent retries.", { + cause: lastError, + }); +} + +async function waitForProtocolV3SmokeReady(sessionId: string): Promise { + const deadline = Date.now() + CUTOVER_SMOKE_TIMEOUT_MS; + + while (true) { + const status = parseProtocolV3SmokeStatus( + executeProdD1Json(protocolV3SmokeStatusSql(sessionId)), + ); + + if (isProtocolV3SmokeReady(status)) { + writeStdout( + ` live Driver ${status.driverVersion} completed protocol v3 boot, hello, and ready`, + ); + return; + } + if (status.driverStatus === "failed" || status.driverStatus === "stopped") { + throw new Error( + `Protocol v3 smoke Driver became ${status.driverStatus} before completing hello and ready.`, + ); + } + if (Date.now() >= deadline) { + throw new Error("Protocol v3 live Driver smoke did not reach ready within five minutes."); + } + + await sleep(POLL_INTERVAL_MS); + } +} + +async function deleteProtocolV3SmokeThread( + config: ProtocolV3SmokeConfig, + threadId: string, +): Promise { + let lastError: unknown = new Error("Protocol v3 smoke cleanup was not attempted."); + + for (let attempt = 1; attempt <= 3; attempt += 1) { + try { + const response = await fetchProdMutation( + `${PROD_PUBLIC_API_URL}/threads/${encodeURIComponent(threadId)}`, + { + headers: { + Accept: "application/json", + Authorization: `Bearer ${config.token}`, + }, + method: "DELETE", + signal: AbortSignal.timeout(30_000), + }, + ); + if (response.ok || response.status === 404) return; + throw new Error(`Protocol v3 smoke DELETE returned HTTP ${response.status}.`); + } catch (error) { + lastError = error; + if (attempt < 3) await sleep(POLL_INTERVAL_MS); + } + } + + throw new Error(`Protocol v3 smoke Thread ${threadId} cleanup failed after retries.`, { + cause: lastError, + }); +} + +function readStoredProtocolV3SmokeSession(): string | null { + return parseStoredProtocolV3SmokeSession(executeProdD1Json(PROTOCOL_V3_SMOKE_SESSION_SQL)); +} + +function readStoredProtocolV3SmokeRequestKey(): string | null { + return parseStoredProtocolV3SmokeRequestKey(executeProdD1Json(PROTOCOL_V3_SMOKE_REQUEST_KEY_SQL)); +} + +async function cleanInterruptedProtocolV3Smoke(config: ProtocolV3SmokeConfig): Promise { + let sessionId = readStoredProtocolV3SmokeSession(); + const requestKey = readStoredProtocolV3SmokeRequestKey(); + + if (sessionId === null && requestKey !== null) { + writeStdout(` recovering interrupted protocol v3 smoke request ${requestKey}`); + sessionId = await createProtocolV3SmokeThread(config, requestKey); + executeProdD1(storeProtocolV3SmokeSessionSql(sessionId)); + } + if (sessionId === null) return; + + writeStdout(` cleaning interrupted protocol v3 smoke Session ${sessionId}`); + await deleteProtocolV3SmokeThread(config, sessionId); + executeProdD1(CLOSE_PROTOCOL_V3_SMOKE_WINDOW_SQL); +} + +async function runProtocolV3LiveSmoke( + config: ProtocolV3SmokeConfig, + closedGate: boolean, +): Promise { + let threadId: string | null = null; + let failure: unknown = null; + const cleanupFailures: unknown[] = []; + + if (closedGate) { + const accountId = await readProtocolV3SmokeAccountId(config); + executeProdD1(openProtocolV3SmokeWindowSql(accountId)); + if ( + readStoredProtocolV3SmokeSession() !== null || + readStoredProtocolV3SmokeRequestKey() !== null + ) { + await cleanInterruptedProtocolV3Smoke(config); + executeProdD1(openProtocolV3SmokeWindowSql(accountId)); + } + } + + try { + const requestKey = `protocol-v3-cutover-${crypto.randomUUID()}`; + if (closedGate) executeProdD1(storeProtocolV3SmokeRequestKeySql(requestKey)); + threadId = await createProtocolV3SmokeThread(config, requestKey); + if (closedGate) executeProdD1(storeProtocolV3SmokeSessionSql(threadId)); + await waitForProtocolV3SmokeReady(threadId); + } catch (error) { + failure = error; + } finally { + if (threadId !== null) { + try { + await deleteProtocolV3SmokeThread(config, threadId); + } catch (error) { + cleanupFailures.push(error); + } + } + + if (closedGate && cleanupFailures.length === 0 && threadId !== null) { + try { + executeProdD1(CLOSE_PROTOCOL_V3_SMOKE_WINDOW_SQL); + if (!parseProtocolV3CommandFreeze(executeProdD1Json(PROTOCOL_V3_COMMAND_FREEZE_SQL))) { + cleanupFailures.push( + new Error("Protocol v3 command freeze did not close after live smoke."), + ); + } + } catch (error) { + cleanupFailures.push(error); + } + } else if (closedGate) { + writeStdout( + `✗ Retaining the smoke allowance for cleanup recovery${threadId === null ? "." : ` of Session ${threadId}.`}`, + ); + } + } + + if (failure !== null || cleanupFailures.length > 0) { + throw new AggregateError( + [...(failure === null ? [] : [failure]), ...cleanupFailures], + "Protocol v3 live Driver smoke or cleanup failed.", + ); + } +} + +async function deployWorkerAndVerify( + smokeConfig: ProtocolV3SmokeConfig, + queueApiConfig: ProdQueueApiConfig, + closedGate: boolean, + releaseTreeOid: string, + cutoverState: ProtocolV3CutoverState | null, + workflowMode: ProdWfpWorkflowMode, +): Promise { + let state = cutoverState; + let workerVersionId = state?.workerVersionId ?? null; + let containerApplicationVersion = state?.containerApplicationVersion ?? null; + let containerImageDigest = state?.containerImageDigest ?? null; + const cloudflare = new Cloudflare({ apiToken: queueApiConfig.apiToken }); + + if ( + workerVersionId === null || + containerApplicationVersion === null || + containerImageDigest === null + ) { + const previousApplication = await readProdContainerApplication(queueApiConfig); + writeStdout("▶ Deploying protocol v3 Worker and Driver image"); + assertReleaseTreeUnchanged(releaseTreeOid); + const namespace = await readProdWfpNamespace(cloudflare, queueApiConfig); + assertProdWfpNamespaceExact(namespace, PROD_APP_DISPATCH_NAMESPACE, false); + const workflowBoundary = await readProdAppDeploymentWorkflowBoundary( + cloudflare, + queueApiConfig, + ); + const currentWorkflowMode = assertProdWfpWorkflowBoundary( + workflowBoundary.workflow, + workflowBoundary.bindings, + prodWfpInfrastructureConfig(queueApiConfig, false, workflowMode === "repair"), + ); + if (currentWorkflowMode !== workflowMode) { + throw new Error("Production App deployment Workflow changed before Worker deployment."); + } + assertProdWfpWorkflowDeployGate(workflowMode, closedGate); + runProdMutation(apiWorkerDeployArgs(protocolV3ReleaseTag(releaseTreeOid))); + assertReleaseTreeUnchanged(releaseTreeOid); + workerVersionId = readExactProdWorkerVersion(releaseTreeOid); + containerImageDigest = readTaggedContainerImageDigest(previousApplication, workerVersionId); + const application = await waitForContainerRollout( + queueApiConfig, + containerImageDigest, + previousApplication.version, + null, + ); + containerApplicationVersion = application.version; + } else { + if (state === null) throw new Error("Stored protocol v3 rollout state is missing."); + writeStdout("▶ Reusing the exact protocol v3 Worker and Container rollout"); + await assertPublishedProtocolV3Release(state, releaseTreeOid, queueApiConfig); + await waitForContainerRollout( + queueApiConfig, + containerImageDigest, + containerApplicationVersion, + containerApplicationVersion, + ); + } + + const publishedWorkflowBoundary = await readProdAppDeploymentWorkflowBoundary( + cloudflare, + queueApiConfig, + ); + if ( + assertProdWfpWorkflowBoundary( + publishedWorkflowBoundary.workflow, + publishedWorkflowBoundary.bindings, + prodWfpInfrastructureConfig(queueApiConfig, false), + ) !== "exact" + ) { + throw new Error("Published Worker does not expose the exact App deployment Workflow."); + } + + writeStdout("▶ Verifying production Worker, D1, and local protocol contract"); + assertProtocolV3Artifacts(); + await verifyProdHealth(); + + writeStdout("▶ Re-probing the production WfP HTTP, streaming, cancellation, and WebSocket path"); + await verifyProdWfpGatewayProbe(PROD_APP_DEPLOYMENT_DOMAIN, { + fetch: globalThis.fetch, + probeWebSocket: probeProdWfpWebSocket, + }); + + writeStdout("▶ Running a live protocol v3 Driver boot, hello, and ready smoke"); + await runProtocolV3LiveSmoke(smokeConfig, closedGate); + + if (readExactProdWorkerVersion(releaseTreeOid) !== workerVersionId) { + throw new Error("Production Worker moved away from the published protocol v3 version."); + } + const application = await readProdContainerApplication(queueApiConfig); + if ( + application.version !== containerApplicationVersion || + application.imageDigest !== containerImageDigest || + readTaggedContainerImageDigest(application, workerVersionId) !== containerImageDigest + ) { + throw new Error("Production Container rollout changed before release binding."); + } + + if (state !== null && state.workerVersionId === null) { + executeProdD1( + storeProtocolV3RolloutSql( + releaseTreeOid, + workerVersionId, + containerApplicationVersion, + containerImageDigest, + ), + ); + state = readProtocolV3CutoverState(); + } + if (state !== null) { + await assertPublishedProtocolV3Release(state, releaseTreeOid, queueApiConfig); + } + return state; +} + +async function runProtocolV3Cutover( + initialPendingMigrations: readonly string[], + localMigrationNames: readonly string[], + smokeConfig: ProtocolV3SmokeConfig, + queueApiConfig: ProdQueueApiConfig, + expectedProdSchema: ProdSchemaCatalog, + releaseTreeOid: string, + workflowMode: ProdWfpWorkflowMode, +): Promise { + let bookmark: string | null = null; + let migrationStarted = true; + let queuesVerified = false; + const durableMcpMigrationPending = initialPendingMigrations.includes(PROTOCOL_V3_MIGRATION); + const sessionEventMigrationPending = initialPendingMigrations.includes( + PROTOCOL_V3_SESSION_EVENT_MIGRATION, + ); + const postRuntimeAuthorityMigration = !initialPendingMigrations.includes( + PROTOCOL_V3_RUNTIME_AUTHORITY_MIGRATION, + ); + + try { + writeStdout("▶ Installing the one-shot D1 admission gate"); + let cutoverState = installProtocolV3CutoverGate(releaseTreeOid, postRuntimeAuthorityMigration); + migrationStarted = cutoverState.migrationStarted; + if (cutoverState.phase === "queues_resuming") { + if (initialPendingMigrations.length > 0) { + throw new Error( + "Protocol v3 queue-resume phase exists before every migration was applied.", + ); + } + await assertPublishedProtocolV3Release(cutoverState, releaseTreeOid, queueApiConfig); + queuesVerified = !cutoverState.enabled; + writeStdout( + `▶ Recovering the durable production queue-resume ${cutoverState.enabled ? "pre-acceptance" : "accepted"} phase`, + ); + await completeProtocolV3QueueResume(cutoverState, { + commitAcceptance: acceptProtocolV3QueueResume, + removeMarker: removeProtocolV3CutoverGate, + resumeAndVerifyQueues: async () => { + await resumeAndVerifyProdQueues(queueApiConfig); + queuesVerified = true; + }, + }); + return; + } + + writeStdout("▶ Pausing production queue delivery for protocol v3 cutover"); + await pauseAndVerifyProdQueues(queueApiConfig); + + if (initialPendingMigrations.length > 0) { + enterProtocolV3Drain(); + + writeStdout("▶ Letting every already-admitted API command lane reach a terminal state"); + await resumeAndVerifyProdQueues(queueApiConfig); + + writeStdout("▶ Draining protocol v2 runtime state while control commands remain available"); + await waitForProtocolV3Drain(postRuntimeAuthorityMigration, false); + + writeStdout("▶ Freezing all new Driver commands at the final zero-state boundary"); + enableProtocolV3CommandFreeze(); + await waitForProtocolV3Drain(postRuntimeAuthorityMigration); + + writeStdout("▶ Re-pausing every API command lane after all admitted work is terminal"); + await pauseAndVerifyProdQueues(queueApiConfig); + await waitForProtocolV3Drain(postRuntimeAuthorityMigration); + } else { + writeStdout("▶ Recovering any interrupted protocol v3 smoke before roll-forward"); + await cleanInterruptedProtocolV3Smoke(smokeConfig); + executeProdD1(CLOSE_PROTOCOL_V3_SMOKE_WINDOW_SQL); + enableProtocolV3CommandFreeze(); + writeStdout("▶ Re-proving the complete closed runtime boundary before roll-forward"); + await waitForProtocolV3Drain(postRuntimeAuthorityMigration); + } + + if (durableMcpMigrationPending) { + writeStdout("▶ Preflighting migration 0013 for lossy historical rewrites"); + verifyProdLossyMigrationInventory(); + } + let legacyRewriteProof: LegacyRewriteProof | null = null; + if (sessionEventMigrationPending) { + writeStdout("▶ Preflighting legacy terminal integrity before migration 0014"); + legacyRewriteProof = preflightProdLegacyTerminalIntegrity(); + } + if (initialPendingMigrations.includes(PROTOCOL_V3_RUNTIME_AUTHORITY_MIGRATION)) { + writeStdout("▶ Preflighting runtime authority identities and terminal backups"); + assertProtocolV3RuntimeAuthorityPreflight( + parseProtocolV3RuntimeAuthorityPreflight( + executeProdD1Json( + protocolV3RuntimeAuthorityPreflightSql( + initialPendingMigrations.includes(PROTOCOL_V3_SESSION_CLEANUP_MIGRATION), + ), + ), + ), + ); + } + + bookmark = + initialPendingMigrations.length > 0 + ? (readStoredProtocolV3Bookmark() ?? ensureProtocolV3Bookmark()) + : readStoredProtocolV3Bookmark(); + if (bookmark === null) { + writeStdout( + " No emergency backup bookmark is available; this does not block roll-forward of the exact v3 release.", + ); + } else { + printProtocolV3EmergencyBookmark(bookmark); + } + + writeStdout("▶ Re-pausing and independently verifying every production queue before migration"); + await pauseAndVerifyProdQueues(queueApiConfig); + writeStdout("▶ Re-verifying the exact admission gate before migration"); + assertProtocolV3CutoverGateExact(); + assertReleaseTreeUnchanged(releaseTreeOid); + + if (initialPendingMigrations.length > 0) { + writeStdout("▶ Persisting the irreversible D1 migration intent"); + migrationStarted = true; + cutoverState = beginProtocolV3Migration(releaseTreeOid); + } + if (sessionEventMigrationPending) { + if (bookmark === null) { + throw new Error("Migration 0014 requires a persisted pre-migration D1 bookmark."); + } + writeStdout("▶ Authorizing the exact drained legacy terminal rewrite set"); + if (legacyRewriteProof === null) { + throw new Error("Migration 0014 legacy rewrite preflight evidence is missing."); + } + authorizeProdLegacyTerminalRewrite(legacyRewriteProof, releaseTreeOid, bookmark); + } + writeStdout("▶ Applying pending D1 migrations behind the closed gate"); + applyD1Migrations(); + assertReleaseTreeUnchanged(releaseTreeOid); + const remainingMigrations = readPendingProdMigrations(localMigrationNames); + if (remainingMigrations.length > 0) { + throw new Error( + `Production D1 migrations remain pending behind the closed gate: ${remainingMigrations.join(", ")}.`, + ); + } + + assertProtocolV3CutoverGateExact(); + + writeStdout("▶ Verifying prod D1 schema matches the latest migration snapshot"); + assertProdSchemaMatchesSnapshot(expectedProdSchema); + + cutoverState = + (await deployWorkerAndVerify( + smokeConfig, + queueApiConfig, + true, + releaseTreeOid, + readProtocolV3CutoverState(), + workflowMode, + )) ?? cutoverState; + + writeStdout("▶ Rechecking the closed runtime boundary"); + await waitForProtocolV3Drain(true); + + writeStdout("▶ Persisting the durable production queue-resume phase"); + const queueResumeState = enterProtocolV3QueuesResuming(); + + await completeProtocolV3QueueResume(queueResumeState, { + commitAcceptance: acceptProtocolV3QueueResume, + removeMarker: removeProtocolV3CutoverGate, + resumeAndVerifyQueues: async () => { + writeStdout("▶ Resuming and verifying production queues"); + await resumeAndVerifyProdQueues(queueApiConfig); + queuesVerified = true; + }, + }); + } catch (originalError) { + return recoverProtocolV3CutoverFailure( + { bookmark, initialPendingMigrations, migrationStarted, originalError, queuesVerified }, + { + commitQueueAcceptance: acceptProtocolV3QueueResume, + pauseAndVerifyQueues: () => pauseAndVerifyProdQueues(queueApiConfig), + printBookmark: printProtocolV3EmergencyBookmark, + probe: probeProtocolV3Cutover, + readBookmark: readStoredProtocolV3Bookmark, + readPendingMigrations: () => readPendingProdMigrations(localMigrationNames), + removeMarker: removeProtocolV3CutoverGate, + resumeAndVerifyQueues: () => resumeAndVerifyProdQueues(queueApiConfig), + write: writeStdout, + }, + ); + } +} + +async function deployProduction(): Promise { + const expectedProdSchema = loadExpectedProdSchema(); + const localMigrationNames = readLocalMigrationNames(); + assertCutoverMigrationJournalAudited(localMigrationNames); + const releaseTreeOid = readCleanReleaseTreeOid(); + + runLocalPreflight(releaseTreeOid); + const smokeConfig = readProtocolV3SmokeConfig(); + const queueApiConfig = readProdQueueApiConfig(); + const initialCutover = probeProtocolV3Cutover(); + if (initialCutover.gatePresent) { + assertProtocolV3CutoverGateExact(); + assertProtocolV3Release(readProtocolV3CutoverState(), releaseTreeOid); + } + + writeStdout("▶ Verifying pre-provisioned production WfP, DNS, route, binding, and TLS"); + const appDeploymentScriptLedgerPresent = readAppDeploymentScriptLedgerPresent(); + const cloudflare = new Cloudflare({ apiToken: queueApiConfig.apiToken }); + const workflowMode = assertProdWfpReadOnlyInfrastructure( + await readProdWfpReadOnlyInventory(cloudflare, queueApiConfig), + prodWfpInfrastructureConfig( + queueApiConfig, + !appDeploymentScriptLedgerPresent, + initialCutover.gatePresent, + ), + Date.now(), + ); + if (workflowMode !== "exact") { + writeStdout( + `▶ App deployment Workflow needs closed-gate ${workflowMode}; forcing exact roll-forward`, + ); + } + + if (shouldProbeProdWfpBeforeMutation(workflowMode)) { + writeStdout("▶ Probing the pre-existing production WfP data path before any mutation"); + await verifyProdWfpGatewayProbe(PROD_APP_DEPLOYMENT_DOMAIN, { + fetch: globalThis.fetch, + probeWebSocket: probeProdWfpWebSocket, + }); + } else { + writeStdout("▶ Deferring the first WfP data-path probe until the closed-gate candidate exists"); + } -writeStdout("▶ Applying pending D1 migrations"); -applyD1Migrations(); + writeStdout("▶ Verifying the dedicated production smoke Agent is published cattle"); + assertProtocolV3SmokeAgent(executeProdD1Json(protocolV3SmokeAgentSql(smokeConfig.agentId))); -writeStdout("▶ Verifying prod D1 schema matches the latest migration snapshot"); -await assertProdSchemaMatchesMigrations(expectedMigrationTables).catch((error: unknown) => { - writeStdout(error instanceof Error ? error.message : String(error)); - process.exit(1); -}); + writeStdout("▶ Acquiring the durable production deploy lease"); + const deployOwner = crypto.randomUUID(); + writeStdout(` production deploy owner: ${deployOwner}`); + try { + acquireProdDeployLease(deployOwner, executeProdDeployLease); + } catch (error) { + writeStdout( + `✗ Production deploy lease ownership was not proven for ${deployOwner}; a timed-out acquisition may still commit, so verify it is quiescent before any exact-owner manual release.`, + ); + throw error; + } + + try { + writeStdout("▶ Proving Workers Scripts Write with the fixed unrouted WfP canary"); + await verifyProdWfpWriteCanary( + createProdWfpWriteCanaryClient(cloudflare, queueApiConfig), + PROD_APP_DISPATCH_NAMESPACE, + PROD_WFP_WRITE_CANARY_SCRIPT, + ); -writeStdout("▶ Ensuring required production queues exist"); -ensureRequiredProdQueues(); + const cutover = probeProtocolV3Cutover(); + const pendingMigrations = readPendingProdMigrations(localMigrationNames); + if (cutover.gatePresent) { + assertProtocolV3CutoverGateExact(); + } -writeStdout("▶ Building driver"); -runVp(["run", "--filter", "agent-driver", "build"]); + writeStdout("▶ Ensuring required production queues exist"); + ensureRequiredProdQueues(); -writeStdout("▶ Deploying worker"); -run(["deploy", "--env", ENV, "--minify", "--containers-rollout", "immediate"]); + if (cutover.gatePresent || pendingMigrations.length > 0 || workflowMode !== "exact") { + await runProtocolV3Cutover( + pendingMigrations, + localMigrationNames, + smokeConfig, + queueApiConfig, + expectedProdSchema, + releaseTreeOid, + workflowMode, + ); + } else { + writeStdout("▶ Verifying prod D1 schema matches the latest migration snapshot"); + assertProdSchemaMatchesSnapshot(expectedProdSchema); + + await deployWorkerAndVerify( + smokeConfig, + queueApiConfig, + false, + releaseTreeOid, + null, + workflowMode, + ); + } + } catch (error) { + writeStdout( + `✗ Retaining production deploy lease ${deployOwner}; verify that every remote mutation is quiescent before an exact-owner manual release.`, + ); + throw error; + } -writeStdout("✓ deploy complete"); + releaseProdDeployLease(); + writeStdout("✓ deploy complete"); +} + +async function main(): Promise { + const args = process.argv.slice(2); + if (args.length === 0) { + return deployProduction(); + } + if (args.length === 1 && args[0] === "--protocol-v3-legacy-inventory") { + writeStdout("▶ Auditing legacy production terminal source identities (read-only)"); + verifyProdLegacyTerminalSourceInventory(); + writeStdout("✓ legacy production terminal sources can be normalized deterministically"); + return; + } + if (args.length === 1 && args[0] === "--protocol-v3-lossy-migration-inventory") { + writeStdout("▶ Auditing migration 0013 production history (read-only)"); + verifyProdLossyMigrationInventory(); + writeStdout("✓ migration 0013 has no lossy production candidates"); + return; + } + throw new Error( + "Usage: bun apps/api/bin/deploy-prod.ts [--protocol-v3-legacy-inventory|--protocol-v3-lossy-migration-inventory]", + ); +} + +if (import.meta.main) { + await main().catch((error: unknown) => { + writeStdout(error instanceof Error ? error.message : String(error)); + process.exitCode = 1; + }); +} diff --git a/apps/api/bin/prod-deploy-lease.ts b/apps/api/bin/prod-deploy-lease.ts new file mode 100644 index 00000000..59d5f80c --- /dev/null +++ b/apps/api/bin/prod-deploy-lease.ts @@ -0,0 +1,79 @@ +import { parseD1JsonResults } from "./d1-json"; + +export const PROD_DEPLOY_LEASE_TABLE = "__production_deploy_lease"; + +export const PROD_DEPLOY_LEASE_TABLE_SQL = `CREATE TABLE "${PROD_DEPLOY_LEASE_TABLE}" ( + "id" integer PRIMARY KEY CHECK ("id" = 1), + "owner" text NOT NULL +)`; + +function ownerSql(owner: string): string { + if (!/^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/u.test(owner)) { + throw new Error("Production deploy lease owner must be a UUID v4."); + } + return `'${owner}'`; +} + +const LEASE_READBACK_SQL = `(SELECT "sql" FROM "sqlite_master" + WHERE "type" = 'table' AND "name" = '${PROD_DEPLOY_LEASE_TABLE}')`; +const LEASE_TRIGGER_COUNT_SQL = `(SELECT count(*) FROM "sqlite_master" + WHERE "type" = 'trigger' AND "tbl_name" = '${PROD_DEPLOY_LEASE_TABLE}' COLLATE BINARY)`; +const LEASE_SCHEMA_PREDICATE_SQL = `${LEASE_READBACK_SQL} = '${PROD_DEPLOY_LEASE_TABLE_SQL.replaceAll("'", "''")}' COLLATE BINARY + AND ${LEASE_TRIGGER_COUNT_SQL} = 0`; +const LEASE_FINAL_READ_SQL = `SELECT + "lease"."owner", + ${LEASE_READBACK_SQL} AS "table_sql", + ${LEASE_TRIGGER_COUNT_SQL} AS "trigger_count" +FROM (SELECT 1) AS "singleton" +LEFT JOIN "${PROD_DEPLOY_LEASE_TABLE}" AS "lease" ON "lease"."id" = 1`; + +export function acquireProdDeployLeaseStatements(owner: string): readonly string[] { + const quotedOwner = ownerSql(owner); + return [ + PROD_DEPLOY_LEASE_TABLE_SQL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS"), + `INSERT INTO "${PROD_DEPLOY_LEASE_TABLE}" ("id", "owner") +SELECT 1, ${quotedOwner} +WHERE ${LEASE_SCHEMA_PREDICATE_SQL} +ON CONFLICT ("id") DO NOTHING`, + LEASE_FINAL_READ_SQL, + ]; +} + +export function verifyProdDeployLeaseStatements(): readonly string[] { + return [LEASE_FINAL_READ_SQL]; +} + +export function releaseProdDeployLeaseStatements(owner: string): readonly string[] { + return [ + `DELETE FROM "${PROD_DEPLOY_LEASE_TABLE}" +WHERE "id" = 1 AND "owner" = ${ownerSql(owner)} + AND ${LEASE_SCHEMA_PREDICATE_SQL}`, + `SELECT "owner" FROM "${PROD_DEPLOY_LEASE_TABLE}" WHERE "id" = 1`, + ]; +} + +export function assertProdDeployLeaseOwned(raw: string, owner: string): void { + ownerSql(owner); + const statements = parseD1JsonResults(raw); + const readback = statements.length === 1 || statements.length === 3 ? statements.at(-1) : null; + const row = readback?.length === 1 ? readback[0] : undefined; + if (row?.table_sql !== PROD_DEPLOY_LEASE_TABLE_SQL) { + throw new Error("Production deploy lease table schema is invalid."); + } + if (row.trigger_count !== 0) { + throw new Error("Production deploy lease table must not have triggers."); + } + if (row.owner !== owner) throw new Error("Production deploy lease is held by another owner."); +} + +function releaseReadback(raw: string) { + const statements = parseD1JsonResults(raw); + return statements.length === 2 ? statements[1] : undefined; +} + +export function assertProdDeployLeaseReleased(raw: string): void { + const remaining = releaseReadback(raw); + if (!remaining || remaining.length !== 0) { + throw new Error("Production deploy lease release did not remove this owner."); + } +} diff --git a/apps/api/bin/prod-schema-guard.ts b/apps/api/bin/prod-schema-guard.ts index 243bacbf..bc63d170 100644 --- a/apps/api/bin/prod-schema-guard.ts +++ b/apps/api/bin/prod-schema-guard.ts @@ -1,97 +1 @@ -/** - * Fail-fast guard against the DEPLOY-D1-001 hazard. - * - * `wrangler d1 migrations apply` records applied migrations by filename. This - * guard detects damage from skipped, rewritten, or incomplete migrations by - * comparing the live database with the latest Drizzle schema snapshot. Applied - * migrations are immutable; every production schema change must use a new file. - * - * These functions are pure (no I/O) so the deploy script can stay thin and the - * detection logic stays unit-testable. Table-level only: this catches a missing - * table — the catastrophic case where every query against it fails — not an - * added column on an existing table. - */ - -interface DrizzleJournal { - entries?: unknown; -} - -interface DrizzleSnapshot { - tables?: unknown; -} - -/** Latest snapshot filename recorded by Drizzle's migration journal. */ -export function getLatestSnapshotFilename(rawJournal: string): string { - const journal = JSON.parse(rawJournal) as DrizzleJournal | null; - const entries: readonly unknown[] = Array.isArray(journal?.entries) ? journal.entries : []; - - if (entries.length === 0) { - throw new Error("Drizzle migration journal has no valid latest entry."); - } - - for (const [expectedIndex, entry] of entries.entries()) { - const index = - typeof entry === "object" && entry !== null && "idx" in entry ? entry.idx : undefined; - - if (index !== expectedIndex) { - throw new Error("Drizzle migration journal indexes must be contiguous from zero."); - } - } - - return `${String(entries.length - 1).padStart(4, "0")}_snapshot.json`; -} - -/** Table names declared by the latest Drizzle schema snapshot. */ -export function parseExpectedTableNames(rawSnapshot: string): string[] { - const snapshot = JSON.parse(rawSnapshot) as DrizzleSnapshot | null; - - if ( - snapshot === null || - typeof snapshot.tables !== "object" || - snapshot.tables === null || - Array.isArray(snapshot.tables) - ) { - throw new Error("Drizzle schema snapshot has no valid tables object."); - } - - const tableNames = Object.keys(snapshot.tables).toSorted(); - - if (tableNames.length === 0) { - throw new Error("Drizzle schema snapshot has no tables."); - } - - return tableNames; -} - -/** Expected tables that are not present in the live database. */ -export function findMissingProdTables( - expectedTableNames: readonly string[], - liveTableNames: readonly string[], -): string[] { - const live = new Set(liveTableNames); - return expectedTableNames.filter((name) => !live.has(name)); -} - -interface D1ExecuteResult { - results?: ReadonlyArray<{ name?: unknown }>; -} - -/** - * Extract `name` rows from `wrangler d1 execute --json` output. The `--json` - * flag prints the result array on stdout; we slice from the first `[` to the - * last `]` to tolerate any leading log lines, then fail loudly on an - * unrecognized shape rather than silently reporting zero tables. - */ -export function extractTableNames(rawStdout: string): string[] { - const start = rawStdout.indexOf("["); - const end = rawStdout.lastIndexOf("]"); - - if (start === -1 || end === -1 || end < start) { - throw new Error(`Could not locate JSON array in \`d1 execute\` output:\n${rawStdout}`); - } - - const parsed = JSON.parse(rawStdout.slice(start, end + 1)) as D1ExecuteResult[]; - const rows = parsed[0]?.results ?? []; - - return rows.map((row) => row.name).filter((name): name is string => typeof name === "string"); -} +export * from "@mosoo/db/deploy-schema-guard"; diff --git a/apps/api/bin/prod-wfp-preflight.ts b/apps/api/bin/prod-wfp-preflight.ts new file mode 100644 index 00000000..59850749 --- /dev/null +++ b/apps/api/bin/prod-wfp-preflight.ts @@ -0,0 +1,622 @@ +import { + APP_DEPLOYMENT_PROBE_HOST_LABEL, + APP_DEPLOYMENT_PROBE_PATH_PREFIX, + APP_DEPLOYMENT_PROBE_SCRIPT_NAME, +} from "../src/modules/apps/application/app-deployment-gateway"; + +export const PROD_APP_DEPLOYMENT_DOMAIN = "apps.mosoo.ai"; +export const PROD_APP_DEPLOYMENT_WORKFLOW = "mosoo-app-deployment-prod"; +export const PROD_APP_DEPLOYMENT_WORKFLOW_BINDING = "APP_DEPLOYMENT_WORKFLOW"; +export const PROD_APP_DEPLOYMENT_WORKFLOW_CLASS = "AppDeploymentWorkflow"; +export const PROD_APP_DISPATCH_NAMESPACE = "mosoo-app-deployments-prod"; +export const PROD_APP_WILDCARD_DNS = `*.${PROD_APP_DEPLOYMENT_DOMAIN}`; +export const PROD_APP_WILDCARD_ROUTE = `${PROD_APP_WILDCARD_DNS}/*`; +export const PROD_API_WORKER_NAME = "mosoo-api-prod"; +export const PROD_WFP_WRITE_CANARY_SCRIPT = "mosoo-release-write-canary"; +export const WRANGLER_DISABLE_AUTO_PROVISION = "--experimental-provision=false"; + +const CERTIFICATE_MINIMUM_VALIDITY_MS = 30 * 24 * 60 * 60 * 1_000; +const PROBE_TIMEOUT_MS = 10_000; + +export interface WfpNamespaceInventory { + id: string | null; + name: string | null; + scriptCount: number | null; + trustedWorkers: boolean | null; +} + +export interface WfpZoneInventory { + accountId: string | null; + id: string | null; + name: string | null; + status: string | null; +} + +export interface WfpDnsRecordInventory { + name: string | null; + proxied: boolean | null; + type: string | null; +} + +export interface WfpCertificateInventory { + certificates: readonly { + expiresOn: string | null; + hosts: readonly string[]; + status: string | null; + }[]; + hosts: readonly string[]; + status: string | null; +} + +export interface WfpWorkflowInventory { + className: string | null; + id: string | null; + name: string | null; + scriptName: string | null; +} + +export interface WfpScriptInventory { + id: string | null; + namespace: string | null; +} + +export interface WfpWorkerBindingInventory { + className: string | null; + name: string | null; + namespace: string | null; + scriptName: string | null; + type: string | null; + workflowName: string | null; +} + +export type ProdWfpWorkflowMode = "bootstrap" | "exact" | "repair"; + +export interface ProdWfpReadOnlyInventory { + apiWorkerBindings: readonly WfpWorkerBindingInventory[]; + certificates: readonly WfpCertificateInventory[]; + dnsRecords: readonly WfpDnsRecordInventory[]; + legacyResources: readonly string[]; + namespace: WfpNamespaceInventory; + probeScript: WfpScriptInventory; + routes: readonly { + pattern: string | null; + script: string | null; + }[]; + workflow: WfpWorkflowInventory; + zone: WfpZoneInventory; +} + +export interface ProdWfpInfrastructureConfig { + accountId: string; + apiWorkerName: string; + dispatchBinding: string; + dispatchNamespace: string; + requireProbeOnlyNamespace: boolean; + wildcardDns: string; + wildcardRoute: string; + workflowBinding: string; + workflowClass: string; + workflowName: string; + workflowRepairAllowed: boolean; + zoneId: string; + zoneName: string; +} + +export interface WfpWriteCanaryClient { + deleteScript(scriptName: string): Promise; + readScript(scriptName: string): Promise<{ id: string | null; namespace: string | null } | null>; + uploadScript(scriptName: string): Promise<{ id: string | null }>; +} + +export function apiWorkerDryRunArgs(): string[] { + return [ + "deploy", + "--env", + "prod", + "--minify", + "--strict", + "--dry-run", + WRANGLER_DISABLE_AUTO_PROVISION, + ]; +} + +export function apiWorkerDeployArgs(releaseTag: string): string[] { + return [ + "deploy", + "--env", + "prod", + "--minify", + "--strict", + WRANGLER_DISABLE_AUTO_PROVISION, + "--containers-rollout", + "immediate", + "--tag", + releaseTag, + ]; +} + +export function assertProdWfpReadOnlyInfrastructure( + inventory: ProdWfpReadOnlyInventory, + config: ProdWfpInfrastructureConfig, + nowMs: number, +): ProdWfpWorkflowMode { + if ( + inventory.zone.id !== config.zoneId || + inventory.zone.name !== config.zoneName || + inventory.zone.status !== "active" || + inventory.zone.accountId !== config.accountId + ) { + throw new Error("Cloudflare production zone identity or status is not exact."); + } + + assertProdWfpNamespaceExact( + inventory.namespace, + config.dispatchNamespace, + config.requireProbeOnlyNamespace, + ); + assertProdWfpProbeScriptExact(inventory.probeScript, config.dispatchNamespace); + const workflowMode = assertProdWfpWorkflowBoundary( + inventory.workflow, + inventory.apiWorkerBindings, + config, + ); + + const wildcardRoutes = inventory.routes.filter((route) => route.pattern === config.wildcardRoute); + if (wildcardRoutes.length !== 1 || wildcardRoutes[0]?.script !== config.apiWorkerName) { + throw new Error( + `Cloudflare must route ${config.wildcardRoute} exactly once to ${config.apiWorkerName}.`, + ); + } + const deploymentDomain = readWildcardDomain(config.wildcardDns); + const shadowingRoutes = inventory.routes.filter( + (route) => + route.pattern !== config.wildcardRoute && + routeMayMatchDeploymentDomain(route.pattern, deploymentDomain), + ); + if (shadowingRoutes.length > 0) { + throw new Error( + `Cloudflare routes must not shadow ${config.wildcardRoute}: ${shadowingRoutes + .map((route) => `${String(route.pattern)} -> ${route.script ?? "no script"}`) + .join(", ")}.`, + ); + } + + const dispatcherBindings = inventory.apiWorkerBindings.filter( + (binding) => binding.name === config.dispatchBinding, + ); + if ( + dispatcherBindings.length !== 1 || + dispatcherBindings[0]?.type !== "dispatch_namespace" || + dispatcherBindings[0]?.namespace !== config.dispatchNamespace + ) { + throw new Error( + `${config.apiWorkerName} must bind ${config.dispatchBinding} to the exact ${config.dispatchNamespace} dispatch namespace.`, + ); + } + + if (inventory.dnsRecords.some((record) => record.name === null)) { + throw new Error("Cloudflare returned a DNS record without a hostname."); + } + const shadowingRecords = inventory.dnsRecords.filter( + (record) => + record.name !== config.wildcardDns && + record.name !== null && + normalizeDnsName(record.name).endsWith(`.${deploymentDomain}`), + ); + if (shadowingRecords.length > 0) { + throw new Error( + `Cloudflare DNS records must not shadow ${config.wildcardDns}: ${shadowingRecords + .map((record) => `${String(record.type)} ${String(record.name)}`) + .join(", ")}.`, + ); + } + const trafficRecords = inventory.dnsRecords.filter( + (record) => + record.name === config.wildcardDns && + (record.type === "A" || record.type === "AAAA" || record.type === "CNAME"), + ); + if (trafficRecords.length === 0 || trafficRecords.some((record) => record.proxied !== true)) { + throw new Error( + `Cloudflare DNS must contain only proxied traffic records for ${config.wildcardDns}.`, + ); + } + + const validAfter = nowMs + CERTIFICATE_MINIMUM_VALIDITY_MS; + const hasExactCertificate = inventory.certificates.some( + (pack) => + pack.status === "active" && + pack.hosts.includes(config.wildcardDns) && + pack.certificates.some( + (certificate) => + certificate.status === "active" && + certificate.hosts.includes(config.wildcardDns) && + parseExpiry(certificate.expiresOn) >= validAfter, + ), + ); + if (!hasExactCertificate) { + throw new Error( + `Cloudflare requires an active certificate for ${config.wildcardDns} valid for at least 30 more days.`, + ); + } + + if (inventory.legacyResources.length > 0) { + throw new Error( + `Legacy Mosoo Pages/classic Worker resources must be removed manually before WfP cutover:\n${inventory.legacyResources.map((resource) => `- ${resource}`).join("\n")}`, + ); + } + + return workflowMode; +} + +function readWildcardDomain(wildcardDns: string): string { + if (!wildcardDns.startsWith("*.") || wildcardDns.length === 2) { + throw new Error("Cloudflare WfP wildcard DNS configuration is invalid."); + } + return wildcardDns.slice(2).toLowerCase(); +} + +function normalizeDnsName(name: string): string { + return name.toLowerCase().replace(/\.$/u, ""); +} + +function routeMayMatchDeploymentDomain(pattern: string | null, domain: string): boolean { + if (pattern === null) return true; + const withoutScheme = pattern.replace(/^https?:\/\//iu, ""); + const slash = withoutScheme.indexOf("/"); + const hostname = (slash === -1 ? withoutScheme : withoutScheme.slice(0, slash)) + .toLowerCase() + .replace(/\.$/u, ""); + if (hostname.length === 0 || /[:?#\s]/u.test(hostname)) return true; + if (hostname.endsWith(`.${domain}`)) return true; + return hostname.startsWith("*") && `probe.${domain}`.endsWith(hostname.slice(1)); +} + +export function assertProdWfpNamespaceExact( + namespace: WfpNamespaceInventory, + expectedName: string, + requireProbeOnly: boolean, +): void { + if ( + namespace.id === null || + namespace.id.length === 0 || + namespace.name !== expectedName || + namespace.scriptCount === null || + !Number.isSafeInteger(namespace.scriptCount) || + namespace.scriptCount < 1 || + namespace.trustedWorkers !== false + ) { + throw new Error( + `Workers for Platforms namespace ${expectedName} is missing, mismatched, or trusted. Provision the exact untrusted namespace outside this deploy.`, + ); + } + + if (requireProbeOnly && namespace.scriptCount !== 1) { + throw new Error( + `Bootstrap requires ${expectedName} to contain only the pre-provisioned ${APP_DEPLOYMENT_PROBE_SCRIPT_NAME} script; found ${namespace.scriptCount} scripts. Inspect and clean it manually before retrying.`, + ); + } +} + +export function assertProdWfpProbeScriptExact( + script: WfpScriptInventory, + expectedNamespace: string, +): void { + if (script.id !== APP_DEPLOYMENT_PROBE_SCRIPT_NAME || script.namespace !== expectedNamespace) { + throw new Error( + `Cloudflare requires the pre-provisioned ${APP_DEPLOYMENT_PROBE_SCRIPT_NAME} operational probe in the exact ${expectedNamespace} namespace.`, + ); + } +} + +export function assertProdWfpWorkflowExact( + workflow: WfpWorkflowInventory, + expectedName: string, + expectedClass: string, + expectedScript: string, +): void { + if ( + workflow.id === null || + workflow.id.length === 0 || + workflow.name !== expectedName || + workflow.className !== expectedClass || + workflow.scriptName !== expectedScript + ) { + throw new Error( + `Cloudflare Workflow ${expectedName} is missing or mismatched. Bootstrap the exact ${expectedClass} Workflow on ${expectedScript} outside this deploy.`, + ); + } +} + +export function assertProdWfpWorkflowBoundary( + workflow: WfpWorkflowInventory, + bindings: readonly WfpWorkerBindingInventory[], + expected: Pick< + ProdWfpInfrastructureConfig, + "apiWorkerName" | "workflowBinding" | "workflowClass" | "workflowName" | "workflowRepairAllowed" + >, +): ProdWfpWorkflowMode { + const workflowBindings = bindings.filter((binding) => binding.type === "workflow"); + const workflowAbsent = + workflow.className === null && + workflow.id === null && + workflow.name === null && + workflow.scriptName === null; + + if (workflowAbsent && workflowBindings.length === 0) { + return "bootstrap"; + } + + const workflowExact = + workflow.id !== null && + workflow.id.length > 0 && + workflow.name === expected.workflowName && + workflow.className === expected.workflowClass && + workflow.scriptName === expected.apiWorkerName; + if (workflowExact && workflowBindings.length === 0 && expected.workflowRepairAllowed) { + return "repair"; + } + + assertProdWfpWorkflowExact( + workflow, + expected.workflowName, + expected.workflowClass, + expected.apiWorkerName, + ); + + const exactBindings = workflowBindings.filter( + (binding) => binding.name === expected.workflowBinding, + ); + if ( + workflowBindings.length !== 1 || + exactBindings.length !== 1 || + exactBindings[0]?.workflowName !== expected.workflowName || + exactBindings[0]?.className !== expected.workflowClass || + exactBindings[0]?.scriptName !== expected.apiWorkerName + ) { + throw new Error( + `${expected.apiWorkerName} must bind ${expected.workflowBinding} to the exact ${expected.workflowName} Workflow class ${expected.workflowClass}.`, + ); + } + + return "exact"; +} + +export function assertProdWfpWorkflowDeployGate( + mode: ProdWfpWorkflowMode, + closedGate: boolean, +): void { + if (mode !== "exact" && !closedGate) { + throw new Error( + "App deployment Workflow bootstrap or repair requires the closed production gate.", + ); + } +} + +export function shouldProbeProdWfpBeforeMutation(mode: ProdWfpWorkflowMode): boolean { + return mode === "exact"; +} + +function parseExpiry(value: string | null): number { + if (value === null) return Number.NaN; + return Date.parse(value); +} + +export async function verifyProdWfpWriteCanary( + client: WfpWriteCanaryClient, + dispatchNamespace: string, + scriptName: string, +): Promise { + try { + await client.deleteScript(scriptName); + if ((await client.readScript(scriptName)) !== null) { + throw new Error("The previous WfP write canary still exists after deletion."); + } + } catch (error) { + throw canaryFailure(scriptName, dispatchNamespace, [error]); + } + + let verificationError: unknown = null; + + try { + const uploaded = await client.uploadScript(scriptName); + if (uploaded.id !== scriptName) { + throw new Error( + `WfP write canary uploaded as ${String(uploaded.id)}, expected ${scriptName}.`, + ); + } + + const stored = await client.readScript(scriptName); + if (stored?.id !== scriptName || stored.namespace !== dispatchNamespace) { + throw new Error("WfP write canary readback did not match its exact script and namespace."); + } + } catch (error) { + verificationError = error; + } + + let cleanupError: unknown = null; + try { + await client.deleteScript(scriptName); + if ((await client.readScript(scriptName)) !== null) { + throw new Error("WfP write canary still exists after deletion."); + } + } catch (error) { + cleanupError = error; + } + + if (verificationError !== null || cleanupError !== null) { + throw canaryFailure(scriptName, dispatchNamespace, [ + ...(verificationError === null ? [] : [verificationError]), + ...(cleanupError === null ? [] : [cleanupError]), + ]); + } +} + +function canaryFailure( + scriptName: string, + dispatchNamespace: string, + errors: readonly unknown[], +): AggregateError { + const endpoint = + `https://api.cloudflare.com/client/v4/accounts/$CLOUDFLARE_ACCOUNT_ID/workers/dispatch/` + + `namespaces/${dispatchNamespace}/scripts/${scriptName}`; + return new AggregateError( + errors, + `WfP write canary ${scriptName} failed. Remove it manually with: curl -X DELETE '${endpoint}' -H 'Authorization: Bearer $CLOUDFLARE_API_TOKEN'`, + ); +} + +interface WfpGatewayProbeDependencies { + fetch: typeof fetch; + probeWebSocket(url: string, nonce: string): Promise; +} + +export async function verifyProdWfpGatewayProbe( + appDomain: string, + dependencies: WfpGatewayProbeDependencies, + createNonce: () => string = () => crypto.randomUUID(), +): Promise { + const hostname = `${APP_DEPLOYMENT_PROBE_HOST_LABEL}.${appDomain}`; + const nonce = createNonce(); + + await verifyHttpBodyProbe(hostname, nonce, dependencies.fetch); + await verifyStreamProbe(hostname, nonce, dependencies.fetch); + await verifyCancellationProbe(hostname, nonce, dependencies.fetch); + await verifyHttpBodyProbe(hostname, `${nonce}-after-cancel`, dependencies.fetch); + await dependencies.probeWebSocket( + `wss://${hostname}${APP_DEPLOYMENT_PROBE_PATH_PREFIX}websocket?nonce=${encodeURIComponent(nonce)}`, + nonce, + ); +} + +async function verifyHttpBodyProbe( + hostname: string, + nonce: string, + fetcher: typeof fetch, +): Promise { + const pathname = `${APP_DEPLOYMENT_PROBE_PATH_PREFIX}http`; + const body = `body:${nonce}`; + const response = await fetcher( + `https://${hostname}${pathname}?nonce=${encodeURIComponent(nonce)}`, + { + body, + headers: { "X-Mosoo-Wfp-Probe-Nonce": nonce }, + method: "POST", + signal: AbortSignal.timeout(PROBE_TIMEOUT_MS), + }, + ); + if (!response.ok || response.headers.get("Cache-Control") !== "no-store") { + throw new Error("WfP HTTP/body probe did not return an uncached success response."); + } + + const actual: unknown = await response.json(); + if ( + !isJsonRecord(actual) || + actual["body"] !== body || + actual["hostname"] !== hostname || + actual["method"] !== "POST" || + actual["nonce"] !== nonce || + actual["pathname"] !== pathname || + actual["search"] !== `?nonce=${encodeURIComponent(nonce)}` + ) { + throw new Error("WfP HTTP/body probe did not preserve the original request."); + } +} + +function isJsonRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +async function verifyStreamProbe( + hostname: string, + nonce: string, + fetcher: typeof fetch, +): Promise { + const response = await fetcher( + `https://${hostname}${APP_DEPLOYMENT_PROBE_PATH_PREFIX}stream?nonce=${encodeURIComponent(nonce)}`, + { signal: AbortSignal.timeout(PROBE_TIMEOUT_MS) }, + ); + const reader = response.body?.getReader(); + if (!response.ok || reader === undefined) { + throw new Error("WfP stream probe did not return a readable success response."); + } + + const decoder = new TextDecoder(); + const first = await reader.read(); + if (first.done || decoder.decode(first.value) !== `start:${nonce}\n`) { + throw new Error("WfP stream probe buffered or changed its first chunk."); + } + + let remainder = ""; + while (true) { + const part = await reader.read(); + if (part.done) break; + remainder += decoder.decode(part.value, { stream: true }); + } + remainder += decoder.decode(); + if (remainder !== `end:${nonce}\n`) { + throw new Error("WfP stream probe changed its final chunk."); + } +} + +async function verifyCancellationProbe( + hostname: string, + nonce: string, + fetcher: typeof fetch, +): Promise { + const response = await fetcher( + `https://${hostname}${APP_DEPLOYMENT_PROBE_PATH_PREFIX}cancel?nonce=${encodeURIComponent(nonce)}`, + { signal: AbortSignal.timeout(PROBE_TIMEOUT_MS) }, + ); + const reader = response.body?.getReader(); + if (!response.ok || reader === undefined) { + throw new Error("WfP cancellation probe did not return a readable success response."); + } + + const first = await reader.read(); + if (first.done || new TextDecoder().decode(first.value) !== `start:${nonce}\n`) { + throw new Error("WfP cancellation probe did not stream its first chunk."); + } + await reader.cancel(`cancel:${nonce}`); +} + +export function probeProdWfpWebSocket( + url: string, + nonce: string, + createSocket: (url: string) => WebSocket = (socketUrl) => new WebSocket(socketUrl), +): Promise { + return new Promise((resolve, reject) => { + const socket = createSocket(url); + let settled = false; + const timeout = setTimeout(() => { + socket.close(); + finish(new Error("WfP WebSocket probe timed out.")); + }, PROBE_TIMEOUT_MS); + const finish = (error?: Error): void => { + if (settled) return; + settled = true; + clearTimeout(timeout); + if (error === undefined) resolve(); + else reject(error); + }; + + socket.addEventListener("open", () => { + socket.send(`echo:${nonce}`); + }); + socket.addEventListener("message", (event) => { + if (event.data !== `echo:${nonce}`) { + socket.close(); + finish(new Error("WfP WebSocket probe changed the echo payload.")); + } + }); + socket.addEventListener("close", (event) => { + if (event.code !== 1000 || event.reason !== `verified:${nonce}`) { + finish(new Error("WfP WebSocket probe did not close normally.")); + return; + } + finish(); + }); + socket.addEventListener("error", () => { + finish(new Error("WfP WebSocket probe failed.")); + }); + }); +} diff --git a/apps/api/bin/protocol-v3-cutover.ts b/apps/api/bin/protocol-v3-cutover.ts new file mode 100644 index 00000000..37eea6ea --- /dev/null +++ b/apps/api/bin/protocol-v3-cutover.ts @@ -0,0 +1,3138 @@ +import type { D1JsonRow } from "./d1-json"; +import { parseD1JsonResults, requireSingleD1Row } from "./d1-json"; +import { PROD_DEPLOY_LEASE_TABLE } from "./prod-deploy-lease"; +import { + MANAGED_PROD_SCHEMA_TRIGGERS, + PROTOCOL_V3_MIGRATION_INTENT_TABLE, + PROTOCOL_V3_MIGRATION_INTENT_TABLE_SQL, +} from "./prod-schema-guard"; + +export { PROTOCOL_V3_MIGRATION_INTENT_TABLE } from "./prod-schema-guard"; + +export const PROTOCOL_V3_MIGRATION = "0013_durable-mcp-effect-v3.sql"; +export const PROTOCOL_V3_SESSION_EVENT_MIGRATION = "0014_session-event-stream-identity.sql"; +export const PROTOCOL_V3_SESSION_CLEANUP_MIGRATION = "0015_session-cleanup-operation.sql"; +export const PROTOCOL_V3_RUNTIME_AUTHORITY_MIGRATION = + "0019_runtime-subject-operation-authority.sql"; +export const LAST_CUTOVER_AUDITED_MIGRATION = "0020_sandbox-backup-object-authority.sql"; +export const AUDITED_MIGRATION_NAMES = [ + "0000_baseline.sql", + "0001_bound-capability-run-provenance.sql", + "0002_bound-agent-call-idempotency.sql", + "0003_collapse-sandbox-error-state.sql", + "0004_usage-rollup-receipt.sql", + "0005_public-thread-end-user.sql", + "0006_runtime_subject_quota_scope.sql", + "0007_app-deployment-secrets.sql", + "0008_public-thread-tool-call-identity.sql", + "0009_terminal_event_lookup_index.sql", + "0010_external-tool-effects.sql", + "0011_cattle-terminal-checkpoints.sql", + "0012_agent-task-snapshot-state.sql", + PROTOCOL_V3_MIGRATION, + PROTOCOL_V3_SESSION_EVENT_MIGRATION, + PROTOCOL_V3_SESSION_CLEANUP_MIGRATION, + "0016_durable-event-side-effects.sql", + "0017_terminal-reconciliation-scheduling.sql", + "0018_runtime-operation-ready-authority.sql", + PROTOCOL_V3_RUNTIME_AUTHORITY_MIGRATION, + LAST_CUTOVER_AUDITED_MIGRATION, +] as const; +export const PROTOCOL_V3_CUTOVER_TABLE = "__protocol_v3_cutover"; +export const PROD_APPLIED_MIGRATIONS_SQL = `SELECT "name" FROM "d1_migrations" ORDER BY "id";`; + +const GIT_TREE_OID_PATTERN = /^(?:[0-9a-f]{40}|[0-9a-f]{64})$/u; +const SHA256_DIGEST_PATTERN = /^[0-9a-f]{64}$/u; +const WORKER_VERSION_ID_PATTERN = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/u; + +export type ProtocolV3CutoverPhase = "draining" | "queues_resuming"; + +const ACTIVE_RUN_STATUSES = "'queued', 'booting', 'running', 'waiting_input'"; +const ACTIVE_APP_DEPLOYMENT_RUN_STATUSES = + "'queued', 'preparing', 'building', 'submitting', 'submitted', 'activating'"; +const CUTOVER_BLOCKED_API_COMMAND_KINDS = + "'session_run_dispatch', 'app_deployment_run_dispatch', 'environment_package_artifact_build'"; + +export const PROTOCOL_V3_CUTOVER_QUEUE_NAMES = [ + "api-command", + "api-command-dlq", + "environment-artifact-build", +] as const; + +export interface ProtocolV3QueueDeliveryControl { + list(): Promise; + mutate(queueName: string, action: "pause" | "resume"): Promise | void; + read(queueId: string): Promise<{ + readonly deliveryPaused: boolean | undefined; + readonly name: string; + }>; +} + +export async function updateAndVerifyProtocolV3QueueDelivery( + control: ProtocolV3QueueDeliveryControl, + action: "pause" | "resume", +): Promise { + const queueIds = new Map((await control.list()).map(({ id, name }) => [name, id])); + const failures: unknown[] = []; + + for (const queueName of PROTOCOL_V3_CUTOVER_QUEUE_NAMES) { + const queueId = queueIds.get(queueName); + if (queueId === undefined) { + failures.push(new Error(`Production queue ${queueName} was not found.`)); + continue; + } + + try { + await control.mutate(queueName, action); + const queue = await control.read(queueId); + if (queue.name !== queueName || queue.deliveryPaused !== (action === "pause")) { + throw new Error(`Production queue ${queueName} delivery did not ${action}.`); + } + } catch (error) { + failures.push(error); + } + } + + if (failures.length > 0) { + throw new AggregateError(failures, `Production queue ${action} and readback failed.`); + } +} +const LIVE_DRIVER_STATUSES = "'provisioning', 'connecting', 'ready', 'stopping'"; + +function exactSmokeSessionSql(sessionId: string, sandboxId?: string): string { + const sandboxJoin = + sandboxId === undefined + ? "" + : ` + INNER JOIN "sandbox" AS "smoke_sandbox" + ON "smoke_sandbox"."id" = ${sandboxId} + AND "smoke_sandbox"."subject_kind" = 'session' + AND "smoke_sandbox"."subject_id" = "smoke_session"."id"`; + return `EXISTS ( + SELECT 1 + FROM "${PROTOCOL_V3_CUTOVER_TABLE}" AS "gate" + INNER JOIN "session" AS "smoke_session" + ON "smoke_session"."id" = ${sessionId} + AND "smoke_session"."creator_account_id" = "gate"."smoke_account_id" + AND "smoke_session"."end_user_id" = "gate"."smoke_request_key"${sandboxJoin} + WHERE "gate"."enabled" = 1 + AND "gate"."smoke_request_key" IS NOT NULL + AND ("gate"."smoke_session_id" IS NULL OR "gate"."smoke_session_id" = "smoke_session"."id") + )`; +} + +const EXACT_NEW_SESSION_SMOKE_SQL = `EXISTS ( + SELECT 1 + FROM "${PROTOCOL_V3_CUTOVER_TABLE}" AS "gate" + WHERE "gate"."enabled" = 1 + AND "gate"."smoke_request_key" IS NOT NULL + AND "gate"."smoke_session_id" IS NULL + AND "gate"."smoke_account_id" = NEW."creator_account_id" + AND "gate"."smoke_request_key" = NEW."end_user_id" + )`; + +function postMigrationSessionStaticSql(row: string): string { + return `${row}."status" IN ('IDLE', 'TERMINATED') + AND ${row}."status_operation_id" IS NULL + AND (${row}."cleanup_operation_kind" IS NULL + OR (${row}."cleanup_operation_kind" = 'archive' + AND ${row}."status" = 'IDLE' + AND ${row}."archived_at" IS NOT NULL)) + AND ${row}."runtime_provisioning_operation_id" IS NULL + AND ${row}."runtime_provisioning_run_id" IS NULL + AND ${row}."runtime_provisioning_sandbox_id" IS NULL + AND ${row}."runtime_provisioning_sandbox_session_id" IS NULL + AND ${row}."runtime_provisioning_sandbox_incarnation" IS NULL + AND ${row}."runtime_provisioning_heartbeat_at" IS NULL`; +} + +interface ProtocolV3CutoverObject { + readonly name: string; + readonly sql: string; + readonly tableName: string; + readonly type: "table" | "trigger"; +} + +const PROTOCOL_V3_CUTOVER_OBJECTS: readonly ProtocolV3CutoverObject[] = [ + { + name: PROTOCOL_V3_CUTOVER_TABLE, + sql: `CREATE TABLE "${PROTOCOL_V3_CUTOVER_TABLE}" ( + "id" integer PRIMARY KEY CHECK ("id" = 1), + "command_freeze" integer NOT NULL DEFAULT 0 CHECK ("command_freeze" IN (0, 1)), + "enabled" integer NOT NULL DEFAULT 1 CHECK ("enabled" IN (0, 1)), + "phase" text NOT NULL DEFAULT 'draining' CHECK ("phase" IN ('draining', 'queues_resuming')), + "pre_migration_bookmark" text, + "release_tree_oid" text NOT NULL CHECK ((length("release_tree_oid") = 40 OR length("release_tree_oid") = 64) AND "release_tree_oid" = lower("release_tree_oid") AND "release_tree_oid" NOT GLOB '*[^0-9a-f]*'), + "smoke_account_id" text, + "smoke_request_key" text, + "smoke_session_id" text, + "target_container_application_version" integer, + "target_container_image_digest" text, + "target_worker_version_id" text, + "started_at" integer NOT NULL DEFAULT (CAST(unixepoch('subsec') * 1000 AS INTEGER)), + CONSTRAINT "protocol_v3_cutover_phase_check" CHECK ( + ("phase" = 'draining' AND "enabled" = 1) + OR ("phase" = 'queues_resuming' AND "command_freeze" = 1) + ), + CONSTRAINT "protocol_v3_cutover_rollout_check" CHECK ( + ("target_container_application_version" IS NULL AND "target_container_image_digest" IS NULL AND "target_worker_version_id" IS NULL) + OR ("target_container_application_version" >= 0 AND length("target_container_image_digest") = 64 AND "target_container_image_digest" = lower("target_container_image_digest") AND "target_container_image_digest" NOT GLOB '*[^0-9a-f]*' AND length(trim("target_worker_version_id")) > 0) + ) +)`.trim(), + tableName: PROTOCOL_V3_CUTOVER_TABLE, + type: "table", + }, + { + name: "__protocol_v3_cutover_session_run_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_session_run_insert" +BEFORE INSERT ON "session_run" +WHEN NEW."status" IN (${ACTIVE_RUN_STATUSES}) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT ${exactSmokeSessionSql('NEW."session_id"')} +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new active Session Runs'); +END`, + tableName: "session_run", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_session_run_update", + sql: `CREATE TRIGGER "__protocol_v3_cutover_session_run_update" +BEFORE UPDATE OF "status" ON "session_run" +WHEN NEW."status" IN (${ACTIVE_RUN_STATUSES}) + AND OLD."status" NOT IN (${ACTIVE_RUN_STATUSES}) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT ${exactSmokeSessionSql('NEW."session_id"')} +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks Session Run reactivation'); +END`, + tableName: "session_run", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_app_deployment_run_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_app_deployment_run_insert" +BEFORE INSERT ON "app_deployment_run" +WHEN NEW."status" IN (${ACTIVE_APP_DEPLOYMENT_RUN_STATUSES}) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new active App deployment Runs'); +END`, + tableName: "app_deployment_run", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_app_deployment_run_update", + sql: `CREATE TRIGGER "__protocol_v3_cutover_app_deployment_run_update" +BEFORE UPDATE OF "status" ON "app_deployment_run" +WHEN NEW."status" IN (${ACTIVE_APP_DEPLOYMENT_RUN_STATUSES}) + AND OLD."status" NOT IN (${ACTIVE_APP_DEPLOYMENT_RUN_STATUSES}) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks App deployment Run reactivation'); +END`, + tableName: "app_deployment_run", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_driver_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_driver_insert" +BEFORE INSERT ON "driver_instance" +WHEN NEW."status" IN (${LIVE_DRIVER_STATUSES}) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT ${exactSmokeSessionSql('NEW."sandbox_session_id"', 'NEW."sandbox_id"')} +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new live Driver instances'); +END`, + tableName: "driver_instance", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_driver_update", + sql: `CREATE TRIGGER "__protocol_v3_cutover_driver_update" +BEFORE UPDATE OF "status" ON "driver_instance" +WHEN NEW."status" IN (${LIVE_DRIVER_STATUSES}) + AND OLD."status" NOT IN (${LIVE_DRIVER_STATUSES}) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT ${exactSmokeSessionSql('NEW."sandbox_session_id"', 'NEW."sandbox_id"')} +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks Driver reactivation'); +END`, + tableName: "driver_instance", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_command_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_command_insert" +BEFORE INSERT ON "driver_command" +WHEN EXISTS ( + SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" + WHERE "enabled" = 1 + AND ("command_freeze" = 1 OR NEW."kind" IN ('input.start', 'mcp.execute')) + ) + AND NOT ( + NEW."kind" = 'session.stop' + AND EXISTS ( + SELECT 1 + FROM "driver_instance" AS "smoke_driver" + WHERE "smoke_driver"."id" = NEW."driver_instance_id" + AND ${exactSmokeSessionSql('"smoke_driver"."sandbox_session_id"', '"smoke_driver"."sandbox_id"')} + ) + ) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new Driver commands'); +END`, + tableName: "driver_command", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_sandbox_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_sandbox_insert" +BEFORE INSERT ON "sandbox" +WHEN (NEW."status" <> 'cold' + OR NEW."status_operation_id" IS NOT NULL + OR NEW."claim_owner" IS NOT NULL + OR NEW."claim_expires_at" IS NOT NULL) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT (NEW."subject_kind" = 'session' AND ${exactSmokeSessionSql('NEW."subject_id"')}) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new active sandboxes'); +END`, + tableName: "sandbox", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_sandbox_update", + sql: `CREATE TRIGGER "__protocol_v3_cutover_sandbox_update" +BEFORE UPDATE OF "status", "status_operation_id", "claim_owner", "claim_expires_at" ON "sandbox" +WHEN OLD."status" = 'cold' + AND OLD."status_operation_id" IS NULL + AND OLD."claim_owner" IS NULL + AND OLD."claim_expires_at" IS NULL + AND (NEW."status" <> 'cold' + OR NEW."status_operation_id" IS NOT NULL + OR NEW."claim_owner" IS NOT NULL + OR NEW."claim_expires_at" IS NOT NULL) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT (NEW."subject_kind" = 'session' AND ${exactSmokeSessionSql('NEW."subject_id"')}) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks sandbox activation'); +END`, + tableName: "sandbox", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_sandbox_session_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_sandbox_session_insert" +BEFORE INSERT ON "sandbox_session" +WHEN NEW."status" NOT IN ('closed', 'error') + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT ${exactSmokeSessionSql('NEW."session_id"', 'NEW."sandbox_id"')} +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new active sandbox Sessions'); +END`, + tableName: "sandbox_session", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_sandbox_session_update", + sql: `CREATE TRIGGER "__protocol_v3_cutover_sandbox_session_update" +BEFORE UPDATE OF "status" ON "sandbox_session" +WHEN OLD."status" IN ('closed', 'error') + AND NEW."status" NOT IN ('closed', 'error') + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT ${exactSmokeSessionSql('NEW."session_id"', 'NEW."sandbox_id"')} +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks sandbox Session reactivation'); +END`, + tableName: "sandbox_session", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_sandbox_backup_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_sandbox_backup_insert" +BEFORE INSERT ON "sandbox_backup" +WHEN NEW."status" NOT IN ('ready', 'pruned') + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new sandbox backup work'); +END`, + tableName: "sandbox_backup", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_sandbox_backup_update", + sql: `CREATE TRIGGER "__protocol_v3_cutover_sandbox_backup_update" +BEFORE UPDATE OF "status" ON "sandbox_backup" +WHEN OLD."status" IN ('ready', 'pruned') + AND NEW."status" NOT IN ('ready', 'pruned') + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks sandbox backup reactivation'); +END`, + tableName: "sandbox_backup", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_session_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_session_insert" +BEFORE INSERT ON "session" +WHEN (NEW."status" NOT IN ('IDLE', 'TERMINATED') OR NEW."status_operation_id" IS NOT NULL) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT ${EXACT_NEW_SESSION_SMOKE_SQL} +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new Session operations'); +END`, + tableName: "session", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_session_update", + sql: `CREATE TRIGGER "__protocol_v3_cutover_session_update" +BEFORE UPDATE OF "status", "status_operation_id" ON "session" +WHEN OLD."status" IN ('IDLE', 'TERMINATED') + AND OLD."status_operation_id" IS NULL + AND (NEW."status" NOT IN ('IDLE', 'TERMINATED') OR NEW."status_operation_id" IS NOT NULL) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT ${exactSmokeSessionSql('NEW."id"')} +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks Session operation acquisition'); +END`, + tableName: "session", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_api_command_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_api_command_insert" +BEFORE INSERT ON "api_command" +WHEN NEW."status" IN ('queued', 'running') + AND EXISTS ( + SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" + WHERE "enabled" = 1 + AND ("command_freeze" = 1 OR NEW."kind" IN (${CUTOVER_BLOCKED_API_COMMAND_KINDS})) + ) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new nonterminal API commands'); +END`, + tableName: "api_command", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_api_command_update", + sql: `CREATE TRIGGER "__protocol_v3_cutover_api_command_update" +BEFORE UPDATE OF "kind", "status", "claim_owner", "claim_expires_at" ON "api_command" +WHEN NEW."status" IN ('queued', 'running') + AND (OLD."status" NOT IN ('queued', 'running') OR NEW."kind" IS NOT OLD."kind") + AND EXISTS ( + SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" + WHERE "enabled" = 1 + AND ("command_freeze" = 1 OR NEW."kind" IN (${CUTOVER_BLOCKED_API_COMMAND_KINDS})) + ) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks API command admission'); +END`, + tableName: "api_command", + type: "trigger", + }, +]; + +export const PROTOCOL_V3_CUTOVER_OBJECT_COUNT = PROTOCOL_V3_CUTOVER_OBJECTS.length; + +const POST_MIGRATION_CUTOVER_REPLACEMENTS: readonly ProtocolV3CutoverObject[] = [ + { + name: "__protocol_v3_cutover_sandbox_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_sandbox_insert" +BEFORE INSERT ON "sandbox" +WHEN (NEW."status" <> 'cold' + OR NEW."operation_kind" IS NOT NULL + OR NEW."status_operation_id" IS NOT NULL + OR NEW."claim_owner" IS NOT NULL + OR NEW."claim_expires_at" IS NOT NULL) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT (NEW."subject_kind" = 'session' AND ${exactSmokeSessionSql('NEW."subject_id"')}) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new active sandboxes'); +END`, + tableName: "sandbox", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_sandbox_update", + sql: `CREATE TRIGGER "__protocol_v3_cutover_sandbox_update" +BEFORE UPDATE OF "status", "operation_kind", "status_operation_id", "claim_owner", "claim_expires_at" ON "sandbox" +WHEN OLD."status" = 'cold' + AND OLD."operation_kind" IS NULL + AND OLD."status_operation_id" IS NULL + AND OLD."claim_owner" IS NULL + AND OLD."claim_expires_at" IS NULL + AND (NEW."status" <> 'cold' + OR NEW."operation_kind" IS NOT NULL + OR NEW."status_operation_id" IS NOT NULL + OR NEW."claim_owner" IS NOT NULL + OR NEW."claim_expires_at" IS NOT NULL) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT (NEW."subject_kind" = 'session' AND ${exactSmokeSessionSql('NEW."subject_id"')}) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks sandbox activation'); +END`, + tableName: "sandbox", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_session_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_session_insert" +BEFORE INSERT ON "session" +WHEN NOT (${postMigrationSessionStaticSql("NEW")}) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT ${EXACT_NEW_SESSION_SMOKE_SQL} +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new Session operations'); +END`, + tableName: "session", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_session_update", + sql: `CREATE TRIGGER "__protocol_v3_cutover_session_update" +BEFORE UPDATE OF "status", "status_operation_id", "archived_at", "cleanup_operation_kind", "runtime_provisioning_operation_id", "runtime_provisioning_run_id", "runtime_provisioning_sandbox_id", "runtime_provisioning_sandbox_session_id", "runtime_provisioning_sandbox_incarnation", "runtime_provisioning_heartbeat_at" ON "session" +WHEN (${postMigrationSessionStaticSql("OLD")}) + AND NOT (${postMigrationSessionStaticSql("NEW")}) + AND EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT ${exactSmokeSessionSql('NEW."id"')} +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks Session operation acquisition'); +END`, + tableName: "session", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_api_command_update", + sql: `CREATE TRIGGER "__protocol_v3_cutover_api_command_update" +BEFORE UPDATE OF "kind", "status", "claim_owner", "claim_expires_at", "delivery_generation" ON "api_command" +WHEN EXISTS ( + SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" + WHERE "enabled" = 1 + AND ( + ("command_freeze" = 1 + AND ( + NEW."delivery_generation" IS NOT OLD."delivery_generation" + OR (NEW."status" IN ('queued', 'running') + AND (OLD."status" NOT IN ('queued', 'running') OR NEW."kind" IS NOT OLD."kind")) + )) + OR ("command_freeze" = 0 + AND NEW."kind" IN (${CUTOVER_BLOCKED_API_COMMAND_KINDS}) + AND NEW."status" IN ('queued', 'running') + AND (OLD."status" NOT IN ('queued', 'running') OR NEW."kind" IS NOT OLD."kind")) + ) + ) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks API command admission'); +END`, + tableName: "api_command", + type: "trigger", + }, +]; + +const POST_MIGRATION_CUTOVER_REPLACEMENT_BY_NAME = new Map( + POST_MIGRATION_CUTOVER_REPLACEMENTS.map((object) => [object.name, object]), +); +const PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECTS: readonly ProtocolV3CutoverObject[] = [ + ...PROTOCOL_V3_CUTOVER_OBJECTS.map( + (object) => POST_MIGRATION_CUTOVER_REPLACEMENT_BY_NAME.get(object.name) ?? object, + ), + { + name: "__protocol_v3_cutover_sandbox_backup_staging_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_sandbox_backup_staging_insert" +BEFORE INSERT ON "sandbox_backup_staging" +WHEN EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT ${exactSmokeSessionSql('NEW."workspace_session_id"', 'NEW."sandbox_id"')} +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new sandbox backup staging'); +END`, + tableName: "sandbox_backup_staging", + type: "trigger", + }, + { + name: "__protocol_v3_cutover_environment_artifact_backup_staging_insert", + sql: `CREATE TRIGGER "__protocol_v3_cutover_environment_artifact_backup_staging_insert" +BEFORE INSERT ON "environment_package_artifact_backup_staging" +WHEN EXISTS (SELECT 1 FROM "${PROTOCOL_V3_CUTOVER_TABLE}" WHERE "enabled" = 1) + AND NOT EXISTS ( + SELECT 1 + FROM "${PROTOCOL_V3_CUTOVER_TABLE}" AS "gate" + INNER JOIN "api_command" AS "command" + ON "command"."id" = NEW."command_id" + AND "command"."created_at" <= "gate"."started_at" + AND "command"."kind" = 'environment_package_artifact_build' + AND "command"."status" = 'running' + AND "command"."delivery_generation" = NEW."delivery_generation" + AND "command"."attempt_count" = NEW."attempt_count" + AND "command"."claim_owner" = NEW."claim_owner" + AND typeof("command"."claim_expires_at") = 'integer' + AND "command"."claim_expires_at" > unixepoch('subsec') * 1000 + AND json_valid("command"."payload_json") = 1 + AND json_extract("command"."payload_json", '$.appId') = NEW."app_id" + AND json_extract("command"."payload_json", '$.inputDigest') = NEW."input_digest" + WHERE "gate"."enabled" = 1 + AND "gate"."command_freeze" = 0 + ) +BEGIN + SELECT RAISE(ABORT, 'protocol v3 cutover blocks new environment artifact backup staging'); +END`, + tableName: "environment_package_artifact_backup_staging", + type: "trigger", + }, +]; + +export const PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT = + PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECTS.length; + +function installableCutoverObjectSql(object: ProtocolV3CutoverObject): string { + const create = object.type === "table" ? "CREATE TABLE" : "CREATE TRIGGER"; + return object.sql.replace(create, `${create} IF NOT EXISTS`); +} + +function sqlString(value: string): string { + return `'${value.replaceAll("'", "''")}'`; +} + +const CUTOVER_RESERVED_NAMES_SQL = [ + ...new Set(PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECTS.map((object) => object.name)), +] + .map(sqlString) + .join(", "); +const PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_TABLE = "__protocol_v3_legacy_rewrite_authorization"; +const PROTOCOL_V3_LEGACY_REWRITE_GATE_UPDATE_TRIGGER = "__protocol_v3_legacy_rewrite_gate_update"; +export const PROTOCOL_V3_LEGACY_REWRITE_GATE_UPDATE_TRIGGER_SQL = `CREATE TRIGGER "${PROTOCOL_V3_LEGACY_REWRITE_GATE_UPDATE_TRIGGER}" +AFTER UPDATE ON "${PROTOCOL_V3_CUTOVER_TABLE}" +BEGIN + DELETE FROM "${PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_TABLE}" WHERE "id" = 1; +END`; +const PROTOCOL_V3_LEGACY_REWRITE_GATE_UPDATE_TRIGGER_MATCH_SQL = `( + "type" = 'trigger' + AND "name" = ${sqlString(PROTOCOL_V3_LEGACY_REWRITE_GATE_UPDATE_TRIGGER)} COLLATE BINARY + AND "tbl_name" = ${sqlString(PROTOCOL_V3_CUTOVER_TABLE)} COLLATE BINARY + AND "sql" = ${sqlString(PROTOCOL_V3_LEGACY_REWRITE_GATE_UPDATE_TRIGGER_SQL)} COLLATE BINARY +)`; +const CUTOVER_PROTECTED_TABLES = [ + PROD_DEPLOY_LEASE_TABLE, + PROTOCOL_V3_CUTOVER_TABLE, + PROTOCOL_V3_MIGRATION_INTENT_TABLE, + PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_TABLE, + "api_command", + "app_deployment_run", + "driver_command", + "driver_instance", + "environment_package_artifact_backup", + "environment_package_artifact_backup_staging", + "sandbox", + "sandbox_backup", + "sandbox_backup_delete_intent", + "sandbox_backup_staging", + "sandbox_session", + "session", + "session_run", +] as const; +const CUTOVER_PROTECTED_TABLES_SQL = CUTOVER_PROTECTED_TABLES.map(sqlString).join(", "); +const CUTOVER_MANAGED_TRIGGER_NAMES = [ + "app_deployment_run_target_script_insert_authority", + "app_deployment_run_target_script_retire", + "app_deployment_run_target_script_update_authority", + "app_deployment_run_terminal_script_retire", + "environment_package_artifact_backup_staging_authority", + "environment_package_artifact_backup_staging_immutable", + "environment_package_artifact_backup_authority", + "environment_package_artifact_backup_retirement_authority", + "environment_package_artifact_backup_retirement_tombstone", + "environment_package_artifact_backup_rotation_authority", + "environment_package_artifact_backup_rotation_tombstone", + "sandbox_backup_delete_intent_attempt_clock", + "sandbox_backup_delete_intent_authority", + "sandbox_backup_delete_intent_blocks_environment_commit", + "sandbox_backup_delete_intent_blocks_environment_stage_insert", + "sandbox_backup_delete_intent_blocks_environment_stage_update", + "sandbox_backup_delete_intent_blocks_runtime_record", + "sandbox_backup_delete_intent_blocks_runtime_record_update", + "sandbox_backup_delete_intent_blocks_runtime_stage_insert", + "sandbox_backup_delete_intent_blocks_runtime_stage_update", + "sandbox_backup_delete_intent_completion_monotonic", + "sandbox_backup_delete_intent_identity_immutable", + "sandbox_backup_delete_intent_permanent", + "sandbox_backup_identity_immutable", + "sandbox_backup_permanent", + "sandbox_backup_staging_identity_immutable", + "sandbox_backup_status_monotonic", + "sandbox_identity_immutable", +] as const; +const CUTOVER_MANAGED_TRIGGER_MATCH_SQL = CUTOVER_MANAGED_TRIGGER_NAMES.map((name) => { + const trigger = MANAGED_PROD_SCHEMA_TRIGGERS.find((candidate) => candidate.name === name); + if ( + trigger === undefined || + !CUTOVER_PROTECTED_TABLES.some((tableName) => tableName === trigger.tableName) + ) { + throw new Error(`Managed cutover trigger ${name} is missing or protects the wrong table.`); + } + const canonicalSql = trigger.sql.replaceAll("`", '"'); + return `( + "type" = 'trigger' + AND "name" = ${sqlString(trigger.name)} COLLATE BINARY + AND "tbl_name" = ${sqlString(trigger.tableName)} COLLATE BINARY + AND replace("sql", char(96), '"') = ${sqlString(canonicalSql)} COLLATE BINARY + )`; +}).join("\n OR "); + +export function installProtocolV3CutoverSql(releaseTreeOid: string): string { + const release = requireGitTreeOid(releaseTreeOid, "Protocol v3 release tree OID"); + return ` +${PROTOCOL_V3_CUTOVER_OBJECTS.map((object) => `${installableCutoverObjectSql(object)};`).join("\n")} +${PROTOCOL_V3_MIGRATION_INTENT_TABLE_SQL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS")}; +INSERT INTO "${PROTOCOL_V3_CUTOVER_TABLE}" ("id", "enabled", "release_tree_oid") +VALUES (1, 1, '${release}') +ON CONFLICT ("id") DO NOTHING; +`; +} + +const PROTOCOL_V3_CUTOVER_TRIGGER_NAMES = [ + ...new Set( + PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECTS.filter((object) => object.type === "trigger").map( + (object) => object.name, + ), + ), +]; + +export const DROP_PROTOCOL_V3_CUTOVER_TRIGGERS_SQL = PROTOCOL_V3_CUTOVER_TRIGGER_NAMES.toReversed() + .map((name) => `DROP TRIGGER IF EXISTS "${name}";`) + .join("\n"); + +export const INSTALL_PROTOCOL_V3_POST_MIGRATION_CUTOVER_SQL = ` +${installableCutoverObjectSql(PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECTS[0])}; +${PROTOCOL_V3_MIGRATION_INTENT_TABLE_SQL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS")}; +${DROP_PROTOCOL_V3_CUTOVER_TRIGGERS_SQL} +${PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECTS.slice(1) + .map((object) => `${installableCutoverObjectSql(object)};`) + .join("\n")} +`; + +export function installProtocolV3PostMigrationCutoverSql(releaseTreeOid: string): string { + const release = requireGitTreeOid(releaseTreeOid, "Protocol v3 release tree OID"); + return `${INSTALL_PROTOCOL_V3_POST_MIGRATION_CUTOVER_SQL} +INSERT INTO "${PROTOCOL_V3_CUTOVER_TABLE}" ("id", "enabled", "release_tree_oid") +VALUES (1, 1, '${release}') +ON CONFLICT ("id") DO NOTHING; +`; +} + +export const REMOVE_PROTOCOL_V3_CUTOVER_SQL = ` +${DROP_PROTOCOL_V3_CUTOVER_TRIGGERS_SQL} +DROP TABLE IF EXISTS "${PROTOCOL_V3_MIGRATION_INTENT_TABLE}"; +DROP TABLE IF EXISTS "${PROTOCOL_V3_CUTOVER_TABLE}"; +`; + +export const PROTOCOL_V3_CUTOVER_PROBE_SQL = ` +SELECT EXISTS( + SELECT 1 FROM "sqlite_master" + WHERE ("name" COLLATE NOCASE IN (${CUTOVER_RESERVED_NAMES_SQL}) + OR "name" = '${PROTOCOL_V3_MIGRATION_INTENT_TABLE}' COLLATE NOCASE) + OR ( + "type" = 'trigger' + AND "tbl_name" COLLATE NOCASE IN (${CUTOVER_PROTECTED_TABLES_SQL}) + AND NOT ( + ${PROTOCOL_V3_LEGACY_REWRITE_GATE_UPDATE_TRIGGER_MATCH_SQL} + OR ${CUTOVER_MANAGED_TRIGGER_MATCH_SQL} + ) + ) + ) AS "gate_present"; +`; + +export const PROTOCOL_V3_LEGACY_TERMINAL_SOURCE_INVENTORY_SQL = ` +WITH "terminal_events" AS ( + SELECT + "event_type", + "id", + "run_id", + "session_id", + "source_event_id", + 'session-run-terminal:' || "run_id" || ':' || "event_type" AS "canonical_source_event_id" + FROM "session_event" + WHERE "event_type" IN ('run.cancelled', 'run.completed', 'run.failed') +), +"source_inventory" AS ( + SELECT + "terminal"."event_type", + count(*) AS "total", + sum("terminal"."source_event_id" = "terminal"."canonical_source_event_id") AS "canonical", + sum("terminal"."source_event_id" <> "terminal"."canonical_source_event_id") AS "noncanonical", + sum(EXISTS ( + SELECT 1 + FROM "session_event" AS "other" + WHERE "other"."session_id" = "terminal"."session_id" + AND "other"."source_event_id" = "terminal"."canonical_source_event_id" + AND "other"."id" <> "terminal"."id" + )) AS "canonical_target_collisions" + FROM "terminal_events" AS "terminal" + GROUP BY "terminal"."event_type" +), +"multiple_terminal_runs" AS ( + SELECT 1 + FROM "terminal_events" + WHERE "run_id" IS NOT NULL + GROUP BY "session_id", "run_id" + HAVING count(*) > 1 +) +SELECT + coalesce(max(CASE WHEN "event_type" = 'run.cancelled' THEN "total" END), 0) AS "cancelled_total", + coalesce(max(CASE WHEN "event_type" = 'run.cancelled' THEN "canonical" END), 0) AS "cancelled_canonical", + coalesce(max(CASE WHEN "event_type" = 'run.cancelled' THEN "noncanonical" END), 0) AS "cancelled_noncanonical", + coalesce(max(CASE WHEN "event_type" = 'run.cancelled' THEN "canonical_target_collisions" END), 0) AS "cancelled_canonical_target_collisions", + coalesce(max(CASE WHEN "event_type" = 'run.completed' THEN "total" END), 0) AS "completed_total", + coalesce(max(CASE WHEN "event_type" = 'run.completed' THEN "canonical" END), 0) AS "completed_canonical", + coalesce(max(CASE WHEN "event_type" = 'run.completed' THEN "noncanonical" END), 0) AS "completed_noncanonical", + coalesce(max(CASE WHEN "event_type" = 'run.completed' THEN "canonical_target_collisions" END), 0) AS "completed_canonical_target_collisions", + coalesce(max(CASE WHEN "event_type" = 'run.failed' THEN "total" END), 0) AS "failed_total", + coalesce(max(CASE WHEN "event_type" = 'run.failed' THEN "canonical" END), 0) AS "failed_canonical", + coalesce(max(CASE WHEN "event_type" = 'run.failed' THEN "noncanonical" END), 0) AS "failed_noncanonical", + coalesce(max(CASE WHEN "event_type" = 'run.failed' THEN "canonical_target_collisions" END), 0) AS "failed_canonical_target_collisions", + ( + SELECT count(*) + FROM "terminal_events" AS "event" + LEFT JOIN "session_run" AS "run" ON "run"."id" = "event"."run_id" + WHERE "event"."run_id" IS NULL + OR "run"."id" IS NULL + OR "run"."session_id" <> "event"."session_id" + ) AS "invalid_terminal_links", + ( + SELECT count(*) + FROM "terminal_events" AS "event" + INNER JOIN "session_run" AS "run" + ON "run"."id" = "event"."run_id" + AND "run"."session_id" = "event"."session_id" + WHERE NOT ( + ("run"."status" = 'completed' AND "event"."event_type" = 'run.completed') + OR ("run"."status" = 'failed' AND "event"."event_type" = 'run.failed') + OR ("run"."status" IN ('cancelled', 'expired') AND "event"."event_type" = 'run.cancelled') + ) + ) AS "mismatched_terminal_events", + (SELECT count(*) FROM "multiple_terminal_runs") AS "multiple_terminal_runs" +FROM "source_inventory"; +`; + +export const PROTOCOL_V3_LOSSY_MIGRATION_INVENTORY_SQL = ` +WITH "effect_source" AS ( + SELECT + "effect"."id" AS "effect_id", + "effect"."status" AS "effect_status", + "effect"."provider_receipt_json" AS "effect_provider_receipt_json", + "effect"."result_json" AS "effect_result_json", + "command"."error_json" AS "command_error_json", + "command"."payload_json" AS "command_payload_json", + "command"."result_json" AS "command_result_json", + "command"."status" AS "command_status", + max( + coalesce(length(CAST("effect"."result_json" AS BLOB)), 0), + coalesce(( + SELECT max(length(CAST("attempt"."result_json" AS BLOB))) + FROM "external_tool_effect_attempt" AS "attempt" + WHERE "attempt"."effect_id" = "effect"."id" + AND "attempt"."status" = 'succeeded' + ), 0), + coalesce(length(CAST("command"."result_json" AS BLOB)), 0) + ) AS "original_result_bytes" + FROM "external_tool_effect" AS "effect" + INNER JOIN "driver_command" AS "command" ON "command"."id" = "effect"."command_id" +), +"effect_result_classified" AS ( + SELECT + "effect_source".*, + "effect_status" = 'succeeded' AND ( + "original_result_bytes" > 1044480 + OR length(CAST('{"kind":"succeeded","result":' || "effect_result_json" || '}' AS BLOB)) > 1044480 + ) AS "result_omitted" + FROM "effect_source" +), +"effect_result_target" AS ( + SELECT + "effect_result_classified".*, + CASE + WHEN "result_omitted" THEN + '{"isError":true,"outputText":' || json_quote('Stored MCP result omitted because it contained ' || "original_result_bytes" || ' UTF-8 bytes.') || + ',"requestId":' || json_quote(json_extract("command_payload_json", '$.requestId')) || + ',"serverId":' || json_quote(json_extract("command_payload_json", '$.serverId')) || + ',"toolName":' || json_quote(json_extract("command_payload_json", '$.toolName')) || '}' + WHEN "effect_status" = 'succeeded' THEN "effect_result_json" + ELSE NULL + END AS "normalized_result_json" + FROM "effect_result_classified" +), +"effect_target" AS ( + SELECT + "effect_result_target".*, + CASE + WHEN "effect_status" <> 'succeeded' THEN NULL + WHEN "effect_provider_receipt_json" IS NULL THEN NULL + WHEN length(CAST( + '{"kind":"succeeded","providerReceiptJson":' || json_quote("effect_provider_receipt_json") || + ',"result":' || "normalized_result_json" || '}' + AS BLOB)) > 1044480 THEN NULL + ELSE "effect_provider_receipt_json" + END AS "normalized_provider_receipt_json" + FROM "effect_result_target" +), +"loss_candidates" ("category", "id") AS ( + SELECT "category"."value", "command"."id" + FROM "driver_command" AS "command" + CROSS JOIN json_each(json_array( + 'command_payload_conflict', + 'mcp_argument_omission', + 'input_text_omission', + 'input_start_result_omission', + 'control_reason_omission', + 'permission_payload_rewrite', + 'command_error_omission' + )) AS "category" + WHERE CASE "category"."value" + WHEN 'command_payload_conflict' THEN + EXISTS ( + SELECT "key" + FROM json_each("command"."payload_json") + GROUP BY "key" + HAVING count(*) > 1 + ) + OR ( + "command"."kind" = 'input.start' + AND EXISTS ( + SELECT "key" + FROM json_each("command"."payload_json", '$.input') + GROUP BY "key" + HAVING count(*) > 1 + ) + ) + WHEN 'mcp_argument_omission' THEN + "command"."kind" = 'mcp.execute' + AND length(CAST(json_set( + "command"."payload_json", + '$.commandId', "command"."id", + '$.runId', ( + SELECT "effect"."session_run_id" + FROM "external_tool_effect" AS "effect" + WHERE "effect"."command_id" = "command"."id" + LIMIT 1 + ) + ) AS BLOB)) > 824448 + WHEN 'input_text_omission' THEN + "command"."kind" = 'input.start' + AND length(CAST("command"."payload_json" AS BLOB)) > 824448 + WHEN 'input_start_result_omission' THEN + "command"."kind" = 'input.start' + AND "command"."result_json" IS NOT NULL + AND json_type("command"."result_json") <> 'null' + AND ( + length(CAST("command"."payload_json" AS BLOB)) > 824448 + OR length(CAST("command"."result_json" AS BLOB)) > 1044480 + ) + WHEN 'control_reason_omission' THEN + "command"."kind" IN ('turn.cancel', 'session.stop') + AND length(CAST("command"."payload_json" AS BLOB)) > 824448 + WHEN 'permission_payload_rewrite' THEN + "command"."kind" = 'permission.resolve' + AND length(CAST("command"."payload_json" AS BLOB)) > 824448 + WHEN 'command_error_omission' THEN + "command"."error_json" IS NOT NULL + AND length(CAST("command"."error_json" AS BLOB)) > 1044480 + AND NOT EXISTS ( + SELECT 1 + FROM "external_tool_effect" AS "effect" + WHERE "effect"."command_id" = "command"."id" AND "effect"."status" = 'succeeded' + ) + ELSE 0 + END + UNION ALL + SELECT "category"."value", "target"."effect_id" + FROM "effect_target" AS "target" + CROSS JOIN json_each(json_array( + 'mcp_result_omission', + 'mcp_result_conflict', + 'provider_receipt_loss', + 'mcp_command_terminal_conflict' + )) AS "category" + WHERE CASE "category"."value" + WHEN 'mcp_result_omission' THEN "target"."result_omitted" + WHEN 'mcp_result_conflict' THEN + ( + NOT "target"."result_omitted" + AND ( + ("target"."effect_status" <> 'succeeded' AND "target"."effect_result_json" IS NOT NULL) + OR ( + "target"."effect_status" = 'succeeded' + AND "target"."command_result_json" IS NOT NULL + AND json_type("target"."command_result_json") <> 'null' + AND "target"."command_result_json" IS NOT "target"."effect_result_json" + ) + OR EXISTS ( + SELECT 1 + FROM "external_tool_effect_attempt" AS "attempt" + WHERE "attempt"."effect_id" = "target"."effect_id" + AND "attempt"."result_json" IS NOT NULL + AND ( + "attempt"."status" <> 'succeeded' + OR "target"."normalized_result_json" IS NULL + OR "attempt"."result_json" IS NOT "target"."effect_result_json" + ) + ) + ) + ) + OR ( + "target"."effect_result_json" IS NOT NULL + AND EXISTS ( + SELECT "key" + FROM json_each("target"."effect_result_json") + GROUP BY "key" + HAVING count(*) > 1 + ) + ) + OR ( + "target"."command_result_json" IS NOT NULL + AND EXISTS ( + SELECT "key" + FROM json_each("target"."command_result_json") + GROUP BY "key" + HAVING count(*) > 1 + ) + ) + OR EXISTS ( + SELECT 1 + FROM "external_tool_effect_attempt" AS "attempt" + WHERE "attempt"."effect_id" = "target"."effect_id" + AND "attempt"."result_json" IS NOT NULL + AND EXISTS ( + SELECT "key" + FROM json_each("attempt"."result_json") + GROUP BY "key" + HAVING count(*) > 1 + ) + ) + WHEN 'provider_receipt_loss' THEN + "target"."effect_provider_receipt_json" IS NOT "target"."normalized_provider_receipt_json" + OR EXISTS ( + SELECT 1 + FROM "external_tool_effect_attempt" AS "attempt" + WHERE "attempt"."effect_id" = "target"."effect_id" + AND "attempt"."provider_receipt_json" IS NOT NULL + AND "attempt"."provider_receipt_json" IS NOT CASE + WHEN "attempt"."status" = 'succeeded' THEN "target"."normalized_provider_receipt_json" + ELSE NULL + END + ) + WHEN 'mcp_command_terminal_conflict' THEN + "target"."effect_status" = 'succeeded' + AND ( + "target"."command_status" <> 'completed' OR "target"."command_error_json" IS NOT NULL + ) + ELSE 0 + END + UNION ALL + SELECT 'orphan_effect', "effect"."id" + FROM "external_tool_effect" AS "effect" + WHERE NOT EXISTS ( + SELECT 1 FROM "driver_command" AS "command" WHERE "command"."id" = "effect"."command_id" + ) + UNION ALL + SELECT DISTINCT 'attempt_completion_time_fabrication', "attempt"."effect_id" + FROM "external_tool_effect_attempt" AS "attempt" + WHERE "attempt"."status" IN ('succeeded', 'unknown') + AND "attempt"."completed_at" IS NULL + UNION ALL + SELECT 'session_run_error_omission', "run"."id" + FROM "session_run" AS "run" + WHERE "run"."error_code" IS NOT NULL + AND "run"."error_message" IS NOT NULL + AND length(CAST( + '{"code":' || json_quote("run"."error_code") || + ',"details":' || coalesce(nullif("run"."error_details_json", ''), '{}') || + ',"message":' || json_quote("run"."error_message") || + ',"retryable":false}' + AS BLOB)) > 1044480 +) +SELECT + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'attempt_completion_time_fabrication') AS "attempt_completion_time_fabrications", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'command_payload_conflict') AS "command_payload_conflicts", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'mcp_argument_omission') AS "mcp_argument_omissions", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'input_text_omission') AS "input_text_omissions", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'input_start_result_omission') AS "input_start_result_omissions", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'control_reason_omission') AS "control_reason_omissions", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'permission_payload_rewrite') AS "permission_payload_rewrites", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'mcp_result_omission') AS "mcp_result_omissions", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'orphan_effect') AS "orphan_effects", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'mcp_result_conflict') AS "mcp_result_conflicts", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'provider_receipt_loss') AS "provider_receipt_losses", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'mcp_command_terminal_conflict') AS "mcp_command_terminal_conflicts", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'command_error_omission') AS "command_error_omissions", + (SELECT count(*) FROM "loss_candidates" WHERE "category" = 'session_run_error_omission') AS "session_run_error_omissions", + (SELECT count(*) FROM "loss_candidates") AS "total_candidates", + coalesce(( + SELECT json_group_array(json_array("category", "id")) + FROM ( + SELECT "category", "id" + FROM "loss_candidates" + ORDER BY "category", "id" + LIMIT 50 + ) + ), json('[]')) AS "candidate_ids_json"; +`; + +export const PROTOCOL_V3_LEGACY_TERMINAL_INTEGRITY_SQL = ` +WITH "terminal_events" AS ( + SELECT "event_type", "id", "run_id", "seq", "session_id", "source_event_id" + FROM "session_event" + WHERE "event_type" IN ('run.cancelled', 'run.completed', 'run.failed') +), +"assistant_runs" AS ( + SELECT "session_id", "session_run_id", count(*) AS "assistant_count" + FROM "session_message" + WHERE "role" = 'assistant' AND "session_run_id" IS NOT NULL + GROUP BY "session_id", "session_run_id" +) +SELECT + ( + SELECT count(*) + FROM ( + SELECT 1 + FROM "terminal_events" + WHERE "run_id" IS NOT NULL + GROUP BY "session_id", "run_id" + HAVING count(*) > 1 + ) + ) AS "duplicate_terminal_runs", + ( + SELECT count(*) + FROM "terminal_events" AS "event" + INNER JOIN "session_run" AS "run" + ON "run"."id" = "event"."run_id" + AND "run"."session_id" = "event"."session_id" + WHERE NOT ( + ("run"."status" = 'completed' AND "event"."event_type" = 'run.completed') + OR ("run"."status" = 'failed' AND "event"."event_type" = 'run.failed') + OR ("run"."status" IN ('cancelled', 'expired') AND "event"."event_type" = 'run.cancelled') + ) + ) AS "mismatched_terminal_events", + ( + SELECT count(*) + FROM "terminal_events" AS "event" + LEFT JOIN "session_run" AS "run" ON "run"."id" = "event"."run_id" + WHERE "event"."run_id" IS NULL + OR "run"."id" IS NULL + OR "run"."session_id" <> "event"."session_id" + ) AS "invalid_terminal_links", + ( + SELECT count(*) + FROM "terminal_events" + WHERE "run_id" IS NULL + OR "source_event_id" <> 'session-run-terminal:' || "run_id" || ':' || "event_type" + ) AS "noncanonical_terminal_sources", + ( + SELECT json_group_array(json_array( + "id", "session_id", "run_id", "event_type", "source_event_id", "seq" + )) + FROM ( + SELECT "id", "session_id", "run_id", "event_type", "source_event_id", "seq" + FROM "terminal_events" + WHERE "run_id" IS NULL + OR "source_event_id" <> + 'session-run-terminal:' || "run_id" || ':' || "event_type" + ORDER BY "id" COLLATE BINARY + ) + ) AS "rewrite_candidate_manifest_json", + ( + SELECT count(*) + FROM "session_run" AS "run" + WHERE "run"."status" IN ('cancelled', 'completed', 'expired', 'failed') + AND NOT EXISTS ( + SELECT 1 + FROM "terminal_events" AS "event" + WHERE "event"."session_id" = "run"."session_id" + AND "event"."run_id" = "run"."id" + ) + ) AS "missing_terminal_events", + ( + SELECT count(*) + FROM "terminal_events" AS "event" + INNER JOIN "session_run" AS "run" + ON "run"."id" = "event"."run_id" + AND "run"."session_id" = "event"."session_id" + LEFT JOIN "session" ON "session"."id" = "run"."session_id" + WHERE "session"."id" IS NULL + OR "run"."completed_at" IS NULL + OR "run"."status_event" <> CASE "run"."status" + WHEN 'completed' THEN 'run.complete' + WHEN 'failed' THEN 'run.fail' + WHEN 'cancelled' THEN 'run.cancel' + WHEN 'expired' THEN 'run.expire' + END + OR "event"."seq" > "session"."runtime_event_seq_cursor" + OR ( + "session"."last_run_id" = "run"."id" + AND ( + "session"."status" NOT IN ('IDLE', 'TERMINATED') + OR "session"."status_operation_id" IS NOT NULL + ) + ) + OR EXISTS ( + SELECT 1 + FROM "session_permission_request" AS "permission" + WHERE "permission"."session_id" = "run"."session_id" + AND "permission"."run_id" = "run"."id" + ) + ) AS "partial_terminal_projections", + ( + SELECT count(*) + FROM "assistant_runs" AS "assistant" + INNER JOIN "session_run" AS "run" + ON "run"."id" = "assistant"."session_run_id" + AND "run"."session_id" = "assistant"."session_id" + WHERE "run"."status" = 'completed' AND "assistant"."assistant_count" > 1 + ) AS "ambiguous_assistant_runs", + ( + SELECT count(*) + FROM "session_message" AS "message" + LEFT JOIN "session_run" AS "run" ON "run"."id" = "message"."session_run_id" + LEFT JOIN "session" ON "session"."id" = "message"."session_id" + WHERE "message"."role" = 'assistant' + AND "message"."session_run_id" IS NOT NULL + AND ( + "run"."id" IS NULL + OR "run"."session_id" <> "message"."session_id" + OR "session"."id" IS NULL + OR "message"."seq" > "session"."message_seq_cursor" + OR "run"."status" IN ('cancelled', 'expired', 'failed') + OR CASE + WHEN "message"."plan_json" IS NULL OR "message"."plan_json" = '' THEN 0 + WHEN json_valid("message"."plan_json") = 0 THEN 1 + WHEN json_type("message"."plan_json") <> 'array' THEN 1 + ELSE 0 + END = 1 + OR CASE + WHEN "message"."segments_json" IS NULL OR "message"."segments_json" = '' THEN 0 + WHEN json_valid("message"."segments_json") = 0 THEN 1 + WHEN json_type("message"."segments_json") <> 'array' THEN 1 + ELSE 0 + END = 1 + ) + ) AS "partial_assistant_projections", + ( + SELECT count(*) + FROM "session_run" + WHERE "status" = 'failed' + AND ( + "error_code" IS NULL + OR trim("error_code") = '' + OR "error_message" IS NULL + OR trim("error_message") = '' + OR CASE + WHEN "error_details_json" IS NULL THEN 0 + WHEN json_valid("error_details_json") = 0 THEN 1 + WHEN json_type("error_details_json") <> 'object' THEN 1 + ELSE 0 + END = 1 + ) + ) AS "invalid_failed_runs", + ( + SELECT count(*) + FROM "session_run" + WHERE "status" IN ('cancelled', 'completed', 'expired') + AND ( + "error_code" IS NOT NULL + OR "error_details_json" IS NOT NULL + OR "error_message" IS NOT NULL + ) + ) AS "invalid_nonfailed_run_errors", + ( + SELECT count(*) + FROM "session_run" + WHERE "status" = 'failed' + AND "error_code" IS NOT NULL + AND trim("error_code") <> '' + AND "error_message" IS NOT NULL + AND trim("error_message") <> '' + AND CASE + WHEN "error_details_json" IS NULL THEN 1 + WHEN json_valid("error_details_json") = 0 THEN 0 + WHEN json_type("error_details_json") = 'object' THEN 1 + ELSE 0 + END = 1 + ) AS "repairable_failed_runs", + ( + SELECT count(*) + FROM "driver_command" + WHERE "status" IN ('queued', 'delivered', 'accepted') + ) AS "nonterminal_commands", + ( + SELECT count(*) + FROM "external_tool_effect" + WHERE "status" IN ('executing', 'claimed') + ) AS "unsettled_effects", + (SELECT count(*) FROM "session_message") AS "legacy_materialized_messages", + (SELECT count(*) FROM "terminal_events") AS "legacy_terminal_events", + ( + SELECT count(*) + FROM "session_event" + WHERE "event_type" LIKE 'message.%' OR "event_type" LIKE 'thought.%' + ) AS "legacy_stream_rows"; +`; + +export const PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_TABLE_SQL = `CREATE TABLE "${PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_TABLE}" ( + "id" integer PRIMARY KEY CHECK ("id" = 1), + "gate_id" integer NOT NULL CHECK ("gate_id" = 1) REFERENCES "${PROTOCOL_V3_CUTOVER_TABLE}" ("id") ON DELETE CASCADE, + "bookmark" text NOT NULL CHECK (length(trim("bookmark")) > 0), + "candidate_count" integer NOT NULL CHECK ("candidate_count" >= 0), + "candidate_manifest_json" text NOT NULL CHECK (json_valid("candidate_manifest_json") AND json_type("candidate_manifest_json") = 'array'), + "deploy_owner" text NOT NULL, + "expires_at" integer NOT NULL, + "release_tree_oid" text NOT NULL CHECK ((length("release_tree_oid") = 40 OR length("release_tree_oid") = 64) AND "release_tree_oid" = lower("release_tree_oid") AND "release_tree_oid" NOT GLOB '*[^0-9a-f]*') +)`; + +const PROTOCOL_V3_LEGACY_REWRITE_MANIFEST_SQL = ` +SELECT + count(*) AS "candidate_count", + json_group_array(json_array( + "id", "session_id", "run_id", "event_type", "source_event_id", "seq" + )) AS "candidate_manifest_json" +FROM ( + SELECT "id", "session_id", "run_id", "event_type", "source_event_id", "seq" + FROM "session_event" + WHERE "event_type" IN ('run.cancelled', 'run.completed', 'run.failed') + AND "source_event_id" <> + 'session-run-terminal:' || "run_id" || ':' || "event_type" + ORDER BY "id" COLLATE BINARY +)`; + +export function authorizeProtocolV3LegacyRewriteSql( + owner: string, + candidateCount: number, + candidateManifestJson: string, +): string { + if (!/^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/u.test(owner)) { + throw new Error("Protocol v3 legacy rewrite owner must be a UUID v4."); + } + let candidateManifest: unknown; + try { + candidateManifest = JSON.parse(candidateManifestJson); + } catch { + throw new Error("Protocol v3 legacy rewrite manifest must be valid JSON."); + } + if ( + !Number.isSafeInteger(candidateCount) || + candidateCount < 0 || + !Array.isArray(candidateManifest) || + candidateManifest.length !== candidateCount + ) { + throw new Error("Protocol v3 legacy rewrite manifest count is invalid."); + } + const quotedOwner = sqlString(owner); + const quotedManifest = sqlString(candidateManifestJson); + return ` +${PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_TABLE_SQL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS")}; +${PROTOCOL_V3_LEGACY_REWRITE_GATE_UPDATE_TRIGGER_SQL.replace("CREATE TRIGGER", "CREATE TRIGGER IF NOT EXISTS")}; +DELETE FROM "${PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_TABLE}" WHERE "id" = 1; +WITH "manifest" AS (${PROTOCOL_V3_LEGACY_REWRITE_MANIFEST_SQL}) +INSERT INTO "${PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_TABLE}" ( + "id", "gate_id", "bookmark", "candidate_count", "candidate_manifest_json", "deploy_owner", "expires_at", "release_tree_oid" +) +SELECT + 1, + "gate"."id", + "gate"."pre_migration_bookmark", + "manifest"."candidate_count", + "manifest"."candidate_manifest_json", + "lease"."owner", + unixepoch() + 600, + "gate"."release_tree_oid" +FROM "manifest" +INNER JOIN "${PROTOCOL_V3_CUTOVER_TABLE}" AS "gate" + ON "gate"."id" = 1 + AND "gate"."enabled" = 1 + AND "gate"."command_freeze" = 1 + AND "gate"."phase" = 'draining' + AND length(trim("gate"."pre_migration_bookmark")) > 0 + AND "gate"."smoke_account_id" IS NULL + AND "gate"."smoke_request_key" IS NULL + AND "gate"."smoke_session_id" IS NULL +INNER JOIN "${PROD_DEPLOY_LEASE_TABLE}" AS "lease" + ON "lease"."id" = 1 + AND "lease"."owner" = ${quotedOwner} +WHERE NOT EXISTS ( + SELECT 1 FROM "session_run" + WHERE "status" IN (${ACTIVE_RUN_STATUSES}) +) +AND NOT EXISTS ( + SELECT 1 FROM "driver_instance" + WHERE "status" IN (${LIVE_DRIVER_STATUSES}) +) +AND NOT EXISTS ( + SELECT 1 FROM "driver_command" + WHERE "status" IN ('queued', 'delivered', 'accepted') +) +AND NOT EXISTS ( + SELECT 1 FROM "external_tool_effect" + WHERE "status" IN ('executing', 'claimed') +) +AND NOT EXISTS ( + SELECT 1 FROM "api_command" + WHERE "status" IN ('queued', 'running') +) +AND NOT EXISTS ( + SELECT 1 FROM "app_deployment_run" + WHERE "status" IN (${ACTIVE_APP_DEPLOYMENT_RUN_STATUSES}) +) +AND NOT EXISTS ( + SELECT 1 FROM "sandbox" + WHERE "status" <> 'cold' + OR "status_operation_id" IS NOT NULL + OR "claim_owner" IS NOT NULL + OR "claim_expires_at" IS NOT NULL +) +AND NOT EXISTS ( + SELECT 1 FROM "sandbox_session" + WHERE "status" NOT IN ('closed', 'error') +) +AND NOT EXISTS ( + SELECT 1 FROM "sandbox_backup" + WHERE "status" NOT IN ('ready', 'pruned') OR "error_message" IS NOT NULL +) +AND NOT EXISTS ( + SELECT 1 FROM "session" + WHERE "status" NOT IN ('IDLE', 'TERMINATED') OR "status_operation_id" IS NOT NULL +) +AND "manifest"."candidate_count" = ${candidateCount} +AND "manifest"."candidate_manifest_json" = ${quotedManifest} COLLATE BINARY; +`; +} + +export const PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_SQL = ` +SELECT + "authorization"."bookmark", + "authorization"."candidate_count", + "authorization"."candidate_manifest_json", + "authorization"."deploy_owner", + "authorization"."expires_at", + "authorization"."release_tree_oid", + (SELECT "sql" FROM "sqlite_master" + WHERE "type" = 'table' AND "name" = '${PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_TABLE}') AS "authorization_table_sql", + (SELECT "sql" FROM "sqlite_master" + WHERE "type" = 'trigger' AND "name" = '${PROTOCOL_V3_LEGACY_REWRITE_GATE_UPDATE_TRIGGER}') AS "gate_update_trigger_sql" +FROM "${PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_TABLE}" AS "authorization" +WHERE "authorization"."id" = 1; +`; + +function protocolV3CutoverDrainSql(postRuntimeAuthorityMigration: boolean): string { + const sandboxOperation = postRuntimeAuthorityMigration + ? '\n OR "operation_kind" IS NOT NULL' + : ""; + const unsafeSessions = postRuntimeAuthorityMigration + ? `NOT (${postMigrationSessionStaticSql('"session"')})` + : `"status" NOT IN ('IDLE', 'TERMINATED') OR "status_operation_id" IS NOT NULL`; + const stagingCount = postRuntimeAuthorityMigration + ? '(SELECT count(*) FROM "sandbox_backup_staging")' + : "0"; + const environmentArtifactStagingCount = postRuntimeAuthorityMigration + ? '(SELECT count(*) FROM "environment_package_artifact_backup_staging")' + : "0"; + return ` +SELECT + (SELECT count(*) FROM "session_run" + WHERE "status" IN (${ACTIVE_RUN_STATUSES})) AS "active_runs", + (SELECT count(*) FROM "app_deployment_run" + WHERE "status" IN (${ACTIVE_APP_DEPLOYMENT_RUN_STATUSES})) AS "active_app_deployment_runs", + (SELECT count(*) FROM "driver_instance" + WHERE "status" IN (${LIVE_DRIVER_STATUSES})) AS "live_drivers", + (SELECT count(*) FROM "external_tool_effect" + WHERE "status" IN ('executing', 'claimed')) AS "unsettled_effects", + (SELECT count(*) FROM "driver_command" + WHERE "status" IN ('queued', 'delivered', 'accepted')) AS "nonterminal_commands", + (SELECT count(*) FROM "api_command" + WHERE "status" IN ('queued', 'running')) AS "nonterminal_api_commands", + (SELECT count(*) FROM "sandbox" + WHERE "status" <> 'cold' + ${sandboxOperation} + OR "status_operation_id" IS NOT NULL + OR "claim_owner" IS NOT NULL + OR "claim_expires_at" IS NOT NULL) AS "unsafe_sandboxes", + (SELECT count(*) FROM "sandbox_session" + WHERE "status" NOT IN ('closed', 'error')) AS "unsafe_sandbox_sessions", + (SELECT count(*) FROM "sandbox_backup" + WHERE "status" NOT IN ('ready', 'pruned')) AS "unsafe_sandbox_backups", + ${stagingCount} AS "unsafe_sandbox_backup_staging", + ${environmentArtifactStagingCount} AS "unsafe_environment_artifact_backup_staging", + (SELECT count(*) FROM "session" WHERE ${unsafeSessions}) AS "unsafe_sessions"; +`; +} + +export const PROTOCOL_V3_CUTOVER_DRAIN_SQL = protocolV3CutoverDrainSql(false); +export const PROTOCOL_V3_POST_MIGRATION_CUTOVER_DRAIN_SQL = protocolV3CutoverDrainSql(true); + +function protocolV3UnsafeSandboxesSql(postRuntimeAuthorityMigration: boolean): string { + const operationColumn = postRuntimeAuthorityMigration ? ', "operation_kind"' : ""; + const operationPredicate = postRuntimeAuthorityMigration + ? '\n OR "operation_kind" IS NOT NULL' + : ""; + return ` +SELECT "id", "status", "status_operation_id", "claim_owner", "claim_expires_at"${operationColumn} +FROM "sandbox" +WHERE "status" <> 'cold' + ${operationPredicate} + OR "status_operation_id" IS NOT NULL + OR "claim_owner" IS NOT NULL + OR "claim_expires_at" IS NOT NULL +ORDER BY "id" +LIMIT 50; +`; +} + +export const PROTOCOL_V3_UNSAFE_SANDBOXES_SQL = protocolV3UnsafeSandboxesSql(false); +export const PROTOCOL_V3_POST_MIGRATION_UNSAFE_SANDBOXES_SQL = protocolV3UnsafeSandboxesSql(true); + +export function protocolV3RuntimeAuthorityPreflightSql( + sessionCleanupMigrationPending: boolean, +): string { + const nonstaticSessionsSql = sessionCleanupMigrationPending + ? `(SELECT count(*) + FROM "session" + WHERE "status" NOT IN ('IDLE', 'TERMINATED') + OR "status_operation_id" IS NOT NULL + )` + : `(SELECT count(*) + FROM "session" + WHERE "status" NOT IN ('IDLE', 'TERMINATED') + OR "status_operation_id" IS NOT NULL + OR NOT ( + "cleanup_operation_kind" IS NULL + OR ( + "cleanup_operation_kind" = 'archive' + AND "status" = 'IDLE' + AND "archived_at" IS NOT NULL + ) + ) + OR "runtime_provisioning_operation_id" IS NOT NULL + OR "runtime_provisioning_run_id" IS NOT NULL + OR "runtime_provisioning_sandbox_id" IS NOT NULL + OR "runtime_provisioning_heartbeat_at" IS NOT NULL + )`; + return ` +WITH "runtime_subject_identity" ( + "sandbox_id", "kind", "subject_kind", "subject_id", "agent_id", "app_id", "owner_account_id" +) AS ( + SELECT + "sandbox"."id", 'pet', 'agent', "sandbox"."subject_id", + "agent"."id", "agent"."app_id", "agent"."owner_account_id" + FROM "sandbox" + INNER JOIN "agent" + ON "sandbox"."kind" = 'pet' + AND "sandbox"."subject_kind" = 'agent' + AND "agent"."id" = "sandbox"."subject_id" + AND "agent"."kind" = 'pet' + INNER JOIN "app" ON "app"."id" = "agent"."app_id" + UNION ALL + SELECT + "sandbox"."id", 'cattle', 'session', "sandbox"."subject_id", + "agent"."id", "session"."app_id", "agent"."owner_account_id" + FROM "sandbox" + INNER JOIN "session" + ON "sandbox"."kind" = 'cattle' + AND "sandbox"."subject_kind" = 'session' + AND "session"."id" = "sandbox"."subject_id" + AND "session"."kind" = 'cattle' + INNER JOIN "agent" + ON "agent"."id" = "session"."agent_id" + AND "agent"."app_id" = "session"."app_id" + AND "agent"."kind" = 'cattle' + INNER JOIN "app" ON "app"."id" = "session"."app_id" +) +SELECT + (SELECT count(*) + FROM "sandbox" AS "sandbox" + LEFT JOIN "runtime_subject_identity" AS "identity" + ON "identity"."sandbox_id" = "sandbox"."id" + WHERE "identity"."sandbox_id" IS NULL + OR NOT ( + ("sandbox"."agent_id" IS NULL + AND "sandbox"."app_id" IS NULL + AND "sandbox"."owner_account_id" IS NULL) + OR ("sandbox"."agent_id" IS "identity"."agent_id" + AND "sandbox"."app_id" IS "identity"."app_id" + AND "sandbox"."owner_account_id" IS "identity"."owner_account_id") + ) + ) AS "invalid_sandbox_identities", + (SELECT count(*) + FROM "sandbox_session" AS "workspace" + LEFT JOIN "runtime_subject_identity" AS "identity" + ON "identity"."sandbox_id" = "workspace"."sandbox_id" + LEFT JOIN "session" AS "session" ON "session"."id" = "workspace"."session_id" + WHERE "identity"."sandbox_id" IS NULL + OR "session"."id" IS NULL + OR "session"."kind" IS NOT "identity"."kind" + OR "session"."agent_id" IS NOT "identity"."agent_id" + OR "session"."app_id" IS NOT "identity"."app_id" + OR ("identity"."subject_kind" = 'session' + AND "workspace"."session_id" IS NOT "identity"."subject_id") + OR ("identity"."subject_kind" = 'agent' + AND "session"."agent_id" IS NOT "identity"."subject_id") + ) AS "invalid_sandbox_session_authorities", + (SELECT count(*) + FROM "driver_instance" + WHERE typeof("generation") <> 'integer' + OR "generation" NOT BETWEEN 0 AND 9007199254740991 + ) AS "invalid_driver_generations", + ((SELECT count(*) + FROM "sandbox" + WHERE "last_backup_id" IS NOT NULL + AND NOT EXISTS ( + SELECT 1 FROM "sandbox_backup" + WHERE "sandbox_backup"."id" = "sandbox"."last_backup_id" + AND "sandbox_backup"."sandbox_id" = "sandbox"."id" + AND "sandbox_backup"."status" = 'ready' + )) + + + (SELECT count(*) + FROM "sandbox" + WHERE "last_restore_backup_id" IS NOT NULL + AND NOT EXISTS ( + SELECT 1 FROM "sandbox_backup" + WHERE "sandbox_backup"."id" = "sandbox"."last_restore_backup_id" + AND "sandbox_backup"."sandbox_id" = "sandbox"."id" + AND "sandbox_backup"."status" IN ('ready', 'pruned') + ))) AS "invalid_sandbox_backup_pointers", + (SELECT count(*) + FROM "app_deployment" + WHERE "last_successful_url" IS NOT NULL + ) AS "legacy_app_deployment_traffic", + ${nonstaticSessionsSql} AS "nonstatic_sessions", + (SELECT count(*) FROM pragma_foreign_key_check) AS "foreign_key_violations", + (SELECT count(*) + FROM "sandbox_backup" AS "backup" + LEFT JOIN "runtime_subject_identity" AS "identity" + ON "identity"."sandbox_id" = "backup"."sandbox_id" + LEFT JOIN "session_run" AS "run" ON "run"."id" = "backup"."session_run_id" + LEFT JOIN "session" AS "run_session" ON "run_session"."id" = "run"."session_id" + WHERE "backup"."status" NOT IN ('ready', 'pruned') + OR "backup"."error_message" IS NOT NULL + OR typeof("backup"."dir") <> 'text' + OR length("backup"."dir") = 0 + OR typeof("backup"."keep") <> 'integer' + OR "backup"."keep" NOT IN (0, 1) + OR typeof("backup"."ttl_seconds") <> 'integer' + OR "backup"."ttl_seconds" NOT BETWEEN 1 AND 9007199254740991 + OR typeof("backup"."created_at") <> 'integer' + OR "backup"."created_at" NOT BETWEEN 0 AND 9007199254740991 + OR typeof("backup"."updated_at") <> 'integer' + OR "backup"."updated_at" NOT BETWEEN "backup"."created_at" AND 9007199254740991 + OR ("backup"."session_run_id" IS NOT NULL AND ( + "identity"."sandbox_id" IS NULL + OR "run"."id" IS NULL + OR "run"."status" IS NOT 'completed' + OR "run"."agent_id" IS NOT "identity"."agent_id" + OR "run_session"."id" IS NULL + OR "run_session"."kind" IS NOT "identity"."kind" + OR NOT EXISTS ( + SELECT 1 + FROM "sandbox_session" AS "workspace" + WHERE "workspace"."session_id" = "run"."session_id" + AND "workspace"."sandbox_id" = "backup"."sandbox_id" + AND "workspace"."cwd" = "backup"."dir" + ) + OR ("identity"."subject_kind" = 'session' + AND "run_session"."id" IS NOT "identity"."subject_id") + OR ("identity"."subject_kind" = 'agent' AND ( + "run_session"."agent_id" IS NOT "identity"."agent_id" + OR "run_session"."app_id" IS NOT "identity"."app_id" + )) + )) + ) AS "invalid_sandbox_backups", + (SELECT count(*) + FROM ( + SELECT 1 + FROM "sandbox_backup" + WHERE "session_run_id" IS NOT NULL + GROUP BY "sandbox_id", "dir", "session_run_id" + HAVING count(*) > 1 + ) + ) AS "duplicate_sandbox_backups"; +`; +} + +export const PROTOCOL_V3_RUNTIME_AUTHORITY_PREFLIGHT_SQL = + protocolV3RuntimeAuthorityPreflightSql(false); + +export const PROTOCOL_V3_CUTOVER_BOOKMARK_SQL = ` +SELECT "pre_migration_bookmark" AS "bookmark" +FROM "${PROTOCOL_V3_CUTOVER_TABLE}" +WHERE "id" = 1; +`; + +export const PROTOCOL_V3_SMOKE_SESSION_SQL = ` +SELECT "smoke_session_id" AS "smoke_session_id" +FROM "${PROTOCOL_V3_CUTOVER_TABLE}" +WHERE "id" = 1; +`; + +export const PROTOCOL_V3_SMOKE_REQUEST_KEY_SQL = ` +SELECT "smoke_request_key" AS "smoke_request_key" +FROM "${PROTOCOL_V3_CUTOVER_TABLE}" +WHERE "id" = 1; +`; + +export const ENTER_PROTOCOL_V3_DRAIN_SQL = ` +UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" +SET + "command_freeze" = 0, + "smoke_account_id" = NULL, + "smoke_request_key" = NULL, + "smoke_session_id" = NULL +WHERE "id" = 1 AND "enabled" = 1 AND "phase" = 'draining'; +`; + +export const ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL = ` +UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" +SET "command_freeze" = 1 +WHERE "id" = 1 + AND "enabled" = 1 + AND "phase" = 'draining' + AND NOT EXISTS ( + SELECT 1 FROM "api_command" WHERE "status" IN ('queued', 'running') + ) + AND NOT EXISTS ( + SELECT 1 FROM "app_deployment_run" + WHERE "status" IN (${ACTIVE_APP_DEPLOYMENT_RUN_STATUSES}) + ); +`; + +export const PROTOCOL_V3_COMMAND_FREEZE_SQL = ` +SELECT + "command_freeze" AS "command_freeze", + "release_tree_oid" AS "release_tree_oid", + "enabled" AS "enabled", + "phase" AS "phase", + "target_container_application_version" AS "target_container_application_version", + "target_container_image_digest" AS "target_container_image_digest", + "target_worker_version_id" AS "target_worker_version_id", + EXISTS ( + SELECT 1 + FROM "${PROTOCOL_V3_MIGRATION_INTENT_TABLE}" AS "intent" + WHERE "intent"."id" = "gate"."id" + ) AS "migration_started", + (SELECT "sql" FROM "sqlite_master" + WHERE "type" = 'table' AND "name" = '${PROTOCOL_V3_MIGRATION_INTENT_TABLE}') AS "migration_intent_table_sql", + (SELECT count(*) FROM "sqlite_master" + WHERE "type" = 'trigger' AND "tbl_name" = '${PROTOCOL_V3_MIGRATION_INTENT_TABLE}' COLLATE BINARY) AS "migration_intent_trigger_count" +FROM "${PROTOCOL_V3_CUTOVER_TABLE}" AS "gate" +WHERE "gate"."id" = 1; +`; + +export function beginProtocolV3MigrationSql(releaseTreeOid: string): string { + const release = requireGitTreeOid(releaseTreeOid, "Protocol v3 release tree OID"); + return `INSERT INTO "${PROTOCOL_V3_MIGRATION_INTENT_TABLE}" ("id") +SELECT "id" +FROM "${PROTOCOL_V3_CUTOVER_TABLE}" +WHERE "id" = 1 + AND "enabled" = 1 + AND "command_freeze" = 1 + AND "phase" = 'draining' + AND "release_tree_oid" = '${release}' +ON CONFLICT ("id") DO NOTHING;`; +} + +export function storeProtocolV3RolloutSql( + releaseTreeOid: string, + workerVersionId: string, + containerApplicationVersion: number, + containerImageDigest: string, +): string { + const release = requireGitTreeOid(releaseTreeOid, "Protocol v3 release tree OID"); + const workerVersion = requireWorkerVersionId(workerVersionId); + if (!Number.isSafeInteger(containerApplicationVersion) || containerApplicationVersion < 0) { + throw new Error("Protocol v3 Container application version must be a non-negative integer."); + } + const imageDigest = requireSha256Digest( + containerImageDigest, + "Protocol v3 Container image digest", + ); + return `UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" +SET "target_container_application_version" = ${containerApplicationVersion}, + "target_container_image_digest" = '${imageDigest}', + "target_worker_version_id" = '${workerVersion}' +WHERE "id" = 1 + AND "enabled" = 1 + AND "phase" = 'draining' + AND "release_tree_oid" = '${release}' + AND ( + ("target_container_application_version" IS NULL AND "target_container_image_digest" IS NULL AND "target_worker_version_id" IS NULL) + OR ( + "target_container_application_version" = ${containerApplicationVersion} + AND "target_container_image_digest" = '${imageDigest}' + AND "target_worker_version_id" = '${workerVersion}' + ) + );`; +} + +export const ENTER_PROTOCOL_V3_QUEUES_RESUMING_SQL = ` +UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" +SET + "command_freeze" = 1, + "enabled" = 1, + "phase" = 'queues_resuming', + "smoke_account_id" = NULL, + "smoke_request_key" = NULL, + "smoke_session_id" = NULL +WHERE "id" = 1 AND "enabled" = 1 AND "phase" = 'draining'; +`; + +export const ACCEPT_PROTOCOL_V3_QUEUE_RESUME_SQL = ` +UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" +SET "enabled" = 0 +WHERE "id" = 1 + AND "enabled" = 1 + AND "command_freeze" = 1 + AND "phase" = 'queues_resuming'; +`; + +export const CLOSE_PROTOCOL_V3_SMOKE_WINDOW_SQL = ` +UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" +SET + "command_freeze" = 1, + "smoke_account_id" = NULL, + "smoke_request_key" = NULL, + "smoke_session_id" = NULL +WHERE "id" = 1 AND "enabled" = 1 AND "phase" = 'draining'; +`; + +export const PROTOCOL_V3_CUTOVER_OBJECTS_SQL = ` +WITH "post_schema" ("value") AS ( + SELECT EXISTS ( + SELECT 1 FROM "sqlite_master" + WHERE "type" = 'table' AND "name" = 'sandbox_backup_staging' COLLATE BINARY + ) +), +"pre_migration_expected" ("type", "name", "table_name", "sql") AS ( + VALUES ${PROTOCOL_V3_CUTOVER_OBJECTS.map( + (object) => + `(${sqlString(object.type)}, ${sqlString(object.name)}, ${sqlString(object.tableName)}, ${sqlString(object.sql)})`, + ).join(",\n ")} +), +"post_migration_expected" ("type", "name", "table_name", "sql") AS ( + VALUES ${PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECTS.map( + (object) => + `(${sqlString(object.type)}, ${sqlString(object.name)}, ${sqlString(object.tableName)}, ${sqlString(object.sql)})`, + ).join(",\n ")} +), +"expected" AS ( + SELECT * FROM "pre_migration_expected" WHERE (SELECT "value" FROM "post_schema") = 0 + UNION ALL + SELECT * FROM "post_migration_expected" WHERE (SELECT "value" FROM "post_schema") = 1 +), +"actual" AS ( + SELECT "type", "name", "tbl_name" AS "table_name", "sql" + FROM "sqlite_master" + WHERE "name" COLLATE NOCASE IN (${CUTOVER_RESERVED_NAMES_SQL}) + OR ( + "type" = 'trigger' + AND "tbl_name" COLLATE NOCASE IN (${CUTOVER_PROTECTED_TABLES_SQL}) + AND NOT ( + ${PROTOCOL_V3_LEGACY_REWRITE_GATE_UPDATE_TRIGGER_MATCH_SQL} + OR ${CUTOVER_MANAGED_TRIGGER_MATCH_SQL} + ) + ) +) +SELECT + count(*) AS "object_count", + CASE WHEN count(*) = (SELECT count(*) FROM "expected") + THEN coalesce(sum(EXISTS ( + SELECT 1 + FROM "expected" + WHERE "expected"."type" = "actual"."type" COLLATE BINARY + AND "expected"."name" = "actual"."name" COLLATE BINARY + AND "expected"."table_name" = "actual"."table_name" COLLATE BINARY + AND "expected"."sql" = "actual"."sql" COLLATE BINARY + )), 0) + ELSE 0 + END AS "exact_object_count" +FROM "actual"; +`; + +export interface ProtocolV3CutoverProbe { + readonly gatePresent: boolean; +} + +export function findPendingProdMigrations( + raw: string, + localMigrationNames: readonly string[], +): string[] { + const validateName = (name: unknown): string => { + if (typeof name !== "string" || !/^\d{4}_[a-z0-9_-]+\.sql$/u.test(name)) { + throw new Error("Production D1 migration name is invalid."); + } + return name; + }; + const local = localMigrationNames.map(validateName); + if (new Set(local).size !== local.length) { + throw new Error("Local D1 migration names must be unique."); + } + + const applied = parseD1JsonResults(raw).flatMap((result) => + result.map((row) => validateName(row.name)), + ); + if (new Set(applied).size !== applied.length) { + throw new Error("Production D1 migration ledger contains duplicate names."); + } + const appliedSet = new Set(applied); + if ( + applied.length > local.length || + local.some((name, index) => appliedSet.has(name) !== index < applied.length) + ) { + throw new Error("Production D1 migration ledger is not an exact prefix of local history."); + } + + return local.filter((name) => !appliedSet.has(name)); +} + +export function assertCutoverMigrationJournalAudited(localMigrationNames: readonly string[]): void { + if ( + localMigrationNames.length !== AUDITED_MIGRATION_NAMES.length || + AUDITED_MIGRATION_NAMES.some((name, index) => localMigrationNames[index] !== name) + ) { + throw new Error( + `Production cutover safety is audited only through ${LAST_CUTOVER_AUDITED_MIGRATION}; review every new migration against the gate and drain before updating that boundary.`, + ); + } +} + +export interface ProtocolV3CutoverObjects { + readonly exactObjectCount: number; + readonly objectCount: number; +} + +export interface ProtocolV3LegacyTerminalIntegrity { + readonly ambiguousAssistantRuns: number; + readonly duplicateTerminalRuns: number; + readonly invalidFailedRuns: number; + readonly invalidNonfailedRunErrors: number; + readonly invalidTerminalLinks: number; + readonly legacyMaterializedMessages: number; + readonly legacyStreamRows: number; + readonly legacyTerminalEvents: number; + readonly mismatchedTerminalEvents: number; + readonly missingTerminalEvents: number; + readonly noncanonicalTerminalSources: number; + readonly nonterminalCommands: number; + readonly partialAssistantProjections: number; + readonly partialTerminalProjections: number; + readonly repairableFailedRuns: number; + readonly rewriteCandidateManifestJson: string; + readonly unsettledEffects: number; +} + +export interface ProtocolV3LegacyRewriteAuthorization { + readonly bookmark: string; + readonly candidateCount: number; + readonly candidateManifestJson: string; + readonly deployOwner: string; + readonly expiresAt: number; + readonly releaseTreeOid: string; +} + +export interface ProtocolV3LegacyTerminalSourceKindInventory { + readonly canonical: number; + readonly canonicalTargetCollisions: number; + readonly noncanonical: number; + readonly total: number; +} + +export interface ProtocolV3LegacyTerminalSourceInventory { + readonly cancelled: ProtocolV3LegacyTerminalSourceKindInventory; + readonly completed: ProtocolV3LegacyTerminalSourceKindInventory; + readonly failed: ProtocolV3LegacyTerminalSourceKindInventory; + readonly invalidTerminalLinks: number; + readonly mismatchedTerminalEvents: number; + readonly multipleTerminalRuns: number; +} + +export type ProtocolV3LossyMigrationCategory = + | "attempt_completion_time_fabrication" + | "command_error_omission" + | "command_payload_conflict" + | "control_reason_omission" + | "input_text_omission" + | "input_start_result_omission" + | "mcp_argument_omission" + | "mcp_command_terminal_conflict" + | "mcp_result_conflict" + | "mcp_result_omission" + | "orphan_effect" + | "provider_receipt_loss" + | "permission_payload_rewrite" + | "session_run_error_omission"; + +export interface ProtocolV3LossyMigrationCandidateId { + readonly category: ProtocolV3LossyMigrationCategory; + readonly id: string; +} + +export interface ProtocolV3LossyMigrationInventory { + readonly attemptCompletionTimeFabrications: number; + readonly candidateIds: readonly ProtocolV3LossyMigrationCandidateId[]; + readonly commandErrorOmissions: number; + readonly commandPayloadConflicts: number; + readonly controlReasonOmissions: number; + readonly inputTextOmissions: number; + readonly inputStartResultOmissions: number; + readonly mcpArgumentOmissions: number; + readonly mcpCommandTerminalConflicts: number; + readonly mcpResultConflicts: number; + readonly mcpResultOmissions: number; + readonly orphanEffects: number; + readonly providerReceiptLosses: number; + readonly permissionPayloadRewrites: number; + readonly sessionRunErrorOmissions: number; + readonly totalCandidates: number; +} + +export interface ProtocolV3CutoverDrain { + readonly activeAppDeploymentRuns: number; + readonly activeRuns: number; + readonly liveDrivers: number; + readonly nonterminalApiCommands: number; + readonly nonterminalCommands: number; + readonly unsafeEnvironmentArtifactBackupStaging: number; + readonly unsafeSandboxBackups: number; + readonly unsafeSandboxBackupStaging: number; + readonly unsafeSandboxes: number; + readonly unsafeSandboxSessions: number; + readonly unsafeSessions: number; + readonly unsettledEffects: number; +} + +export interface ProtocolV3CutoverState { + readonly commandFreeze: boolean; + readonly containerApplicationVersion: number | null; + readonly containerImageDigest: string | null; + readonly enabled: boolean; + readonly migrationStarted: boolean; + readonly phase: ProtocolV3CutoverPhase; + readonly releaseTreeOid: string; + readonly workerVersionId: string | null; +} + +export interface ProtocolV3RuntimeAuthorityPreflight { + readonly duplicateSandboxBackups: number; + readonly foreignKeyViolations: number; + readonly invalidDriverGenerations: number; + readonly invalidSandboxBackupPointers: number; + readonly invalidSandboxBackups: number; + readonly invalidSandboxIdentities: number; + readonly invalidSandboxSessionAuthorities: number; + readonly legacyAppDeploymentTraffic: number; + readonly nonstaticSessions: number; +} + +export interface ProtocolV3WorkerDeployment { + readonly versionId: string; +} + +export interface ProtocolV3SmokeStatus { + readonly bootTokenUsedAt: number | null; + readonly connectionId: string | null; + readonly driverPid: number | null; + readonly driverStartedAt: number | null; + readonly driverStatus: string | null; + readonly driverVersion: string | null; + readonly protocolVersion: number | null; + readonly statusEvent: string | null; +} + +export interface ProtocolV3ContainerInstance { + readonly state: string; + readonly version: string | number | null; +} + +export interface ProtocolV3ContainerApplication { + readonly id: string; + readonly imageDigest: string; + readonly imageRepository: string; + readonly name: string; + readonly state: "active" | "degraded" | "provisioning" | "ready"; + readonly version: number; +} + +interface ProtocolV3ContainerPage { + readonly items: readonly T[]; + readonly nextPageToken: string | null; + readonly pageToken: string | null; +} + +export interface ProtocolV3CutoverRecoveryState { + readonly bookmark: string | null; + readonly initialPendingMigrations: readonly string[]; + readonly migrationStarted: boolean; + readonly originalError: unknown; + readonly queuesVerified: boolean; +} + +export interface ProtocolV3CutoverRecoveryEffects { + readonly commitQueueAcceptance: () => void; + readonly pauseAndVerifyQueues: () => Promise; + readonly printBookmark: (bookmark: string) => void; + readonly probe: () => ProtocolV3CutoverProbe; + readonly readBookmark: () => string | null; + readonly readPendingMigrations: () => readonly string[]; + readonly removeMarker: () => void; + readonly resumeAndVerifyQueues: () => Promise; + readonly write: (message: string) => void; +} + +export interface ProtocolV3QueueResumeEffects { + readonly commitAcceptance: () => void; + readonly removeMarker: () => void; + readonly resumeAndVerifyQueues: () => Promise; +} + +type JsonRow = D1JsonRow; +const MAX_PROTOCOL_V3_CONTAINER_PAGES = 100; + +function requireNonNegativeInteger(row: JsonRow, key: string): number { + const value = row[key]; + + if (!Number.isSafeInteger(value) || (value as number) < 0) { + throw new Error(`D1 JSON field ${key} must be a non-negative safe integer.`); + } + + return value as number; +} + +function requireNullableNonNegativeInteger(row: JsonRow, key: string): number | null { + return row[key] === null ? null : requireNonNegativeInteger(row, key); +} + +function requireNullableString(row: JsonRow, key: string): string | null { + const value = row[key]; + if (value !== null && typeof value !== "string") { + throw new Error(`D1 JSON field ${key} must be a string or null.`); + } + return value; +} + +function requireJsonArrayString(row: JsonRow, key: string): string { + const value = row[key]; + if (typeof value !== "string") throw new Error(`D1 JSON field ${key} must be a string.`); + let parsed: unknown; + try { + parsed = JSON.parse(value); + } catch { + throw new Error(`D1 JSON field ${key} must contain valid JSON.`); + } + if (!Array.isArray(parsed)) throw new Error(`D1 JSON field ${key} must contain an array.`); + return value; +} + +function requirePageToken(value: unknown, label: string): string | null { + if (value === null) return null; + if ( + typeof value !== "string" || + value.length === 0 || + value.length > 4_096 || + [...value].some((character) => { + const codePoint = character.codePointAt(0) ?? 0; + return codePoint < 32 || codePoint === 127; + }) + ) { + throw new Error(`${label} must be a valid opaque token or null.`); + } + return value; +} + +function requirePlatformId(value: string, label: string): string { + if (!/^[0-7][0-9A-HJKMNP-TV-Z]{25}$/iu.test(value)) { + throw new Error(`${label} must be a ULID.`); + } + return value; +} + +function requireGitTreeOid(value: string, label: string): string { + if (!GIT_TREE_OID_PATTERN.test(value)) throw new Error(`${label} must be a Git tree OID.`); + return value; +} + +function requireSha256Digest(value: string, label: string): string { + if (!SHA256_DIGEST_PATTERN.test(value)) throw new Error(`${label} must be a SHA-256 digest.`); + return value; +} + +function requireWorkerVersionId(value: string): string { + if (!WORKER_VERSION_ID_PATTERN.test(value)) { + throw new Error("Protocol v3 Worker version ID must be a UUID."); + } + return value; +} + +export function parseCleanGitTreeOid( + treeOutput: string, + statusOutput: string, + indexOutput = "", + ignoredBuildInputOutput = "", + environmentKeys: readonly string[] = [], +): string { + if (statusOutput.trim().length > 0 || /^(?:S|[a-z]) /mu.test(indexOutput)) { + throw new Error("Production deploy requires a clean Git worktree and clean submodules."); + } + if ( + ignoredBuildInputOutput.trim().length > 0 || + environmentKeys.some((key) => key.startsWith("VITE_")) + ) { + throw new Error("Production deploy rejects local or process-level Vite build inputs."); + } + return requireGitTreeOid(treeOutput.trim(), "Production release tree OID"); +} + +export function protocolV3ReleaseTag(releaseTreeOid: string): string { + return `protocol-v3-${requireGitTreeOid(releaseTreeOid, "Protocol v3 release tree OID")}`; +} + +export function assertProtocolV3Release( + state: ProtocolV3CutoverState, + releaseTreeOid: string, +): void { + if (state.releaseTreeOid !== requireGitTreeOid(releaseTreeOid, "Protocol v3 release tree OID")) { + throw new Error( + `Protocol v3 cutover belongs to release tree ${state.releaseTreeOid}, not ${releaseTreeOid}.`, + ); + } +} + +export function parseProtocolV3WorkerDeployment(raw: string): ProtocolV3WorkerDeployment { + const deployment = JSON.parse(raw) as { versions?: unknown }; + if (!Array.isArray(deployment.versions) || deployment.versions.length !== 1) { + throw new Error("Protocol v3 requires exactly one deployed Worker version."); + } + const [version] = deployment.versions as Array>; + if (version?.percentage !== 100 || typeof version.version_id !== "string") { + throw new Error("Protocol v3 Worker release must receive 100% of production traffic."); + } + return { versionId: requireWorkerVersionId(version.version_id) }; +} + +export function assertProtocolV3WorkerVersion( + raw: string, + expectedVersionId: string, + releaseTreeOid: string, +): void { + const version = JSON.parse(raw) as { + annotations?: Record; + id?: unknown; + }; + const versionId = requireWorkerVersionId(expectedVersionId); + if ( + version.id !== versionId || + version.annotations?.["workers/tag"] !== protocolV3ReleaseTag(releaseTreeOid) + ) { + throw new Error("Production Worker version is not the exact protocol v3 release."); + } +} + +export function protocolV3ContainerImageTag( + imageRepository: string, + workerVersionId: string, +): string { + if (imageRepository.trim() !== imageRepository || !imageRepository.includes("/")) { + throw new Error("Protocol v3 Container image repository is invalid."); + } + return `${imageRepository}:${requireWorkerVersionId(workerVersionId).slice(0, 8)}`; +} + +export function parseProtocolV3ContainerManifestDigest(raw: string): string { + const manifest = JSON.parse(raw) as { Descriptor?: { digest?: unknown } }; + const digest = manifest.Descriptor?.digest; + if (typeof digest !== "string" || !digest.startsWith("sha256:")) { + throw new Error("Protocol v3 Container manifest is missing its SHA-256 digest."); + } + return requireSha256Digest(digest.slice("sha256:".length), "Protocol v3 Container manifest"); +} + +function requireSmokeRequestKey(value: string): string { + if ( + !/^protocol-v3-cutover-[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/u.test( + value, + ) + ) { + throw new Error("Protocol v3 smoke request key must contain a UUID v4."); + } + return value; +} + +export function parseProtocolV3CutoverProbe(raw: string): ProtocolV3CutoverProbe { + const row = requireSingleD1Row(raw); + const gatePresent = requireNonNegativeInteger(row, "gate_present"); + if (gatePresent > 1) throw new Error("D1 cutover probe field must be zero or one."); + return { gatePresent: gatePresent === 1 }; +} + +export function parseProtocolV3LegacyTerminalIntegrity( + raw: string, +): ProtocolV3LegacyTerminalIntegrity { + const row = requireSingleD1Row(raw); + const noncanonicalTerminalSources = requireNonNegativeInteger( + row, + "noncanonical_terminal_sources", + ); + const rewriteCandidateManifestJson = requireJsonArrayString( + row, + "rewrite_candidate_manifest_json", + ); + if ( + (JSON.parse(rewriteCandidateManifestJson) as unknown[]).length !== noncanonicalTerminalSources + ) { + throw new Error("Legacy rewrite candidate count and manifest disagree."); + } + + return { + ambiguousAssistantRuns: requireNonNegativeInteger(row, "ambiguous_assistant_runs"), + duplicateTerminalRuns: requireNonNegativeInteger(row, "duplicate_terminal_runs"), + invalidFailedRuns: requireNonNegativeInteger(row, "invalid_failed_runs"), + invalidNonfailedRunErrors: requireNonNegativeInteger(row, "invalid_nonfailed_run_errors"), + invalidTerminalLinks: requireNonNegativeInteger(row, "invalid_terminal_links"), + legacyMaterializedMessages: requireNonNegativeInteger(row, "legacy_materialized_messages"), + legacyStreamRows: requireNonNegativeInteger(row, "legacy_stream_rows"), + legacyTerminalEvents: requireNonNegativeInteger(row, "legacy_terminal_events"), + mismatchedTerminalEvents: requireNonNegativeInteger(row, "mismatched_terminal_events"), + missingTerminalEvents: requireNonNegativeInteger(row, "missing_terminal_events"), + noncanonicalTerminalSources, + nonterminalCommands: requireNonNegativeInteger(row, "nonterminal_commands"), + partialAssistantProjections: requireNonNegativeInteger(row, "partial_assistant_projections"), + partialTerminalProjections: requireNonNegativeInteger(row, "partial_terminal_projections"), + repairableFailedRuns: requireNonNegativeInteger(row, "repairable_failed_runs"), + rewriteCandidateManifestJson, + unsettledEffects: requireNonNegativeInteger(row, "unsettled_effects"), + }; +} + +export function parseProtocolV3LegacyRewriteAuthorization( + raw: string, +): ProtocolV3LegacyRewriteAuthorization { + const row = requireSingleD1Row(raw); + if (row.authorization_table_sql !== PROTOCOL_V3_LEGACY_REWRITE_AUTHORIZATION_TABLE_SQL) { + throw new Error("Protocol v3 legacy rewrite authorization table schema is invalid."); + } + if (row.gate_update_trigger_sql !== PROTOCOL_V3_LEGACY_REWRITE_GATE_UPDATE_TRIGGER_SQL) { + throw new Error("Protocol v3 legacy rewrite gate update guard is invalid."); + } + const candidateCount = requireNonNegativeInteger(row, "candidate_count"); + const candidateManifestJson = requireJsonArrayString(row, "candidate_manifest_json"); + if ((JSON.parse(candidateManifestJson) as unknown[]).length !== candidateCount) { + throw new Error("Legacy rewrite authorization count and manifest disagree."); + } + const bookmark = row.bookmark; + if (typeof bookmark !== "string" || !/^[0-9a-z-]{1,256}$/iu.test(bookmark)) { + throw new Error("Protocol v3 legacy rewrite authorization has an invalid bookmark."); + } + const deployOwner = row.deploy_owner; + if ( + typeof deployOwner !== "string" || + !/^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/u.test(deployOwner) + ) { + throw new Error("Protocol v3 legacy rewrite authorization has an invalid owner."); + } + return { + bookmark, + candidateCount, + candidateManifestJson, + deployOwner, + expiresAt: requireNonNegativeInteger(row, "expires_at"), + releaseTreeOid: requireGitTreeOid( + typeof row.release_tree_oid === "string" ? row.release_tree_oid : "", + "Protocol v3 legacy rewrite release tree OID", + ), + }; +} + +export function parseProtocolV3LegacyTerminalSourceInventory( + raw: string, +): ProtocolV3LegacyTerminalSourceInventory { + const row = requireSingleD1Row(raw); + const readKind = (kind: "cancelled" | "completed" | "failed") => { + const inventory = { + canonical: requireNonNegativeInteger(row, `${kind}_canonical`), + canonicalTargetCollisions: requireNonNegativeInteger( + row, + `${kind}_canonical_target_collisions`, + ), + noncanonical: requireNonNegativeInteger(row, `${kind}_noncanonical`), + total: requireNonNegativeInteger(row, `${kind}_total`), + }; + if ( + inventory.canonical + inventory.noncanonical !== inventory.total || + inventory.canonicalTargetCollisions > inventory.total + ) { + throw new Error(`Legacy ${kind} terminal source inventory is inconsistent.`); + } + return inventory; + }; + + return { + cancelled: readKind("cancelled"), + completed: readKind("completed"), + failed: readKind("failed"), + invalidTerminalLinks: requireNonNegativeInteger(row, "invalid_terminal_links"), + mismatchedTerminalEvents: requireNonNegativeInteger(row, "mismatched_terminal_events"), + multipleTerminalRuns: requireNonNegativeInteger(row, "multiple_terminal_runs"), + }; +} + +const PROTOCOL_V3_LOSSY_MIGRATION_CATEGORIES = new Set([ + "attempt_completion_time_fabrication", + "command_error_omission", + "command_payload_conflict", + "control_reason_omission", + "input_text_omission", + "input_start_result_omission", + "mcp_argument_omission", + "mcp_command_terminal_conflict", + "mcp_result_conflict", + "mcp_result_omission", + "orphan_effect", + "provider_receipt_loss", + "permission_payload_rewrite", + "session_run_error_omission", +]); + +export function parseProtocolV3LossyMigrationInventory( + raw: string, +): ProtocolV3LossyMigrationInventory { + const row = requireSingleD1Row(raw); + const inventory = { + attemptCompletionTimeFabrications: requireNonNegativeInteger( + row, + "attempt_completion_time_fabrications", + ), + commandErrorOmissions: requireNonNegativeInteger(row, "command_error_omissions"), + commandPayloadConflicts: requireNonNegativeInteger(row, "command_payload_conflicts"), + controlReasonOmissions: requireNonNegativeInteger(row, "control_reason_omissions"), + inputTextOmissions: requireNonNegativeInteger(row, "input_text_omissions"), + inputStartResultOmissions: requireNonNegativeInteger(row, "input_start_result_omissions"), + mcpArgumentOmissions: requireNonNegativeInteger(row, "mcp_argument_omissions"), + mcpCommandTerminalConflicts: requireNonNegativeInteger(row, "mcp_command_terminal_conflicts"), + mcpResultConflicts: requireNonNegativeInteger(row, "mcp_result_conflicts"), + mcpResultOmissions: requireNonNegativeInteger(row, "mcp_result_omissions"), + orphanEffects: requireNonNegativeInteger(row, "orphan_effects"), + providerReceiptLosses: requireNonNegativeInteger(row, "provider_receipt_losses"), + permissionPayloadRewrites: requireNonNegativeInteger(row, "permission_payload_rewrites"), + sessionRunErrorOmissions: requireNonNegativeInteger(row, "session_run_error_omissions"), + totalCandidates: requireNonNegativeInteger(row, "total_candidates"), + }; + const countedCandidates = + inventory.attemptCompletionTimeFabrications + + inventory.commandErrorOmissions + + inventory.commandPayloadConflicts + + inventory.controlReasonOmissions + + inventory.inputTextOmissions + + inventory.inputStartResultOmissions + + inventory.mcpArgumentOmissions + + inventory.mcpCommandTerminalConflicts + + inventory.mcpResultConflicts + + inventory.mcpResultOmissions + + inventory.orphanEffects + + inventory.providerReceiptLosses + + inventory.permissionPayloadRewrites + + inventory.sessionRunErrorOmissions; + if (countedCandidates !== inventory.totalCandidates) { + throw new Error("Lossy migration inventory category counts are inconsistent."); + } + + const parsedIds = JSON.parse(requireJsonArrayString(row, "candidate_ids_json")) as unknown[]; + if (parsedIds.length !== Math.min(inventory.totalCandidates, 50)) { + throw new Error("Lossy migration inventory candidate ID count is inconsistent."); + } + let previousKey: string | null = null; + const candidateIds = parsedIds.map((entry): ProtocolV3LossyMigrationCandidateId => { + if (!Array.isArray(entry) || entry.length !== 2) { + throw new Error("Lossy migration inventory candidate IDs are invalid."); + } + const [category, id] = entry; + if ( + typeof category !== "string" || + !PROTOCOL_V3_LOSSY_MIGRATION_CATEGORIES.has(category as ProtocolV3LossyMigrationCategory) || + typeof id !== "string" + ) { + throw new Error("Lossy migration inventory candidate IDs are invalid."); + } + const key = `${category}\0${requirePlatformId(id, "Lossy migration candidate ID")}`; + if (previousKey !== null && previousKey >= key) { + throw new Error("Lossy migration inventory candidate IDs are not stable and unique."); + } + previousKey = key; + return { category: category as ProtocolV3LossyMigrationCategory, id }; + }); + + return { ...inventory, candidateIds }; +} + +export function assertProtocolV3LossyMigrationInventory( + inventory: ProtocolV3LossyMigrationInventory, +): void { + if (inventory.totalCandidates === 0) return; + throw new Error( + `Migration 0013 would discard or overwrite production history: ${[ + ["attemptCompletionTimeFabrications", inventory.attemptCompletionTimeFabrications], + ["mcpArgumentOmissions", inventory.mcpArgumentOmissions], + ["commandPayloadConflicts", inventory.commandPayloadConflicts], + ["inputTextOmissions", inventory.inputTextOmissions], + ["inputStartResultOmissions", inventory.inputStartResultOmissions], + ["controlReasonOmissions", inventory.controlReasonOmissions], + ["mcpResultOmissions", inventory.mcpResultOmissions], + ["orphanEffects", inventory.orphanEffects], + ["mcpResultConflicts", inventory.mcpResultConflicts], + ["providerReceiptLosses", inventory.providerReceiptLosses], + ["permissionPayloadRewrites", inventory.permissionPayloadRewrites], + ["mcpCommandTerminalConflicts", inventory.mcpCommandTerminalConflicts], + ["commandErrorOmissions", inventory.commandErrorOmissions], + ["sessionRunErrorOmissions", inventory.sessionRunErrorOmissions], + ] + .filter(([, count]) => count !== 0) + .map(([category, count]) => `${category}=${count}`) + .join(" ")}. No lossy migration candidates are authorized.`, + ); +} + +export function assertProtocolV3LegacyTerminalSourceInventory( + inventory: ProtocolV3LegacyTerminalSourceInventory, +): void { + if ( + inventory.invalidTerminalLinks > 0 || + inventory.mismatchedTerminalEvents > 0 || + inventory.multipleTerminalRuns > 0 || + [inventory.cancelled, inventory.completed, inventory.failed].some( + (entry) => entry.canonicalTargetCollisions > 0, + ) + ) { + throw new Error( + "Legacy terminal source inventory is not ready for the protocol v3 production cutover.", + ); + } +} + +export function assertProtocolV3LegacyTerminalIntegrity( + integrity: ProtocolV3LegacyTerminalIntegrity, +): void { + const blockers = { + ambiguousAssistantRuns: integrity.ambiguousAssistantRuns, + duplicateTerminalRuns: integrity.duplicateTerminalRuns, + invalidFailedRuns: integrity.invalidFailedRuns, + invalidNonfailedRunErrors: integrity.invalidNonfailedRunErrors, + invalidTerminalLinks: integrity.invalidTerminalLinks, + mismatchedTerminalEvents: integrity.mismatchedTerminalEvents, + missingTerminalEvents: integrity.missingTerminalEvents, + nonterminalCommands: integrity.nonterminalCommands, + partialAssistantProjections: integrity.partialAssistantProjections, + partialTerminalProjections: integrity.partialTerminalProjections, + unsettledEffects: integrity.unsettledEffects, + }; + + if (Object.values(blockers).some((count) => count > 0)) { + throw new Error( + `Legacy terminal integrity is ambiguous and cannot be migrated: ${Object.entries(blockers) + .map(([name, count]) => `${name}=${count}`) + .join(" ")}.`, + ); + } +} + +function parseProtocolV3ContainerApplicationPage( + value: unknown, +): ProtocolV3ContainerPage { + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new Error("Container applications API output must be an object."); + } + + const applications = "result" in value ? value.result : null; + const resultInfo = "result_info" in value ? value.result_info : null; + if ( + !("success" in value) || + value.success !== true || + !Array.isArray(applications) || + typeof resultInfo !== "object" || + resultInfo === null + ) { + throw new Error("Container applications API output is missing pagination fields."); + } + + return { + items: applications.map((application) => { + if ( + typeof application !== "object" || + application === null || + !("id" in application) || + typeof application.id !== "string" || + !/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/iu.test(application.id) || + !("name" in application) || + typeof application.name !== "string" || + application.name.length === 0 || + !("version" in application) || + !("health" in application) || + typeof application.health !== "object" || + application.health === null || + !("instances" in application.health) || + typeof application.health.instances !== "object" || + application.health.instances === null + ) { + throw new Error("Container application API output is invalid."); + } + + const configuration = "configuration" in application ? application.configuration : null; + const image = + typeof configuration === "object" && + configuration !== null && + "image" in configuration && + typeof configuration.image === "string" + ? configuration.image + : ""; + const imageMatch = /^(.+)@sha256:([0-9a-f]{64})$/u.exec(image); + if (imageMatch === null || imageMatch[1]?.trim() !== imageMatch[1]) { + throw new Error("Container application image is not content-addressed."); + } + + const health = application.health.instances as JsonRow; + const active = requireNonNegativeInteger(health, "active"); + const failed = requireNonNegativeInteger(health, "failed"); + requireNonNegativeInteger(health, "healthy"); + const scheduling = requireNonNegativeInteger(health, "scheduling"); + const starting = requireNonNegativeInteger(health, "starting"); + const version = requireNonNegativeInteger(application as JsonRow, "version"); + const state = + failed > 0 + ? "degraded" + : starting > 0 || scheduling > 0 + ? "provisioning" + : active > 0 + ? "active" + : "ready"; + + return { + id: application.id, + imageDigest: requireSha256Digest(imageMatch[2] ?? "", "Container application image"), + imageRepository: imageMatch[1] ?? "", + name: application.name, + state, + version, + }; + }), + nextPageToken: requirePageToken( + "next_page_token" in resultInfo ? resultInfo.next_page_token : undefined, + "Container application next page token", + ), + pageToken: requirePageToken( + "page_token" in resultInfo ? resultInfo.page_token : undefined, + "Container application page token", + ), + }; +} + +function parseProtocolV3ContainerInstancePage( + raw: string, +): ProtocolV3ContainerPage { + const value = JSON.parse(raw) as unknown; + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new Error("Container instances JSON output must be a paginated object."); + } + + const instances = "instances" in value ? value.instances : null; + const resultInfo = "result_info" in value ? value.result_info : null; + if (!Array.isArray(instances) || typeof resultInfo !== "object" || resultInfo === null) { + throw new Error("Container instances JSON output is missing pagination fields."); + } + + return { + items: instances.map((instance) => { + if ( + typeof instance !== "object" || + instance === null || + !("state" in instance) || + typeof instance.state !== "string" || + !("version" in instance) || + (instance.version !== null && + typeof instance.version !== "string" && + typeof instance.version !== "number") + ) { + throw new Error("Container instance JSON output is invalid."); + } + return { state: instance.state, version: instance.version }; + }), + nextPageToken: requirePageToken( + "next_page_token" in resultInfo ? resultInfo.next_page_token : undefined, + "Container next page token", + ), + pageToken: requirePageToken( + "page_token" in resultInfo ? resultInfo.page_token : undefined, + "Container page token", + ), + }; +} + +async function collectPages( + readPage: (pageToken: string | null) => Promise | Raw, + parsePage: (raw: Raw) => ProtocolV3ContainerPage, + label: string, +): Promise { + const items: T[] = []; + const seenPageTokens = new Set(); + let pageToken: string | null = null; + let pageCount = 0; + + while (true) { + pageCount += 1; + if (pageCount > MAX_PROTOCOL_V3_CONTAINER_PAGES) { + throw new Error(`${label} pagination exceeded its safety limit.`); + } + if (pageToken !== null && seenPageTokens.has(pageToken)) { + throw new Error(`${label} pagination returned a repeated page token.`); + } + if (pageToken !== null) seenPageTokens.add(pageToken); + + const page = parsePage(await readPage(pageToken)); + if (page.pageToken !== pageToken) { + throw new Error(`${label} pagination returned the wrong page token.`); + } + items.push(...page.items); + if (page.nextPageToken === null) return items; + pageToken = page.nextPageToken; + } +} + +export function collectProtocolV3ContainerInstances( + readPage: (pageToken: string | null) => Promise | string, +): Promise { + return collectPages(readPage, parseProtocolV3ContainerInstancePage, "Container instance"); +} + +export function collectProtocolV3ContainerApplications( + readPage: (pageToken: string | null) => Promise | unknown, +): Promise { + return collectPages(readPage, parseProtocolV3ContainerApplicationPage, "Container application"); +} + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +export async function recoverProtocolV3CutoverFailure( + state: ProtocolV3CutoverRecoveryState, + effects: ProtocolV3CutoverRecoveryEffects, +): Promise { + if (state.queuesVerified) { + try { + await effects.resumeAndVerifyQueues(); + if (!effects.probe().gatePresent) { + effects.write(" protocol v3 queue and marker recovery was already complete"); + return; + } + effects.commitQueueAcceptance(); + effects.removeMarker(); + effects.write(" protocol v3 queue and marker recovery completed after acceptance"); + return; + } catch (recoveryError) { + effects.write( + `✗ Protocol v3 was accepted, but final queue/marker confirmation failed: ${errorMessage(recoveryError)}`, + ); + effects.write( + " Queue delivery was already verified open and will not be re-paused after this commit point.", + ); + effects.write( + " Rerun this exact v3 release to repeat queue verification and marker cleanup.", + ); + throw new AggregateError( + [state.originalError, recoveryError], + "Protocol v3 post-acceptance cleanup remains incomplete.", + { cause: recoveryError }, + ); + } + } + + let migrationMayHaveCommitted = + state.initialPendingMigrations.length === 0 || state.migrationStarted; + if (!migrationMayHaveCommitted) { + try { + const stillPending = new Set(effects.readPendingMigrations()); + migrationMayHaveCommitted = state.initialPendingMigrations.some( + (name) => !stillPending.has(name), + ); + } catch (migrationProbeError) { + migrationMayHaveCommitted = true; + effects.write( + `✗ Could not prove every initially pending migration remained unapplied; keeping production closed: ${errorMessage(migrationProbeError)}`, + ); + } + } + + const reportFailures = (label: string, failures: readonly string[]) => { + if (failures.length === 0) return; + effects.write(`✗ ${label}:`); + for (const failure of failures) effects.write(` ${failure}`); + }; + + if (!migrationMayHaveCommitted) { + const failures: string[] = []; + try { + await effects.resumeAndVerifyQueues(); + effects.removeMarker(); + } catch (recoveryError) { + failures.push(`queue or gate recovery: ${errorMessage(recoveryError)}`); + try { + await effects.pauseAndVerifyQueues(); + } catch (pauseError) { + failures.push(`queue re-pause verification: ${errorMessage(pauseError)}`); + } + try { + if (!effects.probe().gatePresent) { + try { + await effects.resumeAndVerifyQueues(); + failures.length = 0; + } catch (resumeError) { + failures.push(`old service queue resume verification: ${errorMessage(resumeError)}`); + try { + await effects.pauseAndVerifyQueues(); + } catch (pauseError) { + failures.push(`final queue re-pause verification: ${errorMessage(pauseError)}`); + } + } + } + } catch (probeError) { + failures.push(`gate cleanup readback: ${errorMessage(probeError)}`); + } + } + reportFailures("Old service recovery was incomplete", failures); + if (failures.length === 0) { + effects.write(" no migration committed; the old service admission and queues were restored"); + } + } else { + try { + await effects.pauseAndVerifyQueues(); + } catch (pauseError) { + reportFailures("Failed to re-pause and verify every queue", [errorMessage(pauseError)]); + } + let bookmark = state.bookmark; + if (bookmark === null) { + try { + bookmark = effects.readBookmark(); + } catch { + // The bookmark is an emergency destructive backup, not a roll-forward prerequisite. + } + } + effects.write( + "✗ A migration request may have committed; queue delivery remains paused. Do not roll back only the Worker.", + ); + effects.write(" Repair only by rolling forward this exact v3 release."); + if (bookmark !== null) effects.printBookmark(bookmark); + } + + throw state.originalError; +} + +export async function completeProtocolV3QueueResume( + state: ProtocolV3CutoverState, + effects: ProtocolV3QueueResumeEffects, +): Promise { + if (state.phase !== "queues_resuming" || !state.commandFreeze) { + throw new Error("Protocol v3 queue resume requires its durable frozen phase."); + } + await effects.resumeAndVerifyQueues(); + if (state.enabled) effects.commitAcceptance(); + effects.removeMarker(); +} + +export function isProtocolV3ContainerRolloutConverged( + application: { readonly state: string; readonly version: string | number }, + instances: readonly ProtocolV3ContainerInstance[], +): boolean { + if (application.state !== "active" && application.state !== "ready") return false; + + return instances.every( + (instance) => + instance.state === "inactive" || + (instance.state === "running" && String(instance.version) === String(application.version)), + ); +} + +export function parseProtocolV3CutoverDrain(raw: string): ProtocolV3CutoverDrain { + const row = requireSingleD1Row(raw); + + return { + activeAppDeploymentRuns: requireNonNegativeInteger(row, "active_app_deployment_runs"), + activeRuns: requireNonNegativeInteger(row, "active_runs"), + liveDrivers: requireNonNegativeInteger(row, "live_drivers"), + nonterminalApiCommands: requireNonNegativeInteger(row, "nonterminal_api_commands"), + nonterminalCommands: requireNonNegativeInteger(row, "nonterminal_commands"), + unsafeEnvironmentArtifactBackupStaging: requireNonNegativeInteger( + row, + "unsafe_environment_artifact_backup_staging", + ), + unsafeSandboxBackups: requireNonNegativeInteger(row, "unsafe_sandbox_backups"), + unsafeSandboxBackupStaging: requireNonNegativeInteger(row, "unsafe_sandbox_backup_staging"), + unsafeSandboxes: requireNonNegativeInteger(row, "unsafe_sandboxes"), + unsafeSandboxSessions: requireNonNegativeInteger(row, "unsafe_sandbox_sessions"), + unsafeSessions: requireNonNegativeInteger(row, "unsafe_sessions"), + unsettledEffects: requireNonNegativeInteger(row, "unsettled_effects"), + }; +} + +export function parseProtocolV3RuntimeAuthorityPreflight( + raw: string, +): ProtocolV3RuntimeAuthorityPreflight { + const row = requireSingleD1Row(raw); + return { + duplicateSandboxBackups: requireNonNegativeInteger(row, "duplicate_sandbox_backups"), + foreignKeyViolations: requireNonNegativeInteger(row, "foreign_key_violations"), + invalidDriverGenerations: requireNonNegativeInteger(row, "invalid_driver_generations"), + invalidSandboxBackupPointers: requireNonNegativeInteger(row, "invalid_sandbox_backup_pointers"), + invalidSandboxBackups: requireNonNegativeInteger(row, "invalid_sandbox_backups"), + invalidSandboxIdentities: requireNonNegativeInteger(row, "invalid_sandbox_identities"), + invalidSandboxSessionAuthorities: requireNonNegativeInteger( + row, + "invalid_sandbox_session_authorities", + ), + legacyAppDeploymentTraffic: requireNonNegativeInteger(row, "legacy_app_deployment_traffic"), + nonstaticSessions: requireNonNegativeInteger(row, "nonstatic_sessions"), + }; +} + +export function assertProtocolV3RuntimeAuthorityPreflight( + state: ProtocolV3RuntimeAuthorityPreflight, +): void { + const violations = Object.entries(state).filter(([, count]) => count !== 0); + if (violations.length === 0) return; + throw new Error( + `Runtime authority migration preflight failed: ${violations + .map(([name, count]) => `${name}=${count}`) + .join( + ", ", + )}. Migration 0019 remains the atomic authority and will not apply until these rows are repaired through supported lifecycle paths.`, + ); +} + +export function isProtocolV3CutoverDrained(state: ProtocolV3CutoverDrain): boolean { + return isProtocolV3RuntimeDrained(state) && state.nonterminalCommands === 0; +} + +export function isProtocolV3RuntimeDrained(state: ProtocolV3CutoverDrain): boolean { + return ( + state.activeAppDeploymentRuns === 0 && + state.activeRuns === 0 && + state.liveDrivers === 0 && + state.nonterminalApiCommands === 0 && + state.unsafeEnvironmentArtifactBackupStaging === 0 && + state.unsafeSandboxBackups === 0 && + state.unsafeSandboxBackupStaging === 0 && + state.unsafeSandboxes === 0 && + state.unsafeSandboxSessions === 0 && + state.unsafeSessions === 0 && + state.unsettledEffects === 0 + ); +} + +export function parseTimeTravelBookmark(raw: string): string { + const parsed = JSON.parse(raw) as { bookmark?: unknown } | null; + const bookmark = parsed?.bookmark; + + if (typeof bookmark !== "string" || !/^[0-9a-z-]{1,256}$/iu.test(bookmark)) { + throw new Error("D1 Time Travel returned an invalid bookmark."); + } + + return bookmark; +} + +export function parseStoredProtocolV3CutoverBookmark(raw: string): string | null { + const bookmark = requireSingleD1Row(raw).bookmark; + + if (bookmark === null) { + return null; + } + if (typeof bookmark !== "string" || !/^[0-9a-z-]{1,256}$/iu.test(bookmark)) { + throw new Error("The stored D1 Time Travel bookmark is invalid."); + } + + return bookmark; +} + +export function parseProtocolV3CutoverObjects(raw: string): ProtocolV3CutoverObjects { + const row = requireSingleD1Row(raw); + const objectCount = requireNonNegativeInteger(row, "object_count"); + const exactObjectCount = requireNonNegativeInteger(row, "exact_object_count"); + if (exactObjectCount > objectCount) { + throw new Error("D1 cutover exact object count exceeds its reserved object count."); + } + return { exactObjectCount, objectCount }; +} + +export function parseProtocolV3CutoverState(raw: string): ProtocolV3CutoverState { + const row = requireSingleD1Row(raw); + const enabled = requireNonNegativeInteger(row, "enabled"); + const freeze = requireNonNegativeInteger(row, "command_freeze"); + const containerApplicationVersion = requireNullableNonNegativeInteger( + row, + "target_container_application_version", + ); + const containerImage = requireNullableString(row, "target_container_image_digest"); + const workerVersion = requireNullableString(row, "target_worker_version_id"); + const migrationStarted = requireNonNegativeInteger(row, "migration_started"); + const migrationIntentTriggerCount = requireNonNegativeInteger( + row, + "migration_intent_trigger_count", + ); + const phase = row.phase; + if (enabled > 1) throw new Error("D1 cutover gate enabled field must be zero or one."); + if (freeze > 1) throw new Error("D1 command freeze field must be zero or one."); + if (migrationStarted > 1) { + throw new Error("D1 migration intent field must be zero or one."); + } + if ( + row.migration_intent_table_sql !== PROTOCOL_V3_MIGRATION_INTENT_TABLE_SQL || + migrationIntentTriggerCount !== 0 + ) { + throw new Error("D1 migration intent authority is invalid."); + } + if (phase !== "draining" && phase !== "queues_resuming") { + throw new Error("D1 cutover phase is invalid."); + } + if ( + (phase === "draining" && enabled !== 1) || + (phase === "queues_resuming" && freeze !== 1) || + new Set([containerApplicationVersion === null, containerImage === null, workerVersion === null]) + .size !== 1 + ) { + throw new Error("D1 cutover phase state is inconsistent."); + } + + return { + commandFreeze: freeze === 1, + containerApplicationVersion, + containerImageDigest: + containerImage === null + ? null + : requireSha256Digest(containerImage, "Stored protocol v3 Container image"), + enabled: enabled === 1, + migrationStarted: migrationStarted === 1, + phase, + releaseTreeOid: requireGitTreeOid( + typeof row.release_tree_oid === "string" ? row.release_tree_oid : "", + "Stored protocol v3 release tree OID", + ), + workerVersionId: workerVersion === null ? null : requireWorkerVersionId(workerVersion), + }; +} + +export function parseProtocolV3CommandFreeze(raw: string): boolean { + return parseProtocolV3CutoverState(raw).commandFreeze; +} + +export function parseProtocolV3SmokeStatus(raw: string): ProtocolV3SmokeStatus { + const row = requireSingleD1Row(raw); + + return { + bootTokenUsedAt: requireNullableNonNegativeInteger(row, "boot_token_used_at"), + connectionId: requireNullableString(row, "connection_id"), + driverPid: requireNullableNonNegativeInteger(row, "driver_pid"), + driverStartedAt: requireNullableNonNegativeInteger(row, "driver_started_at"), + driverStatus: requireNullableString(row, "driver_status"), + driverVersion: requireNullableString(row, "driver_version"), + protocolVersion: requireNullableNonNegativeInteger(row, "protocol_version"), + statusEvent: requireNullableString(row, "status_event"), + }; +} + +export function parseStoredProtocolV3SmokeSession(raw: string): string | null { + const sessionId = requireSingleD1Row(raw).smoke_session_id; + if (sessionId === null) return null; + if (typeof sessionId !== "string") { + throw new Error("The stored protocol v3 smoke Session ID is invalid."); + } + return requirePlatformId(sessionId, "Stored protocol v3 smoke Session ID"); +} + +export function parseStoredProtocolV3SmokeRequestKey(raw: string): string | null { + const requestKey = requireSingleD1Row(raw).smoke_request_key; + if (requestKey === null) return null; + if (typeof requestKey !== "string") { + throw new Error("The stored protocol v3 smoke request key is invalid."); + } + return requireSmokeRequestKey(requestKey); +} + +export function protocolV3SmokeAgentSql(agentId: string): string { + const validated = requirePlatformId(agentId, "Protocol v3 smoke Agent ID"); + return `SELECT + (SELECT "kind" FROM "agent" WHERE "id" = '${validated}') AS "kind", + (SELECT "status" FROM "agent" WHERE "id" = '${validated}') AS "status";`; +} + +export function assertProtocolV3SmokeAgent(raw: string): void { + const row = requireSingleD1Row(raw); + if (row.kind !== "cattle" || row.status !== "published") { + throw new Error( + "Protocol v3 smoke configuration must identify an existing published cattle Agent.", + ); + } +} + +export function isProtocolV3SmokeReady(status: ProtocolV3SmokeStatus): boolean { + return ( + status.driverStatus === "ready" && + status.statusEvent === "driver.ready" && + status.protocolVersion === 3 && + status.bootTokenUsedAt !== null && + status.connectionId !== null && + status.connectionId.length > 0 && + status.driverPid !== null && + status.driverPid > 0 && + status.driverStartedAt !== null && + status.driverVersion !== null && + status.driverVersion.length > 0 + ); +} + +export function openProtocolV3SmokeWindowSql(accountId: string): string { + const validated = requirePlatformId(accountId, "Protocol v3 smoke account ID"); + return `UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" SET "command_freeze" = 1, "smoke_account_id" = '${validated}' WHERE "id" = 1 AND "enabled" = 1 AND "phase" = 'draining';`; +} + +export function storeProtocolV3SmokeRequestKeySql(requestKey: string): string { + const validated = requireSmokeRequestKey(requestKey); + return `UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" SET "smoke_request_key" = '${validated}' WHERE "id" = 1 AND "enabled" = 1 AND "phase" = 'draining';`; +} + +export function storeProtocolV3SmokeSessionSql(sessionId: string): string { + const validated = requirePlatformId(sessionId, "Protocol v3 smoke Session ID"); + return `UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" SET "smoke_session_id" = '${validated}' WHERE "id" = 1 AND "enabled" = 1 AND "phase" = 'draining';`; +} + +export function protocolV3SmokeStatusSql(sessionId: string): string { + const validated = requirePlatformId(sessionId, "Protocol v3 smoke Session ID"); + return ` +WITH "latest" AS ( + SELECT * + FROM "driver_instance" + WHERE "sandbox_session_id" = '${validated}' + ORDER BY "created_at" DESC + LIMIT 1 +) +SELECT + "latest"."status" AS "driver_status", + "latest"."status_event" AS "status_event", + "latest"."protocol_version" AS "protocol_version", + "latest"."driver_pid" AS "driver_pid", + "latest"."driver_started_at" AS "driver_started_at", + "latest"."driver_version" AS "driver_version", + "latest"."connection_id" AS "connection_id", + "latest"."boot_token_used_at" AS "boot_token_used_at" +FROM (SELECT 1) AS "singleton" +LEFT JOIN "latest" ON 1 = 1; +`; +} + +export function storeProtocolV3CutoverBookmarkSql(bookmark: string): string { + if (!/^[0-9a-z-]{1,256}$/iu.test(bookmark)) { + throw new Error("Cannot persist an invalid D1 Time Travel bookmark."); + } + + return `UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" SET "pre_migration_bookmark" = '${bookmark}' WHERE "id" = 1 AND "enabled" = 1 AND "phase" = 'draining' AND "pre_migration_bookmark" IS NULL;`; +} diff --git a/apps/api/package.json b/apps/api/package.json index 80fa6125..c2f55db8 100644 --- a/apps/api/package.json +++ b/apps/api/package.json @@ -10,10 +10,10 @@ "db:migrate:local": "vp exec wrangler d1 migrations apply DB --local", "lint": "vp lint .", "tc": "vp run cf:types && vp exec tsc --noEmit", - "test": "vp exec bun test tests" + "test": "vp exec bun test tests --isolate" }, "dependencies": { - "@cloudflare/sandbox": "0.12.6", + "@cloudflare/sandbox": "0.12.9", "@mosoo/ag-ui-session": "workspace:*", "@mosoo/agent-driver": "workspace:*", "@mosoo/agent-package": "workspace:*", @@ -30,20 +30,20 @@ "@orpc/server": "^1.15.0", "arktype": "^2.2.0", "better-auth": "1.6.28", - "cloudflare": "7.0.0", + "cloudflare": "7.1.0", "drizzle-orm": "^0.45.2", "fflate": "^0.8.3", "graphql": "^17.0.2", "graphql-yoga": "^5.21.3", "hono": "^4.12.32", + "ignore": "7.0.7", "jsonc-parser": "3.3.1", - "smol-toml": "1.8.0", - "xstate": "^5.32.0" + "smol-toml": "1.8.0" }, "devDependencies": { "@types/node": "^25.8.0", "typescript": "^6.0.3", - "vite-plus": "^0.1.23", - "wrangler": "^4.123.0" + "vite-plus": "^0.3.0", + "wrangler": "^4.127.1" } } diff --git a/apps/api/src/adapters/durable-objects/sandbox.do.ts b/apps/api/src/adapters/durable-objects/sandbox.do.ts index 7f40dbb0..d475dbcf 100644 --- a/apps/api/src/adapters/durable-objects/sandbox.do.ts +++ b/apps/api/src/adapters/durable-objects/sandbox.do.ts @@ -1,5 +1,7 @@ +import { isAsyncTimeoutError, promiseWithTimeout } from "@mosoo/effects"; import { DurableObject } from "cloudflare:workers"; +import { hashSandboxNetworkConstraints } from "../../modules/runtime/domain/sandbox-network-constraints"; import type { ApiBindings } from "../../platform/cloudflare/worker-types"; import { configureSandboxNetworkConstraints, @@ -12,7 +14,31 @@ import type { SandboxRpcForwardMethod } from "./sandbox-rpc-methods"; interface SandboxDelegate extends SandboxNetworkDelegate { alarm(alarmProps?: { isRetry: boolean; retryCount: number }): Promise; + createBackup(options: { + dir: string; + excludes?: string[]; + localBucket?: boolean; + name?: string; + ttl?: number; + }): Promise<{ dir: string; id: string }>; + createSession(options: { id: string }): Promise<{ + readFile(path: string, options?: { encoding?: "utf8" }): Promise<{ content: string }>; + terminal(request: Request, options?: unknown): Promise; + }>; + deleteSession(sessionId: string): Promise; + destroy(): Promise; + exists(path: string): Promise<{ + exists: boolean; + path: string; + success: boolean; + timestamp: string; + }>; fetch(request: Request): Promise; + getContainerPlacementId(): Promise; + mkdir(path: string, options?: { recursive?: boolean }): Promise; + readFile(path: string, options?: { encoding?: "utf8" }): Promise<{ content: string }>; + setKeepAlive(keepAlive: boolean): Promise; + writeFile(path: string, content: string): Promise; } type SandboxContainerState = DurableObjectState<{}> & { @@ -22,8 +48,35 @@ type SandboxContainerState = DurableObjectState<{}> & { }; const FORWARD_SANDBOX_METHOD = Symbol("forwardSandboxMethod"); +const RUNTIME_SUBJECT_INCARNATION_STORAGE_KEY = "mosooRuntimeSubjectIncarnation"; +const RUNTIME_SUBJECT_NETWORK_CONSTRAINTS_HASH_STORAGE_KEY = + "mosooRuntimeSubjectNetworkConstraintsHash"; +const RUNTIME_SUBJECT_READY_PLACEMENT_STORAGE_KEY = "mosooRuntimeSubjectReadyPlacement"; +const RUNTIME_SUBJECT_RETIRED_STORAGE_KEY = "mosooRuntimeSubjectRetiredAt"; +const RUNTIME_SUBJECT_SENTINEL_PATH = "/tmp/.mosoo-runtime-subject-incarnation"; +const RUNTIME_SUBJECT_RETIRED_RETRY_DELAY_MS = 5_000; +const RUNTIME_SUBJECT_RETIRED_DESTROY_TIMEOUT_MS = 15_000; + +function assertRuntimeSubjectIncarnation(value: number): void { + if (!Number.isSafeInteger(value) || value < 0) { + throw new TypeError("Runtime subject incarnation must be a non-negative safe integer."); + } +} + +function assertRuntimeSubjectNetworkConstraintsHash(value: string): void { + if (!/^[0-9a-f]{64}$/u.test(value)) { + throw new TypeError("Runtime subject network constraints hash must be a SHA-256 digest."); + } +} + +function isRuntimeSubjectSentinelMissing(error: unknown): boolean { + return ( + typeof error === "object" && error !== null && Reflect.get(error, "code") === "FILE_NOT_FOUND" + ); +} export class Sandbox extends DurableObject { + readonly #backupCreations = new Map>(); readonly #delegatePromise: Promise; readonly #httpsInterceptionDisabled: boolean; readonly #networkRestorePromise: Promise; @@ -46,38 +99,590 @@ export class Sandbox extends DurableObject { } async configureNetworkConstraints(constraints: unknown): Promise { - await this.#networkRestorePromise; - const delegate = await this.#delegatePromise; + await this.#runUnlessRuntimeSubjectRetired(async () => { + const networkConstraintsHash = await hashSandboxNetworkConstraints(constraints); + const admittedHash = await this.ctx.storage.get( + RUNTIME_SUBJECT_NETWORK_CONSTRAINTS_HASH_STORAGE_KEY, + ); + if (admittedHash !== networkConstraintsHash) { + throw new Error("Sandbox network constraints do not match the admitted incarnation."); + } + await this.#networkRestorePromise; + const delegate = await this.#delegatePromise; - await configureSandboxNetworkConstraints(this.ctx.storage, delegate, constraints, { - containerRunning: (this.ctx as SandboxContainerState).container?.running === true, - httpsInterceptionDisabled: this.#httpsInterceptionDisabled, + await configureSandboxNetworkConstraints(this.ctx.storage, delegate, constraints, { + containerRunning: (this.ctx as SandboxContainerState).container?.running === true, + httpsInterceptionDisabled: this.#httpsInterceptionDisabled, + }); }); } override async fetch(request: Request): Promise { - await this.#networkRestorePromise; - return (await this.#delegatePromise).fetch(request); + return this.#runUnlessRuntimeSubjectRetired(async () => { + await this.#networkRestorePromise; + return (await this.#delegatePromise).fetch(request); + }); } override async alarm(alarmProps?: { isRetry: boolean; retryCount: number }): Promise { - await this.#networkRestorePromise; - await (await this.#delegatePromise).alarm(alarmProps); + try { + if (await this.#isRuntimeSubjectRetired()) { + return; + } + + await this.#networkRestorePromise; + await (await this.#delegatePromise).alarm(alarmProps); + } finally { + await this.#destroyRuntimeSubjectContainerIfRetired(); + } + } + + async activateRuntimeSubjectIncarnation( + incarnation: number, + networkConstraintsHash: string, + ): Promise { + assertRuntimeSubjectIncarnation(incarnation); + assertRuntimeSubjectNetworkConstraintsHash(networkConstraintsHash); + + await this.ctx.blockConcurrencyWhile(async () => { + this.#assertRuntimeSubjectNotRetired( + await this.ctx.storage.get(RUNTIME_SUBJECT_RETIRED_STORAGE_KEY), + ); + const [current, currentNetworkConstraintsHash] = await Promise.all([ + this.ctx.storage.get(RUNTIME_SUBJECT_INCARNATION_STORAGE_KEY), + this.ctx.storage.get(RUNTIME_SUBJECT_NETWORK_CONSTRAINTS_HASH_STORAGE_KEY), + ]); + + if (current === incarnation && currentNetworkConstraintsHash === networkConstraintsHash) { + return; + } + if (current === undefined && currentNetworkConstraintsHash === undefined) { + await this.ctx.storage.put({ + [RUNTIME_SUBJECT_INCARNATION_STORAGE_KEY]: incarnation, + [RUNTIME_SUBJECT_NETWORK_CONSTRAINTS_HASH_STORAGE_KEY]: networkConstraintsHash, + }); + return; + } + throw new Error("Runtime subject incarnation identity does not match this Durable Object."); + }); + } + + async createRuntimeSubjectBackup( + incarnation: number, + options: { + dir: string; + excludes?: string[]; + forbiddenPaths?: string[]; + localBucket?: boolean; + name: string; + ttl?: number; + }, + ): Promise<{ dir: string; id: string }> { + assertRuntimeSubjectIncarnation(incarnation); + if (options.name.length === 0) { + throw new TypeError("Runtime subject backup name must not be empty."); + } + + const key = `${incarnation}:${options.name}`; + let creation = this.#backupCreations.get(key); + if (creation === undefined) { + creation = this.#runUnlessRuntimeSubjectRetired(async () => { + const current = await this.ctx.storage.get(RUNTIME_SUBJECT_INCARNATION_STORAGE_KEY); + if (current !== incarnation) { + throw new Error("Runtime subject backup targeted a stale incarnation."); + } + await this.#networkRestorePromise; + const delegate = await this.#delegatePromise; + for (const path of options.forbiddenPaths ?? []) { + const result = await delegate.exists(path); + if (!result.success) { + throw new Error("Runtime subject backup secret preflight failed."); + } + if (result.exists) { + throw new Error("Runtime subject backup contains legacy persistent credentials."); + } + } + await delegate.mkdir(options.dir, { recursive: true }); + const { forbiddenPaths: _, ...backupOptions } = options; + return delegate.createBackup(backupOptions); + }); + this.#backupCreations.set(key, creation); + } + try { + const backup = await creation; + if (backup.dir !== options.dir) { + throw new Error("Runtime subject backup name was reused for a different directory."); + } + return backup; + } finally { + if (this.#backupCreations.get(key) === creation) { + this.#backupCreations.delete(key); + } + } + } + + async inspectRuntimeSubjectIncarnation( + incarnation: number, + networkConstraintsHash: string, + ): Promise<{ kind: "healthy" | "missing" | "retired" | "stale" | "unknown" }> { + assertRuntimeSubjectIncarnation(incarnation); + assertRuntimeSubjectNetworkConstraintsHash(networkConstraintsHash); + const before = await this.#readRuntimeSubjectReadyState(); + if (before.retired) { + return { kind: "retired" }; + } + if ( + before.incarnation !== incarnation || + before.networkConstraintsHash !== networkConstraintsHash + ) { + return { kind: "stale" }; + } + if (before.placement === undefined) { + return { kind: "unknown" }; + } + + let health: "healthy" | "missing" | "unknown"; + try { + health = await this.#readyRuntimeSubjectHealth( + incarnation, + networkConstraintsHash, + before.placement, + ); + } finally { + await this.#destroyRuntimeSubjectContainerIfRetired(); + } + const after = await this.#readRuntimeSubjectReadyState(); + if (after.retired) { + return { kind: "retired" }; + } + if ( + after.incarnation !== incarnation || + after.networkConstraintsHash !== networkConstraintsHash + ) { + return { kind: "stale" }; + } + return { kind: after.placement === before.placement ? health : "unknown" }; + } + + async markRuntimeSubjectIncarnationReady( + incarnation: number, + networkConstraintsHash: string, + ): Promise { + assertRuntimeSubjectIncarnation(incarnation); + assertRuntimeSubjectNetworkConstraintsHash(networkConstraintsHash); + + const before = await this.#readRuntimeSubjectReadyState(); + if (before.retired) { + throw new Error("Runtime subject incarnation is retired."); + } + if ( + before.incarnation !== incarnation || + before.networkConstraintsHash !== networkConstraintsHash + ) { + throw new Error("Runtime subject readiness targeted a stale incarnation."); + } + if ((this.ctx as SandboxContainerState).container?.running !== true) { + throw new Error("Runtime subject container stopped before readiness was recorded."); + } + + try { + await this.#networkRestorePromise; + const delegate = await this.#delegatePromise; + await delegate.writeFile( + RUNTIME_SUBJECT_SENTINEL_PATH, + `${incarnation}:${networkConstraintsHash}`, + ); + const placement = await delegate.getContainerPlacementId(); + if (placement === undefined) { + throw new Error("Runtime subject container placement is not available."); + } + + await this.ctx.blockConcurrencyWhile(async () => { + this.#assertRuntimeSubjectNotRetired( + await this.ctx.storage.get(RUNTIME_SUBJECT_RETIRED_STORAGE_KEY), + ); + const [current, currentNetworkConstraintsHash] = await Promise.all([ + this.ctx.storage.get(RUNTIME_SUBJECT_INCARNATION_STORAGE_KEY), + this.ctx.storage.get(RUNTIME_SUBJECT_NETWORK_CONSTRAINTS_HASH_STORAGE_KEY), + ]); + if (current !== incarnation || currentNetworkConstraintsHash !== networkConstraintsHash) { + throw new Error("Runtime subject readiness targeted a stale incarnation."); + } + if ((this.ctx as SandboxContainerState).container?.running !== true) { + throw new Error("Runtime subject container stopped before readiness was recorded."); + } + await this.ctx.storage.put(RUNTIME_SUBJECT_READY_PLACEMENT_STORAGE_KEY, placement); + }); + } finally { + await this.#destroyRuntimeSubjectContainerIfRetired(); + } + } + + async destroyRuntimeSubjectIncarnation( + incarnation: number, + ): Promise<{ kind: "destroyed" | "stale" }> { + assertRuntimeSubjectIncarnation(incarnation); + + const stale = await this.ctx.blockConcurrencyWhile(async () => { + const current = await this.ctx.storage.get(RUNTIME_SUBJECT_INCARNATION_STORAGE_KEY); + + if (current !== undefined && current !== incarnation) { + return true; + } + if (current === undefined) { + await this.ctx.storage.put(RUNTIME_SUBJECT_INCARNATION_STORAGE_KEY, incarnation); + } + + await this.ctx.storage.put(RUNTIME_SUBJECT_RETIRED_STORAGE_KEY, Date.now()); + await this.ctx.storage.setAlarm(Date.now() + RUNTIME_SUBJECT_RETIRED_RETRY_DELAY_MS); + return false; + }); + if (stale) { + return { kind: "stale" }; + } + + await this.#destroyRetiredRuntimeSubjectContainer(); + return { kind: "destroyed" }; + } + + async #isRuntimeSubjectRetired(): Promise { + return (await this.ctx.storage.get(RUNTIME_SUBJECT_RETIRED_STORAGE_KEY)) !== undefined; + } + + async #readyRuntimeSubjectHealth( + incarnation: number, + networkConstraintsHash: string, + readyPlacement: string | null, + ): Promise<"healthy" | "missing" | "unknown"> { + if ((this.ctx as SandboxContainerState).container?.running !== true) { + return "missing"; + } + + const delegate = await this.#delegatePromise; + try { + await this.#networkRestorePromise; + // The SDK placement id is cached in DO storage and survives onStop. + // A sentinel read forces a real container handshake before we trust it. + const sentinel = await delegate.readFile(RUNTIME_SUBJECT_SENTINEL_PATH, { + encoding: "utf8", + }); + const placement = await delegate.getContainerPlacementId(); + if ((this.ctx as SandboxContainerState).container?.running !== true) { + return "missing"; + } + if (sentinel.content !== `${incarnation}:${networkConstraintsHash}`) { + return "missing"; + } + if (typeof placement === "string" && typeof readyPlacement === "string") { + return placement === readyPlacement ? "healthy" : "missing"; + } + if (placement !== null || readyPlacement !== null) { + return "unknown"; + } + return "healthy"; + } catch (error) { + if ((this.ctx as SandboxContainerState).container?.running !== true) { + return "missing"; + } + try { + const placement = await delegate.getContainerPlacementId(); + if ( + typeof placement === "string" && + typeof readyPlacement === "string" && + placement !== readyPlacement + ) { + return "missing"; + } + } catch { + // Preserve the original, more informative health failure below. + } + return isRuntimeSubjectSentinelMissing(error) ? "missing" : "unknown"; + } + } + + #assertRuntimeSubjectNotRetired(retiredAt: number | undefined): void { + if (retiredAt !== undefined) { + throw new Error("Runtime subject incarnation is retired."); + } + } + + async #destroyRetiredRuntimeSubjectContainer(): Promise { + try { + const delegate = await this.#delegatePromise; + await promiseWithTimeout( + (async () => { + await delegate.setKeepAlive(false); + await delegate.destroy(); + })(), + { + label: "Runtime subject retired container destroy", + timeoutMs: RUNTIME_SUBJECT_RETIRED_DESTROY_TIMEOUT_MS, + }, + ); + } catch (error) { + await this.ctx.storage.setAlarm(Date.now() + RUNTIME_SUBJECT_RETIRED_RETRY_DELAY_MS); + if (isAsyncTimeoutError(error)) { + this.ctx.abort("Runtime subject retired container destroy timed out.", { + retryAlarm: true, + }); + } + throw error; + } + } + + async #destroyRuntimeSubjectContainerIfRetired(): Promise { + if (await this.#isRuntimeSubjectRetired()) { + await this.#destroyRetiredRuntimeSubjectContainer(); + } + } + + async #assertReadyRuntimeSubjectContainerHealthy(): Promise<{ + readonly incarnation: number | undefined; + readonly networkConstraintsHash: string | undefined; + readonly placement: string | null | undefined; + readonly retired: boolean; + }> { + const before = await this.#readRuntimeSubjectReadyState(); + if (before.retired) { + throw new Error("Runtime subject incarnation is retired."); + } + if (before.incarnation === undefined || before.placement === undefined) { + return before; + } + let health: "healthy" | "missing" | "unknown"; + try { + if (before.networkConstraintsHash === undefined) { + throw new Error("Runtime subject ready container has no network identity."); + } + health = await this.#readyRuntimeSubjectHealth( + before.incarnation, + before.networkConstraintsHash, + before.placement, + ); + } finally { + await this.#destroyRuntimeSubjectContainerIfRetired(); + } + const after = await this.#readRuntimeSubjectReadyState(); + if ( + after.retired || + after.incarnation !== before.incarnation || + after.networkConstraintsHash !== before.networkConstraintsHash || + after.placement !== before.placement + ) { + throw new Error("Runtime subject incarnation changed during its health probe."); + } + if (health === "healthy") { + return after; + } + if (health === "unknown") { + throw new Error("Runtime subject ready container health is unknown."); + } + + await this.ctx.blockConcurrencyWhile(async () => { + const current = await this.#readRuntimeSubjectReadyStateWithoutGate(); + if ( + current.retired || + current.incarnation !== before.incarnation || + current.networkConstraintsHash !== before.networkConstraintsHash || + current.placement !== before.placement + ) { + throw new Error("Runtime subject incarnation changed before retirement."); + } + await this.ctx.storage.put(RUNTIME_SUBJECT_RETIRED_STORAGE_KEY, Date.now()); + await this.ctx.storage.setAlarm(Date.now() + RUNTIME_SUBJECT_RETIRED_RETRY_DELAY_MS); + }); + throw new Error("Runtime subject ready container was replaced or stopped."); + } + + async #readRuntimeSubjectReadyState(): Promise<{ + readonly incarnation: number | undefined; + readonly networkConstraintsHash: string | undefined; + readonly placement: string | null | undefined; + readonly retired: boolean; + }> { + return this.ctx.blockConcurrencyWhile(() => this.#readRuntimeSubjectReadyStateWithoutGate()); + } + + async #readRuntimeSubjectReadyStateWithoutGate(): Promise<{ + readonly incarnation: number | undefined; + readonly networkConstraintsHash: string | undefined; + readonly placement: string | null | undefined; + readonly retired: boolean; + }> { + const [incarnation, networkConstraintsHash, placement, retiredAt] = await Promise.all([ + this.ctx.storage.get(RUNTIME_SUBJECT_INCARNATION_STORAGE_KEY), + this.ctx.storage.get(RUNTIME_SUBJECT_NETWORK_CONSTRAINTS_HASH_STORAGE_KEY), + this.ctx.storage.get(RUNTIME_SUBJECT_READY_PLACEMENT_STORAGE_KEY), + this.ctx.storage.get(RUNTIME_SUBJECT_RETIRED_STORAGE_KEY), + ]); + return { incarnation, networkConstraintsHash, placement, retired: retiredAt !== undefined }; + } + + #guardRuntimeSubjectProcess(process: unknown): unknown { + if (typeof process !== "object" || process === null) { + throw new TypeError("Cloudflare Sandbox process handle is not an object."); + } + return this.#guardRuntimeSubjectRpcTarget(process); + } + + #guardRuntimeSubjectForwardResult(method: string, result: unknown): unknown { + if (method === "createSession" || method === "getSession") { + return this.#guardRuntimeSubjectExecutionSession(result); + } + if (method === "startProcess") { + return this.#guardRuntimeSubjectProcess(result); + } + if (method === "getProcess") { + return result === null ? null : this.#guardRuntimeSubjectProcess(result); + } + if (method === "listProcesses") { + if (!Array.isArray(result)) { + throw new TypeError("Cloudflare Sandbox process list is not an array."); + } + return result.map((process) => this.#guardRuntimeSubjectProcess(process)); + } + return result; + } + + #guardRuntimeSubjectExecutionSession(session: unknown): unknown { + if (typeof session !== "object" || session === null) { + throw new TypeError("Cloudflare Sandbox execution session is not an object."); + } + return this.#guardRuntimeSubjectRpcTarget(session); + } + + #guardRuntimeSubjectRpcTarget(target: object): object { + const facade: Record = {}; + for (const key of Reflect.ownKeys(target)) { + const value = Reflect.get(target, key); + Reflect.set( + facade, + key, + typeof value === "function" + ? (...args: unknown[]) => + this.#runUnlessRuntimeSubjectRetired(async () => { + const result = await (Reflect.apply( + value, + target, + this.#guardRuntimeSubjectInvocationArgs(String(key), args), + ) as Promise); + return this.#guardRuntimeSubjectForwardResult(String(key), result); + }) + : value, + ); + } + return facade; + } + + #guardRuntimeSubjectInvocationArgs(method: string, args: readonly unknown[]): readonly unknown[] { + if (method !== "startProcess") { + return args; + } + const options = args[1]; + if (typeof options !== "object" || options === null) { + return args; + } + const onStart = Reflect.get(options, "onStart"); + if (typeof onStart !== "function") { + return args; + } + const guardedOptions = { + ...options, + onStart: (process: unknown) => + Reflect.apply(onStart, options, [this.#guardRuntimeSubjectProcess(process)]), + }; + return [args[0], guardedOptions, ...args.slice(2)]; + } + + async #runUnlessRuntimeSubjectRetired(task: () => Promise): Promise { + const failures: unknown[] = []; + let retired = false; + let result: T | undefined; + try { + const before = await this.#assertReadyRuntimeSubjectContainerHealthy(); + + try { + result = await task(); + } catch (error) { + failures.push(error); + } + + try { + retired = await this.#isRuntimeSubjectRetired(); + if (!retired) { + const after = await this.#assertReadyRuntimeSubjectContainerHealthy(); + if ( + before.incarnation !== after.incarnation || + before.networkConstraintsHash !== after.networkConstraintsHash || + before.placement !== after.placement + ) { + await this.ctx.blockConcurrencyWhile(async () => { + await this.ctx.storage.put(RUNTIME_SUBJECT_RETIRED_STORAGE_KEY, Date.now()); + await this.ctx.storage.setAlarm(Date.now() + RUNTIME_SUBJECT_RETIRED_RETRY_DELAY_MS); + }); + throw new Error("Runtime subject incarnation changed during the operation."); + } + } + } catch (error) { + failures.push(error); + } + } catch (error) { + failures.push(error); + } + + if (!retired) { + try { + retired = await this.#isRuntimeSubjectRetired(); + } catch (error) { + failures.push(error); + } + } + if (retired) { + if (failures.length === 0) { + failures.push(new Error("Runtime subject incarnation was retired during the operation.")); + } + try { + await this.#destroyRetiredRuntimeSubjectContainer(); + } catch (error) { + failures.push(error); + } + } + + if (failures.length === 1) { + throw failures[0]; + } + if (failures.length > 1) { + throw new AggregateError(failures, "Runtime subject operation failed in multiple phases."); + } + return result as T; } async [FORWARD_SANDBOX_METHOD]( method: SandboxRpcForwardMethod, args: readonly unknown[], ): Promise { - await waitForSandboxNetworkRestore(this.#networkRestorePromise, method, args); - const delegate = await this.#delegatePromise; - const action = Reflect.get(delegate, method); + return this.#runUnlessRuntimeSubjectRetired(async () => { + await waitForSandboxNetworkRestore(this.#networkRestorePromise, method, args); + const delegate = await this.#delegatePromise; + const action = Reflect.get(delegate, method); - if (typeof action !== "function") { - throw new TypeError(`Cloudflare Sandbox delegate is missing ${method}.`); - } + if (typeof action !== "function") { + throw new TypeError(`Cloudflare Sandbox delegate is missing ${method}.`); + } - return await (Reflect.apply(action, delegate, args) as Promise); + const result = await (Reflect.apply( + action, + delegate, + this.#guardRuntimeSubjectInvocationArgs(method, args), + ) as Promise); + const runtimeSubjectIncarnation = await this.ctx.storage.get( + RUNTIME_SUBJECT_INCARNATION_STORAGE_KEY, + ); + if (runtimeSubjectIncarnation !== undefined) { + return this.#guardRuntimeSubjectForwardResult(method, result); + } + return result; + }); } } diff --git a/apps/api/src/adapters/durable-objects/session.do.ts b/apps/api/src/adapters/durable-objects/session.do.ts index c92f6eb3..d56dce8c 100644 --- a/apps/api/src/adapters/durable-objects/session.do.ts +++ b/apps/api/src/adapters/durable-objects/session.do.ts @@ -8,7 +8,20 @@ interface SessionDelegate { closeViewers(sessionId: string, reason: string): Promise; destroy(sessionId: string, reason: string): Promise; fetch(request: Request): Promise; - publishEvents(sessionId: string, events: AgUiSessionEvent[]): Promise; + publishEventBatches( + sessionId: string, + batches: Array<{ + events: AgUiSessionEvent[]; + previousRuntimeEventSeqCursor: number | null; + runtimeEventSeqCursor: number | null; + }>, + ): Promise; + publishEvents( + sessionId: string, + events: AgUiSessionEvent[], + runtimeEventSeqCursor?: number | null, + previousRuntimeEventSeqCursor?: number | null, + ): Promise; syncViewers(sessionId: string): Promise; webSocketClose(ws: WebSocket, code: number, reason: string): Promise; webSocketError(ws: WebSocket, error: unknown): Promise | void; @@ -34,8 +47,26 @@ export class Session extends DurableObject { await (await this.#delegatePromise).alarm(); } - async publishEvents(sessionId: string, events: AgUiSessionEvent[]): Promise { - await (await this.#delegatePromise).publishEvents(sessionId, events); + async publishEvents( + sessionId: string, + events: AgUiSessionEvent[], + runtimeEventSeqCursor: number | null = null, + previousRuntimeEventSeqCursor: number | null = null, + ): Promise { + await ( + await this.#delegatePromise + ).publishEvents(sessionId, events, runtimeEventSeqCursor, previousRuntimeEventSeqCursor); + } + + async publishEventBatches( + sessionId: string, + batches: Array<{ + events: AgUiSessionEvent[]; + previousRuntimeEventSeqCursor: number | null; + runtimeEventSeqCursor: number | null; + }>, + ): Promise { + await (await this.#delegatePromise).publishEventBatches(sessionId, batches); } async syncViewers(sessionId: string): Promise { diff --git a/apps/api/src/adapters/graphql/schema.generated.graphql b/apps/api/src/adapters/graphql/schema.generated.graphql index bf4b9281..6938fc2b 100644 --- a/apps/api/src/adapters/graphql/schema.generated.graphql +++ b/apps/api/src/adapters/graphql/schema.generated.graphql @@ -379,6 +379,7 @@ type AgentSessionRetrieve { capabilities: [AgentSessionActionCapability!]! recoverability: AgentSessionRecoverability! session: Session! + taskSnapshot: AgentTaskSnapshot } type AgentSessionRetrieveConnection { @@ -419,6 +420,18 @@ type AgentSummary { visibility: AgentVisibility! } +type AgentTask { + taskId: String! + taskType: String + title: String +} + +type AgentTaskSnapshot { + driverInstanceId: ULID! + runId: ULID! + tasks: [AgentTask!]! +} + type AgentToolSummary { enabled: Boolean! iconUrl: String @@ -496,7 +509,7 @@ enum AppDeploymentRunStatus { } enum AppDeploymentTargetKind { - cloudflare_pages + cloudflare_static_assets cloudflare_worker } diff --git a/apps/api/src/adapters/graphql/schema/app-schema.ts b/apps/api/src/adapters/graphql/schema/app-schema.ts index 76625b02..666bfa55 100644 --- a/apps/api/src/adapters/graphql/schema/app-schema.ts +++ b/apps/api/src/adapters/graphql/schema/app-schema.ts @@ -19,7 +19,7 @@ export const appSchema = /* GraphQL */ ` } enum AppDeploymentTargetKind { - cloudflare_pages + cloudflare_static_assets cloudflare_worker } diff --git a/apps/api/src/adapters/graphql/schema/session-schema.ts b/apps/api/src/adapters/graphql/schema/session-schema.ts index 72d67df1..eabf43f9 100644 --- a/apps/api/src/adapters/graphql/schema/session-schema.ts +++ b/apps/api/src/adapters/graphql/schema/session-schema.ts @@ -287,10 +287,23 @@ export const sessionSchema = /* GraphQL */ ` status: AgentSessionActionCapabilityStatus! } + type AgentTask { + taskId: String! + taskType: String + title: String + } + + type AgentTaskSnapshot { + driverInstanceId: ULID! + runId: ULID! + tasks: [AgentTask!]! + } + type AgentSessionRetrieve { capabilities: [AgentSessionActionCapability!]! recoverability: AgentSessionRecoverability! session: Session! + taskSnapshot: AgentTaskSnapshot } type AgentSessionRetrieveConnection { diff --git a/apps/api/src/index.ts b/apps/api/src/index.ts index b9941fa0..c3faa401 100644 --- a/apps/api/src/index.ts +++ b/apps/api/src/index.ts @@ -4,6 +4,7 @@ export { ContainerProxy } from "./adapters/durable-objects/sandbox-container-pro export { DriverConnection } from "./adapters/durable-objects/driver-connection.do"; export { Sandbox } from "./adapters/durable-objects/sandbox.do"; export { Session } from "./adapters/durable-objects/session.do"; +export { AppDeploymentWorkflow } from "./platform/cloudflare/app-deployment-workflow"; void arktypeWorkerConfigInitialized; diff --git a/apps/api/src/modules/api-command/application/api-command-enqueue.ts b/apps/api/src/modules/api-command/application/api-command-enqueue.ts index 50ae37a0..e1d05048 100644 --- a/apps/api/src/modules/api-command/application/api-command-enqueue.ts +++ b/apps/api/src/modules/api-command/application/api-command-enqueue.ts @@ -1,11 +1,13 @@ import type { AppDeploymentRunId } from "@mosoo/id"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; -import { admitApiCommand, enqueueApiCommand } from "./api-command-ledger"; -import type { ApiCommandAdmission, EnqueueApiCommandInput } from "./api-command-ledger"; +import { enqueueApiCommand } from "./api-command-ledger"; +import type { EnqueueApiCommandInput } from "./api-command-ledger"; import type { AppDeploymentRunDispatchCommandPayload, + AppDeploymentScriptReconciliationCommandPayload, CostLedgerReconciliationCommandPayload, + SandboxBackupReconciliationCommandPayload, ScheduledMaintenanceCommandPayload, SessionRunDispatchCommandPayload, } from "./api-command-payload"; @@ -17,7 +19,7 @@ export function createAppDeploymentRunDispatchDedupeKey(runId: AppDeploymentRunI } export async function enqueueAppDeploymentRunDispatchCommand( - bindings: Pick, + bindings: Pick, payload: AppDeploymentRunDispatchCommandPayload, ): Promise { await enqueueApiCommand(bindings, { @@ -27,6 +29,17 @@ export async function enqueueAppDeploymentRunDispatchCommand( }); } +export async function enqueueAppDeploymentScriptReconciliationCommand( + bindings: Pick, + payload: AppDeploymentScriptReconciliationCommandPayload, +): Promise { + await enqueueApiCommand(bindings, { + dedupeKey: `app-deployment-script-reconciliation:${payload.scheduledTime}:${payload.cursor ?? "root"}`, + kind: "app_deployment_script_reconciliation", + payload, + }); +} + export async function enqueueCostLedgerReconciliationCommand( bindings: Pick, payload: CostLedgerReconciliationCommandPayload, @@ -54,11 +67,20 @@ export async function enqueueScheduledMaintenanceCommand( }); } -export async function admitSessionRunDispatchCommand( +export async function enqueueSandboxBackupReconciliationCommand( bindings: Pick, - payload: SessionRunDispatchCommandPayload, -): Promise { - return admitApiCommand(bindings, createSessionRunDispatchApiCommandInput(payload)); + payload: SandboxBackupReconciliationCommandPayload, +): Promise { + await enqueueApiCommand(bindings, { + dedupeKey: [ + "sandbox_backup_reconciliation", + payload.scheduledTime, + payload.databasePage, + payload.cursor ?? "start", + ].join(":"), + kind: "sandbox_backup_reconciliation", + payload, + }); } export function createSessionRunDispatchApiCommandInput( diff --git a/apps/api/src/modules/api-command/application/api-command-ledger.ts b/apps/api/src/modules/api-command/application/api-command-ledger.ts index 4cb37e4f..b6fc6f3b 100644 --- a/apps/api/src/modules/api-command/application/api-command-ledger.ts +++ b/apps/api/src/modules/api-command/application/api-command-ledger.ts @@ -1,7 +1,7 @@ import { apiCommandsTable } from "@mosoo/db"; import type { ApiCommandId, ApiCommandKind, ApiCommandRow } from "@mosoo/db"; import { createPlatformId } from "@mosoo/id"; -import { and, asc, eq, inArray, lt, or, sql } from "drizzle-orm"; +import { and, asc, eq, gt, inArray, isNull, lte, or, sql } from "drizzle-orm"; import { createErrorLogContext, logError } from "../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; @@ -23,6 +23,8 @@ const API_COMMAND_QUEUE_SEND_FAILED_MESSAGE = "API command queue send failed."; const API_COMMAND_QUEUE_REDRIVE_LIMIT = 100; +const APP_DEPLOYMENT_WORKFLOW_RECONCILE_LIMIT = 100; + export interface EnqueueApiCommandInput { dedupeKey: string; kind: ApiCommandKind; @@ -37,12 +39,26 @@ export interface PreparedApiCommand { export interface ApiCommandClaim { attemptCount: number; + claimOwner: string; commandId: ApiCommandId; dedupeKey: string; + deliveryGeneration: number; kind: ApiCommandKind; + lastErrorCode: string | null; + lastErrorMessage: string | null; payloadJson: string; } +export type ApiCommandClaimAuthority = Pick< + ApiCommandClaim, + "attemptCount" | "claimOwner" | "commandId" | "deliveryGeneration" +>; + +export type ApiCommandClaimResult = + | { readonly claim: ApiCommandClaim; readonly kind: "claimed" } + | { readonly claimExpiresAt: number | null; readonly kind: "busy" } + | { readonly kind: "missing" | "stale" | "terminal" }; + export interface ApiCommandAdmission { readonly commandId: ApiCommandId; readonly kind: ApiCommandKind; @@ -50,7 +66,7 @@ export interface ApiCommandAdmission { } type ApiCommandDeliveryBindings = Pick & - Partial>; + Partial>; function normalizeDedupeKey(value: string): string { const dedupeKey = value.trim(); @@ -62,18 +78,42 @@ function normalizeDedupeKey(value: string): string { return dedupeKey; } -function toQueueMessage(commandId: ApiCommandId): ApiCommandMessage { - return { commandId }; +function toQueueMessage(commandId: ApiCommandId, deliveryGeneration: number): ApiCommandMessage { + return { commandId, deliveryGeneration }; +} + +export function getAppDeploymentWorkflowInstanceId( + commandId: ApiCommandId, + deliveryGeneration: number, +): string { + return `app-${commandId.toLowerCase()}-g${deliveryGeneration.toString(36)}`; +} + +async function createAppDeploymentWorkflowInstance( + workflow: Workflow, + message: ApiCommandMessage, +): Promise { + const instanceId = getAppDeploymentWorkflowInstanceId( + message.commandId, + message.deliveryGeneration, + ); + + await workflow.createBatch([{ id: instanceId, params: message }]); } export async function findApiCommandByDedupeKey( database: D1Database, dedupeKey: string, -): Promise | null> { +): Promise | null> { return ( (await getAppDatabase(database) .select({ id: apiCommandsTable.id, + deliveryGeneration: apiCommandsTable.deliveryGeneration, + kind: apiCommandsTable.kind, lastErrorCode: apiCommandsTable.lastErrorCode, lastErrorMessage: apiCommandsTable.lastErrorMessage, status: apiCommandsTable.status, @@ -88,6 +128,7 @@ export async function findApiCommandByDedupeKey( async function markApiCommandQueueSendFailed(input: { commandId: ApiCommandId; database: D1Database; + deliveryGeneration: number; }): Promise { await getAppDatabase(input.database) .update(apiCommandsTable) @@ -96,13 +137,20 @@ async function markApiCommandQueueSendFailed(input: { lastErrorMessage: API_COMMAND_QUEUE_SEND_FAILED_MESSAGE, updatedAt: currentTimestampMs(), }) - .where(and(eq(apiCommandsTable.id, input.commandId), eq(apiCommandsTable.status, "queued"))) + .where( + and( + eq(apiCommandsTable.id, input.commandId), + eq(apiCommandsTable.deliveryGeneration, input.deliveryGeneration), + eq(apiCommandsTable.status, "queued"), + ), + ) .run(); } async function clearApiCommandQueueSendFailure(input: { commandId: ApiCommandId; database: D1Database; + deliveryGeneration: number; }): Promise { await getAppDatabase(input.database) .update(apiCommandsTable) @@ -114,6 +162,7 @@ async function clearApiCommandQueueSendFailure(input: { .where( and( eq(apiCommandsTable.id, input.commandId), + eq(apiCommandsTable.deliveryGeneration, input.deliveryGeneration), eq(apiCommandsTable.status, "queued"), inArray(apiCommandsTable.lastErrorCode, [ API_COMMAND_QUEUE_DELIVERY_PENDING_CODE, @@ -127,22 +176,35 @@ async function clearApiCommandQueueSendFailure(input: { async function sendApiCommandMessage( bindings: ApiCommandDeliveryBindings, commandId: ApiCommandId, + deliveryGeneration: number, kind: ApiCommandKind, -): Promise { - const queue = - kind === "environment_package_artifact_build" - ? bindings.ENVIRONMENT_ARTIFACT_BUILD_QUEUE - : bindings.API_COMMAND_QUEUE; - if (!queue) { - throw new Error("Environment artifact build queue binding is required."); - } +): Promise { try { - await queue.send(toQueueMessage(commandId)); + const message = toQueueMessage(commandId, deliveryGeneration); + if (kind === "app_deployment_run_dispatch") { + if (!bindings.APP_DEPLOYMENT_WORKFLOW) { + throw new Error("App deployment Workflow binding is required."); + } + await createAppDeploymentWorkflowInstance(bindings.APP_DEPLOYMENT_WORKFLOW, message); + } else { + const queue = + kind === "environment_package_artifact_build" + ? bindings.ENVIRONMENT_ARTIFACT_BUILD_QUEUE + : bindings.API_COMMAND_QUEUE; + if (!queue) { + throw new Error("Environment artifact build queue binding is required."); + } + await queue.send(message); + } } catch (error) { // A rejected producer response does not prove that Queue discarded the message. // The durable outbox record remains eligible for scheduled redrive either way. try { - await markApiCommandQueueSendFailed({ commandId, database: bindings.DB }); + await markApiCommandQueueSendFailed({ + commandId, + database: bindings.DB, + deliveryGeneration, + }); } catch (markError) { logError("api-command.enqueue_failure_mark_failed", { ...createErrorLogContext(markError), @@ -154,11 +216,15 @@ async function sendApiCommandMessage( ...createErrorLogContext(error), commandId, }); - return; + return false; } try { - await clearApiCommandQueueSendFailure({ commandId, database: bindings.DB }); + await clearApiCommandQueueSendFailure({ + commandId, + database: bindings.DB, + deliveryGeneration, + }); } catch (error) { // Queue accepted the command. Leaving its delivery marker intact is safe: // a later redrive may send a duplicate, and consumer claiming is idempotent. @@ -167,13 +233,56 @@ async function sendApiCommandMessage( commandId, }); } + return true; +} + +export type ApiCommandQueueHandoffDisposition = "continue" | "finished" | "retry"; + +export async function handoffAppDeploymentQueueMessage( + bindings: ApiCommandDeliveryBindings, + message: ApiCommandMessage, +): Promise { + const command = await getAppDatabase(bindings.DB) + .select({ + deliveryGeneration: apiCommandsTable.deliveryGeneration, + kind: apiCommandsTable.kind, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, message.commandId)) + .limit(1) + .get(); + + if ( + command === undefined || + command.deliveryGeneration !== message.deliveryGeneration || + (command.status !== "queued" && command.status !== "running") + ) { + return "finished"; + } + if (command.kind !== "app_deployment_run_dispatch") { + return "continue"; + } + + return (await sendApiCommandMessage( + bindings, + message.commandId, + message.deliveryGeneration, + command.kind, + )) + ? "finished" + : "retry"; } export async function redriveFailedApiCommandEnqueues( bindings: ApiCommandDeliveryBindings, ): Promise { const commands = await getAppDatabase(bindings.DB) - .select({ id: apiCommandsTable.id, kind: apiCommandsTable.kind }) + .select({ + deliveryGeneration: apiCommandsTable.deliveryGeneration, + id: apiCommandsTable.id, + kind: apiCommandsTable.kind, + }) .from(apiCommandsTable) .where( and( @@ -189,7 +298,156 @@ export async function redriveFailedApiCommandEnqueues( .all(); for (const command of commands) { - await sendApiCommandMessage(bindings, command.id, command.kind); + await sendApiCommandMessage(bindings, command.id, command.deliveryGeneration, command.kind); + } +} + +function recoverableApiCommandPredicate(nowMs: number) { + return or( + eq(apiCommandsTable.status, "queued"), + and( + eq(apiCommandsTable.status, "running"), + or(isNull(apiCommandsTable.claimExpiresAt), lte(apiCommandsTable.claimExpiresAt, nowMs)), + ), + ); +} + +async function markAppDeploymentWorkflowReconciled(input: { + commandId: ApiCommandId; + database: D1Database; + deliveryGeneration: number; + nowMs: number; + payloadJson: string; +}): Promise { + await getAppDatabase(input.database) + .update(apiCommandsTable) + .set({ + updatedAt: sql`CASE WHEN ${apiCommandsTable.updatedAt} = ${Number.MAX_SAFE_INTEGER} THEN ${apiCommandsTable.updatedAt} ELSE max(${apiCommandsTable.updatedAt} + 1, ${input.nowMs}) END`, + }) + .where( + and( + eq(apiCommandsTable.id, input.commandId), + eq(apiCommandsTable.deliveryGeneration, input.deliveryGeneration), + eq(apiCommandsTable.kind, "app_deployment_run_dispatch"), + eq(apiCommandsTable.payloadJson, input.payloadJson), + recoverableApiCommandPredicate(input.nowMs), + ), + ) + .run(); +} + +export async function reconcileAppDeploymentWorkflowDeliveries( + bindings: ApiCommandDeliveryBindings & Pick, + nowMs = currentTimestampMs(), +): Promise { + const commands = await getAppDatabase(bindings.DB) + .select({ + deliveryGeneration: apiCommandsTable.deliveryGeneration, + id: apiCommandsTable.id, + payloadJson: apiCommandsTable.payloadJson, + }) + .from(apiCommandsTable) + .where( + and( + eq(apiCommandsTable.kind, "app_deployment_run_dispatch"), + recoverableApiCommandPredicate(nowMs), + ), + ) + .orderBy(asc(apiCommandsTable.updatedAt), asc(apiCommandsTable.id)) + .limit(APP_DEPLOYMENT_WORKFLOW_RECONCILE_LIMIT) + .all(); + + for (const command of commands) { + try { + const instanceId = getAppDeploymentWorkflowInstanceId(command.id, command.deliveryGeneration); + let instance: WorkflowInstance; + try { + instance = await bindings.APP_DEPLOYMENT_WORKFLOW.get(instanceId); + } catch { + await sendApiCommandMessage( + bindings, + command.id, + command.deliveryGeneration, + "app_deployment_run_dispatch", + ); + continue; + } + + let status: InstanceStatus["status"]; + try { + status = (await instance.status()).status; + } catch (error) { + logError("api-command.app_deployment_workflow_status_failed", { + ...createErrorLogContext(error), + commandId: command.id, + deliveryGeneration: command.deliveryGeneration, + }); + await sendApiCommandMessage( + bindings, + command.id, + command.deliveryGeneration, + "app_deployment_run_dispatch", + ); + continue; + } + + try { + switch (status) { + case "paused": + case "waitingForPause": { + await instance.resume(); + break; + } + case "complete": + case "errored": + case "terminated": { + await instance.restart(); + break; + } + case "queued": + case "running": + case "waiting": { + break; + } + case "unknown": { + logError("api-command.app_deployment_workflow_status_unknown", { + commandId: command.id, + deliveryGeneration: command.deliveryGeneration, + }); + await sendApiCommandMessage( + bindings, + command.id, + command.deliveryGeneration, + "app_deployment_run_dispatch", + ); + break; + } + } + } catch (error) { + logError("api-command.app_deployment_workflow_recovery_failed", { + ...createErrorLogContext(error), + commandId: command.id, + deliveryGeneration: command.deliveryGeneration, + status, + }); + } + } finally { + try { + await markAppDeploymentWorkflowReconciled({ + commandId: command.id, + database: bindings.DB, + deliveryGeneration: command.deliveryGeneration, + nowMs, + payloadJson: command.payloadJson, + }); + } catch (error) { + logError("api-command.app_deployment_workflow_cursor_update_failed", { + ...createErrorLogContext(error), + commandId: command.id, + deliveryGeneration: command.deliveryGeneration, + }); + } + } } } @@ -209,6 +467,7 @@ export function prepareApiCommand( completedAt: null, createdAt: timestampMs, dedupeKey: normalizeDedupeKey(input.dedupeKey), + deliveryGeneration: 1, id: commandId, kind: input.kind, lastErrorCode: API_COMMAND_QUEUE_DELIVERY_PENDING_CODE, @@ -242,15 +501,23 @@ export async function admitApiCommand( if (current === null) { throw new Error("API command enqueue could not confirm the ledger row."); } + if (current.kind !== input.kind) { + throw new Error("API command dedupe key is already used by a different command kind."); + } if (input.retryTerminal === true && current.status !== "queued" && current.status !== "running") { - await database + if (current.deliveryGeneration >= Number.MAX_SAFE_INTEGER) { + throw new Error("API command delivery generation is exhausted."); + } + + const retried = await database .update(apiCommandsTable) .set({ attemptCount: 0, claimExpiresAt: null, claimOwner: null, completedAt: null, + deliveryGeneration: sql`${apiCommandsTable.deliveryGeneration} + 1`, lastErrorCode: API_COMMAND_QUEUE_DELIVERY_PENDING_CODE, lastErrorMessage: API_COMMAND_QUEUE_DELIVERY_PENDING_MESSAGE, payloadJson: prepared.record.payloadJson, @@ -260,11 +527,13 @@ export async function admitApiCommand( .where( and( eq(apiCommandsTable.id, current.id), + eq(apiCommandsTable.deliveryGeneration, current.deliveryGeneration), inArray(apiCommandsTable.status, ["dead_lettered", "failed", "succeeded"]), ), ) - .run(); - return { commandId: current.id, kind: input.kind, shouldDeliver: true }; + .returning({ id: apiCommandsTable.id }) + .get(); + return { commandId: current.id, kind: input.kind, shouldDeliver: retried !== undefined }; } if ( @@ -286,7 +555,24 @@ export async function deliverApiCommand( return; } - await sendApiCommandMessage(bindings, admission.commandId, admission.kind); + const command = await getAppDatabase(bindings.DB) + .select({ + deliveryGeneration: apiCommandsTable.deliveryGeneration, + kind: apiCommandsTable.kind, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, admission.commandId)) + .limit(1) + .get(); + if (command === undefined) { + throw new Error("API command delivery could not find its durable ledger row."); + } + await sendApiCommandMessage( + bindings, + admission.commandId, + command.deliveryGeneration, + command.kind, + ); } export async function enqueueApiCommand( @@ -301,9 +587,13 @@ export async function enqueueApiCommand( export async function claimApiCommand(input: { commandId: ApiCommandId; database: D1Database; + deliveryGeneration: number; nowMs?: number; - ownerId: string; -}): Promise { + claimOwner: string; +}): Promise { + if (input.claimOwner.trim().length === 0) { + throw new Error("API command claim owner is required."); + } const nowMs = input.nowMs ?? currentTimestampMs(); const row = (await getAppDatabase(input.database) @@ -311,16 +601,24 @@ export async function claimApiCommand(input: { .set({ attemptCount: sql`${apiCommandsTable.attemptCount} + 1`, claimExpiresAt: nowMs + API_COMMAND_LEASE_MS, - claimOwner: input.ownerId, + claimOwner: input.claimOwner, status: "running", updatedAt: nowMs, }) .where( and( eq(apiCommandsTable.id, input.commandId), + eq(apiCommandsTable.deliveryGeneration, input.deliveryGeneration), + sql`typeof(${apiCommandsTable.attemptCount}) = 'integer' AND ${apiCommandsTable.attemptCount} BETWEEN 0 AND ${Number.MAX_SAFE_INTEGER - 1}`, or( eq(apiCommandsTable.status, "queued"), - and(eq(apiCommandsTable.status, "running"), lt(apiCommandsTable.claimExpiresAt, nowMs)), + and( + eq(apiCommandsTable.status, "running"), + or( + isNull(apiCommandsTable.claimExpiresAt), + lte(apiCommandsTable.claimExpiresAt, nowMs), + ), + ), ), ), ) @@ -328,19 +626,62 @@ export async function claimApiCommand(input: { attemptCount: apiCommandsTable.attemptCount, commandId: apiCommandsTable.id, dedupeKey: apiCommandsTable.dedupeKey, + deliveryGeneration: apiCommandsTable.deliveryGeneration, kind: apiCommandsTable.kind, + lastErrorCode: apiCommandsTable.lastErrorCode, + lastErrorMessage: apiCommandsTable.lastErrorMessage, payloadJson: apiCommandsTable.payloadJson, }) .get()) ?? null; - return row; + if (row !== null) { + return { claim: { ...row, claimOwner: input.claimOwner }, kind: "claimed" }; + } + + const state = await getAppDatabase(input.database) + .select({ + attemptCount: apiCommandsTable.attemptCount, + claimExpiresAt: apiCommandsTable.claimExpiresAt, + deliveryGeneration: apiCommandsTable.deliveryGeneration, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, input.commandId)) + .limit(1) + .get(); + if (state === undefined) { + return { kind: "missing" }; + } + if (state.deliveryGeneration !== input.deliveryGeneration) { + return { kind: "stale" }; + } + if (state.status !== "queued" && state.status !== "running") { + return { kind: "terminal" }; + } + if (!Number.isSafeInteger(state.attemptCount) || state.attemptCount < 0) { + throw new Error("API command attempt count is corrupt."); + } + if (state.attemptCount === Number.MAX_SAFE_INTEGER) { + throw new Error("API command attempt count is exhausted."); + } + return { claimExpiresAt: state.claimExpiresAt, kind: "busy" }; +} + +export function exactApiCommandClaimPredicate(claim: ApiCommandClaimAuthority, nowMs: number) { + return and( + eq(apiCommandsTable.id, claim.commandId), + eq(apiCommandsTable.deliveryGeneration, claim.deliveryGeneration), + eq(apiCommandsTable.attemptCount, claim.attemptCount), + eq(apiCommandsTable.status, "running"), + eq(apiCommandsTable.claimOwner, claim.claimOwner), + gt(apiCommandsTable.claimExpiresAt, nowMs), + ); } export async function renewApiCommandClaim(input: { - commandId: ApiCommandId; + claim: ApiCommandClaim; database: D1Database; nowMs?: number; - ownerId: string; }): Promise { const nowMs = input.nowMs ?? currentTimestampMs(); const result = await getAppDatabase(input.database) @@ -349,27 +690,20 @@ export async function renewApiCommandClaim(input: { claimExpiresAt: nowMs + API_COMMAND_LEASE_MS, updatedAt: nowMs, }) - .where( - and( - eq(apiCommandsTable.id, input.commandId), - eq(apiCommandsTable.status, "running"), - eq(apiCommandsTable.claimOwner, input.ownerId), - ), - ) + .where(exactApiCommandClaimPredicate(input.claim, nowMs)) .run(); return getD1ChangeCount(result) > 0; } export async function completeApiCommand(input: { - commandId: ApiCommandId; + claim: ApiCommandClaim; database: D1Database; nowMs?: number; - ownerId: string; -}): Promise { +}): Promise { const nowMs = input.nowMs ?? currentTimestampMs(); - await getAppDatabase(input.database) + const result = await getAppDatabase(input.database) .update(apiCommandsTable) .set({ claimExpiresAt: null, @@ -380,27 +714,21 @@ export async function completeApiCommand(input: { status: "succeeded", updatedAt: nowMs, }) - .where( - and( - eq(apiCommandsTable.id, input.commandId), - eq(apiCommandsTable.status, "running"), - eq(apiCommandsTable.claimOwner, input.ownerId), - ), - ) + .where(exactApiCommandClaimPredicate(input.claim, nowMs)) .run(); + return getD1ChangeCount(result) > 0; } export async function releaseApiCommandForRetry(input: { - commandId: ApiCommandId; + claim: ApiCommandClaim; database: D1Database; errorCode: string; errorMessage: string; nowMs?: number; - ownerId: string; -}): Promise { +}): Promise { const nowMs = input.nowMs ?? currentTimestampMs(); - await getAppDatabase(input.database) + const result = await getAppDatabase(input.database) .update(apiCommandsTable) .set({ claimExpiresAt: null, @@ -410,27 +738,21 @@ export async function releaseApiCommandForRetry(input: { status: "queued", updatedAt: nowMs, }) - .where( - and( - eq(apiCommandsTable.id, input.commandId), - eq(apiCommandsTable.status, "running"), - eq(apiCommandsTable.claimOwner, input.ownerId), - ), - ) + .where(exactApiCommandClaimPredicate(input.claim, nowMs)) .run(); + return getD1ChangeCount(result) > 0; } export async function markApiCommandFailed(input: { - commandId: ApiCommandId; + claim: ApiCommandClaim; database: D1Database; errorCode: string; errorMessage: string; nowMs?: number; - ownerId: string; -}): Promise { +}): Promise { const nowMs = input.nowMs ?? currentTimestampMs(); - await getAppDatabase(input.database) + const result = await getAppDatabase(input.database) .update(apiCommandsTable) .set({ claimExpiresAt: null, @@ -441,26 +763,21 @@ export async function markApiCommandFailed(input: { status: "failed", updatedAt: nowMs, }) - .where( - and( - eq(apiCommandsTable.id, input.commandId), - eq(apiCommandsTable.status, "running"), - eq(apiCommandsTable.claimOwner, input.ownerId), - ), - ) + .where(exactApiCommandClaimPredicate(input.claim, nowMs)) .run(); + return getD1ChangeCount(result) > 0; } export async function markApiCommandDeadLettered(input: { - commandId: ApiCommandId; + claim: ApiCommandClaim; database: D1Database; errorCode: string; errorMessage: string; nowMs?: number; -}): Promise { +}): Promise { const nowMs = input.nowMs ?? currentTimestampMs(); - await getAppDatabase(input.database) + const result = await getAppDatabase(input.database) .update(apiCommandsTable) .set({ claimExpiresAt: null, @@ -471,6 +788,7 @@ export async function markApiCommandDeadLettered(input: { status: "dead_lettered", updatedAt: nowMs, }) - .where(eq(apiCommandsTable.id, input.commandId)) + .where(exactApiCommandClaimPredicate(input.claim, nowMs)) .run(); + return getD1ChangeCount(result) > 0; } diff --git a/apps/api/src/modules/api-command/application/api-command-message.ts b/apps/api/src/modules/api-command/application/api-command-message.ts index f6345241..03bf7efc 100644 --- a/apps/api/src/modules/api-command/application/api-command-message.ts +++ b/apps/api/src/modules/api-command/application/api-command-message.ts @@ -3,6 +3,7 @@ import { parsePlatformId } from "@mosoo/id"; export interface ApiCommandMessage { commandId: ApiCommandId; + deliveryGeneration: number; } export function parseApiCommandMessage(value: unknown): ApiCommandMessage { @@ -11,8 +12,20 @@ export function parseApiCommandMessage(value: unknown): ApiCommandMessage { } const commandId = (value as Record)["commandId"]; + const deliveryGeneration = (value as Record)["deliveryGeneration"]; + + if ( + typeof deliveryGeneration !== "number" || + !Number.isSafeInteger(deliveryGeneration) || + deliveryGeneration <= 0 + ) { + throw new Error( + "API command queue message deliveryGeneration must be a positive safe integer.", + ); + } return { commandId: parsePlatformId(commandId, "API command queue message commandId"), + deliveryGeneration, }; } diff --git a/apps/api/src/modules/api-command/application/api-command-payload.ts b/apps/api/src/modules/api-command/application/api-command-payload.ts index 80fe461c..e4aee160 100644 --- a/apps/api/src/modules/api-command/application/api-command-payload.ts +++ b/apps/api/src/modules/api-command/application/api-command-payload.ts @@ -17,8 +17,10 @@ import type { type ApiCommandPayload = | AppDeploymentRunDispatchCommandPayload + | AppDeploymentScriptReconciliationCommandPayload | CostLedgerReconciliationCommandPayload | EnvironmentPackageArtifactBuildCommandPayload + | SandboxBackupReconciliationCommandPayload | ScheduledMaintenanceCommandPayload | SessionRunDispatchCommandPayload; @@ -34,6 +36,17 @@ export interface CostLedgerReconciliationCommandPayload { scheduledTime: number; } +export interface SandboxBackupReconciliationCommandPayload { + cursor: string | null; + databasePage: number; + scheduledTime: number; +} + +export interface AppDeploymentScriptReconciliationCommandPayload { + cursor: string | null; + scheduledTime: number; +} + export interface AppDeploymentRunDispatchCommandPayload { appDeploymentRunId: AppDeploymentRunId; } @@ -228,6 +241,42 @@ function parseCostLedgerReconciliationPayload( }; } +function parseSandboxBackupReconciliationPayload( + value: unknown, +): SandboxBackupReconciliationCommandPayload { + const label = "sandbox_backup_reconciliation payload"; + const record = requireRecord(value, label); + const cursor = readOptionalString(record, "cursor", label); + if (cursor !== null && cursor.length === 0) { + throw new ApiCommandPayloadError(`${label}.cursor must not be empty.`); + } + const databasePage = readInteger(record, "databasePage", label); + if (databasePage < 0) { + throw new ApiCommandPayloadError(`${label}.databasePage must not be negative.`); + } + const scheduledTime = readInteger(record, "scheduledTime", label); + if (scheduledTime < 0 || !Number.isFinite(new Date(scheduledTime).getTime())) { + throw new ApiCommandPayloadError(`${label}.scheduledTime must be a valid timestamp.`); + } + return { cursor, databasePage, scheduledTime }; +} + +function parseAppDeploymentScriptReconciliationPayload( + value: unknown, +): AppDeploymentScriptReconciliationCommandPayload { + const label = "app_deployment_script_reconciliation payload"; + const record = requireRecord(value, label); + const cursor = readOptionalString(record, "cursor", label); + if (cursor !== null && cursor.length === 0) { + throw new ApiCommandPayloadError(`${label}.cursor must not be empty.`); + } + const scheduledTime = readInteger(record, "scheduledTime", label); + if (scheduledTime < 0 || !Number.isFinite(new Date(scheduledTime).getTime())) { + throw new ApiCommandPayloadError(`${label}.scheduledTime must be a valid timestamp.`); + } + return { cursor, scheduledTime }; +} + function parseAppDeploymentRunDispatchPayload( value: unknown, ): AppDeploymentRunDispatchCommandPayload { @@ -291,12 +340,18 @@ export function parseApiCommandPayload( case "app_deployment_run_dispatch": { return parseAppDeploymentRunDispatchPayload(parsed); } + case "app_deployment_script_reconciliation": { + return parseAppDeploymentScriptReconciliationPayload(parsed); + } case "cost_ledger_reconciliation": { return parseCostLedgerReconciliationPayload(parsed); } case "environment_package_artifact_build": { return parseEnvironmentPackageArtifactBuildPayload(parsed); } + case "sandbox_backup_reconciliation": { + return parseSandboxBackupReconciliationPayload(parsed); + } case "scheduled_maintenance": { return parseScheduledMaintenancePayload(parsed); } diff --git a/apps/api/src/modules/api-command/application/api-command-processor.ts b/apps/api/src/modules/api-command/application/api-command-processor.ts index 1e6ab2fc..50aa448e 100644 --- a/apps/api/src/modules/api-command/application/api-command-processor.ts +++ b/apps/api/src/modules/api-command/application/api-command-processor.ts @@ -1,14 +1,25 @@ -import { apiCommandsTable, appDeploymentRunsTable } from "@mosoo/db"; -import type { ApiCommandId } from "@mosoo/db"; +import { apiCommandsTable, appDeploymentRunsTable, appDeploymentsTable } from "@mosoo/db"; import type { AppDeploymentRunId } from "@mosoo/id"; import { parsePlatformId } from "@mosoo/id"; -import { and, eq, inArray } from "drizzle-orm"; +import { and, eq, exists, inArray, isNull, notExists, or, sql } from "drizzle-orm"; import { createErrorLogContext, logError, logInfo } from "../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; -import { getAppDatabase } from "../../../platform/db/drizzle"; +import { + getAppDatabase, + getD1ChangeCount, + runAppDatabaseBatch, +} from "../../../platform/db/drizzle"; +import type { AppDatabase } from "../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../time"; +import { createCloudflareDeploymentClient } from "../../apps/application/app-deployment-cloudflare-client"; import { dispatchAppDeploymentRun } from "../../apps/application/app-deployment-executor.service"; +import type { AppDeploymentDispatchOutcome } from "../../apps/application/app-deployment-executor.service"; +import { reconcileAppDeploymentScriptPage } from "../../apps/application/app-deployment-script-reconciliation.service"; +import type { + AppDeploymentScriptCleanupAction, + AppDeploymentScriptReconciliationAuthority, +} from "../../apps/application/app-deployment-script-reconciliation.service"; import { ACTIVE_APP_DEPLOYMENT_RUN_STATUSES } from "../../apps/domain/app-deployment-lifecycle"; import { parseCostLedgerReconciliationActivationMode, @@ -17,15 +28,23 @@ import { import { runUsageDailyRollup } from "../../cost/application/cost-rollup.service"; import { buildEnvironmentPackageArtifact } from "../../environments/application/environment-package-artifact-build.service"; import { dispatchQueuedSessionRun } from "../../runtime/application/session-runs/dispatch-queued-run.service"; +import { createLeaseOwnershipRenewal } from "../../runtime/infrastructure/runtime-subject-lifecycle/lease-ownership-renewal"; import { runSandboxMaintenance } from "../../runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance.service"; +import { destroyUnversionedSandboxContainer } from "../../runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-platform"; +import { reconcileSandboxBackupPage } from "../../runtime/infrastructure/sandbox-backup-reconciliation.service"; import { APP_DEPLOYMENT_RUN_DISPATCH_DEDUPE_PREFIX, + enqueueAppDeploymentScriptReconciliationCommand, enqueueCostLedgerReconciliationCommand, + enqueueSandboxBackupReconciliationCommand, } from "./api-command-enqueue"; import { + API_COMMAND_LEASE_MS, API_COMMAND_LEASE_RENEWAL_INTERVAL_MS, claimApiCommand, completeApiCommand, + exactApiCommandClaimPredicate, + handoffAppDeploymentQueueMessage, markApiCommandDeadLettered, markApiCommandFailed, releaseApiCommandForRetry, @@ -37,8 +56,10 @@ import type { ApiCommandMessage } from "./api-command-message"; import { ApiCommandPayloadError, parseApiCommandPayload } from "./api-command-payload"; import type { AppDeploymentRunDispatchCommandPayload, + AppDeploymentScriptReconciliationCommandPayload, CostLedgerReconciliationCommandPayload, EnvironmentPackageArtifactBuildCommandPayload, + SandboxBackupReconciliationCommandPayload, ScheduledMaintenanceCommandPayload, SessionRunDispatchCommandPayload, } from "./api-command-payload"; @@ -50,9 +71,25 @@ import { const API_COMMAND_RETRY_DELAY_SECONDS = 30; -function createClaimOwnerId(message: Message): string { - const normalized = message.id.replaceAll(":", "_").trim(); - return normalized || "api-command-worker"; +const APP_DEPLOYMENT_SCRIPT_CLEANUP_TIMEOUT_MS = 2 * 60 * 1_000; + +export type ApiCommandDeliveryDisposition = + | { readonly kind: "finished" } + | { readonly delaySeconds: number; readonly kind: "retry" }; + +export interface ProcessApiCommandDeliveryOptions { + readonly appDeploymentDeadlineMs?: number; + readonly dispatchAppDeploymentRun?: typeof dispatchAppDeploymentRun; +} + +const FINISHED_DISPOSITION = { kind: "finished" } as const; + +function retryDisposition(delaySeconds = API_COMMAND_RETRY_DELAY_SECONDS) { + return { delaySeconds, kind: "retry" } as const; +} + +function createClaimOwnerId(): string { + return crypto.randomUUID(); } function getErrorMessage(error: unknown): string { @@ -79,6 +116,134 @@ function shouldStartCostLedgerReconciliation(now: Date): boolean { return now.getUTCHours() === 1 && now.getUTCMinutes() === 0; } +function shouldStartSandboxBackupReconciliation(now: Date): boolean { + return now.getUTCHours() === 3 && now.getUTCMinutes() === 0; +} + +function shouldStartAppDeploymentScriptReconciliation(now: Date): boolean { + return now.getUTCHours() === 4 && now.getUTCMinutes() === 0; +} + +async function withAppDeploymentCleanupTimeout( + operation: Promise, + label: string, +): Promise { + let timeout: ReturnType | undefined; + const timeoutPromise = new Promise((_resolve, reject) => { + timeout = setTimeout( + () => reject(new Error(`${label} exceeded its cleanup timeout.`)), + APP_DEPLOYMENT_SCRIPT_CLEANUP_TIMEOUT_MS, + ); + }); + try { + return await Promise.race([operation, timeoutPromise]); + } finally { + clearTimeout(timeout); + } +} + +export async function deleteAppDeploymentArtifactPrefix( + bucket: R2Bucket, + prefix: string, +): Promise { + const keys: string[] = []; + let page = await bucket.list({ limit: 1_000, prefix }); + while (page.truncated) { + keys.push(...page.objects.map(({ key }) => key)); + if (!page.cursor) { + throw new Error("App deployment artifact cleanup returned no continuation cursor."); + } + page = await bucket.list({ cursor: page.cursor, limit: 1_000, prefix }); + } + keys.push(...page.objects.map(({ key }) => key)); + for (let offset = 0; offset < keys.length; offset += 1_000) { + await bucket.delete(keys.slice(offset, offset + 1_000)); + } +} + +function assertAppDeploymentCleanupSucceeded( + results: readonly PromiseSettledResult[], + scriptName: string, +): void { + const errors = results.flatMap((result) => (result.status === "rejected" ? [result.reason] : [])); + if (errors.length > 0) { + throw new AggregateError(errors, `App deployment script cleanup failed for ${scriptName}.`); + } +} + +export async function cleanupAppDeploymentScript( + bindings: ApiBindings, + action: AppDeploymentScriptCleanupAction, +): Promise { + const producerResults = await Promise.allSettled([ + destroyUnversionedSandboxContainer( + bindings, + action.buildSandboxId, + APP_DEPLOYMENT_SCRIPT_CLEANUP_TIMEOUT_MS, + ), + destroyUnversionedSandboxContainer( + bindings, + action.deploySandboxId, + APP_DEPLOYMENT_SCRIPT_CLEANUP_TIMEOUT_MS, + ), + ]); + assertAppDeploymentCleanupSucceeded(producerResults, action.scriptName); + + const outputResults = await Promise.allSettled([ + withAppDeploymentCleanupTimeout( + deleteAppDeploymentArtifactPrefix(bindings.FILE_BUCKET, action.artifactPrefix), + "App deployment artifact cleanup", + ), + action.uploadStartedAt === null + ? Promise.resolve() + : (async () => + createCloudflareDeploymentClient(bindings).deleteWorkerScript({ + scriptName: action.scriptName, + timeoutMs: APP_DEPLOYMENT_SCRIPT_CLEANUP_TIMEOUT_MS, + }))(), + ]); + assertAppDeploymentCleanupSucceeded(outputResults, action.scriptName); +} + +async function processAppDeploymentScriptReconciliationCommand( + bindings: ApiBindings, + payload: AppDeploymentScriptReconciliationCommandPayload, + processedAtMs: number, + authority: AppDeploymentScriptReconciliationAuthority, +): Promise { + const result = await reconcileAppDeploymentScriptPage( + bindings.DB, + { authority, cursor: payload.cursor }, + (action) => cleanupAppDeploymentScript(bindings, action), + ); + for (const failure of result.failures) { + logError("app-deployment.script_reconciliation_cleanup_failed", { + ...createErrorLogContext(failure.error), + scriptName: failure.scriptName, + }); + } + logInfo("app-deployment.script_reconciliation_page_completed", { + armed: result.armed, + cleaned: result.cleaned, + deferred: result.deferred, + failureCount: result.failures.length, + processed: result.processed, + processedAtMs, + scheduledTime: payload.scheduledTime, + }); + if (!result.hasMore) { + return; + } + if (result.nextCursor === null) { + throw new Error("App deployment script reconciliation returned no continuation cursor."); + } + await authority.requireOwnership(); + await enqueueAppDeploymentScriptReconciliationCommand(bindings, { + cursor: result.nextCursor, + scheduledTime: payload.scheduledTime, + }); +} + async function processScheduledMaintenanceCommand( bindings: ApiBindings, payload: ScheduledMaintenanceCommandPayload, @@ -106,9 +271,54 @@ async function processScheduledMaintenanceCommand( } } + if (shouldStartSandboxBackupReconciliation(scheduledAt)) { + tasks.push( + enqueueSandboxBackupReconciliationCommand(bindings, { + cursor: null, + databasePage: 0, + scheduledTime: payload.scheduledTime, + }), + ); + } + + if (shouldStartAppDeploymentScriptReconciliation(scheduledAt)) { + tasks.push( + enqueueAppDeploymentScriptReconciliationCommand(bindings, { + cursor: null, + scheduledTime: payload.scheduledTime, + }), + ); + } + await Promise.all(tasks); } +async function processSandboxBackupReconciliationCommand( + bindings: ApiBindings, + payload: SandboxBackupReconciliationCommandPayload, + processedAtMs: number, +): Promise { + const result = await reconcileSandboxBackupPage(bindings, { + cursor: payload.cursor, + }); + logInfo("runtime.sandbox_backup.reconciliation_page_completed", { + ...result, + processedAtMs, + scheduledTime: payload.scheduledTime, + }); + if (!result.hasMore) { + return; + } + if (payload.databasePage === Number.MAX_SAFE_INTEGER) { + throw new Error("Sandbox backup reconciliation exhausted its database page identity."); + } + await enqueueSandboxBackupReconciliationCommand(bindings, { + cursor: result.nextCursor, + databasePage: payload.databasePage + 1, + scheduledTime: payload.scheduledTime, + }); +} + async function processCostLedgerReconciliationCommand( bindings: ApiBindings, payload: CostLedgerReconciliationCommandPayload, @@ -162,59 +372,6 @@ async function processSessionRunDispatchCommand( }); } -async function failActiveAppDeploymentRun( - bindings: ApiBindings, - runId: AppDeploymentRunId, - input: { errorCode: string; errorMessage: string; nowMs: number }, -): Promise { - await getAppDatabase(bindings.DB) - .update(appDeploymentRunsTable) - .set({ - errorCode: input.errorCode, - errorMessage: input.errorMessage, - status: "failed", - updatedAt: input.nowMs, - }) - .where( - and( - eq(appDeploymentRunsTable.id, runId), - inArray(appDeploymentRunsTable.status, ACTIVE_APP_DEPLOYMENT_RUN_STATUSES), - ), - ) - .run(); -} - -async function failAppDeploymentRunFromPayloadJson( - bindings: ApiBindings, - input: { - errorCode: string; - errorMessage: string; - fallbackDedupeKey?: string; - nowMs: number; - payloadJson: string; - }, - logEventName?: string, -): Promise { - try { - const runId = readAppDeploymentRunIdFromPayload(input); - - if (runId === null) { - return; - } - - await failActiveAppDeploymentRun(bindings, runId, input); - } catch (error) { - if (logEventName === undefined) { - throw error; - } - - logError(logEventName, { - ...createErrorLogContext(error), - errorCode: getErrorCode(error), - }); - } -} - function readAppDeploymentRunIdFromPayload(input: { fallbackDedupeKey?: string; payloadJson: string; @@ -255,17 +412,341 @@ function readAppDeploymentRunIdFromPayload(input: { } } +function claimedApiCommandExists(database: AppDatabase, claim: ApiCommandClaim, nowMs: number) { + return exists( + database + .select({ id: apiCommandsTable.id }) + .from(apiCommandsTable) + .where(exactClaimedApiCommand(claim, nowMs)), + ); +} + +function exactClaimedApiCommand(claim: ApiCommandClaim, nowMs: number) { + return and( + exactApiCommandClaimPredicate(claim, nowMs), + eq(apiCommandsTable.kind, claim.kind), + eq(apiCommandsTable.payloadJson, claim.payloadJson), + ); +} + +async function completeAppDeploymentCommand( + bindings: ApiBindings, + claim: ApiCommandClaim, + outcome: Extract, + nowMs: number, +): Promise { + const results = await runAppDatabaseBatch(bindings.DB, (database) => { + const deploymentIsCurrent = exists( + database + .select({ id: appDeploymentsTable.id }) + .from(appDeploymentsTable) + .where( + and( + eq(appDeploymentsTable.id, outcome.deploymentId), + eq(appDeploymentsTable.latestRunId, outcome.runId), + isNull(appDeploymentsTable.deletedAt), + ), + ), + ); + const runIsActivating = exists( + database + .select({ id: appDeploymentRunsTable.id }) + .from(appDeploymentRunsTable) + .where( + and( + eq(appDeploymentRunsTable.id, outcome.runId), + eq(appDeploymentRunsTable.deploymentId, outcome.deploymentId), + eq(appDeploymentRunsTable.status, "activating"), + eq(appDeploymentRunsTable.targetScriptName, outcome.activeScriptName), + ), + ), + ); + const deploymentMatchesOutcome = exists( + database + .select({ id: appDeploymentsTable.id }) + .from(appDeploymentsTable) + .where( + and( + eq(appDeploymentsTable.id, outcome.deploymentId), + eq(appDeploymentsTable.activeScriptName, outcome.activeScriptName), + eq(appDeploymentsTable.latestRunId, outcome.runId), + eq(appDeploymentsTable.lastSuccessfulUrl, outcome.url), + isNull(appDeploymentsTable.deletedAt), + ), + ), + ); + const runMatchesOutcome = exists( + database + .select({ id: appDeploymentRunsTable.id }) + .from(appDeploymentRunsTable) + .where( + and( + eq(appDeploymentRunsTable.id, outcome.runId), + eq(appDeploymentRunsTable.deploymentId, outcome.deploymentId), + eq(appDeploymentRunsTable.status, "success"), + eq(appDeploymentRunsTable.targetScriptName, outcome.activeScriptName), + eq(appDeploymentRunsTable.url, outcome.url), + outcome.externalDeploymentId === null + ? isNull(appDeploymentRunsTable.externalDeploymentId) + : eq(appDeploymentRunsTable.externalDeploymentId, outcome.externalDeploymentId), + outcome.externalProjectId === null + ? isNull(appDeploymentRunsTable.externalProjectId) + : eq(appDeploymentRunsTable.externalProjectId, outcome.externalProjectId), + outcome.externalVersionId === null + ? isNull(appDeploymentRunsTable.externalVersionId) + : eq(appDeploymentRunsTable.externalVersionId, outcome.externalVersionId), + ), + ), + ); + + return [ + database + .update(appDeploymentsTable) + .set({ + activeScriptName: outcome.activeScriptName, + lastSuccessfulUrl: outcome.url, + updatedAt: nowMs, + }) + .where( + and( + eq(appDeploymentsTable.id, outcome.deploymentId), + eq(appDeploymentsTable.latestRunId, outcome.runId), + isNull(appDeploymentsTable.deletedAt), + runIsActivating, + claimedApiCommandExists(database, claim, nowMs), + ), + ), + database + .update(appDeploymentRunsTable) + .set({ + errorCode: null, + errorMessage: null, + externalDeploymentId: outcome.externalDeploymentId, + externalProjectId: outcome.externalProjectId, + externalVersionId: outcome.externalVersionId, + status: "success", + updatedAt: nowMs, + url: outcome.url, + }) + .where( + and( + eq(appDeploymentRunsTable.id, outcome.runId), + eq(appDeploymentRunsTable.deploymentId, outcome.deploymentId), + eq(appDeploymentRunsTable.status, "activating"), + eq(appDeploymentRunsTable.targetScriptName, outcome.activeScriptName), + deploymentIsCurrent, + claimedApiCommandExists(database, claim, nowMs), + ), + ), + database + .update(apiCommandsTable) + .set({ + claimExpiresAt: null, + claimOwner: null, + completedAt: nowMs, + lastErrorCode: null, + lastErrorMessage: null, + status: "succeeded", + updatedAt: nowMs, + }) + .where( + and(exactClaimedApiCommand(claim, nowMs), deploymentMatchesOutcome, runMatchesOutcome), + ), + ]; + }); + + return getD1ChangeCount(results[2]) > 0; +} + +async function failAppDeploymentCommand( + bindings: ApiBindings, + claim: ApiCommandClaim, + input: { + errorCode: string; + errorMessage: string; + nowMs: number; + runId: AppDeploymentRunId; + status: "dead_lettered" | "failed"; + }, +): Promise { + const results = await runAppDatabaseBatch(bindings.DB, (database) => { + const runExists = database + .select({ id: appDeploymentRunsTable.id }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, input.runId)); + const runIsFailed = database + .select({ id: appDeploymentRunsTable.id }) + .from(appDeploymentRunsTable) + .where( + and( + eq(appDeploymentRunsTable.id, input.runId), + eq(appDeploymentRunsTable.status, "failed"), + ), + ); + + return [ + database + .update(appDeploymentRunsTable) + .set({ + errorCode: input.errorCode, + errorMessage: input.errorMessage, + status: "failed", + updatedAt: input.nowMs, + }) + .where( + and( + eq(appDeploymentRunsTable.id, input.runId), + inArray(appDeploymentRunsTable.status, ACTIVE_APP_DEPLOYMENT_RUN_STATUSES), + claimedApiCommandExists(database, claim, input.nowMs), + ), + ), + database + .update(apiCommandsTable) + .set({ + claimExpiresAt: null, + claimOwner: null, + completedAt: input.nowMs, + lastErrorCode: sql`coalesce( + (select ${appDeploymentRunsTable.errorCode} from ${appDeploymentRunsTable} + where ${appDeploymentRunsTable.id} = ${input.runId}), + ${input.errorCode} + )`, + lastErrorMessage: sql`coalesce( + (select ${appDeploymentRunsTable.errorMessage} from ${appDeploymentRunsTable} + where ${appDeploymentRunsTable.id} = ${input.runId}), + ${input.errorMessage} + )`, + status: input.status, + updatedAt: input.nowMs, + }) + .where( + and( + exactClaimedApiCommand(claim, input.nowMs), + or(exists(runIsFailed), notExists(runExists)), + ), + ), + ]; + }); + + return getD1ChangeCount(results[1]) > 0; +} + +async function reconcileSkippedAppDeploymentCommand( + bindings: ApiBindings, + claim: ApiCommandClaim, + nowMs: number, +): Promise { + const runId = readAppDeploymentRunIdFromPayload({ + fallbackDedupeKey: claim.dedupeKey, + payloadJson: claim.payloadJson, + }); + if (runId === null) { + return markApiCommandFailed({ + claim, + database: bindings.DB, + errorCode: "app_deployment_run_missing", + errorMessage: "App deployment command has no valid deployment run.", + nowMs, + }); + } + + const run = await getAppDatabase(bindings.DB) + .select({ + deploymentId: appDeploymentRunsTable.deploymentId, + errorCode: appDeploymentRunsTable.errorCode, + errorMessage: appDeploymentRunsTable.errorMessage, + status: appDeploymentRunsTable.status, + }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, runId)) + .limit(1) + .get(); + if (run === undefined) { + return markApiCommandFailed({ + claim, + database: bindings.DB, + errorCode: "app_deployment_run_missing", + errorMessage: "App deployment run no longer exists.", + nowMs, + }); + } + if (run.status === "success") { + return completeApiCommand({ claim, database: bindings.DB, nowMs }); + } + if (run.status === "failed") { + return markApiCommandFailed({ + claim, + database: bindings.DB, + errorCode: run.errorCode ?? "app_deployment_run_failed", + errorMessage: run.errorMessage ?? "App deployment run failed.", + nowMs, + }); + } + + const currentDeployment = await getAppDatabase(bindings.DB) + .select({ id: appDeploymentsTable.id }) + .from(appDeploymentsTable) + .where( + and( + eq(appDeploymentsTable.id, run.deploymentId), + eq(appDeploymentsTable.latestRunId, runId), + isNull(appDeploymentsTable.deletedAt), + ), + ) + .limit(1) + .get(); + if (currentDeployment !== undefined) { + return false; + } + + return failAppDeploymentCommand(bindings, claim, { + errorCode: "app_deployment_context_lost", + errorMessage: "App deployment is no longer current or active.", + nowMs, + runId, + status: "failed", + }); +} + async function processClaimedApiCommand( bindings: ApiBindings, claim: ApiCommandClaim, processedAtMs: number, -): Promise { + requireOwnership: () => Promise, + dispatchDeployment: typeof dispatchAppDeploymentRun, + appDeploymentDeadlineMs?: number, +): Promise< + { kind: "completed" } | { kind: "app_deployment"; outcome: AppDeploymentDispatchOutcome } +> { const payload = parseApiCommandPayload(claim.kind, claim.payloadJson); + const authority = { + attemptCount: claim.attemptCount, + claimOwner: claim.claimOwner, + commandId: claim.commandId, + deliveryGeneration: claim.deliveryGeneration, + requireOwnership, + }; + + await requireOwnership(); switch (claim.kind) { case "app_deployment_run_dispatch": { - await dispatchAppDeploymentRun(bindings, payload as AppDeploymentRunDispatchCommandPayload); - return; + const outcome = await dispatchDeployment( + bindings, + payload as AppDeploymentRunDispatchCommandPayload, + authority, + appDeploymentDeadlineMs === undefined ? {} : { deadlineMs: appDeploymentDeadlineMs }, + ); + return { kind: "app_deployment", outcome }; + } + case "app_deployment_script_reconciliation": { + await processAppDeploymentScriptReconciliationCommand( + bindings, + payload as AppDeploymentScriptReconciliationCommandPayload, + processedAtMs, + authority, + ); + break; } case "cost_ledger_reconciliation": { await processCostLedgerReconciliationCommand( @@ -273,148 +754,229 @@ async function processClaimedApiCommand( payload as CostLedgerReconciliationCommandPayload, processedAtMs, ); - return; + break; } case "environment_package_artifact_build": { await buildEnvironmentPackageArtifact( bindings, payload as EnvironmentPackageArtifactBuildCommandPayload, + authority, ); - return; + break; + } + case "sandbox_backup_reconciliation": { + await processSandboxBackupReconciliationCommand( + bindings, + payload as SandboxBackupReconciliationCommandPayload, + processedAtMs, + ); + break; } case "scheduled_maintenance": { await processScheduledMaintenanceCommand( bindings, payload as ScheduledMaintenanceCommandPayload, ); - return; + break; } case "session_run_dispatch": { await processSessionRunDispatchCommand(bindings, payload as SessionRunDispatchCommandPayload); - return; + break; } } + + await requireOwnership(); + return { kind: "completed" }; } async function processClaimedApiCommandWithLeaseRenewal( bindings: ApiBindings, claim: ApiCommandClaim, - ownerId: string, - processedAtMs: number, -): Promise { + nowMs: () => number, + dispatchDeployment: typeof dispatchAppDeploymentRun, + appDeploymentDeadlineMs?: number, +): Promise>> { let stopped = false; - let renewal = Promise.resolve(); + let lossEvent: "api-command.claim_lost" | "api-command.claim_renew_failed" | null = null; + let lossLogged = false; + let pendingHeartbeat = Promise.resolve(); + const requireOwnership = createLeaseOwnershipRenewal(async () => { + try { + const renewed = await renewApiCommandClaim({ + claim, + database: bindings.DB, + nowMs: nowMs(), + }); + if (!renewed) { + lossEvent ??= "api-command.claim_lost"; + } + return renewed; + } catch (error) { + lossEvent ??= "api-command.claim_renew_failed"; + throw error; + } + }, "API command lease ownership was lost."); + const logOwnershipLoss = (error: unknown): void => { + if (lossLogged || lossEvent === null) { + return; + } + lossLogged = true; + logError(lossEvent, { + ...createErrorLogContext(error), + attemptCount: claim.attemptCount, + commandId: claim.commandId, + deliveryGeneration: claim.deliveryGeneration, + kind: claim.kind, + }); + }; const timer = setInterval(() => { if (stopped) { return; } - renewal = renewal - .then(() => - renewApiCommandClaim({ - commandId: claim.commandId, - database: bindings.DB, - ownerId, - }), - ) - .then((renewed) => { - if (renewed) { - return; - } - - stopped = true; - logError("api-command.claim_lost", { - commandId: claim.commandId, - kind: claim.kind, - }); - }) - .catch((error: unknown) => { - logError("api-command.claim_renew_failed", { - ...createErrorLogContext(error), - commandId: claim.commandId, - kind: claim.kind, - }); - }); + pendingHeartbeat = requireOwnership().catch(logOwnershipLoss); }, API_COMMAND_LEASE_RENEWAL_INTERVAL_MS); try { - await processClaimedApiCommand(bindings, claim, processedAtMs); - stopped = true; - await renewal; + await requireOwnership(); + const result = await processClaimedApiCommand( + bindings, + claim, + nowMs(), + requireOwnership, + dispatchDeployment, + appDeploymentDeadlineMs, + ); + await requireOwnership(); + return result; + } catch (error) { + logOwnershipLoss(error); + throw error; } finally { stopped = true; clearInterval(timer); + await pendingHeartbeat; } } -export async function processApiCommandMessage( - bindings: ApiBindings, - message: Message, - nowMs: () => number = currentTimestampMs, -): Promise { - let commandId: ApiCommandId; +function busyRetryDelaySeconds(claimExpiresAt: number | null, nowMs: number): number { + if (claimExpiresAt === null) { + return API_COMMAND_RETRY_DELAY_SECONDS; + } + return Math.min( + API_COMMAND_LEASE_MS / 1_000, + Math.max(1, Math.ceil((claimExpiresAt - nowMs) / 1_000)), + ); +} +function retryMessage(message: Message, delaySeconds: number): void { + message.retry({ delaySeconds }); +} + +async function handoffLegacyAppDeploymentQueueMessage( + bindings: ApiBindings, + message: ApiCommandMessage, +): Promise<"continue" | ApiCommandDeliveryDisposition> { try { - commandId = parseApiCommandMessage(message.body).commandId; + const handoff = await handoffAppDeploymentQueueMessage(bindings, message); + if (handoff === "continue") { + return "continue"; + } + return handoff === "finished" ? FINISHED_DISPOSITION : retryDisposition(); } catch (error) { - logError("api-command.message_invalid", { + logError("api-command.app_deployment_workflow_handoff_failed", { ...createErrorLogContext(error), - errorCode: getErrorCode(error), + commandId: message.commandId, + deliveryGeneration: message.deliveryGeneration, }); - message.ack(); - return; + return retryDisposition(); } +} - const ownerId = createClaimOwnerId(message); - const startMs = nowMs(); - const claim = await claimApiCommand({ - commandId, - database: bindings.DB, - nowMs: startMs, - ownerId, +async function finalizeProcessedCommand( + bindings: ApiBindings, + claim: ApiCommandClaim, + result: Awaited>, + nowMs: number, +): Promise { + if (result.kind === "completed") { + return completeApiCommand({ claim, database: bindings.DB, nowMs }); + } + if (result.outcome.kind === "skipped") { + const reconciled = await reconcileSkippedAppDeploymentCommand(bindings, claim, nowMs); + if (!reconciled) { + await releaseApiCommandForRetry({ + claim, + database: bindings.DB, + errorCode: "app_deployment_dispatch_incomplete", + errorMessage: "App deployment dispatch did not reach a terminal state.", + nowMs, + }); + } + return reconciled; + } + if (result.outcome.kind === "succeeded") { + return completeAppDeploymentCommand(bindings, claim, result.outcome, nowMs); + } + return failAppDeploymentCommand(bindings, claim, { + errorCode: result.outcome.errorCode, + errorMessage: result.outcome.errorMessage, + nowMs, + runId: result.outcome.runId, + status: "failed", }); +} - if (!claim) { - message.ack(); - return; +async function terminalizeFailedCommand( + bindings: ApiBindings, + claim: ApiCommandClaim, + input: { errorCode: string; errorMessage: string; nowMs: number }, +): Promise { + if (claim.kind !== "app_deployment_run_dispatch") { + return markApiCommandFailed({ claim, database: bindings.DB, ...input }); } - try { - await processClaimedApiCommandWithLeaseRenewal(bindings, claim, ownerId, nowMs()); - await completeApiCommand({ - commandId, - database: bindings.DB, - nowMs: nowMs(), - ownerId, - }); - message.ack(); - } catch (error) { - const errorCode = getErrorCode(error); - const errorMessage = getErrorMessage(error); + const runId = readAppDeploymentRunIdFromPayload({ + fallbackDedupeKey: claim.dedupeKey, + payloadJson: claim.payloadJson, + }); + if (runId === null) { + return markApiCommandFailed({ claim, database: bindings.DB, ...input }); + } + return failAppDeploymentCommand(bindings, claim, { + ...input, + runId, + status: "failed", + }); +} - logError("api-command.failed", { - ...createErrorLogContext(error), - attemptCount: claim.attemptCount, - commandId, - errorCode, - kind: claim.kind, - }); +async function handleClaimedCommandFailure( + bindings: ApiBindings, + claim: ApiCommandClaim, + error: unknown, + nowMs: () => number, +): Promise { + const errorCode = getErrorCode(error); + const errorMessage = getErrorMessage(error); + + logError("api-command.failed", { + ...createErrorLogContext(error), + attemptCount: claim.attemptCount, + commandId: claim.commandId, + deliveryGeneration: claim.deliveryGeneration, + errorCode, + kind: claim.kind, + }); - const appDeploymentRunHasTerminalError = - claim.kind === "app_deployment_run_dispatch" && - (error instanceof ApiCommandPayloadError || - (error instanceof Error && - (error.name === "AppDeploymentDetectionError" || - error.name === "AppDeploymentNonRetryableError"))); - const appDeploymentRunRetryExhausted = - claim.kind === "app_deployment_run_dispatch" && - !appDeploymentRunHasTerminalError && - claim.attemptCount >= APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS; - const shouldFailAppDeploymentRun = - appDeploymentRunHasTerminalError || appDeploymentRunRetryExhausted; - - if (shouldFailAppDeploymentRun) { - const failedAtMs = nowMs(); + const appDeploymentPayloadInvalid = + claim.kind === "app_deployment_run_dispatch" && error instanceof ApiCommandPayloadError; + const appDeploymentRunRetryExhausted = + claim.kind === "app_deployment_run_dispatch" && + !appDeploymentPayloadInvalid && + claim.attemptCount >= APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS; + + try { + if (appDeploymentPayloadInvalid || appDeploymentRunRetryExhausted) { const failureCode = appDeploymentRunRetryExhausted ? APP_DEPLOYMENT_RUN_DISPATCH_RETRY_EXHAUSTED_CODE : errorCode; @@ -424,113 +986,242 @@ export async function processApiCommandMessage( lastErrorMessage: errorMessage, }) : errorMessage; - - await failAppDeploymentRunFromPayloadJson( - bindings, - { - errorCode: failureCode, - errorMessage: failureMessage, - fallbackDedupeKey: claim.dedupeKey, - nowMs: failedAtMs, - payloadJson: claim.payloadJson, - }, - "api-command.app_deployment_run_fail_failed", - ); - - await markApiCommandFailed({ - commandId, - database: bindings.DB, + const terminalized = await terminalizeFailedCommand(bindings, claim, { errorCode: failureCode, errorMessage: failureMessage, - nowMs: failedAtMs, - ownerId, + nowMs: nowMs(), }); - message.ack(); - return; + return terminalized ? FINISHED_DISPOSITION : retryDisposition(); } if (error instanceof ApiCommandPayloadError) { - await markApiCommandFailed({ - commandId, + const terminalized = await markApiCommandFailed({ + claim, database: bindings.DB, errorCode, errorMessage, nowMs: nowMs(), - ownerId, }); - message.ack(); - return; + return terminalized ? FINISHED_DISPOSITION : retryDisposition(); } await releaseApiCommandForRetry({ - commandId, + claim, database: bindings.DB, errorCode, errorMessage, nowMs: nowMs(), - ownerId, }); - message.retry({ delaySeconds: API_COMMAND_RETRY_DELAY_SECONDS }); + } catch (persistenceError) { + logError("api-command.failure_persist_failed", { + ...createErrorLogContext(persistenceError), + commandId: claim.commandId, + deliveryGeneration: claim.deliveryGeneration, + }); + } + + return retryDisposition(); +} + +export async function processApiCommandDelivery( + bindings: ApiBindings, + body: unknown, + nowMs: () => number = currentTimestampMs, + options: ProcessApiCommandDeliveryOptions = {}, +): Promise { + let queueMessage: ApiCommandMessage; + + try { + queueMessage = parseApiCommandMessage(body); + } catch (error) { + logError("api-command.message_invalid", { + ...createErrorLogContext(error), + errorCode: getErrorCode(error), + }); + return FINISHED_DISPOSITION; + } + + const claimStartedAtMs = nowMs(); + let claimResult: Awaited>; + try { + claimResult = await claimApiCommand({ + claimOwner: createClaimOwnerId(), + commandId: queueMessage.commandId, + database: bindings.DB, + deliveryGeneration: queueMessage.deliveryGeneration, + nowMs: claimStartedAtMs, + }); + } catch (error) { + logError("api-command.claim_failed", { + ...createErrorLogContext(error), + commandId: queueMessage.commandId, + deliveryGeneration: queueMessage.deliveryGeneration, + }); + return retryDisposition(); + } + + if (claimResult.kind === "busy") { + return retryDisposition(busyRetryDelaySeconds(claimResult.claimExpiresAt, claimStartedAtMs)); + } + if (claimResult.kind !== "claimed") { + return FINISHED_DISPOSITION; + } + + const claim = claimResult.claim; + if ( + claim.kind === "app_deployment_run_dispatch" && + claim.attemptCount > APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS + ) { + try { + const terminalized = await terminalizeFailedCommand(bindings, claim, { + errorCode: APP_DEPLOYMENT_RUN_DISPATCH_RETRY_EXHAUSTED_CODE, + errorMessage: createAppDeploymentDispatchRetryExhaustedMessage({ + attemptCount: APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS, + lastErrorMessage: claim.lastErrorMessage, + }), + nowMs: nowMs(), + }); + return terminalized ? FINISHED_DISPOSITION : retryDisposition(); + } catch (error) { + logError("api-command.retry_exhaustion_persist_failed", { + ...createErrorLogContext(error), + attemptCount: claim.attemptCount, + commandId: claim.commandId, + deliveryGeneration: claim.deliveryGeneration, + }); + return retryDisposition(); + } + } + + let result: Awaited>; + try { + result = await processClaimedApiCommandWithLeaseRenewal( + bindings, + claim, + nowMs, + options.dispatchAppDeploymentRun ?? dispatchAppDeploymentRun, + options.appDeploymentDeadlineMs, + ); + } catch (error) { + return handleClaimedCommandFailure(bindings, claim, error, nowMs); + } + + try { + const finalized = await finalizeProcessedCommand(bindings, claim, result, nowMs()); + return finalized ? FINISHED_DISPOSITION : retryDisposition(); + } catch (error) { + logError("api-command.completion_failed", { + ...createErrorLogContext(error), + commandId: claim.commandId, + deliveryGeneration: claim.deliveryGeneration, + }); + return retryDisposition(); } } +export async function processApiCommandMessage( + bindings: ApiBindings, + message: Message, + nowMs: () => number = currentTimestampMs, + options: ProcessApiCommandDeliveryOptions = {}, +): Promise { + let queueMessage: ApiCommandMessage; + try { + queueMessage = parseApiCommandMessage(message.body); + } catch (error) { + logError("api-command.message_invalid", { + ...createErrorLogContext(error), + errorCode: getErrorCode(error), + }); + message.ack(); + return; + } + const handoff = await handoffLegacyAppDeploymentQueueMessage(bindings, queueMessage); + const disposition = + handoff === "continue" + ? await processApiCommandDelivery(bindings, queueMessage, nowMs, options) + : handoff; + if (disposition.kind === "finished") { + message.ack(); + return; + } + retryMessage(message, disposition.delaySeconds); +} + export async function processApiCommandDeadLetterMessage( bindings: ApiBindings, message: Message, nowMs: () => number = currentTimestampMs, ): Promise { + let queueMessage: ApiCommandMessage; try { - const { commandId } = parseApiCommandMessage(message.body); - const deadLetteredAtMs = nowMs(); - const command = - (await getAppDatabase(bindings.DB) - .select({ - dedupeKey: apiCommandsTable.dedupeKey, - kind: apiCommandsTable.kind, - lastErrorCode: apiCommandsTable.lastErrorCode, - lastErrorMessage: apiCommandsTable.lastErrorMessage, - payloadJson: apiCommandsTable.payloadJson, - }) - .from(apiCommandsTable) - .where(eq(apiCommandsTable.id, commandId)) - .limit(1) - .get()) ?? null; + queueMessage = parseApiCommandMessage(message.body); + } catch (error) { + logError("api-command.dead_letter_invalid", { + ...createErrorLogContext(error), + errorCode: getErrorCode(error), + }); + message.ack(); + return; + } - if (command?.kind === "app_deployment_run_dispatch") { - await failAppDeploymentRunFromPayloadJson( - bindings, - { - errorCode: "queue_dead_lettered", - errorMessage: "Deployment dispatch reached the queue dead-letter consumer.", - fallbackDedupeKey: command.dedupeKey, - nowMs: deadLetteredAtMs, - payloadJson: command.payloadJson, - }, - "api-command.app_deployment_run_dead_letter_fail_failed", - ); + const handoff = await handoffLegacyAppDeploymentQueueMessage(bindings, queueMessage); + if (handoff !== "continue") { + if (handoff.kind === "finished") { + message.ack(); + } else { + retryMessage(message, handoff.delaySeconds); } + return; + } - const preserveArtifactFailure = command?.kind === "environment_package_artifact_build"; + const claimStartedAtMs = nowMs(); + try { + const claimResult = await claimApiCommand({ + claimOwner: createClaimOwnerId(), + commandId: queueMessage.commandId, + database: bindings.DB, + deliveryGeneration: queueMessage.deliveryGeneration, + nowMs: claimStartedAtMs, + }); + if (claimResult.kind === "busy") { + retryMessage(message, busyRetryDelaySeconds(claimResult.claimExpiresAt, claimStartedAtMs)); + return; + } + if (claimResult.kind !== "claimed") { + message.ack(); + return; + } - await markApiCommandDeadLettered({ - commandId, + const claim = claimResult.claim; + const preserveArtifactFailure = claim.kind === "environment_package_artifact_build"; + const errorCode = + preserveArtifactFailure && claimResult.claim.lastErrorCode + ? claimResult.claim.lastErrorCode + : "queue_dead_lettered"; + const errorMessage = + preserveArtifactFailure && claimResult.claim.lastErrorMessage + ? claimResult.claim.lastErrorMessage + : "API command reached the queue dead-letter consumer."; + const deadLettered = await markApiCommandDeadLettered({ + claim, database: bindings.DB, - errorCode: - preserveArtifactFailure && command.lastErrorCode - ? command.lastErrorCode - : "queue_dead_lettered", - errorMessage: - preserveArtifactFailure && command.lastErrorMessage - ? command.lastErrorMessage - : "API command reached the queue dead-letter consumer.", - nowMs: deadLetteredAtMs, + errorCode, + errorMessage, + nowMs: nowMs(), }); + + if (deadLettered) { + message.ack(); + } else { + retryMessage(message, API_COMMAND_RETRY_DELAY_SECONDS); + } } catch (error) { - logError("api-command.dead_letter_invalid", { + logError("api-command.dead_letter_failed", { ...createErrorLogContext(error), - errorCode: getErrorCode(error), + commandId: queueMessage.commandId, + deliveryGeneration: queueMessage.deliveryGeneration, }); + retryMessage(message, API_COMMAND_RETRY_DELAY_SECONDS); } - - message.ack(); } diff --git a/apps/api/src/modules/api-command/application/api-command-workflow.ts b/apps/api/src/modules/api-command/application/api-command-workflow.ts new file mode 100644 index 00000000..3ba52f15 --- /dev/null +++ b/apps/api/src/modules/api-command/application/api-command-workflow.ts @@ -0,0 +1,43 @@ +import type { WorkflowStep } from "cloudflare:workers"; + +import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; +import { currentTimestampMs } from "../../../time"; +import type { ApiCommandMessage } from "./api-command-message"; +import { processApiCommandDelivery } from "./api-command-processor"; +import type { ProcessApiCommandDeliveryOptions } from "./api-command-processor"; + +export const APP_DEPLOYMENT_WORKFLOW_EXECUTION_BUDGET_MS = 29 * 60 * 1_000; + +export class ApiCommandWorkflowRetryError extends Error { + constructor(delaySeconds: number) { + super(`API command delivery requested retry after ${delaySeconds} seconds.`); + this.name = "ApiCommandWorkflowRetryError"; + } +} + +export async function runAppDeploymentWorkflow( + bindings: ApiBindings, + message: ApiCommandMessage, + step: WorkflowStep, + nowMs: () => number = currentTimestampMs, + options: ProcessApiCommandDeliveryOptions = {}, +): Promise { + await step.do( + "dispatch", + { + retries: { backoff: "constant", delay: "30 seconds", limit: 1_000 }, + timeout: "30 minutes", + }, + async () => { + const deadlineMs = nowMs() + APP_DEPLOYMENT_WORKFLOW_EXECUTION_BUDGET_MS; + const disposition = await processApiCommandDelivery(bindings, message, nowMs, { + ...options, + appDeploymentDeadlineMs: deadlineMs, + }); + if (disposition.kind === "retry") { + throw new ApiCommandWorkflowRetryError(disposition.delaySeconds); + } + return null; + }, + ); +} diff --git a/apps/api/src/modules/apps/application/app-deployment-cloudflare-client.ts b/apps/api/src/modules/apps/application/app-deployment-cloudflare-client.ts index 87d6670b..d0392fb0 100644 --- a/apps/api/src/modules/apps/application/app-deployment-cloudflare-client.ts +++ b/apps/api/src/modules/apps/application/app-deployment-cloudflare-client.ts @@ -1,28 +1,15 @@ import Cloudflare from "cloudflare"; -import { createErrorLogContext, logError } from "../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; -type CloudflareWorkerVersion = Awaited< - ReturnType ->; - -export interface CloudflarePagesProjectInput { - branch: string; - projectName: string; -} - -export interface CloudflarePagesDomainInput { - hostname: string; - projectName: string; -} - export interface CloudflareWorkerModuleInput { compatibilityDate: string; mainModuleName: string; scriptContent: string; scriptName: string; - /** Plain-text env vars injected into the Worker (e.g. agent thread URLs). */ + tags: string[]; + timeoutMs?: number; + /** Bearer capability URLs injected as Worker secrets. */ vars: Record; } @@ -31,376 +18,378 @@ export interface CloudflareWorkerDeploymentResult { versionId: string | null; } -export interface CloudflarePagesDomainResult { - status: string | null; +export interface CloudflareStaticAssetManifestEntry { + hash: string; + size: number; } -export type CloudflareDeploymentResourceTargetKind = - | "cloudflare_pages" - | "cloudflare_pages_domain" - | "cloudflare_worker" - | "cloudflare_worker_domain" - | "cloudflare_worker_route"; - -export interface CloudflareDeploymentResourceDeleteFailure { - error: unknown; - resourceName: string; - targetKind: CloudflareDeploymentResourceTargetKind; +export type CloudflareStaticAssetManifest = Record; + +export interface CloudflareStaticAssetsUploadSession { + buckets: string[][]; + uploadToken: string; } export interface CloudflareDeploymentClient { - deletePagesDomain(input: CloudflarePagesDomainInput): Promise; - deletePagesProject(input: { projectName: string }): Promise; - deleteWorkerDomain(input: { hostname: string }): Promise; - deleteWorkerRoute(input: { hostname: string }): Promise; - deleteWorkerScript(input: { scriptName: string }): Promise; + createStaticAssetsUploadSession(input: { + manifest: CloudflareStaticAssetManifest; + scriptName: string; + timeoutMs?: number; + }): Promise; + deleteWorkerScript(input: { scriptName: string; timeoutMs?: number }): Promise; + deployStaticAssets(input: { + compatibilityDate: string; + completionToken: string; + headers: string | null; + redirects: string | null; + scriptName: string; + tags: string[]; + timeoutMs?: number; + }): Promise; deployWorkerModule(input: CloudflareWorkerModuleInput): Promise; - ensurePagesDomain(input: CloudflarePagesDomainInput): Promise; - ensurePagesProject(input: CloudflarePagesProjectInput): Promise<{ projectId: string | null }>; - ensureWorkerDomain(input: { hostname: string; scriptName: string }): Promise; - ensureWorkerRoute(input: { hostname: string; scriptName: string }): Promise; - getLatestPagesDeployment(input: { - projectName: string; - }): Promise<{ deploymentId: string | null; url: string | null }>; } export type CloudflareClientBindings = Pick< ApiBindings, - "CLOUDFLARE_ACCOUNT_ID" | "CLOUDFLARE_API_TOKEN" | "CLOUDFLARE_ZONE_ID" + "CLOUDFLARE_ACCOUNT_ID" | "CLOUDFLARE_API_TOKEN" >; -function toStatus(error: unknown, status: number): boolean { - return ( - typeof error === "object" && - error !== null && - "status" in error && - Reflect.get(error, "status") === status - ); +interface WorkerModuleUpload { + files: File[]; + metadata: { + bindings: Array<{ name: string; text: string; type: "secret_text" }>; + compatibility_date: string; + main_module: string; + tags: string[]; + }; } -export function logCloudflareDeploymentResourceDeleteFailures( - eventName: string, - failures: readonly CloudflareDeploymentResourceDeleteFailure[], -): void { - for (const failure of failures) { - logError(eventName, { - ...createErrorLogContext(failure.error), - resourceName: failure.resourceName, - targetKind: failure.targetKind, - }); - } +interface CloudflareScriptUploadResponse { + etag?: string; + id?: string; + startup_time_ms: number; + tag?: string; } -export function createCloudflareDeploymentClient( - bindings: CloudflareClientBindings, -): CloudflareDeploymentClient { - const client = new Cloudflare({ apiToken: bindings.CLOUDFLARE_API_TOKEN }); - const accountId = bindings.CLOUDFLARE_ACCOUNT_ID; - const zoneId = bindings.CLOUDFLARE_ZONE_ID; +interface CloudflareApiEnvelope { + result?: T; + success?: boolean; +} - return { - async deletePagesDomain(input) { - try { - await client.pages.projects.domains.delete(input.hostname, { - account_id: accountId, - project_name: input.projectName, - }); - } catch (error) { - if (!toStatus(error, 404)) { - throw error; - } - } - }, - async deletePagesProject(input) { - try { - await client.pages.projects.delete(input.projectName, { account_id: accountId }); - } catch (error) { - if (!toStatus(error, 404)) { - throw error; - } - } - }, - async deleteWorkerDomain(input) { - const domain = await findWorkerDomain(client, accountId, input.hostname); +const CLOUDFLARE_API_BASE_URL = "https://api.cloudflare.com/client/v4"; +const CLOUDFLARE_APP_DEPLOYMENT_REQUEST_TIMEOUT_MS = 10 * 60 * 1000; +const STATIC_ASSET_HASH_PATTERN = /^[a-f0-9]{32}$/u; - if (domain?.id === undefined) { - return; - } +function requiredStringBinding(bindings: CloudflareClientBindings, name: string): string { + const value = Reflect.get(bindings, name); - await client.workers.domains.delete(domain.id, { account_id: accountId }); - }, - async deleteWorkerRoute(input) { - const pattern = workerRoutePattern(input.hostname); - const route = await findWorkerRoute(client, zoneId, pattern); + if (typeof value !== "string" || value.trim().length === 0) { + throw new Error(`${name} is required for App deployment.`); + } - if (route?.id === undefined) { - return; - } + return value; +} - await client.workers.routes.delete(route.id, { zone_id: zoneId }); - }, - async deleteWorkerScript(input) { - try { - await client.workers.scripts.delete(input.scriptName, { account_id: accountId }); - } catch (error) { - if (!toStatus(error, 404)) { - throw error; - } - } - }, - async deployWorkerModule(input) { - const scriptPath = `/accounts/${accountId}/workers/scripts/${encodeURIComponent(input.scriptName)}`; - const createVersion = async (): Promise => - ( - await client.post<{ result: CloudflareWorkerVersion }>(`${scriptPath}/versions`, { - body: createWorkerModuleUpload(input), - }) - ).result; - let version; +function requestTimeout(timeoutMs: number | undefined): number { + if (timeoutMs === undefined) { + return CLOUDFLARE_APP_DEPLOYMENT_REQUEST_TIMEOUT_MS; + } - try { - version = await createVersion(); - } catch (error) { - if (!toCloudflareCode(error, 10007)) { - throw error; - } + if (!Number.isSafeInteger(timeoutMs) || timeoutMs <= 0) { + throw new Error("Cloudflare App deployment request timeout must be a positive integer."); + } - await client.put(scriptPath, { - body: createWorkerModuleUpload(input), - }); - version = await createVersion(); - } - const versionId = version.id ?? null; + return Math.min(timeoutMs, CLOUDFLARE_APP_DEPLOYMENT_REQUEST_TIMEOUT_MS); +} - if (versionId === null) { - throw new Error("Cloudflare Worker version response did not include an id."); - } +function remainingRequestTimeout(deadlineMs: number): number { + const remaining = Math.floor(deadlineMs - Date.now()); - const deployment = await client.workers.scripts.deployments.create(input.scriptName, { - account_id: accountId, - strategy: "percentage", - versions: [{ percentage: 100, version_id: versionId }], - }); + if (remaining <= 0) { + throw new Error("Cloudflare App deployment request exceeded its execution budget."); + } - return { - deploymentId: deployment.id ?? null, - versionId, - }; - }, - async ensurePagesDomain(input) { - try { - const domain = await client.pages.projects.domains.create(input.projectName, { - account_id: accountId, - name: input.hostname, - }); + return remaining; +} - return { status: domain.status ?? null }; - } catch (error) { - if (!toStatus(error, 409)) { - throw error; - } +function toStatus(error: unknown, status: number): boolean { + return ( + typeof error === "object" && + error !== null && + "status" in error && + Reflect.get(error, "status") === status + ); +} - const domain = await client.pages.projects.domains.get(input.hostname, { - account_id: accountId, - project_name: input.projectName, - }); +export function validateStaticAssetsManifest(manifest: CloudflareStaticAssetManifest): Set { + const hashes = new Set(); - return { status: domain.status ?? null }; - } - }, - async ensurePagesProject(input) { - try { - const project = await client.pages.projects.create({ - account_id: accountId, - name: input.projectName, - production_branch: input.branch, - }); + for (const [path, entry] of Object.entries(manifest)) { + const parts = path.slice(1).split("/"); - return { projectId: project.id ?? null }; - } catch (error) { - if (!toStatus(error, 409)) { - throw error; - } + if ( + !path.startsWith("/") || + path.includes("\0") || + path.includes("\\") || + parts.some((part) => part.length === 0 || part === "." || part === "..") + ) { + throw new Error("Cloudflare Static Assets manifest contains an invalid path."); + } - const project = await client.pages.projects.get(input.projectName, { - account_id: accountId, - }); + if ( + !STATIC_ASSET_HASH_PATTERN.test(entry.hash) || + !Number.isSafeInteger(entry.size) || + entry.size < 0 + ) { + throw new Error(`Cloudflare Static Assets manifest entry is invalid for ${path}.`); + } - return { projectId: project.id ?? null }; - } - }, - async ensureWorkerDomain(input) { - const existingDomain = await findWorkerDomain(client, accountId, input.hostname); + hashes.add(entry.hash); + } - if (existingDomain !== null && existingDomain.service === input.scriptName) { - return; - } + return hashes; +} - await client.workers.domains.update({ - account_id: accountId, - hostname: input.hostname, - service: input.scriptName, - zone_id: zoneId, - }); - }, - async ensureWorkerRoute(input) { - const pattern = workerRoutePattern(input.hostname); - const existingRoute = await findWorkerRoute(client, zoneId, pattern); - - if (existingRoute !== null) { - if (existingRoute.script !== input.scriptName && existingRoute.id !== undefined) { - await client.workers.routes.update(existingRoute.id, { - pattern, - script: input.scriptName, - zone_id: zoneId, - }); - } +function normalizeAssetUploadSession( + response: { buckets?: string[][]; jwt?: string }, + manifestHashes: ReadonlySet, +): CloudflareStaticAssetsUploadSession { + if (typeof response.jwt !== "string" || response.jwt.length === 0) { + throw new Error("Cloudflare did not return a Static Assets upload token."); + } - return; - } + const buckets = response.buckets ?? []; + const requested = new Set(); - await client.workers.routes.create({ - pattern, - script: input.scriptName, - zone_id: zoneId, - }); - }, - async getLatestPagesDeployment(input) { - const deployments = client.pages.projects.deployments.list(input.projectName, { - account_id: accountId, - per_page: 1, - }); + for (const bucket of buckets) { + if (!Array.isArray(bucket) || bucket.length === 0) { + throw new Error("Cloudflare returned an invalid Static Assets upload bucket."); + } - for await (const deployment of deployments) { - return { deploymentId: deployment.id ?? null, url: deployment.url ?? null }; + for (const hash of bucket) { + if (typeof hash !== "string" || !manifestHashes.has(hash) || requested.has(hash)) { + throw new Error("Cloudflare requested an invalid Static Assets upload hash."); } - return { deploymentId: null, url: null }; - }, - }; -} - -export async function deleteCloudflareDeploymentResources( - cloudflareClient: CloudflareDeploymentClient, - input: { hostname: string; resourceName: string }, -): Promise { - const failures = await Promise.all([ - deleteCloudflareDeploymentResource("cloudflare_pages_domain", input.resourceName, () => - cloudflareClient.deletePagesDomain({ - hostname: input.hostname, - projectName: input.resourceName, - }), - ), - deleteCloudflareDeploymentResource("cloudflare_pages", input.resourceName, () => - cloudflareClient.deletePagesProject({ projectName: input.resourceName }), - ), - deleteCloudflareDeploymentResource("cloudflare_worker_domain", input.hostname, () => - cloudflareClient.deleteWorkerDomain({ hostname: input.hostname }), - ), - deleteCloudflareDeploymentResource("cloudflare_worker_route", input.resourceName, () => - cloudflareClient.deleteWorkerRoute({ hostname: input.hostname }), - ), - deleteCloudflareDeploymentResource("cloudflare_worker", input.resourceName, () => - cloudflareClient.deleteWorkerScript({ scriptName: input.resourceName }), - ), - ]); - - return failures.filter( - (failure): failure is CloudflareDeploymentResourceDeleteFailure => failure !== null, - ); -} - -async function deleteCloudflareDeploymentResource( - targetKind: CloudflareDeploymentResourceTargetKind, - resourceName: string, - runDelete: () => Promise, -): Promise { - try { - await runDelete(); - return null; - } catch (error) { - return { error, resourceName, targetKind }; + requested.add(hash); + } } -} -function workerRoutePattern(hostname: string): string { - return `${hostname}/*`; + return { buckets, uploadToken: response.jwt }; } -export function createWorkerModuleUpload(input: CloudflareWorkerModuleInput): FormData { - const upload = new FormData(); - const file = new File([input.scriptContent], input.mainModuleName, { - type: "application/javascript+module", - }); - const metadata = { - bindings: Object.entries(input.vars).map(([name, text]) => ({ - name, - text, - type: "plain_text" as const, - })), - compatibility_date: input.compatibilityDate, - main_module: input.mainModuleName, - }; - - upload.append("metadata", JSON.stringify(metadata)); - upload.append(input.mainModuleName, file); - - return upload; +function staticAssetsScriptUrl(input: { + accountId: string; + dispatchNamespace: string; + scriptName: string; +}): string { + return `${CLOUDFLARE_API_BASE_URL}/accounts/${encodeURIComponent( + input.accountId, + )}/workers/dispatch/namespaces/${encodeURIComponent( + input.dispatchNamespace, + )}/scripts/${encodeURIComponent(input.scriptName)}`; } -function toCloudflareCode(error: unknown, code: number): boolean { - if (typeof error !== "object" || error === null) { - return false; +async function readCloudflareApiResult(response: Response, operation: string): Promise { + let envelope: CloudflareApiEnvelope; + + try { + envelope = (await response.json()) as CloudflareApiEnvelope; + } catch { + throw new Error(`Cloudflare ${operation} returned a non-JSON response (${response.status}).`); } - if ("code" in error && Reflect.get(error, "code") === code) { - return true; + if (!response.ok || envelope.success !== true || envelope.result === undefined) { + throw new Error(`Cloudflare ${operation} failed (${response.status}).`); } - const cause = Reflect.get(error, "error"); + return envelope.result; +} - if (typeof cause === "object" && cause !== null && Reflect.get(cause, "code") === code) { - return true; +export function createStaticAssetsUploadForm(input: { + compatibilityDate: string; + completionToken: string; + headers: string | null; + redirects: string | null; + tags: string[]; +}): FormData { + const config: { _headers?: string; _redirects?: string } = {}; + + if (input.headers !== null) { + config["_headers"] = input.headers; } - const errors = Reflect.get(error, "errors"); + if (input.redirects !== null) { + config["_redirects"] = input.redirects; + } - return ( - Array.isArray(errors) && - errors.some( - (entry) => typeof entry === "object" && entry !== null && Reflect.get(entry, "code") === code, - ) + const form = new FormData(); + form.set( + "metadata", + JSON.stringify({ + assets: { + config, + jwt: input.completionToken, + }, + compatibility_date: input.compatibilityDate, + tags: input.tags, + }), ); + return form; } -async function findWorkerDomain( - client: Cloudflare, - accountId: string, - hostname: string, -): Promise<{ id?: string; service?: string } | null> { - const domains = client.workers.domains.list({ account_id: accountId }); +export function createCloudflareDeploymentClient( + bindings: CloudflareClientBindings, +): CloudflareDeploymentClient { + const apiToken = requiredStringBinding(bindings, "CLOUDFLARE_API_TOKEN"); + const client = new Cloudflare({ + apiToken, + timeout: CLOUDFLARE_APP_DEPLOYMENT_REQUEST_TIMEOUT_MS, + }); + const accountId = requiredStringBinding(bindings, "CLOUDFLARE_ACCOUNT_ID"); + const dispatchNamespace = requiredStringBinding(bindings, "MOSOO_APP_DISPATCH_NAMESPACE"); + + async function verifyWorkerScriptTags( + scriptName: string, + expectedTags: readonly string[], + deadlineMs: number, + ): Promise { + const actualTags: string[] = []; + + for await (const tag of client.workersForPlatforms.dispatch.namespaces.scripts.tags.list( + scriptName, + { + account_id: accountId, + dispatch_namespace: dispatchNamespace, + }, + { timeout: remainingRequestTimeout(deadlineMs) }, + )) { + actualTags.push(tag); + } - for await (const domain of domains) { - if (domain.hostname === hostname) { - return domain; + if (!sameStringSet(actualTags, expectedTags)) { + throw new Error(`Cloudflare did not retain the expected tags for ${scriptName}.`); } } - return null; + return { + async createStaticAssetsUploadSession(input) { + const manifestHashes = validateStaticAssetsManifest(input.manifest); + const response = + await client.workersForPlatforms.dispatch.namespaces.scripts.assetUpload.create( + input.scriptName, + { + account_id: accountId, + dispatch_namespace: dispatchNamespace, + manifest: input.manifest, + }, + { timeout: requestTimeout(input.timeoutMs) }, + ); + + return normalizeAssetUploadSession(response, manifestHashes); + }, + async deleteWorkerScript(input) { + try { + await client.workersForPlatforms.dispatch.namespaces.scripts.delete( + input.scriptName, + { + account_id: accountId, + dispatch_namespace: dispatchNamespace, + }, + { timeout: requestTimeout(input.timeoutMs) }, + ); + } catch (error) { + if (!toStatus(error, 404)) { + throw error; + } + } + }, + async deployStaticAssets(input) { + const timeoutMs = requestTimeout(input.timeoutMs); + const deadlineMs = Date.now() + timeoutMs; + const response = await fetch( + staticAssetsScriptUrl({ accountId, dispatchNamespace, scriptName: input.scriptName }), + { + body: createStaticAssetsUploadForm(input), + headers: { Authorization: `Bearer ${apiToken}` }, + method: "PUT", + signal: AbortSignal.timeout(remainingRequestTimeout(deadlineMs)), + }, + ); + const worker = await readCloudflareApiResult( + response, + "Static Assets deployment", + ); + + if (worker.id !== undefined && worker.id !== input.scriptName) { + throw new Error( + `Cloudflare uploaded App deployment as ${worker.id}, expected ${input.scriptName}.`, + ); + } + + await verifyWorkerScriptTags(input.scriptName, input.tags, deadlineMs); + + return { + deploymentId: worker.tag ?? null, + versionId: worker.etag ?? null, + }; + }, + async deployWorkerModule(input) { + const timeoutMs = requestTimeout(input.timeoutMs); + const deadlineMs = Date.now() + timeoutMs; + const upload = createWorkerModuleUpload(input); + const worker = await client.workersForPlatforms.dispatch.namespaces.scripts.update( + input.scriptName, + { + account_id: accountId, + dispatch_namespace: dispatchNamespace, + ...upload, + }, + { timeout: remainingRequestTimeout(deadlineMs) }, + ); + + if (worker.id !== undefined && worker.id !== input.scriptName) { + throw new Error( + `Cloudflare uploaded App deployment as ${worker.id}, expected ${input.scriptName}.`, + ); + } + + await verifyWorkerScriptTags(input.scriptName, input.tags, deadlineMs); + + return { + deploymentId: worker.tag ?? null, + versionId: worker.etag ?? null, + }; + }, + }; } -async function findWorkerRoute( - client: Cloudflare, - zoneId: string, - pattern: string, -): Promise<{ id?: string; script?: string } | null> { - const routes = client.workers.routes.list({ zone_id: zoneId }); +function sameStringSet(left: readonly string[], right: readonly string[]): boolean { + const sortedLeft = left.toSorted(); + const sortedRight = right.toSorted(); - for await (const route of routes) { - if (route.pattern === pattern) { - return route; - } - } + return ( + sortedLeft.length === sortedRight.length && + sortedLeft.every((value, index) => value === sortedRight[index]) + ); +} - return null; +export function createWorkerModuleUpload(input: CloudflareWorkerModuleInput): WorkerModuleUpload { + return { + files: [ + new File([input.scriptContent], input.mainModuleName, { + type: "application/javascript+module", + }), + ], + metadata: { + bindings: Object.entries(input.vars).map(([name, text]) => ({ + name, + text, + type: "secret_text", + })), + compatibility_date: input.compatibilityDate, + main_module: input.mainModuleName, + tags: input.tags, + }, + }; } diff --git a/apps/api/src/modules/apps/application/app-deployment-detector.ts b/apps/api/src/modules/apps/application/app-deployment-detector.ts index 80de1a29..bbfae33b 100644 --- a/apps/api/src/modules/apps/application/app-deployment-detector.ts +++ b/apps/api/src/modules/apps/application/app-deployment-detector.ts @@ -1,10 +1,10 @@ -import type { AppDeploymentTargetKind } from "@mosoo/db"; import type { ParseError } from "jsonc-parser"; import { parse as parseJsonc } from "jsonc-parser"; import { parse as parseToml, stringify } from "smol-toml"; export type AppDeploymentPackageManager = "bun" | "none" | "npm" | "pnpm" | "yarn"; -export type AppDeploymentTargetMode = "static_assets" | "worker_module" | "worker_with_assets"; +export type AppDeploymentTargetKind = "cloudflare_static_assets" | "cloudflare_worker"; +export type AppDeploymentTargetMode = "static_assets" | "worker_module"; export type AppDeploymentDetectionErrorCode = | "deployment_config_required" | "deployment_shape_unsupported"; @@ -35,10 +35,6 @@ export interface AppDeploymentRepositorySnapshot { files: Readonly>; } -export interface AppDeploymentDetectionOptions { - resourceName: string; -} - interface PackageJson { dependencies: Readonly>; devDependencies: Readonly>; @@ -62,6 +58,7 @@ interface MosooConfig { interface RepositoryFiles { has(path: string): boolean; + paths(): readonly string[]; read(path: string): string | null; } @@ -80,24 +77,18 @@ export class AppDeploymentDetectionError extends Error { export function detectAppDeploymentPlan( snapshot: AppDeploymentRepositorySnapshot, - options: AppDeploymentDetectionOptions, ): AppDeploymentPlan { const files = createRepositoryFiles(snapshot.files); const mosooConfig = files.read(".mosoo.toml"); - const resourceName = normalizeResourceName(options.resourceName); if (mosooConfig !== null) { - return detectFromMosooConfig(files, mosooConfig, resourceName); + return detectFromMosooConfig(files, mosooConfig); } - return detectFromRepository(files, ".", resourceName); + return detectFromRepository(files, "."); } -function detectFromMosooConfig( - files: RepositoryFiles, - source: string, - resourceName: string, -): AppDeploymentPlan { +function detectFromMosooConfig(files: RepositoryFiles, source: string): AppDeploymentPlan { const config = parseMosooConfig(source); const packageJson = readPackageJson(files, config.rootDir); const packageManager = detectPackageManager(files, config.rootDir, packageJson); @@ -117,14 +108,15 @@ function detectFromMosooConfig( config.outputDir ?? fail("deployment_config_required", "static deployment requires build.output"); - return pagesPlan({ + assertStaticAssetsCompatible(files, config.rootDir); + + return staticAssetsPlan({ agentBindings: config.agents, buildCommand, installCommand, mosooConfigPath: ".mosoo.toml", outputDir, packageManager, - resourceName, routesFallback: config.routesFallback, rootDir: config.rootDir, }); @@ -148,17 +140,12 @@ function detectFromMosooConfig( installCommand, mosooConfigPath: ".mosoo.toml", packageManager, - resourceName, rootDir: config.rootDir, workerEntry, }); } -function detectFromRepository( - files: RepositoryFiles, - rootDir: string, - resourceName: string, -): AppDeploymentPlan { +function detectFromRepository(files: RepositoryFiles, rootDir: string): AppDeploymentPlan { const packageJson = readPackageJson(files, rootDir); const packageManager = detectPackageManager(files, rootDir, packageJson); const wranglerMain = readWranglerMain(files, rootDir); @@ -170,7 +157,6 @@ function detectFromRepository( installCommand: installCommandFor(packageManager, files, rootDir), mosooConfigPath: null, packageManager, - resourceName, rootDir, workerEntry: wranglerMain, }); @@ -178,14 +164,15 @@ function detectFromRepository( if (packageJson === null) { if (files.has("index.html")) { - return pagesPlan({ + assertStaticAssetsCompatible(files, rootDir); + + return staticAssetsPlan({ agentBindings: [], buildCommand: null, installCommand: null, mosooConfigPath: null, outputDir: ".", packageManager: "none", - resourceName, routesFallback: null, rootDir, }); @@ -198,20 +185,20 @@ function detectFromRepository( } if (hasDependency(packageJson, "vite")) { - return packagePagesPlan(files, rootDir, packageJson, packageManager, "dist", resourceName); + return packageStaticAssetsPlan(files, rootDir, packageJson, packageManager, "dist"); } if (hasDependency(packageJson, "astro")) { - return packagePagesPlan(files, rootDir, packageJson, packageManager, "dist", resourceName); + return packageStaticAssetsPlan(files, rootDir, packageJson, packageManager, "dist"); } if (hasDependency(packageJson, "@docusaurus/core")) { - return packagePagesPlan(files, rootDir, packageJson, packageManager, "build", resourceName); + return packageStaticAssetsPlan(files, rootDir, packageJson, packageManager, "build"); } if (hasDependency(packageJson, "next")) { if (isNextStaticExport(files, rootDir, packageJson)) { - return packagePagesPlan(files, rootDir, packageJson, packageManager, "out", resourceName); + return packageStaticAssetsPlan(files, rootDir, packageJson, packageManager, "out"); } throw new AppDeploymentDetectionError( @@ -221,14 +208,15 @@ function detectFromRepository( } if (files.has(pathInRoot(rootDir, "index.html")) && packageJson.scripts["build"] === undefined) { - return pagesPlan({ + assertStaticAssetsCompatible(files, rootDir); + + return staticAssetsPlan({ agentBindings: [], buildCommand: null, installCommand: null, mosooConfigPath: null, outputDir: ".", packageManager: "none", - resourceName, routesFallback: null, rootDir, }); @@ -240,39 +228,38 @@ function detectFromRepository( ); } -function packagePagesPlan( +function packageStaticAssetsPlan( files: RepositoryFiles, rootDir: string, packageJson: PackageJson, packageManager: AppDeploymentPackageManager, outputDir: string, - resourceName: string, ): AppDeploymentPlan { const buildCommand = buildCommandFor(packageManager, packageJson) ?? fail("deployment_config_required", "static framework deployment requires scripts.build"); - return pagesPlan({ + assertStaticAssetsCompatible(files, rootDir); + + return staticAssetsPlan({ agentBindings: [], buildCommand, installCommand: installCommandFor(packageManager, files, rootDir), mosooConfigPath: null, outputDir, packageManager, - resourceName, routesFallback: null, rootDir, }); } -function pagesPlan(input: { +function staticAssetsPlan(input: { agentBindings: AppDeploymentAgentBinding[]; buildCommand: string | null; installCommand: string | null; mosooConfigPath: ".mosoo.toml" | null; outputDir: string; packageManager: AppDeploymentPackageManager; - resourceName: string; routesFallback: string | null; rootDir: string; }): AppDeploymentPlan { @@ -281,8 +268,7 @@ function pagesPlan(input: { buildCommand: input.buildCommand, generatedWranglerConfig: stringify({ compatibility_date: APP_DEPLOYMENT_COMPATIBILITY_DATE, - name: input.resourceName, - pages_build_output_dir: input.outputDir, + assets: { directory: "artifact" }, }), installCommand: input.installCommand, mosooConfigPath: input.mosooConfigPath, @@ -290,7 +276,7 @@ function pagesPlan(input: { packageManager: input.packageManager, routesFallback: input.routesFallback, rootDir: input.rootDir, - targetKind: "cloudflare_pages", + targetKind: "cloudflare_static_assets", targetMode: "static_assets", warnings: [], workerEntry: null, @@ -303,7 +289,6 @@ function workerPlan(input: { installCommand: string | null; mosooConfigPath: ".mosoo.toml" | null; packageManager: AppDeploymentPackageManager; - resourceName: string; rootDir: string; workerEntry: string; }): AppDeploymentPlan { @@ -320,7 +305,6 @@ function workerPlan(input: { generatedWranglerConfig: stringify({ compatibility_date: APP_DEPLOYMENT_COMPATIBILITY_DATE, main: input.workerEntry, - name: input.resourceName, }), installCommand: input.installCommand, mosooConfigPath: input.mosooConfigPath, @@ -346,12 +330,38 @@ function createRepositoryFiles(files: Readonly>): Reposit has(path) { return normalized.has(normalizePath(path)); }, + paths() { + return [...normalized.keys()]; + }, read(path) { return normalized.get(normalizePath(path)) ?? null; }, }; } +function assertStaticAssetsCompatible(files: RepositoryFiles, rootDir: string): void { + const prefix = rootDir === "." ? "" : `${rootDir}/`; + const unsupportedPath = files.paths().find((path) => { + if (!path.startsWith(prefix)) { + return false; + } + + const relativePath = path.slice(prefix.length); + return ( + relativePath === "_routes.json" || + relativePath === "_worker.js" || + relativePath.startsWith("functions/") + ); + }); + + if (unsupportedPath !== undefined) { + throw new AppDeploymentDetectionError( + "deployment_shape_unsupported", + `${unsupportedPath} is not supported by isolated Workers Static Assets; migrate it to a Worker module.`, + ); + } +} + function parseMosooConfig(source: string): MosooConfig { const value = parseTomlObject(source, ".mosoo.toml"); requireAllowedKeys( @@ -888,19 +898,6 @@ function normalizeOptionalRelativePath(path: string | null, field: string): stri return normalizeRelativePath(path, field); } -function normalizeResourceName(value: string): string { - const name = value.trim(); - - if (name.length === 0) { - throw new AppDeploymentDetectionError( - "deployment_config_required", - "deployment resource name is required", - ); - } - - return name; -} - function normalizeRelativePath(path: string, field: string): string { const rawPath = path.replaceAll("\\", "/"); const parts = rawPath.split("/").filter((part) => part !== "" && part !== "."); diff --git a/apps/api/src/modules/apps/application/app-deployment-executor.service.ts b/apps/api/src/modules/apps/application/app-deployment-executor.service.ts index 87ca0bb5..43517e6b 100644 --- a/apps/api/src/modules/apps/application/app-deployment-executor.service.ts +++ b/apps/api/src/modules/apps/application/app-deployment-executor.service.ts @@ -1,19 +1,27 @@ import type { AppDeploymentRunStatus } from "@mosoo/contracts/app"; import type { AppDeploymentRunRow, AppDeploymentRow } from "@mosoo/db"; -import { appDeploymentRunsTable, appDeploymentsTable } from "@mosoo/db"; +import { apiCommandsTable, appDeploymentRunsTable, appDeploymentsTable } from "@mosoo/db"; import { parsePlatformId } from "@mosoo/id"; import type { AgentId, AppDeploymentRunId } from "@mosoo/id"; -import { and, eq, inArray, isNotNull, isNull } from "drizzle-orm"; +import { and, eq, exists, inArray, isNull, or, sql } from "drizzle-orm"; +import createIgnore from "ignore"; import { createErrorLogContext, logError } from "../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; import { getAppDatabase, getD1ChangeCount } from "../../../platform/db/drizzle"; +import { quoteShellArg } from "../../../shared/shell"; import { currentTimestampMs } from "../../../time"; import { listAppOwnerAgentRows } from "../../agents/application/agent-repository"; +import type { ApiCommandClaimAuthority } from "../../api-command/application/api-command-ledger"; +import { exactApiCommandClaimPredicate } from "../../api-command/application/api-command-ledger"; import { boundAgentUrl, mintAppAgentCapabilityToken } from "../../public-api/app-agent-capability"; import { - destroyRuntimeSubjectContainer, - getRuntimeSubjectKeepAliveHandle, + createRuntimeSandboxBucketMountOptions, + resolveRuntimeSandboxBucketMountTarget, +} from "../../runtime/infrastructure/runtime-sandbox-bucket-mount"; +import { + destroyUnversionedSandboxContainer, + getEphemeralUnversionedSandboxHandle, } from "../../runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-platform"; import type { ExecutionSessionHandle, @@ -25,17 +33,21 @@ import { resolveAppAgentBindings, } from "./app-agent-binding-resolution"; import type { ResolvableAppAgent } from "./app-agent-binding-resolution"; -import type { CloudflareDeploymentClient } from "./app-deployment-cloudflare-client"; -import { - createCloudflareDeploymentClient, - deleteCloudflareDeploymentResources, - logCloudflareDeploymentResourceDeleteFailures, +import type { + CloudflareDeploymentClient, + CloudflareStaticAssetManifest, } from "./app-deployment-cloudflare-client"; +import { createCloudflareDeploymentClient } from "./app-deployment-cloudflare-client"; import { APP_DEPLOYMENT_COMPATIBILITY_DATE, + AppDeploymentDetectionError, detectAppDeploymentPlan, } from "./app-deployment-detector"; import type { AppDeploymentPlan, AppDeploymentRepositorySnapshot } from "./app-deployment-detector"; +import { + markAppDeploymentScriptUploadStarted, + registerAppDeploymentScriptCandidate, +} from "./app-deployment-script-reconciliation.service"; interface AppDeploymentDispatchContext { deployment: AppDeploymentRow; @@ -47,13 +59,34 @@ interface PreparedAppDeploymentRepository { snapshot: AppDeploymentRepositorySnapshot; } -interface AppDeploymentDeployResult { +export interface AppDeploymentDeployResult { + activeScriptName: string; externalDeploymentId: string | null; externalProjectId: string | null; externalVersionId: string | null; url: string; } +type AppDeploymentUploadedResult = Omit; + +export interface AppDeploymentAttemptAuthority extends ApiCommandClaimAuthority { + requireOwnership: () => Promise; +} + +export type AppDeploymentDispatchOutcome = + | { kind: "skipped" } + | ({ + deploymentId: AppDeploymentRow["id"]; + kind: "succeeded"; + runId: AppDeploymentRunId; + } & AppDeploymentDeployResult) + | { + errorCode: string; + errorMessage: string; + kind: "terminal_failure"; + runId: AppDeploymentRunId; + }; + export interface AppDeploymentBuildRunner { build(input: { plan: AppDeploymentPlan; @@ -61,12 +94,13 @@ export interface AppDeploymentBuildRunner { }): Promise; cleanup?(): Promise; deploy(input: { + activeScriptName: string; deployment: AppDeploymentRow; envVars: Record; plan: AppDeploymentPlan; prepared: PreparedAppDeploymentRepository; run: AppDeploymentRunRow; - }): Promise; + }): Promise; prepare(input: { deployment: AppDeploymentRow; run: AppDeploymentRunRow; @@ -75,6 +109,7 @@ export interface AppDeploymentBuildRunner { export interface DispatchAppDeploymentRunOptions { cloudflareClient?: CloudflareDeploymentClient; + deadlineMs?: number; runner?: AppDeploymentBuildRunner; } @@ -103,16 +138,156 @@ const SNAPSHOT_FILE_NAMES = new Set([ "yarn.lock", ]); const WORKER_JS_ENTRY_PATTERN = /\.(?:mjs|js)$/u; -export function appDeploymentBuildSandboxId(runId: AppDeploymentRunId): string { - return `${runId}-build`; +const APP_DEPLOYMENT_SANDBOX_SLEEP_AFTER_SECONDS = 15 * 60; +const APP_DEPLOYMENT_SANDBOX_COMMAND_TIMEOUT_MS = 10 * 60 * 1000; +const APP_DEPLOYMENT_EXECUTION_TIMEOUT_MS = 29 * 60 * 1000; +const APP_DEPLOYMENT_RPC_TIMEOUT_MS = 10 * 60 * 1000; +const APP_DEPLOYMENT_MANAGED_TAG = "mosoo-managed"; +const APP_DEPLOYMENT_MAX_SCRIPT_TAGS = 8; +const APP_DEPLOYMENT_ARTIFACT_MOUNT_PATH = "/mnt/mosoo-app-deployment-artifact"; +const APP_DEPLOYMENT_ARTIFACT_FILE_NAME = "artifact.tar"; +const APP_DEPLOYMENT_STATIC_ASSET_TOOL_FILE_NAME = "static-assets.mjs"; +const APP_DEPLOYMENT_STATIC_ASSET_INVENTORY_FILE_NAME = "inventory.json"; +const APP_DEPLOYMENT_STATIC_ASSET_SELECTION_FILE_NAME = "selection.json"; +const APP_DEPLOYMENT_STATIC_ASSET_MANIFEST_FILE_NAME = "manifest.json"; +const APP_DEPLOYMENT_STATIC_ASSET_BUCKETS_FILE_NAME = "buckets.json"; +const APP_DEPLOYMENT_STATIC_ASSET_COMPLETION_FILE_NAME = "completion-token"; +export type AppDeploymentAttemptIdentity = Pick< + AppDeploymentAttemptAuthority, + "attemptCount" | "deliveryGeneration" +>; +type AppDeploymentCandidateIdentity = Pick< + AppDeploymentAttemptAuthority, + "attemptCount" | "commandId" | "deliveryGeneration" +>; + +interface AppDeploymentDispatcher { + get(scriptName: string): { fetch(request: Request): Promise }; +} + +class AppDeploymentExecutionBudget { + readonly #deadlineMs: number; + + constructor(deadlineMs: number) { + if (!Number.isFinite(deadlineMs) || deadlineMs <= Date.now()) { + throw new Error("App deployment execution deadline has already expired."); + } + + this.#deadlineMs = deadlineMs; + } + + remaining(maximumMs = APP_DEPLOYMENT_RPC_TIMEOUT_MS): number { + const remainingMs = Math.floor(this.#deadlineMs - Date.now()); + + if (remainingMs <= 0) { + throw new Error("App deployment exceeded its execution deadline."); + } + + return Math.max(1, Math.min(remainingMs, maximumMs)); + } + + async run( + label: string, + effect: (timeoutMs: number) => Promise, + maximumMs = APP_DEPLOYMENT_RPC_TIMEOUT_MS, + ): Promise { + const timeoutMs = this.remaining(maximumMs); + let timeout: ReturnType | undefined; + + try { + return await Promise.race([ + effect(timeoutMs), + new Promise((_, reject) => { + timeout = setTimeout( + () => reject(new Error(`${label} exceeded the App deployment execution budget.`)), + timeoutMs, + ); + }), + ]); + } finally { + if (timeout !== undefined) { + clearTimeout(timeout); + } + } + } +} + +function appDeploymentAttemptSuffix(attempt: AppDeploymentAttemptIdentity): string { + return `g${attempt.deliveryGeneration.toString(36)}-a${attempt.attemptCount.toString(36)}`; +} + +export function appDeploymentCandidateScriptName(attempt: AppDeploymentCandidateIdentity): string { + const scriptName = `app-${attempt.commandId.toLowerCase()}-${appDeploymentAttemptSuffix(attempt)}`; + + if (!/^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$/u.test(scriptName)) { + throw new Error("App deployment attempt cannot be represented as a Cloudflare script name."); + } + + return scriptName; +} + +function appDeploymentCandidateScriptTags(input: { + deploymentId: AppDeploymentRow["id"]; + runId: AppDeploymentRunId; +}): string[] { + const tags = [ + APP_DEPLOYMENT_MANAGED_TAG, + `d-${input.deploymentId.toLowerCase()}`, + `r-${input.runId.toLowerCase()}`, + ]; + + if ( + tags.length > APP_DEPLOYMENT_MAX_SCRIPT_TAGS || + tags.some((tag) => !/^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$/u.test(tag)) + ) { + throw new Error("App deployment identifiers cannot be represented as Cloudflare tags."); + } + + return tags; +} + +function appDeploymentArtifactPrefix(scriptName: string): string { + return `app-deployments/${scriptName}`; +} + +function mountAppDeploymentArtifactBucket( + bindings: ApiBindings, + sandbox: SandboxHandle, + scriptName: string, + readOnly: boolean, +): Promise { + return sandbox.mountBucket( + resolveRuntimeSandboxBucketMountTarget(bindings), + APP_DEPLOYMENT_ARTIFACT_MOUNT_PATH, + createRuntimeSandboxBucketMountOptions(bindings, { + prefix: appDeploymentArtifactPrefix(scriptName), + readOnly, + }), + ); } -export function appDeploymentDeploySandboxId(runId: AppDeploymentRunId): string { - return `${runId}-deploy`; +export function appDeploymentBuildSandboxId( + runId: AppDeploymentRunId, + attempt: AppDeploymentAttemptIdentity, +): string { + return `${runId}-${appDeploymentAttemptSuffix(attempt)}-build`; } -function quoteShellArg(value: string): string { - return `'${value.replaceAll("'", `'"'"'`)}'`; +export function appDeploymentDeploySandboxId( + runId: AppDeploymentRunId, + attempt: AppDeploymentAttemptIdentity, +): string { + return `${runId}-${appDeploymentAttemptSuffix(attempt)}-deploy`; +} + +async function runOwnedEffect( + requireOwnership: () => Promise, + effect: () => Promise, +): Promise { + await requireOwnership(); + const result = await effect(); + await requireOwnership(); + return result; } function deploymentHostname(deployment: AppDeploymentRow, domain: string): string { @@ -123,6 +298,30 @@ function deploymentUrl(deployment: AppDeploymentRow, domain: string): string { return `https://${deploymentHostname(deployment, domain)}`; } +function requiredBinding(bindings: ApiBindings, name: string): unknown { + const value = Reflect.get(bindings, name); + + if (value === undefined || value === null || value === "") { + throw new Error(`${name} is required for App deployment.`); + } + + return value; +} + +function appDeploymentDispatcher(bindings: ApiBindings): AppDeploymentDispatcher { + const dispatcher = requiredBinding(bindings, "APP_DEPLOYMENT_DISPATCHER"); + + if ( + typeof dispatcher !== "object" || + dispatcher === null || + typeof Reflect.get(dispatcher, "get") !== "function" + ) { + throw new Error("APP_DEPLOYMENT_DISPATCHER is invalid."); + } + + return dispatcher as AppDeploymentDispatcher; +} + function isActiveRunStatus(status: AppDeploymentRunStatus): boolean { return (ACTIVE_APP_DEPLOYMENT_RUN_STATUSES as readonly AppDeploymentRunStatus[]).includes(status); } @@ -153,10 +352,15 @@ async function execChecked( session: ExecutionSessionHandle, command: string, label: string, + budget: AppDeploymentExecutionBudget, options: { retryable?: boolean } = {}, ): Promise { const message = commandFailureMessage( - await session.exec(`sh -lc ${quoteShellArg(command)}`), + await budget.run(label, (timeoutMs) => + session.exec(`sh -lc ${quoteShellArg(command)}`, { + timeout: Math.min(timeoutMs, APP_DEPLOYMENT_SANDBOX_COMMAND_TIMEOUT_MS), + }), + ), label, ); @@ -171,14 +375,6 @@ async function execChecked( throw new Error(message); } -function assertSelfContainedWorkerModule(scriptContent: string): void { - if (/^\s*import\s/mu.test(scriptContent) || /\bimport\s*\(/u.test(scriptContent)) { - throw new AppDeploymentNonRetryableError( - "Worker deployment only supports self-contained JavaScript modules in the first cut.", - ); - } -} - function assertRequestedMosooConfigPresent( run: AppDeploymentRunRow, snapshot: AppDeploymentRepositorySnapshot, @@ -208,28 +404,25 @@ function assertRequestedMosooConfigPresent( } } -async function assertControlledWranglerAvailable(sandbox: SandboxHandle): Promise { - await execChecked( - sandbox, - "command -v wrangler >/dev/null && wrangler --version >/dev/null", - "Controlled Wrangler availability", - { retryable: false }, - ); -} - export async function destroyAppDeploymentRunSandboxesBestEffort( bindings: ApiBindings, runId: AppDeploymentRunId, + attempt?: AppDeploymentAttemptIdentity, ): Promise { + const buildSandboxId = + attempt === undefined ? `${runId}-build` : appDeploymentBuildSandboxId(runId, attempt); + const deploySandboxId = + attempt === undefined ? `${runId}-deploy` : appDeploymentDeploySandboxId(runId, attempt); + await Promise.all([ destroyDeploymentSandboxBestEffort( bindings, - appDeploymentBuildSandboxId(runId), + buildSandboxId, "app-deployment.build_sandbox_destroy_failed", ), destroyDeploymentSandboxBestEffort( bindings, - appDeploymentDeploySandboxId(runId), + deploySandboxId, "app-deployment.deploy_sandbox_destroy_failed", ), ]); @@ -241,7 +434,7 @@ async function destroyDeploymentSandboxBestEffort( eventName: string, ): Promise { try { - await destroyRuntimeSubjectContainer(bindings, sandboxId); + await destroyUnversionedSandboxContainer(bindings, sandboxId); } catch (error) { logError(eventName, { ...createErrorLogContext(error), @@ -283,193 +476,163 @@ async function readCurrentDispatchContext( return deployment === null ? null : { deployment, run }; } -async function updateRunStatus( - database: D1Database, - runId: AppDeploymentRunId, - status: Extract< - AppDeploymentRunStatus, - "activating" | "building" | "preparing" | "submitted" | "submitting" - >, -): Promise { - const result = await getAppDatabase(database) - .update(appDeploymentRunsTable) - .set({ status, updatedAt: currentTimestampMs() }) - .where( - and( - eq(appDeploymentRunsTable.id, runId), - inArray(appDeploymentRunsTable.status, ACTIVE_APP_DEPLOYMENT_RUN_STATUSES), - ), - ) - .run(); +type AppDeploymentProgressStatus = Extract< + AppDeploymentRunStatus, + "activating" | "building" | "preparing" | "queued" | "submitted" | "submitting" +>; - return getD1ChangeCount(result) > 0; -} - -async function storeDeploymentPlan(input: { - database: D1Database; - plan: AppDeploymentPlan; - runId: AppDeploymentRunId; - targetName: string; -}): Promise { - const targetKind = input.plan.targetKind; - const result = await getAppDatabase(input.database) - .update(appDeploymentRunsTable) - .set({ - generatedWranglerConfigJson: JSON.stringify({ toml: input.plan.generatedWranglerConfig }), - planJson: JSON.stringify(input.plan), - targetKind, - targetProjectName: targetKind === "cloudflare_pages" ? input.targetName : null, - targetScriptName: targetKind === "cloudflare_worker" ? input.targetName : null, - updatedAt: currentTimestampMs(), - }) - .where( - and( - eq(appDeploymentRunsTable.id, input.runId), - inArray(appDeploymentRunsTable.status, ACTIVE_APP_DEPLOYMENT_RUN_STATUSES), - ), - ) - .run(); +const APP_DEPLOYMENT_STATUS_INDEX = new Map( + ACTIVE_APP_DEPLOYMENT_RUN_STATUSES.map((status, index) => [status, index]), +); - return getD1ChangeCount(result) > 0; +function activeRunStatusIndex(status: AppDeploymentRunStatus): number | null { + return APP_DEPLOYMENT_STATUS_INDEX.get(status as AppDeploymentProgressStatus) ?? null; } -async function failDeploymentRunIfActive(input: { - database: D1Database; - errorCode: string; - errorMessage: string; - runId: AppDeploymentRunId; -}): Promise { - await getAppDatabase(input.database) - .update(appDeploymentRunsTable) - .set({ - errorCode: input.errorCode, - errorMessage: input.errorMessage, - status: "failed", - updatedAt: currentTimestampMs(), - }) - .where( - and( - eq(appDeploymentRunsTable.id, input.runId), - inArray(appDeploymentRunsTable.status, ACTIVE_APP_DEPLOYMENT_RUN_STATUSES), +function exactAppDeploymentAttempt( + database: D1Database, + authority: AppDeploymentAttemptAuthority, + runId: AppDeploymentRunId, + nowMs: number, +) { + return exists( + getAppDatabase(database) + .select({ id: apiCommandsTable.id }) + .from(apiCommandsTable) + .where( + and( + exactApiCommandClaimPredicate(authority, nowMs), + eq(apiCommandsTable.kind, "app_deployment_run_dispatch"), + sql`json_extract(${apiCommandsTable.payloadJson}, '$.appDeploymentRunId') = ${runId}`, + ), ), - ) - .run(); + ); } -async function completeDeploymentRun(input: { +async function advanceRunStatus(input: { + authority: AppDeploymentAttemptAuthority; database: D1Database; - deployment: AppDeploymentRow; - result: AppDeploymentDeployResult; - run: AppDeploymentRunRow; + expected: AppDeploymentProgressStatus; + next: Exclude; + runId: AppDeploymentRunId; }): Promise { - const nowMs = currentTimestampMs(); - const deploymentUpdate = await getAppDatabase(input.database) - .update(appDeploymentsTable) - .set({ - lastSuccessfulUrl: input.result.url, - updatedAt: nowMs, - }) - .where( - and( - eq(appDeploymentsTable.id, input.deployment.id), - eq(appDeploymentsTable.latestRunId, input.run.id), - isNull(appDeploymentsTable.deletedAt), - ), - ) - .run(); + const changed = await runOwnedEffect(input.authority.requireOwnership, async () => { + const nowMs = currentTimestampMs(); + const result = await getAppDatabase(input.database) + .update(appDeploymentRunsTable) + .set({ status: input.next, updatedAt: nowMs }) + .where( + and( + eq(appDeploymentRunsTable.id, input.runId), + eq(appDeploymentRunsTable.status, input.expected), + exactAppDeploymentAttempt(input.database, input.authority, input.runId, nowMs), + ), + ) + .run(); - if (getD1ChangeCount(deploymentUpdate) === 0) { - return false; + return getD1ChangeCount(result) > 0; + }); + + if (changed) { + return true; } - const runUpdate = await getAppDatabase(input.database) - .update(appDeploymentRunsTable) - .set({ - errorCode: null, - errorMessage: null, - externalDeploymentId: input.result.externalDeploymentId, - externalProjectId: input.result.externalProjectId, - externalVersionId: input.result.externalVersionId, - status: "success", - updatedAt: nowMs, - url: input.result.url, - }) - .where( - and( - eq(appDeploymentRunsTable.id, input.run.id), - inArray(appDeploymentRunsTable.status, ACTIVE_APP_DEPLOYMENT_RUN_STATUSES), - ), - ) - .run(); + const currentStatus = await runOwnedEffect(input.authority.requireOwnership, async () => { + const row = await getAppDatabase(input.database) + .select({ status: appDeploymentRunsTable.status }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, input.runId)) + .limit(1) + .get(); + return row?.status ?? null; + }); + const currentIndex = currentStatus === null ? null : activeRunStatusIndex(currentStatus); + const nextIndex = APP_DEPLOYMENT_STATUS_INDEX.get(input.next); - return getD1ChangeCount(runUpdate) > 0; + return currentIndex !== null && nextIndex !== undefined && currentIndex >= nextIndex; } -async function shouldCompensateDeletedDeployment(input: { +async function storeDeploymentPlan(input: { + authority: AppDeploymentAttemptAuthority; database: D1Database; - deployment: AppDeploymentRow; + plan: AppDeploymentPlan; + runId: AppDeploymentRunId; + targetScriptName: string; }): Promise { - const deletedDeployment = - (await getAppDatabase(input.database) - .select({ id: appDeploymentsTable.id }) - .from(appDeploymentsTable) + const targetKind = input.plan.targetKind; + const planJson = JSON.stringify(input.plan); + const generatedWranglerConfigJson = JSON.stringify({ + toml: input.plan.generatedWranglerConfig, + }); + const changed = await runOwnedEffect(input.authority.requireOwnership, async () => { + const nowMs = currentTimestampMs(); + const result = await getAppDatabase(input.database) + .update(appDeploymentRunsTable) + .set({ + generatedWranglerConfigJson, + planJson, + targetKind, + targetProjectName: null, + targetScriptName: input.targetScriptName, + updatedAt: nowMs, + }) .where( and( - eq(appDeploymentsTable.id, input.deployment.id), - isNotNull(appDeploymentsTable.deletedAt), + eq(appDeploymentRunsTable.id, input.runId), + inArray(appDeploymentRunsTable.status, ACTIVE_APP_DEPLOYMENT_RUN_STATUSES), + exactAppDeploymentAttempt(input.database, input.authority, input.runId, nowMs), + or( + isNull(appDeploymentRunsTable.planJson), + and( + eq(appDeploymentRunsTable.planJson, planJson), + eq(appDeploymentRunsTable.generatedWranglerConfigJson, generatedWranglerConfigJson), + eq(appDeploymentRunsTable.targetKind, targetKind), + isNull(appDeploymentRunsTable.targetProjectName), + ), + ), ), ) - .limit(1) - .get()) ?? null; + .run(); + return getD1ChangeCount(result) > 0; + }); - if (deletedDeployment === null) { - return false; + if (changed) { + return true; } - const replacement = - (await getAppDatabase(input.database) - .select({ id: appDeploymentsTable.id }) - .from(appDeploymentsTable) - .where( - and( - eq(appDeploymentsTable.appId, input.deployment.appId), - isNull(appDeploymentsTable.deletedAt), - ), - ) + const run = await runOwnedEffect(input.authority.requireOwnership, async () => + getAppDatabase(input.database) + .select({ + generatedWranglerConfigJson: appDeploymentRunsTable.generatedWranglerConfigJson, + planJson: appDeploymentRunsTable.planJson, + status: appDeploymentRunsTable.status, + targetKind: appDeploymentRunsTable.targetKind, + targetProjectName: appDeploymentRunsTable.targetProjectName, + targetScriptName: appDeploymentRunsTable.targetScriptName, + }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, input.runId)) .limit(1) - .get()) ?? null; + .get(), + ); - return replacement === null; -} + if (run === undefined || !isActiveRunStatus(run.status)) { + return false; + } -async function compensateDeletedDeploymentResources(input: { - bindings: ApiBindings; - cloudflareClient: CloudflareDeploymentClient | null; - deployment: AppDeploymentRow; -}): Promise { if ( - !(await shouldCompensateDeletedDeployment({ - database: input.bindings.DB, - deployment: input.deployment, - })) + run.planJson !== planJson || + run.generatedWranglerConfigJson !== generatedWranglerConfigJson || + run.targetKind !== targetKind || + run.targetProjectName !== null || + run.targetScriptName !== input.targetScriptName ) { - return; - } - - const deleteFailures = await deleteCloudflareDeploymentResources( - input.cloudflareClient ?? createCloudflareDeploymentClient(input.bindings), - { - hostname: deploymentHostname(input.deployment, input.bindings.MOSOO_APP_DEPLOYMENT_DOMAIN), - resourceName: input.deployment.mosooSubdomain, - }, - ); - - if (deleteFailures.length > 0) { - logCloudflareDeploymentResourceDeleteFailures( - "app-deployment.cloudflare_delete_after_deletion_failed", - deleteFailures, + throw new AppDeploymentNonRetryableError( + "Deployment plan changed while retrying the same source commit.", ); } + + return true; } function shouldIncludeSnapshotPath(path: string): boolean { @@ -478,79 +641,103 @@ function shouldIncludeSnapshotPath(path: string): boolean { return SNAPSHOT_FILE_NAMES.has(fileName); } +function isUnsupportedStaticAssetsPath(path: string): boolean { + const parts = path.split("/"); + const fileName = parts.at(-1); + + return fileName === "_routes.json" || fileName === "_worker.js" || parts.includes("functions"); +} + async function readRepositorySnapshot( sandbox: SandboxHandle, repoDir: string, + budget: AppDeploymentExecutionBudget, ): Promise { - const listResult = await sandbox.exec( - `sh -lc ${quoteShellArg(`cd ${quoteShellArg(repoDir)} && find . -type f -print | sort`)}`, + const listResult = await budget.run("Repository file listing", (timeoutMs) => + sandbox.exec( + `sh -lc ${quoteShellArg(`cd ${quoteShellArg(repoDir)} && find . -type f -print | sort`)}`, + { timeout: Math.min(timeoutMs, APP_DEPLOYMENT_SANDBOX_COMMAND_TIMEOUT_MS) }, + ), ); assertSuccessfulCommand(listResult, "Repository file listing"); const files: Record = {}; - const paths = listResult.stdout + const repositoryPaths = listResult.stdout .split("\n") .map((line) => line.trim().replace(/^\.\//u, "")) - .filter((path) => path.length > 0 && shouldIncludeSnapshotPath(path)); + .filter((path) => path.length > 0); + const paths = repositoryPaths.filter(shouldIncludeSnapshotPath); await Promise.all( paths.map(async (path) => { - files[path] = (await sandbox.readFile(`${repoDir}/${path}`, { encoding: "utf8" })).content; + files[path] = ( + await budget.run(`Repository snapshot read (${path})`, () => + sandbox.readFile(`${repoDir}/${path}`, { encoding: "utf8" }), + ) + ).content; }), ); + for (const path of repositoryPaths.filter(isUnsupportedStaticAssetsPath)) { + files[path] ??= ""; + } + return { files }; } -function pagesRoutesFallbackCommands(plan: AppDeploymentPlan, outputDir: string): string[] { +function staticAssetsRoutesFallbackCommands(plan: AppDeploymentPlan): string[] { if (plan.routesFallback === null) { return []; } return [ - `printf '%s\\n' ${quoteShellArg(`/* /${plan.routesFallback} 200`)} > ${quoteShellArg( - `${outputDir}/_redirects`, - )}`, + `printf '\\n%s\\n' ${quoteShellArg(`/* /${plan.routesFallback} 200`)} >> ${quoteShellArg("_redirects")}`, ]; } -async function createPagesArtifactArchive(input: { +async function createStaticAssetsArtifactArchive(input: { + budget: AppDeploymentExecutionBudget; plan: AppDeploymentPlan; prepared: PreparedAppDeploymentRepository; buildSandbox: SandboxHandle; - workDir: string; -}): Promise { +}): Promise { if (input.plan.outputDir === null) { - throw new AppDeploymentNonRetryableError("Pages deployment plan is missing outputDir."); + throw new AppDeploymentNonRetryableError("Static Assets deployment plan is missing outputDir."); } - const archivePath = `${input.workDir}/artifact.tar`; + const archivePath = `${APP_DEPLOYMENT_ARTIFACT_MOUNT_PATH}/${APP_DEPLOYMENT_ARTIFACT_FILE_NAME}`; + const repoDir = input.prepared.repoDir; const outputDir = `${input.prepared.repoDir}/${input.plan.rootDir}/${input.plan.outputDir}`; await execChecked( input.buildSandbox, [ `rm -f ${quoteShellArg(archivePath)}`, - ...pagesRoutesFallbackCommands(input.plan, outputDir), + `repo_root=$(cd ${quoteShellArg(repoDir)} && pwd -P)`, + `case "$repo_root" in ${quoteShellArg(repoDir)}) ;; *) echo 'Static Assets repository root is not canonical.' >&2; exit 64 ;; esac`, `cd ${quoteShellArg(outputDir)}`, - `find . -type f -print0 | tar --null --no-recursion -cf ${quoteShellArg(archivePath)} -T -`, + `output_root=$(pwd -P)`, + `case "$output_root" in "$repo_root"|"$repo_root"/*) ;; *) echo 'Static Assets output escapes repository root.' >&2; exit 64 ;; esac`, + `test ! -e ${quoteShellArg("_worker.js")} || { echo '_worker.js must be migrated to a Worker module.' >&2; exit 64; }`, + `test ! -e ${quoteShellArg("_routes.json")} || { echo '_routes.json must be migrated to Static Assets routing.' >&2; exit 64; }`, + `test ! -e ${quoteShellArg("functions")} || { echo 'Pages functions must be migrated to a Worker module.' >&2; exit 64; }`, + ...staticAssetsRoutesFallbackCommands(input.plan), + `find . -name .git -prune -o -type f -print0 | tar --null --no-recursion -cf ${quoteShellArg(archivePath)} -T -`, ].join(" && "), - "Pages artifact archive", + "Static Assets artifact archive", + input.budget, { retryable: false }, ); - - return (await input.buildSandbox.readFile(archivePath, { encoding: "base64" })).content; } -async function extractPagesArtifactArchive(input: { - archiveBase64: string; +async function extractStaticAssetsArtifactArchive(input: { + budget: AppDeploymentExecutionBudget; deploySandbox: SandboxHandle; workDir: string; }): Promise<{ artifactDir: string; deployDir: string }> { - const artifactDir = `${input.workDir}/artifact`; - const archiveBase64Path = `${input.workDir}/artifact.tar.b64`; - const archivePath = `${input.workDir}/artifact.tar`; + const archivePath = `${APP_DEPLOYMENT_ARTIFACT_MOUNT_PATH}/${APP_DEPLOYMENT_ARTIFACT_FILE_NAME}`; const deployDir = `${input.workDir}/deploy`; + const artifactDir = `${deployDir}/artifact`; await execChecked( input.deploySandbox, @@ -558,42 +745,391 @@ async function extractPagesArtifactArchive(input: { `rm -rf ${quoteShellArg(input.workDir)}`, `mkdir -p ${quoteShellArg(artifactDir)} ${quoteShellArg(deployDir)}`, ].join(" && "), - "Pages deploy workspace", + "Static Assets deploy workspace", + input.budget, ); - await input.deploySandbox.writeFile(archiveBase64Path, input.archiveBase64); await execChecked( input.deploySandbox, - [ - `base64 -d ${quoteShellArg(archiveBase64Path)} > ${quoteShellArg(archivePath)}`, - `tar -xf ${quoteShellArg(archivePath)} -C ${quoteShellArg(artifactDir)}`, - ].join(" && "), - "Pages artifact extraction", + `tar -xf ${quoteShellArg(archivePath)} -C ${quoteShellArg(artifactDir)}`, + "Static Assets artifact extraction", + input.budget, ); return { artifactDir, deployDir }; } +interface StaticAssetsInventory { + assetsIgnore: string | null; + headers: string | null; + paths: string[]; + redirects: string | null; +} + +interface PreparedStaticAssetsArtifact { + artifactDir: string; + headers: string | null; + manifest: CloudflareStaticAssetManifest; + redirects: string | null; + toolPath: string; + workDir: string; +} + +export const APP_DEPLOYMENT_STATIC_ASSET_TOOL_SOURCE = String.raw` +import { createHash } from "node:crypto"; +import { extname, join } from "node:path"; +import { readFile, readdir, writeFile } from "node:fs/promises"; + +const MAX_ASSET_SIZE = 25 * 1024 * 1024; +const MIME_TYPES = { + avif: "image/avif", css: "text/css", csv: "text/csv", gif: "image/gif", + htm: "text/html", html: "text/html", ico: "image/x-icon", jpeg: "image/jpeg", + jpg: "image/jpeg", js: "application/javascript", json: "application/json", + m4a: "audio/mp4", map: "application/json", mjs: "application/javascript", + mp3: "audio/mpeg", mp4: "video/mp4", ogg: "audio/ogg", otf: "font/otf", + pdf: "application/pdf", png: "image/png", svg: "image/svg+xml", ttf: "font/ttf", + txt: "text/plain", wasm: "application/wasm", wav: "audio/wav", webm: "video/webm", + webmanifest: "application/manifest+json", webp: "image/webp", woff: "font/woff", + woff2: "font/woff2", xml: "application/xml" +}; + +function checkedRelativePath(value) { + if (typeof value !== "string" || value.length === 0 || value.startsWith("/") || value.includes("\\0")) { + throw new Error("Static Assets artifact contains an invalid path."); + } + const parts = value.split("/"); + if (parts.some((part) => part.length === 0 || part === "." || part === "..")) { + throw new Error("Static Assets artifact contains an invalid path."); + } + return value; +} + +async function readOptional(path) { + try { + return await readFile(path, "utf8"); + } catch (error) { + if (error && typeof error === "object" && error.code === "ENOENT") return null; + throw error; + } +} + +async function listFiles(root, relative = "") { + const entries = await readdir(join(root, relative), { withFileTypes: true }); + const paths = []; + for (const entry of entries.sort((left, right) => left.name.localeCompare(right.name))) { + const child = relative.length === 0 ? entry.name : relative + "/" + entry.name; + if (entry.isDirectory()) paths.push(...await listFiles(root, child)); + else if (entry.isFile()) paths.push(child); + } + return paths; +} + +async function createInventory(root, outputPath) { + const paths = await listFiles(root); + await writeFile(outputPath, JSON.stringify({ + assetsIgnore: await readOptional(join(root, ".assetsignore")), + headers: await readOptional(join(root, "_headers")), + paths, + redirects: await readOptional(join(root, "_redirects")) + })); +} + +async function createManifest(root, selectionPath, outputPath, scriptName) { + const selected = JSON.parse(await readFile(selectionPath, "utf8")); + if (!Array.isArray(selected) || typeof scriptName !== "string" || scriptName.length === 0) { + throw new Error("Static Assets manifest input is invalid."); + } + const manifest = {}; + for (const item of selected) { + const relative = checkedRelativePath(item); + const contents = await readFile(join(root, relative)); + if (contents.byteLength > MAX_ASSET_SIZE) { + throw new Error("Static asset exceeds the 25 MiB Cloudflare limit: " + relative); + } + const hash = createHash("sha256") + .update(scriptName) + .update(Buffer.from([0])) + .update(extname(relative).slice(1).toLowerCase()) + .update(Buffer.from([0])) + .update(contents) + .digest("hex") + .slice(0, 32); + manifest["/" + relative] = { hash, size: contents.byteLength }; + } + await writeFile(outputPath, JSON.stringify(manifest)); +} + +function jwtPayload(token) { + const encoded = token.split(".")[1]; + if (!encoded) throw new Error("Static Assets upload token is invalid."); + return JSON.parse(Buffer.from(encoded, "base64url").toString("utf8")); +} + +function contentType(path) { + const type = MIME_TYPES[extname(path).slice(1).toLowerCase()] ?? "application/null"; + return type.startsWith("text/") ? type + "; charset=utf-8" : type; +} + +async function uploadRequest(url, init) { + const response = await fetch(url, init); + let envelope; + try { + envelope = await response.json(); + } catch { + throw new Error("Cloudflare asset upload returned non-JSON (" + response.status + ")."); + } + if (!response.ok || envelope?.success !== true || !envelope.result || typeof envelope.result !== "object") { + throw new Error("Cloudflare asset upload failed (" + response.status + ")."); + } + const completion = envelope.result.jwt; + if (completion !== undefined && (typeof completion !== "string" || completion.length === 0)) { + throw new Error("Cloudflare asset upload returned an invalid completion token."); + } + return completion; +} + +async function uploadAssets(root, manifestPath, bucketsPath, completionPath) { + const accountId = process.env.CLOUDFLARE_ACCOUNT_ID; + const uploadToken = process.env.CLOUDFLARE_ASSET_UPLOAD_JWT; + if (!accountId || !uploadToken) throw new Error("Static Assets upload environment is incomplete."); + const manifest = JSON.parse(await readFile(manifestPath, "utf8")); + const buckets = JSON.parse(await readFile(bucketsPath, "utf8")); + if (!Array.isArray(buckets)) throw new Error("Static Assets upload buckets are invalid."); + const hashPaths = new Map(); + for (const [assetPath, entry] of Object.entries(manifest)) { + if (!entry || typeof entry !== "object" || typeof entry.hash !== "string") { + throw new Error("Static Assets manifest is invalid."); + } + hashPaths.set(entry.hash, checkedRelativePath(assetPath.slice(1))); + } + const single = jwtPayload(uploadToken).wrangler_single_asset_uploads === true; + const baseUrl = "https://api.cloudflare.com/client/v4/accounts/" + encodeURIComponent(accountId) + "/workers/assets/upload"; + let completion; + for (const bucket of buckets) { + if (!Array.isArray(bucket) || bucket.length === 0) throw new Error("Static Assets upload bucket is invalid."); + if (single) { + for (const hash of bucket) { + const relative = hashPaths.get(hash); + if (!relative) throw new Error("Cloudflare requested an unknown Static Assets hash."); + completion = await uploadRequest(baseUrl + "/" + encodeURIComponent(hash), { + body: await readFile(join(root, relative)), + headers: { Authorization: "Bearer " + uploadToken, "Content-Type": contentType(relative) }, + method: "POST" + }) ?? completion; + } + } else { + const form = new FormData(); + for (const hash of bucket) { + const relative = hashPaths.get(hash); + if (!relative) throw new Error("Cloudflare requested an unknown Static Assets hash."); + form.append(hash, (await readFile(join(root, relative))).toString("base64")); + } + completion = await uploadRequest(baseUrl + "?base64=true", { + body: form, + headers: { Authorization: "Bearer " + uploadToken }, + method: "POST" + }) ?? completion; + } + } + if (!completion) throw new Error("Cloudflare did not return a Static Assets completion token."); + await writeFile(completionPath, completion, { mode: 0o600 }); +} + +const [mode, ...args] = process.argv.slice(2); +if (mode === "inventory" && args.length === 2) await createInventory(args[0], args[1]); +else if (mode === "manifest" && args.length === 4) await createManifest(args[0], args[1], args[2], args[3]); +else if (mode === "upload" && args.length === 4) await uploadAssets(args[0], args[1], args[2], args[3]); +else throw new Error("Static Assets tool invocation is invalid."); +`; + +function parseStaticAssetsInventory(source: string): StaticAssetsInventory { + const value = JSON.parse(source) as unknown; + + if (typeof value !== "object" || value === null) { + throw new Error("Static Assets inventory is invalid."); + } + + const assetsIgnore = Reflect.get(value, "assetsIgnore"); + const headers = Reflect.get(value, "headers"); + const paths = Reflect.get(value, "paths"); + const redirects = Reflect.get(value, "redirects"); + + if ( + (assetsIgnore !== null && typeof assetsIgnore !== "string") || + (headers !== null && typeof headers !== "string") || + !Array.isArray(paths) || + !paths.every((path) => typeof path === "string" && path.length > 0) || + (redirects !== null && typeof redirects !== "string") + ) { + throw new Error("Static Assets inventory is invalid."); + } + + return { assetsIgnore, headers, paths, redirects }; +} + +function selectStaticAssetPaths(inventory: StaticAssetsInventory): string[] { + const patterns = ["/.assetsignore", "/_redirects", "/_headers"]; + + if (inventory.assetsIgnore !== null) { + patterns.push(...inventory.assetsIgnore.split("\n")); + } + + const matcher = createIgnore().add(patterns); + return inventory.paths.filter((path) => !matcher.test(path).ignored).toSorted(); +} + +function parseStaticAssetsManifest(source: string): CloudflareStaticAssetManifest { + const value = JSON.parse(source) as unknown; + + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new Error("Static Assets manifest is invalid."); + } + + return value as CloudflareStaticAssetManifest; +} + +async function prepareStaticAssetsArtifact(input: { + activeScriptName: string; + budget: AppDeploymentExecutionBudget; + deploySandbox: SandboxHandle; + workDir: string; +}): Promise { + const { artifactDir, deployDir } = await extractStaticAssetsArtifactArchive(input); + const toolPath = `${deployDir}/${APP_DEPLOYMENT_STATIC_ASSET_TOOL_FILE_NAME}`; + const inventoryPath = `${deployDir}/${APP_DEPLOYMENT_STATIC_ASSET_INVENTORY_FILE_NAME}`; + const selectionPath = `${deployDir}/${APP_DEPLOYMENT_STATIC_ASSET_SELECTION_FILE_NAME}`; + const manifestPath = `${deployDir}/${APP_DEPLOYMENT_STATIC_ASSET_MANIFEST_FILE_NAME}`; + + await input.budget.run("Static Assets tool write", () => + input.deploySandbox.writeFile(toolPath, APP_DEPLOYMENT_STATIC_ASSET_TOOL_SOURCE), + ); + await execChecked( + input.deploySandbox, + `node ${quoteShellArg(toolPath)} inventory ${quoteShellArg(artifactDir)} ${quoteShellArg(inventoryPath)}`, + "Static Assets inventory", + input.budget, + { retryable: false }, + ); + const inventory = parseStaticAssetsInventory( + ( + await input.budget.run("Static Assets inventory read", () => + input.deploySandbox.readFile(inventoryPath, { encoding: "utf8" }), + ) + ).content, + ); + await input.budget.run("Static Assets selection write", () => + input.deploySandbox.writeFile(selectionPath, JSON.stringify(selectStaticAssetPaths(inventory))), + ); + await execChecked( + input.deploySandbox, + [ + "node", + quoteShellArg(toolPath), + "manifest", + quoteShellArg(artifactDir), + quoteShellArg(selectionPath), + quoteShellArg(manifestPath), + quoteShellArg(input.activeScriptName), + ].join(" "), + "Static Assets manifest", + input.budget, + { retryable: false }, + ); + const manifest = parseStaticAssetsManifest( + ( + await input.budget.run("Static Assets manifest read", () => + input.deploySandbox.readFile(manifestPath, { encoding: "utf8" }), + ) + ).content, + ); + + return { + artifactDir, + headers: inventory.headers, + manifest, + redirects: inventory.redirects, + toolPath, + workDir: deployDir, + }; +} + +async function bundleWorkerModule(input: { + budget: AppDeploymentExecutionBudget; + plan: AppDeploymentPlan; + prepared: PreparedAppDeploymentRepository; + sandbox: SandboxHandle; +}): Promise<{ mainModuleName: string; scriptContent: string }> { + if (input.plan.workerEntry === null) { + throw new AppDeploymentNonRetryableError("Worker deployment plan is missing workerEntry."); + } + + if (!WORKER_JS_ENTRY_PATTERN.test(input.plan.workerEntry)) { + throw new AppDeploymentNonRetryableError( + "Worker deployment requires a JavaScript module entry.", + ); + } + + const rootDir = `${input.prepared.repoDir}/${input.plan.rootDir}`; + const bundlePath = `${rootDir}/.mosoo-worker-bundle.mjs`; + await execChecked( + input.sandbox, + [ + `cd ${quoteShellArg(rootDir)} &&`, + "bun build", + quoteShellArg(input.plan.workerEntry), + "--target=browser", + "--format=esm", + `--outfile=${quoteShellArg(bundlePath)}`, + `--external=${quoteShellArg("cloudflare:*")}`, + ].join(" "), + "Worker module bundle", + input.budget, + { retryable: false }, + ); + + return { + mainModuleName: "worker.mjs", + scriptContent: ( + await input.budget.run("Worker bundle read", () => + input.sandbox.readFile(bundlePath, { encoding: "utf8" }), + ) + ).content, + }; +} + class SandboxAppDeploymentBuildRunner implements AppDeploymentBuildRunner { + readonly #attempt: AppDeploymentAttemptIdentity; readonly #bindings: ApiBindings; + readonly #budget: AppDeploymentExecutionBudget; readonly #cloudflareClient: CloudflareDeploymentClient; + readonly #requireOwnership: () => Promise; #buildSandbox: SandboxHandle | null = null; - #buildWorkDir: string | null = null; #runId: AppDeploymentRunId | null = null; - constructor(bindings: ApiBindings, cloudflareClient: CloudflareDeploymentClient) { + constructor( + bindings: ApiBindings, + cloudflareClient: CloudflareDeploymentClient, + authority: AppDeploymentAttemptAuthority, + budget: AppDeploymentExecutionBudget, + ) { + this.#attempt = authority; this.#bindings = bindings; + this.#budget = budget; this.#cloudflareClient = cloudflareClient; + this.#requireOwnership = authority.requireOwnership; } async prepare(input: { deployment: AppDeploymentRow; run: AppDeploymentRunRow; }): Promise { - const sandbox = await getRuntimeSubjectKeepAliveHandle( - this.#bindings, - appDeploymentBuildSandboxId(input.run.id), + const sandbox = await this.#budget.run("Build sandbox acquisition", () => + getEphemeralUnversionedSandboxHandle( + this.#bindings, + appDeploymentBuildSandboxId(input.run.id, this.#attempt), + APP_DEPLOYMENT_SANDBOX_SLEEP_AFTER_SECONDS, + ), ); - const workDir = `/tmp/mosoo-app-deployment-build-${input.run.id}`; + const workDir = `/tmp/mosoo-app-deployment-${input.run.id}-${appDeploymentAttemptSuffix(this.#attempt)}-build`; const repoDir = `${workDir}/repo`; const cloneCommand = [ `rm -rf ${quoteShellArg(workDir)}`, @@ -604,16 +1140,14 @@ class SandboxAppDeploymentBuildRunner implements AppDeploymentBuildRunner { `git checkout --detach ${quoteShellArg(input.run.sourceCommitSha)}`, ].join(" && "); - await sandbox.setKeepAlive(true); - await execChecked(sandbox, cloneCommand, "Repository clone"); + await execChecked(sandbox, cloneCommand, "Repository clone", this.#budget); this.#buildSandbox = sandbox; - this.#buildWorkDir = workDir; this.#runId = input.run.id; return { repoDir, - snapshot: await readRepositorySnapshot(sandbox, repoDir), + snapshot: await readRepositorySnapshot(sandbox, repoDir, this.#budget), }; } @@ -630,9 +1164,11 @@ class SandboxAppDeploymentBuildRunner implements AppDeploymentBuildRunner { return; } - const buildSession = await sandbox.createSession({ - cwd: `${input.prepared.repoDir}/${input.plan.rootDir}`, - }); + const buildSession = await this.#budget.run("Build session creation", () => + sandbox.createSession({ + cwd: `${input.prepared.repoDir}/${input.plan.rootDir}`, + }), + ); await execChecked( buildSession, @@ -640,151 +1176,185 @@ class SandboxAppDeploymentBuildRunner implements AppDeploymentBuildRunner { " && ", ), "App deployment build", + this.#budget, { retryable: false }, ); } async deploy(input: { + activeScriptName: string; deployment: AppDeploymentRow; envVars: Record; plan: AppDeploymentPlan; prepared: PreparedAppDeploymentRepository; run: AppDeploymentRunRow; - }): Promise { + }): Promise { const buildSandbox = this.#requireBuildSandbox(); - const buildWorkDir = this.#requireBuildWorkDir(); - const targetName = input.deployment.mosooSubdomain; - const domain = this.#bindings.MOSOO_APP_DEPLOYMENT_DOMAIN; - const hostname = deploymentHostname(input.deployment, domain); - - if (input.plan.targetKind === "cloudflare_pages") { - const project = await this.#cloudflareClient.ensurePagesProject({ - branch: input.run.sourceBranch, - projectName: targetName, - }); - const archiveBase64 = await createPagesArtifactArchive({ + const tags = appDeploymentCandidateScriptTags({ + deploymentId: input.deployment.id, + runId: input.run.id, + }); + + if (input.plan.targetKind === "cloudflare_static_assets") { + await this.#owned(() => + this.#budget.run("Build artifact bucket mount", () => + mountAppDeploymentArtifactBucket( + this.#bindings, + buildSandbox, + input.activeScriptName, + false, + ), + ), + ); + await createStaticAssetsArtifactArchive({ + budget: this.#budget, buildSandbox, plan: input.plan, prepared: input.prepared, - workDir: buildWorkDir, }); + await this.#owned(() => + this.#budget.run("Build artifact bucket unmount", () => + buildSandbox.unmountBucket(APP_DEPLOYMENT_ARTIFACT_MOUNT_PATH), + ), + ); await this.#destroyBuildSandbox(); - const deploySandbox = await getRuntimeSubjectKeepAliveHandle( - this.#bindings, - appDeploymentDeploySandboxId(input.run.id), + const deploySandbox = await this.#budget.run("Deploy sandbox acquisition", () => + getEphemeralUnversionedSandboxHandle( + this.#bindings, + appDeploymentDeploySandboxId(input.run.id, this.#attempt), + APP_DEPLOYMENT_SANDBOX_SLEEP_AFTER_SECONDS, + ), ); - const deployWorkDir = `/tmp/mosoo-app-deployment-deploy-${input.run.id}`; - await deploySandbox.setKeepAlive(true); - await assertControlledWranglerAvailable(deploySandbox); - const { artifactDir, deployDir } = await extractPagesArtifactArchive({ - archiveBase64, + const deployWorkDir = `/tmp/mosoo-app-deployment-${input.run.id}-${appDeploymentAttemptSuffix(this.#attempt)}-deploy`; + await this.#owned(() => + this.#budget.run("Deploy artifact bucket mount", () => + mountAppDeploymentArtifactBucket( + this.#bindings, + deploySandbox, + input.activeScriptName, + true, + ), + ), + ); + const artifact = await prepareStaticAssetsArtifact({ + activeScriptName: input.activeScriptName, + budget: this.#budget, deploySandbox, workDir: deployWorkDir, }); - const deploySession = await deploySandbox.createSession({ - cwd: deployDir, - env: { - CLOUDFLARE_ACCOUNT_ID: this.#bindings.CLOUDFLARE_ACCOUNT_ID, - CLOUDFLARE_API_TOKEN: this.#bindings.CLOUDFLARE_API_TOKEN, - }, - }); - - await execChecked( - deploySession, - [ - "wrangler", - "pages", - "deploy", - quoteShellArg(artifactDir), - "--project-name", - quoteShellArg(targetName), - "--branch", - quoteShellArg(input.run.sourceBranch), - ].join(" "), - "Cloudflare Pages deploy", + const uploadSession = await this.#owned(() => + this.#cloudflareClient.createStaticAssetsUploadSession({ + manifest: artifact.manifest, + scriptName: input.activeScriptName, + timeoutMs: this.#budget.remaining(), + }), ); + let completionToken = uploadSession.uploadToken; + + if (uploadSession.buckets.length > 0) { + const bucketsPath = `${artifact.workDir}/${APP_DEPLOYMENT_STATIC_ASSET_BUCKETS_FILE_NAME}`; + const manifestPath = `${artifact.workDir}/${APP_DEPLOYMENT_STATIC_ASSET_MANIFEST_FILE_NAME}`; + const completionPath = `${artifact.workDir}/${APP_DEPLOYMENT_STATIC_ASSET_COMPLETION_FILE_NAME}`; + await this.#budget.run("Static Assets upload bucket write", () => + deploySandbox.writeFile(bucketsPath, JSON.stringify(uploadSession.buckets)), + ); + const deploySession = await this.#budget.run("Static Assets upload session creation", () => + deploySandbox.createSession({ + cwd: artifact.workDir, + env: { + CLOUDFLARE_ACCOUNT_ID: this.#bindings.CLOUDFLARE_ACCOUNT_ID, + CLOUDFLARE_ASSET_UPLOAD_JWT: uploadSession.uploadToken, + }, + }), + ); + + await this.#owned(() => + execChecked( + deploySession, + [ + "node", + quoteShellArg(artifact.toolPath), + "upload", + quoteShellArg(artifact.artifactDir), + quoteShellArg(manifestPath), + quoteShellArg(bucketsPath), + quoteShellArg(completionPath), + ].join(" "), + "Cloudflare Static Assets upload", + this.#budget, + ), + ); + completionToken = ( + await this.#budget.run("Static Assets completion token read", () => + deploySession.readFile(completionPath, { encoding: "utf8" }), + ) + ).content; + + if (completionToken.length === 0) { + throw new Error("Cloudflare did not return a Static Assets completion token."); + } + + await execChecked( + deploySession, + `rm -f ${quoteShellArg(completionPath)}`, + "Static Assets completion token cleanup", + this.#budget, + ); + } - const [latestDeployment, domainResult] = await Promise.all([ - this.#cloudflareClient.getLatestPagesDeployment({ - projectName: targetName, + const deployed = await this.#owned(() => + this.#cloudflareClient.deployStaticAssets({ + compatibilityDate: APP_DEPLOYMENT_COMPATIBILITY_DATE, + completionToken, + headers: artifact.headers, + redirects: artifact.redirects, + scriptName: input.activeScriptName, + tags, + timeoutMs: this.#budget.remaining(), }), - this.#cloudflareClient.ensurePagesDomain({ - hostname, - projectName: targetName, - }), - ]); - const url = - domainResult.status === "active" - ? deploymentUrl(input.deployment, domain) - : latestDeployment.url; - - if (url === null) { - throw new Error("Cloudflare Pages deployment response did not include a live URL."); - } + ); return { - externalDeploymentId: latestDeployment.deploymentId, - externalProjectId: project.projectId, - externalVersionId: null, - url, + externalDeploymentId: deployed.deploymentId, + externalProjectId: null, + externalVersionId: deployed.versionId, }; } - if (input.plan.workerEntry === null) { - throw new AppDeploymentNonRetryableError("Worker deployment plan is missing workerEntry."); - } - - if (!WORKER_JS_ENTRY_PATTERN.test(input.plan.workerEntry)) { - throw new AppDeploymentNonRetryableError( - "Worker deployment requires a JavaScript module entry.", - ); - } - - const mainModuleName = input.plan.workerEntry.split("/").at(-1) ?? input.plan.workerEntry; - const scriptContent = ( - await buildSandbox.readFile( - `${input.prepared.repoDir}/${input.plan.rootDir}/${input.plan.workerEntry}`, - { - encoding: "utf8", - }, - ) - ).content; - assertSelfContainedWorkerModule(scriptContent); + const { mainModuleName, scriptContent } = await bundleWorkerModule({ + budget: this.#budget, + plan: input.plan, + prepared: input.prepared, + sandbox: buildSandbox, + }); await this.#destroyBuildSandbox(); - const worker = await this.#cloudflareClient.deployWorkerModule({ - compatibilityDate: APP_DEPLOYMENT_COMPATIBILITY_DATE, - mainModuleName, - scriptContent, - scriptName: targetName, - vars: input.envVars, - }); - await this.#cloudflareClient.ensureWorkerRoute({ - hostname, - scriptName: targetName, - }); - await this.#cloudflareClient.ensureWorkerDomain({ - hostname, - scriptName: targetName, - }); + const worker = await this.#owned(() => + this.#cloudflareClient.deployWorkerModule({ + compatibilityDate: APP_DEPLOYMENT_COMPATIBILITY_DATE, + mainModuleName, + scriptContent, + scriptName: input.activeScriptName, + tags, + timeoutMs: this.#budget.remaining(), + vars: input.envVars, + }), + ); return { externalDeploymentId: worker.deploymentId, externalProjectId: null, externalVersionId: worker.versionId, - url: deploymentUrl(input.deployment, domain), }; } async cleanup(): Promise { - if (this.#runId === null) { - return; + if (this.#runId !== null) { + await destroyAppDeploymentRunSandboxesBestEffort(this.#bindings, this.#runId, this.#attempt); } - await destroyAppDeploymentRunSandboxesBestEffort(this.#bindings, this.#runId); this.#buildSandbox = null; - this.#buildWorkDir = null; } async #destroyBuildSandbox(): Promise { @@ -792,13 +1362,15 @@ class SandboxAppDeploymentBuildRunner implements AppDeploymentBuildRunner { return; } - await destroyDeploymentSandboxBestEffort( - this.#bindings, - appDeploymentBuildSandboxId(this.#runId), - "app-deployment.build_sandbox_destroy_failed", + const runId = this.#runId; + await this.#budget.run("Build sandbox destruction", () => + destroyDeploymentSandboxBestEffort( + this.#bindings, + appDeploymentBuildSandboxId(runId, this.#attempt), + "app-deployment.build_sandbox_destroy_failed", + ), ); this.#buildSandbox = null; - this.#buildWorkDir = null; } #requireBuildSandbox(): SandboxHandle { @@ -809,12 +1381,8 @@ class SandboxAppDeploymentBuildRunner implements AppDeploymentBuildRunner { return this.#buildSandbox; } - #requireBuildWorkDir(): string { - if (this.#buildWorkDir === null) { - throw new Error("App deployment work directory was not prepared."); - } - - return this.#buildWorkDir; + #owned(effect: () => Promise): Promise { + return runOwnedEffect(this.#requireOwnership, effect); } } @@ -868,19 +1436,109 @@ async function resolveDeploymentEnvVars( return envVars; } +function toAppDeploymentTerminalFailure( + runId: AppDeploymentRunId, + error: unknown, +): Extract | null { + if ( + error instanceof AppAgentBindingResolutionError || + error instanceof AppDeploymentDetectionError + ) { + return { + errorCode: error.code, + errorMessage: error.message, + kind: "terminal_failure", + runId, + }; + } + + if (error instanceof AppDeploymentNonRetryableError) { + return { + errorCode: error.name, + errorMessage: error.message, + kind: "terminal_failure", + runId, + }; + } + + return null; +} + +async function admitAppDeploymentCandidate(input: { + bindings: ApiBindings; + budget: AppDeploymentExecutionBudget; + scriptName: string; + url: string; +}): Promise { + const response = await input.budget.run( + "App deployment candidate admission", + (timeoutMs) => + appDeploymentDispatcher(input.bindings) + .get(input.scriptName) + .fetch( + new Request(input.url, { + signal: AbortSignal.timeout(timeoutMs), + }), + ), + 30_000, + ); + + if (!(response instanceof Response)) { + throw new Error("App deployment candidate admission did not return a Response."); + } + + await input.budget.run( + "App deployment candidate response disposal", + () => response.body?.cancel().catch(() => undefined) ?? Promise.resolve(), + 30_000, + ); +} + export async function dispatchAppDeploymentRun( bindings: ApiBindings, input: { appDeploymentRunId: AppDeploymentRunId }, + authority: AppDeploymentAttemptAuthority, options: DispatchAppDeploymentRunOptions = {}, -): Promise { - let context = await readCurrentDispatchContext(bindings.DB, input.appDeploymentRunId); +): Promise { + const startedAt = Date.now(); + const budget = new AppDeploymentExecutionBudget( + Math.min( + options.deadlineMs ?? startedAt + APP_DEPLOYMENT_EXECUTION_TIMEOUT_MS, + startedAt + APP_DEPLOYMENT_EXECUTION_TIMEOUT_MS, + ), + ); + let context = await budget.run("App deployment context read", () => + runOwnedEffect(authority.requireOwnership, () => + readCurrentDispatchContext(bindings.DB, input.appDeploymentRunId), + ), + ); if (context === null) { - return; + return { kind: "skipped" }; } - if (!(await updateRunStatus(bindings.DB, input.appDeploymentRunId, "preparing"))) { - return; + const activeScriptName = appDeploymentCandidateScriptName(authority); + const deploymentId = context.deployment.id; + const runId = context.run.id; + await budget.run("App deployment candidate registration", () => + runOwnedEffect(authority.requireOwnership, () => + registerAppDeploymentScriptCandidate(bindings.DB, { + authority, + deploymentId, + nowMs: currentTimestampMs(), + runId, + scriptName: activeScriptName, + }), + ), + ); + context = await budget.run("Registered App deployment context read", () => + runOwnedEffect(authority.requireOwnership, () => + readCurrentDispatchContext(bindings.DB, input.appDeploymentRunId), + ), + ); + + if (context === null) { + return { kind: "skipped" }; } const cloudflareClient = @@ -891,125 +1549,176 @@ export async function dispatchAppDeploymentRun( new SandboxAppDeploymentBuildRunner( bindings, cloudflareClient ?? createCloudflareDeploymentClient(bindings), + authority, + budget, ); - let externallyAttemptedDeployment: AppDeploymentRow | null = null; try { - const prepared = await runner.prepare(context); - const targetName = context.deployment.mosooSubdomain; + const prepared = await budget.run("App deployment preparation", () => + runOwnedEffect(authority.requireOwnership, () => + runner.prepare(context as AppDeploymentDispatchContext), + ), + ); assertRequestedMosooConfigPresent(context.run, prepared.snapshot); - const plan = detectAppDeploymentPlan(prepared.snapshot, { resourceName: targetName }); + const plan = detectAppDeploymentPlan(prepared.snapshot); if ( - !(await storeDeploymentPlan({ - database: bindings.DB, - plan, - runId: context.run.id, - targetName, - })) + !(await budget.run("App deployment plan persistence", () => + storeDeploymentPlan({ + authority, + database: bindings.DB, + plan, + runId: context!.run.id, + targetScriptName: activeScriptName, + }), + )) ) { - return; + await authority.requireOwnership(); + return { kind: "skipped" }; } - let envVars: Record; - try { - envVars = await resolveDeploymentEnvVars(bindings, context.deployment, context.run, plan); - } catch (error) { - if (error instanceof AppAgentBindingResolutionError) { - await failDeploymentRunIfActive({ - database: bindings.DB, - errorCode: error.code, - errorMessage: error.message, - runId: context.run.id, - }); - return; - } - throw error; - } + const envVars = await budget.run("App deployment binding resolution", () => + runOwnedEffect(authority.requireOwnership, () => + resolveDeploymentEnvVars(bindings, context!.deployment, context!.run, plan), + ), + ); - if (!(await updateRunStatus(bindings.DB, input.appDeploymentRunId, "building"))) { - return; + if ( + !(await budget.run("App deployment building transition", () => + advanceRunStatus({ + authority, + database: bindings.DB, + expected: "preparing", + next: "building", + runId: input.appDeploymentRunId, + }), + )) + ) { + await authority.requireOwnership(); + return { kind: "skipped" }; } - await runner.build({ plan, prepared }); + await budget.run("App deployment build", () => + runOwnedEffect(authority.requireOwnership, () => runner.build({ plan, prepared })), + ); - if (!(await updateRunStatus(bindings.DB, input.appDeploymentRunId, "submitting"))) { - await failDeploymentRunIfActive({ - database: bindings.DB, - errorCode: "deployment_submission_lost", - errorMessage: "Deployment built but the deployment run changed.", - runId: context.run.id, - }); - return; + if ( + !(await budget.run("App deployment submitting transition", () => + advanceRunStatus({ + authority, + database: bindings.DB, + expected: "building", + next: "submitting", + runId: input.appDeploymentRunId, + }), + )) + ) { + await authority.requireOwnership(); + return { kind: "skipped" }; } - context = await readCurrentDispatchContext(bindings.DB, input.appDeploymentRunId); + context = await budget.run("Submitting App deployment context read", () => + runOwnedEffect(authority.requireOwnership, () => + readCurrentDispatchContext(bindings.DB, input.appDeploymentRunId), + ), + ); if (context === null) { - await failDeploymentRunIfActive({ - database: bindings.DB, - errorCode: "deployment_context_lost", - errorMessage: "Deployment context was lost after build.", - runId: input.appDeploymentRunId, - }); - return; + return { kind: "skipped" }; } + const submittingContext = context; + + await budget.run("App deployment upload ledger mark", () => + runOwnedEffect(authority.requireOwnership, () => + markAppDeploymentScriptUploadStarted(bindings.DB, { + authority, + deploymentId: submittingContext.deployment.id, + nowMs: currentTimestampMs(), + runId: submittingContext.run.id, + scriptName: activeScriptName, + }), + ), + ); + const uploaded = await budget.run("App deployment upload", () => + runOwnedEffect(authority.requireOwnership, () => + runner.deploy({ activeScriptName, ...submittingContext, envVars, plan, prepared }), + ), + ); + const plannedUrl = deploymentUrl( + submittingContext.deployment, + bindings.MOSOO_APP_DEPLOYMENT_DOMAIN, + ); + const result: AppDeploymentDeployResult = { + ...uploaded, + activeScriptName, + url: plannedUrl, + }; - externallyAttemptedDeployment = context.deployment; - const result = await runner.deploy({ ...context, envVars, plan, prepared }); + await budget.run("App deployment admission ownership", () => + runOwnedEffect(authority.requireOwnership, () => + admitAppDeploymentCandidate({ + bindings, + budget, + scriptName: activeScriptName, + url: plannedUrl, + }), + ), + ); - if (!(await updateRunStatus(bindings.DB, input.appDeploymentRunId, "submitted"))) { - await failDeploymentRunIfActive({ - database: bindings.DB, - errorCode: "deployment_submission_lost", - errorMessage: "Deployment submitted externally but the deployment run changed.", - runId: context.run.id, - }); - return; + if ( + !(await budget.run("App deployment submitted transition", () => + advanceRunStatus({ + authority, + database: bindings.DB, + expected: "submitting", + next: "submitted", + runId: input.appDeploymentRunId, + }), + )) + ) { + await authority.requireOwnership(); + return { kind: "skipped" }; } - if (!(await updateRunStatus(bindings.DB, input.appDeploymentRunId, "activating"))) { - await failDeploymentRunIfActive({ - database: bindings.DB, - errorCode: "deployment_activation_lost", - errorMessage: "Deployment activated externally but the deployment run changed.", - runId: context.run.id, - }); - return; + if ( + !(await budget.run("App deployment activating transition", () => + advanceRunStatus({ + authority, + database: bindings.DB, + expected: "submitted", + next: "activating", + runId: input.appDeploymentRunId, + }), + )) + ) { + await authority.requireOwnership(); + return { kind: "skipped" }; } - const completed = await completeDeploymentRun({ - database: bindings.DB, - deployment: context.deployment, - result, - run: context.run, - }); + await budget.run("App deployment final ownership check", () => authority.requireOwnership()); + return { + ...result, + deploymentId: context.deployment.id, + kind: "succeeded", + runId: context.run.id, + }; + } catch (error) { + const terminalFailure = toAppDeploymentTerminalFailure(input.appDeploymentRunId, error); - if (!completed) { - await failDeploymentRunIfActive({ - database: bindings.DB, - errorCode: "deployment_completion_lost", - errorMessage: "Deployment completed externally but the App deployment row changed.", - runId: context.run.id, - }); + if (terminalFailure === null) { + throw error; } + + await authority.requireOwnership(); + return terminalFailure; } finally { - if (externallyAttemptedDeployment !== null) { - try { - await compensateDeletedDeploymentResources({ - bindings, - cloudflareClient, - deployment: externallyAttemptedDeployment, - }); - } catch (error) { - logError("app-deployment.cloudflare_delete_after_deletion_check_failed", { - ...createErrorLogContext(error), - deploymentId: externallyAttemptedDeployment.id, - runId: input.appDeploymentRunId, - }); - } + try { + await runner.cleanup?.(); + } catch (error) { + logError("app-deployment.sandbox_cleanup_failed", { + ...createErrorLogContext(error), + runId: input.appDeploymentRunId, + }); } - - await runner.cleanup?.(); } } diff --git a/apps/api/src/modules/apps/application/app-deployment-gateway.ts b/apps/api/src/modules/apps/application/app-deployment-gateway.ts new file mode 100644 index 00000000..85cacb0c --- /dev/null +++ b/apps/api/src/modules/apps/application/app-deployment-gateway.ts @@ -0,0 +1,150 @@ +import { PLATFORM_ID_PATTERN } from "@mosoo/id"; + +import { createErrorLogContext, logError } from "../../../platform/cloudflare/logger"; + +export const APP_DEPLOYMENT_PROBE_HOST_LABEL = "probe"; +export const APP_DEPLOYMENT_PROBE_PATH_PREFIX = "/.well-known/mosoo-wfp-probe/"; +export const APP_DEPLOYMENT_PROBE_SCRIPT_NAME = "mosoo-wfp-probe"; + +interface AppDeploymentUserWorker { + fetch(request: Request): Promise | Response; +} + +export interface AppDeploymentDispatchNamespace { + get(scriptName: string): AppDeploymentUserWorker; +} + +export interface AppDeploymentGatewayOptions { + appDeploymentDomain: string; + dispatcher: AppDeploymentDispatchNamespace | undefined; + resolveActiveScriptName(mosooSubdomain: string): Promise; +} + +export async function resolveActiveAppDeploymentScriptName( + database: D1Database, + mosooSubdomain: string, +): Promise { + const deployment = await database + .prepare( + `SELECT active_script_name AS activeScriptName + FROM app_deployment + WHERE mosoo_subdomain = ? + AND deleted_at IS NULL + AND active_script_name IS NOT NULL + LIMIT 1`, + ) + .bind(mosooSubdomain) + .first<{ activeScriptName: string }>(); + + return deployment?.activeScriptName ?? null; +} + +const platformIdPattern = new RegExp(PLATFORM_ID_PATTERN, "u"); + +function normalizeHostname(hostname: string): string { + return hostname.toLowerCase().replace(/\.$/u, ""); +} + +function readDeploymentSubdomain(hostname: string, domain: string): string | null { + const normalizedDomain = normalizeHostname(domain); + const normalizedHostname = normalizeHostname(hostname); + const suffix = `.${normalizedDomain}`; + + if (normalizedDomain.length === 0 || !normalizedHostname.endsWith(suffix)) { + return null; + } + + return normalizedHostname.slice(0, -suffix.length); +} + +export function parseAppDeploymentSubdomain(hostname: string, domain: string): string | null { + const subdomain = readDeploymentSubdomain(hostname, domain); + + if ( + subdomain === null || + subdomain.includes(".") || + !subdomain.startsWith("app-") || + !platformIdPattern.test(subdomain.slice(4).toUpperCase()) + ) { + return null; + } + + return subdomain; +} + +export async function dispatchAppDeploymentGatewayRequest( + request: Request, + options: AppDeploymentGatewayOptions, +): Promise { + const url = new URL(request.url); + const hostname = url.hostname; + const domainSubdomain = readDeploymentSubdomain(hostname, options.appDeploymentDomain); + + if (domainSubdomain === null) { + return null; + } + + const isProbe = + domainSubdomain === APP_DEPLOYMENT_PROBE_HOST_LABEL && + url.pathname.startsWith(APP_DEPLOYMENT_PROBE_PATH_PREFIX); + const mosooSubdomain = isProbe + ? null + : parseAppDeploymentSubdomain(hostname, options.appDeploymentDomain); + + if (!isProbe && mosooSubdomain === null) { + return notFound(); + } + + if (options.dispatcher === undefined) { + logError("app-deployment.gateway_dispatcher_missing", { + hostname, + mosooSubdomain, + }); + return serviceUnavailable(); + } + + let activeScriptName: string | null = APP_DEPLOYMENT_PROBE_SCRIPT_NAME; + + if (mosooSubdomain !== null) { + try { + activeScriptName = await options.resolveActiveScriptName(mosooSubdomain); + } catch (error) { + logError("app-deployment.gateway_lookup_failed", { + ...createErrorLogContext(error), + hostname, + mosooSubdomain, + }); + return serviceUnavailable(); + } + + if (activeScriptName === null) { + return notFound(); + } + } + + try { + return await options.dispatcher.get(activeScriptName).fetch(request); + } catch (error) { + logError("app-deployment.gateway_dispatch_failed", { + ...createErrorLogContext(error), + activeScriptName, + hostname, + mosooSubdomain, + }); + return serviceUnavailable(); + } +} + +function notFound(): Response { + return new Response(null, { + headers: { "Cache-Control": "no-store" }, + status: 404, + }); +} + +function serviceUnavailable(): Response { + return new Response(null, { + headers: { "Cache-Control": "no-store" }, + status: 503, + }); +} diff --git a/apps/api/src/modules/apps/application/app-deployment-script-reconciliation.service.ts b/apps/api/src/modules/apps/application/app-deployment-script-reconciliation.service.ts new file mode 100644 index 00000000..04e8efba --- /dev/null +++ b/apps/api/src/modules/apps/application/app-deployment-script-reconciliation.service.ts @@ -0,0 +1,629 @@ +import { + apiCommandsTable, + appDeploymentRunsTable, + appDeploymentScriptsTable, + appDeploymentsTable, +} from "@mosoo/db"; +import type { AppDeploymentScriptRow } from "@mosoo/db"; +import type { AppDeploymentId, AppDeploymentRunId } from "@mosoo/id"; +import { and, asc, eq, exists, gt, inArray, isNotNull, isNull, lte, or, sql } from "drizzle-orm"; + +import { getAppDatabase, getD1ChangeCount } from "../../../platform/db/drizzle"; +import type { ApiCommandClaimAuthority } from "../../api-command/application/api-command-ledger"; +import { + API_COMMAND_LEASE_MS, + exactApiCommandClaimPredicate, +} from "../../api-command/application/api-command-ledger"; +import { ACTIVE_APP_DEPLOYMENT_RUN_STATUSES } from "../domain/app-deployment-lifecycle"; + +export const APP_DEPLOYMENT_SCRIPT_GRACE_MS = 24 * 60 * 60_000; + +// ponytail: Keep one external cleanup per Queue command until measured backlog proves +// that bounded parallelism is needed; four remote effects cannot fit a serial 64-row page. +const RECONCILIATION_PAGE_SIZE = 1; +const TOMBSTONE_BACKOFF_DAYS = [1, 2, 4, 8, 16, 30] as const; +const DATABASE_NOW_MS = sql`CAST(unixepoch('subsec') * 1000 AS INTEGER)`; + +function databaseDeadlineMs(delayMs: number) { + return sql`min(${DATABASE_NOW_MS} + ${delayMs}, ${Number.MAX_SAFE_INTEGER})`; +} + +export interface AppDeploymentScriptCandidateAuthority extends ApiCommandClaimAuthority { + readonly requireOwnership: () => Promise; +} + +export interface AppDeploymentScriptCleanupAction { + readonly artifactPrefix: string; + readonly attemptCount: number; + readonly buildSandboxId: string; + readonly commandId: AppDeploymentScriptRow["commandId"]; + readonly deliveryGeneration: number; + readonly deploySandboxId: string; + readonly deploymentId: AppDeploymentId; + readonly reconcileCount: number; + readonly runId: AppDeploymentRunId; + readonly scriptName: string; + readonly uploadStartedAt: number | null; +} + +export interface AppDeploymentScriptReconciliationFailure { + readonly error: unknown; + readonly scriptName: string; +} + +export interface AppDeploymentScriptReconciliationResult { + readonly armed: number; + readonly cleaned: number; + readonly deferred: number; + readonly failures: readonly AppDeploymentScriptReconciliationFailure[]; + readonly hasMore: boolean; + readonly nextCursor: string | null; + readonly processed: number; +} + +export type AppDeploymentScriptReconciliationAuthority = AppDeploymentScriptCandidateAuthority; + +type ReconciliationDisposition = + | { readonly kind: "armed" } + | { readonly kind: "deferred" } + | { + readonly action: AppDeploymentScriptCleanupAction; + readonly kind: "cleanup"; + readonly reconcileOwner: string; + }; + +function exactAppDeploymentCommand( + database: D1Database, + authority: ApiCommandClaimAuthority, + nowMs: number, + runId: AppDeploymentRunId, +) { + return exists( + getAppDatabase(database) + .select({ id: apiCommandsTable.id }) + .from(apiCommandsTable) + .where( + and( + exactApiCommandClaimPredicate(authority, nowMs), + eq(apiCommandsTable.kind, "app_deployment_run_dispatch"), + gt(apiCommandsTable.claimExpiresAt, DATABASE_NOW_MS), + sql`json_valid(${apiCommandsTable.payloadJson}) = 1`, + sql`json_extract(${apiCommandsTable.payloadJson}, '$.appDeploymentRunId') = ${runId}`, + ), + ), + ); +} + +function exactReconciliationCommand(database: D1Database, authority: ApiCommandClaimAuthority) { + return exists( + getAppDatabase(database) + .select({ id: apiCommandsTable.id }) + .from(apiCommandsTable) + .where( + and( + eq(apiCommandsTable.id, authority.commandId), + eq(apiCommandsTable.deliveryGeneration, authority.deliveryGeneration), + eq(apiCommandsTable.attemptCount, authority.attemptCount), + eq(apiCommandsTable.status, "running"), + eq(apiCommandsTable.claimOwner, authority.claimOwner), + eq(apiCommandsTable.kind, "app_deployment_script_reconciliation"), + gt(apiCommandsTable.claimExpiresAt, DATABASE_NOW_MS), + ), + ), + ); +} + +function currentDeploymentRun( + database: D1Database, + deploymentId: AppDeploymentId, + runId: AppDeploymentRunId, +) { + return exists( + getAppDatabase(database) + .select({ id: appDeploymentsTable.id }) + .from(appDeploymentsTable) + .where( + and( + eq(appDeploymentsTable.id, deploymentId), + eq(appDeploymentsTable.latestRunId, runId), + isNull(appDeploymentsTable.deletedAt), + ), + ), + ); +} + +function scriptIsActive(database: D1Database, scriptName: string) { + return exists( + getAppDatabase(database) + .select({ id: appDeploymentsTable.id }) + .from(appDeploymentsTable) + .where(eq(appDeploymentsTable.activeScriptName, scriptName)), + ); +} + +function scriptHasTrafficAuthority(database: D1Database, scriptName: string) { + return or( + scriptIsActive(database, scriptName), + exists( + getAppDatabase(database) + .select({ scriptName: appDeploymentScriptsTable.scriptName }) + .from(appDeploymentScriptsTable) + .innerJoin( + appDeploymentRunsTable, + and( + eq(appDeploymentRunsTable.id, appDeploymentScriptsTable.runId), + eq(appDeploymentRunsTable.deploymentId, appDeploymentScriptsTable.deploymentId), + eq(appDeploymentRunsTable.targetScriptName, appDeploymentScriptsTable.scriptName), + ), + ) + .innerJoin( + appDeploymentsTable, + and( + eq(appDeploymentsTable.id, appDeploymentScriptsTable.deploymentId), + eq(appDeploymentsTable.latestRunId, appDeploymentScriptsTable.runId), + isNull(appDeploymentsTable.deletedAt), + ), + ) + .innerJoin(apiCommandsTable, eq(apiCommandsTable.id, appDeploymentScriptsTable.commandId)) + .where( + and( + eq(appDeploymentScriptsTable.scriptName, scriptName), + isNull(appDeploymentScriptsTable.retireAfter), + inArray(appDeploymentRunsTable.status, ACTIVE_APP_DEPLOYMENT_RUN_STATUSES), + eq(apiCommandsTable.kind, "app_deployment_run_dispatch"), + eq(apiCommandsTable.status, "running"), + eq(apiCommandsTable.deliveryGeneration, appDeploymentScriptsTable.deliveryGeneration), + eq(apiCommandsTable.attemptCount, appDeploymentScriptsTable.attemptCount), + eq(apiCommandsTable.claimOwner, appDeploymentScriptsTable.registeredClaimOwner), + gt(apiCommandsTable.claimExpiresAt, DATABASE_NOW_MS), + sql`json_valid(${apiCommandsTable.payloadJson}) = 1`, + sql`json_extract(${apiCommandsTable.payloadJson}, '$.appDeploymentRunId') = ${appDeploymentScriptsTable.runId}`, + ), + ), + ), + ); +} + +function attemptSuffix(attempt: { attemptCount: number; deliveryGeneration: number }): string { + return `g${attempt.deliveryGeneration.toString(36)}-a${attempt.attemptCount.toString(36)}`; +} + +function cleanupAction( + row: Omit< + AppDeploymentScriptCleanupAction, + "artifactPrefix" | "buildSandboxId" | "deploySandboxId" + >, +): AppDeploymentScriptCleanupAction { + const suffix = attemptSuffix(row); + return { + ...row, + artifactPrefix: `app-deployments/${row.scriptName}/`, + buildSandboxId: `${row.runId}-${suffix}-build`, + deploySandboxId: `${row.runId}-${suffix}-deploy`, + }; +} + +export async function registerAppDeploymentScriptCandidate( + database: D1Database, + input: { + readonly authority: AppDeploymentScriptCandidateAuthority; + readonly deploymentId: AppDeploymentId; + readonly nowMs: number; + readonly runId: AppDeploymentRunId; + readonly scriptName: string; + }, +): Promise { + await input.authority.requireOwnership(); + + const row: AppDeploymentScriptRow = { + attemptCount: input.authority.attemptCount, + commandId: input.authority.commandId, + deliveryGeneration: input.authority.deliveryGeneration, + deploymentId: input.deploymentId, + externalDeletedAt: null, + lastReconciledAt: null, + nextReconcileAt: input.nowMs + APP_DEPLOYMENT_SCRIPT_GRACE_MS, + reconcileCount: 0, + reconcileExpiresAt: null, + reconcileOwner: null, + registeredAt: input.nowMs, + registeredClaimOwner: input.authority.claimOwner, + retireAfter: null, + runId: input.runId, + scriptName: input.scriptName, + uploadStartedAt: null, + }; + + await getAppDatabase(database).insert(appDeploymentScriptsTable).values(row).run(); + + const registered = await getAppDatabase(database) + .select() + .from(appDeploymentScriptsTable) + .where(eq(appDeploymentScriptsTable.scriptName, input.scriptName)) + .limit(1) + .get(); + + if ( + registered === undefined || + registered.commandId !== row.commandId || + registered.deliveryGeneration !== row.deliveryGeneration || + registered.attemptCount !== row.attemptCount || + registered.registeredClaimOwner !== row.registeredClaimOwner || + registered.deploymentId !== row.deploymentId || + registered.runId !== row.runId || + registered.retireAfter !== null + ) { + throw new Error( + "App deployment script candidate registration conflicts with durable authority.", + ); + } + + return registered; +} + +export async function markAppDeploymentScriptUploadStarted( + database: D1Database, + input: { + readonly authority: AppDeploymentScriptCandidateAuthority; + readonly deploymentId: AppDeploymentId; + readonly nowMs: number; + readonly runId: AppDeploymentRunId; + readonly scriptName: string; + }, +): Promise { + await input.authority.requireOwnership(); + + const result = await getAppDatabase(database) + .update(appDeploymentScriptsTable) + .set({ + uploadStartedAt: sql`coalesce(${appDeploymentScriptsTable.uploadStartedAt}, ${input.nowMs})`, + }) + .where( + and( + eq(appDeploymentScriptsTable.scriptName, input.scriptName), + eq(appDeploymentScriptsTable.deploymentId, input.deploymentId), + eq(appDeploymentScriptsTable.runId, input.runId), + eq(appDeploymentScriptsTable.commandId, input.authority.commandId), + eq(appDeploymentScriptsTable.deliveryGeneration, input.authority.deliveryGeneration), + eq(appDeploymentScriptsTable.attemptCount, input.authority.attemptCount), + eq(appDeploymentScriptsTable.registeredClaimOwner, input.authority.claimOwner), + isNull(appDeploymentScriptsTable.retireAfter), + exactAppDeploymentCommand(database, input.authority, input.nowMs, input.runId), + currentDeploymentRun(database, input.deploymentId, input.runId), + exists( + getAppDatabase(database) + .select({ id: appDeploymentRunsTable.id }) + .from(appDeploymentRunsTable) + .where( + and( + eq(appDeploymentRunsTable.id, input.runId), + eq(appDeploymentRunsTable.targetScriptName, input.scriptName), + inArray(appDeploymentRunsTable.status, ACTIVE_APP_DEPLOYMENT_RUN_STATUSES), + ), + ), + ), + ), + ) + .run(); + + if (getD1ChangeCount(result) !== 1) { + throw new Error("App deployment script upload lost durable authority."); + } +} + +function decodeCursor( + cursor: string | null, +): { nextReconcileAt: number; scriptName: string } | null { + if (cursor === null) { + return null; + } + + const separator = cursor.indexOf(":"); + const nextReconcileAt = Number(cursor.slice(0, separator)); + const scriptName = cursor.slice(separator + 1); + + if ( + separator < 1 || + !Number.isSafeInteger(nextReconcileAt) || + nextReconcileAt < 0 || + !/^app-[0-9a-z-]+$/u.test(scriptName) + ) { + throw new Error("App deployment script reconciliation cursor is invalid."); + } + + return { nextReconcileAt, scriptName }; +} + +function encodeCursor(row: { nextReconcileAt: number | null; scriptName: string }): string { + if (row.nextReconcileAt === null) { + throw new Error("App deployment script reconciliation row has no cursor timestamp."); + } + return `${row.nextReconcileAt}:${row.scriptName}`; +} + +async function prepareReconciliation( + database: D1Database, + input: { + readonly authority: AppDeploymentScriptReconciliationAuthority; + readonly scriptName: string; + }, +): Promise { + const db = getAppDatabase(database); + const exactCommand = exactReconciliationCommand(database, input.authority); + const due = and( + eq(appDeploymentScriptsTable.scriptName, input.scriptName), + isNotNull(appDeploymentScriptsTable.nextReconcileAt), + lte(appDeploymentScriptsTable.nextReconcileAt, DATABASE_NOW_MS), + exactCommand, + ); + const protectedScript = scriptHasTrafficAuthority(database, input.scriptName); + const activeScript = scriptIsActive(database, input.scriptName); + + const protectedResult = await db + .update(appDeploymentScriptsTable) + .set({ + lastReconciledAt: DATABASE_NOW_MS, + nextReconcileAt: sql`CASE + WHEN ${activeScript} THEN NULL + ELSE ${databaseDeadlineMs(APP_DEPLOYMENT_SCRIPT_GRACE_MS)} + END`, + reconcileExpiresAt: null, + reconcileOwner: null, + }) + .where(and(due, protectedScript)) + .run(); + if (getD1ChangeCount(protectedResult) === 1) { + return { kind: "deferred" }; + } + + const retireAfter = sql`CAST(unixepoch('subsec') * 1000 AS INTEGER) + ${APP_DEPLOYMENT_SCRIPT_GRACE_MS}`; + const armedResult = await db + .update(appDeploymentScriptsTable) + .set({ + lastReconciledAt: DATABASE_NOW_MS, + nextReconcileAt: retireAfter, + reconcileExpiresAt: null, + reconcileOwner: null, + retireAfter, + }) + .where(and(due, isNull(appDeploymentScriptsTable.retireAfter), sql`NOT ${protectedScript}`)) + .run(); + if (getD1ChangeCount(armedResult) === 1) { + return { kind: "armed" }; + } + + const graceResult = await db + .update(appDeploymentScriptsTable) + .set({ nextReconcileAt: appDeploymentScriptsTable.retireAfter }) + .where( + and( + due, + isNotNull(appDeploymentScriptsTable.retireAfter), + gt(appDeploymentScriptsTable.retireAfter, DATABASE_NOW_MS), + ), + ) + .run(); + if (getD1ChangeCount(graceResult) === 1) { + return { kind: "deferred" }; + } + + const reconcileOwner = `${input.authority.commandId}:${crypto.randomUUID()}`; + const claimResult = await db + .update(appDeploymentScriptsTable) + .set({ + reconcileExpiresAt: databaseDeadlineMs(API_COMMAND_LEASE_MS), + reconcileOwner, + }) + .where( + and( + due, + isNotNull(appDeploymentScriptsTable.retireAfter), + lte(appDeploymentScriptsTable.retireAfter, DATABASE_NOW_MS), + sql`NOT ${protectedScript}`, + or( + isNull(appDeploymentScriptsTable.reconcileOwner), + lte(appDeploymentScriptsTable.reconcileExpiresAt, DATABASE_NOW_MS), + ), + ), + ) + .run(); + if (getD1ChangeCount(claimResult) !== 1) { + return { kind: "deferred" }; + } + + const row = await db + .select({ + attemptCount: appDeploymentScriptsTable.attemptCount, + commandId: appDeploymentScriptsTable.commandId, + deliveryGeneration: appDeploymentScriptsTable.deliveryGeneration, + deploymentId: appDeploymentScriptsTable.deploymentId, + reconcileCount: appDeploymentScriptsTable.reconcileCount, + runId: appDeploymentScriptsTable.runId, + scriptName: appDeploymentScriptsTable.scriptName, + uploadStartedAt: appDeploymentScriptsTable.uploadStartedAt, + }) + .from(appDeploymentScriptsTable) + .where( + and( + eq(appDeploymentScriptsTable.scriptName, input.scriptName), + eq(appDeploymentScriptsTable.reconcileOwner, reconcileOwner), + ), + ) + .limit(1) + .get(); + if (row === undefined) { + throw new Error("Claimed App deployment script reconciliation row disappeared."); + } + + return { action: cleanupAction(row), kind: "cleanup", reconcileOwner }; +} + +function tombstoneBackoffMs(reconcileCount: number): number { + const index = Math.min(reconcileCount, TOMBSTONE_BACKOFF_DAYS.length - 1); + return (TOMBSTONE_BACKOFF_DAYS[index] ?? 30) * APP_DEPLOYMENT_SCRIPT_GRACE_MS; +} + +async function completeReconciliation( + database: D1Database, + input: { + readonly action: AppDeploymentScriptCleanupAction; + readonly authority: AppDeploymentScriptReconciliationAuthority; + readonly reconcileOwner: string; + }, +): Promise { + const result = await getAppDatabase(database) + .update(appDeploymentScriptsTable) + .set({ + externalDeletedAt: + input.action.uploadStartedAt === null + ? appDeploymentScriptsTable.externalDeletedAt + : sql`coalesce(${appDeploymentScriptsTable.externalDeletedAt}, ${DATABASE_NOW_MS})`, + lastReconciledAt: DATABASE_NOW_MS, + nextReconcileAt: databaseDeadlineMs(tombstoneBackoffMs(input.action.reconcileCount)), + reconcileCount: sql`min(${appDeploymentScriptsTable.reconcileCount} + 1, 9007199254740991)`, + reconcileExpiresAt: null, + reconcileOwner: null, + }) + .where( + and( + eq(appDeploymentScriptsTable.scriptName, input.action.scriptName), + eq(appDeploymentScriptsTable.reconcileOwner, input.reconcileOwner), + gt(appDeploymentScriptsTable.reconcileExpiresAt, DATABASE_NOW_MS), + isNotNull(appDeploymentScriptsTable.retireAfter), + sql`NOT ${scriptHasTrafficAuthority(database, input.action.scriptName)}`, + exactReconciliationCommand(database, input.authority), + ), + ) + .run(); + + return getD1ChangeCount(result) === 1; +} + +async function failReconciliation( + database: D1Database, + input: { + readonly authority: AppDeploymentScriptReconciliationAuthority; + readonly reconcileOwner: string; + readonly scriptName: string; + }, +): Promise { + const result = await getAppDatabase(database) + .update(appDeploymentScriptsTable) + .set({ + lastReconciledAt: DATABASE_NOW_MS, + nextReconcileAt: databaseDeadlineMs(APP_DEPLOYMENT_SCRIPT_GRACE_MS), + reconcileExpiresAt: null, + reconcileOwner: null, + }) + .where( + and( + eq(appDeploymentScriptsTable.scriptName, input.scriptName), + eq(appDeploymentScriptsTable.reconcileOwner, input.reconcileOwner), + gt(appDeploymentScriptsTable.reconcileExpiresAt, DATABASE_NOW_MS), + exactReconciliationCommand(database, input.authority), + ), + ) + .run(); + + return getD1ChangeCount(result) === 1; +} + +export async function reconcileAppDeploymentScriptPage( + database: D1Database, + input: { + readonly authority: AppDeploymentScriptReconciliationAuthority; + readonly cursor: string | null; + }, + cleanup: (action: AppDeploymentScriptCleanupAction) => Promise, +): Promise { + const cursor = decodeCursor(input.cursor); + await input.authority.requireOwnership(); + + const rows = await getAppDatabase(database) + .select({ + nextReconcileAt: appDeploymentScriptsTable.nextReconcileAt, + scriptName: appDeploymentScriptsTable.scriptName, + }) + .from(appDeploymentScriptsTable) + .where( + and( + isNotNull(appDeploymentScriptsTable.nextReconcileAt), + lte(appDeploymentScriptsTable.nextReconcileAt, DATABASE_NOW_MS), + cursor === null + ? undefined + : or( + gt(appDeploymentScriptsTable.nextReconcileAt, cursor.nextReconcileAt), + and( + eq(appDeploymentScriptsTable.nextReconcileAt, cursor.nextReconcileAt), + gt(appDeploymentScriptsTable.scriptName, cursor.scriptName), + ), + ), + ), + ) + .orderBy( + asc(appDeploymentScriptsTable.nextReconcileAt), + asc(appDeploymentScriptsTable.scriptName), + ) + .limit(RECONCILIATION_PAGE_SIZE + 1) + .all(); + + const page = rows.slice(0, RECONCILIATION_PAGE_SIZE); + const failures: AppDeploymentScriptReconciliationFailure[] = []; + let armed = 0; + let cleaned = 0; + let deferred = 0; + + for (const row of page) { + await input.authority.requireOwnership(); + const disposition = await prepareReconciliation(database, { + authority: input.authority, + scriptName: row.scriptName, + }); + + if (disposition.kind === "armed") { + armed += 1; + continue; + } + if (disposition.kind === "deferred") { + deferred += 1; + continue; + } + + try { + await cleanup(disposition.action); + await input.authority.requireOwnership(); + const completed = await completeReconciliation(database, { + action: disposition.action, + authority: input.authority, + reconcileOwner: disposition.reconcileOwner, + }); + if (!completed) { + throw new Error("App deployment script reconciliation lost durable authority."); + } + cleaned += 1; + } catch (error) { + await input.authority.requireOwnership(); + const released = await failReconciliation(database, { + authority: input.authority, + reconcileOwner: disposition.reconcileOwner, + scriptName: disposition.action.scriptName, + }); + if (!released) { + throw error; + } + failures.push({ error, scriptName: disposition.action.scriptName }); + } + } + + const hasMore = rows.length > RECONCILIATION_PAGE_SIZE; + const lastRow = page.at(-1); + return { + armed, + cleaned, + deferred, + failures, + hasMore, + nextCursor: hasMore && lastRow !== undefined ? encodeCursor(lastRow) : null, + processed: page.length, + }; +} diff --git a/apps/api/src/modules/apps/application/app-deployment.service.ts b/apps/api/src/modules/apps/application/app-deployment.service.ts index b87cc9fa..61a75477 100644 --- a/apps/api/src/modules/apps/application/app-deployment.service.ts +++ b/apps/api/src/modules/apps/application/app-deployment.service.ts @@ -4,7 +4,7 @@ import type { DeleteAppDeploymentInput, DeployAppInput, } from "@mosoo/contracts/app"; -import type { ApiCommandId, AppDeploymentRunRow, AppDeploymentRow } from "@mosoo/db"; +import type { AppDeploymentRunRow, AppDeploymentRow } from "@mosoo/db"; import { apiCommandsTable, appDeploymentRunsTable, appDeploymentsTable } from "@mosoo/db"; import type { AppDeploymentId, AppDeploymentRunId, AppId } from "@mosoo/id"; import { createPlatformId } from "@mosoo/id"; @@ -12,42 +12,25 @@ import { and, desc, eq, inArray, isNull } from "drizzle-orm"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; import { getAppDatabase, getD1ChangeCount } from "../../../platform/db/drizzle"; -import { API_ERROR_CODE, createApiError, validationError } from "../../../platform/errors"; +import { validationError } from "../../../platform/errors"; import { currentTimestampMs, toIsoString } from "../../../time"; import { createAppDeploymentRunDispatchDedupeKey, enqueueAppDeploymentRunDispatchCommand, } from "../../api-command/application/api-command-enqueue"; import { API_COMMAND_LEASE_MS } from "../../api-command/application/api-command-ledger"; -import { - APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS, - APP_DEPLOYMENT_RUN_DISPATCH_RETRY_EXHAUSTED_CODE, - createAppDeploymentDispatchRetryExhaustedMessage, -} from "../../api-command/application/api-command-policy"; import type { AuthenticatedViewer } from "../../auth/application/viewer-auth.service"; import { ACTIVE_APP_DEPLOYMENT_RUN_STATUSES } from "../domain/app-deployment-lifecycle"; -import { - createCloudflareDeploymentClient, - deleteCloudflareDeploymentResources, - logCloudflareDeploymentResourceDeleteFailures, -} from "./app-deployment-cloudflare-client"; -import type { - CloudflareClientBindings, - CloudflareDeploymentClient, -} from "./app-deployment-cloudflare-client"; +import type { AppDeploymentAttemptIdentity } from "./app-deployment-executor.service"; import { destroyAppDeploymentRunSandboxesBestEffort } from "./app-deployment-executor.service"; import { ensureAppOwnership } from "./app.service"; import { normalizeLimit } from "./normalize-limit"; type AppDeploymentBindings = Pick< ApiBindings, - "API_COMMAND_QUEUE" | "DB" | "MOSOO_APP_DEPLOYMENT_DOMAIN" + "API_COMMAND_QUEUE" | "APP_DEPLOYMENT_WORKFLOW" | "DB" | "MOSOO_APP_DEPLOYMENT_DOMAIN" >; -type AppDeploymentDeleteBindings = Pick< - AppDeploymentBindings, - "DB" | "MOSOO_APP_DEPLOYMENT_DOMAIN" -> & - CloudflareClientBindings & +type AppDeploymentDeleteBindings = Pick & Partial>; export type AppDeploymentReadBindings = Pick< @@ -56,7 +39,6 @@ export type AppDeploymentReadBindings = Pick< >; interface AppDeploymentServiceOptions { - cloudflareClient?: CloudflareDeploymentClient; fetch?: typeof fetch; nowMs?: () => number; } @@ -75,7 +57,7 @@ export async function readAppDeploymentForOwnedApp( return null; } - await recoverStaleActiveDeploymentRun(bindings.DB, appId); + await recoverActiveDeploymentRun(bindings.DB, appId); const latestRun = await readLatestDeploymentRun(bindings.DB, appId); @@ -97,7 +79,7 @@ export async function getAppDeploymentStatus( appId: AppId, ): Promise { await ensureAppOwnership(bindings.DB, viewer.id, appId); - await recoverStaleActiveDeploymentRun(bindings.DB, appId); + await recoverActiveDeploymentRun(bindings.DB, appId); const run = await readLatestDeploymentRun(bindings.DB, appId); @@ -117,7 +99,7 @@ export async function listAppDeploymentRuns( limit?: number | null, ): Promise { await ensureAppOwnership(bindings.DB, viewer.id, appId); - await recoverStaleActiveDeploymentRun(bindings.DB, appId); + await recoverActiveDeploymentRun(bindings.DB, appId); const runLimit = normalizeLimit(limit, "limit", RUN_LIST_LIMITS); const runs = await getAppDatabase(bindings.DB) @@ -159,9 +141,7 @@ export async function deployApp( input.repoUrl, options.fetch ?? globalThis.fetch, ); - const activeRun = await readActiveDeploymentRun(bindings.DB, input.appId, { - recoverMissingDispatch: true, - }); + const activeRun = await recoverActiveDeploymentRun(bindings.DB, input.appId); if (activeRun !== null) { throw validationError("An App deployment run is already active."); @@ -172,6 +152,7 @@ export async function deployApp( const deployment = existingDeployment ?? ({ + activeScriptName: null, appId: input.appId, createdAt: nowMs, defaultBranch: repository.defaultBranch, @@ -199,18 +180,6 @@ export async function deployApp( if (getD1ChangeCount(insertDeploymentResult) === 0) { throw validationError("An App deployment is already active."); } - } else { - await getAppDatabase(bindings.DB) - .update(appDeploymentsTable) - .set({ - defaultBranch: repository.defaultBranch, - repoName: repository.repoName, - repoOwner: repository.repoOwner, - repoUrl: repository.repoUrl, - updatedAt: nowMs, - }) - .where(eq(appDeploymentsTable.id, deployment.id)) - .run(); } const insertRunResult = await getAppDatabase(bindings.DB) @@ -249,7 +218,14 @@ export async function deployApp( try { linkRunResult = await getAppDatabase(bindings.DB) .update(appDeploymentsTable) - .set({ latestRunId: runId, updatedAt: nowMs }) + .set({ + defaultBranch: repository.defaultBranch, + latestRunId: runId, + repoName: repository.repoName, + repoOwner: repository.repoOwner, + repoUrl: repository.repoUrl, + updatedAt: nowMs, + }) .where(and(eq(appDeploymentsTable.id, deployment.id), isNull(appDeploymentsTable.deletedAt))) .run(); } catch (error) { @@ -326,70 +302,79 @@ export async function deleteAppDeployment( return { ok: true }; } - const activeRunIds = await readActiveDeploymentRunIds(bindings.DB, input.appId); - const nowMs = currentTimestampMs(); - - await getAppDatabase(bindings.DB) - .update(appDeploymentRunsTable) - .set({ - errorCode: "deployment_deleted", - errorMessage: "Deployment was deleted.", - status: "failed", - updatedAt: nowMs, - }) - .where( - and( - eq(appDeploymentRunsTable.appId, input.appId), - inArray(appDeploymentRunsTable.status, ACTIVE_APP_DEPLOYMENT_RUN_STATUSES), - ), - ) - .run(); - - await destroyActiveDeploymentRunSandboxes(bindings, activeRunIds); - - const deleteFailures = await deleteCloudflareDeploymentResources( - options.cloudflareClient ?? createCloudflareDeploymentClient(bindings), - { - hostname: createPlannedHost(deployment.mosooSubdomain, bindings.MOSOO_APP_DEPLOYMENT_DOMAIN), - resourceName: deployment.mosooSubdomain, - }, - ); + const activeRunAttempts = await readActiveDeploymentRunAttempts(bindings.DB, deployment.id); + const nowMs = options.nowMs?.() ?? currentTimestampMs(); - if (deleteFailures.length > 0) { - logCloudflareDeploymentResourceDeleteFailures( - "app-deployment.cloudflare_delete_failed", - deleteFailures, - ); - throw createApiError( - API_ERROR_CODE.appDeploymentCleanupFailed, - "Cloudflare deployment cleanup failed. Retry deletion.", - ); - } + await tombstoneAppDeployment(bindings.DB, deployment.id, nowMs); - await getAppDatabase(bindings.DB) - .update(appDeploymentsTable) - .set({ - deletedAt: nowMs, - lastSuccessfulUrl: null, - updatedAt: nowMs, - }) - .where(eq(appDeploymentsTable.id, deployment.id)) - .run(); + await destroyActiveDeploymentRunSandboxes(bindings, activeRunAttempts); return { ok: true }; } +async function tombstoneAppDeployment( + database: D1Database, + deploymentId: AppDeploymentId, + nowMs: number, +): Promise { + await database.batch([ + database + .prepare( + `UPDATE app_deployment + SET active_script_name = NULL, + deleted_at = ?, + last_successful_url = NULL, + updated_at = ? + WHERE id = ? AND deleted_at IS NULL`, + ) + .bind(nowMs, nowMs, deploymentId), + database + .prepare( + `UPDATE api_command + SET claim_expires_at = NULL, + claim_owner = NULL, + completed_at = ?, + last_error_code = 'deployment_deleted', + last_error_message = 'Deployment was deleted.', + status = 'failed', + updated_at = ? + WHERE kind = 'app_deployment_run_dispatch' + AND status IN ('queued', 'running') + AND EXISTS ( + SELECT 1 + FROM app_deployment_run AS run + WHERE run.deployment_id = ? + AND run.status IN ('queued', 'preparing', 'building', 'submitting', 'submitted', 'activating') + AND json_valid(api_command.payload_json) = 1 + AND json_extract(api_command.payload_json, '$.appDeploymentRunId') = run.id + )`, + ) + .bind(nowMs, nowMs, deploymentId), + database + .prepare( + `UPDATE app_deployment_run + SET error_code = 'deployment_deleted', + error_message = 'Deployment was deleted.', + status = 'failed', + updated_at = ? + WHERE deployment_id = ? + AND status IN ('queued', 'preparing', 'building', 'submitting', 'submitted', 'activating')`, + ) + .bind(nowMs, deploymentId), + ]); +} + async function destroyActiveDeploymentRunSandboxes( bindings: AppDeploymentDeleteBindings, - runIds: readonly AppDeploymentRunId[], + runs: ReadonlyArray<{ attempt: AppDeploymentAttemptIdentity; runId: AppDeploymentRunId }>, ): Promise { if (!hasRuntimeSubjectDestroyBinding(bindings)) { return; } await Promise.all( - runIds.map((runId) => - destroyAppDeploymentRunSandboxesBestEffort(bindings as ApiBindings, runId), + runs.map(({ attempt, runId }) => + destroyAppDeploymentRunSandboxesBestEffort(bindings as ApiBindings, runId, attempt), ), ); } @@ -412,22 +397,38 @@ async function readActiveDeployment( ); } -async function readActiveDeploymentRunIds( +async function readActiveDeploymentRunAttempts( database: D1Database, - appId: AppId, -): Promise { - const rows = await getAppDatabase(database) - .select({ id: appDeploymentRunsTable.id }) - .from(appDeploymentRunsTable) - .where( - and( - eq(appDeploymentRunsTable.appId, appId), - inArray(appDeploymentRunsTable.status, ACTIVE_APP_DEPLOYMENT_RUN_STATUSES), - ), + deploymentId: AppDeploymentId, +): Promise> { + const { results } = await database + .prepare( + `SELECT command.attempt_count AS attemptCount, + command.delivery_generation AS deliveryGeneration, + run.id AS runId + FROM app_deployment_run AS run + JOIN api_command AS command + ON command.kind = 'app_deployment_run_dispatch' + AND json_valid(command.payload_json) = 1 + AND json_extract(command.payload_json, '$.appDeploymentRunId') = run.id + WHERE run.deployment_id = ? + AND run.status IN ('queued', 'preparing', 'building', 'submitting', 'submitted', 'activating') + AND command.attempt_count > 0`, ) - .all(); - - return rows.map((row) => row.id); + .bind(deploymentId) + .all<{ + attemptCount: number; + deliveryGeneration: number; + runId: AppDeploymentRunId; + }>(); + + return results.map((row) => ({ + attempt: { + attemptCount: row.attemptCount, + deliveryGeneration: row.deliveryGeneration, + }, + runId: row.runId, + })); } async function readDeploymentById( @@ -477,10 +478,9 @@ async function readLatestDeploymentRun( ); } -async function readActiveDeploymentRun( +async function recoverActiveDeploymentRun( database: D1Database, appId: AppId, - options: { recoverMissingDispatch?: boolean } = {}, ): Promise | null> { const run = (await getAppDatabase(database) @@ -503,19 +503,10 @@ async function readActiveDeploymentRun( return null; } - if (options.recoverMissingDispatch !== true) { - return run; - } - const nowMs = currentTimestampMs(); const dispatchCommand = (await getAppDatabase(database) .select({ - attemptCount: apiCommandsTable.attemptCount, - claimExpiresAt: apiCommandsTable.claimExpiresAt, - id: apiCommandsTable.id, - lastErrorCode: apiCommandsTable.lastErrorCode, - lastErrorMessage: apiCommandsTable.lastErrorMessage, status: apiCommandsTable.status, }) .from(apiCommandsTable) @@ -523,40 +514,7 @@ async function readActiveDeploymentRun( .limit(1) .get()) ?? null; - const dispatchRetryExhausted = - dispatchCommand !== null && - (dispatchCommand.status === "queued" || dispatchCommand.status === "running") && - dispatchCommand.attemptCount >= APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS && - dispatchCommand.lastErrorCode !== null; - - if (dispatchRetryExhausted) { - const errorMessage = createAppDeploymentDispatchRetryExhaustedMessage({ - attemptCount: dispatchCommand.attemptCount, - lastErrorMessage: dispatchCommand.lastErrorMessage ?? dispatchCommand.lastErrorCode, - }); - - await markDeploymentRunFailed( - database, - run.id, - APP_DEPLOYMENT_RUN_DISPATCH_RETRY_EXHAUSTED_CODE, - new Error(errorMessage), - nowMs, - ); - await markDeploymentDispatchCommandFailed(database, dispatchCommand.id, { - errorCode: APP_DEPLOYMENT_RUN_DISPATCH_RETRY_EXHAUSTED_CODE, - errorMessage, - nowMs, - }); - - return null; - } - - if ( - dispatchCommand?.status === "queued" || - (dispatchCommand?.status === "running" && - dispatchCommand.claimExpiresAt !== null && - dispatchCommand.claimExpiresAt > nowMs) - ) { + if (dispatchCommand?.status === "queued" || dispatchCommand?.status === "running") { return run; } @@ -564,55 +522,17 @@ async function readActiveDeploymentRun( return run; } - const staleDispatch = - dispatchCommand?.status === "running" && - dispatchCommand.claimExpiresAt !== null && - dispatchCommand.claimExpiresAt <= nowMs; - await markDeploymentRunFailed( database, run.id, - staleDispatch ? "deployment_dispatch_expired" : "deployment_dispatch_missing", - new Error( - staleDispatch - ? "Deployment dispatch claim expired before completion." - : "Deployment dispatch command is missing.", - ), + "deployment_dispatch_missing", + new Error("Deployment dispatch command is missing."), nowMs, ); return null; } -async function markDeploymentDispatchCommandFailed( - database: D1Database, - commandId: ApiCommandId, - input: { errorCode: string; errorMessage: string; nowMs: number }, -): Promise { - await getAppDatabase(database) - .update(apiCommandsTable) - .set({ - claimExpiresAt: null, - claimOwner: null, - completedAt: input.nowMs, - lastErrorCode: input.errorCode, - lastErrorMessage: input.errorMessage, - status: "failed", - updatedAt: input.nowMs, - }) - .where( - and( - eq(apiCommandsTable.id, commandId), - inArray(apiCommandsTable.status, ["queued", "running"]), - ), - ) - .run(); -} - -async function recoverStaleActiveDeploymentRun(database: D1Database, appId: AppId): Promise { - await readActiveDeploymentRun(database, appId, { recoverMissingDispatch: true }); -} - async function markDeploymentRunFailed( database: D1Database, runId: AppDeploymentRunId, diff --git a/apps/api/src/modules/auth/application/personal-access-token.service.ts b/apps/api/src/modules/auth/application/personal-access-token.service.ts index fce8273b..fcaa6cef 100644 --- a/apps/api/src/modules/auth/application/personal-access-token.service.ts +++ b/apps/api/src/modules/auth/application/personal-access-token.service.ts @@ -62,7 +62,7 @@ export function isPersonalAccessTokenValue(tokenValue: string): boolean { export async function hashTokenValue(tokenValue: string): Promise { const encoded = new TextEncoder().encode(tokenValue); const digest = await crypto.subtle.digest("SHA-256", encoded); - return [...new Uint8Array(digest)].map((byte) => byte.toString(16).padStart(2, "0")).join(""); + return new Uint8Array(digest).toHex(); } function toTokenSummary(row: PersonalAccessTokenListRow): PersonalAccessTokenSummary { diff --git a/apps/api/src/modules/cost/application/cost-rollup.service.ts b/apps/api/src/modules/cost/application/cost-rollup.service.ts index b8b9b2c7..4b710ba0 100644 --- a/apps/api/src/modules/cost/application/cost-rollup.service.ts +++ b/apps/api/src/modules/cost/application/cost-rollup.service.ts @@ -1,4 +1,9 @@ -import { usageDailyRollupsTable, usageEventRollupReceiptsTable, usageEventsTable } from "@mosoo/db"; +import { + sessionModelCallsTable, + usageDailyRollupsTable, + usageEventRollupReceiptsTable, + usageEventsTable, +} from "@mosoo/db"; import { lt, sql } from "drizzle-orm"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; @@ -55,6 +60,26 @@ export async function runUsageDailyRollup(env: ApiBindings, now = new Date()): P // Drizzle's D1 batch cannot prepare parameterized db.run(sql) queries. const cutoffSql = sql.raw(String(cutoffMs)); const rolledUpAtSql = sql.raw(String(rolledUpAtMs)); + const eligibleUsage = sql` + ${usageEventsTable.createdAt} < ${cutoffSql} + AND ( + ${usageEventsTable.source} <> 'runtime_driver' + OR EXISTS ( + SELECT 1 + FROM ${sessionModelCallsTable} AS rollup_model_call + WHERE rollup_model_call.session_id = ${usageEventsTable.sessionId} + AND rollup_model_call.session_run_id = ${usageEventsTable.sessionRunId} + AND ${usageEventsTable.sourceEventId} = + rollup_model_call.driver_instance_id || ':' || + CASE + WHEN trim(COALESCE(rollup_model_call.native_call_id, '')) <> '' + THEN trim(rollup_model_call.native_call_id) + ELSE rollup_model_call.session_run_id || ':' || rollup_model_call.call_key + END + AND rollup_model_call.source_event_seq >= ${usageEventsTable.sourceEventSeq} + ) + ) + `; await runAppDatabaseBatch(env.DB, (db) => [ db.run(sql` @@ -98,7 +123,7 @@ export async function runUsageDailyRollup(env: ApiBindings, now = new Date()): P SUM(CASE WHEN ${usageEventsTable.pricingStatus} = 'unknown' THEN 1 ELSE 0 END) AS unpriced_request_count FROM ${usageEventsTable} - WHERE ${usageEventsTable.createdAt} < ${cutoffSql} + WHERE ${eligibleUsage} GROUP BY ${usageEventsTable.organizationId}, ${usageEventsTable.appId}, @@ -140,10 +165,10 @@ export async function runUsageDailyRollup(env: ApiBindings, now = new Date()): P ${usageEventsTable.sourceEventId}, ${rolledUpAtSql} FROM ${usageEventsTable} - WHERE ${usageEventsTable.createdAt} < ${cutoffSql} + WHERE ${eligibleUsage} ON CONFLICT(source, source_event_id) DO NOTHING `), - db.delete(usageEventsTable).where(lt(usageEventsTable.createdAt, cutoffMs)), + db.delete(usageEventsTable).where(eligibleUsage), db.delete(usageDailyRollupsTable).where(createDailyRollupRetentionPredicate(now)), db .delete(usageEventRollupReceiptsTable) diff --git a/apps/api/src/modules/cost/application/cost-usage-event.service.ts b/apps/api/src/modules/cost/application/cost-usage-event.service.ts index 344397d6..96c1e0df 100644 --- a/apps/api/src/modules/cost/application/cost-usage-event.service.ts +++ b/apps/api/src/modules/cost/application/cost-usage-event.service.ts @@ -12,7 +12,8 @@ import type { SessionId, SessionRunId, } from "@mosoo/id"; -import { and, eq, sql } from "drizzle-orm"; +import { and, eq, exists, gte, notExists, or, sql } from "drizzle-orm"; +import type { SQL } from "drizzle-orm"; import { getAppDatabase } from "../../../platform/db/drizzle"; import type { AppDatabase } from "../../../platform/db/drizzle"; @@ -46,9 +47,26 @@ export interface RecordRuntimeUsageEventInput { driverInstanceId: DriverInstanceId; nativeCallId: string | null; run: RuntimeUsageRunContext; + sourceEventSeq?: number; usage: SessionUsageSummary; } +function selectedValue(value: Value, alias: string) { + return sql`${value}`.as(alias); +} + +function normalizeSourceEventSeq(value: number | undefined): number { + if (value === undefined) { + return 0; + } + + if (!Number.isSafeInteger(value) || value < 0) { + throw new Error("Runtime usage source event seq must be a non-negative safe integer."); + } + + return value; +} + function toTokenCount(value: number | null | undefined): number { if (typeof value !== "number" || !Number.isFinite(value) || value < 0) { return 0; @@ -57,6 +75,10 @@ function toTokenCount(value: number | null | undefined): number { return Math.round(value); } +function isProvidedTokenCount(value: number | null | undefined): boolean { + return typeof value === "number" && Number.isFinite(value) && value >= 0; +} + function isUsageContract(value: string | null | undefined): value is UsageContract { return ( value === "anthropic_bucketed" || @@ -127,6 +149,17 @@ export function hasRecordableRuntimeUsage(usage: SessionUsageSummary): boolean { ); } +function hasRuntimeUsageMutation(input: RecordRuntimeUsageEventInput): boolean { + return ( + hasRecordableRuntimeUsage(input.usage) || + (input.sourceEventSeq !== undefined && + (isProvidedTokenCount(input.usage.inputTokens) || + isProvidedTokenCount(input.usage.outputTokens) || + isProvidedTokenCount(input.usage.cachedReadTokens) || + isProvidedTokenCount(input.usage.cachedWriteTokens))) + ); +} + const RUNTIME_USAGE_SOURCE = "runtime_driver"; function resolveUsageEventIdentity(input: RecordRuntimeUsageEventInput): { @@ -158,14 +191,79 @@ async function isUsageEventAlreadyRolledUp( return existing.length > 0; } -function createRuntimeUsageEventInsert(database: AppDatabase, input: RecordRuntimeUsageEventInput) { +export async function hasRuntimeUsageEventRollupReceipt( + database: D1Database, + input: RecordRuntimeUsageEventInput, +): Promise { + return isUsageEventAlreadyRolledUp(getAppDatabase(database), resolveUsageEventIdentity(input)); +} + +export function createRuntimeUsageEventUnrolledPredicate( + database: AppDatabase, + input: RecordRuntimeUsageEventInput, +) { + const identity = resolveUsageEventIdentity(input); + + return notExists( + database + .select({ source: usageEventRollupReceiptsTable.source }) + .from(usageEventRollupReceiptsTable) + .where( + and( + eq(usageEventRollupReceiptsTable.source, identity.source), + eq(usageEventRollupReceiptsTable.sourceEventId, identity.sourceEventId), + ), + ), + ); +} + +export function createRuntimeUsageEventConvergencePredicate( + database: AppDatabase, + input: RecordRuntimeUsageEventInput, +) { + const identity = resolveUsageEventIdentity(input); + const identityPredicate = and( + eq(usageEventsTable.source, identity.source), + eq(usageEventsTable.sourceEventId, identity.sourceEventId), + ); + const converged = exists( + database + .select({ source: usageEventsTable.source }) + .from(usageEventsTable) + .where( + and( + identityPredicate, + eq(usageEventsTable.sessionId, input.run.sessionId), + eq(usageEventsTable.sessionRunId, input.run.sessionRunId), + eq(usageEventsTable.model, input.run.model), + eq(usageEventsTable.provider, input.run.provider), + eq(usageEventsTable.usageContract, requireUsageContract(input.usage)), + gte(usageEventsTable.sourceEventSeq, normalizeSourceEventSeq(input.sourceEventSeq)), + ), + ), + ); + + return hasRecordableRuntimeUsage(input.usage) + ? converged + : or( + notExists( + database + .select({ source: usageEventsTable.source }) + .from(usageEventsTable) + .where(identityPredicate), + ), + converged, + ); +} + +function prepareRuntimeUsageEventValues(input: RecordRuntimeUsageEventInput) { const rawInputTokens = toTokenCount(input.usage.inputTokens); const rawOutputTokens = toTokenCount(input.usage.outputTokens); const rawCacheReadTokens = toTokenCount(input.usage.cachedReadTokens); const rawCacheCreationTokens = toTokenCount(input.usage.cachedWriteTokens); const providedCostUsd = toProvidedUsdCost(input.usage); - if (!hasRecordableRuntimeUsage(input.usage)) { + if (!hasRuntimeUsageMutation(input)) { return null; } @@ -190,8 +288,7 @@ function createRuntimeUsageEventInsert(database: AppDatabase, input: RecordRunti provider, }); const { source, sourceEventId } = resolveUsageEventIdentity(input); - - return database.insert(usageEventsTable).values({ + const values = { actorUserId: input.run.actorUserId, agentId: input.run.agentId, agentOwnerUserId: input.run.agentOwnerUserId, @@ -215,34 +312,189 @@ function createRuntimeUsageEventInsert(database: AppDatabase, input: RecordRunti sessionRunId: input.run.sessionRunId, source, sourceEventId, + sourceEventSeq: normalizeSourceEventSeq(input.sourceEventSeq), totalCostUsdMicros: toUsdMicros(cost.totalCostUsd), usageContract, - }); + } satisfies typeof usageEventsTable.$inferInsert; + + return { providedCostUsd, values }; +} + +function createRuntimeUsageEventInsert( + database: AppDatabase, + input: RecordRuntimeUsageEventInput, + writeFence?: SQL, +) { + const prepared = prepareRuntimeUsageEventValues(input); + + if (prepared === null) { + return null; + } + + const { values } = prepared; + const identity = resolveUsageEventIdentity(input); + const mayCreate = hasRecordableRuntimeUsage(input.usage) + ? undefined + : exists( + database + .select({ source: usageEventsTable.source }) + .from(usageEventsTable) + .where( + and( + eq(usageEventsTable.source, identity.source), + eq(usageEventsTable.sourceEventId, identity.sourceEventId), + ), + ), + ); + + return { + prepared, + query: database.insert(usageEventsTable).select( + database + .select({ + actorUserId: selectedValue(values.actorUserId, "actor_user_id"), + agentId: selectedValue(values.agentId, "agent_id"), + agentOwnerUserId: selectedValue(values.agentOwnerUserId, "agent_owner_user_id"), + agentPublicationStateAtRun: selectedValue( + values.agentPublicationStateAtRun, + "agent_publication_state_at_run", + ), + agentRevisionId: selectedValue(values.agentRevisionId, "agent_revision_id"), + cacheCreationTokens: selectedValue(values.cacheCreationTokens, "cache_creation_tokens"), + cacheReadTokens: selectedValue(values.cacheReadTokens, "cache_read_tokens"), + createdAt: selectedValue(values.createdAt, "created_at"), + id: selectedValue(values.id, "id"), + inputTokens: selectedValue(values.inputTokens, "input_tokens"), + model: selectedValue(values.model, "model"), + organizationId: selectedValue(values.organizationId, "organization_id"), + appId: selectedValue(values.appId, "app_id"), + outputTokens: selectedValue(values.outputTokens, "output_tokens"), + priceSnapshotJson: selectedValue(values.priceSnapshotJson, "price_snapshot_json"), + pricingStatus: selectedValue(values.pricingStatus, "pricing_status"), + provider: selectedValue(values.provider, "provider"), + runPurpose: selectedValue(values.runPurpose, "run_purpose"), + runtimeId: selectedValue(values.runtimeId, "runtime_id"), + sessionId: selectedValue(values.sessionId, "session_id"), + sessionRunId: selectedValue(values.sessionRunId, "session_run_id"), + source: selectedValue(values.source, "source"), + sourceEventId: selectedValue(values.sourceEventId, "source_event_id"), + sourceEventSeq: selectedValue(values.sourceEventSeq, "source_event_seq"), + totalCostUsdMicros: selectedValue(values.totalCostUsdMicros, "total_cost_usd_micros"), + usageContract: selectedValue(values.usageContract, "usage_contract"), + }) + .from(sql`(SELECT 1)`) + .where( + and(createRuntimeUsageEventUnrolledPredicate(database, input), writeFence, mayCreate), + ), + ), + }; } export function createRuntimeUsageEventUpsert( database: AppDatabase, input: RecordRuntimeUsageEventInput, + writeFence?: SQL, ) { - const query = createRuntimeUsageEventInsert(database, input); + const insert = createRuntimeUsageEventInsert(database, input, writeFence); - if (query === null) { + if (insert === null) { return null; } + const { prepared, query } = insert; + + const sourceEventSeqFence = + input.sourceEventSeq === undefined + ? eq(usageEventsTable.sourceEventSeq, 0) + : sql`${usageEventsTable.sourceEventSeq} < excluded.source_event_seq`; + const cacheCreationTokensProvided = isProvidedTokenCount(input.usage.cachedWriteTokens); + const cacheReadTokensProvided = isProvidedTokenCount(input.usage.cachedReadTokens); + const inputTokensProvided = isProvidedTokenCount(input.usage.inputTokens); + const outputTokensProvided = isProvidedTokenCount(input.usage.outputTokens); + const mergedCacheCreationTokens = cacheCreationTokensProvided + ? sql`excluded.cache_creation_tokens` + : sql`${usageEventsTable.cacheCreationTokens}`; + const mergedCacheReadTokens = cacheReadTokensProvided + ? sql`excluded.cache_read_tokens` + : sql`${usageEventsTable.cacheReadTokens}`; + const mergedInputTokens = + input.usage.usageContract !== "anthropic_bucketed" + ? inputTokensProvided + ? sql`excluded.input_tokens` + : sql`${usageEventsTable.inputTokens}` + : inputTokensProvided + ? cacheReadTokensProvided + ? sql`excluded.input_tokens` + : sql`excluded.input_tokens + ${usageEventsTable.cacheReadTokens}` + : cacheReadTokensProvided + ? sql`${usageEventsTable.inputTokens} - ${usageEventsTable.cacheReadTokens} + excluded.cache_read_tokens` + : sql`${usageEventsTable.inputTokens}`; + const mergedOutputTokens = outputTokensProvided + ? sql`excluded.output_tokens` + : sql`${usageEventsTable.outputTokens}`; + const mergedHasTokens = sql`( + ${mergedCacheCreationTokens} > 0 OR + ${mergedCacheReadTokens} > 0 OR + ${mergedInputTokens} > 0 OR + ${mergedOutputTokens} > 0 + )`; + const preserveRuntimeReportedCost = sql`( + ${prepared.providedCostUsd === null} AND + NOT ${mergedHasTokens} AND + json_extract(${usageEventsTable.priceSnapshotJson}, '$.source') = 'runtime_reported_usd' + )`; + const billableInputTokens = sql`MAX(0, ${mergedInputTokens} - ${mergedCacheReadTokens})`; + const pricingRateSnapshot = sql`CASE + WHEN ${inputTokensProvided} OR + json_extract(${usageEventsTable.priceSnapshotJson}, '$.source') IS NOT 'mosoo_seed_2026_07_10' + THEN excluded.price_snapshot_json + ELSE ${usageEventsTable.priceSnapshotJson} + END`; + const priceSnapshotJson = sql`CASE + WHEN excluded.pricing_status <> 'priced' THEN NULL + WHEN ${preserveRuntimeReportedCost} THEN ${usageEventsTable.priceSnapshotJson} + WHEN ${mergedHasTokens} THEN + json_set(${pricingRateSnapshot}, '$.billableInputTokens', ${billableInputTokens}) + ELSE excluded.price_snapshot_json + END`; + const totalCostUsdMicros = sql`CASE + WHEN excluded.pricing_status = 'priced' AND ${preserveRuntimeReportedCost} + THEN ${usageEventsTable.totalCostUsdMicros} + WHEN excluded.pricing_status = 'priced' AND ${mergedHasTokens} THEN + CAST(ROUND( + ${billableInputTokens} * json_extract(${pricingRateSnapshot}, '$.inputUsdPerMillion') + + ${mergedOutputTokens} * json_extract(${pricingRateSnapshot}, '$.outputUsdPerMillion') + + ${mergedCacheReadTokens} * json_extract(${pricingRateSnapshot}, '$.cacheReadUsdPerMillion') + + ${mergedCacheCreationTokens} * json_extract(${pricingRateSnapshot}, '$.cacheWriteUsdPerMillion') + ) AS INTEGER) + WHEN excluded.pricing_status = 'unknown' AND ${prepared.providedCostUsd === null} + THEN ${usageEventsTable.totalCostUsdMicros} + ELSE excluded.total_cost_usd_micros + END`; + const identityFence = sql` + ${usageEventsTable.sessionId} IS excluded.session_id AND + ${usageEventsTable.sessionRunId} IS excluded.session_run_id AND + ${usageEventsTable.model} = excluded.model AND + ${usageEventsTable.provider} = excluded.provider AND + ${usageEventsTable.usageContract} = excluded.usage_contract + `; + const updateFence = sql`${sourceEventSeqFence} AND ${identityFence}`; + return query.onConflictDoUpdate({ set: { - cacheCreationTokens: sql`excluded.cache_creation_tokens`, - cacheReadTokens: sql`excluded.cache_read_tokens`, - inputTokens: sql`excluded.input_tokens`, + cacheCreationTokens: mergedCacheCreationTokens, + cacheReadTokens: mergedCacheReadTokens, + inputTokens: mergedInputTokens, model: sql`excluded.model`, - outputTokens: sql`excluded.output_tokens`, - priceSnapshotJson: sql`excluded.price_snapshot_json`, + outputTokens: mergedOutputTokens, + priceSnapshotJson, pricingStatus: sql`excluded.pricing_status`, provider: sql`excluded.provider`, - totalCostUsdMicros: sql`excluded.total_cost_usd_micros`, + sourceEventSeq: sql`excluded.source_event_seq`, + totalCostUsdMicros, usageContract: sql`excluded.usage_contract`, }, + setWhere: writeFence === undefined ? updateFence : sql`${updateFence} AND ${writeFence}`, target: [usageEventsTable.source, usageEventsTable.sourceEventId], }); } @@ -251,28 +503,59 @@ export function createRuntimeUsageEventInsertIfMissing( database: AppDatabase, input: RecordRuntimeUsageEventInput, ) { - const query = createRuntimeUsageEventInsert(database, input); + const insert = createRuntimeUsageEventInsert(database, input); - if (query === null) { + if (insert === null) { return null; } - return query.onConflictDoNothing({ + return insert.query.onConflictDoNothing({ target: [usageEventsTable.source, usageEventsTable.sourceEventId], }); } +function readRuntimeReportedUsd(priceSnapshotJson: string | null): number | null { + if (priceSnapshotJson === null) { + return null; + } + + try { + const snapshot = JSON.parse(priceSnapshotJson) as unknown; + + if ( + typeof snapshot !== "object" || + snapshot === null || + !("source" in snapshot) || + snapshot.source !== "runtime_reported_usd" || + !("reportedCostUsd" in snapshot) || + typeof snapshot.reportedCostUsd !== "number" || + !Number.isFinite(snapshot.reportedCostUsd) || + snapshot.reportedCostUsd < 0 + ) { + return null; + } + + return snapshot.reportedCostUsd; + } catch { + return null; + } +} + export async function recordRuntimeUsageEvent( database: D1Database, input: RecordRuntimeUsageEventInput, ): Promise { - if (!hasRecordableRuntimeUsage(input.usage)) { + if (!hasRuntimeUsageMutation(input)) { return; } const appDatabase = getAppDatabase(database); if (await isUsageEventAlreadyRolledUp(appDatabase, resolveUsageEventIdentity(input))) { + if (input.sourceEventSeq !== undefined) { + throw new Error("Runtime usage event was already rolled up and cannot be replaced safely."); + } + return; } @@ -283,4 +566,116 @@ export async function recordRuntimeUsageEvent( } await query.run(); + + if ( + input.sourceEventSeq !== undefined && + (await isUsageEventAlreadyRolledUp(appDatabase, resolveUsageEventIdentity(input))) + ) { + throw new Error("Runtime usage event was rolled up before its durable replay was verified."); + } + + if (input.sourceEventSeq !== undefined) { + const identity = resolveUsageEventIdentity(input); + const prepared = prepareRuntimeUsageEventValues(input); + + if (prepared === null) { + return; + } + const { values: expected } = prepared; + + const row = + (await appDatabase + .select({ + cacheCreationTokens: usageEventsTable.cacheCreationTokens, + cacheReadTokens: usageEventsTable.cacheReadTokens, + inputTokens: usageEventsTable.inputTokens, + model: usageEventsTable.model, + outputTokens: usageEventsTable.outputTokens, + priceSnapshotJson: usageEventsTable.priceSnapshotJson, + pricingStatus: usageEventsTable.pricingStatus, + provider: usageEventsTable.provider, + sessionId: usageEventsTable.sessionId, + sessionRunId: usageEventsTable.sessionRunId, + sourceEventSeq: usageEventsTable.sourceEventSeq, + totalCostUsdMicros: usageEventsTable.totalCostUsdMicros, + usageContract: usageEventsTable.usageContract, + }) + .from(usageEventsTable) + .where( + and( + eq(usageEventsTable.source, identity.source), + eq(usageEventsTable.sourceEventId, identity.sourceEventId), + ), + ) + .limit(1) + .get()) ?? null; + + if (row === null) { + if (!hasRecordableRuntimeUsage(input.usage)) { + return; + } + throw new Error("Runtime usage event CAS did not persist the durable event."); + } + + if (row.sourceEventSeq < input.sourceEventSeq) { + throw new Error("Runtime usage event CAS did not persist the durable event."); + } + + if ( + row.model !== expected.model || + row.provider !== expected.provider || + row.sessionId !== expected.sessionId || + row.sessionRunId !== expected.sessionRunId || + row.usageContract !== expected.usageContract + ) { + throw new Error("Runtime usage event identity conflicts with its durable event."); + } + + if (row.sourceEventSeq > input.sourceEventSeq) { + return; + } + + const tokensUnavailable = + row.cacheCreationTokens === 0 && + row.cacheReadTokens === 0 && + row.inputTokens === 0 && + row.outputTokens === 0; + const providedCostUsd = + prepared.providedCostUsd ?? + (tokensUnavailable ? readRuntimeReportedUsd(row.priceSnapshotJson) : null); + const mergedCost = calculateUsageCost({ + cacheCreationTokens: row.cacheCreationTokens, + cacheReadTokens: row.cacheReadTokens, + inputTokens: row.inputTokens, + model: row.model, + outputTokens: row.outputTokens, + pricedAtMs: input.run.createdAtMs, + providedCostUsd, + provider: row.provider, + }); + const costMustMatch = + (mergedCost.pricing !== null || providedCostUsd !== null) && + (tokensUnavailable || + isProvidedTokenCount(input.usage.inputTokens) || + mergedCost.pricing === null); + + if ( + (isProvidedTokenCount(input.usage.cachedWriteTokens) && + row.cacheCreationTokens !== expected.cacheCreationTokens) || + (isProvidedTokenCount(input.usage.cachedReadTokens) && + row.cacheReadTokens !== expected.cacheReadTokens) || + (isProvidedTokenCount(input.usage.inputTokens) && + (input.usage.usageContract === "anthropic_bucketed" && + !isProvidedTokenCount(input.usage.cachedReadTokens) + ? row.inputTokens - row.cacheReadTokens !== expected.inputTokens + : row.inputTokens !== expected.inputTokens)) || + (isProvidedTokenCount(input.usage.outputTokens) && + row.outputTokens !== expected.outputTokens) || + (costMustMatch && row.priceSnapshotJson !== mergedCost.priceSnapshotJson) || + (costMustMatch && row.pricingStatus !== mergedCost.pricingStatus) || + (costMustMatch && row.totalCostUsdMicros !== toUsdMicros(mergedCost.totalCostUsd)) + ) { + throw new Error("Runtime usage event seq was replayed with conflicting content."); + } + } } diff --git a/apps/api/src/modules/environments/application/environment-package-artifact-backup-store.ts b/apps/api/src/modules/environments/application/environment-package-artifact-backup-store.ts new file mode 100644 index 00000000..8a2f05eb --- /dev/null +++ b/apps/api/src/modules/environments/application/environment-package-artifact-backup-store.ts @@ -0,0 +1,826 @@ +import { + environmentPackageArtifactBackupsTable, + environmentPackageArtifactBackupStagingTable, +} from "@mosoo/db"; +import type { ApiCommandId } from "@mosoo/db"; +import type { AppId, SandboxBackupId } from "@mosoo/id"; +import { parsePlatformId } from "@mosoo/id"; +import { and, eq } from "drizzle-orm"; + +import { getAppDatabase } from "../../../platform/db/drizzle"; +import type { + EnvironmentPackageArtifactKey, + EnvironmentPackageArtifactPaths, +} from "../domain/environment-package-artifact"; +import { + environmentPackageArtifactDir, + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_REFRESH_WINDOW_MS, + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS, + parseEnvironmentPackageArtifactPaths, +} from "../domain/environment-package-artifact"; + +const DATABASE_NOW_MS_SQL = "CAST(unixepoch('subsec') * 1000 AS INTEGER)"; + +export interface EnvironmentPackageArtifactBackupStage { + readonly actualBackupId: SandboxBackupId | null; + readonly appId: AppId; + readonly attemptCount: number; + readonly claimOwner: string; + readonly commandId: ApiCommandId; + readonly createdAt: number; + readonly deliveryGeneration: number; + readonly dir: string; + readonly inputDigest: string; + readonly paths: EnvironmentPackageArtifactPaths; + readonly updatedAt: number; +} + +export interface EnvironmentPackageArtifactBackupManifest extends EnvironmentPackageArtifactKey { + readonly attemptCount: number; + readonly backupId: SandboxBackupId; + readonly commandId: ApiCommandId; + readonly committedAt: number; + readonly deliveryGeneration: number; + readonly expiresAt: number; + readonly manifestGeneration: number; + readonly paths: EnvironmentPackageArtifactPaths; +} + +export interface EnvironmentPackageArtifactCommandAuthority extends EnvironmentPackageArtifactKey { + readonly attemptCount: number; + readonly commandId: ApiCommandId; + readonly deliveryGeneration: number; +} + +export type EnvironmentPackageArtifactBackupStageAuthority = Pick< + EnvironmentPackageArtifactBackupStage, + "attemptCount" | "claimOwner" | "commandId" | "deliveryGeneration" +>; + +function mapStage( + row: typeof environmentPackageArtifactBackupStagingTable.$inferSelect, +): EnvironmentPackageArtifactBackupStage { + const artifactDir = environmentPackageArtifactDir(row); + const paths = parseEnvironmentPackageArtifactPaths(JSON.parse(row.pathsJson), artifactDir); + if (row.dir !== artifactDir || paths === null) { + throw new Error("Environment package artifact backup stage paths are invalid."); + } + return { ...row, paths }; +} + +function mapManifest( + row: typeof environmentPackageArtifactBackupsTable.$inferSelect, +): EnvironmentPackageArtifactBackupManifest { + const paths = parseEnvironmentPackageArtifactPaths( + JSON.parse(row.pathsJson), + environmentPackageArtifactDir(row), + ); + if (paths === null) { + throw new Error("Environment package artifact backup manifest paths are invalid."); + } + return { ...row, paths }; +} + +export async function getEnvironmentPackageArtifactBackupManifest( + database: D1Database, + key: EnvironmentPackageArtifactKey, +): Promise { + const row = await getAppDatabase(database) + .select() + .from(environmentPackageArtifactBackupsTable) + .where( + and( + eq(environmentPackageArtifactBackupsTable.appId, key.appId), + eq(environmentPackageArtifactBackupsTable.inputDigest, key.inputDigest), + ), + ) + .limit(1) + .get(); + return row === undefined ? null : mapManifest(row); +} + +export async function retireExpiredEnvironmentPackageArtifactBackups( + database: D1Database, + limit: number, +): Promise { + if (!Number.isSafeInteger(limit) || limit <= 0) { + throw new TypeError("Environment package artifact retirement limit must be positive."); + } + const rows = await database + .prepare( + `DELETE FROM environment_package_artifact_backup + WHERE backup_id IN ( + SELECT artifact.backup_id + FROM environment_package_artifact_backup AS artifact + WHERE artifact.expires_at <= ${DATABASE_NOW_MS_SQL} + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_delete_intent AS deletion + WHERE deletion.backup_id = artifact.backup_id + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup AS backup + WHERE backup.id = artifact.backup_id AND backup.status = 'ready' + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_staging AS stage + WHERE stage.actual_backup_id = artifact.backup_id + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup_staging AS stage + WHERE stage.actual_backup_id = artifact.backup_id + ) + ORDER BY artifact.expires_at, artifact.backup_id + LIMIT ? + ) + AND expires_at <= ${DATABASE_NOW_MS_SQL} + RETURNING backup_id`, + ) + .bind(limit) + .all<{ backup_id: SandboxBackupId }>(); + return rows.results.map((row) => row.backup_id); +} + +function commandMatchesStage(alias: string): string { + return `EXISTS ( + SELECT 1 FROM api_command AS command + WHERE command.id = ${alias}.command_id + AND command.kind = 'environment_package_artifact_build' + AND json_valid(command.payload_json) = 1 + AND json_extract(command.payload_json, '$.appId') = ${alias}.app_id + AND json_extract(command.payload_json, '$.inputDigest') = ${alias}.input_digest + AND command.delivery_generation = ${alias}.delivery_generation + AND command.attempt_count = ${alias}.attempt_count + AND command.claim_owner = ${alias}.claim_owner + )`; +} + +export async function getEnvironmentPackageArtifactBackupStage( + database: D1Database, + commandId: ApiCommandId, +): Promise { + const row = await getAppDatabase(database) + .select() + .from(environmentPackageArtifactBackupStagingTable) + .where(eq(environmentPackageArtifactBackupStagingTable.commandId, commandId)) + .limit(1) + .get(); + return row === undefined ? null : mapStage(row); +} + +function mapCommandAuthority(row: { + readonly attempt_count: number; + readonly delivery_generation: number; + readonly id: string; + readonly payload_json: string; +}): EnvironmentPackageArtifactCommandAuthority | null { + try { + const payload = JSON.parse(row.payload_json) as unknown; + if (typeof payload !== "object" || payload === null) { + return null; + } + const appId = Reflect.get(payload, "appId"); + const inputDigest = Reflect.get(payload, "inputDigest"); + if ( + typeof appId !== "string" || + typeof inputDigest !== "string" || + !Number.isSafeInteger(row.attempt_count) || + row.attempt_count <= 0 || + !Number.isSafeInteger(row.delivery_generation) || + row.delivery_generation <= 0 + ) { + return null; + } + return { + appId: parsePlatformId(appId, "environment artifact command app ID"), + attemptCount: row.attempt_count, + commandId: parsePlatformId(row.id, "environment artifact command ID"), + deliveryGeneration: row.delivery_generation, + inputDigest, + }; + } catch { + return null; + } +} + +export async function getEnvironmentPackageArtifactCommandIntent( + database: D1Database, + commandId: ApiCommandId, +): Promise { + const row = await database + .prepare( + `SELECT attempt_count, delivery_generation, id, payload_json + FROM api_command + WHERE id = ? AND kind = 'environment_package_artifact_build'`, + ) + .bind(commandId) + .first<{ + attempt_count: number; + delivery_generation: number; + id: string; + payload_json: string; + }>(); + return row === null ? null : mapCommandAuthority(row); +} + +export async function findEnvironmentPackageArtifactCommandAuthority( + database: D1Database, + key: EnvironmentPackageArtifactKey, +): Promise { + const rows = ( + await database + .prepare( + `SELECT attempt_count, delivery_generation, id, payload_json + FROM api_command + WHERE kind = 'environment_package_artifact_build' + AND json_valid(payload_json) = 1 + AND json_extract(payload_json, '$.appId') = ? + AND json_extract(payload_json, '$.inputDigest') = ? + ORDER BY id + LIMIT 2`, + ) + .bind(key.appId, key.inputDigest) + .all<{ + attempt_count: number; + delivery_generation: number; + id: string; + payload_json: string; + }>() + ).results; + return rows.length === 1 ? mapCommandAuthority(rows[0]!) : null; +} + +function manifestMatches( + manifest: EnvironmentPackageArtifactBackupManifest, + input: { + readonly backupId: SandboxBackupId; + readonly commandId: ApiCommandId; + readonly attemptCount: number; + readonly deliveryGeneration: number; + readonly expiresAt: number; + readonly key: EnvironmentPackageArtifactKey; + readonly paths: EnvironmentPackageArtifactPaths; + }, +): boolean { + return ( + manifest.backupId === input.backupId && + manifest.commandId === input.commandId && + manifest.attemptCount === input.attemptCount && + manifest.deliveryGeneration === input.deliveryGeneration && + manifest.expiresAt === input.expiresAt && + manifest.appId === input.key.appId && + manifest.inputDigest === input.key.inputDigest && + JSON.stringify(manifest.paths) === JSON.stringify(input.paths) + ); +} + +export async function commitEnvironmentPackageArtifactBackup( + database: D1Database, + input: EnvironmentPackageArtifactBackupStageAuthority & { + readonly actualBackupId: SandboxBackupId; + readonly expiresAt: number; + readonly key: EnvironmentPackageArtifactKey; + readonly paths: EnvironmentPackageArtifactPaths; + }, +): Promise { + const paths = parseEnvironmentPackageArtifactPaths( + input.paths, + environmentPackageArtifactDir(input.key), + ); + if (paths === null) { + throw new Error("Environment package artifact backup paths are invalid."); + } + const pathsJson = JSON.stringify(paths); + const previous = await getEnvironmentPackageArtifactBackupManifest(database, input.key); + if ( + previous !== null && + manifestMatches(previous, { ...input, backupId: input.actualBackupId }) + ) { + return true; + } + if (previous?.backupId === input.actualBackupId) { + return false; + } + const authoritySql = `EXISTS ( + SELECT 1 + FROM environment_package_artifact_backup_staging AS stage + JOIN api_command AS command ON command.id = stage.command_id + WHERE stage.command_id = ? AND stage.actual_backup_id = ? + AND stage.app_id = ? AND stage.input_digest = ? AND stage.dir = ? AND stage.paths_json = ? + AND stage.delivery_generation = ? AND stage.attempt_count = ? + AND stage.claim_owner = ? AND command.status = 'running' + AND command.claim_expires_at > ${DATABASE_NOW_MS_SQL} + AND ${commandMatchesStage("stage")} + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_delete_intent AS deletion + WHERE deletion.backup_id = stage.actual_backup_id + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup + WHERE id = stage.actual_backup_id + UNION ALL + SELECT 1 FROM sandbox_backup_staging + WHERE actual_backup_id = stage.actual_backup_id + ) + )`; + const bindings = [ + input.commandId, + input.actualBackupId, + input.key.appId, + input.key.inputDigest, + environmentPackageArtifactDir(input.key), + pathsJson, + input.deliveryGeneration, + input.attemptCount, + input.claimOwner, + ] as const; + const insert = async (): Promise => { + await database + .prepare( + `INSERT INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) + SELECT ?, ?, ?, ?, ${DATABASE_NOW_MS_SQL}, ?, ?, ?, 1, ? + WHERE ${authoritySql} + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup + WHERE backup_id = ? OR (app_id = ? AND input_digest = ?) + )`, + ) + .bind( + input.key.appId, + input.attemptCount, + input.actualBackupId, + input.commandId, + input.deliveryGeneration, + input.expiresAt, + input.key.inputDigest, + pathsJson, + ...bindings, + input.actualBackupId, + input.key.appId, + input.key.inputDigest, + ) + .run(); + }; + if (previous === null) { + await insert(); + } else { + await database + .prepare( + `UPDATE environment_package_artifact_backup + SET attempt_count = ?, backup_id = ?, command_id = ?, committed_at = ${DATABASE_NOW_MS_SQL}, + delivery_generation = ?, expires_at = ?, manifest_generation = manifest_generation + 1, + paths_json = ? + WHERE app_id = ? AND input_digest = ? AND manifest_generation = ? + AND backup_id = ? AND ${authoritySql} + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup AS other + WHERE other.backup_id = ? + AND (other.app_id <> ? OR other.input_digest <> ?) + )`, + ) + .bind( + input.attemptCount, + input.actualBackupId, + input.commandId, + input.deliveryGeneration, + input.expiresAt, + pathsJson, + input.key.appId, + input.key.inputDigest, + previous.manifestGeneration, + previous.backupId, + ...bindings, + input.actualBackupId, + input.key.appId, + input.key.inputDigest, + ) + .run(); + } + let manifest = await getEnvironmentPackageArtifactBackupManifest(database, input.key); + if (previous !== null && manifest === null) { + await insert(); + manifest = await getEnvironmentPackageArtifactBackupManifest(database, input.key); + } + return ( + manifest !== null && manifestMatches(manifest, { ...input, backupId: input.actualBackupId }) + ); +} + +export async function adoptLegacyEnvironmentPackageArtifactBackup( + database: D1Database, + input: EnvironmentPackageArtifactCommandAuthority & { + readonly actualBackupId: SandboxBackupId; + readonly expiresAt: number; + readonly key: EnvironmentPackageArtifactKey; + readonly paths: EnvironmentPackageArtifactPaths; + }, +): Promise { + const paths = parseEnvironmentPackageArtifactPaths( + input.paths, + environmentPackageArtifactDir(input.key), + ); + if (paths === null) { + throw new Error("Environment package artifact backup paths are invalid."); + } + const pathsJson = JSON.stringify(paths); + await database + .prepare( + `INSERT INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) + SELECT ?, command.attempt_count, ?, command.id, ${DATABASE_NOW_MS_SQL}, + command.delivery_generation, ?, ?, 1, ? + FROM api_command AS command + WHERE command.id = ? AND command.kind = 'environment_package_artifact_build' + AND command.status = 'succeeded' + AND command.completed_at IS NOT NULL + AND command.claim_owner IS NULL AND command.claim_expires_at IS NULL + AND command.attempt_count = ? AND command.delivery_generation = ? + AND ? > ${DATABASE_NOW_MS_SQL} + ? + AND ? <= ${DATABASE_NOW_MS_SQL} + ? + AND json_valid(command.payload_json) = 1 + AND json_extract(command.payload_json, '$.appId') = ? + AND json_extract(command.payload_json, '$.inputDigest') = ? + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_delete_intent AS deletion + WHERE deletion.backup_id = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup AS artifact + WHERE artifact.backup_id = ? + OR (artifact.app_id = ? AND artifact.input_digest = ?) + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup_staging AS stage + WHERE stage.actual_backup_id = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup + WHERE id = ? + UNION ALL + SELECT 1 FROM sandbox_backup_staging + WHERE actual_backup_id = ? + )`, + ) + .bind( + input.key.appId, + input.actualBackupId, + input.expiresAt, + input.key.inputDigest, + pathsJson, + input.commandId, + input.attemptCount, + input.deliveryGeneration, + input.expiresAt, + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_REFRESH_WINDOW_MS, + input.expiresAt, + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS * 1_000, + input.key.appId, + input.key.inputDigest, + input.actualBackupId, + input.actualBackupId, + input.actualBackupId, + input.key.appId, + input.key.inputDigest, + input.actualBackupId, + input.actualBackupId, + ) + .run(); + const manifest = await getEnvironmentPackageArtifactBackupManifest(database, input.key); + return ( + manifest !== null && manifestMatches(manifest, { ...input, backupId: input.actualBackupId }) + ); +} + +export async function stageEnvironmentPackageArtifactBackup( + database: D1Database, + input: { + readonly attemptCount: number; + readonly claimOwner: string; + readonly commandId: ApiCommandId; + readonly deliveryGeneration: number; + readonly dir: string; + readonly key: EnvironmentPackageArtifactKey; + readonly paths: EnvironmentPackageArtifactPaths; + }, +): Promise { + if ( + !Number.isSafeInteger(input.deliveryGeneration) || + input.deliveryGeneration <= 0 || + !Number.isSafeInteger(input.attemptCount) || + input.attemptCount <= 0 + ) { + throw new Error("Environment package artifact attempt must be a positive safe integer."); + } + const artifactDir = environmentPackageArtifactDir(input.key); + const paths = parseEnvironmentPackageArtifactPaths(input.paths, artifactDir); + if (input.dir !== artifactDir || paths === null) { + throw new Error("Environment package artifact backup paths are invalid."); + } + const pathsJson = JSON.stringify(paths); + await database + .prepare( + `DELETE FROM environment_package_artifact_backup_staging AS stage + WHERE stage.command_id = ? + AND ( + stage.delivery_generation <> ? OR stage.attempt_count <> ? OR stage.claim_owner <> ? + ) + AND EXISTS ( + SELECT 1 FROM api_command AS command + WHERE command.id = stage.command_id + AND command.kind = 'environment_package_artifact_build' + AND command.status = 'running' AND command.claim_owner = ? + AND command.delivery_generation = ? AND command.attempt_count = ? + AND command.claim_expires_at > ${DATABASE_NOW_MS_SQL} + AND json_valid(command.payload_json) = 1 + AND json_extract(command.payload_json, '$.appId') = ? + AND json_extract(command.payload_json, '$.inputDigest') = ? + ) + `, + ) + .bind( + input.commandId, + input.deliveryGeneration, + input.attemptCount, + input.claimOwner, + input.claimOwner, + input.deliveryGeneration, + input.attemptCount, + input.key.appId, + input.key.inputDigest, + ) + .run(); + await database + .prepare( + `INSERT INTO environment_package_artifact_backup_staging ( + actual_backup_id, app_id, attempt_count, claim_owner, command_id, created_at, + delivery_generation, dir, input_digest, paths_json, updated_at + ) + SELECT NULL, ?, ?, ?, ?, ${DATABASE_NOW_MS_SQL}, ?, ?, ?, ?, ${DATABASE_NOW_MS_SQL} + WHERE EXISTS ( + SELECT 1 FROM api_command AS command + WHERE command.id = ? AND command.kind = 'environment_package_artifact_build' + AND command.status = 'running' AND command.claim_owner = ? + AND command.delivery_generation = ? AND command.attempt_count = ? + AND command.claim_expires_at > ${DATABASE_NOW_MS_SQL} + AND json_valid(command.payload_json) = 1 + AND json_extract(command.payload_json, '$.appId') = ? + AND json_extract(command.payload_json, '$.inputDigest') = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup_staging AS existing + WHERE existing.command_id = ? + OR (existing.app_id = ? AND existing.input_digest = ?) + ) + ON CONFLICT DO NOTHING`, + ) + .bind( + input.key.appId, + input.attemptCount, + input.claimOwner, + input.commandId, + input.deliveryGeneration, + input.dir, + input.key.inputDigest, + pathsJson, + input.commandId, + input.claimOwner, + input.deliveryGeneration, + input.attemptCount, + input.key.appId, + input.key.inputDigest, + input.commandId, + input.key.appId, + input.key.inputDigest, + ) + .run(); + const stage = await getEnvironmentPackageArtifactBackupStage(database, input.commandId); + if ( + stage === null || + stage.appId !== input.key.appId || + stage.attemptCount !== input.attemptCount || + stage.claimOwner !== input.claimOwner || + stage.deliveryGeneration !== input.deliveryGeneration || + stage.inputDigest !== input.key.inputDigest || + stage.dir !== input.dir || + JSON.stringify(stage.paths) !== pathsJson + ) { + throw new Error("Environment package artifact backup lost its immutable stage."); + } + const owned = await database + .prepare( + `SELECT 1 + FROM environment_package_artifact_backup_staging AS stage + JOIN api_command AS command ON command.id = stage.command_id + WHERE stage.command_id = ? AND command.status = 'running' + AND command.claim_owner = ? AND command.delivery_generation = ? + AND command.attempt_count = ? + AND command.claim_expires_at > ${DATABASE_NOW_MS_SQL} + AND ${commandMatchesStage("stage")}`, + ) + .bind(input.commandId, input.claimOwner, input.deliveryGeneration, input.attemptCount) + .first(); + if (owned === null) { + throw new Error("Environment package artifact build lost its API command lease."); + } + return stage; +} + +export async function claimEnvironmentPackageArtifactBackupActual( + database: D1Database, + input: { + readonly actualBackupId: SandboxBackupId; + readonly authority: { + readonly attemptCount: number; + readonly claimOwner: string; + readonly deliveryGeneration: number; + }; + readonly commandId: ApiCommandId; + readonly dir: string; + }, +): Promise<{ readonly actualBackupId: SandboxBackupId } | null> { + const authorityBindings = [ + input.authority.claimOwner, + input.authority.deliveryGeneration, + input.authority.attemptCount, + ] as const; + const authorityPredicate = `command.status = 'running' AND command.claim_owner = ? AND command.delivery_generation = ? AND command.attempt_count = ? AND command.claim_expires_at > ${DATABASE_NOW_MS_SQL}`; + const claimed = await database + .prepare( + `UPDATE environment_package_artifact_backup_staging AS stage + SET actual_backup_id = ?, updated_at = ${DATABASE_NOW_MS_SQL} + WHERE stage.command_id = ? AND stage.dir = ? AND stage.actual_backup_id IS NULL + AND EXISTS ( + SELECT 1 FROM api_command AS command + WHERE command.id = stage.command_id AND ${authorityPredicate} + ) + AND ${commandMatchesStage("stage")} + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_delete_intent AS deletion + WHERE deletion.backup_id = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup AS artifact + WHERE artifact.backup_id = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup + WHERE id = ? + UNION ALL + SELECT 1 FROM sandbox_backup_staging + WHERE actual_backup_id = ? + ) + RETURNING actual_backup_id`, + ) + .bind( + input.actualBackupId, + input.commandId, + input.dir, + ...authorityBindings, + input.actualBackupId, + input.actualBackupId, + input.actualBackupId, + input.actualBackupId, + ) + .first<{ actual_backup_id: SandboxBackupId }>(); + if (claimed !== null) { + return { actualBackupId: claimed.actual_backup_id }; + } + const winner = await database + .prepare( + `SELECT stage.actual_backup_id + FROM environment_package_artifact_backup_staging AS stage + JOIN api_command AS command ON command.id = stage.command_id + WHERE stage.command_id = ? AND stage.dir = ? AND stage.actual_backup_id IS NOT NULL + AND ${authorityPredicate} AND ${commandMatchesStage("stage")} + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_delete_intent AS deletion + WHERE deletion.backup_id = stage.actual_backup_id + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup AS artifact + WHERE artifact.backup_id = stage.actual_backup_id + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup + WHERE id = stage.actual_backup_id + UNION ALL + SELECT 1 FROM sandbox_backup_staging + WHERE actual_backup_id = stage.actual_backup_id + )`, + ) + .bind(input.commandId, input.dir, ...authorityBindings) + .first<{ actual_backup_id: SandboxBackupId }>(); + return winner === null ? null : { actualBackupId: winner.actual_backup_id }; +} + +export async function clearMissingEnvironmentPackageArtifactBackupActual( + database: D1Database, + input: EnvironmentPackageArtifactBackupStageAuthority & { + readonly actualBackupId: SandboxBackupId; + }, +): Promise { + const result = await database + .prepare( + `UPDATE environment_package_artifact_backup_staging + SET actual_backup_id = NULL, updated_at = ${DATABASE_NOW_MS_SQL} + WHERE command_id = ? AND delivery_generation = ? AND attempt_count = ? + AND claim_owner = ? AND actual_backup_id = ?`, + ) + .bind( + input.commandId, + input.deliveryGeneration, + input.attemptCount, + input.claimOwner, + input.actualBackupId, + ) + .run(); + return (result.meta.changes ?? 0) === 1; +} + +export async function completeEnvironmentPackageArtifactBackupStage( + database: D1Database, + input: EnvironmentPackageArtifactBackupStageAuthority & { + readonly actualBackupId: SandboxBackupId | null; + }, +): Promise { + const result = await database + .prepare( + `DELETE FROM environment_package_artifact_backup_staging + WHERE command_id = ? AND delivery_generation = ? AND attempt_count = ? + AND claim_owner = ? AND actual_backup_id IS ? + AND EXISTS ( + SELECT 1 FROM environment_package_artifact_backup AS artifact + WHERE artifact.app_id = environment_package_artifact_backup_staging.app_id + AND artifact.input_digest = environment_package_artifact_backup_staging.input_digest + AND artifact.backup_id = environment_package_artifact_backup_staging.actual_backup_id + AND artifact.command_id = environment_package_artifact_backup_staging.command_id + AND artifact.delivery_generation = environment_package_artifact_backup_staging.delivery_generation + AND artifact.attempt_count = environment_package_artifact_backup_staging.attempt_count + AND artifact.paths_json = environment_package_artifact_backup_staging.paths_json + )`, + ) + .bind( + input.commandId, + input.deliveryGeneration, + input.attemptCount, + input.claimOwner, + input.actualBackupId, + ) + .run(); + return (result.meta.changes ?? 0) === 1; +} + +export async function revokeTerminalEnvironmentPackageArtifactBackupStage( + database: D1Database, + commandId: ApiCommandId, +): Promise { + const row = await database + .prepare( + `DELETE FROM environment_package_artifact_backup_staging AS stage + WHERE stage.command_id = ? + AND NOT EXISTS ( + SELECT 1 FROM api_command AS command + WHERE command.id = stage.command_id AND command.status = 'running' + AND command.claim_expires_at > ${DATABASE_NOW_MS_SQL} + AND ${commandMatchesStage("stage")} + ) + RETURNING actual_backup_id`, + ) + .bind(commandId) + .first<{ actual_backup_id: SandboxBackupId | null }>(); + return row?.actual_backup_id ?? null; +} + +export async function revokeTerminalEnvironmentPackageArtifactBackupStages( + database: D1Database, +): Promise { + const result = await database + .prepare( + `DELETE FROM environment_package_artifact_backup_staging AS stage + WHERE stage.command_id IN ( + SELECT candidate.command_id + FROM environment_package_artifact_backup_staging AS candidate + WHERE NOT EXISTS ( + SELECT 1 FROM api_command AS command + WHERE command.id = candidate.command_id + AND command.status = 'running' + AND command.claim_expires_at > ${DATABASE_NOW_MS_SQL} + AND command.kind = 'environment_package_artifact_build' + AND json_valid(command.payload_json) = 1 + AND json_extract(command.payload_json, '$.appId') = candidate.app_id + AND json_extract(command.payload_json, '$.inputDigest') = candidate.input_digest + AND command.delivery_generation = candidate.delivery_generation + AND command.attempt_count = candidate.attempt_count + AND command.claim_owner = candidate.claim_owner + ) + ORDER BY candidate.updated_at, candidate.command_id + LIMIT 64 + )`, + ) + .run(); + return result.meta.changes ?? 0; +} diff --git a/apps/api/src/modules/environments/application/environment-package-artifact-backup.ts b/apps/api/src/modules/environments/application/environment-package-artifact-backup.ts new file mode 100644 index 00000000..dcefc1c3 --- /dev/null +++ b/apps/api/src/modules/environments/application/environment-package-artifact-backup.ts @@ -0,0 +1,304 @@ +import type { ApiCommandId } from "@mosoo/db"; +import type { SandboxBackupId } from "@mosoo/id"; +import { parsePlatformId } from "@mosoo/id"; + +import { logWarn } from "../../../platform/cloudflare/logger"; +import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; +import { + decodeSandboxBackupIdForPlatform, + encodeSandboxBackupIdForStorage, +} from "../../runtime/infrastructure/sandbox-backup-id"; +import { + getSandboxBackupObjectKeys, + parseSandboxBackupMetadata, +} from "../../runtime/infrastructure/sandbox-backup-platform"; +import { authorizeSandboxBackupDeletion } from "../../runtime/infrastructure/sandbox-backup-store"; +import { + createEnvironmentPackageArtifactBackupName, + environmentPackageArtifactDir, + environmentPackageArtifactMetadataKey, + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_REFRESH_WINDOW_MS, + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS, + parseEnvironmentPackageArtifactBackupName, + parseEnvironmentPackageArtifactMetadata, +} from "../domain/environment-package-artifact"; +import type { + EnvironmentPackageArtifactKey, + EnvironmentPackageArtifactMetadata, + EnvironmentPackageArtifactPaths, +} from "../domain/environment-package-artifact"; +import { + adoptLegacyEnvironmentPackageArtifactBackup, + commitEnvironmentPackageArtifactBackup, + findEnvironmentPackageArtifactCommandAuthority, + getEnvironmentPackageArtifactBackupManifest, + getEnvironmentPackageArtifactCommandIntent, +} from "./environment-package-artifact-backup-store"; + +export function environmentPackageArtifactMetadataBackupId( + metadata: EnvironmentPackageArtifactMetadata, +): SandboxBackupId | null { + try { + return encodeSandboxBackupIdForStorage(metadata.backupId); + } catch { + return null; + } +} + +interface EnvironmentPackageArtifactBackupVerification { + readonly authority: { + readonly attemptCount: number; + readonly commandId: ApiCommandId; + readonly deliveryGeneration: number; + } | null; + readonly expiresAt: number; +} + +async function readEnvironmentPackageArtifactBackupVerification( + bindings: Pick, + input: { readonly backupId: SandboxBackupId; readonly dir: string }, +): Promise { + const [dataKey, metadataKey] = getSandboxBackupObjectKeys(input.backupId); + const [data, storedMetadata] = await Promise.all([ + bindings.SANDBOX_STATE_BUCKET.head(dataKey), + bindings.SANDBOX_STATE_BUCKET.get(metadataKey), + ]); + if (data === null || storedMetadata === null) { + return null; + } + let metadata = null; + let createdAt: unknown; + let ttl: unknown; + try { + const value = JSON.parse(await storedMetadata.text()) as unknown; + metadata = parseSandboxBackupMetadata(value); + createdAt = + typeof value === "object" && value !== null ? Reflect.get(value, "createdAt") : null; + ttl = typeof value === "object" && value !== null ? Reflect.get(value, "ttl") : null; + } catch { + return null; + } + if ( + metadata?.id !== decodeSandboxBackupIdForPlatform(input.backupId) || + metadata.dir !== input.dir + ) { + return null; + } + const createdAtMs = typeof createdAt === "string" ? Date.parse(createdAt) : Number.NaN; + const expiresAt = createdAtMs + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS * 1_000; + if ( + !Number.isSafeInteger(createdAtMs) || + new Date(createdAtMs).toISOString() !== createdAt || + ttl !== ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS || + !Number.isSafeInteger(expiresAt) + ) { + return null; + } + if (metadata.name === null) { + return { authority: null, expiresAt }; + } + const authority = parseEnvironmentPackageArtifactBackupName(metadata.name); + if (authority === null) { + return null; + } + try { + return { + authority: { + attemptCount: authority.attemptCount, + commandId: parsePlatformId( + authority.commandId, + "environment artifact command ID", + ), + deliveryGeneration: authority.deliveryGeneration, + }, + expiresAt, + }; + } catch { + return null; + } +} + +export async function isEnvironmentPackageArtifactBackupReady( + bindings: Pick, + input: { + readonly attemptCount: number; + readonly backupId: SandboxBackupId; + readonly commandId: ApiCommandId; + readonly deliveryGeneration: number; + readonly dir: string; + }, +): Promise { + const verified = await readEnvironmentPackageArtifactBackupVerification(bindings, input); + return ( + verified?.authority?.commandId === input.commandId && + verified.authority.deliveryGeneration === input.deliveryGeneration && + verified.authority.attemptCount === input.attemptCount + ); +} + +async function readLegacyEnvironmentPackageArtifactMetadata( + bindings: Pick, + key: EnvironmentPackageArtifactKey, +): Promise { + const object = await bindings.SANDBOX_STATE_BUCKET.get( + environmentPackageArtifactMetadataKey(key), + ); + if (object === null) { + return null; + } + try { + return parseEnvironmentPackageArtifactMetadata( + JSON.parse(await object.text()), + environmentPackageArtifactDir(key), + ); + } catch { + return null; + } +} + +export async function resolveEnvironmentPackageArtifactBackup( + bindings: Pick, + key: EnvironmentPackageArtifactKey, +): Promise { + const manifest = await getEnvironmentPackageArtifactBackupManifest(bindings.DB, key); + const clock = await bindings.DB.prepare( + "SELECT CAST(unixepoch('subsec') * 1000 AS INTEGER) AS now_ms", + ).first<{ now_ms: number }>(); + if (clock === null) { + return null; + } + if (manifest !== null) { + if ( + manifest.expiresAt <= + clock.now_ms + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_REFRESH_WINDOW_MS + ) { + return null; + } + const verified = await readEnvironmentPackageArtifactBackupVerification(bindings, { + backupId: manifest.backupId, + dir: environmentPackageArtifactDir(key), + }); + const command = await getEnvironmentPackageArtifactCommandIntent( + bindings.DB, + manifest.commandId, + ); + if ( + verified === null || + command === null || + command.appId !== key.appId || + command.inputDigest !== key.inputDigest || + command.commandId !== manifest.commandId || + verified.expiresAt !== manifest.expiresAt || + (verified.authority !== null && + (verified.authority.commandId !== manifest.commandId || + verified.authority.deliveryGeneration !== manifest.deliveryGeneration || + verified.authority.attemptCount !== manifest.attemptCount)) + ) { + return null; + } + return { + backupId: decodeSandboxBackupIdForPlatform(manifest.backupId), + paths: manifest.paths, + }; + } + + const projection = await readLegacyEnvironmentPackageArtifactMetadata(bindings, key); + if (projection === null) { + return null; + } + const backupId = environmentPackageArtifactMetadataBackupId(projection); + if (backupId === null) { + return null; + } + const verified = await readEnvironmentPackageArtifactBackupVerification(bindings, { + backupId, + dir: environmentPackageArtifactDir(key), + }); + if ( + verified === null || + verified.expiresAt <= clock.now_ms + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_REFRESH_WINDOW_MS || + verified.expiresAt > clock.now_ms + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS * 1_000 + ) { + return null; + } + const command = + verified.authority === null + ? await findEnvironmentPackageArtifactCommandAuthority(bindings.DB, key) + : await getEnvironmentPackageArtifactCommandIntent(bindings.DB, verified.authority.commandId); + if ( + command === null || + command.appId !== key.appId || + command.inputDigest !== key.inputDigest || + (verified.authority !== null && + (verified.authority.commandId !== command.commandId || + verified.authority.deliveryGeneration !== command.deliveryGeneration || + verified.authority.attemptCount !== command.attemptCount)) || + !(await adoptLegacyEnvironmentPackageArtifactBackup(bindings.DB, { + actualBackupId: backupId, + expiresAt: verified.expiresAt, + ...command, + key, + paths: projection.paths, + })) + ) { + logWarn("runtime.environment_artifact.legacy_projection_adoption_failed", { + backupId, + inputDigest: key.inputDigest, + }); + return null; + } + return projection; +} + +export async function publishEnvironmentPackageArtifactBackup( + bindings: Pick, + input: { + readonly attemptCount: number; + readonly backupId: SandboxBackupId; + readonly claimOwner: string; + readonly commandId: ApiCommandId; + readonly deliveryGeneration: number; + readonly key: EnvironmentPackageArtifactKey; + readonly paths: EnvironmentPackageArtifactPaths; + }, +): Promise { + const verified = await readEnvironmentPackageArtifactBackupVerification(bindings, { + backupId: input.backupId, + dir: environmentPackageArtifactDir(input.key), + }); + if ( + verified?.authority?.commandId !== input.commandId || + verified.authority.deliveryGeneration !== input.deliveryGeneration || + verified.authority.attemptCount !== input.attemptCount + ) { + throw new Error("Environment package artifact backup objects are incomplete."); + } + const committed = await commitEnvironmentPackageArtifactBackup(bindings.DB, { + actualBackupId: input.backupId, + attemptCount: input.attemptCount, + claimOwner: input.claimOwner, + commandId: input.commandId, + deliveryGeneration: input.deliveryGeneration, + expiresAt: verified.expiresAt, + key: input.key, + paths: input.paths, + }); + const manifest = await getEnvironmentPackageArtifactBackupManifest(bindings.DB, input.key); + if (!committed || manifest?.backupId !== input.backupId) { + if (manifest !== null && manifest.backupId !== input.backupId) { + await authorizeSandboxBackupDeletion(bindings.DB, { + authority: { + attemptCount: input.attemptCount, + commandId: input.commandId, + deliveryGeneration: input.deliveryGeneration, + kind: "environment_invalid", + }, + backupId: input.backupId, + }); + return; + } + throw new Error("Environment package artifact backup lost D1 commit authority."); + } +} + +export { createEnvironmentPackageArtifactBackupName }; diff --git a/apps/api/src/modules/environments/application/environment-package-artifact-build.service.ts b/apps/api/src/modules/environments/application/environment-package-artifact-build.service.ts index ad9a30f8..a9afdecd 100644 --- a/apps/api/src/modules/environments/application/environment-package-artifact-build.service.ts +++ b/apps/api/src/modules/environments/application/environment-package-artifact-build.service.ts @@ -1,29 +1,44 @@ +import type { ApiCommandId } from "@mosoo/db"; +import type { SandboxBackupId } from "@mosoo/id"; + import { disposeRpcResource, withDisposedRpcResult, } from "../../../platform/cloudflare/rpc-disposal"; -import { requireCloudflareSandboxBinding } from "../../../platform/cloudflare/sandbox-binding"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; +import { quoteShellArg } from "../../../shared/shell"; import type { EnvironmentPackageArtifactBuildCommandPayload } from "../../api-command/application/api-command-payload"; import { isRuntimeSandboxLocalBucketEnabled } from "../../runtime/infrastructure/runtime-sandbox-bucket-mount"; -import { deleteSandboxBackupObjects } from "../../runtime/infrastructure/sandbox-backup-platform"; +import { getEphemeralUnversionedSandboxHandle } from "../../runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-platform"; +import { encodeSandboxBackupIdForStorage } from "../../runtime/infrastructure/sandbox-backup-id"; +import { authorizeSandboxBackupDeletion } from "../../runtime/infrastructure/sandbox-backup-store"; import type { SandboxHandle } from "../../runtime/infrastructure/sandbox-handles"; -import { toSandboxHandle } from "../../runtime/infrastructure/sandbox-handles"; import { createEnvironmentPackageArtifactKey, + environmentPackageArtifactBuildSandboxId, environmentPackageArtifactDir, - environmentPackageArtifactMetadataKey, - environmentPackageArtifactSandboxId, ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS, ENVIRONMENT_PACKAGE_ARTIFACT_MAX_BUILD_MS, } from "../domain/environment-package-artifact"; import type { EnvironmentPackageArtifactPaths } from "../domain/environment-package-artifact"; import { normalizePackages } from "./environment-config"; -import { readEnvironmentPackageArtifactMetadata } from "./environment-package-artifact.service"; +import { + createEnvironmentPackageArtifactBackupName, + environmentPackageArtifactMetadataBackupId, + isEnvironmentPackageArtifactBackupReady, + publishEnvironmentPackageArtifactBackup, + resolveEnvironmentPackageArtifactBackup, +} from "./environment-package-artifact-backup"; +import { + claimEnvironmentPackageArtifactBackupActual, + clearMissingEnvironmentPackageArtifactBackupActual, + completeEnvironmentPackageArtifactBackupStage, + getEnvironmentPackageArtifactBackupStage, + stageEnvironmentPackageArtifactBackup, +} from "./environment-package-artifact-backup-store"; -function quoteShellArg(value: string): string { - return `'${value.replaceAll("'", `'"'"'`)}'`; -} +const ENVIRONMENT_ARTIFACT_BUILD_SANDBOX_SLEEP_AFTER_SECONDS = + ENVIRONMENT_PACKAGE_ARTIFACT_MAX_BUILD_MS / 1_000 + 5 * 60; export function createEnvironmentPackageArtifactBuildScript(input: { npmRoot: string; @@ -62,9 +77,112 @@ async function closeBuildSandbox(sandbox: SandboxHandle): Promise { } } +export interface EnvironmentPackageArtifactBuildAuthority { + readonly attemptCount: number; + readonly claimOwner: string; + readonly commandId: ApiCommandId; + readonly deliveryGeneration: number; + requireOwnership(): Promise; +} + +async function withBuildOwnership( + authority: EnvironmentPackageArtifactBuildAuthority, + effect: () => Promise, +): Promise { + await authority.requireOwnership(); + try { + return await effect(); + } finally { + await authority.requireOwnership(); + } +} + +async function settlePublishedArtifact( + bindings: ApiBindings, + input: { + readonly authority: EnvironmentPackageArtifactBuildAuthority; + readonly commandId: ApiCommandId; + readonly key: Awaited>; + }, +): Promise { + const metadata = await resolveEnvironmentPackageArtifactBackup(bindings, input.key); + if (metadata === null) { + return null; + } + const backupId = environmentPackageArtifactMetadataBackupId(metadata); + if (backupId === null) { + throw new Error("Environment package artifact metadata has an invalid backup ID."); + } + const stage = await getEnvironmentPackageArtifactBackupStage(bindings.DB, input.commandId); + if (stage === null) { + return backupId; + } + if (JSON.stringify(stage.paths) !== JSON.stringify(metadata.paths)) { + throw new Error("Environment package artifact manifest changed its immutable paths."); + } + await withBuildOwnership(input.authority, () => + completeEnvironmentPackageArtifactBackupStage(bindings.DB, { + actualBackupId: stage.actualBackupId, + attemptCount: stage.attemptCount, + claimOwner: stage.claimOwner, + commandId: stage.commandId, + deliveryGeneration: stage.deliveryGeneration, + }), + ); + return backupId; +} + +async function publishClaimedArtifact( + bindings: ApiBindings, + input: { + readonly attemptCount: number; + readonly authority: EnvironmentPackageArtifactBuildAuthority; + readonly backupId: SandboxBackupId; + readonly commandId: ApiCommandId; + readonly deliveryGeneration: number; + readonly dir: string; + readonly key: Awaited>; + readonly paths: EnvironmentPackageArtifactPaths; + }, +): Promise { + if ( + (await settlePublishedArtifact(bindings, { + authority: input.authority, + commandId: input.commandId, + key: input.key, + })) !== null + ) { + return; + } + if (!(await isEnvironmentPackageArtifactBackupReady(bindings, input))) { + throw new Error("Environment package artifact backup objects are incomplete."); + } + await withBuildOwnership(input.authority, () => + publishEnvironmentPackageArtifactBackup(bindings, { + attemptCount: input.attemptCount, + backupId: input.backupId, + claimOwner: input.authority.claimOwner, + commandId: input.commandId, + deliveryGeneration: input.deliveryGeneration, + key: input.key, + paths: input.paths, + }), + ); + if ( + (await settlePublishedArtifact(bindings, { + authority: input.authority, + commandId: input.commandId, + key: input.key, + })) === null + ) { + throw new Error("Environment package artifact manifest was not committed."); + } +} + export async function buildEnvironmentPackageArtifact( bindings: ApiBindings, payload: EnvironmentPackageArtifactBuildCommandPayload, + authority: EnvironmentPackageArtifactBuildAuthority, ): Promise { const packages = normalizePackages(payload.packages); const key = await createEnvironmentPackageArtifactKey({ @@ -75,7 +193,13 @@ export async function buildEnvironmentPackageArtifact( if (key.inputDigest !== payload.inputDigest) { throw new Error("Environment package artifact digest does not match its payload."); } - if ((await readEnvironmentPackageArtifactMetadata(bindings, key)) !== null) { + if ( + (await settlePublishedArtifact(bindings, { + authority, + commandId: authority.commandId, + key, + })) !== null + ) { return; } @@ -85,32 +209,36 @@ export async function buildEnvironmentPackageArtifact( const tempRoot = `/tmp/mosoo-environment-artifact-${key.inputDigest}`; const npmRoot = `${dir}/npm`; const pipRoot = `${dir}/python`; - let backupId: string | null = null; - const { getSandbox } = await import("@cloudflare/sandbox"); - const sandbox = toSandboxHandle( - getSandbox( - requireCloudflareSandboxBinding(bindings), - environmentPackageArtifactSandboxId(key), - { keepAlive: true, normalizeId: true }, + const sandbox = await getEphemeralUnversionedSandboxHandle( + bindings, + environmentPackageArtifactBuildSandboxId( + authority.commandId, + authority.deliveryGeneration, + authority.attemptCount, ), + ENVIRONMENT_ARTIFACT_BUILD_SANDBOX_SLEEP_AFTER_SECONDS, ); try { - const reset = await sandbox.exec( - `rm -rf ${quoteShellArg(dir)} ${quoteShellArg(tempRoot)} && mkdir -p ${quoteShellArg(dir)} ${quoteShellArg(tempRoot)}`, + const reset = await withBuildOwnership(authority, () => + sandbox.exec( + `rm -rf ${quoteShellArg(dir)} ${quoteShellArg(tempRoot)} && mkdir -p ${quoteShellArg(dir)} ${quoteShellArg(tempRoot)}`, + ), ); if (!reset.success) { throw new Error("Environment package build directory could not be prepared."); } - const result = await sandbox.exec( - createEnvironmentPackageArtifactBuildScript({ - npmRoot, - npmSpecs, - pipRoot, - pipSpecs, - tempRoot, - }), - { timeout: ENVIRONMENT_PACKAGE_ARTIFACT_MAX_BUILD_MS }, + const result = await withBuildOwnership(authority, () => + sandbox.exec( + createEnvironmentPackageArtifactBuildScript({ + npmRoot, + npmSpecs, + pipRoot, + pipSpecs, + tempRoot, + }), + { timeout: ENVIRONMENT_PACKAGE_ARTIFACT_MAX_BUILD_MS }, + ), ); if (!result.success) { const tail = `${result.stdout}\n${result.stderr}`.trim().slice(-4096); @@ -129,29 +257,154 @@ export async function buildEnvironmentPackageArtifact( node: npmSpecs.length === 0 ? [] : [`${npmRoot}/node_modules`], python: pipSpecs.length === 0 ? [] : [pipSite], }; - const backup = await withDisposedRpcResult( - sandbox.createBackup({ + let stage = await withBuildOwnership(authority, () => + stageEnvironmentPackageArtifactBackup(bindings.DB, { + attemptCount: authority.attemptCount, + claimOwner: authority.claimOwner, + commandId: authority.commandId, + deliveryGeneration: authority.deliveryGeneration, dir, - localBucket: isRuntimeSandboxLocalBucketEnabled(bindings), - ttl: ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS, + key, + paths, }), - (backupResult) => ({ dir: backupResult.dir, id: backupResult.id }), ); - backupId = backup.id; - if (!backupId || backup.dir !== dir) { + if ( + stage.actualBackupId !== null && + !(await isEnvironmentPackageArtifactBackupReady(bindings, { + attemptCount: authority.attemptCount, + backupId: stage.actualBackupId, + commandId: authority.commandId, + deliveryGeneration: authority.deliveryGeneration, + dir, + })) + ) { + const cleared = await withBuildOwnership(authority, () => + clearMissingEnvironmentPackageArtifactBackupActual(bindings.DB, { + actualBackupId: stage.actualBackupId as SandboxBackupId, + attemptCount: stage.attemptCount, + claimOwner: stage.claimOwner, + commandId: stage.commandId, + deliveryGeneration: stage.deliveryGeneration, + }), + ); + if (!cleared) { + if ( + (await settlePublishedArtifact(bindings, { + authority, + commandId: authority.commandId, + key, + })) !== null + ) { + return; + } + throw new Error("Environment package artifact build lost its immutable stage."); + } + stage = { ...stage, actualBackupId: null }; + } + if (stage.actualBackupId !== null) { + const claimed = await withBuildOwnership(authority, () => + claimEnvironmentPackageArtifactBackupActual(bindings.DB, { + actualBackupId: stage.actualBackupId as SandboxBackupId, + authority, + commandId: authority.commandId, + dir, + }), + ); + if (claimed?.actualBackupId !== stage.actualBackupId) { + if ( + (await settlePublishedArtifact(bindings, { + authority, + commandId: authority.commandId, + key, + })) !== null + ) { + return; + } + throw new Error("Environment package artifact build lost its API command lease."); + } + await publishClaimedArtifact(bindings, { + attemptCount: authority.attemptCount, + authority, + backupId: stage.actualBackupId, + commandId: authority.commandId, + deliveryGeneration: authority.deliveryGeneration, + dir, + key, + paths, + }); + return; + } + + const backup = await withBuildOwnership(authority, () => + withDisposedRpcResult( + sandbox.createBackup({ + dir, + localBucket: isRuntimeSandboxLocalBucketEnabled(bindings), + name: createEnvironmentPackageArtifactBackupName(authority), + ttl: ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS, + }), + (backupResult) => ({ dir: backupResult.dir, id: backupResult.id }), + ), + ); + if (!backup.id || backup.dir !== dir) { throw new Error("Environment package artifact backup is invalid."); } - await bindings.SANDBOX_STATE_BUCKET.put( - environmentPackageArtifactMetadataKey(key), - JSON.stringify({ backupId, paths }), - { httpMetadata: { contentType: "application/json" } }, + const candidateId = encodeSandboxBackupIdForStorage(backup.id); + const claimed = await withBuildOwnership(authority, () => + claimEnvironmentPackageArtifactBackupActual(bindings.DB, { + actualBackupId: candidateId, + authority, + commandId: authority.commandId, + dir, + }), ); - backupId = null; - } catch (error) { - if (backupId !== null) { - await deleteSandboxBackupObjects(bindings, [backupId]).catch(() => undefined); + if (claimed?.actualBackupId !== candidateId) { + const published = await settlePublishedArtifact(bindings, { + authority, + commandId: authority.commandId, + key, + }); + if (claimed !== null && (published === null || published !== candidateId)) { + await withBuildOwnership(authority, () => + authorizeSandboxBackupDeletion(bindings.DB, { + authority: { + attemptCount: authority.attemptCount, + commandId: authority.commandId, + deliveryGeneration: authority.deliveryGeneration, + kind: "environment_candidate", + }, + backupId: candidateId, + }), + ); + } + if (published !== null) { + return; + } + if (claimed === null) { + throw new Error("Environment package artifact build lost its API command lease."); + } + await publishClaimedArtifact(bindings, { + attemptCount: authority.attemptCount, + authority, + backupId: claimed.actualBackupId, + commandId: authority.commandId, + deliveryGeneration: authority.deliveryGeneration, + dir, + key, + paths, + }); + return; } - throw error; + await publishClaimedArtifact(bindings, { + attemptCount: authority.attemptCount, + authority, + backupId: candidateId, + commandId: authority.commandId, + deliveryGeneration: authority.deliveryGeneration, + dir, + key, + paths, + }); } finally { await closeBuildSandbox(sandbox); } diff --git a/apps/api/src/modules/environments/application/environment-package-artifact.service.ts b/apps/api/src/modules/environments/application/environment-package-artifact.service.ts index a750ef23..31b9fbc1 100644 --- a/apps/api/src/modules/environments/application/environment-package-artifact.service.ts +++ b/apps/api/src/modules/environments/application/environment-package-artifact.service.ts @@ -10,7 +10,6 @@ import { import { createEnvironmentPackageArtifactKey, environmentPackageArtifactDir, - environmentPackageArtifactMetadataKey, ENVIRONMENT_PACKAGE_ARTIFACT_ABI, } from "../domain/environment-package-artifact"; import type { @@ -18,34 +17,14 @@ import type { EnvironmentPackageArtifactMetadata, } from "../domain/environment-package-artifact"; import { normalizePackages, parsePackagesJson } from "./environment-config"; +import { resolveEnvironmentPackageArtifactBackup } from "./environment-package-artifact-backup"; +import { getEnvironmentPackageArtifactBackupManifest } from "./environment-package-artifact-backup-store"; type ArtifactBindings = Pick< ApiBindings, "API_COMMAND_QUEUE" | "DB" | "ENVIRONMENT_ARTIFACT_BUILD_QUEUE" | "SANDBOX_STATE_BUCKET" >; -export async function readEnvironmentPackageArtifactMetadata( - bindings: Pick, - key: EnvironmentPackageArtifactKey, -): Promise { - const object = await bindings.SANDBOX_STATE_BUCKET.get( - environmentPackageArtifactMetadataKey(key), - ); - if (object === null) { - return null; - } - const metadata = JSON.parse(await object.text()) as Partial; - const paths = metadata.paths; - if ( - typeof metadata.backupId !== "string" || - paths === undefined || - ![paths.executable, paths.node, paths.python].every(Array.isArray) - ) { - throw new Error("Environment package artifact metadata is invalid."); - } - return { backupId: metadata.backupId, paths }; -} - export async function resolveEnvironmentPackageArtifact( bindings: ArtifactBindings, appId: AppId, @@ -64,13 +43,19 @@ export async function resolveEnvironmentPackageArtifact( artifactAbi: ENVIRONMENT_PACKAGE_ARTIFACT_ABI, packages: normalized, }); - const metadata = await readEnvironmentPackageArtifactMetadata(bindings, key); + const metadata = await resolveEnvironmentPackageArtifactBackup(bindings, key); if (metadata === null) { + const refreshCurrentManifest = + (await getEnvironmentPackageArtifactBackupManifest(bindings.DB, key)) !== null; const dedupeKey = `environment_package_artifact_build:${key.appId}:${key.inputDigest}`; + const existingCommand = await findApiCommandByDedupeKey(bindings.DB, dedupeKey); await enqueueApiCommand(bindings, { dedupeKey, kind: "environment_package_artifact_build", - retryTerminal: options.retryFailed === true, + retryTerminal: + refreshCurrentManifest || + existingCommand?.status === "succeeded" || + options.retryFailed === true, payload: { ...key, artifactAbi: ENVIRONMENT_PACKAGE_ARTIFACT_ABI, @@ -80,7 +65,7 @@ export async function resolveEnvironmentPackageArtifact( if (options.retryFailed !== true) { const command = await findApiCommandByDedupeKey(bindings.DB, dedupeKey); if (command !== null && command.status !== "queued" && command.status !== "running") { - const completedMetadata = await readEnvironmentPackageArtifactMetadata(bindings, key); + const completedMetadata = await resolveEnvironmentPackageArtifactBackup(bindings, key); if (completedMetadata !== null) { return { key, metadata: completedMetadata }; } diff --git a/apps/api/src/modules/environments/domain/environment-package-artifact.ts b/apps/api/src/modules/environments/domain/environment-package-artifact.ts index 6e6bebbf..4ff67976 100644 --- a/apps/api/src/modules/environments/domain/environment-package-artifact.ts +++ b/apps/api/src/modules/environments/domain/environment-package-artifact.ts @@ -1,10 +1,11 @@ import type { EnvironmentPackageSpec } from "@mosoo/contracts/environment"; import type { AppId } from "@mosoo/id"; -export const ENVIRONMENT_PACKAGE_ARTIFACT_ROOT = "/workspace/.mosoo/environment-artifacts"; export const ENVIRONMENT_PACKAGE_ARTIFACT_ABI = "environment-artifact-v1"; export const ENVIRONMENT_PACKAGE_ARTIFACT_MAX_BUILD_MS = 10 * 60 * 1000; export const ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS = 10 * 365 * 24 * 60 * 60; +export const ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_REFRESH_WINDOW_MS = 24 * 60 * 60 * 1_000; +const ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_NAME_PREFIX = "mosoo:environment-artifact:v1:"; export interface EnvironmentPackageArtifactPaths { executable: string[]; @@ -22,8 +23,114 @@ export interface EnvironmentPackageArtifactMetadata { paths: EnvironmentPackageArtifactPaths; } -function bytesToHex(bytes: Uint8Array): string { - return [...bytes].map((byte) => byte.toString(16).padStart(2, "0")).join(""); +export interface EnvironmentPackageArtifactBackupAuthorityRef { + readonly attemptCount: number; + readonly commandId: string; + readonly deliveryGeneration: number; +} + +function isPositiveSafeInteger(value: number): boolean { + return Number.isSafeInteger(value) && value > 0; +} + +function parseArtifactPathArray( + value: unknown, + artifactDir: string, + seen: Set, +): string[] | null { + if (!Array.isArray(value)) { + return null; + } + for (const path of value) { + if ( + typeof path !== "string" || + path.includes("\0") || + path.includes(":") || + !path.startsWith(`${artifactDir}/`) || + path + .slice(artifactDir.length + 1) + .split("/") + .some((segment) => segment.length === 0 || segment === "." || segment === "..") || + seen.has(path) + ) { + return null; + } + seen.add(path); + } + return value; +} + +export function parseEnvironmentPackageArtifactPaths( + value: unknown, + artifactDir: string, +): EnvironmentPackageArtifactPaths | null { + if ( + typeof value !== "object" || + value === null || + Object.keys(value).toSorted().join(",") !== "executable,node,python" + ) { + return null; + } + const seen = new Set(); + const executable = parseArtifactPathArray(Reflect.get(value, "executable"), artifactDir, seen); + const node = parseArtifactPathArray(Reflect.get(value, "node"), artifactDir, seen); + const python = parseArtifactPathArray(Reflect.get(value, "python"), artifactDir, seen); + return executable === null || node === null || python === null + ? null + : { executable, node, python }; +} + +export function parseEnvironmentPackageArtifactMetadata( + value: unknown, + artifactDir: string, +): EnvironmentPackageArtifactMetadata | null { + if ( + typeof value !== "object" || + value === null || + Object.keys(value).toSorted().join(",") !== "backupId,paths" + ) { + return null; + } + const backupId = Reflect.get(value, "backupId"); + const paths = parseEnvironmentPackageArtifactPaths(Reflect.get(value, "paths"), artifactDir); + return typeof backupId === "string" && backupId.length > 0 && paths !== null + ? { backupId, paths } + : null; +} + +export function createEnvironmentPackageArtifactBackupName( + authority: EnvironmentPackageArtifactBackupAuthorityRef, +): string { + if ( + authority.commandId.length === 0 || + !isPositiveSafeInteger(authority.deliveryGeneration) || + !isPositiveSafeInteger(authority.attemptCount) + ) { + throw new Error("Environment package artifact command ID is required."); + } + return `${ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_NAME_PREFIX}${authority.commandId}:${authority.deliveryGeneration}:${authority.attemptCount}`; +} + +export function parseEnvironmentPackageArtifactBackupName( + value: string | null, +): EnvironmentPackageArtifactBackupAuthorityRef | null { + if (value === null || !value.startsWith(ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_NAME_PREFIX)) { + return null; + } + const [commandId, deliveryGenerationValue, attemptCountValue, extra] = value + .slice(ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_NAME_PREFIX.length) + .split(":"); + const deliveryGeneration = Number(deliveryGenerationValue); + const attemptCount = Number(attemptCountValue); + return commandId !== undefined && + commandId.length > 0 && + extra === undefined && + /^\d+$/u.test(deliveryGenerationValue ?? "") && + /^\d+$/u.test(attemptCountValue ?? "") && + isPositiveSafeInteger(deliveryGeneration) && + isPositiveSafeInteger(attemptCount) + ? { attemptCount, commandId, deliveryGeneration } + : null; } export async function createEnvironmentPackageArtifactKey(input: { @@ -39,17 +146,28 @@ export async function createEnvironmentPackageArtifactKey(input: { "SHA-256", new TextEncoder().encode(JSON.stringify({ artifactAbi, packages: input.packages })), ); - return { appId: input.appId, inputDigest: bytesToHex(new Uint8Array(digest)) }; + return { appId: input.appId, inputDigest: new Uint8Array(digest).toHex() }; } export function environmentPackageArtifactDir(key: EnvironmentPackageArtifactKey): string { - return `${ENVIRONMENT_PACKAGE_ARTIFACT_ROOT}/${key.inputDigest}`; + return `/workspace/.mosoo/environment-artifacts/${key.inputDigest}`; } export function environmentPackageArtifactMetadataKey(key: EnvironmentPackageArtifactKey): string { return `environment-artifacts/${key.appId}/${key.inputDigest}.json`; } -export function environmentPackageArtifactSandboxId(key: EnvironmentPackageArtifactKey): string { - return `envpkg-${key.appId}-${key.inputDigest}`.toLowerCase().slice(0, 63); +export function environmentPackageArtifactBuildSandboxId( + commandId: string, + deliveryGeneration: number, + attemptCount: number, +): string { + if ( + commandId.length === 0 || + !isPositiveSafeInteger(deliveryGeneration) || + !isPositiveSafeInteger(attemptCount) + ) { + throw new Error("Environment package artifact attempt must be a positive safe integer."); + } + return `envpkg-${commandId}-${deliveryGeneration.toString(36)}-${attemptCount.toString(36)}`.toLowerCase(); } diff --git a/apps/api/src/modules/files/application/file-store.ts b/apps/api/src/modules/files/application/file-store.ts index 4747294a..cec71eca 100644 --- a/apps/api/src/modules/files/application/file-store.ts +++ b/apps/api/src/modules/files/application/file-store.ts @@ -20,13 +20,13 @@ import type { SessionResource, } from "@mosoo/contracts/session"; import { fileRecordsTable, sessionsTable } from "@mosoo/db"; -import { createPlatformId, parsePlatformId } from "@mosoo/id"; +import { parsePlatformId } from "@mosoo/id"; import type { AccountId, AppId, FileId, SessionId } from "@mosoo/id"; -import { and, asc, desc, eq, inArray, or } from "drizzle-orm"; +import { and, asc, desc, eq, inArray, or, sql } from "drizzle-orm"; import type { SQL } from "drizzle-orm"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; -import { getAppDatabase } from "../../../platform/db/drizzle"; +import { getAppDatabase, parameterizedSql } from "../../../platform/db/drizzle"; import { toArrayBuffer } from "../../../shared/bytes"; import { currentTimestampMs } from "../../../time"; import { ensureAppOwnership } from "../../apps/application/app.service"; @@ -45,23 +45,19 @@ import { createFileNotFoundError, createUnexpectedFileError, FileControlError, + isRetryableFileControlError, } from "../infrastructure/file-errors"; -import { - createFinalObjectKey, - createSessionArtifactPath, - normalizeContentType, - normalizeFileName, -} from "../infrastructure/file-paths"; +import { normalizeFileName } from "../infrastructure/file-paths"; import { ensureFileAccess, fileRecordRowColumns, listFileRecords, listFileRecordsById, - parseRuntimeOutputSourcePath, toFileEntry, toFileRecord, toSessionFile, } from "../infrastructure/file-record-store"; +import type { FileRecordRow } from "../infrastructure/file-record-store"; import { updateFile } from "../infrastructure/file-update"; import { completeFileUpload as completeFileUploadRecord } from "../infrastructure/file-upload-complete"; import { createFileUpload, getFileUpload } from "../infrastructure/file-upload-create"; @@ -70,7 +66,7 @@ import { uploadFileContent, uploadFilePart, } from "../infrastructure/file-upload-transfer"; -import { getObjectBody, putObject } from "../infrastructure/r2-s3-client"; +import { deleteObject, getObjectBody, putObject } from "../infrastructure/r2-s3-client"; import { normalizeR2Etag } from "../infrastructure/r2-s3-client"; import { ensureAppSessionFileAccess, @@ -88,16 +84,6 @@ export interface CompleteFileUploadCommand { viewer: AuthenticatedViewer; } -export interface RuntimeOutputFileInput { - bindings: ApiBindings; - body: Uint8Array; - contentSha256?: string; - contentType?: string | null; - createdBy: AccountId; - path: string; - sessionId: SessionId; -} - export interface AgentPackageFileAdmissionInput { appId: AppId; fileId: FileId; @@ -193,7 +179,6 @@ export interface FileStore { database: D1Database, sessionId: SessionId, ): Promise; - listReadySessionArtifactKeys(database: D1Database, sessionId: SessionId): Promise; listReadySessionFiles(database: D1Database, sessionId: SessionId): Promise; listSessionResourcePathEntries( database: D1Database, @@ -215,7 +200,6 @@ export interface FileStore { body: ContentBody, ): Promise; readSessionArtifactBytes(bindings: ApiBindings, objectKey: string): Promise; - recordRuntimeOutput(input: RuntimeOutputFileInput): Promise; streamContent( bindings: ApiBindings, viewer: AuthenticatedViewer, @@ -263,13 +247,50 @@ export function getRuntimeOutputName(path: string): string { export async function createRuntimeOutputContentSha256(body: Uint8Array): Promise { const digest = await crypto.subtle.digest("SHA-256", toArrayBuffer(body)); - return [...new Uint8Array(digest)].map((byte) => byte.toString(16).padStart(2, "0")).join(""); + return new Uint8Array(digest).toHex(); } export function createRuntimeOutputParentPath(path: string, contentSha256: string): string { return ["runtime-output", ...readRuntimeOutputPathSegments(path), contentSha256].join("/"); } +export async function putRuntimeArtifactObject(input: { + bindings: ApiBindings; + body: Uint8Array; + contentSha256: string; + contentType: string; + objectKey: string; + sourcePath: string; + attemptId: string; +}): Promise<{ contentLength: number; contentType: string | null; etag: string } | null> { + try { + return await putObject({ + bindings: input.bindings, + body: input.body, + contentType: input.contentType, + customMetadata: { + attemptId: input.attemptId, + contentSha256: input.contentSha256, + sourcePath: input.sourcePath, + }, + objectKey: input.objectKey, + options: { ifNoneMatch: "*" }, + }); + } catch (error) { + if (isRetryableFileControlError(error)) { + return null; + } + throw error; + } +} + +export async function deleteRuntimeArtifactObject( + bindings: ApiBindings, + objectKey: string, +): Promise { + await deleteObject(bindings, objectKey); +} + function toSessionResource(file: FileRecord): SessionResource { return { createdAt: file.createdAt, @@ -609,93 +630,116 @@ async function listVisibleFileRecords( .all(); } -async function listReadySessionFiles( +type CurrentSessionArtifactFileRow = Pick< + FileRecordRow, + "committed" | "created_at" | "id" | "mime_type" | "name" | "session_kind" | "size" +>; + +interface CurrentSessionArtifactSourceRow { + readonly authority_source_path: string; + readonly object_key: string; + readonly size: number; +} + +const CURRENT_SESSION_ARTIFACT_AUTHORITY_CTE = ` +WITH input(session_id) AS ( + VALUES (?) +), +artifact_authority AS ( + SELECT artifact_head.source_path, artifact_head.file_id + FROM session_artifact_head AS artifact_head + INNER JOIN input ON input.session_id = artifact_head.session_id +)`; + +const CURRENT_SESSION_ARTIFACT_FILES_SQL = `${CURRENT_SESSION_ARTIFACT_AUTHORITY_CTE} +SELECT + file_record.committed, + file_record.created_at, + file_record.id, + file_record.mime_type, + file_record.name, + file_record.session_kind, + file_record.size, + artifact_authority.source_path +FROM artifact_authority +INNER JOIN file_record ON file_record.id = artifact_authority.file_id +INNER JOIN input ON input.session_id = file_record.scope_id +WHERE file_record.scope_kind = 'session' + AND file_record.session_kind = 'artifact' + AND file_record.status = 'ready' +ORDER BY file_record.created_at DESC, file_record.id DESC`; + +const CURRENT_SESSION_ARTIFACT_SOURCES_SQL = `${CURRENT_SESSION_ARTIFACT_AUTHORITY_CTE} +SELECT + file_record.object_key, + file_record.size, + artifact_authority.source_path AS authority_source_path +FROM artifact_authority +INNER JOIN file_record ON file_record.id = artifact_authority.file_id +INNER JOIN input ON input.session_id = file_record.scope_id +WHERE file_record.scope_kind = 'session' + AND file_record.session_kind = 'artifact' + AND file_record.status = 'ready' +ORDER BY artifact_authority.source_path`; + +async function listCurrentSessionArtifactFiles( database: D1Database, sessionId: SessionId, -): Promise { - const rows = await getAppDatabase(database) - .select(fileRecordRowColumns) - .from(fileRecordsTable) - .where( - and( - eq(fileRecordsTable.scopeKind, "session"), - eq(fileRecordsTable.scopeId, sessionId), - eq(fileRecordsTable.status, "ready"), - ), - ) - .orderBy(desc(fileRecordsTable.createdAt)) - .all(); +): Promise { + return getAppDatabase(database).all( + parameterizedSql(CURRENT_SESSION_ARTIFACT_FILES_SQL, [sessionId]), + ); +} - return rows.map(toSessionFile); +async function listCurrentSessionArtifactSources( + database: D1Database, + sessionId: SessionId, +): Promise { + return getAppDatabase(database).all( + parameterizedSql(CURRENT_SESSION_ARTIFACT_SOURCES_SQL, [sessionId]), + ); } -async function listReadySessionArtifactKeys( +async function listReadySessionFiles( database: D1Database, sessionId: SessionId, -): Promise { - const rows = await getAppDatabase(database) - .select({ - parentPath: fileRecordsTable.parentPath, - }) - .from(fileRecordsTable) - .where( - and( - eq(fileRecordsTable.scopeKind, "session"), - eq(fileRecordsTable.scopeId, sessionId), - eq(fileRecordsTable.status, "ready"), - eq(fileRecordsTable.sessionKind, "artifact"), - ), - ) - .all(); +): Promise { + const db = getAppDatabase(database); + const [rows, artifacts] = await Promise.all([ + db + .select(fileRecordRowColumns) + .from(fileRecordsTable) + .where( + and( + eq(fileRecordsTable.scopeKind, "session"), + eq(fileRecordsTable.scopeId, sessionId), + eq(fileRecordsTable.status, "ready"), + sql`${fileRecordsTable.sessionKind} IS NOT 'artifact'`, + ), + ) + .all(), + listCurrentSessionArtifactFiles(database, sessionId), + ]); - return rows.map((row) => row.parentPath); + return [...rows, ...artifacts] + .toSorted( + (left, right) => right.created_at - left.created_at || right.id.localeCompare(left.id), + ) + .map(toSessionFile); } async function listLatestReadySessionArtifactSources( database: D1Database, sessionId: SessionId, ): Promise { - const rows = await getAppDatabase(database) - .select({ - id: fileRecordsTable.id, - createdAt: fileRecordsTable.createdAt, - objectKey: fileRecordsTable.objectKey, - parentPath: fileRecordsTable.parentPath, - size: fileRecordsTable.size, - }) - .from(fileRecordsTable) - .where( - and( - eq(fileRecordsTable.scopeKind, "session"), - eq(fileRecordsTable.scopeId, sessionId), - eq(fileRecordsTable.status, "ready"), - eq(fileRecordsTable.sessionKind, "artifact"), - ), - ) - .orderBy(asc(fileRecordsTable.createdAt), asc(fileRecordsTable.id)) - .all(); - - // Ascending scan + map overwrite keeps the newest record per source path - // (ULID ids break created-at ties in creation order). - const latestBySourcePath = new Map(); - - for (const row of rows) { - const sourcePath = parseRuntimeOutputSourcePath(row.parentPath); - - if (sourcePath === null) { - continue; - } - - latestBySourcePath.set(sourcePath, { - objectKey: row.objectKey, + const rows = await listCurrentSessionArtifactSources(database, sessionId); + return rows + .map((row) => ({ + objectKey: row.object_key, size: row.size, - sourcePath, - }); - } - - return [...latestBySourcePath.values()].toSorted((left, right) => - left.sourcePath.localeCompare(right.sourcePath), - ); + sourcePath: row.authority_source_path, + })) + .toSorted((left, right) => left.sourcePath.localeCompare(right.sourcePath)); } async function listSessionResources( @@ -899,68 +943,6 @@ async function readSessionArtifactBytes( return body === null ? null : new Uint8Array(await body.arrayBuffer()); } -async function recordRuntimeOutput(input: RuntimeOutputFileInput): Promise { - const fileId = createPlatformId(); - const name = getRuntimeOutputName(input.path); - const contentType = normalizeContentType(input.contentType ?? "application/octet-stream"); - const contentSha256 = input.contentSha256 ?? (await createRuntimeOutputContentSha256(input.body)); - const path = createSessionArtifactPath(fileId, name); - const timestampMs = currentTimestampMs(); - const objectKey = createFinalObjectKey({ - created_by_account_id: input.createdBy, - id: fileId, - name, - path, - scope_id: input.sessionId, - scope_kind: "session", - session_kind: "artifact", - }); - const object = await putObject({ - bindings: input.bindings, - body: input.body, - contentType, - objectKey, - }); - - await getAppDatabase(input.bindings.DB) - .insert(fileRecordsTable) - .values({ - committed: true, - createdAt: timestampMs, - createdByAccountId: input.createdBy, - etag: object.etag, - expiresAt: null, - id: fileId, - mimeType: object.contentType ?? contentType, - name, - objectKey, - ownerId: input.sessionId, - ownerKind: "session", - parentPath: createRuntimeOutputParentPath(input.path, contentSha256), - path, - purpose: "session_artifact", - scopeId: input.sessionId, - scopeKind: "session", - sessionKind: "artifact", - size: object.contentLength, - status: "ready", - updatedAt: timestampMs, - version: 1, - }) - .run(); - - const createdRows = await listFileRecordsById(input.bindings.DB, [fileId]); - const createdRow = createdRows[0]; - - if (createdRow === undefined) { - throw createFileNotFoundError("Runtime output file was not created."); - } - - const file = toFileRecord(createdRow); - await publishSessionResourceUpsert(input.bindings, file); - return file; -} - export { createFileErrorResponse, createUnexpectedFileError, @@ -986,14 +968,12 @@ export const fileStore: FileStore = { getUpload, list, listLatestReadySessionArtifactSources, - listReadySessionArtifactKeys, listReadySessionFiles, listSessionResourcePathEntries, listSessionResources, putContent, putPart, readSessionArtifactBytes, - recordRuntimeOutput, streamContent, update, }; diff --git a/apps/api/src/modules/files/infrastructure/file-record-model.ts b/apps/api/src/modules/files/infrastructure/file-record-model.ts index f9011cc0..bae51647 100644 --- a/apps/api/src/modules/files/infrastructure/file-record-model.ts +++ b/apps/api/src/modules/files/infrastructure/file-record-model.ts @@ -97,7 +97,7 @@ const SHA256_PATTERN = /^[a-f0-9]{64}$/; // path an artifact was recorded from. Path segments the writer could never // produce (empty, ".", "..") mark the record malformed, which keeps the value // safe to use as a workspace-relative write target. -export function parseRuntimeOutputSourcePath(parentPath: string): string | null { +function parseRuntimeOutputSourcePath(parentPath: string): string | null { const segments = parentPath.split("/"); const contentSha256 = segments.at(-1); @@ -114,8 +114,9 @@ export function parseRuntimeOutputSourcePath(parentPath: string): string | null const hasUnsafeSegment = pathSegments.some( (segment) => segment.length === 0 || segment === "." || segment === "..", ); + const sourcePath = pathSegments.join("/"); - return hasUnsafeSegment ? null : pathSegments.join("/"); + return hasUnsafeSegment || !sourcePath.startsWith("outputs/") ? null : sourcePath; } function toRuntimeOutputSourcePath(row: FileRecordRow): string | null { @@ -255,7 +256,12 @@ export function toUploadSummary(upload: FileUploadRow, file: FileRecordRow): Fil }; } -export function toSessionFile(row: FileRecordRow): SessionFile { +type SessionFileRow = Pick< + FileRecordRow, + "committed" | "created_at" | "id" | "mime_type" | "name" | "session_kind" | "size" +>; + +export function toSessionFile(row: SessionFileRow): SessionFile { return { committed: row.committed === 1, createdAt: toIsoString(row.created_at), diff --git a/apps/api/src/modules/files/infrastructure/file-record-store.ts b/apps/api/src/modules/files/infrastructure/file-record-store.ts index 84eccc4e..b4970340 100644 --- a/apps/api/src/modules/files/infrastructure/file-record-store.ts +++ b/apps/api/src/modules/files/infrastructure/file-record-store.ts @@ -11,7 +11,6 @@ export { } from "./file-record-mutations"; export { fileRecordRowColumns, - parseRuntimeOutputSourcePath, toFileEntry, toFileRecord, toSessionFile, diff --git a/apps/api/src/modules/files/infrastructure/r2-s3-client-types.ts b/apps/api/src/modules/files/infrastructure/r2-s3-client-types.ts index fd39430d..3fc20a19 100644 --- a/apps/api/src/modules/files/infrastructure/r2-s3-client-types.ts +++ b/apps/api/src/modules/files/infrastructure/r2-s3-client-types.ts @@ -55,6 +55,7 @@ export interface PutObjectInput { bindings: ApiBindings; body: ReadableStream | ArrayBuffer | ArrayBufferView | string | Blob | null; contentType: string; + customMetadata?: Record | undefined; objectKey: string; options?: PutObjectOptions | undefined; } diff --git a/apps/api/src/modules/files/infrastructure/r2-s3-object-client.ts b/apps/api/src/modules/files/infrastructure/r2-s3-object-client.ts index 3ba78a6c..2259430f 100644 --- a/apps/api/src/modules/files/infrastructure/r2-s3-object-client.ts +++ b/apps/api/src/modules/files/infrastructure/r2-s3-object-client.ts @@ -233,6 +233,7 @@ export async function copyObject(input: CopyObjectInput): Promise { const options = input.options ?? {}; const putOptions: R2PutOptions = { + ...(input.customMetadata === undefined ? {} : { customMetadata: input.customMetadata }), httpMetadata: { contentType: input.contentType, }, diff --git a/apps/api/src/modules/mcp/application/mcp-mappers.ts b/apps/api/src/modules/mcp/application/mcp-mappers.ts index dc269ecc..356d5cdf 100644 --- a/apps/api/src/modules/mcp/application/mcp-mappers.ts +++ b/apps/api/src/modules/mcp/application/mcp-mappers.ts @@ -120,7 +120,7 @@ function toCredentialSummary(row: CredentialRow): McpCredentialSummary { id: row.id, scope: row.scope, scopeValues: decodeJsonArray(row.scopeValuesJson), - status: getCredentialStatus(row) as Exclude, + status: getCredentialStatus(row), subjectLabel: row.subjectLabel, updatedAt: toIsoString(row.updatedAt), }; @@ -192,9 +192,7 @@ export function toAgentBinding( export function toOAuthFlowState(flow: OAuthFlowRow, server: ServerRow): McpOAuthFlowState { return { authorizationState: - flow.status === "succeeded" - ? ((server.enabled === 1 ? "active" : "disabled") as McpAuthorizationState) - : null, + flow.status === "succeeded" ? (server.enabled === 1 ? "active" : "disabled") : null, errorMessage: flow.errorMessage, flowId: flow.id, serverId: flow.serverId, diff --git a/apps/api/src/modules/public-api/app-agent-bound-ask.service.ts b/apps/api/src/modules/public-api/app-agent-bound-ask.service.ts index 218b769f..d3625d51 100644 --- a/apps/api/src/modules/public-api/app-agent-bound-ask.service.ts +++ b/apps/api/src/modules/public-api/app-agent-bound-ask.service.ts @@ -350,7 +350,6 @@ async function startBoundAgentRun(input: { if (createdSessionId !== null && input.recoverableSessionId === undefined) { await cleanupFailedThreadCreation({ bindings: input.bindings, - fileIds: [], sessionId: createdSessionId, }).catch((cleanupError: unknown) => { logError("public-api.bound_agent_call.cleanup_failed", { diff --git a/apps/api/src/modules/public-api/app-agent-capability.ts b/apps/api/src/modules/public-api/app-agent-capability.ts index 35d93d99..c8ed0535 100644 --- a/apps/api/src/modules/public-api/app-agent-capability.ts +++ b/apps/api/src/modules/public-api/app-agent-capability.ts @@ -15,6 +15,8 @@ import type { AgentId, AppDeploymentId, AppDeploymentRunId, AppId } from "@mosoo/id"; import { isPlatformId } from "@mosoo/id"; +import { fromBase64Url, toBase64Url } from "../../shared/bytes"; + export type AppAgentCapabilityExpose = "public_thread"; export interface AppAgentCapabilityBinding { @@ -56,24 +58,6 @@ export function boundAgentUrl(apiOrigin: string, token: string): string { return `${stripTrailingSlashes(apiOrigin)}${APP_AGENT_BOUND_PATH_PREFIX}/${token}`; } -function bytesToBase64Url(bytes: Uint8Array): string { - let binary = ""; - for (const byte of bytes) { - binary += String.fromCharCode(byte); - } - return btoa(binary).replaceAll("+", "-").replaceAll("/", "_").replaceAll("=", ""); -} - -function base64UrlToBytes(value: string): Uint8Array { - const normalized = value.replaceAll("-", "+").replaceAll("_", "/"); - const binary = atob(normalized); - const bytes = new Uint8Array(binary.length); - for (let index = 0; index < binary.length; index += 1) { - bytes[index] = binary.charCodeAt(index); - } - return bytes; -} - async function importSigningKey(secret: string): Promise { return crypto.subtle.importKey("raw", new TextEncoder().encode(secret), HMAC_PARAMS, false, [ "sign", @@ -112,14 +96,14 @@ export async function mintAppAgentCapabilityToken( secret: string, claims: AppAgentCapabilityClaims, ): Promise { - const payload = bytesToBase64Url(new TextEncoder().encode(JSON.stringify(claims))); + const payload = toBase64Url(new TextEncoder().encode(JSON.stringify(claims))); const key = await importSigningKey(secret); const signature = await crypto.subtle.sign( HMAC_PARAMS.name, key, new TextEncoder().encode(payload), ); - return `${payload}.${bytesToBase64Url(new Uint8Array(signature))}`; + return `${payload}.${toBase64Url(new Uint8Array(signature))}`; } /** @@ -159,7 +143,7 @@ export async function inspectAppAgentCapabilityToken( signatureValid = await crypto.subtle.verify( HMAC_PARAMS.name, key, - base64UrlToBytes(signaturePart), + fromBase64Url(signaturePart), new TextEncoder().encode(payload), ); } catch { @@ -171,7 +155,7 @@ export async function inspectAppAgentCapabilityToken( let parsed: unknown; try { - parsed = JSON.parse(new TextDecoder().decode(base64UrlToBytes(payload))); + parsed = JSON.parse(new TextDecoder().decode(fromBase64Url(payload))); } catch { return { status: "invalid" }; } diff --git a/apps/api/src/modules/public-api/public-api-idempotency.service.ts b/apps/api/src/modules/public-api/public-api-idempotency.service.ts index ae7cf869..5c63ec93 100644 --- a/apps/api/src/modules/public-api/public-api-idempotency.service.ts +++ b/apps/api/src/modules/public-api/public-api-idempotency.service.ts @@ -78,7 +78,7 @@ export async function hashPublicApiIdempotencyBody(value: unknown): Promise byte.toString(16).padStart(2, "0")).join(""); + return new Uint8Array(digest).toHex(); } function enforceSameIdempotentRequest( diff --git a/apps/api/src/modules/public-api/public-thread-create.ts b/apps/api/src/modules/public-api/public-thread-create.ts index c43024a1..93d40d64 100644 --- a/apps/api/src/modules/public-api/public-thread-create.ts +++ b/apps/api/src/modules/public-api/public-thread-create.ts @@ -162,7 +162,6 @@ export async function createPublicThread( if (createdSessionId !== null) { await cleanupFailedThreadCreation({ bindings: request.bindings, - fileIds: request.input.fileIds, sessionId: createdSessionId, }).catch((cleanupError: unknown) => { logError("public-api.thread.cleanup_failed", { diff --git a/apps/api/src/modules/public-api/public-thread-events.ts b/apps/api/src/modules/public-api/public-thread-events.ts index 0e250ba4..7af58ebe 100644 --- a/apps/api/src/modules/public-api/public-thread-events.ts +++ b/apps/api/src/modules/public-api/public-thread-events.ts @@ -12,16 +12,30 @@ import { import type { SessionProcessEvent } from "@mosoo/contracts/session"; import { parseJsonObject } from "@mosoo/contracts/validation"; import type { JsonObject } from "@mosoo/contracts/validation"; -import { sessionEventsTable, sessionMessagesTable } from "@mosoo/db"; +import { sessionEventsTable, sessionMessagesTable, sessionRunsTable } from "@mosoo/db"; import { parsePlatformId } from "@mosoo/id"; -import type { RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; -import { and, asc, desc, eq, gt, lt } from "drizzle-orm"; +import type { RuntimeEventId, SessionId, SessionMessageId, SessionRunId } from "@mosoo/id"; +import { createSessionRunTerminalSourceId } from "@mosoo/runtime-events"; +import { and, asc, desc, eq, gt, inArray, lt } from "drizzle-orm"; import type { SQL } from "drizzle-orm"; import { createErrorLogContext, logWarn } from "../../platform/cloudflare/logger"; import { getAppDatabase } from "../../platform/db/drizzle"; import { createSessionProcessEventsFromSessionEventRows } from "../sessions/application/session-process-events.service"; import type { SessionEventProcessRow } from "../sessions/application/session-process-events.service"; +import { + excludeSessionEventStreams, + findLeftIncompleteSessionEventStreamKeys, + foldStreamedSessionEventRows, + getMessageStreamTextFragment, + getSessionEventStreamKey, +} from "../sessions/domain/session-event-stream-fold"; +import { readTerminalEventSemanticAuthority } from "../sessions/domain/session-terminal-event-authority"; +import type { SessionMessageEventStreamCursor } from "../sessions/infrastructure/session-message-event-stream.repository"; +import { + iteratePublicSessionMessageEventRows, + readSealedPublicSessionMessage, +} from "../sessions/infrastructure/session-message-event-stream.repository"; import { connectSessionPublicEventWebSocket } from "../sessions/infrastructure/session/client"; import { publicInternalError, publicInvalidRequest, toPublicApiError } from "./public-api-errors"; import { sanitizePublicOutput } from "./public-output-sanitization"; @@ -44,8 +58,10 @@ const SSE_TEXT_ENCODER = new TextEncoder(); interface PublicThreadEventWindow { events: PublicThreadEventLogEntry[]; + leftIncompleteStreamKeys: Set; latestSeq: number | null; rows: PublicThreadEventProcessRow[]; + trustUnknownMessageDeltas: boolean; truncated: boolean; } @@ -54,130 +70,276 @@ interface PublicThreadEventProcessRow extends SessionEventProcessRow { tool_call_id: string | null; tool_input_json: string | null; tool_name: string | null; + tool_status: "cancelled" | "completed" | "failed" | "running" | null; +} + +interface LiveMessageSnapshot { + endSeq: number; + row: PublicThreadEventProcessRow; + sealed: boolean; + startSeq: number; } interface LiveMessageState { - filter: OpenAiPrivateCitationStreamFilter; - text: string; + draftEndSeq: number; + draftFilter: OpenAiPrivateCitationStreamFilter; + draftProjectedChars: number; + draftStartSeq: number; + emittedCursor: SessionMessageEventStreamCursor | null; + firstOccurredAt: number; + firstSeq: number; + messageTerminal: boolean; + runId: SessionRunId | null; + runTerminal: boolean; + snapshot: LiveMessageSnapshot | null; + trusted: boolean; } -class PublicLiveEventRowProjector { - readonly #messages = new Map(); - readonly #messageKeysByRun = new Map>(); +interface PublicMessageReconciliationRequest { + candidateCursor: SessionMessageEventStreamCursor; + emittedCursor: SessionMessageEventStreamCursor | null; + key: string; + outputRow: PublicThreadEventProcessRow; + sourceRow: PublicThreadEventProcessRow; +} - #getMessageRunKey(row: PublicThreadEventProcessRow): string { - return row.run_id ?? ""; - } +interface PublicMessageReconciliationResult { + compatible: boolean; + emitted: boolean; +} - #getMessageKey(row: PublicThreadEventProcessRow): string { - return `${this.#getMessageRunKey(row)}:${row.process_type}`; - } +function projectPublicMessageText(text: string, finish: boolean): string { + const filter = new OpenAiPrivateCitationStreamFilter(); + const projected = filter.push(text).text; + return finish ? projected + filter.finish().text : projected; +} - #setMessage(row: PublicThreadEventProcessRow, message: LiveMessageState): void { - const key = this.#getMessageKey(row); - const runKey = this.#getMessageRunKey(row); - const runKeys = this.#messageKeysByRun.get(runKey); +class PublicLiveEventRowProjector { + #isSeeding = false; + readonly #messages = new Map(); + readonly #trustedUnknownDeltaRuns = new Set(); + #trustUnknownMessageDeltas = true; + + seedFromCanonicalEvents( + rows: readonly PublicThreadEventProcessRow[], + events: readonly PublicThreadEventLogEntry[], + leftIncompleteStreamKeys: ReadonlySet = new Set(), + trustUnknownMessageDeltas = true, + ): PublicThreadEventLogEntry[] { + this.#trustUnknownMessageDeltas = trustUnknownMessageDeltas; + for (const row of rows) { + if (row.process_type !== "agent.message.delta" || !row.event_type.startsWith("message.")) { + continue; + } + const key = this.#getMessageKey(row); + if (leftIncompleteStreamKeys.has(key) && !this.#messages.has(key)) { + this.#messages.set(key, this.#createMessage(row, false)); + } + } - this.#messages.set(key, message); - if (runKeys === undefined) { - this.#messageKeysByRun.set(runKey, new Set([key])); - } else { - runKeys.add(key); + this.#isSeeding = true; + try { + this.project(rows); + } finally { + this.#isSeeding = false; } - } - #deleteMessage(row: PublicThreadEventProcessRow): void { - const key = this.#getMessageKey(row); - const runKey = this.#getMessageRunKey(row); - const runKeys = this.#messageKeysByRun.get(runKey); + const rowsByEventId = new Map(rows.map((row) => [row.id, row])); + return events.map((event) => { + const row = rowsByEventId.get(event.id); + if (row?.event_type.startsWith("message.") !== true) { + return event; + } - this.#messages.delete(key); - runKeys?.delete(key); - if (runKeys?.size === 0) { - this.#messageKeysByRun.delete(runKey); - } - } + const message = this.#messages.get(this.#getMessageKey(row)); + if (message === undefined) { + return event; + } - #deleteRunMessages(runId: SessionRunId | null): void { - const runKey = runId ?? ""; - const messageKeys = this.#messageKeysByRun.get(runKey); + const cursor = this.#currentCursor(message); + message.emittedCursor = cursor; + const content = projectPublicMessageText(event.content, cursor.finish); + return content === event.content ? event : { ...event, content }; + }); + } - if (messageKeys === undefined) { - return; + getReconciliationRequests( + row: PublicThreadEventProcessRow, + ): PublicMessageReconciliationRequest[] { + if (row.event_type === "message.completed") { + const key = this.#getMessageKey(row); + const message = this.#messages.get(key); + if ( + message?.snapshot === null || + message?.snapshot === undefined || + message.runTerminal || + message.snapshot.sealed || + !message.trusted + ) { + return []; + } + return [this.#createReconciliationRequest(key, message, row, true)]; } - for (const messageKey of messageKeys) { - this.#messages.delete(messageKey); + if ( + row.event_type !== "run.cancelled" && + row.event_type !== "run.completed" && + row.event_type !== "run.failed" + ) { + return []; } - this.#messageKeysByRun.delete(runKey); + const requests: PublicMessageReconciliationRequest[] = []; + for (const [key, message] of this.#messages) { + if ( + message.runId === row.run_id && + message.snapshot !== null && + !message.snapshot.sealed && + !message.runTerminal && + message.trusted + ) { + requests.push(this.#createReconciliationRequest(key, message, row, false)); + } + } + return requests; } - project(rows: readonly PublicThreadEventProcessRow[]): PublicThreadEventProcessRow[] { + project( + rows: readonly PublicThreadEventProcessRow[], + reconciled: ReadonlyMap = new Map(), + ): PublicThreadEventProcessRow[] { const output: PublicThreadEventProcessRow[] = []; for (const row of rows) { - const key = this.#getMessageKey(row); - if (row.event_type === "message.started") { - this.#setMessage(row, { - filter: new OpenAiPrivateCitationStreamFilter(), - text: "", - }); + const key = this.#getMessageKey(row); + if (!this.#messages.has(key)) { + this.#messages.set(key, this.#createMessage(row)); + } continue; } if (row.event_type === "message.delta") { + const key = this.#getMessageKey(row); let message = this.#messages.get(key); - if (message === undefined) { - message = { filter: new OpenAiPrivateCitationStreamFilter(), text: "" }; - this.#setMessage(row, message); + const runBoundaryIsTrusted = + row.run_id !== null && this.#trustedUnknownDeltaRuns.has(row.run_id); + if ( + !this.#trustUnknownMessageDeltas && + !runBoundaryIsTrusted && + row.stream_id !== row.id + ) { + continue; + } + message = this.#createMessage(row); + this.#messages.set(key, message); + } + if (!message.trusted) { + continue; + } + if (message.snapshot !== null) { + message.snapshot.endSeq = row.seq; + message.snapshot.sealed = false; + message.messageTerminal = false; + continue; } - const contentText = message.filter.push(row.content_text).text; - message.text += contentText; - + const projectedBefore = message.draftProjectedChars; + const contentText = message.draftFilter.push(row.content_text).text; + message.draftEndSeq = row.seq; + message.draftProjectedChars += contentText.length; + message.messageTerminal = false; if (contentText.length > 0) { + this.#recordDraftOutput(message, row.seq, projectedBefore); output.push({ ...row, content_text: contentText }); } continue; } - if (row.event_type === "message.completed") { - const message = this.#messages.get(key); - const contentText = message?.filter.finish().text ?? ""; - - if (message !== undefined) { - message.text += contentText; + if (row.event_type === "message.added") { + const key = this.#getMessageKey(row); + let message = this.#messages.get(key); + if (message === undefined && row.process_type === "agent.message.delta") { + message = this.#createMessage(row); + this.#messages.set(key, message); } - if (contentText.length > 0) { - output.push({ ...row, content_text: contentText }); + if (message === undefined) { + output.push({ ...row, content_text: sanitizePublicOutput(row.content_text).text }); + continue; } + if (message.runTerminal && !this.#isSeeding) { + message.trusted = false; + continue; + } + message.trusted = true; + message.messageTerminal = false; + message.snapshot = { + endSeq: row.seq, + row: { ...row, content_text: "" }, + sealed: false, + startSeq: row.seq, + }; continue; } - if (row.event_type === "message.added") { + if ( + row.event_type === "message.cancelled" || + row.event_type === "message.completed" || + row.event_type === "message.failed" + ) { + const key = this.#getMessageKey(row); const message = this.#messages.get(key); - this.#deleteMessage(row); - - if (message !== undefined) { - const snapshotText = sanitizePublicOutput(row.content_text).text; - - // SSE is append-only; a divergent snapshot stays canonical through - // run.finalOutput because emitted text cannot be retracted. - if (snapshotText.startsWith(message.text)) { - const suffix = snapshotText.slice(message.text.length); - - if (suffix.length > 0) { - output.push({ ...row, content_text: suffix }); - } + if (message === undefined || message.runTerminal) { + if (message === undefined) { + output.push({ ...row, content_text: "" }); } continue; } + if (!message.trusted && message.snapshot === null) { + message.messageTerminal = true; + continue; + } + const wasTerminal = message.messageTerminal; + let contentText = ""; + if (row.event_type === "message.completed" && message.snapshot !== null) { + message.snapshot.endSeq = row.seq; + message.snapshot.sealed = true; + if (!this.#isSeeding) { + const compatible = reconciled.get(key)?.compatible === true; + message.trusted = compatible; + if (compatible) { + message.emittedCursor = this.#snapshotCursor(message.snapshot, true); + } + } + } else { + message.snapshot = null; + if (!this.#isSeeding) { + contentText = message.draftFilter.finish().text; + const projectedBefore = message.draftProjectedChars; + message.draftProjectedChars += contentText.length; + message.draftEndSeq = row.seq; + if (contentText.length > 0) { + this.#recordDraftOutput(message, row.seq, projectedBefore); + } + if (message.emittedCursor?.startSeq === message.draftStartSeq) { + message.emittedCursor = { ...message.emittedCursor, endSeq: row.seq, finish: true }; + } + } + message.trusted = row.event_type === "message.completed"; + } + message.messageTerminal = true; + if ((!wasTerminal || contentText.length > 0) && reconciled.get(key)?.emitted !== true) { + output.push({ ...row, content_text: contentText }); + } + continue; } - if (row.event_type === "thought.started" || row.event_type === "thought.completed") { + if ( + row.event_type === "thought.cancelled" || + row.event_type === "thought.completed" || + row.event_type === "thought.started" + ) { continue; } @@ -186,7 +348,41 @@ class PublicLiveEventRowProjector { row.event_type === "run.completed" || row.event_type === "run.failed" ) { - this.#deleteRunMessages(row.run_id); + for (const [key, message] of this.#messages) { + if (message.runId !== row.run_id) { + continue; + } + if (message.snapshot !== null && !message.snapshot.sealed && !this.#isSeeding) { + const compatible = reconciled.get(key)?.compatible === true; + message.trusted = compatible; + if (compatible) { + message.emittedCursor = this.#snapshotCursor(message.snapshot, false); + } + } + message.runTerminal = true; + } + if (row.run_id !== null) { + this.#trustedUnknownDeltaRuns.delete(row.run_id); + } + output.push(row); + continue; + } + + if (row.event_type === "run.started") { + let released = false; + for (const [key, message] of this.#messages) { + if (message.runId !== row.run_id) { + this.#messages.delete(key); + released = true; + } + } + if (released) { + this.#trustUnknownMessageDeltas = false; + } + if (row.run_id !== null) { + this.#trustedUnknownDeltaRuns.clear(); + this.#trustedUnknownDeltaRuns.add(row.run_id); + } } output.push(row); @@ -194,6 +390,261 @@ class PublicLiveEventRowProjector { return output; } + + #createMessage(row: PublicThreadEventProcessRow, trusted = true): LiveMessageState { + return { + draftEndSeq: row.seq, + draftFilter: new OpenAiPrivateCitationStreamFilter(), + draftProjectedChars: 0, + draftStartSeq: row.seq, + emittedCursor: null, + firstOccurredAt: row.occurred_at, + firstSeq: row.seq, + messageTerminal: false, + runId: row.run_id, + runTerminal: false, + snapshot: null, + trusted, + }; + } + + #getMessageKey(row: PublicThreadEventProcessRow): string { + const key = getSessionEventStreamKey(row); + if (key === null) { + throw new Error(`Persisted ${row.event_type} event is missing its stream ID.`); + } + return key; + } + + #currentCursor(message: LiveMessageState): SessionMessageEventStreamCursor { + return message.snapshot === null + ? { + endSeq: message.draftEndSeq, + finish: message.messageTerminal, + outputOffset: 0, + startSeq: message.draftStartSeq, + } + : this.#snapshotCursor(message.snapshot, message.snapshot.sealed); + } + + #snapshotCursor(snapshot: LiveMessageSnapshot, finish: boolean): SessionMessageEventStreamCursor { + return { + endSeq: snapshot.endSeq, + finish, + outputOffset: 0, + startSeq: snapshot.startSeq, + }; + } + + #createReconciliationRequest( + key: string, + message: LiveMessageState, + boundaryRow: PublicThreadEventProcessRow, + finish: boolean, + ): PublicMessageReconciliationRequest { + const snapshot = message.snapshot; + if (snapshot === null) { + throw new Error("Cannot reconcile a message without an authoritative snapshot."); + } + return { + candidateCursor: { + ...this.#snapshotCursor(snapshot, finish), + endSeq: finish ? boundaryRow.seq : snapshot.endSeq, + }, + emittedCursor: message.emittedCursor, + key, + outputRow: { + ...(finish ? boundaryRow : snapshot.row), + event_type: "message.delta", + occurred_at: message.firstOccurredAt, + seq: message.firstSeq, + }, + sourceRow: snapshot.row, + }; + } + + #recordDraftOutput(message: LiveMessageState, endSeq: number, outputOffset: number): void { + if (message.emittedCursor === null) { + message.emittedCursor = { + endSeq, + finish: false, + outputOffset, + startSeq: message.draftStartSeq, + }; + } else if (message.emittedCursor.startSeq === message.draftStartSeq) { + message.emittedCursor = { ...message.emittedCursor, endSeq, finish: false }; + } + } +} + +interface ProjectedMessageChunk { + row: PublicThreadEventProcessRow; + text: string; +} + +async function* iterateProjectedMessageChunks(input: { + cursor: SessionMessageEventStreamCursor; + database: D1Database; + row: PublicThreadEventProcessRow; + sessionId: SessionId; +}): AsyncGenerator { + if (input.row.stream_id === null) { + throw new Error(`Persisted ${input.row.event_type} event is missing its stream ID.`); + } + + let filter = new OpenAiPrivateCitationStreamFilter(); + let lastContentRow: PublicThreadEventProcessRow | null = null; + let outputOffset = input.cursor.outputOffset; + let pending: ProjectedMessageChunk | null = null; + + for await (const row of iteratePublicSessionMessageEventRows(input.database, { + cursor: input.cursor, + processType: input.row.process_type, + runId: input.row.run_id, + sessionId: input.sessionId, + streamId: input.row.stream_id, + })) { + const fragment = getMessageStreamTextFragment(row); + if (fragment === null) { + continue; + } + if (fragment.kind === "reset") { + filter = new OpenAiPrivateCitationStreamFilter(); + } + + const projectedRow: PublicThreadEventProcessRow = { + ...row, + tool_call_id: null, + tool_input_json: null, + tool_name: null, + tool_status: null, + }; + lastContentRow = projectedRow; + let text = filter.push(fragment.text).text; + if (outputOffset >= text.length) { + outputOffset -= text.length; + text = ""; + } else if (outputOffset > 0) { + text = text.slice(outputOffset); + outputOffset = 0; + } + if (text.length === 0) { + continue; + } + if (pending !== null) { + yield pending; + } + pending = { row: projectedRow, text }; + } + + if (input.cursor.finish) { + let tail = filter.finish().text; + if (outputOffset >= tail.length) { + outputOffset -= tail.length; + tail = ""; + } else if (outputOffset > 0) { + tail = tail.slice(outputOffset); + outputOffset = 0; + } + if (tail.length > 0) { + if (pending === null) { + if (lastContentRow === null) { + throw new Error("Cannot finish a projected message without a content row."); + } + pending = { row: lastContentRow, text: tail }; + } else { + pending = { ...pending, text: pending.text + tail }; + } + } + } + + if (pending !== null) { + yield pending; + } +} + +async function reconcilePublicMessage(input: { + database: D1Database; + emit: (chunk: ProjectedMessageChunk) => void; + request: PublicMessageReconciliationRequest; + sessionId: SessionId; +}): Promise { + const suffix: string[] = []; + const emitSuffix = () => { + if (suffix.length === 0) { + return false; + } + input.emit({ row: input.request.outputRow, text: suffix.join("") }); + return true; + }; + const candidate = iterateProjectedMessageChunks({ + cursor: input.request.candidateCursor, + database: input.database, + row: input.request.sourceRow, + sessionId: input.sessionId, + })[Symbol.asyncIterator](); + + if (input.request.emittedCursor === null) { + let chunk = await candidate.next(); + while (!chunk.done) { + suffix.push(chunk.value.text); + chunk = await candidate.next(); + } + return { compatible: true, emitted: emitSuffix() }; + } + + const emitted = iterateProjectedMessageChunks({ + cursor: input.request.emittedCursor, + database: input.database, + row: input.request.sourceRow, + sessionId: input.sessionId, + })[Symbol.asyncIterator](); + let candidateChunk = await candidate.next(); + let candidateOffset = 0; + + for (;;) { + const emittedChunk = await emitted.next(); + if (emittedChunk.done) { + break; + } + + let emittedOffset = 0; + while (emittedOffset < emittedChunk.value.text.length) { + if (candidateChunk.done) { + await emitted.return?.(undefined); + return { compatible: false, emitted: false }; + } + const comparedLength = Math.min( + emittedChunk.value.text.length - emittedOffset, + candidateChunk.value.text.length - candidateOffset, + ); + if ( + emittedChunk.value.text.slice(emittedOffset, emittedOffset + comparedLength) !== + candidateChunk.value.text.slice(candidateOffset, candidateOffset + comparedLength) + ) { + await candidate.return?.(undefined); + await emitted.return?.(undefined); + return { compatible: false, emitted: false }; + } + emittedOffset += comparedLength; + candidateOffset += comparedLength; + if (candidateOffset === candidateChunk.value.text.length) { + candidateChunk = await candidate.next(); + candidateOffset = 0; + } + } + } + + if (!candidateChunk.done && candidateOffset < candidateChunk.value.text.length) { + suffix.push(candidateChunk.value.text.slice(candidateOffset)); + } + for (;;) { + candidateChunk = await candidate.next(); + if (candidateChunk.done) { + return { compatible: true, emitted: emitSuffix() }; + } + suffix.push(candidateChunk.value.text); + } } class PublicThreadEventWakeup { @@ -327,7 +778,11 @@ function toPublicThreadEventLogEntry(input: { let toolInput: JsonObject | undefined; if (input.row?.tool_input_json !== null && input.row?.tool_input_json !== undefined) { - toolInput = parseJsonObject(JSON.parse(input.row.tool_input_json), "Persisted tool input"); + try { + toolInput = parseJsonObject(JSON.parse(input.row.tool_input_json), "Persisted tool input"); + } catch { + toolInput = undefined; + } } return { @@ -351,13 +806,12 @@ function toPublicThreadEventLogEntry(input: { function toPublicThreadEventLogEntries( rows: PublicThreadEventProcessRow[], - options: { foldStreamedRows?: boolean } = {}, ): PublicThreadEventLogEntry[] { const rowsByEventId = new Map( rows.map((row) => [row.id, row]), ); - return createSessionProcessEventsFromSessionEventRows(rows, options).flatMap((event) => { + return createSessionProcessEventsFromSessionEventRows(rows).flatMap((event) => { const publicEvent = toPublicThreadEventLogEntry({ event, row: rowsByEventId.get(event.id), @@ -366,6 +820,15 @@ function toPublicThreadEventLogEntries( }); } +function toCanonicalPublicThreadEventLogEntries( + rows: PublicThreadEventProcessRow[], +): PublicThreadEventLogEntry[] { + return new PublicLiveEventRowProjector().seedFromCanonicalEvents( + rows, + toPublicThreadEventLogEntries(rows), + ); +} + function selectPublicThreadEventRows(input: { database: D1Database; filters: SQL[]; @@ -383,9 +846,11 @@ function selectPublicThreadEventRows(input: { process_type: sessionEventsTable.processType, run_id: sessionEventsTable.runId, seq: sessionEventsTable.seq, + stream_id: sessionEventsTable.streamId, tool_call_id: sessionEventsTable.toolCallId, tool_input_json: sessionEventsTable.toolInputJson, tool_name: sessionEventsTable.toolName, + tool_status: sessionEventsTable.toolStatus, tokens: sessionEventsTable.tokens, }) .from(sessionEventsTable) @@ -407,7 +872,8 @@ async function readPublicThreadEventWindow(input: { while (scannedRows.length < THREAD_EVENT_RAW_ROW_SCAN_LIMIT) { const remainingRows = THREAD_EVENT_RAW_ROW_SCAN_LIMIT - scannedRows.length; - const pageSize = Math.min(THREAD_EVENT_ROW_PAGE_SIZE, remainingRows); + const rowCapacity = Math.min(THREAD_EVENT_ROW_PAGE_SIZE, remainingRows); + const querySize = rowCapacity + 1; const filters = [ eq(sessionEventsTable.sessionId, input.sessionId), eq(sessionEventsTable.visibility, "all_consumers"), @@ -421,7 +887,7 @@ async function readPublicThreadEventWindow(input: { database: input.database, filters, order: desc(sessionEventsTable.seq), - pageSize, + pageSize: querySize, }); if (page.length === 0) { @@ -429,34 +895,60 @@ async function readPublicThreadEventWindow(input: { break; } - latestSeq ??= page[0]?.seq ?? null; - scannedRows.push(...page); - beforeSeq = page[page.length - 1]?.seq ?? beforeSeq; - - const events = toPublicThreadEventLogEntries(scannedRows.toReversed()); + const scannedPage = page.slice(0, rowCapacity); + latestSeq ??= scannedPage[0]?.seq ?? null; + scannedRows.push(...scannedPage); + beforeSeq = scannedPage[scannedPage.length - 1]?.seq ?? beforeSeq; + const chronologicalRows = scannedRows.toReversed(); + const foldedRows = foldStreamedSessionEventRows(chronologicalRows); + const events = toCanonicalPublicThreadEventLogEntries(chronologicalRows); if (events.length > input.limit) { + const reachedDatabaseStart = page.length <= rowCapacity; + const incompleteStreams = reachedDatabaseStart + ? new Set() + : findLeftIncompleteSessionEventStreamKeys(chronologicalRows); + const rowsByEventId = new Map(foldedRows.map((row) => [row.id, row])); + const retainedStreamsAreComplete = events.slice(-input.limit).every((event) => { + const row = rowsByEventId.get(event.id); + const key = row === undefined ? null : getSessionEventStreamKey(row); + return key === null || !incompleteStreams.has(key); + }); + + if (!retainedStreamsAreComplete) { + continue; + } + return { events: events.slice(-input.limit), + leftIncompleteStreamKeys: incompleteStreams, latestSeq, - rows: scannedRows.toReversed(), + rows: chronologicalRows, + trustUnknownMessageDeltas: reachedDatabaseStart, truncated: true, }; } - if (page.length < pageSize) { + if (page.length <= rowCapacity) { reachedStart = true; break; } } - const events = toPublicThreadEventLogEntries(scannedRows.toReversed()); + const chronologicalRows = scannedRows.toReversed(); + const incompleteStreams = reachedStart + ? new Set() + : findLeftIncompleteSessionEventStreamKeys(chronologicalRows); + const completeRows = excludeSessionEventStreams(chronologicalRows, incompleteStreams); + const events = toCanonicalPublicThreadEventLogEntries(completeRows); const truncated = !reachedStart || events.length > input.limit; return { events: truncated ? events.slice(-input.limit) : events, + leftIncompleteStreamKeys: incompleteStreams, latestSeq, - rows: scannedRows.toReversed(), + rows: chronologicalRows, + trustUnknownMessageDeltas: reachedStart, truncated, }; } @@ -483,10 +975,155 @@ export async function readPublicThreadRunFinalOutput(input: { runId: SessionRunId; sessionId: SessionId; }): Promise { + const database = getAppDatabase(input.database); + const snapshotRows = await database + .select({ + eventType: sessionEventsTable.eventType, + runStatus: sessionRunsTable.status, + semanticHash: sessionEventsTable.semanticHash, + seq: sessionEventsTable.seq, + sourceEventId: sessionEventsTable.sourceEventId, + streamId: sessionEventsTable.streamId, + terminalEventJson: sessionEventsTable.terminalEventJson, + }) + .from(sessionRunsTable) + .leftJoin( + sessionEventsTable, + and( + eq(sessionEventsTable.sessionId, sessionRunsTable.sessionId), + eq(sessionEventsTable.runId, sessionRunsTable.id), + inArray(sessionEventsTable.eventType, ["run.cancelled", "run.completed", "run.failed"]), + ), + ) + .where( + and(eq(sessionRunsTable.id, input.runId), eq(sessionRunsTable.sessionId, input.sessionId)), + ) + .limit(3) + .all(); + const [snapshot] = snapshotRows; + if (snapshot?.runStatus !== "completed") { + return null; + } + const terminalRows = snapshotRows.flatMap((row) => + row.eventType === null + ? [] + : [ + { + eventType: row.eventType, + semanticHash: row.semanticHash, + seq: row.seq, + sourceEventId: row.sourceEventId, + streamId: row.streamId, + terminalEventJson: row.terminalEventJson, + }, + ], + ); + if (terminalRows.length !== 1) { + throw new Error(`Final assistant reference for run ${input.runId} is not immutable.`); + } + + const terminal = terminalRows[0]; + if (terminal === undefined) { + throw new Error(`Final assistant reference for run ${input.runId} is not immutable.`); + } + if (terminal.semanticHash !== null) { + if ( + terminal.eventType !== "run.completed" || + terminal.seq === null || + terminal.sourceEventId !== createSessionRunTerminalSourceId(input.runId, "run.completed") + ) { + throw new Error(`Final assistant reference for run ${input.runId} is invalid.`); + } + const semanticAuthority = await readTerminalEventSemanticAuthority({ + eventJson: terminal.terminalEventJson, + eventType: terminal.eventType, + runId: input.runId, + semanticHash: terminal.semanticHash, + sessionId: input.sessionId, + sourceEventId: terminal.sourceEventId, + streamId: terminal.streamId, + }); + if (semanticAuthority.finalMessageId === null) { + return null; + } + const finalMessageId = parsePlatformId( + semanticAuthority.finalMessageId, + "Final assistant message ID", + ); + const message = + (await database + .select({ + content_text: sessionMessagesTable.contentText, + id: sessionMessagesTable.id, + plan_json: sessionMessagesTable.planJson, + projection_format: sessionMessagesTable.projectionFormat, + segments_json: sessionMessagesTable.segmentsJson, + }) + .from(sessionMessagesTable) + .where( + and( + eq(sessionMessagesTable.id, finalMessageId), + eq(sessionMessagesTable.sessionId, input.sessionId), + eq(sessionMessagesTable.sessionRunId, input.runId), + eq(sessionMessagesTable.role, "assistant"), + ), + ) + .limit(1) + .get()) ?? null; + if ( + message === null || + message.content_text !== "" || + message.plan_json !== null || + message.projection_format !== "event_stream_v3" || + message.segments_json !== null + ) { + throw new Error(`Final assistant reference ${finalMessageId} is not lightweight.`); + } + const sealed = await readSealedPublicSessionMessage(input.database, { + endSeq: terminal.seq, + processType: "agent.message.delta", + runId: input.runId, + sessionId: input.sessionId, + streamId: finalMessageId, + }); + if (sealed === null) { + throw new Error(`Final assistant reference ${finalMessageId} is not sealed.`); + } + const sanitizedOutput = sanitizePublicOutput(sealed.text); + return { + text: sanitizedOutput.text, + ...(sanitizedOutput.warnings.length === 0 ? {} : { warnings: sanitizedOutput.warnings }), + }; + } + + if ( + terminal.eventType !== "run.completed" || + terminal.sourceEventId !== createSessionRunTerminalSourceId(input.runId, "run.completed") + ) { + throw new Error(`Legacy final assistant projection for run ${input.runId} is invalid.`); + } + + const eventStreamReference = await database + .select({ id: sessionMessagesTable.id }) + .from(sessionMessagesTable) + .where( + and( + eq(sessionMessagesTable.sessionId, input.sessionId), + eq(sessionMessagesTable.sessionRunId, input.runId), + eq(sessionMessagesTable.role, "assistant"), + eq(sessionMessagesTable.projectionFormat, "event_stream_v3"), + ), + ) + .limit(1) + .get(); + if (eventStreamReference !== undefined) { + throw new Error(`Final assistant reference for run ${input.runId} has no terminal pointer.`); + } + const message = - (await getAppDatabase(input.database) + (await database .select({ - content: sessionMessagesTable.contentText, + content_text: sessionMessagesTable.contentText, }) .from(sessionMessagesTable) .where( @@ -494,6 +1131,7 @@ export async function readPublicThreadRunFinalOutput(input: { eq(sessionMessagesTable.sessionId, input.sessionId), eq(sessionMessagesTable.sessionRunId, input.runId), eq(sessionMessagesTable.role, "assistant"), + eq(sessionMessagesTable.projectionFormat, "materialized"), ), ) .orderBy(desc(sessionMessagesTable.seq)) @@ -503,8 +1141,7 @@ export async function readPublicThreadRunFinalOutput(input: { if (message === null) { return null; } - - const sanitizedOutput = sanitizePublicOutput(message.content); + const sanitizedOutput = sanitizePublicOutput(message.content_text); return { text: sanitizedOutput.text, @@ -553,21 +1190,62 @@ function enqueueSseText(controller: ReadableStreamDefaultController, function enqueueThreadEvents(input: { controller: ReadableStreamDefaultController; - emittedEventIds: Set; events: PublicThreadEventLogEntry[]; }): boolean { - let enqueued = false; - for (const event of input.events) { - if (input.emittedEventIds.has(event.id)) { - continue; + enqueueSseText(input.controller, encodeSseThreadEvent(event)); + } + + return input.events.length > 0; +} + +async function enqueueLiveThreadRows(input: { + controller: ReadableStreamDefaultController; + database: D1Database; + projector: PublicLiveEventRowProjector; + rows: readonly PublicThreadEventProcessRow[]; + sessionId: SessionId; +}): Promise { + let enqueued = false; + let pendingRows: PublicThreadEventProcessRow[] = []; + const flushPendingRows = () => { + if (pendingRows.length === 0) { + return; } + enqueued = + enqueueThreadEvents({ + controller: input.controller, + events: toPublicThreadEventLogEntries(pendingRows), + }) || enqueued; + pendingRows = []; + }; - input.emittedEventIds.add(event.id); - enqueueSseText(input.controller, encodeSseThreadEvent(event)); - enqueued = true; + for (const row of input.rows) { + const reconciled = new Map(); + const requests = input.projector.getReconciliationRequests(row); + if (requests.length > 0) { + flushPendingRows(); + } + for (const request of requests) { + const result = await reconcilePublicMessage({ + database: input.database, + emit: (chunk) => { + enqueued = + enqueueThreadEvents({ + controller: input.controller, + events: toPublicThreadEventLogEntries([{ ...chunk.row, content_text: chunk.text }]), + }) || enqueued; + }, + request, + sessionId: input.sessionId, + }); + reconciled.set(request.key, result); + } + pendingRows.push(...input.projector.project([row], reconciled)); } + flushPendingRows(); + return enqueued; } @@ -607,15 +1285,17 @@ export async function createPublicThreadEventStream( wakeup.close(); throw error; } - const emittedEventIds = new Set(); const state = { cancelled: false, }; let lastSeenSeq = initialWindow.latestSeq ?? 0; const liveProjector = new PublicLiveEventRowProjector(); - const initialEvents = toPublicThreadEventLogEntries( - liveProjector.project(initialWindow.rows), - ).slice(-limit); + const initialEvents = liveProjector.seedFromCanonicalEvents( + initialWindow.rows, + initialWindow.events, + initialWindow.leftIncompleteStreamKeys, + initialWindow.trustUnknownMessageDeltas, + ); return new ReadableStream({ async start(controller) { @@ -625,7 +1305,6 @@ export async function createPublicThreadEventStream( enqueueSseText(controller, encodeSseComment("connected")); enqueueThreadEvents({ controller, - emittedEventIds, events: initialEvents, }); @@ -651,11 +1330,13 @@ export async function createPublicThreadEventStream( lastSeenSeq = rows[rows.length - 1]?.seq ?? lastSeenSeq; enqueuedEvents = - enqueueThreadEvents({ + (await enqueueLiveThreadRows({ controller, - emittedEventIds, - events: toPublicThreadEventLogEntries(liveProjector.project(rows)), - }) || enqueuedEvents; + database: request.database, + projector: liveProjector, + rows, + sessionId, + })) || enqueuedEvents; if (rows.length < THREAD_EVENT_ROW_PAGE_SIZE) { break; diff --git a/apps/api/src/modules/public-api/public-thread-store.ts b/apps/api/src/modules/public-api/public-thread-store.ts index 93549cf2..a464f19c 100644 --- a/apps/api/src/modules/public-api/public-thread-store.ts +++ b/apps/api/src/modules/public-api/public-thread-store.ts @@ -1,19 +1,14 @@ import type { SessionSummary } from "@mosoo/contracts/session"; -import { - sessionEventsTable, - sessionMessagesTable, - sessionRunsTable, - sessionsTable, -} from "@mosoo/db"; +import { sessionRunsTable, sessionsTable } from "@mosoo/db"; import { parsePlatformId } from "@mosoo/id"; -import type { AccountId, AgentId, FileId, PublicThreadId, SessionId } from "@mosoo/id"; +import type { AccountId, AgentId, PublicThreadId, SessionId } from "@mosoo/id"; import { and, eq, sql } from "drizzle-orm"; import type { SQL } from "drizzle-orm"; import type { ApiBindings } from "../../platform/cloudflare/worker-types"; import { getAppDatabase } from "../../platform/db/drizzle"; import { currentTimestampMs, toIsoString } from "../../time"; -import { fileStore } from "../files/application/file-store"; +import { deleteSessionCascade } from "../sessions/application/session-cleanup.service"; import { buildSessionSummaryFromJoinedRow, sessionSummaryWithLastRunColumns, @@ -46,35 +41,9 @@ export interface ThreadSnapshot { export async function cleanupFailedThreadCreation(input: { bindings: ApiBindings; - fileIds: FileId[]; sessionId: SessionId; }): Promise { - if (input.fileIds.length > 0) { - await fileStore.deleteScope(input.bindings, { - id: input.sessionId, - kind: "session", - }); - } - - await getAppDatabase(input.bindings.DB) - .delete(sessionEventsTable) - .where(eq(sessionEventsTable.sessionId, input.sessionId)) - .run(); - - await getAppDatabase(input.bindings.DB) - .delete(sessionMessagesTable) - .where(eq(sessionMessagesTable.sessionId, input.sessionId)) - .run(); - - await getAppDatabase(input.bindings.DB) - .delete(sessionRunsTable) - .where(eq(sessionRunsTable.sessionId, input.sessionId)) - .run(); - - await getAppDatabase(input.bindings.DB) - .delete(sessionsTable) - .where(eq(sessionsTable.id, input.sessionId)) - .run(); + await deleteSessionCascade(input.bindings, input.sessionId); } export async function getThreadSnapshot( diff --git a/apps/api/src/modules/runtime/application/execution-plane/execution-plane-adapter.ts b/apps/api/src/modules/runtime/application/execution-plane/execution-plane-adapter.ts index 04a024ad..772b1247 100644 --- a/apps/api/src/modules/runtime/application/execution-plane/execution-plane-adapter.ts +++ b/apps/api/src/modules/runtime/application/execution-plane/execution-plane-adapter.ts @@ -1,4 +1,3 @@ -import type { RunError, SessionRunStatus } from "@mosoo/contracts/session-run"; import type { AgentId, DriverInstanceId, @@ -16,6 +15,7 @@ import type { RuntimeTimingSnapshot } from "../session-runs/session-runtime-timi import type { DriverBootPayloadPreparedHandler } from "./driver-boot-payload-prepared"; export interface RuntimeExecutionPlaneRunLease { + driverGeneration: number; driverInstanceId: DriverInstanceId; timing: RuntimeTimingSnapshot; readiness(): Promise; @@ -39,6 +39,7 @@ export interface PrepareRuntimeRunInput { export interface DispatchRuntimeTurnInput { attachmentIds: FileId[]; + driverGeneration: number; driverInstanceId: DriverInstanceId; prompt: string; sessionRunId: SessionRunId; @@ -47,13 +48,9 @@ export interface DispatchRuntimeTurnInput { export interface StopRuntimeSubjectDriversInput { operationId?: RuntimeOperationId; runtimeSubjectId: SandboxId; - preserveSessionLifecycle?: boolean; + sandboxIncarnation?: number; reason: string; targets?: readonly RuntimeSubjectOperationSessionTarget[]; - terminalRun?: { - error?: RunError | null; - status: Extract; - }; } export interface RuntimeSubjectOperationInput { @@ -61,10 +58,6 @@ export interface RuntimeSubjectOperationInput { runtimeSubjectId: SandboxId; reason: string; targets: readonly RuntimeSubjectOperationSessionTarget[]; - terminalRun: { - error?: RunError | null; - status: Extract; - }; } export interface RuntimeSubjectOperationSessionTarget { diff --git a/apps/api/src/modules/runtime/application/owner-debug-terminal.service.ts b/apps/api/src/modules/runtime/application/owner-debug-terminal.service.ts index 05cc9160..98c2c641 100644 --- a/apps/api/src/modules/runtime/application/owner-debug-terminal.service.ts +++ b/apps/api/src/modules/runtime/application/owner-debug-terminal.service.ts @@ -2,6 +2,7 @@ import type { PtyOptions } from "@cloudflare/sandbox"; import { parsePlatformId } from "@mosoo/id"; import type { AccountId, AgentId } from "@mosoo/id"; +import { withDisposedRpcResource } from "../../../platform/cloudflare/rpc-disposal"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; import { API_ERROR_CODE, ApiError, createApiError } from "../../../platform/errors"; import { ensureAgentOwner } from "../../agents/application/agent-access.service"; @@ -88,9 +89,11 @@ export async function connectOwnerDebugTerminalWebSocket( networkConstraints: { allowedHosts: [], networkPolicy: "full" }, }); - return connectPreparedSandboxTerminal(activation.subject, { - options, - request: input.request, - terminalSessionId: target.terminalSessionId, - }); + return withDisposedRpcResource(activation.subject, (subject) => + connectPreparedSandboxTerminal(subject, { + options, + request: input.request, + terminalSessionId: target.terminalSessionId, + }), + ); } diff --git a/apps/api/src/modules/runtime/application/runtime-diagnostic-events.ts b/apps/api/src/modules/runtime/application/runtime-diagnostic-events.ts index 0451dde4..6afb350f 100644 --- a/apps/api/src/modules/runtime/application/runtime-diagnostic-events.ts +++ b/apps/api/src/modules/runtime/application/runtime-diagnostic-events.ts @@ -30,6 +30,7 @@ export interface RuntimeDiagnosticEventInput< TName extends RuntimeDiagnosticEventName = RuntimeDiagnosticEventName, > { eventName: TName; + sourceEventId?: string; value: RuntimeDiagnosticEventValue; } @@ -76,11 +77,18 @@ export async function appendRuntimeDiagnosticEvent; }, ): Promise { return appendRuntimeDiagnosticEvents(bindings, { - events: [{ eventName: input.eventName, value: input.value }], + events: [ + { + eventName: input.eventName, + ...(input.sourceEventId === undefined ? {} : { sourceEventId: input.sourceEventId }), + value: input.value, + }, + ], sessionId: input.sessionId, }); } @@ -104,6 +112,7 @@ export async function appendRuntimeDiagnosticEvents( createRuntimeDiagnosticSessionEvent({ eventName: event.eventName, sessionId: input.sessionId, + ...(event.sourceEventId === undefined ? {} : { sourceEventId: event.sourceEventId }), value: event.value, }), ), @@ -172,6 +181,7 @@ function createRuntimeDiagnosticSessionEvent([ + "recreateSandbox", + "resetAgentState", + "restartDriver", +]); + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function canonicalError(input: { + readonly operationId: RuntimeOperationId; + readonly sessionId: SessionId; + readonly status: RuntimeOperationEventStatus; +}): Error { + return new Error( + `Runtime operation ${input.status} source runtime-operation:${input.operationId}:${input.sessionId}:${input.status} is not canonical.`, + ); +} + +export function readRuntimeOperationEventIdentity( + event: RuntimeEventEnvelope, + input: { + readonly agentId: AgentId; + readonly operationId: RuntimeOperationId; + readonly sessionId: SessionId; + readonly status: RuntimeOperationEventStatus; + }, +): RuntimeOperationEventIdentity { + const sourceEventId = `runtime-operation:${input.operationId}:${input.sessionId}:${input.status}`; + if ( + event.actor !== "api" || + event.delivery !== "lossless" || + event.kind !== "agent.task.updated" || + event.origin !== "api" || + event.sessionId !== input.sessionId || + event.sourceEventId !== sourceEventId || + event.visibility !== "participant" || + event.context !== undefined || + event.correlationId !== undefined || + event.driverInstanceId !== undefined || + event.native !== undefined || + event.receivedAt !== undefined || + event.runId !== undefined || + event.runtimeId !== undefined || + event.seq !== undefined || + event.traceId !== undefined || + !isRecord(event.payload) + ) { + throw canonicalError(input); + } + + const payload = event.payload; + const timingField = input.status === "ready" ? "readyAt" : "startedAt"; + const allowedKeys = new Set([ + "agentId", + "deploymentVersionId", + "deploymentVersionNumber", + "operation", + "operationId", + "status", + timingField, + ]); + if (Object.keys(payload).some((key) => !allowedKeys.has(key))) { + throw canonicalError(input); + } + + let agentId: AgentId; + let deploymentVersionId: AgentDeploymentVersionId | null = null; + let operationId: RuntimeOperationId; + try { + agentId = parsePlatformId(payload["agentId"], "Runtime operation Agent ID"); + operationId = parsePlatformId( + payload["operationId"], + "Runtime operation ID", + ); + if (payload["deploymentVersionId"] !== undefined) { + deploymentVersionId = parsePlatformId( + payload["deploymentVersionId"], + "Runtime operation deployment version ID", + ); + } + } catch { + throw canonicalError(input); + } + + const operation = payload["operation"]; + const observedAt = payload[timingField]; + const hasDeploymentVersion = payload["deploymentVersionId"] !== undefined; + const deploymentVersionNumber = payload["deploymentVersionNumber"]; + if ( + agentId !== input.agentId || + operationId !== input.operationId || + typeof operation !== "string" || + !runtimeStateOperationNames.has(operation as RuntimeStateOperationName) || + payload["status"] !== input.status || + typeof observedAt !== "string" || + observedAt !== event.occurredAt || + !Number.isFinite(Date.parse(observedAt)) || + hasDeploymentVersion !== (deploymentVersionNumber !== undefined) || + (hasDeploymentVersion && + (!Number.isSafeInteger(deploymentVersionNumber) || (deploymentVersionNumber as number) <= 0)) + ) { + throw canonicalError(input); + } + + return { + agentId, + deploymentVersionId, + deploymentVersionNumber: hasDeploymentVersion ? (deploymentVersionNumber as number) : null, + operation: operation as RuntimeStateOperationName, + }; +} + +export function createRuntimeOperationEventAuthorityJson(input: { + readonly agentId: AgentId; + readonly event: RuntimeEventEnvelope; + readonly operationId: RuntimeOperationId; + readonly sessionId: SessionId; + readonly status: RuntimeOperationEventStatus; +}): string { + readRuntimeOperationEventIdentity(input.event, input); + return stringifyRuntimeEventSemanticValue(input.event); +} + +export async function readRuntimeOperationEventAuthority(input: { + readonly agentId: AgentId; + readonly eventId: string; + readonly eventJson: string | null; + readonly eventType: string; + readonly occurredAt: number; + readonly operationId: RuntimeOperationId; + readonly rowAgentId: AgentId; + readonly semanticHash: string | null; + readonly sessionId: SessionId; + readonly source: string; + readonly sourceEventId: string; + readonly status: RuntimeOperationEventStatus; + readonly visibility: string; +}): Promise<{ + readonly agentId: AgentId; + readonly deploymentVersionId: AgentDeploymentVersionId | null; + readonly deploymentVersionNumber: number | null; + readonly event: RuntimeEventEnvelope; + readonly occurredAt: number; + readonly operation: RuntimeStateOperationName; +}> { + const sourceEventId = `runtime-operation:${input.operationId}:${input.sessionId}:${input.status}`; + const event = await readSessionRuntimeEventSemanticAuthority({ + eventJson: input.eventJson, + invalidMessage: `Runtime operation ${input.status} source ${sourceEventId} is not canonical.`, + missingMessage: `Runtime operation ${input.status} source ${sourceEventId} has no semantic authority.`, + semanticHash: input.semanticHash, + }); + const identity = readRuntimeOperationEventIdentity(event, input); + const projection = createSessionRuntimeEventProjection(event); + + if ( + input.rowAgentId !== input.agentId || + event.id !== input.eventId || + event.kind !== input.eventType || + Date.parse(event.occurredAt) !== input.occurredAt || + event.sourceEventId !== input.sourceEventId || + input.sourceEventId !== sourceEventId || + projection.eventType !== input.eventType || + projection.runId !== null || + projection.source !== input.source || + projection.source !== "api" || + projection.traceId !== null || + projection.visibility !== input.visibility || + projection.visibility !== "all_consumers" + ) { + throw canonicalError(input); + } + + return { ...identity, event, occurredAt: input.occurredAt }; +} diff --git a/apps/api/src/modules/runtime/application/runtime-state-operation-events.ts b/apps/api/src/modules/runtime/application/runtime-state-operation-events.ts index dff85f1e..ec6b44a1 100644 --- a/apps/api/src/modules/runtime/application/runtime-state-operation-events.ts +++ b/apps/api/src/modules/runtime/application/runtime-state-operation-events.ts @@ -1,6 +1,7 @@ import type { RuntimeStateOperationName } from "@mosoo/contracts/agent"; -import type { AgentDeploymentVersionId, AgentId } from "@mosoo/id"; +import type { AgentDeploymentVersionId, AgentId, RuntimeOperationId, SessionId } from "@mosoo/id"; +import { createSessionRuntimeEvent } from "../../sessions/application/session-event-write.service"; import type { RuntimeOperationTargetVersion } from "./runtime-state-operation-version"; export interface RuntimeOperationEvent { @@ -12,6 +13,36 @@ export interface RuntimeOperationEvent { status: "ready" | "updating"; } +export function createRuntimeOperationSessionEvent(input: { + readonly event: RuntimeOperationEvent; + readonly operationId: RuntimeOperationId; + readonly sessionId: SessionId; +}) { + const occurredAtMs = Date.parse(input.event.observedAt); + + return createSessionRuntimeEvent({ + kind: "agent.task.updated", + ...(Number.isFinite(occurredAtMs) ? { occurredAtMs } : {}), + payload: { + agentId: input.event.agentId, + ...(input.event.deploymentVersionId + ? { + deploymentVersionId: input.event.deploymentVersionId, + deploymentVersionNumber: input.event.deploymentVersionNumber, + } + : {}), + operation: input.event.operation, + operationId: input.operationId, + ...(input.event.status === "ready" + ? { readyAt: input.event.observedAt } + : { startedAt: input.event.observedAt }), + status: input.event.status, + }, + sessionId: input.sessionId, + sourceEventId: `runtime-operation:${input.operationId}:${input.sessionId}:${input.event.status}`, + }); +} + export function buildRuntimeStateOperationEvents(input: { agentId: AgentId; operation: RuntimeStateOperationName; diff --git a/apps/api/src/modules/runtime/application/runtime-state-operation-execution.ts b/apps/api/src/modules/runtime/application/runtime-state-operation-execution.ts index c367874b..d3efd33c 100644 --- a/apps/api/src/modules/runtime/application/runtime-state-operation-execution.ts +++ b/apps/api/src/modules/runtime/application/runtime-state-operation-execution.ts @@ -3,7 +3,6 @@ import type { RuntimeOperationId } from "@mosoo/id"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; import type { RuntimeExecutionPlaneAdapter } from "./execution-plane/execution-plane-adapter"; -import { RUNTIME_STATE_OPERATION_INTERRUPTED_ERROR } from "./runtime-state-operation-errors"; import type { RuntimeOperationSubject } from "./runtime-state-operation-subjects"; export type RuntimeStateOperationExecutionPlane = Pick< @@ -19,10 +18,6 @@ function operationInput(input: RuntimeOperationSubject & { operationId: RuntimeO runtimeSubjectId: input.runtimeSubjectId, reason: "agent.runtime_state_operation", targets: input.targets, - terminalRun: { - error: RUNTIME_STATE_OPERATION_INTERRUPTED_ERROR, - status: "cancelled" as const, - }, }; } @@ -36,10 +31,7 @@ async function executeRuntimeStateOperationSubject( ): Promise { switch (input.operation) { case "restartDriver": { - await executionPlane.stopSubjectDrivers(bindings, { - ...operationInput(input), - preserveSessionLifecycle: true, - }); + await executionPlane.stopSubjectDrivers(bindings, operationInput(input)); return; } case "recreateSandbox": { @@ -70,7 +62,7 @@ export async function executeRuntimeStateOperationSubjects( index < input.subjects.length; index += RUNTIME_OPERATION_SUBJECT_CONCURRENCY ) { - await Promise.all( + const outcomes = await Promise.allSettled( input.subjects.slice(index, index + RUNTIME_OPERATION_SUBJECT_CONCURRENCY).map((subject) => executeRuntimeStateOperationSubject(input.executionPlane, bindings, { operationId: input.operationId, @@ -80,5 +72,11 @@ export async function executeRuntimeStateOperationSubjects( }), ), ); + const failure = outcomes.find( + (outcome): outcome is PromiseRejectedResult => outcome.status === "rejected", + ); + if (failure !== undefined) { + throw failure.reason; + } } } diff --git a/apps/api/src/modules/runtime/application/runtime-state-operation-phases.ts b/apps/api/src/modules/runtime/application/runtime-state-operation-phases.ts index f6aa5ffb..bf6be9c2 100644 --- a/apps/api/src/modules/runtime/application/runtime-state-operation-phases.ts +++ b/apps/api/src/modules/runtime/application/runtime-state-operation-phases.ts @@ -4,16 +4,15 @@ import type { AgentId, RuntimeOperationId } from "@mosoo/id"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; import { buildRuntimeStateOperationEvents } from "./runtime-state-operation-events"; -import type { RuntimeOperationEvent } from "./runtime-state-operation-events"; import { - broadcastRuntimeOperationEvent, - writeRuntimeOperationInterruptedSnapshots, + publishRuntimeOperationEvent, writeRuntimeOperationTimedOutSnapshots, } from "./runtime-state-operation-target-events"; import { restoreRuntimeOperationFailedTargets } from "./runtime-state-operation-target-recovery"; import { + claimRuntimeOperationTargets, expireStaleRuntimeOperationTargets, - transitionRuntimeTargetSessionStatus, + listRuntimeOperationTargets, } from "./runtime-state-operation-target-store"; import type { RuntimeSessionTarget, @@ -46,31 +45,6 @@ function listAdmissibleOperationTargets( ); } -async function broadcastOperationPhase( - bindings: ApiBindings, - input: { - readonly event: RuntimeOperationEvent; - readonly expectedStatus?: RuntimeSessionTarget["sessionStatus"]; - readonly operationId: RuntimeOperationId; - readonly status: RuntimeSessionTarget["sessionStatus"]; - readonly targets: RuntimeSessionTarget[]; - }, -): Promise { - const transitions = await transitionRuntimeTargetSessionStatus(bindings.DB, { - ...(input.expectedStatus ? { expectedStatus: input.expectedStatus } : {}), - expectedOperationId: input.expectedStatus === undefined ? null : input.operationId, - operationId: input.operationId, - status: input.status, - targets: input.targets, - }); - await broadcastRuntimeOperationEvent(bindings, { - event: input.event, - operationId: input.operationId, - targets: listCurrentTargets(transitions), - }); - return transitions; -} - export async function startRuntimeStateOperationPhase( bindings: ApiBindings, input: { @@ -89,11 +63,51 @@ export async function startRuntimeStateOperationPhase( startedAt, targetVersion: input.targetVersion, }); - const reschedulingTargets = await broadcastOperationPhase(bindings, { + const admissibleTargets = listAdmissibleOperationTargets(input.targets); + let reschedulingTargets: RuntimeSessionTargetTransition[]; + + try { + reschedulingTargets = await claimRuntimeOperationTargets(bindings.DB, { + event: updatingEvent, + operationId, + targets: admissibleTargets, + }); + } catch (error) { + const claimedTargets = await listRuntimeOperationTargets(bindings.DB, { + operationId, + targets: admissibleTargets, + }); + const partialTransitions = claimedTargets.map((current) => ({ current })); + const [, readyEvent] = buildRuntimeStateOperationEvents({ + agentId: input.agentId, + operation: input.operation, + readyAt: new Date().toISOString(), + startedAt, + targetVersion: input.targetVersion, + }); + + try { + await restoreRuntimeOperationFailedTargets(bindings, { + operationId, + readyEvent, + targets: partialTransitions, + terminalTimestampMs: Date.parse(startedAt), + }); + } catch (recoveryError) { + throw new AggregateError( + [error, recoveryError], + "Runtime operation admission and partial-target recovery both failed.", + { cause: recoveryError }, + ); + } + throw new Error("Runtime operation admission failed after partial-target recovery.", { + cause: error, + }); + } + await publishRuntimeOperationEvent(bindings, { event: updatingEvent, operationId, - status: "RESCHEDULING", - targets: listAdmissibleOperationTargets(input.targets), + targets: listCurrentTargets(reschedulingTargets), }); return { @@ -104,31 +118,7 @@ export async function startRuntimeStateOperationPhase( }; } -export async function failRuntimeStateOperationPhase( - bindings: ApiBindings, - input: { - readonly agentId: AgentId; - readonly operation: RuntimeStateOperationName; - readonly phase: RuntimeStateOperationPhase; - }, -): Promise { - const failedAt = new Date().toISOString(); - const [, failureReadyEvent] = buildRuntimeStateOperationEvents({ - agentId: input.agentId, - operation: input.operation, - readyAt: failedAt, - startedAt: input.phase.startedAt, - targetVersion: input.phase.targetVersion, - }); - - await restoreRuntimeOperationFailedTargets(bindings, { - operationId: input.phase.operationId, - readyEvent: failureReadyEvent, - targets: input.phase.reschedulingTargets, - }); -} - -export async function completeRuntimeStateOperationPhase( +async function finishRuntimeStateOperationPhase( bindings: ApiBindings, input: { readonly agentId: AgentId; @@ -144,28 +134,26 @@ export async function completeRuntimeStateOperationPhase( startedAt: input.phase.startedAt, targetVersion: input.phase.targetVersion, }); - + const phaseTargets = listRuntimeStateOperationPhaseTargets(input.phase); const timedOutTargets = await expireStaleRuntimeOperationTargets(bindings.DB, { operationId: input.phase.operationId, - targets: listRuntimeStateOperationPhaseTargets(input.phase), + targets: phaseTargets, }); await writeRuntimeOperationTimedOutSnapshots(bindings, { operationId: input.phase.operationId, targets: timedOutTargets, }); - const readyTransitions = await transitionRuntimeTargetSessionStatus(bindings.DB, { - expectedOperationId: input.phase.operationId, - expectedStatus: "RESCHEDULING", - status: "IDLE", - targets: listRuntimeStateOperationPhaseTargets(input.phase), - }); - await writeRuntimeOperationInterruptedSnapshots(bindings, { - operationId: input.phase.operationId, - targets: readyTransitions.map((transition) => transition.previous), - }); - await broadcastRuntimeOperationEvent(bindings, { - event: readyEvent, + const timedOutIds = new Set(timedOutTargets.map((target) => target.sessionId)); + + await restoreRuntimeOperationFailedTargets(bindings, { operationId: input.phase.operationId, - targets: listCurrentTargets(readyTransitions), + readyEvent, + targets: input.phase.reschedulingTargets.filter( + (target) => !timedOutIds.has(target.current.sessionId), + ), + terminalTimestampMs: Date.parse(input.phase.startedAt), }); } + +export const failRuntimeStateOperationPhase = finishRuntimeStateOperationPhase; +export const completeRuntimeStateOperationPhase = finishRuntimeStateOperationPhase; diff --git a/apps/api/src/modules/runtime/application/runtime-state-operation-subjects.ts b/apps/api/src/modules/runtime/application/runtime-state-operation-subjects.ts index 3cd2ec2e..a07c24b4 100644 --- a/apps/api/src/modules/runtime/application/runtime-state-operation-subjects.ts +++ b/apps/api/src/modules/runtime/application/runtime-state-operation-subjects.ts @@ -71,9 +71,11 @@ async function resolveLeaseScopedRuntimeOperationScope( runtimeSubjectId: sandboxesTable.id, sandboxId: sandboxSessionsTable.sandboxId, sessionId: sessionsTable.id, + sessionRuntimeEventSeqCursor: sessionsTable.runtimeEventSeqCursor, sessionStatusOperationId: sessionsTable.statusOperationId, sessionStatusSeq: sessionsTable.statusSeq, sessionStatus: sql`${sessionsTable.status}`, + sessionUpdatedAt: sessionsTable.updatedAt, }) .from(sessionsTable) .innerJoin(sandboxSessionsTable, eq(sandboxSessionsTable.sessionId, sessionsTable.id)) @@ -97,9 +99,11 @@ async function resolveLeaseScopedRuntimeOperationScope( lastRunId: row.lastRunId, sandboxId: row.sandboxId, sessionId: row.sessionId, + sessionRuntimeEventSeqCursor: row.sessionRuntimeEventSeqCursor, sessionStatusOperationId: row.sessionStatusOperationId, sessionStatusSeq: row.sessionStatusSeq, sessionStatus: row.sessionStatus, + sessionUpdatedAt: row.sessionUpdatedAt, })), }; } diff --git a/apps/api/src/modules/runtime/application/runtime-state-operation-target-events.ts b/apps/api/src/modules/runtime/application/runtime-state-operation-target-events.ts index d12b6920..cc076764 100644 --- a/apps/api/src/modules/runtime/application/runtime-state-operation-target-events.ts +++ b/apps/api/src/modules/runtime/application/runtime-state-operation-target-events.ts @@ -1,19 +1,20 @@ -import type { RunError, SessionRunSummary } from "@mosoo/contracts/session-run"; +import type { SessionRunSummary } from "@mosoo/contracts/session-run"; +import { sessionEventsTable } from "@mosoo/db"; import { createPlatformId } from "@mosoo/id"; import type { AgentId, RuntimeOperationId, SandboxId, SessionId, SessionRunId } from "@mosoo/id"; import { RUNTIME_DIAGNOSTIC_EVENT } from "@mosoo/runtime-events"; +import { and, inArray } from "drizzle-orm"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; +import { getAppDatabase } from "../../../platform/db/drizzle"; import { isTruthy } from "../../../shared/truthiness"; -import { - appendOneSessionRuntimeEventPerSession, - createSessionRuntimeEvent, -} from "../../sessions/application/session-event-write.service"; +import { publishPersistedSessionRuntimeEvents } from "../../sessions/application/session-event-write.service"; +import { RESCHEDULING_RECONNECT_WINDOW_MS } from "../../sessions/domain/session-lifecycle"; +import { readTerminalEventSemanticAuthority } from "../../sessions/domain/session-terminal-event-authority"; +import type { TerminalEventSemanticAuthority } from "../../sessions/domain/session-terminal-event-authority"; +import { adoptTerminalRunProjection } from "../infrastructure/driver-instance/completed-run-commit.repository"; import { listLiveDriverInstanceRefsForSandboxSessions } from "../infrastructure/driver-instance/live-driver-instance.repository"; -import { - cancelActiveSessionRunsForRuntimeOperation, - getSessionRunSummariesByIds, -} from "../infrastructure/session-runs/session-run-store.repository"; +import { getSessionRunSummariesByIds } from "../infrastructure/session-runs/session-run-store.repository"; import { appendOneRuntimeDiagnosticEventPerSession, toRuntimeDiagnosticBaseValue, @@ -22,72 +23,180 @@ import { RUNTIME_STATE_OPERATION_INTERRUPTED_ERROR, RUNTIME_STATE_OPERATION_TIMEOUT_ERROR, } from "./runtime-state-operation-errors"; +import { createRuntimeOperationSessionEvent } from "./runtime-state-operation-events"; import type { RuntimeOperationEvent } from "./runtime-state-operation-events"; +import { + adoptRuntimeOperationReadyReceipt, + commitSessionLifecycleEventProjection, + listRuntimeOperationTargets, + transitionRuntimeTargetSessionStatus, +} from "./runtime-state-operation-target-store"; import type { RuntimeSessionTarget } from "./runtime-state-operation-target-store"; import type { RuntimeOperationTargetVersion } from "./runtime-state-operation-version"; -import { - createCancelledSessionRunRuntimeEvent, - createSessionLifecycleTerminatedEvent, -} from "./session-runs/session-run-view-events.service"; +import { recordCanonicalSessionRunTerminal } from "./session-runs/session-run-terminal-failure.service"; +import { createSessionLifecycleTerminatedEvent } from "./session-runs/session-run-view-events.service"; function listTargetRunIds(targets: readonly RuntimeSessionTarget[]): SessionRunId[] { - const runIds: SessionRunId[] = []; + return [...new Set(targets.flatMap((target) => (target.lastRunId ? [target.lastRunId] : [])))]; +} - for (const target of targets) { - if (target.sessionStatus !== "IDLE" && isTruthy(target.lastRunId)) { - runIds.push(target.lastRunId); - } +function runtimeOperationSessionEventId(input: { + readonly kind: "interrupted" | "timed_out"; + readonly operationId: RuntimeOperationId; + readonly sessionId: SessionId; +}): string { + return `runtime-operation:${input.operationId}:${input.sessionId}:${input.kind}`; +} + +async function getTargetRuns( + database: D1Database, + targets: readonly RuntimeSessionTarget[], +): Promise> { + const runIds = listTargetRunIds(targets); + + if (runIds.length === 0) { + return new Map(); } - return runIds; + return getSessionRunSummariesByIds(database, runIds); } -function toObservedAtMs(value: string): number | undefined { - const observedAtMs = Date.parse(value); - - return Number.isFinite(observedAtMs) ? observedAtMs : undefined; +function isTerminalRun(run: SessionRunSummary): boolean { + return ["cancelled", "completed", "expired", "failed"].includes(run.status); } -function runtimeOperationRunEventId(input: { - readonly kind: "interrupted" | "timed_out"; - readonly operationId: RuntimeOperationId; - readonly runId: SessionRunId; -}): string { - return `runtime-operation:${input.operationId}:${input.runId}:${input.kind}`; +async function throwFirstFailure( + outcomes: readonly PromiseSettledResult[], +): Promise { + const failed = outcomes.find( + (outcome): outcome is PromiseRejectedResult => outcome.status === "rejected", + ); + + if (failed !== undefined) { + throw failed.reason; + } } -function runtimeOperationSessionEventId(input: { - readonly kind: "timed_out"; - readonly operationId: RuntimeOperationId; - readonly sessionId: SessionId; -}): string { - return `runtime-operation:${input.operationId}:${input.sessionId}:${input.kind}`; +async function adoptCanonicalTerminalRuns( + database: D1Database, + targets: readonly { + readonly run: SessionRunSummary; + readonly sessionId: SessionId; + }[], +): Promise> { + const outcomes = await Promise.allSettled( + targets.map(async ({ run, sessionId }) => { + const outcome = await adoptTerminalRunProjection(database, { + runId: run.id, + sessionId, + }); + if (outcome.kind === "missing" || outcome.kind === "stale") { + throw new Error( + `Runtime operation found an incomplete terminal projection for run ${run.id}.`, + ); + } + }), + ); + await throwFirstFailure(outcomes); + + if (targets.length === 0) { + return new Map(); + } + const rows = await getAppDatabase(database) + .select({ + eventType: sessionEventsTable.eventType, + runId: sessionEventsTable.runId, + semanticHash: sessionEventsTable.semanticHash, + sessionId: sessionEventsTable.sessionId, + sourceEventId: sessionEventsTable.sourceEventId, + streamId: sessionEventsTable.streamId, + terminalEventJson: sessionEventsTable.terminalEventJson, + }) + .from(sessionEventsTable) + .where( + and( + inArray( + sessionEventsTable.runId, + targets.map(({ run }) => run.id), + ), + inArray(sessionEventsTable.eventType, ["run.cancelled", "run.completed", "run.failed"]), + ), + ) + .all(); + const rowsByRunId = Map.groupBy( + rows.filter((row) => row.runId !== null), + (row) => row.runId!, + ); + const authorities = new Map(); + for (const { run, sessionId } of targets) { + const receipts = rowsByRunId.get(run.id) ?? []; + const [receipt] = receipts; + if (receipts.length !== 1 || receipt === undefined || receipt.semanticHash === null) { + throw new Error(`Runtime operation found no v3 lifecycle authority for run ${run.id}.`); + } + authorities.set( + run.id, + await readTerminalEventSemanticAuthority({ + eventJson: receipt.terminalEventJson, + eventType: receipt.eventType, + runId: run.id, + semanticHash: receipt.semanticHash, + sessionId, + sourceEventId: receipt.sourceEventId, + streamId: receipt.streamId, + }), + ); + } + return authorities; } -async function cancelRuntimeOperationTargetRuns( +export async function commitRuntimeOperationReadySnapshots( bindings: ApiBindings, input: { - readonly error: RunError; + readonly event: RuntimeOperationEvent; readonly operationId: RuntimeOperationId; readonly targets: readonly RuntimeSessionTarget[]; }, -): Promise> { - const runIds = listTargetRunIds(input.targets); - - if (runIds.length === 0) { - return new Map(); +): Promise { + if (input.event.status !== "ready") { + throw new Error("Runtime operation release requires one ready event."); + } + const timestampMs = Date.parse(input.event.observedAt); + if (!Number.isFinite(timestampMs)) { + throw new Error("Runtime operation ready time must be a valid ISO timestamp."); } - const updated = await cancelActiveSessionRunsForRuntimeOperation(bindings.DB, { - error: input.error, - operationId: input.operationId, - runIds, - }); - - return getSessionRunSummariesByIds(bindings.DB, [...updated.runIds]); + const outcomes = await Promise.allSettled( + input.targets.map(async (target) => { + const event = createRuntimeOperationSessionEvent({ + event: input.event, + operationId: input.operationId, + sessionId: target.sessionId, + }); + const outcome = await commitSessionLifecycleEventProjection(bindings.DB, { + event, + runtimeOperation: { operationId: input.operationId, status: "ready" }, + status: "IDLE", + target, + timestampMs, + }); + + return outcome.kind === "stale" ? null : { event, sessionId: target.sessionId }; + }), + ); + await throwFirstFailure(outcomes); + const committed = outcomes.flatMap((outcome) => + outcome.status === "fulfilled" && outcome.value !== null ? [outcome.value] : [], + ); + const deliveries = await Promise.allSettled( + committed.map(({ event, sessionId }) => + publishPersistedSessionRuntimeEvents({ bindings, events: [event], sessionId }), + ), + ); + await throwFirstFailure(deliveries); } -export async function broadcastRuntimeOperationEvent( +export async function publishRuntimeOperationEvent( bindings: ApiBindings, input: { readonly event: RuntimeOperationEvent; @@ -95,84 +204,133 @@ export async function broadcastRuntimeOperationEvent( readonly targets: readonly RuntimeSessionTarget[]; }, ): Promise { - await appendOneSessionRuntimeEventPerSession({ - bindings, - records: input.targets.map((target) => { - const occurredAtMs = toObservedAtMs(input.event.observedAt); - - return { - event: createSessionRuntimeEvent({ - kind: "agent.task.updated", - ...(occurredAtMs === undefined ? {} : { occurredAtMs }), - payload: { - agentId: input.event.agentId, - ...(input.event.deploymentVersionId - ? { - deploymentVersionId: input.event.deploymentVersionId, - deploymentVersionNumber: input.event.deploymentVersionNumber, - } - : {}), - operation: input.event.operation, + const outcomes = await Promise.allSettled( + input.targets.map((target) => + publishPersistedSessionRuntimeEvents({ + bindings, + events: [ + createRuntimeOperationSessionEvent({ + event: input.event, operationId: input.operationId, - ...(input.event.status === "ready" - ? { readyAt: input.event.observedAt } - : { startedAt: input.event.observedAt }), - status: input.event.status, - }, - sessionId: target.sessionId, - }), + sessionId: target.sessionId, + }), + ], sessionId: target.sessionId, - }; - }), - }); + }), + ), + ); + await throwFirstFailure(outcomes); } export async function writeRuntimeOperationInterruptedSnapshots( bindings: ApiBindings, input: { readonly operationId: RuntimeOperationId; + readonly timestampMs: number; readonly targets: readonly RuntimeSessionTarget[]; }, ): Promise { - const runsById = await cancelRuntimeOperationTargetRuns(bindings, { - error: RUNTIME_STATE_OPERATION_INTERRUPTED_ERROR, + const runsById = await getTargetRuns(bindings.DB, input.targets); + const outcomes = await Promise.allSettled( + input.targets.flatMap((target) => { + const run = target.lastRunId === null ? null : (runsById.get(target.lastRunId) ?? null); + + return run === null || isTerminalRun(run) + ? [] + : [ + recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + error: RUNTIME_STATE_OPERATION_INTERRUPTED_ERROR, + expectedSessionOperationId: input.operationId, + lifecycle: "IDLE", + runId: run.id, + sessionId: target.sessionId, + source: "runtime_operation", + status: "cancelled", + timestampMs: input.timestampMs, + }), + ]; + }), + ); + await throwFirstFailure(outcomes); + + const latestRuns = await getTargetRuns(bindings.DB, input.targets); + const terminalRuns = input.targets.flatMap((target) => { + const run = target.lastRunId === null ? null : (latestRuns.get(target.lastRunId) ?? null); + + if (run === null) { + return []; + } + if (!isTerminalRun(run)) { + throw new Error(`Runtime operation did not terminalize active run ${run.id}.`); + } + return [{ run, sessionId: target.sessionId }]; + }); + const terminalAuthorities = await adoptCanonicalTerminalRuns(bindings.DB, terminalRuns); + const terminatedTargets = input.targets.filter((target) => { + const run = target.lastRunId === null ? null : (latestRuns.get(target.lastRunId) ?? null); + return run !== null && terminalAuthorities.get(run.id)?.lifecycle === "TERMINATED"; + }); + const currentTerminatedTargets = await listRuntimeOperationTargets(bindings.DB, { operationId: input.operationId, - targets: input.targets, + targets: terminatedTargets, }); - - const records = input.targets - .map((target) => { - if (!isTruthy(target.lastRunId) || target.sessionStatus === "IDLE") { + const lifecycleOutcomes = await Promise.allSettled( + currentTerminatedTargets.map(async (target) => { + const ready = await adoptRuntimeOperationReadyReceipt(bindings.DB, { + operationId: input.operationId, + target, + }); + if (ready !== "missing") { return null; } - - const run = runsById.get(target.lastRunId); - - if (!run) { - return null; + if (target.sessionUpdatedAt > input.timestampMs) { + throw new Error("Runtime operation target advanced beyond its interruption time."); } - - return { - event: createCancelledSessionRunRuntimeEvent({ - eventId: createPlatformId(), - sourceEventId: runtimeOperationRunEventId({ - kind: "interrupted", - operationId: input.operationId, - runId: target.lastRunId, - }), - run, - runError: RUNTIME_STATE_OPERATION_INTERRUPTED_ERROR, + const event = createSessionLifecycleTerminatedEvent({ + eventId: createPlatformId(), + sourceEventId: runtimeOperationSessionEventId({ + kind: "interrupted", + operationId: input.operationId, sessionId: target.sessionId, }), + lastSeen: new Date(input.timestampMs).toISOString(), + message: RUNTIME_STATE_OPERATION_INTERRUPTED_ERROR.message, + occurredAtMs: input.timestampMs, + reason: RUNTIME_STATE_OPERATION_INTERRUPTED_ERROR.code, sessionId: target.sessionId, - }; - }) - .filter(isTruthy); - - await appendOneSessionRuntimeEventPerSession({ - bindings, - records, + }); + const outcome = await commitSessionLifecycleEventProjection(bindings.DB, { + event, + status: "TERMINATED", + target, + timestampMs: input.timestampMs, + }); + return outcome.kind === "stale" ? null : { event, sessionId: target.sessionId }; + }), + ); + await throwFirstFailure(lifecycleOutcomes); + const lifecycleDeliveries = await Promise.allSettled( + lifecycleOutcomes.flatMap((outcome) => + outcome.status === "fulfilled" && outcome.value !== null + ? [ + publishPersistedSessionRuntimeEvents({ + bindings, + events: [outcome.value.event], + sessionId: outcome.value.sessionId, + }), + ] + : [], + ), + ); + await throwFirstFailure(lifecycleDeliveries); + const unreleased = await listRuntimeOperationTargets(bindings.DB, { + operationId: input.operationId, + targets: terminatedTargets, }); + if (unreleased.length > 0) { + throw new Error("Runtime operation interruption did not preserve terminal Session lifecycle."); + } } export async function appendRuntimeDriverRestartAttemptedEvents( @@ -271,72 +429,186 @@ export async function writeRuntimeOperationTimedOutSnapshots( readonly targets: readonly RuntimeSessionTarget[]; }, ): Promise { - const runsById = await cancelRuntimeOperationTargetRuns(bindings, { - error: RUNTIME_STATE_OPERATION_TIMEOUT_ERROR, - operationId: input.operationId, - targets: input.targets, - }); + await writeRuntimeOperationTimedOutSnapshotsOnce( + bindings, + input, + new Map( + input.targets.map((target) => [ + target.sessionId, + target.sessionUpdatedAt + RESCHEDULING_RECONNECT_WINDOW_MS, + ]), + ), + true, + ); +} - const records = input.targets - .map((target) => { - if (!isTruthy(target.lastRunId) || target.sessionStatus === "IDLE") { - return { - event: createSessionLifecycleTerminatedEvent({ - eventId: createPlatformId(), - sourceEventId: runtimeOperationSessionEventId({ - kind: "timed_out", - operationId: input.operationId, - sessionId: target.sessionId, - }), - lastSeen: new Date().toISOString(), - message: RUNTIME_STATE_OPERATION_TIMEOUT_ERROR.message, - reason: RUNTIME_STATE_OPERATION_TIMEOUT_ERROR.code, - sessionId: target.sessionId, - }), - sessionId: target.sessionId, - }; +async function writeRuntimeOperationTimedOutSnapshotsOnce( + bindings: ApiBindings, + input: { + readonly operationId: RuntimeOperationId; + readonly targets: readonly RuntimeSessionTarget[]; + }, + deadlines: ReadonlyMap, + retryStaleTargets: boolean, +): Promise { + const deadlineFor = (target: RuntimeSessionTarget): number => { + const deadline = deadlines.get(target.sessionId); + if (deadline === undefined) { + throw new Error("Runtime operation timeout lost its stable target deadline."); + } + return deadline; + }; + const runsById = await getTargetRuns(bindings.DB, input.targets); + const activeTargets = input.targets.filter((target) => { + const run = target.lastRunId === null ? null : (runsById.get(target.lastRunId) ?? null); + return run !== null && !isTerminalRun(run); + }); + const outcomes = await Promise.allSettled( + activeTargets.map((target) => { + const run = runsById.get(target.lastRunId!); + if (run === undefined) { + throw new Error("Runtime operation timeout lost its active run snapshot."); } - const run = runsById.get(target.lastRunId); + return recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + error: RUNTIME_STATE_OPERATION_TIMEOUT_ERROR, + expectedSessionOperationId: input.operationId, + lifecycle: "TERMINATED", + runId: run.id, + sessionId: target.sessionId, + source: "runtime_operation", + status: "expired", + timestampMs: deadlineFor(target), + }); + }), + ); + await throwFirstFailure(outcomes); + + const latestRuns = await getTargetRuns(bindings.DB, input.targets); + const terminalRuns = input.targets.flatMap((target) => { + const run = target.lastRunId === null ? null : (latestRuns.get(target.lastRunId) ?? null); + if (run === null) { + return []; + } + if (!isTerminalRun(run)) { + throw new Error(`Runtime operation timeout did not terminalize active run ${run.id}.`); + } + return [{ run, sessionId: target.sessionId }]; + }); + const terminalAuthorities = await adoptCanonicalTerminalRuns(bindings.DB, terminalRuns); - if (!run) { - return { - event: createSessionLifecycleTerminatedEvent({ - eventId: createPlatformId(), - sourceEventId: runtimeOperationSessionEventId({ - kind: "timed_out", - operationId: input.operationId, - sessionId: target.sessionId, - }), - lastSeen: new Date().toISOString(), - message: RUNTIME_STATE_OPERATION_TIMEOUT_ERROR.message, - reason: RUNTIME_STATE_OPERATION_TIMEOUT_ERROR.code, - sessionId: target.sessionId, - }), - sessionId: target.sessionId, - }; + const noRunLifecycleTargets = input.targets.filter((target) => { + const run = target.lastRunId === null ? null : (latestRuns.get(target.lastRunId) ?? null); + return run === null; + }); + const terminatedTargets = input.targets.filter((target) => { + const run = target.lastRunId === null ? null : (latestRuns.get(target.lastRunId) ?? null); + return run !== null && terminalAuthorities.get(run.id)?.lifecycle === "TERMINATED"; + }); + const currentTerminatedTargets = await listRuntimeOperationTargets(bindings.DB, { + operationId: input.operationId, + targets: terminatedTargets, + }); + const originalTargetsById = new Map(input.targets.map((target) => [target.sessionId, target])); + const lifecycleTargets = [ + ...noRunLifecycleTargets.map((target) => { + const timestampMs = deadlineFor(target); + if (target.sessionUpdatedAt > timestampMs) { + throw new Error("Runtime operation target advanced beyond its stable timeout deadline."); + } + return { target, timestampMs }; + }), + ...currentTerminatedTargets.map((target) => { + const original = originalTargetsById.get(target.sessionId); + if (original === undefined) { + throw new Error("Runtime operation timeout lost its original target."); + } + const timestampMs = deadlineFor(original); + if (target.sessionUpdatedAt > timestampMs) { + throw new Error("Runtime operation target advanced beyond its stable timeout deadline."); + } + return { target, timestampMs }; + }), + ]; + const lifecycleOutcomes = await Promise.allSettled( + lifecycleTargets.map(async ({ target, timestampMs }) => { + const ready = await adoptRuntimeOperationReadyReceipt(bindings.DB, { + operationId: input.operationId, + target, + }); + if (ready !== "missing") { + return null; } - return { - event: createCancelledSessionRunRuntimeEvent({ - eventId: createPlatformId(), - sourceEventId: runtimeOperationRunEventId({ - kind: "timed_out", - operationId: input.operationId, - runId: target.lastRunId, - }), - lifecycle: "TERMINATED", - run, - runError: RUNTIME_STATE_OPERATION_TIMEOUT_ERROR, + const event = createSessionLifecycleTerminatedEvent({ + eventId: createPlatformId(), + sourceEventId: runtimeOperationSessionEventId({ + kind: "timed_out", + operationId: input.operationId, sessionId: target.sessionId, }), + lastSeen: new Date(timestampMs).toISOString(), + message: RUNTIME_STATE_OPERATION_TIMEOUT_ERROR.message, + occurredAtMs: timestampMs, + reason: RUNTIME_STATE_OPERATION_TIMEOUT_ERROR.code, sessionId: target.sessionId, - }; - }) - .filter(isTruthy); + }); + const outcome = await commitSessionLifecycleEventProjection(bindings.DB, { + event, + status: "TERMINATED", + target, + timestampMs, + }); + + return outcome.kind === "stale" ? null : { event, sessionId: target.sessionId }; + }), + ); + await throwFirstFailure(lifecycleOutcomes); + const lifecycleDeliveries = await Promise.allSettled( + lifecycleOutcomes.flatMap((outcome) => + outcome.status === "fulfilled" && outcome.value !== null + ? [ + publishPersistedSessionRuntimeEvents({ + bindings, + events: [outcome.value.event], + sessionId: outcome.value.sessionId, + }), + ] + : [], + ), + ); + await throwFirstFailure(lifecycleDeliveries); - await appendOneSessionRuntimeEventPerSession({ - bindings, - records, + const adoptedTargets = input.targets.filter((target) => { + const run = target.lastRunId === null ? null : (latestRuns.get(target.lastRunId) ?? null); + return run !== null && terminalAuthorities.get(run.id)?.lifecycle === "IDLE"; }); + const currentAdoptedTargets = await listRuntimeOperationTargets(bindings.DB, { + operationId: input.operationId, + targets: adoptedTargets, + }); + await transitionRuntimeTargetSessionStatus(bindings.DB, { + expectedOperationId: input.operationId, + operationId: null, + status: "IDLE", + targets: currentAdoptedTargets, + }); + + const unreleased = await listRuntimeOperationTargets(bindings.DB, { + operationId: input.operationId, + targets: input.targets, + }); + if (unreleased.length > 0) { + if (retryStaleTargets) { + await writeRuntimeOperationTimedOutSnapshotsOnce( + bindings, + { operationId: input.operationId, targets: unreleased }, + deadlines, + false, + ); + return; + } + throw new Error("Runtime operation timeout did not release every owned Session."); + } } diff --git a/apps/api/src/modules/runtime/application/runtime-state-operation-target-recovery.ts b/apps/api/src/modules/runtime/application/runtime-state-operation-target-recovery.ts index c31d6d10..7f064848 100644 --- a/apps/api/src/modules/runtime/application/runtime-state-operation-target-recovery.ts +++ b/apps/api/src/modules/runtime/application/runtime-state-operation-target-recovery.ts @@ -1,79 +1,13 @@ -import type { RuntimeOperationId, SessionRunId } from "@mosoo/id"; +import type { RuntimeOperationId } from "@mosoo/id"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; -import { isTruthy } from "../../../shared/truthiness"; -import { getSessionRunSummariesByIds } from "../infrastructure/session-runs/session-run-store.repository"; import type { RuntimeOperationEvent } from "./runtime-state-operation-events"; import { - broadcastRuntimeOperationEvent, + commitRuntimeOperationReadySnapshots, writeRuntimeOperationInterruptedSnapshots, } from "./runtime-state-operation-target-events"; -import { - isTerminalRunStatus, - transitionRuntimeTargetSessionStatus, -} from "./runtime-state-operation-target-store"; -import type { - RuntimeSessionTarget, - RuntimeSessionTargetTransition, -} from "./runtime-state-operation-target-store"; - -type RuntimeTargetGroups = Record; -type RuntimeTargetTransitionGroups = Record< - RuntimeSessionTarget["sessionStatus"], - RuntimeSessionTargetTransition[] ->; -const RUNTIME_TARGET_SESSION_STATUSES = ["IDLE", "RESCHEDULING", "RUNNING"] as const; - -function createRuntimeTargetGroups(): RuntimeTargetGroups { - return { - IDLE: [], - RESCHEDULING: [], - RUNNING: [], - }; -} - -function createRuntimeTargetTransitionGroups(): RuntimeTargetTransitionGroups { - return { - IDLE: [], - RESCHEDULING: [], - RUNNING: [], - }; -} - -async function groupFailureRestoreTargets( - bindings: ApiBindings, - targets: readonly RuntimeSessionTargetTransition[], -): Promise { - const runIds: SessionRunId[] = []; - - for (const { previous: target } of targets) { - if (target.sessionStatus === "RUNNING" && isTruthy(target.lastRunId)) { - runIds.push(target.lastRunId); - } - } - - const runsById = await getSessionRunSummariesByIds(bindings.DB, runIds); - const groups = createRuntimeTargetTransitionGroups(); - - for (const target of targets) { - const previous = target.previous; - - if (previous.sessionStatus === "IDLE" || previous.sessionStatus === "RESCHEDULING") { - groups[previous.sessionStatus].push(target); - continue; - } - - if (!isTruthy(previous.lastRunId)) { - groups.IDLE.push(target); - continue; - } - - const run = runsById.get(previous.lastRunId); - groups[run && !isTerminalRunStatus(run.status) ? "RUNNING" : "IDLE"].push(target); - } - - return groups; -} +import { listRuntimeOperationTargets } from "./runtime-state-operation-target-store"; +import type { RuntimeSessionTargetTransition } from "./runtime-state-operation-target-store"; export async function restoreRuntimeOperationFailedTargets( bindings: ApiBindings, @@ -81,52 +15,36 @@ export async function restoreRuntimeOperationFailedTargets( readonly operationId: RuntimeOperationId; readonly readyEvent: RuntimeOperationEvent; readonly targets: readonly RuntimeSessionTargetTransition[]; + readonly terminalTimestampMs: number; }, -): Promise { - const restoreGroups = await groupFailureRestoreTargets(bindings, input.targets); - const restoredTargets: RuntimeSessionTarget[] = []; - const restoredGroups = createRuntimeTargetGroups(); - - const restoredEntries = await Promise.all( - RUNTIME_TARGET_SESSION_STATUSES.map(async (status) => { - const transitions = await transitionRuntimeTargetSessionStatus(bindings.DB, { - expectedOperationId: input.operationId, - expectedStatus: "RESCHEDULING", - status, - targets: restoreGroups[status].map((target) => target.current), - }); - const updatedTargets = transitions.map((transition) => transition.current); - - return [status, updatedTargets] as const; - }), - ); - - for (const [status, updatedTargets] of restoredEntries) { - restoredGroups[status] = updatedTargets; - restoredTargets.push(...updatedTargets); - } +): Promise { + const phaseTargets = input.targets.map((target) => target.current); + const ownedTargets = await listRuntimeOperationTargets(bindings.DB, { + operationId: input.operationId, + targets: phaseTargets, + }); await writeRuntimeOperationInterruptedSnapshots(bindings, { operationId: input.operationId, - targets: input.targets.map((target) => target.previous), + targets: ownedTargets, + timestampMs: input.terminalTimestampMs, }); - await broadcastRuntimeOperationEvent(bindings, { + const readyTargets = await listRuntimeOperationTargets(bindings.DB, { + operationId: input.operationId, + targets: phaseTargets, + }); + await commitRuntimeOperationReadySnapshots(bindings, { event: input.readyEvent, operationId: input.operationId, - targets: restoredGroups.IDLE, + targets: readyTargets, }); - const runningEvent: RuntimeOperationEvent = { - ...input.readyEvent, - observedAt: new Date().toISOString(), - status: "ready", - }; - await broadcastRuntimeOperationEvent(bindings, { - event: runningEvent, + const unreleased = await listRuntimeOperationTargets(bindings.DB, { operationId: input.operationId, - targets: restoredGroups.RUNNING, + targets: phaseTargets, }); - - return restoredTargets; + if (unreleased.length > 0) { + throw new Error("Runtime operation ready projection did not release every owned Session."); + } } diff --git a/apps/api/src/modules/runtime/application/runtime-state-operation-target-store.ts b/apps/api/src/modules/runtime/application/runtime-state-operation-target-store.ts index 55bec015..b8908824 100644 --- a/apps/api/src/modules/runtime/application/runtime-state-operation-target-store.ts +++ b/apps/api/src/modules/runtime/application/runtime-state-operation-target-store.ts @@ -1,19 +1,31 @@ -import type { SessionRunStatus } from "@mosoo/contracts/session-run"; -import { sandboxSessionsTable, sessionsTable } from "@mosoo/db"; +import { sandboxSessionsTable, sessionEventsTable, sessionsTable } from "@mosoo/db"; +import { createPlatformId } from "@mosoo/id"; import type { AgentId, PlatformId, + RuntimeEventId, RuntimeOperationId, SandboxId, SessionId, SessionRunId, } from "@mosoo/id"; -import { and, eq, inArray, isNull, lte, or, sql } from "drizzle-orm"; +import { createRuntimeEventSemanticHash } from "@mosoo/runtime-events"; +import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; +import { and, asc, eq, inArray, isNotNull, isNull, lte, or, sql } from "drizzle-orm"; -import { getAppDatabase } from "../../../platform/db/drizzle"; +import { getAppDatabase, getD1ChangeCount } from "../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../time"; import { RESCHEDULING_RECONNECT_WINDOW_MS } from "../../sessions/domain/session-lifecycle"; +import { createSessionRuntimeEventProjection } from "../../sessions/domain/session-runtime-event-projection"; import { createSessionStatusTransitionPatch } from "../infrastructure/session-runs/session-lifecycle-projection.repository"; +import { + createRuntimeOperationEventAuthorityJson, + readRuntimeOperationEventIdentity, + readRuntimeOperationEventAuthority, +} from "./runtime-state-operation-event-authority"; +import type { RuntimeOperationEventStatus } from "./runtime-state-operation-event-authority"; +import { createRuntimeOperationSessionEvent } from "./runtime-state-operation-events"; +import type { RuntimeOperationEvent } from "./runtime-state-operation-events"; export interface RuntimeSessionTarget { readonly agentId: AgentId | null; @@ -21,11 +33,28 @@ export interface RuntimeSessionTarget { readonly lastRunId: SessionRunId | null; readonly sandboxId: SandboxId; readonly sessionId: SessionId; + readonly sessionRuntimeEventSeqCursor: number; readonly sessionStatusOperationId: RuntimeOperationId | null; readonly sessionStatusSeq: number; readonly sessionStatus: "IDLE" | "RUNNING" | "RESCHEDULING"; + readonly sessionUpdatedAt: number; } +export type RuntimeSessionLifecycleTarget = Pick< + RuntimeSessionTarget, + | "lastRunId" + | "sessionId" + | "sessionRuntimeEventSeqCursor" + | "sessionStatus" + | "sessionStatusOperationId" + | "sessionStatusSeq" + | "sessionUpdatedAt" +> & { readonly agentId?: AgentId | null }; + +export type SessionLifecycleEventProjectionOutcome = + | { readonly kind: "applied" | "duplicate" } + | { readonly kind: "stale" }; + export const RUNTIME_TARGET_SESSION_STATUSES: RuntimeSessionTarget["sessionStatus"][] = [ "IDLE", "RUNNING", @@ -35,13 +64,10 @@ const RUNTIME_TARGET_STATUS_WRITE_BATCH_SIZE = 25; export interface RuntimeSessionTargetTransition { readonly current: RuntimeSessionTarget; - readonly previous: RuntimeSessionTarget; } -function isTerminalRunStatus(status: SessionRunStatus): boolean { - return ( - status === "completed" || status === "failed" || status === "cancelled" || status === "expired" - ); +export interface StaleRuntimeOperationTarget extends RuntimeSessionTarget { + readonly operationId: RuntimeOperationId; } function sessionStatusOperationCondition(operationId: RuntimeOperationId | null) { @@ -72,11 +98,223 @@ function sessionTargetFreshnessCondition( eq(sessionsTable.id, target.sessionId), eq(sessionsTable.status, input.expectedStatus ?? target.sessionStatus), eq(sessionsTable.statusSeq, target.sessionStatusSeq), + eq(sessionsTable.runtimeEventSeqCursor, target.sessionRuntimeEventSeqCursor), sessionLastRunCondition(target.lastRunId), ...sessionStatusOperationCondition(expectedOperationId), ); } +interface PreparedSessionLifecycleEvent { + readonly event: RuntimeEventEnvelope; + readonly eventType: string; + readonly occurredAtMs: number; + readonly projection: ReturnType; + readonly semanticHash: string; + readonly sourceEventId: string; + readonly runtimeOperationEventJson: string | null; +} + +async function prepareSessionLifecycleEvent( + event: RuntimeEventEnvelope, + runtimeOperation?: { + readonly agentId: AgentId; + readonly operationId: RuntimeOperationId; + readonly status: RuntimeOperationEventStatus; + }, +): Promise { + const occurredAtMs = Date.parse(event.occurredAt); + if (!Number.isFinite(occurredAtMs)) { + throw new Error("Session lifecycle event time must be a valid ISO timestamp."); + } + if (event.sourceEventId === undefined) { + throw new Error("Atomic Session lifecycle events require one stable source event ID."); + } + + const projection = createSessionRuntimeEventProjection(event); + return { + event, + eventType: projection.eventType, + occurredAtMs, + projection, + semanticHash: await createRuntimeEventSemanticHash(event), + sourceEventId: event.sourceEventId, + runtimeOperationEventJson: + runtimeOperation === undefined + ? null + : createRuntimeOperationEventAuthorityJson({ + event, + ...runtimeOperation, + sessionId: event.sessionId, + }), + }; +} + +function createSessionLifecycleEventInsertStatements( + database: D1Database, + input: { + readonly prepared: PreparedSessionLifecycleEvent; + readonly target: { + readonly lastRunId: SessionRunId | null; + readonly operationId: RuntimeOperationId | null; + readonly runtimeEventSeqCursor: number; + readonly sessionId: SessionId; + readonly status: "IDLE" | "RESCHEDULING" | "TERMINATED"; + readonly statusSeq: number; + readonly updatedAt: number; + }; + }, +): [D1PreparedStatement, D1PreparedStatement] { + const { event, projection } = input.prepared; + const target = input.target; + + return [ + database + .prepare( + `INSERT INTO session_event ( + agent_id, content_text, created_at, ended_at, event_type, family, id, + occurred_at, process_status, process_type, run_id, semantic_hash, + runtime_operation_event_json, seq, + session_id, source_event_id, source, stream_id, tool_call_id, + tool_input_delta_json, tool_input_json, tool_name, tool_output_delta_text, + tool_output_text, tool_parent_message_id, tool_result_message_id, tool_status, + tokens, trace_id, visibility + ) + SELECT + s.agent_id, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, s.id, ?, ?, ?, ?, ?, ?, ?, + ?, ?, ?, ?, ?, ?, ?, ? + FROM session AS s + WHERE s.id = ? + AND s.archived_at IS NULL + AND s.cleanup_operation_kind IS NULL + AND s.last_run_id IS ? + AND s.runtime_event_seq_cursor = ? + AND s.status = ? + AND s.status_operation_id IS ? + AND s.status_seq = ? + AND s.updated_at = ? + ON CONFLICT(session_id, source_event_id) DO NOTHING`, + ) + .bind( + projection.contentText, + input.prepared.occurredAtMs, + input.prepared.occurredAtMs, + projection.eventType, + projection.family, + event.id, + input.prepared.occurredAtMs, + projection.processStatus, + projection.processType, + projection.runId, + input.prepared.semanticHash, + input.prepared.runtimeOperationEventJson, + target.runtimeEventSeqCursor, + input.prepared.sourceEventId, + projection.source, + projection.streamId, + projection.toolCallId, + projection.toolInputDeltaJson, + projection.toolInputJson, + projection.toolName, + projection.toolOutputDeltaText, + projection.toolOutputText, + projection.toolParentMessageId, + projection.toolResultMessageId, + projection.toolStatus, + projection.tokens, + projection.traceId, + projection.visibility, + target.sessionId, + target.lastRunId, + target.runtimeEventSeqCursor, + target.status, + target.operationId, + target.statusSeq, + target.updatedAt, + ), + database + .prepare( + `INSERT INTO session_event (id) + SELECT ? + WHERE EXISTS ( + SELECT 1 + FROM session AS s + WHERE s.id = ? + AND s.archived_at IS NULL + AND s.cleanup_operation_kind IS NULL + AND s.last_run_id IS ? + AND s.runtime_event_seq_cursor = ? + AND s.status = ? + AND s.status_operation_id IS ? + AND s.status_seq = ? + AND s.updated_at = ? + ) + AND NOT EXISTS ( + SELECT 1 + FROM session_event AS event + WHERE event.session_id = ? + AND event.source_event_id = ? + AND event.event_type = ? + AND event.semantic_hash = ? + AND event.runtime_operation_event_json IS ? + AND event.seq = ? + )`, + ) + .bind( + createPlatformId(), + target.sessionId, + target.lastRunId, + target.runtimeEventSeqCursor, + target.status, + target.operationId, + target.statusSeq, + target.updatedAt, + target.sessionId, + input.prepared.sourceEventId, + input.prepared.eventType, + input.prepared.semanticHash, + input.prepared.runtimeOperationEventJson, + target.runtimeEventSeqCursor, + ), + ]; +} + +async function getSessionLifecycleEventReceipt( + database: D1Database, + input: Pick< + PreparedSessionLifecycleEvent, + "eventType" | "runtimeOperationEventJson" | "semanticHash" | "sourceEventId" + > & { + readonly sessionId: SessionId; + }, +): Promise<{ readonly exact: boolean; readonly seq: number } | null> { + const receipt = await getAppDatabase(database) + .select({ + eventType: sessionEventsTable.eventType, + runtimeOperationEventJson: sessionEventsTable.runtimeOperationEventJson, + semanticHash: sessionEventsTable.semanticHash, + seq: sessionEventsTable.seq, + }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.sourceEventId, input.sourceEventId), + ), + ) + .limit(1) + .get(); + + return receipt === undefined + ? null + : { + exact: + receipt.eventType === input.eventType && + receipt.runtimeOperationEventJson === input.runtimeOperationEventJson && + receipt.semanticHash === input.semanticHash, + seq: receipt.seq, + }; +} + export async function listRuntimeSessionTargetsForSandboxIds( database: D1Database, sandboxIds: readonly SandboxId[], @@ -94,9 +332,11 @@ export async function listRuntimeSessionTargetsForSandboxIds( lastRunId: sessionsTable.lastRunId, sandboxId: sandboxSessionsTable.sandboxId, sessionId: sessionsTable.id, + sessionRuntimeEventSeqCursor: sessionsTable.runtimeEventSeqCursor, sessionStatusOperationId: sessionsTable.statusOperationId, sessionStatusSeq: sessionsTable.statusSeq, sessionStatus: sql`${sessionsTable.status}`, + sessionUpdatedAt: sessionsTable.updatedAt, }) .from(sessionsTable) .innerJoin(sandboxSessionsTable, eq(sandboxSessionsTable.sessionId, sessionsTable.id)) @@ -105,6 +345,7 @@ export async function listRuntimeSessionTargetsForSandboxIds( inArray(sandboxSessionsTable.sandboxId, uniqueSandboxIds), eq(sandboxSessionsTable.status, "active"), isNull(sessionsTable.archivedAt), + isNull(sessionsTable.runtimeProvisioningOperationId), inArray(sessionsTable.status, RUNTIME_TARGET_SESSION_STATUSES), ), ) @@ -194,43 +435,612 @@ async function transitionRuntimeTargetSessionStatusBatch( sessionStatus: input.status, sessionStatusOperationId: updated.status_operation_id, sessionStatusSeq: updated.status_seq, + sessionUpdatedAt: timestampMs, }, - previous: target, }, ]; }); } -export async function expireStaleRuntimeOperationTargets( +export async function commitSessionLifecycleEventProjection( + database: D1Database, + input: { + readonly event: RuntimeEventEnvelope; + readonly runtimeOperation?: { + readonly operationId: RuntimeOperationId; + readonly status: RuntimeOperationEventStatus; + }; + readonly status: "IDLE" | "TERMINATED"; + readonly target: RuntimeSessionLifecycleTarget; + readonly timestampMs: number; + }, +): Promise { + if (input.runtimeOperation !== undefined && input.target.agentId == null) { + throw new Error("Runtime operation lifecycle events require an Agent-bound Session."); + } + const prepared = await prepareSessionLifecycleEvent( + input.event, + input.runtimeOperation === undefined + ? undefined + : { agentId: input.target.agentId!, ...input.runtimeOperation }, + ); + if (input.runtimeOperation?.status === "ready") { + const claim = await getRuntimeOperationEventReceipt(database, { + agentId: input.target.agentId!, + operationId: input.runtimeOperation.operationId, + sessionId: input.target.sessionId, + status: "updating", + }); + const ready = readRuntimeOperationEventIdentity(input.event, { + agentId: input.target.agentId!, + operationId: input.runtimeOperation.operationId, + sessionId: input.target.sessionId, + status: "ready", + }); + if (claim === null) { + throw new Error( + `Runtime operation ready source ${prepared.sourceEventId} has no canonical claim.`, + ); + } + if (!runtimeOperationEventIdentitiesEqual(claim, ready)) { + throw new Error( + `Runtime operation ready source ${prepared.sourceEventId} conflicts with its claim.`, + ); + } + } + if (prepared.event.sessionId !== input.target.sessionId) { + throw new Error("Session lifecycle event does not match its target Session."); + } + if (prepared.occurredAtMs !== input.timestampMs) { + throw new Error("Session lifecycle event time does not match its projection timestamp."); + } + + const receiptInput = { + eventType: prepared.eventType, + runtimeOperationEventJson: prepared.runtimeOperationEventJson, + semanticHash: prepared.semanticHash, + sessionId: input.target.sessionId, + sourceEventId: prepared.sourceEventId, + }; + const existingReceipt = await getSessionLifecycleEventReceipt(database, receiptInput); + if (existingReceipt !== null && !existingReceipt.exact) { + throw new Error( + `Session lifecycle source ${prepared.sourceEventId} conflicts with its durable receipt.`, + ); + } + + let changed = 0; + if (existingReceipt !== null) { + const result = await database + .prepare( + `UPDATE session + SET status = ?, + status_operation_id = NULL, + status_seq = status_seq + 1, + updated_at = ? + WHERE id = ? + AND archived_at IS NULL + AND cleanup_operation_kind IS NULL + AND last_run_id IS ? + AND runtime_event_seq_cursor = ? + AND status = ? + AND status_operation_id IS ? + AND status_seq = ? + AND updated_at = ? + AND EXISTS ( + SELECT 1 + FROM session_event AS event + WHERE event.session_id = ? + AND event.source_event_id = ? + AND event.event_type = ? + AND event.runtime_operation_event_json IS ? + AND event.semantic_hash = ? + )`, + ) + .bind( + input.status, + input.timestampMs, + input.target.sessionId, + input.target.lastRunId, + input.target.sessionRuntimeEventSeqCursor, + input.target.sessionStatus, + input.target.sessionStatusOperationId, + input.target.sessionStatusSeq, + input.target.sessionUpdatedAt, + input.target.sessionId, + prepared.sourceEventId, + prepared.eventType, + prepared.runtimeOperationEventJson, + prepared.semanticHash, + ) + .run(); + changed = getD1ChangeCount(result); + } else { + const nextEventSeq = input.target.sessionRuntimeEventSeqCursor + 1; + const nextStatusSeq = input.target.sessionStatusSeq + 1; + const eventStatements = createSessionLifecycleEventInsertStatements(database, { + prepared, + target: { + lastRunId: input.target.lastRunId, + operationId: null, + runtimeEventSeqCursor: nextEventSeq, + sessionId: input.target.sessionId, + status: input.status, + statusSeq: nextStatusSeq, + updatedAt: input.timestampMs, + }, + }); + const results = await database.batch([ + database + .prepare( + `UPDATE session + SET runtime_event_seq_cursor = runtime_event_seq_cursor + 1, + status = ?, + status_operation_id = NULL, + status_seq = status_seq + 1, + updated_at = ? + WHERE id = ? + AND archived_at IS NULL + AND cleanup_operation_kind IS NULL + AND last_run_id IS ? + AND runtime_event_seq_cursor = ? + AND status = ? + AND status_operation_id IS ? + AND status_seq = ? + AND updated_at = ?`, + ) + .bind( + input.status, + input.timestampMs, + input.target.sessionId, + input.target.lastRunId, + input.target.sessionRuntimeEventSeqCursor, + input.target.sessionStatus, + input.target.sessionStatusOperationId, + input.target.sessionStatusSeq, + input.target.sessionUpdatedAt, + ), + ...eventStatements, + ]); + changed = getD1ChangeCount(results[0]); + } + + if (changed > 0) { + return { kind: "applied" }; + } + + const [receipt, session] = await Promise.all([ + getSessionLifecycleEventReceipt(database, receiptInput), + database + .prepare( + `SELECT archived_at, cleanup_operation_kind, last_run_id, + runtime_event_seq_cursor, status, status_operation_id + FROM session + WHERE id = ?`, + ) + .bind(input.target.sessionId) + .first<{ + archived_at: number | null; + cleanup_operation_kind: string | null; + last_run_id: string | null; + runtime_event_seq_cursor: number; + status: string; + status_operation_id: string | null; + }>(), + ]); + if ( + receipt?.exact === true && + session?.archived_at === null && + session.cleanup_operation_kind === null && + session.last_run_id === input.target.lastRunId && + session.runtime_event_seq_cursor >= receipt.seq && + session.status === input.status && + session.status_operation_id === null + ) { + return { kind: "duplicate" }; + } + + return { kind: "stale" }; +} + +interface RuntimeOperationEventReceipt { + readonly agentId: AgentId; + readonly deploymentVersionId: RuntimeOperationEvent["deploymentVersionId"] | null; + readonly deploymentVersionNumber: RuntimeOperationEvent["deploymentVersionNumber"] | null; + readonly eventId: string; + readonly eventJson: string; + readonly eventType: string; + readonly occurredAt: number; + readonly operation: RuntimeOperationEvent["operation"]; + readonly semanticHash: string; + readonly source: string; + readonly sourceEventId: string; + readonly visibility: string; + readonly seq: number; +} + +async function getRuntimeOperationEventReceipt( + database: D1Database, + input: { + readonly agentId: AgentId; + readonly operationId: RuntimeOperationId; + readonly sessionId: SessionId; + readonly status: RuntimeOperationEventStatus; + }, +): Promise { + const sourceEventId = `runtime-operation:${input.operationId}:${input.sessionId}:${input.status}`; + const receipt = await getAppDatabase(database) + .select({ + agentId: sessionEventsTable.agentId, + eventId: sessionEventsTable.id, + eventJson: sessionEventsTable.runtimeOperationEventJson, + eventType: sessionEventsTable.eventType, + occurredAt: sessionEventsTable.occurredAt, + semanticHash: sessionEventsTable.semanticHash, + seq: sessionEventsTable.seq, + source: sessionEventsTable.source, + sourceEventId: sessionEventsTable.sourceEventId, + visibility: sessionEventsTable.visibility, + }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.sourceEventId, sourceEventId), + ), + ) + .limit(1) + .get(); + + if (receipt === undefined) { + return null; + } + const authority = await readRuntimeOperationEventAuthority({ + agentId: input.agentId, + eventId: receipt.eventId, + eventJson: receipt.eventJson, + eventType: receipt.eventType, + occurredAt: receipt.occurredAt, + operationId: input.operationId, + rowAgentId: receipt.agentId, + semanticHash: receipt.semanticHash, + sessionId: input.sessionId, + source: receipt.source, + sourceEventId: receipt.sourceEventId, + status: input.status, + visibility: receipt.visibility, + }); + if (receipt.eventJson === null || receipt.semanticHash === null) { + throw new Error( + `Runtime operation ${input.status} source ${sourceEventId} has no semantic authority.`, + ); + } + + return { + ...receipt, + agentId: authority.agentId, + deploymentVersionId: authority.deploymentVersionId, + deploymentVersionNumber: authority.deploymentVersionNumber, + eventJson: receipt.eventJson, + operation: authority.operation, + semanticHash: receipt.semanticHash, + }; +} + +function runtimeOperationEventReceiptsEqual( + left: RuntimeOperationEventReceipt, + right: RuntimeOperationEventReceipt, +): boolean { + return ( + left.deploymentVersionId === right.deploymentVersionId && + left.deploymentVersionNumber === right.deploymentVersionNumber && + left.eventId === right.eventId && + left.eventJson === right.eventJson && + left.eventType === right.eventType && + left.occurredAt === right.occurredAt && + left.operation === right.operation && + left.semanticHash === right.semanticHash && + left.seq === right.seq && + left.source === right.source && + left.sourceEventId === right.sourceEventId && + left.visibility === right.visibility + ); +} + +function runtimeOperationEventIdentitiesEqual( + left: Pick< + RuntimeOperationEventReceipt, + "agentId" | "deploymentVersionId" | "deploymentVersionNumber" | "operation" + >, + right: Pick< + RuntimeOperationEventReceipt, + "agentId" | "deploymentVersionId" | "deploymentVersionNumber" | "operation" + >, +): boolean { + return ( + left.agentId === right.agentId && + left.deploymentVersionId === right.deploymentVersionId && + left.deploymentVersionNumber === right.deploymentVersionNumber && + left.operation === right.operation + ); +} + +export async function adoptRuntimeOperationReadyReceipt( + database: D1Database, + input: { + readonly operationId: RuntimeOperationId; + readonly target: RuntimeSessionLifecycleTarget; + }, +): Promise<"applied" | "duplicate" | "missing" | "stale"> { + if (input.target.agentId == null) { + throw new Error("Runtime operation ready adoption requires an Agent-bound Session."); + } + const claimSourceEventId = `runtime-operation:${input.operationId}:${input.target.sessionId}:updating`; + const sourceEventId = `runtime-operation:${input.operationId}:${input.target.sessionId}:ready`; + const [claim, receipt] = await Promise.all([ + getRuntimeOperationEventReceipt(database, { + agentId: input.target.agentId, + operationId: input.operationId, + sessionId: input.target.sessionId, + status: "updating", + }), + getRuntimeOperationEventReceipt(database, { + agentId: input.target.agentId, + operationId: input.operationId, + sessionId: input.target.sessionId, + status: "ready", + }), + ]); + + if (receipt === null) { + return "missing"; + } + if (claim === null) { + throw new Error(`Runtime operation ready source ${sourceEventId} has no canonical claim.`); + } + if (!runtimeOperationEventIdentitiesEqual(claim, receipt)) { + throw new Error(`Runtime operation ready source ${sourceEventId} conflicts with its claim.`); + } + + const result = await database + .prepare( + `UPDATE session + SET status = 'IDLE', + status_operation_id = NULL, + status_seq = status_seq + 1, + updated_at = ? + WHERE id = ? + AND archived_at IS NULL + AND cleanup_operation_kind IS NULL + AND last_run_id IS ? + AND runtime_event_seq_cursor = ? + AND status = ? + AND status_operation_id = ? + AND status_seq = ? + AND updated_at = ? + AND EXISTS ( + SELECT 1 + FROM session_event AS event + WHERE event.session_id = ? + AND event.source_event_id = ? + AND event.id = ? + AND event.event_type = ? + AND event.occurred_at = ? + AND event.runtime_operation_event_json = ? + AND event.semantic_hash = ? + AND event.seq = ? + AND event.source = ? + AND event.visibility = ? + ) + AND EXISTS ( + SELECT 1 + FROM session_event AS claim + WHERE claim.session_id = ? + AND claim.source_event_id = ? + AND claim.id = ? + AND claim.event_type = ? + AND claim.occurred_at = ? + AND claim.runtime_operation_event_json = ? + AND claim.semantic_hash = ? + AND claim.seq = ? + AND claim.source = ? + AND claim.visibility = ? + )`, + ) + .bind( + receipt.occurredAt, + input.target.sessionId, + input.target.lastRunId, + input.target.sessionRuntimeEventSeqCursor, + input.target.sessionStatus, + input.operationId, + input.target.sessionStatusSeq, + input.target.sessionUpdatedAt, + input.target.sessionId, + sourceEventId, + receipt.eventId, + receipt.eventType, + receipt.occurredAt, + receipt.eventJson, + receipt.semanticHash, + receipt.seq, + receipt.source, + receipt.visibility, + input.target.sessionId, + claimSourceEventId, + claim.eventId, + claim.eventType, + claim.occurredAt, + claim.eventJson, + claim.semanticHash, + claim.seq, + claim.source, + claim.visibility, + ) + .run(); + if (getD1ChangeCount(result) > 0) { + return "applied"; + } + + const [currentClaim, currentReceipt, session] = await Promise.all([ + getRuntimeOperationEventReceipt(database, { + agentId: input.target.agentId, + operationId: input.operationId, + sessionId: input.target.sessionId, + status: "updating", + }), + getRuntimeOperationEventReceipt(database, { + agentId: input.target.agentId, + operationId: input.operationId, + sessionId: input.target.sessionId, + status: "ready", + }), + database + .prepare( + `SELECT archived_at, cleanup_operation_kind, last_run_id, + runtime_event_seq_cursor, status, status_operation_id + FROM session + WHERE id = ?`, + ) + .bind(input.target.sessionId) + .first<{ + archived_at: number | null; + cleanup_operation_kind: string | null; + last_run_id: string | null; + runtime_event_seq_cursor: number; + status: string; + status_operation_id: string | null; + }>(), + ]); + + return currentClaim !== null && + currentReceipt !== null && + runtimeOperationEventReceiptsEqual(currentClaim, claim) && + runtimeOperationEventReceiptsEqual(currentReceipt, receipt) && + session?.archived_at === null && + session.cleanup_operation_kind === null && + session.last_run_id === input.target.lastRunId && + session.runtime_event_seq_cursor >= receipt.seq && + session.status === "IDLE" && + session.status_operation_id === null + ? "duplicate" + : "stale"; +} + +export async function claimRuntimeOperationTargets( database: D1Database, input: { + readonly event: RuntimeOperationEvent; readonly operationId: RuntimeOperationId; readonly targets: readonly RuntimeSessionTarget[]; }, -): Promise { - if (input.targets.length === 0) { - return []; +): Promise { + const transitions: RuntimeSessionTargetTransition[] = []; + if (input.event.status !== "updating") { + throw new Error("Runtime operation target claims require one updating event."); } + const occurredAtMs = Date.parse(input.event.observedAt); - const expired: RuntimeSessionTarget[] = []; + if (!Number.isFinite(occurredAtMs)) { + throw new Error("Runtime operation start time must be a valid ISO timestamp."); + } for ( let index = 0; index < input.targets.length; index += RUNTIME_TARGET_STATUS_WRITE_BATCH_SIZE ) { - expired.push( - ...(await expireStaleRuntimeOperationTargetsBatch(database, { - operationId: input.operationId, - targets: input.targets.slice(index, index + RUNTIME_TARGET_STATUS_WRITE_BATCH_SIZE), - })), + const targets = input.targets.slice(index, index + RUNTIME_TARGET_STATUS_WRITE_BATCH_SIZE); + const prepared = await Promise.all( + targets.map(async (target) => { + if (target.agentId === null) { + throw new Error("Runtime operation target claims require an Agent-bound Session."); + } + const event = createRuntimeOperationSessionEvent({ + event: input.event, + operationId: input.operationId, + sessionId: target.sessionId, + }); + const preparedEvent = await prepareSessionLifecycleEvent(event, { + agentId: target.agentId, + operationId: input.operationId, + status: "updating", + }); + const nextStatusSeq = target.sessionStatusSeq + 1; + const nextEventSeq = target.sessionRuntimeEventSeqCursor + 1; + + return { + target, + statements: [ + database + .prepare( + `UPDATE session + SET runtime_event_seq_cursor = runtime_event_seq_cursor + 1, + status = 'RESCHEDULING', + status_operation_id = ?, + status_seq = status_seq + 1, + updated_at = ? + WHERE id = ? + AND archived_at IS NULL + AND cleanup_operation_kind IS NULL + AND runtime_provisioning_operation_id IS NULL + AND last_run_id IS ? + AND runtime_event_seq_cursor = ? + AND status = ? + AND status_operation_id IS ? + AND status_seq = ? + AND updated_at = ?`, + ) + .bind( + input.operationId, + occurredAtMs, + target.sessionId, + target.lastRunId, + target.sessionRuntimeEventSeqCursor, + target.sessionStatus, + target.sessionStatusOperationId, + target.sessionStatusSeq, + target.sessionUpdatedAt, + ), + ...createSessionLifecycleEventInsertStatements(database, { + prepared: preparedEvent, + target: { + lastRunId: target.lastRunId, + operationId: input.operationId, + runtimeEventSeqCursor: nextEventSeq, + sessionId: target.sessionId, + status: "RESCHEDULING", + statusSeq: nextStatusSeq, + updatedAt: occurredAtMs, + }, + }), + ], + }; + }), ); + const results = await database.batch(prepared.flatMap((record) => record.statements)); + + for (const [targetIndex, record] of prepared.entries()) { + if (getD1ChangeCount(results[targetIndex * 3]) === 0) { + continue; + } + transitions.push({ + current: { + ...record.target, + sessionRuntimeEventSeqCursor: record.target.sessionRuntimeEventSeqCursor + 1, + sessionStatus: "RESCHEDULING", + sessionStatusOperationId: input.operationId, + sessionStatusSeq: record.target.sessionStatusSeq + 1, + sessionUpdatedAt: occurredAtMs, + }, + }); + } } - return expired; + return transitions; } -async function expireStaleRuntimeOperationTargetsBatch( +export async function listRuntimeOperationTargets( database: D1Database, input: { readonly operationId: RuntimeOperationId; @@ -241,34 +1051,98 @@ async function expireStaleRuntimeOperationTargetsBatch( return []; } - const now = currentTimestampMs(); - const whereClause = and( - eq(sessionsTable.status, "RESCHEDULING"), - eq(sessionsTable.statusOperationId, input.operationId), - lte(sessionsTable.updatedAt, now - RESCHEDULING_RECONNECT_WINDOW_MS), - or( - ...input.targets.map((target) => - and( - eq(sessionsTable.id, target.sessionId), - eq(sessionsTable.statusSeq, target.sessionStatusSeq), + const targetsById = new Map(input.targets.map((target) => [target.sessionId, target])); + const rows = await getAppDatabase(database) + .select({ + lastRunId: sessionsTable.lastRunId, + sessionId: sessionsTable.id, + sessionRuntimeEventSeqCursor: sessionsTable.runtimeEventSeqCursor, + sessionStatusOperationId: sessionsTable.statusOperationId, + sessionStatusSeq: sessionsTable.statusSeq, + sessionStatus: sql`${sessionsTable.status}`, + sessionUpdatedAt: sessionsTable.updatedAt, + }) + .from(sessionsTable) + .where( + and( + isNull(sessionsTable.archivedAt), + isNull(sessionsTable.runtimeProvisioningOperationId), + inArray( + sessionsTable.id, + input.targets.map((target) => target.sessionId), ), + inArray(sessionsTable.status, ["IDLE", "RESCHEDULING"]), + eq(sessionsTable.statusOperationId, input.operationId), ), - ), - ); - const results = await getAppDatabase(database) - .update(sessionsTable) - .set( - createSessionStatusTransitionPatch({ - status: "TERMINATED", - timestampMs: now, - }), ) - .where(whereClause) - .returning({ id: sessionsTable.id }) .all(); - const expiredIds = new Set(results.map((row) => row.id)); - return input.targets.filter((target) => expiredIds.has(target.sessionId)); + return rows.flatMap((row) => { + const target = targetsById.get(row.sessionId); + + return target === undefined ? [] : [{ ...target, ...row }]; + }); } -export { isTerminalRunStatus }; +export async function listStaleRuntimeOperationTargets( + database: D1Database, + input: { + readonly limit: number; + readonly staleUpdatedAtLte: number; + }, +): Promise { + if (!Number.isSafeInteger(input.limit) || input.limit <= 0) { + throw new Error("Stale runtime operation target limit must be a positive integer."); + } + + const rows = await getAppDatabase(database) + .select({ + agentId: sessionsTable.agentId, + creatorAccountId: sessionsTable.creatorAccountId, + lastRunId: sessionsTable.lastRunId, + operationId: sessionsTable.statusOperationId, + sandboxId: sandboxSessionsTable.sandboxId, + sessionId: sessionsTable.id, + sessionRuntimeEventSeqCursor: sessionsTable.runtimeEventSeqCursor, + sessionStatusOperationId: sessionsTable.statusOperationId, + sessionStatusSeq: sessionsTable.statusSeq, + sessionStatus: sql`${sessionsTable.status}`, + sessionUpdatedAt: sessionsTable.updatedAt, + }) + .from(sessionsTable) + .innerJoin(sandboxSessionsTable, eq(sandboxSessionsTable.sessionId, sessionsTable.id)) + .where( + and( + isNull(sessionsTable.archivedAt), + isNull(sessionsTable.cleanupOperationKind), + isNull(sessionsTable.runtimeProvisioningOperationId), + inArray(sessionsTable.status, ["IDLE", "RESCHEDULING"]), + isNotNull(sessionsTable.statusOperationId), + lte(sessionsTable.updatedAt, input.staleUpdatedAtLte), + ), + ) + .orderBy(asc(sessionsTable.updatedAt), asc(sessionsTable.id)) + .limit(input.limit) + .all(); + + return rows.flatMap((row) => + row.operationId === null ? [] : [{ ...row, operationId: row.operationId }], + ); +} + +export async function expireStaleRuntimeOperationTargets( + database: D1Database, + input: { + readonly operationId: RuntimeOperationId; + readonly targets: readonly RuntimeSessionTarget[]; + }, +): Promise { + if (input.targets.length === 0) { + return []; + } + + const currentTargets = await listRuntimeOperationTargets(database, input); + const staleBeforeMs = currentTimestampMs() - RESCHEDULING_RECONNECT_WINDOW_MS; + + return currentTargets.filter((target) => target.sessionUpdatedAt <= staleBeforeMs); +} diff --git a/apps/api/src/modules/runtime/application/session-lifecycle-transition.service.ts b/apps/api/src/modules/runtime/application/session-lifecycle-transition.service.ts deleted file mode 100644 index 9f2f1910..00000000 --- a/apps/api/src/modules/runtime/application/session-lifecycle-transition.service.ts +++ /dev/null @@ -1,23 +0,0 @@ -import type { RunError, SessionRunStatus } from "@mosoo/contracts/session-run"; -import type { SessionRunId } from "@mosoo/id"; - -import { setSessionRunStatus } from "../infrastructure/session-runs/session-run-store.repository"; - -export { createSessionStatusTransitionPatch } from "../infrastructure/session-runs/session-lifecycle-projection.repository"; -export type { SessionRunTransitionOutcome } from "../infrastructure/session-runs/session-run-store.repository"; - -export async function setSystemSessionRunStatus( - database: D1Database, - input: { - error?: RunError | null; - runId: SessionRunId; - status: SessionRunStatus; - }, -) { - return setSessionRunStatus(database, { - ...(input.error !== undefined ? { error: input.error } : {}), - runId: input.runId, - source: "system", - status: input.status, - }); -} diff --git a/apps/api/src/modules/runtime/application/session-runs/cancel-run.service.ts b/apps/api/src/modules/runtime/application/session-runs/cancel-run.service.ts index 7c392a3f..4426cc9b 100644 --- a/apps/api/src/modules/runtime/application/session-runs/cancel-run.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/cancel-run.service.ts @@ -1,13 +1,13 @@ import type { RuntimeCommand } from "@mosoo/contracts/runtime-command"; import type { SessionRunSummary } from "@mosoo/contracts/session-run"; -import { sessionRunsTable, sessionsTable } from "@mosoo/db"; +import { driverInstancesTable, sessionRunsTable, sessionsTable } from "@mosoo/db"; +import { sleepPromise } from "@mosoo/effects"; import { createPlatformId, parsePlatformId } from "@mosoo/id"; import type { AccountId, DriverCommandId, DriverInstanceId, AppId, - RuntimeEventId, SessionId, SessionRunId, } from "@mosoo/id"; @@ -19,18 +19,18 @@ import { getAppDatabase } from "../../../../platform/db/drizzle"; import { isTruthy } from "../../../../shared/truthiness"; import { ensureAppOwnership } from "../../../apps/application/app.service"; import type { AuthenticatedViewer } from "../../../auth/application/viewer-auth.service"; -import { appendSessionRuntimeEvents } from "../../../sessions/application/session-event-write.service"; import { sessionParticipantCondition } from "../../../sessions/domain/session-access.policy"; +import { RUNTIME_SOCKET_TIMEOUT_MS } from "../../domain/runtime-config"; import { sendDriverInstanceCommand } from "../../infrastructure/driver-instance/client"; import { isDriverControlSocketMissingError } from "../../infrastructure/driver-session-stop-errors"; +import { stopDriverSession } from "../../infrastructure/driver-session-stop.service"; import { expireUndeliveredInputStartCommandsForRun } from "../../infrastructure/session-runs/runtime-command-store.repository"; import { toSessionRunSummary } from "../../infrastructure/session-runs/session-run-row.mapper"; import type { SessionRunRow } from "../../infrastructure/session-runs/session-run-row.mapper"; -import { - getSessionRunSummary, - setSessionRunStatus, -} from "../../infrastructure/session-runs/session-run-store.repository"; -import { createCancelledSessionRunRuntimeEvent } from "./session-run-view-events.service"; +import { getSessionRunSummary } from "../../infrastructure/session-runs/session-run-store.repository"; +import { recordCanonicalSessionRunTerminal } from "./session-run-terminal-failure.service"; + +const RUN_CANCEL_POLL_MS = 100; interface CancelSessionRunInput { appId: AppId; runId: SessionRunId; @@ -42,7 +42,12 @@ async function getOwnedSessionRun( viewerId: AccountId, input: CancelSessionRunInput, ): Promise<{ + driverConnectionId: string | null; + driverGeneration: number | null; driverInstanceId: DriverInstanceId | null; + driverLastHeartbeatAt: number | null; + driverStatus: string | null; + driverUpdatedAt: number | null; run: SessionRunSummary; sessionId: SessionId; } | null> { @@ -53,10 +58,16 @@ async function getOwnedSessionRun( created_at: sessionRunsTable.createdAt, deployment_version_id: sessionRunsTable.deploymentVersionId, deployment_version_number: sessionRunsTable.deploymentVersionNumber, + driver_connection_id: driverInstancesTable.connectionId, + driver_generation: driverInstancesTable.generation, driver_instance_id: sessionRunsTable.driverInstanceId, + driver_last_heartbeat_at: driverInstancesTable.lastHeartbeatAt, + driver_status: driverInstancesTable.status, + driver_updated_at: driverInstancesTable.updatedAt, error_code: sessionRunsTable.errorCode, error_details_json: sessionRunsTable.errorDetailsJson, error_message: sessionRunsTable.errorMessage, + error_retryable: sessionRunsTable.errorRetryable, id: sessionRunsTable.id, model: sessionRunsTable.model, provider: sessionRunsTable.provider, @@ -68,6 +79,10 @@ async function getOwnedSessionRun( updated_at: sessionRunsTable.updatedAt, }) .from(sessionRunsTable) + .leftJoin( + driverInstancesTable, + eq(driverInstancesTable.id, sessionRunsTable.driverInstanceId), + ) .innerJoin(sessionsTable, eq(sessionsTable.id, sessionRunsTable.sessionId)) .where( and( @@ -85,12 +100,43 @@ async function getOwnedSessionRun( } return { + driverConnectionId: row.driver_connection_id, + driverGeneration: row.driver_generation, driverInstanceId: row.driver_instance_id, + driverLastHeartbeatAt: row.driver_last_heartbeat_at, + driverStatus: row.driver_status, + driverUpdatedAt: row.driver_updated_at, run: toSessionRunSummary(row satisfies SessionRunRow), sessionId: row.session_id, }; } +async function waitForDriverTerminalRun( + database: D1Database, + runId: SessionRunId, +): Promise { + const deadline = Date.now() + RUNTIME_SOCKET_TIMEOUT_MS; + + while (true) { + const run = await getSessionRunSummary(database, runId); + if (run === null) { + throw new Error("Session run disappeared while waiting for cancellation."); + } + if ( + run.status === "cancelled" || + run.status === "completed" || + run.status === "expired" || + run.status === "failed" + ) { + return run; + } + if (Date.now() >= deadline) { + throw new Error("Driver did not settle the run cancellation before the control timeout."); + } + await sleepPromise(RUN_CANCEL_POLL_MS); + } +} + export async function cancelRun( bindings: ApiBindings, viewer: AuthenticatedViewer, @@ -137,55 +183,70 @@ export async function cancelRun( }; } - if (isTruthy(run.driverInstanceId)) { + let driverTerminalRun: SessionRunSummary | null = null; + let requiresSyntheticCancellation = + !isTruthy(run.driverInstanceId) || run.driverGeneration === null; + + if (!requiresSyntheticCancellation && run.driverInstanceId && run.driverGeneration !== null) { const command: RuntimeCommand = { commandId: createPlatformId(), kind: "turn.cancel", reason: "viewer.cancelled", + runId, }; try { - await sendDriverInstanceCommand(bindings, run.driverInstanceId, command); + await sendDriverInstanceCommand( + bindings, + run.driverInstanceId, + run.driverGeneration, + command, + ); + driverTerminalRun = await waitForDriverTerminalRun(database, runId); } catch (error) { if (!isDriverControlSocketMissingError(error)) { throw error; } + requiresSyntheticCancellation = true; } } - const outcome = await setSessionRunStatus(database, { - runId, - source: "viewer", - status: "cancelled", - }); - - if (outcome.kind === "repair_needed") { - throw new Error("Session lifecycle projection needs repair."); - } - - if (outcome.kind === "duplicate") { - if (isTruthy(run.driverInstanceId)) { - await expireUndeliveredInputStartCommandsForRun(database, { - driverInstanceId: run.driverInstanceId, + const outcome = requiresSyntheticCancellation + ? await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + error: null, + ...(run.driverGeneration === null || run.driverInstanceId === null + ? {} + : { + expectedDriverObservation: { + connectionId: run.driverConnectionId, + driverInstanceId: run.driverInstanceId, + generation: run.driverGeneration, + lastHeartbeatAt: run.driverLastHeartbeatAt, + status: run.driverStatus, + updatedAt: run.driverUpdatedAt, + }, + }), runId, - }); - } - - return { - run: outcome.run, - }; - } - - if (outcome.kind === "rejected" || outcome.kind === "stale") { - const latestRun = await getSessionRunSummary(database, runId); + sessionId: run.sessionId, + source: "viewer", + status: "cancelled", + }) + : null; - return { - run: latestRun ?? currentRun, - }; + if ( + outcome?.kind === "committed" && + run.driverInstanceId !== null && + run.driverGeneration !== null + ) { + await stopDriverSession(bindings, { + driverInstanceId: run.driverInstanceId, + expectedDriverGeneration: run.driverGeneration, + expectedSessionRunId: runId, + reason: "viewer.cancelled", + }); } - const updatedRun = outcome.run; - if (isTruthy(run.driverInstanceId)) { await expireUndeliveredInputStartCommandsForRun(database, { driverInstanceId: run.driverInstanceId, @@ -193,17 +254,9 @@ export async function cancelRun( }); } - const cancelledEvent = createCancelledSessionRunRuntimeEvent({ - eventId: createPlatformId(), - run: updatedRun, - sessionId: run.sessionId, - sourceEventId: `viewer-cancel:${runId}:cancelled`, - }); - await appendSessionRuntimeEvents({ - bindings, - events: [cancelledEvent], - sessionId: run.sessionId, - }); + if (outcome?.kind === "stale") { + return { run: outcome.run }; + } logInfo("session.turn.cancelled", { driverInstanceId: run.driverInstanceId, @@ -214,6 +267,6 @@ export async function cancelRun( }); return { - run: updatedRun, + run: outcome?.run ?? driverTerminalRun ?? currentRun, }; } diff --git a/apps/api/src/modules/runtime/application/session-runs/dispatch-queued-run.service.ts b/apps/api/src/modules/runtime/application/session-runs/dispatch-queued-run.service.ts index b04bfe64..c9ffc96a 100644 --- a/apps/api/src/modules/runtime/application/session-runs/dispatch-queued-run.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/dispatch-queued-run.service.ts @@ -6,14 +6,13 @@ import { logError, logInfo, logWarn } from "../../../../platform/cloudflare/logg import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; import type { AuthenticatedViewer } from "../../../auth/application/viewer-auth.service"; import { fileStore } from "../../../files/application/file-store"; -import { appendSessionRuntimeEvents } from "../../../sessions/application/session-event-write.service"; import { getSupportedRuntimeId } from "../../domain/runtime-config"; import { hydrateCachedRunContextFromSession } from "../session-definition/hydrate-run-context.service"; import { appendSessionResourceContextToPrompt } from "../session-resources/session-resource-prompt.service"; import { dispatchSessionRun } from "./dispatch-run.service"; import { describeRunError } from "./run-error-message"; -import { getSessionRunState, updateSessionRunStatusIfActive } from "./session-run-state.repository"; -import { createFailedSessionRunRuntimeEvent } from "./session-run-view-events.service"; +import { getSessionRunState } from "./session-run-state.repository"; +import { recordCanonicalSessionRunTerminal } from "./session-run-terminal-failure.service"; import { appendSessionRuntimeTimingEventBestEffort, createRuntimeTimingRecorder, @@ -35,42 +34,28 @@ async function failQueuedSessionRunBeforeDispatch( message, retryable: false, } as const; - // Only fail the run while it is still queued. Once another dispatcher CASed - // it to booting, this path is a losing contender and its hydration error - // must not tear down the run the winner is provisioning. - const failedRun = await updateSessionRunStatusIfActive(bindings.DB, { + const outcome = await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, error: runError, - expectedCurrentStatus: "queued", + expectedRunStatus: "queued", runId: input.sessionRunId, + sessionId: input.sessionId, + source: "api", status: "failed", }); - if (!failedRun) { - const state = await getSessionRunState(bindings.DB, input.sessionRunId); - + if (outcome.kind === "stale") { logWarn("session.run.context_hydration.failed.run-not-queued", { message, runId: input.sessionRunId, sessionId: input.sessionId, - status: state?.status ?? null, + status: outcome.run.status, traceId: input.traceId, }); return; } - await appendSessionRuntimeEvents({ - bindings, - events: [ - createFailedSessionRunRuntimeEvent({ - run: failedRun, - runError, - sessionId: input.sessionId, - }), - ], - sessionId: input.sessionId, - }); - logError("session.run.context_hydration.failed", { message, runId: input.sessionRunId, diff --git a/apps/api/src/modules/runtime/application/session-runs/dispatch-run-cleanup.service.ts b/apps/api/src/modules/runtime/application/session-runs/dispatch-run-cleanup.service.ts index 36dcf031..4e282d56 100644 --- a/apps/api/src/modules/runtime/application/session-runs/dispatch-run-cleanup.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/dispatch-run-cleanup.service.ts @@ -9,6 +9,7 @@ import { expireUndeliveredInputStartCommandsForRun } from "../../infrastructure/ export async function cleanupDispatchedDriver( bindings: ApiBindings, input: { + driverGeneration: number; driverInstanceId: DriverInstanceId; reason: string; runId: SessionRunId; @@ -23,6 +24,8 @@ export async function cleanupDispatchedDriver( }); await stopDriverSession(bindings, { driverInstanceId: input.driverInstanceId, + expectedDriverGeneration: input.driverGeneration, + expectedSessionRunId: input.runId, reason: input.reason, }); } catch (error) { @@ -32,6 +35,7 @@ export async function cleanupDispatchedDriver( try { await createRuntimeSubjectLifecycleService(bindings).releaseRunLease({ driverInstanceId: input.driverInstanceId, + expectedDriverGeneration: input.driverGeneration, expectedSessionRunId: input.runId, }); } catch (releaseError) { diff --git a/apps/api/src/modules/runtime/application/session-runs/dispatch-run.service.ts b/apps/api/src/modules/runtime/application/session-runs/dispatch-run.service.ts index e812ae35..a2c1fdfc 100644 --- a/apps/api/src/modules/runtime/application/session-runs/dispatch-run.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/dispatch-run.service.ts @@ -118,6 +118,7 @@ export async function dispatchSessionRun( }, ): Promise { const sandboxId = input.profile.sandbox.id; + let driverGeneration: number | null = null; let driverInstanceId: DriverInstanceId | null = null; let prepareTimingEventPromise: Promise = Promise.resolve(); // Boot-payload config traces are owner-debug telemetry; they persist off the @@ -208,6 +209,7 @@ export async function dispatchSessionRun( traceId: input.traceId, }); runLease = preparedRunLease; + driverGeneration = preparedRunLease.driverGeneration; driverInstanceId = preparedRunLease.driverInstanceId; const preparedDriverInstanceId = preparedRunLease.driverInstanceId; logInfo("session.run.prepared", { @@ -232,6 +234,7 @@ export async function dispatchSessionRun( await dispatchTiming.measure("dispatchDriverTurn", () => executionPlane.dispatchTurn(bindings, { attachmentIds: input.attachmentIds, + driverGeneration: preparedRunLease.driverGeneration, driverInstanceId: preparedDriverInstanceId, prompt: input.prompt, sessionRunId: input.sessionRunId, @@ -273,8 +276,9 @@ export async function dispatchSessionRun( traceId: input.traceId, }); - if (isTruthy(driverInstanceId)) { + if (isTruthy(driverInstanceId) && driverGeneration !== null) { await cleanupDispatchedDriver(bindings, { + driverGeneration, driverInstanceId, reason: "session.run.pre-ready-retry", runId: input.sessionRunId, @@ -285,6 +289,7 @@ export async function dispatchSessionRun( runLease?.release(); runLease = null; + driverGeneration = null; driverInstanceId = null; // Stop retrying if the run was cancelled or failed elsewhere while @@ -310,8 +315,9 @@ export async function dispatchSessionRun( await pendingBootPayloadEvents; if (error instanceof SessionRunNoLongerActiveError) { - if (isTruthy(driverInstanceId)) { + if (isTruthy(driverInstanceId) && driverGeneration !== null) { await cleanupDispatchedDriver(bindings, { + driverGeneration, driverInstanceId, reason: `session.run.${error.status}`, runId: input.sessionRunId, @@ -340,8 +346,9 @@ export async function dispatchSessionRun( retryable: false, } as const; - if (isTruthy(driverInstanceId)) { + if (isTruthy(driverInstanceId) && driverGeneration !== null) { await cleanupDispatchedDriver(bindings, { + driverGeneration, driverInstanceId, reason: "session.run.provision-failed", runId: input.sessionRunId, @@ -356,12 +363,6 @@ export async function dispatchSessionRun( sessionId: input.sessionId, source: "api", }); - if (failureOutcome.kind === "repair_needed") { - throw new Error("Session lifecycle projection needs repair.", { - cause: error, - }); - } - if (failureOutcome.kind === "not_failed") { const state = await getSessionRunState(bindings.DB, input.sessionRunId); diff --git a/apps/api/src/modules/runtime/application/session-runs/prewarm-agent-session-runtime.service.ts b/apps/api/src/modules/runtime/application/session-runs/prewarm-agent-session-runtime.service.ts index ba2eb7f4..af1ec072 100644 --- a/apps/api/src/modules/runtime/application/session-runs/prewarm-agent-session-runtime.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/prewarm-agent-session-runtime.service.ts @@ -73,25 +73,27 @@ export async function prewarmAgentSessionRuntime( } const sandboxId = hydrated.value.profile.sandbox.id; - const { subject: sandbox } = await timing.measure("activateRuntimeSubject", () => - createRuntimeSubjectLifecycleService(bindings).activate({ - agentId: hydrated.value.profile.agentId, - executionOwnerUserId: hydrated.value.profile.session.origin.executionOwnerUserId, - kind: hydrated.value.profile.kind, - networkConstraints: resolveRuntimeSubjectNetworkConstraints(bindings, { - envVars: hydrated.value.profile.envVars, + const { incarnation: sandboxIncarnation, subject: sandbox } = await timing.measure( + "activateRuntimeSubject", + () => + createRuntimeSubjectLifecycleService(bindings).activate({ + agentId: hydrated.value.profile.agentId, + executionOwnerUserId: hydrated.value.profile.session.origin.executionOwnerUserId, kind: hydrated.value.profile.kind, - network: hydrated.value.profile.network, - requestUrl: request.requestUrl, + networkConstraints: resolveRuntimeSubjectNetworkConstraints(bindings, { + envVars: hydrated.value.profile.envVars, + kind: hydrated.value.profile.kind, + network: hydrated.value.profile.network, + requestUrl: request.requestUrl, + subjectKind: hydrated.value.profile.sandbox.subjectKind, + }), + runtimeSubjectId: sandboxId, + appId: session.appId, + purpose: "prewarm", + subjectId: hydrated.value.profile.sandbox.subjectId, subjectKind: hydrated.value.profile.sandbox.subjectKind, + timing, }), - runtimeSubjectId: sandboxId, - appId: session.appId, - purpose: "prewarm", - subjectId: hydrated.value.profile.sandbox.subjectId, - subjectKind: hydrated.value.profile.sandbox.subjectKind, - timing, - }), ); handles.subject = sandbox; @@ -103,6 +105,7 @@ export async function prewarmAgentSessionRuntime( origin: hydrated.value.profile.session.origin, sandbox, sandboxId, + sandboxIncarnation, sessionId: session.id, timing, }), @@ -136,6 +139,7 @@ export async function prewarmAgentSessionRuntime( resolvedSkillCatalog: hydrated.value.skillCatalog, resolvedSkills: hydrated.value.skills, sandbox, + sandboxIncarnation, sandboxSessionId: session.id, sessionId: session.id, }), diff --git a/apps/api/src/modules/runtime/application/session-runs/resolve-permission-request.service.ts b/apps/api/src/modules/runtime/application/session-runs/resolve-permission-request.service.ts index 7a9509b8..7a0dfa0b 100644 --- a/apps/api/src/modules/runtime/application/session-runs/resolve-permission-request.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/resolve-permission-request.service.ts @@ -1,6 +1,6 @@ import { driverInstancesTable, sessionRunsTable, sessionsTable } from "@mosoo/db"; import { createPlatformId, parsePlatformId } from "@mosoo/id"; -import type { AccountId, DriverInstanceId, AppId, SessionId } from "@mosoo/id"; +import type { AccountId, DriverInstanceId, AppId, SessionId, SessionRunId } from "@mosoo/id"; import { and, eq, inArray } from "drizzle-orm"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; @@ -16,6 +16,7 @@ interface ResolveDriverPermissionInput { driverInstanceId: DriverInstanceId; appId: AppId; requestId: string; + runId: SessionRunId; sessionId: SessionId; } @@ -29,6 +30,8 @@ export async function resolvePermissionRequest( const row = (await getAppDatabase(bindings.DB) .select({ + driverGeneration: driverInstancesTable.generation, + runId: sessionRunsTable.id, sessionId: sessionsTable.id, }) .from(driverInstancesTable) @@ -43,6 +46,7 @@ export async function resolvePermissionRequest( .where( and( eq(driverInstancesTable.id, input.driverInstanceId), + eq(sessionRunsTable.id, input.runId), eq(sessionsTable.id, input.sessionId), eq(sessionsTable.appId, input.appId), sessionParticipantCondition(viewerId), @@ -55,10 +59,11 @@ export async function resolvePermissionRequest( throw new Error("Driver instance not found."); } - await sendDriverInstanceCommand(bindings, input.driverInstanceId, { + await sendDriverInstanceCommand(bindings, input.driverInstanceId, row.driverGeneration, { commandId: createPlatformId(), decision: input.decision, kind: "permission.resolve", requestId: input.requestId, + runId: row.runId, }); } diff --git a/apps/api/src/modules/runtime/application/session-runs/send-agent-session-events.service.ts b/apps/api/src/modules/runtime/application/session-runs/send-agent-session-events.service.ts index b01c269a..15df7863 100644 --- a/apps/api/src/modules/runtime/application/session-runs/send-agent-session-events.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/send-agent-session-events.service.ts @@ -258,7 +258,7 @@ async function handleAgentSessionEvent(input: { if (updated) { await appendSessionRuntimeEvents({ bindings: input.bindings, - events: [updated.event], + events: updated.events, sessionId: input.sessionId, }); } diff --git a/apps/api/src/modules/runtime/application/session-runs/session-permission-decision.service.ts b/apps/api/src/modules/runtime/application/session-runs/session-permission-decision.service.ts index ad1dc011..3b600dce 100644 --- a/apps/api/src/modules/runtime/application/session-runs/session-permission-decision.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/session-permission-decision.service.ts @@ -1,5 +1,5 @@ import { parsePlatformId } from "@mosoo/id"; -import type { DriverInstanceId, AppId, SessionId } from "@mosoo/id"; +import type { DriverInstanceId, AppId, SessionId, SessionRunId } from "@mosoo/id"; import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; @@ -17,7 +17,7 @@ import { resolvePermissionRequest } from "./resolve-permission-request.service"; type PermissionDecision = "allow_once" | "reject_once"; export interface SessionPermissionStateUpdate { - event: RuntimeEventEnvelope; + events: RuntimeEventEnvelope[]; state: SessionLiveState; } @@ -68,27 +68,31 @@ function requirePermissionRequestDriverInstanceId( async function createPermissionStateUpdate(input: { currentState: SessionLiveState; + events: RuntimeEventEnvelope[]; +}): Promise { + return { + events: input.events, + state: input.events.reduce(applyRuntimeEventToSessionLiveState, input.currentState), + }; +} + +async function createPermissionResolvedEvent(input: { outcome?: PermissionDecision; - permissionRequests: SessionLiveState["permissionRequests"]; - requestId?: string; + requestId: string; + runId: SessionRunId; sessionId: SessionId; -}): Promise { - const event = createSessionRuntimeEvent({ +}): Promise { + return createSessionRuntimeEvent({ actor: "user", kind: "permission.resolved", origin: "viewer", payload: { ...(input.outcome === undefined ? {} : { outcome: input.outcome }), - permissionRequests: input.permissionRequests, - ...(input.requestId === undefined ? {} : { requestId: input.requestId }), + requestId: input.requestId, }, + runId: input.runId, sessionId: input.sessionId, }); - - return { - event, - state: applyRuntimeEventToSessionLiveState(input.currentState, event), - }; } export async function resolveSessionPermissionDecision( @@ -112,19 +116,22 @@ export async function resolveSessionPermissionDecision( driverInstanceId: requirePermissionRequestDriverInstanceId(request), appId: input.appId, requestId: input.requestId, + runId: parsePlatformId(request.runId, "permission request run id"), sessionId: input.sessionId, }); - const permissionRequests = currentState.permissionRequests.filter( - (candidate) => candidate.requestId !== input.requestId, - ); + const runId = parsePlatformId(request.runId, "permission request run id"); return createPermissionStateUpdate({ currentState, - outcome: input.decision, - permissionRequests, - requestId: input.requestId, - sessionId: input.sessionId, + events: [ + await createPermissionResolvedEvent({ + outcome: input.decision, + requestId: input.requestId, + runId, + sessionId: input.sessionId, + }), + ], }); } @@ -141,8 +148,6 @@ export async function rejectSessionPermissionRequests( return null; } - const remainingRequests: SessionLiveState["permissionRequests"] = []; - const cleanupResults = await runOrderedAsyncTasks( currentState.permissionRequests.map((request) => async () => { try { @@ -151,6 +156,7 @@ export async function rejectSessionPermissionRequests( driverInstanceId: requirePermissionRequestDriverInstanceId(request), appId: input.appId, requestId: request.requestId, + runId: parsePlatformId(request.runId, "permission request run id"), sessionId: input.sessionId, }); return { rejected: true, request }; @@ -161,19 +167,25 @@ export async function rejectSessionPermissionRequests( }), ); - for (const result of cleanupResults) { - if (!result.rejected) { - remainingRequests.push(result.request); - } - } + const rejectedRequests = cleanupResults.flatMap((result) => + result.rejected ? [result.request] : [], + ); - if (remainingRequests.length === currentState.permissionRequests.length) { + if (rejectedRequests.length === 0) { return null; } return createPermissionStateUpdate({ currentState, - permissionRequests: remainingRequests, - sessionId: input.sessionId, + events: await Promise.all( + rejectedRequests.map((request) => + createPermissionResolvedEvent({ + outcome: "reject_once", + requestId: request.requestId, + runId: parsePlatformId(request.runId, "permission request run id"), + sessionId: input.sessionId, + }), + ), + ), }); } diff --git a/apps/api/src/modules/runtime/application/session-runs/session-run-state.repository.ts b/apps/api/src/modules/runtime/application/session-runs/session-run-state.repository.ts index dd282636..8ca8dc5e 100644 --- a/apps/api/src/modules/runtime/application/session-runs/session-run-state.repository.ts +++ b/apps/api/src/modules/runtime/application/session-runs/session-run-state.repository.ts @@ -1,4 +1,4 @@ -import type { RunError, SessionRunSummary } from "@mosoo/contracts/session-run"; +import type { SessionRunSummary } from "@mosoo/contracts/session-run"; import { sessionRunsTable } from "@mosoo/db"; import type { DriverInstanceId, SessionRunId } from "@mosoo/id"; import { eq } from "drizzle-orm"; @@ -59,40 +59,6 @@ export async function getSessionRunState( return row; } -export async function updateSessionRunStatusIfActive( - database: D1Database, - input: { - error?: RunError | null; - expectedCurrentStatus?: SessionRunSummary["status"]; - runId: SessionRunId; - status: SessionRunSummary["status"]; - }, -): Promise { - const outcome = await setSessionRunStatus(database, { - ...(input.error !== undefined ? { error: input.error } : {}), - ...(input.expectedCurrentStatus !== undefined - ? { expectedCurrentStatus: input.expectedCurrentStatus } - : {}), - runId: input.runId, - source: "api", - status: input.status, - }); - - switch (outcome.kind) { - case "applied": - case "duplicate": { - return outcome.run; - } - case "repair_needed": { - throw new Error("Session lifecycle projection needs repair."); - } - case "rejected": - case "stale": { - return null; - } - } -} - export async function acquireSessionRunDispatch( database: D1Database, runId: SessionRunId, diff --git a/apps/api/src/modules/runtime/application/session-runs/session-run-terminal-failure.service.ts b/apps/api/src/modules/runtime/application/session-runs/session-run-terminal-failure.service.ts index 2c8cc17b..4775e8db 100644 --- a/apps/api/src/modules/runtime/application/session-runs/session-run-terminal-failure.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/session-run-terminal-failure.service.ts @@ -1,28 +1,167 @@ -import type { RunError } from "@mosoo/contracts/session-run"; -import type { SessionId, SessionRunId } from "@mosoo/id"; +import type { RunError, SessionRunSummary } from "@mosoo/contracts/session-run"; +import { + sessionEventsTable, + sessionMessagesTable, + sessionRunsTable, + sessionsTable, +} from "@mosoo/db"; +import { createPlatformId } from "@mosoo/id"; +import type { + DriverInstanceId, + RuntimeOperationId, + RuntimeEventId, + SessionId, + SessionRunId, +} from "@mosoo/id"; +import { and, eq, inArray, sql } from "drizzle-orm"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; -import { appendSessionRuntimeEvents } from "../../../sessions/application/session-event-write.service"; -import { createSessionRunTerminalFailureSourceId } from "../../domain/session-run-terminal-event-id"; +import { getAppDatabase } from "../../../../platform/db/drizzle"; +import { currentTimestampMs, toIsoString } from "../../../../time"; +import { publishPersistedSessionRuntimeEvents } from "../../../sessions/application/session-event-write.service"; +import { readTerminalEventSemanticAuthority } from "../../../sessions/domain/session-terminal-event-authority"; +import { createSessionRunTerminalSourceId } from "../../domain/session-run-terminal-event-id"; +import type { PreparedAssistantMessageProjection } from "../../infrastructure/driver-instance/assistant-message-projection"; +import { commitTerminalRunProjection } from "../../infrastructure/driver-instance/completed-run-commit.repository"; +import type { + ExpectedTerminalDriverObservation, + ExpectedTerminalSessionObservation, + HostTerminalRunStatus, + TerminalRunProjectionSource, +} from "../../infrastructure/driver-instance/completed-run-commit.repository"; +import { createCanonicalDriverRunFailedEvent } from "../../infrastructure/driver-instance/driver-event-canonicalization"; +import { getSessionRunSummary } from "../../infrastructure/session-runs/session-run-store.repository"; import { - getSessionRunSummary, - setSessionRunStatus, -} from "../../infrastructure/session-runs/session-run-store.repository"; -import type { SessionRunTransitionOutcome } from "../../infrastructure/session-runs/session-run-store.repository"; -import { createFailedSessionRunRuntimeEvent } from "./session-run-view-events.service"; + createCompletedSessionRunRuntimeEvent, + createFailedSessionRunRuntimeEvent, + createSessionRunUpdatedEvent, +} from "./session-run-view-events.service"; export type CanonicalSessionRunFailureOutcome = + | { kind: "failed" } + | { kind: "not_failed"; status: SessionRunSummary["status"] }; + +export type CanonicalSessionRunTerminalOutcome = | { - kind: "failed"; + commitKind: "applied" | "duplicate"; + kind: "committed"; + run: SessionRunSummary; } - | { - kind: "not_failed"; - transition: SessionRunTransitionOutcome; + | { kind: "stale"; run: SessionRunSummary }; + +function terminalRunSummary( + current: SessionRunSummary, + input: { + error: RunError | null; + status: HostTerminalRunStatus; + }, + timestamp: string, +): SessionRunSummary { + return { + ...current, + completedAt: current.completedAt ?? timestamp, + error: input.error, + startedAt: current.startedAt ?? timestamp, + status: input.status, + updatedAt: current.status === input.status ? current.updatedAt : timestamp, + }; +} + +async function readRunExecutionIdentity( + database: D1Database, + runId: SessionRunId, +): Promise<{ driverInstanceId: DriverInstanceId; runtimeId: string } | null> { + const row = await getAppDatabase(database) + .select({ + driverInstanceId: sessionRunsTable.driverInstanceId, + runtimeId: sql< + string | null + >`coalesce(${sessionRunsTable.runtimeId}, ${sessionsTable.runtimeId})`, + }) + .from(sessionRunsTable) + .innerJoin(sessionsTable, eq(sessionsTable.id, sessionRunsTable.sessionId)) + .where(eq(sessionRunsTable.id, runId)) + .limit(1) + .get(); + + return row?.driverInstanceId === null || row?.runtimeId === null || row === undefined + ? null + : { driverInstanceId: row.driverInstanceId, runtimeId: row.runtimeId }; +} + +async function hasCompletedWithoutAssistantAuthority( + database: D1Database, + input: { readonly runId: SessionRunId; readonly sessionId: SessionId }, +): Promise { + const db = getAppDatabase(database); + const [receipts, assistantMessages] = await Promise.all([ + db + .select({ + eventType: sessionEventsTable.eventType, + semanticHash: sessionEventsTable.semanticHash, + sourceEventId: sessionEventsTable.sourceEventId, + streamId: sessionEventsTable.streamId, + terminalEventJson: sessionEventsTable.terminalEventJson, + }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.runId, input.runId), + inArray(sessionEventsTable.eventType, ["run.cancelled", "run.completed", "run.failed"]), + ), + ) + .all(), + db + .select({ + id: sessionMessagesTable.id, + projectionFormat: sessionMessagesTable.projectionFormat, + }) + .from(sessionMessagesTable) + .where( + and( + eq(sessionMessagesTable.sessionId, input.sessionId), + eq(sessionMessagesTable.sessionRunId, input.runId), + eq(sessionMessagesTable.role, "assistant"), + ), + ) + .all(), + ]); + const [receipt] = receipts; + const finalAssistantMessages = assistantMessages.filter( + (message) => message.id.toString() !== input.runId, + ); + + if (receipts.length === 0) { + return finalAssistantMessages.length === 0; + } + if ( + receipts.length !== 1 || + receipt?.eventType !== "run.completed" || + receipt.sourceEventId !== createSessionRunTerminalSourceId(input.runId, "run.completed") + ) { + return false; + } + if (receipt.semanticHash === null) { + if (receipt.terminalEventJson !== null) { + throw new Error(`Session run ${input.runId} has an invalid legacy terminal authority.`); } - | { - kind: "repair_needed"; - transition: Extract; - }; + return ( + finalAssistantMessages.length === 1 && + finalAssistantMessages[0]?.projectionFormat === "materialized" + ); + } + const authority = await readTerminalEventSemanticAuthority({ + eventJson: receipt.terminalEventJson, + eventType: receipt.eventType, + runId: input.runId, + semanticHash: receipt.semanticHash, + sessionId: input.sessionId, + sourceEventId: receipt.sourceEventId, + streamId: receipt.streamId, + }); + return authority.finalMessageId === null && finalAssistantMessages.length === 0; +} export async function recordCanonicalSessionRunFailure( bindings: ApiBindings, @@ -30,41 +169,164 @@ export async function recordCanonicalSessionRunFailure( error: RunError; runId: SessionRunId; sessionId: SessionId; - source: "api" | "driver"; + source: TerminalRunProjectionSource; }, ): Promise { - const outcome = await setSessionRunStatus(bindings.DB, { + const outcome = await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, error: input.error, runId: input.runId, + sessionId: input.sessionId, source: input.source, status: "failed", }); - if (outcome.kind === "repair_needed") { - return { kind: "repair_needed", transition: outcome }; - } + return outcome.kind === "committed" + ? { kind: "failed" } + : { kind: "not_failed", status: outcome.run.status }; +} - const run = - outcome.kind === "applied" || outcome.kind === "duplicate" - ? outcome.run - : await getSessionRunSummary(bindings.DB, input.runId); +export async function recordCanonicalSessionRunTerminal( + bindings: ApiBindings, + input: { + assistantMessage: PreparedAssistantMessageProjection | null; + deliver?: boolean; + error: RunError | null; + expectedDriverObservation?: ExpectedTerminalDriverObservation; + expectedRunStatus?: SessionRunSummary["status"]; + expectedSessionObservation?: ExpectedTerminalSessionObservation; + expectedSessionOperationId?: RuntimeOperationId | null; + lifecycle?: "IDLE" | "TERMINATED"; + runId: SessionRunId; + sessionId: SessionId; + source: TerminalRunProjectionSource; + sourceEventId?: string; + status: HostTerminalRunStatus; + timestampMs?: number; + }, +): Promise { + const current = await getSessionRunSummary(bindings.DB, input.runId); + if (current === null) { + throw new Error("Session Run was not found while recording its terminal projection."); + } + if ( + current.status !== input.status && + (current.status === "cancelled" || + current.status === "completed" || + current.status === "expired" || + current.status === "failed") + ) { + return { kind: "stale", run: current }; + } + if (input.expectedRunStatus !== undefined && current.status !== input.expectedRunStatus) { + return { kind: "stale", run: current }; + } - if (run?.status !== "failed" || run.error === null) { - return { kind: "not_failed", transition: outcome }; + const error = current.status === input.status ? current.error : input.error; + if (input.status === "failed" && error === null) { + throw new Error("A failed Session Run requires one durable RunError."); + } + const completedReceiptWithoutAssistant = + input.status === "completed" && input.assistantMessage === null + ? await hasCompletedWithoutAssistantAuthority(bindings.DB, { + runId: input.runId, + sessionId: input.sessionId, + }) + : false; + if ( + input.status === "completed" && + (error !== null || (input.assistantMessage === null && !completedReceiptWithoutAssistant)) + ) { + throw new Error("A completed Session Run requires one sealed final assistant reference."); + } + if (input.status !== "completed" && input.assistantMessage !== null) { + throw new Error("Only a completed Session Run can reference a final assistant message."); } - await appendSessionRuntimeEvents({ - bindings, - events: [ - createFailedSessionRunRuntimeEvent({ - run, - runError: run.error, - sessionId: input.sessionId, - sourceEventId: createSessionRunTerminalFailureSourceId(input.runId), - }), - ], + const timestampMs = input.timestampMs ?? currentTimestampMs(); + const timestamp = toIsoString(timestampMs); + const run = terminalRunSummary(current, { error, status: input.status }, timestamp); + const kind = input.status === "expired" ? "run.cancelled" : (`run.${input.status}` as const); + const sourceEventId = input.sourceEventId ?? createSessionRunTerminalSourceId(input.runId, kind); + const lifecycle = input.lifecycle ?? "IDLE"; + const event = await (async () => { + if (input.status === "completed") { + return input.assistantMessage === null + ? createSessionRunUpdatedEvent(run, input.sessionId, lifecycle, sourceEventId) + : createCompletedSessionRunRuntimeEvent({ + finalMessageId: input.assistantMessage.id, + lifecycle, + run, + sessionId: input.sessionId, + sourceEventId, + }); + } + if (input.status !== "failed") { + return createSessionRunUpdatedEvent(run, input.sessionId, lifecycle, sourceEventId); + } + + const execution = + lifecycle === "IDLE" && input.source === "driver" && input.sourceEventId === undefined + ? await readRunExecutionIdentity(bindings.DB, input.runId) + : null; + return execution === null + ? createFailedSessionRunRuntimeEvent({ + lifecycle, + run, + runError: error!, + sessionId: input.sessionId, + sourceEventId, + }) + : createCanonicalDriverRunFailedEvent({ + driverInstanceId: execution.driverInstanceId, + error: error!, + id: createPlatformId(), + occurredAt: timestamp, + runId: input.runId, + runtimeId: execution.runtimeId, + sessionId: input.sessionId, + traceId: run.traceId, + }); + })(); + const committed = await commitTerminalRunProjection(bindings.DB, { + assistantMessage: input.assistantMessage, + error, + ...(input.expectedDriverObservation === undefined + ? {} + : { expectedDriverObservation: input.expectedDriverObservation }), + ...(input.expectedRunStatus === undefined + ? {} + : { expectedRunStatus: input.expectedRunStatus }), + ...(input.expectedSessionObservation === undefined + ? {} + : { expectedSessionObservation: input.expectedSessionObservation }), + ...(input.expectedSessionOperationId === undefined + ? {} + : { expectedSessionOperationId: input.expectedSessionOperationId }), + runId: input.runId, sessionId: input.sessionId, + source: input.source, + targetStatus: input.status, + terminalEvent: { event, occurredAt: timestampMs, sourceEventId }, + timestampMs, }); + const persisted = await getSessionRunSummary(bindings.DB, input.runId); + if (persisted === null) { + throw new Error("Session Run disappeared after its terminal projection."); + } - return { kind: "failed" }; + if (committed.kind === "stale") { + return { kind: "stale", run: persisted }; + } + if (persisted.status !== input.status) { + throw new Error("Atomic terminal projection returned without its exact Session Run status."); + } + if (input.deliver !== false) { + await publishPersistedSessionRuntimeEvents({ + bindings, + events: [event], + sessionId: input.sessionId, + }); + } + return { commitKind: committed.kind, kind: "committed", run: persisted }; } diff --git a/apps/api/src/modules/runtime/application/session-runs/session-run-view-events.service.ts b/apps/api/src/modules/runtime/application/session-runs/session-run-view-events.service.ts index 8c7f5ca1..0fdae9b0 100644 --- a/apps/api/src/modules/runtime/application/session-runs/session-run-view-events.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/session-run-view-events.service.ts @@ -120,60 +120,54 @@ export function createQueuedSessionRunRuntimeEvents(input: { ]; } -export function createCancelledSessionRunRuntimeEvent(input: { +export function createFailedSessionRunRuntimeEvent(input: { eventId?: RuntimeEventId; lifecycle?: Extract; run: SessionRunSummary; - runError?: RunError | null; + runError: RunError; sessionId: SessionId; sourceEventId?: string; }): RuntimeEventEnvelope { - const run: SessionRunView = { - ...toSessionRunView(input.run), - error: input.runError - ? { - ...input.runError, - details: toPrimitiveRecord(input.runError.details), - } - : input.run.error, - status: "cancelled", - }; - return createSessionRuntimeEvent({ ...(input.eventId === undefined ? {} : { id: input.eventId }), - ...(input.sourceEventId === undefined ? {} : { sourceEventId: input.sourceEventId }), - kind: "run.cancelled", + kind: "run.failed", payload: { + error: { + code: input.runError.code, + details: toPrimitiveRecord(input.runError.details), + message: input.runError.message, + retryable: input.runError.retryable, + }, + recoverable: input.runError.retryable, lifecycle: input.lifecycle ?? "IDLE", - run, + run: toSessionRunView(input.run), }, runId: input.run.id, sessionId: input.sessionId, + ...(input.sourceEventId === undefined ? {} : { sourceEventId: input.sourceEventId }), traceId: input.run.traceId, }); } -export function createFailedSessionRunRuntimeEvent(input: { +export function createCompletedSessionRunRuntimeEvent(input: { + eventId?: RuntimeEventId; + finalMessageId: SessionMessageId; + lifecycle?: Extract; run: SessionRunSummary; - runError: RunError; sessionId: SessionId; sourceEventId?: string; }): RuntimeEventEnvelope { return createSessionRuntimeEvent({ - kind: "run.failed", + ...(input.eventId === undefined ? {} : { id: input.eventId }), + ...(input.sourceEventId === undefined ? {} : { sourceEventId: input.sourceEventId }), + kind: "run.completed", payload: { - error: { - code: input.runError.code, - details: toPrimitiveRecord(input.runError.details), - message: input.runError.message, - retryable: input.runError.retryable, - }, - lifecycle: "IDLE", + finalMessageId: input.finalMessageId, + lifecycle: input.lifecycle ?? "IDLE", run: toSessionRunView(input.run), }, runId: input.run.id, sessionId: input.sessionId, - ...(input.sourceEventId === undefined ? {} : { sourceEventId: input.sourceEventId }), traceId: input.run.traceId, }); } @@ -182,12 +176,14 @@ export function createSessionLifecycleTerminatedEvent(input: { eventId?: RuntimeEventId; lastSeen: string; message: string; + occurredAtMs?: number; reason: string; sessionId: SessionId; sourceEventId?: string; }): RuntimeEventEnvelope { return createSessionRuntimeEvent({ ...(input.eventId === undefined ? {} : { id: input.eventId }), + ...(input.occurredAtMs === undefined ? {} : { occurredAtMs: input.occurredAtMs }), ...(input.sourceEventId === undefined ? {} : { sourceEventId: input.sourceEventId }), kind: "session.lifecycle.updated", payload: { diff --git a/apps/api/src/modules/runtime/application/session-runs/stale-run-reconciliation.service.ts b/apps/api/src/modules/runtime/application/session-runs/stale-run-reconciliation.service.ts index 6d60ead7..39f4dd17 100644 --- a/apps/api/src/modules/runtime/application/session-runs/stale-run-reconciliation.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/stale-run-reconciliation.service.ts @@ -1,10 +1,11 @@ import type { RunError } from "@mosoo/contracts/session-run"; -import { driverInstancesTable, sessionRunsTable } from "@mosoo/db"; +import { driverInstancesTable, sessionRunsTable, sessionsTable } from "@mosoo/db"; import type { DriverInstanceId, SessionId, SessionRunId } from "@mosoo/id"; import { and, asc, desc, eq, inArray, isNull, lte, notInArray, or, sql } from "drizzle-orm"; import { alias } from "drizzle-orm/sqlite-core"; import { logWarn } from "../../../../platform/cloudflare/logger"; +import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; import { getAppDatabase } from "../../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../../time"; import { @@ -13,15 +14,17 @@ import { } from "../../domain/runtime-config"; import { classifyReclaim } from "../../domain/session-run-reclaim-recovery"; import { recordRuntimeRunLeaseReleasedOutcome } from "../../infrastructure/runtime-subject-lifecycle/runtime-run-lease-store"; -import { setSessionRunStatus } from "../../infrastructure/session-runs/session-run-store.repository"; +import { recordCanonicalSessionRunTerminal } from "./session-run-terminal-failure.service"; export interface ActiveRunDriverRow { driver_error_message: string | null; + driver_generation: number | null; driver_instance_id: DriverInstanceId | null; driver_last_heartbeat_at: number | null; driver_status: string | null; driver_updated_at: number | null; run_id: SessionRunId; + run_status: "queued" | "booting" | "running" | "waiting_input"; session_id: SessionId; run_trace_id: string | null; run_updated_at: number; @@ -74,11 +77,13 @@ const ACTIVE_SESSION_RUN_STATUSES = ["queued", "booting", "running", "waiting_in function activeRunDriverColumns() { return { driver_error_message: runDriverInstancesTable.errorMessage, + driver_generation: runDriverInstancesTable.generation, driver_instance_id: sessionRunsTable.driverInstanceId, driver_last_heartbeat_at: runDriverInstancesTable.lastHeartbeatAt, driver_status: runDriverInstancesTable.status, driver_updated_at: runDriverInstancesTable.updatedAt, run_id: sessionRunsTable.id, + run_status: sql`${sessionRunsTable.status}`, run_trace_id: sessionRunsTable.traceId, run_updated_at: sessionRunsTable.updatedAt, session_id: sessionRunsTable.sessionId, @@ -118,6 +123,7 @@ async function findStaleActiveRun( (await getAppDatabase(database) .select(activeRunDriverColumns()) .from(sessionRunsTable) + .innerJoin(sessionsTable, eq(sessionsTable.id, sessionRunsTable.sessionId)) .leftJoin( runDriverInstancesTable, eq(runDriverInstancesTable.id, sessionRunsTable.driverInstanceId), @@ -125,6 +131,7 @@ async function findStaleActiveRun( .where( and( eq(sessionRunsTable.sessionId, sessionId), + isNull(sessionsTable.archivedAt), inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), ), ) @@ -152,6 +159,7 @@ async function findStaleActiveRuns( return getAppDatabase(database) .select(activeRunDriverColumns()) .from(sessionRunsTable) + .innerJoin(sessionsTable, eq(sessionsTable.id, sessionRunsTable.sessionId)) .leftJoin( runDriverInstancesTable, eq(runDriverInstancesTable.id, sessionRunsTable.driverInstanceId), @@ -159,6 +167,7 @@ async function findStaleActiveRuns( .where( and( inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + isNull(sessionsTable.archivedAt), staleActiveRunPredicate(input.nowMs), ), ) @@ -169,29 +178,29 @@ async function findStaleActiveRuns( async function failStaleActiveRun(database: D1Database, staleRun: ActiveRunDriverRow) { const error = staleRunError(staleRun); - const outcome = await setSessionRunStatus(database, { + const outcome = await recordCanonicalSessionRunTerminal({ DB: database } as ApiBindings, { + assistantMessage: null, + deliver: false, error, + expectedDriverObservation: { + driverInstanceId: staleRun.driver_instance_id, + lastHeartbeatAt: staleRun.driver_last_heartbeat_at, + status: staleRun.driver_status, + updatedAt: staleRun.driver_updated_at, + }, + expectedRunStatus: staleRun.run_status, runId: staleRun.run_id, + sessionId: staleRun.session_id, source: "maintenance", status: "failed", }); - switch (outcome.kind) { - case "applied": - case "duplicate": { - await releaseStaleRunLease(database, staleRun); - return true; - } - case "repair_needed": { - throw new Error( - "Stale session run reconciliation left the session lifecycle projection stale.", - ); - } - case "rejected": - case "stale": { - return false; - } + if (outcome.kind === "stale") { + return false; } + + await releaseStaleRunLease(database, staleRun); + return true; } // Failing the run ends the lease, but only the release write re-arms the @@ -201,12 +210,13 @@ async function releaseStaleRunLease( database: D1Database, staleRun: ActiveRunDriverRow, ): Promise { - if (staleRun.driver_instance_id === null) { + if (staleRun.driver_instance_id === null || staleRun.driver_generation === null) { return; } const outcome = await recordRuntimeRunLeaseReleasedOutcome(database, { driverInstanceId: staleRun.driver_instance_id, + expectedDriverGeneration: staleRun.driver_generation, expectedSessionRunId: staleRun.run_id, }); diff --git a/apps/api/src/modules/runtime/application/session-runs/terminal-run-reconciliation.service.ts b/apps/api/src/modules/runtime/application/session-runs/terminal-run-reconciliation.service.ts index 264d4311..a689eb5b 100644 --- a/apps/api/src/modules/runtime/application/session-runs/terminal-run-reconciliation.service.ts +++ b/apps/api/src/modules/runtime/application/session-runs/terminal-run-reconciliation.service.ts @@ -1,103 +1,106 @@ import type { SessionStatus } from "@mosoo/contracts/session"; -import type { SessionRunStatus, SessionRunSummary } from "@mosoo/contracts/session-run"; +import type { SessionRunStatus } from "@mosoo/contracts/session-run"; import { driverInstancesTable, sessionEventsTable, + sessionMessagesTable, sessionRunsTable, sessionsTable, } from "@mosoo/db"; -import type { SessionId, SessionRunId } from "@mosoo/id"; -import { and, asc, eq, inArray, isNull, notExists, or, sql } from "drizzle-orm"; +import type { PlatformId, RuntimeOperationId, SessionId, SessionRunId } from "@mosoo/id"; +import { and, asc, eq, gt, inArray, isNull, lte, notExists, or, sql } from "drizzle-orm"; +import { createErrorLogContext, logWarn } from "../../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; -import { getAppDatabase } from "../../../../platform/db/drizzle"; -import { appendSessionRuntimeEvents } from "../../../sessions/application/session-event-write.service"; +import { getAppDatabase, getD1ChangeCount } from "../../../../platform/db/drizzle"; +import { currentTimestampMs } from "../../../../time"; +import { readTerminalEventSemanticAuthority } from "../../../sessions/domain/session-terminal-event-authority"; +import { isSealedPublicSessionMessageStream } from "../../../sessions/infrastructure/session-message-event-stream.repository"; import { createSessionRunTerminalSourceId } from "../../domain/session-run-terminal-event-id"; -import { - getSessionRunSummariesByIds, - setSessionRunStatus, -} from "../../infrastructure/session-runs/session-run-store.repository"; -import type { SessionRunTransitionOutcome } from "../../infrastructure/session-runs/session-run-store.repository"; -import { - createFailedSessionRunRuntimeEvent, - createSessionRunUpdatedEvent, -} from "./session-run-view-events.service"; +import { prepareAssistantMessageProjection } from "../../infrastructure/driver-instance/assistant-message-projection"; +import type { PreparedAssistantMessageProjection } from "../../infrastructure/driver-instance/assistant-message-projection"; +import type { + HostTerminalRunStatus, + TerminalRunProjectionSource, +} from "../../infrastructure/driver-instance/completed-run-commit.repository"; +import { adoptTerminalRunProjection } from "../../infrastructure/driver-instance/completed-run-commit.repository"; +import { getSessionRunSummary } from "../../infrastructure/session-runs/session-run-store.repository"; +import { recordCanonicalSessionRunTerminal } from "./session-run-terminal-failure.service"; const TERMINAL_RUN_STATUSES = ["cancelled", "completed", "expired", "failed"] as const; const TERMINAL_DRIVER_STATUSES = ["failed", "stopped"] as const; +const TERMINAL_RECONCILIATION_RETRY_AFTER_MS = 10 * 60_000; +const TERMINAL_SOURCES = new Set([ + "api", + "driver", + "maintenance", + "runtime_operation", + "system", + "viewer", +]); interface TerminalRunCandidate { readonly runId: SessionRunId; + readonly runStatus: SessionRunStatus; + readonly runStatusOperationId: RuntimeOperationId | null; + readonly runStatusSeq: number; + readonly runStatusSource: string; + readonly runTerminalReconciliationAttemptedAt: number | null; + readonly runUpdatedAt: number; readonly sessionId: SessionId; readonly sessionLastRunId: SessionRunId | null; readonly sessionStatus: SessionStatus; } +interface TerminalEventReceipt { + readonly eventType: string; + readonly runId: SessionRunId | null; + readonly semanticHash: string | null; + readonly sourceEventId: string; + readonly streamId: string | null; + readonly terminalEventJson: string | null; +} + +interface CompletedAssistantRow { + readonly contentText: string; + readonly createdByAccountId: PlatformId; + readonly id: string; + readonly planJson: string | null; + readonly projectionFormat: "event_stream_v3" | "materialized"; + readonly segmentsJson: string | null; +} + export interface TerminalRunReconciliationResult { + readonly failures: readonly { + readonly message: string; + readonly runId: SessionRunId; + readonly sessionId: SessionId; + }[]; readonly reconciledRunIds: readonly SessionRunId[]; readonly reconciledSessionIds: readonly SessionId[]; } -function assertTerminalRunProjection(outcome: SessionRunTransitionOutcome): void { - switch (outcome.kind) { - case "applied": - case "duplicate": { - return; - } - case "repair_needed": { - throw new Error("Terminal run reconciliation left the session lifecycle projection stale."); - } - case "rejected": - case "stale": { - throw new Error("Terminal run reconciliation lost a concurrent run transition."); - } - } -} - function terminalEventKind( status: SessionRunStatus, ): "run.cancelled" | "run.completed" | "run.failed" { switch (status) { - case "completed": { + case "completed": return "run.completed"; - } - case "failed": { + case "failed": return "run.failed"; - } case "cancelled": - case "expired": { + case "expired": return "run.cancelled"; - } case "queued": case "booting": case "running": - case "waiting_input": { + case "waiting_input": throw new Error(`Expected terminal Session Run status, received ${status}.`); - } } } -function createTerminalRunRecoveryEvent(input: { - readonly kind: "run.cancelled" | "run.completed" | "run.failed"; - readonly run: SessionRunSummary; - readonly sessionId: SessionId; - readonly sourceEventId: string; -}) { - if (input.kind !== "run.failed") { - return createSessionRunUpdatedEvent(input.run, input.sessionId, "IDLE", input.sourceEventId); - } - - return createFailedSessionRunRuntimeEvent({ - run: input.run, - runError: input.run.error ?? { - code: "runtime.terminal_error_missing", - details: {}, - message: "The run failed without a persisted error.", - retryable: false, - }, - sessionId: input.sessionId, - sourceEventId: input.sourceEventId, - }); +function isHostTerminalRunStatus(status: SessionRunStatus): status is HostTerminalRunStatus { + return TERMINAL_RUN_STATUSES.some((terminalStatus) => terminalStatus === status); } async function findTerminalRunCandidates( @@ -128,10 +131,17 @@ async function findTerminalRunCandidates( eq(sessionsTable.lastRunId, sessionRunsTable.id), eq(sessionsTable.status, "RUNNING"), ); + const retryBefore = currentTimestampMs() - TERMINAL_RECONCILIATION_RETRY_AFTER_MS; return database .select({ runId: sessionRunsTable.id, + runStatus: sessionRunsTable.status, + runStatusOperationId: sessionRunsTable.statusOperationId, + runStatusSeq: sessionRunsTable.statusSeq, + runStatusSource: sessionRunsTable.statusSource, + runTerminalReconciliationAttemptedAt: sessionRunsTable.terminalReconciliationAttemptedAt, + runUpdatedAt: sessionRunsTable.updatedAt, sessionId: sessionRunsTable.sessionId, sessionLastRunId: sessionsTable.lastRunId, sessionStatus: sessionsTable.status, @@ -142,6 +152,12 @@ async function findTerminalRunCandidates( .where( and( inArray(sessionRunsTable.status, TERMINAL_RUN_STATUSES), + or( + isNull(sessionRunsTable.terminalReconciliationAttemptedAt), + lte(sessionRunsTable.terminalReconciliationAttemptedAt, retryBefore), + gt(sessionRunsTable.updatedAt, sessionRunsTable.terminalReconciliationAttemptedAt), + gt(sessionsTable.updatedAt, sessionRunsTable.terminalReconciliationAttemptedAt), + ), isNull(sessionsTable.archivedAt), inArray(sessionsTable.status, ["IDLE", "RESCHEDULING", "RUNNING"]), or( @@ -152,103 +168,346 @@ async function findTerminalRunCandidates( or(staleSessionProjection, missingTerminalEvent), ), ) - .orderBy(asc(sessionRunsTable.updatedAt), asc(sessionRunsTable.id)) + .orderBy( + asc( + sql`coalesce(${sessionRunsTable.terminalReconciliationAttemptedAt}, ${sessionRunsTable.updatedAt})`, + ), + asc(sessionRunsTable.id), + ) .limit(limit) .all(); } -async function readPersistedTerminalEventKeys( +async function recordTerminalReconciliationAttempt( + bindings: ApiBindings, + candidate: TerminalRunCandidate, + attemptedAt: number, +): Promise { + const result = await getAppDatabase(bindings.DB) + .update(sessionRunsTable) + .set({ terminalReconciliationAttemptedAt: attemptedAt }) + .where( + and( + eq(sessionRunsTable.id, candidate.runId), + eq(sessionRunsTable.status, candidate.runStatus), + eq(sessionRunsTable.statusSeq, candidate.runStatusSeq), + eq(sessionRunsTable.statusSource, candidate.runStatusSource), + eq(sessionRunsTable.updatedAt, candidate.runUpdatedAt), + candidate.runStatusOperationId === null + ? isNull(sessionRunsTable.statusOperationId) + : eq(sessionRunsTable.statusOperationId, candidate.runStatusOperationId), + candidate.runTerminalReconciliationAttemptedAt === null + ? isNull(sessionRunsTable.terminalReconciliationAttemptedAt) + : eq( + sessionRunsTable.terminalReconciliationAttemptedAt, + candidate.runTerminalReconciliationAttemptedAt, + ), + ), + ) + .run(); + return getD1ChangeCount(result) > 0; +} + +async function readTerminalEventReceipts( bindings: ApiBindings, runIds: readonly SessionRunId[], -): Promise> { +): Promise> { if (runIds.length === 0) { - return new Set(); + return new Map(); } - const rows = await getAppDatabase(bindings.DB) .select({ eventType: sessionEventsTable.eventType, runId: sessionEventsTable.runId, + semanticHash: sessionEventsTable.semanticHash, + sourceEventId: sessionEventsTable.sourceEventId, + streamId: sessionEventsTable.streamId, + terminalEventJson: sessionEventsTable.terminalEventJson, }) .from(sessionEventsTable) .where( and( - inArray(sessionEventsTable.runId, [...runIds]), + inArray(sessionEventsTable.runId, runIds), inArray(sessionEventsTable.eventType, ["run.cancelled", "run.completed", "run.failed"]), ), ) .all(); - return new Set( - rows.flatMap((row) => (row.runId === null ? [] : [`${row.runId}:${row.eventType}`])), + return new Map( + [...Map.groupBy(rows, (row) => row.runId)].flatMap(([runId, receipts]) => + runId === null ? [] : [[runId, receipts]], + ), ); } -/** - * Repairs terminal Run projections after the Driver can no longer replay its - * final event. The terminal Run row itself is the durable, idempotent repair - * obligation: a missing matching terminal session_event is reconstructed with - * a stable source id, while a duplicate status transition repairs the owning - * Session lifecycle projection. - */ +export async function assertCanonicalTerminalSessionRunProjection( + bindings: ApiBindings, + input: { + readonly runId: SessionRunId; + readonly sessionId: SessionId; + readonly status: HostTerminalRunStatus; + }, +): Promise { + const run = await getSessionRunSummary(bindings.DB, input.runId); + if (run?.status !== input.status) { + throw new Error(`Terminal run ${input.runId} has no unique canonical terminal receipt.`); + } + const outcome = await adoptTerminalRunProjection(bindings.DB, { + runId: input.runId, + sessionId: input.sessionId, + }); + if (outcome.kind === "missing" || outcome.kind === "stale") { + throw new Error(`Terminal run ${input.runId} has no unique canonical terminal receipt.`); + } +} + +async function resolveCompletedAssistant( + bindings: ApiBindings, + input: { + readonly legacy: boolean; + readonly runId: SessionRunId; + readonly sessionId: SessionId; + readonly streamId?: string | null; + }, +): Promise { + const rows: CompletedAssistantRow[] = await getAppDatabase(bindings.DB) + .select({ + contentText: sessionMessagesTable.contentText, + createdByAccountId: sessionMessagesTable.createdByAccountId, + id: sessionMessagesTable.id, + planJson: sessionMessagesTable.planJson, + projectionFormat: sessionMessagesTable.projectionFormat, + segmentsJson: sessionMessagesTable.segmentsJson, + }) + .from(sessionMessagesTable) + .where( + and( + eq(sessionMessagesTable.sessionId, input.sessionId), + eq(sessionMessagesTable.sessionRunId, input.runId), + eq(sessionMessagesTable.role, "assistant"), + ), + ) + .all(); + + const isEventStreamReference = (row: CompletedAssistantRow): boolean => + row.contentText === "" && + row.planJson === null && + row.projectionFormat === "event_stream_v3" && + row.segmentsJson === null; + const withoutToolCarrier = (): CompletedAssistantRow[] => { + const carrier = rows.find((row) => row.id === input.runId); + if (carrier !== undefined && !isEventStreamReference(carrier)) { + throw new Error(`Completed run ${input.runId} has an invalid parentless tool carrier.`); + } + return rows.filter((row) => row.id !== input.runId); + }; + + if (input.legacy) { + const authorityRows = withoutToolCarrier(); + if (authorityRows.length !== 1 || authorityRows[0]?.projectionFormat !== "materialized") { + throw new Error(`Legacy completed run ${input.runId} has ambiguous assistant authority.`); + } + return null; + } + + if (input.streamId === null) { + if (withoutToolCarrier().length !== 0) { + throw new Error( + `Completed run ${input.runId} has assistant rows despite declaring no final stream.`, + ); + } + return null; + } + + const authorityRows = + input.streamId === undefined + ? withoutToolCarrier() + : rows.filter((row) => row.id === input.streamId); + const otherRows = + input.streamId === undefined ? [] : rows.filter((row) => row.id !== input.streamId); + if (input.streamId !== undefined) { + const carrier = otherRows.find((row) => row.id === input.runId); + if ( + otherRows.length > 1 || + (carrier === undefined && otherRows.length !== 0) || + (carrier !== undefined && !isEventStreamReference(carrier)) + ) { + throw new Error( + `Completed run ${input.runId} has no unique event-stream assistant reference.`, + ); + } + } + + if (authorityRows.length === 0 && input.streamId === undefined) { + return null; + } + + const [row] = authorityRows; + if (authorityRows.length !== 1 || row === undefined || !isEventStreamReference(row)) { + throw new Error(`Completed run ${input.runId} has no unique event-stream assistant reference.`); + } + const sealed = await isSealedPublicSessionMessageStream(bindings.DB, { + processType: "agent.message.delta", + runId: input.runId, + sessionId: input.sessionId, + streamId: row.id, + }); + if (!sealed) { + throw new Error(`Completed run ${input.runId} has no sealed authoritative assistant stream.`); + } + + return prepareAssistantMessageProjection({ + createdByAccountId: row.createdByAccountId, + messageId: row.id, + sessionId: input.sessionId, + sessionRunId: input.runId, + }); +} + export async function reconcileTerminalSessionRuns( bindings: ApiBindings, input: { readonly limit: number; }, ): Promise { + const database = getAppDatabase(bindings.DB); const candidates = await findTerminalRunCandidates(bindings, input.limit); const runIds = candidates.map((candidate) => candidate.runId); - const [runsById, persistedTerminalEventKeys] = await Promise.all([ - getSessionRunSummariesByIds(bindings.DB, runIds), - readPersistedTerminalEventKeys(bindings, runIds), - ]); + const receiptsByRunId = await readTerminalEventReceipts(bindings, runIds); + const failures: { + message: string; + runId: SessionRunId; + sessionId: SessionId; + }[] = []; const reconciledRunIds: SessionRunId[] = []; const reconciledSessionIds = new Set(); for (const candidate of candidates) { - const run = runsById.get(candidate.runId); - - if (run === undefined) { - continue; - } - - if (candidate.sessionLastRunId === run.id && candidate.sessionStatus === "RUNNING") { - const projection = await setSessionRunStatus(bindings.DB, { - runId: run.id, - source: "maintenance", - status: run.status, - }); - assertTerminalRunProjection(projection); - } - - const kind = terminalEventKind(run.status); - const eventKey = `${run.id}:${kind}`; - - if (!persistedTerminalEventKeys.has(eventKey)) { - const sourceEventId = createSessionRunTerminalSourceId(run.id, kind); - const persisted = await appendSessionRuntimeEvents({ - bindings, - events: [ - createTerminalRunRecoveryEvent({ - kind, - run, + try { + if (!(await recordTerminalReconciliationAttempt(bindings, candidate, currentTimestampMs()))) { + continue; + } + const run = await getSessionRunSummary(bindings.DB, candidate.runId); + if (run === null) { + continue; + } + if (!isHostTerminalRunStatus(run.status)) { + continue; + } + const terminalStatus = run.status; + const source = TERMINAL_SOURCES.has(candidate.runStatusSource as TerminalRunProjectionSource) + ? (candidate.runStatusSource as TerminalRunProjectionSource) + : null; + if (source === null) { + throw new Error(`Terminal run ${run.id} has an unknown durable source.`); + } + if (source === "runtime_operation" && candidate.runStatusOperationId === null) { + throw new Error(`Runtime-operation terminal run ${run.id} has no operation identity.`); + } + const receipts = receiptsByRunId.get(run.id) ?? []; + const expectedKind = terminalEventKind(terminalStatus); + const matchingReceipts = receipts.filter((receipt) => receipt.eventType === expectedKind); + const [matchingReceipt] = matchingReceipts; + const legacy = matchingReceipts.length === 1 && matchingReceipt?.semanticHash === null; + if (receipts.length > 1 || matchingReceipts.length > 1) { + throw new Error(`Terminal run ${run.id} has ambiguous durable terminal receipts.`); + } + if ( + matchingReceipt !== undefined && + matchingReceipt.sourceEventId !== createSessionRunTerminalSourceId(run.id, expectedKind) + ) { + throw new Error(`Terminal run ${run.id} has a noncanonical terminal receipt identity.`); + } + if (legacy && matchingReceipt?.terminalEventJson !== null) { + throw new Error(`Terminal run ${run.id} has invalid legacy semantic authority.`); + } + const semanticAuthority = + matchingReceipt === undefined || matchingReceipt.semanticHash === null + ? null + : await readTerminalEventSemanticAuthority({ + eventJson: matchingReceipt.terminalEventJson, + eventType: matchingReceipt.eventType, + runId: run.id, + semanticHash: matchingReceipt.semanticHash, + sessionId: candidate.sessionId, + sourceEventId: matchingReceipt.sourceEventId, + streamId: matchingReceipt.streamId, + }); + const assistantMessage = + terminalStatus === "completed" && matchingReceipt === undefined + ? await resolveCompletedAssistant(bindings, { + legacy, + runId: run.id, + sessionId: candidate.sessionId, + }) + : null; + const commitKind = await (async () => { + if (matchingReceipt !== undefined) { + const outcome = await adoptTerminalRunProjection(bindings.DB, { + runId: run.id, sessionId: candidate.sessionId, - sourceEventId, - }), - ], - sessionId: candidate.sessionId, - }); + }); + if (outcome.kind === "missing") { + throw new Error(`Terminal run ${run.id} lost its durable terminal receipt.`); + } + return outcome.kind === "stale" ? null : outcome.kind; + } - if (persisted.persistedCount > 0) { + const outcome = await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage, + deliver: false, + error: run.error, + ...(source === "runtime_operation" + ? { expectedSessionOperationId: candidate.runStatusOperationId } + : {}), + runId: run.id, + sessionId: candidate.sessionId, + source, + status: terminalStatus, + }); + return outcome.kind === "stale" ? null : outcome.commitKind; + })(); + if (commitKind === null) { + continue; + } + if (commitKind === "applied") { reconciledRunIds.push(run.id); } + const session = await database + .select({ + lastRunId: sessionsTable.lastRunId, + status: sessionsTable.status, + }) + .from(sessionsTable) + .where(eq(sessionsTable.id, candidate.sessionId)) + .limit(1) + .get(); + const expectedSessionStatus = semanticAuthority?.lifecycle ?? "IDLE"; + if (session?.lastRunId !== run.id) { + continue; + } + if (session.status !== expectedSessionStatus) { + throw new Error( + `Terminal run ${run.id} did not converge its current Session to ${expectedSessionStatus}.`, + ); + } + reconciledSessionIds.add(candidate.sessionId); + } catch (error) { + failures.push({ + message: error instanceof Error ? error.message : "Unknown terminal reconciliation error.", + runId: candidate.runId, + sessionId: candidate.sessionId, + }); + logWarn("runtime.terminal_run.reconciliation_failed", { + ...createErrorLogContext(error), + runId: candidate.runId, + sessionId: candidate.sessionId, + }); } - - reconciledSessionIds.add(candidate.sessionId); } return { + failures, reconciledRunIds, reconciledSessionIds: [...reconciledSessionIds], }; diff --git a/apps/api/src/modules/runtime/domain/runtime-kind-policy.ts b/apps/api/src/modules/runtime/domain/runtime-kind-policy.ts index dd6c5684..dfeb2ddf 100644 --- a/apps/api/src/modules/runtime/domain/runtime-kind-policy.ts +++ b/apps/api/src/modules/runtime/domain/runtime-kind-policy.ts @@ -16,7 +16,6 @@ export type RuntimeCheckpointRule = readonly updateSubjectCheckpoint: true; } | { - readonly sanitizeTransientState: boolean; readonly type: "session_workspaces"; readonly updateSubjectCheckpoint: false; }; @@ -89,16 +88,10 @@ const SUBJECT_MEMORY_CHECKPOINT = { } as const satisfies RuntimeCheckpointRule; const SESSION_WORKSPACES_CHECKPOINT = { - sanitizeTransientState: false, type: "session_workspaces", updateSubjectCheckpoint: false, } as const satisfies RuntimeCheckpointRule; -const CATTLE_SESSION_WORKSPACE_CHECKPOINT = { - ...SESSION_WORKSPACES_CHECKPOINT, - sanitizeTransientState: true, -} as const satisfies RuntimeCheckpointRule; - const SUBJECT_MEMORY_CLEAR = { path: SANDBOX_MEMORY_PATH, type: "subject_memory", @@ -115,8 +108,8 @@ export const RUNTIME_KIND_POLICIES = { createOnHibernate: [], createOnRecreate: [], createOnReset: [], - createOnTerminal: [CATTLE_SESSION_WORKSPACE_CHECKPOINT], - restoreOnActivate: [CATTLE_SESSION_WORKSPACE_CHECKPOINT], + createOnTerminal: [SESSION_WORKSPACES_CHECKPOINT], + restoreOnActivate: [SESSION_WORKSPACES_CHECKPOINT], }, // Cattle commits the complete session workspace, including provider-native // state, before terminal lease release. Platform history remains the diff --git a/apps/api/src/modules/runtime/domain/runtime-subject-lifecycle.machine.ts b/apps/api/src/modules/runtime/domain/runtime-subject-lifecycle.machine.ts index c1fbb050..b260e29e 100644 --- a/apps/api/src/modules/runtime/domain/runtime-subject-lifecycle.machine.ts +++ b/apps/api/src/modules/runtime/domain/runtime-subject-lifecycle.machine.ts @@ -1,5 +1,4 @@ import type { SandboxStatus } from "@mosoo/contracts/sandbox"; -import { createMachine, transition } from "xstate"; export const RUNTIME_SUBJECT_CLAIMABLE_STATUSES = [ "active", @@ -11,7 +10,14 @@ export const RUNTIME_SUBJECT_OPERATION_STATUSES = [ "destroying", ] as const satisfies readonly SandboxStatus[]; +export const RUNTIME_SUBJECT_RECOVERABLE_OPERATION_STATUSES = [ + "restoring", + ...RUNTIME_SUBJECT_OPERATION_STATUSES, +] as const satisfies readonly SandboxStatus[]; + export type RuntimeSubjectOperationStatus = (typeof RUNTIME_SUBJECT_OPERATION_STATUSES)[number]; +export type RuntimeSubjectRecoverableOperationStatus = + (typeof RUNTIME_SUBJECT_RECOVERABLE_OPERATION_STATUSES)[number]; // There is no dedicated failure state. A failed lifecycle step makes the // container untrustworthy, so activation first enters `destroying`. Successful @@ -24,14 +30,6 @@ export type RuntimeSubjectLifecycleEvent = | { type: "runtime_subject.cold" } | { type: "runtime_subject.destroy" }; -const RUNTIME_SUBJECT_STATUS_BY_EVENT = { - "runtime_subject.activate": "restoring", - "runtime_subject.active": "active", - "runtime_subject.back_up": "backing_up", - "runtime_subject.cold": "cold", - "runtime_subject.destroy": "destroying", -} as const satisfies Record; - const RUNTIME_SUBJECT_EVENT_BY_STATUS = { active: { type: "runtime_subject.active" }, backing_up: { type: "runtime_subject.back_up" }, @@ -40,49 +38,13 @@ const RUNTIME_SUBJECT_EVENT_BY_STATUS = { restoring: { type: "runtime_subject.activate" }, } as const satisfies Record; -const runtimeSubjectLifecycleMachine = createMachine({ - id: "runtimeSubjectLifecycle", - initial: "cold", - states: { - active: { - on: { - "runtime_subject.active": "active", - "runtime_subject.back_up": "backing_up", - "runtime_subject.cold": "cold", - "runtime_subject.destroy": "destroying", - }, - }, - backing_up: { - on: { - "runtime_subject.active": "active", - "runtime_subject.cold": "cold", - "runtime_subject.destroy": "destroying", - }, - }, - cold: { - on: { - "runtime_subject.activate": "restoring", - "runtime_subject.back_up": "backing_up", - "runtime_subject.destroy": "destroying", - }, - }, - destroying: { - on: { - "runtime_subject.cold": "cold", - }, - }, - restoring: { - on: { - "runtime_subject.active": "active", - "runtime_subject.cold": "cold", - "runtime_subject.destroy": "destroying", - }, - }, - }, - types: {} as { - events: RuntimeSubjectLifecycleEvent; - }, -}); +const RUNTIME_SUBJECT_TRANSITIONS: Record = { + active: ["backing_up", "cold", "destroying"], + backing_up: ["active", "cold", "destroying"], + cold: ["restoring", "backing_up", "destroying"], + destroying: ["cold"], + restoring: ["active", "cold", "destroying"], +}; export type RuntimeSubjectTransitionDecision = | { @@ -126,11 +88,7 @@ export function decideRuntimeSubjectTransition(input: { }; } - const snapshot = runtimeSubjectLifecycleMachine.resolveState({ value: input.currentStatus }); - const [nextSnapshot] = transition(runtimeSubjectLifecycleMachine, snapshot, event); - const nextStatus = readRuntimeSubjectSnapshotValue(nextSnapshot.value); - - if (nextStatus === input.currentStatus) { + if (!RUNTIME_SUBJECT_TRANSITIONS[input.currentStatus].includes(input.targetStatus)) { return { currentStatus: input.currentStatus, event, @@ -143,17 +101,7 @@ export function decideRuntimeSubjectTransition(input: { return { event, kind: "accepted", - nextStatus, + nextStatus: input.targetStatus, previousStatus: input.currentStatus, }; } - -function readRuntimeSubjectSnapshotValue(value: unknown): SandboxStatus { - if (typeof value !== "string" || !(value in RUNTIME_SUBJECT_EVENT_BY_STATUS)) { - throw new Error("Runtime subject lifecycle machine returned an unknown state."); - } - - return RUNTIME_SUBJECT_STATUS_BY_EVENT[ - toRuntimeSubjectLifecycleEvent(value as SandboxStatus).type - ]; -} diff --git a/apps/api/src/modules/runtime/domain/sandbox-network-constraints.ts b/apps/api/src/modules/runtime/domain/sandbox-network-constraints.ts index 9445daf1..24f3b19b 100644 --- a/apps/api/src/modules/runtime/domain/sandbox-network-constraints.ts +++ b/apps/api/src/modules/runtime/domain/sandbox-network-constraints.ts @@ -15,6 +15,18 @@ export interface SandboxNetworkConstraints { readonly networkPolicy: EnvironmentNetworkPolicy; } +export async function hashSandboxNetworkConstraints(input: unknown): Promise { + const constraints = parseSandboxNetworkConstraints(input); + const digest = await crypto.subtle.digest( + "SHA-256", + new TextEncoder().encode( + JSON.stringify([constraints.networkPolicy, ...constraints.allowedHosts]), + ), + ); + + return new Uint8Array(digest).toHex(); +} + export function normalizeSandboxNetworkHost(host: string): string { const normalized = host.trim().toLowerCase().replace(/\.$/u, ""); diff --git a/apps/api/src/modules/runtime/domain/session-run-lifecycle.machine.ts b/apps/api/src/modules/runtime/domain/session-run-lifecycle.machine.ts index 26a12636..5273df95 100644 --- a/apps/api/src/modules/runtime/domain/session-run-lifecycle.machine.ts +++ b/apps/api/src/modules/runtime/domain/session-run-lifecycle.machine.ts @@ -1,5 +1,4 @@ import type { SessionRunStatus } from "@mosoo/contracts/session-run"; -import { createMachine, transition } from "xstate"; const TERMINAL_SESSION_RUN_STATUSES = [ "cancelled", @@ -27,17 +26,6 @@ export type SessionRunLifecycleEvent = | { type: "run.start" } | { type: "run.wait_for_input" }; -const SESSION_RUN_STATUS_BY_EVENT = { - "run.boot": "booting", - "run.cancel": "cancelled", - "run.complete": "completed", - "run.expire": "expired", - "run.fail": "failed", - "run.queue": "queued", - "run.start": "running", - "run.wait_for_input": "waiting_input", -} as const satisfies Record; - const SESSION_RUN_EVENT_BY_STATUS = { booting: { type: "run.boot" }, cancelled: { type: "run.cancel" }, @@ -49,56 +37,16 @@ const SESSION_RUN_EVENT_BY_STATUS = { waiting_input: { type: "run.wait_for_input" }, } as const satisfies Record; -const sessionRunLifecycleMachine = createMachine({ - id: "sessionRunLifecycle", - initial: "queued", - states: { - booting: { - on: { - "run.cancel": "cancelled", - "run.complete": "completed", - "run.expire": "expired", - "run.fail": "failed", - "run.start": "running", - "run.wait_for_input": "waiting_input", - }, - }, - cancelled: {}, - completed: {}, - expired: {}, - failed: {}, - queued: { - on: { - "run.boot": "booting", - "run.cancel": "cancelled", - "run.expire": "expired", - "run.fail": "failed", - "run.start": "running", - }, - }, - running: { - on: { - "run.cancel": "cancelled", - "run.complete": "completed", - "run.expire": "expired", - "run.fail": "failed", - "run.wait_for_input": "waiting_input", - }, - }, - waiting_input: { - on: { - "run.cancel": "cancelled", - "run.complete": "completed", - "run.expire": "expired", - "run.fail": "failed", - "run.start": "running", - }, - }, - }, - types: {} as { - events: SessionRunLifecycleEvent; - }, -}); +const previousStatusesByTarget: Readonly> = { + booting: ["queued"], + cancelled: ["booting", "queued", "running", "waiting_input"], + completed: ["booting", "running", "waiting_input"], + expired: ["booting", "queued", "running", "waiting_input"], + failed: ["booting", "queued", "running", "waiting_input"], + queued: [], + running: ["booting", "queued", "waiting_input"], + waiting_input: ["booting", "running"], +}; export type SessionRunTransitionDecision = | { @@ -167,11 +115,7 @@ export function decideSessionRunTransition(input: { }; } - const snapshot = sessionRunLifecycleMachine.resolveState({ value: input.currentStatus }); - const [nextSnapshot] = transition(sessionRunLifecycleMachine, snapshot, event); - const nextStatus = readSessionRunSnapshotValue(nextSnapshot.value); - - if (nextStatus === input.currentStatus) { + if (!previousStatusesByTarget[input.targetStatus].includes(input.currentStatus)) { return { currentStatus: input.currentStatus, event, @@ -184,15 +128,7 @@ export function decideSessionRunTransition(input: { return { event, kind: "accepted", - nextStatus, + nextStatus: input.targetStatus, previousStatus: input.currentStatus, }; } - -function readSessionRunSnapshotValue(value: unknown): SessionRunStatus { - if (typeof value !== "string" || !(value in SESSION_RUN_EVENT_BY_STATUS)) { - throw new Error("Session run lifecycle machine returned an unknown state."); - } - - return SESSION_RUN_STATUS_BY_EVENT[toSessionRunLifecycleEvent(value as SessionRunStatus).type]; -} diff --git a/apps/api/src/modules/runtime/domain/session-run-terminal-event-id.ts b/apps/api/src/modules/runtime/domain/session-run-terminal-event-id.ts index 2faef709..a0dbfdd7 100644 --- a/apps/api/src/modules/runtime/domain/session-run-terminal-event-id.ts +++ b/apps/api/src/modules/runtime/domain/session-run-terminal-event-id.ts @@ -1,17 +1,7 @@ import type { SessionRunId } from "@mosoo/id"; -import type { RuntimeEventKind } from "@mosoo/runtime-events"; +import { createSessionRunTerminalSourceId } from "@mosoo/runtime-events"; -type TerminalSessionRunEventKind = Extract< - RuntimeEventKind, - "run.cancelled" | "run.completed" | "run.failed" ->; - -export function createSessionRunTerminalSourceId( - runId: SessionRunId, - kind: TerminalSessionRunEventKind, -): string { - return `session-run-terminal:${runId}:${kind}`; -} +export { createSessionRunTerminalSourceId }; export function createSessionRunTerminalFailureSourceId(runId: SessionRunId): string { return createSessionRunTerminalSourceId(runId, "run.failed"); diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/assistant-message-projection.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/assistant-message-projection.ts index 6e49a1a5..88550b18 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/assistant-message-projection.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/assistant-message-projection.ts @@ -1,161 +1,48 @@ -import type { SessionMessageSegment } from "@mosoo/contracts/session"; -import { sessionMessagesTable } from "@mosoo/db"; import { parsePlatformId } from "@mosoo/id"; -import type { - DriverInstanceId, - PlatformId, - SessionId, - SessionMessageId, - SessionRunId, -} from "@mosoo/id"; -import { and, desc, eq } from "drizzle-orm"; - -import { logInfo, logWarn } from "../../../../platform/cloudflare/logger"; -import { getAppDatabase } from "../../../../platform/db/drizzle"; -import type { - SessionLiveState, - SessionLiveStateMessage, - SessionViewSegment, -} from "../../../sessions/application/session-live-state.service"; -import { insertSessionMessage } from "../../../sessions/application/session-message-write.service"; - -function toSessionMessageSegments(segments: SessionViewSegment[]): SessionMessageSegment[] { - const result: SessionMessageSegment[] = []; - - for (const segment of segments) { - switch (segment.kind) { - case "text": { - result.push({ kind: "text", text: segment.text }); - break; - } - case "reasoning": { - break; - } - case "tool_use": { - result.push({ - argsText: segment.argsText, - kind: "tool_use", - path: segment.path, - tool: segment.tool, - toolCallId: segment.toolCallId, - }); - break; - } - case "tool_result": { - result.push({ - kind: "tool_result", - output: segment.output, - tool: segment.tool, - toolCallId: segment.toolCallId, - }); - break; - } - default: { - const exhaustiveSegment: never = segment; - - throw new Error(`Unsupported session segment kind: ${String(exhaustiveSegment)}`); - } - } - } - - return result; +import type { PlatformId, SessionId, SessionMessageId, SessionRunId } from "@mosoo/id"; + +export interface PreparedAssistantMessageProjection { + contentText: string; + createdByAccountId: PlatformId; + id: SessionMessageId; + planJson: string | null; + projectionFormat: "event_stream_v3"; + segmentsJson: string | null; + sessionId: SessionId; + sessionRunId: SessionRunId; } -function findAssistantMessage( - messages: SessionLiveStateMessage[], - messageId: string, -): SessionLiveStateMessage | null { - return ( - messages.find((message) => message.id === messageId && message.role === "assistant") ?? null - ); -} - -export async function persistAssistantMessageProjection( - database: D1Database, - input: { - createdByAccountId: PlatformId; - driverInstanceId: DriverInstanceId; - messageId: string; - messageText: string; - sessionId: SessionId; - sessionRunId: SessionRunId; - state: SessionLiveState; - }, -): Promise { - const message = findAssistantMessage(input.state.messages, input.messageId); - const useStructuredProjection = message?.content === input.messageText; - - if (message !== null && !useStructuredProjection) { - logWarn("runtime.assistant.message.snapshot_mismatch", { - driverInstanceId: input.driverInstanceId, - finalMessageId: input.messageId, - projectedTextLength: message.content.length, - reason: "RUN_FINISHED final assistant snapshot did not match the live projection", - sessionId: input.sessionId, - sessionRunId: input.sessionRunId, - snapshotTextLength: input.messageText.length, - }); - } - +export function prepareAssistantMessageProjection(input: { + createdByAccountId: PlatformId; + messageId: string; + sessionId: SessionId; + sessionRunId: SessionRunId; +}): PreparedAssistantMessageProjection { const messageId = parsePlatformId(input.messageId, "assistant message id"); - const plan = useStructuredProjection && message !== null ? message.plan : []; - const segments = - useStructuredProjection && message !== null - ? toSessionMessageSegments(message.segments) - : [{ kind: "text" as const, text: input.messageText }]; - const persistedMessage = await readPersistedAssistantMessage(database, input.sessionRunId); - - if (persistedMessage !== null) { - if (persistedMessage.content !== input.messageText) { - throw new Error( - `Canonical final assistant message conflicts with the persisted projection for run ${input.sessionRunId}.`, - ); - } - - // Provider reconnects can replay the same canonical snapshot after the - // driver process has generated a new message id. The run-level text is the - // durable identity here: preserve the first canonical transcript row. - return; - } - - logInfo("runtime.assistant.message.persisting", { - driverInstanceId: input.driverInstanceId, - planEntries: plan.length, - segmentCount: segments.length, + return { + // The authoritative aggregate lives in paginated session_event rows. This + // row is only the transcript identity/order reference; materializing the + // same unbounded body here would exceed D1's per-row limit. + contentText: "", + createdByAccountId: input.createdByAccountId, + id: messageId, + planJson: null, + projectionFormat: "event_stream_v3", + segmentsJson: null, sessionId: input.sessionId, sessionRunId: input.sessionRunId, - textLength: input.messageText.length, - }); + }; +} - await insertSessionMessage(database, { - content: input.messageText, +export function prepareToolCarrierProjection(input: { + createdByAccountId: PlatformId; + sessionId: SessionId; + sessionRunId: SessionRunId; +}): PreparedAssistantMessageProjection { + return prepareAssistantMessageProjection({ createdByAccountId: input.createdByAccountId, - id: messageId, - plan, - role: "assistant", - segments, + messageId: input.sessionRunId, sessionId: input.sessionId, sessionRunId: input.sessionRunId, }); } - -async function readPersistedAssistantMessage( - database: D1Database, - sessionRunId: SessionRunId, -): Promise<{ content: string; id: SessionMessageId } | null> { - const row = - (await getAppDatabase(database) - .select({ content: sessionMessagesTable.contentText, id: sessionMessagesTable.id }) - .from(sessionMessagesTable) - .where( - and( - eq(sessionMessagesTable.sessionRunId, sessionRunId), - eq(sessionMessagesTable.role, "assistant"), - ), - ) - .orderBy(desc(sessionMessagesTable.seq)) - .limit(1) - .get()) ?? null; - - return row; -} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/client.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/client.ts index eeb3b54f..154a32d4 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/client.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/client.ts @@ -3,7 +3,6 @@ import type { DriverInstanceId } from "@mosoo/id"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; import type { - DriverInstanceHelloResult, DriverInstanceReadyResult, DriverInstanceSnapshot, DriverInstanceWaitForCloseResult, @@ -87,26 +86,18 @@ export async function upgradeDriverInstanceSocket( ); } -export async function waitForDriverInstanceHello( - env: ApiBindings, - driverInstanceId: DriverInstanceId, - timeoutMs: number, -): Promise { - return expectJson( - await getDriverConnectionStub(env, driverInstanceId).fetch( - createDoRequest(driverInstanceId, `/wait/hello?timeoutMs=${timeoutMs}`), - ), - ); -} - export async function waitForDriverInstanceReady( env: ApiBindings, driverInstanceId: DriverInstanceId, + generation: number, timeoutMs: number, ): Promise { return expectJson( await getDriverConnectionStub(env, driverInstanceId).fetch( - createDoRequest(driverInstanceId, `/wait/ready?timeoutMs=${timeoutMs}`), + createDoRequest( + driverInstanceId, + `/wait/ready?generation=${generation}&timeoutMs=${timeoutMs}`, + ), ), ); } @@ -114,11 +105,15 @@ export async function waitForDriverInstanceReady( export async function waitForDriverInstanceClose( env: ApiBindings, driverInstanceId: DriverInstanceId, + generation: number, timeoutMs: number, ): Promise { return expectJson( await getDriverConnectionStub(env, driverInstanceId).fetch( - createDoRequest(driverInstanceId, `/wait/close?timeoutMs=${timeoutMs}`), + createDoRequest( + driverInstanceId, + `/wait/close?generation=${generation}&timeoutMs=${timeoutMs}`, + ), ), ); } @@ -137,12 +132,13 @@ export async function getDriverInstanceSnapshot( export async function sendDriverInstanceCommand( env: ApiBindings, driverInstanceId: DriverInstanceId, + generation: number, command: RuntimeCommand, ): Promise { await expectJson<{ ok: true }>( await getDriverConnectionStub(env, driverInstanceId).fetch( createDoRequest(driverInstanceId, "/control/send", { - body: JSON.stringify(command), + body: JSON.stringify({ command, generation }), headers: { "content-type": "application/json", }, @@ -155,12 +151,13 @@ export async function sendDriverInstanceCommand( export async function failDriverInstance( env: ApiBindings, driverInstanceId: DriverInstanceId, + generation: number, message: string, ): Promise { await expectJson<{ ok: true }>( await getDriverConnectionStub(env, driverInstanceId).fetch( createDoRequest(driverInstanceId, "/control/fail", { - body: JSON.stringify({ message }), + body: JSON.stringify({ generation, message }), headers: { "content-type": "application/json", }, @@ -173,12 +170,13 @@ export async function failDriverInstance( export async function destroyDriverInstanceDurableObject( env: ApiBindings, driverInstanceId: DriverInstanceId, + generation: number, reason: string, ): Promise { await expectJson<{ ok: true }>( await getDriverConnectionStub(env, driverInstanceId).fetch( createDoRequest(driverInstanceId, "/control/destroy", { - body: JSON.stringify({ reason }), + body: JSON.stringify({ generation, reason }), headers: { "content-type": "application/json", }, diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/commands.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/commands.ts index 4dec7623..58a9334f 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/commands.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/commands.ts @@ -12,7 +12,7 @@ export interface DriverInstanceCommandState { } interface CommandOptions { - markCommandDelivered: (command: RuntimeCommand) => Promise; + markCommandDelivered: (command: RuntimeCommand) => Promise<"delivered" | "discarded" | "retry">; persistCommandQueue: () => Promise; } @@ -37,17 +37,17 @@ async function dispatchQueuedRuntimeCommands( try { waiter.assertActiveConnection(); await options.persistCommandQueue(); - const delivered = await waiter.markCommandDelivered(command); + const delivery = await waiter.markCommandDelivered(command); waiter.assertActiveConnection(); - if (!delivered) { + if (delivery === "retry") { state.commandQueue.unshift(command); await options.persistCommandQueue(); waiter.deferred.resolve(null); continue; } - waiter.deferred.resolve(command); + waiter.deferred.resolve(delivery === "delivered" ? command : null); } catch { state.commandQueue.unshift(command); await options.persistCommandQueue(); @@ -89,16 +89,16 @@ export async function nextRuntimeCommand( } options.assertActiveConnection(); - const delivered = await options.markCommandDelivered(command); + const delivery = await options.markCommandDelivered(command); options.assertActiveConnection(); - if (!delivered) { + if (delivery === "retry") { state.commandQueue.unshift(command); await options.persistCommandQueue(); return null; } - return command; + return delivery === "delivered" ? command : null; } if (state.terminalized) { diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/completed-run-commit.repository.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/completed-run-commit.repository.ts new file mode 100644 index 00000000..90e84470 --- /dev/null +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/completed-run-commit.repository.ts @@ -0,0 +1,2864 @@ +import type { AgentKind } from "@mosoo/contracts/agent"; +import type { SessionStatus } from "@mosoo/contracts/session"; +import { DurableRunError } from "@mosoo/contracts/session-run"; +import type { RunError, SessionRunStatus } from "@mosoo/contracts/session-run"; +import { parseSchemaValue } from "@mosoo/contracts/validation"; +import { createPlatformId } from "@mosoo/id"; +import type { + DriverInstanceId, + PlatformId, + RuntimeEventId, + RuntimeOperationId, + SessionId, + SessionRunId, +} from "@mosoo/id"; +import { + createSessionRunTerminalSourceId, + createRuntimeEventSemanticHash, + readRuntimeEventPayload, + readRuntimeEventString, + readRuntimeRunPayload, + stringifyRuntimeEventSemanticValue, +} from "@mosoo/runtime-events"; +import type { RuntimeRunView } from "@mosoo/runtime-events"; + +import { getD1ChangeCount } from "../../../../platform/db/drizzle"; +import { currentTimestampMs, toIsoString } from "../../../../time"; +import { createSessionRuntimeEventProjection } from "../../../sessions/domain/session-runtime-event-projection"; +import { readTerminalEventSemanticAuthority } from "../../../sessions/domain/session-terminal-event-authority"; +import type { SessionRuntimeEventInput } from "../../../sessions/infrastructure/session-runtime-event-store.repository"; +import { decideSessionRunTransition } from "../../domain/session-run-lifecycle.machine"; +import { + prepareAssistantMessageProjection, + prepareToolCarrierProjection, +} from "./assistant-message-projection"; +import type { PreparedAssistantMessageProjection } from "./assistant-message-projection"; +import { prepareRuntimeArtifactPromotion } from "./runtime-artifact-attempt.repository"; + +interface CompletionStateRow { + archived_at: number | null; + cleanup_operation_kind: "archive" | "delete" | null; + completed_at: number | null; + created_at: number; + error_code: string | null; + error_details_json: string | null; + error_message: string | null; + error_retryable: number | null; + driver_last_heartbeat_at: number | null; + driver_connection_id: string | null; + driver_generation: number | null; + driver_status: string | null; + driver_status_operation_id: string | null; + driver_updated_at: number | null; + last_run_id: SessionRunId | null; + message_seq_cursor: number; + tool_carrier_anchor_json: string | null; + permission_request_count: number; + run_status: SessionRunStatus; + run_status_operation_id: string | null; + run_status_seq: number; + run_status_source: string; + run_trace_id: string; + run_updated_at: number; + run_created_by_account_id: PlatformId; + run_driver_instance_id: DriverInstanceId | null; + runtime_event_seq_cursor: number; + session_kind: AgentKind; + session_status: SessionStatus; + session_status_operation_id: string | null; + session_status_seq: number; + session_updated_at: number; + started_at: number | null; + final_message_artifact_json: string | null; + unfinalized_model_call_count: number; +} + +interface PersistedAssistantMessageRow { + content_text: string; + created_at: number; + created_by_account_id: string; + id: string; + plan_json: string | null; + projection_format: "event_stream_v3" | "materialized"; + segments_json: string | null; + seq: number; + session_id: string; + session_run_id: string; +} + +interface PreparedTerminalAssistantMessage { + createdAt: number; + firstEventSeq: number; + message: PreparedAssistantMessageProjection; +} + +interface GuardedAssistantMessage { + contentText: string; + createdAt: number; + createdByAccountId: string; + id: string; + planJson: string | null; + projectionFormat: "event_stream_v3" | "materialized"; + segmentsJson: string | null; + seq: number; + sessionId: SessionId; +} + +interface TranscriptAnchor { + occurredAt: number; + seq: number; +} + +interface PersistedTerminalEventRow { + artifact_attempt_id: string | null; + artifact_manifest_json: string | null; + artifact_manifest_sha256: string | null; + created_at: number; + ended_at: number | null; + event_type: string; + family: string; + id: RuntimeEventId; + occurred_at: number; + process_status: string; + process_type: string; + run_id: string | null; + semantic_hash: string | null; + terminal_event_json: string | null; + seq: number; + session_id: string; + source: string; + source_event_id: string; + stream_id: string | null; + tool_call_id: string | null; + tool_input_json: string | null; + tool_name: string | null; + tokens: number | null; + trace_id: string | null; + visibility: string; +} + +interface PersistedTerminalAuthorityRow extends PersistedTerminalEventRow { + agent_id: string; + content_text: string; + mcp_command_id: string | null; + tool_input_delta_json: string | null; + tool_output_delta_text: string | null; + tool_output_text: string | null; + tool_parent_message_id: string | null; + tool_result_message_id: string | null; + tool_status: string | null; +} + +interface CompletionSnapshot { + assistantMessages: PersistedAssistantMessageRow[]; + state: CompletionStateRow | null; + terminalEvents: PersistedTerminalEventRow[]; +} + +export type CompletedRunCommitResult = + | { + kind: "applied" | "duplicate"; + persistedSourceEventIds: readonly string[]; + runDurationMs: number | null; + } + | { + currentStatus: SessionRunStatus; + kind: "stale"; + persistedSourceEventIds: readonly []; + runDurationMs: null; + }; + +export type AdoptTerminalRunProjectionResult = + | CompletedRunCommitResult + | { + kind: "missing"; + persistedSourceEventIds: readonly []; + runDurationMs: null; + }; + +export type DriverTerminalRunStatus = "cancelled" | "completed" | "failed"; +export type HostTerminalRunStatus = DriverTerminalRunStatus | "expired"; +export type TerminalRunProjectionSource = + | "api" + | "driver" + | "maintenance" + | "runtime_operation" + | "system" + | "viewer"; + +export interface ExpectedTerminalSessionObservation { + lastRunId: SessionRunId | null; + status: SessionStatus; + statusSeq: number; + updatedAt: number; +} + +export interface ExpectedTerminalDriverObservation { + connectionId?: string | null; + driverInstanceId: DriverInstanceId | null; + generation?: number; + lastHeartbeatAt?: number | null; + status?: string | null; + updatedAt?: number | null; +} + +interface AtomicTerminalDriverReleaseObservation { + readonly connectionId: string | null; + readonly driverInstanceId: DriverInstanceId; + readonly generation: number; +} + +function readAtomicTerminalDriverReleaseObservation( + observation: ExpectedTerminalDriverObservation | undefined, +): AtomicTerminalDriverReleaseObservation | null { + if ( + observation?.driverInstanceId === null || + observation?.driverInstanceId === undefined || + observation.connectionId === undefined || + observation.generation === undefined + ) { + return null; + } + + return { + connectionId: observation.connectionId, + driverInstanceId: observation.driverInstanceId, + generation: observation.generation, + }; +} + +function expectedDriverObservationMatchesState( + state: CompletionStateRow, + observation: ExpectedTerminalDriverObservation, +): boolean { + return ( + state.run_driver_instance_id === observation.driverInstanceId && + (observation.connectionId === undefined || + state.driver_connection_id === observation.connectionId) && + (observation.generation === undefined || state.driver_generation === observation.generation) && + (observation.status === undefined || state.driver_status === observation.status) && + (observation.updatedAt === undefined || state.driver_updated_at === observation.updatedAt) && + (observation.lastHeartbeatAt === undefined || + state.driver_last_heartbeat_at === observation.lastHeartbeatAt) + ); +} + +function terminalEventKind( + status: HostTerminalRunStatus, +): "run.cancelled" | "run.completed" | "run.failed" { + return status === "expired" ? "run.cancelled" : `run.${status}`; +} + +function terminalLifecycleEvent(status: HostTerminalRunStatus) { + switch (status) { + case "cancelled": + return "run.cancel"; + case "completed": + return "run.complete"; + case "failed": + return "run.fail"; + case "expired": + return "run.expire"; + } +} + +function readTerminalRunOperationId(input: { + expectedSessionOperationId?: RuntimeOperationId | null; + source: TerminalRunProjectionSource; +}): RuntimeOperationId | null { + if (input.source !== "runtime_operation") { + return null; + } + if (input.expectedSessionOperationId === undefined || input.expectedSessionOperationId === null) { + throw new Error("A runtime-operation terminal transition requires its operation identity."); + } + return input.expectedSessionOperationId; +} + +interface PreparedTerminalEvent { + artifactAttemptId: string | null; + artifactManifestJson: string | null; + artifactManifestSha256: string | null; + contentText: string; + createdAt: number; + endedAt: number; + eventType: string; + explicitRunView: RuntimeRunView | null; + family: string; + id: RuntimeEventId; + occurredAt: number; + processStatus: string; + processType: string; + semanticHash: string; + terminalEventJson: string; + sessionStatus: Extract; + source: string; + sourceEventId: string; + streamId: string | null; + toolCallId: string | null; + toolInputJson: string | null; + toolName: string | null; + tokens: number | null; + traceId: string | null; + visibility: string; +} + +async function readTerminalAuthorityReceipt( + database: D1Database, + runId: SessionRunId, +): Promise { + const { results } = await database + .prepare( + `SELECT agent_id, artifact_attempt_id, artifact_manifest_json, + artifact_manifest_sha256, content_text, created_at, ended_at, event_type, + family, id, mcp_command_id, occurred_at, process_status, process_type, run_id, + semantic_hash, terminal_event_json, seq, session_id, source, source_event_id, + stream_id, tool_call_id, + tool_input_delta_json, tool_input_json, tool_name, tool_output_delta_text, + tool_output_text, tool_parent_message_id, tool_result_message_id, tool_status, + tokens, trace_id, visibility + FROM session_event + WHERE run_id = ? + AND event_type IN ('run.cancelled', 'run.completed', 'run.failed') + ORDER BY seq`, + ) + .bind(runId) + .all(); + + if (results.length > 1) { + throw new Error(`Session run ${runId} has conflicting durable terminal events.`); + } + + return results[0] ?? null; +} + +const PERSISTED_TERMINAL_EVENT_KEYS = [ + "artifact_attempt_id", + "artifact_manifest_json", + "artifact_manifest_sha256", + "created_at", + "ended_at", + "event_type", + "family", + "id", + "occurred_at", + "process_status", + "process_type", + "run_id", + "semantic_hash", + "terminal_event_json", + "seq", + "session_id", + "source", + "source_event_id", + "stream_id", + "tool_call_id", + "tool_input_json", + "tool_name", + "tokens", + "trace_id", + "visibility", +] as const satisfies readonly (keyof PersistedTerminalEventRow)[]; + +const PERSISTED_TERMINAL_AUTHORITY_KEYS = [ + ...PERSISTED_TERMINAL_EVENT_KEYS, + "agent_id", + "content_text", + "mcp_command_id", + "tool_input_delta_json", + "tool_output_delta_text", + "tool_output_text", + "tool_parent_message_id", + "tool_result_message_id", + "tool_status", +] as const satisfies readonly (keyof PersistedTerminalAuthorityRow)[]; + +function hasSameTerminalAuthority( + left: PersistedTerminalAuthorityRow, + right: PersistedTerminalAuthorityRow, +): boolean { + return PERSISTED_TERMINAL_AUTHORITY_KEYS.every((key) => left[key] === right[key]); +} + +function assertSnapshotTerminalAuthority( + snapshot: CompletionSnapshot, + expected: PersistedTerminalAuthorityRow, + runId: SessionRunId, +): void { + const [actual] = snapshot.terminalEvents; + if ( + snapshot.terminalEvents.length !== 1 || + actual === undefined || + PERSISTED_TERMINAL_EVENT_KEYS.some((key) => actual[key] !== expected[key]) + ) { + throw new Error(`Session run ${runId} terminal authority changed during adoption.`); + } +} + +async function assertDurableTerminalAuthority( + database: D1Database, + expected: PersistedTerminalAuthorityRow, + runId: SessionRunId, +): Promise { + const actual = await readTerminalAuthorityReceipt(database, runId); + if (actual === null || !hasSameTerminalAuthority(actual, expected)) { + throw new Error(`Session run ${runId} terminal authority changed during adoption.`); + } +} + +async function readCompletionSnapshot( + database: D1Database, + input: { + finalMessageId: string | null; + runId: SessionRunId; + sessionId: SessionId; + sourceEventId: string; + }, +): Promise { + const [state, assistantMessages, terminalEvents] = await Promise.all([ + database + .prepare( + `SELECT + s.archived_at, + s.cleanup_operation_kind, + r.completed_at, + r.created_at, + r.created_by_account_id AS run_created_by_account_id, + r.error_code, + r.error_details_json, + r.error_message, + r.error_retryable, + driver.connection_id AS driver_connection_id, + driver.generation AS driver_generation, + driver.last_heartbeat_at AS driver_last_heartbeat_at, + driver.status AS driver_status, + driver.status_operation_id AS driver_status_operation_id, + driver.updated_at AS driver_updated_at, + s.last_run_id, + s.message_seq_cursor, + ( + SELECT json_object('occurredAt', output.occurred_at, 'seq', output.seq) + FROM session_event AS output + WHERE output.session_id = s.id + AND output.run_id = r.id + AND output.event_type = 'tool.call.updated' + AND output.visibility = 'all_consumers' + AND output.tool_call_id IS NOT NULL + AND output.seq < COALESCE( + ( + SELECT terminal.seq + FROM session_event AS terminal + WHERE terminal.session_id = s.id + AND terminal.run_id = r.id + AND terminal.event_type IN ('run.cancelled', 'run.completed', 'run.failed') + ORDER BY terminal.seq + LIMIT 1 + ), + s.runtime_event_seq_cursor + 1 + ) + AND ( + ( + output.tool_parent_message_id IS NOT NULL + AND (? IS NULL OR output.tool_parent_message_id <> ?) + ) + OR ( + ( + output.tool_output_delta_text IS NOT NULL + OR output.tool_output_text IS NOT NULL + ) + AND NOT EXISTS ( + SELECT 1 + FROM session_event AS identity + WHERE identity.session_id = output.session_id + AND identity.run_id = output.run_id + AND identity.event_type = 'tool.call.updated' + AND identity.visibility = 'all_consumers' + AND identity.tool_call_id = output.tool_call_id + AND identity.tool_parent_message_id IS NOT NULL + AND identity.seq <= output.seq + ) + ) + ) + ORDER BY output.seq + LIMIT 1 + ) AS tool_carrier_anchor_json, + ( + SELECT COUNT(*) + FROM session_permission_request AS permission + WHERE permission.session_id = r.session_id AND permission.run_id = r.id + ) AS permission_request_count, + r.status AS run_status, + r.status_operation_id AS run_status_operation_id, + r.status_seq AS run_status_seq, + r.status_source AS run_status_source, + r.trace_id AS run_trace_id, + r.updated_at AS run_updated_at, + r.driver_instance_id AS run_driver_instance_id, + ( + SELECT COUNT(*) + FROM session_model_call AS model_call + WHERE model_call.session_id = r.session_id + AND model_call.session_run_id = r.id + AND ( + model_call.completed_at IS NULL + OR model_call.status <> CASE + WHEN r.status = 'completed' THEN 'completed' + ELSE 'failed' + END + ) + ) AS unfinalized_model_call_count, + s.runtime_event_seq_cursor, + s.kind AS session_kind, + s.status AS session_status, + s.status_operation_id AS session_status_operation_id, + s.status_seq AS session_status_seq, + s.updated_at AS session_updated_at, + r.started_at, + ( + SELECT json_object('occurredAt', message.occurred_at, 'seq', message.seq) + FROM session_event AS message + WHERE ? IS NOT NULL + AND message.session_id = s.id + AND message.run_id = r.id + AND message.stream_id = ? + AND message.process_type = 'agent.message.delta' + AND message.visibility = 'all_consumers' + AND message.event_type = 'message.added' + AND message.seq < COALESCE( + ( + SELECT terminal.seq + FROM session_event AS terminal + WHERE terminal.session_id = s.id + AND terminal.run_id = r.id + AND terminal.event_type IN ('run.cancelled', 'run.completed', 'run.failed') + ORDER BY terminal.seq + LIMIT 1 + ), + s.runtime_event_seq_cursor + 1 + ) + ORDER BY message.seq DESC + LIMIT 1 + ) AS final_message_artifact_json + FROM session_run AS r + INNER JOIN session AS s ON s.id = r.session_id + LEFT JOIN driver_instance AS driver ON driver.id = r.driver_instance_id + WHERE r.id = ? AND r.session_id = ? + LIMIT 1`, + ) + .bind( + input.finalMessageId, + input.finalMessageId, + input.finalMessageId, + input.finalMessageId, + input.runId, + input.sessionId, + ) + .first(), + database + .prepare( + `SELECT content_text, created_at, created_by_account_id, id, plan_json, projection_format, + segments_json, seq, session_id, session_run_id + FROM session_message + WHERE session_run_id = ? AND role = 'assistant' + ORDER BY seq`, + ) + .bind(input.runId) + .all(), + database + .prepare( + `SELECT artifact_attempt_id, artifact_manifest_json, artifact_manifest_sha256, + created_at, ended_at, event_type, family, id, occurred_at, + process_status, process_type, run_id, semantic_hash, terminal_event_json, seq, + session_id, source, source_event_id, stream_id, tool_call_id, tool_input_json, + tool_name, tokens, trace_id, visibility + FROM session_event + WHERE session_id = ? + AND ( + source_event_id = ? + OR (run_id = ? AND event_type IN ('run.cancelled', 'run.completed', 'run.failed')) + ) + ORDER BY seq`, + ) + .bind(input.sessionId, input.sourceEventId, input.runId) + .all(), + ]); + + return { + assistantMessages: assistantMessages.results, + state, + terminalEvents: terminalEvents.results, + }; +} + +function parseTranscriptAnchor(value: string | null, label: string): TranscriptAnchor | null { + if (value === null) { + return null; + } + const parsed: unknown = JSON.parse(value); + if ( + typeof parsed !== "object" || + parsed === null || + !("occurredAt" in parsed) || + typeof parsed.occurredAt !== "number" || + !("seq" in parsed) || + typeof parsed.seq !== "number" + ) { + throw new Error(`Canonical ${label} transcript anchor is invalid.`); + } + return { occurredAt: parsed.occurredAt, seq: parsed.seq }; +} + +function prepareTerminalAssistantMessages( + snapshot: CompletionSnapshot, + input: { + finalMessage: PreparedAssistantMessageProjection | null; + runId: SessionRunId; + sessionId: SessionId; + }, +): PreparedTerminalAssistantMessage[] { + const { state } = snapshot; + if (state === null) { + return []; + } + const finalArtifact = parseTranscriptAnchor(state.final_message_artifact_json, "final message"); + const toolArtifact = parseTranscriptAnchor(state.tool_carrier_anchor_json, "tool carrier"); + + if (input.finalMessage !== null && finalArtifact === null) { + throw new Error( + `Canonical final assistant ${input.finalMessage.id} has no pre-terminal message stream.`, + ); + } + + const final = + input.finalMessage === null || finalArtifact === null + ? null + : { + createdAt: finalArtifact.occurredAt, + firstEventSeq: finalArtifact.seq, + message: input.finalMessage, + }; + const carrier = + toolArtifact === null + ? null + : { + createdAt: toolArtifact.occurredAt, + firstEventSeq: toolArtifact.seq, + message: prepareToolCarrierProjection({ + createdByAccountId: state.run_created_by_account_id, + sessionId: input.sessionId, + sessionRunId: input.runId, + }), + }; + + if (final !== null && carrier !== null && final.message.id === carrier.message.id) { + return [ + final.firstEventSeq <= carrier.firstEventSeq + ? final + : { + ...final, + createdAt: carrier.createdAt, + firstEventSeq: carrier.firstEventSeq, + }, + ]; + } + + return [final, carrier] + .filter((message): message is PreparedTerminalAssistantMessage => message !== null) + .toSorted((left, right) => left.firstEventSeq - right.firstEventSeq); +} + +function isCanonicalAssistantMessage( + row: PersistedAssistantMessageRow, + expected: PreparedTerminalAssistantMessage, +): boolean { + return ( + row.content_text === expected.message.contentText && + row.created_at === expected.createdAt && + row.created_by_account_id === expected.message.createdByAccountId && + row.id === expected.message.id && + row.plan_json === expected.message.planJson && + row.projection_format === expected.message.projectionFormat && + row.segments_json === expected.message.segmentsJson && + row.session_id === expected.message.sessionId && + row.session_run_id === expected.message.sessionRunId + ); +} + +function assertCanonicalAssistantMessages( + rows: readonly PersistedAssistantMessageRow[], + expected: readonly PreparedTerminalAssistantMessage[], + runId: SessionRunId, +): void { + if ( + rows.length === expected.length && + rows.every((row, index) => { + const message = expected[index]; + return message !== undefined && isCanonicalAssistantMessage(row, message); + }) + ) { + return; + } + + throw new Error( + `Canonical assistant messages conflict with the persisted projection for run ${runId}.`, + ); +} + +function assertRepairableAssistantMessages( + rows: readonly PersistedAssistantMessageRow[], + expected: readonly PreparedTerminalAssistantMessage[], + runId: SessionRunId, +): void { + if ( + rows.length <= expected.length && + rows.every((row, index) => { + const message = expected[index]; + return message !== undefined && isCanonicalAssistantMessage(row, message); + }) + ) { + return; + } + + throw new Error(`Canonical assistant messages are not safely repairable for run ${runId}.`); +} + +function assertCanonicalTerminalRunError( + state: CompletionStateRow, + expected: RunError | null, + runId: SessionRunId, +): void { + const expectedDetails = + expected === null ? null : stringifyRuntimeEventSemanticValue(expected.details); + const persistedDetails = + state.error_details_json === null + ? null + : stringifyRuntimeEventSemanticValue(JSON.parse(state.error_details_json)); + + if ( + state.error_code !== (expected?.code ?? null) || + persistedDetails !== expectedDetails || + state.error_message !== (expected?.message ?? null) || + state.error_retryable !== (expected === null ? null : Number(expected.retryable)) + ) { + throw new Error(`Canonical terminal error conflicts with the persisted run ${runId}.`); + } +} + +function assertCanonicalTerminalEvent( + rows: readonly PersistedTerminalEventRow[], + expected: { + eventType: ReturnType; + runId: SessionRunId; + sourceEventId: string; + semanticHash: string; + }, +): PersistedTerminalEventRow { + const [row] = rows; + + if ( + rows.length !== 1 || + row === undefined || + row.event_type !== expected.eventType || + row.run_id !== expected.runId || + row.semantic_hash !== expected.semanticHash || + row.source_event_id !== expected.sourceEventId + ) { + throw new Error( + `Canonical ${expected.eventType} receipt conflicts with the persisted projection for run ${expected.runId}.`, + ); + } + + return row; +} + +function readRunDurationMs(state: CompletionStateRow): number | null { + if (state.completed_at === null) { + return null; + } + + return Math.max(0, state.completed_at - (state.started_at ?? state.created_at)); +} + +function projectedSessionOperationId( + state: CompletionStateRow, + sessionStatus: PreparedTerminalEvent["sessionStatus"], +): string | null { + return sessionStatus === "IDLE" ? state.session_status_operation_id : null; +} + +function isLegacyTerminalProjection(snapshot: CompletionSnapshot): boolean { + return snapshot.terminalEvents.length === 1 && snapshot.terminalEvents[0]?.semantic_hash === null; +} + +function assertLegacyAssistantMessages( + rows: readonly PersistedAssistantMessageRow[], + expected: readonly PreparedTerminalAssistantMessage[], + runId: SessionRunId, + sessionId: SessionId, +): void { + const materialized = rows.filter((row) => row.projection_format === "materialized"); + const carriers = rows.filter((row) => row.projection_format === "event_stream_v3"); + const expectedCarrier = expected.find(({ message }) => String(message.id) === String(runId)); + if ( + materialized.length > 1 || + carriers.length > 1 || + materialized.length + carriers.length !== rows.length || + rows.some((row) => row.session_id !== sessionId || row.session_run_id !== runId) || + carriers.some( + (row) => + expectedCarrier === undefined || + row.id !== runId || + !isCanonicalAssistantMessage(row, expectedCarrier), + ) + ) { + throw new Error(`Legacy terminal assistant messages conflict with run ${runId}.`); + } +} + +function staleTerminalResult(state: CompletionStateRow): CompletedRunCommitResult { + return { + currentStatus: state.run_status, + kind: "stale", + persistedSourceEventIds: [], + runDurationMs: null, + }; +} + +function hasSafelyRepairableStaleSessionProjection( + state: CompletionStateRow, + input: { runId: SessionRunId }, +): boolean { + return ( + state.archived_at === null && + state.cleanup_operation_kind === null && + state.last_run_id === input.runId && + state.session_status === "RUNNING" && + state.session_status_operation_id === null && + state.session_updated_at <= state.run_updated_at + ); +} + +function explicitRunViewMatchesFreshState( + state: CompletionStateRow, + input: { + createdAt: number; + error: RunError | null; + explicitRunView: RuntimeRunView | null; + runId: SessionRunId; + targetStatus: HostTerminalRunStatus; + }, +): boolean { + const run = input.explicitRunView; + if (run === null) { + return true; + } + + return ( + run.completedAt === toIsoString(state.completed_at ?? input.createdAt) && + run.id === input.runId && + run.startedAt === toIsoString(state.started_at ?? input.createdAt) && + run.status === input.targetStatus && + run.traceId === state.run_trace_id && + stringifyRuntimeEventSemanticValue(run.error) === + stringifyRuntimeEventSemanticValue(input.error) + ); +} + +function classifyTerminalSnapshot( + snapshot: CompletionSnapshot, + input: { + assistantMessage: PreparedAssistantMessageProjection | null; + assistantMessages: readonly PreparedTerminalAssistantMessage[]; + createdAt: number; + error: RunError | null; + explicitRunView: RuntimeRunView | null; + expectedDriverObservation?: ExpectedTerminalDriverObservation; + expectedRunStatus?: SessionRunStatus; + expectedSessionObservation?: ExpectedTerminalSessionObservation; + expectedSessionOperationId?: RuntimeOperationId | null; + eventId: RuntimeEventId; + runId: SessionRunId; + runStatusOperationId: RuntimeOperationId | null; + semanticHash: string; + sessionId: SessionId; + sessionStatus: PreparedTerminalEvent["sessionStatus"]; + sourceEventId: string; + source: TerminalRunProjectionSource; + targetStatus: HostTerminalRunStatus; + }, +): CompletedRunCommitResult | null { + const { state } = snapshot; + + if (state === null) { + throw new Error(`Session run ${input.runId} was not found for atomic completion.`); + } + const terminalDriverRelease = readAtomicTerminalDriverReleaseObservation( + input.expectedDriverObservation, + ); + if (state.run_status === input.targetStatus) { + if ( + input.expectedDriverObservation !== undefined && + !expectedDriverObservationMatchesState(state, input.expectedDriverObservation) + ) { + return staleTerminalResult(state); + } + if ( + terminalDriverRelease !== null && + state.driver_status_operation_id !== null && + state.driver_status_operation_id !== input.runId && + state.driver_status !== "stopping" + ) { + return staleTerminalResult(state); + } + if (isLegacyTerminalProjection(snapshot)) { + if (input.assistantMessage !== null) { + throw new Error( + "A legacy materialized terminal run cannot adopt an event-stream reference.", + ); + } + assertLegacyAssistantMessages( + snapshot.assistantMessages, + input.assistantMessages, + input.runId, + input.sessionId, + ); + assertCanonicalTerminalRunError(state, input.error, input.runId); + const [legacyEvent] = snapshot.terminalEvents; + if ( + legacyEvent === undefined || + legacyEvent.event_type !== terminalEventKind(input.targetStatus) || + legacyEvent.run_id !== input.runId || + legacyEvent.source_event_id !== + createSessionRunTerminalSourceId(input.runId, terminalEventKind(input.targetStatus)) + ) { + throw new Error(`Legacy terminal projection conflicts with run ${input.runId}.`); + } + const cursorCovered = + legacyEvent.seq <= state.runtime_event_seq_cursor && + snapshot.assistantMessages.every((message) => message.seq <= state.message_seq_cursor); + if ( + cursorCovered && + state.permission_request_count === 0 && + state.unfinalized_model_call_count === 0 + ) { + return { + kind: "duplicate", + persistedSourceEventIds: [], + runDurationMs: readRunDurationMs(state), + }; + } + if ( + state.archived_at !== null || + state.last_run_id !== input.runId || + state.session_status_operation_id !== null || + (state.session_status !== "IDLE" && state.session_status !== "TERMINATED") + ) { + throw new Error( + `Legacy terminal projection for run ${input.runId} is not safely repairable.`, + ); + } + return null; + } + if ( + state.run_status_operation_id !== input.runStatusOperationId || + state.run_status_source !== input.source + ) { + return staleTerminalResult(state); + } + if (!explicitRunViewMatchesFreshState(state, input)) { + return staleTerminalResult(state); + } + assertCanonicalTerminalRunError(state, input.error, input.runId); + assertRepairableAssistantMessages( + snapshot.assistantMessages, + input.assistantMessages, + input.runId, + ); + + if (snapshot.terminalEvents.length === 0) { + if ( + state.archived_at !== null || + state.last_run_id !== input.runId || + (state.session_status !== "RUNNING" && state.session_status !== input.sessionStatus) || + state.session_status_operation_id !== null || + state.session_updated_at > state.run_updated_at + ) { + throw new Error( + `Missing terminal projection for run ${input.runId} is not safely repairable.`, + ); + } + return null; + } + + const hasMissingAssistantMessages = + snapshot.assistantMessages.length < input.assistantMessages.length; + if (!hasMissingAssistantMessages) { + assertCanonicalAssistantMessages( + snapshot.assistantMessages, + input.assistantMessages, + input.runId, + ); + } + const terminalEvent = assertCanonicalTerminalEvent(snapshot.terminalEvents, { + ...input, + eventType: terminalEventKind(input.targetStatus), + semanticHash: input.semanticHash, + }); + if (terminalEvent.seq > state.runtime_event_seq_cursor) { + throw new Error( + `Canonical terminal event is ahead of the Session cursor for run ${input.runId}.`, + ); + } + if ( + hasMissingAssistantMessages || + snapshot.assistantMessages.some((message) => message.seq > state.message_seq_cursor) || + state.permission_request_count !== 0 + ) { + return null; + } + + if (state.unfinalized_model_call_count > 0) { + return null; + } + + if (hasSafelyRepairableStaleSessionProjection(state, input)) { + return null; + } + + return { + kind: terminalEvent.id === input.eventId ? "applied" : "duplicate", + persistedSourceEventIds: terminalEvent.id === input.eventId ? [input.sourceEventId] : [], + runDurationMs: readRunDurationMs(state), + }; + } + + if (["cancelled", "completed", "expired", "failed"].includes(state.run_status)) { + return staleTerminalResult(state); + } + + if (input.expectedRunStatus !== undefined && state.run_status !== input.expectedRunStatus) { + return staleTerminalResult(state); + } + if ( + input.expectedSessionObservation !== undefined && + (state.last_run_id !== input.expectedSessionObservation.lastRunId || + state.session_status !== input.expectedSessionObservation.status || + state.session_status_seq !== input.expectedSessionObservation.statusSeq || + state.session_updated_at !== input.expectedSessionObservation.updatedAt) + ) { + return staleTerminalResult(state); + } + if ( + input.expectedDriverObservation !== undefined && + (!expectedDriverObservationMatchesState(state, input.expectedDriverObservation) || + (terminalDriverRelease !== null && + state.driver_status_operation_id !== null && + state.driver_status !== "stopping")) + ) { + return staleTerminalResult(state); + } + + if ( + input.expectedSessionOperationId !== undefined && + state.session_status_operation_id !== input.expectedSessionOperationId + ) { + return staleTerminalResult(state); + } + if (!explicitRunViewMatchesFreshState(state, input)) { + return staleTerminalResult(state); + } + + return null; +} + +async function prepareTerminalEvent( + record: SessionRuntimeEventInput, + input: { + assistantMessage: PreparedAssistantMessageProjection | null; + error: RunError | null; + runId: SessionRunId; + sessionId: SessionId; + targetStatus: HostTerminalRunStatus; + timestampMs?: number; + }, +): Promise { + const expectedKind = terminalEventKind(input.targetStatus); + const expectedSourceEventId = createSessionRunTerminalSourceId(input.runId, expectedKind); + + if ( + record.event.kind !== expectedKind || + record.event.runId !== input.runId || + record.event.sessionId !== input.sessionId || + record.sourceEventId !== expectedSourceEventId || + record.event.sourceEventId !== expectedSourceEventId + ) { + throw new Error( + `Atomic run terminal projection requires one exact canonical ${expectedKind} identity.`, + ); + } + + const eventPayload = readRuntimeEventPayload(record.event); + const runtimeRunPayload = readRuntimeRunPayload(record.event); + const runPayload = runtimeRunPayload.run; + if ( + runPayload === null || + runPayload.id !== input.runId || + runPayload.status !== input.targetStatus + ) { + throw new Error(`Atomic ${expectedKind} payload does not match its terminal run identity.`); + } + if (runtimeRunPayload.lifecycle !== "IDLE" && runtimeRunPayload.lifecycle !== "TERMINATED") { + throw new Error(`Atomic ${expectedKind} projection requires a terminal Session lifecycle.`); + } + if (input.targetStatus === "failed") { + if (input.error === null || runPayload.error === null) { + throw new Error("Atomic run.failed projection requires one exact RunError."); + } + if ( + runPayload.error.code !== input.error.code || + stringifyRuntimeEventSemanticValue(runPayload.error.details) !== + stringifyRuntimeEventSemanticValue(input.error.details) || + runPayload.error.message !== input.error.message || + runPayload.error.retryable !== input.error.retryable + ) { + throw new Error("Atomic run.failed payload conflicts with its persisted RunError."); + } + } else if (input.targetStatus === "completed") { + if (input.error !== null || runPayload.error !== null) { + throw new Error(`Atomic ${expectedKind} projection cannot persist a RunError.`); + } + } else if ((input.error === null) !== (runPayload.error === null)) { + throw new Error(`Atomic ${expectedKind} RunError must be present together.`); + } else if ( + input.error !== null && + runPayload.error !== null && + (runPayload.error.code !== input.error.code || + stringifyRuntimeEventSemanticValue(runPayload.error.details) !== + stringifyRuntimeEventSemanticValue(input.error.details) || + runPayload.error.message !== input.error.message || + runPayload.error.retryable !== input.error.retryable) + ) { + throw new Error(`Atomic ${expectedKind} projection cannot persist a RunError.`); + } + + const finalMessageId = readRuntimeEventString(eventPayload, "finalMessageId"); + if ( + input.targetStatus === "completed" && + (finalMessageId === null) !== (input.assistantMessage === null) + ) { + throw new Error("run.completed final message identity and reference must be present together."); + } + if (input.assistantMessage !== null) { + if ( + input.targetStatus !== "completed" || + input.assistantMessage.sessionId !== input.sessionId || + input.assistantMessage.sessionRunId !== input.runId || + input.assistantMessage.id !== finalMessageId + ) { + throw new Error("Atomic final assistant reference conflicts with run.completed."); + } + } + + const timestampMs = input.timestampMs ?? currentTimestampMs(); + const occurredAt = record.occurredAt ?? timestampMs; + const parsedEndedAt = Date.parse(record.event.occurredAt); + const projection = createSessionRuntimeEventProjection(record.event); + const artifactAttemptId = record.artifactAttemptId ?? null; + const artifactManifestJson = record.artifactManifestJson ?? null; + const artifactManifestSha256 = record.artifactManifestSha256 ?? null; + if ( + (artifactAttemptId === null) !== (artifactManifestJson === null) || + (artifactAttemptId === null) !== (artifactManifestSha256 === null) || + (artifactAttemptId !== null && expectedKind !== "run.completed") + ) { + throw new Error("Atomic terminal artifact projection is incomplete or out of scope."); + } + + if ( + projection.eventType !== expectedKind || + projection.runId !== input.runId || + projection.visibility !== "all_consumers" + ) { + throw new Error("Atomic run terminal event projection is not public and run-scoped."); + } + + return { + artifactAttemptId, + artifactManifestJson, + artifactManifestSha256, + contentText: projection.contentText, + createdAt: timestampMs, + endedAt: + Number.isFinite(parsedEndedAt) && parsedEndedAt >= occurredAt ? parsedEndedAt : occurredAt, + eventType: projection.eventType, + explicitRunView: Object.hasOwn(eventPayload, "run") ? runPayload : null, + family: projection.family, + id: createPlatformId(), + occurredAt, + processStatus: projection.processStatus, + processType: projection.processType, + semanticHash: await createRuntimeEventSemanticHash(record.event), + terminalEventJson: stringifyRuntimeEventSemanticValue(record.event), + sessionStatus: runtimeRunPayload.lifecycle, + source: projection.source, + sourceEventId: expectedSourceEventId, + streamId: projection.streamId, + toolCallId: projection.toolCallId, + toolInputJson: projection.toolInputJson, + toolName: projection.toolName, + tokens: projection.tokens, + traceId: projection.traceId, + visibility: projection.visibility, + }; +} + +function prepareTerminalEventInsert( + database: D1Database, + input: { + assistantMessageCount: number; + expectedDriverObservation?: ExpectedTerminalDriverObservation; + event: PreparedTerminalEvent; + runId: SessionRunId; + sessionId: SessionId; + state: CompletionStateRow; + }, +): D1PreparedStatement { + const { event, state } = input; + + return database + .prepare( + `/* completed-run:event */ + INSERT INTO session_event ( + agent_id, artifact_attempt_id, artifact_manifest_json, artifact_manifest_sha256, + content_text, created_at, ended_at, event_type, family, id, + occurred_at, process_status, process_type, run_id, semantic_hash, terminal_event_json, + seq, session_id, source_event_id, source, stream_id, tool_call_id, tool_input_json, + tool_name, tokens, trace_id, visibility + ) + SELECT + s.agent_id, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, r.id, ?, ?, + s.runtime_event_seq_cursor + 1, s.id, ?, ?, ?, ?, ?, ?, ?, ?, ? + FROM session AS s + INNER JOIN session_run AS r ON r.session_id = s.id + WHERE s.id = ? + AND r.id = ? + AND r.status = ? + AND r.status_seq = ? + AND r.status_operation_id IS ? + AND r.status_source = ? + AND (? = 0 OR r.driver_instance_id IS ?) + AND ( + ? = 0 OR ( + SELECT driver.connection_id FROM driver_instance AS driver + WHERE driver.id = r.driver_instance_id + ) IS ? + ) + AND ( + ? = 0 OR ( + SELECT driver.generation FROM driver_instance AS driver + WHERE driver.id = r.driver_instance_id + ) IS ? + ) + AND ( + ? = 0 OR ( + SELECT driver.status FROM driver_instance AS driver + WHERE driver.id = r.driver_instance_id + ) IS ? + ) + AND ( + ? = 0 OR ( + SELECT driver.updated_at FROM driver_instance AS driver + WHERE driver.id = r.driver_instance_id + ) IS ? + ) + AND ( + ? = 0 OR ( + SELECT driver.last_heartbeat_at FROM driver_instance AS driver + WHERE driver.id = r.driver_instance_id + ) IS ? + ) + AND s.archived_at IS ? + AND s.cleanup_operation_kind IS ? + AND s.last_run_id IS ? + AND s.status = ? + AND s.status_operation_id IS ? + AND s.status_seq = ? + AND s.updated_at = ? + AND s.message_seq_cursor = ? + AND s.runtime_event_seq_cursor = ? + AND ( + SELECT COUNT(*) + FROM session_message AS m + WHERE m.session_run_id = r.id AND m.role = 'assistant' + ) = ? + AND NOT EXISTS ( + SELECT 1 + FROM session_event AS existing + WHERE existing.session_id = s.id + AND ( + existing.source_event_id = ? + OR ( + existing.run_id = r.id + AND existing.event_type IN ('run.cancelled', 'run.completed', 'run.failed') + ) + ) + ) + AND NOT EXISTS ( + SELECT 1 + FROM driver_command AS command + WHERE command.kind = 'mcp.execute' + AND command.status = 'accepted' + AND json_extract(command.payload_json, '$.runId') = r.id + ) + AND NOT EXISTS ( + SELECT 1 + FROM external_tool_effect AS effect + INNER JOIN driver_command AS command ON command.id = effect.command_id + WHERE effect.session_run_id = r.id + AND effect.status IN ('claimed', 'succeeded', 'unknown') + AND command.status IN ('queued', 'delivered', 'accepted') + )`, + ) + .bind( + event.artifactAttemptId, + event.artifactManifestJson, + event.artifactManifestSha256, + event.contentText, + event.createdAt, + event.endedAt, + event.eventType, + event.family, + event.id, + event.occurredAt, + event.processStatus, + event.processType, + event.semanticHash, + event.terminalEventJson, + event.sourceEventId, + event.source, + event.streamId, + event.toolCallId, + event.toolInputJson, + event.toolName, + event.tokens, + event.traceId, + event.visibility, + input.sessionId, + input.runId, + state.run_status, + state.run_status_seq, + state.run_status_operation_id, + state.run_status_source, + Number(input.expectedDriverObservation !== undefined), + input.expectedDriverObservation?.driverInstanceId ?? null, + Number(input.expectedDriverObservation?.connectionId !== undefined), + input.expectedDriverObservation?.connectionId ?? null, + Number(input.expectedDriverObservation?.generation !== undefined), + input.expectedDriverObservation?.generation ?? null, + Number(input.expectedDriverObservation?.status !== undefined), + input.expectedDriverObservation?.status ?? null, + Number(input.expectedDriverObservation?.updatedAt !== undefined), + input.expectedDriverObservation?.updatedAt ?? null, + Number(input.expectedDriverObservation?.lastHeartbeatAt !== undefined), + input.expectedDriverObservation?.lastHeartbeatAt ?? null, + state.archived_at, + state.cleanup_operation_kind, + state.last_run_id, + state.session_status, + state.session_status_operation_id, + state.session_status_seq, + state.session_updated_at, + state.message_seq_cursor, + state.runtime_event_seq_cursor, + input.assistantMessageCount, + event.sourceEventId, + ); +} + +function prepareTerminalEventAdoptionFence( + database: D1Database, + row: PersistedTerminalAuthorityRow, +): D1PreparedStatement { + return database + .prepare( + `/* completed-run:adopt-event */ + INSERT INTO session_event (id) + SELECT ? + WHERE NOT EXISTS ( + SELECT 1 + FROM session_event + WHERE id = ? + AND agent_id = ? + AND artifact_attempt_id IS ? + AND artifact_manifest_json IS ? + AND artifact_manifest_sha256 IS ? + AND content_text = ? + AND created_at = ? + AND ended_at IS ? + AND event_type = ? + AND family = ? + AND mcp_command_id IS ? + AND occurred_at = ? + AND process_status = ? + AND process_type = ? + AND run_id IS ? + AND semantic_hash IS ? + AND terminal_event_json IS ? + AND seq = ? + AND session_id = ? + AND source = ? + AND source_event_id = ? + AND stream_id IS ? + AND tool_call_id IS ? + AND tool_input_delta_json IS ? + AND tool_input_json IS ? + AND tool_name IS ? + AND tool_output_delta_text IS ? + AND tool_output_text IS ? + AND tool_parent_message_id IS ? + AND tool_result_message_id IS ? + AND tool_status IS ? + AND tokens IS ? + AND trace_id IS ? + AND visibility = ? + )`, + ) + .bind( + row.id, + row.id, + row.agent_id, + row.artifact_attempt_id, + row.artifact_manifest_json, + row.artifact_manifest_sha256, + row.content_text, + row.created_at, + row.ended_at, + row.event_type, + row.family, + row.mcp_command_id, + row.occurred_at, + row.process_status, + row.process_type, + row.run_id, + row.semantic_hash, + row.terminal_event_json, + row.seq, + row.session_id, + row.source, + row.source_event_id, + row.stream_id, + row.tool_call_id, + row.tool_input_delta_json, + row.tool_input_json, + row.tool_name, + row.tool_output_delta_text, + row.tool_output_text, + row.tool_parent_message_id, + row.tool_result_message_id, + row.tool_status, + row.tokens, + row.trace_id, + row.visibility, + ); +} + +function prepareAssistantMessageInsert( + database: D1Database, + input: { + eventId: RuntimeEventId; + message: PreparedTerminalAssistantMessage; + seqOffset: number; + }, +): D1PreparedStatement { + return database + .prepare( + `/* completed-run:message */ + INSERT INTO session_message ( + content_text, created_at, created_by_account_id, id, plan_json, role, + projection_format, segments_json, seq, session_id, session_run_id + ) + SELECT ?, ?, ?, ?, ?, 'assistant', ?, ?, s.message_seq_cursor + ?, s.id, ? + FROM session AS s + INNER JOIN session_event AS terminal + ON terminal.session_id = s.id AND terminal.id = ? + WHERE s.id = ? + AND terminal.event_type IN ('run.cancelled', 'run.completed', 'run.failed')`, + ) + .bind( + input.message.message.contentText, + input.message.createdAt, + input.message.message.createdByAccountId, + input.message.message.id, + input.message.message.planJson, + input.message.message.projectionFormat, + input.message.message.segmentsJson, + input.seqOffset, + input.message.message.sessionRunId, + input.eventId, + input.message.message.sessionId, + ); +} + +function prepareRunTerminalUpdate( + database: D1Database, + input: { + error: RunError | null; + eventId: RuntimeEventId; + runId: SessionRunId; + sessionId: SessionId; + source: TerminalRunProjectionSource; + statusOperationId: RuntimeOperationId | null; + state: CompletionStateRow; + targetStatus: HostTerminalRunStatus; + timestampMs: number; + }, +): D1PreparedStatement { + return database + .prepare( + `/* completed-run:run */ + UPDATE session_run + SET completed_at = ?, + error_code = ?, + error_details_json = ?, + error_message = ?, + error_retryable = ?, + started_at = COALESCE(started_at, ?), + status = ?, + status_changed_at = ?, + status_event = ?, + status_operation_id = ?, + status_seq = status_seq + 1, + status_source = ?, + updated_at = ? + WHERE id = ? + AND session_id = ? + AND status = ? + AND status_seq = ? + AND EXISTS ( + SELECT 1 + FROM session_event AS terminal + WHERE terminal.id = ? + AND terminal.session_id = ? + AND terminal.run_id = ? + AND terminal.event_type = ? + )`, + ) + .bind( + input.timestampMs, + input.error?.code ?? null, + input.error === null ? null : stringifyRuntimeEventSemanticValue(input.error.details), + input.error?.message ?? null, + input.error === null ? null : Number(input.error.retryable), + input.timestampMs, + input.targetStatus, + input.timestampMs, + terminalLifecycleEvent(input.targetStatus), + input.statusOperationId, + input.source, + input.timestampMs, + input.runId, + input.sessionId, + input.state.run_status, + input.state.run_status_seq, + input.eventId, + input.sessionId, + input.runId, + terminalEventKind(input.targetStatus), + ); +} + +function prepareSessionModelCallsTerminalUpdate( + database: D1Database, + input: { + eventId: RuntimeEventId; + runId: SessionRunId; + semanticHash: string | null; + sessionId: SessionId; + targetStatus: HostTerminalRunStatus; + timestampMs: number; + }, +): D1PreparedStatement { + const status = input.targetStatus === "completed" ? "completed" : "failed"; + + return database + .prepare( + `/* completed-run:model-calls */ + UPDATE session_model_call + SET completed_at = COALESCE(completed_at, ?), + status = ?, + updated_at = MAX(updated_at, ?) + WHERE session_id = ? + AND session_run_id = ? + AND EXISTS ( + SELECT 1 + FROM session_event AS terminal + WHERE terminal.id = ? + AND terminal.session_id = ? + AND terminal.run_id = ? + AND terminal.event_type = ? + AND terminal.semantic_hash IS ? + )`, + ) + .bind( + input.timestampMs, + status, + input.timestampMs, + input.sessionId, + input.runId, + input.eventId, + input.sessionId, + input.runId, + terminalEventKind(input.targetStatus), + input.semanticHash, + ); +} + +function prepareSessionTerminalUpdate( + database: D1Database, + input: { + assistantMessageCount: number; + eventId: RuntimeEventId; + lastMessageAt: number; + runId: SessionRunId; + sessionId: SessionId; + sessionOperationId: string | null; + sessionStatus: PreparedTerminalEvent["sessionStatus"]; + state: CompletionStateRow; + targetStatus: HostTerminalRunStatus; + timestampMs: number; + }, +): D1PreparedStatement { + return database + .prepare( + `/* completed-run:session */ + UPDATE session + SET last_message_at = CASE WHEN ? > 0 THEN ? ELSE last_message_at END, + message_seq_cursor = message_seq_cursor + ?, + runtime_event_seq_cursor = runtime_event_seq_cursor + 1, + status = ?, + status_operation_id = ?, + status_seq = status_seq + 1, + updated_at = MAX(updated_at, ?), + workspace_checkpoint_required = CASE + WHEN ? = 'completed' AND kind = 'cattle' THEN 1 + ELSE workspace_checkpoint_required + END + WHERE id = ? + AND archived_at IS ? + AND cleanup_operation_kind IS ? + AND last_run_id = ? + AND status = ? + AND status_operation_id IS ? + AND status_seq = ? + AND updated_at = ? + AND message_seq_cursor = ? + AND runtime_event_seq_cursor = ? + AND EXISTS ( + SELECT 1 + FROM session_event AS terminal + WHERE terminal.id = ? + AND terminal.session_id = ? + AND terminal.run_id = ? + AND terminal.event_type = ? + ) + AND EXISTS ( + SELECT 1 + FROM session_run AS completed_run + WHERE completed_run.id = ? + AND completed_run.session_id = ? + AND completed_run.status = ? + AND completed_run.status_seq = ? + )`, + ) + .bind( + input.assistantMessageCount, + input.lastMessageAt, + input.assistantMessageCount, + input.sessionStatus, + input.sessionOperationId, + input.timestampMs, + input.targetStatus, + input.sessionId, + input.state.archived_at, + input.state.cleanup_operation_kind, + input.runId, + input.state.session_status, + input.state.session_status_operation_id, + input.state.session_status_seq, + input.state.session_updated_at, + input.state.message_seq_cursor, + input.state.runtime_event_seq_cursor, + input.eventId, + input.sessionId, + input.runId, + terminalEventKind(input.targetStatus), + input.runId, + input.sessionId, + input.targetStatus, + input.state.run_status_seq + 1, + ); +} + +function prepareSessionTerminalRepair( + database: D1Database, + input: { + eventId: RuntimeEventId; + lastMessageAt: number; + messageCursorIncrement: number; + runId: SessionRunId; + runtimeEventCursorIncrement: 0 | 1; + sessionId: SessionId; + sessionOperationId: string | null; + sessionStatus: PreparedTerminalEvent["sessionStatus"]; + state: CompletionStateRow; + targetStatus: HostTerminalRunStatus; + timestampMs: number; + }, +): D1PreparedStatement { + const touchLastMessage = Number(input.messageCursorIncrement > 0); + + return database + .prepare( + `/* completed-run:session-repair */ + UPDATE session + SET last_message_at = CASE WHEN ? = 1 THEN ? ELSE last_message_at END, + message_seq_cursor = message_seq_cursor + ?, + runtime_event_seq_cursor = runtime_event_seq_cursor + ?, + status = CASE WHEN last_run_id = ? THEN ? ELSE status END, + status_operation_id = CASE + WHEN last_run_id = ? THEN ? + ELSE status_operation_id + END, + status_seq = status_seq + CASE + WHEN last_run_id = ? AND (status <> ? OR status_operation_id IS NOT ?) + THEN 1 + ELSE 0 + END, + updated_at = MAX(updated_at, ?), + workspace_checkpoint_required = CASE + WHEN ? = 'completed' AND kind = 'cattle' THEN 1 + ELSE workspace_checkpoint_required + END + WHERE id = ? + AND archived_at IS ? + AND cleanup_operation_kind IS ? + AND last_run_id IS ? + AND status = ? + AND status_operation_id IS ? + AND status_seq = ? + AND updated_at = ? + AND message_seq_cursor = ? + AND runtime_event_seq_cursor = ? + AND EXISTS ( + SELECT 1 + FROM session_event AS terminal + WHERE terminal.id = ? + AND terminal.session_id = ? + AND terminal.run_id = ? + AND terminal.event_type = ? + ) + AND EXISTS ( + SELECT 1 + FROM session_run AS terminal_run + WHERE terminal_run.id = ? + AND terminal_run.session_id = ? + AND terminal_run.status = ? + AND terminal_run.status_seq = ? + )`, + ) + .bind( + touchLastMessage, + input.lastMessageAt, + input.messageCursorIncrement, + input.runtimeEventCursorIncrement, + input.runId, + input.sessionStatus, + input.runId, + input.sessionOperationId, + input.runId, + input.sessionStatus, + input.sessionOperationId, + input.timestampMs, + input.targetStatus, + input.sessionId, + input.state.archived_at, + input.state.cleanup_operation_kind, + input.state.last_run_id, + input.state.session_status, + input.state.session_status_operation_id, + input.state.session_status_seq, + input.state.session_updated_at, + input.state.message_seq_cursor, + input.state.runtime_event_seq_cursor, + input.eventId, + input.sessionId, + input.runId, + terminalEventKind(input.targetStatus), + input.runId, + input.sessionId, + input.targetStatus, + input.state.run_status_seq, + ); +} + +function preparePermissionRequestDelete( + database: D1Database, + input: { + eventId: RuntimeEventId; + runId: SessionRunId; + sessionId: SessionId; + targetStatus: HostTerminalRunStatus; + }, +): D1PreparedStatement { + return database + .prepare( + `/* completed-run:permissions */ + DELETE FROM session_permission_request + WHERE session_id = ? + AND run_id = ? + AND EXISTS ( + SELECT 1 + FROM session_event AS terminal + INNER JOIN session_run AS completed_run + ON completed_run.id = terminal.run_id + WHERE terminal.id = ? + AND terminal.session_id = ? + AND terminal.run_id = ? + AND terminal.event_type = ? + AND completed_run.status = ? + )`, + ) + .bind( + input.sessionId, + input.runId, + input.eventId, + input.sessionId, + input.runId, + terminalEventKind(input.targetStatus), + input.targetStatus, + ); +} + +function prepareTerminalDriverReleaseClaim( + database: D1Database, + input: { + eventId: RuntimeEventId; + observation: AtomicTerminalDriverReleaseObservation; + runId: SessionRunId; + sessionId: SessionId; + targetStatus: HostTerminalRunStatus; + }, +): D1PreparedStatement { + return database + .prepare( + `/* completed-run:driver-release-claim */ + UPDATE driver_instance + SET status_operation_id = ? + WHERE id = ? + AND generation = ? + AND connection_id IS ? + AND status IN ('provisioning', 'connecting', 'ready', 'stopped', 'failed') + AND status_operation_id IS NULL + AND NOT EXISTS ( + SELECT 1 + FROM session_run AS active_successor + WHERE active_successor.driver_instance_id = driver_instance.id + AND active_successor.id <> ? + AND active_successor.status IN ('queued', 'booting', 'running', 'waiting_input') + ) + AND EXISTS ( + SELECT 1 + FROM session_event AS terminal + INNER JOIN session_run AS terminal_run + ON terminal_run.id = terminal.run_id + AND terminal_run.session_id = terminal.session_id + WHERE terminal.id = ? + AND terminal.session_id = ? + AND terminal.run_id = ? + AND terminal.event_type = ? + AND terminal_run.driver_instance_id = driver_instance.id + AND terminal_run.status = ? + )`, + ) + .bind( + input.runId, + input.observation.driverInstanceId, + input.observation.generation, + input.observation.connectionId, + input.runId, + input.eventId, + input.sessionId, + input.runId, + terminalEventKind(input.targetStatus), + input.targetStatus, + ); +} + +async function claimTerminalDriverForAdoption( + database: D1Database, + input: { + fence: PersistedTerminalAuthorityRow; + observation: AtomicTerminalDriverReleaseObservation; + result: Extract; + runId: SessionRunId; + sessionId: SessionId; + state: CompletionStateRow; + targetStatus: HostTerminalRunStatus; + }, +): Promise { + const results = await database.batch([ + prepareTerminalEventAdoptionFence(database, input.fence), + prepareTerminalDriverReleaseClaim(database, { + eventId: input.fence.id, + observation: input.observation, + runId: input.runId, + sessionId: input.sessionId, + targetStatus: input.targetStatus, + }), + ]); + + return getD1ChangeCount(results[1]) === 1 ? input.result : staleTerminalResult(input.state); +} + +function prepareTerminalCommitGuard( + database: D1Database, + input: { + assistantMessages: readonly GuardedAssistantMessage[]; + driverRelease: AtomicTerminalDriverReleaseObservation | null; + event: PreparedTerminalEvent; + error: RunError | null; + messageCursorIncrement: number; + repairingExistingTerminal: boolean; + runId: SessionRunId; + runtimeEventCursorIncrement: 0 | 1; + sessionId: SessionId; + source: string; + state: CompletionStateRow; + statusOperationId: string | null; + targetStatus: HostTerminalRunStatus; + terminalEventId: RuntimeEventId; + terminalSemanticHash: string | null; + terminalSourceEventId: string; + }, +): D1PreparedStatement { + const sameCurrentRun = input.state.last_run_id === input.runId; + const repairedLifecycle = + input.repairingExistingTerminal && + sameCurrentRun && + (input.state.session_status !== input.event.sessionStatus || + input.state.session_status_operation_id !== + projectedSessionOperationId(input.state, input.event.sessionStatus)); + const expectedSessionStatus = sameCurrentRun + ? input.event.sessionStatus + : input.state.session_status; + const expectedSessionOperationId = sameCurrentRun + ? projectedSessionOperationId(input.state, input.event.sessionStatus) + : input.state.session_status_operation_id; + const expectedSessionStatusSeq = + input.state.session_status_seq + + (input.repairingExistingTerminal ? Number(repairedLifecycle) : 1); + const expectedRunStatusSeq = + input.state.run_status_seq + Number(!input.repairingExistingTerminal); + const assistantMessagesJson = JSON.stringify(input.assistantMessages); + const expectedErrorDetailsJson = input.repairingExistingTerminal + ? input.state.error_details_json + : input.error === null + ? null + : stringifyRuntimeEventSemanticValue(input.error.details); + + return database + .prepare( + `/* completed-run:guard */ + INSERT INTO session_event (id) + SELECT ? + WHERE NOT EXISTS ( + SELECT 1 + FROM session_event AS terminal + INNER JOIN session_run AS terminal_run + ON terminal_run.id = terminal.run_id AND terminal_run.session_id = terminal.session_id + INNER JOIN session AS terminal_session ON terminal_session.id = terminal.session_id + WHERE terminal.id = ? + AND terminal.session_id = ? + AND terminal.run_id = ? + AND terminal.event_type = ? + AND terminal.source_event_id = ? + AND terminal.semantic_hash IS ? + AND terminal_run.status = ? + AND terminal_run.status_seq = ? + AND terminal_run.status_operation_id IS ? + AND terminal_run.status_source = ? + AND terminal_run.error_code IS ? + AND terminal_run.error_details_json IS ? + AND terminal_run.error_message IS ? + AND terminal_run.error_retryable IS ? + AND terminal_session.archived_at IS ? + AND terminal_session.cleanup_operation_kind IS ? + AND terminal_session.last_run_id IS ? + AND terminal_session.status = ? + AND terminal_session.status_operation_id IS ? + AND terminal_session.status_seq = ? + AND terminal_session.message_seq_cursor = ? + AND terminal_session.runtime_event_seq_cursor = ? + AND ( + ? = 0 OR EXISTS ( + SELECT 1 + FROM driver_instance AS terminal_driver + WHERE terminal_driver.id = ? + AND terminal_driver.generation = ? + AND terminal_driver.connection_id IS ? + AND terminal_driver.status_operation_id = ? + AND terminal_run.driver_instance_id = terminal_driver.id + ) + ) + AND NOT EXISTS ( + SELECT 1 FROM session_permission_request AS permission + WHERE permission.session_id = ? AND permission.run_id = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM session_model_call AS model_call + WHERE model_call.session_id = ? + AND model_call.session_run_id = ? + AND ( + model_call.completed_at IS NULL + OR model_call.status <> ? + ) + ) + AND ? = ( + SELECT COUNT(*) + FROM session_message AS message + WHERE message.session_run_id = ? AND message.role = 'assistant' + ) + AND NOT EXISTS ( + SELECT 1 + FROM session_message AS message + WHERE message.session_run_id = ? + AND message.role = 'assistant' + AND NOT EXISTS ( + SELECT 1 + FROM json_each(?) AS expected + WHERE message.id = json_extract(expected.value, '$.id') + AND message.session_id = json_extract(expected.value, '$.sessionId') + AND message.content_text = json_extract(expected.value, '$.contentText') + AND message.created_at = json_extract(expected.value, '$.createdAt') + AND message.plan_json IS json_extract(expected.value, '$.planJson') + AND message.projection_format = json_extract( + expected.value, + '$.projectionFormat' + ) + AND message.segments_json IS json_extract(expected.value, '$.segmentsJson') + AND message.created_by_account_id = json_extract( + expected.value, + '$.createdByAccountId' + ) + AND message.seq = json_extract(expected.value, '$.seq') + ) + ) + )`, + ) + .bind( + input.event.id, + input.terminalEventId, + input.sessionId, + input.runId, + terminalEventKind(input.targetStatus), + input.terminalSourceEventId, + input.terminalSemanticHash, + input.targetStatus, + expectedRunStatusSeq, + input.statusOperationId, + input.source, + input.error?.code ?? null, + expectedErrorDetailsJson, + input.error?.message ?? null, + input.error === null ? null : Number(input.error.retryable), + input.state.archived_at, + input.state.cleanup_operation_kind, + input.state.last_run_id, + expectedSessionStatus, + expectedSessionOperationId, + expectedSessionStatusSeq, + input.state.message_seq_cursor + input.messageCursorIncrement, + input.state.runtime_event_seq_cursor + input.runtimeEventCursorIncrement, + Number(input.driverRelease !== null), + input.driverRelease?.driverInstanceId ?? "", + input.driverRelease?.generation ?? -1, + input.driverRelease === null ? "" : input.driverRelease.connectionId, + input.runId, + input.sessionId, + input.runId, + input.sessionId, + input.runId, + input.targetStatus === "completed" ? "completed" : "failed", + input.assistantMessages.length, + input.runId, + input.runId, + assistantMessagesJson, + ); +} + +async function classifyAfterBatch( + database: D1Database, + input: { + assistantMessage: PreparedAssistantMessageProjection | null; + error: RunError | null; + event: PreparedTerminalEvent; + expectedDriverObservation?: ExpectedTerminalDriverObservation; + expectedRunStatus?: SessionRunStatus; + expectedSessionObservation?: ExpectedTerminalSessionObservation; + expectedSessionOperationId?: RuntimeOperationId | null; + runId: SessionRunId; + runStatusOperationId: RuntimeOperationId | null; + sessionId: SessionId; + source: TerminalRunProjectionSource; + targetStatus: HostTerminalRunStatus; + terminalEventFence?: PersistedTerminalAuthorityRow; + }, +): Promise { + const snapshot = await readCompletionSnapshot(database, { + finalMessageId: input.assistantMessage?.id ?? null, + runId: input.runId, + sessionId: input.sessionId, + sourceEventId: input.event.sourceEventId, + }); + if (input.terminalEventFence !== undefined) { + assertSnapshotTerminalAuthority(snapshot, input.terminalEventFence, input.runId); + await assertDurableTerminalAuthority(database, input.terminalEventFence, input.runId); + } + const assistantMessages = prepareTerminalAssistantMessages(snapshot, { + finalMessage: input.assistantMessage, + runId: input.runId, + sessionId: input.sessionId, + }); + + return classifyTerminalSnapshot(snapshot, { + assistantMessage: input.assistantMessage, + assistantMessages, + createdAt: input.event.createdAt, + error: input.error, + explicitRunView: input.event.explicitRunView, + ...(input.expectedDriverObservation === undefined + ? {} + : { expectedDriverObservation: input.expectedDriverObservation }), + ...(input.expectedRunStatus === undefined + ? {} + : { expectedRunStatus: input.expectedRunStatus }), + ...(input.expectedSessionObservation === undefined + ? {} + : { expectedSessionObservation: input.expectedSessionObservation }), + ...(input.expectedSessionOperationId === undefined + ? {} + : { expectedSessionOperationId: input.expectedSessionOperationId }), + eventId: input.event.id, + runId: input.runId, + runStatusOperationId: input.runStatusOperationId, + semanticHash: input.event.semanticHash, + sessionId: input.sessionId, + sessionStatus: input.event.sessionStatus, + sourceEventId: input.event.sourceEventId, + source: input.source, + targetStatus: input.targetStatus, + }); +} + +function readPersistedTerminalStatus( + status: SessionRunStatus, + runId: SessionRunId, +): HostTerminalRunStatus { + switch (status) { + case "cancelled": + case "completed": + case "expired": + case "failed": + return status; + case "booting": + case "queued": + case "running": + case "waiting_input": + throw new Error(`Session run ${runId} has a terminal receipt without a terminal status.`); + } +} + +function readPersistedTerminalSource( + source: string, + runId: SessionRunId, +): TerminalRunProjectionSource { + switch (source) { + case "api": + case "driver": + case "maintenance": + case "runtime_operation": + case "system": + case "viewer": + return source; + default: + throw new Error(`Session run ${runId} has an invalid terminal status source.`); + } +} + +function readPersistedTerminalError( + state: CompletionStateRow, + runId: SessionRunId, +): RunError | null { + if ( + state.error_code === null && + state.error_details_json === null && + state.error_message === null && + state.error_retryable === null + ) { + return null; + } + if ( + state.error_code === null || + state.error_details_json === null || + state.error_message === null || + (state.error_retryable !== 0 && state.error_retryable !== 1) + ) { + throw new Error(`Session run ${runId} has an invalid durable terminal error.`); + } + + try { + return parseSchemaValue(DurableRunError, { + code: state.error_code, + details: JSON.parse(state.error_details_json), + message: state.error_message, + retryable: state.error_retryable === 1, + }); + } catch { + throw new Error(`Session run ${runId} has an invalid durable terminal error.`); + } +} + +function assertAdoptableTerminalReceipt( + row: PersistedTerminalAuthorityRow, + input: { + runId: SessionRunId; + sessionId: SessionId; + targetStatus: HostTerminalRunStatus; + }, +): void { + const eventType = terminalEventKind(input.targetStatus); + const hasArtifact = row.artifact_attempt_id !== null; + if ( + row.session_id !== input.sessionId || + row.run_id !== input.runId || + row.event_type !== eventType || + row.source_event_id !== createSessionRunTerminalSourceId(input.runId, eventType) || + row.visibility !== "all_consumers" || + row.mcp_command_id !== null || + row.tool_call_id !== null || + row.tool_input_delta_json !== null || + row.tool_input_json !== null || + row.tool_name !== null || + row.tool_output_delta_text !== null || + row.tool_output_text !== null || + row.tool_parent_message_id !== null || + row.tool_result_message_id !== null || + row.tool_status !== null || + (row.semantic_hash !== null && !/^[0-9a-f]{64}$/.test(row.semantic_hash)) || + (row.semantic_hash === null) !== (row.terminal_event_json === null) || + (eventType !== "run.completed" && row.stream_id !== null) || + hasArtifact !== (row.artifact_manifest_json !== null) || + hasArtifact !== (row.artifact_manifest_sha256 !== null) || + (hasArtifact && eventType !== "run.completed") + ) { + throw new Error(`Session run ${input.runId} has an invalid canonical terminal receipt.`); + } +} + +function assertPreparedTerminalEventMatchesAuthority( + event: PreparedTerminalEvent, + row: PersistedTerminalAuthorityRow, + runId: SessionRunId, +): void { + if ( + event.artifactAttemptId !== row.artifact_attempt_id || + event.artifactManifestJson !== row.artifact_manifest_json || + event.artifactManifestSha256 !== row.artifact_manifest_sha256 || + event.contentText !== row.content_text || + event.createdAt !== row.created_at || + event.endedAt !== row.ended_at || + event.eventType !== row.event_type || + event.family !== row.family || + event.occurredAt !== row.occurred_at || + event.processStatus !== row.process_status || + event.processType !== row.process_type || + event.semanticHash !== row.semantic_hash || + event.source !== row.source || + event.sourceEventId !== row.source_event_id || + event.streamId !== row.stream_id || + event.terminalEventJson !== row.terminal_event_json || + event.toolCallId !== row.tool_call_id || + event.toolInputJson !== row.tool_input_json || + event.toolName !== row.tool_name || + event.tokens !== row.tokens || + event.traceId !== row.trace_id || + event.visibility !== row.visibility + ) { + throw new Error(`Session run ${runId} terminal semantic authority conflicts with its receipt.`); + } +} + +function prepareAdoptedTerminalEvent( + row: PersistedTerminalAuthorityRow, + state: CompletionStateRow, +): PreparedTerminalEvent { + return { + artifactAttemptId: row.artifact_attempt_id, + artifactManifestJson: row.artifact_manifest_json, + artifactManifestSha256: row.artifact_manifest_sha256, + contentText: row.content_text, + createdAt: row.created_at, + endedAt: row.ended_at ?? row.occurred_at, + eventType: row.event_type, + explicitRunView: null, + family: row.family, + id: createPlatformId(), + occurredAt: row.occurred_at, + processStatus: row.process_status, + processType: row.process_type, + semanticHash: row.semantic_hash ?? "", + sessionStatus: state.session_status === "TERMINATED" ? "TERMINATED" : "IDLE", + source: row.source, + sourceEventId: row.source_event_id, + streamId: row.stream_id, + terminalEventJson: row.terminal_event_json ?? "", + toolCallId: row.tool_call_id, + toolInputJson: row.tool_input_json, + toolName: row.tool_name, + tokens: row.tokens, + traceId: row.trace_id, + visibility: row.visibility, + }; +} + +export async function adoptTerminalRunProjection( + database: D1Database, + input: { + expectedDriverObservation?: AtomicTerminalDriverReleaseObservation; + expectedTargetStatus?: DriverTerminalRunStatus; + runId: SessionRunId; + sessionId: SessionId; + }, +): Promise { + const authority = await readTerminalAuthorityReceipt(database, input.runId); + if (authority === null) { + return { kind: "missing", persistedSourceEventIds: [], runDurationMs: null }; + } + if (authority.session_id !== input.sessionId) { + throw new Error(`Session run ${input.runId} terminal receipt belongs to another session.`); + } + const semanticAuthority = + authority.semantic_hash === null + ? null + : await readTerminalEventSemanticAuthority({ + eventJson: authority.terminal_event_json, + eventType: authority.event_type, + runId: input.runId, + semanticHash: authority.semantic_hash, + sessionId: input.sessionId, + sourceEventId: authority.source_event_id, + streamId: authority.stream_id, + }); + const persistedEvent = semanticAuthority?.event ?? null; + const persistedFinalMessageId = semanticAuthority?.finalMessageId ?? null; + + const snapshot = await readCompletionSnapshot(database, { + finalMessageId: persistedFinalMessageId, + runId: input.runId, + sessionId: input.sessionId, + sourceEventId: authority.source_event_id, + }); + const { state } = snapshot; + if (state === null) { + throw new Error(`Session run ${input.runId} was not found for terminal adoption.`); + } + assertSnapshotTerminalAuthority(snapshot, authority, input.runId); + + const targetStatus = readPersistedTerminalStatus(state.run_status, input.runId); + const error = readPersistedTerminalError(state, input.runId); + assertAdoptableTerminalReceipt(authority, { + runId: input.runId, + sessionId: input.sessionId, + targetStatus, + }); + if ( + input.expectedTargetStatus !== undefined && + terminalEventKind(input.expectedTargetStatus) !== terminalEventKind(targetStatus) + ) { + return staleTerminalResult(state); + } + + const source = readPersistedTerminalSource(state.run_status_source, input.runId); + const expectedDriverObservation = + input.expectedDriverObservation === undefined + ? undefined + : { + connectionId: input.expectedDriverObservation.connectionId, + driverInstanceId: input.expectedDriverObservation.driverInstanceId, + generation: input.expectedDriverObservation.generation, + }; + const assistantMessage = + persistedEvent === null || persistedFinalMessageId === null + ? null + : prepareAssistantMessageProjection({ + createdByAccountId: state.run_created_by_account_id, + messageId: persistedFinalMessageId, + sessionId: input.sessionId, + sessionRunId: input.runId, + }); + const event = + persistedEvent === null + ? prepareAdoptedTerminalEvent(authority, state) + : await prepareTerminalEvent( + { + artifactAttemptId: authority.artifact_attempt_id, + artifactManifestJson: authority.artifact_manifest_json, + artifactManifestSha256: authority.artifact_manifest_sha256, + event: persistedEvent, + occurredAt: authority.occurred_at, + sourceEventId: authority.source_event_id, + }, + { + assistantMessage, + error, + runId: input.runId, + sessionId: input.sessionId, + targetStatus, + timestampMs: authority.created_at, + }, + ); + if (persistedEvent !== null) { + assertPreparedTerminalEventMatchesAuthority(event, authority, input.runId); + } + + return commitPreparedTerminalRunProjection(database, { + assistantMessage, + error, + event, + ...(expectedDriverObservation === undefined ? {} : { expectedDriverObservation }), + runId: input.runId, + runStatusOperationId: state.run_status_operation_id as RuntimeOperationId | null, + sessionId: input.sessionId, + source, + targetStatus, + terminalEventFence: authority, + }); +} + +export async function commitTerminalRunProjection( + database: D1Database, + input: { + assistantMessage: PreparedAssistantMessageProjection | null; + error: RunError | null; + expectedDriverObservation?: ExpectedTerminalDriverObservation; + expectedRunStatus?: SessionRunStatus; + expectedSessionObservation?: ExpectedTerminalSessionObservation; + expectedSessionOperationId?: RuntimeOperationId | null; + runId: SessionRunId; + sessionId: SessionId; + source: TerminalRunProjectionSource; + targetStatus: HostTerminalRunStatus; + terminalEvent: SessionRuntimeEventInput; + timestampMs?: number; + }, +): Promise { + if (input.targetStatus !== "completed" && input.assistantMessage !== null) { + throw new Error("Only a completed run can project a final assistant message."); + } + + const runStatusOperationId = readTerminalRunOperationId(input); + const event = await prepareTerminalEvent(input.terminalEvent, input); + + return commitPreparedTerminalRunProjection(database, { + ...input, + event, + runStatusOperationId, + }); +} + +async function commitPreparedTerminalRunProjection( + database: D1Database, + input: { + assistantMessage: PreparedAssistantMessageProjection | null; + error: RunError | null; + event: PreparedTerminalEvent; + expectedDriverObservation?: ExpectedTerminalDriverObservation; + expectedRunStatus?: SessionRunStatus; + expectedSessionObservation?: ExpectedTerminalSessionObservation; + expectedSessionOperationId?: RuntimeOperationId | null; + runId: SessionRunId; + runStatusOperationId: RuntimeOperationId | null; + sessionId: SessionId; + source: TerminalRunProjectionSource; + targetStatus: HostTerminalRunStatus; + terminalEventFence?: PersistedTerminalAuthorityRow; + }, +): Promise { + const { event, runStatusOperationId } = input; + const snapshot = await readCompletionSnapshot(database, { + finalMessageId: input.assistantMessage?.id ?? null, + runId: input.runId, + sessionId: input.sessionId, + sourceEventId: event.sourceEventId, + }); + if (input.terminalEventFence !== undefined) { + assertSnapshotTerminalAuthority(snapshot, input.terminalEventFence, input.runId); + await assertDurableTerminalAuthority(database, input.terminalEventFence, input.runId); + } + const assistantMessages = prepareTerminalAssistantMessages(snapshot, { + finalMessage: input.assistantMessage, + runId: input.runId, + sessionId: input.sessionId, + }); + const terminal = classifyTerminalSnapshot(snapshot, { + assistantMessage: input.assistantMessage, + assistantMessages, + createdAt: event.createdAt, + error: input.error, + explicitRunView: event.explicitRunView, + ...(input.expectedDriverObservation === undefined + ? {} + : { expectedDriverObservation: input.expectedDriverObservation }), + ...(input.expectedRunStatus === undefined + ? {} + : { expectedRunStatus: input.expectedRunStatus }), + ...(input.expectedSessionObservation === undefined + ? {} + : { expectedSessionObservation: input.expectedSessionObservation }), + ...(input.expectedSessionOperationId === undefined + ? {} + : { expectedSessionOperationId: input.expectedSessionOperationId }), + eventId: event.id, + runId: input.runId, + runStatusOperationId, + semanticHash: event.semanticHash, + sessionId: input.sessionId, + sessionStatus: event.sessionStatus, + sourceEventId: event.sourceEventId, + source: input.source, + targetStatus: input.targetStatus, + }); + const { state } = snapshot; + + if (terminal !== null) { + const driverRelease = readAtomicTerminalDriverReleaseObservation( + input.expectedDriverObservation, + ); + if ( + (terminal.kind === "applied" || terminal.kind === "duplicate") && + input.terminalEventFence !== undefined && + state !== null && + driverRelease !== null && + state.driver_status_operation_id === null + ) { + return claimTerminalDriverForAdoption(database, { + fence: input.terminalEventFence, + observation: driverRelease, + result: terminal, + runId: input.runId, + sessionId: input.sessionId, + state, + targetStatus: input.targetStatus, + }); + } + return terminal; + } + + if (state === null) { + throw new Error(`Session run ${input.runId} was not found for atomic completion.`); + } + + const repairingExistingTerminal = state.run_status === input.targetStatus; + const observedDriverRelease = readAtomicTerminalDriverReleaseObservation( + input.expectedDriverObservation, + ); + const driverRelease = + observedDriverRelease !== null && state.driver_status_operation_id === null + ? observedDriverRelease + : null; + + if (!repairingExistingTerminal) { + const decision = decideSessionRunTransition({ + currentStatus: state.run_status, + targetStatus: input.targetStatus, + }); + + if (decision.kind !== "accepted") { + throw new Error(`Driver terminal run projection was rejected: ${decision.kind}.`); + } + } + + if ( + state.session_status === "TERMINATED" && + (!repairingExistingTerminal || event.sessionStatus !== "TERMINATED") + ) { + throw new Error("Session is not writable for an atomic terminal run projection."); + } + + if ( + (!repairingExistingTerminal && snapshot.terminalEvents.length > 0) || + (!repairingExistingTerminal && snapshot.assistantMessages.length > 0) + ) { + throw new Error("Atomic terminal run projection found a partial canonical projection."); + } + + if (!repairingExistingTerminal && state.last_run_id !== input.runId) { + throw new Error("Active terminal run is no longer the Session's current run."); + } + + const persistedTerminalEvent = repairingExistingTerminal + ? (snapshot.terminalEvents[0] ?? null) + : null; + const terminalEventId = persistedTerminalEvent?.id ?? event.id; + const terminalSemanticHash = + persistedTerminalEvent === null ? event.semanticHash : persistedTerminalEvent.semantic_hash; + const terminalSourceEventId = persistedTerminalEvent?.source_event_id ?? event.sourceEventId; + const repairingLegacyTerminal = persistedTerminalEvent?.semantic_hash === null; + const guardedRunSource = repairingLegacyTerminal ? state.run_status_source : input.source; + const guardedRunStatusOperationId = repairingLegacyTerminal + ? state.run_status_operation_id + : runStatusOperationId; + const sessionOperationId = projectedSessionOperationId(state, event.sessionStatus); + const runtimeEventCursorIncrement = (() => { + if (persistedTerminalEvent === null) { + return 1; + } + if (persistedTerminalEvent.seq <= state.runtime_event_seq_cursor) { + return 0; + } + if (persistedTerminalEvent.seq === state.runtime_event_seq_cursor + 1) { + return 1; + } + throw new Error("Canonical terminal event is separated from the Session cursor by a gap."); + })() satisfies 0 | 1; + const assistantMessageInserts: { + message: PreparedTerminalAssistantMessage; + seqOffset: number; + }[] = []; + const guardedAssistantMessages: GuardedAssistantMessage[] = []; + let messageCursorIncrement = 0; + + if (repairingLegacyTerminal) { + for (const persisted of snapshot.assistantMessages) { + if (persisted.seq > state.message_seq_cursor) { + messageCursorIncrement += 1; + if (persisted.seq !== state.message_seq_cursor + messageCursorIncrement) { + throw new Error( + "Legacy assistant message is separated from the Session cursor by a gap.", + ); + } + } + guardedAssistantMessages.push({ + contentText: persisted.content_text, + createdAt: persisted.created_at, + createdByAccountId: persisted.created_by_account_id, + id: persisted.id, + planJson: persisted.plan_json, + projectionFormat: persisted.projection_format, + segmentsJson: persisted.segments_json, + seq: persisted.seq, + sessionId: input.sessionId, + }); + } + } else { + assertRepairableAssistantMessages(snapshot.assistantMessages, assistantMessages, input.runId); + for (const [index, message] of assistantMessages.entries()) { + const persisted = snapshot.assistantMessages[index]; + let seq: number; + if (persisted === undefined) { + messageCursorIncrement += 1; + seq = state.message_seq_cursor + messageCursorIncrement; + assistantMessageInserts.push({ + message, + seqOffset: messageCursorIncrement, + }); + } else { + seq = persisted.seq; + if (seq > state.message_seq_cursor) { + messageCursorIncrement += 1; + if (seq !== state.message_seq_cursor + messageCursorIncrement) { + throw new Error( + "Canonical assistant messages are separated from the Session cursor by a gap.", + ); + } + } + } + guardedAssistantMessages.push({ + contentText: message.message.contentText, + createdAt: message.createdAt, + createdByAccountId: message.message.createdByAccountId, + id: message.message.id, + planJson: message.message.planJson, + projectionFormat: message.message.projectionFormat, + segmentsJson: message.message.segmentsJson, + seq, + sessionId: message.message.sessionId, + }); + } + } + const lastMessageAt = guardedAssistantMessages.at(-1)?.createdAt ?? event.createdAt; + + const statements = [ + input.terminalEventFence === undefined + ? prepareTerminalEventInsert(database, { + assistantMessageCount: snapshot.assistantMessages.length, + event, + ...(input.expectedDriverObservation === undefined + ? {} + : { expectedDriverObservation: input.expectedDriverObservation }), + runId: input.runId, + sessionId: input.sessionId, + state, + }) + : prepareTerminalEventAdoptionFence(database, input.terminalEventFence), + ...(!repairingExistingTerminal && + event.artifactAttemptId !== null && + event.artifactManifestJson !== null && + event.artifactManifestSha256 !== null + ? prepareRuntimeArtifactPromotion(database, { + attemptId: event.artifactAttemptId, + eventId: event.id, + manifestJson: event.artifactManifestJson, + manifestSha256: event.artifactManifestSha256, + timestampMs: event.createdAt, + }) + : []), + ...assistantMessageInserts.map(({ message, seqOffset }) => + prepareAssistantMessageInsert(database, { + eventId: terminalEventId, + message, + seqOffset, + }), + ), + ...(repairingExistingTerminal + ? [ + prepareSessionTerminalRepair(database, { + eventId: terminalEventId, + lastMessageAt, + messageCursorIncrement, + runId: input.runId, + runtimeEventCursorIncrement, + sessionId: input.sessionId, + sessionOperationId, + sessionStatus: event.sessionStatus, + state, + targetStatus: input.targetStatus, + timestampMs: event.createdAt, + }), + ] + : [ + prepareRunTerminalUpdate(database, { + error: input.error, + eventId: event.id, + runId: input.runId, + sessionId: input.sessionId, + source: input.source, + state, + statusOperationId: runStatusOperationId, + targetStatus: input.targetStatus, + timestampMs: event.createdAt, + }), + prepareSessionTerminalUpdate(database, { + assistantMessageCount: assistantMessages.length, + eventId: event.id, + lastMessageAt, + runId: input.runId, + sessionId: input.sessionId, + sessionOperationId, + sessionStatus: event.sessionStatus, + state, + targetStatus: input.targetStatus, + timestampMs: event.createdAt, + }), + ]), + prepareSessionModelCallsTerminalUpdate(database, { + eventId: terminalEventId, + runId: input.runId, + semanticHash: terminalSemanticHash, + sessionId: input.sessionId, + targetStatus: input.targetStatus, + timestampMs: event.createdAt, + }), + ...(driverRelease === null + ? [] + : [ + prepareTerminalDriverReleaseClaim(database, { + eventId: terminalEventId, + observation: driverRelease, + runId: input.runId, + sessionId: input.sessionId, + targetStatus: input.targetStatus, + }), + ]), + preparePermissionRequestDelete(database, { + eventId: terminalEventId, + runId: input.runId, + sessionId: input.sessionId, + targetStatus: input.targetStatus, + }), + prepareTerminalCommitGuard(database, { + assistantMessages: guardedAssistantMessages, + driverRelease, + event, + error: input.error, + messageCursorIncrement, + repairingExistingTerminal, + runId: input.runId, + source: guardedRunSource, + statusOperationId: guardedRunStatusOperationId, + runtimeEventCursorIncrement, + sessionId: input.sessionId, + state, + targetStatus: input.targetStatus, + terminalEventId, + terminalSemanticHash, + terminalSourceEventId, + }), + ]; + + try { + const results = await database.batch(statements); + + if (input.terminalEventFence === undefined && getD1ChangeCount(results[0]) === 0) { + const raced = await classifyAfterBatch(database, { + ...input, + event, + runStatusOperationId, + }); + + if (raced !== null) { + return raced; + } + + throw new Error("Atomic terminal run projection lost a concurrent session mutation."); + } + } catch (error) { + try { + const committed = await classifyAfterBatch(database, { + ...input, + event, + runStatusOperationId, + }); + + if (committed !== null) { + return committed; + } + } catch { + // Preserve the batch failure when its commit outcome cannot be proven. + } + + throw error; + } + + const committed = await classifyAfterBatch(database, { + ...input, + event, + runStatusOperationId, + }); + + if (committed === null) { + throw new Error("Atomic terminal run projection did not commit its canonical projection."); + } + + return committed; +} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/connections.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/connections.ts index 1c970e78..8dde0ee4 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/connections.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/connections.ts @@ -17,6 +17,7 @@ function isDriverInstanceConflictMessage(message: string): boolean { message === "Driver hello is required before pushLogs." || message === "Driver instance id does not match the active Durable Object." || message === "Driver instance id mismatch." || + message === "Driver generation is no longer current." || message.includes("already closed.") || message.includes("closed before hello.") ); diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/debug-resume-snapshot.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/debug-resume-snapshot.ts index 5a936c41..5bdb6213 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/debug-resume-snapshot.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/debug-resume-snapshot.ts @@ -3,7 +3,6 @@ import type { SandboxId } from "@mosoo/id"; export type DriverDebugRecoveryMode = "fresh" | "ready" | "disconnected" | "turn_interrupted"; export interface DriverDebugResumeSnapshot { - readonly lastEventSeq: number; readonly recoveryMode: DriverDebugRecoveryMode; readonly sandboxId: SandboxId | null; } @@ -11,9 +10,5 @@ export interface DriverDebugResumeSnapshot { export function createDriverDebugResumeSnapshot( input: DriverDebugResumeSnapshot, ): DriverDebugResumeSnapshot { - if (!Number.isInteger(input.lastEventSeq) || input.lastEventSeq < 0) { - throw new Error("Driver debug resume lastEventSeq must be a non-negative integer."); - } - return input; } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/do.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/do.ts index e956128d..2a43b68d 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/do.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/do.ts @@ -24,11 +24,14 @@ import { import { decodeAndHashBootToken } from "../runtime-boot-token"; import { toDriverInstanceRequestErrorStatus } from "./connections"; import { json, toErrorMessage } from "./driver-instance-support"; -import { claimDriverInstanceByBootTokenHash } from "./driver-instance-token.repository"; +import { + claimDriverInstanceByBootTokenHash, + validateDriverInstanceBootTokenHash, +} from "./driver-instance-token.repository"; import { runtimeSessionLinkNeedsRefresh } from "./event-types"; import { handleDriverInstanceRequest } from "./http"; import type { DriverInstanceHttpHandler } from "./http"; -import { getDriverInstanceStatus, markDriverInstanceConnected } from "./lifecycle"; +import { getDriverInstanceLifecycleIdentity, markDriverInstanceConnected } from "./lifecycle"; import { createDriverInstanceRpcContext } from "./rpc"; import type { DriverInstanceRpcContext } from "./rpc"; import { DriverInstanceRpcController } from "./rpc-controller"; @@ -39,8 +42,7 @@ import { SessionViewerEventDeliveryBuffer } from "./session-viewer-event-deliver import { DriverInstanceSocketRegistry } from "./sockets"; import type { DriverInstanceCloseSnapshot, - DriverInstanceHeartbeatResult, - DriverInstanceHelloResult, + DriverInstanceConnectionEpoch, DriverInstanceReadyResult, DriverInstanceSnapshot, DriverInstanceWaitForCloseResult, @@ -48,7 +50,8 @@ import type { import { DriverInstanceTerminalStateCoordinator } from "./terminal-state-coordinator"; export class DriverInstance extends DurableObject implements DriverInstanceHttpHandler { - #destroyed = false; + #destroyedGeneration: number | null = null; + #destroyTask: Promise | null = null; readonly #identity = new DurableObjectIdentity({ mismatchMessage: "Driver instance id does not match the active Durable Object.", requiredMessage: "Driver instance id is required.", @@ -87,7 +90,7 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH }); this.#rpcController = new DriverInstanceRpcController({ env, - finalizeTerminalState: async () => this.#terminalState.finalize(), + finalizeTerminalState: async (epoch) => this.#terminalState.finalize(epoch), sockets: this.#sockets, state: this.#state, viewCache: this.#viewCache, @@ -95,16 +98,40 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH waitUntil: (task) => this.ctx.waitUntil(task), withRuntimeLogContext: (fn) => this.#withRuntimeLogContext(fn), }); - void this.ctx.blockConcurrencyWhile(async () => this.#state.load()); + void this.ctx.blockConcurrencyWhile(async () => { + await this.#state.load(); + + if (this.#state.driverInstanceId !== null && this.#state.driverGeneration === null) { + const identity = await getDriverInstanceLifecycleIdentity( + this.env, + this.#state.driverInstanceId, + ); + + if (identity !== null) { + await this.#state.setDriverGeneration(identity.generation); + } + } + + if ( + (this.#state.terminalized || this.#state.errorMessage !== null) && + !this.#state.terminalCleanupComplete + ) { + const epoch = this.#state.connectionEpoch(); + + if (epoch !== null) { + await this.#finalizeTerminalState(epoch); + } + } + }); } override async fetch(request: Request): Promise { try { const url = new URL(request.url); - if (this.#destroyed) { + if (this.#destroyedGeneration !== null) { if (request.method === "POST" && url.pathname === "/control/destroy") { - return json({ ok: true }); + return handleDriverInstanceRequest(this, request); } return json({ error: "Driver instance Durable Object was destroyed." }, { status: 410 }); @@ -127,28 +154,28 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH } override async webSocketClose(ws: WebSocket, code: number, reason: string): Promise { - if (this.#destroyed) { + if (this.#destroyedGeneration !== null) { return; } this.#rpcHandler?.close(ws); + const epoch = this.#sockets.getSocketEpoch(ws); - // A socket replaced by a newer accepted connection must not finalize the - // state that now belongs to its successor. - if (this.#sockets.isSupersededDriverSocket(ws)) { - this.#sockets.releaseDriverSocket(ws); + if (epoch === null || !this.#isCurrentSocketEpoch(ws, epoch)) { return; } - this.#sockets.releaseDriverSocket(ws); - const close: DriverInstanceCloseSnapshot = { at: new Date().toISOString(), code, reason, }; - await this.#state.persistClose(close); + await this.#state.persistClose(close, epoch); + + if (!this.#isCurrentSocketEpoch(ws, epoch)) { + return; + } this.#withRuntimeLogContext(() => { logInfo("runtime.socket.closed", { @@ -157,48 +184,54 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH driverInstanceId: this.#state.driverInstanceId, }); }); - await this.#appendTransportWsDisconnectedEvent(close); + await this.#appendTransportWsDisconnectedEvent(close, epoch); - if (!this.#state.hello) { - this.#state.rejectHelloWaiters( - new Error(`Driver instance socket closed before hello: ${reason || code}.`), - ); + if (!this.#isCurrentSocketEpoch(ws, epoch)) { + return; } - await this.#terminalState.finalize(); + await this.#finalizeTerminalState(epoch); } override async webSocketMessage(ws: WebSocket, message: ArrayBuffer | string): Promise { - if (this.#destroyed) { + if (this.#destroyedGeneration !== null) { return; } - if (!this.#sockets.isActiveDriverSocket(ws)) { + const epoch = this.#sockets.getSocketEpoch(ws); + + if (epoch === null || !this.#isCurrentSocketEpoch(ws, epoch)) { return; } try { const rpcHandler = await this.#getRpcHandler(); - if (!this.#sockets.isActiveDriverSocket(ws)) { + if (!this.#isCurrentSocketEpoch(ws, epoch)) { return; } - const connectionId = this.#state.requireConnectionId(); await rpcHandler.message(ws, message, { context: createDriverInstanceRpcContext(this.#rpcController, { assertActiveConnection: () => { - if ( - this.#state.connectionId !== connectionId || - !this.#sockets.isActiveDriverSocket(ws) - ) { + if (!this.#isCurrentSocketEpoch(ws, epoch)) { throw new Error("Driver connection is no longer current."); } }, - connectionId, + connectionId: epoch.connectionId, + epoch, }), }); + + if (!this.#isCurrentSocketEpoch(ws, epoch)) { + this.#closeSocket(ws, 1000, "runtime.socket.superseded"); + } } catch (error) { + if (!this.#isCurrentSocketEpoch(ws, epoch)) { + this.#closeSocket(ws, 1000, "runtime.socket.superseded"); + return; + } + this.#withRuntimeLogContext(() => { logError("runtime.socket.message.failed", { ...createErrorLogContext(error), @@ -206,15 +239,27 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH }); }); - await this.#state.setErrorMessage( + await this.#state.setConnectionErrorMessage( + epoch, toErrorMessage(error, "Driver instance WebSocket message failed."), ); - await this.#appendTransportRpcErrorEvent(error); + + if (!this.#isCurrentSocketEpoch(ws, epoch)) { + this.#closeSocket(ws, 1000, "runtime.socket.superseded"); + return; + } + + await this.#appendTransportRpcErrorEvent(error, epoch); + + if (!this.#isCurrentSocketEpoch(ws, epoch)) { + this.#closeSocket(ws, 1000, "runtime.socket.superseded"); + return; + } if (ws.readyState === WebSocket.OPEN) { ws.close(1003, "runtime.invalid-message"); } else { - await this.#terminalState.finalize(); + await this.#finalizeTerminalState(epoch); } } } @@ -236,11 +281,14 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH } override async webSocketError(ws: WebSocket, _error: unknown): Promise { - if (this.#destroyed) { + if (this.#destroyedGeneration !== null) { return; } - if (this.#sockets.isSupersededDriverSocket(ws)) { + const epoch = this.#sockets.getSocketEpoch(ws); + + if (epoch === null || !this.#isCurrentSocketEpoch(ws, epoch)) { + this.#closeSocket(ws, 1000, "runtime.socket.superseded"); return; } @@ -250,15 +298,24 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH }); }); - await this.#state.setErrorMessage("Driver instance WebSocket error."); - await this.#appendTransportRpcErrorEvent("Driver instance WebSocket error."); + await this.#state.setConnectionErrorMessage(epoch, "Driver instance WebSocket error."); + + if (!this.#isCurrentSocketEpoch(ws, epoch)) { + this.#closeSocket(ws, 1000, "runtime.socket.superseded"); + return; + } + + await this.#appendTransportRpcErrorEvent("Driver instance WebSocket error.", epoch); - const socket = this.#sockets.getDriverSocket(); + if (!this.#isCurrentSocketEpoch(ws, epoch)) { + this.#closeSocket(ws, 1000, "runtime.socket.superseded"); + return; + } - if (socket && socket.readyState === WebSocket.OPEN) { - socket.close(1011, "runtime.socket.error"); + if (ws.readyState === WebSocket.OPEN) { + ws.close(1011, "runtime.socket.error"); } else { - await this.#terminalState.finalize(); + await this.#finalizeTerminalState(epoch); } } @@ -282,59 +339,97 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH return json({ error: "Boot token is invalid." }, { status: 401 }); } - if (this.#state.terminalized) { - await this.#terminalState.resetForReuse(); - } + return this.ctx.blockConcurrencyWhile(async () => { + const validation = await validateDriverInstanceBootTokenHash(this.env, bootTokenHash); + const driverInstanceId = this.#state.requireDriverInstanceId(); - const claim = await claimDriverInstanceByBootTokenHash(this.env, bootTokenHash); + if ( + validation.driverInstanceId === null || + validation.generation === null || + validation.driverInstanceId !== driverInstanceId + ) { + return json({ error: validation.error ?? "Boot token is invalid." }, { status: 401 }); + } - if ( - claim.driverInstanceId === null || - claim.generation === null || - claim.driverInstanceId !== this.#state.requireDriverInstanceId() - ) { - return json({ error: claim.error ?? "Boot token is invalid." }, { status: 401 }); - } + let shouldReset = this.#state.terminalized; - const connectedAt = Date.now(); - const connectionId = createPlatformId(); - const connected = await markDriverInstanceConnected(this.env, { - bootTokenHash, - connectedAt, - connectionId, - driverInstanceId: this.#state.requireDriverInstanceId(), - generation: claim.generation, - }); + if (shouldReset) { + if (!this.#state.terminalCleanupComplete) { + const epoch = this.#state.connectionEpoch(); - if (!connected) { - return json({ error: "Driver connection is no longer current." }, { status: 409 }); - } + if (epoch !== null) { + await this.#finalizeTerminalState(epoch); + } + } - const traceparent = url.searchParams.get("traceparent"); - const parsedTraceparent = isTruthy(traceparent) ? parseTraceparent(traceparent) : null; - const pair = new WebSocketPair(); - const [clientSocket, serverSocket] = [pair[0], pair[1]]; + if (validation.generation <= this.#state.requireDriverGeneration()) { + return json({ error: "Driver generation is no longer current." }, { status: 409 }); + } - this.#sockets.replaceDriverSockets(); - this.#sockets.acceptDriverSocket(serverSocket); + await this.#terminalState.prepareForReuse(); + } else if (this.#state.requireDriverGeneration() !== validation.generation) { + if (this.#state.connectionId !== null) { + return json({ error: "Driver generation is no longer current." }, { status: 409 }); + } - await this.#state.recordAcceptedConnection({ - connectedAt, - connectionId, - driverGeneration: claim.generation, - traceId: parsedTraceparent?.traceId ?? null, - }); + shouldReset = true; + await this.#terminalState.prepareForReuse(); + } - this.#withRuntimeLogContext(() => { - logInfo("runtime.socket.accepted", { + const claim = await claimDriverInstanceByBootTokenHash(this.env, bootTokenHash); + + if ( + claim.driverInstanceId !== driverInstanceId || + claim.generation !== validation.generation + ) { + return json({ error: claim.error ?? "Boot token is invalid." }, { status: 401 }); + } + + if (shouldReset) { + await this.#terminalState.resetForReuse(validation.generation); + } + + const connectedAt = Date.now(); + const connectionId = createPlatformId(); + const connected = await markDriverInstanceConnected(this.env, { + bootTokenHash, + connectedAt, connectionId, - driverInstanceId: this.#state.requireDriverInstanceId(), + driverInstanceId, + generation: validation.generation, }); - }); - return new Response(null, { - status: 101, - webSocket: clientSocket, + if (!connected) { + return json({ error: "Driver connection is no longer current." }, { status: 409 }); + } + + const traceparent = url.searchParams.get("traceparent"); + const parsedTraceparent = isTruthy(traceparent) ? parseTraceparent(traceparent) : null; + const pair = new WebSocketPair(); + const [clientSocket, serverSocket] = [pair[0], pair[1]]; + const epoch = { connectionId, generation: validation.generation }; + + this.#sockets.replaceDriverSockets(); + this.#sockets.acceptDriverSocket(serverSocket, epoch); + + await this.#state.recordAcceptedConnection({ + connectedAt, + connectionId, + driverGeneration: validation.generation, + traceId: parsedTraceparent?.traceId ?? null, + }); + + this.#withRuntimeLogContext(() => { + logInfo("runtime.socket.accepted", { + connectionId, + driverInstanceId, + }); + }); + + return new Response(null, { + status: 101, + webSocket: clientSocket, + }); }); } @@ -346,12 +441,17 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH this.#identity.ensure(candidate); } - if (this.#state.terminalized) { - const status = await getDriverInstanceStatus(this.env, this.#state.driverInstanceId); + if (this.#state.driverGeneration === null) { + const identity = await getDriverInstanceLifecycleIdentity( + this.env, + this.#state.driverInstanceId, + ); - if (status === "provisioning" || status === "connecting" || status === "ready") { - await this.#terminalState.resetForReuse(); + if (identity === null) { + throw new Error("Driver instance record was not found."); } + + await this.#state.setDriverGeneration(identity.generation); } return this.#state.driverInstanceId; @@ -361,77 +461,117 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH this.#identity.ensure(candidate), "driver instance id", ); - await this.#state.setDriverInstanceId(driverInstanceId); + const identity = await getDriverInstanceLifecycleIdentity(this.env, driverInstanceId); + + if (identity === null) { + throw new Error("Driver instance record was not found."); + } + + await this.#state.initializeDriverInstance(driverInstanceId, identity.generation); return driverInstanceId; } - async sendControlCommand(command: RuntimeCommand): Promise { - const socket = this.#sockets.getDriverSocket(); - - if (!socket || socket.readyState !== WebSocket.OPEN) { - const message = "Runtime driver control socket is not connected."; - await this.#state.setErrorMessage(message); - await this.#terminalState.finalize(); - throw new Error(message); - } + async sendControlCommand(generation: number, command: RuntimeCommand): Promise { + return this.ctx.blockConcurrencyWhile(async () => { + await this.#assertCurrentGeneration(generation); + const epoch = this.#state.requireConnectionEpoch(); + const socket = this.#sockets.getDriverSocket(epoch); + + if (!socket || socket.readyState !== WebSocket.OPEN) { + const message = "Runtime driver control socket is not connected."; + await this.#state.setConnectionErrorMessage(epoch, message); + await this.#finalizeTerminalState(epoch); + throw new Error(message); + } - await this.#rpcController.enqueueCommand(command); + await this.#rpcController.enqueueCommand(generation, command); + }); } - async destroy(reason: string): Promise { - if (this.#destroyed) { - return; - } + async destroy(generation: number, reason: string): Promise { + return this.ctx.blockConcurrencyWhile(async () => { + if (this.#destroyedGeneration !== null) { + if (this.#destroyedGeneration !== generation) { + throw new Error("Driver generation is no longer current."); + } - this.#destroyed = true; - this.#identity.clear(); - const socket = this.#sockets.getDriverSocket(); + return; + } + + await this.#assertCurrentGeneration(generation); + + return (this.#destroyTask ??= this.#destroy(generation, reason).finally(() => { + if (this.#destroyedGeneration === null) { + this.#destroyTask = null; + } + })); + }); + } + + async #destroy(generation: number, reason: string): Promise { + const socket = this.#sockets.getDriverSocket(this.#state.connectionEpoch()); if (socket?.readyState === WebSocket.OPEN) { socket.close(1000, reason); } - await this.#terminalState.destroy(reason); + await this.#rpcController.runAfterPendingEvents(() => this.#terminalState.destroy(reason)); + this.#identity.clear(); + this.#destroyedGeneration = generation; } - async fail(message: string): Promise { - await this.#state.setErrorMessage(message); + async fail(generation: number, message: string): Promise { + return this.ctx.blockConcurrencyWhile(async () => { + await this.#assertCurrentGeneration(generation); + const epoch = this.#state.requireConnectionEpoch(); + await this.#state.setConnectionErrorMessage(epoch, message); - const socket = this.#sockets.getDriverSocket(); + const socket = this.#sockets.getDriverSocket(epoch); - if (socket && socket.readyState === WebSocket.OPEN) { - socket.close(1011, "runtime.failed"); - return; - } + if (socket && socket.readyState === WebSocket.OPEN) { + socket.close(1011, "runtime.failed"); + return; + } - await this.#terminalState.finalize(); + await this.#finalizeTerminalState(epoch); + }); } - snapshot(): DriverInstanceSnapshot { - const socket = this.#sockets.getDriverSocket(); - return this.#state.snapshot(Boolean(socket && socket.readyState === WebSocket.OPEN)); + async #finalizeTerminalState(epoch: DriverInstanceConnectionEpoch): Promise { + await this.#rpcController.runAfterPendingEvents(() => this.#terminalState.finalize(epoch)); } - async waitForClose(timeoutMs: number): Promise { - return this.#state.waitForClose(timeoutMs); + async #assertCurrentGeneration(generation: number): Promise { + if (this.#state.requireDriverGeneration() !== generation) { + throw new Error("Driver generation is no longer current."); + } + + const driverInstanceId = this.#state.requireDriverInstanceId(); + const identity = await getDriverInstanceLifecycleIdentity(this.env, driverInstanceId); + + if (identity === null || identity.generation !== generation) { + throw new Error("Driver generation is no longer current."); + } } - async waitForHeartbeat( - afterCount: number, - timeoutMs: number, - ): Promise { - return this.#state.waitForHeartbeat(afterCount, timeoutMs); + snapshot(): DriverInstanceSnapshot { + const socket = this.#sockets.getDriverSocket(this.#state.connectionEpoch()); + return this.#state.snapshot(Boolean(socket && socket.readyState === WebSocket.OPEN)); } - async waitForHello(timeoutMs: number): Promise { - return this.#state.waitForHello(timeoutMs); + async waitForClose( + generation: number, + timeoutMs: number, + ): Promise { + return this.#state.waitForClose(generation, timeoutMs); } - async waitForReady(timeoutMs: number): Promise { - return this.#state.waitForReady(timeoutMs); + async waitForReady(generation: number, timeoutMs: number): Promise { + return this.#state.waitForReady(generation, timeoutMs); } - async #getRuntimeSessionLink() { + async #getRuntimeSessionLink(epoch: DriverInstanceConnectionEpoch) { + this.#state.assertConnectionEpoch(epoch); const existing = this.#state.runtimeSessionLink; if (existing !== null && !runtimeSessionLinkNeedsRefresh(existing)) { @@ -439,13 +579,18 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH } const link = await getRuntimeSessionLink(this.env.DB, this.#state.requireDriverInstanceId()); + this.#state.assertConnectionEpoch(epoch); this.#state.setRuntimeSessionLink(link); return link; } - async #appendTransportRpcErrorEvent(error: unknown): Promise { + async #appendTransportRpcErrorEvent( + error: unknown, + epoch: DriverInstanceConnectionEpoch, + ): Promise { try { - const link = await this.#getRuntimeSessionLink(); + const link = await this.#getRuntimeSessionLink(epoch); + this.#state.assertConnectionEpoch(epoch); if (!isTruthy(link.agentId) || !isTruthy(link.sessionId)) { return; @@ -475,9 +620,13 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH } } - async #appendTransportWsDisconnectedEvent(close: DriverInstanceCloseSnapshot): Promise { + async #appendTransportWsDisconnectedEvent( + close: DriverInstanceCloseSnapshot, + epoch: DriverInstanceConnectionEpoch, + ): Promise { try { - const link = await this.#getRuntimeSessionLink(); + const link = await this.#getRuntimeSessionLink(epoch); + this.#state.assertConnectionEpoch(epoch); if (!isTruthy(link.agentId) || !isTruthy(link.sessionId)) { return; @@ -507,6 +656,16 @@ export class DriverInstance extends DurableObject implements DriverInstanceHttpH } } + #isCurrentSocketEpoch(ws: WebSocket, epoch: DriverInstanceConnectionEpoch): boolean { + return this.#sockets.isCurrentDriverSocket(ws, epoch, this.#state.connectionEpoch()); + } + + #closeSocket(ws: WebSocket, code: number, reason: string): void { + if (ws.readyState === WebSocket.OPEN) { + ws.close(code, reason); + } + } + #withRuntimeLogContext(fn: () => T): T { return runWithApiLogContext( { diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-event-canonicalization.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-event-canonicalization.ts new file mode 100644 index 00000000..a2098107 --- /dev/null +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-event-canonicalization.ts @@ -0,0 +1,127 @@ +import { DurableRunError } from "@mosoo/contracts/session-run"; +import { parseSchemaValue } from "@mosoo/contracts/validation"; +import type { DriverInstanceId, RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; +import { + createRuntimeEvent, + parseRuntimeEventEnvelope, + readRuntimeEventPayload, +} from "@mosoo/runtime-events"; +import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; + +import { createSessionRunTerminalSourceId } from "../../domain/session-run-terminal-event-id"; +import type { CanonicalDriverEventEnvelope } from "./event-types"; + +interface DriverEventEnvelopeInput { + readonly event: unknown; + readonly eventId: string; + readonly occurredAt?: string | null | undefined; +} + +function isTerminalRunEventKind( + kind: RuntimeEventEnvelope["kind"], +): kind is "run.cancelled" | "run.completed" | "run.failed" { + return kind === "run.cancelled" || kind === "run.completed" || kind === "run.failed"; +} + +export function createCanonicalDriverRunFailedEvent(input: { + driverInstanceId: DriverInstanceId; + error: typeof DurableRunError.infer; + id: RuntimeEventId; + occurredAt: string; + runId: SessionRunId; + runtimeId: string; + sessionId: SessionId; + traceId: string; +}): RuntimeEventEnvelope { + return createRuntimeEvent({ + actor: "driver", + delivery: "lossless", + driverInstanceId: input.driverInstanceId, + id: input.id, + kind: "run.failed", + occurredAt: input.occurredAt, + origin: "driver", + payload: { + error: input.error, + lifecycle: "IDLE", + recoverable: input.error.retryable, + }, + runId: input.runId, + runtimeId: input.runtimeId, + sessionId: input.sessionId, + sourceEventId: createSessionRunTerminalSourceId(input.runId, "run.failed"), + traceId: input.traceId, + visibility: "participant", + }); +} + +export function canonicalizeDriverEventEnvelope( + envelope: DriverEventEnvelopeInput, + input: { traceId: string | null }, +): CanonicalDriverEventEnvelope { + const parsedEvent = parseRuntimeEventEnvelope(envelope.event); + + if (parsedEvent.sourceEventId !== undefined && parsedEvent.sourceEventId !== envelope.eventId) { + throw new Error("Runtime driver event source id does not match the driver envelope."); + } + + const event = (() => { + if (parsedEvent.runId === undefined) { + return parsedEvent; + } + + if (input.traceId === null) { + throw new Error("Runtime driver event run is missing its authoritative trace identity."); + } + + if (parsedEvent.traceId !== undefined && parsedEvent.traceId !== input.traceId) { + throw new Error("Runtime driver event trace id does not match its Session Run."); + } + + return parsedEvent.traceId === undefined + ? { ...parsedEvent, traceId: input.traceId } + : parsedEvent; + })(); + + if (!isTerminalRunEventKind(event.kind) || event.runId === undefined) { + return { ...envelope, event }; + } + + if (event.kind !== "run.failed") { + return { + ...envelope, + event: { + ...event, + payload: { + ...readRuntimeEventPayload(event), + lifecycle: "IDLE", + }, + sourceEventId: createSessionRunTerminalSourceId(event.runId, event.kind), + }, + }; + } + + if ( + event.driverInstanceId === undefined || + event.runtimeId === undefined || + event.traceId === undefined + ) { + throw new Error("Runtime driver failure event is missing its canonical execution identity."); + } + + const error = parseSchemaValue(DurableRunError, readRuntimeEventPayload(event)["error"]); + + return { + ...envelope, + event: createCanonicalDriverRunFailedEvent({ + driverInstanceId: event.driverInstanceId, + error, + id: event.id, + occurredAt: event.occurredAt, + runId: event.runId, + runtimeId: event.runtimeId, + sessionId: event.sessionId, + traceId: event.traceId, + }), + }; +} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-event-receipts.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-event-receipts.ts deleted file mode 100644 index 014184ce..00000000 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-event-receipts.ts +++ /dev/null @@ -1,90 +0,0 @@ -import type { DriverEventEnvelope } from "@mosoo/agent-driver/events"; -import type { DriverEventReceipt } from "@mosoo/agent-driver/orpc"; - -const MAX_PROCESSED_DRIVER_EVENT_RECEIPTS = 8192; - -export function filterNewDriverEvents(input: { - events: readonly DriverEventEnvelope[]; - processedReceipts: Map; -}): DriverEventEnvelope[] { - const seenEventIds = new Set(); - - return input.events.filter((envelope) => { - if (envelope.eventId.length === 0) { - return true; - } - - if (seenEventIds.has(envelope.eventId)) { - return false; - } - - seenEventIds.add(envelope.eventId); - - return !input.processedReceipts.has(envelope.eventId); - }); -} - -export function createReceiptsForDriverEvents(input: { - events: readonly DriverEventEnvelope[]; - nextSeq: number; -}): { - nextSeq: number; - receipts: DriverEventReceipt[]; -} { - let nextSeq = input.nextSeq; - const receipts = input.events.map((envelope) => { - nextSeq += 1; - - return { - ...(envelope.eventId.length > 0 ? { eventId: envelope.eventId } : {}), - seq: nextSeq, - type: envelope.event.kind, - }; - }); - - return { nextSeq, receipts }; -} - -export function readReceiptsForProcessedDriverEvents(input: { - events: readonly DriverEventEnvelope[]; - processedReceipts: Map; -}): DriverEventReceipt[] { - const receipts: DriverEventReceipt[] = []; - - for (const envelope of input.events) { - if (envelope.eventId.length === 0) { - continue; - } - - const receipt = input.processedReceipts.get(envelope.eventId); - - if (receipt !== undefined) { - receipts.push(receipt); - } - } - - return receipts; -} - -export function rememberDriverEventReceipts(input: { - processedReceipts: Map; - receipts: DriverEventReceipt[]; -}): void { - for (const receipt of input.receipts) { - if (typeof receipt.eventId !== "string" || receipt.eventId.length === 0) { - continue; - } - - input.processedReceipts.set(receipt.eventId, receipt); - } - - while (input.processedReceipts.size > MAX_PROCESSED_DRIVER_EVENT_RECEIPTS) { - const oldest = input.processedReceipts.keys().next().value; - - if (typeof oldest !== "string") { - return; - } - - input.processedReceipts.delete(oldest); - } -} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-instance-record.repository.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-instance-record.repository.ts index 900438ae..6add8cbf 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-instance-record.repository.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-instance-record.repository.ts @@ -1,12 +1,37 @@ import { DRIVER_PROTOCOL_VERSION } from "@mosoo/agent-driver/boot"; import type { DriverRuntime } from "@mosoo/agent-driver/runtime"; -import { driverCommandsTable, driverInstanceMcpGrantsTable, driverInstancesTable } from "@mosoo/db"; +import { + driverCommandsTable, + driverInstanceMcpGrantsTable, + driverInstancesTable, + externalToolEffectsTable, + sandboxesTable, + sandboxSessionsTable, + sessionRunsTable, + sessionsTable, +} from "@mosoo/db"; import type { DriverInstanceId, SandboxId, SessionId } from "@mosoo/id"; -import { and, desc, eq, gt, inArray, isNotNull, notInArray, or, sql } from "drizzle-orm"; +import { + and, + desc, + eq, + exists, + gt, + inArray, + isNotNull, + isNull, + notExists, + or, + sql, +} from "drizzle-orm"; import type { SQL } from "drizzle-orm"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; -import { getAppDatabase, runAppDatabaseBatch } from "../../../../platform/db/drizzle"; +import { + getAppDatabase, + getD1ChangeCount, + runAppDatabaseBatch, +} from "../../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../../time"; import { REUSABLE_DRIVER_INSTANCE_STATUSES, @@ -18,6 +43,8 @@ import { DRIVER_COLD_READY_TIMEOUT_MS, RUNTIME_SOCKET_TIMEOUT_MS, } from "../../domain/runtime-config"; +import { ACTIVE_SESSION_RUN_STATUSES } from "../../domain/session-run-lifecycle.machine"; +import type { RuntimeRunProvisioningLease } from "../runtime-subject-lifecycle/runtime-provisioning-lease-store"; import type { DriverInstanceMcpGrantRecord } from "./mcp-grants.repository"; import { driverInstanceExpiresAt } from "./status"; import type { DriverInstanceStatus } from "./status"; @@ -35,6 +62,186 @@ export type CreateDriverInstanceRecordResult = status: "skipped"; }; +function selectedValue(value: Value, alias: string) { + return sql`${value}`.as(alias); +} + +async function insertProvisioningOwnedDriverRecord( + database: D1Database, + input: { + readonly driverRecord: typeof driverInstancesTable.$inferInsert; + readonly lease: RuntimeRunProvisioningLease; + }, +): Promise<{ bootTokenExpiresAt: number; generation: number } | null> { + const { driverRecord: record, lease } = input; + if ( + lease.sandboxIncarnation === null || + lease.sandboxSessionId === null || + record.bootTokenExpiresAt === undefined || + record.generation === undefined || + record.sandboxId !== lease.sandboxId || + record.sandboxIncarnation !== lease.sandboxIncarnation || + record.sandboxSessionId !== lease.sessionId + ) { + throw new Error("Driver provisioning requires a complete immutable sandbox target."); + } + + const inserted = await database + .prepare( + ` + INSERT INTO driver_instance ( + boot_token_expires_at, boot_token_hash, created_at, expires_at, + generation, heartbeat_count, id, protocol, protocol_version, + restart_count, runtime, sandbox_id, sandbox_incarnation, + sandbox_session_id, status, status_changed_at, status_event, + status_seq, status_source, updated_at + ) + SELECT ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? + FROM session AS provisioning + INNER JOIN session_run AS run + ON run.id = provisioning.runtime_provisioning_run_id + AND run.session_id = provisioning.id + INNER JOIN sandbox AS subject + ON subject.id = provisioning.runtime_provisioning_sandbox_id + INNER JOIN sandbox_session AS conversation + ON conversation.session_id = provisioning.id + AND conversation.sandbox_id = subject.id + WHERE provisioning.id = ? + AND provisioning.runtime_provisioning_operation_id = ? + AND provisioning.runtime_provisioning_run_id = ? + AND provisioning.runtime_provisioning_sandbox_id = ? + AND provisioning.runtime_provisioning_sandbox_incarnation = ? + AND provisioning.runtime_provisioning_sandbox_session_id = ? + AND provisioning.last_run_id = run.id + AND provisioning.status = 'RUNNING' + AND provisioning.archived_at IS NULL + AND provisioning.cleanup_operation_kind IS NULL + AND provisioning.status_operation_id IS NULL + AND run.status IN ('queued', 'booting', 'running', 'waiting_input') + AND subject.incarnation = ? + AND subject.status = 'active' + AND subject.claim_owner IS NULL + AND subject.operation_kind IS NULL + AND subject.status_operation_id IS NULL + AND conversation.sandbox_incarnation = ? + AND conversation.cloudflare_session_id = ? + AND conversation.status = 'active' + ON CONFLICT DO NOTHING + `, + ) + .bind( + record.bootTokenExpiresAt, + record.bootTokenHash, + record.createdAt, + record.expiresAt, + record.generation, + record.heartbeatCount, + record.id, + record.protocol, + record.protocolVersion, + record.restartCount, + record.runtime, + record.sandboxId, + record.sandboxIncarnation, + record.sandboxSessionId, + record.status, + record.statusChangedAt, + record.statusEvent, + record.statusSeq, + record.statusSource, + record.updatedAt, + lease.sessionId, + lease.operationId, + lease.runId, + lease.sandboxId, + lease.sandboxIncarnation, + lease.sandboxSessionId, + lease.sandboxIncarnation, + lease.sandboxIncarnation, + lease.sandboxSessionId, + ) + .run(); + + return getD1ChangeCount(inserted) === 1 + ? { bootTokenExpiresAt: record.bootTokenExpiresAt, generation: record.generation } + : null; +} + +export async function runtimeProvisioningDriverLaunchIsOwned( + database: D1Database, + input: { + readonly bootTokenHash: Uint8Array; + readonly driverGeneration: number; + readonly driverInstanceId: DriverInstanceId; + readonly lease: RuntimeRunProvisioningLease; + }, +): Promise { + const { lease } = input; + if (lease.sandboxIncarnation === null || lease.sandboxSessionId === null) { + return false; + } + + const row = await getAppDatabase(database) + .select({ id: driverInstancesTable.id }) + .from(sessionsTable) + .innerJoin( + sessionRunsTable, + and( + eq(sessionRunsTable.id, sessionsTable.runtimeProvisioningRunId), + eq(sessionRunsTable.sessionId, sessionsTable.id), + ), + ) + .innerJoin(sandboxesTable, eq(sandboxesTable.id, sessionsTable.runtimeProvisioningSandboxId)) + .innerJoin( + sandboxSessionsTable, + and( + eq(sandboxSessionsTable.sessionId, sessionsTable.id), + eq(sandboxSessionsTable.sandboxId, sandboxesTable.id), + ), + ) + .innerJoin( + driverInstancesTable, + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.sandboxSessionId, sessionsTable.id), + ), + ) + .where( + and( + eq(sessionsTable.id, lease.sessionId), + eq(sessionsTable.runtimeProvisioningOperationId, lease.operationId), + eq(sessionsTable.runtimeProvisioningRunId, lease.runId), + eq(sessionsTable.runtimeProvisioningSandboxId, lease.sandboxId), + eq(sessionsTable.runtimeProvisioningSandboxIncarnation, lease.sandboxIncarnation), + eq(sessionsTable.runtimeProvisioningSandboxSessionId, lease.sandboxSessionId), + eq(sessionsTable.lastRunId, lease.runId), + eq(sessionsTable.status, "RUNNING"), + isNull(sessionsTable.archivedAt), + isNull(sessionsTable.cleanupOperationKind), + isNull(sessionsTable.statusOperationId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + eq(sandboxesTable.incarnation, lease.sandboxIncarnation), + eq(sandboxesTable.status, "active"), + isNull(sandboxesTable.claimOwner), + isNull(sandboxesTable.operationKind), + isNull(sandboxesTable.statusOperationId), + eq(sandboxSessionsTable.sandboxIncarnation, lease.sandboxIncarnation), + eq(sandboxSessionsTable.sandboxSessionId, lease.sandboxSessionId), + eq(sandboxSessionsTable.status, "active"), + eq(driverInstancesTable.generation, input.driverGeneration), + eq(driverInstancesTable.bootTokenHash, input.bootTokenHash), + eq(driverInstancesTable.sandboxId, lease.sandboxId), + eq(driverInstancesTable.sandboxIncarnation, lease.sandboxIncarnation), + inArray(driverInstancesTable.status, REUSABLE_DRIVER_INSTANCE_STATUSES), + isNull(driverInstancesTable.statusOperationId), + ), + ) + .limit(1) + .get(); + + return row !== undefined; +} + export async function createDriverInstanceRecord( bindings: ApiBindings, input: { @@ -43,8 +250,10 @@ export async function createDriverInstanceRecord( driverInstanceId: DriverInstanceId; runtime: DriverRuntime; sandboxId: SandboxId; + sandboxIncarnation: number; sandboxSessionId: SessionId; mcpGrants?: DriverInstanceMcpGrantRecord[]; + runtimeProvisioningLease?: RuntimeRunProvisioningLease; }, ): Promise { const now = currentTimestampMs(); @@ -84,6 +293,7 @@ export async function createDriverInstanceRecord( restartCount: 0, runtime: input.runtime, sandboxId: input.sandboxId, + sandboxIncarnation: input.sandboxIncarnation, sandboxSessionId: input.sandboxSessionId, status: "provisioning", statusChangedAt: now, @@ -96,16 +306,20 @@ export async function createDriverInstanceRecord( if (input.conflictStrategy === "insert-only") { const database = getAppDatabase(bindings.DB); - const inserted = - (await database - .insert(driverInstancesTable) - .values(driverRecord) - .onConflictDoNothing() - .returning({ - bootTokenExpiresAt: driverInstancesTable.bootTokenExpiresAt, - generation: driverInstancesTable.generation, + const inserted = input.runtimeProvisioningLease + ? await insertProvisioningOwnedDriverRecord(bindings.DB, { + driverRecord, + lease: input.runtimeProvisioningLease, }) - .get()) ?? null; + : ((await database + .insert(driverInstancesTable) + .values(driverRecord) + .onConflictDoNothing() + .returning({ + bootTokenExpiresAt: driverInstancesTable.bootTokenExpiresAt, + generation: driverInstancesTable.generation, + }) + .get()) ?? null); if (inserted === null) { return { @@ -127,65 +341,141 @@ export async function createDriverInstanceRecord( }; } + const [replacement] = await runAppDatabaseBatch(bindings.DB, (batchDb) => { + const acceptedCommand = batchDb + .select({ id: driverCommandsTable.id }) + .from(driverCommandsTable) + .where( + and( + eq(driverCommandsTable.id, externalToolEffectsTable.commandId), + eq(driverCommandsTable.status, "accepted"), + ), + ); + const protectedEffect = batchDb + .select({ id: externalToolEffectsTable.id }) + .from(externalToolEffectsTable) + .where( + and( + eq(externalToolEffectsTable.driverInstanceId, input.driverInstanceId), + or( + inArray(externalToolEffectsTable.status, ["claimed", "unknown"]), + exists(acceptedCommand), + ), + ), + ); + const replacementAllowed = and(sql`status_operation_id IS NULL`, notExists(protectedEffect))!; + const replacementDriver = and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.bootTokenHash, input.bootTokenHash), + ); + const replacementCommitted = exists( + batchDb + .select({ id: driverInstancesTable.id }) + .from(driverInstancesTable) + .where(replacementDriver), + ); + + return [ + batchDb + .insert(driverInstancesTable) + .values(driverRecord) + .onConflictDoUpdate({ + set: { + bootTokenExpiresAt: sql`excluded.boot_token_expires_at`, + bootTokenHash: sql`excluded.boot_token_hash`, + bootTokenUsedAt: null, + closeCode: null, + closeReason: null, + connectionId: null, + createdAt: sql`excluded.created_at`, + driverPid: null, + driverStartedAt: null, + driverVersion: null, + errorMessage: null, + expiresAt: sql`excluded.expires_at`, + generation: sql`${driverInstancesTable.generation} + 1`, + heartbeatCount: 0, + lastHeartbeatAt: null, + processId: null, + protocol: sql`excluded.protocol`, + protocolVersion: sql`excluded.protocol_version`, + restartCount: sql`${driverInstancesTable.restartCount} + 1`, + runtime: sql`excluded.runtime`, + sandboxId: sql`excluded.sandbox_id`, + sandboxIncarnation: sql`excluded.sandbox_incarnation`, + sandboxSessionId: sql`excluded.sandbox_session_id`, + status: sql`excluded.status`, + statusChangedAt: sql`excluded.status_changed_at`, + statusEvent: sql`excluded.status_event`, + statusOperationId: null, + statusSeq: sql`${driverInstancesTable.statusSeq} + 1`, + statusSource: sql`excluded.status_source`, + updatedAt: sql`excluded.updated_at`, + }, + setWhere: replacementAllowed, + target: driverInstancesTable.id, + }), + batchDb + .delete(driverCommandsTable) + .where( + and( + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + replacementCommitted, + ), + ), + batchDb + .delete(driverInstanceMcpGrantsTable) + .where( + and( + eq(driverInstanceMcpGrantsTable.driverInstanceId, input.driverInstanceId), + replacementCommitted, + ), + ), + ...mcpGrantRows.map((grant) => + batchDb.insert(driverInstanceMcpGrantsTable).select( + batchDb + .select({ + authType: selectedValue(grant.authType, "auth_type"), + authorizationState: selectedValue(grant.authorizationState, "authorization_state"), + canInvalidate: selectedValue(grant.canInvalidate, "can_invalidate"), + canRefresh: selectedValue(grant.canRefresh, "can_refresh"), + createdAt: selectedValue(grant.createdAt, "created_at"), + credentialId: selectedValue(grant.credentialId, "credential_id"), + driverInstanceId: driverInstancesTable.id, + appId: selectedValue(grant.appId, "app_id"), + serverId: selectedValue(grant.serverId, "server_id"), + updatedAt: selectedValue(grant.updatedAt, "updated_at"), + }) + .from(driverInstancesTable) + .where(replacementDriver), + ), + ), + ]; + }); + + if (getD1ChangeCount(replacement) === 0) { + throw new Error("Driver instance replacement is blocked by a protected external effect."); + } + const database = getAppDatabase(bindings.DB); - await runAppDatabaseBatch(bindings.DB, (batchDb) => [ - batchDb - .delete(driverCommandsTable) - .where(eq(driverCommandsTable.driverInstanceId, input.driverInstanceId)), - batchDb - .delete(driverInstanceMcpGrantsTable) - .where(eq(driverInstanceMcpGrantsTable.driverInstanceId, input.driverInstanceId)), - ]); const upserted = (await database - .insert(driverInstancesTable) - .values(driverRecord) - .onConflictDoUpdate({ - set: { - bootTokenExpiresAt: sql`excluded.boot_token_expires_at`, - bootTokenHash: sql`excluded.boot_token_hash`, - bootTokenUsedAt: null, - closeCode: null, - closeReason: null, - connectionId: null, - createdAt: sql`excluded.created_at`, - driverPid: null, - driverStartedAt: null, - driverVersion: null, - errorMessage: null, - expiresAt: sql`excluded.expires_at`, - generation: sql`${driverInstancesTable.generation} + 1`, - heartbeatCount: 0, - lastHeartbeatAt: null, - processId: null, - protocol: sql`excluded.protocol`, - protocolVersion: sql`excluded.protocol_version`, - restartCount: sql`${driverInstancesTable.restartCount} + 1`, - runtime: sql`excluded.runtime`, - sandboxId: sql`excluded.sandbox_id`, - sandboxSessionId: sql`excluded.sandbox_session_id`, - status: sql`excluded.status`, - statusChangedAt: sql`excluded.status_changed_at`, - statusEvent: sql`excluded.status_event`, - statusOperationId: null, - statusSeq: sql`${driverInstancesTable.statusSeq} + 1`, - statusSource: sql`excluded.status_source`, - updatedAt: sql`excluded.updated_at`, - }, - target: driverInstancesTable.id, - }) - .returning({ + .select({ bootTokenExpiresAt: driverInstancesTable.bootTokenExpiresAt, generation: driverInstancesTable.generation, }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.bootTokenHash, input.bootTokenHash), + ), + ) + .limit(1) .get()) ?? null; - if (mcpGrantRows.length > 0) { - await database.insert(driverInstanceMcpGrantsTable).values(mcpGrantRows).run(); - } - if (upserted === null) { - throw new Error("Driver instance record was not created."); + throw new Error("Driver instance record was replaced before provisioning could claim it."); } return { @@ -229,6 +519,7 @@ export async function getDriverInstanceRecord( ): Promise<{ generation: number; sandboxId: SandboxId; + sandboxIncarnation: number; sandboxSessionId: SessionId; status: DriverInstanceStatus; } | null> { @@ -237,6 +528,7 @@ export async function getDriverInstanceRecord( .select({ generation: driverInstancesTable.generation, sandboxId: driverInstancesTable.sandboxId, + sandboxIncarnation: driverInstancesTable.sandboxIncarnation, sandboxSessionId: driverInstancesTable.sandboxSessionId, status: driverInstancesTable.status, }) @@ -303,6 +595,7 @@ export async function getReusableDriverInstanceRecord( database: D1Database, input: { sandboxId: SandboxId; + sandboxIncarnation: number; sandboxSessionId: SessionId; }, ): Promise<{ @@ -321,6 +614,7 @@ export async function getReusableDriverInstanceRecord( .where( and( eq(driverInstancesTable.sandboxId, input.sandboxId), + eq(driverInstancesTable.sandboxIncarnation, input.sandboxIncarnation), eq(driverInstancesTable.sandboxSessionId, input.sandboxSessionId), inArray(driverInstancesTable.status, REUSABLE_DRIVER_INSTANCE_STATUSES), ), @@ -343,7 +637,8 @@ export async function recordRuntimeProcessStarted( const now = currentTimestampMs(); const conditions: SQL[] = [ eq(driverInstancesTable.id, driverInstanceId), - notInArray(driverInstancesTable.status, ["stopped", "failed"]), + inArray(driverInstancesTable.status, REUSABLE_DRIVER_INSTANCE_STATUSES), + isNull(driverInstancesTable.statusOperationId), ]; if (options.expectedBootTokenHash !== undefined) { diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-instance-token.repository.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-instance-token.repository.ts index 06ad29a8..c3db90f3 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-instance-token.repository.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/driver-instance-token.repository.ts @@ -11,19 +11,32 @@ import type { DriverInstanceStatus } from "./status"; interface DriverInstanceTokenRow { boot_token_expires_at: number; boot_token_used_at: number | null; + connection_id: string | null; generation: number; id: DriverInstanceId; status: DriverInstanceStatus; } -export async function claimDriverInstanceByBootTokenHash( - bindings: ApiBindings, - bootTokenHash: Uint8Array, -): Promise<{ +interface DriverInstanceTokenValidation { driverInstanceId: DriverInstanceId | null; error: string | null; generation: number | null; -}> { +} + +export async function validateDriverInstanceBootTokenHash( + bindings: ApiBindings, + bootTokenHash: Uint8Array, +): Promise { + return validateDriverInstanceTokenRow( + await readDriverInstanceTokenRow(bindings, bootTokenHash), + currentTimestampMs(), + ); +} + +export async function claimDriverInstanceByBootTokenHash( + bindings: ApiBindings, + bootTokenHash: Uint8Array, +): Promise { const now = currentTimestampMs(); const claimed = (await getAppDatabase(bindings.DB) @@ -59,8 +72,16 @@ export async function claimDriverInstanceByBootTokenHash( }; } - const row = await readDriverInstanceTokenRow(bindings, bootTokenHash); + return validateDriverInstanceTokenRow( + await readDriverInstanceTokenRow(bindings, bootTokenHash), + now, + ); +} +function validateDriverInstanceTokenRow( + row: DriverInstanceTokenRow | null, + now: number, +): DriverInstanceTokenValidation { if (!row) { return { driverInstanceId: null, @@ -77,6 +98,14 @@ export async function claimDriverInstanceByBootTokenHash( }; } + if (row.status === "connecting" && row.boot_token_used_at !== null) { + return { + driverInstanceId: row.id, + error: null, + generation: row.generation, + }; + } + if (row.boot_token_used_at !== null || row.status !== "provisioning") { return { driverInstanceId: null, @@ -86,9 +115,9 @@ export async function claimDriverInstanceByBootTokenHash( } return { - driverInstanceId: null, - error: "Boot token is invalid.", - generation: null, + driverInstanceId: row.id, + error: null, + generation: row.generation, }; } @@ -109,6 +138,7 @@ async function readDriverInstanceTokenRow( .select({ bootTokenExpiresAt: driverInstancesTable.bootTokenExpiresAt, bootTokenUsedAt: driverInstancesTable.bootTokenUsedAt, + connectionId: driverInstancesTable.connectionId, generation: driverInstancesTable.generation, id: driverInstancesTable.id, status: driverInstancesTable.status, @@ -125,6 +155,7 @@ async function readDriverInstanceTokenRow( return { boot_token_expires_at: row.bootTokenExpiresAt, boot_token_used_at: row.bootTokenUsedAt, + connection_id: row.connectionId, generation: row.generation, id: row.id, status: row.status, diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/event-link-assertion.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/event-link-assertion.ts index 8917df0c..06a4a48c 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/event-link-assertion.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/event-link-assertion.ts @@ -1,8 +1,11 @@ -import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; -import type { RuntimeEventKind } from "@mosoo/runtime-events"; +import { readRuntimeAgentTaskSnapshot } from "@mosoo/runtime-events"; +import type { RuntimeEventEnvelope, RuntimeEventKind } from "@mosoo/runtime-events"; +import { ACTIVE_SESSION_RUN_STATUSES } from "../../domain/session-run-lifecycle.machine"; import type { RuntimeSessionLink } from "./event-types"; +const activeSessionRunStatuses = new Set(ACTIVE_SESSION_RUN_STATUSES); + const runBoundRuntimeEventDomains = new Set([ "image", "item", @@ -16,12 +19,15 @@ const runBoundRuntimeEventDomains = new Set([ ]); const runBoundRuntimeEventKinds = new Set([ + "agent.tasks.replaced", "file.change.updated", "mcp.tool.updated", "permission.requested", "permission.resolved", "permission.review.completed", "permission.review.started", + "runtime.resume.updated", + "usage.updated", ]); function runtimeEventRequiresRunLink(kind: RuntimeEventKind): boolean { @@ -49,6 +55,14 @@ export function assertRuntimeEventMatchesDriverLink( throw new Error("Runtime driver event driver instance id does not match the request."); } + if ((event.runtimeId ?? null) !== input.link.runtimeId) { + throw new Error("Runtime driver event runtime id does not match the driver session link."); + } + + if (event.runId !== undefined && (event.traceId ?? null) !== input.link.traceId) { + throw new Error("Runtime driver event trace id does not match the driver session link."); + } + if (event.runId === undefined && !runtimeEventRequiresRunLink(event.kind)) { return; } @@ -56,15 +70,12 @@ export function assertRuntimeEventMatchesDriverLink( if (event.runId !== input.link.sessionRunId) { throw new Error("Runtime driver event run id does not match the driver session link."); } -} -export function assertRuntimeEventMatchesDriverEnvelope( - event: RuntimeEventEnvelope, - input: { - eventId: string; - }, -): void { - if (event.sourceEventId !== undefined && event.sourceEventId !== input.eventId) { - throw new Error("Runtime driver event source id does not match the driver envelope."); + if ( + event.kind === "agent.tasks.replaced" && + !activeSessionRunStatuses.has(input.link.sessionRunStatus ?? "") && + readRuntimeAgentTaskSnapshot(event).tasks.length > 0 + ) { + throw new Error("Runtime agent task snapshot requires an active session run."); } } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/event-persistence.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/event-persistence.ts index f2d5dae3..6ea29f15 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/event-persistence.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/event-persistence.ts @@ -1,20 +1,21 @@ -import { sessionEventsTable, sessionsTable } from "@mosoo/db"; import type { DriverInstanceId } from "@mosoo/id"; -import { and, eq, isNull } from "drizzle-orm"; import { captureServerProductEvent, SERVER_PRODUCT_ANALYTICS_EVENTS, } from "../../../../platform/analytics/product-analytics"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; -import { getAppDatabase } from "../../../../platform/db/drizzle"; -import { currentTimestampMs } from "../../../../time"; -import { createSessionRuntimeEvent } from "../../../sessions/application/session-event-write.service"; -import { upsertSessionModelCallUsage } from "../../../sessions/application/session-model-call.service"; +import { reduceMessageStreamLifecycle } from "../../../sessions/domain/session-event-stream-fold"; +import { createSessionRuntimeEventProjection } from "../../../sessions/domain/session-runtime-event-projection"; +import { + isSealedPublicSessionMessageStream, + readPublicSessionMessageStreamSealState, +} from "../../../sessions/infrastructure/session-message-event-stream.repository"; import { persistSessionRuntimeEvents } from "../../../sessions/infrastructure/session-runtime-event-store.repository"; -import { setSessionRunStatus } from "../session-runs/session-run-store.repository"; -import type { SessionRunTransitionOutcome } from "../session-runs/session-run-store.repository"; -import { persistAssistantMessageProjection } from "./assistant-message-projection"; +import { prepareAssistantMessageProjection } from "./assistant-message-projection"; +import type { PreparedAssistantMessageProjection } from "./assistant-message-projection"; +import { commitTerminalRunProjection } from "./completed-run-commit.repository"; +import type { DriverTerminalRunStatus } from "./completed-run-commit.repository"; import { compactRuntimeDriverRunTransitions } from "./event-projection"; import type { AppRuntimeDriverEventsResult, @@ -22,122 +23,181 @@ import type { RuntimeSessionLink, SessionLiveState, } from "./event-types"; -import { hasTerminalRuntimeDriverRunTransition } from "./run-transitions"; - -async function loadTerminalRunRelease() { - return import("./terminal-run-release"); -} -function getModelCallStatus( - transitions: ReturnType, -): "completed" | "failed" | "started" { - for (const transition of transitions) { - if (transition.status === "completed") { - return "completed"; - } +async function resolveFinalAssistantMessage( + database: D1Database, + input: { + messageId: string; + runId: NonNullable; + sessionId: NonNullable; + }, +): Promise<{ id: string }> { + const sealed = await isSealedPublicSessionMessageStream(database, { + processType: "agent.message.delta", + runId: input.runId, + sessionId: input.sessionId, + streamId: input.messageId, + }); - if (transition.status === "cancelled" || transition.status === "failed") { - return "failed"; - } + if (!sealed) { + throw new Error( + `Completed run final message ${input.messageId} has no sealed authoritative snapshot.`, + ); } - return "started"; + return { id: input.messageId }; } -type DriverProjectedSessionRunStatusInput = Parameters[1]; +async function isFinalAssistantMessageSealedAfterProjection( + database: D1Database, + input: { + messageId: string; + projection: AppRuntimeDriverEventsResult; + runId: NonNullable; + sessionId: NonNullable; + }, +): Promise { + let state = await readPublicSessionMessageStreamSealState(database, { + processType: "agent.message.delta", + runId: input.runId, + sessionId: input.sessionId, + streamId: input.messageId, + }); -function assertDriverProjectedSessionRunTransition(outcome: SessionRunTransitionOutcome): void { - switch (outcome.kind) { - case "applied": - case "duplicate": { - return; + for (const { event } of input.projection.runtimeEvents) { + if ( + event.kind !== "message.added" && + event.kind !== "message.cancelled" && + event.kind !== "message.completed" && + event.kind !== "message.delta" && + event.kind !== "message.failed" && + event.kind !== "message.started" + ) { + continue; } - case "stale": { - if (outcome.reason === "terminal_run") { - return; - } - throw new Error("Driver run transition lost a concurrent status race."); + const row = createSessionRuntimeEventProjection(event); + if (row.streamId !== input.messageId) { + continue; } - case "repair_needed": { - throw new Error("Driver run transition left the session lifecycle projection stale."); + if ( + event.sessionId !== input.sessionId || + row.runId !== input.runId || + row.processType !== "agent.message.delta" + ) { + throw new Error(`Session message stream ${input.messageId} has conflicting identity rows.`); } - case "rejected": { - throw new Error(`Driver run transition was rejected: ${outcome.reason}.`); + if (row.visibility !== "all_consumers") { + throw new Error(`Session message stream ${input.messageId} has mixed visibility.`); } + + state = reduceMessageStreamLifecycle(state, event.kind); } -} -async function setDriverProjectedSessionRunStatus( - database: D1Database, - input: DriverProjectedSessionRunStatusInput, -): Promise { - const outcome = await setSessionRunStatus(database, input); - assertDriverProjectedSessionRunTransition(outcome); - return outcome; + return state.sealed; } -function isStaleTerminalRunTransition(outcome: SessionRunTransitionOutcome | null): boolean { - return outcome?.kind === "stale" && outcome.reason === "terminal_run"; +function isTerminalRunTransition( + transition: RuntimeDriverRunTransition | undefined, +): transition is RuntimeDriverRunTransition & { status: DriverTerminalRunStatus } { + return transition !== undefined && transition.status !== "running"; } -function isMatchingStaleTerminalRunTransition(input: { - readonly outcome: SessionRunTransitionOutcome | null; - readonly transition: RuntimeDriverRunTransition | undefined; -}): boolean { - if (input.outcome?.kind !== "stale" || input.outcome.reason !== "terminal_run") { - return true; - } +function terminalRuntimeEventKind(status: DriverTerminalRunStatus) { + return `run.${status}` as const; +} - return input.transition !== undefined && input.outcome.currentStatus === input.transition.status; +export interface PersistProjectedRuntimeDriverEventsResult { + liveState: SessionLiveState | null; + persistedSourceEventIds: readonly string[]; } -function getRunDurationMs(outcome: SessionRunTransitionOutcome): number | null { +export async function preflightProjectedRuntimeDriverEvents( + database: D1Database, + projection: AppRuntimeDriverEventsResult, +): Promise { + const { link } = projection; + if (link.sessionId === null) { + return; + } + const [runTransition] = compactRuntimeDriverRunTransitions(projection.transitions); + const terminalRunTransition = isTerminalRunTransition(runTransition) ? runTransition : undefined; + + if (terminalRunTransition !== undefined && link.sessionRunId === null) { + throw new Error("Terminal run projection is missing its run scope."); + } if ( - outcome.kind !== "applied" || - outcome.run.startedAt === null || - outcome.run.completedAt === null + terminalRunTransition?.status === "completed" && + projection.finalAssistantMessageId !== null ) { - return null; + if (link.sessionRunId === null) { + throw new Error("Completed run final message is missing its run scope."); + } + if ( + !(await isFinalAssistantMessageSealedAfterProjection(database, { + messageId: projection.finalAssistantMessageId, + projection, + runId: link.sessionRunId, + sessionId: link.sessionId, + })) + ) { + throw new Error( + `Completed run final message ${projection.finalAssistantMessageId} has no sealed authoritative snapshot.`, + ); + } + } + if (terminalRunTransition !== undefined) { + const terminalKind = terminalRuntimeEventKind(terminalRunTransition.status); + if (projection.runtimeEvents.filter(({ event }) => event.kind === terminalKind).length !== 1) { + throw new Error("Terminal run projection requires exactly one terminal runtime event."); + } } - - const durationMs = Date.parse(outcome.run.completedAt) - Date.parse(outcome.run.startedAt); - return Number.isFinite(durationMs) ? Math.max(0, durationMs) : null; } -async function autoTitleRuntimeSession( +export async function persistProjectedRuntimeDriverEventPrerequisites( database: D1Database, - link: RuntimeSessionLink, - title: string, + input: { + driverConnectionId: string; + driverGeneration: number; + driverInstanceId: DriverInstanceId; + projection: AppRuntimeDriverEventsResult; + }, ): Promise { - if (link.creatorId === null || link.sessionId === null) { + const { link, runtimeEvents } = input.projection; + if (link.sessionId === null) { return; } - - await getAppDatabase(database) - .update(sessionsTable) - .set({ - title, - updatedAt: currentTimestampMs(), - }) - .where( - and( - eq(sessionsTable.id, link.sessionId), - eq(sessionsTable.creatorAccountId, link.creatorId), - isNull(sessionsTable.title), - eq(sessionsTable.renamed, false), - ), - ) - .run(); -} - -export interface PersistProjectedRuntimeDriverEventsResult { - liveState: SessionLiveState | null; - persistedSourceEventIds: readonly string[]; + const [runTransition] = compactRuntimeDriverRunTransitions(input.projection.transitions); + const terminalRunTransition = isTerminalRunTransition(runTransition) ? runTransition : undefined; + if (terminalRunTransition === undefined) { + return; + } + const terminalKind = terminalRuntimeEventKind(terminalRunTransition.status); + const prerequisites = runtimeEvents.filter( + ({ event }) => + event.kind !== terminalKind && + event.kind !== "file.change.updated" && + event.kind !== "file.changed", + ); + if (prerequisites.length === 0) { + return; + } + await persistSessionRuntimeEvents(database, { + driverFence: { + connectionId: input.driverConnectionId, + driverInstanceId: input.driverInstanceId, + generation: input.driverGeneration, + sessionRunId: link.sessionRunId, + }, + records: prerequisites, + sessionId: link.sessionId, + }); } export async function persistProjectedRuntimeDriverEvents( bindings: ApiBindings, input: { + driverConnectionId: string; + driverGeneration: number; driverInstanceId: DriverInstanceId; projection: AppRuntimeDriverEventsResult; }, @@ -150,8 +210,7 @@ export async function persistProjectedRuntimeDriverEvents( }; const transitions = compactRuntimeDriverRunTransitions(projection.transitions); const [runTransition] = transitions; - const deferCompletedRunTransition = runTransition?.status === "completed"; - let runTransitionOutcome: SessionRunTransitionOutcome | null = null; + const terminalRunTransition = isTerminalRunTransition(runTransition) ? runTransition : undefined; if (link.sessionId === null) { return { @@ -159,119 +218,88 @@ export async function persistProjectedRuntimeDriverEvents( persistedSourceEventIds: [], }; } + const driverFence = { + connectionId: input.driverConnectionId, + driverInstanceId: input.driverInstanceId, + generation: input.driverGeneration, + sessionRunId: link.sessionRunId, + }; - if (runTransition !== undefined && link.sessionRunId !== null && !deferCompletedRunTransition) { - if (runTransition.status === "running") { - runTransitionOutcome = await setDriverProjectedSessionRunStatus(database, { - runId: link.sessionRunId, - source: "driver", - status: "running", - }); - } else if (runTransition.status === "cancelled") { - runTransitionOutcome = await setDriverProjectedSessionRunStatus(database, { - runId: link.sessionRunId, - source: "driver", - status: "cancelled", - }); - } else { - runTransitionOutcome = await setDriverProjectedSessionRunStatus(database, { - error: runTransition.error ?? null, - runId: link.sessionRunId, - source: "driver", - status: "failed", - }); - } + const completedTransition = terminalRunTransition?.status === "completed"; + const terminalEventKind = + terminalRunTransition === undefined + ? null + : terminalRuntimeEventKind(terminalRunTransition.status); + const preTerminalRuntimeEvents = + terminalEventKind === null + ? [] + : projection.runtimeEvents.filter((record) => record.event.kind !== terminalEventKind); + const terminalRuntimeEvents = + terminalEventKind === null + ? projection.runtimeEvents + : projection.runtimeEvents.filter((record) => record.event.kind === terminalEventKind); + const persistedSourceEventIds: string[] = []; + + if (preTerminalRuntimeEvents.length > 0) { + const persisted = await persistSessionRuntimeEvents(database, { + driverFence, + records: preTerminalRuntimeEvents, + sessionId: link.sessionId, + }); + persistedSourceEventIds.push(...persisted.persistedSourceEventIds); } - const shouldReleaseDriverRun = hasTerminalRuntimeDriverRunTransition(transitions); - let staleTerminalRunTransition = isStaleTerminalRunTransition(runTransitionOutcome); + let finalAssistantMessage: { id: string } | null = null; + let preparedAssistantMessage: PreparedAssistantMessageProjection | null = null; - if ( - staleTerminalRunTransition && - !isMatchingStaleTerminalRunTransition({ - outcome: runTransitionOutcome, - transition: runTransition, - }) - ) { - if (shouldReleaseDriverRun && link.sessionRunId !== null) { - const { releaseTerminalDriverInstanceSessionRun } = await loadTerminalRunRelease(); - await releaseTerminalDriverInstanceSessionRun(bindings, { - driverInstanceId: input.driverInstanceId, - sessionRunId: link.sessionRunId, - }); + if (completedTransition && projection.finalAssistantMessageId !== null) { + if (link.sessionRunId === null) { + throw new Error("Completed run final message is missing its run scope."); } - return { - liveState: null, - persistedSourceEventIds: [], - }; - } - - if (projection.sessionTitle !== null && projection.sessionTitle.length > 0) { - await autoTitleRuntimeSession(database, link, projection.sessionTitle); - } - - const traceId = link.traceId ?? link.sessionRunId ?? link.sessionId; + finalAssistantMessage = await resolveFinalAssistantMessage(database, { + messageId: projection.finalAssistantMessageId, + runId: link.sessionRunId, + sessionId: link.sessionId, + }); - if (projection.usage && link.sessionRunId !== null) { - await upsertSessionModelCallUsage(database, { - driverInstanceId: input.driverInstanceId, + preparedAssistantMessage = prepareAssistantMessageProjection({ + createdByAccountId: link.callerId ?? link.creatorId ?? input.driverInstanceId, + messageId: finalAssistantMessage.id, sessionId: link.sessionId, sessionRunId: link.sessionRunId, - status: getModelCallStatus(transitions), - traceId, - usage: projection.usage, }); } - const completedTransition = transitions.find((transition) => transition.status === "completed"); - const preCompletionRuntimeEvents = - completedTransition === undefined - ? [] - : projection.runtimeEvents.filter((record) => record.event.kind !== "run.completed"); - const terminalRuntimeEvents = - completedTransition === undefined - ? projection.runtimeEvents - : projection.runtimeEvents.filter((record) => record.event.kind === "run.completed"); - const finalAssistantRuntimeEvents: typeof terminalRuntimeEvents = []; - const persistedSourceEventIds: string[] = []; + let terminalRunCommit: Awaited> | null = null; - if (preCompletionRuntimeEvents.length > 0) { - const persisted = await persistSessionRuntimeEvents(database, { - records: preCompletionRuntimeEvents, - sessionId: link.sessionId, - }); - persistedSourceEventIds.push(...persisted.persistedSourceEventIds); - } + if (terminalRunTransition !== undefined) { + if (link.sessionRunId === null) { + throw new Error("Terminal run projection is missing its run scope."); + } - // Completion is a retry boundary. Arbitrate the terminal Run status before - // writing canonical output, so a concurrent failure/cancellation cannot - // leave a final assistant row on a non-completed Run. The terminal receipt - // remains last: a crash after the CAS is repaired by replay against the exact - // completed Run link. - if (deferCompletedRunTransition && link.sessionRunId !== null) { - runTransitionOutcome = await setDriverProjectedSessionRunStatus(database, { + const [terminalEvent] = terminalRuntimeEvents; + + if (terminalRuntimeEvents.length !== 1 || terminalEvent === undefined) { + throw new Error("Terminal run projection requires exactly one terminal runtime event."); + } + + terminalRunCommit = await commitTerminalRunProjection(database, { + assistantMessage: preparedAssistantMessage, + error: terminalRunTransition.error ?? null, + expectedDriverObservation: { + connectionId: input.driverConnectionId, + driverInstanceId: input.driverInstanceId, + generation: input.driverGeneration, + }, runId: link.sessionRunId, + sessionId: link.sessionId, source: "driver", - status: "completed", + targetStatus: terminalRunTransition.status, + terminalEvent, }); - staleTerminalRunTransition = isStaleTerminalRunTransition(runTransitionOutcome); - - if ( - staleTerminalRunTransition && - !isMatchingStaleTerminalRunTransition({ - outcome: runTransitionOutcome, - transition: runTransition, - }) - ) { - if (shouldReleaseDriverRun) { - const { releaseTerminalDriverInstanceSessionRun } = await loadTerminalRunRelease(); - await releaseTerminalDriverInstanceSessionRun(bindings, { - driverInstanceId: input.driverInstanceId, - sessionRunId: link.sessionRunId, - }); - } - + persistedSourceEventIds.push(...terminalRunCommit.persistedSourceEventIds); + if (terminalRunCommit.kind === "stale") { return { liveState: null, persistedSourceEventIds: [], @@ -279,73 +307,15 @@ export async function persistProjectedRuntimeDriverEvents( } } - if ( - completedTransition !== undefined && - projection.finalAssistantMessage !== null && - nextLiveState !== null && - link.sessionRunId !== null && - nextLiveState.run.id === link.sessionRunId - ) { - await persistAssistantMessageProjection(database, { - createdByAccountId: link.callerId ?? link.creatorId ?? input.driverInstanceId, - driverInstanceId: input.driverInstanceId, - messageId: projection.finalAssistantMessage.id, - messageText: projection.finalAssistantMessage.text, + if (terminalRunTransition === undefined) { + const persistedTerminalEvents = await persistSessionRuntimeEvents(database, { + driverFence, + records: terminalRuntimeEvents, sessionId: link.sessionId, - sessionRunId: link.sessionRunId, - state: nextLiveState, }); - - const terminalRuntimeEvent = terminalRuntimeEvents[0]; - const finalSnapshotAlreadyPersisted = - (await getAppDatabase(database) - .select({ id: sessionEventsTable.id }) - .from(sessionEventsTable) - .where( - and( - eq(sessionEventsTable.sessionId, link.sessionId), - eq(sessionEventsTable.runId, link.sessionRunId), - eq(sessionEventsTable.eventType, "message.added"), - eq(sessionEventsTable.processType, "agent.message.delta"), - eq(sessionEventsTable.contentText, projection.finalAssistantMessage.text), - ), - ) - .limit(1) - .get()) !== undefined; - - if (terminalRuntimeEvent !== undefined && !finalSnapshotAlreadyPersisted) { - const sourceEventId = `session-run:${link.sessionRunId}:final-assistant`; - finalAssistantRuntimeEvents.push({ - event: createSessionRuntimeEvent({ - actor: terminalRuntimeEvent.event.actor, - kind: "message.added", - ...(terminalRuntimeEvent.occurredAt === null - ? {} - : { occurredAtMs: terminalRuntimeEvent.occurredAt }), - origin: terminalRuntimeEvent.event.origin, - payload: { - content: projection.finalAssistantMessage.text, - messageId: projection.finalAssistantMessage.id, - role: "agent", - }, - runId: link.sessionRunId, - sessionId: link.sessionId, - sourceEventId, - traceId: terminalRuntimeEvent.event.traceId ?? link.traceId, - visibility: terminalRuntimeEvent.event.visibility, - }), - occurredAt: terminalRuntimeEvent.occurredAt, - sourceEventId, - }); - } + persistedSourceEventIds.push(...persistedTerminalEvents.persistedSourceEventIds); } - const persistedTerminalEvents = await persistSessionRuntimeEvents(database, { - records: [...finalAssistantRuntimeEvents, ...terminalRuntimeEvents], - sessionId: link.sessionId, - }); - persistedSourceEventIds.push(...persistedTerminalEvents.persistedSourceEventIds); - let committedLiveState: SessionLiveState | null = null; if (projection.liveStateChanged && nextLiveState !== null) { @@ -353,28 +323,10 @@ export async function persistProjectedRuntimeDriverEvents( } if ( - runTransition === undefined && - !staleTerminalRunTransition && - link.sessionRunId !== null && - projection.liveStateChanged && - nextLiveState !== null && - nextLiveState.run.id === link.sessionRunId && - nextLiveState.run.status === "waiting_input" - ) { - await setDriverProjectedSessionRunStatus(database, { - runId: link.sessionRunId, - source: "driver", - status: "waiting_input", - }); - } - - if ( - completedTransition !== undefined && - runTransitionOutcome?.kind === "applied" && + completedTransition && + terminalRunCommit?.kind === "applied" && link.executionOwnerId !== null ) { - const runDurationMs = getRunDurationMs(runTransitionOutcome); - await captureServerProductEvent(bindings, { distinctId: link.executionOwnerId, event: SERVER_PRODUCT_ANALYTICS_EVENTS.taskSucceeded, @@ -382,7 +334,7 @@ export async function persistProjectedRuntimeDriverEvents( agent_id: link.agentId, app_id: link.appId, run_id: link.sessionRunId, - run_duration_ms: runDurationMs, + run_duration_ms: terminalRunCommit.runDurationMs, sandbox_id: link.sandboxId, sandbox_kind: link.sandboxKind, sandbox_subject_kind: link.sandboxSubjectKind, @@ -392,14 +344,6 @@ export async function persistProjectedRuntimeDriverEvents( }); } - if (shouldReleaseDriverRun && link.sessionRunId !== null) { - const { releaseTerminalDriverInstanceSessionRun } = await loadTerminalRunRelease(); - await releaseTerminalDriverInstanceSessionRun(bindings, { - driverInstanceId: input.driverInstanceId, - sessionRunId: link.sessionRunId, - }); - } - return { liveState: committedLiveState, persistedSourceEventIds, diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/event-projection.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/event-projection.ts index bd76d6a1..997433e0 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/event-projection.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/event-projection.ts @@ -6,7 +6,6 @@ import type { SessionLiveState, SessionPermissionRequestView, } from "../../../sessions/application/session-live-state.service"; -import { normalizeSessionTitle } from "../../../sessions/domain/session-title"; import type { RuntimeDriverRunTransition, RuntimeSessionLink } from "./event-types"; export function compactRuntimeDriverRunTransitions( @@ -92,23 +91,6 @@ export function createBaseLiveState( }); } -export function normalizeRuntimeSessionInfoTitle(title: string | null | undefined): string | null { - if (typeof title !== "string" || title.trim().length === 0) { - return null; - } - - return normalizeSessionTitle(title); -} - -export function upsertPermissionRequest( - current: SessionPermissionRequestView[], - next: SessionPermissionRequestView, -): SessionPermissionRequestView[] { - const requests = current.filter((request) => request.requestId !== next.requestId); - requests.push(next); - return requests; -} - export function removePermissionRequest( current: SessionPermissionRequestView[], requestId: string, diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/event-types.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/event-types.ts index af005f42..41c787c4 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/event-types.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/event-types.ts @@ -1,4 +1,3 @@ -import type { SessionUsageSummary } from "@mosoo/ag-ui-session"; import type { AgentKind } from "@mosoo/contracts/agent"; import type { SandboxSubjectKind } from "@mosoo/contracts/sandbox"; import type { SessionType } from "@mosoo/contracts/session"; @@ -7,6 +6,7 @@ import type { AccountId, AgentId, AppId, + DriverCommandId, PlatformId, SandboxId, SessionId, @@ -21,6 +21,17 @@ import type { export type { SessionLiveState }; +export interface CanonicalDriverEventEnvelope { + readonly event: RuntimeEventEnvelope; + readonly eventId: string; + readonly occurredAt?: string | null | undefined; +} + +export interface HostDriverEventBatchInput { + readonly driverInstanceId: string; + readonly events: readonly CanonicalDriverEventEnvelope[]; +} + export interface RuntimeSessionLink { agentId: AgentId | null; appId: AppId | null; @@ -29,6 +40,7 @@ export interface RuntimeSessionLink { executionOwnerId: AccountId | null; sandboxId: SandboxId | null; sandboxKind: AgentKind | null; + runtimeId: string | null; sessionId: SessionId | null; sessionRunId: SessionRunId | null; sessionRunStatus: SessionRunStatus | null; @@ -42,8 +54,12 @@ export function runtimeSessionLinkNeedsRefresh(link: RuntimeSessionLink | null): } export interface ProjectedRuntimeEventRecord { + artifactAttemptId?: string | null; + artifactManifestJson?: string | null; + artifactManifestSha256?: string | null; event: RuntimeEventEnvelope; occurredAt: number | null; + provenMcpCommandId?: DriverCommandId | null; sourceEventId: string | null; } @@ -64,13 +80,11 @@ export interface RuntimeDriverRunTransition { } export interface AppRuntimeDriverEventsResult { - finalAssistantMessage: { id: string; text: string } | null; + finalAssistantMessageId: string | null; link: RuntimeSessionLink; liveStateChanged: boolean; nextLiveState: SessionLiveState | null; - sessionTitle: string | null; transitions: RuntimeDriverRunTransition[]; - usage: SessionUsageSummary | null; runtimeEvents: ProjectedRuntimeEventRecord[]; sessionDeliveryEvents: ProjectedSessionDeliveryEvent[]; } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/events.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/events.ts index 9e5a504d..621692ff 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/events.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/events.ts @@ -1,32 +1,13 @@ -import { - EventType, - MOSOO_CUSTOM_EVENT, - createServerCustomEvent, - parseNullableSessionUsageSummary, -} from "@mosoo/ag-ui-session"; -import type { DriverEventEnvelope } from "@mosoo/agent-driver/events"; -import { parsePlatformId } from "@mosoo/id"; -import type { DriverInstanceId } from "@mosoo/id"; -import type { AccountId, SessionId } from "@mosoo/id"; +import type { DriverCommandId, DriverInstanceId } from "@mosoo/id"; import { parseRuntimeEventEnvelope, - readRuntimeEventFileChanges, readRuntimeEventPayload, - readRuntimeEventPermissionRequest, readRuntimeEventString, - readRuntimeRunPayload, } from "@mosoo/runtime-events"; import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; -import { createErrorLogContext, logInfo, logWarn } from "../../../../platform/cloudflare/logger"; -import { withDisposedRpcResource } from "../../../../platform/cloudflare/rpc-disposal"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; import { isTruthy } from "../../../../shared/truthiness"; -import { - createRuntimeOutputContentSha256, - createRuntimeOutputParentPath, - fileStore, -} from "../../../files/application/file-store"; import { applyAgUiEventToSessionLiveState, loadSessionViewerState, @@ -36,364 +17,43 @@ import type { SessionDeliveryEvent, SessionLiveState, } from "../../../sessions/application/session-live-state.service"; -import { getRuntimeKindPolicy } from "../../domain/runtime-kind-policy"; -import { createSessionRunTerminalFailureSourceId } from "../../domain/session-run-terminal-event-id"; -import { upsertNativeResumeRef } from "../native-resume-ref.repository"; -import { getRuntimeSubjectKeepAliveHandle } from "../runtime-subject-lifecycle/runtime-subject-lifecycle.service"; -import { getRuntimeConversationSession } from "../runtime-subject-lifecycle/runtime-subject-store"; -import { readSandboxFileBytes } from "../sandbox-file-bytes"; -import type { ExecutionSessionHandle } from "../sandbox-handles"; -import { - assertRuntimeEventMatchesDriverEnvelope, - assertRuntimeEventMatchesDriverLink, -} from "./event-link-assertion"; -import { - createBaseLiveState, - normalizeRuntimeSessionInfoTitle, - readPermissionRequestViews, - readRuntimeDriverRunTransition, - removePermissionRequest, - upsertPermissionRequest, -} from "./event-projection"; +import { createSessionRunTerminalSourceId } from "../../domain/session-run-terminal-event-id"; +import { assertRuntimeEventMatchesDriverLink } from "./event-link-assertion"; +import { createBaseLiveState, readRuntimeDriverRunTransition } from "./event-projection"; import type { - ProjectedRuntimeEventRecord, AppRuntimeDriverEventsResult, + CanonicalDriverEventEnvelope, + ProjectedRuntimeEventRecord, RuntimeDriverRunTransition, RuntimeSessionLink, } from "./event-types"; -import { readNativeResumeRef } from "./native-resume-ref-event"; -import { - RUNTIME_SESSION_OUTPUT_DIR_NAME, - RUNTIME_SESSION_OUTPUT_SCAN_MAX_FILES, - getRuntimeSessionOutputDirectory, - guessRuntimeSessionOutputContentType, - readRuntimeSessionOutputListing, - toRuntimeSessionOutputArtifactPath, - toRuntimeSessionOutputFile, -} from "./runtime-session-outputs"; import { getRuntimeSessionLink } from "./session-link.repository"; export type { AppRuntimeDriverEventsResult, RuntimeDriverRunTransition, RuntimeSessionLink, } from "./event-types"; -export { persistProjectedRuntimeDriverEvents } from "./event-persistence"; +export { + persistProjectedRuntimeDriverEventPrerequisites, + persistProjectedRuntimeDriverEvents, + preflightProjectedRuntimeDriverEvents, +} from "./event-persistence"; export { getRuntimeSessionLink } from "./session-link.repository"; export { recordDriverInstanceCompletion, recordDriverInstanceFailure, } from "./terminal-driver-events"; -function quoteShellArg(value: string): string { - return `'${value.replaceAll("'", `'"'"'`)}'`; -} - -function resolveRuntimeOutputCreator(link: RuntimeSessionLink): AccountId | null { - const actorId = link.executionOwnerId ?? link.callerId ?? link.creatorId; - - if (!isTruthy(actorId)) { - return null; - } - - return parsePlatformId(actorId, "runtime output creator account ID"); -} - -function readRuntimeFileChangeContentType( - metadata: Record | undefined, -): string | null { - const contentType = metadata?.["contentType"] ?? metadata?.["mimeType"]; - return typeof contentType === "string" && contentType.trim().length > 0 ? contentType : null; -} - -function createRuntimeSessionOutputListCommand(outputDir: string): string { - const quotedOutputDir = quoteShellArg(outputDir); - const command = [ - `if [ ! -d ${quotedOutputDir} ]; then exit 0; fi`, - `cd ${quotedOutputDir}`, - `find . -type f -print | sed 's#^\\./##' | sort | head -n ${RUNTIME_SESSION_OUTPUT_SCAN_MAX_FILES}`, - ].join(" && "); - - return `sh -lc ${quoteShellArg(command)}`; -} - -async function listRuntimeSessionOutputFiles( - handle: ExecutionSessionHandle, - outputDir: string, -): Promise { - const result = await handle.exec(createRuntimeSessionOutputListCommand(outputDir)); - - if (!result.success || result.exitCode !== 0) { - throw new Error( - result.stderr.trim() || - result.stdout.trim() || - `Failed to list runtime session outputs in ${outputDir}.`, - ); - } - - return readRuntimeSessionOutputListing(result.stdout); -} - -async function recordRuntimeSessionOutputFile(input: { - bindings: ApiBindings; - body: Uint8Array; - contentType: string | null; - createdBy: AccountId; - existingArtifacts: Set; - path: string; - recordedArtifacts: Set; - sessionId: SessionId; -}): Promise { - const contentSha256 = await createRuntimeOutputContentSha256(input.body); - const artifactKey = createRuntimeOutputParentPath(input.path, contentSha256); - - if (input.recordedArtifacts.has(artifactKey) || input.existingArtifacts.has(artifactKey)) { - return; - } - - await fileStore.recordRuntimeOutput({ - bindings: input.bindings, - body: input.body, - contentSha256, - contentType: input.contentType, - createdBy: input.createdBy, - path: input.path, - sessionId: input.sessionId, - }); - input.recordedArtifacts.add(artifactKey); - input.existingArtifacts.add(artifactKey); -} - -function readTerminalPendingToolResult(event: RuntimeEventEnvelope): string | null { - if (event.kind === "run.failed") { - const run = readRuntimeRunPayload(event).run; - const message = run?.error?.message ?? "Run failed before the tool returned a result."; - return `Tool failed before returning a result: ${message}`; - } - - if (event.kind === "run.cancelled") { - return "Tool was cancelled before returning a result."; - } - - return null; -} - -function createPendingToolResultEvents( - state: SessionLiveState, - event: RuntimeEventEnvelope, -): SessionDeliveryEvent[] { - const content = readTerminalPendingToolResult(event); - - if (content === null) { - return []; - } - - return state.messages.flatMap((message) => { - const completedToolCallIds = new Set( - message.segments.flatMap((segment) => - segment.kind === "tool_result" ? [segment.toolCallId] : [], - ), - ); - - return message.segments.flatMap((segment) => { - if (segment.kind !== "tool_use" || completedToolCallIds.has(segment.toolCallId)) { - return []; - } - - return [ - { - content, - messageId: message.id, - toolCallId: segment.toolCallId, - type: EventType.TOOL_CALL_RESULT, - }, - { - toolCallId: segment.toolCallId, - type: EventType.TOOL_CALL_END, - }, - ]; - }); - }); -} - -async function recordRuntimeFileChanges(input: { - bindings: ApiBindings; - event: RuntimeEventEnvelope; - link: RuntimeSessionLink; -}): Promise { - const sessionId = input.link.sessionId; - const sandboxId = input.link.sandboxId; - const createdBy = resolveRuntimeOutputCreator(input.link); - const changes = readRuntimeEventFileChanges(input.event).filter( - (change) => change.change === "upsert", - ); - - if (changes.length === 0) { - return; - } - - if (sessionId === null || sandboxId === null || createdBy === null) { - logWarn("runtime.file_artifact.record_skipped", { - driverInstanceId: input.event.driverInstanceId ?? null, - hasCreatedBy: createdBy !== null, - sandboxId, - sessionId, - }); - return; - } - - const conversation = await getRuntimeConversationSession(input.bindings.DB, sessionId); - - if (conversation === null) { - logWarn("runtime.file_artifact.record_skipped.missing_session", { - sandboxId, - sessionId, - }); - return; - } - - const outputChanges = changes.flatMap((change) => { - const outputFile = toRuntimeSessionOutputFile({ - contentType: readRuntimeFileChangeContentType(change.metadata), - cwd: conversation.cwd, - path: change.path, - }); - - return outputFile === null ? [] : [outputFile]; - }); - - if (outputChanges.length === 0) { - return; - } - - const parsedSessionId = parsePlatformId(sessionId, "runtime output session ID"); - const existingArtifacts = new Set( - await fileStore.listReadySessionArtifactKeys(input.bindings.DB, parsedSessionId), - ); - const recordedArtifacts = new Set(); - - await withDisposedRpcResource( - await getRuntimeSubjectKeepAliveHandle(input.bindings, sandboxId), - async (sandbox) => { - const sandboxSession = await sandbox.getSession(conversation.sandboxSessionId); - - for (const outputFile of outputChanges) { - try { - await recordRuntimeSessionOutputFile({ - bindings: input.bindings, - body: await readSandboxFileBytes(sandboxSession, outputFile.readPath), - contentType: outputFile.contentType, - createdBy, - existingArtifacts, - path: outputFile.artifactPath, - recordedArtifacts, - sessionId: parsedSessionId, - }); - } catch (error) { - logWarn("runtime.file_artifact.record_failed", { - ...createErrorLogContext(error), - path: outputFile.artifactPath, - sandboxId, - sessionId, - }); - } - } - }, - ); -} - -async function recordRuntimeSessionOutputDirectory(input: { - bindings: ApiBindings; - event: RuntimeEventEnvelope; - link: RuntimeSessionLink; -}): Promise { - const sessionId = input.link.sessionId; - const sandboxId = input.link.sandboxId; - const createdBy = resolveRuntimeOutputCreator(input.link); - - if (sessionId === null || sandboxId === null || createdBy === null) { - return; - } - - const parsedSessionId = parsePlatformId(sessionId, "runtime output session ID"); - let conversation; - - try { - conversation = await getRuntimeConversationSession(input.bindings.DB, parsedSessionId); - } catch (error) { - logWarn("runtime.file_artifact.output_scan_session_lookup_failed", { - ...createErrorLogContext(error), - driverInstanceId: input.event.driverInstanceId ?? null, - sandboxId, - sessionId, - }); - return; - } - - if (conversation === null) { - return; - } - - try { - await withDisposedRpcResource( - await getRuntimeSubjectKeepAliveHandle(input.bindings, sandboxId), - async (sandbox) => { - const sandboxSession = await sandbox.getSession(conversation.sandboxSessionId); - const outputDir = getRuntimeSessionOutputDirectory(conversation.cwd); - const outputPaths = await listRuntimeSessionOutputFiles(sandboxSession, outputDir); - - if (outputPaths.length === 0) { - return; - } - - const existingArtifacts = new Set( - await fileStore.listReadySessionArtifactKeys(input.bindings.DB, parsedSessionId), - ); - const recordedArtifacts = new Set(); - - for (const outputPath of outputPaths) { - const artifactPath = toRuntimeSessionOutputArtifactPath(outputPath); - - try { - await recordRuntimeSessionOutputFile({ - bindings: input.bindings, - body: await readSandboxFileBytes(sandboxSession, `${outputDir}/${outputPath}`), - contentType: guessRuntimeSessionOutputContentType(outputPath), - createdBy, - existingArtifacts, - path: artifactPath, - recordedArtifacts, - sessionId: parsedSessionId, - }); - } catch (error) { - logWarn("runtime.file_artifact.output_record_failed", { - ...createErrorLogContext(error), - path: artifactPath, - sandboxId, - sessionId, - }); - } - } - }, - ); - } catch (error) { - logWarn("runtime.file_artifact.output_scan_failed", { - ...createErrorLogContext(error), - driverInstanceId: input.event.driverInstanceId ?? null, - outputDir: `${RUNTIME_SESSION_OUTPUT_DIR_NAME}/`, - sandboxId, - sessionId, - }); - } -} - export async function appRuntimeDriverEvents( bindings: ApiBindings, input: { assertCurrentConnection?: () => void; currentLiveState?: SessionLiveState | null; - events: readonly DriverEventEnvelope[]; + projectLiveState?: boolean; + events: readonly CanonicalDriverEventEnvelope[]; driverInstanceId: DriverInstanceId; link?: RuntimeSessionLink | null; + provenMcpCommandIds?: ReadonlyMap; }, ): Promise { const database = bindings.DB; @@ -410,32 +70,41 @@ export async function appRuntimeDriverEvents( throw new Error("Runtime driver run event is missing a session run id."); } - const currentLiveState = - input.currentLiveState ?? - (await loadStoredRuntimeLiveState(database, { - driverInstanceId: input.driverInstanceId, - link, - })); + const projectLiveState = input.projectLiveState !== false; + const currentLiveState = !projectLiveState + ? createBaseLiveState({ + callerId: link.callerId, + creatorId: link.creatorId, + driverInstanceId: input.driverInstanceId, + sessionId: link.sessionId, + }) + : (input.currentLiveState ?? + (await loadStoredRuntimeLiveState(database, { + driverInstanceId: input.driverInstanceId, + link, + }))); let nextLiveState = currentLiveState; let liveStateChanged = false; - let finalAssistantMessage: AppRuntimeDriverEventsResult["finalAssistantMessage"] = null; - let sessionTitle: string | null = null; - let usage: AppRuntimeDriverEventsResult["usage"] = null; + let finalAssistantMessageId: AppRuntimeDriverEventsResult["finalAssistantMessageId"] = null; const runtimeEvents: ProjectedRuntimeEventRecord[] = []; const sessionDeliveryEvents: AppRuntimeDriverEventsResult["sessionDeliveryEvents"] = []; const transitions: RuntimeDriverRunTransition[] = []; - function appendCanonicalEvent(source: DriverEventEnvelope, event: RuntimeEventEnvelope): void { + function appendCanonicalEvent( + source: CanonicalDriverEventEnvelope, + event: RuntimeEventEnvelope, + ): void { runtimeEvents.push({ event, occurredAt: toDriverEventOccurredAtMs(source.occurredAt), + provenMcpCommandId: input.provenMcpCommandIds?.get(source.eventId) ?? null, sourceEventId: resolveDriverEventPersistenceSourceId(source, event), }); } function appendSessionDeliveryEvent( - source: DriverEventEnvelope, + source: CanonicalDriverEventEnvelope, event: RuntimeEventEnvelope, deliveryEvent: SessionDeliveryEvent, ): void { @@ -453,145 +122,32 @@ export async function appRuntimeDriverEvents( driverInstanceId: input.driverInstanceId, link, }); - assertRuntimeEventMatchesDriverEnvelope(event, { - eventId: envelope.eventId, - }); appendCanonicalEvent(envelope, event); if (event.kind === "runtime.resume.updated") { - const nativeResumeRef = readNativeResumeRef(event); - - if (nativeResumeRef === null) { - continue; - } - - const policy = link.sandboxKind === null ? null : getRuntimeKindPolicy(link.sandboxKind); - - if (policy?.nativeResume.persistence !== "platform") { - logInfo("runtime.native_resume_ref.ignored", { - driverInstanceId: input.driverInstanceId, - kind: nativeResumeRef.kind, - runtimeId: nativeResumeRef.runtimeId, - sandboxKind: link.sandboxKind, - sandboxSubjectKind: link.sandboxSubjectKind, - sessionId: link.sessionId, - sessionRunId: link.sessionRunId, - }); - continue; - } - - if (link.sessionRunId === null) { - logInfo("runtime.native_resume_ref.deferred", { - driverInstanceId: input.driverInstanceId, - kind: nativeResumeRef.kind, - runtimeId: nativeResumeRef.runtimeId, - sessionId: link.sessionId, - }); - continue; - } - - input.assertCurrentConnection?.(); - await upsertNativeResumeRef(database, { - driverInstanceId: input.driverInstanceId, - nativeResumeRef, - sessionId: link.sessionId, - sessionRunId: link.sessionRunId, - }); - logInfo("runtime.native_resume_ref.observed", { - driverInstanceId: input.driverInstanceId, - kind: nativeResumeRef.kind, - runtimeId: nativeResumeRef.runtimeId, - sessionId: link.sessionId, - sessionRunId: link.sessionRunId, - }); continue; } if (event.kind === "file.change.updated" || event.kind === "file.changed") { - await recordRuntimeFileChanges({ - bindings, - event, - link, - }); continue; } if (event.kind === "run.completed") { const payload = readRuntimeEventPayload(event); - const finalMessageId = readRuntimeEventString(payload, "finalMessageId"); - const finalMessageText = readRuntimeEventString(payload, "finalMessageText"); - finalAssistantMessage = - finalMessageId === null || finalMessageText === null - ? null - : { id: finalMessageId, text: finalMessageText }; - await recordRuntimeSessionOutputDirectory({ - bindings, - event, - link, - }); + finalAssistantMessageId = readRuntimeEventString(payload, "finalMessageId"); } - if (event.kind === "permission.requested") { - const request = readRuntimeEventPermissionRequest(event); - - if (request) { - const permissionsUpdatedEvent = createServerCustomEvent( - MOSOO_CUSTOM_EVENT.sessionPermissionsUpdated.name, - { - permissionRequests: upsertPermissionRequest(nextLiveState.permissionRequests, request), - }, - ); - - nextLiveState = applyAgUiEventToSessionLiveState(nextLiveState, permissionsUpdatedEvent); - appendSessionDeliveryEvent(envelope, event, permissionsUpdatedEvent); - liveStateChanged = true; - } + const transition = readRuntimeDriverRunTransition(event); - continue; + if (transition !== null) { + transitions.push(transition); } - if (event.kind === "permission.resolved") { - const payload = readRuntimeEventPayload(event); - const requestId = readRuntimeEventString(payload, "requestId"); - const permissionRequests = - readPermissionRequestViews(payload["permissionRequests"]) ?? - (requestId === null - ? null - : removePermissionRequest(nextLiveState.permissionRequests, requestId)); - - if (permissionRequests !== null) { - const permissionsUpdatedEvent = createServerCustomEvent( - MOSOO_CUSTOM_EVENT.sessionPermissionsUpdated.name, - { - permissionRequests, - }, - ); - - nextLiveState = applyAgUiEventToSessionLiveState(nextLiveState, permissionsUpdatedEvent); - appendSessionDeliveryEvent(envelope, event, permissionsUpdatedEvent); - liveStateChanged = true; - } - + if (!projectLiveState) { continue; } - const liveEvents = [ - ...createPendingToolResultEvents(nextLiveState, event), - ...appRuntimeEventToSessionDeliveryEvents(event), - ]; - - const setSessionTitle = (title: string | null): void => { - sessionTitle = title; - }; - const setUsage = (nextUsage: AppRuntimeDriverEventsResult["usage"]): void => { - usage = nextUsage; - }; - - appendRuntimeDriverCanonicalSideEffects(event, { - setSessionTitle, - setUsage, - transitions, - }); + const liveEvents = appRuntimeEventToSessionDeliveryEvents(event); for (const liveEvent of liveEvents) { nextLiveState = applyAgUiEventToSessionLiveState(nextLiveState, liveEvent); @@ -601,21 +157,21 @@ export async function appRuntimeDriverEvents( } return { - finalAssistantMessage, + finalAssistantMessageId, link, liveStateChanged, nextLiveState, runtimeEvents, - sessionTitle, sessionDeliveryEvents, transitions, - usage, }; } -// Driver Contract v2 carries envelope occurredAt as an ISO 8601 string; +// Driver Contract v3 carries envelope occurredAt as an ISO 8601 string; // persistence keeps epoch milliseconds. Unparsable or absent values stay null. -function toDriverEventOccurredAtMs(occurredAt: DriverEventEnvelope["occurredAt"]): number | null { +function toDriverEventOccurredAtMs( + occurredAt: CanonicalDriverEventEnvelope["occurredAt"], +): number | null { if (typeof occurredAt !== "string") { return null; } @@ -625,45 +181,21 @@ function toDriverEventOccurredAtMs(occurredAt: DriverEventEnvelope["occurredAt"] } function resolveDriverEventPersistenceSourceId( - source: DriverEventEnvelope, + source: CanonicalDriverEventEnvelope, event: RuntimeEventEnvelope, ): string | null { - if (event.kind === "run.failed" && event.runId !== undefined) { - return createSessionRunTerminalFailureSourceId(event.runId); + if ( + (event.kind === "run.cancelled" || + event.kind === "run.completed" || + event.kind === "run.failed") && + event.runId !== undefined + ) { + return createSessionRunTerminalSourceId(event.runId, event.kind); } return source.eventId.trim().length > 0 ? source.eventId : null; } -function appendRuntimeDriverCanonicalSideEffects( - event: RuntimeEventEnvelope, - output: { - setSessionTitle: (title: string | null) => void; - setUsage: (usage: AppRuntimeDriverEventsResult["usage"]) => void; - transitions: RuntimeDriverRunTransition[]; - }, -): void { - if (event.kind === "session.info.updated") { - output.setSessionTitle( - normalizeRuntimeSessionInfoTitle( - readRuntimeEventString(readRuntimeEventPayload(event), "title"), - ), - ); - return; - } - - if (event.kind === "usage.updated") { - output.setUsage(parseNullableSessionUsageSummary(event.payload)); - return; - } - - const transition = readRuntimeDriverRunTransition(event); - - if (transition !== null) { - output.transitions.push(transition); - } -} - async function loadStoredRuntimeLiveState( database: D1Database, input: { diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/http.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/http.ts index ef7df8ec..6da81594 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/http.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/http.ts @@ -1,9 +1,8 @@ +import { parseRuntimeCommand } from "@mosoo/contracts/runtime-command"; import type { RuntimeCommand } from "@mosoo/contracts/runtime-command"; import { json, readPositiveTimeout, toErrorMessage } from "./driver-instance-support"; import type { - DriverInstanceHeartbeatResult, - DriverInstanceHelloResult, DriverInstanceReadyResult, DriverInstanceSnapshot, DriverInstanceWaitForCloseResult, @@ -11,24 +10,29 @@ import type { export interface DriverInstanceHttpHandler { acceptDriverSocket(request: Request): Promise; - destroy(reason: string): Promise; - fail(message: string): Promise; - sendControlCommand(command: RuntimeCommand): Promise; + destroy(generation: number, reason: string): Promise; + fail(generation: number, message: string): Promise; + sendControlCommand(generation: number, command: RuntimeCommand): Promise; snapshot(): DriverInstanceSnapshot; - waitForClose(timeoutMs: number): Promise; - waitForHeartbeat(afterCount: number, timeoutMs: number): Promise; - waitForHello(timeoutMs: number): Promise; - waitForReady(timeoutMs: number): Promise; + waitForClose(generation: number, timeoutMs: number): Promise; + waitForReady(generation: number, timeoutMs: number): Promise; } interface RuntimeFailRequest { + generation: number; message?: string; } interface RuntimeCloseRequest { + generation: number; reason?: string; } +interface RuntimeSendRequest { + command: RuntimeCommand; + generation: number; +} + function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); } @@ -62,25 +66,46 @@ function readOptionalString( function parseFailRequest(value: unknown): RuntimeFailRequest { const requestName = "RuntimeFailRequest"; const record = readObject(value, requestName); + const generation = record["generation"]; const message = readOptionalString(record, "message", requestName); - return message === undefined ? {} : { message }; + if (!Number.isSafeInteger(generation) || (generation as number) < 0) { + throw new TypeError(`${requestName}.generation must be a non-negative safe integer.`); + } + + return message === undefined + ? { generation: generation as number } + : { generation: generation as number, message }; } function parseCloseRequest(value: unknown): RuntimeCloseRequest { const requestName = "RuntimeCloseRequest"; const record = readObject(value, requestName); + const generation = record["generation"]; const reason = readOptionalString(record, "reason", requestName); - return reason === undefined ? {} : { reason }; + if (!Number.isSafeInteger(generation) || (generation as number) < 0) { + throw new TypeError(`${requestName}.generation must be a non-negative safe integer.`); + } + + return reason === undefined + ? { generation: generation as number } + : { generation: generation as number, reason }; } -function parseRuntimeCommand(value: unknown): RuntimeCommand { - if (!isRecord(value) || typeof value["kind"] !== "string") { - throw new TypeError("Runtime command must be an object with a string kind."); +function parseSendRequest(value: unknown): RuntimeSendRequest { + const requestName = "RuntimeSendRequest"; + const record = readObject(value, requestName); + const generation = record["generation"]; + + if (!Number.isSafeInteger(generation) || (generation as number) < 0) { + throw new TypeError(`${requestName}.generation must be a non-negative safe integer.`); } - return value as RuntimeCommand; + return { + command: parseRuntimeCommand(record["command"]), + generation: generation as number, + }; } async function readOptionalJsonBody(request: Request): Promise { @@ -88,6 +113,17 @@ async function readOptionalJsonBody(request: Request): Promise { return text.trim().length === 0 ? {} : JSON.parse(text); } +function readGeneration(url: URL): number { + const value = url.searchParams.get("generation"); + const generation = value === null || value.trim().length === 0 ? Number.NaN : Number(value); + + if (!Number.isSafeInteger(generation) || generation < 0) { + throw new TypeError("generation must be a non-negative safe integer."); + } + + return generation; +} + export async function handleDriverInstanceRequest( handler: DriverInstanceHttpHandler, request: Request, @@ -98,27 +134,12 @@ export async function handleDriverInstanceRequest( return handler.acceptDriverSocket(request); } - if (request.method === "GET" && url.pathname === "/wait/hello") { - return json(await handler.waitForHello(readPositiveTimeout(url, "hello"))); - } - if (request.method === "GET" && url.pathname === "/wait/ready") { - return json(await handler.waitForReady(readPositiveTimeout(url, "ready"))); - } - - if (request.method === "GET" && url.pathname === "/wait/heartbeat") { - const timeoutMs = readPositiveTimeout(url, "heartbeat"); - const afterCount = Number(url.searchParams.get("afterCount") ?? "0"); - - if (!Number.isInteger(afterCount) || afterCount < 0) { - return json({ error: "afterCount must be a non-negative integer." }, { status: 400 }); - } - - return json(await handler.waitForHeartbeat(afterCount, timeoutMs)); + return json(await handler.waitForReady(readGeneration(url), readPositiveTimeout(url, "ready"))); } if (request.method === "GET" && url.pathname === "/wait/close") { - return json(await handler.waitForClose(readPositiveTimeout(url, "close"))); + return json(await handler.waitForClose(readGeneration(url), readPositiveTimeout(url, "close"))); } if (request.method === "GET" && url.pathname === "/snapshot") { @@ -126,10 +147,10 @@ export async function handleDriverInstanceRequest( } if (request.method === "POST" && url.pathname === "/control/send") { - let command: RuntimeCommand; + let body: RuntimeSendRequest; try { - command = parseRuntimeCommand(await request.json()); + body = parseSendRequest(await request.json()); } catch (error) { return json( { @@ -139,7 +160,7 @@ export async function handleDriverInstanceRequest( ); } - await handler.sendControlCommand(command); + await handler.sendControlCommand(body.generation, body.command); return json({ ok: true }); } @@ -162,7 +183,7 @@ export async function handleDriverInstanceRequest( ? body.message : "Driver instance failed."; - await handler.fail(message); + await handler.fail(body.generation, message); return json({ ok: true }); } @@ -181,6 +202,7 @@ export async function handleDriverInstanceRequest( } await handler.destroy( + body.generation, typeof body.reason === "string" && body.reason.trim().length > 0 ? body.reason : "runtime.driver_instance.destroyed", diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/lifecycle.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/lifecycle.ts index 4ac96984..f3cafa5e 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/lifecycle.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/lifecycle.ts @@ -5,7 +5,7 @@ import type { } from "@mosoo/agent-driver/orpc"; import { driverInstancesTable } from "@mosoo/db"; import type { DriverInstanceId } from "@mosoo/id"; -import { and, eq, inArray, sql } from "drizzle-orm"; +import { and, eq, inArray, isNull, sql } from "drizzle-orm"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; import { getAppDatabase } from "../../../../platform/db/drizzle"; @@ -17,6 +17,8 @@ import { import { parseDriverTimestampMs, driverInstanceExpiresAt } from "./status"; import type { DriverInstanceStatus } from "./status"; +export type DriverInstanceProjectionOutcome = "applied" | "replay" | "conflict"; + export async function markDriverInstanceConnected( bindings: ApiBindings, input: { @@ -58,7 +60,7 @@ export async function recordDriverInstanceHello( generation: number; hello: DriverHelloInput; }, -): Promise { +): Promise { const now = currentTimestampMs(); if (input.hello === undefined) { @@ -81,12 +83,40 @@ export async function recordDriverInstanceHello( eq(driverInstancesTable.connectionId, input.connectionId), eq(driverInstancesTable.generation, input.generation), eq(driverInstancesTable.status, "connecting"), + isNull(driverInstancesTable.driverPid), + isNull(driverInstancesTable.driverStartedAt), + isNull(driverInstancesTable.driverVersion), ), ) .returning({ id: driverInstancesTable.id }) .get()) ?? null; - return row !== null; + if (row !== null) { + return "applied"; + } + + const startedAt = parseDriverTimestampMs(input.hello.startedAt, "Driver hello startedAt"); + const existing = + (await getAppDatabase(bindings.DB) + .select({ id: driverInstancesTable.id }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.connectionId, input.connectionId), + eq(driverInstancesTable.generation, input.generation), + inArray(driverInstancesTable.status, ["connecting", "ready"]), + eq(driverInstancesTable.driverPid, input.hello.pid), + eq(driverInstancesTable.driverStartedAt, startedAt), + eq(driverInstancesTable.driverVersion, input.hello.driverVersion), + eq(driverInstancesTable.protocolVersion, input.hello.protocolVersion), + eq(driverInstancesTable.runtime, input.hello.runtime), + ), + ) + .limit(1) + .get()) ?? null; + + return existing === null ? "conflict" : "replay"; } export async function markDriverInstanceReady( @@ -96,7 +126,7 @@ export async function markDriverInstanceReady( driverInstanceId: DriverInstanceId; generation: number; }, -): Promise { +): Promise { const now = currentTimestampMs(); const row = @@ -122,7 +152,27 @@ export async function markDriverInstanceReady( .returning({ id: driverInstancesTable.id }) .get()) ?? null; - return row !== null; + if (row !== null) { + return "applied"; + } + + const existing = + (await getAppDatabase(bindings.DB) + .select({ id: driverInstancesTable.id }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.connectionId, input.connectionId), + eq(driverInstancesTable.generation, input.generation), + eq(driverInstancesTable.status, "ready"), + eq(driverInstancesTable.driverPid, input.pid), + ), + ) + .limit(1) + .get()) ?? null; + + return existing === null ? "conflict" : "replay"; } export async function recordDriverInstanceHeartbeat( @@ -165,7 +215,7 @@ export async function finalizeDriverInstance( input: { closeCode?: number | null; closeReason?: string | null; - connectionId: string; + connectionId: string | null; connectedAt?: number | null; driverPid?: number | null; driverStartedAt?: string | null; @@ -175,8 +225,10 @@ export async function finalizeDriverInstance( lastHeartbeatAt?: string | null; status: Extract; }, -): Promise { +): Promise | null> { const completedAt = currentTimestampMs(); + const connectionMatches = + input.connectionId === null ? [] : [eq(driverInstancesTable.connectionId, input.connectionId)]; const driverStartedAt = typeof input.driverStartedAt === "string" && input.driverStartedAt.length > 0 @@ -209,7 +261,7 @@ export async function finalizeDriverInstance( .where( and( eq(driverInstancesTable.id, driverInstanceId), - eq(driverInstancesTable.connectionId, input.connectionId), + ...connectionMatches, eq(driverInstancesTable.generation, input.generation), inArray(driverInstancesTable.status, LIVE_DRIVER_INSTANCE_STATUSES), ), @@ -217,20 +269,42 @@ export async function finalizeDriverInstance( .returning({ id: driverInstancesTable.id }) .get()) ?? null; - return row !== null; + if (row !== null) { + return input.status; + } + + const existing = + (await getAppDatabase(bindings.DB) + .select({ status: driverInstancesTable.status }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, driverInstanceId), + ...connectionMatches, + eq(driverInstancesTable.generation, input.generation), + inArray(driverInstancesTable.status, ["stopped", "failed"]), + ), + ) + .limit(1) + .get()) ?? null; + + return existing?.status === "stopped" || existing?.status === "failed" ? existing.status : null; } -export async function getDriverInstanceStatus( +export async function getDriverInstanceLifecycleIdentity( bindings: ApiBindings, driverInstanceId: DriverInstanceId, -): Promise { +): Promise<{ generation: number; status: DriverInstanceStatus } | null> { const row = (await getAppDatabase(bindings.DB) - .select({ status: driverInstancesTable.status }) + .select({ + generation: driverInstancesTable.generation, + status: driverInstancesTable.status, + }) .from(driverInstancesTable) .where(eq(driverInstancesTable.id, driverInstanceId)) .limit(1) .get()) ?? null; - return row?.status ?? null; + return row; } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/live-driver-instance.repository.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/live-driver-instance.repository.ts index 18973423..af11dd4e 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/live-driver-instance.repository.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/live-driver-instance.repository.ts @@ -10,6 +10,7 @@ export async function listLiveDriverInstanceRefsForSandboxSessions( sandboxSessionIds: readonly SessionId[], ): Promise< { + generation: number; id: DriverInstanceId; sandboxSessionId: SessionId; }[] @@ -22,6 +23,7 @@ export async function listLiveDriverInstanceRefsForSandboxSessions( return getAppDatabase(database) .select({ + generation: driverInstancesTable.generation, id: driverInstancesTable.id, sandboxSessionId: driverInstancesTable.sandboxSessionId, }) @@ -34,11 +36,3 @@ export async function listLiveDriverInstanceRefsForSandboxSessions( ) .all(); } - -export async function listLiveDriverInstanceIdsForSandboxSessions( - database: D1Database, - sandboxSessionIds: readonly SessionId[], -): Promise { - const rows = await listLiveDriverInstanceRefsForSandboxSessions(database, sandboxSessionIds); - return rows.map((row) => row.id); -} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/maintenance.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/maintenance.ts index c9e4d52a..444556e8 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/maintenance.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/maintenance.ts @@ -1,5 +1,5 @@ -import { driverInstancesTable, externalToolEffectsTable } from "@mosoo/db"; -import { and, eq, inArray, isNull, lte, notExists, sql } from "drizzle-orm"; +import { driverInstancesTable, externalToolEffectsTable, sessionsTable } from "@mosoo/db"; +import { and, eq, inArray, isNotNull, isNull, lte, notExists, sql } from "drizzle-orm"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; import { getAppDatabase } from "../../../../platform/db/drizzle"; @@ -9,7 +9,9 @@ import { DRIVER_COLD_READY_TIMEOUT_MS, RUNTIME_SOCKET_TIMEOUT_MS, } from "../../domain/runtime-config"; +import { repairClaimedDriverStopsGlobally } from "../driver-session-stop.service"; import { driverInstanceExpiresAt } from "./status"; +import { repairTerminalDriverRuntimeCommandsGlobally } from "./terminal-run-release"; export async function cleanupDriverInstances(bindings: ApiBindings): Promise { const database = getAppDatabase(bindings.DB); @@ -77,11 +79,33 @@ export async function cleanupDriverInstances(bindings: ApiBindings): Promise["result"] + >, ): RuntimeCommandResult { return parseSchemaValue(RuntimeCommandResult, result); } @@ -51,11 +54,16 @@ export class DriverInstanceRpcCommandController { this.#dependencies = dependencies; } - async enqueueCommand(command: RuntimeCommand): Promise { + async enqueueCommand(driverGeneration: number, command: RuntimeCommand): Promise { const { env, state } = this.#dependencies; + if (state.requireDriverGeneration() !== driverGeneration) { + throw new Error("Driver generation is no longer current."); + } + await createRuntimeCommandRecord(env.DB, { command, + driverGeneration, driverInstanceId: state.requireDriverInstanceId(), expiresAt: currentTimestampPlus(COMMAND_LEASE_MS), }); @@ -76,19 +84,31 @@ export class DriverInstanceRpcCommandController { throw new Error("Driver instance id mismatch."); } const driverInstanceId = state.requireDriverInstanceId(); + const driverGeneration = state.requireDriverGeneration(); context.assertActiveConnection(); const commandId = parsePlatformId(input.commandId, "driver command id"); - const command = await getRuntimeCommandRecord(env.DB, driverInstanceId, commandId); + const command = await getRuntimeCommandRecord( + env.DB, + driverInstanceId, + driverGeneration, + commandId, + ); context.assertActiveConnection(); + const terminalPayload = + input.status === "failed" + ? { error: input.error } + : input.status === "completed" && input.result !== undefined + ? { result: toStoredRuntimeCommandResult(input.result) } + : {}; const updateOutcome = await updateRuntimeCommandRecord(env.DB, { commandId, deliveryConnectionId: context.connectionId, + driverGeneration, driverInstanceId, - ...(input.error === undefined ? {} : { error: input.error }), + ...terminalPayload, status: input.status, - ...(input.result === undefined ? {} : { result: toStoredRuntimeCommandResult(input.result) }), }); context.assertActiveConnection(); @@ -96,24 +116,9 @@ export class DriverInstanceRpcCommandController { throw new Error(`Runtime command status update rejected: ${updateOutcome.reason}.`); } - if ( - command?.payload.kind === "input.start" && - (input.status === "completed" || - input.status === "failed" || - input.status === "cancelled" || - input.status === "expired") - ) { - const release = releaseLinkedTerminalDriverInstanceSessionRun(env, driverInstanceId).catch( - (error: unknown) => { - this.#dependencies.withRuntimeLogContext(() => { - logError("runtime.terminal.lease_release.failed", { - ...createErrorLogContext(error), - driverInstanceId, - }); - }); - }, - ); - this.#dependencies.waitUntil(release); + if (command?.payload.kind === "input.start" && input.status !== "accepted") { + await releaseLinkedTerminalDriverInstanceSessionRun(env, driverInstanceId, driverGeneration); + context.assertActiveConnection(); } return { ok: true }; @@ -130,50 +135,57 @@ export class DriverInstanceRpcCommandController { } context.assertActiveConnection(); - return claimExternalToolEffect(env.DB, { + const claim = await claimExternalToolEffect(env.DB, { + claimToken: input.claimToken, commandId: parsePlatformId(input.commandId, "driver command id"), + driverGeneration: state.requireDriverGeneration(), driverInstanceId: state.requireDriverInstanceId(), }); + context.assertActiveConnection(); + return claim; } - async handleCompleteExternalToolEffect( - input: DriverExternalToolEffectCompleteInput, + async handleObserveExternalToolEffect( + input: DriverExternalToolEffectObserveInput, context: DriverInstanceRpcOperationContext, - ): Promise<{ ok: true }> { + ): Promise { const { env, state } = this.#dependencies; if (input.driverInstanceId !== state.requireDriverInstanceId()) { throw new Error("Driver instance id mismatch."); } context.assertActiveConnection(); - await completeExternalToolEffect(env.DB, { + + const observation = await observeExternalToolEffect(env.DB, { commandId: parsePlatformId(input.commandId, "driver command id"), + driverGeneration: state.requireDriverGeneration(), driverInstanceId: state.requireDriverInstanceId(), - ...(input.providerReceiptJson === undefined - ? {} - : { providerReceiptJson: input.providerReceiptJson }), - result: parseSchemaValue(McpExecuteCommandResult, input.result), }); context.assertActiveConnection(); - return { ok: true }; + return observation; } - async handleMarkExternalToolEffectUnknown( - input: DriverExternalToolEffectUnknownInput, + async handleSettleExternalToolEffect( + input: DriverExternalToolEffectSettleInput, context: DriverInstanceRpcOperationContext, - ): Promise<{ ok: true }> { + ): Promise { const { env, state } = this.#dependencies; if (input.driverInstanceId !== state.requireDriverInstanceId()) { throw new Error("Driver instance id mismatch."); } context.assertActiveConnection(); - await markExternalToolEffectUnknown(env.DB, { + + const settlement = await settleExternalToolEffect(env.DB, { + claimToken: input.claimToken, commandId: parsePlatformId(input.commandId, "driver command id"), + driverGeneration: state.requireDriverGeneration(), driverInstanceId: state.requireDriverInstanceId(), + effectId: parsePlatformId(input.effectId, "external tool effect id"), + settlement: parseSchemaValue(ExternalToolEffectSettlement, input.settlement), }); context.assertActiveConnection(); - return { ok: true }; + return settlement; } async handleNextCommand( @@ -195,6 +207,7 @@ export class DriverInstanceRpcCommandController { const record = await claimNextQueuedRuntimeCommandRecord( env.DB, driverInstanceId, + state.requireDriverGeneration(), context.connectionId, ); @@ -220,7 +233,7 @@ export class DriverInstanceRpcCommandController { async #markCommandDelivered( command: RuntimeCommand, context: DriverInstanceRpcOperationContext, - ): Promise { + ): Promise<"delivered" | "discarded" | "retry"> { const { env, state } = this.#dependencies; context.assertActiveConnection(); @@ -228,11 +241,18 @@ export class DriverInstanceRpcCommandController { const deliveryOutcome = await markRuntimeCommandRecordDelivered(env.DB, { commandId, connectionId: context.connectionId, + driverGeneration: state.requireDriverGeneration(), driverInstanceId: state.requireDriverInstanceId(), }); context.assertActiveConnection(); - return deliveryOutcome.kind === "applied"; + if (deliveryOutcome.kind === "applied") { + return "delivered"; + } + + return deliveryOutcome.kind === "rejected" && deliveryOutcome.reason === "inactive_session_run" + ? "discarded" + : "retry"; } async #nextCommand(context: DriverInstanceRpcOperationContext): Promise { diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-controller-dependencies.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-controller-dependencies.ts index 49b6c3de..451a149a 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-controller-dependencies.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-controller-dependencies.ts @@ -3,10 +3,11 @@ import type { RuntimeSessionViewCache } from "./runtime-session-view-cache"; import type { DriverInstanceRuntimeState } from "./runtime-state"; import type { SessionViewerEventDeliveryBuffer } from "./session-viewer-event-delivery-buffer"; import type { DriverInstanceSocketRegistry } from "./sockets"; +import type { DriverInstanceConnectionEpoch } from "./state"; export interface DriverInstanceRpcControllerDependencies { env: ApiBindings; - finalizeTerminalState: () => Promise; + finalizeTerminalState: (epoch: DriverInstanceConnectionEpoch) => Promise; sockets: DriverInstanceSocketRegistry; state: DriverInstanceRuntimeState; viewCache: RuntimeSessionViewCache; diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-controller.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-controller.ts index ab2db5af..3cd45f3c 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-controller.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-controller.ts @@ -1,12 +1,12 @@ import type { DriverCommandUpdateInput, DriverCompletionInput, - DriverEventBatchInput, DriverEventBatchOutput, DriverExternalToolEffectClaimInput, DriverExternalToolEffectClaimOutput, - DriverExternalToolEffectCompleteInput, - DriverExternalToolEffectUnknownInput, + DriverExternalToolEffectObserveInput, + DriverExternalToolEffectSettleInput, + DriverExternalToolEffectState, DriverFailureInput, DriverHeartbeatInput, DriverHelloInput, @@ -19,6 +19,7 @@ import type { } from "@mosoo/agent-driver/orpc"; import type { RuntimeCommand } from "@mosoo/contracts/runtime-command"; +import type { HostDriverEventBatchInput } from "./event-types"; import type { DriverInstanceRpcHandler, DriverInstanceRpcOperationContext } from "./rpc"; import { DriverInstanceRpcCommandController } from "./rpc-command-controller"; import type { DriverInstanceRpcControllerDependencies } from "./rpc-controller-dependencies"; @@ -39,8 +40,8 @@ export class DriverInstanceRpcController implements DriverInstanceRpcHandler { this.#terminal = new DriverInstanceRpcRunTerminalController(dependencies); } - async enqueueCommand(command: RuntimeCommand): Promise { - await this.#commands.enqueueCommand(command); + async enqueueCommand(driverGeneration: number, command: RuntimeCommand): Promise { + await this.#commands.enqueueCommand(driverGeneration, command); } async handleCommandUpdate( @@ -57,11 +58,11 @@ export class DriverInstanceRpcController implements DriverInstanceRpcHandler { return this.#commands.handleClaimExternalToolEffect(input, context); } - async handleCompleteExternalToolEffect( - input: DriverExternalToolEffectCompleteInput, + async handleObserveExternalToolEffect( + input: DriverExternalToolEffectObserveInput, context: DriverInstanceRpcOperationContext, - ): Promise<{ ok: true }> { - return this.#commands.handleCompleteExternalToolEffect(input, context); + ): Promise { + return this.#commands.handleObserveExternalToolEffect(input, context); } async handleCompleteRun( @@ -106,7 +107,7 @@ export class DriverInstanceRpcController implements DriverInstanceRpcHandler { } async handlePushEvents( - input: DriverEventBatchInput, + input: HostDriverEventBatchInput, context: DriverInstanceRpcOperationContext, ): Promise { return this.#events.handlePushEvents(input, context); @@ -119,11 +120,15 @@ export class DriverInstanceRpcController implements DriverInstanceRpcHandler { return this.#events.handlePushLogs(input, context); } - async handleMarkExternalToolEffectUnknown( - input: DriverExternalToolEffectUnknownInput, + async runAfterPendingEvents(operation: () => Promise): Promise { + return this.#events.runAfterPendingEvents(operation); + } + + async handleSettleExternalToolEffect( + input: DriverExternalToolEffectSettleInput, context: DriverInstanceRpcOperationContext, - ): Promise<{ ok: true }> { - return this.#commands.handleMarkExternalToolEffectUnknown(input, context); + ): Promise { + return this.#commands.handleSettleExternalToolEffect(input, context); } async handleReady( diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-event-ingestion-controller.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-event-ingestion-controller.ts index c226a97f..4969e5b5 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-event-ingestion-controller.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-event-ingestion-controller.ts @@ -1,33 +1,70 @@ -import type { DriverEventEnvelope } from "@mosoo/agent-driver/events"; +import { createMcpExecuteFailedEventIdentity } from "@mosoo/agent-driver/events"; import type { - DriverEventBatchInput, DriverEventBatchOutput, DriverEventReceipt, DriverLogBatchInput, DriverLogBatchOutput, } from "@mosoo/agent-driver/orpc"; +import { McpExecuteCommandResult } from "@mosoo/contracts/runtime-command"; +import { parseSchemaValue } from "@mosoo/contracts/validation"; import { parsePlatformId } from "@mosoo/id"; -import type { SessionRunId } from "@mosoo/id"; +import type { DriverCommandId, DriverInstanceId, SessionRunId } from "@mosoo/id"; +import { + createRuntimeEventSemanticHash, + isRuntimeEventRecord, + readRuntimeEventToolCallUpdate, +} from "@mosoo/runtime-events"; import { createErrorLogContext, logError } from "../../../../platform/cloudflare/logger"; -import type { SessionDeliveryEvent } from "../../../sessions/application/session-live-state.service"; +import { loadSessionViewerState } from "../../../sessions/application/session-live-state.service"; +import type { + SessionDeliveryEvent, + SessionLiveState, +} from "../../../sessions/application/session-live-state.service"; import { getSessionRuntimeEventSourceReceipts } from "../../../sessions/infrastructure/session-runtime-event-store.repository"; -import { createSessionRunTerminalFailureSourceId } from "../../domain/session-run-terminal-event-id"; +import type { SessionRuntimeEventSourceReceipt } from "../../../sessions/infrastructure/session-runtime-event-store.repository"; +import { createSessionRunTerminalSourceId } from "../../domain/session-run-terminal-event-id"; +import { getExternalToolEffectForCommand } from "../session-runs/external-tool-effect-store.repository"; +import { getRuntimeCommandRecord } from "../session-runs/runtime-command-store.repository"; import { EVENT_BATCH_MAX_SIZE, LOG_BATCH_MAX_SIZE } from "./connections"; +import { canonicalizeDriverEventEnvelope } from "./driver-event-canonicalization"; import { DriverEventTerminalGate } from "./driver-event-terminal-gate"; import { publishDriverLogBatch } from "./driver-log-batch-publisher"; import { runtimeSessionLinkNeedsRefresh } from "./event-types"; +import type { CanonicalDriverEventEnvelope, HostDriverEventBatchInput } from "./event-types"; import { getRuntimeSessionLink, + persistProjectedRuntimeDriverEventPrerequisites, persistProjectedRuntimeDriverEvents, + preflightProjectedRuntimeDriverEvents, appRuntimeDriverEvents, } from "./events"; import type { RuntimeSessionLink } from "./events"; import type { DriverInstanceRpcOperationContext } from "./rpc"; import type { DriverInstanceRpcControllerDependencies } from "./rpc-controller-dependencies"; -import { filterDurablyAcceptedRuntimeStreamReplays } from "./runtime-event-replay-filter"; +import { stageRuntimeArtifactEvents } from "./runtime-artifact-staging"; +import { assertActiveRuntimeSessionRun } from "./session-link.repository"; + +interface HashedDriverEvent { + readonly envelope: CanonicalDriverEventEnvelope; + readonly persistenceSourceId: string; + readonly semanticHash: string; +} -function summarizeDriverEvents(events: readonly DriverEventEnvelope[]) { +interface PreparedDriverEventBatch { + readonly uniqueOuterEvents: readonly HashedDriverEvent[]; + readonly uniquePersistenceEvents: readonly HashedDriverEvent[]; +} + +function toRuntimeArtifactEvent(event: HashedDriverEvent) { + return { + event: event.envelope.event, + semanticHash: event.semanticHash, + sourceEventId: event.persistenceSourceId, + }; +} + +function summarizeDriverEvents(events: readonly CanonicalDriverEventEnvelope[]) { return { eventCount: events.length, eventKinds: events.map((event) => event.event.kind).slice(0, 24), @@ -36,7 +73,7 @@ function summarizeDriverEvents(events: readonly DriverEventEnvelope[]) { } function resolveEventSessionRunId( - events: readonly DriverEventEnvelope[], + events: readonly { readonly event: { readonly runId?: unknown } }[], ): SessionRunId | undefined { let eventRunId: string | undefined; @@ -46,6 +83,9 @@ function resolveEventSessionRunId( if (candidateRunId === undefined) { continue; } + if (typeof candidateRunId !== "string") { + throw new TypeError("Driver event run id must be a string."); + } if (eventRunId !== undefined && eventRunId !== candidateRunId) { throw new Error("Event batch cannot contain events from multiple runs."); @@ -59,10 +99,16 @@ function resolveEventSessionRunId( : parsePlatformId(eventRunId, "driver event run id"); } -function resolveDriverEventPersistenceSourceId(event: DriverEventEnvelope): string { - if (event.event.kind === "run.failed" && event.event.runId !== undefined) { - return createSessionRunTerminalFailureSourceId( - parsePlatformId(event.event.runId, "driver failed event run id"), +function resolveDriverEventPersistenceSourceId(event: CanonicalDriverEventEnvelope): string { + if ( + (event.event.kind === "run.cancelled" || + event.event.kind === "run.completed" || + event.event.kind === "run.failed") && + event.event.runId !== undefined + ) { + return createSessionRunTerminalSourceId( + parsePlatformId(event.event.runId, "driver terminal event run id"), + event.event.kind, ); } @@ -82,7 +128,7 @@ export class DriverInstanceRpcEventIngestionController { } public async handlePushEvents( - input: DriverEventBatchInput, + input: HostDriverEventBatchInput, context: DriverInstanceRpcOperationContext, ): Promise { const { state } = this.#dependencies; @@ -107,7 +153,10 @@ export class DriverInstanceRpcEventIngestionController { const cachedLink = state.runtimeSessionLink; const eventSessionRunId = resolveEventSessionRunId(input.events); const shouldRefreshLink = - input.events.some((envelope) => envelope.event.kind === "run.started") || + input.events.some( + (envelope) => + envelope.event.kind === "run.started" || envelope.event.kind === "agent.tasks.replaced", + ) || runtimeSessionLinkNeedsRefresh(cachedLink) || (eventSessionRunId !== undefined && cachedLink?.sessionRunId !== eventSessionRunId); const link = await this.#getRuntimeSessionLink({ @@ -115,23 +164,68 @@ export class DriverInstanceRpcEventIngestionController { ...(eventSessionRunId === undefined ? {} : { sessionRunId: eventSessionRunId }), }); context.assertActiveConnection(); - const replayedReceipts = state.readProcessedDriverEventReceipts(input.events); - const candidateEvents = state.filterUnprocessedDriverEvents(input.events); - const durableReceipts = await this.#readPersistedEventReceipts(link, candidateEvents); + const canonicalInputEvents = input.events.map((event) => + canonicalizeDriverEventEnvelope(event, { traceId: link.traceId }), + ); + const preparedEvents = await prepareDriverEventBatch(canonicalInputEvents); + assertDriverEventBatchTerminalOrder( + preparedEvents.uniquePersistenceEvents.map((event) => event.envelope), + ); + const provenMcpCommandIds = await proveMcpDriverEvents(env.DB, { + driverGeneration: state.requireDriverGeneration(), + driverInstanceId, + events: preparedEvents.uniquePersistenceEvents, + }); + const durableReceiptsBySource = await this.#readPersistedEventReceipts( + link, + preparedEvents.uniquePersistenceEvents, + ); + const durableReceipts = projectDriverEventReceipts( + preparedEvents.uniqueOuterEvents, + durableReceiptsBySource, + ); context.assertActiveConnection(); - const durableEventIds = new Set( - durableReceipts.flatMap((receipt) => - receipt.eventId === undefined ? [] : [receipt.eventId], - ), + const pendingEvents = preparedEvents.uniquePersistenceEvents.filter( + (event) => !durableReceiptsBySource.has(event.persistenceSourceId), ); - const events = filterDurablyAcceptedRuntimeStreamReplays(candidateEvents, durableEventIds); - const replayedAccepted = [...replayedReceipts, ...durableReceipts]; - + const events = pendingEvents.map((event) => event.envelope); + const containsTerminalEvent = preparedEvents.uniquePersistenceEvents.some( + ({ envelope: { event } }) => + event.kind === "run.cancelled" || + event.kind === "run.completed" || + event.kind === "run.failed", + ); + let requiresDurableStateSync = durableReceipts.length > 0 || containsTerminalEvent; + let requiresDriverStateReload = requiresDurableStateSync && !containsTerminalEvent; if (events.length === 0) { - state.rememberProcessedDriverEventReceipts(replayedAccepted); - return { accepted: replayedAccepted }; + const replayState = await this.#readDurableStateForSync(link, requiresDriverStateReload); + context.assertActiveConnection(); + if (replayState !== null) { + viewCache.update(replayState); + } else if (containsTerminalEvent) { + // Terminal state is sealed in D1. Do not rebuild an unbounded final + // transcript in this hibernatable Driver DO merely to discard it. + viewCache.reset(); + } + if (requiresDurableStateSync) { + viewerEventDelivery.requestStateSync(link.sessionId); + } + return { accepted: durableReceipts }; } + const assertActiveRun = async (): Promise => { + if (eventSessionRunId === undefined || link.sessionId === null) { + return; + } + + await assertActiveRuntimeSessionRun(env.DB, { + driverInstanceId, + sessionId: link.sessionId, + sessionRunId: eventSessionRunId, + }); + }; + + await assertActiveRun(); const projection = await (async () => { try { return await appRuntimeDriverEvents(env, { @@ -139,7 +233,9 @@ export class DriverInstanceRpcEventIngestionController { currentLiveState: viewCache.currentState, driverInstanceId, events, + projectLiveState: !containsTerminalEvent, link, + provenMcpCommandIds, }); } catch (error) { logError("runtime.driver.events.projection_failed", { @@ -150,14 +246,50 @@ export class DriverInstanceRpcEventIngestionController { throw error; } })(); + await persistProjectedRuntimeDriverEventPrerequisites(env.DB, { + driverConnectionId: context.connectionId, + driverGeneration: state.requireDriverGeneration(), + driverInstanceId, + projection, + }); + await preflightProjectedRuntimeDriverEvents(env.DB, projection); + context.assertActiveConnection(); + const artifactProjections = await stageRuntimeArtifactEvents(env, { + driverFence: { + connectionId: context.connectionId, + driverInstanceId, + generation: state.requireDriverGeneration(), + sessionRunId: link.sessionRunId, + }, + events: pendingEvents.map(toRuntimeArtifactEvent), + link, + }); + if (artifactProjections.size > 0) { + requiresDurableStateSync = true; + requiresDriverStateReload = !containsTerminalEvent; + } + context.assertActiveConnection(); // An accepted source identity must already be durable. Buffering stream // fragments only in this hibernatable DO would acknowledge text that a // fresh instance cannot reconstruct, so persist every canonical event. - const persistenceRuntimeEvents = projection.runtimeEvents; + const persistenceRuntimeEvents = projection.runtimeEvents.map((record) => { + const artifact = + record.sourceEventId === null ? undefined : artifactProjections.get(record.sourceEventId); + return artifact === undefined + ? record + : { + ...record, + artifactAttemptId: artifact.attemptId, + artifactManifestJson: artifact.manifestJson, + artifactManifestSha256: artifact.manifestSha256, + }; + }); const commit = await (async () => { try { return await persistProjectedRuntimeDriverEvents(env, { + driverConnectionId: context.connectionId, + driverGeneration: state.requireDriverGeneration(), driverInstanceId, projection: { ...projection, @@ -181,22 +313,59 @@ export class DriverInstanceRpcEventIngestionController { })(); context.assertActiveConnection(); - if (commit.liveState) { - viewCache.update(commit.liveState); + const committedReceiptsBySource = await this.#readPersistedEventReceipts( + link, + preparedEvents.uniquePersistenceEvents, + ); + if (committedReceiptsBySource.size !== preparedEvents.uniquePersistenceEvents.length) { + throw new Error("Driver event commit did not produce every durable source receipt."); } - - viewerEventDelivery.enqueue( - projection.link.sessionId, - filterDurablyCommittedDeliveryEvents({ - persistenceEvents: persistenceRuntimeEvents, - persistedSourceEventIds: commit.persistedSourceEventIds, - sessionDeliveryEvents: projection.sessionDeliveryEvents, - }), + const persistedRuntimeEventSeqs = commit.persistedSourceEventIds.flatMap( + (sourceEventId): number[] => { + const seq = committedReceiptsBySource.get(sourceEventId)?.seq; + return seq === undefined ? [] : [seq]; + }, ); + const deliveryRuntimeEventSeqCursor = + persistedRuntimeEventSeqs.length === 0 ? null : Math.max(...persistedRuntimeEventSeqs); + const previousDeliveryRuntimeEventSeqCursor = + deliveryRuntimeEventSeqCursor === null || + persistedRuntimeEventSeqs.length === 0 || + deliveryRuntimeEventSeqCursor - Math.min(...persistedRuntimeEventSeqs) + 1 !== + persistedRuntimeEventSeqs.length + ? null + : Math.min(...persistedRuntimeEventSeqs) - 1; + + const replayState = await this.#readDurableStateForSync(link, requiresDriverStateReload); + context.assertActiveConnection(); + if (containsTerminalEvent) { + viewCache.reset(); + } else { + const committedLiveState = replayState ?? commit.liveState; + if (committedLiveState !== null) { + viewCache.update(committedLiveState); + } + } - const accepted = [...replayedAccepted, ...state.createDriverEventReceipts(events)]; + if (!requiresDurableStateSync) { + viewerEventDelivery.enqueue( + projection.link.sessionId, + filterDurablyCommittedDeliveryEvents({ + persistenceEvents: persistenceRuntimeEvents, + persistedSourceEventIds: commit.persistedSourceEventIds, + sessionDeliveryEvents: projection.sessionDeliveryEvents, + }), + deliveryRuntimeEventSeqCursor, + previousDeliveryRuntimeEventSeqCursor, + ); + } else { + viewerEventDelivery.requestStateSync(link.sessionId); + } - state.rememberProcessedDriverEventReceipts(accepted); + const accepted = projectDriverEventReceipts( + preparedEvents.uniqueOuterEvents, + committedReceiptsBySource, + ); return { accepted }; }); @@ -289,18 +458,16 @@ export class DriverInstanceRpcEventIngestionController { async #readPersistedEventReceipts( link: RuntimeSessionLink, - events: readonly DriverEventEnvelope[], - ): Promise { + events: readonly HashedDriverEvent[], + ): Promise> { if (link.sessionId === null) { - return []; + return new Map(); } - const sourceEventIds = events - .map(resolveDriverEventPersistenceSourceId) - .filter((eventId) => eventId.length > 0); + const sourceEventIds = events.map((event) => event.persistenceSourceId); if (sourceEventIds.length === 0) { - return []; + return new Map(); } const receiptsByEventId = await getSessionRuntimeEventSourceReceipts( @@ -310,27 +477,310 @@ export class DriverInstanceRpcEventIngestionController { sourceEventIds, }, ); - const receipts: DriverEventReceipt[] = []; - const seenEventIds = new Set(); - for (const event of events) { - if (event.eventId.length === 0 || seenEventIds.has(event.eventId)) { + const receipt = receiptsByEventId.get(event.persistenceSourceId); + + if (receipt === undefined) { continue; } - seenEventIds.add(event.eventId); + if ( + receipt.semanticHash === null || + receipt.semanticHash !== event.semanticHash || + receipt.type !== event.envelope.event.kind + ) { + throw new Error( + `Driver event source ${event.persistenceSourceId} conflicts with its durable receipt.`, + ); + } + } - const receipt = receiptsByEventId.get(resolveDriverEventPersistenceSourceId(event)); + return receiptsByEventId; + } - if (receipt === undefined) { - continue; + async #readDurableStateForSync( + link: RuntimeSessionLink, + required: boolean, + ): Promise { + if (link.sessionId === null || !required) { + return null; + } + const viewerId = link.callerId ?? link.creatorId; + if (viewerId === null) { + throw new Error("Durable Driver event replay is missing its Session viewer identity."); + } + + return loadSessionViewerState(this.#dependencies.env.DB, { + sessionId: link.sessionId, + viewerId, + }); + } +} + +function projectDriverEventReceipts( + events: readonly HashedDriverEvent[], + receiptsBySource: ReadonlyMap, +): DriverEventReceipt[] { + return events.flatMap((event) => { + const receipt = receiptsBySource.get(event.persistenceSourceId); + + return receipt === undefined + ? [] + : [ + { + eventId: event.envelope.eventId, + seq: receipt.seq, + type: event.envelope.event.kind, + }, + ]; + }); +} + +async function prepareDriverEventBatch( + events: readonly CanonicalDriverEventEnvelope[], +): Promise { + const hashedEvents = await Promise.all( + events.map(async (envelope) => ({ + envelope, + persistenceSourceId: resolveDriverEventPersistenceSourceId(envelope), + semanticHash: await createRuntimeEventSemanticHash(envelope.event), + })), + ); + const outerIdentities = new Map(); + const persistenceIdentities = new Map(); + + for (const event of hashedEvents) { + const outer = outerIdentities.get(event.envelope.eventId); + + if (outer !== undefined) { + assertMatchingDriverEventIdentity(outer, event, "outer"); + } else { + outerIdentities.set(event.envelope.eventId, event); + } + + const persisted = persistenceIdentities.get(event.persistenceSourceId); + + if (persisted !== undefined) { + assertMatchingDriverEventIdentity(persisted, event, "durable"); + } else { + persistenceIdentities.set(event.persistenceSourceId, event); + } + } + + return { + uniqueOuterEvents: [...outerIdentities.values()], + uniquePersistenceEvents: [...persistenceIdentities.values()], + }; +} + +function assertMatchingDriverEventIdentity( + expected: HashedDriverEvent, + actual: HashedDriverEvent, + identityKind: "durable" | "outer", +): void { + if ( + expected.persistenceSourceId !== actual.persistenceSourceId || + expected.semanticHash !== actual.semanticHash + ) { + throw new Error(`Driver event batch contains a conflicting ${identityKind} source identity.`); + } +} + +export function assertDriverEventBatchTerminalOrder( + events: readonly CanonicalDriverEventEnvelope[], +): void { + const terminalIndexes = events.flatMap((event, index) => + event.event.kind === "run.cancelled" || + event.event.kind === "run.completed" || + event.event.kind === "run.failed" + ? [index] + : [], + ); + + if (terminalIndexes.length > 1) { + throw new Error("Driver event batch cannot contain multiple run terminal events."); + } + + if (terminalIndexes[0] !== undefined && terminalIndexes[0] !== events.length - 1) { + throw new Error("Driver event batch run terminal event must be last."); + } +} + +function readMcpDriverEventIdentity(event: CanonicalDriverEventEnvelope): { + commandId: DriverCommandId; + sourceEventId: string; + status: "cancelled" | "completed" | "failed" | "running"; +} | null { + const sourceEventId = event.event.sourceEventId ?? event.eventId; + const reservedSource = sourceEventId.startsWith("mcp.execute."); + const mcpPayload = + event.event.kind === "tool.call.updated" && + readRuntimeEventToolCallUpdate(event.event).kind === "mcp"; + + if (!reservedSource && !mcpPayload) { + return null; + } + if (!reservedSource || !mcpPayload) { + throw new Error("MCP tool events require their reserved command source identity."); + } + + const commandId = parsePlatformId( + event.event.correlationId, + "MCP event command correlation ID", + ); + + for (const status of ["running", "completed", "failed", "cancelled"] as const) { + const prefix = `mcp.execute.${status}:`; + + if (!sourceEventId.startsWith(prefix)) { + continue; + } + + if (event.event.sourceEventId !== sourceEventId || event.eventId !== sourceEventId) { + throw new Error("MCP event transport identity does not match its command."); + } + + if ( + status !== "failed" && + commandId !== + parsePlatformId( + sourceEventId.slice(prefix.length), + "MCP event source command ID", + ) + ) { + throw new Error("MCP event transport identity does not match its command."); + } + + return { commandId, sourceEventId, status }; + } + + throw new Error("MCP event source identity has an unsupported terminal status."); +} + +async function proveMcpDriverEvents( + database: D1Database, + input: { + driverGeneration: number; + driverInstanceId: DriverInstanceId; + events: readonly HashedDriverEvent[]; + }, +): Promise> { + const proven = new Map(); + + for (const { envelope } of input.events) { + const identity = readMcpDriverEventIdentity(envelope); + + if (identity === null) { + continue; + } + + if (envelope.event.kind !== "tool.call.updated") { + throw new Error("MCP command source identity requires a tool.call.updated event."); + } + + const record = await getRuntimeCommandRecord( + database, + input.driverInstanceId, + input.driverGeneration, + identity.commandId, + ); + + if (record === null || record.kind !== "mcp.execute") { + throw new Error("MCP event does not have an immutable command in this Driver generation."); + } + + if ( + record.status !== "accepted" && + (identity.status === "running" || record.status !== identity.status) + ) { + throw new Error("MCP event conflicts with the command's durable lifecycle status."); + } + + const command = record.payload; + const toolCall = readRuntimeEventToolCallUpdate(envelope.event); + const effect = await getExternalToolEffectForCommand(database, { + commandId: identity.commandId, + driverGeneration: input.driverGeneration, + driverInstanceId: input.driverInstanceId, + }); + + if ( + effect === null || + envelope.event.driverInstanceId !== input.driverInstanceId || + envelope.event.runId !== command.runId || + toolCall.content !== null || + toolCall.kind !== "mcp" || + toolCall.messageId !== null || + toolCall.parentMessageId !== null || + toolCall.rawInput !== command.argumentsJson || + toolCall.rawInputDelta !== null || + toolCall.rawOutputDelta !== null || + toolCall.status !== identity.status || + toolCall.title !== command.toolName || + toolCall.toolCallId !== command.toolCallId + ) { + throw new Error("MCP event does not match its immutable command intent."); + } + + if (identity.status === "completed") { + if (effect.status !== "succeeded" || effect.resultJson === null) { + throw new Error("MCP completion event has no succeeded durable external effect."); } - receipts.push({ ...receipt, eventId: event.eventId }); + const result = parseSchemaValue(McpExecuteCommandResult, JSON.parse(effect.resultJson)); + + if ( + result.requestId !== command.requestId || + result.serverId !== command.serverId || + result.toolName !== command.toolName || + toolCall.rawOutput !== result.outputText || + (record.status === "completed" && + (record.result === null || + record.result.outputText !== result.outputText || + record.result.requestId !== result.requestId || + record.result.serverId !== result.serverId || + record.result.toolName !== result.toolName || + (record.result.isError ?? false) !== (result.isError ?? false))) + ) { + throw new Error("MCP completion event conflicts with its durable external effect result."); + } + } else if (identity.status === "failed") { + if ( + (effect.status !== "intent" && effect.status !== "unknown") || + toolCall.rawOutput === null || + toolCall.rawOutput.length === 0 || + (record.status === "failed" && record.error?.message !== toolCall.rawOutput) + ) { + throw new Error("MCP failure event conflicts with its durable external effect state."); + } + + const failedEvent = createMcpExecuteFailedEventIdentity({ + commandId: identity.commandId, + rawInput: command.argumentsJson, + rawOutput: toolCall.rawOutput, + title: command.toolName, + toolCallId: command.toolCallId, + }); + const payload = envelope.event.payload; + + if ( + identity.sourceEventId !== failedEvent.sourceEventId || + !isRuntimeEventRecord(payload) || + Object.keys(payload).length !== Object.keys(failedEvent.payload).length || + Object.entries(failedEvent.payload).some(([key, value]) => payload[key] !== value) + ) { + throw new Error("MCP failure event does not have its canonical content identity."); + } + } else if (toolCall.rawOutput !== null || effect.status !== "intent") { + throw new Error("Unsettled MCP event contains an impossible external effect output."); } - return receipts; + if (identity.status !== "running") { + proven.set(envelope.eventId, identity.commandId); + } } + + return proven; } function filterDurablyCommittedDeliveryEvents(input: { diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-handshake-controller.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-handshake-controller.ts index 0d4b892b..d6b6e81b 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-handshake-controller.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-handshake-controller.ts @@ -5,7 +5,6 @@ import type { DriverReadyInput, } from "@mosoo/agent-driver/orpc"; import { SANDBOX_ORGANIZATION_ROOT } from "@mosoo/agent-driver/paths"; -import { createPlatformId } from "@mosoo/id"; import { logInfo } from "../../../../platform/cloudflare/logger"; import { DRIVER_HEARTBEAT_INTERVAL_MS } from "../../domain/runtime-config"; @@ -67,31 +66,57 @@ export class DriverInstanceRpcHandshakeController { context: DriverInstanceRpcOperationContext, ): Promise { const { env, state, withRuntimeLogContext } = this.#dependencies; + context.assertActiveConnection(); + + let link: RuntimeSessionLink | null = null; + let output = state.helloOutput ?? state.pendingHello?.output ?? null; - if (state.hello) { - throw new Error("Driver hello has already been received."); + if (output === null) { + link = await this.#getRuntimeSessionLink(context); + context.assertActiveConnection(); + output = { + acceptedCapabilities: input.capabilities, + connectionId: context.connectionId, + driverInstanceId: state.requireDriverInstanceId(), + heartbeatIntervalMs: DRIVER_HEARTBEAT_INTERVAL_MS, + runConfig: { + commandLeaseMs: COMMAND_LEASE_MS, + envPolicy: "strict", + eventBatchMaxSize: EVENT_BATCH_MAX_SIZE, + organizationPath: SANDBOX_ORGANIZATION_ROOT, + }, + runId: link.sessionRunId, + }; } + + const staged = await state.stageHello(context.epoch, input, output); context.assertActiveConnection(); - const recorded = await recordDriverInstanceHello(env, { + if (staged === "replay") { + return output; + } + + const projection = await recordDriverInstanceHello(env, { connectionId: context.connectionId, driverInstanceId: state.requireDriverInstanceId(), - generation: state.requireDriverGeneration(), + generation: context.epoch.generation, hello: input, }); - if (!recorded) { - throw new Error("Driver connection is no longer current."); + if (projection === "conflict") { + throw new Error("Driver hello conflicts with the lifecycle projection."); } context.assertActiveConnection(); - const result = await state.recordHello(input); - state.resolveHelloWaiters(result); + const committedOutput = await state.commitHello(context.epoch); + context.assertActiveConnection(); - const link = await this.#getRuntimeSessionLink(); + link ??= await this.#getRuntimeSessionLink(context); + context.assertActiveConnection(); if (state.traceId === null && link.traceId !== null) { - await state.setTraceId(link.traceId); + await state.setTraceId(link.traceId, context.epoch); + context.assertActiveConnection(); } withRuntimeLogContext(() => { @@ -105,19 +130,7 @@ export class DriverInstanceRpcHandshakeController { }); }); - return { - acceptedCapabilities: input.capabilities, - connectionId: state.connectionId ?? createPlatformId(), - driverInstanceId: state.requireDriverInstanceId(), - heartbeatIntervalMs: DRIVER_HEARTBEAT_INTERVAL_MS, - runConfig: { - commandLeaseMs: COMMAND_LEASE_MS, - envPolicy: "strict", - eventBatchMaxSize: EVENT_BATCH_MAX_SIZE, - organizationPath: SANDBOX_ORGANIZATION_ROOT, - }, - runId: link.sessionRunId, - }; + return committedOutput; } async handleReady( @@ -134,25 +147,29 @@ export class DriverInstanceRpcHandshakeController { throw new Error("Driver hello is required before ready."); } - if (state.ready) { - throw new Error("Driver ready has already been received."); - } context.assertActiveConnection(); - const markedReady = await markDriverInstanceReady(env, { + const staged = await state.stageReady(context.epoch, input); + context.assertActiveConnection(); + + if (staged === "replay") { + return { ok: true }; + } + + const projection = await markDriverInstanceReady(env, { ...input, connectionId: context.connectionId, driverInstanceId: state.requireDriverInstanceId(), - generation: state.requireDriverGeneration(), + generation: context.epoch.generation, }); - if (!markedReady) { - throw new Error("Driver connection is no longer current."); + if (projection === "conflict") { + throw new Error("Driver ready conflicts with the lifecycle projection."); } context.assertActiveConnection(); - const result = await state.recordReady(input); - state.resolveReadyWaiters(result); + const result = await state.commitReady(context.epoch); + state.resolveReadyWaiters(result, context.epoch.generation); withRuntimeLogContext(() => { logInfo("runtime.driver.ready.received", { @@ -165,7 +182,10 @@ export class DriverInstanceRpcHandshakeController { return { ok: true }; } - async #getRuntimeSessionLink(options: { refresh?: boolean } = {}): Promise { + async #getRuntimeSessionLink( + context: DriverInstanceRpcOperationContext, + options: { refresh?: boolean } = {}, + ): Promise { const { env, state } = this.#dependencies; if (options.refresh !== true && state.runtimeSessionLink !== null) { @@ -173,6 +193,7 @@ export class DriverInstanceRpcHandshakeController { } const link = await getRuntimeSessionLink(env.DB, state.requireDriverInstanceId()); + context.assertActiveConnection(); state.setRuntimeSessionLink(link); return link; } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-run-terminal-controller.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-run-terminal-controller.ts index 126fd5f5..b94a9b2e 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-run-terminal-controller.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-run-terminal-controller.ts @@ -1,7 +1,8 @@ import type { DriverCompletionInput, DriverFailureInput } from "@mosoo/agent-driver/orpc"; +import { parsePlatformId } from "@mosoo/id"; +import type { SessionRunId } from "@mosoo/id"; import { logError, logInfo } from "../../../../platform/cloudflare/logger"; -import { syncSessionViewerState } from "../../../sessions/application/session-viewer-events.service"; import { runtimeSessionLinkNeedsRefresh } from "./event-types"; import { getRuntimeSessionLink, @@ -36,14 +37,23 @@ export class DriverInstanceRpcRunTerminalController { throw new Error("Driver instance id mismatch."); } const driverInstanceId = state.requireDriverInstanceId(); + const sessionRunId = readTerminalRunId(input); context.assertActiveConnection(); - await viewerEventDelivery.flushSafely(); - context.assertActiveConnection(); await recordDriverInstanceCompletion(env, { + driverConnectionId: context.connectionId, + driverGeneration: context.epoch.generation, driverInstanceId, - driverReady: state.hello !== null, + sessionRunId, + }); + + await state.setTerminalSessionRunId(sessionRunId, context.epoch); + + const link = await this.#getRuntimeSessionLink(context, { + refresh: runtimeSessionLinkNeedsRefresh(state.runtimeSessionLink), + sessionRunId, }); + viewerEventDelivery.requestStateSync(link.sessionId); context.assertActiveConnection(); withRuntimeLogContext(() => { @@ -54,24 +64,22 @@ export class DriverInstanceRpcRunTerminalController { }); }); - const socket = sockets.getDriverSocket(); + const socket = sockets.getDriverSocket(context.epoch); if (socket && socket.readyState === WebSocket.OPEN) { - sockets.scheduleDriverSocketClose(1000, "runtime.completed"); + sockets.scheduleDriverSocketClose(context.epoch, 1000, "runtime.completed"); } else { - await state.persistClose({ - at: new Date().toISOString(), - code: 1000, - reason: "runtime.completed", - }); - await finalizeTerminalState(); + await state.persistClose( + { + at: new Date().toISOString(), + code: 1000, + reason: "runtime.completed", + }, + context.epoch, + ); + await finalizeTerminalState(context.epoch); } - const link = await this.#getRuntimeSessionLink({ - refresh: runtimeSessionLinkNeedsRefresh(state.runtimeSessionLink), - }); - await syncSessionViewerState(env, link.sessionId); - return { ok: true }; } @@ -92,18 +100,25 @@ export class DriverInstanceRpcRunTerminalController { throw new Error("Driver instance id mismatch."); } const driverInstanceId = state.requireDriverInstanceId(); + const sessionRunId = readTerminalRunId(input); context.assertActiveConnection(); - await viewerEventDelivery.flushSafely(); - context.assertActiveConnection(); - const link = await this.#getRuntimeSessionLink({ + const link = await this.#getRuntimeSessionLink(context, { refresh: runtimeSessionLinkNeedsRefresh(state.runtimeSessionLink), + sessionRunId, }); await recordDriverInstanceFailure(env, { + driverConnectionId: context.connectionId, + driverGeneration: context.epoch.generation, driverInstanceId, error: input.error, link, + sessionRunId, }); + + await state.setTerminalSessionRunId(sessionRunId, context.epoch); + + viewerEventDelivery.requestStateSync(link.sessionId); context.assertActiveConnection(); withRuntimeLogContext(() => { @@ -117,30 +132,47 @@ export class DriverInstanceRpcRunTerminalController { }); }); - await state.setErrorMessage(input.error.message); + await state.setConnectionErrorMessage(context.epoch, input.error.message); + context.assertActiveConnection(); - const socket = sockets.getDriverSocket(); + const socket = sockets.getDriverSocket(context.epoch); if (socket && socket.readyState === WebSocket.OPEN) { - sockets.scheduleDriverSocketClose(1011, "runtime.failed"); + sockets.scheduleDriverSocketClose(context.epoch, 1011, "runtime.failed"); } else { - await finalizeTerminalState(); + await finalizeTerminalState(context.epoch); } - await syncSessionViewerState(env, link.sessionId); - return { ok: true }; } - async #getRuntimeSessionLink(options: { refresh?: boolean } = {}): Promise { + async #getRuntimeSessionLink( + context: DriverInstanceRpcOperationContext, + options: { refresh?: boolean; sessionRunId: SessionRunId }, + ): Promise { const { env, state } = this.#dependencies; - if (options.refresh !== true && state.runtimeSessionLink !== null) { + if ( + options.refresh !== true && + state.runtimeSessionLink?.sessionRunId === options.sessionRunId + ) { return state.runtimeSessionLink; } - const link = await getRuntimeSessionLink(env.DB, state.requireDriverInstanceId()); + const link = await getRuntimeSessionLink(env.DB, state.requireDriverInstanceId(), { + sessionRunId: options.sessionRunId, + }); + context.assertActiveConnection(); + + if (link.sessionRunId !== options.sessionRunId) { + throw new Error("Terminal Driver Session Run identity does not match the request."); + } + state.setRuntimeSessionLink(link); return link; } } + +function readTerminalRunId(input: DriverCompletionInput | DriverFailureInput): SessionRunId { + return parsePlatformId(input.runId, "terminal Driver Session Run id"); +} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-wire.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-wire.ts index afbde39b..a4cc7d6b 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-wire.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc-wire.ts @@ -1,13 +1,12 @@ -import { parseDriverEventEnvelope } from "@mosoo/agent-driver/events"; import type { DriverCommandUpdateInput, DriverCompletionInput, - DriverEventBatchInput, DriverEventBatchOutput, DriverExternalToolEffectClaimInput, DriverExternalToolEffectClaimOutput, - DriverExternalToolEffectCompleteInput, - DriverExternalToolEffectUnknownInput, + DriverExternalToolEffectObserveInput, + DriverExternalToolEffectSettleInput, + DriverExternalToolEffectState, DriverFailureInput, DriverHeartbeatInput, DriverHeartbeatOutput, @@ -20,156 +19,235 @@ import type { DriverReadyInput, } from "@mosoo/agent-driver/orpc"; import { DriverCapability } from "@mosoo/contracts/driver-instance"; -import { ExternalToolEffectClaim } from "@mosoo/contracts/external-tool-effect"; import { - RuntimeCommand, + ExternalToolEffectClaim, + ExternalToolEffectClaimToken, + ExternalToolEffectSettlement, + ExternalToolEffectState, +} from "@mosoo/contracts/external-tool-effect"; +import { + InputStartCommandResult, McpExecuteCommandResult, - RuntimeCommandResult, - RuntimeCommandStatus, + RuntimeCommand, } from "@mosoo/contracts/runtime-command"; -import { RunError } from "@mosoo/contracts/session-run"; +import { DurableRunError } from "@mosoo/contracts/session-run"; import { NonEmptyString, PrimitiveRecord, parseSchemaValue } from "@mosoo/contracts/validation"; +import { + RUNTIME_EVENT_KINDS, + RUNTIME_EVENT_SCHEMA_VERSION, + parseRuntimeEventEnvelope, +} from "@mosoo/runtime-events"; import { eventIterator, os } from "@orpc/server"; import { type } from "arktype"; +import { EVENT_BATCH_MAX_SIZE, LOG_BATCH_MAX_SIZE } from "./connections"; +import type { HostDriverEventBatchInput } from "./event-types"; + +const PositiveSafeInteger = type("number.integer >= 1 & number.safe"); +const NonNegativeSafeInteger = type("number.integer >= 0 & number.safe"); +const DriverEventBatchMaxSizeWire = PositiveSafeInteger.narrow((size, context) => + size <= EVENT_BATCH_MAX_SIZE + ? true + : context.reject({ expected: `at most ${EVENT_BATCH_MAX_SIZE}` }), +); +const DriverCapabilitiesWire = DriverCapability.array().narrow((capabilities, context) => + new Set(capabilities.map(({ id }) => id)).size === capabilities.length + ? true + : context.reject({ expected: "unique capability ids" }), +); + const DriverHelloInputWire = type({ - capabilities: DriverCapability.array(), + capabilities: DriverCapabilitiesWire, driverVersion: NonEmptyString, - pid: "number", - protocolVersion: "2", + pid: PositiveSafeInteger, + protocolVersion: "3", runtime: '"openai-runtime" | "claude-agent-sdk" | "acp-fallback"', - startedAt: "string", -}); + startedAt: NonEmptyString, +}).onUndeclaredKey("reject"); + +const DriverRunConfigWire = type({ + commandLeaseMs: NonNegativeSafeInteger, + envPolicy: '"strict"', + eventBatchMaxSize: DriverEventBatchMaxSizeWire, + organizationPath: NonEmptyString, +}).onUndeclaredKey("reject"); const DriverHelloOutputWire = type({ - acceptedCapabilities: DriverCapability.array(), + acceptedCapabilities: DriverCapabilitiesWire, connectionId: NonEmptyString, driverInstanceId: NonEmptyString, - heartbeatIntervalMs: "number >= 250", - runConfig: { - commandLeaseMs: "number >= 0", - envPolicy: '"strict"', - eventBatchMaxSize: "number >= 0", - organizationPath: NonEmptyString, - }, - runId: "string | null", -}); + heartbeatIntervalMs: "number.integer >= 250 & number.safe", + runConfig: DriverRunConfigWire, + runId: type("null").or(NonEmptyString), +}).onUndeclaredKey("reject"); const DriverHeartbeatInputWire = type({ - at: "string", - pid: "number", + at: NonEmptyString, + pid: PositiveSafeInteger, reason: '"interval" | "ping"', -}); +}).onUndeclaredKey("reject"); const DriverHeartbeatOutputWire = type({ - heartbeatCount: "number >= 0", + heartbeatCount: NonNegativeSafeInteger, ok: "true", -}); +}).onUndeclaredKey("reject"); const DriverReadyInputWire = type({ at: NonEmptyString, driverInstanceId: NonEmptyString, - pid: "number", -}); + pid: PositiveSafeInteger, +}).onUndeclaredKey("reject"); + +const RuntimeEventKindWire = type.enumerated(...RUNTIME_EVENT_KINDS); +const DriverRuntimeEventWire = type({ + actor: '"agent" | "api" | "driver" | "system" | "tool" | "user"', + "correlationId?": "string", + delivery: '"best_effort" | "lossless"', + "driverInstanceId?": "string", + id: NonEmptyString, + kind: RuntimeEventKindWire, + "native?": "unknown", + occurredAt: NonEmptyString, + origin: '"api" | "driver" | "file" | "runtime" | "system" | "viewer"', + payload: "unknown", + "receivedAt?": "string", + "runId?": "string", + "runtimeId?": "string", + schemaVersion: `'${RUNTIME_EVENT_SCHEMA_VERSION}'`, + sessionId: NonEmptyString, + "sourceEventId?": "string", + "traceId?": "string", + visibility: '"owner_debug" | "participant" | "public" | "system_internal"', +}).onUndeclaredKey("reject"); const DriverEventEnvelopeWire = type({ - event: "unknown", + event: DriverRuntimeEventWire, eventId: NonEmptyString, "occurredAt?": "string | null | undefined", -}); +}).onUndeclaredKey("reject"); const DriverEventReceiptWire = type({ - "eventId?": "string | undefined", - seq: "number >= 0", - type: NonEmptyString, -}); + eventId: NonEmptyString, + seq: NonNegativeSafeInteger, + type: RuntimeEventKindWire, +}).onUndeclaredKey("reject"); const DriverEventBatchInputWire = type({ driverInstanceId: NonEmptyString, - events: DriverEventEnvelopeWire.array(), -}); + events: DriverEventEnvelopeWire.array().atMostLength(EVENT_BATCH_MAX_SIZE), +}).onUndeclaredKey("reject"); const DriverEventBatchOutputWire = type({ accepted: DriverEventReceiptWire.array(), -}); +}).onUndeclaredKey("reject"); const DriverLogContextWire = type({ "parentSpanId?": "string", "requestId?": "string", "sandboxId?": "string", "sessionId?": "string", - "spanId?": NonEmptyString, - "traceId?": NonEmptyString, -}); + "spanId?": "string", + "traceId?": "string", +}).onUndeclaredKey("reject"); const DriverLogErrorWire = type({ "code?": "string | number", - message: NonEmptyString, - name: NonEmptyString, + message: "string", + name: "string", "stack?": "string | null", -}); +}).onUndeclaredKey("reject"); const DriverLogEntryWire = type({ "context?": DriverLogContextWire, "error?": DriverLogErrorWire, "fields?": PrimitiveRecord, level: '"debug" | "error" | "info" | "trace" | "warn"', - message: NonEmptyString, + message: "string", "namespace?": "string | null", - seq: "number >= 0", + seq: NonNegativeSafeInteger, timestamp: NonEmptyString, -}); +}).onUndeclaredKey("reject"); const DriverLogBatchInputWire = type({ driverInstanceId: NonEmptyString, - logs: DriverLogEntryWire.array(), -}); + logs: DriverLogEntryWire.array().atMostLength(LOG_BATCH_MAX_SIZE), +}).onUndeclaredKey("reject"); const DriverLogBatchOutputWire = type({ ok: "true", -}); +}).onUndeclaredKey("reject"); -const DriverCommandUpdateInputWire = type({ +const OkOutputWire = type({ ok: "true" }).onUndeclaredKey("reject"); + +const DriverCommandAcceptedInputWire = type({ commandId: NonEmptyString, driverInstanceId: NonEmptyString, - "error?": RunError, - "result?": RuntimeCommandResult, - status: RuntimeCommandStatus, -}); + status: '"accepted"', +}).onUndeclaredKey("reject"); + +const DriverCommandCancelledInputWire = type({ + commandId: NonEmptyString, + driverInstanceId: NonEmptyString, + status: '"cancelled"', +}).onUndeclaredKey("reject"); + +const DriverCommandCompletedInputWire = type({ + commandId: NonEmptyString, + driverInstanceId: NonEmptyString, + "result?": InputStartCommandResult.or(McpExecuteCommandResult), + status: '"completed"', +}).onUndeclaredKey("reject"); + +const DriverCommandFailedInputWire = type({ + commandId: NonEmptyString, + driverInstanceId: NonEmptyString, + error: DurableRunError, + status: '"failed"', +}).onUndeclaredKey("reject"); + +const DriverCommandUpdateInputWire = DriverCommandAcceptedInputWire.or( + DriverCommandCancelledInputWire, +) + .or(DriverCommandCompletedInputWire) + .or(DriverCommandFailedInputWire); const DriverExternalToolEffectClaimInputWire = type({ + claimToken: ExternalToolEffectClaimToken, commandId: NonEmptyString, driverInstanceId: NonEmptyString, -}); +}).onUndeclaredKey("reject"); -const DriverExternalToolEffectCompleteInputWire = type({ +const DriverExternalToolEffectObserveInputWire = type({ commandId: NonEmptyString, driverInstanceId: NonEmptyString, - "providerReceiptJson?": "string | null | undefined", - result: McpExecuteCommandResult, -}); +}).onUndeclaredKey("reject"); -const DriverExternalToolEffectUnknownInputWire = type({ +const DriverExternalToolEffectSettleInputWire = type({ + claimToken: ExternalToolEffectClaimToken, commandId: NonEmptyString, driverInstanceId: NonEmptyString, -}); + effectId: NonEmptyString, + settlement: ExternalToolEffectSettlement, +}).onUndeclaredKey("reject"); const DriverNextCommandInputWire = type({ driverInstanceId: NonEmptyString, -}); +}).onUndeclaredKey("reject"); const DriverNextCommandOutputWire = type({ command: type("null").or(RuntimeCommand), -}); +}).onUndeclaredKey("reject"); const DriverCompletionInputWire = type({ driverInstanceId: NonEmptyString, -}); + runId: NonEmptyString, +}).onUndeclaredKey("reject"); const DriverFailureInputWire = type({ driverInstanceId: NonEmptyString, - error: RunError, -}); + error: DurableRunError, + runId: NonEmptyString, +}).onUndeclaredKey("reject"); type DriverEventBatchOutputWireValue = typeof DriverEventBatchOutputWire.infer; type DriverHelloOutputWireValue = typeof DriverHelloOutputWire.infer; @@ -180,65 +258,87 @@ export interface RuntimeOrpcContext { onClaimExternalToolEffect( input: DriverExternalToolEffectClaimInput, ): Promise; - onCompleteExternalToolEffect(input: DriverExternalToolEffectCompleteInput): Promise<{ ok: true }>; onCompleteRun(input: DriverCompletionInput): Promise<{ ok: true }>; onFailRun(input: DriverFailureInput): Promise<{ ok: true }>; onHeartbeat(input: DriverHeartbeatInput): Promise; onHello(input: DriverHelloInput): Promise; onNextCommand(input: DriverNextCommandInput): Promise; - onPushEvents(input: DriverEventBatchInput): Promise; + onObserveExternalToolEffect( + input: DriverExternalToolEffectObserveInput, + ): Promise; + onPushEvents(input: HostDriverEventBatchInput): Promise; onPushLogs(input: DriverLogBatchInput): Promise; - onMarkExternalToolEffectUnknown( - input: DriverExternalToolEffectUnknownInput, - ): Promise<{ ok: true }>; onReady(input: DriverReadyInput): Promise<{ ok: true }>; + onSettleExternalToolEffect( + input: DriverExternalToolEffectSettleInput, + ): Promise; onWatchCommands(): AsyncIteratorObject; } -function parseDriverCommandUpdateInput(input: unknown): DriverCommandUpdateInput { +export function parseDriverCommandUpdateInput(input: unknown): DriverCommandUpdateInput { return parseSchemaValue(DriverCommandUpdateInputWire, input); } -function parseDriverExternalToolEffectClaimInput( +export function parseDriverExternalToolEffectClaimInput( input: unknown, ): DriverExternalToolEffectClaimInput { return parseSchemaValue(DriverExternalToolEffectClaimInputWire, input); } -function parseDriverExternalToolEffectCompleteInput( +export function parseDriverExternalToolEffectObserveInput( input: unknown, -): DriverExternalToolEffectCompleteInput { - return parseSchemaValue(DriverExternalToolEffectCompleteInputWire, input); +): DriverExternalToolEffectObserveInput { + return parseSchemaValue(DriverExternalToolEffectObserveInputWire, input); } -function parseDriverExternalToolEffectUnknownInput( +export function parseDriverExternalToolEffectSettleInput( input: unknown, -): DriverExternalToolEffectUnknownInput { - return parseSchemaValue(DriverExternalToolEffectUnknownInputWire, input); +): DriverExternalToolEffectSettleInput { + return parseSchemaValue(DriverExternalToolEffectSettleInputWire, input); } -function parseDriverCompletionInput(input: unknown): DriverCompletionInput { +export function parseDriverCompletionInput(input: unknown): DriverCompletionInput { return parseSchemaValue(DriverCompletionInputWire, input); } -export function parseDriverEventBatchInput(input: unknown): DriverEventBatchInput { +export function parseDriverEventBatchInput(input: unknown): HostDriverEventBatchInput { const batch = parseSchemaValue(DriverEventBatchInputWire, input); return { driverInstanceId: batch.driverInstanceId, - events: batch.events.map(parseDriverEventEnvelope), + events: batch.events.map((envelope) => ({ + event: parseRuntimeEventEnvelope(envelope.event), + eventId: envelope.eventId, + occurredAt: envelope.occurredAt, + })), }; } -function parseDriverFailureInput(input: unknown): DriverFailureInput { +export function parseDriverFailureInput(input: unknown): DriverFailureInput { return parseSchemaValue(DriverFailureInputWire, input); } -function parseDriverLogBatchInput(input: unknown): DriverLogBatchInput { +export function parseDriverHeartbeatInput(input: unknown): DriverHeartbeatInput { + return parseSchemaValue(DriverHeartbeatInputWire, input); +} + +export function parseDriverHeartbeatOutput(input: unknown): DriverHeartbeatOutput { + return parseSchemaValue(DriverHeartbeatOutputWire, input); +} + +export function parseDriverHelloInput(input: unknown): DriverHelloInput { + return parseSchemaValue(DriverHelloInputWire, input); +} + +export function parseDriverLogBatchInput(input: unknown): DriverLogBatchInput { return parseSchemaValue(DriverLogBatchInputWire, input); } -function parseDriverNextCommandInput(input: unknown): DriverNextCommandInput { +export function parseDriverLogBatchOutput(input: unknown): DriverLogBatchOutput { + return parseSchemaValue(DriverLogBatchOutputWire, input); +} + +export function parseDriverNextCommandInput(input: unknown): DriverNextCommandInput { return parseSchemaValue(DriverNextCommandInputWire, input); } @@ -246,20 +346,22 @@ export function parseDriverReadyInput(input: unknown): DriverReadyInput { return parseSchemaValue(DriverReadyInputWire, input); } -function toDriverEventBatchOutputWire( - output: DriverEventBatchOutput, -): DriverEventBatchOutputWireValue { - return parseSchemaValue(DriverEventBatchOutputWire, output); +export function parseDriverEventBatchOutput(input: unknown): DriverEventBatchOutputWireValue { + return parseSchemaValue(DriverEventBatchOutputWire, input); } -function toDriverHelloOutputWire(output: DriverHelloOutput): DriverHelloOutputWireValue { - return parseSchemaValue(DriverHelloOutputWire, output); +export function parseDriverHelloOutput(input: unknown): DriverHelloOutputWireValue { + return parseSchemaValue(DriverHelloOutputWire, input); +} + +export function parseDriverNextCommandOutput(input: unknown): DriverNextCommandOutputWireValue { + return parseSchemaValue(DriverNextCommandOutputWire, input); } function toDriverNextCommandOutputWire( output: DriverNextCommandOutput, ): DriverNextCommandOutputWireValue { - return parseSchemaValue(DriverNextCommandOutputWire, output); + return parseDriverNextCommandOutput(output); } const base = os.$context(); @@ -270,43 +372,55 @@ export const runtimeOrpcRouter = { .input(DriverExternalToolEffectClaimInputWire) .output(ExternalToolEffectClaim) .handler(async ({ context, input }) => - context.onClaimExternalToolEffect(parseDriverExternalToolEffectClaimInput(input)), + parseSchemaValue( + ExternalToolEffectClaim, + await context.onClaimExternalToolEffect(parseDriverExternalToolEffectClaimInput(input)), + ), ), commandUpdate: base .input(DriverCommandUpdateInputWire) - .output(type({ ok: "true" })) + .output(OkOutputWire) .handler(async ({ context, input }) => context.onCommandUpdate(parseDriverCommandUpdateInput(input)), ), - completeExternalToolEffect: base - .input(DriverExternalToolEffectCompleteInputWire) - .output(type({ ok: "true" })) - .handler(async ({ context, input }) => - context.onCompleteExternalToolEffect(parseDriverExternalToolEffectCompleteInput(input)), - ), completeRun: base .input(DriverCompletionInputWire) - .output(type({ ok: "true" })) + .output(OkOutputWire) .handler(async ({ context, input }) => context.onCompleteRun(parseDriverCompletionInput(input)), ), failRun: base .input(DriverFailureInputWire) - .output(type({ ok: "true" })) + .output(OkOutputWire) .handler(async ({ context, input }) => context.onFailRun(parseDriverFailureInput(input))), heartbeat: base .input(DriverHeartbeatInputWire) .output(DriverHeartbeatOutputWire) - .handler(async ({ context, input }) => context.onHeartbeat(input)), + .handler(async ({ context, input }) => + parseDriverHeartbeatOutput(await context.onHeartbeat(parseDriverHeartbeatInput(input))), + ), hello: base .input(DriverHelloInputWire) .output(DriverHelloOutputWire) - .handler(async ({ context, input }) => toDriverHelloOutputWire(await context.onHello(input))), + .handler(async ({ context, input }) => + parseDriverHelloOutput(await context.onHello(parseDriverHelloInput(input))), + ), + observeExternalToolEffect: base + .input(DriverExternalToolEffectObserveInputWire) + .output(ExternalToolEffectState) + .handler(async ({ context, input }) => + parseSchemaValue( + ExternalToolEffectState, + await context.onObserveExternalToolEffect( + parseDriverExternalToolEffectObserveInput(input), + ), + ), + ), pushEvents: base .input(DriverEventBatchInputWire) .output(DriverEventBatchOutputWire) .handler(async ({ context, input }) => - toDriverEventBatchOutputWire(await context.onPushEvents(parseDriverEventBatchInput(input))), + parseDriverEventBatchOutput(await context.onPushEvents(parseDriverEventBatchInput(input))), ), pushLogs: base .input(DriverLogBatchInputWire) @@ -314,13 +428,16 @@ export const runtimeOrpcRouter = { .handler(async ({ context, input }) => context.onPushLogs(parseDriverLogBatchInput(input))), ready: base .input(DriverReadyInputWire) - .output(type({ ok: "true" })) + .output(OkOutputWire) .handler(async ({ context, input }) => context.onReady(parseDriverReadyInput(input))), - markExternalToolEffectUnknown: base - .input(DriverExternalToolEffectUnknownInputWire) - .output(type({ ok: "true" })) + settleExternalToolEffect: base + .input(DriverExternalToolEffectSettleInputWire) + .output(ExternalToolEffectState) .handler(async ({ context, input }) => - context.onMarkExternalToolEffectUnknown(parseDriverExternalToolEffectUnknownInput(input)), + parseSchemaValue( + ExternalToolEffectState, + await context.onSettleExternalToolEffect(parseDriverExternalToolEffectSettleInput(input)), + ), ), }, driverInstance: { diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc.ts index 268ec76b..f17ffb62 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/rpc.ts @@ -1,12 +1,12 @@ import type { DriverCommandUpdateInput, DriverCompletionInput, - DriverEventBatchInput, DriverEventBatchOutput, DriverExternalToolEffectClaimInput, DriverExternalToolEffectClaimOutput, - DriverExternalToolEffectCompleteInput, - DriverExternalToolEffectUnknownInput, + DriverExternalToolEffectObserveInput, + DriverExternalToolEffectSettleInput, + DriverExternalToolEffectState, DriverFailureInput, DriverHeartbeatInput, DriverHelloInput, @@ -19,12 +19,15 @@ import type { } from "@mosoo/agent-driver/orpc"; import type { RuntimeCommand } from "@mosoo/contracts/runtime-command"; +import type { HostDriverEventBatchInput } from "./event-types"; import type { RuntimeOrpcContext } from "./rpc-wire"; +import type { DriverInstanceConnectionEpoch } from "./state"; export type DriverInstanceRpcContext = RuntimeOrpcContext; export interface DriverInstanceRpcOperationContext { readonly connectionId: string; + readonly epoch: DriverInstanceConnectionEpoch; assertActiveConnection(): void; } @@ -37,10 +40,10 @@ export interface DriverInstanceRpcHandler { input: DriverExternalToolEffectClaimInput, context: DriverInstanceRpcOperationContext, ): Promise; - handleCompleteExternalToolEffect( - input: DriverExternalToolEffectCompleteInput, + handleObserveExternalToolEffect( + input: DriverExternalToolEffectObserveInput, context: DriverInstanceRpcOperationContext, - ): Promise<{ ok: true }>; + ): Promise; handleCompleteRun( input: DriverCompletionInput, context: DriverInstanceRpcOperationContext, @@ -62,17 +65,17 @@ export interface DriverInstanceRpcHandler { context: DriverInstanceRpcOperationContext, ): Promise; handlePushEvents( - input: DriverEventBatchInput, + input: HostDriverEventBatchInput, context: DriverInstanceRpcOperationContext, ): Promise; handlePushLogs( input: DriverLogBatchInput, context: DriverInstanceRpcOperationContext, ): Promise; - handleMarkExternalToolEffectUnknown( - input: DriverExternalToolEffectUnknownInput, + handleSettleExternalToolEffect( + input: DriverExternalToolEffectSettleInput, context: DriverInstanceRpcOperationContext, - ): Promise<{ ok: true }>; + ): Promise; handleReady( input: DriverReadyInput, context: DriverInstanceRpcOperationContext, @@ -88,18 +91,18 @@ export function createDriverInstanceRpcContext( onClaimExternalToolEffect: async (input) => handler.handleClaimExternalToolEffect(input, context), onCommandUpdate: async (input) => handler.handleCommandUpdate(input, context), - onCompleteExternalToolEffect: async (input) => - handler.handleCompleteExternalToolEffect(input, context), onCompleteRun: async (input) => handler.handleCompleteRun(input, context), onFailRun: async (input) => handler.handleFailRun(input, context), onHeartbeat: async (input) => handler.handleHeartbeat(input, context), onHello: async (input) => handler.handleHello(input, context), onNextCommand: async (input) => handler.handleNextCommand(input, context), + onObserveExternalToolEffect: async (input) => + handler.handleObserveExternalToolEffect(input, context), onPushEvents: async (input) => handler.handlePushEvents(input, context), onPushLogs: async (input) => handler.handlePushLogs(input, context), - onMarkExternalToolEffectUnknown: async (input) => - handler.handleMarkExternalToolEffectUnknown(input, context), onReady: async (input) => handler.handleReady(input, context), + onSettleExternalToolEffect: async (input) => + handler.handleSettleExternalToolEffect(input, context), onWatchCommands: () => handler.watchCommands(context)[Symbol.asyncIterator]() as ReturnType< RuntimeOrpcContext["onWatchCommands"] diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/run-transitions.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/run-transitions.ts deleted file mode 100644 index 6c204579..00000000 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/run-transitions.ts +++ /dev/null @@ -1,12 +0,0 @@ -import type { RuntimeDriverRunTransition } from "./event-types"; - -export function hasTerminalRuntimeDriverRunTransition( - transitions: readonly RuntimeDriverRunTransition[], -): boolean { - return transitions.some( - (transition) => - transition.status === "cancelled" || - transition.status === "completed" || - transition.status === "failed", - ); -} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-artifact-attempt.repository.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-artifact-attempt.repository.ts new file mode 100644 index 00000000..34ef7b96 --- /dev/null +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-artifact-attempt.repository.ts @@ -0,0 +1,896 @@ +import { fileRecordsTable, runtimeArtifactAttemptsTable } from "@mosoo/db"; +import type { AccountId, FileId, RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; +import { stringifyRuntimeEventSemanticValue } from "@mosoo/runtime-events"; +import { and, eq, inArray, isNull, lte, or, sql } from "drizzle-orm"; + +import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; +import { getAppDatabase } from "../../../../platform/db/drizzle"; +import { currentTimestampMs } from "../../../../time"; +import { + createRuntimeOutputContentSha256, + deleteRuntimeArtifactObject, +} from "../../../files/application/file-store"; +import type { DriverRuntimeEventFence } from "../../../sessions/infrastructure/session-runtime-event-store.types"; +import { + normalizeRuntimeSessionOutputRelativePath, + toRuntimeSessionOutputArtifactPath, +} from "./runtime-session-outputs"; + +const RUNTIME_ARTIFACT_STAGE_TTL_MS = 24 * 60 * 60_000; +const RUNTIME_ARTIFACT_DELETE_QUARANTINE_MS = 7 * 24 * 60 * 60_000; +const RUNTIME_ARTIFACT_DELETE_LEASE_MS = 5 * 60_000; +const RUNTIME_ARTIFACT_CLEANUP_BATCH_SIZE = 25; + +export interface RuntimeArtifactCapturePlanFile { + readonly contentType: string | null; + readonly expectedSize: number; + readonly fileId: FileId; + readonly name: string; + readonly objectKey: string; + readonly operation: "upsert"; + readonly readPath: string; + readonly sourcePath: string; +} + +export type RuntimeArtifactCaptureStatus = + | "complete" + | "omitted_file_limit" + | "omitted_runtime_unavailable" + | "omitted_size_limit" + | "omitted_source_changed" + | "omitted_source_missing"; + +export interface RuntimeArtifactCapturePlan { + readonly captureStatus: RuntimeArtifactCaptureStatus; + readonly files: readonly ( + | RuntimeArtifactCapturePlanFile + | { readonly operation: "delete"; readonly sourcePath: string } + )[]; + readonly mode: "delta" | "snapshot"; + readonly version: 1; +} + +export interface RuntimeArtifactUpsertManifestFile { + readonly contentSha256: string; + readonly contentType: string | null; + readonly disposition: "create" | "reuse"; + readonly etag: string; + readonly fileId: FileId; + readonly name: string; + readonly objectKey: string; + readonly operation: "upsert"; + readonly parentPath: string; + readonly path: string; + readonly size: number; + readonly sourcePath: string; +} + +export interface RuntimeArtifactDeleteManifestFile { + readonly operation: "delete"; + readonly sourcePath: string; +} + +export type RuntimeArtifactManifestFile = + | RuntimeArtifactDeleteManifestFile + | RuntimeArtifactUpsertManifestFile; + +export interface RuntimeArtifactManifest { + readonly captureStatus: RuntimeArtifactCaptureStatus; + readonly files: readonly RuntimeArtifactManifestFile[]; + readonly mode: "delta" | "snapshot"; + readonly semanticHash: string; + readonly sourceEventId: string; + readonly version: 1; +} + +export interface StagedRuntimeArtifactProjection { + readonly attemptId: string; + readonly manifestJson: string; + readonly manifestSha256: string; +} + +export interface ReadyRuntimeArtifactRecord { + readonly contentType: string | null; + readonly etag: string; + readonly fileId: FileId; + readonly name: string; + readonly objectKey: string; + readonly parentPath: string; + readonly path: string; + readonly size: number; +} + +interface RuntimeArtifactAttemptIdentity { + readonly createdByAccountId: AccountId; + readonly driverFence: DriverRuntimeEventFence; + readonly eventType: string; + readonly runId: SessionRunId; + readonly semanticHash: string; + readonly sessionId: SessionId; + readonly sourceEventId: string; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function readRequiredString(record: Record, key: string): string { + const value = record[key]; + if (typeof value !== "string" || value.length === 0) { + throw new Error(`Runtime artifact manifest ${key} must be a non-empty string.`); + } + return value; +} + +function parseManifestFile(value: unknown): RuntimeArtifactManifestFile { + if (!isRecord(value)) { + throw new Error("Runtime artifact manifest file must be an object."); + } + const operation = value["operation"]; + const sourcePath = readRequiredString(value, "sourcePath"); + if (operation === "delete") { + return { operation, sourcePath }; + } + const contentType = value["contentType"]; + const disposition = value["disposition"]; + const size = value["size"]; + if ( + operation !== "upsert" || + (contentType !== null && typeof contentType !== "string") || + (disposition !== "create" && disposition !== "reuse") || + !Number.isSafeInteger(size) || + (size as number) < 0 + ) { + throw new Error("Runtime artifact manifest file has invalid storage metadata."); + } + const contentSha256 = readRequiredString(value, "contentSha256"); + if (!/^[0-9a-f]{64}$/.test(contentSha256)) { + throw new Error("Runtime artifact manifest content hash is invalid."); + } + return { + contentSha256, + contentType, + disposition, + etag: readRequiredString(value, "etag"), + fileId: readRequiredString(value, "fileId") as FileId, + name: readRequiredString(value, "name"), + objectKey: readRequiredString(value, "objectKey"), + operation, + parentPath: readRequiredString(value, "parentPath"), + path: readRequiredString(value, "path"), + size: size as number, + sourcePath, + }; +} + +export function parseRuntimeArtifactManifest(value: string): RuntimeArtifactManifest { + const parsed: unknown = JSON.parse(value); + if ( + !isRecord(parsed) || + parsed["version"] !== 1 || + (parsed["captureStatus"] !== "complete" && + parsed["captureStatus"] !== "omitted_file_limit" && + parsed["captureStatus"] !== "omitted_runtime_unavailable" && + parsed["captureStatus"] !== "omitted_size_limit" && + parsed["captureStatus"] !== "omitted_source_changed" && + parsed["captureStatus"] !== "omitted_source_missing") || + (parsed["mode"] !== "delta" && parsed["mode"] !== "snapshot") || + !Array.isArray(parsed["files"]) + ) { + throw new Error("Runtime artifact manifest is invalid."); + } + const sourceEventId = readRequiredString(parsed, "sourceEventId"); + const semanticHash = readRequiredString(parsed, "semanticHash"); + if (!/^[0-9a-f]{64}$/.test(semanticHash)) { + throw new Error("Runtime artifact manifest semantic hash is invalid."); + } + const files = parsed["files"].map(parseManifestFile); + if (parsed["captureStatus"] !== "complete" && files.length !== 0) { + throw new Error("An omitted runtime artifact capture must be empty."); + } + const identities = new Set(); + for (const file of files) { + for (const identity of [ + file.sourcePath, + ...(file.operation === "upsert" ? [file.fileId, file.path, file.objectKey] : []), + ]) { + if (identities.has(identity)) { + throw new Error("Runtime artifact manifest contains a duplicate file identity."); + } + identities.add(identity); + } + } + return { + captureStatus: parsed["captureStatus"], + files, + mode: parsed["mode"], + semanticHash, + sourceEventId, + version: 1, + }; +} + +export async function createRuntimeArtifactManifest( + input: Omit, +): Promise<{ manifestJson: string; manifestSha256: string }> { + const manifest: RuntimeArtifactManifest = { + ...input, + files: input.files.toSorted((left, right) => + left.sourcePath === right.sourcePath + ? (left.operation === "upsert" ? left.fileId : "").localeCompare( + right.operation === "upsert" ? right.fileId : "", + ) + : left.sourcePath.localeCompare(right.sourcePath), + ), + version: 1, + }; + const manifestJson = stringifyRuntimeEventSemanticValue(manifest); + return { + manifestJson, + manifestSha256: await createRuntimeOutputContentSha256(new TextEncoder().encode(manifestJson)), + }; +} + +export async function createRuntimeArtifactAttempt( + database: D1Database, + input: RuntimeArtifactAttemptIdentity & { + readonly attemptId: string; + readonly timestampMs?: number; + }, +): Promise { + if (input.driverFence.sessionRunId !== input.runId) { + throw new Error("Runtime artifact attempt does not match its fenced Session Run."); + } + const timestampMs = input.timestampMs ?? currentTimestampMs(); + const result = await database + .prepare( + `INSERT INTO runtime_artifact_attempt ( + accepted_event_id, created_at, created_by_account_id, + delete_after, driver_connection_id, driver_generation, driver_instance_id, + event_type, expires_at, id, manifest_json, manifest_sha256, + owned_object_keys_json, run_id, semantic_hash, session_id, source_event_id, + status, updated_at + ) + SELECT NULL, ?, ?, NULL, ?, ?, ?, ?, ?, ?, NULL, NULL, '[]', ?, ?, ?, ?, + 'staging', ? + FROM session AS session + INNER JOIN session_run AS run + ON run.id = ? AND run.session_id = session.id + INNER JOIN driver_instance AS driver + ON driver.id = run.driver_instance_id + WHERE session.id = ? + AND session.archived_at IS NULL + AND session.cleanup_operation_kind IS NULL + AND session.last_run_id = run.id + AND session.status = 'RUNNING' + AND session.status_operation_id IS NULL + AND run.status IN ('queued', 'booting', 'running', 'waiting_input') + AND driver.id = ? + AND driver.generation = ? + AND driver.connection_id = ? + AND driver.sandbox_session_id = session.id + RETURNING id`, + ) + .bind( + timestampMs, + input.createdByAccountId, + input.driverFence.connectionId, + input.driverFence.generation, + input.driverFence.driverInstanceId, + input.eventType, + timestampMs + RUNTIME_ARTIFACT_STAGE_TTL_MS, + input.attemptId, + input.runId, + input.semanticHash, + input.sessionId, + input.sourceEventId, + timestampMs, + input.runId, + input.sessionId, + input.driverFence.driverInstanceId, + input.driverFence.generation, + input.driverFence.connectionId, + ) + .first<{ id: string }>(); + + if (result?.id !== input.attemptId) { + throw new Error("Runtime artifact attempt lost its active Driver fence."); + } +} + +function exactActiveAttemptFenceSql(): string { + return `EXISTS ( + SELECT 1 + FROM session AS session + INNER JOIN session_run AS run + ON run.id = runtime_artifact_attempt.run_id + AND run.session_id = session.id + INNER JOIN driver_instance AS driver + ON driver.id = run.driver_instance_id + WHERE session.id = runtime_artifact_attempt.session_id + AND session.archived_at IS NULL + AND session.cleanup_operation_kind IS NULL + AND session.last_run_id = run.id + AND session.status = 'RUNNING' + AND session.status_operation_id IS NULL + AND run.status IN ('queued', 'booting', 'running', 'waiting_input') + AND driver.id = runtime_artifact_attempt.driver_instance_id + AND driver.generation = runtime_artifact_attempt.driver_generation + AND driver.connection_id = runtime_artifact_attempt.driver_connection_id + AND driver.sandbox_session_id = session.id + )`; +} + +export async function getReadyRuntimeArtifact( + database: D1Database, + input: { readonly parentPath: string; readonly sessionId: SessionId }, +): Promise { + const row = await getAppDatabase(database) + .select({ + contentType: fileRecordsTable.mimeType, + etag: fileRecordsTable.etag, + fileId: fileRecordsTable.id, + name: fileRecordsTable.name, + objectKey: fileRecordsTable.objectKey, + parentPath: fileRecordsTable.parentPath, + path: fileRecordsTable.path, + size: fileRecordsTable.size, + }) + .from(fileRecordsTable) + .where( + and( + eq(fileRecordsTable.scopeKind, "session"), + eq(fileRecordsTable.scopeId, input.sessionId), + eq(fileRecordsTable.sessionKind, "artifact"), + eq(fileRecordsTable.status, "ready"), + eq(fileRecordsTable.parentPath, input.parentPath), + ), + ) + .get(); + if (row === undefined) { + return null; + } + if (row.etag === null) { + throw new Error("Ready runtime artifact is missing its object etag."); + } + return { ...row, etag: row.etag }; +} + +export async function claimRuntimeArtifactObjectKey( + database: D1Database, + input: { readonly attemptId: string; readonly objectKey: string; readonly timestampMs?: number }, +): Promise { + if (!input.objectKey.startsWith(`runtime-artifact-attempts/v1/${input.attemptId}/files/`)) { + throw new Error("Runtime artifact object key is outside its attempt namespace."); + } + const timestampMs = input.timestampMs ?? currentTimestampMs(); + const result = await database + .prepare( + `UPDATE runtime_artifact_attempt + SET owned_object_keys_json = json_insert(owned_object_keys_json, '$[#]', ?), + updated_at = ? + WHERE id = ? + AND status = 'staging' + AND expires_at > ? + AND NOT EXISTS ( + SELECT 1 FROM json_each(owned_object_keys_json) WHERE value = ? + ) + AND ${exactActiveAttemptFenceSql()} + RETURNING id`, + ) + .bind(input.objectKey, timestampMs, input.attemptId, timestampMs, input.objectKey) + .first<{ id: string }>(); + if (result?.id !== input.attemptId) { + throw new Error("Runtime artifact object lost its staging ownership fence."); + } +} + +export async function sealRuntimeArtifactAttempt( + database: D1Database, + input: StagedRuntimeArtifactProjection & { readonly timestampMs?: number }, +): Promise { + const timestampMs = input.timestampMs ?? currentTimestampMs(); + const manifest = parseRuntimeArtifactManifest(input.manifestJson); + for (const file of manifest.files) { + const outputPrefix = "outputs/"; + const relativePath = file.sourcePath.startsWith(outputPrefix) + ? normalizeRuntimeSessionOutputRelativePath(file.sourcePath.slice(outputPrefix.length)) + : null; + if ( + relativePath === null || + file.sourcePath !== toRuntimeSessionOutputArtifactPath(relativePath) + ) { + throw new Error("Runtime artifact manifest source path is not canonical."); + } + if ( + file.operation === "upsert" && + file.disposition === "create" && + !file.objectKey.startsWith(`runtime-artifact-attempts/v1/${input.attemptId}/files/`) + ) { + throw new Error("Runtime artifact manifest object is outside its attempt namespace."); + } + } + const expectedManifestSha256 = await createRuntimeOutputContentSha256( + new TextEncoder().encode(input.manifestJson), + ); + if (input.manifestSha256 !== expectedManifestSha256) { + throw new Error("Runtime artifact manifest hash does not match its canonical bytes."); + } + const result = await database + .prepare( + `UPDATE runtime_artifact_attempt + SET manifest_json = ?, manifest_sha256 = ?, status = 'staged', updated_at = ? + WHERE id = ? + AND status = 'staging' + AND expires_at > ? + AND source_event_id = ? + AND semantic_hash = ? + AND ( + (event_type = 'run.completed' + AND json_extract(?, '$.mode') = 'snapshot' + AND NOT EXISTS ( + SELECT 1 FROM json_each(?, '$.files') AS file + WHERE json_extract(file.value, '$.operation') = 'delete' + )) + OR (event_type IN ('file.change.updated', 'file.changed') + AND json_extract(?, '$.mode') = 'delta') + ) + AND ${exactActiveAttemptFenceSql()} + RETURNING id`, + ) + .bind( + input.manifestJson, + input.manifestSha256, + timestampMs, + input.attemptId, + timestampMs, + manifest.sourceEventId, + manifest.semanticHash, + input.manifestJson, + input.manifestJson, + input.manifestJson, + ) + .first<{ id: string }>(); + if (result?.id !== input.attemptId) { + throw new Error("Runtime artifact manifest lost its staging ownership fence."); + } +} + +export function prepareRuntimeArtifactPromotion( + database: D1Database, + input: StagedRuntimeArtifactProjection & { + readonly eventId: RuntimeEventId; + readonly timestampMs: number; + }, +): D1PreparedStatement[] { + return [ + database + .prepare( + `INSERT INTO file_record ( + committed, created_at, created_by_account_id, etag, expires_at, id, + mime_type, name, object_key, owner_id, owner_kind, parent_path, path, + purpose, runtime_event_seq, scope_id, scope_kind, session_kind, size, + status, updated_at, version + ) + SELECT 1, event.created_at, attempt.created_by_account_id, + json_extract(file.value, '$.etag'), NULL, + json_extract(file.value, '$.fileId'), + json_extract(file.value, '$.contentType'), + json_extract(file.value, '$.name'), + json_extract(file.value, '$.objectKey'), + event.session_id, 'session', + json_extract(file.value, '$.parentPath'), + json_extract(file.value, '$.path'), + 'session_artifact', event.seq, event.session_id, 'session', 'artifact', + json_extract(file.value, '$.size'), 'ready', event.created_at, 1 + FROM runtime_artifact_attempt AS attempt + INNER JOIN session_event AS event + ON event.id = ? + AND event.session_id = attempt.session_id + AND event.run_id = attempt.run_id + AND event.source_event_id = attempt.source_event_id + AND event.semantic_hash = attempt.semantic_hash + AND event.artifact_attempt_id = attempt.id + AND event.artifact_manifest_sha256 = attempt.manifest_sha256 + INNER JOIN json_each(attempt.manifest_json, '$.files') AS file + ON json_extract(file.value, '$.operation') = 'upsert' + AND json_extract(file.value, '$.disposition') = 'create' + WHERE attempt.id = ? + AND attempt.status = 'staged' + AND attempt.expires_at > ? + AND attempt.manifest_sha256 = ?`, + ) + .bind(input.eventId, input.attemptId, input.timestampMs, input.manifestSha256), + database + .prepare( + `UPDATE file_record + SET runtime_event_seq = MAX(COALESCE(runtime_event_seq, 0), ( + SELECT event.seq FROM session_event AS event WHERE event.id = ? + )), + updated_at = MAX(updated_at, ( + SELECT event.created_at FROM session_event AS event WHERE event.id = ? + )) + WHERE EXISTS ( + SELECT 1 + FROM runtime_artifact_attempt AS attempt + INNER JOIN session_event AS event + ON event.id = ? + AND event.session_id = attempt.session_id + AND event.run_id = attempt.run_id + AND event.source_event_id = attempt.source_event_id + AND event.semantic_hash = attempt.semantic_hash + AND event.artifact_attempt_id = attempt.id + AND event.artifact_manifest_sha256 = attempt.manifest_sha256 + INNER JOIN json_each(attempt.manifest_json, '$.files') AS file + ON json_extract(file.value, '$.operation') = 'upsert' + AND json_extract(file.value, '$.disposition') = 'reuse' + WHERE attempt.id = ? + AND attempt.status = 'staged' + AND attempt.expires_at > ? + AND attempt.manifest_sha256 = ? + AND file_record.id = json_extract(file.value, '$.fileId') + AND file_record.scope_kind = 'session' + AND file_record.scope_id = event.session_id + AND file_record.session_kind = 'artifact' + AND file_record.status = 'ready' + AND file_record.parent_path = json_extract(file.value, '$.parentPath') + AND file_record.path = json_extract(file.value, '$.path') + AND file_record.name = json_extract(file.value, '$.name') + AND file_record.object_key = json_extract(file.value, '$.objectKey') + AND file_record.etag = json_extract(file.value, '$.etag') + AND file_record.size = json_extract(file.value, '$.size') + AND file_record.mime_type IS json_extract(file.value, '$.contentType') + )`, + ) + .bind( + input.eventId, + input.eventId, + input.eventId, + input.attemptId, + input.timestampMs, + input.manifestSha256, + ), + database + .prepare( + `UPDATE session_artifact_head + SET file_id = NULL, + runtime_event_seq = (SELECT seq FROM session_event WHERE id = ?), + source_event_id = (SELECT source_event_id FROM session_event WHERE id = ?), + updated_at = (SELECT created_at FROM session_event WHERE id = ?) + WHERE session_id = (SELECT session_id FROM session_event WHERE id = ?) + AND runtime_event_seq < (SELECT seq FROM session_event WHERE id = ?) + AND EXISTS ( + SELECT 1 + FROM runtime_artifact_attempt AS attempt + INNER JOIN session_event AS event + ON event.id = ? + AND event.artifact_attempt_id = attempt.id + AND event.artifact_manifest_sha256 = attempt.manifest_sha256 + WHERE attempt.id = ? + AND attempt.status = 'staged' + AND attempt.expires_at > ? + AND attempt.manifest_sha256 = ? + AND json_extract(attempt.manifest_json, '$.mode') = 'snapshot' + AND json_extract(attempt.manifest_json, '$.captureStatus') = 'complete' + )`, + ) + .bind( + input.eventId, + input.eventId, + input.eventId, + input.eventId, + input.eventId, + input.eventId, + input.attemptId, + input.timestampMs, + input.manifestSha256, + ), + database + .prepare( + `INSERT INTO session_artifact_head ( + file_id, runtime_event_seq, session_id, source_event_id, source_path, updated_at + ) + SELECT CASE + WHEN json_extract(file.value, '$.operation') = 'upsert' + THEN json_extract(file.value, '$.fileId') + ELSE NULL + END, + event.seq, event.session_id, event.source_event_id, + json_extract(file.value, '$.sourcePath'), event.created_at + FROM runtime_artifact_attempt AS attempt + INNER JOIN session_event AS event + ON event.id = ? + AND event.session_id = attempt.session_id + AND event.run_id = attempt.run_id + AND event.source_event_id = attempt.source_event_id + AND event.semantic_hash = attempt.semantic_hash + AND event.artifact_attempt_id = attempt.id + AND event.artifact_manifest_sha256 = attempt.manifest_sha256 + INNER JOIN json_each(attempt.manifest_json, '$.files') AS file ON 1 = 1 + WHERE attempt.id = ? + AND attempt.status = 'staged' + AND attempt.expires_at > ? + AND attempt.manifest_sha256 = ? + ON CONFLICT(session_id, source_path) DO UPDATE SET + file_id = excluded.file_id, + runtime_event_seq = excluded.runtime_event_seq, + source_event_id = excluded.source_event_id, + updated_at = excluded.updated_at + WHERE excluded.runtime_event_seq > session_artifact_head.runtime_event_seq + OR ( + excluded.runtime_event_seq = session_artifact_head.runtime_event_seq + AND excluded.source_event_id = session_artifact_head.source_event_id + )`, + ) + .bind(input.eventId, input.attemptId, input.timestampMs, input.manifestSha256), + database + .prepare( + `UPDATE runtime_artifact_attempt AS attempt + SET accepted_event_id = ?, expires_at = NULL, owned_object_keys_json = '[]', + status = 'accepted', updated_at = ? + WHERE attempt.id = ? + AND attempt.status = 'staged' + AND attempt.expires_at > ? + AND attempt.manifest_json = ? + AND attempt.manifest_sha256 = ? + AND NOT EXISTS ( + SELECT 1 + FROM json_each(attempt.owned_object_keys_json) AS owned + WHERE NOT EXISTS ( + SELECT 1 + FROM json_each(attempt.manifest_json, '$.files') AS file + WHERE json_extract(file.value, '$.operation') = 'upsert' + AND json_extract(file.value, '$.disposition') = 'create' + AND json_extract(file.value, '$.objectKey') = owned.value + ) + ) + AND NOT EXISTS ( + SELECT 1 + FROM json_each(attempt.manifest_json, '$.files') AS file + WHERE json_extract(file.value, '$.operation') = 'upsert' + AND json_extract(file.value, '$.disposition') = 'create' + AND NOT EXISTS ( + SELECT 1 + FROM json_each(attempt.owned_object_keys_json) AS owned + WHERE owned.value = json_extract(file.value, '$.objectKey') + ) + ) + AND ( + SELECT COUNT(DISTINCT json_extract(file.value, '$.objectKey')) + FROM json_each(attempt.manifest_json, '$.files') AS file + WHERE json_extract(file.value, '$.operation') = 'upsert' + AND json_extract(file.value, '$.disposition') = 'create' + ) = json_array_length(attempt.owned_object_keys_json) + AND EXISTS ( + SELECT 1 FROM session_event AS event + WHERE event.id = ? + AND event.session_id = attempt.session_id + AND event.run_id = attempt.run_id + AND event.source_event_id = attempt.source_event_id + AND event.semantic_hash = attempt.semantic_hash + AND event.artifact_attempt_id = attempt.id + AND event.artifact_manifest_json = attempt.manifest_json + AND event.artifact_manifest_sha256 = attempt.manifest_sha256 + ) + AND NOT EXISTS ( + SELECT 1 + FROM json_each(attempt.manifest_json, '$.files') AS file + INNER JOIN session_event AS event ON event.id = ? + WHERE NOT EXISTS ( + SELECT 1 + FROM session_artifact_head AS head + WHERE head.session_id = attempt.session_id + AND head.source_path = json_extract(file.value, '$.sourcePath') + AND head.runtime_event_seq = event.seq + AND head.source_event_id = event.source_event_id + AND head.file_id IS CASE + WHEN json_extract(file.value, '$.operation') = 'upsert' + THEN json_extract(file.value, '$.fileId') + ELSE NULL + END + ) + ) + AND NOT EXISTS ( + SELECT 1 + FROM session_artifact_head AS head + INNER JOIN session_event AS event ON event.id = ? + WHERE json_extract(attempt.manifest_json, '$.mode') = 'snapshot' + AND json_extract(attempt.manifest_json, '$.captureStatus') = 'complete' + AND head.session_id = attempt.session_id + AND head.runtime_event_seq < event.seq + ) + AND NOT EXISTS ( + SELECT 1 + FROM json_each(attempt.manifest_json, '$.files') AS file + WHERE json_extract(file.value, '$.operation') = 'upsert' + AND NOT EXISTS ( + SELECT 1 FROM file_record AS stored + INNER JOIN session_event AS event ON event.id = ? + WHERE stored.id = json_extract(file.value, '$.fileId') + AND stored.scope_kind = 'session' + AND stored.scope_id = attempt.session_id + AND stored.session_kind = 'artifact' + AND stored.status = 'ready' + AND stored.parent_path = json_extract(file.value, '$.parentPath') + AND stored.path = json_extract(file.value, '$.path') + AND stored.name = json_extract(file.value, '$.name') + AND stored.object_key = json_extract(file.value, '$.objectKey') + AND stored.etag = json_extract(file.value, '$.etag') + AND stored.size = json_extract(file.value, '$.size') + AND stored.mime_type IS json_extract(file.value, '$.contentType') + AND stored.runtime_event_seq >= event.seq + ) + )`, + ) + .bind( + input.eventId, + input.timestampMs, + input.attemptId, + input.timestampMs, + input.manifestJson, + input.manifestSha256, + input.eventId, + input.eventId, + input.eventId, + input.eventId, + ), + database + .prepare( + `INSERT INTO session_event (id) + SELECT ? + WHERE NOT EXISTS ( + SELECT 1 + FROM runtime_artifact_attempt AS attempt + INNER JOIN session_event AS event ON event.id = attempt.accepted_event_id + WHERE attempt.id = ? + AND attempt.status = 'accepted' + AND attempt.accepted_event_id = ? + AND attempt.manifest_json = event.artifact_manifest_json + AND attempt.manifest_sha256 = event.artifact_manifest_sha256 + AND event.artifact_attempt_id = attempt.id + )`, + ) + .bind(crypto.randomUUID(), input.attemptId, input.eventId), + ]; +} + +interface CleanupAttempt { + readonly id: string; + readonly keys: readonly string[]; + readonly deleteAfter: number; +} + +function parseOwnedObjectKeys(value: string): string[] { + const parsed: unknown = JSON.parse(value); + if (!Array.isArray(parsed) || parsed.some((key) => typeof key !== "string" || key.length === 0)) { + throw new Error("Runtime artifact attempt owned object keys are invalid."); + } + return [...new Set(parsed)]; +} + +async function claimRuntimeArtifactCleanupAttempts( + database: D1Database, + nowMs: number, +): Promise { + await getAppDatabase(database) + .delete(runtimeArtifactAttemptsTable) + .where( + and( + eq(runtimeArtifactAttemptsTable.status, "accepted"), + sql`NOT EXISTS ( + SELECT 1 FROM session_event + WHERE session_event.id = ${runtimeArtifactAttemptsTable.acceptedEventId} + AND session_event.artifact_attempt_id = ${runtimeArtifactAttemptsTable.id} + )`, + ), + ) + .run(); + + const candidates = await getAppDatabase(database) + .select({ + deleteAfter: runtimeArtifactAttemptsTable.deleteAfter, + id: runtimeArtifactAttemptsTable.id, + ownedObjectKeysJson: runtimeArtifactAttemptsTable.ownedObjectKeysJson, + status: runtimeArtifactAttemptsTable.status, + updatedAt: runtimeArtifactAttemptsTable.updatedAt, + }) + .from(runtimeArtifactAttemptsTable) + .where( + or( + and( + inArray(runtimeArtifactAttemptsTable.status, ["staging", "staged"]), + lte(runtimeArtifactAttemptsTable.expiresAt, nowMs), + ), + and( + eq(runtimeArtifactAttemptsTable.status, "deleting"), + lte(runtimeArtifactAttemptsTable.updatedAt, nowMs - RUNTIME_ARTIFACT_DELETE_LEASE_MS), + ), + ), + ) + .orderBy(runtimeArtifactAttemptsTable.id) + .limit(RUNTIME_ARTIFACT_CLEANUP_BATCH_SIZE) + .all(); + const claimed: CleanupAttempt[] = []; + + for (const candidate of candidates) { + const deleteAfter = candidate.deleteAfter ?? nowMs + RUNTIME_ARTIFACT_DELETE_QUARANTINE_MS; + const row = await getAppDatabase(database) + .update(runtimeArtifactAttemptsTable) + .set({ + acceptedEventId: null, + deleteAfter, + expiresAt: null, + status: "deleting", + updatedAt: nowMs, + }) + .where( + and( + eq(runtimeArtifactAttemptsTable.id, candidate.id), + eq(runtimeArtifactAttemptsTable.status, candidate.status), + candidate.deleteAfter === null + ? isNull(runtimeArtifactAttemptsTable.deleteAfter) + : eq(runtimeArtifactAttemptsTable.deleteAfter, candidate.deleteAfter), + eq(runtimeArtifactAttemptsTable.updatedAt, candidate.updatedAt), + ), + ) + .returning({ id: runtimeArtifactAttemptsTable.id }) + .get(); + if (row !== undefined) { + claimed.push({ + deleteAfter, + id: candidate.id, + keys: parseOwnedObjectKeys(candidate.ownedObjectKeysJson), + }); + } + } + return claimed; +} + +async function runtimeArtifactObjectIsOwned( + database: D1Database, + input: { readonly attemptId: string; readonly objectKey: string }, +): Promise { + const row = await database + .prepare( + `SELECT 1 AS owned + WHERE EXISTS (SELECT 1 FROM file_record WHERE object_key = ?) + OR EXISTS ( + SELECT 1 + FROM runtime_artifact_attempt AS other, + json_each(other.owned_object_keys_json) AS owned_key + WHERE other.id <> ? + AND other.status IN ('staging', 'staged', 'accepted') + AND owned_key.value = ? + ) + LIMIT 1`, + ) + .bind(input.objectKey, input.attemptId, input.objectKey) + .first<{ owned: number }>(); + return row !== null; +} + +export async function cleanupRuntimeArtifactAttempts(bindings: ApiBindings): Promise { + const nowMs = currentTimestampMs(); + const attempts = await claimRuntimeArtifactCleanupAttempts(bindings.DB, nowMs); + for (const attempt of attempts) { + for (const objectKey of attempt.keys) { + if (await runtimeArtifactObjectIsOwned(bindings.DB, { attemptId: attempt.id, objectKey })) { + continue; + } + await deleteRuntimeArtifactObject(bindings, objectKey); + } + if (attempt.deleteAfter > nowMs) { + continue; + } + await getAppDatabase(bindings.DB) + .delete(runtimeArtifactAttemptsTable) + .where( + and( + eq(runtimeArtifactAttemptsTable.id, attempt.id), + eq(runtimeArtifactAttemptsTable.status, "deleting"), + eq(runtimeArtifactAttemptsTable.deleteAfter, attempt.deleteAfter), + ), + ) + .run(); + } +} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-artifact-staging.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-artifact-staging.ts new file mode 100644 index 00000000..aa9ffc87 --- /dev/null +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-artifact-staging.ts @@ -0,0 +1,665 @@ +import { createSessionArtifactPath, normalizeContentType } from "@mosoo/contracts/file"; +import { createPlatformId, parsePlatformId } from "@mosoo/id"; +import type { AccountId, FileId, SessionId, SessionRunId } from "@mosoo/id"; +import { readRuntimeEventFileChanges } from "@mosoo/runtime-events"; +import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; + +import { withDisposedRpcResource } from "../../../../platform/cloudflare/rpc-disposal"; +import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; +import { quoteShellArg } from "../../../../shared/shell"; +import { + createRuntimeOutputContentSha256, + createRuntimeOutputParentPath, + getRuntimeOutputName, + putRuntimeArtifactObject, +} from "../../../files/application/file-store"; +import type { DriverRuntimeEventFence } from "../../../sessions/infrastructure/session-runtime-event-store.types"; +import { getRuntimeSubjectKeepAliveHandle } from "../runtime-subject-lifecycle/runtime-subject-lifecycle.service"; +import { getRuntimeConversationSession } from "../runtime-subject-lifecycle/runtime-subject-store"; +import { readSandboxFileBytes } from "../sandbox-file-bytes"; +import type { ExecutionSessionHandle } from "../sandbox-handles"; +import type { RuntimeSessionLink } from "./event-types"; +import { + claimRuntimeArtifactObjectKey, + createRuntimeArtifactAttempt, + createRuntimeArtifactManifest, + getReadyRuntimeArtifact, + sealRuntimeArtifactAttempt, +} from "./runtime-artifact-attempt.repository"; +import type { + RuntimeArtifactCapturePlan, + RuntimeArtifactCapturePlanFile, + RuntimeArtifactManifestFile, + StagedRuntimeArtifactProjection, +} from "./runtime-artifact-attempt.repository"; +import { + RUNTIME_SESSION_OUTPUT_MAX_FILE_BYTES, + RUNTIME_SESSION_OUTPUT_MAX_TOTAL_BYTES, + RUNTIME_SESSION_OUTPUT_SCAN_MAX_FILES, + getRuntimeSessionOutputDirectory, + guessRuntimeSessionOutputContentType, + readRuntimeSessionOutputInventory, + toRuntimeSessionOutputArtifactPath, + toRuntimeSessionOutputFile, +} from "./runtime-session-outputs"; + +export interface DurableRuntimeArtifactEvent { + readonly event: RuntimeEventEnvelope; + readonly semanticHash: string; + readonly sourceEventId: string; +} + +interface RuntimeArtifactCaptureUpsert { + readonly contentType: string | null; + readonly expectedSize: number | null; + readonly operation: "upsert"; + readonly readPath: string; + readonly sourcePath: string; +} + +type RuntimeArtifactCaptureFile = + | RuntimeArtifactCaptureUpsert + | { readonly operation: "delete"; readonly sourcePath: string }; + +interface RuntimeArtifactCaptureSources { + readonly captureStatus: RuntimeArtifactCapturePlan["captureStatus"]; + readonly files: readonly RuntimeArtifactCaptureFile[]; + readonly mode: "delta" | "snapshot"; + readonly sandboxSessionId: string | null; + readonly sourcePaths: readonly string[]; +} + +class RuntimeArtifactCaptureUnavailable extends Error { + constructor(message: string, cause?: unknown) { + super(message, { cause }); + this.name = "RuntimeArtifactCaptureUnavailable"; + } +} + +function omittedRuntimeArtifactSources(mode: "delta" | "snapshot"): RuntimeArtifactCaptureSources { + return { + captureStatus: "omitted_runtime_unavailable", + files: [], + mode, + sandboxSessionId: null, + sourcePaths: [], + }; +} + +async function captureRuntimeArtifactSourcesOrOmit( + mode: "delta" | "snapshot", + capture: () => Promise, +): Promise { + try { + return await capture(); + } catch (error) { + if (error instanceof RuntimeArtifactCaptureUnavailable) { + return omittedRuntimeArtifactSources(mode); + } + throw error; + } +} + +function readRuntimeFileChangeContentType( + metadata: Record | undefined, +): string | null { + const contentType = metadata?.["contentType"] ?? metadata?.["mimeType"]; + return typeof contentType === "string" && contentType.trim().length > 0 ? contentType : null; +} + +function resolveRuntimeOutputCreator(link: RuntimeSessionLink): AccountId | null { + const actorId = link.executionOwnerId ?? link.callerId ?? link.creatorId; + return actorId === null ? null : parsePlatformId(actorId, "runtime output creator"); +} + +function createRuntimeSessionOutputListCommand(outputDir: string): string { + const quotedOutputDir = quoteShellArg(outputDir); + const command = [ + `if [ ! -d ${quotedOutputDir} ]; then exit 0; fi`, + `cd ${quotedOutputDir}`, + "runtime_output_scan_file=$(mktemp)", + "trap 'rm -f \"$runtime_output_scan_file\"' EXIT", + `find . -type f -print0 | head -z -n ${RUNTIME_SESSION_OUTPUT_SCAN_MAX_FILES + 1} > "$runtime_output_scan_file"; runtime_output_pipe_status=("\${PIPESTATUS[@]}"); runtime_output_find_status="\${runtime_output_pipe_status[0]}"; runtime_output_head_status="\${runtime_output_pipe_status[1]}"; if { [ "$runtime_output_find_status" -ne 0 ] && [ "$runtime_output_find_status" -ne 141 ]; } || [ "$runtime_output_head_status" -ne 0 ]; then exit 1; fi`, + 'LC_ALL=C sort -z -o "$runtime_output_scan_file" "$runtime_output_scan_file"', + "xargs -0 -r stat --printf='%n\\0%s\\0' -- < \"$runtime_output_scan_file\"", + ].join(" && "); + return `bash -lc ${quoteShellArg(command)}`; +} + +async function listRuntimeSessionOutputFiles( + handle: ExecutionSessionHandle, + outputDir: string, +): Promise> { + const result = await handle.exec(createRuntimeSessionOutputListCommand(outputDir)); + if (!result.success || result.exitCode !== 0) { + throw new Error( + result.stderr.trim() || + result.stdout.trim() || + `Failed to list runtime session outputs in ${outputDir}.`, + ); + } + const inventory = readRuntimeSessionOutputInventory(result.stdout); + if (inventory.length > RUNTIME_SESSION_OUTPUT_SCAN_MAX_FILES) { + return { captureStatus: "omitted_file_limit", files: [] }; + } + let totalBytes = 0; + for (const file of inventory) { + if ( + file.size > RUNTIME_SESSION_OUTPUT_MAX_FILE_BYTES || + totalBytes > RUNTIME_SESSION_OUTPUT_MAX_TOTAL_BYTES - file.size + ) { + return { captureStatus: "omitted_size_limit", files: [] }; + } + totalBytes += file.size; + } + return { + captureStatus: "complete", + files: inventory.map((file) => ({ + contentType: guessRuntimeSessionOutputContentType(file.relativePath), + expectedSize: file.size, + operation: "upsert", + readPath: `${outputDir}/${file.relativePath}`, + sourcePath: toRuntimeSessionOutputArtifactPath(file.relativePath), + })), + }; +} + +async function readRuntimeArtifactFileSize( + handle: ExecutionSessionHandle, + path: string, +): Promise { + const quotedPath = quoteShellArg(path); + const result = await handle.exec( + `sh -lc ${quoteShellArg( + `if runtime_output_size=$(stat --printf='%s' -- ${quotedPath}); then printf '%s' "$runtime_output_size"; elif [ ! -e ${quotedPath} ]; then exit 44; else exit 1; fi`, + )}`, + ); + if (result.exitCode === 44) { + return null; + } + if (!result.success || result.exitCode !== 0) { + throw new Error( + result.stderr.trim() || result.stdout.trim() || `Failed to stat runtime output ${path}.`, + ); + } + const output = result.stdout.trim(); + const size = Number(output); + if (!/^\d+$/.test(output) || !Number.isSafeInteger(size)) { + throw new Error(`Runtime output size is invalid for ${path}.`); + } + return size; +} + +function getRuntimeArtifactCaptureStatus( + files: readonly RuntimeArtifactCaptureFile[], +): RuntimeArtifactCaptureSources["captureStatus"] { + if (files.length > RUNTIME_SESSION_OUTPUT_SCAN_MAX_FILES) { + return "omitted_file_limit"; + } + let totalBytes = 0; + for (const file of files) { + if (file.operation === "delete") { + continue; + } + if (file.expectedSize === null) { + throw new Error(`Runtime output size is missing for ${file.sourcePath}.`); + } + if ( + file.expectedSize > RUNTIME_SESSION_OUTPUT_MAX_FILE_BYTES || + totalBytes > RUNTIME_SESSION_OUTPUT_MAX_TOTAL_BYTES - file.expectedSize + ) { + return "omitted_size_limit"; + } + totalBytes += file.expectedSize; + } + return "complete"; +} + +function isRuntimeArtifactEvent(event: RuntimeEventEnvelope): boolean { + return ( + event.kind === "file.change.updated" || + event.kind === "file.changed" || + event.kind === "run.completed" + ); +} + +async function captureRuntimeFileChangeSources(input: { + bindings: ApiBindings; + event: RuntimeEventEnvelope; + excludedSourcePaths: ReadonlySet; + link: RuntimeSessionLink; +}): Promise { + if (input.link.sessionId === null || input.link.sandboxId === null) { + throw new RuntimeArtifactCaptureUnavailable( + "Runtime file artifact event is missing its active sandbox identity.", + ); + } + const conversation = await getRuntimeConversationSession(input.bindings.DB, input.link.sessionId); + if (conversation === null) { + throw new RuntimeArtifactCaptureUnavailable( + "Runtime file artifact event is missing its sandbox conversation.", + ); + } + const filesByPath = new Map(); + for (const change of readRuntimeEventFileChanges(input.event)) { + const outputFile = toRuntimeSessionOutputFile({ + contentType: readRuntimeFileChangeContentType(change.metadata), + cwd: conversation.cwd, + path: change.path, + }); + if (outputFile === null) { + continue; + } + filesByPath.set( + outputFile.artifactPath, + change.change === "delete" + ? { operation: "delete", sourcePath: outputFile.artifactPath } + : { + contentType: outputFile.contentType, + expectedSize: null, + operation: "upsert", + readPath: outputFile.readPath, + sourcePath: outputFile.artifactPath, + }, + ); + } + const sourcePaths = [...filesByPath.keys()]; + const files = [...filesByPath.values()] + .filter((file) => !input.excludedSourcePaths.has(file.sourcePath)) + .toSorted((left, right) => left.sourcePath.localeCompare(right.sourcePath)); + if (files.length > RUNTIME_SESSION_OUTPUT_SCAN_MAX_FILES) { + return { + captureStatus: "omitted_file_limit", + files: [], + mode: "delta", + sandboxSessionId: conversation.sandboxSessionId, + sourcePaths, + }; + } + let sizedFiles = files; + if (files.some((file) => file.operation === "upsert")) { + try { + sizedFiles = await withDisposedRpcResource( + await getRuntimeSubjectKeepAliveHandle( + input.bindings, + input.link.sandboxId, + conversation.sandboxIncarnation, + ), + async (sandbox) => { + const sandboxSession = await sandbox.getSession(conversation.sandboxSessionId); + return Promise.all( + files.map(async (file) => + file.operation === "delete" + ? file + : { + contentType: file.contentType, + expectedSize: await readRuntimeArtifactFileSize(sandboxSession, file.readPath), + operation: file.operation, + readPath: file.readPath, + sourcePath: file.sourcePath, + }, + ), + ); + }, + ); + } catch (error) { + throw new RuntimeArtifactCaptureUnavailable( + "Runtime file artifact sources are unavailable.", + error, + ); + } + } + const captureStatus = sizedFiles.some( + (file) => file.operation === "upsert" && file.expectedSize === null, + ) + ? "omitted_source_missing" + : getRuntimeArtifactCaptureStatus(sizedFiles); + return { + captureStatus, + files: captureStatus === "complete" ? sizedFiles : [], + mode: "delta", + sandboxSessionId: conversation.sandboxSessionId, + sourcePaths, + }; +} + +async function captureRuntimeOutputSnapshotSources(input: { + bindings: ApiBindings; + link: RuntimeSessionLink; +}): Promise { + if (input.link.sessionId === null || input.link.sandboxId === null) { + throw new RuntimeArtifactCaptureUnavailable( + "Runtime output snapshot is missing its active sandbox identity.", + ); + } + const conversation = await getRuntimeConversationSession(input.bindings.DB, input.link.sessionId); + if (conversation === null) { + throw new RuntimeArtifactCaptureUnavailable( + "Runtime output snapshot is missing its sandbox conversation.", + ); + } + const outputDir = getRuntimeSessionOutputDirectory(conversation.cwd); + let capture: Pick; + try { + capture = await withDisposedRpcResource( + await getRuntimeSubjectKeepAliveHandle( + input.bindings, + input.link.sandboxId, + conversation.sandboxIncarnation, + ), + async (sandbox) => { + const sandboxSession = await sandbox.getSession(conversation.sandboxSessionId); + return listRuntimeSessionOutputFiles(sandboxSession, outputDir); + }, + ); + } catch (error) { + throw new RuntimeArtifactCaptureUnavailable("Runtime output snapshot is unavailable.", error); + } + return { + ...capture, + mode: "snapshot", + sandboxSessionId: conversation.sandboxSessionId, + sourcePaths: capture.files.map((file) => file.sourcePath), + }; +} + +async function stageRuntimeArtifactOmission( + bindings: ApiBindings, + input: { + readonly createdByAccountId: AccountId; + readonly driverFence: DriverRuntimeEventFence; + readonly event: DurableRuntimeArtifactEvent; + readonly mode: "delta" | "snapshot"; + readonly runId: SessionRunId; + readonly sessionId: SessionId; + }, +): Promise { + const attemptId = crypto.randomUUID(); + await createRuntimeArtifactAttempt(bindings.DB, { + attemptId, + createdByAccountId: input.createdByAccountId, + driverFence: input.driverFence, + eventType: input.event.event.kind, + runId: input.runId, + semanticHash: input.event.semanticHash, + sessionId: input.sessionId, + sourceEventId: input.event.sourceEventId, + }); + const manifest = await createRuntimeArtifactManifest({ + captureStatus: "omitted_runtime_unavailable", + files: [], + mode: input.mode, + semanticHash: input.event.semanticHash, + sourceEventId: input.event.sourceEventId, + }); + const projection = { attemptId, ...manifest }; + await sealRuntimeArtifactAttempt(bindings.DB, projection); + return projection; +} + +export async function stageRuntimeArtifactEvents( + bindings: ApiBindings, + input: { + readonly driverFence: DriverRuntimeEventFence; + readonly events: readonly DurableRuntimeArtifactEvent[]; + readonly link: RuntimeSessionLink; + }, +): Promise> { + const staged = new Map(); + if (input.link.sessionId === null || input.link.sessionRunId === null) { + if (input.events.some(({ event }) => isRuntimeArtifactEvent(event))) { + throw new Error("Runtime artifact events require a Session Run identity."); + } + return staged; + } + const runId = input.link.sessionRunId; + const createdBy = resolveRuntimeOutputCreator(input.link); + if (createdBy === null) { + throw new Error("Runtime artifact events require a creator identity."); + } + const sessionId = parsePlatformId(input.link.sessionId, "runtime output session ID"); + const artifactEvents = input.events.filter(({ event }) => isRuntimeArtifactEvent(event)); + const completedEvent = artifactEvents.find(({ event }) => event.kind === "run.completed"); + const stageUnavailableArtifacts = async (): Promise< + ReadonlyMap + > => { + const projections = new Map(); + for (const event of completedEvent === undefined ? artifactEvents : [completedEvent]) { + projections.set( + event.sourceEventId, + await stageRuntimeArtifactOmission(bindings, { + createdByAccountId: createdBy, + driverFence: input.driverFence, + event, + mode: event.event.kind === "run.completed" ? "snapshot" : "delta", + runId, + sessionId, + }), + ); + } + return projections; + }; + const captures: { event: DurableRuntimeArtifactEvent; sources: RuntimeArtifactCaptureSources }[] = + []; + if (completedEvent !== undefined) { + captures.push({ + event: completedEvent, + sources: await captureRuntimeArtifactSourcesOrOmit("snapshot", () => + captureRuntimeOutputSnapshotSources({ bindings, link: input.link }), + ), + }); + } else { + const laterPaths = new Set(); + for (let index = artifactEvents.length - 1; index >= 0; index -= 1) { + const event = artifactEvents[index]; + if (event === undefined) { + continue; + } + const sources = await captureRuntimeArtifactSourcesOrOmit("delta", () => + captureRuntimeFileChangeSources({ + bindings, + event: event.event, + excludedSourcePaths: laterPaths, + link: input.link, + }), + ); + for (const sourcePath of sources.sourcePaths) { + laterPaths.add(sourcePath); + } + captures.unshift({ + event, + sources, + }); + } + } + if (captures.some(({ sources }) => sources.captureStatus === "omitted_runtime_unavailable")) { + return stageUnavailableArtifacts(); + } + + for (const { event: durableEvent, sources } of captures) { + if ( + sources.mode === "delta" && + sources.captureStatus === "complete" && + sources.files.length === 0 + ) { + continue; + } + const proposedAttemptId = crypto.randomUUID(); + const proposedCapturePlan: RuntimeArtifactCapturePlan = { + captureStatus: sources.captureStatus, + files: sources.files.map((file) => { + if (file.operation === "delete") { + return file; + } + if (file.expectedSize === null) { + throw new Error(`Runtime output size is missing for ${file.sourcePath}.`); + } + const fileId = createPlatformId(); + return { + contentType: file.contentType, + expectedSize: file.expectedSize, + fileId, + name: getRuntimeOutputName(file.sourcePath), + objectKey: `runtime-artifact-attempts/v1/${proposedAttemptId}/files/${fileId}`, + operation: file.operation, + readPath: file.readPath, + sourcePath: file.sourcePath, + }; + }), + mode: sources.mode, + version: 1, + }; + await createRuntimeArtifactAttempt(bindings.DB, { + attemptId: proposedAttemptId, + createdByAccountId: createdBy, + driverFence: input.driverFence, + eventType: durableEvent.event.kind, + runId, + semanticHash: durableEvent.semanticHash, + sessionId, + sourceEventId: durableEvent.sourceEventId, + }); + const attemptId = proposedAttemptId; + const capturePlan = proposedCapturePlan; + + let captureStatus: RuntimeArtifactCapturePlan["captureStatus"] = capturePlan.captureStatus; + const manifestFiles: RuntimeArtifactManifestFile[] = capturePlan.files.flatMap((file) => + file.operation === "delete" ? [file] : [], + ); + const upserts = capturePlan.files.filter( + (file): file is RuntimeArtifactCapturePlanFile => file.operation === "upsert", + ); + const capturedBodies = new Map(); + if (captureStatus === "complete" && upserts.length > 0) { + const sandboxSessionId = sources.sandboxSessionId; + if (input.link.sandboxId === null || sandboxSessionId === null) { + return stageUnavailableArtifacts(); + } + const conversation = await getRuntimeConversationSession(bindings.DB, sessionId); + if ( + conversation === null || + conversation.sandboxSessionId !== sandboxSessionId || + conversation.sandboxId !== input.link.sandboxId + ) { + return stageUnavailableArtifacts(); + } + try { + await withDisposedRpcResource( + await getRuntimeSubjectKeepAliveHandle( + bindings, + input.link.sandboxId, + conversation.sandboxIncarnation, + ), + async (sandbox) => { + const sandboxSession = await sandbox.getSession(sandboxSessionId); + let totalBytes = 0; + for (const file of upserts) { + const body = await readSandboxFileBytes( + sandboxSession, + file.readPath, + RUNTIME_SESSION_OUTPUT_MAX_FILE_BYTES, + ); + if ( + body.byteLength > RUNTIME_SESSION_OUTPUT_MAX_FILE_BYTES || + totalBytes > RUNTIME_SESSION_OUTPUT_MAX_TOTAL_BYTES - body.byteLength + ) { + captureStatus = "omitted_size_limit"; + capturedBodies.clear(); + return; + } + if (body.byteLength !== file.expectedSize) { + captureStatus = "omitted_source_changed"; + capturedBodies.clear(); + return; + } + totalBytes += body.byteLength; + capturedBodies.set(file.objectKey, body); + } + }, + ); + } catch { + return stageUnavailableArtifacts(); + } + } + if (captureStatus !== "complete") { + manifestFiles.length = 0; + } else { + let storageUnavailable = false; + for (const file of upserts) { + const body = capturedBodies.get(file.objectKey); + if (body === undefined) { + throw new Error(`Runtime output content is missing for ${file.sourcePath}.`); + } + const contentSha256 = await createRuntimeOutputContentSha256(body); + const parentPath = createRuntimeOutputParentPath(file.sourcePath, contentSha256); + const existing = await getReadyRuntimeArtifact(bindings.DB, { + parentPath, + sessionId, + }); + if (existing !== null) { + manifestFiles.push({ + contentSha256, + contentType: existing.contentType, + disposition: "reuse", + etag: existing.etag, + fileId: existing.fileId, + name: existing.name, + objectKey: existing.objectKey, + operation: "upsert", + parentPath, + path: existing.path, + size: existing.size, + sourcePath: file.sourcePath, + }); + continue; + } + await claimRuntimeArtifactObjectKey(bindings.DB, { + attemptId, + objectKey: file.objectKey, + }); + const contentType = normalizeContentType(file.contentType ?? "application/octet-stream"); + const stored = await putRuntimeArtifactObject({ + attemptId, + bindings, + body, + contentSha256, + contentType, + objectKey: file.objectKey, + sourcePath: file.sourcePath, + }); + if (stored === null) { + storageUnavailable = true; + break; + } + manifestFiles.push({ + contentSha256, + contentType: stored.contentType ?? contentType, + disposition: "create", + etag: stored.etag, + fileId: file.fileId, + name: file.name, + objectKey: file.objectKey, + operation: "upsert", + parentPath, + path: createSessionArtifactPath(file.fileId, file.name), + size: stored.contentLength, + sourcePath: file.sourcePath, + }); + } + if (storageUnavailable) { + return stageUnavailableArtifacts(); + } + } + const manifest = await createRuntimeArtifactManifest({ + captureStatus, + files: manifestFiles, + mode: sources.mode, + semanticHash: durableEvent.semanticHash, + sourceEventId: durableEvent.sourceEventId, + }); + const projection = { attemptId, ...manifest }; + await sealRuntimeArtifactAttempt(bindings.DB, projection); + staged.set(durableEvent.sourceEventId, projection); + } + return staged; +} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-event-compaction.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-event-compaction.ts deleted file mode 100644 index 12e6bd6e..00000000 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-event-compaction.ts +++ /dev/null @@ -1,197 +0,0 @@ -import { - readRuntimeEventMessageContent, - readRuntimeEventPayload, - readRuntimeEventString, -} from "@mosoo/runtime-events"; -import type { - RuntimeEventEnvelope, - RuntimeEventKind, - RuntimeEventToolStatus, -} from "@mosoo/runtime-events"; - -import type { ProjectedRuntimeEventRecord } from "./event-types"; - -export type RuntimeEventPayloadRecord = Record; - -export type TextStreamKind = "message" | "thought"; - -export interface StreamAccumulator { - content: string; - firstRecord: ProjectedRuntimeEventRecord; - key: string; - lastRecord: ProjectedRuntimeEventRecord; - runId: string | null; - sessionId: string; -} - -export interface TextStreamAccumulator extends StreamAccumulator { - kind: TextStreamKind; - role: "agent" | "user"; -} - -export interface ToolCallAccumulator extends StreamAccumulator { - itemCompleted: boolean; - name: string | null; - payload: RuntimeEventPayloadRecord; - rawInput: string | null; - rawOutput: string | null; - status: RuntimeEventToolStatus; - toolCallId: string; -} - -const terminalRunKinds = new Set([ - "run.cancelled", - "run.completed", - "run.failed", -]); - -export function toStreamKey(input: { - id: string; - kind: TextStreamKind | "tool"; - runId: string | null; - sessionId: string; -}): string { - return `${input.sessionId}:${input.runId ?? ""}:${input.kind}:${input.id}`; -} - -function readRecordOccurredAt(record: ProjectedRuntimeEventRecord): number | null { - if (record.occurredAt !== null) { - return record.occurredAt; - } - - const occurredAt = Date.parse(record.event.occurredAt); - return Number.isFinite(occurredAt) ? occurredAt : null; -} - -export function readPayloadString( - payload: RuntimeEventPayloadRecord, - field: string, -): string | null { - const value = payload[field]; - return typeof value === "string" ? value : null; -} - -function appendStreamText(current: string | null, next: string | null): string | null { - if (next === null || next.length === 0) { - return current; - } - - if (current === null || current.length === 0) { - return next; - } - - return `${current}${next}`; -} - -function mergeSnapshotText(current: string | null, next: string | null): string | null { - if (next === null || next.length === 0) { - return current; - } - - if (current === null || current.length === 0) { - return next; - } - - if (next.length > current.length && next.startsWith(current)) { - return next; - } - - if (current.length >= next.length && current.startsWith(next)) { - return current; - } - - return `${current}${next}`; -} - -export function mergeDeliveredText( - current: string | null, - next: string | null, - delivery: RuntimeEventEnvelope["delivery"], -): string | null { - return delivery === "best_effort" - ? appendStreamText(current, next) - : mergeSnapshotText(current, next); -} - -export function mergeTextEventContent( - accumulator: TextStreamAccumulator, - event: RuntimeEventEnvelope, -): void { - const delta = readRuntimeEventString(readRuntimeEventPayload(event), "contentDelta"); - - if (delta !== null) { - accumulator.content = appendStreamText(accumulator.content, delta) ?? ""; - return; - } - - accumulator.content = - mergeSnapshotText(accumulator.content, readRuntimeEventMessageContent(event)) ?? ""; -} - -export function readRuntimeEventMessageRoleUpdate( - event: RuntimeEventEnvelope, -): TextStreamAccumulator["role"] | null { - const role = readRuntimeEventString(readRuntimeEventPayload(event), "role"); - - return role === "agent" || role === "user" ? role : null; -} - -function createEventWithPayload( - source: RuntimeEventEnvelope, - input: { - delivery?: RuntimeEventEnvelope["delivery"]; - kind: RuntimeEventKind; - payload: RuntimeEventPayloadRecord; - sourceEventId: string | null; - }, -): RuntimeEventEnvelope { - const { sourceEventId: _sourceEventId, ...rest } = source; - - return { - ...rest, - delivery: input.delivery ?? "lossless", - kind: input.kind, - payload: input.payload, - ...(input.sourceEventId === null ? {} : { sourceEventId: input.sourceEventId }), - }; -} - -export function createCompactedRecord( - accumulator: StreamAccumulator, - input: { - kind: RuntimeEventKind; - payload: RuntimeEventPayloadRecord; - }, -): ProjectedRuntimeEventRecord { - return { - event: createEventWithPayload(accumulator.lastRecord.event, { - kind: input.kind, - payload: input.payload, - sourceEventId: accumulator.lastRecord.sourceEventId, - }), - occurredAt: readRecordOccurredAt(accumulator.firstRecord), - sourceEventId: accumulator.lastRecord.sourceEventId, - }; -} - -export function readToolItemId(event: RuntimeEventEnvelope): string | null { - if (event.kind !== "item.completed" && event.kind !== "item.started") { - return null; - } - - const payload = readRuntimeEventPayload(event); - - if (readRuntimeEventString(payload, "itemType") !== "tool_call") { - return null; - } - - return readRuntimeEventString(payload, "itemId") ?? event.id; -} - -export function isTerminalRunEvent(event: RuntimeEventEnvelope): boolean { - return terminalRunKinds.has(event.kind); -} - -export function toTerminalToolStatus(event: RuntimeEventEnvelope): RuntimeEventToolStatus { - return event.kind === "run.failed" || event.kind === "run.cancelled" ? "failed" : "completed"; -} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-event-persistence-compactor.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-event-persistence-compactor.ts deleted file mode 100644 index 5e523ed2..00000000 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-event-persistence-compactor.ts +++ /dev/null @@ -1,440 +0,0 @@ -import { - readRuntimeEventMessageKey, - readRuntimeEventMessageRole, - readRuntimeEventPayload, - readRuntimeEventString, - readRuntimeEventToolCallId, - readRuntimeEventToolName, - readRuntimeEventToolStatusFromEvent, - readRuntimeRunPayload, -} from "@mosoo/runtime-events"; -import type { RuntimeEventEnvelope, RuntimeEventToolStatus } from "@mosoo/runtime-events"; - -import type { ProjectedRuntimeEventRecord } from "./event-types"; -import { - createCompactedRecord, - isTerminalRunEvent, - mergeDeliveredText, - mergeTextEventContent, - readPayloadString, - readRuntimeEventMessageRoleUpdate, - readToolItemId, - toStreamKey, - toTerminalToolStatus, -} from "./runtime-event-compaction"; -import type { - RuntimeEventPayloadRecord, - TextStreamAccumulator, - TextStreamKind, - ToolCallAccumulator, -} from "./runtime-event-compaction"; - -export class RuntimeEventPersistenceCompactor { - readonly #messages = new Map(); - readonly #thoughts = new Map(); - readonly #tools = new Map(); - - compact(records: ProjectedRuntimeEventRecord[]): ProjectedRuntimeEventRecord[] { - const output: ProjectedRuntimeEventRecord[] = []; - - for (const record of records) { - if (this.#applyMessage(record, output)) { - continue; - } - - if (this.#applyThought(record, output)) { - continue; - } - - if (this.#applyTool(record)) { - continue; - } - - if (isTerminalRunEvent(record.event)) { - output.push(...this.#flushRun(record.event, toTerminalToolStatus(record.event))); - output.push(record); - continue; - } - - output.push(record); - } - - output.push(...this.#flushReadyTools()); - return output; - } - - #applyMessage( - record: ProjectedRuntimeEventRecord, - output: ProjectedRuntimeEventRecord[], - ): boolean { - const event = record.event; - - if ( - event.kind !== "message.added" && - event.kind !== "message.completed" && - event.kind !== "message.delta" && - event.kind !== "message.started" - ) { - return false; - } - - const key = readRuntimeEventMessageKey(event); - - if (key === null) { - return true; - } - - const streamKey = toStreamKey({ - id: key, - kind: "message", - runId: event.runId ?? null, - sessionId: event.sessionId, - }); - const accumulator = this.#upsertTextAccumulator(this.#messages, { - key: streamKey, - kind: "message", - record, - }); - mergeTextEventContent(accumulator, event); - - if (event.kind === "message.added" || event.kind === "message.completed") { - this.#messages.delete(streamKey); - const compacted = this.#createMessageRecord(accumulator); - - if (compacted !== null) { - output.push(compacted); - } else if (event.kind === "message.added") { - output.push(record); - } - } - - return true; - } - - #applyThought( - record: ProjectedRuntimeEventRecord, - output: ProjectedRuntimeEventRecord[], - ): boolean { - const event = record.event; - - if ( - event.kind !== "thought.completed" && - event.kind !== "thought.delta" && - event.kind !== "thought.started" - ) { - return false; - } - - const key = readRuntimeEventMessageKey(event); - - if (key === null) { - return true; - } - - const streamKey = toStreamKey({ - id: key, - kind: "thought", - runId: event.runId ?? null, - sessionId: event.sessionId, - }); - const accumulator = this.#upsertTextAccumulator(this.#thoughts, { - key: streamKey, - kind: "thought", - record, - }); - mergeTextEventContent(accumulator, event); - - if (event.kind === "thought.completed") { - this.#thoughts.delete(streamKey); - const compacted = this.#createThoughtRecord(accumulator); - - if (compacted !== null) { - output.push(compacted); - } - } - - return true; - } - - #applyTool(record: ProjectedRuntimeEventRecord): boolean { - const event = record.event; - const toolCallId = readRuntimeEventToolCallId(event) ?? readToolItemId(event); - - if (toolCallId === null) { - return false; - } - - const key = toStreamKey({ - id: toolCallId, - kind: "tool", - runId: event.runId ?? null, - sessionId: event.sessionId, - }); - const accumulator = this.#upsertToolAccumulator(key, record, toolCallId); - const payload = readRuntimeEventPayload(event); - - accumulator.lastRecord = record; - this.#mergeToolPayload(accumulator, payload); - accumulator.content = - mergeDeliveredText( - accumulator.content, - readPayloadString(payload, "content"), - event.delivery, - ) ?? ""; - accumulator.rawInput = mergeDeliveredText( - accumulator.rawInput, - readPayloadString(payload, "rawInput"), - event.delivery, - ); - accumulator.rawOutput = mergeDeliveredText( - accumulator.rawOutput, - readPayloadString(payload, "rawOutput"), - event.delivery, - ); - accumulator.status = readRuntimeEventToolStatusFromEvent(event); - - const name = readRuntimeEventToolName(event); - - if (name !== null) { - accumulator.name = name; - } - - if (event.kind === "item.completed") { - accumulator.itemCompleted = true; - const status = readRuntimeEventString(payload, "status"); - accumulator.status = status === "failed" ? "failed" : "completed"; - } - - return true; - } - - #upsertTextAccumulator( - map: Map, - input: { - key: string; - kind: TextStreamKind; - record: ProjectedRuntimeEventRecord; - }, - ): TextStreamAccumulator { - const existing = map.get(input.key); - - if (existing !== undefined) { - existing.lastRecord = input.record; - existing.role = readRuntimeEventMessageRoleUpdate(input.record.event) ?? existing.role; - return existing; - } - - const accumulator: TextStreamAccumulator = { - content: "", - firstRecord: input.record, - key: input.key, - kind: input.kind, - lastRecord: input.record, - role: readRuntimeEventMessageRole(input.record.event), - runId: input.record.event.runId ?? null, - sessionId: input.record.event.sessionId, - }; - map.set(input.key, accumulator); - return accumulator; - } - - #upsertToolAccumulator( - key: string, - record: ProjectedRuntimeEventRecord, - toolCallId: string, - ): ToolCallAccumulator { - const existing = this.#tools.get(key); - - if (existing !== undefined) { - return existing; - } - - const accumulator: ToolCallAccumulator = { - content: "", - firstRecord: record, - itemCompleted: false, - key, - lastRecord: record, - name: readRuntimeEventToolName(record.event), - payload: {}, - rawInput: null, - rawOutput: null, - runId: record.event.runId ?? null, - sessionId: record.event.sessionId, - status: readRuntimeEventToolStatusFromEvent(record.event), - toolCallId, - }; - this.#tools.set(key, accumulator); - return accumulator; - } - - #createMessageRecord(accumulator: TextStreamAccumulator): ProjectedRuntimeEventRecord | null { - if (accumulator.content.length === 0) { - return null; - } - - const payload: RuntimeEventPayloadRecord = { - ...readRuntimeEventPayload(accumulator.lastRecord.event), - content: accumulator.content, - messageId: readRuntimeEventMessageKey(accumulator.lastRecord.event) ?? accumulator.key, - role: accumulator.role, - }; - delete payload["contentDelta"]; - - return createCompactedRecord(accumulator, { - kind: "message.added", - payload, - }); - } - - #createThoughtRecord(accumulator: TextStreamAccumulator): ProjectedRuntimeEventRecord | null { - if (accumulator.content.length === 0) { - return null; - } - - const payload: RuntimeEventPayloadRecord = { - ...readRuntimeEventPayload(accumulator.lastRecord.event), - content: accumulator.content, - thoughtId: readRuntimeEventMessageKey(accumulator.lastRecord.event) ?? accumulator.key, - }; - delete payload["contentDelta"]; - - return createCompactedRecord(accumulator, { - kind: "thought.completed", - payload, - }); - } - - #createToolRecord(accumulator: ToolCallAccumulator): ProjectedRuntimeEventRecord { - const payload = { - ...this.#mergeToolPayloads(accumulator), - status: accumulator.status, - toolCallId: accumulator.toolCallId, - }; - - return createCompactedRecord(accumulator, { - kind: "tool.call.updated", - payload, - }); - } - - #mergeToolPayloads(accumulator: ToolCallAccumulator): RuntimeEventPayloadRecord { - const merged: RuntimeEventPayloadRecord = { ...accumulator.payload }; - - if (accumulator.name !== null) { - merged["title"] ??= accumulator.name; - } - - if (accumulator.content.length > 0) { - merged["content"] = accumulator.content; - } - - if (accumulator.rawInput !== null) { - merged["rawInput"] = accumulator.rawInput; - } - - if (accumulator.rawOutput !== null) { - merged["rawOutput"] = accumulator.rawOutput; - } - - return merged; - } - - #mergeToolPayload(accumulator: ToolCallAccumulator, payload: RuntimeEventPayloadRecord): void { - for (const [key, value] of Object.entries(payload)) { - if (key !== "content" && key !== "rawInput" && key !== "rawOutput") { - accumulator.payload[key] = value; - } - } - } - - #flushReadyTools(): ProjectedRuntimeEventRecord[] { - const output: ProjectedRuntimeEventRecord[] = []; - - for (const [key, accumulator] of this.#tools) { - if (!accumulator.itemCompleted) { - continue; - } - - this.#tools.delete(key); - output.push(this.#createToolRecord(accumulator)); - } - - return output; - } - - #flushRun( - terminalEvent: RuntimeEventEnvelope, - terminalToolStatus: RuntimeEventToolStatus, - ): ProjectedRuntimeEventRecord[] { - const output: ProjectedRuntimeEventRecord[] = []; - const runId = terminalEvent.runId ?? null; - const terminalToolResult = readTerminalToolResult(terminalEvent); - - for (const [key, accumulator] of this.#messages) { - if (accumulator.sessionId !== terminalEvent.sessionId || accumulator.runId !== runId) { - continue; - } - - this.#messages.delete(key); - const compacted = this.#createMessageRecord(accumulator); - - if (compacted !== null) { - output.push(compacted); - } - } - - for (const [key, accumulator] of this.#thoughts) { - if (accumulator.sessionId !== terminalEvent.sessionId || accumulator.runId !== runId) { - continue; - } - - this.#thoughts.delete(key); - const compacted = this.#createThoughtRecord(accumulator); - - if (compacted !== null) { - output.push(compacted); - } - } - - for (const [key, accumulator] of this.#tools) { - if (accumulator.sessionId !== terminalEvent.sessionId || accumulator.runId !== runId) { - continue; - } - - this.#tools.delete(key); - - if (accumulator.status === "running") { - accumulator.status = terminalToolStatus; - } - - if ( - terminalToolResult !== null && - accumulator.status === "failed" && - accumulator.rawOutput === null && - accumulator.content.length === 0 - ) { - accumulator.content = terminalToolResult; - accumulator.rawOutput = terminalToolResult; - } - - output.push(this.#createToolRecord(accumulator)); - } - - return output; - } -} - -function readTerminalToolResult(event: RuntimeEventEnvelope): string | null { - if (event.kind === "run.failed") { - const run = readRuntimeRunPayload(event).run; - const message = run?.error?.message ?? "Run failed before the tool returned a result."; - return `Tool failed before returning a result: ${message}`; - } - - if (event.kind === "run.cancelled") { - return "Tool was cancelled before returning a result."; - } - - return null; -} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-event-replay-filter.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-event-replay-filter.ts deleted file mode 100644 index 86c7554a..00000000 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-event-replay-filter.ts +++ /dev/null @@ -1,208 +0,0 @@ -import type { DriverEventEnvelope } from "@mosoo/agent-driver/events"; - -type ReplayRuntimeEvent = DriverEventEnvelope["event"]; -type ReplayRuntimeEventPayload = Record; - -export function filterDurablyAcceptedRuntimeStreamReplays( - events: readonly DriverEventEnvelope[], - durableEventIds: ReadonlySet, -): DriverEventEnvelope[] { - const durableStreams = createDurableRuntimeStreamReplayIndex(events, durableEventIds); - - return events.filter((envelope) => { - if (envelope.eventId.length > 0 && durableEventIds.has(envelope.eventId)) { - return false; - } - - return !isDurablyAcceptedRuntimeStreamReplay(envelope.event, durableStreams); - }); -} - -interface DurableRuntimeStreamReplayIndex { - messageKeys: Set; - runKeys: Set; - thoughtKeys: Set; - toolKeys: Set; -} - -function createDurableRuntimeStreamReplayIndex( - events: readonly DriverEventEnvelope[], - durableEventIds: ReadonlySet, -): DurableRuntimeStreamReplayIndex { - const index: DurableRuntimeStreamReplayIndex = { - messageKeys: new Set(), - runKeys: new Set(), - thoughtKeys: new Set(), - toolKeys: new Set(), - }; - - for (const envelope of events) { - if (envelope.eventId.length === 0 || !durableEventIds.has(envelope.eventId)) { - continue; - } - - const event = envelope.event; - - if (event.kind === "message.added" || event.kind === "message.completed") { - const key = readRuntimeMessageReplayKey(event); - - if (key !== null) { - index.messageKeys.add(key); - } - continue; - } - - if (event.kind === "thought.completed") { - const key = readRuntimeThoughtReplayKey(event); - - if (key !== null) { - index.thoughtKeys.add(key); - } - continue; - } - - if (event.kind === "item.completed" || event.kind === "tool.call.updated") { - const key = readRuntimeToolReplayKey(event); - - if (key !== null) { - index.toolKeys.add(key); - } - continue; - } - - if ( - event.kind === "run.cancelled" || - event.kind === "run.completed" || - event.kind === "run.failed" - ) { - const key = readRuntimeRunReplayKey(event); - - if (key !== null) { - index.runKeys.add(key); - } - } - } - - return index; -} - -function isDurablyAcceptedRuntimeStreamReplay( - event: ReplayRuntimeEvent, - index: DurableRuntimeStreamReplayIndex, -): boolean { - if (isRunBoundRuntimeStreamEvent(event)) { - const runKey = readRuntimeRunReplayKey(event); - - if (runKey !== null && index.runKeys.has(runKey)) { - return true; - } - } - - const messageKey = readRuntimeMessageReplayKey(event); - - if (messageKey !== null && index.messageKeys.has(messageKey)) { - return true; - } - - const thoughtKey = readRuntimeThoughtReplayKey(event); - - if (thoughtKey !== null && index.thoughtKeys.has(thoughtKey)) { - return true; - } - - const toolKey = readRuntimeToolReplayKey(event); - - return toolKey !== null && index.toolKeys.has(toolKey); -} - -function isRunBoundRuntimeStreamEvent(event: ReplayRuntimeEvent): boolean { - return ( - event.kind === "item.completed" || - event.kind === "item.started" || - event.kind === "message.added" || - event.kind === "message.completed" || - event.kind === "message.delta" || - event.kind === "message.started" || - event.kind === "thought.completed" || - event.kind === "thought.delta" || - event.kind === "thought.started" || - event.kind === "tool.call.updated" - ); -} - -function readRuntimeRunReplayKey(event: ReplayRuntimeEvent): string | null { - return event.runId === undefined ? null : `${event.sessionId}:${event.runId}`; -} - -function readRuntimeMessageReplayKey(event: ReplayRuntimeEvent): string | null { - if ( - event.kind !== "message.added" && - event.kind !== "message.completed" && - event.kind !== "message.delta" && - event.kind !== "message.started" - ) { - return null; - } - - const payload = readRuntimeEventPayload(event); - const messageKey = readRuntimeEventString(payload, "messageId") ?? event.id; - return messageKey === null ? null : `${event.sessionId}:${event.runId ?? ""}:${messageKey}`; -} - -function readRuntimeThoughtReplayKey(event: ReplayRuntimeEvent): string | null { - if ( - event.kind !== "thought.completed" && - event.kind !== "thought.delta" && - event.kind !== "thought.started" - ) { - return null; - } - - const payload = readRuntimeEventPayload(event); - const thoughtKey = readRuntimeEventString(payload, "thoughtId") ?? event.id; - return thoughtKey === null ? null : `${event.sessionId}:${event.runId ?? ""}:${thoughtKey}`; -} - -function readRuntimeToolReplayKey(event: ReplayRuntimeEvent): string | null { - const toolCallId = readRuntimeEventToolCallId(event) ?? readRuntimeToolItemReplayId(event); - return toolCallId === null ? null : `${event.sessionId}:${event.runId ?? ""}:${toolCallId}`; -} - -function readRuntimeToolItemReplayId(event: ReplayRuntimeEvent): string | null { - if (event.kind !== "item.completed" && event.kind !== "item.started") { - return null; - } - - const payload = readRuntimeEventPayload(event); - - if (readRuntimeEventString(payload, "itemType") !== "tool_call") { - return null; - } - - return readRuntimeEventString(payload, "itemId") ?? event.id; -} - -function readRuntimeEventToolCallId(event: ReplayRuntimeEvent): string | null { - if (event.kind !== "tool.call.updated") { - return null; - } - - return readRuntimeEventString(readRuntimeEventPayload(event), "toolCallId") ?? event.id; -} - -function readRuntimeEventPayload(event: ReplayRuntimeEvent): ReplayRuntimeEventPayload { - return isRuntimeEventPayload(event.payload) ? event.payload : {}; -} - -function isRuntimeEventPayload(value: unknown): value is ReplayRuntimeEventPayload { - return typeof value === "object" && value !== null && !Array.isArray(value); -} - -function readRuntimeEventString(value: unknown, field: string): string | null { - if (!isRuntimeEventPayload(value)) { - return null; - } - - const entry = value[field]; - return typeof entry === "string" && entry.length > 0 ? entry : null; -} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-session-outputs.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-session-outputs.ts index 2ccf1d22..a6741faf 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-session-outputs.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-session-outputs.ts @@ -1,5 +1,14 @@ +import { normalizeLibraryFilePath } from "@mosoo/contracts/file"; + export const RUNTIME_SESSION_OUTPUT_DIR_NAME = "outputs"; export const RUNTIME_SESSION_OUTPUT_SCAN_MAX_FILES = 100; +export const RUNTIME_SESSION_OUTPUT_MAX_FILE_BYTES = 8 * 1024 * 1024; +export const RUNTIME_SESSION_OUTPUT_MAX_TOTAL_BYTES = 32 * 1024 * 1024; + +export interface RuntimeSessionOutputInventoryFile { + readonly relativePath: string; + readonly size: number; +} export interface RuntimeSessionOutputFile { readonly artifactPath: string; @@ -27,10 +36,6 @@ const contentTypesByExtension = new Map([ ["zip", "application/zip"], ]); -function normalizeSandboxPath(value: string): string { - return value.trim().replaceAll("\\", "/").replace(/\/+$/, ""); -} - function joinSandboxPath(parent: string, child: string): string { return `${parent.replace(/\/+$/, "")}/${child.replace(/^\/+/, "")}`; } @@ -44,31 +49,12 @@ export function normalizeRuntimeSessionOutputRelativePath(value: unknown): strin return null; } - const normalizedPath = value.trim().replaceAll("\\", "/"); - - if ( - normalizedPath.length === 0 || - normalizedPath.includes("\0") || - normalizedPath.startsWith("/") - ) { + try { + const normalizedPath = normalizeLibraryFilePath(value); + return normalizedPath === value ? normalizedPath : null; + } catch { return null; } - - const segments: string[] = []; - - for (const segment of normalizedPath.split("/")) { - if (segment.length === 0 || segment === ".") { - continue; - } - - if (segment === "..") { - return null; - } - - segments.push(segment); - } - - return segments.length === 0 ? null : segments.join("/"); } export function toRuntimeSessionOutputArtifactPath(relativePath: string): string { @@ -91,18 +77,16 @@ export function toRuntimeSessionOutputFile(input: { readonly path: string; }): RuntimeSessionOutputFile | null { const outputDir = getRuntimeSessionOutputDirectory(input.cwd); - const normalizedPath = normalizeSandboxPath(input.path); + const normalizedPath = input.path; let relativePath: string | null; if (normalizedPath.startsWith("/")) { - const normalizedOutputDir = normalizeSandboxPath(outputDir); - - if (!normalizedPath.startsWith(`${normalizedOutputDir}/`)) { + if (!normalizedPath.startsWith(`${outputDir}/`)) { return null; } relativePath = normalizeRuntimeSessionOutputRelativePath( - normalizedPath.slice(normalizedOutputDir.length + 1), + normalizedPath.slice(outputDir.length + 1), ); } else { const normalizedRelativePath = normalizeRuntimeSessionOutputRelativePath(normalizedPath); @@ -140,24 +124,41 @@ export function toRuntimeSessionOutputFile(input: { }; } -export function readRuntimeSessionOutputListing(stdout: string): string[] { +export function readRuntimeSessionOutputInventory( + stdout: string, +): RuntimeSessionOutputInventoryFile[] { + const files: RuntimeSessionOutputInventoryFile[] = []; const seen = new Set(); - const paths: string[] = []; + const fields = stdout.split("\0"); - for (const line of stdout.split("\n")) { - const normalizedPath = normalizeRuntimeSessionOutputRelativePath(line); + if (fields.at(-1) === "") { + fields.pop(); + } + if (fields.length % 2 !== 0) { + throw new Error("Runtime output inventory is invalid."); + } - if (normalizedPath === null || seen.has(normalizedPath)) { - continue; + for (let index = 0; index < fields.length; index += 2) { + const inventoryPath = fields[index]; + const relativePath = inventoryPath?.startsWith("./") + ? normalizeRuntimeSessionOutputRelativePath(inventoryPath.slice(2)) + : null; + const sizeText = fields[index + 1]; + const size = Number(sizeText); + + if ( + relativePath === null || + seen.has(relativePath) || + !/^\d+$/.test(sizeText ?? "") || + !Number.isSafeInteger(size) || + size < 0 + ) { + throw new Error("Runtime output inventory is invalid."); } - seen.add(normalizedPath); - paths.push(normalizedPath); - - if (paths.length >= RUNTIME_SESSION_OUTPUT_SCAN_MAX_FILES) { - break; - } + seen.add(relativePath); + files.push({ relativePath, size }); } - return paths; + return files; } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-state-store.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-state-store.ts index 12ab0f27..1bfbdae2 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-state-store.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-state-store.ts @@ -6,14 +6,16 @@ import { import type { DriverHeartbeatInput, DriverHelloInput, + DriverHelloOutput, DriverReadyInput, } from "@mosoo/agent-driver/orpc"; import { parseRuntimeCommand } from "@mosoo/contracts/runtime-command"; import type { RuntimeCommand } from "@mosoo/contracts/runtime-command"; import { parsePlatformId } from "@mosoo/id"; -import type { DriverInstanceId } from "@mosoo/id"; +import type { DriverInstanceId, SessionRunId } from "@mosoo/id"; -import type { DriverInstanceCloseSnapshot } from "./state"; +import { parseDriverHelloOutput } from "./rpc-wire"; +import type { DriverInstanceCloseSnapshot, DriverInstanceConnectionEpoch } from "./state"; export const HEARTBEAT_STATE_PERSIST_INTERVAL_MS = 10_000; export const DRIVER_INSTANCE_STATE_STORAGE_KEY = "driverInstanceState"; @@ -28,11 +30,27 @@ export interface DriverInstanceStoredState { errorMessage: string | null; heartbeatCount: number; hello: DriverHelloInput | null; + helloOutput: DriverHelloOutput | null; lastHeartbeat: DriverHeartbeatInput | null; + pendingHello: DriverInstancePendingHello | null; + pendingReady: DriverInstancePendingReady | null; ready: DriverReadyInput | null; + terminalCleanupComplete: boolean; + terminalSessionRunId: SessionRunId | null; traceId: string | null; } +export interface DriverInstancePendingHello { + epoch: DriverInstanceConnectionEpoch; + input: DriverHelloInput; + output: DriverHelloOutput; +} + +export interface DriverInstancePendingReady { + epoch: DriverInstanceConnectionEpoch; + input: DriverReadyInput; +} + interface DriverInstanceRuntimeStorage { deleteAll(): Promise; get(key: string): Promise; @@ -72,8 +90,13 @@ export function parseStoredState(value: unknown): DriverInstanceStoredState { errorMessage: readNullableString(value, "errorMessage"), heartbeatCount: readRequiredNumber(value, "heartbeatCount"), hello: parseNullableHello(value["hello"]), + helloOutput: parseOptionalNullableHelloOutput(value["helloOutput"]), lastHeartbeat: parseNullableHeartbeat(value["lastHeartbeat"]), + pendingHello: parseOptionalPendingHello(value["pendingHello"]), + pendingReady: parseOptionalPendingReady(value["pendingReady"]), ready: parseNullableReady(value["ready"]), + terminalCleanupComplete: readOptionalBoolean(value, "terminalCleanupComplete", false), + terminalSessionRunId: readOptionalNullableSessionRunId(value, "terminalSessionRunId"), traceId: readNullableString(value, "traceId"), }; } @@ -89,12 +112,43 @@ export function createEmptyStoredState(): DriverInstanceStoredState { errorMessage: null, heartbeatCount: 0, hello: null, + helloOutput: null, lastHeartbeat: null, + pendingHello: null, + pendingReady: null, ready: null, + terminalCleanupComplete: false, + terminalSessionRunId: null, traceId: null, }; } +function readOptionalBoolean( + value: Record, + field: string, + defaultValue: boolean, +): boolean { + const entry = value[field]; + + if (entry === undefined) { + return defaultValue; + } + if (typeof entry !== "boolean") { + throw new TypeError(`Driver instance stored state ${field} must be a boolean.`); + } + + return entry; +} + +function readOptionalNullableSessionRunId( + value: Record, + field: string, +): SessionRunId | null { + const entry = value[field]; + + return entry === undefined || entry === null ? null : parsePlatformId(entry, field); +} + function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); } @@ -188,6 +242,54 @@ function parseNullableHello(value: unknown): DriverHelloInput | null { return value === null ? null : parseDriverHelloInput(value); } +function parseOptionalNullableHelloOutput(value: unknown): DriverHelloOutput | null { + return value === undefined || value === null ? null : parseDriverHelloOutput(value); +} + +function parseConnectionEpoch(value: unknown): DriverInstanceConnectionEpoch { + if (!isRecord(value)) { + throw new TypeError("Driver instance connection epoch must be an object."); + } + + const connectionId = readRequiredString(value, "connectionId"); + const generation = readRequiredNumber(value, "generation"); + + if (connectionId.length === 0 || !Number.isSafeInteger(generation)) { + throw new TypeError("Driver instance connection epoch is invalid."); + } + + return { connectionId, generation }; +} + +function parseOptionalPendingHello(value: unknown): DriverInstancePendingHello | null { + if (value === undefined || value === null) { + return null; + } + if (!isRecord(value)) { + throw new TypeError("Pending Driver hello must be an object."); + } + + return { + epoch: parseConnectionEpoch(value["epoch"]), + input: parseDriverHelloInput(value["input"]), + output: parseDriverHelloOutput(value["output"]), + }; +} + +function parseOptionalPendingReady(value: unknown): DriverInstancePendingReady | null { + if (value === undefined || value === null) { + return null; + } + if (!isRecord(value)) { + throw new TypeError("Pending Driver ready must be an object."); + } + + return { + epoch: parseConnectionEpoch(value["epoch"]), + input: parseDriverReadyInput(value["input"]), + }; +} + function parseNullableReady(value: unknown): DriverReadyInput | null { if (value === null) { return null; diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-state.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-state.ts index d500dc11..afe6cc0a 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-state.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/runtime-state.ts @@ -1,23 +1,16 @@ -import type { DriverEventEnvelope } from "@mosoo/agent-driver/events"; import type { - DriverEventReceipt, DriverHeartbeatInput, DriverHelloInput, + DriverHelloOutput, DriverReadyInput, } from "@mosoo/agent-driver/orpc"; import type { RuntimeCommand } from "@mosoo/contracts/runtime-command"; -import type { DriverInstanceId } from "@mosoo/id"; +import type { DriverInstanceId, SessionRunId } from "@mosoo/id"; import { isTruthy } from "../../../../shared/truthiness"; import type { DriverInstanceCommandState, RuntimeCommandWaiter } from "./commands"; import { createDriverDebugResumeSnapshot } from "./debug-resume-snapshot"; import type { DriverDebugRecoveryMode } from "./debug-resume-snapshot"; -import { - createReceiptsForDriverEvents, - filterNewDriverEvents, - readReceiptsForProcessedDriverEvents, - rememberDriverEventReceipts, -} from "./driver-event-receipts"; import { createDeferred, withTimeout } from "./driver-instance-support"; import type { Deferred } from "./driver-instance-support"; import type { RuntimeSessionLink } from "./event-types"; @@ -29,21 +22,33 @@ import { parseStoredState, } from "./runtime-state-store"; import type { + DriverInstancePendingHello, + DriverInstancePendingReady, DriverInstanceRuntimeStateContext, DriverInstanceStoredState, } from "./runtime-state-store"; import type { DriverInstanceCloseSnapshot, - DriverInstanceHeartbeatResult, - DriverInstanceHelloResult, + DriverInstanceConnectionEpoch, DriverInstanceReadyResult, DriverInstanceSnapshot, DriverInstanceWaitForCloseResult, - HeartbeatWaiter, } from "./state"; +interface GenerationWaiter { + deferred: Deferred; + generation: number; +} + +export type DriverInstanceHandshakeStageOutcome = "applied" | "replay" | "resume"; + +function isExactJson(left: unknown, right: unknown): boolean { + return JSON.stringify(left) === JSON.stringify(right); +} + interface DriverInstanceResetOptions { beforeReset: () => Promise; + driverGeneration: number; } export interface DriverInstanceHeartbeatRecord { @@ -52,25 +57,26 @@ export interface DriverInstanceHeartbeatRecord { export class DriverInstanceRuntimeState { close: DriverInstanceCloseSnapshot | null = null; - readonly closeWaiters: Deferred[] = []; + readonly closeWaiters: GenerationWaiter[] = []; commandQueue: RuntimeCommand[] = []; readonly commandWaiters: RuntimeCommandWaiter[] = []; connectedAt: number | null = null; connectionId: string | null = null; driverGeneration: number | null = null; driverInstanceId: DriverInstanceId | null = null; - driverEventReceiptSeq = 0; errorMessage: string | null = null; heartbeatCount = 0; - readonly heartbeatWaiters: HeartbeatWaiter[] = []; hello: DriverHelloInput | null = null; - readonly helloWaiters: Deferred[] = []; + helloOutput: DriverHelloOutput | null = null; lastHeartbeat: DriverHeartbeatInput | null = null; lastPersistedHeartbeatAtMs: number | null = null; - readonly processedDriverEventReceipts = new Map(); + pendingHello: DriverInstancePendingHello | null = null; + pendingReady: DriverInstancePendingReady | null = null; ready: DriverReadyInput | null = null; - readonly readyWaiters: Deferred[] = []; + readonly readyWaiters: GenerationWaiter[] = []; runtimeSessionLink: RuntimeSessionLink | null = null; + terminalCleanupComplete = false; + terminalSessionRunId: SessionRunId | null = null; terminalized = false; traceId: string | null = null; readonly #ctx: DriverInstanceRuntimeStateContext; @@ -86,17 +92,20 @@ export class DriverInstanceRuntimeState { this.connectionId = snapshot.connectionId; this.driverGeneration = snapshot.driverGeneration; this.driverInstanceId = snapshot.driverInstanceId; - this.driverEventReceiptSeq = 0; this.errorMessage = snapshot.errorMessage; this.heartbeatCount = snapshot.heartbeatCount; this.hello = snapshot.hello; + this.helloOutput = snapshot.helloOutput; this.lastHeartbeat = snapshot.lastHeartbeat; this.lastPersistedHeartbeatAtMs = snapshot.lastHeartbeat ? parseHeartbeatTimestampMs(snapshot.lastHeartbeat.at) : null; - this.processedDriverEventReceipts.clear(); + this.pendingHello = snapshot.pendingHello; + this.pendingReady = snapshot.pendingReady; this.ready = snapshot.ready ?? null; this.runtimeSessionLink = null; + this.terminalCleanupComplete = snapshot.terminalCleanupComplete; + this.terminalSessionRunId = snapshot.terminalSessionRunId; this.terminalized = snapshot.close !== null; this.traceId = snapshot.traceId; } @@ -116,8 +125,13 @@ export class DriverInstanceRuntimeState { errorMessage: this.errorMessage, heartbeatCount: this.heartbeatCount, hello: this.hello, + helloOutput: this.helloOutput, lastHeartbeat: this.lastHeartbeat, + pendingHello: this.pendingHello, + pendingReady: this.pendingReady, ready: this.ready, + terminalCleanupComplete: this.terminalCleanupComplete, + terminalSessionRunId: this.terminalSessionRunId, traceId: this.traceId, }; } @@ -136,19 +150,37 @@ export class DriverInstanceRuntimeState { ); } - async persistClose(close: DriverInstanceCloseSnapshot): Promise { + async persistClose( + close: DriverInstanceCloseSnapshot, + epoch: DriverInstanceConnectionEpoch, + ): Promise { + this.assertConnectionEpoch(epoch); + await this.#ctx.storage.put(DRIVER_INSTANCE_STATE_STORAGE_KEY, { + ...this.#toStoredState(), + close, + terminalCleanupComplete: false, + }); + this.assertConnectionEpoch(epoch); this.close = close; - await this.#persistState(); + this.terminalCleanupComplete = false; + this.terminalized = true; } async persistCommandQueue(): Promise { await this.#persistState(); } - async persistTerminalSnapshot(): Promise { + async persistTerminalSnapshot(epoch: DriverInstanceConnectionEpoch): Promise { + this.assertConnectionEpoch(epoch); this.requireDriverInstanceId(); + await this.#ctx.storage.put(DRIVER_INSTANCE_STATE_STORAGE_KEY, { + ...this.#toStoredState(), + commandQueue: [], + terminalCleanupComplete: true, + }); + this.assertConnectionEpoch(epoch); this.commandQueue = []; - await this.#persistState(); + this.terminalCleanupComplete = true; } async recordAcceptedConnection(input: { @@ -157,15 +189,30 @@ export class DriverInstanceRuntimeState { driverGeneration: number; traceId: string | null; }): Promise { - this.connectedAt = input.connectedAt; - this.connectionId = input.connectionId; - this.driverGeneration = input.driverGeneration; - - if (input.traceId !== null) { - this.traceId = input.traceId; - } + const isNewConnection = + this.connectionId !== input.connectionId || this.driverGeneration !== input.driverGeneration; + const next: DriverInstanceStoredState = { + ...this.#toStoredState(), + connectedAt: input.connectedAt, + connectionId: input.connectionId, + driverGeneration: input.driverGeneration, + errorMessage: null, + ...(isNewConnection + ? { + heartbeatCount: 0, + hello: null, + helloOutput: null, + lastHeartbeat: null, + pendingHello: null, + pendingReady: null, + ready: null, + } + : {}), + traceId: input.traceId ?? this.traceId, + }; - await this.#persistState(); + await this.#ctx.storage.put(DRIVER_INSTANCE_STATE_STORAGE_KEY, next); + this.#applyStoredState(next); } async recordHeartbeat(payload: DriverHeartbeatInput): Promise { @@ -179,71 +226,120 @@ export class DriverInstanceRuntimeState { this.lastPersistedHeartbeatAtMs = heartbeatAtMs; } - const result: DriverInstanceHeartbeatResult = { - heartbeat: payload, - heartbeatCount: this.heartbeatCount, - lastHeartbeatAt: payload.at, - }; - - for (const waiter of this.heartbeatWaiters.splice(0)) { - if (result.heartbeatCount > waiter.afterCount) { - waiter.deferred.resolve(result); - continue; - } - - this.heartbeatWaiters.push(waiter); - } - return { shouldPersistCanonical }; } - async recordHello(input: DriverHelloInput): Promise { - if (this.hello) { - throw new Error("Driver hello has already been received."); + async stageHello( + epoch: DriverInstanceConnectionEpoch, + input: DriverHelloInput, + output: DriverHelloOutput, + ): Promise { + this.assertConnectionEpoch(epoch); + + if (this.hello !== null || this.helloOutput !== null) { + if (isExactJson(this.hello, input) && isExactJson(this.helloOutput, output)) { + return "replay"; + } + throw new Error("Driver hello conflicts with the canonical receipt."); } - this.hello = input; - await this.#persistState(); + const pending = { epoch, input, output } satisfies DriverInstancePendingHello; - const result: DriverInstanceHelloResult = { - heartbeatCount: this.heartbeatCount, - hello: input, - lastHeartbeatAt: this.lastHeartbeat?.at ?? null, - }; + if (this.pendingHello !== null) { + if (isExactJson(this.pendingHello, pending)) { + return "resume"; + } + throw new Error("Driver hello conflicts with the pending receipt."); + } - return result; + await this.#ctx.storage.put(DRIVER_INSTANCE_STATE_STORAGE_KEY, { + ...this.#toStoredState(), + pendingHello: pending, + }); + this.assertConnectionEpoch(epoch); + this.pendingHello = pending; + return "applied"; } - async recordReady(input: DriverReadyInput): Promise { - if (this.ready) { - throw new Error("Driver ready has already been received."); + async commitHello(epoch: DriverInstanceConnectionEpoch): Promise { + this.assertConnectionEpoch(epoch); + const pending = this.pendingHello; + + if (pending === null || !isExactJson(pending.epoch, epoch)) { + throw new Error("Pending Driver hello receipt was lost."); } - this.ready = input; - await this.#persistState(); + await this.#ctx.storage.put(DRIVER_INSTANCE_STATE_STORAGE_KEY, { + ...this.#toStoredState(), + hello: pending.input, + helloOutput: pending.output, + pendingHello: null, + }); + this.assertConnectionEpoch(epoch); + this.hello = pending.input; + this.helloOutput = pending.output; + this.pendingHello = null; + return pending.output; + } + + async stageReady( + epoch: DriverInstanceConnectionEpoch, + input: DriverReadyInput, + ): Promise { + this.assertConnectionEpoch(epoch); + + if (this.ready !== null) { + if (isExactJson(this.ready, input)) { + return "replay"; + } + throw new Error("Driver ready conflicts with the canonical receipt."); + } - return { - heartbeatCount: this.heartbeatCount, - lastHeartbeatAt: this.lastHeartbeat?.at ?? null, - ready: input, - }; - } + const pending = { epoch, input } satisfies DriverInstancePendingReady; - rejectHeartbeatWaiters(error: Error): void { - for (const waiter of this.heartbeatWaiters.splice(0)) { - waiter.deferred.reject(error); + if (this.pendingReady !== null) { + if (isExactJson(this.pendingReady, pending)) { + return "resume"; + } + throw new Error("Driver ready conflicts with the pending receipt."); } + + await this.#ctx.storage.put(DRIVER_INSTANCE_STATE_STORAGE_KEY, { + ...this.#toStoredState(), + pendingReady: pending, + }); + this.assertConnectionEpoch(epoch); + this.pendingReady = pending; + return "applied"; } - rejectHelloWaiters(error: Error): void { - for (const waiter of this.helloWaiters.splice(0)) { - waiter.reject(error); + async commitReady(epoch: DriverInstanceConnectionEpoch): Promise { + this.assertConnectionEpoch(epoch); + const pending = this.pendingReady; + + if (pending === null || !isExactJson(pending.epoch, epoch)) { + throw new Error("Pending Driver ready receipt was lost."); } + + await this.#ctx.storage.put(DRIVER_INSTANCE_STATE_STORAGE_KEY, { + ...this.#toStoredState(), + pendingReady: null, + ready: pending.input, + }); + this.assertConnectionEpoch(epoch); + this.pendingReady = null; + this.ready = pending.input; + + return this.readyResult(); } - rejectReadyWaiters(error: Error): void { + rejectReadyWaiters(error: Error, generation: number): void { for (const waiter of this.readyWaiters.splice(0)) { - waiter.reject(error); + if (waiter.generation === generation) { + waiter.deferred.reject(error); + } else { + this.readyWaiters.push(waiter); + } } } @@ -271,12 +367,47 @@ export class DriverInstanceRuntimeState { return this.driverGeneration; } + connectionEpoch(): DriverInstanceConnectionEpoch | null { + return this.connectionId === null || this.driverGeneration === null + ? null + : { connectionId: this.connectionId, generation: this.driverGeneration }; + } + + requireConnectionEpoch(): DriverInstanceConnectionEpoch { + const epoch = this.connectionEpoch(); + + if (epoch === null) { + throw new Error("Driver connection epoch was not initialized."); + } + + return epoch; + } + + matchesConnectionEpoch(epoch: DriverInstanceConnectionEpoch): boolean { + return this.connectionId === epoch.connectionId && this.driverGeneration === epoch.generation; + } + + assertConnectionEpoch(epoch: DriverInstanceConnectionEpoch): void { + if (!this.matchesConnectionEpoch(epoch)) { + throw new Error("Driver connection is no longer current."); + } + } + async resetForReuse(options: DriverInstanceResetOptions): Promise { await options.beforeReset(); + const staleError = new Error("Driver generation is no longer current."); + for (const waiter of this.closeWaiters.splice(0)) { + waiter.deferred.reject(staleError); + } + for (const waiter of this.readyWaiters.splice(0)) { + waiter.deferred.reject(staleError); + } + const driverInstanceId = this.requireDriverInstanceId(); this.#applyStoredState({ ...createEmptyStoredState(), + driverGeneration: options.driverGeneration, driverInstanceId, }); await this.#ctx.storage.deleteAll(); @@ -287,49 +418,80 @@ export class DriverInstanceRuntimeState { const error = new Error(reason); for (const waiter of this.closeWaiters.splice(0)) { - waiter.reject(error); + waiter.deferred.reject(error); } for (const waiter of this.commandWaiters.splice(0)) { waiter.deferred.resolve(null); } - for (const waiter of this.helloWaiters.splice(0)) { - waiter.reject(error); - } - - for (const waiter of this.heartbeatWaiters.splice(0)) { - waiter.deferred.reject(error); - } - for (const waiter of this.readyWaiters.splice(0)) { - waiter.reject(error); + waiter.deferred.reject(error); } this.#applyStoredState(createEmptyStoredState()); } - resolveCloseWaiters(result: DriverInstanceWaitForCloseResult): void { + resolveCloseWaiters(result: DriverInstanceWaitForCloseResult, generation: number): void { for (const waiter of this.closeWaiters.splice(0)) { - waiter.resolve(result); - } - } - - resolveHelloWaiters(result: DriverInstanceHelloResult): void { - for (const waiter of this.helloWaiters.splice(0)) { - waiter.resolve(result); + if (waiter.generation === generation) { + waiter.deferred.resolve(result); + } else { + this.closeWaiters.push(waiter); + } } } - resolveReadyWaiters(result: DriverInstanceReadyResult): void { + resolveReadyWaiters(result: DriverInstanceReadyResult, generation: number): void { for (const waiter of this.readyWaiters.splice(0)) { - waiter.resolve(result); + if (waiter.generation === generation) { + waiter.deferred.resolve(result); + } else { + this.readyWaiters.push(waiter); + } } } - async setDriverInstanceId(driverInstanceId: DriverInstanceId): Promise { + async initializeDriverInstance( + driverInstanceId: DriverInstanceId, + driverGeneration: number, + ): Promise { + await this.#ctx.storage.put(DRIVER_INSTANCE_STATE_STORAGE_KEY, { + ...this.#toStoredState(), + driverGeneration, + driverInstanceId, + }); + this.driverGeneration = driverGeneration; this.driverInstanceId = driverInstanceId; - await this.#persistState(); + } + + async setDriverGeneration(driverGeneration: number): Promise { + await this.#ctx.storage.put(DRIVER_INSTANCE_STATE_STORAGE_KEY, { + ...this.#toStoredState(), + driverGeneration, + }); + this.driverGeneration = driverGeneration; + } + + async setTerminalSessionRunId( + sessionRunId: SessionRunId, + epoch?: DriverInstanceConnectionEpoch, + ): Promise { + if (epoch !== undefined) { + this.assertConnectionEpoch(epoch); + } + if (this.terminalSessionRunId !== null && this.terminalSessionRunId !== sessionRunId) { + throw new Error("Terminal Session Run identity is already fixed."); + } + + await this.#ctx.storage.put(DRIVER_INSTANCE_STATE_STORAGE_KEY, { + ...this.#toStoredState(), + terminalSessionRunId: sessionRunId, + }); + if (epoch !== undefined) { + this.assertConnectionEpoch(epoch); + } + this.terminalSessionRunId = sessionRunId; } async setErrorMessage(message: string): Promise { @@ -339,9 +501,26 @@ export class DriverInstanceRuntimeState { this.errorMessage = message; await this.#persistState(); - this.rejectHelloWaiters(new Error(message)); - this.rejectHeartbeatWaiters(new Error(message)); - this.rejectReadyWaiters(new Error(message)); + this.rejectReadyWaiters(new Error(message), this.requireDriverGeneration()); + } + + async setConnectionErrorMessage( + epoch: DriverInstanceConnectionEpoch, + message: string, + ): Promise { + this.assertConnectionEpoch(epoch); + + if (isTruthy(this.errorMessage)) { + return; + } + + await this.#ctx.storage.put(DRIVER_INSTANCE_STATE_STORAGE_KEY, { + ...this.#toStoredState(), + errorMessage: message, + }); + this.assertConnectionEpoch(epoch); + this.errorMessage = message; + this.rejectReadyWaiters(new Error(message), epoch.generation); } setRuntimeSessionLink(link: RuntimeSessionLink): void { @@ -355,43 +534,19 @@ export class DriverInstanceRuntimeState { ); } - filterUnprocessedDriverEvents(events: readonly DriverEventEnvelope[]): DriverEventEnvelope[] { - return filterNewDriverEvents({ - events, - processedReceipts: this.processedDriverEventReceipts, - }); - } - - createDriverEventReceipts(events: readonly DriverEventEnvelope[]): DriverEventReceipt[] { - const result = createReceiptsForDriverEvents({ - events, - nextSeq: this.driverEventReceiptSeq, - }); - this.driverEventReceiptSeq = result.nextSeq; - return result.receipts; - } - - readProcessedDriverEventReceipts(events: readonly DriverEventEnvelope[]): DriverEventReceipt[] { - return readReceiptsForProcessedDriverEvents({ - events, - processedReceipts: this.processedDriverEventReceipts, - }); - } - - rememberProcessedDriverEventReceipts(receipts: DriverEventReceipt[]): void { - rememberDriverEventReceipts({ - processedReceipts: this.processedDriverEventReceipts, - receipts, - }); - } - - async setTraceId(traceId: string): Promise { + async setTraceId(traceId: string, epoch?: DriverInstanceConnectionEpoch): Promise { + if (epoch !== undefined) { + this.assertConnectionEpoch(epoch); + } if (this.traceId === traceId) { return; } this.traceId = traceId; await this.#persistState(); + if (epoch !== undefined) { + this.assertConnectionEpoch(epoch); + } } snapshot(driverSocketConnected: boolean): DriverInstanceSnapshot { @@ -399,7 +554,6 @@ export class DriverInstanceRuntimeState { return { close: this.close, debugResume: createDriverDebugResumeSnapshot({ - lastEventSeq: this.driverEventReceiptSeq, recoveryMode, sandboxId: this.runtimeSessionLink?.sandboxId ?? null, }), @@ -426,103 +580,88 @@ export class DriverInstanceRuntimeState { return "fresh"; } - async waitForClose(timeoutMs: number): Promise { + async waitForClose( + generation: number, + timeoutMs: number, + ): Promise { + this.assertGeneration(generation); + if (this.close) { return this.closeResult(); } - const deferred = createDeferred(); - this.closeWaiters.push(deferred); - return withTimeout( - deferred.promise, - timeoutMs, - `Driver instance ${this.requireDriverInstanceId()} close`, - ); + const waiter: GenerationWaiter = { + deferred: createDeferred(), + generation, + }; + this.closeWaiters.push(waiter); + try { + return await withTimeout( + waiter.deferred.promise, + timeoutMs, + `Driver instance ${this.requireDriverInstanceId()} close`, + ); + } finally { + this.#removeWaiter(this.closeWaiters, waiter); + } } - async waitForHeartbeat( - afterCount: number, - timeoutMs: number, - ): Promise { - if (this.lastHeartbeat && this.heartbeatCount > afterCount) { - return { - heartbeat: this.lastHeartbeat, - heartbeatCount: this.heartbeatCount, - lastHeartbeatAt: this.lastHeartbeat.at, - }; - } + async waitForReady(generation: number, timeoutMs: number): Promise { + this.assertGeneration(generation); if (isTruthy(this.errorMessage)) { throw new Error(this.errorMessage); } if (this.close) { - throw new Error(`Driver instance ${this.requireDriverInstanceId()} is already closed.`); + throw new Error(`Driver instance ${this.requireDriverInstanceId()} closed before ready.`); + } + + if (this.ready) { + return this.readyResult(); } - const waiter: HeartbeatWaiter = { - afterCount, - deferred: createDeferred(), + const waiter: GenerationWaiter = { + deferred: createDeferred(), + generation, }; - this.heartbeatWaiters.push(waiter); - return withTimeout( - waiter.deferred.promise, - timeoutMs, - `Driver instance ${this.requireDriverInstanceId()} heartbeat`, - ); + this.readyWaiters.push(waiter); + try { + return await withTimeout( + waiter.deferred.promise, + timeoutMs, + `Driver instance ${this.requireDriverInstanceId()} ready`, + ); + } finally { + this.#removeWaiter(this.readyWaiters, waiter); + } } - async waitForHello(timeoutMs: number): Promise { - if (this.hello) { - return { - heartbeatCount: this.heartbeatCount, - hello: this.hello, - lastHeartbeatAt: this.lastHeartbeat?.at ?? null, - }; - } + #removeWaiter(waiters: GenerationWaiter[], waiter: GenerationWaiter): void { + const index = waiters.indexOf(waiter); - if (isTruthy(this.errorMessage)) { - throw new Error(this.errorMessage); + if (index !== -1) { + waiters.splice(index, 1); } - - if (this.close) { - throw new Error(`Driver instance ${this.requireDriverInstanceId()} closed before hello.`); - } - - const deferred = createDeferred(); - this.helloWaiters.push(deferred); - return withTimeout( - deferred.promise, - timeoutMs, - `Driver instance ${this.requireDriverInstanceId()} hello`, - ); } - async waitForReady(timeoutMs: number): Promise { - if (isTruthy(this.errorMessage)) { - throw new Error(this.errorMessage); + assertGeneration(generation: number): void { + if (this.requireDriverGeneration() !== generation) { + throw new Error("Driver generation is no longer current."); } + } - if (this.close) { - throw new Error(`Driver instance ${this.requireDriverInstanceId()} closed before ready.`); + readyResult(): DriverInstanceReadyResult { + if (this.ready === null) { + throw new Error(`Driver instance ${this.requireDriverInstanceId()} is not ready yet.`); } - if (this.ready) { - return { - heartbeatCount: this.heartbeatCount, - lastHeartbeatAt: this.lastHeartbeat?.at ?? null, - ready: this.ready, - }; - } - - const deferred = createDeferred(); - this.readyWaiters.push(deferred); - return withTimeout( - deferred.promise, - timeoutMs, - `Driver instance ${this.requireDriverInstanceId()} ready`, - ); + return { + heartbeatCount: this.heartbeatCount, + lastHeartbeatAt: this.lastHeartbeat?.at ?? null, + ready: this.ready, + }; } closeResult(): DriverInstanceWaitForCloseResult { diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/sandbox-binding.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/sandbox-binding.ts deleted file mode 100644 index f8dafe04..00000000 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/sandbox-binding.ts +++ /dev/null @@ -1,6 +0,0 @@ -import { requireCloudflareSandboxBinding } from "../../../../platform/cloudflare/sandbox-binding"; -import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; - -export function requireSandboxBinding(env: ApiBindings) { - return requireCloudflareSandboxBinding(env); -} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/session-link.repository.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/session-link.repository.ts index e3437ff4..34d8f922 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/session-link.repository.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/session-link.repository.ts @@ -21,7 +21,7 @@ import type { SessionRunId, } from "@mosoo/id"; import { parsePlatformId } from "@mosoo/id"; -import { and, eq, inArray, sql } from "drizzle-orm"; +import { and, desc, eq, inArray, sql } from "drizzle-orm"; import { getAppDatabase } from "../../../../platform/db/drizzle"; import { ACTIVE_SESSION_RUN_STATUSES } from "../../domain/session-run-lifecycle.machine"; @@ -35,6 +35,7 @@ interface RuntimeSessionLinkRow { caller_account_id: AccountId | null; creator_account_id: PlatformId | null; origin_json: string | null; + runtime_id: string | null; sandbox_id: SandboxId | null; sandbox_kind: AgentKind | null; sandbox_subject_kind: SandboxSubjectKind | null; @@ -46,6 +47,7 @@ interface RuntimeSessionLinkRow { } export interface GetRuntimeSessionLinkOptions { + latestTerminalRun?: boolean; sessionRunId?: SessionRunId; } @@ -82,14 +84,19 @@ export async function getRuntimeSessionLink( ): Promise { const linkedSessionId = sql`coalesce(${sessionRunsTable.sessionId}, ${sandboxSessionsTable.sessionId})`; const linkedSessionRun = - options.sessionRunId === undefined + options.sessionRunId !== undefined ? and( eq(sessionRunsTable.driverInstanceId, driverInstancesTable.id), - inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + eq(sessionRunsTable.id, options.sessionRunId), ) : and( eq(sessionRunsTable.driverInstanceId, driverInstancesTable.id), - eq(sessionRunsTable.id, options.sessionRunId), + inArray( + sessionRunsTable.status, + options.latestTerminalRun === true + ? ["cancelled", "completed", "expired", "failed"] + : ACTIVE_SESSION_RUN_STATUSES, + ), ); const row = (await getAppDatabase(database) @@ -100,6 +107,9 @@ export async function getRuntimeSessionLink( caller_account_id: sessionRunsTable.createdByAccountId, creator_account_id: sessionsTable.creatorAccountId, origin_json: sandboxSessionsTable.originJson, + runtime_id: sql< + string | null + >`coalesce(${sessionRunsTable.runtimeId}, ${sessionsTable.runtimeId})`, sandbox_id: driverInstancesTable.sandboxId, sandbox_kind: sandboxesTable.kind, sandbox_subject_kind: sandboxesTable.subjectKind, @@ -119,6 +129,11 @@ export async function getRuntimeSessionLink( .leftJoin(agentsTable, eq(agentsTable.id, sessionsTable.agentId)) .leftJoin(sandboxesTable, eq(sandboxesTable.id, driverInstancesTable.sandboxId)) .where(eq(driverInstancesTable.id, driverInstanceId)) + .orderBy( + desc(sessionRunsTable.updatedAt), + desc(sessionRunsTable.createdAt), + desc(sessionRunsTable.id), + ) .limit(1) .get()) ?? null; const principals = resolveRuntimeSessionPrincipalIds(row ?? null); @@ -132,6 +147,7 @@ export async function getRuntimeSessionLink( sandboxId: row?.sandbox_id ?? null, sandboxKind: row?.sandbox_kind ?? null, sandboxSubjectKind: row?.sandbox_subject_kind ?? null, + runtimeId: row?.runtime_id ?? null, sessionId: row?.session_id ?? null, sessionRunId: row?.session_run_id ?? null, sessionRunStatus: row?.session_run_status ?? null, @@ -139,3 +155,31 @@ export async function getRuntimeSessionLink( traceId: row?.trace_id ?? null, }; } + +export async function assertActiveRuntimeSessionRun( + database: D1Database, + input: { + driverInstanceId: DriverInstanceId; + sessionId: SessionId; + sessionRunId: SessionRunId; + }, +): Promise { + const row = + (await getAppDatabase(database) + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, input.sessionRunId), + eq(sessionRunsTable.sessionId, input.sessionId), + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ) + .limit(1) + .get()) ?? null; + + if (row === null) { + throw new Error("Fresh runtime driver event requires its exact active Session Run."); + } +} diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/session-viewer-event-delivery-buffer.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/session-viewer-event-delivery-buffer.ts index 3ed2826e..de31b9a8 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/session-viewer-event-delivery-buffer.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/session-viewer-event-delivery-buffer.ts @@ -5,21 +5,28 @@ import { isAgUiSessionRunStartedEvent, isAgUiSessionRunTerminalEvent, } from "@mosoo/ag-ui-session"; -import { discardPromiseResult, ignorePromiseRejection } from "@mosoo/effects"; import type { DriverInstanceId, SessionId } from "@mosoo/id"; import { createErrorLogContext, logError } from "../../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; import type { SessionDeliveryEvent } from "../../../sessions/application/session-live-state.service"; -import { publishSessionViewerEvents } from "../../../sessions/application/session-viewer-events.service"; +import { + publishSessionViewerEventBatches, + syncSessionViewerState, +} from "../../../sessions/infrastructure/session/client"; const SESSION_VIEWER_EVENT_DELIVERY_FLUSH_MS = 150; const SESSION_VIEWER_EVENT_DELIVERY_MAX_DELTA_BYTES = 4 * 1024; const SESSION_VIEWER_EVENT_DELIVERY_MAX_EVENTS = 64; +const SESSION_VIEWER_EVENT_DELIVERY_MAX_SERIALIZED_BYTES = 1_536 * 1024; +const sessionViewerEventEncoder = new TextEncoder(); interface BufferedSessionViewerEvents { deltaBytes: number; events: SessionDeliveryEvent[]; + previousRuntimeEventSeqCursor: number | null; + requiresStateSync: boolean; + runtimeEventSeqCursor: number | null; sessionId: SessionId | null; } @@ -42,18 +49,20 @@ function hasDeltaEvent(events: SessionDeliveryEvent[]): boolean { return events.some((event) => getAgUiSessionEventDeltaLength(event) > 0); } -function estimateDeltaBytes(events: SessionDeliveryEvent[]): number { - return events.reduce((bytes, event) => bytes + getAgUiSessionEventDeltaLength(event), 0); +function measureSerializedBytes(events: SessionDeliveryEvent[]): number { + return sessionViewerEventEncoder.encode(JSON.stringify(events)).byteLength; } export class SessionViewerEventDeliveryBuffer { #buffer: BufferedSessionViewerEvents | null = null; + #delivery: Promise | null = null; #pendingFirstDelta = false; + readonly #pendingBatches: BufferedSessionViewerEvents[] = []; readonly #ctx: DurableObjectState; readonly #env: ApiBindings; - #gate: Promise = Promise.resolve(); readonly #getDriverInstanceId: () => DriverInstanceId | null; #timer: ReturnType | null = null; + #timerDone: (() => void) | null = null; readonly #withRuntimeLogContext: (fn: () => T) => T; constructor(options: SessionViewerEventDeliveryBufferOptions) { @@ -63,25 +72,69 @@ export class SessionViewerEventDeliveryBuffer { this.#withRuntimeLogContext = options.withRuntimeLogContext; } - enqueue(sessionId: SessionId | null, events: SessionDeliveryEvent[]): void { + enqueue( + sessionId: SessionId | null, + events: SessionDeliveryEvent[], + runtimeEventSeqCursor: number | null = null, + previousRuntimeEventSeqCursor: number | null = null, + ): void { const compactedEvents = compactAgUiSessionEvents(events); if (compactedEvents.length === 0) { return; } - const buffered = this.#buffer; - const incomingDeltaBytes = estimateDeltaBytes(compactedEvents); - const nextEvents = buffered - ? appendCompactedAgUiSessionEvents(buffered.events, compactedEvents) - : compactedEvents; - const nextDeltaBytes = (buffered?.deltaBytes ?? 0) + incomingDeltaBytes; + if (runtimeEventSeqCursor !== null) { + this.#queueBuffer(); + this.#pendingBatches.push({ + deltaBytes: compactedEvents.reduce( + (total, event) => total + getAgUiSessionEventDeltaLength(event), + 0, + ), + events: + measureSerializedBytes(compactedEvents) <= + SESSION_VIEWER_EVENT_DELIVERY_MAX_SERIALIZED_BYTES + ? compactedEvents + : [], + previousRuntimeEventSeqCursor, + requiresStateSync: + previousRuntimeEventSeqCursor === null || + measureSerializedBytes(compactedEvents) > + SESSION_VIEWER_EVENT_DELIVERY_MAX_SERIALIZED_BYTES, + runtimeEventSeqCursor, + sessionId, + }); - this.#buffer = { - deltaBytes: nextDeltaBytes, - events: nextEvents, - sessionId: buffered?.sessionId ?? sessionId, - }; + if (this.#delivery !== null) { + this.#replacePendingWithStateSync(sessionId, runtimeEventSeqCursor); + } + + if (hasRunStartedEvent(compactedEvents)) { + this.#pendingFirstDelta = true; + } + const isFirstDeltaOfRun = this.#pendingFirstDelta && hasDeltaEvent(compactedEvents); + if (isFirstDeltaOfRun) { + this.#pendingFirstDelta = false; + } + if (isFirstDeltaOfRun || hasTerminalEvent(compactedEvents)) { + void this.#startFlush(); + } else { + this.#scheduleFlush(); + } + return; + } + + if ( + this.#buffer !== null && + (this.#buffer.sessionId !== sessionId || + this.#buffer.runtimeEventSeqCursor !== runtimeEventSeqCursor) + ) { + this.#queueBuffer(); + } + + for (const event of compactedEvents) { + this.#appendEvent(sessionId, event, runtimeEventSeqCursor); + } // O2: cut first-token latency. Arm on RUN_STARTED, then flush the first // delta of the run immediately instead of waiting out the 150ms timer. @@ -95,71 +148,78 @@ export class SessionViewerEventDeliveryBuffer { this.#pendingFirstDelta = false; } - if ( - nextEvents.length >= SESSION_VIEWER_EVENT_DELIVERY_MAX_EVENTS || - nextDeltaBytes >= SESSION_VIEWER_EVENT_DELIVERY_MAX_DELTA_BYTES || - isFirstDeltaOfRun || - hasTerminalEvent(compactedEvents) - ) { - this.#startFlush(); + if (isFirstDeltaOfRun || hasTerminalEvent(compactedEvents)) { + void this.#startFlush(); return; } - this.#scheduleFlush(); - } - - async flush(): Promise { - if (this.#timer !== null) { - clearTimeout(this.#timer); - this.#timer = null; + if (this.#delivery !== null) { + this.#replacePendingWithStateSync(sessionId, runtimeEventSeqCursor); } - const buffered = this.#buffer; - - if (!buffered) { - await this.#gate; - return; - } - - this.#buffer = null; - - const task = this.#deliverAfterGate(buffered); - - this.#gate = SessionViewerEventDeliveryBuffer.#discardDeliveryResult(task); - - try { - await task; - } catch (error) { - const currentBuffer = this.#getBuffer(); - this.#buffer = { - deltaBytes: buffered.deltaBytes + (currentBuffer?.deltaBytes ?? 0), - events: currentBuffer - ? appendCompactedAgUiSessionEvents(buffered.events, currentBuffer.events) - : buffered.events, - sessionId: buffered.sessionId ?? currentBuffer?.sessionId ?? null, - }; - throw error; + if (this.#buffer) { + this.#scheduleFlush(); } } - async #deliverAfterGate(buffered: BufferedSessionViewerEvents): Promise { - try { - await this.#gate; - } catch (error) { - ignorePromiseRejection(error); - } - - await publishSessionViewerEvents(this.#env, buffered.sessionId, buffered.events); + async flush(): Promise { + this.#clearTimer(); + this.#queueBuffer(); + await this.#getOrStartDelivery(); } - static async #discardDeliveryResult(task: Promise): Promise { + async #deliverPendingBatches(): Promise { try { - await task; - } catch (error) { - ignorePromiseRejection(error); + while (this.#pendingBatches.length > 0) { + const first = this.#pendingBatches.shift(); + + if (!first) { + return; + } + if (first.requiresStateSync) { + try { + await syncSessionViewerState(this.#env, first.sessionId); + } catch (error) { + this.#pendingBatches.unshift(first); + throw error; + } + continue; + } + const batches = [first]; + let serializedBytes = measureSerializedBytes(first.events); + for (;;) { + const next = this.#pendingBatches[0]; + if ( + next === undefined || + next.requiresStateSync || + next.sessionId !== first.sessionId || + serializedBytes + measureSerializedBytes(next.events) > + SESSION_VIEWER_EVENT_DELIVERY_MAX_SERIALIZED_BYTES + ) { + break; + } + batches.push(this.#pendingBatches.shift()!); + serializedBytes += measureSerializedBytes(next.events); + } + + try { + await publishSessionViewerEventBatches( + this.#env, + first.sessionId, + batches.map((batch) => ({ + events: batch.events, + previousRuntimeEventSeqCursor: batch.previousRuntimeEventSeqCursor, + runtimeEventSeqCursor: batch.runtimeEventSeqCursor, + })), + ); + } catch (error) { + this.#pendingBatches.unshift(...batches); + throw error; + } + } + } finally { + this.#delivery = null; } - - discardPromiseResult(); } async flushSafely(): Promise { @@ -167,55 +227,255 @@ export class SessionViewerEventDeliveryBuffer { await this.flush(); } catch (error) { this.#reportDeliveryError(error); - - if (this.#buffer && this.#timer === null) { - this.#scheduleFlush(); - } + this.resetAfterFlush(); } } + requestStateSync(sessionId: SessionId | null): void { + this.#replacePendingWithStateSync(sessionId, null); + void this.#startFlush(); + } + resetAfterFlush(): void { this.#buffer = null; this.#pendingFirstDelta = false; - this.#gate = Promise.resolve(); + this.#pendingBatches.length = 0; + this.#clearTimer(); + } + + #appendEvent( + sessionId: SessionId | null, + event: SessionDeliveryEvent, + runtimeEventSeqCursor: number | null, + ): void { + let buffered = this.#buffer; + let events = buffered ? appendCompactedAgUiSessionEvents(buffered.events, [event]) : [event]; + let serializedBytes = measureSerializedBytes(events); + + if (buffered && serializedBytes > SESSION_VIEWER_EVENT_DELIVERY_MAX_SERIALIZED_BYTES) { + void this.#startFlush(); + buffered = null; + events = [event]; + serializedBytes = measureSerializedBytes(events); + } + + this.#buffer = { + deltaBytes: (buffered?.deltaBytes ?? 0) + getAgUiSessionEventDeltaLength(event), + events, + previousRuntimeEventSeqCursor: null, + requiresStateSync: false, + runtimeEventSeqCursor: buffered?.runtimeEventSeqCursor ?? runtimeEventSeqCursor, + sessionId: buffered?.sessionId ?? sessionId, + }; + + if ( + events.length >= SESSION_VIEWER_EVENT_DELIVERY_MAX_EVENTS || + this.#buffer.deltaBytes >= SESSION_VIEWER_EVENT_DELIVERY_MAX_DELTA_BYTES || + serializedBytes >= SESSION_VIEWER_EVENT_DELIVERY_MAX_SERIALIZED_BYTES + ) { + void this.#startFlush(); + } + } + + #clearTimer(): void { if (this.#timer !== null) { clearTimeout(this.#timer); this.#timer = null; } + this.#timerDone?.(); + this.#timerDone = null; } - #getBuffer(): BufferedSessionViewerEvents | null { - return this.#buffer; + #getOrStartDelivery(): Promise { + if (this.#delivery) { + this.#compactPendingBatches(); + return this.#delivery; + } + + if (this.#pendingBatches.length === 0) { + return Promise.resolve(); + } + + this.#compactPendingBatches(); + const delivery = this.#deliverPendingBatches(); + this.#delivery = delivery; + return delivery; } - #scheduleFlush(): void { - if (this.#timer !== null) { + #compactPendingBatches(): void { + if (this.#pendingBatches.length < 2) { return; } - this.#timer = setTimeout(() => { - this.#timer = null; - this.#startFlush(); - }, SESSION_VIEWER_EVENT_DELIVERY_FLUSH_MS); + const queued = this.#pendingBatches.splice(0); + let events: SessionDeliveryEvent[] = []; + let sessionId = queued[0]?.sessionId ?? null; + let previousRuntimeEventSeqCursor = queued[0]?.previousRuntimeEventSeqCursor ?? null; + let requiresStateSync = queued[0]?.requiresStateSync ?? false; + let runtimeEventSeqCursor = queued[0]?.runtimeEventSeqCursor ?? null; + + for (const batch of queued) { + if ( + batch.sessionId !== sessionId || + batch.runtimeEventSeqCursor !== runtimeEventSeqCursor || + batch.previousRuntimeEventSeqCursor !== previousRuntimeEventSeqCursor || + batch.requiresStateSync !== requiresStateSync + ) { + this.#queueCompactedEvents( + sessionId, + events, + runtimeEventSeqCursor, + previousRuntimeEventSeqCursor, + requiresStateSync, + ); + events = []; + sessionId = batch.sessionId; + previousRuntimeEventSeqCursor = batch.previousRuntimeEventSeqCursor; + requiresStateSync = batch.requiresStateSync; + runtimeEventSeqCursor = batch.runtimeEventSeqCursor; + } + + events = appendCompactedAgUiSessionEvents(events, batch.events); + } + + this.#queueCompactedEvents( + sessionId, + events, + runtimeEventSeqCursor, + previousRuntimeEventSeqCursor, + requiresStateSync, + ); + } + + #queueCompactedEvents( + sessionId: SessionId | null, + events: SessionDeliveryEvent[], + runtimeEventSeqCursor: number | null, + previousRuntimeEventSeqCursor: number | null, + requiresStateSync: boolean, + ): void { + if (requiresStateSync) { + this.#pendingBatches.push({ + deltaBytes: 0, + events: [], + previousRuntimeEventSeqCursor: null, + requiresStateSync: true, + runtimeEventSeqCursor, + sessionId, + }); + return; + } + + if (runtimeEventSeqCursor !== null) { + const compactedEvents = compactAgUiSessionEvents(events); + const mustSync = + measureSerializedBytes(compactedEvents) > + SESSION_VIEWER_EVENT_DELIVERY_MAX_SERIALIZED_BYTES; + this.#pendingBatches.push({ + deltaBytes: compactedEvents.reduce( + (total, event) => total + getAgUiSessionEventDeltaLength(event), + 0, + ), + events: mustSync ? [] : compactedEvents, + previousRuntimeEventSeqCursor, + requiresStateSync: mustSync, + runtimeEventSeqCursor, + sessionId, + }); + return; + } + for (const event of events) { + const previous = this.#pendingBatches.at(-1); + const nextEvents = previous + ? appendCompactedAgUiSessionEvents(previous.events, [event]) + : [event]; + const nextDeltaBytes = (previous?.deltaBytes ?? 0) + getAgUiSessionEventDeltaLength(event); + + if ( + previous && + previous.sessionId === sessionId && + !previous.requiresStateSync && + previous.runtimeEventSeqCursor === runtimeEventSeqCursor && + nextEvents.length <= SESSION_VIEWER_EVENT_DELIVERY_MAX_EVENTS && + nextDeltaBytes <= SESSION_VIEWER_EVENT_DELIVERY_MAX_DELTA_BYTES && + measureSerializedBytes(nextEvents) <= SESSION_VIEWER_EVENT_DELIVERY_MAX_SERIALIZED_BYTES + ) { + previous.deltaBytes = nextDeltaBytes; + previous.events = nextEvents; + continue; + } + + this.#pendingBatches.push({ + deltaBytes: getAgUiSessionEventDeltaLength(event), + events: [event], + previousRuntimeEventSeqCursor, + requiresStateSync: false, + runtimeEventSeqCursor, + sessionId, + }); + } + } + + #queueBuffer(): void { + if (!this.#buffer) { + return; + } + + this.#pendingBatches.push(this.#buffer); + this.#buffer = null; } - #startFlush(): void { - const task = this.#flushAndReportDeliveryErrors(); + #replacePendingWithStateSync( + sessionId: SessionId | null, + runtimeEventSeqCursor: number | null, + ): void { + this.#buffer = null; + this.#pendingBatches.length = 0; + this.#clearTimer(); + this.#pendingBatches.push({ + deltaBytes: 0, + events: [], + previousRuntimeEventSeqCursor: null, + requiresStateSync: true, + runtimeEventSeqCursor, + sessionId, + }); + } + + #scheduleFlush(): void { + if (this.#timer !== null) { + return; + } + const task = new Promise((resolve) => { + this.#timerDone = resolve; + this.#timer = setTimeout(() => { + this.#timer = null; + this.#timerDone = null; + void this.#startFlush().finally(resolve); + }, SESSION_VIEWER_EVENT_DELIVERY_FLUSH_MS); + }); this.#ctx.waitUntil(task); } - async #flushAndReportDeliveryErrors(): Promise { - try { - await this.flush(); - } catch (error) { - this.#reportDeliveryError(error); + #startFlush(): Promise { + this.#clearTimer(); + this.#queueBuffer(); - if (this.#buffer && this.#timer === null) { - this.#scheduleFlush(); - } + if (this.#delivery) { + this.#compactPendingBatches(); + return this.#delivery; + } + + if (this.#pendingBatches.length === 0) { + return Promise.resolve(); } + + const task = this.flushSafely(); + + this.#ctx.waitUntil(task); + return task; } #reportDeliveryError(error: unknown): void { diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/sockets.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/sockets.ts index 46848dbc..8c5d8de4 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/sockets.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/sockets.ts @@ -1,68 +1,102 @@ import { sleepPromise } from "@mosoo/effects"; +import type { DriverInstanceConnectionEpoch } from "./state"; + const DRIVER_SOCKET_TAG = "driver"; +function parseDriverSocketEpoch(socket: WebSocket): DriverInstanceConnectionEpoch | null { + const attachment: unknown = socket.deserializeAttachment(); + + if ( + typeof attachment !== "object" || + attachment === null || + !("connectionId" in attachment) || + !("generation" in attachment) || + typeof attachment.connectionId !== "string" || + attachment.connectionId.length === 0 || + !Number.isSafeInteger(attachment.generation) || + (attachment.generation as number) < 0 + ) { + return null; + } + + return { + connectionId: attachment.connectionId, + generation: attachment.generation as number, + }; +} + +function epochsMatch( + left: DriverInstanceConnectionEpoch, + right: DriverInstanceConnectionEpoch, +): boolean { + return left.connectionId === right.connectionId && left.generation === right.generation; +} + export class DriverInstanceSocketRegistry { readonly #ctx: DurableObjectState; - #activeDriverSocket: WebSocket | null = null; constructor(ctx: DurableObjectState) { this.#ctx = ctx; } - acceptDriverSocket(socket: WebSocket): void { - // Hibernation accept: the socket survives Durable Object eviction, and - // driver messages re-instantiate this object. Command delivery does not - // depend on in-memory waiters — the driver polls nextCommand, which - // claims from D1 — so waking into a fresh instance is safe. + acceptDriverSocket(socket: WebSocket, epoch: DriverInstanceConnectionEpoch): void { this.#ctx.acceptWebSocket(socket, [DRIVER_SOCKET_TAG]); - this.#activeDriverSocket = socket; + socket.serializeAttachment(epoch); } - getDriverSocket(): WebSocket | null { - if (this.#activeDriverSocket && this.#activeDriverSocket.readyState !== WebSocket.CLOSED) { - return this.#activeDriverSocket; - } - - const [socket] = this.#ctx.getWebSockets(DRIVER_SOCKET_TAG); - return socket ?? null; + getSocketEpoch(socket: WebSocket): DriverInstanceConnectionEpoch | null { + return parseDriverSocketEpoch(socket); } - isActiveDriverSocket(socket: WebSocket): boolean { - return this.getDriverSocket() === socket; + socketMatchesEpoch(socket: WebSocket, epoch: DriverInstanceConnectionEpoch): boolean { + const socketEpoch = parseDriverSocketEpoch(socket); + return socketEpoch !== null && epochsMatch(socketEpoch, epoch); } - /** - * True when a different, still-open driver socket has superseded this one. - * Close/error events from superseded sockets must not finalize the state - * that now belongs to the successor connection. - */ - isSupersededDriverSocket(socket: WebSocket): boolean { - const current = this.getDriverSocket(); - return current !== null && current !== socket; + isCurrentDriverSocket( + socket: WebSocket, + capturedEpoch: DriverInstanceConnectionEpoch, + currentEpoch: DriverInstanceConnectionEpoch | null, + ): boolean { + return ( + currentEpoch !== null && + epochsMatch(capturedEpoch, currentEpoch) && + this.socketMatchesEpoch(socket, capturedEpoch) + ); } - releaseDriverSocket(socket: WebSocket): void { - if (this.#activeDriverSocket === socket) { - this.#activeDriverSocket = null; + getDriverSocket(epoch: DriverInstanceConnectionEpoch | null): WebSocket | null { + if (epoch === null) { + return null; } + + return ( + this.#ctx + .getWebSockets(DRIVER_SOCKET_TAG) + .find( + (socket) => + socket.readyState === WebSocket.OPEN && this.socketMatchesEpoch(socket, epoch), + ) ?? null + ); } replaceDriverSockets(): void { - if (this.#activeDriverSocket && this.#activeDriverSocket.readyState !== WebSocket.CLOSED) { - this.#activeDriverSocket.close(1012, "runtime.socket.replaced"); - this.#activeDriverSocket = null; - } - - for (const existingSocket of this.#ctx.getWebSockets(DRIVER_SOCKET_TAG)) { - existingSocket.close(1012, "runtime.socket.replaced"); + for (const socket of this.#ctx.getWebSockets(DRIVER_SOCKET_TAG)) { + if (socket.readyState === WebSocket.OPEN) { + socket.close(1012, "runtime.socket.replaced"); + } } } - scheduleDriverSocketClose(code: number, reason: string): void { - const socket = this.getDriverSocket(); + scheduleDriverSocketClose( + epoch: DriverInstanceConnectionEpoch, + code: number, + reason: string, + ): void { + const socket = this.getDriverSocket(epoch); - if (!socket || socket.readyState !== WebSocket.OPEN) { + if (socket === null) { return; } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/state.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/state.ts index 866a3721..8ece68fa 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/state.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/state.ts @@ -1,11 +1,11 @@ -import type { - DriverHeartbeatInput, - DriverHelloInput, - DriverReadyInput, -} from "@mosoo/agent-driver/orpc"; +import type { DriverHelloInput, DriverReadyInput } from "@mosoo/agent-driver/orpc"; import type { DriverDebugResumeSnapshot } from "./debug-resume-snapshot"; -import type { Deferred } from "./driver-instance-support"; + +export interface DriverInstanceConnectionEpoch { + readonly connectionId: string; + readonly generation: number; +} export interface DriverInstanceCloseSnapshot { at: string; @@ -13,23 +13,6 @@ export interface DriverInstanceCloseSnapshot { reason: string; } -export interface DriverInstanceHeartbeatResult { - heartbeat: DriverHeartbeatInput; - heartbeatCount: number; - lastHeartbeatAt: string; -} - -export interface HeartbeatWaiter { - afterCount: number; - deferred: Deferred; -} - -export interface DriverInstanceHelloResult { - heartbeatCount: number; - hello: DriverHelloInput; - lastHeartbeatAt: string | null; -} - export interface DriverInstanceReadyResult { heartbeatCount: number; lastHeartbeatAt: string | null; diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-driver-events.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-driver-events.ts index a5a336b9..e0fe8f48 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-driver-events.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-driver-events.ts @@ -1,218 +1,233 @@ import type { DriverFailureInput } from "@mosoo/agent-driver/orpc"; +import type { RunError, SessionRunSummary } from "@mosoo/contracts/session-run"; import { createPlatformId } from "@mosoo/id"; import type { DriverInstanceId, RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; -import { createRuntimeEvent } from "@mosoo/runtime-events"; -import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; -import { appendSessionRuntimeEvents } from "../../../sessions/application/session-event-write.service"; -import { appRuntimeEventToSessionDeliveryEvents } from "../../../sessions/application/session-live-state.service"; -import { recordCanonicalSessionRunFailure } from "../../application/session-runs/session-run-terminal-failure.service"; -import { isTerminalSessionRunStatus } from "../../domain/session-run-status"; -import { setSessionRunStatus } from "../session-runs/session-run-store.repository"; -import type { SessionRunTransitionOutcome } from "../session-runs/session-run-store.repository"; +import { currentTimestampMs, toIsoString } from "../../../../time"; +import { createSessionRunTerminalSourceId } from "../../domain/session-run-terminal-event-id"; +import { getSessionRunSummary } from "../session-runs/session-run-store.repository"; +import { + adoptTerminalRunProjection, + commitTerminalRunProjection, +} from "./completed-run-commit.repository"; +import type { DriverTerminalRunStatus } from "./completed-run-commit.repository"; +import { createCanonicalDriverRunFailedEvent } from "./driver-event-canonicalization"; import type { RuntimeSessionLink } from "./event-types"; +import { getDriverInstanceLifecycleIdentity } from "./lifecycle"; import { getRuntimeSessionLink } from "./session-link.repository"; import { releaseTerminalDriverInstanceSessionRun } from "./terminal-run-release"; -function assertTerminalDriverSessionRunTransition(outcome: SessionRunTransitionOutcome): void { - switch (outcome.kind) { - case "applied": - case "duplicate": { - return; - } - case "stale": { - if (outcome.reason === "terminal_run") { - return; - } - throw new Error("Terminal driver event lost a concurrent run transition."); - } - case "repair_needed": { - throw new Error("Terminal driver event left the session lifecycle projection stale."); - } - case "rejected": { - throw new Error(`Terminal driver event run transition was rejected: ${outcome.reason}.`); - } +function terminalTargetFromRun( + status: SessionRunSummary["status"], +): DriverTerminalRunStatus | null { + switch (status) { + case "cancelled": + return "cancelled"; + case "completed": + return "completed"; + case "failed": + return "failed"; + case "expired": + return "cancelled"; + case "booting": + case "queued": + case "running": + case "waiting_input": + return null; } } -function isStaleTerminalRunTransition(outcome: SessionRunTransitionOutcome): boolean { - return outcome.kind === "stale" && outcome.reason === "terminal_run"; -} - -function isStaleTerminalRunStatus( - outcome: SessionRunTransitionOutcome, - status: "completed" | "failed", -): boolean { - return ( - outcome.kind === "stale" && - outcome.reason === "terminal_run" && - outcome.currentStatus === status - ); -} - -function createTerminalDriverEventId(input: { - readonly driverInstanceId: DriverInstanceId; - readonly kind: "run.completed" | "run.failed"; - readonly sessionRunId: SessionRunId; -}): string { - return `driver-terminal:${input.driverInstanceId}:${input.sessionRunId}:${input.kind}`; -} - -export async function recordDriverInstanceCompletion( +async function commitOrAdoptTerminalRun( bindings: ApiBindings, input: { - driverReady: boolean; + driverConnectionId?: string; + driverGeneration?: number; driverInstanceId: DriverInstanceId; + error: RunError | null; + requestedStatus: "completed" | "failed"; + runtimeId: string; + sessionId: SessionId; + sessionRunId: SessionRunId; }, ): Promise { - void input.driverReady; - const database = bindings.DB; - const link = await getRuntimeSessionLink(database, input.driverInstanceId); + const current = await getSessionRunSummary(bindings.DB, input.sessionRunId); + if (current === null) { + throw new Error("Terminal Driver Session Run was not found."); + } - if ( - hasLinkedSessionRun(link) && - link.sessionRunStatus !== null && - (!isTerminalSessionRunStatus(link.sessionRunStatus) || link.sessionRunStatus === "completed") - ) { - await synthesizeDriverRunFinished(database, { - bindings, - driverInstanceId: input.driverInstanceId, - link, - }); + const currentTarget = terminalTargetFromRun(current.status); + const targetStatus = currentTarget ?? input.requestedStatus; + const expectedDriverObservation = + input.driverConnectionId === undefined || input.driverGeneration === undefined + ? undefined + : { + connectionId: input.driverConnectionId, + driverInstanceId: input.driverInstanceId, + generation: input.driverGeneration, + }; + const adopted = await adoptTerminalRunProjection(bindings.DB, { + ...(expectedDriverObservation === undefined ? {} : { expectedDriverObservation }), + expectedTargetStatus: targetStatus, + runId: input.sessionRunId, + sessionId: input.sessionId, + }); + if (adopted.kind !== "missing") { + if (adopted.kind === "stale") { + throw new Error( + `Terminal Driver Session Run lost a concurrent ${adopted.currentStatus} race.`, + ); + } + return; + } + + if (targetStatus !== "failed") { + throw new Error( + `${targetStatus === "completed" ? "Completed" : "Cancelled"} Session Run is missing its canonical terminal event.`, + ); } - await releaseLinkedRunLease(bindings, { + const runError = currentTarget === "failed" ? current.error : input.error; + if (runError === null) { + throw new Error("Failed Session Run is missing its authoritative durable error."); + } + const timestampMs = currentTimestampMs(); + const timestamp = toIsoString(timestampMs); + const sourceEventId = createSessionRunTerminalSourceId(input.sessionRunId, "run.failed"); + const event = createCanonicalDriverRunFailedEvent({ driverInstanceId: input.driverInstanceId, - sessionRunId: link.sessionRunId, + error: runError, + id: createPlatformId(), + occurredAt: timestamp, + runId: input.sessionRunId, + runtimeId: input.runtimeId, + sessionId: input.sessionId, + traceId: current.traceId, }); + + const outcome = await commitTerminalRunProjection(bindings.DB, { + assistantMessage: null, + error: runError, + ...(expectedDriverObservation === undefined ? {} : { expectedDriverObservation }), + runId: input.sessionRunId, + sessionId: input.sessionId, + source: "driver", + targetStatus: "failed", + terminalEvent: { event, occurredAt: timestampMs, sourceEventId }, + timestampMs, + }); + if (outcome.kind === "stale") { + throw new Error(`Terminal Driver Session Run lost a concurrent ${outcome.currentStatus} race.`); + } } -export async function recordDriverInstanceFailure( +function hasLinkedSessionRun(link: RuntimeSessionLink): link is RuntimeSessionLink & { + runtimeId: string; + sessionId: SessionId; + sessionRunId: SessionRunId; +} { + return link.runtimeId !== null && link.sessionId !== null && link.sessionRunId !== null; +} + +async function finishTerminalDriverRun( bindings: ApiBindings, input: { - error: DriverFailureInput["error"]; + driverConnectionId?: string; + driverGeneration?: number; driverInstanceId: DriverInstanceId; + error: RunError | null; link?: RuntimeSessionLink; + requestedStatus: "completed" | "failed"; + sessionRunId: SessionRunId; }, ): Promise { - const database = bindings.DB; - const link = input.link ?? (await getRuntimeSessionLink(database, input.driverInstanceId)); + if ((input.driverConnectionId === undefined) !== (input.driverGeneration === undefined)) { + throw new Error("Terminal Driver connection identity must be provided together."); + } - if (hasLinkedSessionRun(link)) { - const outcome = await recordCanonicalSessionRunFailure(bindings, { - error: input.error, - runId: link.sessionRunId, - sessionId: link.sessionId, - source: "driver", - }); - if (outcome.kind !== "failed") { - assertTerminalDriverSessionRunTransition(outcome.transition); - } - } else if (link.sessionRunId !== null) { - const outcome = await setSessionRunStatus(database, { - error: input.error, - runId: link.sessionRunId, - source: "driver", - status: "failed", - }); - assertTerminalDriverSessionRunTransition(outcome); + const link = + input.link ?? + (await getRuntimeSessionLink(bindings.DB, input.driverInstanceId, { + sessionRunId: input.sessionRunId, + })); + + if (link.sessionRunId !== input.sessionRunId) { + throw new Error("Terminal Driver Session Run identity does not match the request."); + } + if (!hasLinkedSessionRun(link)) { + throw new Error("Terminal Driver Session Run is missing its durable session identity."); + } + + const driverGeneration = + input.driverGeneration ?? + (await getDriverInstanceLifecycleIdentity(bindings, input.driverInstanceId))?.generation; + if (driverGeneration === undefined) { + throw new Error("Terminal Driver instance identity was not found."); } - await releaseLinkedRunLease(bindings, { + await commitOrAdoptTerminalRun(bindings, { + ...(input.driverConnectionId === undefined + ? {} + : { driverConnectionId: input.driverConnectionId }), + ...(input.driverGeneration === undefined ? {} : { driverGeneration: input.driverGeneration }), driverInstanceId: input.driverInstanceId, - sessionRunId: link.sessionRunId, + error: input.error, + requestedStatus: input.requestedStatus, + runtimeId: link.runtimeId, + sessionId: link.sessionId, + sessionRunId: input.sessionRunId, + }); + await releaseTerminalDriverInstanceSessionRun(bindings, { + ...(input.driverConnectionId === undefined + ? {} + : { expectedDriverConnectionId: input.driverConnectionId }), + driverGeneration, + driverInstanceId: input.driverInstanceId, + sessionRunId: input.sessionRunId, }); } -async function synthesizeDriverRunFinished( - database: D1Database, +export async function recordDriverInstanceCompletion( + bindings: ApiBindings, input: { - bindings: ApiBindings; + driverConnectionId?: string; + driverGeneration?: number; driverInstanceId: DriverInstanceId; - link: RuntimeSessionLink & { - sessionId: SessionId; - sessionRunId: SessionRunId; - }; + sessionRunId: SessionRunId; }, ): Promise { - const eventId = createTerminalDriverEventId({ - driverInstanceId: input.driverInstanceId, - kind: "run.completed", - sessionRunId: input.link.sessionRunId, - }); - const runCompletedEvent = createRuntimeEvent({ + await finishTerminalDriverRun(bindings, { + ...(input.driverConnectionId === undefined + ? {} + : { driverConnectionId: input.driverConnectionId }), + ...(input.driverGeneration === undefined ? {} : { driverGeneration: input.driverGeneration }), driverInstanceId: input.driverInstanceId, - id: createPlatformId(), - kind: "run.completed", - occurredAt: new Date().toISOString(), - payload: { - stopReason: "end_turn", - }, - runId: input.link.sessionRunId, - sessionId: input.link.sessionId, - sourceEventId: eventId, - }); - const [runFinishedEvent] = appRuntimeEventToSessionDeliveryEvents(runCompletedEvent); - - if (runFinishedEvent === undefined) { - throw new Error("Run completion event did not app to session delivery."); - } - - // The terminal RPC proves only that execution ended; it carries no final - // assistant item identity. Never guess from the session's last assistant - // message, which may be progress or belong to an earlier run. The canonical - // projection is written only by an ordered runtime run.completed event that - // names finalMessageId. - const outcome = await setSessionRunStatus(database, { - runId: input.link.sessionRunId, - source: "driver", - status: "completed", - }); - assertTerminalDriverSessionRunTransition(outcome); - if (isStaleTerminalRunTransition(outcome) && !isStaleTerminalRunStatus(outcome, "completed")) { - return; - } - await appendCanonicalTerminalDriverEvent({ - bindings: input.bindings, - event: runCompletedEvent, - }); -} - -async function appendCanonicalTerminalDriverEvent(input: { - bindings: ApiBindings; - event: RuntimeEventEnvelope; -}): Promise { - await appendSessionRuntimeEvents({ - bindings: input.bindings, - events: [input.event], - sessionId: input.event.sessionId, - sourceEventId: input.event.sourceEventId ?? input.event.id, + error: null, + requestedStatus: "completed", + sessionRunId: input.sessionRunId, }); } -function hasLinkedSessionRun(link: RuntimeSessionLink): link is RuntimeSessionLink & { - sessionId: SessionId; - sessionRunId: SessionRunId; -} { - return link.sessionId !== null && link.sessionRunId !== null; -} - -async function releaseLinkedRunLease( +export async function recordDriverInstanceFailure( bindings: ApiBindings, input: { - readonly driverInstanceId: DriverInstanceId; - readonly sessionRunId: SessionRunId | null; + driverConnectionId?: string; + driverGeneration?: number; + error: DriverFailureInput["error"]; + driverInstanceId: DriverInstanceId; + link?: RuntimeSessionLink; + sessionRunId: SessionRunId; }, -): Promise { - if (input.sessionRunId === null) { - return false; - } - - const outcome = await releaseTerminalDriverInstanceSessionRun(bindings, { +): Promise { + await finishTerminalDriverRun(bindings, { + ...(input.driverConnectionId === undefined + ? {} + : { driverConnectionId: input.driverConnectionId }), + ...(input.driverGeneration === undefined ? {} : { driverGeneration: input.driverGeneration }), driverInstanceId: input.driverInstanceId, + error: input.error, + ...(input.link === undefined ? {} : { link: input.link }), + requestedStatus: "failed", sessionRunId: input.sessionRunId, }); - - return outcome.released; } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-run-release.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-run-release.ts index 576754f1..353078be 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-run-release.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-run-release.ts @@ -1,40 +1,172 @@ -import type { RunError, SessionRunStatus, SessionRunSummary } from "@mosoo/contracts/session-run"; -import { sessionRunsTable } from "@mosoo/db"; -import type { DriverInstanceId, SessionId, SessionRunId } from "@mosoo/id"; -import { and, eq, inArray } from "drizzle-orm"; +import { createMcpExecuteFailedEventIdentity } from "@mosoo/agent-driver/events"; +import type { RunError, SessionRunSummary } from "@mosoo/contracts/session-run"; +import { + driverInstancesTable, + sessionEventsTable, + sessionRunsTable, + sessionsTable, +} from "@mosoo/db"; +import { createPlatformId, parsePlatformId } from "@mosoo/id"; +import type { + DriverInstanceId, + RuntimeEventId, + RuntimeOperationId, + SessionId, + SessionRunId, +} from "@mosoo/id"; +import { createRuntimeEvent, createRuntimeEventSemanticHash } from "@mosoo/runtime-events"; +import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; +import { and, eq, exists, inArray, isNull, ne, notExists, sql } from "drizzle-orm"; -import { logInfo, logWarn } from "../../../../platform/cloudflare/logger"; +import { logInfo } from "../../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; import { getAppDatabase } from "../../../../platform/db/drizzle"; +import { currentTimestampMs, toIsoString } from "../../../../time"; import { appendSessionRuntimeEvents } from "../../../sessions/application/session-event-write.service"; -import { createFailedSessionRunRuntimeEvent } from "../../application/session-runs/session-run-view-events.service"; +import { createSessionRuntimeEventProjection } from "../../../sessions/domain/session-runtime-event-projection"; import { getRuntimeKindPolicy } from "../../domain/runtime-kind-policy"; +import { ACTIVE_SESSION_RUN_STATUSES } from "../../domain/session-run-lifecycle.machine"; import { classifyReclaim, decideReclaimRecovery } from "../../domain/session-run-reclaim-recovery"; import { isTerminalSessionRunStatus } from "../../domain/session-run-status"; -import { createSessionRunTerminalFailureSourceId } from "../../domain/session-run-terminal-event-id"; import { recordRuntimeRunLeaseReleasedOutcome } from "../runtime-subject-lifecycle/runtime-run-lease-store"; import { createSandboxCheckpoints } from "../sandbox-backup.service"; -import { markExecutingExternalToolEffectsUnknownForDriver } from "../session-runs/external-tool-effect-store.repository"; -import { failAcceptedRuntimeCommandsForTerminalDriver } from "../session-runs/runtime-command-store.repository"; -import { setSessionRunStatus } from "../session-runs/session-run-store.repository"; -import type { SessionRunTransitionOutcome } from "../session-runs/session-run-store.repository"; +import { markClaimedExternalToolEffectsUnknownForDriver } from "../session-runs/external-tool-effect-store.repository"; +import { + listAcceptedInputStartCommandRepairsForTerminalDriver, + listAcceptedMcpCommandRepairsForTerminalDriver, + listTerminalDriversWithPendingRuntimeCommands, + repairAcceptedRuntimeCommandsForTerminalDriver, + updateRuntimeCommandRecord, +} from "../session-runs/runtime-command-store.repository"; +import type { + AcceptedInputStartCommandRepair, + AcceptedMcpCommandRepair, +} from "../session-runs/runtime-command-store.repository"; +import { isCattleTerminalCheckpointReadyForNextRun } from "../session-runs/session-run-admission.repository"; +import { getSessionRunSummary } from "../session-runs/session-run-store.repository"; +import { commitTerminalRunProjection } from "./completed-run-commit.repository"; +import { createCanonicalDriverRunFailedEvent } from "./driver-event-canonicalization"; import type { RuntimeSessionLink } from "./event-types"; +import { getDriverInstanceLifecycleIdentity } from "./lifecycle"; import { getRuntimeSessionLink } from "./session-link.repository"; -import { closeReleasedTerminalRuntimeLeaseIfNeeded } from "./terminal-runtime-lease"; - -interface LinkedSessionRunStatusRow { - readonly sessionRunId: SessionRunId | null; - readonly status: SessionRunStatus | null; -} +import { + closeTerminalRuntimeConversationIfNeeded, + recycleTerminalRuntimeLeaseIfNeeded, +} from "./terminal-runtime-lease"; export interface TerminalDriverInstanceSessionRunReleaseResult { readonly link: RuntimeSessionLink | null; readonly released: boolean; } +async function claimTerminalRunRelease( + database: D1Database, + input: { + readonly expectedDriverConnectionId?: string | null; + readonly driverGeneration: number; + readonly driverInstanceId: DriverInstanceId; + readonly sessionRunId: SessionRunId; + }, +): Promise<{ + operationId: RuntimeOperationId; + phase: "cleanup_required" | "release_committed"; +} | null> { + const operationId = parsePlatformId( + input.sessionRunId, + "terminal Run release operation ID", + ); + const db = getAppDatabase(database); + const exactTerminalRun = db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, input.sessionRunId), + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + inArray(sessionRunsTable.status, ["cancelled", "completed", "expired", "failed"]), + ), + ); + const activeSuccessorRun = db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + ne(sessionRunsTable.id, input.sessionRunId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ); + const claimed = await db + .update(driverInstancesTable) + .set({ statusOperationId: operationId }) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + input.expectedDriverConnectionId === undefined + ? undefined + : input.expectedDriverConnectionId === null + ? isNull(driverInstancesTable.connectionId) + : eq(driverInstancesTable.connectionId, input.expectedDriverConnectionId), + eq(driverInstancesTable.generation, input.driverGeneration), + inArray(driverInstancesTable.status, [ + "provisioning", + "connecting", + "ready", + "stopped", + "failed", + ]), + isNull(driverInstancesTable.statusOperationId), + exists(exactTerminalRun), + notExists(activeSuccessorRun), + ), + ) + .returning({ id: driverInstancesTable.id }) + .get(); + + if (claimed !== undefined) { + return { operationId, phase: "cleanup_required" }; + } + + const adopted = await db + .select({ status: driverInstancesTable.status }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + input.expectedDriverConnectionId === undefined + ? undefined + : input.expectedDriverConnectionId === null + ? isNull(driverInstancesTable.connectionId) + : eq(driverInstancesTable.connectionId, input.expectedDriverConnectionId), + eq(driverInstancesTable.generation, input.driverGeneration), + inArray(driverInstancesTable.status, [ + "provisioning", + "connecting", + "ready", + "stopping", + "stopped", + "failed", + ]), + eq(driverInstancesTable.statusOperationId, operationId), + exists(exactTerminalRun), + notExists(activeSuccessorRun), + ), + ) + .limit(1) + .get(); + + return adopted === undefined + ? null + : { + operationId, + phase: adopted.status === "stopping" ? "release_committed" : "cleanup_required", + }; +} + async function checkpointTerminalRuntimeSessionIfNeeded( bindings: ApiBindings, link: RuntimeSessionLink, + input: { readonly driverGeneration: number; readonly driverInstanceId: DriverInstanceId }, ): Promise { if ( link.sandboxId === null || @@ -46,192 +178,710 @@ async function checkpointTerminalRuntimeSessionIfNeeded( return; } + const lifecycle = await getAppDatabase(bindings.DB) + .select({ + cleanupOperationKind: sessionsTable.cleanupOperationKind, + operationId: sessionsTable.statusOperationId, + }) + .from(sessionsTable) + .where(eq(sessionsTable.id, link.sessionId)) + .limit(1) + .get(); + if (lifecycle?.cleanupOperationKind === "delete" && lifecycle.operationId !== null) { + return; + } + if ( + link.sandboxKind === "cattle" && + (await isCattleTerminalCheckpointReadyForNextRun(bindings.DB, link.sessionId)) + ) { + return; + } + const rules = getRuntimeKindPolicy(link.sandboxKind).checkpoint.createOnTerminal; if (rules.length === 0) { return; } + const driver = await getAppDatabase(bindings.DB) + .select({ sandboxIncarnation: driverInstancesTable.sandboxIncarnation }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.driverGeneration), + eq(driverInstancesTable.sandboxId, link.sandboxId), + ), + ) + .limit(1) + .get(); + if (driver === undefined) { + throw new Error("Terminal sandbox checkpoint lost its exact Driver incarnation."); + } + await createSandboxCheckpoints(bindings, { requiredSessionId: link.sessionId, rules, sandboxId: link.sandboxId, - sessionRunId: link.sessionRunId, + terminalAuthority: { + driverGeneration: input.driverGeneration, + driverInstanceId: input.driverInstanceId, + incarnation: driver.sandboxIncarnation, + sessionId: link.sessionId, + sessionRunId: link.sessionRunId, + }, }); } -function toFinalizedDriverRunTransitionRun( - outcome: SessionRunTransitionOutcome, -): SessionRunSummary | null { - switch (outcome.kind) { - case "applied": - case "duplicate": { - return outcome.run; - } - case "stale": { - if (outcome.reason === "terminal_run") { - return null; - } - throw new Error("Finalized driver repair lost a concurrent run transition."); - } - case "repair_needed": { - throw new Error("Finalized driver repair left the session lifecycle projection stale."); - } - case "rejected": { - throw new Error(`Finalized driver repair run transition was rejected: ${outcome.reason}.`); - } - } -} - -async function appendFinalizedDriverRunEvent( +async function failFinalizedDriverRun( bindings: ApiBindings, input: { - readonly run: SessionRunSummary; + readonly driverInstanceId: DriverInstanceId; + readonly runId: SessionRunId; readonly runError: RunError; + readonly runtimeId: string; readonly sessionId: SessionId; }, -): Promise { - await appendSessionRuntimeEvents({ - bindings, - events: [ - createFailedSessionRunRuntimeEvent({ - run: input.run, - runError: input.runError, - sessionId: input.sessionId, - sourceEventId: createSessionRunTerminalFailureSourceId(input.run.id), - }), - ], +): Promise { + const current = await getSessionRunSummary(bindings.DB, input.runId); + + if (current === null) { + throw new Error("Finalized Driver Session Run was not found."); + } + + if (current.status !== "failed" && isTerminalSessionRunStatus(current.status)) { + return null; + } + + const timestampMs = currentTimestampMs(); + const timestamp = toIsoString(timestampMs); + const runError = current.status === "failed" ? (current.error ?? input.runError) : input.runError; + const failedRun: SessionRunSummary = { + ...current, + completedAt: current.completedAt ?? timestamp, + error: runError, + startedAt: current.startedAt ?? timestamp, + status: "failed", + updatedAt: current.status === "failed" ? current.updatedAt : timestamp, + }; + const event = createCanonicalDriverRunFailedEvent({ + driverInstanceId: input.driverInstanceId, + error: runError, + id: createPlatformId(), + occurredAt: timestamp, + runId: input.runId, + runtimeId: input.runtimeId, sessionId: input.sessionId, + traceId: failedRun.traceId, }); + const outcome = await commitTerminalRunProjection(bindings.DB, { + assistantMessage: null, + error: runError, + runId: input.runId, + sessionId: input.sessionId, + source: "maintenance", + targetStatus: "failed", + terminalEvent: { + event, + occurredAt: timestampMs, + sourceEventId: event.sourceEventId ?? null, + }, + timestampMs, + }); + + if (outcome.kind === "stale") { + return null; + } + + return getSessionRunSummary(bindings.DB, input.runId); } export async function releaseTerminalDriverInstanceSessionRun( bindings: ApiBindings, input: { + expectedDriverConnectionId?: string | null; + driverGeneration: number; driverInstanceId: DriverInstanceId; sessionRunId: SessionRunId; }, ): Promise { const database = bindings.DB; + const claim = await claimTerminalRunRelease(database, input); + if (claim === null) { + throw new Error("Terminal Session Run release lost its exact Driver ownership."); + } const link = await getRuntimeSessionLink(database, input.driverInstanceId, { sessionRunId: input.sessionRunId, }); - await checkpointTerminalRuntimeSessionIfNeeded(bindings, link); + if ( + link.sessionRunId !== input.sessionRunId || + link.sessionRunStatus === null || + !isTerminalSessionRunStatus(link.sessionRunStatus) + ) { + throw new Error("Terminal Session Run no longer belongs to this Driver instance."); + } + + await repairTerminalDriverInputCommands(bindings, input); + + if (claim.phase === "release_committed") { + return { link, released: true }; + } + + await checkpointTerminalRuntimeSessionIfNeeded(bindings, link, input); + await closeTerminalRuntimeConversationIfNeeded(bindings, link); + await recycleTerminalRuntimeLeaseIfNeeded(bindings, link); const outcome = await recordRuntimeRunLeaseReleasedOutcome(database, { driverInstanceId: input.driverInstanceId, + expectedDriverGeneration: input.driverGeneration, + expectedDriverOperationId: claim.operationId, expectedSessionRunId: input.sessionRunId, + retainDriverOperationUntilTerminal: true, }); - const released = outcome.status === "applied"; - if (!released) { - logWarn("runtime.terminal.lease_release_skipped", { - driverInstanceId: input.driverInstanceId, - reason: "reason" in outcome ? outcome.reason : outcome.status, - sessionRunId: input.sessionRunId, - status: outcome.status, - }); + if (outcome.status !== "applied") { + const reason = "reason" in outcome ? outcome.reason : outcome.status; + throw new Error(`Terminal Session Run lease release failed: ${reason}.`); + } + + return { link, released: true }; +} + +function runHasReclaimError(run: SessionRunSummary, runError: RunError): boolean { + return ( + run.error?.code === runError.code && + run.error.details["driverInstanceId"] === runError.details["driverInstanceId"] + ); +} + +type McpTerminalStatus = "cancelled" | "completed" | "failed"; +type CompletedMcpCommandRepair = Extract< + AcceptedMcpCommandRepair["terminal"], + { status: "completed" } +>; +type ReconciledMcpTerminal = + | { result: CompletedMcpCommandRepair["result"]; status: "completed" } + | { error: RunError; status: "failed" } + | { status: "cancelled" }; +interface PersistedMcpTerminalEvent extends Record { + contentText: string; + semanticHash: string | null; + sourceEventId: string; + toolOutputText: string | null; + toolStatus: McpTerminalStatus; +} + +function mcpTerminalSourceEventId( + commandId: string, + status: Exclude, +): string { + return `mcp.execute.${status}:${commandId}`; +} + +function createMcpTerminalEvent( + driverInstanceId: DriverInstanceId, + repair: AcceptedMcpCommandRepair, + terminal: ReconciledMcpTerminal, + traceId: string, +): RuntimeEventEnvelope { + const identity = + terminal.status === "failed" + ? createMcpExecuteFailedEventIdentity({ + commandId: repair.commandId, + rawInput: repair.command.argumentsJson, + rawOutput: terminal.error.message, + title: repair.command.toolName, + toolCallId: repair.command.toolCallId, + }) + : ({ + payload: { + kind: "mcp", + rawInput: repair.command.argumentsJson, + ...(terminal.status === "completed" ? { rawOutput: terminal.result.outputText } : {}), + status: terminal.status, + title: repair.command.toolName, + toolCallId: repair.command.toolCallId, + }, + sourceEventId: mcpTerminalSourceEventId(repair.commandId, terminal.status), + } as const); + + return createRuntimeEvent({ + correlationId: repair.commandId, + driverInstanceId, + id: createPlatformId(), + kind: "tool.call.updated", + occurredAt: new Date().toISOString(), + payload: identity.payload, + runId: parsePlatformId(repair.command.runId, "MCP command Session Run ID"), + runtimeId: repair.runtimeId, + sessionId: repair.sessionId, + sourceEventId: identity.sourceEventId, + traceId, + }); +} + +async function readExistingMcpTerminalEvent(input: { + bindings: ApiBindings; + repair: AcceptedMcpCommandRepair; +}): Promise { + const rows = await getAppDatabase(input.bindings.DB) + .select({ + contentText: sessionEventsTable.contentText, + eventType: sessionEventsTable.eventType, + family: sessionEventsTable.family, + mcpCommandId: sessionEventsTable.mcpCommandId, + processStatus: sessionEventsTable.processStatus, + processType: sessionEventsTable.processType, + runId: sessionEventsTable.runId, + semanticHash: sessionEventsTable.semanticHash, + source: sessionEventsTable.source, + sourceEventId: sessionEventsTable.sourceEventId, + streamId: sessionEventsTable.streamId, + toolCallId: sessionEventsTable.toolCallId, + toolInputDeltaJson: sessionEventsTable.toolInputDeltaJson, + toolInputJson: sessionEventsTable.toolInputJson, + toolName: sessionEventsTable.toolName, + toolOutputDeltaText: sessionEventsTable.toolOutputDeltaText, + toolOutputText: sessionEventsTable.toolOutputText, + toolParentMessageId: sessionEventsTable.toolParentMessageId, + toolResultMessageId: sessionEventsTable.toolResultMessageId, + toolStatus: sessionEventsTable.toolStatus, + tokens: sessionEventsTable.tokens, + traceId: sessionEventsTable.traceId, + visibility: sessionEventsTable.visibility, + }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, input.repair.sessionId), + eq(sessionEventsTable.mcpCommandId, input.repair.commandId), + eq(sessionEventsTable.eventType, "tool.call.updated"), + inArray(sessionEventsTable.toolStatus, ["completed", "cancelled", "failed"]), + ), + ) + .all(); + + if (rows.length > 1) { + throw new Error("MCP command has conflicting durable terminal events."); + } + + return (rows[0] as PersistedMcpTerminalEvent | undefined) ?? null; +} + +function adoptExistingMcpTerminal(input: { + rawOutput: string | null; + repair: AcceptedMcpCommandRepair; + status: McpTerminalStatus; +}): ReconciledMcpTerminal { + if (input.status === "completed") { + if (input.repair.effectStatus !== "succeeded" || input.repair.terminal.status !== "completed") { + throw new Error("MCP completion event conflicts with its durable external effect."); + } + + return { result: input.repair.terminal.result, status: "completed" }; + } + + if (input.status === "cancelled") { + if (input.repair.effectStatus !== "intent") { + throw new Error("MCP cancellation event conflicts with its durable external effect."); + } + + return { status: "cancelled" }; + } + + if (input.repair.effectStatus === "succeeded" || input.repair.terminal.status !== "failed") { + throw new Error("MCP failure event conflicts with its durable external effect."); + } + + if (input.rawOutput === null) { + throw new Error("MCP failure event is missing its durable raw output."); + } + + return { + error: { ...input.repair.terminal.error, message: input.rawOutput }, + status: "failed", + }; +} + +async function assertMcpTerminalEventProjection( + row: Record, + event: RuntimeEventEnvelope, + commandId: AcceptedMcpCommandRepair["commandId"], +): Promise { + const expected = createSessionRuntimeEventProjection(event, { + provenMcpCommandId: commandId, + }); + + if (row["sourceEventId"] !== event.sourceEventId) { + throw new Error("MCP terminal event source identity conflicts with its command."); + } + + for (const [key, value] of Object.entries(expected)) { + if (row[key] !== value) { + throw new Error(`MCP terminal event projection ${key} conflicts with its command.`); + } + } + + if (row["semanticHash"] !== (await createRuntimeEventSemanticHash(event))) { + throw new Error("MCP terminal event semantic hash conflicts with its command."); + } +} + +async function reconcileAcceptedMcpCommand( + bindings: ApiBindings, + input: { + driverGeneration: number; + driverInstanceId: DriverInstanceId; + repair: AcceptedMcpCommandRepair; + }, +): Promise { + const runId = parsePlatformId( + input.repair.command.runId, + "MCP command Session Run ID", + ); + const link = await getRuntimeSessionLink(bindings.DB, input.driverInstanceId, { + sessionRunId: runId, + }); + + if ( + link.sessionId !== input.repair.sessionId || + link.sessionRunId !== runId || + link.runtimeId !== input.repair.runtimeId || + link.traceId === null + ) { + throw new Error("MCP command repair lost its authoritative Session Run identity."); + } + + const traceId = link.traceId; + const existing = await readExistingMcpTerminalEvent({ bindings, repair: input.repair }); + const proposedTerminal = + existing === null + ? input.repair.terminal.status === "completed" + ? ({ result: input.repair.terminal.result, status: "completed" } as const) + : ({ error: input.repair.terminal.error, status: "failed" } as const) + : adoptExistingMcpTerminal({ + rawOutput: existing.toolOutputText, + repair: input.repair, + status: existing.toolStatus, + }); + + if (existing === null) { + try { + await appendSessionRuntimeEvents({ + bindings, + events: [ + createMcpTerminalEvent(input.driverInstanceId, input.repair, proposedTerminal, traceId), + ], + provenMcpCommandId: input.repair.commandId, + sessionId: input.repair.sessionId, + }); + } catch (error) { + // The terminal tool partial-unique index chooses exactly one concurrent + // Driver/repair winner. Only suppress the insert error when that winner + // can now be read and validated below. + if ((await readExistingMcpTerminalEvent({ bindings, repair: input.repair })) === null) { + throw error; + } + } + } + + const persisted = + existing ?? (await readExistingMcpTerminalEvent({ bindings, repair: input.repair })); + + if (persisted === null) { + throw new Error("MCP terminal event was not durably persisted."); + } + + const terminal = adoptExistingMcpTerminal({ + rawOutput: persisted.toolOutputText, + repair: input.repair, + status: persisted.toolStatus, + }); + const event = createMcpTerminalEvent(input.driverInstanceId, input.repair, terminal, traceId); + await assertMcpTerminalEventProjection(persisted, event, input.repair.commandId); + const outcome = await updateRuntimeCommandRecord(bindings.DB, { + commandId: input.repair.commandId, + driverGeneration: input.driverGeneration, + driverInstanceId: input.driverInstanceId, + ...(terminal.status === "completed" + ? { result: terminal.result } + : terminal.status === "failed" + ? { error: terminal.error } + : {}), + status: terminal.status, + }); + + if (outcome.kind === "rejected") { + throw new Error(`MCP command terminal repair was rejected: ${outcome.reason}.`); + } +} + +async function assertInputTerminalEventPersisted( + bindings: ApiBindings, + repair: AcceptedInputStartCommandRepair, +): Promise { + const eventType = + repair.terminal.status === "completed" + ? "run.completed" + : repair.terminal.status === "cancelled" + ? "run.cancelled" + : "run.failed"; + const row = + (await getAppDatabase(bindings.DB) + .select({ id: sessionEventsTable.id }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, repair.sessionId), + eq( + sessionEventsTable.runId, + parsePlatformId(repair.command.runId, "input command Session Run ID"), + ), + eq(sessionEventsTable.eventType, eventType), + ), + ) + .limit(1) + .get()) ?? null; + + if (row === null) { + throw new Error("Input command terminal Session Run event is not durably persisted."); + } +} + +async function reconcileAcceptedInputStartCommand( + bindings: ApiBindings, + input: { + driverGeneration: number; + driverInstanceId: DriverInstanceId; + repair: AcceptedInputStartCommandRepair; + }, +): Promise { + await assertInputTerminalEventPersisted(bindings, input.repair); + const outcome = await updateRuntimeCommandRecord(bindings.DB, { + commandId: input.repair.commandId, + driverGeneration: input.driverGeneration, + driverInstanceId: input.driverInstanceId, + ...(input.repair.terminal.status === "completed" + ? { result: input.repair.terminal.result } + : input.repair.terminal.status === "failed" + ? { error: input.repair.terminal.error } + : {}), + status: input.repair.terminal.status, + }); + + if (outcome.kind === "rejected") { + throw new Error(`Input command terminal repair was rejected: ${outcome.reason}.`); } +} - await closeReleasedTerminalRuntimeLeaseIfNeeded(bindings, { link, released }); +async function repairTerminalDriverMcpCommands( + bindings: ApiBindings, + input: { + driverGeneration: number; + driverInstanceId: DriverInstanceId; + nowMs?: number; + }, +): Promise { + await markClaimedExternalToolEffectsUnknownForDriver(bindings.DB, input); + const mcpRepairs = await listAcceptedMcpCommandRepairsForTerminalDriver(bindings.DB, input); - return { link, released }; + for (const repair of mcpRepairs) { + await reconcileAcceptedMcpCommand(bindings, { ...input, repair }); + } +} + +async function repairTerminalDriverInputCommands( + bindings: ApiBindings, + input: { + driverGeneration: number; + driverInstanceId: DriverInstanceId; + nowMs?: number; + }, +): Promise { + const inputRepairs = await listAcceptedInputStartCommandRepairsForTerminalDriver( + bindings.DB, + input, + ); + + for (const repair of inputRepairs) { + await reconcileAcceptedInputStartCommand(bindings, { ...input, repair }); + } } export async function repairFinalizedTerminalDriverRunState( bindings: ApiBindings, input: { + driverGeneration: number; driverInstanceId: DriverInstanceId; + sessionRunId: SessionRunId | null; status: "failed" | "stopped"; }, ): Promise { - await markExecutingExternalToolEffectsUnknownForDriver(bindings.DB, input.driverInstanceId); - await failAcceptedRuntimeCommandsForTerminalDriver(bindings.DB, { + const commandRepairInput = { + driverGeneration: input.driverGeneration, driverInstanceId: input.driverInstanceId, - }); + }; - const link = await getRuntimeSessionLink(bindings.DB, input.driverInstanceId); + // Close any durable tool UI state before the enclosing run is made terminal. + await repairTerminalDriverMcpCommands(bindings, commandRepairInput); - if (link.sessionRunId === null || link.sessionRunStatus === null) { - return { link, released: false }; + if (input.sessionRunId === null) { + await repairTerminalDriverInputCommands(bindings, commandRepairInput); + await repairAcceptedRuntimeCommandsForTerminalDriver(bindings.DB, commandRepairInput); + return { link: null, released: false }; + } + + const link = await getRuntimeSessionLink(bindings.DB, input.driverInstanceId, { + sessionRunId: input.sessionRunId, + }); + + if (link.sessionRunId !== input.sessionRunId || link.sessionRunStatus === null) { + throw new Error("Finalized Driver Session Run ownership was lost."); } - if (!isTerminalSessionRunStatus(link.sessionRunStatus)) { + const wasActive = !isTerminalSessionRunStatus(link.sessionRunStatus); + + if (wasActive || link.sessionRunStatus === "failed") { const runError = classifyReclaim({ driverInstanceId: input.driverInstanceId, driverTerminalStatus: input.status, reclaimReason: "socket_closed", }); - const outcome = await setSessionRunStatus(bindings.DB, { - error: runError, - runId: link.sessionRunId, - source: "driver", - status: "failed", - }); - const run = toFinalizedDriverRunTransitionRun(outcome); - - if (link.sessionId !== null && run !== null) { - await appendFinalizedDriverRunEvent(bindings, { - run, - runError, - sessionId: link.sessionId, - }); + const run = + link.sessionId === null || link.runtimeId === null + ? null + : await failFinalizedDriverRun(bindings, { + driverInstanceId: input.driverInstanceId, + runError, + runId: link.sessionRunId, + runtimeId: link.runtimeId, + sessionId: link.sessionId, + }); + if (link.sessionId !== null && run !== null && runHasReclaimError(run, runError)) { // Decide recovery for the reclaimed run. v1 records the decision so it is // observable and unit-testable; executing the auto-requeue (a fresh // `resume` run + re-dispatch) is a follow-up because this DO finalize // context lacks the viewer + requestUrl that enqueueSessionRunDispatchCommand // needs to rebuild the sandbox's action-token callback URLs. - const recovery = decideReclaimRecovery({ - driverTerminalStatus: input.status, - priorTrigger: run.trigger, - reclaimReason: "socket_closed", - runStatus: link.sessionRunStatus, - }); - logInfo("runtime.reclaim.recovery.decided", { - action: recovery.kind, - driverInstanceId: input.driverInstanceId, - priorTrigger: run.trigger, - runId: run.id, - sessionId: link.sessionId, - }); + if (wasActive) { + const recovery = decideReclaimRecovery({ + driverTerminalStatus: input.status, + priorTrigger: run.trigger, + reclaimReason: "socket_closed", + runStatus: link.sessionRunStatus, + }); + logInfo("runtime.reclaim.recovery.decided", { + action: recovery.kind, + driverInstanceId: input.driverInstanceId, + priorTrigger: run.trigger, + runId: run.id, + sessionId: link.sessionId, + }); + } } } + // input.start is derived from the authoritative terminal run and its durable + // event, so it must be repaired after the run transition above. + await repairTerminalDriverInputCommands(bindings, commandRepairInput); + await repairAcceptedRuntimeCommandsForTerminalDriver(bindings.DB, commandRepairInput); + return releaseTerminalDriverInstanceSessionRun(bindings, { + driverGeneration: input.driverGeneration, driverInstanceId: input.driverInstanceId, sessionRunId: link.sessionRunId, }); } +export async function repairTerminalDriverRuntimeCommandsGlobally( + bindings: ApiBindings, +): Promise { + const failures: unknown[] = []; + const claimedReleases = await getAppDatabase(bindings.DB) + .select({ + driverGeneration: driverInstancesTable.generation, + driverInstanceId: driverInstancesTable.id, + sessionRunId: sessionRunsTable.id, + }) + .from(driverInstancesTable) + .innerJoin( + sessionRunsTable, + and( + eq(sessionRunsTable.driverInstanceId, driverInstancesTable.id), + sql`${sessionRunsTable.id} = ${driverInstancesTable.statusOperationId}`, + ), + ) + .where( + and( + inArray(driverInstancesTable.status, ["ready", "stopped", "failed"]), + inArray(sessionRunsTable.status, ["cancelled", "completed", "expired", "failed"]), + ), + ) + .all(); + + for (const release of claimedReleases) { + try { + await releaseTerminalDriverInstanceSessionRun(bindings, release); + } catch (error) { + failures.push(error); + } + } + + const drivers = await listTerminalDriversWithPendingRuntimeCommands(bindings.DB); + + for (const driver of drivers) { + try { + const identity = await getDriverInstanceLifecycleIdentity(bindings, driver.id); + + if ( + identity === null || + identity.generation !== driver.generation || + (identity.status !== "failed" && identity.status !== "stopped") + ) { + continue; + } + + let link = await getRuntimeSessionLink(bindings.DB, driver.id); + + if (link.sessionRunId === null) { + link = await getRuntimeSessionLink(bindings.DB, driver.id, { + latestTerminalRun: true, + }); + } + await repairFinalizedTerminalDriverRunState(bindings, { + driverGeneration: driver.generation, + driverInstanceId: driver.id, + sessionRunId: link.sessionRunId, + status: identity.status, + }); + } catch (error) { + failures.push(error); + } + } + + if (failures[0] !== undefined) { + throw failures[0]; + } +} + export async function releaseLinkedTerminalDriverInstanceSessionRun( bindings: ApiBindings, driverInstanceId: DriverInstanceId, + driverGeneration: number, ): Promise { - const database = bindings.DB; - const row: LinkedSessionRunStatusRow | null = - (await getAppDatabase(database) - .select({ - sessionRunId: sessionRunsTable.id, - status: sessionRunsTable.status, - }) - .from(sessionRunsTable) - .where( - and( - eq(sessionRunsTable.driverInstanceId, driverInstanceId), - inArray(sessionRunsTable.status, ["cancelled", "completed", "expired", "failed"]), - ), - ) - .limit(1) - .get()) ?? null; + const link = await getRuntimeSessionLink(bindings.DB, driverInstanceId, { + latestTerminalRun: true, + }); - if (row === null || row.sessionRunId === null || !isTerminalSessionRunStatus(row.status)) { - return { link: null, released: false }; + if (link.sessionRunId === null) { + return { link, released: false }; } return releaseTerminalDriverInstanceSessionRun(bindings, { + driverGeneration, driverInstanceId, - sessionRunId: row.sessionRunId, + sessionRunId: link.sessionRunId, }); } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-runtime-lease.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-runtime-lease.ts index 96d7eaf7..21142797 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-runtime-lease.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-runtime-lease.ts @@ -5,21 +5,6 @@ import { recycleInactiveRuntimeSubjectNow } from "../runtime-subject-lifecycle/r import { closeSandboxConversationSession } from "../sandbox-session.service"; import type { RuntimeSessionLink } from "./event-types"; -export async function closeReleasedTerminalRuntimeLeaseIfNeeded( - bindings: ApiBindings, - input: { - readonly link: RuntimeSessionLink | null; - readonly released: boolean; - }, -): Promise { - if (!input.released || input.link === null) { - return; - } - - await closeTerminalRuntimeConversationIfNeeded(bindings, input.link); - await recycleReleasedTerminalRuntimeLeaseIfNeeded(bindings, input); -} - export async function closeTerminalRuntimeConversationIfNeeded( bindings: ApiBindings, link: RuntimeSessionLink, @@ -40,19 +25,10 @@ export async function closeTerminalRuntimeConversationIfNeeded( }); } -export async function recycleReleasedTerminalRuntimeLeaseIfNeeded( +export async function recycleTerminalRuntimeLeaseIfNeeded( bindings: ApiBindings, - input: { - readonly link: RuntimeSessionLink | null; - readonly released: boolean; - }, + link: RuntimeSessionLink, ): Promise { - if (!input.released || input.link === null) { - return; - } - - const link = input.link; - if (link.sandboxKind === null || link.sandboxId === null || link.sessionId === null) { return; } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-state-coordinator.ts b/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-state-coordinator.ts index e0c45ba9..7ac4cfca 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-state-coordinator.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-instance/terminal-state-coordinator.ts @@ -10,16 +10,21 @@ import { } from "../../application/runtime-diagnostic-events"; import { resolvePendingRuntimeCommands } from "./commands"; import { runtimeSessionLinkNeedsRefresh } from "./event-types"; +import type { RuntimeSessionLink } from "./event-types"; import { finalizeDriverInstance } from "./lifecycle"; import type { RuntimeSessionViewCache } from "./runtime-session-view-cache"; import type { DriverInstanceRuntimeState } from "./runtime-state"; import { getRuntimeSessionLink } from "./session-link.repository"; import type { SessionViewerEventDeliveryBuffer } from "./session-viewer-event-delivery-buffer"; -import type { DriverInstanceCloseSnapshot } from "./state"; +import type { DriverInstanceCloseSnapshot, DriverInstanceConnectionEpoch } from "./state"; import { repairFinalizedTerminalDriverRunState } from "./terminal-run-release"; + interface DriverInstanceTerminalStateCoordinatorOptions { + appendDiagnosticEvent?: typeof appendRuntimeDiagnosticEvent; clearStorage: () => Promise; env: ApiBindings; + finalizeDriver?: typeof finalizeDriverInstance; + repairFinalizedRunState?: typeof repairFinalizedTerminalDriverRunState; state: DriverInstanceRuntimeState; viewCache: RuntimeSessionViewCache; viewerEventDelivery: SessionViewerEventDeliveryBuffer; @@ -27,115 +32,198 @@ interface DriverInstanceTerminalStateCoordinatorOptions { } export class DriverInstanceTerminalStateCoordinator { + readonly #appendDiagnosticEvent: typeof appendRuntimeDiagnosticEvent; readonly #clearStorage: () => Promise; readonly #env: ApiBindings; + readonly #finalizeDriver: typeof finalizeDriverInstance; + #finalizationTask: { epoch: DriverInstanceConnectionEpoch; task: Promise } | null = null; + readonly #repairFinalizedRunState: typeof repairFinalizedTerminalDriverRunState; readonly #state: DriverInstanceRuntimeState; + #resetTask: Promise | null = null; readonly #viewCache: RuntimeSessionViewCache; readonly #viewerEventDelivery: SessionViewerEventDeliveryBuffer; readonly #withRuntimeLogContext: (fn: () => T) => T; constructor(options: DriverInstanceTerminalStateCoordinatorOptions) { + this.#appendDiagnosticEvent = options.appendDiagnosticEvent ?? appendRuntimeDiagnosticEvent; this.#clearStorage = options.clearStorage; this.#env = options.env; + this.#finalizeDriver = options.finalizeDriver ?? finalizeDriverInstance; + this.#repairFinalizedRunState = + options.repairFinalizedRunState ?? repairFinalizedTerminalDriverRunState; this.#state = options.state; this.#viewCache = options.viewCache; this.#viewerEventDelivery = options.viewerEventDelivery; this.#withRuntimeLogContext = options.withRuntimeLogContext; } - async finalize(): Promise { - if (this.#state.terminalized) { + async finalize(epoch: DriverInstanceConnectionEpoch): Promise { + if (!this.#state.matchesConnectionEpoch(epoch)) { + return; + } + + if (this.#finalizationTask !== null && this.#epochsMatch(this.#finalizationTask.epoch, epoch)) { + return this.#finalizationTask.task; + } + + if (this.#state.terminalCleanupComplete) { return; } - this.#state.terminalized = true; - await this.#viewerEventDelivery.flushSafely(); + const task = this.#finalize(epoch).finally(() => { + if (this.#finalizationTask?.task === task) { + this.#finalizationTask = null; + } + }); + this.#finalizationTask = { epoch, task }; + return task; + } + + async #finalize(epoch: DriverInstanceConnectionEpoch): Promise { + if (!this.#state.matchesConnectionEpoch(epoch)) { + return; + } const driverInstanceId = this.#state.requireDriverInstanceId(); - const close = await this.#ensureCloseSnapshot(); - const status = getDriverInstanceTerminalStatus(this.#state.errorMessage, close.code); + const close = await this.#ensureCloseSnapshot(epoch); + + if (!this.#state.matchesConnectionEpoch(epoch)) { + return; + } + + const terminalSessionLink = await this.#captureTerminalSessionLink(driverInstanceId, epoch); + + if (!this.#state.matchesConnectionEpoch(epoch)) { + return; + } + + this.#viewerEventDelivery.requestStateSync(terminalSessionLink.sessionId); + const desiredStatus = getDriverInstanceTerminalStatus(this.#state.errorMessage, close.code); const closeResult = this.#state.closeResult(); - const connectionId = this.#state.connectionId; - const finalized = isTruthy(connectionId) - ? await finalizeDriverInstance(this.#env, driverInstanceId, { - closeCode: close.code, - closeReason: close.reason || null, - connectionId, - connectedAt: this.#state.connectedAt, - driverPid: this.#state.hello?.pid ?? null, - driverStartedAt: this.#state.hello?.startedAt ?? null, - errorMessage: this.#state.errorMessage, - generation: this.#state.requireDriverGeneration(), - heartbeatCount: this.#state.heartbeatCount, - lastHeartbeatAt: this.#state.lastHeartbeat?.at ?? null, - status, - }) - : false; - - if (finalized) { - await this.#repairFinalizedRunState({ - driverInstanceId, - status, - }); - await this.#appendDriverCrashedEventIfNeeded({ - close, - driverInstanceId, - status, - }); + const snapshot = { + connectedAt: this.#state.connectedAt, + driverPid: this.#state.hello?.pid ?? null, + driverStartedAt: this.#state.hello?.startedAt ?? null, + errorMessage: this.#state.errorMessage, + heartbeatCount: this.#state.heartbeatCount, + lastHeartbeatAt: this.#state.lastHeartbeat?.at ?? null, + terminalSessionRunId: this.#state.terminalSessionRunId, + traceId: this.#state.traceId, + }; + const terminalStatus = await this.#finalizeDriver(this.#env, driverInstanceId, { + closeCode: close.code, + closeReason: close.reason || null, + connectionId: epoch.connectionId, + connectedAt: snapshot.connectedAt, + driverPid: snapshot.driverPid, + driverStartedAt: snapshot.driverStartedAt, + errorMessage: snapshot.errorMessage, + generation: epoch.generation, + heartbeatCount: snapshot.heartbeatCount, + lastHeartbeatAt: snapshot.lastHeartbeatAt, + status: desiredStatus, + }); - this.#withRuntimeLogContext(() => { - logInfo("runtime.run.finalized", { - closeCode: close.code, - closeReason: close.reason || null, - connectedAt: this.#state.connectedAt, - connectionId, - driverInstanceId, - driverPid: this.#state.hello?.pid ?? null, - errorMessage: this.#state.errorMessage, - heartbeatCount: this.#state.heartbeatCount, - status, - }); - }); + if (terminalStatus === null || !this.#state.matchesConnectionEpoch(epoch)) { + return; } - this.#state.resolveCloseWaiters(closeResult); - this.#state.rejectHeartbeatWaiters(new Error(`Driver instance ${driverInstanceId} is closed.`)); + await this.#repairFinalizedRunState(this.#env, { + driverGeneration: epoch.generation, + driverInstanceId, + sessionRunId: snapshot.terminalSessionRunId, + status: terminalStatus, + }); - if (!this.#state.hello) { - this.#state.rejectHelloWaiters( - new Error(`Driver instance ${driverInstanceId} closed before hello.`), - ); + if (!this.#state.matchesConnectionEpoch(epoch)) { + return; + } + + await this.#appendDriverCrashedEventIfNeeded({ + close, + driverGeneration: epoch.generation, + driverInstanceId, + link: terminalSessionLink, + status: terminalStatus, + traceId: snapshot.traceId, + }); + + if (!this.#state.matchesConnectionEpoch(epoch)) { + return; } + this.#withRuntimeLogContext(() => { + logInfo("runtime.run.finalized", { + closeCode: close.code, + closeReason: close.reason || null, + connectedAt: snapshot.connectedAt, + connectionId: epoch.connectionId, + driverInstanceId, + driverPid: snapshot.driverPid, + errorMessage: snapshot.errorMessage, + heartbeatCount: snapshot.heartbeatCount, + status: terminalStatus, + }); + }); + + this.#state.resolveCloseWaiters(closeResult, epoch.generation); + if (!this.#state.ready) { this.#state.rejectReadyWaiters( new Error(`Driver instance ${driverInstanceId} closed before ready.`), + epoch.generation, ); } resolvePendingRuntimeCommands(this.#state.commandWaiters); - await this.#state.persistTerminalSnapshot(); + await this.#state.persistTerminalSnapshot(epoch); } - async resetForReuse(): Promise { + async resetForReuse(driverGeneration: number): Promise { + await this.prepareForReuse(); + + return (this.#resetTask ??= this.#resetForReuse(driverGeneration).finally(() => { + this.#resetTask = null; + })); + } + + async prepareForReuse(): Promise { + if (this.#finalizationTask !== null) { + await this.#finalizationTask.task; + } + + this.#viewerEventDelivery.resetAfterFlush(); + } + + async #resetForReuse(driverGeneration: number): Promise { await this.#state.resetForReuse({ beforeReset: async () => { - await this.#viewerEventDelivery.flushSafely(); this.#viewerEventDelivery.resetAfterFlush(); this.#viewCache.reset(); }, + driverGeneration, }); } async destroy(reason: string): Promise { - await this.#viewerEventDelivery.flushSafely(); + if (this.#finalizationTask !== null) { + await this.#finalizationTask.task; + } + + if (this.#resetTask !== null) { + await this.#resetTask; + } + this.#viewerEventDelivery.resetAfterFlush(); this.#viewCache.reset(); await this.#clearStorage(); this.#state.resetAfterDestroy(reason); } - async #ensureCloseSnapshot(): Promise { + async #ensureCloseSnapshot( + epoch: DriverInstanceConnectionEpoch, + ): Promise { + this.#state.assertConnectionEpoch(epoch); const close = this.#state.close ?? ({ @@ -145,7 +233,7 @@ export class DriverInstanceTerminalStateCoordinator { } satisfies DriverInstanceCloseSnapshot); if (!this.#state.close) { - await this.#state.persistClose(close); + await this.#state.persistClose(close, epoch); } return close; @@ -153,33 +241,32 @@ export class DriverInstanceTerminalStateCoordinator { async #appendDriverCrashedEventIfNeeded(input: { close: DriverInstanceCloseSnapshot; + driverGeneration: number; driverInstanceId: DriverInstanceId; + link: RuntimeSessionLink; status: "failed" | "stopped"; + traceId: string | null; }): Promise { if (input.status !== "failed") { return; } try { - const cachedLink = this.#state.runtimeSessionLink; - const link = - cachedLink !== null && !runtimeSessionLinkNeedsRefresh(cachedLink) - ? cachedLink - : await getRuntimeSessionLink(this.#env.DB, input.driverInstanceId); - this.#state.setRuntimeSessionLink(link); + const { link } = input; if (!isTruthy(link.agentId) || !isTruthy(link.sessionId)) { return; } - await appendRuntimeDiagnosticEvent(this.#env, { + await this.#appendDiagnosticEvent(this.#env, { eventName: RUNTIME_DIAGNOSTIC_EVENT.driverCrashed.name, sessionId: link.sessionId, + sourceEventId: `driver-terminal:${input.driverInstanceId}:${String(input.driverGeneration)}:crashed`, value: { ...toRuntimeDiagnosticBaseValue({ agentId: link.agentId, sessionId: link.sessionId, - traceId: this.#state.traceId, + traceId: input.traceId, }), driverInstanceId: input.driverInstanceId, status: input.close.reason || "failed", @@ -195,21 +282,39 @@ export class DriverInstanceTerminalStateCoordinator { } } - async #repairFinalizedRunState(input: { - driverInstanceId: DriverInstanceId; - status: "failed" | "stopped"; - }): Promise { - try { - await repairFinalizedTerminalDriverRunState(this.#env, input); - } catch (error) { - this.#withRuntimeLogContext(() => { - logWarn("runtime.driver.finalize_repair.failed", { - ...createErrorLogContext(error), - driverInstanceId: input.driverInstanceId, - status: input.status, - }); - }); + async #captureTerminalSessionLink( + driverInstanceId: DriverInstanceId, + epoch: DriverInstanceConnectionEpoch, + ): Promise { + this.#state.assertConnectionEpoch(epoch); + const sessionRunId = this.#state.terminalSessionRunId; + const cachedLink = this.#state.runtimeSessionLink; + const link = + sessionRunId !== null + ? cachedLink?.sessionRunId === sessionRunId && !runtimeSessionLinkNeedsRefresh(cachedLink) + ? cachedLink + : await getRuntimeSessionLink(this.#env.DB, driverInstanceId, { sessionRunId }) + : cachedLink !== null && !runtimeSessionLinkNeedsRefresh(cachedLink) + ? cachedLink + : await getRuntimeSessionLink(this.#env.DB, driverInstanceId); + + this.#state.assertConnectionEpoch(epoch); + + if (sessionRunId !== null && link.sessionRunId !== sessionRunId) { + throw new Error("Terminal Session Run ownership was lost."); } + + if (sessionRunId === null && link.sessionRunId !== null) { + await this.#state.setTerminalSessionRunId(link.sessionRunId, epoch); + } + + this.#state.assertConnectionEpoch(epoch); + this.#state.setRuntimeSessionLink(link); + return link; + } + + #epochsMatch(left: DriverInstanceConnectionEpoch, right: DriverInstanceConnectionEpoch): boolean { + return left.connectionId === right.connectionId && left.generation === right.generation; } } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-session-startup.ts b/apps/api/src/modules/runtime/infrastructure/driver-session-startup.ts index f5acdb6e..961f099b 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-session-startup.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-session-startup.ts @@ -26,6 +26,7 @@ async function createDriverStartupExitError( bindings: ApiBindings, input: { driverInstanceId: DriverInstanceId; + driverGeneration: number; eventContext?: DriverRuntimeStartupEventContext; exitCode: number; logContext: Record; @@ -63,7 +64,7 @@ async function createDriverStartupExitError( task: () => input.markStartupFailed ? input.markStartupFailed(message) - : failDriverInstance(bindings, input.driverInstanceId, message), + : failDriverInstance(bindings, input.driverInstanceId, input.driverGeneration, message), }); return new Error(message); @@ -73,6 +74,7 @@ export async function waitForDriverReady( bindings: ApiBindings, input: { driverInstanceId: DriverInstanceId; + driverGeneration: number; eventContext: DriverRuntimeStartupEventContext; getStaleStartupError?: () => Promise; logContext: Record; @@ -85,6 +87,7 @@ export async function waitForDriverReady( const readyPromise = waitForDriverInstanceReady( bindings, input.driverInstanceId, + input.driverGeneration, DRIVER_COLD_READY_TIMEOUT_MS, ).then(() => { ready = true; @@ -102,6 +105,7 @@ export async function waitForDriverReady( throw await createDriverStartupExitError(bindings, { driverInstanceId: input.driverInstanceId, + driverGeneration: input.driverGeneration, eventContext: input.eventContext, exitCode: exit.exitCode, logContext: input.logContext, diff --git a/apps/api/src/modules/runtime/infrastructure/driver-session-state.ts b/apps/api/src/modules/runtime/infrastructure/driver-session-state.ts index 69db540d..a0695fd4 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-session-state.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-session-state.ts @@ -1,5 +1,5 @@ import { driverInstancesTable, sessionRunsTable } from "@mosoo/db"; -import type { DriverInstanceId, SessionRunId } from "@mosoo/id"; +import type { DriverInstanceId, SandboxId, SessionId, SessionRunId } from "@mosoo/id"; import { and, eq, inArray } from "drizzle-orm"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; @@ -11,7 +11,13 @@ import type { DriverInstanceStatus } from "./driver-instance/status"; export async function getDriverUsage( database: D1Database, driverInstanceId: DriverInstanceId, + scope: { + readonly sandboxId: SandboxId; + readonly sandboxIncarnation: number; + readonly sandboxSessionId: SessionId; + }, ): Promise<{ + generation: number; sessionRunId: SessionRunId | null; status: DriverInstanceStatus; } | null> { @@ -19,10 +25,18 @@ export async function getDriverUsage( const row = (await appDb .select({ + generation: driverInstancesTable.generation, status: driverInstancesTable.status, }) .from(driverInstancesTable) - .where(eq(driverInstancesTable.id, driverInstanceId)) + .where( + and( + eq(driverInstancesTable.id, driverInstanceId), + eq(driverInstancesTable.sandboxId, scope.sandboxId), + eq(driverInstancesTable.sandboxIncarnation, scope.sandboxIncarnation), + eq(driverInstancesTable.sandboxSessionId, scope.sandboxSessionId), + ), + ) .limit(1) .get()) ?? null; @@ -44,6 +58,7 @@ export async function getDriverUsage( .get()) ?? null; return { + generation: row.generation, sessionRunId: activeRun?.id ?? null, status: row.status, }; diff --git a/apps/api/src/modules/runtime/infrastructure/driver-session-stop.service.ts b/apps/api/src/modules/runtime/infrastructure/driver-session-stop.service.ts index b0cb6509..09f6185b 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-session-stop.service.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-session-stop.service.ts @@ -1,129 +1,365 @@ -import type { RunError, SessionRunStatus } from "@mosoo/contracts/session-run"; -import { sessionRunsTable } from "@mosoo/db"; +import { driverInstancesTable, sessionRunsTable, sessionsTable } from "@mosoo/db"; import { createPlatformId } from "@mosoo/id"; import type { DriverInstanceId, RuntimeOperationId, SessionRunId } from "@mosoo/id"; -import { and, eq, inArray } from "drizzle-orm"; +import { and, eq, exists, inArray, isNotNull, isNull, ne, notExists, or, sql } from "drizzle-orm"; import { logWarn } from "../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; import { getAppDatabase } from "../../../platform/db/drizzle"; +import { currentTimestampMs } from "../../../time"; +import { + ASSIGNABLE_DRIVER_INSTANCE_STATUSES, + toDriverInstanceStatusLifecycleEventName, +} from "../domain/driver-instance-lifecycle.machine"; import { RUNTIME_SOCKET_TIMEOUT_MS } from "../domain/runtime-config"; import { ACTIVE_SESSION_RUN_STATUSES } from "../domain/session-run-lifecycle.machine"; import { + destroyDriverInstanceDurableObject, failDriverInstance, sendDriverInstanceCommand, waitForDriverInstanceClose, } from "./driver-instance/client"; -import { getDriverInstanceRecord } from "./driver-instance/driver-instance-record.repository"; import { isDriverControlSocketMissingError } from "./driver-session-stop-errors"; import { recordRuntimeRunLeaseReleasedOutcome } from "./runtime-subject-lifecycle/runtime-run-lease-store"; -import { setSessionRunStatus } from "./session-runs/session-run-store.repository"; -import type { SessionRunTransitionOutcome } from "./session-runs/session-run-store.repository"; -function assertStoppedDriverRunTransition(outcome: SessionRunTransitionOutcome): void { - switch (outcome.kind) { - case "applied": - case "duplicate": { - return; - } - case "stale": { - if (outcome.reason === "terminal_run") { - return; - } - throw new Error("Driver stop lost a concurrent run transition."); - } - case "repair_needed": { - throw new Error("Driver stop left session projection stale."); - } - case "rejected": { - throw new Error(`Driver stop run transition was rejected: ${outcome.reason}.`); - } - } +interface DriverStopSnapshot { + readonly generation: number; + readonly status: (typeof driverInstancesTable.$inferSelect)["status"]; + readonly statusOperationId: RuntimeOperationId | null; +} + +async function readDriverStopSnapshot( + database: D1Database, + driverInstanceId: DriverInstanceId, +): Promise { + return ( + (await getAppDatabase(database) + .select({ + generation: driverInstancesTable.generation, + status: driverInstancesTable.status, + statusOperationId: driverInstancesTable.statusOperationId, + }) + .from(driverInstancesTable) + .where(eq(driverInstancesTable.id, driverInstanceId)) + .limit(1) + .get()) ?? null + ); } async function getActiveDriverSessionRunId( database: D1Database, driverInstanceId: DriverInstanceId, ): Promise { - const row = - (await getAppDatabase(database) - .select({ id: sessionRunsTable.id }) + const row = await getAppDatabase(database) + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.driverInstanceId, driverInstanceId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ) + .limit(1) + .get(); + + return row?.id ?? null; +} + +async function claimDriverStop( + database: D1Database, + input: { + readonly driverGeneration: number; + readonly driverInstanceId: DriverInstanceId; + readonly expectedSessionRunId: SessionRunId | null; + readonly operationId: RuntimeOperationId; + }, +): Promise { + const db = getAppDatabase(database); + const unexpectedActiveRun = db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ...(input.expectedSessionRunId === null + ? [] + : [ne(sessionRunsTable.id, input.expectedSessionRunId)]), + ), + ); + const expectedRun = + input.expectedSessionRunId === null + ? null + : db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, input.expectedSessionRunId), + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + ), + ); + const now = currentTimestampMs(); + const claimed = await db + .update(driverInstancesTable) + .set({ + status: "stopping", + statusChangedAt: now, + statusEvent: toDriverInstanceStatusLifecycleEventName("stopping"), + statusOperationId: input.operationId, + statusSeq: sql`${driverInstancesTable.statusSeq} + 1`, + statusSource: "api", + updatedAt: now, + }) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.driverGeneration), + inArray(driverInstancesTable.status, ASSIGNABLE_DRIVER_INSTANCE_STATUSES), + or( + isNull(driverInstancesTable.statusOperationId), + eq(driverInstancesTable.statusOperationId, input.operationId), + ), + notExists(unexpectedActiveRun), + ...(expectedRun === null ? [] : [exists(expectedRun)]), + ), + ) + .returning({ id: driverInstancesTable.id }) + .get(); + + if (claimed !== undefined) { + return true; + } + + const adopted = await db + .select({ id: driverInstancesTable.id }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.driverGeneration), + eq(driverInstancesTable.status, "stopping"), + eq(driverInstancesTable.statusOperationId, input.operationId), + notExists(unexpectedActiveRun), + ...(expectedRun === null ? [] : [exists(expectedRun)]), + ), + ) + .limit(1) + .get(); + + return adopted !== undefined; +} + +async function releaseDriverStopClaim( + bindings: ApiBindings, + input: { + readonly driverGeneration: number; + readonly driverInstanceId: DriverInstanceId; + readonly operationId: RuntimeOperationId; + readonly sessionRunId: SessionRunId | null; + }, +): Promise { + if (input.sessionRunId !== null) { + const terminalRun = await getAppDatabase(bindings.DB) + .select({ status: sessionRunsTable.status }) .from(sessionRunsTable) .where( and( - eq(sessionRunsTable.driverInstanceId, driverInstanceId), - inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + eq(sessionRunsTable.id, input.sessionRunId), + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + inArray(sessionRunsTable.status, ["cancelled", "completed", "expired", "failed"]), ), ) .limit(1) - .get()) ?? null; + .get(); + if (terminalRun !== undefined) { + const terminalOperationId = input.sessionRunId as unknown as RuntimeOperationId; + if (input.operationId !== terminalOperationId) { + const handedOff = await getAppDatabase(bindings.DB) + .update(driverInstancesTable) + .set({ statusOperationId: terminalOperationId }) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.driverGeneration), + inArray(driverInstancesTable.status, ["stopped", "failed"]), + eq(driverInstancesTable.statusOperationId, input.operationId), + ), + ) + .returning({ id: driverInstancesTable.id }) + .get(); + if (handedOff === undefined) { + throw new Error("Driver stop lost ownership before terminal cleanup handoff."); + } + } + const { releaseTerminalDriverInstanceSessionRun } = + await import("./driver-instance/terminal-run-release"); + await releaseTerminalDriverInstanceSessionRun(bindings, { + driverGeneration: input.driverGeneration, + driverInstanceId: input.driverInstanceId, + sessionRunId: input.sessionRunId, + }); + return; + } - return row?.id ?? null; + const outcome = await recordRuntimeRunLeaseReleasedOutcome(bindings.DB, { + driverInstanceId: input.driverInstanceId, + expectedDriverGeneration: input.driverGeneration, + expectedDriverOperationId: input.operationId, + expectedSessionRunId: input.sessionRunId, + }); + + if (outcome.status === "applied") { + return; + } + + logWarn("runtime.driver_stop.lease_release_skipped", { + driverInstanceId: input.driverInstanceId, + reason: "reason" in outcome ? outcome.reason : outcome.status, + sessionRunId: input.sessionRunId, + status: outcome.status, + }); + throw new Error("Driver stop could not release its exact Session Run lease."); + } + + const db = getAppDatabase(bindings.DB); + const activeRun = db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ); + const released = await db + .update(driverInstancesTable) + .set({ statusOperationId: null }) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.driverGeneration), + eq(driverInstancesTable.statusOperationId, input.operationId), + notExists(activeRun), + ), + ) + .returning({ id: driverInstancesTable.id }) + .get(); + + if (released === undefined) { + throw new Error("Driver stop lost its exact operation ownership before release."); + } +} + +async function findCurrentTerminalDriverRunId( + database: D1Database, + driverInstanceId: DriverInstanceId, +): Promise { + const rows = await getAppDatabase(database) + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .innerJoin(sessionsTable, eq(sessionsTable.lastRunId, sessionRunsTable.id)) + .where( + and( + eq(sessionRunsTable.driverInstanceId, driverInstanceId), + inArray(sessionRunsTable.status, ["cancelled", "completed", "expired", "failed"]), + ), + ) + .limit(2) + .all(); + + return rows.length === 1 ? (rows[0]?.id ?? null) : null; } export async function stopDriverSession( bindings: ApiBindings, input: { driverInstanceId: DriverInstanceId; + expectedDriverGeneration?: number; + expectedSessionRunId?: SessionRunId; operationId?: RuntimeOperationId; - preserveSessionLifecycle?: boolean; reason: string; - terminalRun?: { - error?: RunError | null; - status: Extract; - }; }, ): Promise { - const driver = await getDriverInstanceRecord(bindings.DB, input.driverInstanceId); + const driver = await readDriverStopSnapshot(bindings.DB, input.driverInstanceId); if (!driver) { return; } - const activeDriver = driver; - const activeSessionRunId = await getActiveDriverSessionRunId(bindings.DB, input.driverInstanceId); + if ( + input.expectedDriverGeneration !== undefined && + driver.generation !== input.expectedDriverGeneration + ) { + return; + } - async function releaseLinkedRun(): Promise { - if (activeSessionRunId === null) { - return; - } + const observedActiveRunId = await getActiveDriverSessionRunId( + bindings.DB, + input.driverInstanceId, + ); + const terminalOwnedRunId = + input.expectedSessionRunId === undefined && + observedActiveRunId === null && + driver.statusOperationId !== null + ? await findCurrentTerminalDriverRunId(bindings.DB, input.driverInstanceId) + : null; + const expectedSessionRunId = + input.expectedSessionRunId ?? observedActiveRunId ?? terminalOwnedRunId; - if (input.terminalRun !== undefined) { - const outcome = await setSessionRunStatus(bindings.DB, { - error: input.terminalRun.error ?? null, - ...(input.operationId !== undefined ? { operationId: input.operationId } : {}), - preserveSessionLifecycle: input.preserveSessionLifecycle === true, - runId: activeSessionRunId, - source: "runtime_operation", - status: input.terminalRun.status, + if (driver.status === "stopped" || driver.status === "failed") { + if (driver.statusOperationId !== null) { + if (input.operationId !== undefined && input.operationId !== driver.statusOperationId) { + return; + } + await destroyDriverInstanceDurableObject( + bindings, + input.driverInstanceId, + driver.generation, + input.reason, + ); + await releaseDriverStopClaim(bindings, { + driverGeneration: driver.generation, + driverInstanceId: input.driverInstanceId, + operationId: driver.statusOperationId, + sessionRunId: expectedSessionRunId, }); - assertStoppedDriverRunTransition(outcome); + return; } - - const outcome = await recordRuntimeRunLeaseReleasedOutcome(bindings.DB, { - driverInstanceId: input.driverInstanceId, - expectedSessionRunId: activeSessionRunId, - }); - - if (outcome.status !== "applied") { - logWarn("runtime.driver_stop.lease_release_skipped", { + if (observedActiveRunId !== null) { + const outcome = await recordRuntimeRunLeaseReleasedOutcome(bindings.DB, { driverInstanceId: input.driverInstanceId, - reason: "reason" in outcome ? outcome.reason : outcome.status, - sessionRunId: activeSessionRunId, - status: outcome.status, + expectedDriverGeneration: driver.generation, + expectedSessionRunId: observedActiveRunId, }); + if (outcome.status !== "applied") { + throw new Error("Terminal Driver lost its exact Session Run lease before release."); + } } + return; } - if (activeDriver.status === "stopped" || activeDriver.status === "failed") { - await releaseLinkedRun(); - return; + const operationId = + driver.status === "stopping" && driver.statusOperationId !== null + ? driver.statusOperationId + : ((expectedSessionRunId as unknown as RuntimeOperationId | null) ?? + input.operationId ?? + createPlatformId()); + const claimed = await claimDriverStop(bindings.DB, { + driverGeneration: driver.generation, + driverInstanceId: input.driverInstanceId, + expectedSessionRunId, + operationId, + }); + if (!claimed) { + throw new Error("Driver stop lost its exact Driver and Session Run ownership."); } + let stopped = false; try { - if (activeDriver.status === "ready") { + if (driver.status === "ready") { try { - await sendDriverInstanceCommand(bindings, input.driverInstanceId, { + await sendDriverInstanceCommand(bindings, input.driverInstanceId, driver.generation, { commandId: createPlatformId(), kind: "session.stop", reason: input.reason, @@ -131,19 +367,105 @@ export async function stopDriverSession( await waitForDriverInstanceClose( bindings, input.driverInstanceId, + driver.generation, RUNTIME_SOCKET_TIMEOUT_MS, ); } catch (error) { if (!isDriverControlSocketMissingError(error)) { throw error; } + await failDriverInstance(bindings, input.driverInstanceId, driver.generation, input.reason); + await waitForDriverInstanceClose( + bindings, + input.driverInstanceId, + driver.generation, + RUNTIME_SOCKET_TIMEOUT_MS, + ); } - - return; + } else { + await failDriverInstance(bindings, input.driverInstanceId, driver.generation, input.reason); + await waitForDriverInstanceClose( + bindings, + input.driverInstanceId, + driver.generation, + RUNTIME_SOCKET_TIMEOUT_MS, + ); } - - await failDriverInstance(bindings, input.driverInstanceId, input.reason); + stopped = true; } finally { - await releaseLinkedRun(); + if (stopped) { + await releaseDriverStopClaim(bindings, { + driverGeneration: driver.generation, + driverInstanceId: input.driverInstanceId, + operationId, + sessionRunId: expectedSessionRunId, + }); + } + } +} + +export async function repairClaimedDriverStopsGlobally(bindings: ApiBindings): Promise { + const db = getAppDatabase(bindings.DB); + const terminalReleaseRun = db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + sql`${sessionRunsTable.id} = ${driverInstancesTable.statusOperationId}`, + eq(sessionRunsTable.driverInstanceId, driverInstancesTable.id), + inArray(sessionRunsTable.status, ["cancelled", "completed", "expired", "failed"]), + ), + ); + const claims = await db + .select({ + driverGeneration: driverInstancesTable.generation, + driverInstanceId: driverInstancesTable.id, + operationId: driverInstancesTable.statusOperationId, + sessionRunId: sessionRunsTable.id, + }) + .from(driverInstancesTable) + .leftJoin( + sessionRunsTable, + and( + sql`${sessionRunsTable.id} = ${driverInstancesTable.statusOperationId}`, + eq(sessionRunsTable.driverInstanceId, driverInstancesTable.id), + ), + ) + .where( + and( + isNotNull(driverInstancesTable.statusOperationId), + or( + eq(driverInstancesTable.status, "stopping"), + and( + inArray(driverInstancesTable.status, ["stopped", "failed"]), + notExists(terminalReleaseRun), + ), + ), + ), + ) + .all(); + + for (const claim of claims) { + if (claim.operationId === null) { + continue; + } + try { + const sessionRunId = + claim.sessionRunId ?? + (await findCurrentTerminalDriverRunId(bindings.DB, claim.driverInstanceId)); + await stopDriverSession(bindings, { + driverInstanceId: claim.driverInstanceId, + expectedDriverGeneration: claim.driverGeneration, + ...(sessionRunId === null ? {} : { expectedSessionRunId: sessionRunId }), + operationId: claim.operationId, + reason: "runtime.driver_stop.repair", + }); + } catch (error) { + logWarn("runtime.driver_stop.repair_failed", { + driverInstanceId: claim.driverInstanceId, + error: error instanceof Error ? error.message : "Driver stop repair failed.", + operationId: claim.operationId, + }); + } } } diff --git a/apps/api/src/modules/runtime/infrastructure/driver-session.service.ts b/apps/api/src/modules/runtime/infrastructure/driver-session.service.ts index c9b64e3f..b3372a38 100644 --- a/apps/api/src/modules/runtime/infrastructure/driver-session.service.ts +++ b/apps/api/src/modules/runtime/infrastructure/driver-session.service.ts @@ -38,6 +38,7 @@ import { provisionSessionDriver, } from "./runtime-sandbox-provisioner"; import { stopProvisionProcess } from "./runtime-sandbox-provisioning/runtime-driver-process-cleanup"; +import type { RuntimeRunProvisioningLease } from "./runtime-subject-lifecycle/runtime-provisioning-lease-store"; import type { RuntimeRunLeaseTransitionOutcome } from "./runtime-subject-lifecycle/runtime-run-lease-store"; import { createRuntimeSubjectLifecycleService } from "./runtime-subject-lifecycle/runtime-subject-lifecycle.service"; import type { @@ -53,6 +54,7 @@ async function allocateDriverInstanceId( database: D1Database, input: { sandboxId: SandboxId; + sandboxIncarnation: number; sandboxSessionId: SessionId; }, ): Promise { @@ -97,6 +99,7 @@ async function waitForRetryableRunLeaseOutcome( } async function releasePreparedRunLeaseAfterFailure(input: { + driverGeneration: number; driverInstanceId: DriverInstanceId; runtimeSubjectLifecycle: ReturnType; sessionId: SessionId; @@ -106,6 +109,7 @@ async function releasePreparedRunLeaseAfterFailure(input: { try { const released = await input.runtimeSubjectLifecycle.releaseRunLease({ driverInstanceId: input.driverInstanceId, + expectedDriverGeneration: input.driverGeneration, expectedSessionRunId: input.sessionRunId, }); @@ -138,7 +142,9 @@ export async function ensureDriverSessionReady( resolvedMcpServers: DriverResolvedMcpServer[]; resolvedSkillCatalog: DriverSkillCatalogEntry[]; resolvedSkills: Omit[]; + runtimeProvisioningLease: RuntimeRunProvisioningLease; sandbox: SandboxHandle; + sandboxIncarnation: number; sandboxSessionId: SessionId; sessionId: SessionId; sessionRunId: SessionRunId; @@ -146,12 +152,14 @@ export async function ensureDriverSessionReady( onBootPayloadPrepared?: DriverBootPayloadPreparedHandler; }, ): Promise<{ + driverGeneration: number; driverInstanceId: DriverInstanceId; readiness(): Promise; timing: RuntimeTimingSnapshot; }> { let driverInstanceId = await allocateDriverInstanceId(bindings.DB, { sandboxId: input.profile.sandbox.id, + sandboxIncarnation: input.sandboxIncarnation, sandboxSessionId: input.sandboxSessionId, }); const timing = createRuntimeTimingRecorder({ @@ -173,7 +181,11 @@ export async function ensureDriverSessionReady( traceId: input.traceId, }; const usage = await timing.measure("driver.getUsage", () => - getDriverUsage(bindings.DB, driverInstanceId), + getDriverUsage(bindings.DB, driverInstanceId, { + sandboxId: input.profile.sandbox.id, + sandboxIncarnation: input.sandboxIncarnation, + sandboxSessionId: input.sandboxSessionId, + }), ); const usageSessionRunId = usage?.sessionRunId ?? null; @@ -200,9 +212,15 @@ export async function ensureDriverSessionReady( }); try { - await failDriverInstance(bindings, driverInstanceId, DRIVER_SOCKET_MISSING_MESSAGE); + await failDriverInstance( + bindings, + driverInstanceId, + usage.generation, + DRIVER_SOCKET_MISSING_MESSAGE, + ); await runtimeSubjectLifecycle.releaseRunLease({ driverInstanceId, + expectedDriverGeneration: usage.generation, expectedSessionRunId: input.sessionRunId, }); } catch (error) { @@ -219,8 +237,10 @@ export async function ensureDriverSessionReady( const runLeaseOutcome = await timing.measure("driver.bindRun", () => runtimeSubjectLifecycle.acquireRunLease({ + driverGeneration: usage.generation, driverInstanceId, runtimeSubjectId: input.profile.sandbox.id, + runtimeSubjectIncarnation: input.sandboxIncarnation, sessionId: input.sessionId, sessionRunId: input.sessionRunId, }), @@ -239,6 +259,7 @@ export async function ensureDriverSessionReady( const readyTiming = timing.snapshot({ path: "warm" }); return { + driverGeneration: usage.generation, driverInstanceId, readiness: async () => readyTiming, timing: readyTiming, @@ -248,8 +269,10 @@ export async function ensureDriverSessionReady( if (usage && (usage.status === "provisioning" || usage.status === "connecting")) { const runLeaseOutcome = await timing.measure("driver.bindProvisioningRun", () => runtimeSubjectLifecycle.acquireRunLease({ + driverGeneration: usage.generation, driverInstanceId, runtimeSubjectId: input.profile.sandbox.id, + runtimeSubjectIncarnation: input.sandboxIncarnation, sessionId: input.sessionId, sessionRunId: input.sessionRunId, }), @@ -261,12 +284,14 @@ export async function ensureDriverSessionReady( } return { + driverGeneration: usage.generation, driverInstanceId, readiness: async () => { try { await timing.measure("driver.waitProvisioningReady", () => waitForDriverReady(bindings, { driverInstanceId, + driverGeneration: usage.generation, eventContext, logContext: { driverInstanceId, @@ -284,6 +309,7 @@ export async function ensureDriverSessionReady( eventContext, }); await releasePreparedRunLeaseAfterFailure({ + driverGeneration: usage.generation, driverInstanceId, runtimeSubjectLifecycle, sessionId: input.sessionId, @@ -314,13 +340,16 @@ export async function ensureDriverSessionReady( resolvedSkillCatalog: input.resolvedSkillCatalog, resolvedSkills: input.resolvedSkills, runtime: input.profile.runtimeId, + runtimeProvisioningLease: input.runtimeProvisioningLease, sandbox: input.sandbox, + sandboxIncarnation: input.sandboxIncarnation, sandboxSessionId: input.sandboxSessionId, sessionRunId: input.sessionRunId, traceId: input.traceId, }; const { onBootPayloadPrepared } = input; let provisionProcess: RuntimeProcessHandle | null = null; + let provisionDriverGeneration: number | null = null; let reconnectFailureEventContext = eventContext; try { @@ -339,6 +368,7 @@ export async function ensureDriverSessionReady( timing.addPhase(`driver.provision.${phase.name}`, phase.durationMs); } provisionProcess = provision.process; + provisionDriverGeneration = provision.driverGeneration; const provisionEventContext = { ...eventContext, driverInstanceId: provision.driverInstanceId, @@ -347,8 +377,10 @@ export async function ensureDriverSessionReady( const provisioningRunLeaseOutcome = await timing.measure("driver.bindProvisioningRun", () => runtimeSubjectLifecycle.acquireRunLease({ + driverGeneration: provision.driverGeneration, driverInstanceId: provision.driverInstanceId, runtimeSubjectId: input.profile.sandbox.id, + runtimeSubjectIncarnation: input.sandboxIncarnation, sessionId: input.sessionId, sessionRunId: input.sessionRunId, }), @@ -363,12 +395,14 @@ export async function ensureDriverSessionReady( provisionProcess = null; return { + driverGeneration: provision.driverGeneration, driverInstanceId: provision.driverInstanceId, readiness: async () => { try { await timing.measure("driver.waitForReady", () => waitForDriverReady(bindings, { driverInstanceId: provision.driverInstanceId, + driverGeneration: provision.driverGeneration, eventContext: provisionEventContext, logContext: { driverInstanceId: provision.driverInstanceId, @@ -382,6 +416,7 @@ export async function ensureDriverSessionReady( ); } catch (error) { await releasePreparedRunLeaseAfterFailure({ + driverGeneration: provision.driverGeneration, driverInstanceId: provision.driverInstanceId, runtimeSubjectLifecycle, sessionId: input.sessionId, @@ -421,18 +456,22 @@ export async function ensureDriverSessionReady( }); driverInstanceId = await allocateDriverInstanceId(bindings.DB, { sandboxId: input.profile.sandbox.id, + sandboxIncarnation: input.sandboxIncarnation, sandboxSessionId: input.sandboxSessionId, }); reconnectAttempt = null; continue; } - await releasePreparedRunLeaseAfterFailure({ - driverInstanceId, - runtimeSubjectLifecycle, - sessionId: input.sessionId, - sessionRunId: input.sessionRunId, - traceId: input.traceId, - }); + if (provisionDriverGeneration !== null) { + await releasePreparedRunLeaseAfterFailure({ + driverGeneration: provisionDriverGeneration, + driverInstanceId, + runtimeSubjectLifecycle, + sessionId: input.sessionId, + sessionRunId: input.sessionRunId, + traceId: input.traceId, + }); + } await appendDriverSocketReconnectFailedIfNeeded(bindings, { attempt: reconnectAttempt, error, @@ -456,6 +495,7 @@ export async function prewarmDriverSession( resolvedSkillCatalog: DriverSkillCatalogEntry[]; resolvedSkills: Omit[]; sandbox: SandboxHandle; + sandboxIncarnation: number; sandboxSessionId: SessionId; sessionId: SessionId; }, @@ -465,6 +505,7 @@ export async function prewarmDriverSession( } | null> { let driverInstanceId = await allocateDriverInstanceId(bindings.DB, { sandboxId: input.profile.sandbox.id, + sandboxIncarnation: input.sandboxIncarnation, sandboxSessionId: input.sandboxSessionId, }); const timing = createRuntimeTimingRecorder({ @@ -485,7 +526,11 @@ export async function prewarmDriverSession( traceId: null, }; const usage = await timing.measure("driver.getUsage", () => - getDriverUsage(bindings.DB, driverInstanceId), + getDriverUsage(bindings.DB, driverInstanceId, { + sandboxId: input.profile.sandbox.id, + sandboxIncarnation: input.sandboxIncarnation, + sandboxSessionId: input.sandboxSessionId, + }), ); if ((usage?.sessionRunId ?? null) !== null) { @@ -506,7 +551,12 @@ export async function prewarmDriverSession( return { driverInstanceId, timing: timing.snapshot({ path: "warm" }) }; } - await failDriverInstance(bindings, driverInstanceId, DRIVER_SOCKET_MISSING_MESSAGE); + await failDriverInstance( + bindings, + driverInstanceId, + usage.generation, + DRIVER_SOCKET_MISSING_MESSAGE, + ); continue; } @@ -514,6 +564,7 @@ export async function prewarmDriverSession( await timing.measure("driver.waitProvisioningReady", () => waitForDriverReady(bindings, { driverInstanceId, + driverGeneration: usage.generation, eventContext, logContext: { driverInstanceId, @@ -539,6 +590,7 @@ export async function prewarmDriverSession( resolvedSkills: input.resolvedSkills, runtime: input.profile.runtimeId, sandbox: input.sandbox, + sandboxIncarnation: input.sandboxIncarnation, sandboxSessionId: input.sandboxSessionId, sessionRunId: null, traceId: null, @@ -557,6 +609,7 @@ export async function prewarmDriverSession( await timing.measure("driver.waitForReady", () => waitForDriverReady(bindings, { driverInstanceId: provision.driverInstanceId, + driverGeneration: provision.driverGeneration, eventContext: { ...eventContext, driverInstanceId: provision.driverInstanceId, @@ -631,6 +684,7 @@ export async function dispatchDriverTurn( bindings: ApiBindings, input: { attachmentIds: FileId[]; + driverGeneration: number; driverInstanceId: DriverInstanceId; prompt: string; sessionRunId: SessionRunId; @@ -647,6 +701,7 @@ export async function dispatchDriverTurn( requestId: createPlatformId(), runId: input.sessionRunId, }, + driverGeneration: input.driverGeneration, driverInstanceId: input.driverInstanceId, expiresAt: currentTimestampPlus(DRIVER_COLD_READY_TIMEOUT_MS), }); diff --git a/apps/api/src/modules/runtime/infrastructure/execution-plane/sandbox-execution-plane-adapter.ts b/apps/api/src/modules/runtime/infrastructure/execution-plane/sandbox-execution-plane-adapter.ts index 7a087f27..5ebddfb5 100644 --- a/apps/api/src/modules/runtime/infrastructure/execution-plane/sandbox-execution-plane-adapter.ts +++ b/apps/api/src/modules/runtime/infrastructure/execution-plane/sandbox-execution-plane-adapter.ts @@ -23,6 +23,19 @@ import { } from "../../application/runtime-diagnostic-events"; import { createRuntimeTimingRecorder } from "../../application/session-runs/session-runtime-timing"; import { dispatchDriverTurn, ensureDriverSessionReady } from "../driver-session.service"; +import { + cleanupRuntimeProvisioningResources, + retireRuntimeProvisioningIncarnation, +} from "../runtime-subject-lifecycle/runtime-provisioning-cleanup.service"; +import { + claimRuntimeRunProvisioningLease, + heartbeatRuntimeRunProvisioningLease, + recordRuntimeProvisioningSandboxIncarnation, + releaseAbortedRuntimeProvisioningLease, + releaseReadyRuntimeRunProvisioningLease, + renewRuntimeProvisioningLeaseOwnership, +} from "../runtime-subject-lifecycle/runtime-provisioning-lease-store"; +import type { RuntimeRunProvisioningLease } from "../runtime-subject-lifecycle/runtime-provisioning-lease-store"; import { createRuntimeSubjectLifecycleService, getRuntimeSubjectKeepAliveHandle, @@ -33,7 +46,10 @@ import { resetRuntimeSubjectAgentState, stopRuntimeSubjectDrivers, } from "../runtime-subject-lifecycle/runtime-subject-operations.service"; -import { getRuntimeConversationSession } from "../runtime-subject-lifecycle/runtime-subject-store"; +import { + ensureRuntimeSubjectId, + getRuntimeConversationSession, +} from "../runtime-subject-lifecycle/runtime-subject-store"; import type { ExecutionSessionHandle, SandboxHandle } from "../sandbox-handles"; import { ensureSandboxConversationSession } from "../sandbox-session.service"; import { ensureSessionResourcesMounted } from "../session-resources/session-resource-mount.service"; @@ -52,6 +68,8 @@ type TerminalSessionHandle = ExecutionSessionHandle & { terminal(request: Request, options?: PtyOptions): Promise; }; +const RUNTIME_PROVISIONING_HEARTBEAT_MS = 60_000; + function isSessionAlreadyExistsError(error: unknown): boolean { if (!(error instanceof Error)) { return false; @@ -103,7 +121,9 @@ export async function connectPreparedSandboxTerminal( ): Promise { if (input.terminalSessionId) { const terminalSession = await ensureTerminalSession(subject, input.terminalSessionId); - return terminalSession.terminal(input.request, input.options); + return withDisposedRpcResource(terminalSession, (session) => + session.terminal(input.request, input.options), + ); } return subject.terminal(input.request, input.options); @@ -123,6 +143,9 @@ class SandboxExecutionPlaneAdapter implements RuntimeExecutionPlaneAdapter { }); const sandboxProvisioningTimer = createStopwatch(); let sandboxProvisioned = false; + let provisioningLease: RuntimeRunProvisioningLease | null = null; + let provisioningLeaseReleased = false; + let provisioningHeartbeat: ReturnType | null = null; const handles: { executionSession: ExecutionSessionHandle | null; subject: SandboxHandle | null; @@ -144,6 +167,43 @@ class SandboxExecutionPlaneAdapter implements RuntimeExecutionPlaneAdapter { traceId: input.traceId, }); const runtimeSubjectLifecycle = createRuntimeSubjectLifecycleService(bindings); + const ensuredSandboxId = await ensureRuntimeSubjectId(bindings.DB, { + agentId: input.profile.agentId, + appId: input.profile.vendorCredential.appId, + executionOwnerUserId: input.profile.session.origin.executionOwnerUserId, + kind: input.profile.kind, + runtimeSubjectId: sandboxId, + subjectId: input.profile.sandbox.subjectId, + subjectKind: input.profile.sandbox.subjectKind, + }); + if (ensuredSandboxId !== sandboxId) { + throw new Error("Session Run provisioning resolved a different runtime subject."); + } + provisioningLease = await claimRuntimeRunProvisioningLease(bindings.DB, { + runId: input.sessionRunId, + sandboxId, + sessionId: input.sessionId, + }); + if (provisioningLease === null) { + throw new Error("Session Run provisioning could not acquire lifecycle ownership."); + } + const activationProvisioningLease = provisioningLease; + provisioningHeartbeat = setInterval(() => { + if (provisioningLease === null || provisioningLeaseReleased) { + return; + } + void heartbeatRuntimeRunProvisioningLease(bindings.DB, provisioningLease).catch( + () => undefined, + ); + }, RUNTIME_PROVISIONING_HEARTBEAT_MS); + const heartbeatProvisioning = async (): Promise => { + if ( + provisioningLease === null || + !(await heartbeatRuntimeRunProvisioningLease(bindings.DB, provisioningLease)) + ) { + throw new Error("Session Run provisioning lost lifecycle ownership."); + } + }; pendingDiagnostics = appendRuntimeDiagnosticEvent(bindings, { eventName: RUNTIME_DIAGNOSTIC_EVENT.sandboxProvisioningStarted.name, sessionId: input.sessionId, @@ -152,31 +212,46 @@ class SandboxExecutionPlaneAdapter implements RuntimeExecutionPlaneAdapter { sandboxId, }, }); - const { subject: sandbox } = await timing.measure("activateRuntimeSubject", () => - runtimeSubjectLifecycle.activate({ - agentId: input.profile.agentId, - diagnosticContext: { - agentId: input.profile.configRevision.agentId, - sessionId: input.sessionId, - traceId: input.traceId, - }, - executionOwnerUserId: input.profile.session.origin.executionOwnerUserId, - kind: input.profile.kind, - networkConstraints: resolveRuntimeSubjectNetworkConstraints(bindings, { - envVars: input.profile.envVars, + const { incarnation: sandboxIncarnation, subject: sandbox } = await timing.measure( + "activateRuntimeSubject", + () => + runtimeSubjectLifecycle.activate({ + agentId: input.profile.agentId, + diagnosticContext: { + agentId: input.profile.configRevision.agentId, + sessionId: input.sessionId, + traceId: input.traceId, + }, + executionOwnerUserId: input.profile.session.origin.executionOwnerUserId, kind: input.profile.kind, - network: input.profile.network, - requestUrl, + networkConstraints: resolveRuntimeSubjectNetworkConstraints(bindings, { + envVars: input.profile.envVars, + kind: input.profile.kind, + network: input.profile.network, + requestUrl, + subjectKind: input.profile.sandbox.subjectKind, + }), + runtimeSubjectId: sandboxId, + appId: input.profile.vendorCredential.appId, + subjectId: input.profile.sandbox.subjectId, subjectKind: input.profile.sandbox.subjectKind, + timing, + provisioningAuthority: { + operationId: activationProvisioningLease.operationId, + runId: activationProvisioningLease.runId, + sessionId: activationProvisioningLease.sessionId, + }, }), - runtimeSubjectId: sandboxId, - appId: input.profile.vendorCredential.appId, - subjectId: input.profile.sandbox.subjectId, - subjectKind: input.profile.sandbox.subjectKind, - timing, - }), ); handles.subject = sandbox; + provisioningLease = await recordRuntimeProvisioningSandboxIncarnation(bindings.DB, { + lease: provisioningLease, + sandboxIncarnation, + }); + if (provisioningLease === null) { + throw new Error("Session Run provisioning lost ownership before sandbox handoff."); + } + await heartbeatProvisioning(); const provisioningCompletedValue = { ...runtimeBase, coldStartMs: sandboxProvisioningTimer.elapsedMs(), @@ -198,13 +273,22 @@ class SandboxExecutionPlaneAdapter implements RuntimeExecutionPlaneAdapter { kind: input.profile.kind, mountSessionResources: input.attachmentIds.length > 0, origin: input.profile.session.origin, + ...(provisioningLease === null ? {} : { provisioningLease }), + replaceClosedExecutionSession: true, sandbox, sandboxId, + sandboxIncarnation, sessionId: input.sessionId, timing, }), ); handles.executionSession = executionSession.cloudflareSession; + provisioningLease = executionSession.provisioningLease ?? provisioningLease; + await heartbeatProvisioning(); + const driverProvisioningLease = provisioningLease; + if (driverProvisioningLease === null) { + throw new Error("Session Run provisioning lost ownership before Driver handoff."); + } const driverProfile = { ...input.profile, @@ -227,7 +311,9 @@ class SandboxExecutionPlaneAdapter implements RuntimeExecutionPlaneAdapter { resolvedMcpServers: input.resolvedMcpServers, resolvedSkillCatalog: input.resolvedSkillCatalog, resolvedSkills: input.resolvedSkills, + runtimeProvisioningLease: driverProvisioningLease, sandbox, + sandboxIncarnation, sandboxSessionId: input.sessionId, sessionId: input.sessionId, sessionRunId: input.sessionRunId, @@ -238,10 +324,19 @@ class SandboxExecutionPlaneAdapter implements RuntimeExecutionPlaneAdapter { timing.addPhase(phase.name, phase.durationMs); } const initialDriverPhaseCount = driver.timing.phases.length; + provisioningLeaseReleased = await releaseReadyRuntimeRunProvisioningLease(bindings.DB, { + driverGeneration: driver.driverGeneration, + driverInstanceId: driver.driverInstanceId, + lease: provisioningLease, + }); + if (!provisioningLeaseReleased) { + throw new Error("Session Run provisioning could not complete its durable handoff."); + } await pendingDiagnostics; return { + driverGeneration: driver.driverGeneration, driverInstanceId: driver.driverInstanceId, readiness: async () => { const driverTiming = await driver.readiness(); @@ -271,7 +366,29 @@ class SandboxExecutionPlaneAdapter implements RuntimeExecutionPlaneAdapter { }); } releaseRunResources(handles); + if (provisioningLease !== null && !provisioningLeaseReleased) { + const stillOwnsLease = await renewRuntimeProvisioningLeaseOwnership( + bindings.DB, + provisioningLease, + ); + if (stillOwnsLease) { + try { + if (provisioningLease.sandboxIncarnation === null) { + await cleanupRuntimeProvisioningResources(bindings, provisioningLease, "api"); + await releaseAbortedRuntimeProvisioningLease(bindings.DB, provisioningLease); + } else { + await retireRuntimeProvisioningIncarnation(bindings, provisioningLease, "api"); + } + } catch { + // Keep the durable lease so maintenance retries before cleanup can proceed. + } + } + } throw error; + } finally { + if (provisioningHeartbeat !== null) { + clearInterval(provisioningHeartbeat); + } } } @@ -290,7 +407,11 @@ class SandboxExecutionPlaneAdapter implements RuntimeExecutionPlaneAdapter { } await withDisposedRpcResource( - await getRuntimeSubjectKeepAliveHandle(bindings, sandboxSession.sandboxId), + await getRuntimeSubjectKeepAliveHandle( + bindings, + sandboxSession.sandboxId, + sandboxSession.sandboxIncarnation, + ), async (sandbox) => { await ensureSessionResourcesMounted({ bindings, diff --git a/apps/api/src/modules/runtime/infrastructure/native-resume-ref.repository.ts b/apps/api/src/modules/runtime/infrastructure/native-resume-ref.repository.ts index 0c7a307a..f55ca712 100644 --- a/apps/api/src/modules/runtime/infrastructure/native-resume-ref.repository.ts +++ b/apps/api/src/modules/runtime/infrastructure/native-resume-ref.repository.ts @@ -8,11 +8,12 @@ import { parseDriverNativeRuntimeRef, } from "@mosoo/agent-driver/runtime"; import { nativeResumeRefsTable, sessionsTable } from "@mosoo/db"; -import type { DriverInstanceId, SessionId, SessionRunId } from "@mosoo/id"; +import type { DriverInstanceId, RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; import { eq, inArray, sql } from "drizzle-orm"; import { getAppDatabase } from "../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../time"; +import { RUNTIME_KIND_POLICIES } from "../domain/runtime-kind-policy"; interface NativeResumeRefRow { committed_value: string | null; @@ -22,13 +23,141 @@ interface NativeResumeRefRow { value: string; } +interface ObservedNativeResumeRefRow { + kind: string; + observed_driver_instance_id: string | null; + observed_event_seq: number; + observed_session_run_id: string | null; + runtime_id: string; + value: string; +} + export interface NativeResumeRefObservation { driverInstanceId: DriverInstanceId; nativeResumeRef: DriverNativeRuntimeRef; + observedEventSeq: number; sessionId: SessionId; sessionRunId: SessionRunId; } +const PLATFORM_NATIVE_RESUME_KINDS = Object.values(RUNTIME_KIND_POLICIES) + .filter((policy) => policy.nativeResume.persistence === "platform") + .map((policy) => policy.kind); + +export function prepareNativeResumeRefProjection( + database: D1Database, + input: NativeResumeRefObservation & { + createdAt: number; + eventId: RuntimeEventId; + semanticHash: string; + }, +): D1PreparedStatement[] { + enforceNativeRuntimeRefShape(input.nativeResumeRef); + + if (!Number.isSafeInteger(input.observedEventSeq) || input.observedEventSeq < 0) { + throw new Error("Native resume ref event seq must be a non-negative safe integer."); + } + + const platformKinds = PLATFORM_NATIVE_RESUME_KINDS.map(() => "?").join(", "); + const eligibleReceipt = `EXISTS ( + SELECT 1 + FROM session_event AS receipt + JOIN session_run AS run + ON run.id = ? + AND run.session_id = receipt.session_id + AND run.driver_instance_id = ? + JOIN driver_instance AS driver + ON driver.id = ? + AND driver.sandbox_session_id = receipt.session_id + JOIN sandbox AS runtime_sandbox + ON runtime_sandbox.id = driver.sandbox_id + AND runtime_sandbox.kind IN (${platformKinds}) + WHERE receipt.id = ? + AND receipt.session_id = ? + AND receipt.event_type = 'runtime.resume.updated' + AND receipt.run_id = ? + AND receipt.semantic_hash = ? + AND receipt.seq = ? + )`; + const eligibilityBindings = [ + input.sessionRunId, + input.driverInstanceId, + input.driverInstanceId, + ...PLATFORM_NATIVE_RESUME_KINDS, + input.eventId, + input.sessionId, + input.sessionRunId, + input.semanticHash, + input.observedEventSeq, + ]; + const upsert = database + .prepare( + `INSERT INTO native_resume_ref ( + created_at, kind, observed_driver_instance_id, observed_event_seq, + observed_session_run_id, runtime_id, session_id, updated_at, value + ) + SELECT ?, ?, ?, ?, ?, ?, ?, ?, ? + WHERE ${eligibleReceipt} + ON CONFLICT (session_id) DO UPDATE SET + kind = excluded.kind, + observed_driver_instance_id = excluded.observed_driver_instance_id, + observed_event_seq = excluded.observed_event_seq, + observed_session_run_id = excluded.observed_session_run_id, + runtime_id = excluded.runtime_id, + updated_at = excluded.updated_at, + value = excluded.value + WHERE native_resume_ref.observed_event_seq < excluded.observed_event_seq`, + ) + .bind( + input.createdAt, + input.nativeResumeRef.kind, + input.driverInstanceId, + input.observedEventSeq, + input.sessionRunId, + input.nativeResumeRef.runtimeId, + input.sessionId, + input.createdAt, + input.nativeResumeRef.value, + ...eligibilityBindings, + ); + const guard = database + .prepare( + `INSERT INTO session_event (id) + SELECT ? + WHERE ${eligibleReceipt} + AND NOT EXISTS ( + SELECT 1 + FROM native_resume_ref AS stored + WHERE stored.session_id = ? + AND ( + stored.observed_event_seq > ? + OR ( + stored.observed_event_seq = ? + AND stored.kind = ? + AND stored.observed_driver_instance_id = ? + AND stored.observed_session_run_id = ? + AND stored.runtime_id = ? + AND stored.value = ? + ) + ) + )`, + ) + .bind( + input.eventId, + ...eligibilityBindings, + input.sessionId, + input.observedEventSeq, + input.observedEventSeq, + input.nativeResumeRef.kind, + input.driverInstanceId, + input.sessionRunId, + input.nativeResumeRef.runtimeId, + input.nativeResumeRef.value, + ); + + return [upsert, guard]; +} + function expectedNativeRuntimeRefKind( runtimeId: DriverNativeRuntimeRef["runtimeId"], ): DriverNativeRuntimeRefKind { @@ -113,14 +242,20 @@ export async function upsertNativeResumeRef( ): Promise { enforceNativeRuntimeRefShape(observation.nativeResumeRef); + if (!Number.isSafeInteger(observation.observedEventSeq) || observation.observedEventSeq < 0) { + throw new Error("Native resume ref event seq must be a non-negative safe integer."); + } + const timestampMs = currentTimestampMs(); + const appDatabase = getAppDatabase(database); - await getAppDatabase(database) + await appDatabase .insert(nativeResumeRefsTable) .values({ createdAt: timestampMs, kind: observation.nativeResumeRef.kind, observedDriverInstanceId: observation.driverInstanceId, + observedEventSeq: observation.observedEventSeq, observedSessionRunId: observation.sessionRunId, runtimeId: observation.nativeResumeRef.runtimeId, sessionId: observation.sessionId, @@ -131,12 +266,54 @@ export async function upsertNativeResumeRef( set: { kind: sql`excluded.kind`, observedDriverInstanceId: sql`excluded.observed_driver_instance_id`, + observedEventSeq: sql`excluded.observed_event_seq`, observedSessionRunId: sql`excluded.observed_session_run_id`, runtimeId: sql`excluded.runtime_id`, updatedAt: sql`excluded.updated_at`, value: sql`excluded.value`, }, + setWhere: sql`${nativeResumeRefsTable.observedEventSeq} < excluded.observed_event_seq`, target: nativeResumeRefsTable.sessionId, }) .run(); + + const stored = + (await appDatabase + .select({ + kind: nativeResumeRefsTable.kind, + observed_driver_instance_id: nativeResumeRefsTable.observedDriverInstanceId, + observed_event_seq: nativeResumeRefsTable.observedEventSeq, + observed_session_run_id: nativeResumeRefsTable.observedSessionRunId, + runtime_id: nativeResumeRefsTable.runtimeId, + value: nativeResumeRefsTable.value, + }) + .from(nativeResumeRefsTable) + .where(eq(nativeResumeRefsTable.sessionId, observation.sessionId)) + .limit(1) + .get()) ?? null; + + assertNativeResumeRefConverged(stored, observation); +} + +function assertNativeResumeRefConverged( + stored: ObservedNativeResumeRefRow | null, + observation: NativeResumeRefObservation, +): void { + if (stored === null || stored.observed_event_seq < observation.observedEventSeq) { + throw new Error("Native resume ref CAS did not persist the durable event."); + } + + if (stored.observed_event_seq > observation.observedEventSeq) { + return; + } + + if ( + stored.kind !== observation.nativeResumeRef.kind || + stored.observed_driver_instance_id !== observation.driverInstanceId || + stored.observed_session_run_id !== observation.sessionRunId || + stored.runtime_id !== observation.nativeResumeRef.runtimeId || + stored.value !== observation.nativeResumeRef.value + ) { + throw new Error("Native resume ref event seq was replayed with conflicting content."); + } } diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-files.service.ts b/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-files.service.ts index 097f9faa..5b09d57c 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-files.service.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-files.service.ts @@ -1,6 +1,7 @@ import { SANDBOX_CACHE_PATH, SANDBOX_MEMORY_PATH } from "@mosoo/agent-driver/paths"; import { disposeRpcResource } from "../../../../platform/cloudflare/rpc-disposal"; +import { quoteShellArg } from "../../../../shared/shell"; import type { DriverProfileConfig } from "../../domain/driver-snapshot"; import type { ExecutionSessionHandle } from "../sandbox-handles"; import { @@ -32,13 +33,9 @@ const RUNTIME_MEMORY_MOUNTS: RuntimeMemoryMountsByRuntime = { ], }; -function quoteShellArg(value: string): string { - return `'${value.replaceAll("'", `'"'"'`)}'`; -} - async function sha256(value: string): Promise { const digest = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(value)); - return [...new Uint8Array(digest)].map((byte) => byte.toString(16).padStart(2, "0")).join(""); + return new Uint8Array(digest).toHex(); } async function readJsonFile( diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-process-cleanup.ts b/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-process-cleanup.ts index 8ec2fcab..920949b7 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-process-cleanup.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-process-cleanup.ts @@ -1,6 +1,29 @@ +import { disposeRpcResource } from "../../../../platform/cloudflare/rpc-disposal"; import { runBestEffortRuntimeCleanup } from "../runtime-cleanup"; import type { RuntimeProcessHandle } from "../sandbox-handles"; +export async function startProvisionProcessWithOwnershipFence(input: { + assertOwned: () => Promise; + context: Record; + message: string; + startProcess: () => Promise; +}): Promise { + await input.assertOwned(); + const process = await input.startProcess(); + try { + await input.assertOwned(); + return process; + } catch (error) { + await stopProvisionProcess({ + context: input.context, + message: input.message, + process, + }); + disposeRpcResource(process); + throw error; + } +} + export async function stopProvisionProcess(input: { context: Record; message: string; diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-provisioning.service.ts b/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-provisioning.service.ts index 97e8d2fe..c44abe89 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-provisioning.service.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-provisioning.service.ts @@ -28,6 +28,7 @@ import { createDriverInstanceRecord, markDriverInstanceFailedIfBootTokenMatches, recordRuntimeProcessStarted, + runtimeProvisioningDriverLaunchIsOwned, } from "../driver-instance/driver-instance-record.repository"; import { relayDriverProcessLogs } from "../driver-process-log-relay"; import { getNativeResumeRefForRuntime } from "../native-resume-ref.repository"; @@ -44,7 +45,10 @@ import { getLostPrewarmOwnershipError, usesInsertOnlyDriverRecord, } from "./runtime-driver-prewarm-ownership"; -import { stopProvisionProcess } from "./runtime-driver-process-cleanup"; +import { + startProvisionProcessWithOwnershipFence, + stopProvisionProcess, +} from "./runtime-driver-process-cleanup"; import { appendRuntimeEnvironmentInstallFailed, createRuntimeEnvironmentInstallState, @@ -158,6 +162,7 @@ async function provisionDriver( const { driverInstanceId } = input; const processId = sanitizeProcessId(driverInstanceId); + const bootPayloadPath = `/tmp/.mosoo-driver-boot-${processId}.json`; const sandboxId = input.profile.sandbox.id; const runtimeEntry = getRuntimeCatalogEntry(input.runtime); @@ -202,7 +207,11 @@ async function provisionDriver( mcpGrants: input.resolvedMcpServers.map(toDriverInstanceMcpGrantRecord), conflictStrategy: input.driverRecordConflictStrategy ?? "replace", runtime: input.runtime, + ...(input.runtimeProvisioningLease + ? { runtimeProvisioningLease: input.runtimeProvisioningLease } + : {}), sandboxId, + sandboxIncarnation: input.sandboxIncarnation, sandboxSessionId: input.sandboxSessionId, }), ); @@ -306,7 +315,6 @@ async function provisionDriver( sandboxId, traceparent, }); - const bootPayloadPath = `${input.profile.session.homePath}/driver-boot-payload-${processId}.json`; const bootPayloadJson = JSON.stringify(bootPayload); const bootPayloadPreparedPromise = timing.measure("onBootPayloadPrepared", async () => { @@ -329,15 +337,34 @@ async function provisionDriver( await timing.measure("writeBootPayload", () => input.cloudflareSession.writeFile(bootPayloadPath, bootPayloadJson), ); + const assertLaunchOwned = async (): Promise => { + if ( + input.runtimeProvisioningLease !== undefined && + !(await runtimeProvisioningDriverLaunchIsOwned(env.DB, { + bootTokenHash: bootToken.hash, + driverGeneration: activeDriverGeneration, + driverInstanceId, + lease: input.runtimeProvisioningLease, + })) + ) { + throw new Error("Runtime Driver provisioning lost launch ownership."); + } + }; const startedProcess = await timing.measure("startProcess", () => - input.cloudflareSession.startProcess(AGENT_DRIVER_PROCESS_COMMAND, { - autoCleanup: true, - cwd: organizationPath, - env: { - [DRIVER_BOOT_PAYLOAD_FILE_ENV_NAME]: bootPayloadPath, - ...toRuntimeProcessProxyEnv(env, runtimeProfile.network.networkPolicy), - }, - processId, + startProvisionProcessWithOwnershipFence({ + assertOwned: assertLaunchOwned, + context: { driverInstanceId, sandboxId }, + message: "runtime.driver.provision.late_process_cleanup_failed", + startProcess: () => + input.cloudflareSession.startProcess(AGENT_DRIVER_PROCESS_COMMAND, { + autoCleanup: true, + cwd: organizationPath, + env: { + [DRIVER_BOOT_PAYLOAD_FILE_ENV_NAME]: bootPayloadPath, + ...toRuntimeProcessProxyEnv(env, runtimeProfile.network.networkPolicy), + }, + processId, + }), }), ); process = startedProcess; @@ -363,6 +390,7 @@ async function provisionDriver( if (staleError !== null) { throw staleError; } + throw new Error("Runtime Driver process start lost durable ownership."); } }); void processRecordPromise.catch(() => undefined); @@ -432,6 +460,10 @@ async function provisionDriver( message: "runtime.driver.provision.skipped_process_cleanup_failed", process, }); + await removeProvisionBootPayload(input.cloudflareSession, bootPayloadPath, { + driverInstanceId, + sandboxId, + }); disposeRpcResource(process); logInfo("runtime.driver.provision.skipped", { driverInstanceId, @@ -503,6 +535,10 @@ async function provisionDriver( message: "runtime.driver.provision.process_cleanup_failed", process, }); + await removeProvisionBootPayload(input.cloudflareSession, bootPayloadPath, { + driverInstanceId, + sandboxId, + }); disposeRpcResource(process); @@ -524,3 +560,17 @@ async function provisionDriver( throw error; } } + +async function removeProvisionBootPayload( + session: ProvisionDriverInput["cloudflareSession"], + path: string, + context: Record, +): Promise { + await runBestEffortRuntimeCleanup({ + context, + message: "runtime.driver.provision.boot_payload_cleanup_failed", + task: async () => { + await session.exec(`rm -f -- '${path}'`); + }, + }); +} diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-environment-artifact.ts b/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-environment-artifact.ts index dcc6a74b..f065aadd 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-environment-artifact.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-environment-artifact.ts @@ -2,15 +2,12 @@ import { discardPromiseResult } from "@mosoo/effects"; import { withDisposedRpcResult } from "../../../../platform/cloudflare/rpc-disposal"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; +import { quoteShellArg } from "../../../../shared/shell"; import type { DriverEnvironmentArtifactProfile } from "../../domain/driver-snapshot"; import { isRuntimeSandboxLocalBucketEnabled } from "../runtime-sandbox-bucket-mount"; import type { ExecutionSessionHandle, SandboxHandle } from "../sandbox-handles"; import { getParentDirectory } from "./runtime-sandbox-provisioning.paths"; -function quoteShellArg(value: string): string { - return `'${value.replaceAll("'", `'"'"'`)}'`; -} - export async function exposeEnvironmentNodeModules( session: Pick, input: { diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-sandbox-provisioning.types.ts b/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-sandbox-provisioning.types.ts index fb01ceeb..460844bb 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-sandbox-provisioning.types.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-sandbox-provisioning.types.ts @@ -11,6 +11,7 @@ import type { DriverRuntime, DriverSkillCatalogEntry, } from "../../domain/driver-snapshot"; +import type { RuntimeRunProvisioningLease } from "../runtime-subject-lifecycle/runtime-provisioning-lease-store"; import type { ExecutionSessionHandle, RuntimeProcessHandle, @@ -40,7 +41,9 @@ export interface ProvisionDriverInput { resolvedSkillCatalog: DriverSkillCatalogEntry[]; resolvedSkills: Omit[]; runtime: DriverRuntime; + runtimeProvisioningLease?: RuntimeRunProvisioningLease; sandbox: SandboxHandle; + sandboxIncarnation: number; sandboxSessionId: SessionId; sessionRunId?: SessionRunId | null; traceId?: string | null; diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/lease-ownership-renewal.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/lease-ownership-renewal.ts new file mode 100644 index 00000000..18f2e4b4 --- /dev/null +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/lease-ownership-renewal.ts @@ -0,0 +1,34 @@ +export function createLeaseOwnershipRenewal( + renew: () => Promise, + lostMessage: string, +): () => Promise { + let ownershipLoss: Error | null = null; + let pending: Promise | null = null; + + return () => { + if (pending !== null) { + return pending; + } + if (ownershipLoss !== null) { + return Promise.reject(ownershipLoss); + } + + pending = renew() + .then( + (renewed) => { + if (!renewed) { + ownershipLoss = new Error(lostMessage); + throw ownershipLoss; + } + }, + (cause) => { + ownershipLoss = new Error(lostMessage, { cause }); + throw ownershipLoss; + }, + ) + .finally(() => { + pending = null; + }); + return pending; + }; +} diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-conversation-session-store.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-conversation-session-store.ts index 8fa7c722..81f1f7fe 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-conversation-session-store.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-conversation-session-store.ts @@ -7,15 +7,19 @@ import { sessionsTable, } from "@mosoo/db"; import { createPlatformId } from "@mosoo/id"; -import type { SandboxId, SandboxSessionId, SessionId } from "@mosoo/id"; +import type { RuntimeOperationId, SandboxId, SandboxSessionId, SessionId } from "@mosoo/id"; import { and, desc, eq, exists, inArray, isNull, lte, notExists, or, sql } from "drizzle-orm"; -import { getAppDatabase, runAppDatabaseBatch } from "../../../../platform/db/drizzle"; +import { + getAppDatabase, + getD1ChangeCount, + runAppDatabaseBatch, +} from "../../../../platform/db/drizzle"; import { getRuntimeKindPolicy, getRuntimeSubjectInactiveDeadline, } from "../../domain/runtime-kind-policy"; -import { toRuntimeSubjectStatusLifecycleEventName } from "../../domain/runtime-subject-lifecycle.machine"; +import { ACTIVE_SESSION_RUN_STATUSES } from "../../domain/session-run-lifecycle.machine"; import { isCattleTerminalCheckpointReadyForNextRun } from "../session-runs/session-run-admission.repository"; import { activeConversationSessionQuery, @@ -25,6 +29,7 @@ import { runLeaseQueryForListedSubject, } from "./runtime-subject-store-queries"; import type { + PendingRuntimeConversationSessionCleanup, RuntimeConversationSessionRecord, RuntimeConversationSessionState, } from "./runtime-subject-store.types"; @@ -37,6 +42,7 @@ export async function getRuntimeConversationSession( (await getAppDatabase(database) .select({ sandboxSessionId: sandboxSessionsTable.sandboxSessionId, + sandboxIncarnation: sandboxSessionsTable.sandboxIncarnation, cwd: sandboxSessionsTable.cwd, latestReadyBackupDir: readyConversationBackupTable.dir, latestReadyBackupId: readyConversationBackupTable.id, @@ -52,6 +58,7 @@ export async function getRuntimeConversationSession( and( eq(readyConversationBackupTable.sandboxId, sandboxSessionsTable.sandboxId), eq(readyConversationBackupTable.dir, sandboxSessionsTable.cwd), + eq(readyConversationBackupTable.workspaceSessionId, sandboxSessionsTable.sessionId), eq(readyConversationBackupTable.status, "ready"), ), ) @@ -66,6 +73,7 @@ export async function getRuntimeConversationSession( return { sandboxSessionId: row.sandboxSessionId, + sandboxIncarnation: row.sandboxIncarnation, cwd: row.cwd, latestReadyBackup: mapReadyRuntimeSubjectBackup({ dir: row.latestReadyBackupDir, @@ -81,7 +89,10 @@ export async function getRuntimeConversationSession( export async function getRuntimeConversationSessionState( database: D1Database, input: { + readonly expectedProvisioningOperationId?: RuntimeOperationId; + readonly expectedSandboxSessionId?: SandboxSessionId; readonly runtimeSubjectId: SandboxId; + readonly expectedSandboxIncarnation?: number; readonly sessionId: SessionId; }, ): Promise { @@ -89,7 +100,9 @@ export async function getRuntimeConversationSessionState( (await getAppDatabase(database) .select({ agentId: sessionsTable.agentId, + cleanupOperationId: sandboxSessionsTable.cleanupOperationId, sandboxSessionId: sandboxSessionsTable.sandboxSessionId, + sandboxIncarnation: sandboxSessionsTable.sandboxIncarnation, kind: sandboxesTable.kind, status: sandboxSessionsTable.status, }) @@ -100,6 +113,20 @@ export async function getRuntimeConversationSessionState( and( eq(sandboxSessionsTable.sessionId, input.sessionId), eq(sandboxSessionsTable.sandboxId, input.runtimeSubjectId), + ...(input.expectedSandboxSessionId === undefined + ? [] + : [eq(sandboxSessionsTable.sandboxSessionId, input.expectedSandboxSessionId)]), + ...(input.expectedSandboxIncarnation === undefined + ? [] + : [eq(sandboxSessionsTable.sandboxIncarnation, input.expectedSandboxIncarnation)]), + ...(input.expectedProvisioningOperationId === undefined + ? [] + : [ + eq( + sessionsTable.runtimeProvisioningOperationId, + input.expectedProvisioningOperationId, + ), + ]), ), ) .limit(1) @@ -132,9 +159,22 @@ export async function listIdleSessionScopedConversationSessions( .where( and( eq(sandboxSessionsTable.status, "active"), + isNull(sandboxSessionsTable.cleanupOperationId), + isNull(sessionsTable.runtimeProvisioningOperationId), eq(sandboxesTable.kind, "cattle"), sql`${sandboxSessionsTable.updatedAt} <= ${input.idleSinceLte}`, notExists(runLeaseQueryForListedSubject(appDb)), + notExists( + appDb + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.sessionId, sandboxSessionsTable.sessionId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ), + ), or( eq(sessionsTable.workspaceCheckpointRequired, false), isNull(sessionsTable.lastRunId), @@ -185,32 +225,158 @@ export async function claimIdleSessionScopedConversationForClose( readonly idleSinceLte: number; readonly now: number; readonly runtimeSubjectId: SandboxId; + readonly sandboxIncarnation: number; readonly sandboxSessionId: SandboxSessionId; readonly sessionId: SessionId; }, -): Promise { +): Promise { if (!(await isCattleTerminalCheckpointReadyForNextRun(database, input.sessionId))) { - return false; + return null; } const appDb = getAppDatabase(database); + const operationId = createPlatformId(); const claimed = await appDb .update(sandboxSessionsTable) - .set({ status: "closed", updatedAt: input.now }) + .set({ cleanupOperationId: operationId, status: "cleanup_pending", updatedAt: input.now }) .where( and( eq(sandboxSessionsTable.sessionId, input.sessionId), eq(sandboxSessionsTable.sandboxId, input.runtimeSubjectId), + eq(sandboxSessionsTable.sandboxIncarnation, input.sandboxIncarnation), eq(sandboxSessionsTable.sandboxSessionId, input.sandboxSessionId), - eq(sandboxSessionsTable.status, "active"), + inArray(sandboxSessionsTable.status, ["active", "error"]), + isNull(sandboxSessionsTable.cleanupOperationId), lte(sandboxSessionsTable.updatedAt, input.idleSinceLte), notExists(runLeaseQuery(appDb, input.runtimeSubjectId)), + notExists( + appDb + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.sessionId, input.sessionId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ), + ), + exists( + appDb + .select({ id: sessionsTable.id }) + .from(sessionsTable) + .where( + and( + eq(sessionsTable.id, input.sessionId), + isNull(sessionsTable.runtimeProvisioningOperationId), + ), + ), + ), ), ) .returning({ sessionId: sandboxSessionsTable.sessionId }) .get(); - return claimed != null; + return claimed == null ? null : operationId; +} + +export async function claimRuntimeConversationSessionCleanup( + database: D1Database, + input: { + readonly expectedProvisioningOperationId?: RuntimeOperationId; + readonly now: number; + readonly runtimeSubjectId: SandboxId; + readonly sandboxIncarnation: number; + readonly sandboxSessionId: SandboxSessionId; + readonly sessionId: SessionId; + }, +): Promise { + const appDb = getAppDatabase(database); + const operationId = createPlatformId(); + const claimed = await appDb + .update(sandboxSessionsTable) + .set({ cleanupOperationId: operationId, status: "cleanup_pending", updatedAt: input.now }) + .where( + and( + eq(sandboxSessionsTable.sessionId, input.sessionId), + eq(sandboxSessionsTable.sandboxId, input.runtimeSubjectId), + eq(sandboxSessionsTable.sandboxIncarnation, input.sandboxIncarnation), + inArray(sandboxSessionsTable.status, ["active", "closed", "error"]), + isNull(sandboxSessionsTable.cleanupOperationId), + eq(sandboxSessionsTable.sandboxSessionId, input.sandboxSessionId), + ...(input.expectedProvisioningOperationId === undefined + ? [] + : [ + exists( + appDb + .select({ id: sessionsTable.id }) + .from(sessionsTable) + .where( + and( + eq(sessionsTable.id, input.sessionId), + eq( + sessionsTable.runtimeProvisioningOperationId, + input.expectedProvisioningOperationId, + ), + eq(sessionsTable.runtimeProvisioningSandboxId, input.runtimeSubjectId), + eq(sessionsTable.runtimeProvisioningSandboxSessionId, input.sandboxSessionId), + eq( + sessionsTable.runtimeProvisioningSandboxIncarnation, + input.sandboxIncarnation, + ), + ), + ), + ), + ]), + ), + ) + .returning({ id: sandboxSessionsTable.sessionId }) + .get(); + + return claimed === undefined ? null : operationId; +} + +export async function listPendingRuntimeConversationSessionCleanups( + database: D1Database, + limit: number, +): Promise { + return getAppDatabase(database) + .select({ + agentId: sessionsTable.agentId, + cleanupOperationId: sandboxSessionsTable.cleanupOperationId, + kind: sandboxesTable.kind, + sandboxId: sandboxSessionsTable.sandboxId, + sandboxSessionId: sandboxSessionsTable.sandboxSessionId, + sandboxIncarnation: sandboxSessionsTable.sandboxIncarnation, + sessionId: sandboxSessionsTable.sessionId, + status: sandboxSessionsTable.status, + }) + .from(sandboxSessionsTable) + .innerJoin(sessionsTable, eq(sessionsTable.id, sandboxSessionsTable.sessionId)) + .innerJoin(sandboxesTable, eq(sandboxesTable.id, sandboxSessionsTable.sandboxId)) + .where(eq(sandboxSessionsTable.status, "cleanup_pending")) + .limit(limit) + .all() as Promise; +} + +export async function retireRuntimeConversationSessionsForIncarnation( + database: D1Database, + input: { + readonly now: number; + readonly runtimeSubjectId: SandboxId; + readonly sandboxIncarnation: number; + }, +): Promise { + await getAppDatabase(database) + .update(sandboxSessionsTable) + .set({ cleanupOperationId: null, status: "closed", updatedAt: input.now }) + .where( + and( + eq(sandboxSessionsTable.sandboxId, input.runtimeSubjectId), + eq(sandboxSessionsTable.sandboxIncarnation, input.sandboxIncarnation), + inArray(sandboxSessionsTable.status, ["active", "cleanup_pending", "error"]), + ), + ) + .run(); } export async function ensureRuntimeConversationSessionRecord( @@ -220,6 +386,7 @@ export async function ensureRuntimeConversationSessionRecord( readonly now: number; readonly originJson: string; readonly runtimeSubjectId: SandboxId; + readonly sandboxIncarnation: number; readonly sessionId: SessionId; }, ): Promise { @@ -230,6 +397,10 @@ export async function ensureRuntimeConversationSessionRecord( throw new Error("Sandbox session is already bound to a different sandbox."); } + if (existing.status !== "closed" && existing.sandboxIncarnation !== input.sandboxIncarnation) { + throw new Error("Sandbox session belongs to a retired sandbox incarnation."); + } + return existing; } @@ -241,6 +412,7 @@ export async function ensureRuntimeConversationSessionRecord( cwd: input.cwd, originJson: input.originJson, sandboxId: input.runtimeSubjectId, + sandboxIncarnation: input.sandboxIncarnation, sessionId: input.sessionId, status: "closed", updatedAt: input.now, @@ -258,6 +430,10 @@ export async function ensureRuntimeConversationSessionRecord( throw new Error("Sandbox session is already bound to a different sandbox."); } + if (created.status !== "closed" && created.sandboxIncarnation !== input.sandboxIncarnation) { + throw new Error("Sandbox session belongs to a retired sandbox incarnation."); + } + return created; } @@ -265,95 +441,165 @@ export async function recordRuntimeConversationSessionError( database: D1Database, input: { readonly sandboxSessionId: SandboxSessionId; + readonly sandboxIncarnation: number; readonly cwd: string; readonly message: string; readonly errorCode: RuntimeSubjectErrorCode; readonly now: number; readonly originJson: string; + readonly expectedProvisioningOperationId?: RuntimeOperationId; readonly runtimeSubjectId: SandboxId; readonly sessionId: SessionId; }, -): Promise { - await runAppDatabaseBatch(database, (appDb) => [ - appDb - .insert(sandboxSessionsTable) - .values({ - sandboxSessionId: input.sandboxSessionId, - createdAt: input.now, - cwd: input.cwd, - originJson: input.originJson, - sandboxId: input.runtimeSubjectId, - sessionId: input.sessionId, - status: "error", - updatedAt: input.now, - }) - .onConflictDoUpdate({ - set: { - status: "error", - updatedAt: sql`excluded.updated_at`, - }, - target: sandboxSessionsTable.sessionId, - }), - appDb - .update(sandboxesTable) - .set({ - lastError: input.message, - lastErrorCode: input.errorCode, - status: "cold", - statusChangedAt: input.now, - statusEvent: toRuntimeSubjectStatusLifecycleEventName("cold"), - statusOperationId: null, - statusSeq: sql`${sandboxesTable.statusSeq} + 1`, - statusSource: "runtime", - updatedAt: input.now, - }) - .where( - and( - eq(sandboxesTable.id, input.runtimeSubjectId), - inArray(sandboxesTable.status, ["restoring", "active"]), +): Promise { + const appDb = getAppDatabase(database); + const updated = await appDb + .update(sandboxSessionsTable) + .set({ + cleanupOperationId: null, + sandboxSessionId: input.sandboxSessionId, + sandboxIncarnation: input.sandboxIncarnation, + status: "error", + updatedAt: input.now, + }) + .where( + and( + eq(sandboxSessionsTable.sessionId, input.sessionId), + eq(sandboxSessionsTable.sandboxId, input.runtimeSubjectId), + or( + and( + eq(sandboxSessionsTable.status, "active"), + eq(sandboxSessionsTable.sandboxIncarnation, input.sandboxIncarnation), + ), + inArray(sandboxSessionsTable.status, ["closed", "error"]), ), + isNull(sandboxSessionsTable.cleanupOperationId), + exists( + appDb + .select({ id: sandboxesTable.id }) + .from(sandboxesTable) + .where( + and( + eq(sandboxesTable.id, input.runtimeSubjectId), + eq(sandboxesTable.incarnation, input.sandboxIncarnation), + eq(sandboxesTable.status, "active"), + ), + ), + ), + ...(input.expectedProvisioningOperationId === undefined + ? [] + : [ + exists( + appDb + .select({ id: sessionsTable.id }) + .from(sessionsTable) + .where( + and( + eq(sessionsTable.id, input.sessionId), + eq( + sessionsTable.runtimeProvisioningOperationId, + input.expectedProvisioningOperationId, + ), + eq(sessionsTable.runtimeProvisioningSandboxId, input.runtimeSubjectId), + eq(sessionsTable.runtimeProvisioningSandboxSessionId, input.sandboxSessionId), + eq( + sessionsTable.runtimeProvisioningSandboxIncarnation, + input.sandboxIncarnation, + ), + ), + ), + ), + ]), ), - ]); + ) + .returning({ id: sandboxSessionsTable.sessionId }) + .get(); + + return updated !== undefined; } export async function recordRuntimeConversationSessionActive( database: D1Database, input: { readonly sandboxSessionId: SandboxSessionId; + readonly sandboxIncarnation: number; readonly cwd: string; readonly now: number; readonly originJson: string; + readonly expectedProvisioningOperationId?: RuntimeOperationId; readonly runtimeSubjectId: SandboxId; readonly sessionId: SessionId; }, -): Promise { +): Promise { const petInactiveDeadlineAt = getRuntimeSubjectInactiveDeadline( getRuntimeKindPolicy("pet"), input.now, ); - await runAppDatabaseBatch(database, (appDb) => [ + const results = await runAppDatabaseBatch(database, (appDb) => [ appDb - .insert(sandboxSessionsTable) - .values({ - sandboxSessionId: input.sandboxSessionId, - createdAt: input.now, + .update(sandboxSessionsTable) + .set({ + cleanupOperationId: null, cwd: input.cwd, - originJson: input.originJson, - sandboxId: input.runtimeSubjectId, - sessionId: input.sessionId, + sandboxSessionId: input.sandboxSessionId, + sandboxIncarnation: input.sandboxIncarnation, status: "active", updatedAt: input.now, }) - .onConflictDoUpdate({ - set: { - sandboxSessionId: sql`excluded.cloudflare_session_id`, - cwd: sql`excluded.cwd`, - status: "active", - updatedAt: sql`excluded.updated_at`, - }, - target: sandboxSessionsTable.sessionId, - }), + .where( + and( + eq(sandboxSessionsTable.sessionId, input.sessionId), + eq(sandboxSessionsTable.sandboxId, input.runtimeSubjectId), + or( + and( + eq(sandboxSessionsTable.status, "active"), + eq(sandboxSessionsTable.sandboxIncarnation, input.sandboxIncarnation), + ), + inArray(sandboxSessionsTable.status, ["closed", "error"]), + ), + isNull(sandboxSessionsTable.cleanupOperationId), + exists( + appDb + .select({ id: sandboxesTable.id }) + .from(sandboxesTable) + .where( + and( + eq(sandboxesTable.id, input.runtimeSubjectId), + eq(sandboxesTable.incarnation, input.sandboxIncarnation), + eq(sandboxesTable.status, "active"), + ), + ), + ), + ...(input.expectedProvisioningOperationId === undefined + ? [] + : [ + exists( + appDb + .select({ id: sessionsTable.id }) + .from(sessionsTable) + .where( + and( + eq(sessionsTable.id, input.sessionId), + eq( + sessionsTable.runtimeProvisioningOperationId, + input.expectedProvisioningOperationId, + ), + eq(sessionsTable.runtimeProvisioningSandboxId, input.runtimeSubjectId), + eq( + sessionsTable.runtimeProvisioningSandboxSessionId, + input.sandboxSessionId, + ), + eq( + sessionsTable.runtimeProvisioningSandboxIncarnation, + input.sandboxIncarnation, + ), + ), + ), + ), + ]), + ), + ), appDb .update(sandboxesTable) .set({ @@ -366,27 +612,88 @@ export async function recordRuntimeConversationSessionActive( `, updatedAt: input.now, }) - .where(eq(sandboxesTable.id, input.runtimeSubjectId)), + .where( + and( + eq(sandboxesTable.id, input.runtimeSubjectId), + eq(sandboxesTable.incarnation, input.sandboxIncarnation), + eq(sandboxesTable.status, "active"), + exists( + appDb + .select({ id: sandboxSessionsTable.sessionId }) + .from(sandboxSessionsTable) + .where( + and( + eq(sandboxSessionsTable.sessionId, input.sessionId), + eq(sandboxSessionsTable.sandboxSessionId, input.sandboxSessionId), + eq(sandboxSessionsTable.sandboxIncarnation, input.sandboxIncarnation), + eq(sandboxSessionsTable.status, "active"), + ), + ), + ), + ), + ), ]); + + return getD1ChangeCount(results[0]) === 1; } export async function recordRuntimeConversationSessionClosed( database: D1Database, input: { + readonly expectedProvisioningOperationId?: RuntimeOperationId; + readonly cleanupOperationId: RuntimeOperationId; readonly inactiveDeadlineAt: number | null; readonly now: number; readonly runtimeSubjectId: SandboxId; + readonly sandboxIncarnation: number; + readonly sandboxSessionId: SandboxSessionId; readonly sessionId: SessionId; }, -): Promise { - await runAppDatabaseBatch(database, (appDb) => [ +): Promise { + const results = await runAppDatabaseBatch(database, (appDb) => [ appDb .update(sandboxSessionsTable) .set({ status: "closed", + cleanupOperationId: null, updatedAt: input.now, }) - .where(eq(sandboxSessionsTable.sessionId, input.sessionId)), + .where( + and( + eq(sandboxSessionsTable.sessionId, input.sessionId), + eq(sandboxSessionsTable.sandboxSessionId, input.sandboxSessionId), + eq(sandboxSessionsTable.sandboxIncarnation, input.sandboxIncarnation), + eq(sandboxSessionsTable.status, "cleanup_pending"), + eq(sandboxSessionsTable.cleanupOperationId, input.cleanupOperationId), + ...(input.expectedProvisioningOperationId === undefined + ? [] + : [ + exists( + appDb + .select({ id: sessionsTable.id }) + .from(sessionsTable) + .where( + and( + eq(sessionsTable.id, input.sessionId), + eq( + sessionsTable.runtimeProvisioningOperationId, + input.expectedProvisioningOperationId, + ), + eq(sessionsTable.runtimeProvisioningSandboxId, input.runtimeSubjectId), + eq( + sessionsTable.runtimeProvisioningSandboxSessionId, + input.sandboxSessionId, + ), + eq( + sessionsTable.runtimeProvisioningSandboxIncarnation, + input.sandboxIncarnation, + ), + ), + ), + ), + ]), + ), + ), appDb .update(sandboxesTable) .set({ @@ -401,4 +708,6 @@ export async function recordRuntimeConversationSessionClosed( ), ), ]); + + return getD1ChangeCount(results[0]) === 1; } diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-provisioning-cleanup.service.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-provisioning-cleanup.service.ts new file mode 100644 index 00000000..7e1fee07 --- /dev/null +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-provisioning-cleanup.service.ts @@ -0,0 +1,160 @@ +import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; +import { LIVE_DRIVER_INSTANCE_STATUSES } from "../../domain/driver-instance-lifecycle.machine"; +import { destroyDriverInstanceDurableObject } from "../driver-instance/client"; +import { stopDriverSession } from "../driver-session-stop.service"; +import { closeSandboxConversationSession } from "../sandbox-session.service"; +import { deleteActiveSandboxConversationSession } from "../sandbox-session/sandbox-conversation-session-delete"; +import { createLeaseOwnershipRenewal } from "./lease-ownership-renewal"; +import type { RuntimeProvisioningLease } from "./runtime-provisioning-lease-store"; +import { + claimRuntimeProvisioningSubjectRetirement, + claimRuntimeProvisioningDriverCleanup, + readRuntimeProvisioningCleanupTargets, + releaseAbortedRuntimeProvisioningLease, + renewRuntimeProvisioningLeaseOwnership, +} from "./runtime-provisioning-lease-store"; +import { runRuntimeSubjectOperation } from "./runtime-subject-operations.service"; +import { getRuntimeSubject } from "./runtime-subject-store"; + +const PROVISIONING_CLEANUP_HEARTBEAT_MS = 60_000; + +export async function retireRuntimeProvisioningIncarnation( + bindings: ApiBindings, + lease: RuntimeProvisioningLease, + source: "api" | "maintenance", +): Promise<"released" | "stale" | "waiting"> { + if (lease.sandboxIncarnation === null) { + return "stale"; + } + + const retirement = await claimRuntimeProvisioningSubjectRetirement(bindings.DB, { + lease, + source, + }); + if (retirement.kind === "stale") { + return "stale"; + } + if (retirement.kind === "waiting" || retirement.kind === "repairing") { + return "waiting"; + } + if (retirement.kind === "destroying") { + const subject = await getRuntimeSubject(bindings.DB, lease.sandboxId); + if (subject === null) { + throw new Error("Runtime provisioning retirement lost its subject."); + } + await runRuntimeSubjectOperation(bindings, { + kind: subject.kind, + lease: retirement.lease, + reason: "runtime.provisioning_ambiguous", + runtimeSubjectId: lease.sandboxId, + }); + } + + return (await releaseAbortedRuntimeProvisioningLease(bindings.DB, lease)) ? "released" : "stale"; +} + +export async function cleanupRuntimeProvisioningResources( + bindings: ApiBindings, + lease: RuntimeProvisioningLease, + source: "api" | "maintenance", +): Promise { + const requireOwnership = createLeaseOwnershipRenewal( + () => renewRuntimeProvisioningLeaseOwnership(bindings.DB, lease), + "Runtime provisioning cleanup lost its lease ownership.", + ); + const heartbeat = setInterval(() => { + void requireOwnership().catch(() => undefined); + }, PROVISIONING_CLEANUP_HEARTBEAT_MS); + + try { + await requireOwnership(); + const targets = await readRuntimeProvisioningCleanupTargets(bindings.DB, lease); + if (targets === null) { + throw new Error("Runtime provisioning cleanup lost its lease ownership."); + } + + const failures: unknown[] = []; + const liveDrivers = targets.driverInstances.filter((driver) => + LIVE_DRIVER_INSTANCE_STATUSES.some((status) => status === driver.status), + ); + const driverClaims = await Promise.allSettled( + liveDrivers.map(async (driver) => ({ + claimed: await claimRuntimeProvisioningDriverCleanup(bindings.DB, { + driverGeneration: driver.generation, + driverInstanceId: driver.id, + lease, + source, + }), + driver, + })), + ); + failures.push( + ...driverClaims.flatMap((result) => (result.status === "rejected" ? [result.reason] : [])), + ); + const claimedDrivers = driverClaims.flatMap((result) => + result.status === "fulfilled" && result.value.claimed ? [result.value.driver] : [], + ); + await requireOwnership(); + const stopResults = await Promise.allSettled( + claimedDrivers.map((driver) => + stopDriverSession(bindings, { + driverInstanceId: driver.id, + expectedDriverGeneration: driver.generation, + reason: "runtime.provisioning_stale", + }), + ), + ); + failures.push( + ...stopResults.flatMap((result) => (result.status === "rejected" ? [result.reason] : [])), + ); + + await requireOwnership(); + if (targets.conversationSessionId !== null) { + if (lease.sandboxIncarnation === null) { + throw new Error("Runtime provisioning cleanup target has no sandbox incarnation."); + } + try { + await deleteActiveSandboxConversationSession(bindings, { + sandboxId: lease.sandboxId, + sandboxIncarnation: lease.sandboxIncarnation, + sandboxSessionId: targets.conversationSessionId, + }); + await closeSandboxConversationSession(bindings, { + expectedProvisioningOperationId: lease.operationId, + expectedSandboxSessionId: targets.conversationSessionId, + sandboxId: lease.sandboxId, + sessionId: lease.sessionId, + }); + } catch (error) { + failures.push(error); + } + } + + await requireOwnership(); + const liveDriverIds = new Set(liveDrivers.map((driver) => driver.id)); + const destructionTargets = [ + ...targets.driverInstances.filter((driver) => !liveDriverIds.has(driver.id)), + ...claimedDrivers, + ]; + const destroyResults = await Promise.allSettled( + destructionTargets.map((driver) => + destroyDriverInstanceDurableObject( + bindings, + driver.id, + driver.generation, + "runtime.provisioning_stale", + ), + ), + ); + failures.push( + ...destroyResults.flatMap((result) => (result.status === "rejected" ? [result.reason] : [])), + ); + + await requireOwnership(); + if (failures[0] !== undefined) { + throw failures[0]; + } + } finally { + clearInterval(heartbeat); + } +} diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-provisioning-lease-store.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-provisioning-lease-store.ts new file mode 100644 index 00000000..26db433a --- /dev/null +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-provisioning-lease-store.ts @@ -0,0 +1,841 @@ +import { + driverInstancesTable, + sandboxesTable, + sandboxSessionsTable, + sessionRunsTable, + sessionsTable, +} from "@mosoo/db"; +import { createPlatformId } from "@mosoo/id"; +import type { + DriverInstanceId, + RuntimeOperationId, + SandboxId, + SandboxSessionId, + SessionId, + SessionRunId, +} from "@mosoo/id"; +import { + and, + eq, + exists, + inArray, + isNotNull, + isNull, + lte, + ne, + notExists, + or, + sql, +} from "drizzle-orm"; +import { alias } from "drizzle-orm/sqlite-core"; + +import { getAppDatabase } from "../../../../platform/db/drizzle"; +import { currentTimestampMs } from "../../../../time"; +import { + ASSIGNABLE_DRIVER_INSTANCE_STATUSES, + toDriverInstanceStatusLifecycleEventName, +} from "../../domain/driver-instance-lifecycle.machine"; +import { toRuntimeSubjectStatusLifecycleEventName } from "../../domain/runtime-subject-lifecycle.machine"; +import { ACTIVE_SESSION_RUN_STATUSES } from "../../domain/session-run-lifecycle.machine"; +import type { RuntimeSubjectOperationLease } from "./runtime-subject-store.types"; + +export interface RuntimeProvisioningLease { + readonly heartbeatAt: number; + readonly operationId: RuntimeOperationId; + readonly runId: SessionRunId | null; + readonly sandboxId: SandboxId; + readonly sandboxIncarnation: number | null; + readonly sandboxSessionId: SandboxSessionId | null; + readonly sessionId: SessionId; +} + +export interface RuntimeRunProvisioningLease extends RuntimeProvisioningLease { + readonly runId: SessionRunId; +} + +export interface RuntimeProvisioningCleanupTargets { + readonly conversationSessionId: SandboxSessionId | null; + readonly driverInstances: readonly { + readonly generation: number; + readonly id: DriverInstanceId; + readonly status: (typeof driverInstancesTable.$inferSelect)["status"]; + }[]; +} + +export type RuntimeProvisioningSubjectRetirementOutcome = + | { readonly kind: "cold" } + | { readonly kind: "destroying"; readonly lease: RuntimeSubjectOperationLease } + | { readonly kind: "repairing" } + | { readonly kind: "stale" } + | { readonly kind: "waiting" }; + +const RUNTIME_PROVISIONING_RETIRE_CLAIM_TTL_MS = 30 * 60_000; + +function runtimeProvisioningRetireClaimOwner(sessionId: SessionId): string { + return `runtime-provisioning-retire:${sessionId}`; +} + +const otherProvisioningSessionsTable = alias(sessionsTable, "other_runtime_provisioning_session"); +const otherProvisioningRunsTable = alias(sessionRunsTable, "other_runtime_provisioning_run"); +const otherProvisioningDriversTable = alias( + driverInstancesTable, + "other_runtime_provisioning_driver", +); +const retirementOwnerSessionsTable = alias(sessionsTable, "runtime_retirement_owner_session"); + +/** + * The durable subject status closes admission. This predicate only decides + * whether the current worker may advance from draining to physical teardown. + */ +export async function runtimeSubjectActivationRetirementIsDrained( + database: D1Database, + input: { + readonly lease: RuntimeSubjectOperationLease; + readonly runtimeSubjectId: SandboxId; + }, +): Promise { + if (input.lease.kind !== "activate" || input.lease.status !== "destroying") { + return true; + } + + const db = getAppDatabase(database); + const otherProvisioning = db + .select({ id: otherProvisioningSessionsTable.id }) + .from(otherProvisioningSessionsTable) + .where( + and( + eq(otherProvisioningSessionsTable.runtimeProvisioningSandboxId, input.runtimeSubjectId), + isNotNull(otherProvisioningSessionsTable.runtimeProvisioningOperationId), + ne(otherProvisioningSessionsTable.runtimeProvisioningOperationId, input.lease.operationId), + ), + ); + const ownedRun = db + .select({ id: retirementOwnerSessionsTable.id }) + .from(retirementOwnerSessionsTable) + .where( + and( + eq(retirementOwnerSessionsTable.runtimeProvisioningOperationId, input.lease.operationId), + eq(retirementOwnerSessionsTable.runtimeProvisioningSandboxId, input.runtimeSubjectId), + eq(retirementOwnerSessionsTable.runtimeProvisioningRunId, otherProvisioningRunsTable.id), + ), + ); + const otherActiveRun = db + .select({ id: otherProvisioningRunsTable.id }) + .from(otherProvisioningRunsTable) + .innerJoin( + otherProvisioningDriversTable, + eq(otherProvisioningDriversTable.id, otherProvisioningRunsTable.driverInstanceId), + ) + .where( + and( + eq(otherProvisioningDriversTable.sandboxId, input.runtimeSubjectId), + eq(otherProvisioningDriversTable.sandboxIncarnation, input.lease.incarnation), + inArray(otherProvisioningRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + notExists(ownedRun), + ), + ); + const drained = await db + .select({ id: sandboxesTable.id }) + .from(sandboxesTable) + .where( + and( + eq(sandboxesTable.id, input.runtimeSubjectId), + eq(sandboxesTable.incarnation, input.lease.incarnation), + eq(sandboxesTable.status, "destroying"), + eq(sandboxesTable.operationKind, "activate"), + eq(sandboxesTable.statusOperationId, input.lease.operationId), + eq(sandboxesTable.claimOwner, input.lease.claimOwner), + notExists(otherProvisioning), + notExists(otherActiveRun), + ), + ) + .limit(1) + .get(); + + return drained !== undefined; +} + +export async function claimRuntimeProvisioningDriverCleanup( + database: D1Database, + input: { + readonly driverGeneration: number; + readonly driverInstanceId: DriverInstanceId; + readonly lease: RuntimeProvisioningLease; + readonly source: "api" | "maintenance"; + }, +): Promise { + if (input.lease.sandboxIncarnation === null) { + return false; + } + const db = getAppDatabase(database); + const ownedLease = db + .select({ id: sessionsTable.id }) + .from(sessionsTable) + .where(runtimeProvisioningLeaseCondition(input.lease)); + const now = currentTimestampMs(); + const claimed = await db + .update(driverInstancesTable) + .set({ + status: "stopping", + statusChangedAt: now, + statusEvent: toDriverInstanceStatusLifecycleEventName("stopping"), + statusOperationId: input.lease.operationId, + statusSeq: sql`${driverInstancesTable.statusSeq} + 1`, + statusSource: input.source, + updatedAt: now, + }) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.driverGeneration), + eq(driverInstancesTable.sandboxId, input.lease.sandboxId), + eq(driverInstancesTable.sandboxIncarnation, input.lease.sandboxIncarnation), + eq(driverInstancesTable.sandboxSessionId, input.lease.sessionId), + inArray(driverInstancesTable.status, ASSIGNABLE_DRIVER_INSTANCE_STATUSES), + isNull(driverInstancesTable.statusOperationId), + exists(ownedLease), + ), + ) + .returning({ id: driverInstancesTable.id }) + .get(); + if (claimed !== undefined) { + return true; + } + + const adopted = await db + .select({ id: driverInstancesTable.id }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.driverGeneration), + eq(driverInstancesTable.sandboxId, input.lease.sandboxId), + eq(driverInstancesTable.sandboxIncarnation, input.lease.sandboxIncarnation), + eq(driverInstancesTable.sandboxSessionId, input.lease.sessionId), + eq(driverInstancesTable.status, "stopping"), + eq(driverInstancesTable.statusOperationId, input.lease.operationId), + exists(ownedLease), + ), + ) + .limit(1) + .get(); + + return adopted !== undefined; +} + +export async function claimRuntimeRunProvisioningLease( + database: D1Database, + input: { + readonly runId: SessionRunId; + readonly sandboxId: SandboxId; + readonly sessionId: SessionId; + }, +): Promise { + const heartbeatAt = currentTimestampMs(); + const operationId = createPlatformId(); + const db = getAppDatabase(database); + const activeRun = db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, input.runId), + eq(sessionRunsTable.sessionId, input.sessionId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ); + const availableSubject = db + .select({ id: sandboxesTable.id }) + .from(sandboxesTable) + .where( + and( + eq(sandboxesTable.id, input.sandboxId), + inArray(sandboxesTable.status, ["active", "cold"]), + isNull(sandboxesTable.operationKind), + isNull(sandboxesTable.statusOperationId), + or( + isNull(sandboxesTable.claimOwner), + isNull(sandboxesTable.claimExpiresAt), + lte(sandboxesTable.claimExpiresAt, heartbeatAt), + ), + ), + ); + const otherProvisioning = db + .select({ id: otherProvisioningSessionsTable.id }) + .from(otherProvisioningSessionsTable) + .where( + and( + ne(otherProvisioningSessionsTable.id, input.sessionId), + eq(otherProvisioningSessionsTable.runtimeProvisioningSandboxId, input.sandboxId), + isNotNull(otherProvisioningSessionsTable.runtimeProvisioningOperationId), + ), + ); + const claimed = await db + .update(sessionsTable) + .set({ + runtimeProvisioningHeartbeatAt: heartbeatAt, + runtimeProvisioningOperationId: operationId, + runtimeProvisioningRunId: input.runId, + runtimeProvisioningSandboxId: input.sandboxId, + runtimeProvisioningSandboxIncarnation: null, + runtimeProvisioningSandboxSessionId: null, + }) + .where( + and( + eq(sessionsTable.id, input.sessionId), + eq(sessionsTable.lastRunId, input.runId), + eq(sessionsTable.status, "RUNNING"), + isNull(sessionsTable.archivedAt), + isNull(sessionsTable.cleanupOperationKind), + isNull(sessionsTable.statusOperationId), + isNull(sessionsTable.runtimeProvisioningOperationId), + exists(activeRun), + exists(availableSubject), + notExists(otherProvisioning), + ), + ) + .returning({ id: sessionsTable.id }) + .get(); + + return claimed === undefined + ? null + : { + heartbeatAt, + operationId, + runId: input.runId, + sandboxId: input.sandboxId, + sandboxIncarnation: null, + sandboxSessionId: null, + sessionId: input.sessionId, + }; +} + +export async function recordRuntimeProvisioningConversationTarget( + database: D1Database, + input: { + readonly lease: RuntimeRunProvisioningLease; + readonly sandboxIncarnation: number; + readonly sandboxSessionId: SandboxSessionId; + }, +): Promise { + if ( + (input.lease.sandboxIncarnation !== null && + input.lease.sandboxIncarnation !== input.sandboxIncarnation) || + (input.lease.sandboxSessionId !== null && + input.lease.sandboxSessionId !== input.sandboxSessionId) + ) { + throw new Error("Runtime provisioning cannot change its immutable conversation target."); + } + + if (input.lease.sandboxSessionId === input.sandboxSessionId) { + return input.lease; + } + + const heartbeatAt = currentTimestampMs(); + const updated = await getAppDatabase(database) + .update(sessionsTable) + .set({ + runtimeProvisioningHeartbeatAt: heartbeatAt, + runtimeProvisioningSandboxIncarnation: input.sandboxIncarnation, + runtimeProvisioningSandboxSessionId: input.sandboxSessionId, + }) + .where(runtimeProvisioningLeaseCondition(input.lease)) + .returning({ id: sessionsTable.id }) + .get(); + + return updated === undefined + ? null + : { + ...input.lease, + heartbeatAt, + sandboxIncarnation: input.sandboxIncarnation, + sandboxSessionId: input.sandboxSessionId, + }; +} + +export async function recordRuntimeProvisioningSandboxIncarnation( + database: D1Database, + input: { + readonly lease: RuntimeRunProvisioningLease; + readonly sandboxIncarnation: number; + }, +): Promise { + if (!Number.isSafeInteger(input.sandboxIncarnation) || input.sandboxIncarnation < 0) { + throw new Error("Runtime provisioning sandbox incarnation must be non-negative."); + } + if ( + input.lease.sandboxIncarnation !== null && + input.lease.sandboxIncarnation !== input.sandboxIncarnation + ) { + throw new Error("Runtime provisioning cannot change its immutable sandbox incarnation."); + } + if (input.lease.sandboxIncarnation === input.sandboxIncarnation) { + return input.lease; + } + + const heartbeatAt = currentTimestampMs(); + const updated = await getAppDatabase(database) + .update(sessionsTable) + .set({ + runtimeProvisioningHeartbeatAt: heartbeatAt, + runtimeProvisioningSandboxIncarnation: input.sandboxIncarnation, + }) + .where(runtimeProvisioningLeaseCondition(input.lease)) + .returning({ id: sessionsTable.id }) + .get(); + return updated === undefined + ? null + : { + ...input.lease, + heartbeatAt, + sandboxIncarnation: input.sandboxIncarnation, + }; +} + +/** Poison N durably first; the expiring claim only elects its teardown worker. */ +export async function claimRuntimeProvisioningSubjectRetirement( + database: D1Database, + input: { + readonly lease: RuntimeProvisioningLease; + readonly source: "api" | "maintenance"; + }, +): Promise { + const { lease } = input; + if (lease.sandboxIncarnation === null) { + return { kind: "stale" }; + } + + const now = currentTimestampMs(); + const claimExpiresAt = now + RUNTIME_PROVISIONING_RETIRE_CLAIM_TTL_MS; + const claimOwner = runtimeProvisioningRetireClaimOwner(lease.sessionId); + const db = getAppDatabase(database); + const ownedLease = db + .select({ id: sessionsTable.id }) + .from(sessionsTable) + .where(runtimeProvisioningLeaseCondition(lease)); + + await db + .update(sandboxesTable) + .set({ + claimExpiresAt, + claimOwner, + lastError: "Runtime provisioning left an ambiguous physical mutation.", + lastErrorCode: "runtime.subject_activation_failed", + operationKind: "activate", + status: "destroying", + statusChangedAt: now, + statusEvent: toRuntimeSubjectStatusLifecycleEventName("destroying"), + statusOperationId: lease.operationId, + statusSeq: sql`${sandboxesTable.statusSeq} + 1`, + statusSource: input.source, + updatedAt: now, + }) + .where( + and( + eq(sandboxesTable.id, lease.sandboxId), + eq(sandboxesTable.incarnation, lease.sandboxIncarnation), + eq(sandboxesTable.status, "active"), + isNull(sandboxesTable.operationKind), + isNull(sandboxesTable.statusOperationId), + or( + isNull(sandboxesTable.claimOwner), + isNull(sandboxesTable.claimExpiresAt), + lte(sandboxesTable.claimExpiresAt, now), + eq(sandboxesTable.claimOwner, claimOwner), + ), + exists(ownedLease), + ), + ) + .run(); + + // A stale provisioning takeover rotates its operation id. Adopt the durable + // poison only when no different worker still owns the subject claim. + await db + .update(sandboxesTable) + .set({ + claimExpiresAt, + claimOwner, + statusOperationId: lease.operationId, + statusSource: input.source, + updatedAt: now, + }) + .where( + and( + eq(sandboxesTable.id, lease.sandboxId), + eq(sandboxesTable.incarnation, lease.sandboxIncarnation), + eq(sandboxesTable.status, "destroying"), + eq(sandboxesTable.operationKind, "activate"), + or( + eq(sandboxesTable.claimOwner, claimOwner), + isNull(sandboxesTable.claimOwner), + isNull(sandboxesTable.claimExpiresAt), + lte(sandboxesTable.claimExpiresAt, now), + ), + exists(ownedLease), + ), + ) + .run(); + + const state = await db + .select({ + claimExpiresAt: sandboxesTable.claimExpiresAt, + claimOwner: sandboxesTable.claimOwner, + incarnation: sandboxesTable.incarnation, + operationId: sandboxesTable.statusOperationId, + operationKind: sandboxesTable.operationKind, + status: sandboxesTable.status, + }) + .from(sessionsTable) + .innerJoin(sandboxesTable, eq(sandboxesTable.id, sessionsTable.runtimeProvisioningSandboxId)) + .where(runtimeProvisioningLeaseCondition(lease)) + .limit(1) + .get(); + if (state === undefined || state.incarnation !== lease.sandboxIncarnation) { + return { kind: "stale" }; + } + if (state.status === "cold") { + return { kind: "cold" }; + } + if ( + state.status !== "destroying" || + state.operationKind !== "activate" || + state.operationId !== lease.operationId || + state.claimOwner !== claimOwner || + state.claimExpiresAt === null + ) { + return state.status === "destroying" && state.operationKind === "activate" + ? { kind: "repairing" } + : { kind: "waiting" }; + } + + const operationLease: RuntimeSubjectOperationLease = { + claimExpiresAt: state.claimExpiresAt, + claimOwner, + incarnation: lease.sandboxIncarnation, + kind: "activate", + operationId: lease.operationId, + status: "destroying", + }; + return (await runtimeSubjectActivationRetirementIsDrained(database, { + lease: operationLease, + runtimeSubjectId: lease.sandboxId, + })) + ? { kind: "destroying", lease: operationLease } + : { kind: "waiting" }; +} + +export async function heartbeatRuntimeRunProvisioningLease( + database: D1Database, + lease: RuntimeRunProvisioningLease, +): Promise { + const heartbeatAt = currentTimestampMs(); + const db = getAppDatabase(database); + const activeRun = db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, lease.runId), + eq(sessionRunsTable.sessionId, lease.sessionId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ); + const updated = await db + .update(sessionsTable) + .set({ runtimeProvisioningHeartbeatAt: heartbeatAt }) + .where( + and( + runtimeProvisioningLeaseCondition(lease), + eq(sessionsTable.lastRunId, lease.runId), + eq(sessionsTable.status, "RUNNING"), + isNull(sessionsTable.archivedAt), + isNull(sessionsTable.cleanupOperationKind), + isNull(sessionsTable.statusOperationId), + exists(activeRun), + ), + ) + .returning({ id: sessionsTable.id }) + .get(); + + return updated !== undefined; +} + +export async function releaseReadyRuntimeRunProvisioningLease( + database: D1Database, + input: { + readonly driverGeneration: number; + readonly driverInstanceId: DriverInstanceId; + readonly lease: RuntimeRunProvisioningLease; + }, +): Promise { + if (input.lease.sandboxIncarnation === null) { + return false; + } + const db = getAppDatabase(database); + const activeConversation = db + .select({ id: sandboxSessionsTable.sessionId }) + .from(sandboxSessionsTable) + .where( + and( + eq(sandboxSessionsTable.sessionId, input.lease.sessionId), + eq(sandboxSessionsTable.sandboxId, input.lease.sandboxId), + input.lease.sandboxIncarnation === null + ? isNull(sandboxSessionsTable.sandboxIncarnation) + : eq(sandboxSessionsTable.sandboxIncarnation, input.lease.sandboxIncarnation), + input.lease.sandboxSessionId === null + ? isNull(sandboxSessionsTable.sandboxSessionId) + : eq(sandboxSessionsTable.sandboxSessionId, input.lease.sandboxSessionId), + eq(sandboxSessionsTable.status, "active"), + ), + ); + const linkedRun = db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, input.lease.runId), + eq(sessionRunsTable.sessionId, input.lease.sessionId), + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ); + const durableDriver = db + .select({ id: driverInstancesTable.id }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.driverGeneration), + eq(driverInstancesTable.sandboxId, input.lease.sandboxId), + eq(driverInstancesTable.sandboxIncarnation, input.lease.sandboxIncarnation), + eq(driverInstancesTable.sandboxSessionId, input.lease.sessionId), + inArray(driverInstancesTable.status, ASSIGNABLE_DRIVER_INSTANCE_STATUSES), + isNull(driverInstancesTable.statusOperationId), + ), + ); + const released = await db + .update(sessionsTable) + .set(runtimeProvisioningLeaseReleasePatch()) + .where( + and( + runtimeProvisioningLeaseCondition(input.lease), + eq(sessionsTable.lastRunId, input.lease.runId), + eq(sessionsTable.status, "RUNNING"), + isNull(sessionsTable.archivedAt), + isNull(sessionsTable.cleanupOperationKind), + isNull(sessionsTable.statusOperationId), + exists(activeConversation), + exists(linkedRun), + exists(durableDriver), + ), + ) + .returning({ id: sessionsTable.id }) + .get(); + + return released !== undefined; +} + +export async function adoptReadyRuntimeRunProvisioningLease( + database: D1Database, + lease: RuntimeProvisioningLease, +): Promise { + if (lease.runId === null || lease.sandboxIncarnation === null) { + return false; + } + const driver = await getAppDatabase(database) + .select({ + generation: driverInstancesTable.generation, + id: driverInstancesTable.id, + }) + .from(sessionRunsTable) + .innerJoin(driverInstancesTable, eq(driverInstancesTable.id, sessionRunsTable.driverInstanceId)) + .where( + and( + eq(sessionRunsTable.id, lease.runId), + eq(sessionRunsTable.sessionId, lease.sessionId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + eq(driverInstancesTable.sandboxId, lease.sandboxId), + eq(driverInstancesTable.sandboxIncarnation, lease.sandboxIncarnation), + eq(driverInstancesTable.sandboxSessionId, lease.sessionId), + inArray(driverInstancesTable.status, ASSIGNABLE_DRIVER_INSTANCE_STATUSES), + isNull(driverInstancesTable.statusOperationId), + ), + ) + .limit(1) + .get(); + + return driver === undefined + ? false + : releaseReadyRuntimeRunProvisioningLease(database, { + driverGeneration: driver.generation, + driverInstanceId: driver.id, + lease: { ...lease, runId: lease.runId }, + }); +} + +export async function releaseAbortedRuntimeProvisioningLease( + database: D1Database, + lease: RuntimeProvisioningLease, +): Promise { + const released = await getAppDatabase(database) + .update(sessionsTable) + .set(runtimeProvisioningLeaseReleasePatch()) + .where(runtimeProvisioningLeaseCondition(lease)) + .returning({ id: sessionsTable.id }) + .get(); + + return released !== undefined; +} + +export async function renewRuntimeProvisioningLeaseOwnership( + database: D1Database, + lease: RuntimeProvisioningLease, +): Promise { + const renewed = await getAppDatabase(database) + .update(sessionsTable) + .set({ runtimeProvisioningHeartbeatAt: currentTimestampMs() }) + .where(runtimeProvisioningLeaseCondition(lease)) + .returning({ id: sessionsTable.id }) + .get(); + + return renewed !== undefined; +} + +export async function readRuntimeProvisioningCleanupTargets( + database: D1Database, + lease: RuntimeProvisioningLease, +): Promise { + const rows = await getAppDatabase(database) + .select({ + driverGeneration: driverInstancesTable.generation, + driverId: driverInstancesTable.id, + driverStatus: driverInstancesTable.status, + }) + .from(sessionsTable) + .leftJoin( + driverInstancesTable, + and( + eq(driverInstancesTable.sandboxSessionId, lease.sessionId), + eq(driverInstancesTable.sandboxId, lease.sandboxId), + ...(lease.sandboxIncarnation === null + ? [sql`0`] + : [eq(driverInstancesTable.sandboxIncarnation, lease.sandboxIncarnation)]), + ), + ) + .where(runtimeProvisioningLeaseCondition(lease)) + .all(); + + if (rows.length === 0) { + return null; + } + + return { + conversationSessionId: lease.sandboxSessionId, + driverInstances: rows.flatMap((row) => + row.driverId === null || row.driverGeneration === null || row.driverStatus === null + ? [] + : [{ generation: row.driverGeneration, id: row.driverId, status: row.driverStatus }], + ), + }; +} + +export async function claimStaleRuntimeProvisioningLeases( + database: D1Database, + input: { readonly heartbeatAtLte: number; readonly limit: number }, +): Promise { + if (!Number.isSafeInteger(input.heartbeatAtLte) || input.heartbeatAtLte < 0) { + throw new Error("Runtime provisioning stale heartbeat must be a non-negative integer."); + } + if (!Number.isSafeInteger(input.limit) || input.limit <= 0) { + throw new Error("Runtime provisioning repair limit must be a positive integer."); + } + const db = getAppDatabase(database); + const rows = await db + .select({ + heartbeatAt: sessionsTable.runtimeProvisioningHeartbeatAt, + operationId: sessionsTable.runtimeProvisioningOperationId, + runId: sessionsTable.runtimeProvisioningRunId, + sandboxId: sessionsTable.runtimeProvisioningSandboxId, + sandboxIncarnation: sessionsTable.runtimeProvisioningSandboxIncarnation, + sandboxSessionId: sessionsTable.runtimeProvisioningSandboxSessionId, + sessionId: sessionsTable.id, + }) + .from(sessionsTable) + .where( + and( + isNotNull(sessionsTable.runtimeProvisioningOperationId), + isNotNull(sessionsTable.runtimeProvisioningSandboxId), + isNotNull(sessionsTable.runtimeProvisioningHeartbeatAt), + lte(sessionsTable.runtimeProvisioningHeartbeatAt, input.heartbeatAtLte), + ), + ) + .orderBy(sessionsTable.runtimeProvisioningHeartbeatAt, sessionsTable.id) + .limit(input.limit) + .all(); + const claims: RuntimeProvisioningLease[] = []; + + for (const row of rows) { + if (row.heartbeatAt === null || row.operationId === null || row.sandboxId === null) { + continue; + } + const previous: RuntimeProvisioningLease = { + heartbeatAt: row.heartbeatAt, + operationId: row.operationId, + runId: row.runId, + sandboxId: row.sandboxId, + sandboxIncarnation: row.sandboxIncarnation, + sandboxSessionId: row.sandboxSessionId, + sessionId: row.sessionId, + }; + const operationId = createPlatformId(); + const heartbeatAt = currentTimestampMs(); + const claimed = await db + .update(sessionsTable) + .set({ + runtimeProvisioningHeartbeatAt: heartbeatAt, + runtimeProvisioningOperationId: operationId, + }) + .where(runtimeProvisioningLeaseCondition(previous, true)) + .returning({ id: sessionsTable.id }) + .get(); + if (claimed !== undefined) { + claims.push({ ...previous, heartbeatAt, operationId }); + } + } + + return claims; +} + +function runtimeProvisioningLeaseCondition( + lease: RuntimeProvisioningLease, + includeHeartbeat = false, +) { + return and( + eq(sessionsTable.id, lease.sessionId), + eq(sessionsTable.runtimeProvisioningOperationId, lease.operationId), + lease.runId === null + ? isNull(sessionsTable.runtimeProvisioningRunId) + : eq(sessionsTable.runtimeProvisioningRunId, lease.runId), + eq(sessionsTable.runtimeProvisioningSandboxId, lease.sandboxId), + lease.sandboxIncarnation === null + ? isNull(sessionsTable.runtimeProvisioningSandboxIncarnation) + : eq(sessionsTable.runtimeProvisioningSandboxIncarnation, lease.sandboxIncarnation), + lease.sandboxSessionId === null + ? isNull(sessionsTable.runtimeProvisioningSandboxSessionId) + : eq(sessionsTable.runtimeProvisioningSandboxSessionId, lease.sandboxSessionId), + ...(includeHeartbeat + ? [eq(sessionsTable.runtimeProvisioningHeartbeatAt, lease.heartbeatAt)] + : []), + ); +} + +function runtimeProvisioningLeaseReleasePatch() { + return { + runtimeProvisioningHeartbeatAt: null, + runtimeProvisioningOperationId: null, + runtimeProvisioningRunId: null, + runtimeProvisioningSandboxId: null, + runtimeProvisioningSandboxIncarnation: null, + runtimeProvisioningSandboxSessionId: null, + } as const; +} diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-run-lease-store.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-run-lease-store.ts index 49aea3df..1571b1ae 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-run-lease-store.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-run-lease-store.ts @@ -4,11 +4,21 @@ import { sandboxesTable, sessionRunsTable, } from "@mosoo/db"; -import type { DriverInstanceId, SandboxId, SessionId, SessionRunId } from "@mosoo/id"; -import { and, eq, inArray, isNull, ne, notExists, or, sql } from "drizzle-orm"; +import type { + DriverInstanceId, + RuntimeOperationId, + SandboxId, + SessionId, + SessionRunId, +} from "@mosoo/id"; +import { and, eq, exists, inArray, isNull, ne, notExists, or, sql } from "drizzle-orm"; import { alias } from "drizzle-orm/sqlite-core"; -import { getAppDatabase } from "../../../../platform/db/drizzle"; +import { + getAppDatabase, + getD1ChangeCount, + runAppDatabaseBatch, +} from "../../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../../time"; import { ACTIVE_SESSION_RUN_STATUSES } from "../../domain/session-run-lifecycle.machine"; import { @@ -58,16 +68,24 @@ export type RuntimeRunLeaseTransitionOutcome = interface RuntimeRunLeaseAcquireSnapshot { readonly driverActiveSessionRunId: SessionRunId | null; + readonly driverGeneration: number; readonly driverSandboxId: SandboxId; + readonly driverSandboxIncarnation: number; readonly driverSandboxSessionId: SessionId; readonly driverStatus: string; + readonly driverStatusOperationId: RuntimeOperationId | null; readonly runDriverInstanceId: DriverInstanceId | null; readonly runId: SessionRunId | null; readonly runSessionId: SessionId | null; readonly runStatus: string | null; readonly runStatusSeq: number | null; readonly sandboxId: SandboxId; + readonly sandboxSessionIncarnation: number | null; readonly sandboxSessionStatus: string | null; + readonly subjectIncarnation: number | null; + readonly subjectClaimOwner: string | null; + readonly subjectOperationId: RuntimeOperationId | null; + readonly subjectStatus: string | null; } const activeDriverLeaseRunsTable = alias(sessionRunsTable, "active_driver_lease"); @@ -106,7 +124,7 @@ export async function recordRuntimeRunLeaseAcquiredOutcome( return admission; } - const linked = await recordRuntimeRunLeaseLinked(appDb, { + const linked = await recordRuntimeRunLeaseLinked(database, { ...input, now, sandboxId: snapshot.sandboxId, @@ -122,6 +140,14 @@ export async function recordRuntimeRunLeaseAcquiredOutcome( }; } + if (linked === "driver_changed") { + return { + reason: "driver_changed", + status: "stale", + transition: "acquire", + }; + } + return { reason: "run_changed", status: "stale", @@ -147,16 +173,24 @@ async function readRuntimeRunLeaseAcquireSnapshot( const row = (await appDb .select({ + driverGeneration: driverInstancesTable.generation, driverSandboxId: driverInstancesTable.sandboxId, + driverSandboxIncarnation: driverInstancesTable.sandboxIncarnation, driverSandboxSessionId: driverInstancesTable.sandboxSessionId, driverStatus: driverInstancesTable.status, + driverStatusOperationId: driverInstancesTable.statusOperationId, runDriverInstanceId: sessionRunsTable.driverInstanceId, runId: sessionRunsTable.id, runSessionId: sessionRunsTable.sessionId, runStatus: sessionRunsTable.status, runStatusSeq: sessionRunsTable.statusSeq, sandboxId: driverInstancesTable.sandboxId, + sandboxSessionIncarnation: sandboxSessionsTable.sandboxIncarnation, sandboxSessionStatus: sandboxSessionsTable.status, + subjectIncarnation: sandboxesTable.incarnation, + subjectClaimOwner: sandboxesTable.claimOwner, + subjectOperationId: sandboxesTable.statusOperationId, + subjectStatus: sandboxesTable.status, }) .from(driverInstancesTable) .leftJoin(sessionRunsTable, eq(sessionRunsTable.id, input.sessionRunId)) @@ -167,6 +201,7 @@ async function readRuntimeRunLeaseAcquireSnapshot( eq(sandboxSessionsTable.sessionId, input.sessionId), ), ) + .leftJoin(sandboxesTable, eq(sandboxesTable.id, input.runtimeSubjectId)) .where(eq(driverInstancesTable.id, input.driverInstanceId)) .limit(1) .get()) ?? null; @@ -198,6 +233,14 @@ function decideRuntimeRunLeaseAcquire( input: RuntimeRunLeaseInput, snapshot: RuntimeRunLeaseAcquireSnapshot, ): RuntimeRunLeaseTransitionOutcome { + if (snapshot.driverGeneration !== input.driverGeneration) { + return { + reason: "driver_changed", + status: "stale", + transition: "acquire", + }; + } + if (snapshot.runId === null) { return { reason: "run_not_found", @@ -208,6 +251,7 @@ function decideRuntimeRunLeaseAcquire( if ( snapshot.driverSandboxId !== input.runtimeSubjectId || + snapshot.driverSandboxIncarnation !== input.runtimeSubjectIncarnation || snapshot.driverSandboxSessionId !== input.sessionId ) { return { @@ -237,7 +281,14 @@ function decideRuntimeRunLeaseAcquire( }; } - if (snapshot.sandboxSessionStatus !== "active") { + if ( + snapshot.sandboxSessionStatus !== "active" || + snapshot.sandboxSessionIncarnation !== input.runtimeSubjectIncarnation || + snapshot.subjectIncarnation !== input.runtimeSubjectIncarnation || + snapshot.subjectStatus !== "active" || + snapshot.subjectClaimOwner !== null || + snapshot.subjectOperationId !== null + ) { return { reason: "sandbox_session_not_active", status: "rejected", @@ -246,6 +297,7 @@ function decideRuntimeRunLeaseAcquire( } if ( + snapshot.driverStatusOperationId !== null || !ASSIGNABLE_DRIVER_STATUSES.includes( snapshot.driverStatus as (typeof ASSIGNABLE_DRIVER_STATUSES)[number], ) @@ -294,57 +346,154 @@ function decideRuntimeRunLeaseAcquire( } async function recordRuntimeRunLeaseLinked( - appDb: AppDatabase, + database: D1Database, input: RuntimeRunLeaseInput & { readonly now: number; readonly sandboxId: SandboxId; readonly statusSeq: number | null; }, -): Promise<"linked" | "run_changed" | "run_link_conflict"> { +): Promise<"driver_changed" | "linked" | "run_changed" | "run_link_conflict"> { if (input.statusSeq === null) { return "run_changed"; } + const statusSeq = input.statusSeq; - const linked = - (await appDb - .update(sessionRunsTable) - .set({ - driverInstanceId: input.driverInstanceId, - updatedAt: sql` - CASE - WHEN ${sessionRunsTable.driverInstanceId} IS NULL THEN ${input.now} - ELSE ${sessionRunsTable.updatedAt} - END - `, - }) + const [linked] = await runAppDatabaseBatch(database, (db) => { + const activeConversation = db + .select({ id: sandboxSessionsTable.sessionId }) + .from(sandboxSessionsTable) + .where( + and( + eq(sandboxSessionsTable.sandboxId, input.runtimeSubjectId), + eq(sandboxSessionsTable.sandboxIncarnation, input.runtimeSubjectIncarnation), + eq(sandboxSessionsTable.sessionId, input.sessionId), + eq(sandboxSessionsTable.status, "active"), + ), + ); + const activeSubject = db + .select({ id: sandboxesTable.id }) + .from(sandboxesTable) + .where( + and( + eq(sandboxesTable.id, input.runtimeSubjectId), + eq(sandboxesTable.incarnation, input.runtimeSubjectIncarnation), + eq(sandboxesTable.status, "active"), + isNull(sandboxesTable.claimOwner), + isNull(sandboxesTable.operationKind), + isNull(sandboxesTable.statusOperationId), + ), + ); + const assignableDriver = db + .select({ id: driverInstancesTable.id }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.driverGeneration), + eq(driverInstancesTable.sandboxId, input.runtimeSubjectId), + eq(driverInstancesTable.sandboxIncarnation, input.runtimeSubjectIncarnation), + eq(driverInstancesTable.sandboxSessionId, input.sessionId), + inArray(driverInstancesTable.status, ASSIGNABLE_DRIVER_STATUSES), + isNull(driverInstancesTable.statusOperationId), + exists(activeConversation), + exists(activeSubject), + ), + ); + const exactLinkedRun = db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) .where( and( eq(sessionRunsTable.id, input.sessionRunId), eq(sessionRunsTable.sessionId, input.sessionId), - eq(sessionRunsTable.statusSeq, input.statusSeq), + eq(sessionRunsTable.statusSeq, statusSeq), + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), - or( - isNull(sessionRunsTable.driverInstanceId), - eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), - ), - notExists( - appDb - .select({ id: activeDriverLeaseRunsTable.id }) - .from(activeDriverLeaseRunsTable) - .where( - and( - eq(activeDriverLeaseRunsTable.driverInstanceId, input.driverInstanceId), - ne(activeDriverLeaseRunsTable.id, input.sessionRunId), - inArray(activeDriverLeaseRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ); + + return [ + db + .update(sessionRunsTable) + .set({ + driverInstanceId: input.driverInstanceId, + updatedAt: sql` + CASE + WHEN ${sessionRunsTable.driverInstanceId} IS NULL THEN ${input.now} + ELSE ${sessionRunsTable.updatedAt} + END + `, + }) + .where( + and( + eq(sessionRunsTable.id, input.sessionRunId), + eq(sessionRunsTable.sessionId, input.sessionId), + eq(sessionRunsTable.statusSeq, statusSeq), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + or( + isNull(sessionRunsTable.driverInstanceId), + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + ), + exists(assignableDriver), + notExists( + db + .select({ id: activeDriverLeaseRunsTable.id }) + .from(activeDriverLeaseRunsTable) + .where( + and( + eq(activeDriverLeaseRunsTable.driverInstanceId, input.driverInstanceId), + ne(activeDriverLeaseRunsTable.id, input.sessionRunId), + inArray(activeDriverLeaseRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), ), - ), + ), ), ), - ) - .returning({ id: sessionRunsTable.id }) - .get()) ?? null; + db + .update(sandboxesTable) + .set({ + inactiveDeadlineAt: null, + updatedAt: sql` + CASE + WHEN ${sandboxesTable.inactiveDeadlineAt} IS NULL THEN ${sandboxesTable.updatedAt} + ELSE ${input.now} + END + `, + }) + .where( + and( + eq(sandboxesTable.id, input.sandboxId), + eq(sandboxesTable.incarnation, input.runtimeSubjectIncarnation), + eq(sandboxesTable.status, "active"), + isNull(sandboxesTable.claimOwner), + isNull(sandboxesTable.operationKind), + isNull(sandboxesTable.statusOperationId), + exists(exactLinkedRun), + ), + ), + ]; + }); - if (linked === null) { + if (getD1ChangeCount(linked) === 0) { + const appDb = getAppDatabase(database); + const driver = await appDb + .select({ + generation: driverInstancesTable.generation, + sandboxIncarnation: driverInstancesTable.sandboxIncarnation, + statusOperationId: driverInstancesTable.statusOperationId, + }) + .from(driverInstancesTable) + .where(eq(driverInstancesTable.id, input.driverInstanceId)) + .limit(1) + .get(); + if ( + driver === undefined || + driver.generation !== input.driverGeneration || + driver.sandboxIncarnation !== input.runtimeSubjectIncarnation || + driver.statusOperationId !== null + ) { + return "driver_changed"; + } const current = (await appDb .select({ @@ -370,20 +519,6 @@ async function recordRuntimeRunLeaseLinked( return "run_link_conflict"; } - await appDb - .update(sandboxesTable) - .set({ - inactiveDeadlineAt: null, - updatedAt: sql` - CASE - WHEN ${sandboxesTable.inactiveDeadlineAt} IS NULL THEN ${sandboxesTable.updatedAt} - ELSE ${input.now} - END - `, - }) - .where(eq(sandboxesTable.id, input.sandboxId)) - .run(); - return "linked"; } @@ -391,7 +526,10 @@ export async function recordRuntimeRunLeaseReleased( database: D1Database, input: { readonly driverInstanceId: DriverInstanceId; + readonly expectedDriverGeneration: number; + readonly expectedDriverOperationId?: RuntimeOperationId; readonly expectedSessionRunId: SessionRunId; + readonly retainDriverOperationUntilTerminal?: boolean; }, ): Promise { const outcome = await recordRuntimeRunLeaseReleasedOutcome(database, input); @@ -402,14 +540,24 @@ export async function recordRuntimeRunLeaseReleasedOutcome( database: D1Database, input: { readonly driverInstanceId: DriverInstanceId; + readonly expectedDriverGeneration: number; + readonly expectedDriverOperationId?: RuntimeOperationId; readonly expectedSessionRunId: SessionRunId; + readonly retainDriverOperationUntilTerminal?: boolean; }, ): Promise { + if ( + input.retainDriverOperationUntilTerminal === true && + input.expectedDriverOperationId === undefined + ) { + throw new Error("A retained Driver release requires exact operation ownership."); + } const now = currentTimestampMs(); const appDb = getAppDatabase(database); const driver = (await appDb .select({ + generation: driverInstancesTable.generation, sandboxId: driverInstancesTable.sandboxId, }) .from(driverInstancesTable) @@ -425,6 +573,14 @@ export async function recordRuntimeRunLeaseReleasedOutcome( }; } + if (driver.generation !== input.expectedDriverGeneration) { + return { + reason: "driver_changed", + status: "stale", + transition: "release", + }; + } + const currentRun = (await appDb .select({ @@ -449,15 +605,15 @@ export async function recordRuntimeRunLeaseReleasedOutcome( .limit(1) .get()) ?? null; - if (currentRun === null || currentRun.driverInstanceId === null) { - if (activeDriverRun !== null && activeDriverRun.id !== input.expectedSessionRunId) { - return { - reason: "lease_mismatch", - status: "stale", - transition: "release", - }; - } + if (activeDriverRun !== null && activeDriverRun.id !== input.expectedSessionRunId) { + return { + reason: "lease_mismatch", + status: "stale", + transition: "release", + }; + } + if (currentRun === null || currentRun.driverInstanceId === null) { return { reason: "lease_missing", status: "rejected", @@ -473,13 +629,25 @@ export async function recordRuntimeRunLeaseReleasedOutcome( }; } - if ( - ACTIVE_SESSION_RUN_STATUSES.includes( - currentRun.status as (typeof ACTIVE_SESSION_RUN_STATUSES)[number], - ) - ) { - const released = - (await appDb + const currentRunIsActive = ACTIVE_SESSION_RUN_STATUSES.includes( + currentRun.status as (typeof ACTIVE_SESSION_RUN_STATUSES)[number], + ); + const [released, , driverFence] = await runAppDatabaseBatch(database, (db) => { + const exactDriverGeneration = db + .select({ id: driverInstancesTable.id }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.expectedDriverGeneration), + ...(input.expectedDriverOperationId === undefined + ? [] + : [eq(driverInstancesTable.statusOperationId, input.expectedDriverOperationId)]), + ), + ); + + return [ + db .update(sessionRunsTable) .set({ driverInstanceId: null, @@ -490,39 +658,106 @@ export async function recordRuntimeRunLeaseReleasedOutcome( eq(sessionRunsTable.id, input.expectedSessionRunId), eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + exists(exactDriverGeneration), ), - ) - .returning({ - id: sessionRunsTable.id, + ), + db + .update(sandboxesTable) + .set({ + inactiveDeadlineAt: sql`COALESCE( + ${sandboxesTable.inactiveDeadlineAt}, + ${getRuntimeSubjectInactiveDeadlineSql(now)} + )`, + updatedAt: sql`CASE + WHEN ${sandboxesTable.inactiveDeadlineAt} IS NULL THEN ${now} + ELSE ${sandboxesTable.updatedAt} + END`, }) - .get()) ?? null; + .where( + and( + eq(sandboxesTable.id, driver.sandboxId), + exists(exactDriverGeneration), + or( + eq(sandboxesTable.kind, "pet"), + notExists(activeConversationSessionQuery(db, driver.sandboxId)), + ), + notExists(runLeaseQuery(db, driver.sandboxId)), + ), + ), + db + .update(driverInstancesTable) + .set( + input.retainDriverOperationUntilTerminal === true + ? { + status: sql`CASE + WHEN ${driverInstancesTable.status} IN ('provisioning', 'connecting', 'ready') + THEN 'stopping' + ELSE ${driverInstancesTable.status} + END`, + statusChangedAt: sql`CASE + WHEN ${driverInstancesTable.status} IN ('provisioning', 'connecting', 'ready') + THEN ${now} + ELSE ${driverInstancesTable.statusChangedAt} + END`, + statusEvent: sql`CASE + WHEN ${driverInstancesTable.status} IN ('provisioning', 'connecting', 'ready') + THEN 'driver.stopping' + ELSE ${driverInstancesTable.statusEvent} + END`, + statusOperationId: sql`CASE + WHEN ${driverInstancesTable.status} IN ('stopped', 'failed') THEN NULL + ELSE ${driverInstancesTable.statusOperationId} + END`, + statusSeq: sql`CASE + WHEN ${driverInstancesTable.status} IN ('provisioning', 'connecting', 'ready') + THEN ${driverInstancesTable.statusSeq} + 1 + ELSE ${driverInstancesTable.statusSeq} + END`, + statusSource: sql`CASE + WHEN ${driverInstancesTable.status} IN ('provisioning', 'connecting', 'ready') + THEN 'api' + ELSE ${driverInstancesTable.statusSource} + END`, + updatedAt: sql`CASE + WHEN ${driverInstancesTable.status} IN ('provisioning', 'connecting', 'ready') + THEN ${now} + ELSE ${driverInstancesTable.updatedAt} + END`, + } + : { + statusOperationId: + input.expectedDriverOperationId === undefined + ? sql`${driverInstancesTable.statusOperationId}` + : null, + }, + ) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.expectedDriverGeneration), + ...(input.expectedDriverOperationId === undefined + ? [] + : [eq(driverInstancesTable.statusOperationId, input.expectedDriverOperationId)]), + ), + ), + ]; + }); - if (!released) { - return { - reason: "run_changed", - status: "stale", - transition: "release", - }; - } + if (getD1ChangeCount(driverFence) === 0) { + return { + reason: "driver_changed", + status: "stale", + transition: "release", + }; } - await appDb - .update(sandboxesTable) - .set({ - inactiveDeadlineAt: getRuntimeSubjectInactiveDeadlineSql(now), - updatedAt: now, - }) - .where( - and( - eq(sandboxesTable.id, driver.sandboxId), - or( - eq(sandboxesTable.kind, "pet"), - notExists(activeConversationSessionQuery(appDb, driver.sandboxId)), - ), - notExists(runLeaseQuery(appDb, driver.sandboxId)), - ), - ) - .run(); + if (currentRunIsActive && getD1ChangeCount(released) === 0) { + return { + reason: "run_changed", + status: "stale", + transition: "release", + }; + } return { repaired: false, diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-driver-stop.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-driver-stop.ts index 9d4934e7..2976d65e 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-driver-stop.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-driver-stop.ts @@ -2,41 +2,31 @@ import type { DriverInstanceId } from "@mosoo/id"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; import type { StopRuntimeSubjectDriversInput } from "../../application/execution-plane/execution-plane-adapter"; -import { listLiveDriverInstanceIdsForSandboxSessions } from "../driver-instance/live-driver-instance.repository"; import { stopDriverSession } from "../driver-session.service"; import { listRuntimeSubjectDriverIds } from "./runtime-subject-store"; -async function listRuntimeSubjectOperationDriverIds( - bindings: ApiBindings, - input: StopRuntimeSubjectDriversInput, -): Promise { - if (input.targets !== undefined) { - return listLiveDriverInstanceIdsForSandboxSessions( - bindings.DB, - input.targets.map((target) => target.sessionId), - ); - } - - return listRuntimeSubjectDriverIds(bindings.DB, input.runtimeSubjectId); -} - export async function stopRuntimeSubjectDrivers( bindings: ApiBindings, input: StopRuntimeSubjectDriversInput, ): Promise { - const driverIds = await listRuntimeSubjectOperationDriverIds(bindings, input); + const driverIds: DriverInstanceId[] = await listRuntimeSubjectDriverIds( + bindings.DB, + input.runtimeSubjectId, + input.sandboxIncarnation, + ); - await Promise.all( + const outcomes = await Promise.allSettled( driverIds.map((driverInstanceId) => stopDriverSession(bindings, { driverInstanceId, - ...(input.operationId !== undefined ? { operationId: input.operationId } : {}), - ...(input.preserveSessionLifecycle !== undefined - ? { preserveSessionLifecycle: input.preserveSessionLifecycle } - : {}), reason: input.reason, - ...(input.terminalRun ? { terminalRun: input.terminalRun } : {}), }), ), ); + const failure = outcomes.find( + (outcome): outcome is PromiseRejectedResult => outcome.status === "rejected", + ); + if (failure !== undefined) { + throw failure.reason; + } } diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-errors.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-errors.ts index b52e4564..7968f815 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-errors.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-errors.ts @@ -57,6 +57,15 @@ export class RuntimeSubjectRestoreFailedError extends Error { } } +export class RuntimeSubjectPhysicalStateLostError extends Error { + constructor(runtimeSubjectId: string) { + super( + `Runtime subject ${runtimeSubjectId} physical incarnation was lost before its checkpoint completed.`, + ); + this.name = "RuntimeSubjectPhysicalStateLostError"; + } +} + export class RuntimeBucketMountConflictError extends Error { readonly bucket: string | null; readonly mountPath: string; diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-lifecycle.service.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-lifecycle.service.ts index a7d1e6fe..40342bb4 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-lifecycle.service.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-lifecycle.service.ts @@ -21,6 +21,7 @@ import { SERVER_PRODUCT_ANALYTICS_EVENTS, } from "../../../../platform/analytics/product-analytics"; import { createErrorLogContext, logWarn } from "../../../../platform/cloudflare/logger"; +import { disposeRpcResource } from "../../../../platform/cloudflare/rpc-disposal"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; import { currentTimestampMs } from "../../../../time"; import { @@ -32,17 +33,17 @@ import type { RuntimeDiagnosticContext } from "../../application/runtime-diagnos import type { RuntimeTimingRecorder } from "../../application/session-runs/session-runtime-timing"; import { getRuntimeKindPolicy, - getRuntimeSubjectInactiveDeadline, runtimeCheckpointRulesInclude, } from "../../domain/runtime-kind-policy"; import type { SandboxNetworkConstraints } from "../../domain/sandbox-network-constraints"; +import { hashSandboxNetworkConstraints } from "../../domain/sandbox-network-constraints"; import type { SandboxHandle } from "../sandbox-handles"; -import { deleteActiveSandboxConversationSession } from "../sandbox-session/sandbox-conversation-session-delete"; import { recordRuntimeRunLeaseAcquiredOutcome, recordRuntimeRunLeaseReleased, } from "./runtime-run-lease-store"; import type { RuntimeRunLeaseTransitionOutcome } from "./runtime-run-lease-store"; +import { stopRuntimeSubjectDrivers } from "./runtime-subject-driver-stop"; import { getRuntimeSubjectErrorCode, RuntimeSubjectBackupNotReadyError, @@ -51,27 +52,34 @@ import { import { assertRuntimeSubjectNetworkPolicySupported } from "./runtime-subject-network"; import { configureRuntimeSubjectNetwork, + activateRuntimeSubjectIncarnation, destroyRuntimeSubjectContainer, getRuntimeSubjectKeepAliveHandle, + inspectRuntimeSubjectIncarnation, + markRuntimeSubjectIncarnationReady, prepareRuntimeSubjectFilesystem, restoreRuntimeSubjectBackup, } from "./runtime-subject-platform"; import { claimRuntimeSubjectActivation, + claimRuntimeSubjectOperationForRepair, ensureRuntimeSubjectId, - getRuntimeConversationSessionState, getRuntimeSubjectActivationRecord, markRuntimeSubjectActivationDestroying, markRuntimeSubjectActivationFailed, + markRuntimeSubjectActiveDestroying, markRuntimeSubjectActive, + markRuntimeSubjectOperationRepairNeeded, markRuntimeSubjectRestoreApplied, markRuntimeSubjectRestoring, preemptRuntimeSubjectActivationClaim, recordRuntimeConversationSessionActive, - recordRuntimeConversationSessionClosed, recordRuntimeConversationSessionError, + releaseRuntimeSubjectActivationClaim, + retireRuntimeConversationSessionsForIncarnation, } from "./runtime-subject-store"; import type { RuntimeSubjectActivationRecord } from "./runtime-subject-store"; +import type { RuntimeSubjectOperationLease } from "./runtime-subject-store"; import type { ReadyRuntimeSubjectBackupRecord } from "./runtime-subject-store"; const RUNTIME_SUBJECT_ACTIVATION_CLAIM_TTL_MS = 10 * 60_000; @@ -90,6 +98,11 @@ export interface ActivateRuntimeSubjectInput { readonly diagnosticContext?: RuntimeDiagnosticContext; readonly networkConstraints: SandboxNetworkConstraints; readonly purpose?: RuntimeSubjectActivationPurpose; + readonly provisioningAuthority?: { + readonly operationId: RuntimeOperationId; + readonly runId: SessionRunId; + readonly sessionId: SessionId; + }; readonly runtimeSubjectId: SandboxId; readonly appId: AppId; readonly subjectId: PlatformId; @@ -98,6 +111,7 @@ export interface ActivateRuntimeSubjectInput { } export interface ActiveRuntimeSubject { + readonly incarnation: number; readonly subject: SandboxHandle; } @@ -193,8 +207,8 @@ export class RuntimeSubjectLifecycleService { this.#bindings = bindings; } - async getHandle(runtimeSubjectId: SandboxId): Promise { - return getRuntimeSubjectKeepAliveHandle(this.#bindings, runtimeSubjectId); + async getHandle(runtimeSubjectId: SandboxId, incarnation: number): Promise { + return getRuntimeSubjectKeepAliveHandle(this.#bindings, runtimeSubjectId, incarnation); } async activate(input: ActivateRuntimeSubjectInput): Promise { @@ -206,55 +220,122 @@ export class RuntimeSubjectLifecycleService { const purpose = input.purpose ?? "interactive"; const claimOwner = createRuntimeSubjectActivationClaimOwner(purpose); + const networkConstraintsHash = await hashSandboxNetworkConstraints(input.networkConstraints); const record = await measureOptional(input.timing, "runtimeSubject.admitLifecycle", () => this.#admitActivation(input, claimOwner, purpose), ); - const subject = await this.getHandle(input.runtimeSubjectId); const isCold = record === null || record.status === "cold"; + let activationLease: RuntimeSubjectOperationLease | null = null; + let subject: SandboxHandle | null = null; + let subjectTransferred = false; + let reusedHealthyIncarnation = false; try { - // Network constraints must land before the first container-starting RPC - // below; a limited policy that cannot be applied fails the activation - // (and the catch path destroys the container) instead of running open. - // Stable Pet subjects support Full only and keep their existing runtime - // path unchanged. Cattle records Full as well as Limited so the - // session-scoped subject can never switch policy after admission. - if (input.kind === "cattle") { - await measureOptional(input.timing, "runtimeSubject.configureNetwork", () => - configureRuntimeSubjectNetwork(subject, input.networkConstraints), - ); - } - await measureOptional(input.timing, "runtimeSubject.prepareFilesystem", () => - prepareRuntimeSubjectFilesystem(subject), - ); - if (isCold) { - const restoring = await measureOptional(input.timing, "runtimeSubject.markRestoring", () => + activationLease = await measureOptional(input.timing, "runtimeSubject.markRestoring", () => markRuntimeSubjectRestoring(this.#bindings.DB, { claimOwner, + expectedIncarnation: record?.incarnation ?? 0, + expectedStatus: "cold", + networkConstraintsHash, + operationId: createPlatformId(), runtimeSubjectId: input.runtimeSubjectId, }), ); - if (!restoring) { + if (activationLease === null) { throw new Error("Runtime subject activation claim expired before restore."); } + } else if (record !== null) { + if (record.networkConstraintsHash !== networkConstraintsHash) { + const retired = await this.#retireActiveIncarnation({ + claimOwner, + errorCode: "runtime.subject_activation_failed", + message: "Runtime subject network constraints changed.", + record, + runtimeSubjectId: input.runtimeSubjectId, + provisioningAuthority: input.provisioningAuthority, + }); + if (!retired) { + throw new Error("Runtime subject activation claim expired before network retirement."); + } + throw new Error("Runtime subject network constraints changed; retry activation."); + } + subject = await this.getHandle(input.runtimeSubjectId, record.incarnation); + const health = await inspectRuntimeSubjectIncarnation( + subject, + record.incarnation, + networkConstraintsHash, + ); + if (health.kind === "healthy") { + reusedHealthyIncarnation = true; + } else if (health.kind === "unknown") { + throw new Error("Runtime subject active-container health is unknown."); + } else { + const retired = await this.#retireActiveIncarnation({ + claimOwner, + errorCode: "runtime.subject_activation_failed", + message: `Runtime subject active container is ${health.kind}.`, + record, + runtimeSubjectId: input.runtimeSubjectId, + provisioningAuthority: input.provisioningAuthority, + }); + if (!retired) { + throw new Error("Runtime subject activation claim expired before recovery."); + } + throw new Error("Runtime subject active container was retired; retry activation."); + } + } + + const incarnation = activationLease?.incarnation ?? record?.incarnation ?? 0; + subject ??= await this.getHandle(input.runtimeSubjectId, incarnation); + const activeSubject = subject; + if (activationLease !== null) { + await activateRuntimeSubjectIncarnation(activeSubject, incarnation, networkConstraintsHash); + } + + // A healthy active incarnation is already prepared. Re-running global + // container mutations during reuse is both redundant and unsafe: one + // caller timing out must never poison a Pet container used by another + // Run. Fresh incarnations still establish network policy before their + // first container-starting filesystem RPC. + if (!reusedHealthyIncarnation) { + if (input.kind === "cattle") { + await measureOptional(input.timing, "runtimeSubject.configureNetwork", () => + configureRuntimeSubjectNetwork(activeSubject, input.networkConstraints), + ); + } + await measureOptional(input.timing, "runtimeSubject.prepareFilesystem", () => + prepareRuntimeSubjectFilesystem(activeSubject), + ); + } + + if (activationLease !== null) { + const lease = activationLease; await measureOptional(input.timing, "runtimeSubject.restoreBackup", () => this.#restoreLastBackup({ - claimOwner, kind: input.kind, + lease, record, runtimeSubjectId: input.runtimeSubjectId, - subject, + subject: activeSubject, }), ); + await markRuntimeSubjectIncarnationReady( + activeSubject, + incarnation, + networkConstraintsHash, + ); } const activated = await measureOptional(input.timing, "runtimeSubject.markActive", () => markRuntimeSubjectActive(this.#bindings.DB, { claimOwner, + incarnation: activationLease?.incarnation ?? record?.incarnation ?? 0, kind: input.kind, + networkConstraintsHash, + operationId: activationLease?.operationId ?? null, runtimeSubjectId: input.runtimeSubjectId, }), ); @@ -283,25 +364,44 @@ export class RuntimeSubjectLifecycleService { }, }); } + + if (subject === null) { + throw new Error("Runtime subject activation completed without a Sandbox handle."); + } + + subjectTransferred = true; + return { + incarnation: activationLease?.incarnation ?? record?.incarnation ?? 0, + subject, + }; } catch (error) { const message = error instanceof Error ? error.message : "Runtime subject activation failed."; const errorCode = getRuntimeSubjectErrorCode(error); - const operationId = createPlatformId(); + let destroyingRecorded = false; - try { - destroyingRecorded = await markRuntimeSubjectActivationDestroying(this.#bindings.DB, { + if (activationLease === null) { + await releaseRuntimeSubjectActivationClaim(this.#bindings.DB, { claimOwner, errorCode, - message, - operationId, - runtimeSubjectId: input.runtimeSubjectId, - }); - } catch (recordError) { - logWarn("runtime.subject.activation_failure.destroy_record_failed", { - ...createErrorLogContext(recordError), + errorMessage: message, + incarnation: record?.incarnation ?? 0, runtimeSubjectId: input.runtimeSubjectId, }); + } else { + try { + destroyingRecorded = await markRuntimeSubjectActivationDestroying(this.#bindings.DB, { + errorCode, + lease: activationLease, + message, + runtimeSubjectId: input.runtimeSubjectId, + }); + } catch (recordError) { + logWarn("runtime.subject.activation_failure.destroy_record_failed", { + ...createErrorLogContext(recordError), + runtimeSubjectId: input.runtimeSubjectId, + }); + } } // Teardown is bounded by the provision timeout. Only confirmed teardown @@ -311,7 +411,11 @@ export class RuntimeSubjectLifecycleService { if (destroyingRecorded) { try { - await destroyRuntimeSubjectContainer(this.#bindings, input.runtimeSubjectId); + await destroyRuntimeSubjectContainer( + this.#bindings, + input.runtimeSubjectId, + activationLease?.incarnation ?? 0, + ); destroyed = true; } catch (destroyError) { logWarn("runtime.subject.activation_failure.destroy_failed", { @@ -321,12 +425,12 @@ export class RuntimeSubjectLifecycleService { } } - if (destroyingRecorded && destroyed) { + if (activationLease !== null && destroyingRecorded && destroyed) { try { await markRuntimeSubjectActivationFailed(this.#bindings.DB, { errorCode, + lease: activationLease, message, - operationId, runtimeSubjectId: input.runtimeSubjectId, }); } catch (finalizeError) { @@ -346,13 +450,16 @@ export class RuntimeSubjectLifecycleService { }); throw new Error(message, { cause: error }); + } finally { + if (!subjectTransferred) { + disposeRpcResource(subject); + } } - - return { subject }; } async activateConversationSession(input: { readonly sandboxSessionId: SandboxSessionId; + readonly sandboxIncarnation: number; readonly cwd: string; readonly now: number; readonly originJson: string; @@ -364,6 +471,7 @@ export class RuntimeSubjectLifecycleService { async failConversationSession(input: { readonly sandboxSessionId: SandboxSessionId; + readonly sandboxIncarnation: number; readonly cwd: string; readonly errorCode: RuntimeSubjectErrorCode; readonly message: string; @@ -379,45 +487,18 @@ export class RuntimeSubjectLifecycleService { readonly runtimeSubjectId: SandboxId; readonly sessionId: SessionId; }): Promise { - const state = await getRuntimeConversationSessionState(this.#bindings.DB, input); - - if (!state || state.status !== "active") { - return; - } - - const now = currentTimestampMs(); - - await deleteActiveSandboxConversationSession(this.#bindings, { - sandboxSessionId: state.sandboxSessionId, + const { closeSandboxConversationSession } = await import("../sandbox-session.service"); + await closeSandboxConversationSession(this.#bindings, { sandboxId: input.runtimeSubjectId, - }); - - if (state.agentId) { - await appendRuntimeDiagnosticEvent(this.#bindings, { - eventName: RUNTIME_DIAGNOSTIC_EVENT.sandboxSessionDestroyed.name, - sessionId: input.sessionId, - value: { - ...toRuntimeDiagnosticBaseValue({ - agentId: state.agentId, - sessionId: input.sessionId, - }), - reason: "runtime_subject_session_closed", - sandboxId: input.runtimeSubjectId, - }, - }); - } - - await recordRuntimeConversationSessionClosed(this.#bindings.DB, { - inactiveDeadlineAt: getRuntimeSubjectInactiveDeadline(getRuntimeKindPolicy(state.kind), now), - now, - runtimeSubjectId: input.runtimeSubjectId, sessionId: input.sessionId, }); } async acquireRunLease(input: { + readonly driverGeneration: number; readonly driverInstanceId: DriverInstanceId; readonly runtimeSubjectId: SandboxId; + readonly runtimeSubjectIncarnation: number; readonly sessionId: SessionId; readonly sessionRunId: SessionRunId; }): Promise { @@ -426,6 +507,7 @@ export class RuntimeSubjectLifecycleService { async releaseRunLease(input: { readonly driverInstanceId: DriverInstanceId; + readonly expectedDriverGeneration: number; readonly expectedSessionRunId: SessionRunId; }): Promise { return recordRuntimeRunLeaseReleased(this.#bindings.DB, input); @@ -498,8 +580,19 @@ export class RuntimeSubjectLifecycleService { }): Promise { let record = input.record; - if (record.kind !== input.activation.kind) { - throw new Error("Runtime subject kind does not match the requested runtime kind."); + if ( + record.kind !== input.activation.kind || + record.subjectKind !== input.activation.subjectKind || + record.subjectId !== input.activation.subjectId + ) { + throw new Error("Runtime subject identity does not match the activation request."); + } + if ( + record.agentId !== input.activation.agentId || + record.appId !== input.activation.appId || + record.ownerAccountId !== input.activation.executionOwnerUserId + ) { + throw new Error("Runtime subject ownership does not match the activation request."); } if (record.status === "backing_up" || record.status === "destroying") { @@ -543,8 +636,12 @@ export class RuntimeSubjectLifecycleService { throw new Error("Runtime subject activation could not refresh the lifecycle record."); } record = refreshed; - if (record.kind !== input.activation.kind) { - throw new Error("Runtime subject kind does not match the requested runtime kind."); + if ( + record.kind !== input.activation.kind || + record.subjectKind !== input.activation.subjectKind || + record.subjectId !== input.activation.subjectId + ) { + throw new Error("Runtime subject identity changed during activation."); } if (record.status === "backing_up" || record.status === "destroying") { throw new Error("Runtime subject is busy with lifecycle maintenance."); @@ -552,6 +649,43 @@ export class RuntimeSubjectLifecycleService { } if (record.status === "restoring") { + if ( + !hasActiveRuntimeSubjectClaim(record, currentTimestampMs()) && + record.operationId !== null && + record.operationKind === "activate" + ) { + const repairNow = currentTimestampMs(); + const lease = await claimRuntimeSubjectOperationForRepair(this.#bindings.DB, { + candidate: { + claimExpiresAt: record.claimExpiresAt, + claimOwner: record.claimOwner, + id: record.id, + incarnation: record.incarnation, + kind: record.kind, + operationId: record.operationId, + operationKind: "activate", + status: "restoring", + }, + claimExpiresAt: repairNow + RUNTIME_SUBJECT_ACTIVATION_CLAIM_TTL_MS, + claimOwner: `activation-repair-${crypto.randomUUID()}`, + now: repairNow, + }); + if (lease !== null) { + const { runRuntimeSubjectOperation } = + await import("./runtime-subject-operations.service"); + await runRuntimeSubjectOperation(this.#bindings, { + kind: record.kind, + lease, + reason: "runtime_subject.activation_takeover", + runtimeSubjectId: record.id, + }); + const repaired = await getRuntimeSubjectActivationRecord(this.#bindings.DB, record.id); + if (repaired === null) { + throw new Error("Runtime subject activation repair lost its lifecycle record."); + } + return this.#claimExistingActivation({ ...input, record: repaired }); + } + } throw new Error("Runtime subject is busy with lifecycle maintenance."); } @@ -635,9 +769,87 @@ export class RuntimeSubjectLifecycleService { }); } - async #restoreLastBackup(input: { + async #retireActiveIncarnation(input: { readonly claimOwner: string; + readonly errorCode: RuntimeSubjectErrorCode; + readonly message: string; + readonly provisioningAuthority: ActivateRuntimeSubjectInput["provisioningAuthority"]; + readonly record: RuntimeSubjectActivationRecord; + readonly runtimeSubjectId: SandboxId; + }): Promise { + if (input.provisioningAuthority === undefined) { + return false; + } + const lease = await markRuntimeSubjectActiveDestroying(this.#bindings.DB, { + claimOwner: input.claimOwner, + errorCode: input.errorCode, + expectedIncarnation: input.record.incarnation, + message: input.message, + operationId: createPlatformId(), + provisioningOperationId: input.provisioningAuthority.operationId, + provisioningRunId: input.provisioningAuthority.runId, + provisioningSessionId: input.provisioningAuthority.sessionId, + runtimeSubjectId: input.runtimeSubjectId, + }); + if (lease === null) { + return false; + } + + try { + await stopRuntimeSubjectDrivers(this.#bindings, { + operationId: lease.operationId, + reason: "runtime_subject.active_incarnation_retired", + runtimeSubjectId: input.runtimeSubjectId, + sandboxIncarnation: lease.incarnation, + }); + await destroyRuntimeSubjectContainer( + this.#bindings, + input.runtimeSubjectId, + lease.incarnation, + ); + await retireRuntimeConversationSessionsForIncarnation(this.#bindings.DB, { + now: currentTimestampMs(), + runtimeSubjectId: input.runtimeSubjectId, + sandboxIncarnation: lease.incarnation, + }); + if ( + !(await markRuntimeSubjectActivationFailed(this.#bindings.DB, { + errorCode: input.errorCode, + lease, + message: input.message, + runtimeSubjectId: input.runtimeSubjectId, + })) + ) { + throw new Error("Runtime subject active-incarnation retirement lost ownership."); + } + } catch (error) { + try { + await markRuntimeSubjectOperationRepairNeeded(this.#bindings.DB, { + errorCode: getRuntimeSubjectErrorCode(error), + errorMessage: error instanceof Error ? error.message : input.message, + expectedStatus: "destroying", + lease, + runtimeSubjectId: input.runtimeSubjectId, + source: "api", + }); + } catch (recordError) { + logWarn("runtime.subject.active_retire.repair_record_failed", { + ...createErrorLogContext(recordError), + runtimeSubjectId: input.runtimeSubjectId, + }); + } + logWarn("runtime.subject.active_retire.failed", { + ...createErrorLogContext(error), + runtimeSubjectId: input.runtimeSubjectId, + }); + } + + return true; + } + + async #restoreLastBackup(input: { readonly kind: AgentKind; + readonly lease: RuntimeSubjectOperationLease; readonly record: RuntimeSubjectActivationRecord | null; readonly runtimeSubjectId: SandboxId; readonly subject: SandboxHandle; @@ -664,11 +876,14 @@ export class RuntimeSubjectLifecycleService { runtimeSubjectId: input.runtimeSubjectId, }); } - await markRuntimeSubjectRestoreApplied(this.#bindings.DB, { + const recorded = await markRuntimeSubjectRestoreApplied(this.#bindings.DB, { backupId: readyBackup.id, - claimOwner: input.claimOwner, + lease: input.lease, runtimeSubjectId: input.runtimeSubjectId, }); + if (!recorded) { + throw new Error("Runtime subject restore lost lifecycle ownership."); + } } async #appendRestoreFailureDiagnostic(input: { @@ -712,7 +927,4 @@ export function createRuntimeSubjectLifecycleService( return new RuntimeSubjectLifecycleService(bindings); } -export { - getRuntimeSubjectKeepAliveHandle, - prepareRuntimeSubjectFilesystem, -} from "./runtime-subject-platform"; +export { getRuntimeSubjectKeepAliveHandle } from "./runtime-subject-platform"; diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance-store.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance-store.ts index 55b86ace..0d14360d 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance-store.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance-store.ts @@ -8,13 +8,9 @@ import { import type { DriverInstanceId, SandboxId, SessionId } from "@mosoo/id"; import { and, asc, eq, exists, inArray, isNotNull, isNull, lte, notExists, or } from "drizzle-orm"; -import { - getAppDatabase, - getD1ChangeCount, - runAppDatabaseBatch, -} from "../../../../platform/db/drizzle"; +import { getAppDatabase, getD1ChangeCount } from "../../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../../time"; -import { RUNTIME_SUBJECT_OPERATION_STATUSES } from "../../domain/runtime-subject-lifecycle.machine"; +import { RUNTIME_SUBJECT_RECOVERABLE_OPERATION_STATUSES } from "../../domain/runtime-subject-lifecycle.machine"; import { ACTIVE_SESSION_RUN_STATUSES } from "../../domain/session-run-lifecycle.machine"; import { activeConversationSessionQuery, @@ -25,52 +21,28 @@ import { liveDriverInstanceQueryForListedSubject, runLeaseQuery, runLeaseQueryForListedSubject, + runtimeProvisioningQuery, + runtimeProvisioningQueryForListedSubject, } from "./runtime-subject-store-queries"; import type { RuntimeSubjectMaintenanceCandidate, RuntimeSubjectOperationRepairCandidate, + RuntimeSubjectOperationLease, RuntimeSubjectStatus, } from "./runtime-subject-store.types"; function isRuntimeSubjectOperationStatus( status: RuntimeSubjectStatus, ): status is RuntimeSubjectOperationRepairCandidate["status"] { - return RUNTIME_SUBJECT_OPERATION_STATUSES.includes( + return RUNTIME_SUBJECT_RECOVERABLE_OPERATION_STATUSES.includes( status as RuntimeSubjectOperationRepairCandidate["status"], ); } -export async function closeRuntimeSubjectSessionsForRecycle( - database: D1Database, - runtimeSubjectId: SandboxId, -): Promise { - const now = currentTimestampMs(); - - await runAppDatabaseBatch(database, (appDb) => [ - appDb - .update(sandboxSessionsTable) - .set({ - status: "closed", - updatedAt: now, - }) - .where( - and( - eq(sandboxSessionsTable.sandboxId, runtimeSubjectId), - eq(sandboxSessionsTable.status, "active"), - ), - ), - appDb - .update(sandboxesTable) - .set({ - updatedAt: now, - }) - .where(eq(sandboxesTable.id, runtimeSubjectId)), - ]); -} - export async function listRuntimeSubjectDriverIds( database: D1Database, runtimeSubjectId: SandboxId, + sandboxIncarnation?: number, ): Promise { const appDb = getAppDatabase(database); const activeRunLeaseQuery = appDb @@ -88,6 +60,9 @@ export async function listRuntimeSubjectDriverIds( .where( and( eq(driverInstancesTable.sandboxId, runtimeSubjectId), + ...(sandboxIncarnation === undefined + ? [] + : [eq(driverInstancesTable.sandboxIncarnation, sandboxIncarnation)]), or(inArray(driverInstancesTable.status, LIVE_DRIVER_STATUSES), exists(activeRunLeaseQuery)), ), ) @@ -147,6 +122,7 @@ export async function listInactiveRuntimeSubjects( notExists(activeConversationSessionQueryForListedSubject(appDb)), ), notExists(runLeaseQueryForListedSubject(appDb)), + notExists(runtimeProvisioningQueryForListedSubject(appDb)), isNotNull(sandboxesTable.inactiveDeadlineAt), lte(sandboxesTable.inactiveDeadlineAt, input.now), ), @@ -179,6 +155,7 @@ export async function repairStrandedRuntimeSubjectDeadlines( isNull(sandboxesTable.inactiveDeadlineAt), notExists(activeConversationSessionQueryForListedSubject(appDb)), notExists(runLeaseQueryForListedSubject(appDb)), + notExists(runtimeProvisioningQueryForListedSubject(appDb)), ), ) .run(); @@ -203,6 +180,7 @@ export async function repairStrandedRuntimeSubjectDeadlines( notExists(liveDriverInstanceQueryForListedSubject(appDb)), notExists(activeSessionRunQueryForListedSubject(appDb)), notExists(runLeaseQueryForListedSubject(appDb)), + notExists(runtimeProvisioningQueryForListedSubject(appDb)), ), ) .run(); @@ -225,14 +203,23 @@ export async function listStaleRuntimeSubjectOperations( id: sandboxesTable.id, kind: sandboxesTable.kind, operationId: sandboxesTable.statusOperationId, + claimExpiresAt: sandboxesTable.claimExpiresAt, + claimOwner: sandboxesTable.claimOwner, + incarnation: sandboxesTable.incarnation, + operationKind: sandboxesTable.operationKind, status: sandboxesTable.status, }) .from(sandboxesTable) .where( and( - inArray(sandboxesTable.status, RUNTIME_SUBJECT_OPERATION_STATUSES), + inArray(sandboxesTable.status, RUNTIME_SUBJECT_RECOVERABLE_OPERATION_STATUSES), isNotNull(sandboxesTable.statusOperationId), - lte(sandboxesTable.statusChangedAt, input.staleChangedAtLte), + isNotNull(sandboxesTable.operationKind), + or( + isNull(sandboxesTable.claimOwner), + isNull(sandboxesTable.claimExpiresAt), + lte(sandboxesTable.claimExpiresAt, input.staleChangedAtLte), + ), ), ) .orderBy(asc(sandboxesTable.statusChangedAt), asc(sandboxesTable.id)) @@ -240,12 +227,18 @@ export async function listStaleRuntimeSubjectOperations( .all(); return rows.flatMap((row) => - row.operationId === null || !isRuntimeSubjectOperationStatus(row.status) + row.operationId === null || + row.operationKind === null || + !isRuntimeSubjectOperationStatus(row.status) ? [] : [ { + claimExpiresAt: row.claimExpiresAt, + claimOwner: row.claimOwner, id: row.id, + incarnation: row.incarnation, kind: row.kind, + operationKind: row.operationKind, operationId: row.operationId, status: row.status, }, @@ -253,6 +246,57 @@ export async function listStaleRuntimeSubjectOperations( ); } +export async function claimRuntimeSubjectOperationForRepair( + database: D1Database, + input: { + readonly candidate: RuntimeSubjectOperationRepairCandidate; + readonly claimExpiresAt: number; + readonly claimOwner: string; + readonly now: number; + }, +): Promise { + const row = await getAppDatabase(database) + .update(sandboxesTable) + .set({ + claimExpiresAt: input.claimExpiresAt, + claimOwner: input.claimOwner, + updatedAt: input.now, + }) + .where( + and( + eq(sandboxesTable.id, input.candidate.id), + eq(sandboxesTable.incarnation, input.candidate.incarnation), + eq(sandboxesTable.operationKind, input.candidate.operationKind), + eq(sandboxesTable.status, input.candidate.status), + eq(sandboxesTable.statusOperationId, input.candidate.operationId), + ...(input.candidate.claimOwner === null + ? [isNull(sandboxesTable.claimOwner)] + : [eq(sandboxesTable.claimOwner, input.candidate.claimOwner)]), + ...(input.candidate.claimExpiresAt === null + ? [isNull(sandboxesTable.claimExpiresAt)] + : [eq(sandboxesTable.claimExpiresAt, input.candidate.claimExpiresAt)]), + or( + isNull(sandboxesTable.claimOwner), + isNull(sandboxesTable.claimExpiresAt), + lte(sandboxesTable.claimExpiresAt, input.now), + ), + ), + ) + .returning({ id: sandboxesTable.id }) + .get(); + + return row === undefined + ? null + : { + claimExpiresAt: input.claimExpiresAt, + claimOwner: input.claimOwner, + incarnation: input.candidate.incarnation, + kind: input.candidate.operationKind, + operationId: input.candidate.operationId, + status: input.candidate.status, + }; +} + export async function claimInactiveRuntimeSubject( database: D1Database, input: { @@ -280,6 +324,7 @@ export async function claimInactiveRuntimeSubject( notExists(activeConversationSessionQuery(appDb, input.runtimeSubjectId)), ), notExists(runLeaseQuery(appDb, input.runtimeSubjectId)), + notExists(runtimeProvisioningQuery(appDb, input.runtimeSubjectId)), isNotNull(sandboxesTable.inactiveDeadlineAt), lte(sandboxesTable.inactiveDeadlineAt, input.now), or( diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance.service.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance.service.ts index 325dd347..e3a52193 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance.service.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance.service.ts @@ -1,34 +1,51 @@ import { sessionsTable } from "@mosoo/db"; import type { SandboxId, SessionId, SessionRunId } from "@mosoo/id"; -import { and, asc, eq, inArray, isNull, lte } from "drizzle-orm"; +import { and, asc, eq, isNull, lte } from "drizzle-orm"; import { createErrorLogContext, logWarn } from "../../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; import { getAppDatabase } from "../../../../platform/db/drizzle"; -import { isTruthy } from "../../../../shared/truthiness"; import { toIsoString } from "../../../../time"; import { repairStaleSessionDeleteCleanups } from "../../../sessions/application/session-cleanup.service"; -import { appendSessionRuntimeEvents } from "../../../sessions/application/session-event-write.service"; +import { publishPersistedSessionRuntimeEvents } from "../../../sessions/application/session-event-write.service"; +import { repairStaleSessionArchiveCleanups } from "../../../sessions/application/session-lifecycle-mutation.service"; import { syncSessionViewerState } from "../../../sessions/application/session-viewer-events.service"; import { RESCHEDULING_RECONNECT_WINDOW_MS } from "../../../sessions/domain/session-lifecycle"; +import { writeRuntimeOperationTimedOutSnapshots } from "../../application/runtime-state-operation-target-events"; +import { + commitSessionLifecycleEventProjection, + listStaleRuntimeOperationTargets, +} from "../../application/runtime-state-operation-target-store"; +import { recordCanonicalSessionRunTerminal } from "../../application/session-runs/session-run-terminal-failure.service"; import { createSessionLifecycleTerminatedEvent } from "../../application/session-runs/session-run-view-events.service"; import { reconcileStaleActiveSessionRuns } from "../../application/session-runs/stale-run-reconciliation.service"; import { reconcileTerminalSessionRuns } from "../../application/session-runs/terminal-run-reconciliation.service"; import { getRuntimeKindPolicy } from "../../domain/runtime-kind-policy"; import { cleanupDriverInstances } from "../driver-instance/maintenance"; +import { cleanupRuntimeArtifactAttempts } from "../driver-instance/runtime-artifact-attempt.repository"; +import { repairStagedSandboxBackups } from "../sandbox-backup.service"; import { repairRuntimeCommandRecords } from "../session-runs/runtime-command-store.repository"; -import { createSessionStatusTransitionPatch } from "../session-runs/session-lifecycle-projection.repository"; -import { setSessionRunStatus } from "../session-runs/session-run-store.repository"; -import type { SessionRunTransitionOutcome } from "../session-runs/session-run-store.repository"; +import { getSessionRunSummariesByIds } from "../session-runs/session-run-store.repository"; import { listIdleSessionScopedConversationSessions } from "./runtime-conversation-session-store"; +import { + cleanupRuntimeProvisioningResources, + retireRuntimeProvisioningIncarnation, +} from "./runtime-provisioning-cleanup.service"; +import { + adoptReadyRuntimeRunProvisioningLease, + claimStaleRuntimeProvisioningLeases, + releaseAbortedRuntimeProvisioningLease, +} from "./runtime-provisioning-lease-store"; import { repairStrandedRuntimeSubjectDeadlines } from "./runtime-subject-maintenance-store"; import { claimInactiveRuntimeSubject, + claimRuntimeSubjectOperationForRepair, listInactiveRuntimeSubjects, listStaleRuntimeSubjectOperations, } from "./runtime-subject-store"; import type { RuntimeSubjectMaintenanceCandidate, + RuntimeSubjectOperationLease, RuntimeSubjectOperationRepairCandidate, } from "./runtime-subject-store"; @@ -51,16 +68,18 @@ type ResumeRuntimeSubjectRecycleOperation = ( bindings: ApiBindings, input: { readonly kind: RuntimeSubjectOperationRepairCandidate["kind"]; - readonly operationId: RuntimeSubjectOperationRepairCandidate["operationId"]; + readonly lease: RuntimeSubjectOperationLease; readonly reason: string; readonly runtimeSubjectId: RuntimeSubjectOperationRepairCandidate["id"]; - readonly status: RuntimeSubjectOperationRepairCandidate["status"]; }, ) => Promise; interface StaleReschedulingSessionRow { id: SessionId; last_run_id: SessionRunId | null; + runtime_event_seq_cursor: number; + status_seq: number; + updated_at: number; } const RESCHEDULING_TIMEOUT_ERROR = { @@ -70,34 +89,23 @@ const RESCHEDULING_TIMEOUT_ERROR = { retryable: false, } as const; -function assertMaintenanceRunTransition(outcome: SessionRunTransitionOutcome): void { - switch (outcome.kind) { - case "applied": - case "duplicate": { - return; - } - case "stale": { - if (outcome.reason === "terminal_run") { - return; - } - throw new Error("Rescheduling timeout lost a concurrent run transition."); - } - case "repair_needed": { - throw new Error("Rescheduling timeout left session projection stale."); - } - case "rejected": { - throw new Error(`Rescheduling timeout run transition was rejected: ${outcome.reason}.`); - } - } -} - async function processInBatches( items: readonly T[], batchSize: number, task: (item: T) => Promise, + onRejected: (item: T, reason: unknown) => void, ): Promise { for (let index = 0; index < items.length; index += batchSize) { - await Promise.all(items.slice(index, index + batchSize).map(task)); + const batch = items.slice(index, index + batchSize); + const outcomes = await Promise.allSettled(batch.map(task)); + for (const [outcomeIndex, outcome] of outcomes.entries()) { + if (outcome.status === "rejected") { + const item = batch[outcomeIndex]; + if (item !== undefined) { + onRejected(item, outcome.reason); + } + } + } } } @@ -146,13 +154,23 @@ async function repairRuntimeSubjectOperationCandidate( readonly resumeRuntimeSubjectRecycleOperation: ResumeRuntimeSubjectRecycleOperation; }, ): Promise { + const now = Date.now(); + const lease = await claimRuntimeSubjectOperationForRepair(bindings.DB, { + candidate: input.candidate, + claimExpiresAt: now + MAINTENANCE_CLAIM_TTL_MS, + claimOwner: `repair-${crypto.randomUUID()}`, + now, + }); + if (lease === null) { + return; + } + try { await input.resumeRuntimeSubjectRecycleOperation(bindings, { kind: input.candidate.kind, - operationId: input.candidate.operationId, + lease, reason: input.reason, runtimeSubjectId: input.candidate.id, - status: input.candidate.status, }); } catch (error) { logWarn("runtime.subject.maintenance.operation_repair_failed", { @@ -164,23 +182,37 @@ async function repairRuntimeSubjectOperationCandidate( } } -async function publishReschedulingTimeoutEvent( +async function commitReschedulingTimeoutProjection( bindings: ApiBindings, target: StaleReschedulingSessionRow, ): Promise { - const stoppedAt = Date.now(); + const stoppedAt = target.updated_at + RESCHEDULING_RECONNECT_WINDOW_MS; const event = createSessionLifecycleTerminatedEvent({ lastSeen: toIsoString(stoppedAt), message: RESCHEDULING_TIMEOUT_ERROR.message, + occurredAtMs: stoppedAt, reason: RESCHEDULING_TIMEOUT_ERROR.code, sessionId: target.id, + sourceEventId: `maintenance:rescheduling-timeout:${target.id}`, }); - await appendSessionRuntimeEvents({ - bindings, - events: [event], - sessionId: target.id, + const outcome = await commitSessionLifecycleEventProjection(bindings.DB, { + event, + status: "TERMINATED", + target: { + lastRunId: target.last_run_id, + sessionId: target.id, + sessionRuntimeEventSeqCursor: target.runtime_event_seq_cursor, + sessionStatus: "RESCHEDULING", + sessionStatusOperationId: null, + sessionStatusSeq: target.status_seq, + sessionUpdatedAt: target.updated_at, + }, + timestampMs: stoppedAt, }); + if (outcome.kind !== "stale") { + await publishPersistedSessionRuntimeEvents({ bindings, events: [event], sessionId: target.id }); + } } export async function expireStaleReschedulingSessions(bindings: ApiBindings): Promise { @@ -188,11 +220,16 @@ export async function expireStaleReschedulingSessions(bindings: ApiBindings): Pr const staleSessions = await getAppDatabase(bindings.DB) .select({ id: sessionsTable.id, + last_run_id: sessionsTable.lastRunId, + runtime_event_seq_cursor: sessionsTable.runtimeEventSeqCursor, + status_seq: sessionsTable.statusSeq, + updated_at: sessionsTable.updatedAt, }) .from(sessionsTable) .where( and( eq(sessionsTable.status, "RESCHEDULING"), + isNull(sessionsTable.archivedAt), isNull(sessionsTable.statusOperationId), lte(sessionsTable.updatedAt, now - RESCHEDULING_RECONNECT_WINDOW_MS), ), @@ -205,49 +242,77 @@ export async function expireStaleReschedulingSessions(bindings: ApiBindings): Pr return; } - const results = await getAppDatabase(bindings.DB) - .update(sessionsTable) - .set( - createSessionStatusTransitionPatch({ - status: "TERMINATED", - timestampMs: now, - }), - ) - .where( - and( - inArray( - sessionsTable.id, - staleSessions.map((session) => session.id), - ), - eq(sessionsTable.status, "RESCHEDULING"), - isNull(sessionsTable.statusOperationId), - lte(sessionsTable.updatedAt, now - RESCHEDULING_RECONNECT_WINDOW_MS), - ), - ) - .returning({ - id: sessionsTable.id, - last_run_id: sessionsTable.lastRunId, - }) - .all(); + const runIds = staleSessions.flatMap((target) => + target.last_run_id === null ? [] : [target.last_run_id], + ); + const runsById = await getSessionRunSummariesByIds(bindings.DB, runIds); await processInBatches( - results.map((target) => target.last_run_id).filter(isTruthy), + staleSessions, RESCHEDULING_TIMEOUT_IO_BATCH_SIZE, - async (runId) => { - const outcome = await setSessionRunStatus(bindings.DB, { - error: RESCHEDULING_TIMEOUT_ERROR, - preserveSessionLifecycle: true, - runId, - source: "maintenance", - status: "failed", + async (target) => { + const run = target.last_run_id === null ? null : (runsById.get(target.last_run_id) ?? null); + if (run !== null && !["cancelled", "completed", "expired", "failed"].includes(run.status)) { + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + error: RESCHEDULING_TIMEOUT_ERROR, + expectedSessionOperationId: null, + expectedSessionObservation: { + lastRunId: target.last_run_id, + status: "RESCHEDULING", + statusSeq: target.status_seq, + updatedAt: target.updated_at, + }, + lifecycle: "TERMINATED", + runId: run.id, + sessionId: target.id, + source: "maintenance", + status: "failed", + timestampMs: target.updated_at + RESCHEDULING_RECONNECT_WINDOW_MS, + }); + return; + } + + await commitReschedulingTimeoutProjection(bindings, target); + }, + (target, reason) => { + logWarn("runtime.session.rescheduling_timeout_repair_failed", { + ...createErrorLogContext(reason), + sessionId: target.id, }); - assertMaintenanceRunTransition(outcome); }, ); +} + +export async function repairStaleRuntimeOperationTargets( + bindings: ApiBindings, + input: { + readonly limit: number; + readonly staleUpdatedAtLte: number; + }, +): Promise { + const targets = await listStaleRuntimeOperationTargets(bindings.DB, input); - await processInBatches(results, RESCHEDULING_TIMEOUT_IO_BATCH_SIZE, async (target) => - publishReschedulingTimeoutEvent(bindings, target), + await Promise.all( + [...Map.groupBy(targets, (target) => target.operationId)].map( + async ([operationId, operationTargets]) => { + try { + await writeRuntimeOperationTimedOutSnapshots(bindings, { + operationId, + targets: operationTargets, + }); + } catch (error) { + logWarn("runtime.operation.timeout_repair_failed", { + ...createErrorLogContext(error), + operationId, + sessionIds: operationTargets.map((target) => target.sessionId), + }); + } + }, + ), ); + + return targets.length; } // Cattle conversations no longer close on run terminal (the resident driver is @@ -286,11 +351,84 @@ async function closeIdleSessionScopedConversationSessions( } } +async function repairPendingConversationSessionCleanups(bindings: ApiBindings): Promise { + const { repairPendingSandboxConversationSessionCleanups } = + await import("../sandbox-session.service"); + try { + await repairPendingSandboxConversationSessionCleanups(bindings, MAINTENANCE_BATCH_SIZE); + } catch (error) { + logWarn("runtime.conversation.cleanup_repair_failed", createErrorLogContext(error)); + } +} + +export async function repairStaleRuntimeProvisioningLeases( + bindings: ApiBindings, + input: { readonly heartbeatAtLte: number; readonly limit: number }, +): Promise { + const claims = await claimStaleRuntimeProvisioningLeases(bindings.DB, input); + + await Promise.allSettled( + claims.map(async (claim) => { + try { + if (await adoptReadyRuntimeRunProvisioningLease(bindings.DB, claim)) { + return; + } + if (claim.sandboxIncarnation === null) { + await cleanupRuntimeProvisioningResources(bindings, claim, "maintenance"); + if (!(await releaseAbortedRuntimeProvisioningLease(bindings.DB, claim))) { + throw new Error("Runtime provisioning repair lost its maintenance ownership."); + } + } else { + await retireRuntimeProvisioningIncarnation(bindings, claim, "maintenance"); + } + } catch (error) { + logWarn("runtime.provisioning.repair_failed", { + ...createErrorLogContext(error), + operationId: claim.operationId, + runId: claim.runId, + sandboxId: claim.sandboxId, + sessionId: claim.sessionId, + }); + } + }), + ); + + return claims.length; +} + export async function runSandboxMaintenance(bindings: ApiBindings): Promise { const now = Date.now(); - await repairRuntimeCommandRecords(bindings.DB, { nowMs: now }); + await repairStaleRuntimeProvisioningLeases(bindings, { + heartbeatAtLte: now - MAINTENANCE_OPERATION_REPAIR_AFTER_MS, + limit: MAINTENANCE_BATCH_SIZE, + }); + await repairStagedSandboxBackups(bindings, MAINTENANCE_BATCH_SIZE); + const cleanupRepairs = await Promise.allSettled([ + repairStaleSessionArchiveCleanups(bindings, { + limit: MAINTENANCE_BATCH_SIZE, + staleUpdatedAtLte: now - MAINTENANCE_OPERATION_REPAIR_AFTER_MS, + }), + repairStaleSessionDeleteCleanups(bindings, { + limit: MAINTENANCE_BATCH_SIZE, + staleUpdatedAtLte: now - MAINTENANCE_OPERATION_REPAIR_AFTER_MS, + }), + ]); + const cleanupRepairFailure = cleanupRepairs.find( + (result): result is PromiseRejectedResult => result.status === "rejected", + ); + if (cleanupRepairFailure !== undefined) { + logWarn( + "runtime.session.cleanup_repair_failed", + createErrorLogContext(cleanupRepairFailure.reason), + ); + } + await repairStaleRuntimeOperationTargets(bindings, { + limit: MAINTENANCE_BATCH_SIZE, + staleUpdatedAtLte: now - RESCHEDULING_RECONNECT_WINDOW_MS, + }); await cleanupDriverInstances(bindings); + await repairRuntimeCommandRecords(bindings.DB, { nowMs: now }); const staleRunReconciliation = await reconcileStaleActiveSessionRuns(bindings.DB, { limit: MAINTENANCE_BATCH_SIZE, }); @@ -306,12 +444,15 @@ export async function runSandboxMaintenance(bindings: ApiBindings): Promise syncSessionViewerState(bindings, sessionId), + (sessionId, reason) => { + logWarn("runtime.session.viewer_state_repair_failed", { + ...createErrorLogContext(reason), + sessionId, + }); + }, ); await expireStaleReschedulingSessions(bindings); - await repairStaleSessionDeleteCleanups(bindings, { - limit: MAINTENANCE_BATCH_SIZE, - staleUpdatedAtLte: now - MAINTENANCE_OPERATION_REPAIR_AFTER_MS, - }); + await repairPendingConversationSessionCleanups(bindings); await closeIdleSessionScopedConversationSessions(bindings, now); const repairedDeadlines = await repairStrandedRuntimeSubjectDeadlines(bindings.DB, { now }); @@ -332,10 +473,16 @@ export async function runSandboxMaintenance(bindings: ApiBindings): Promise; + +function operationWorkerOwner(prefix: string): string { + return `${prefix}-${crypto.randomUUID()}`; +} + function getRuntimeOperationErrorMessage(error: unknown): string { return error instanceof Error ? error.message : "Runtime state operation failed."; } @@ -53,214 +71,392 @@ async function appendCheckpointFailureDiagnostics( } await appendOneRuntimeDiagnosticEventPerSession(bindings, { - events: input.targets.flatMap((target) => { - if (target.agentId === null) { - return []; - } - - return [ - { - eventName: RUNTIME_DIAGNOSTIC_EVENT.sandboxCheckpointFailed.name, - sessionId: target.sessionId, - value: { - ...toRuntimeDiagnosticBaseValue({ - agentId: target.agentId, + events: input.targets.flatMap((target) => + target.agentId === null + ? [] + : [ + { + eventName: RUNTIME_DIAGNOSTIC_EVENT.sandboxCheckpointFailed.name, sessionId: target.sessionId, - }), - backupId: - input.error instanceof RuntimeSubjectCheckpointFailedError - ? input.error.backupId - : null, - dir: - input.error instanceof RuntimeSubjectCheckpointFailedError ? input.error.dir : null, - errorCode, - reason: toRuntimeDiagnosticReason(input.error, "Runtime subject checkpoint failed."), - sandboxId: input.runtimeSubjectId, - }, - }, - ]; - }), + value: { + ...toRuntimeDiagnosticBaseValue({ + agentId: target.agentId, + sessionId: target.sessionId, + }), + backupId: + input.error instanceof RuntimeSubjectCheckpointFailedError + ? input.error.backupId + : null, + dir: + input.error instanceof RuntimeSubjectCheckpointFailedError + ? input.error.dir + : null, + errorCode, + reason: toRuntimeDiagnosticReason( + input.error, + "Runtime subject checkpoint failed.", + ), + sandboxId: input.runtimeSubjectId, + }, + }, + ], + ), }); } -async function appendTerminatedEventsForRuntimeSubject( +function checkpointRules(kind: AgentKind, operationKind: DestructiveRuntimeSubjectOperationKind) { + const checkpoint = getRuntimeKindPolicy(kind).checkpoint; + + switch (operationKind) { + case "hibernate": + return checkpoint.createOnHibernate; + case "recreate": + return checkpoint.createOnRecreate; + case "reset": + return checkpoint.createOnReset; + } +} + +function clearsBackups(kind: AgentKind, operationKind: SandboxOperationKind): boolean { + if (operationKind === "reset") { + return true; + } + return operationKind === "recreate" + ? getRuntimeKindPolicy(kind).checkpoint.createOnRecreate.length === 0 + : false; +} + +type RuntimeSubjectPhysicalState = "gone" | "healthy" | "unknown"; + +async function inspectOperationRuntimeSubjectPhysicalState( bindings: ApiBindings, input: { - readonly reason: string; + readonly lease: RuntimeSubjectOperationLease; readonly runtimeSubjectId: SandboxId; - readonly targets: RuntimeSubjectOperationInput["targets"]; }, -): Promise { - await appendRuntimeSubjectTerminatedEvents(bindings, { - reason: input.reason, - runtimeSubjectId: input.runtimeSubjectId, - targets: input.targets, - }); +): Promise { + const record = await getRuntimeSubject(bindings.DB, input.runtimeSubjectId); + if ( + record === null || + record.incarnation !== input.lease.incarnation || + record.networkConstraintsHash === null + ) { + return "unknown"; + } + + const subject = await getRuntimeSubjectKeepAliveHandle( + bindings, + input.runtimeSubjectId, + input.lease.incarnation, + ); + try { + const result = await inspectRuntimeSubjectIncarnation( + subject, + input.lease.incarnation, + record.networkConstraintsHash, + ); + switch (result.kind) { + case "healthy": + return "healthy"; + case "missing": + case "retired": + return "gone"; + case "stale": + case "unknown": + return "unknown"; + } + } finally { + disposeRpcResource(subject); + } } -export async function recreateRuntimeSubjectPreservingState( +export async function runRuntimeSubjectOperation( bindings: ApiBindings, - input: RuntimeSubjectOperationInput, + input: { + readonly kind: AgentKind; + readonly lease: RuntimeSubjectOperationLease; + readonly reason: string; + readonly runtimeSubjectId: SandboxId; + readonly targets?: RuntimeSubjectOperationInput["targets"]; + }, ): Promise { - const subject = await getRuntimeSubject(bindings.DB, input.runtimeSubjectId); + let lease = input.lease; + let physicalStateLost = false; + const renew = createLeaseOwnershipRenewal( + () => + renewRuntimeSubjectOperationLease(bindings.DB, { + claimExpiresAt: Date.now() + RUNTIME_SUBJECT_OPERATION_LEASE_TTL_MS, + lease, + runtimeSubjectId: input.runtimeSubjectId, + }), + "Runtime subject operation lost lifecycle ownership.", + ); + const heartbeat = setInterval(() => { + void renew().catch(() => undefined); + }, RUNTIME_SUBJECT_OPERATION_HEARTBEAT_MS); - if (!subject) { - return; - } + try { + await renew(); - let destroyStarted = false; + if (lease.status === "restoring") { + const advanced = await advanceRuntimeSubjectOperationStatus(bindings.DB, { + expectedStatus: "restoring", + lease, + runtimeSubjectId: input.runtimeSubjectId, + source: "maintenance", + status: "destroying", + }); + if (!advanced) { + throw new Error("Runtime subject activation repair lost lifecycle ownership."); + } + lease = { ...lease, status: "destroying" }; + } - const policy = getRuntimeKindPolicy(subject.kind); - const started = await markRuntimeSubjectOperationStarted(bindings.DB, { - operationId: input.operationId, - runtimeSubjectId: input.runtimeSubjectId, - status: "backing_up", - }); + if ( + lease.kind === "activate" && + !(await runtimeSubjectActivationRetirementIsDrained(bindings.DB, { + lease, + runtimeSubjectId: input.runtimeSubjectId, + })) + ) { + throw new Error("Runtime subject incarnation is still draining active Runs."); + } - if (!started) { - throw new Error("Runtime subject is busy with lifecycle maintenance."); - } + if (lease.status === "backing_up") { + if (lease.kind === "activate") { + throw new Error("Activation repair cannot enter the backup phase."); + } - try { - await stopRuntimeSubjectDrivers(bindings, { - operationId: input.operationId, - runtimeSubjectId: input.runtimeSubjectId, - preserveSessionLifecycle: true, - reason: input.reason, - targets: input.targets, - terminalRun: input.terminalRun, - }); - await createSandboxCheckpoints(bindings, { - operationId: input.operationId, - rules: policy.checkpoint.createOnRecreate, - sandboxId: input.runtimeSubjectId, - }); - destroyStarted = await advanceRuntimeSubjectOperationStatus(bindings.DB, { - expectedStatus: "backing_up", - operationId: input.operationId, - runtimeSubjectId: input.runtimeSubjectId, - status: "destroying", - }); - if (!destroyStarted) { - throw new Error("Runtime subject changed before destroy."); + const before = await inspectOperationRuntimeSubjectPhysicalState(bindings, { + lease, + runtimeSubjectId: input.runtimeSubjectId, + }); + if (before === "unknown") { + throw new Error("Runtime subject physical state is unknown before checkpoint."); + } + physicalStateLost = before === "gone"; + + if (!physicalStateLost) { + try { + await renew(); + await stopRuntimeSubjectDrivers(bindings, { + operationId: lease.operationId, + reason: input.reason, + runtimeSubjectId: input.runtimeSubjectId, + sandboxIncarnation: lease.incarnation, + ...(input.targets === undefined ? {} : { targets: input.targets }), + }); + await renew(); + + if (lease.kind === "reset") { + const stateTargets = await listRuntimeSubjectSessionStateTargets(bindings.DB, { + runtimeSubjectId: input.runtimeSubjectId, + }); + await clearRuntimeSubjectAgentState(bindings, { + incarnation: lease.incarnation, + rules: getRuntimeKindPolicy(input.kind).checkpoint.clearOnReset, + runtimeSubjectId: input.runtimeSubjectId, + stateTargets, + }); + await renew(); + } + + await createSandboxCheckpoints(bindings, { + operationLease: lease, + rules: checkpointRules(input.kind, lease.kind), + sandboxId: input.runtimeSubjectId, + }); + await renew(); + } catch (error) { + let after: RuntimeSubjectPhysicalState = "unknown"; + try { + after = await inspectOperationRuntimeSubjectPhysicalState(bindings, { + lease, + runtimeSubjectId: input.runtimeSubjectId, + }); + } catch { + // Preserve the original failure unless the exact incarnation is + // durably known to be gone. + } + physicalStateLost = after === "gone"; + if (!physicalStateLost) { + throw error; + } + } + } + + const advanced = await advanceRuntimeSubjectOperationStatus(bindings.DB, { + expectedStatus: "backing_up", + lease, + runtimeSubjectId: input.runtimeSubjectId, + source: "maintenance", + status: "destroying", + }); + if (!advanced) { + throw new Error("Runtime subject changed before destroy."); + } + lease = { ...lease, status: "destroying" }; + } + + if (lease.kind === "activate") { + await renew(); + await stopRuntimeSubjectDrivers(bindings, { + operationId: lease.operationId, + reason: input.reason, + runtimeSubjectId: input.runtimeSubjectId, + sandboxIncarnation: lease.incarnation, + }); + await renew(); + } + + await renew(); + await destroyRuntimeSubjectContainer(bindings, input.runtimeSubjectId, lease.incarnation); + await renew(); + + if (input.targets !== undefined) { + await appendRuntimeSubjectTerminatedEvents(bindings, { + reason: input.reason, + runtimeSubjectId: input.runtimeSubjectId, + targets: input.targets, + }); } - await destroyRuntimeSubjectContainer(bindings, input.runtimeSubjectId); - await appendTerminatedEventsForRuntimeSubject(bindings, { - reason: input.reason, - runtimeSubjectId: input.runtimeSubjectId, - targets: input.targets, - }); - await closeRuntimeSubjectSessionsForRecycle(bindings.DB, input.runtimeSubjectId); const completed = await markRuntimeSubjectCold(bindings.DB, { - clearBackups: policy.checkpoint.createOnRecreate.length === 0, + clearBackups: !physicalStateLost && clearsBackups(input.kind, lease.kind), + clearNativeResumeRefs: !physicalStateLost && lease.kind === "reset", + ...(physicalStateLost + ? { + errorCode: "runtime.subject_operation_failed" as const, + errorMessage: + "Runtime subject physical incarnation was lost before its checkpoint completed.", + } + : {}), expectedStatus: "destroying", - operationId: input.operationId, + lease, runtimeSubjectId: input.runtimeSubjectId, + source: "maintenance", }); if (!completed) { - throw new Error("Runtime subject changed before recreate completion."); + throw new Error("Runtime subject changed before operation completion."); + } + if (physicalStateLost) { + throw new RuntimeSubjectPhysicalStateLostError(input.runtimeSubjectId); } } catch (error) { - await appendCheckpointFailureDiagnostics(bindings, { - error, - runtimeSubjectId: input.runtimeSubjectId, - targets: input.targets, - }); - await markRuntimeSubjectFailed(bindings.DB, { - errorCode: getRuntimeSubjectOperationErrorCode(error), - errorMessage: getRuntimeOperationErrorMessage(error), - expectedStatus: destroyStarted ? "destroying" : "backing_up", - operationId: input.operationId, - runtimeSubjectId: input.runtimeSubjectId, - status: "cold", - }); + if (!(error instanceof RuntimeSubjectPhysicalStateLostError)) { + await markRuntimeSubjectOperationRepairNeeded(bindings.DB, { + errorCode: getRuntimeSubjectOperationErrorCode(error), + errorMessage: getRuntimeOperationErrorMessage(error), + expectedStatus: lease.status, + lease, + runtimeSubjectId: input.runtimeSubjectId, + source: "maintenance", + }); + } throw error; + } finally { + clearInterval(heartbeat); } } -export async function resetRuntimeSubjectAgentState( +async function startRuntimeSubjectOperation( bindings: ApiBindings, - input: RuntimeSubjectOperationInput, + input: { + readonly operationId: RuntimeOperationId; + readonly operationKind: DestructiveRuntimeSubjectOperationKind; + readonly runtimeSubjectId: SandboxId; + }, +): Promise { + const now = Date.now(); + const lease = await markRuntimeSubjectOperationStarted(bindings.DB, { + claimExpiresAt: now + RUNTIME_SUBJECT_OPERATION_LEASE_TTL_MS, + claimOwner: operationWorkerOwner("runtime-operation"), + now, + operationId: input.operationId, + operationKind: input.operationKind, + runtimeSubjectId: input.runtimeSubjectId, + source: "runtime", + }); + if (lease === null) { + throw new Error("Runtime subject is busy with lifecycle maintenance."); + } + return lease; +} + +async function executeRequestedRuntimeSubjectOperation( + bindings: ApiBindings, + input: RuntimeSubjectOperationInput & { + readonly operationKind: "recreate" | "reset"; + }, ): Promise { const subject = await getRuntimeSubject(bindings.DB, input.runtimeSubjectId); - - if (!subject) { + if (subject === null) { return; } - - const policy = getRuntimeKindPolicy(subject.kind); - - if (!policy.operations.resetSubjectState) { + if ( + input.operationKind === "reset" && + !getRuntimeKindPolicy(subject.kind).operations.resetSubjectState + ) { throw new Error("This runtime kind does not have resettable subject state."); } - const started = await markRuntimeSubjectOperationStarted(bindings.DB, { - operationId: input.operationId, - runtimeSubjectId: input.runtimeSubjectId, - status: "destroying", - }); + if (subject.status === "cold") { + if (input.operationKind === "recreate") { + return; + } + if (subject.agentId === null || subject.appId === null || subject.ownerAccountId === null) { + throw new Error("Cold runtime subject has no complete activation identity."); + } - if (!started) { + const activated = await createRuntimeSubjectLifecycleService(bindings).activate({ + agentId: subject.agentId, + appId: subject.appId, + executionOwnerUserId: subject.ownerAccountId, + kind: subject.kind, + networkConstraints: { allowedHosts: [], networkPolicy: "full" }, + runtimeSubjectId: subject.id, + subjectId: subject.subjectId, + subjectKind: subject.subjectKind, + }); + disposeRpcResource(activated.subject); + } else if (subject.status !== "active") { throw new Error("Runtime subject is busy with lifecycle maintenance."); } + const lease = await startRuntimeSubjectOperation(bindings, input); try { - await stopRuntimeSubjectDrivers(bindings, { - operationId: input.operationId, - runtimeSubjectId: input.runtimeSubjectId, - preserveSessionLifecycle: true, - reason: input.reason, - targets: input.targets, - terminalRun: input.terminalRun, - }); - const stateTargets = await listRuntimeSubjectSessionStateTargets(bindings.DB, { - runtimeSubjectId: input.runtimeSubjectId, - sessionIds: input.targets.map((target) => target.sessionId), - }); - await clearRuntimeSubjectAgentState(bindings, { - runtimeSubjectId: input.runtimeSubjectId, - rules: policy.checkpoint.clearOnReset, - stateTargets, - }); - await deleteNativeResumeRefsForSessions( - bindings.DB, - input.targets.map((target) => target.sessionId), - ); - await createSandboxCheckpoints(bindings, { - operationId: input.operationId, - rules: policy.checkpoint.createOnReset, - sandboxId: input.runtimeSubjectId, - }); - await destroyRuntimeSubjectContainer(bindings, input.runtimeSubjectId); - await appendTerminatedEventsForRuntimeSubject(bindings, { + await runRuntimeSubjectOperation(bindings, { + kind: subject.kind, + lease, reason: input.reason, runtimeSubjectId: input.runtimeSubjectId, targets: input.targets, }); - await closeRuntimeSubjectSessionsForRecycle(bindings.DB, input.runtimeSubjectId); - const completed = await markRuntimeSubjectCold(bindings.DB, { - clearBackups: true, - expectedStatus: "destroying", - operationId: input.operationId, - runtimeSubjectId: input.runtimeSubjectId, - }); - if (!completed) { - throw new Error("Runtime subject changed before reset completion."); - } } catch (error) { await appendCheckpointFailureDiagnostics(bindings, { error, runtimeSubjectId: input.runtimeSubjectId, targets: input.targets, }); - await markRuntimeSubjectFailed(bindings.DB, { - errorCode: getRuntimeSubjectOperationErrorCode(error), - errorMessage: getRuntimeOperationErrorMessage(error), - expectedStatus: "destroying", - operationId: input.operationId, - runtimeSubjectId: input.runtimeSubjectId, - status: "cold", - }); throw error; } } + +export function recreateRuntimeSubjectPreservingState( + bindings: ApiBindings, + input: RuntimeSubjectOperationInput, +): Promise { + return executeRequestedRuntimeSubjectOperation(bindings, { + ...input, + operationKind: "recreate", + }); +} + +export function resetRuntimeSubjectAgentState( + bindings: ApiBindings, + input: RuntimeSubjectOperationInput, +): Promise { + return executeRequestedRuntimeSubjectOperation(bindings, { + ...input, + operationKind: "reset", + }); +} diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-platform.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-platform.ts index dc3968e0..47735af1 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-platform.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-platform.ts @@ -11,27 +11,56 @@ import { } from "../../../../platform/cloudflare/rpc-disposal"; import { requireCloudflareSandboxBinding } from "../../../../platform/cloudflare/sandbox-binding"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; +import { quoteShellArg } from "../../../../shared/shell"; import type { RuntimeStateClearRule } from "../../domain/runtime-kind-policy"; import type { SandboxNetworkConstraints } from "../../domain/sandbox-network-constraints"; import { withRuntimeProvisionTimeout } from "../runtime-provision-timeout"; import { decodeSandboxBackupIdForPlatform } from "../sandbox-backup-id"; -import { toSandboxHandle } from "../sandbox-handles"; +import { toRuntimeSubjectIncarnationHandle, toSandboxHandle } from "../sandbox-handles"; import type { SandboxHandle } from "../sandbox-handles"; import type { ReadyRuntimeSubjectBackupRecord } from "./runtime-subject-store"; -function quoteShellArg(value: string): string { - return `'${value.replaceAll("'", `'"'"'`)}'`; -} - export function getRuntimeSubjectKeepAliveHandle( bindings: ApiBindings, runtimeSubjectId: string, + incarnation: number, +): Promise { + const physicalId = + incarnation === 0 ? runtimeSubjectId : `${runtimeSubjectId}-i${incarnation.toString(36)}`; + if (bindings.runtimeSubjectHandleFactory) { + return Promise.resolve(toSandboxHandle(bindings.runtimeSubjectHandleFactory(physicalId))); + } + + return getCloudflareRuntimeSubjectKeepAliveHandle(bindings, physicalId); +} + +function getUnversionedSandboxHandle( + bindings: ApiBindings, + sandboxId: string, +): Promise { + if (bindings.runtimeSubjectHandleFactory) { + return Promise.resolve(toSandboxHandle(bindings.runtimeSubjectHandleFactory(sandboxId))); + } + return getCloudflareRuntimeSubjectKeepAliveHandle(bindings, sandboxId); +} + +export function getEphemeralUnversionedSandboxHandle( + bindings: ApiBindings, + sandboxId: string, + sleepAfter: string | number, ): Promise { if (bindings.runtimeSubjectHandleFactory) { - return Promise.resolve(toSandboxHandle(bindings.runtimeSubjectHandleFactory(runtimeSubjectId))); + return Promise.resolve(toSandboxHandle(bindings.runtimeSubjectHandleFactory(sandboxId))); } + return getCloudflareEphemeralSandboxHandle(bindings, sandboxId, sleepAfter); +} - return getCloudflareRuntimeSubjectKeepAliveHandle(bindings, runtimeSubjectId); +export function createEphemeralSandboxOptions(sleepAfter: string | number): { + readonly keepAlive: false; + readonly normalizeId: true; + readonly sleepAfter: string | number; +} { + return { keepAlive: false, normalizeId: true, sleepAfter }; } async function getCloudflareRuntimeSubjectKeepAliveHandle( @@ -46,6 +75,21 @@ async function getCloudflareRuntimeSubjectKeepAliveHandle( return toSandboxHandle(sandbox); } +async function getCloudflareEphemeralSandboxHandle( + bindings: ApiBindings, + sandboxId: string, + sleepAfter: string | number, +): Promise { + const { getSandbox } = await import("@cloudflare/sandbox"); + return toSandboxHandle( + getSandbox( + requireCloudflareSandboxBinding(bindings), + sandboxId, + createEphemeralSandboxOptions(sleepAfter), + ), + ); +} + /** * Pushes the environment egress policy into the sandbox Durable Object. Must * run before any container-starting call (mkdir/exec/session create): the @@ -83,6 +127,48 @@ export async function prepareRuntimeSubjectFilesystem(subject: SandboxHandle): P ); } +export async function activateRuntimeSubjectIncarnation( + subject: SandboxHandle, + incarnation: number, + networkConstraintsHash: string, +): Promise { + await withRuntimeProvisionTimeout( + toRuntimeSubjectIncarnationHandle(subject).activateRuntimeSubjectIncarnation( + incarnation, + networkConstraintsHash, + ), + `Runtime subject incarnation ${incarnation} activation`, + ); +} + +export async function inspectRuntimeSubjectIncarnation( + subject: SandboxHandle, + incarnation: number, + networkConstraintsHash: string, +): Promise<{ kind: "healthy" | "missing" | "retired" | "stale" | "unknown" }> { + return withRuntimeProvisionTimeout( + toRuntimeSubjectIncarnationHandle(subject).inspectRuntimeSubjectIncarnation( + incarnation, + networkConstraintsHash, + ), + `Runtime subject incarnation ${incarnation} inspection`, + ); +} + +export async function markRuntimeSubjectIncarnationReady( + subject: SandboxHandle, + incarnation: number, + networkConstraintsHash: string, +): Promise { + await withRuntimeProvisionTimeout( + toRuntimeSubjectIncarnationHandle(subject).markRuntimeSubjectIncarnationReady( + incarnation, + networkConstraintsHash, + ), + `Runtime subject incarnation ${incarnation} readiness`, + ); +} + export async function restoreRuntimeSubjectBackup( subject: SandboxHandle, input: { @@ -105,15 +191,21 @@ export async function restoreRuntimeSubjectBackup( export async function destroyRuntimeSubjectContainer( bindings: ApiBindings, runtimeSubjectId: string, + incarnation: number, timeoutMs?: number, ): Promise { await withRuntimeProvisionTimeout( (async () => withDisposedRpcResource( - await getRuntimeSubjectKeepAliveHandle(bindings, runtimeSubjectId), + await getRuntimeSubjectKeepAliveHandle(bindings, runtimeSubjectId, incarnation), async (subject) => { - await subject.setKeepAlive(false); - await subject.destroy(); + const outcome = + await toRuntimeSubjectIncarnationHandle(subject).destroyRuntimeSubjectIncarnation( + incarnation, + ); + if (outcome.kind === "stale") { + throw new Error("Runtime subject destroy targeted a stale incarnation."); + } }, ))(), `Runtime subject destroy for ${runtimeSubjectId}`, @@ -121,16 +213,35 @@ export async function destroyRuntimeSubjectContainer( ); } +export async function destroyUnversionedSandboxContainer( + bindings: ApiBindings, + sandboxId: string, + timeoutMs?: number, +): Promise { + await withRuntimeProvisionTimeout( + withDisposedRpcResource( + await getUnversionedSandboxHandle(bindings, sandboxId), + async (sandbox) => { + await sandbox.setKeepAlive(false); + await sandbox.destroy(); + }, + ), + `Sandbox destroy for ${sandboxId}`, + timeoutMs, + ); +} + export async function clearRuntimeSubjectAgentState( bindings: ApiBindings, input: { readonly rules: readonly RuntimeStateClearRule[]; + readonly incarnation: number; readonly runtimeSubjectId: string; readonly stateTargets: readonly string[]; }, ): Promise { await withDisposedRpcResource( - await getRuntimeSubjectKeepAliveHandle(bindings, input.runtimeSubjectId), + await getRuntimeSubjectKeepAliveHandle(bindings, input.runtimeSubjectId, input.incarnation), async (subject) => { const commands = input.rules.flatMap((rule) => { switch (rule.type) { diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-record-store.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-record-store.ts index 8f02a2fc..47f2b562 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-record-store.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-record-store.ts @@ -1,6 +1,17 @@ import type { AgentKind } from "@mosoo/contracts/agent"; -import type { RuntimeSubjectErrorCode, SandboxSubjectKind } from "@mosoo/contracts/sandbox"; -import { sandboxesTable } from "@mosoo/db"; +import type { + RuntimeSubjectErrorCode, + SandboxOperationKind, + SandboxSubjectKind, +} from "@mosoo/contracts/sandbox"; +import { + driverInstancesTable, + nativeResumeRefsTable, + sandboxesTable, + sandboxSessionsTable, + sessionRunsTable, + sessionsTable, +} from "@mosoo/db"; import { createPlatformId } from "@mosoo/id"; import type { AccountId, @@ -10,11 +21,30 @@ import type { RuntimeOperationId, SandboxBackupId, SandboxId, + SessionId, + SessionRunId, } from "@mosoo/id"; -import { and, eq, inArray, isNull, lte, notExists, or, sql } from "drizzle-orm"; +import { + and, + eq, + exists, + inArray, + isNotNull, + isNull, + lte, + ne, + notExists, + or, + sql, +} from "drizzle-orm"; import type { SQL } from "drizzle-orm"; +import { alias } from "drizzle-orm/sqlite-core"; -import { getAppDatabase, getD1ChangeCount } from "../../../../platform/db/drizzle"; +import { + getAppDatabase, + getD1ChangeCount, + runAppDatabaseBatch, +} from "../../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../../time"; import { getRuntimeKindPolicy, @@ -25,6 +55,7 @@ import { toRuntimeSubjectStatusLifecycleEventName, } from "../../domain/runtime-subject-lifecycle.machine"; import type { RuntimeSubjectOperationStatus } from "../../domain/runtime-subject-lifecycle.machine"; +import { ACTIVE_SESSION_RUN_STATUSES } from "../../domain/session-run-lifecycle.machine"; import { activeSessionRunQueryForListedSubject, lastBackupTable, @@ -32,9 +63,11 @@ import { mapReadyRuntimeSubjectBackup, readyLastBackupTable, runLeaseQuery, + runtimeProvisioningQuery, } from "./runtime-subject-store-queries"; import type { RuntimeSubjectActivationRecord, + RuntimeSubjectOperationLease, RuntimeSubjectRecord, RuntimeSubjectStatus, } from "./runtime-subject-store.types"; @@ -68,6 +101,7 @@ function runtimeSubjectAccountCapacityPredicate(input: { function runtimeSubjectStatusPatch(input: { readonly now: number; + readonly operationKind: SandboxOperationKind | null; readonly operationId: RuntimeOperationId | null; readonly source: "api" | "maintenance" | "runtime"; readonly status: RuntimeSubjectStatus; @@ -77,24 +111,13 @@ function runtimeSubjectStatusPatch(input: { statusChangedAt: input.now, statusEvent: toRuntimeSubjectStatusLifecycleEventName(input.status), statusOperationId: input.operationId ?? null, + operationKind: input.operationKind, statusSeq: sql`${sandboxesTable.statusSeq} + 1`, statusSource: input.source, updatedAt: input.now, } as const; } -function runtimeSubjectStatusOperationCondition( - operationId: RuntimeOperationId | null | undefined, -): SQL[] { - if (operationId === undefined) { - return []; - } - - return operationId === null - ? [isNull(sandboxesTable.statusOperationId)] - : [eq(sandboxesTable.statusOperationId, operationId)]; -} - export async function getRuntimeSubject( database: D1Database, runtimeSubjectId: SandboxId, @@ -102,9 +125,15 @@ export async function getRuntimeSubject( const row = (await getAppDatabase(database) .select({ + agentId: sandboxesTable.agentId, + appId: sandboxesTable.appId, id: sandboxesTable.id, + incarnation: sandboxesTable.incarnation, kind: sandboxesTable.kind, + networkConstraintsHash: sandboxesTable.networkConstraintsHash, + ownerAccountId: sandboxesTable.ownerAccountId, status: sandboxesTable.status, + subjectId: sandboxesTable.subjectId, subjectKind: sandboxesTable.subjectKind, }) .from(sandboxesTable) @@ -156,6 +185,16 @@ export async function ensureRuntimeSubjectId( const existing = await getRuntimeSubjectIdByTuple(database, input); if (existing !== null) { + const record = await getRuntimeSubject(database, existing); + if ( + record === null || + (input.runtimeSubjectId !== undefined && record.id !== input.runtimeSubjectId) || + record.agentId !== input.agentId || + record.appId !== input.appId || + record.ownerAccountId !== input.executionOwnerUserId + ) { + throw new Error("Runtime subject identity does not match the allocation request."); + } return existing; } @@ -173,8 +212,11 @@ export async function ensureRuntimeSubjectId( globalMountsJson: "[]", id: runtimeSubjectId, inactiveDeadlineAt: getRuntimeSubjectInactiveDeadline(getRuntimeKindPolicy(input.kind), now), + incarnation: 0, kind: input.kind, ownerAccountId: input.executionOwnerUserId, + networkConstraintsHash: null, + operationKind: null, status: "cold", statusChangedAt: now, statusEvent: toRuntimeSubjectStatusLifecycleEventName("cold"), @@ -198,6 +240,17 @@ export async function ensureRuntimeSubjectId( throw new Error("Runtime subject could not be allocated."); } + const concurrentRecord = await getRuntimeSubject(database, createdByConcurrentRequest); + if ( + concurrentRecord === null || + (input.runtimeSubjectId !== undefined && concurrentRecord.id !== input.runtimeSubjectId) || + concurrentRecord.agentId !== input.agentId || + concurrentRecord.appId !== input.appId || + concurrentRecord.ownerAccountId !== input.executionOwnerUserId + ) { + throw new Error("Runtime subject identity does not match the allocation request."); + } + return createdByConcurrentRequest; } @@ -208,9 +261,12 @@ export async function getRuntimeSubjectActivationRecord( const row = (await getAppDatabase(database) .select({ + agentId: sandboxesTable.agentId, + appId: sandboxesTable.appId, claimExpiresAt: sandboxesTable.claimExpiresAt, claimOwner: sandboxesTable.claimOwner, id: sandboxesTable.id, + incarnation: sandboxesTable.incarnation, kind: sandboxesTable.kind, lastError: sandboxesTable.lastError, lastErrorCode: sandboxesTable.lastErrorCode, @@ -219,7 +275,13 @@ export async function getRuntimeSubjectActivationRecord( lastBackupStatus: lastBackupTable.status, lastReadyBackupDir: readyLastBackupTable.dir, lastReadyBackupId: readyLastBackupTable.id, + networkConstraintsHash: sandboxesTable.networkConstraintsHash, + ownerAccountId: sandboxesTable.ownerAccountId, + operationId: sandboxesTable.statusOperationId, + operationKind: sandboxesTable.operationKind, status: sandboxesTable.status, + subjectId: sandboxesTable.subjectId, + subjectKind: sandboxesTable.subjectKind, }) .from(sandboxesTable) .leftJoin( @@ -246,9 +308,12 @@ export async function getRuntimeSubjectActivationRecord( } return { + agentId: row.agentId, + appId: row.appId, claimExpiresAt: row.claimExpiresAt, claimOwner: row.claimOwner, id: row.id, + incarnation: row.incarnation, kind: row.kind, lastError: row.lastError, lastErrorCode: row.lastErrorCode, @@ -261,7 +326,13 @@ export async function getRuntimeSubjectActivationRecord( dir: row.lastReadyBackupDir, id: row.lastReadyBackupId, }), + networkConstraintsHash: row.networkConstraintsHash, + ownerAccountId: row.ownerAccountId, + operationId: row.operationId, + operationKind: row.operationKind, status: row.status, + subjectId: row.subjectId, + subjectKind: row.subjectKind, }; } @@ -280,16 +351,16 @@ export async function claimRuntimeSubjectActivation( (await getAppDatabase(database) .update(sandboxesTable) .set({ - agentId: input.agentId, - appId: input.appId, claimExpiresAt: input.claimExpiresAt, claimOwner: input.claimOwner, - ownerAccountId: input.executionOwnerUserId, updatedAt: input.now, }) .where( and( eq(sandboxesTable.id, input.runtimeSubjectId), + eq(sandboxesTable.agentId, input.agentId), + eq(sandboxesTable.appId, input.appId), + eq(sandboxesTable.ownerAccountId, input.executionOwnerUserId), eq(sandboxesTable.status, input.expectedStatus), inArray(sandboxesTable.status, RUNTIME_SUBJECT_CLAIMABLE_STATUSES), or( @@ -347,18 +418,25 @@ export async function markRuntimeSubjectRestoring( database: D1Database, input: { readonly claimOwner: string; + readonly expectedIncarnation: number; + readonly expectedStatus: "active" | "cold"; + readonly networkConstraintsHash: string; + readonly operationId: RuntimeOperationId; readonly runtimeSubjectId: SandboxId; }, -): Promise { +): Promise { const now = currentTimestampMs(); - const result = await getAppDatabase(database) + const row = await getAppDatabase(database) .update(sandboxesTable) .set({ + incarnation: sql`${sandboxesTable.incarnation} + 1`, lastError: null, lastErrorCode: null, + networkConstraintsHash: input.networkConstraintsHash, ...runtimeSubjectStatusPatch({ now, - operationId: null, + operationId: input.operationId, + operationKind: "activate", source: "api", status: "restoring", }), @@ -367,23 +445,145 @@ export async function markRuntimeSubjectRestoring( and( eq(sandboxesTable.id, input.runtimeSubjectId), eq(sandboxesTable.claimOwner, input.claimOwner), - eq(sandboxesTable.status, "cold"), + eq(sandboxesTable.incarnation, input.expectedIncarnation), + isNotNull(sandboxesTable.claimExpiresAt), + eq(sandboxesTable.status, input.expectedStatus), + isNull(sandboxesTable.operationKind), + isNull(sandboxesTable.statusOperationId), ), ) - .run(); + .returning({ + claimExpiresAt: sandboxesTable.claimExpiresAt, + claimOwner: sandboxesTable.claimOwner, + incarnation: sandboxesTable.incarnation, + operationId: sandboxesTable.statusOperationId, + }) + .get(); - return getD1ChangeCount(result) > 0; + return row?.claimExpiresAt === null || row?.claimOwner === null || row?.operationId === null + ? null + : { + claimExpiresAt: row.claimExpiresAt, + claimOwner: row.claimOwner, + incarnation: row.incarnation, + kind: "activate", + operationId: row.operationId, + status: "restoring", + }; +} + +export async function markRuntimeSubjectActiveDestroying( + database: D1Database, + input: { + readonly claimOwner: string; + readonly errorCode: RuntimeSubjectErrorCode; + readonly expectedIncarnation: number; + readonly message: string; + readonly operationId: RuntimeOperationId; + readonly provisioningOperationId: RuntimeOperationId; + readonly provisioningRunId: SessionRunId; + readonly provisioningSessionId: SessionId; + readonly runtimeSubjectId: SandboxId; + }, +): Promise { + const now = currentTimestampMs(); + const db = getAppDatabase(database); + const retirementRuns = alias(sessionRunsTable, "runtime_subject_retirement_run"); + const retirementDrivers = alias(driverInstancesTable, "runtime_subject_retirement_driver"); + const retirementProvisioning = alias(sessionsTable, "runtime_subject_retirement_provisioning"); + const ownedProvisioning = db + .select({ id: sessionsTable.id }) + .from(sessionsTable) + .where( + and( + eq(sessionsTable.id, input.provisioningSessionId), + eq(sessionsTable.runtimeProvisioningOperationId, input.provisioningOperationId), + eq(sessionsTable.runtimeProvisioningRunId, input.provisioningRunId), + eq(sessionsTable.runtimeProvisioningSandboxId, input.runtimeSubjectId), + or( + isNull(sessionsTable.runtimeProvisioningSandboxIncarnation), + eq(sessionsTable.runtimeProvisioningSandboxIncarnation, input.expectedIncarnation), + ), + ), + ); + const otherProvisioning = db + .select({ id: retirementProvisioning.id }) + .from(retirementProvisioning) + .where( + and( + eq(retirementProvisioning.runtimeProvisioningSandboxId, input.runtimeSubjectId), + isNotNull(retirementProvisioning.runtimeProvisioningOperationId), + ne(retirementProvisioning.runtimeProvisioningOperationId, input.provisioningOperationId), + ), + ); + const otherRun = db + .select({ id: retirementRuns.id }) + .from(retirementRuns) + .innerJoin(retirementDrivers, eq(retirementDrivers.id, retirementRuns.driverInstanceId)) + .where( + and( + ne(retirementRuns.id, input.provisioningRunId), + eq(retirementDrivers.sandboxId, input.runtimeSubjectId), + eq(retirementDrivers.sandboxIncarnation, input.expectedIncarnation), + inArray(retirementRuns.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ); + const row = await db + .update(sandboxesTable) + .set({ + lastError: input.message, + lastErrorCode: input.errorCode, + ...runtimeSubjectStatusPatch({ + now, + operationId: input.operationId, + operationKind: "activate", + source: "api", + status: "destroying", + }), + }) + .where( + and( + eq(sandboxesTable.id, input.runtimeSubjectId), + eq(sandboxesTable.claimOwner, input.claimOwner), + eq(sandboxesTable.incarnation, input.expectedIncarnation), + isNotNull(sandboxesTable.claimExpiresAt), + eq(sandboxesTable.status, "active"), + isNull(sandboxesTable.operationKind), + isNull(sandboxesTable.statusOperationId), + exists(ownedProvisioning), + notExists(otherProvisioning), + notExists(otherRun), + ), + ) + .returning({ + claimExpiresAt: sandboxesTable.claimExpiresAt, + claimOwner: sandboxesTable.claimOwner, + incarnation: sandboxesTable.incarnation, + operationId: sandboxesTable.statusOperationId, + }) + .get(); + + return row?.claimExpiresAt === null || row?.claimOwner === null || row?.operationId === null + ? null + : { + claimExpiresAt: row.claimExpiresAt, + claimOwner: row.claimOwner, + incarnation: row.incarnation, + kind: "activate", + operationId: row.operationId, + status: "destroying", + }; } export async function markRuntimeSubjectRestoreApplied( database: D1Database, input: { readonly backupId: SandboxBackupId; - readonly claimOwner: string; + readonly lease: RuntimeSubjectOperationLease; readonly runtimeSubjectId: SandboxId; }, -): Promise { - await getAppDatabase(database) +): Promise { + const result = await getAppDatabase(database) .update(sandboxesTable) .set({ lastRestoreBackupId: input.backupId, @@ -392,17 +592,26 @@ export async function markRuntimeSubjectRestoreApplied( .where( and( eq(sandboxesTable.id, input.runtimeSubjectId), - eq(sandboxesTable.claimOwner, input.claimOwner), + eq(sandboxesTable.claimOwner, input.lease.claimOwner), + eq(sandboxesTable.incarnation, input.lease.incarnation), + eq(sandboxesTable.operationKind, "activate"), + eq(sandboxesTable.status, "restoring"), + eq(sandboxesTable.statusOperationId, input.lease.operationId), ), ) .run(); + + return getD1ChangeCount(result) === 1; } export async function markRuntimeSubjectActive( database: D1Database, input: { readonly claimOwner: string; + readonly incarnation: number; readonly kind: AgentKind; + readonly networkConstraintsHash: string; + readonly operationId: RuntimeOperationId | null; readonly runtimeSubjectId: SandboxId; }, ): Promise { @@ -417,6 +626,7 @@ export async function markRuntimeSubjectActive( inactiveDeadlineAt: getRuntimeSubjectInactiveDeadline(getRuntimeKindPolicy(input.kind), now), lastError: null, lastErrorCode: null, + operationKind: null, status: "active", statusChangedAt: sql` CASE @@ -449,7 +659,15 @@ export async function markRuntimeSubjectActive( and( eq(sandboxesTable.id, input.runtimeSubjectId), eq(sandboxesTable.claimOwner, input.claimOwner), - inArray(sandboxesTable.status, ["restoring", "active"]), + eq(sandboxesTable.incarnation, input.incarnation), + eq(sandboxesTable.networkConstraintsHash, input.networkConstraintsHash), + ...(input.operationId === null + ? [eq(sandboxesTable.status, "active"), isNull(sandboxesTable.statusOperationId)] + : [ + eq(sandboxesTable.status, "restoring"), + eq(sandboxesTable.operationKind, "activate"), + eq(sandboxesTable.statusOperationId, input.operationId), + ]), ), ) .run(); @@ -460,10 +678,9 @@ export async function markRuntimeSubjectActive( export async function markRuntimeSubjectActivationDestroying( database: D1Database, input: { - readonly claimOwner: string; readonly message: string; readonly errorCode: RuntimeSubjectErrorCode; - readonly operationId: RuntimeOperationId; + readonly lease: RuntimeSubjectOperationLease; readonly runtimeSubjectId: SandboxId; }, ): Promise { @@ -472,13 +689,12 @@ export async function markRuntimeSubjectActivationDestroying( const result = await getAppDatabase(database) .update(sandboxesTable) .set({ - claimExpiresAt: null, - claimOwner: null, lastError: input.message, lastErrorCode: input.errorCode, ...runtimeSubjectStatusPatch({ now, - operationId: input.operationId, + operationId: input.lease.operationId, + operationKind: "activate", source: "api", status: "destroying", }), @@ -486,10 +702,11 @@ export async function markRuntimeSubjectActivationDestroying( .where( and( eq(sandboxesTable.id, input.runtimeSubjectId), - eq(sandboxesTable.claimOwner, input.claimOwner), - // Activation can fail at any point after the claim: still cold (during - // prepareFilesystem), restoring (during restore), or active. - inArray(sandboxesTable.status, ["cold", "restoring", "active"]), + eq(sandboxesTable.claimOwner, input.lease.claimOwner), + eq(sandboxesTable.incarnation, input.lease.incarnation), + eq(sandboxesTable.operationKind, "activate"), + eq(sandboxesTable.status, "restoring"), + eq(sandboxesTable.statusOperationId, input.lease.operationId), ), ) .run(); @@ -502,7 +719,7 @@ export async function markRuntimeSubjectActivationFailed( input: { readonly message: string; readonly errorCode: RuntimeSubjectErrorCode; - readonly operationId: RuntimeOperationId; + readonly lease: RuntimeSubjectOperationLease; readonly runtimeSubjectId: SandboxId; }, ): Promise { @@ -520,6 +737,7 @@ export async function markRuntimeSubjectActivationFailed( ...runtimeSubjectStatusPatch({ now, operationId: null, + operationKind: null, source: "api", status: "cold", }), @@ -527,8 +745,11 @@ export async function markRuntimeSubjectActivationFailed( .where( and( eq(sandboxesTable.id, input.runtimeSubjectId), + eq(sandboxesTable.claimOwner, input.lease.claimOwner), + eq(sandboxesTable.incarnation, input.lease.incarnation), + eq(sandboxesTable.operationKind, "activate"), eq(sandboxesTable.status, "destroying"), - eq(sandboxesTable.statusOperationId, input.operationId), + eq(sandboxesTable.statusOperationId, input.lease.operationId), ), ) .run(); @@ -539,62 +760,72 @@ export async function markRuntimeSubjectActivationFailed( export async function markRuntimeSubjectOperationStarted( database: D1Database, input: { - readonly claimOwner?: string; + readonly claimExpiresAt: number; + readonly claimOwner: string; readonly now?: number; - readonly operationId?: RuntimeOperationId | null; + readonly operationId: RuntimeOperationId; + readonly operationKind: Exclude; readonly runtimeSubjectId: SandboxId; readonly source?: "api" | "maintenance" | "runtime"; - readonly status: RuntimeSubjectOperationStatus; + readonly status?: "backing_up"; }, -): Promise { +): Promise { const now = input.now ?? currentTimestampMs(); const appDb = getAppDatabase(database); - const claimPredicate = - input.claimOwner === undefined - ? or( - isNull(sandboxesTable.claimOwner), - isNull(sandboxesTable.claimExpiresAt), - lte(sandboxesTable.claimExpiresAt, now), - ) - : eq(sandboxesTable.claimOwner, input.claimOwner); - const result = await appDb + const claimPredicate = or( + eq(sandboxesTable.claimOwner, input.claimOwner), + isNull(sandboxesTable.claimOwner), + isNull(sandboxesTable.claimExpiresAt), + lte(sandboxesTable.claimExpiresAt, now), + ); + const row = await appDb .update(sandboxesTable) .set({ - claimExpiresAt: null, - claimOwner: null, + claimExpiresAt: input.claimExpiresAt, + claimOwner: input.claimOwner, inactiveDeadlineAt: null, lastError: null, lastErrorCode: null, ...runtimeSubjectStatusPatch({ now, - operationId: input.operationId ?? null, + operationId: input.operationId, + operationKind: input.operationKind, source: input.source ?? "api", - status: input.status, + status: "backing_up", }), }) .where( and( eq(sandboxesTable.id, input.runtimeSubjectId), - inArray(sandboxesTable.status, RUNTIME_SUBJECT_CLAIMABLE_STATUSES), + eq(sandboxesTable.status, "active"), + isNull(sandboxesTable.operationKind), + isNull(sandboxesTable.statusOperationId), claimPredicate, - ...(input.source === "maintenance" - ? [ - notExists(activeSessionRunQueryForListedSubject(appDb)), - notExists(runLeaseQuery(appDb, input.runtimeSubjectId)), - ] - : []), + notExists(runtimeProvisioningQuery(appDb, input.runtimeSubjectId)), + notExists(activeSessionRunQueryForListedSubject(appDb)), + notExists(runLeaseQuery(appDb, input.runtimeSubjectId)), ), ) - .run(); + .returning({ incarnation: sandboxesTable.incarnation }) + .get(); - return getD1ChangeCount(result) > 0; + return row === undefined + ? null + : { + claimExpiresAt: input.claimExpiresAt, + claimOwner: input.claimOwner, + incarnation: row.incarnation, + kind: input.operationKind, + operationId: input.operationId, + status: "backing_up", + }; } export async function advanceRuntimeSubjectOperationStatus( database: D1Database, input: { - readonly expectedStatus: RuntimeSubjectOperationStatus; - readonly operationId?: RuntimeOperationId | null; + readonly expectedStatus: RuntimeSubjectOperationLease["status"]; + readonly lease: RuntimeSubjectOperationLease; readonly runtimeSubjectId: SandboxId; readonly source?: "api" | "maintenance" | "runtime"; readonly status: RuntimeSubjectOperationStatus; @@ -606,7 +837,8 @@ export async function advanceRuntimeSubjectOperationStatus( .set({ ...runtimeSubjectStatusPatch({ now, - operationId: input.operationId ?? null, + operationId: input.lease.operationId, + operationKind: input.lease.kind, source: input.source ?? "api", status: input.status, }), @@ -615,7 +847,10 @@ export async function advanceRuntimeSubjectOperationStatus( and( eq(sandboxesTable.id, input.runtimeSubjectId), eq(sandboxesTable.status, input.expectedStatus), - ...runtimeSubjectStatusOperationCondition(input.operationId), + eq(sandboxesTable.claimOwner, input.lease.claimOwner), + eq(sandboxesTable.incarnation, input.lease.incarnation), + eq(sandboxesTable.operationKind, input.lease.kind), + eq(sandboxesTable.statusOperationId, input.lease.operationId), ), ) .run(); @@ -627,8 +862,11 @@ export async function markRuntimeSubjectCold( database: D1Database, input: { readonly clearBackups: boolean; - readonly expectedStatus: RuntimeSubjectOperationStatus; - readonly operationId?: RuntimeOperationId | null; + readonly clearNativeResumeRefs?: boolean; + readonly errorCode?: RuntimeSubjectErrorCode; + readonly errorMessage?: string; + readonly expectedStatus: "destroying"; + readonly lease: RuntimeSubjectOperationLease; readonly runtimeSubjectId: SandboxId; readonly source?: "api" | "maintenance" | "runtime"; }, @@ -641,27 +879,101 @@ export async function markRuntimeSubjectCold( } : {}; + const results = await runAppDatabaseBatch(database, (appDb) => { + const ownedOperation = and( + eq(sandboxesTable.id, input.runtimeSubjectId), + eq(sandboxesTable.status, input.expectedStatus), + eq(sandboxesTable.claimOwner, input.lease.claimOwner), + eq(sandboxesTable.incarnation, input.lease.incarnation), + eq(sandboxesTable.operationKind, input.lease.kind), + eq(sandboxesTable.statusOperationId, input.lease.operationId), + ); + const stillOwned = exists( + appDb.select({ id: sandboxesTable.id }).from(sandboxesTable).where(ownedOperation), + ); + + return [ + appDb.delete(nativeResumeRefsTable).where( + and( + input.clearNativeResumeRefs ? sql`TRUE` : sql`FALSE`, + stillOwned, + exists( + appDb + .select({ sessionId: sandboxSessionsTable.sessionId }) + .from(sandboxSessionsTable) + .where( + and( + eq(sandboxSessionsTable.sessionId, nativeResumeRefsTable.sessionId), + eq(sandboxSessionsTable.sandboxId, input.runtimeSubjectId), + eq(sandboxSessionsTable.sandboxIncarnation, input.lease.incarnation), + ), + ), + ), + ), + ), + appDb + .update(sandboxSessionsTable) + .set({ cleanupOperationId: null, status: "closed", updatedAt: now }) + .where( + and( + eq(sandboxSessionsTable.sandboxId, input.runtimeSubjectId), + eq(sandboxSessionsTable.sandboxIncarnation, input.lease.incarnation), + inArray(sandboxSessionsTable.status, ["active", "cleanup_pending", "error"]), + stillOwned, + ), + ), + appDb + .update(sandboxesTable) + .set({ + ...backupFields, + claimExpiresAt: null, + claimOwner: null, + inactiveDeadlineAt: null, + lastError: input.errorMessage ?? null, + lastErrorCode: input.errorCode ?? null, + ...runtimeSubjectStatusPatch({ + now, + operationId: null, + operationKind: null, + source: input.source ?? "api", + status: "cold", + }), + }) + .where(ownedOperation), + ]; + }); + + return getD1ChangeCount(results[2]) > 0; +} + +export async function markRuntimeSubjectOperationRepairNeeded( + database: D1Database, + input: { + readonly errorMessage: string; + readonly errorCode: RuntimeSubjectErrorCode; + readonly expectedStatus: RuntimeSubjectOperationLease["status"]; + readonly lease: RuntimeSubjectOperationLease; + readonly runtimeSubjectId: SandboxId; + readonly source?: "api" | "maintenance" | "runtime"; + }, +): Promise { + const now = currentTimestampMs(); const result = await getAppDatabase(database) .update(sandboxesTable) .set({ - ...backupFields, - claimExpiresAt: null, - claimOwner: null, - inactiveDeadlineAt: null, - lastError: null, - lastErrorCode: null, - ...runtimeSubjectStatusPatch({ - now, - operationId: input.operationId ?? null, - source: input.source ?? "api", - status: "cold", - }), + claimExpiresAt: now, + lastError: input.errorMessage, + lastErrorCode: input.errorCode, + updatedAt: now, }) .where( and( eq(sandboxesTable.id, input.runtimeSubjectId), eq(sandboxesTable.status, input.expectedStatus), - ...runtimeSubjectStatusOperationCondition(input.operationId), + eq(sandboxesTable.claimOwner, input.lease.claimOwner), + eq(sandboxesTable.incarnation, input.lease.incarnation), + eq(sandboxesTable.operationKind, input.lease.kind), + eq(sandboxesTable.statusOperationId, input.lease.operationId), ), ) .run(); @@ -669,37 +981,28 @@ export async function markRuntimeSubjectCold( return getD1ChangeCount(result) > 0; } -export async function markRuntimeSubjectOperationRepairNeeded( +export async function renewRuntimeSubjectOperationLease( database: D1Database, input: { - readonly errorMessage: string; - readonly errorCode: RuntimeSubjectErrorCode; - readonly expectedStatus: RuntimeSubjectOperationStatus; - readonly operationId: RuntimeOperationId; + readonly claimExpiresAt: number; + readonly lease: RuntimeSubjectOperationLease; readonly runtimeSubjectId: SandboxId; - readonly source?: "api" | "maintenance" | "runtime"; }, ): Promise { const now = currentTimestampMs(); const result = await getAppDatabase(database) .update(sandboxesTable) .set({ - claimExpiresAt: null, - claimOwner: null, - lastError: input.errorMessage, - lastErrorCode: input.errorCode, - ...runtimeSubjectStatusPatch({ - now, - operationId: input.operationId, - source: input.source ?? "maintenance", - status: input.expectedStatus, - }), + claimExpiresAt: input.claimExpiresAt, + updatedAt: now, }) .where( and( eq(sandboxesTable.id, input.runtimeSubjectId), - eq(sandboxesTable.status, input.expectedStatus), - eq(sandboxesTable.statusOperationId, input.operationId), + eq(sandboxesTable.claimOwner, input.lease.claimOwner), + eq(sandboxesTable.incarnation, input.lease.incarnation), + eq(sandboxesTable.operationKind, input.lease.kind), + eq(sandboxesTable.statusOperationId, input.lease.operationId), ), ) .run(); @@ -707,29 +1010,16 @@ export async function markRuntimeSubjectOperationRepairNeeded( return getD1ChangeCount(result) > 0; } -export async function markRuntimeSubjectFailed( +export async function releaseRuntimeSubjectActivationClaim( database: D1Database, input: { - readonly errorMessage: string; + readonly claimOwner: string; readonly errorCode: RuntimeSubjectErrorCode; - readonly expectedStatus?: RuntimeSubjectStatus; - readonly operationId?: RuntimeOperationId | null; + readonly errorMessage: string; + readonly incarnation: number; readonly runtimeSubjectId: SandboxId; - readonly source?: "api" | "maintenance" | "runtime"; - readonly status: RuntimeSubjectStatus; }, ): Promise { - const now = currentTimestampMs(); - const conditions: SQL[] = [eq(sandboxesTable.id, input.runtimeSubjectId)]; - - if (input.expectedStatus !== undefined) { - conditions.push(eq(sandboxesTable.status, input.expectedStatus)); - } - - if (input.operationId !== undefined) { - conditions.push(...runtimeSubjectStatusOperationCondition(input.operationId)); - } - const result = await getAppDatabase(database) .update(sandboxesTable) .set({ @@ -737,14 +1027,18 @@ export async function markRuntimeSubjectFailed( claimOwner: null, lastError: input.errorMessage, lastErrorCode: input.errorCode, - ...runtimeSubjectStatusPatch({ - now, - operationId: input.operationId ?? null, - source: input.source ?? "api", - status: input.status, - }), + updatedAt: currentTimestampMs(), }) - .where(and(...conditions)) + .where( + and( + eq(sandboxesTable.id, input.runtimeSubjectId), + eq(sandboxesTable.claimOwner, input.claimOwner), + eq(sandboxesTable.incarnation, input.incarnation), + inArray(sandboxesTable.status, ["active", "cold"]), + isNull(sandboxesTable.operationKind), + isNull(sandboxesTable.statusOperationId), + ), + ) .run(); return getD1ChangeCount(result) > 0; diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-recycle.service.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-recycle.service.ts index 897028dd..8c72a134 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-recycle.service.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-recycle.service.ts @@ -3,91 +3,16 @@ import { createPlatformId } from "@mosoo/id"; import type { RuntimeOperationId, SandboxId } from "@mosoo/id"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; -import { getRuntimeKindPolicy } from "../../domain/runtime-kind-policy"; -import type { RuntimeSubjectOperationStatus } from "../../domain/runtime-subject-lifecycle.machine"; -import { createSandboxCheckpoints } from "../sandbox-backup.service"; -import { stopRuntimeSubjectDrivers } from "./runtime-subject-driver-stop"; -import { getRuntimeSubjectOperationErrorCode } from "./runtime-subject-errors"; -import { destroyRuntimeSubjectContainer } from "./runtime-subject-platform"; +import { runRuntimeSubjectOperation } from "./runtime-subject-operations.service"; import { - advanceRuntimeSubjectOperationStatus, claimInactiveRuntimeSubject, - closeRuntimeSubjectSessionsForRecycle, - markRuntimeSubjectCold, markRuntimeSubjectOperationStarted, - markRuntimeSubjectOperationRepairNeeded, releaseInactiveRuntimeSubjectClaim, } from "./runtime-subject-store"; +import type { RuntimeSubjectOperationLease } from "./runtime-subject-store"; const RECYCLE_CLAIM_TTL_MS = 10 * 60_000; -function getRuntimeSubjectRecycleErrorMessage(error: unknown): string { - return error instanceof Error ? error.message : "Runtime subject recycle failed."; -} - -async function runRuntimeSubjectRecycleOperation( - bindings: ApiBindings, - input: { - readonly kind: AgentKind; - readonly operationId: RuntimeOperationId; - readonly reason: string; - readonly runtimeSubjectId: SandboxId; - readonly startStatus: RuntimeSubjectOperationStatus; - }, -): Promise { - let destroyStarted = input.startStatus === "destroying"; - - try { - const policy = getRuntimeKindPolicy(input.kind); - - if (input.startStatus === "backing_up") { - await stopRuntimeSubjectDrivers(bindings, { - operationId: input.operationId, - reason: input.reason, - runtimeSubjectId: input.runtimeSubjectId, - }); - await createSandboxCheckpoints(bindings, { - operationId: input.operationId, - rules: policy.checkpoint.createOnHibernate, - sandboxId: input.runtimeSubjectId, - }); - destroyStarted = await advanceRuntimeSubjectOperationStatus(bindings.DB, { - expectedStatus: "backing_up", - operationId: input.operationId, - runtimeSubjectId: input.runtimeSubjectId, - source: "maintenance", - status: "destroying", - }); - if (!destroyStarted) { - throw new Error("Runtime subject changed before recycle destroy."); - } - } - - await destroyRuntimeSubjectContainer(bindings, input.runtimeSubjectId); - await closeRuntimeSubjectSessionsForRecycle(bindings.DB, input.runtimeSubjectId); - const completed = await markRuntimeSubjectCold(bindings.DB, { - clearBackups: false, - expectedStatus: "destroying", - operationId: input.operationId, - runtimeSubjectId: input.runtimeSubjectId, - source: "maintenance", - }); - if (!completed) { - throw new Error("Runtime subject changed before recycle completion."); - } - } catch (error) { - await markRuntimeSubjectOperationRepairNeeded(bindings.DB, { - errorCode: getRuntimeSubjectOperationErrorCode(error), - errorMessage: getRuntimeSubjectRecycleErrorMessage(error), - expectedStatus: destroyStarted ? "destroying" : "backing_up", - operationId: input.operationId, - runtimeSubjectId: input.runtimeSubjectId, - source: "maintenance", - }); - throw error; - } -} - export async function recycleRuntimeSubject( bindings: ApiBindings, input: { @@ -99,16 +24,17 @@ export async function recycleRuntimeSubject( }, ): Promise { const operationId = createPlatformId(); - const started = await markRuntimeSubjectOperationStarted(bindings.DB, { + const lease = await markRuntimeSubjectOperationStarted(bindings.DB, { + claimExpiresAt: input.now + RECYCLE_CLAIM_TTL_MS, claimOwner: input.claimOwner, now: input.now, operationId, + operationKind: "hibernate", runtimeSubjectId: input.runtimeSubjectId, source: "maintenance", - status: "backing_up", }); - if (!started) { + if (lease === null) { await releaseInactiveRuntimeSubjectClaim(bindings.DB, { claimOwner: input.claimOwner, runtimeSubjectId: input.runtimeSubjectId, @@ -116,14 +42,12 @@ export async function recycleRuntimeSubject( return false; } - await runRuntimeSubjectRecycleOperation(bindings, { + await runRuntimeSubjectOperation(bindings, { kind: input.kind, - operationId, + lease, reason: input.reason, runtimeSubjectId: input.runtimeSubjectId, - startStatus: "backing_up", }); - return true; } @@ -131,20 +55,12 @@ export async function resumeRuntimeSubjectRecycleOperation( bindings: ApiBindings, input: { readonly kind: AgentKind; - readonly operationId: RuntimeOperationId; + readonly lease: RuntimeSubjectOperationLease; readonly reason: string; readonly runtimeSubjectId: SandboxId; - readonly status: RuntimeSubjectOperationStatus; }, ): Promise { - await runRuntimeSubjectRecycleOperation(bindings, { - kind: input.kind, - operationId: input.operationId, - reason: input.reason, - runtimeSubjectId: input.runtimeSubjectId, - startStatus: input.status, - }); - + await runRuntimeSubjectOperation(bindings, input); return true; } @@ -166,15 +82,13 @@ export async function recycleInactiveRuntimeSubjectNow( runtimeSubjectId: input.runtimeSubjectId, }); - if (!claimed) { - return false; - } - - return recycleRuntimeSubject(bindings, { - claimOwner, - kind: input.kind, - now, - reason: input.reason, - runtimeSubjectId: input.runtimeSubjectId, - }); + return claimed + ? recycleRuntimeSubject(bindings, { + claimOwner, + kind: input.kind, + now, + reason: input.reason, + runtimeSubjectId: input.runtimeSubjectId, + }) + : false; } diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store-queries.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store-queries.ts index b9dda0e0..14b45922 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store-queries.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store-queries.ts @@ -4,6 +4,7 @@ import { sandboxesTable, sandboxSessionsTable, sessionRunsTable, + sessionsTable, } from "@mosoo/db"; import type { SandboxBackupId, SandboxId } from "@mosoo/id"; import { and, eq, inArray, or, sql } from "drizzle-orm"; @@ -31,6 +32,7 @@ const activeRuntimeSubjectRunsTable = alias(sessionRunsTable, "active_runtime_su const liveSubjectDriversTable = alias(driverInstancesTable, "live_runtime_subject_driver"); const runLeaseDriversTable = alias(driverInstancesTable, "runtime_run_lease_driver"); const runLeaseRunsTable = alias(sessionRunsTable, "runtime_run_lease_run"); +const provisioningSessionsTable = alias(sessionsTable, "runtime_provisioning_session"); export const readyConversationBackupTable = alias(sandboxBackupsTable, "ready_conversation_backup"); export const lastBackupTable = alias(sandboxBackupsTable, "last_backup"); export const readyLastBackupTable = alias(sandboxBackupsTable, "ready_last_backup"); @@ -72,7 +74,7 @@ export function activeConversationSessionQuery(appDb: AppDatabase, runtimeSubjec .where( and( eq(activeConversationSessionsTable.sandboxId, runtimeSubjectId), - eq(activeConversationSessionsTable.status, "active"), + inArray(activeConversationSessionsTable.status, ["active", "cleanup_pending"]), ), ); } @@ -84,7 +86,7 @@ export function activeConversationSessionQueryForListedSubject(appDb: AppDatabas .where( and( eq(activeConversationSessionsTable.sandboxId, sandboxesTable.id), - eq(activeConversationSessionsTable.status, "active"), + inArray(activeConversationSessionsTable.status, ["active", "cleanup_pending"]), ), ); } @@ -154,6 +156,20 @@ export function runLeaseQueryForListedSubject(appDb: AppDatabase) { ); } +export function runtimeProvisioningQuery(appDb: AppDatabase, runtimeSubjectId: SandboxId) { + return appDb + .select({ id: provisioningSessionsTable.id }) + .from(provisioningSessionsTable) + .where(eq(provisioningSessionsTable.runtimeProvisioningSandboxId, runtimeSubjectId)); +} + +export function runtimeProvisioningQueryForListedSubject(appDb: AppDatabase) { + return appDb + .select({ id: provisioningSessionsTable.id }) + .from(provisioningSessionsTable) + .where(eq(provisioningSessionsTable.runtimeProvisioningSandboxId, sandboxesTable.id)); +} + export function getRuntimeSubjectInactiveDeadlineSql(now: number) { return sql` CASE ${sandboxesTable.kind} diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store.ts index fc99c5ea..a98f3492 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store.ts @@ -1,15 +1,18 @@ export { claimIdleSessionScopedConversationForClose, + claimRuntimeConversationSessionCleanup, ensureRuntimeConversationSessionRecord, getRuntimeConversationSession, getRuntimeConversationSessionState, + listPendingRuntimeConversationSessionCleanups, recordRuntimeConversationSessionActive, recordRuntimeConversationSessionClosed, recordRuntimeConversationSessionError, + retireRuntimeConversationSessionsForIncarnation, } from "./runtime-conversation-session-store"; export { claimInactiveRuntimeSubject, - closeRuntimeSubjectSessionsForRecycle, + claimRuntimeSubjectOperationForRepair, listInactiveRuntimeSubjects, listRuntimeSubjectDriverIds, listRuntimeSubjectSessionStateTargets, @@ -23,13 +26,15 @@ export { getRuntimeSubject, getRuntimeSubjectActivationRecord, getRuntimeSubjectIdByTuple, + markRuntimeSubjectActiveDestroying, markRuntimeSubjectActivationDestroying, markRuntimeSubjectActivationFailed, markRuntimeSubjectActive, markRuntimeSubjectCold, - markRuntimeSubjectFailed, markRuntimeSubjectOperationStarted, markRuntimeSubjectOperationRepairNeeded, + releaseRuntimeSubjectActivationClaim, + renewRuntimeSubjectOperationLease, preemptRuntimeSubjectActivationClaim, markRuntimeSubjectRestoreApplied, markRuntimeSubjectRestoring, @@ -41,6 +46,5 @@ export type { RuntimeSubjectActivationRecord, RuntimeSubjectMaintenanceCandidate, RuntimeSubjectOperationRepairCandidate, - RuntimeSubjectRecord, - RuntimeSubjectStatus, + RuntimeSubjectOperationLease, } from "./runtime-subject-store.types"; diff --git a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store.types.ts b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store.types.ts index 6161009d..9325a7b6 100644 --- a/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store.types.ts +++ b/apps/api/src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store.types.ts @@ -2,13 +2,17 @@ import type { AgentKind } from "@mosoo/contracts/agent"; import type { RuntimeSubjectErrorCode, SandboxBackupStatus, + SandboxOperationKind, SandboxSessionStatus, SandboxStatus, SandboxSubjectKind, } from "@mosoo/contracts/sandbox"; import type { + AccountId, AgentId, + AppId, DriverInstanceId, + PlatformId, RuntimeOperationId, SandboxBackupId, SandboxId, @@ -17,27 +21,51 @@ import type { SessionRunId, } from "@mosoo/id"; -import type { RuntimeSubjectOperationStatus } from "../../domain/runtime-subject-lifecycle.machine"; +import type { RuntimeSubjectRecoverableOperationStatus } from "../../domain/runtime-subject-lifecycle.machine"; export type RuntimeSubjectStatus = SandboxStatus; export interface RuntimeSubjectRecord { + readonly agentId: AgentId | null; + readonly appId: AppId | null; readonly id: SandboxId; + readonly incarnation: number; readonly kind: AgentKind; + readonly networkConstraintsHash: string | null; + readonly ownerAccountId: AccountId | null; readonly status: RuntimeSubjectStatus; + readonly subjectId: PlatformId; readonly subjectKind: SandboxSubjectKind; } export interface RuntimeSubjectActivationRecord { + readonly agentId: AgentId | null; + readonly appId: AppId | null; readonly claimExpiresAt: number | null; readonly claimOwner: string | null; readonly id: SandboxId; + readonly incarnation: number; readonly kind: AgentKind; readonly lastError: string | null; readonly lastErrorCode: RuntimeSubjectErrorCode | null; readonly lastBackup: RuntimeSubjectBackupRecord | null; readonly lastReadyBackup: ReadyRuntimeSubjectBackupRecord | null; + readonly networkConstraintsHash: string | null; + readonly ownerAccountId: AccountId | null; + readonly operationId: RuntimeOperationId | null; + readonly operationKind: SandboxOperationKind | null; readonly status: RuntimeSubjectStatus; + readonly subjectId: PlatformId; + readonly subjectKind: SandboxSubjectKind; +} + +export interface RuntimeSubjectOperationLease { + readonly claimExpiresAt: number; + readonly claimOwner: string; + readonly incarnation: number; + readonly kind: SandboxOperationKind; + readonly operationId: RuntimeOperationId; + readonly status: RuntimeSubjectRecoverableOperationStatus; } export interface RuntimeSubjectBackupRecord { @@ -52,6 +80,7 @@ export interface ReadyRuntimeSubjectBackupRecord { } export interface RuntimeConversationSessionRecord { + readonly sandboxIncarnation: number; readonly sandboxSessionId: SandboxSessionId; readonly cwd: string; readonly latestReadyBackup: ReadyRuntimeSubjectBackupRecord | null; @@ -63,26 +92,41 @@ export interface RuntimeConversationSessionRecord { export interface RuntimeConversationSessionState { readonly agentId: AgentId | null; + readonly cleanupOperationId: RuntimeOperationId | null; readonly sandboxSessionId: SandboxSessionId; + readonly sandboxIncarnation: number; readonly kind: AgentKind; readonly status: RuntimeConversationSessionRecord["status"]; } +export interface PendingRuntimeConversationSessionCleanup extends RuntimeConversationSessionState { + readonly cleanupOperationId: RuntimeOperationId; + readonly sandboxId: SandboxId; + readonly sessionId: SessionId; + readonly status: "cleanup_pending"; +} + export interface RuntimeSubjectMaintenanceCandidate { readonly id: SandboxId; readonly kind: AgentKind; } export interface RuntimeSubjectOperationRepairCandidate { + readonly claimExpiresAt: number | null; + readonly claimOwner: string | null; readonly id: SandboxId; + readonly incarnation: number; readonly kind: AgentKind; + readonly operationKind: SandboxOperationKind; readonly operationId: RuntimeOperationId; - readonly status: RuntimeSubjectOperationStatus; + readonly status: RuntimeSubjectRecoverableOperationStatus; } export interface RuntimeRunLeaseInput { + readonly driverGeneration: number; readonly driverInstanceId: DriverInstanceId; readonly runtimeSubjectId: SandboxId; + readonly runtimeSubjectIncarnation: number; readonly sessionId: SessionId; readonly sessionRunId: SessionRunId; } diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-backup-platform.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-backup-platform.ts index 660df1d7..4e02adc3 100644 --- a/apps/api/src/modules/runtime/infrastructure/sandbox-backup-platform.ts +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-backup-platform.ts @@ -1,8 +1,5 @@ -import { - getSessionResourceRootPath, - getSessionRuntimeStatePath, - getSessionStateRootPath, -} from "@mosoo/agent-driver/paths"; +import { getSessionRuntimeStatePath } from "@mosoo/agent-driver/paths"; +import type { SandboxBackupId } from "@mosoo/id"; import { withDisposedRpcResource, @@ -13,100 +10,124 @@ import { decodeSandboxBackupIdForPlatform, encodeSandboxBackupIdForStorage, } from "./sandbox-backup-id"; -import type { SandboxHandle } from "./sandbox-handles"; +import { + beginSandboxBackupDeletionAttempt, + completeSandboxBackupDeletion, + isSandboxBackupDeletionAuthorized, +} from "./sandbox-backup-store"; +import { toRuntimeSubjectIncarnationHandle } from "./sandbox-handles"; -interface SandboxBackupObject { +const RUNTIME_SANDBOX_BACKUP_NAME_PREFIX = "mosoo:runtime-backup:v1:"; + +export interface SandboxBackupObject { readonly dir: string; - readonly id: string; + readonly id: SandboxBackupId; } -function isMissingRuntimeBucketMountError(error: unknown): boolean { - return error instanceof Error && error.message.includes("No active mount found at path:"); +export interface SandboxBackupMetadata { + readonly dir: string; + readonly id: string; + readonly name: string | null; } -function quoteShellArg(value: string): string { - return `'${value.replaceAll("'", `'"'"'`)}'`; +export function createRuntimeSandboxBackupName(stagingId: SandboxBackupId): string { + return `${RUNTIME_SANDBOX_BACKUP_NAME_PREFIX}${stagingId}`; } -async function prepareRuntimeSessionWorkspaceCheckpoint( - sandbox: SandboxHandle, - input: { - readonly cwd: string; - readonly sessionId: string; - }, -): Promise { - const resourceRoot = getSessionResourceRootPath(input.sessionId); - const stateRoot = getSessionStateRootPath(input.sessionId); - const openAiAuthPath = `${getSessionRuntimeStatePath(input.sessionId, "openai-runtime")}/auth.json`; - - if (!resourceRoot.startsWith(`${input.cwd}/`) || !stateRoot.startsWith(`${input.cwd}/`)) { - throw new Error("Session checkpoint exclusions must stay inside the session workspace."); +export function parseRuntimeSandboxBackupName(value: string | null): SandboxBackupId | null { + if (value === null || !value.startsWith(RUNTIME_SANDBOX_BACKUP_NAME_PREFIX)) { + return null; } - + const stagingId = value.slice(RUNTIME_SANDBOX_BACKUP_NAME_PREFIX.length); try { - await sandbox.unmountBucket(resourceRoot); - } catch (error) { - if (!isMissingRuntimeBucketMountError(error)) { - throw error; - } + return encodeSandboxBackupIdForStorage(stagingId); + } catch { + return null; } +} - const command = [ - "set -eu", - `cwd=${quoteShellArg(input.cwd)}`, - 'test -d "$cwd"', - `resource_root=${quoteShellArg(resourceRoot)}`, - 'if [ -L "$resource_root" ]; then unlink "$resource_root"; elif mountpoint -q "$resource_root"; then fusermount -u "$resource_root"; fi', - 'if [ -e "$resource_root" ]; then rm -rf "$resource_root"; fi', - `state_root=${quoteShellArg(stateRoot)}`, - 'if [ -d "$state_root" ]; then find "$state_root" -type f -name "driver-boot-payload-*.json" -delete; fi', - `rm -f ${quoteShellArg(openAiAuthPath)}`, - "sync", - ].join("; "); +export function parseSandboxBackupMetadata(value: unknown): SandboxBackupMetadata | null { + if (typeof value !== "object" || value === null) { + return null; + } + const dir = Reflect.get(value, "dir"); + const id = Reflect.get(value, "id"); + const name = Reflect.get(value, "name"); + return typeof dir === "string" && + dir.length > 0 && + typeof id === "string" && + (name === null || typeof name === "string") + ? { dir, id, name } + : null; +} - await withDisposedRpcResult(sandbox.exec(`sh -lc ${quoteShellArg(command)}`), (result) => { - if (!result.success || result.exitCode !== 0) { - const detail = result.stderr.trim(); - throw new Error( - `Session workspace could not be prepared for a credential-safe checkpoint${detail ? `: ${detail}` : "."}`, - ); - } - }); +export function getSandboxBackupObjectKeys(backupId: string): readonly [string, string] { + const platformId = decodeSandboxBackupIdForPlatform(backupId); + return [`backups/${platformId}/data.sqsh`, `backups/${platformId}/meta.json`]; } -function getSandboxBackupObjectKeys(backupId: string): string[] { - const platformBackupId = decodeSandboxBackupIdForPlatform(backupId); +async function readSandboxBackupMetadata( + bindings: Pick, + backupId: SandboxBackupId, +): Promise { + const [, metadataKey] = getSandboxBackupObjectKeys(backupId); + const stored = await bindings.SANDBOX_STATE_BUCKET.get(metadataKey); + if (stored === null) { + return null; + } + try { + return parseSandboxBackupMetadata(JSON.parse(await stored.text())); + } catch { + return null; + } +} - return [`backups/${platformBackupId}/data.sqsh`, `backups/${platformBackupId}/meta.json`]; +export async function isRuntimeSandboxBackupObjectReady( + bindings: Pick, + input: { + readonly backupId: SandboxBackupId; + readonly dir: string; + readonly stagingId: SandboxBackupId; + }, +): Promise { + const [dataKey] = getSandboxBackupObjectKeys(input.backupId); + const [data, metadata] = await Promise.all([ + bindings.SANDBOX_STATE_BUCKET.head(dataKey), + readSandboxBackupMetadata(bindings, input.backupId), + ]); + return ( + data !== null && + metadata?.id === decodeSandboxBackupIdForPlatform(input.backupId) && + metadata.dir === input.dir && + parseRuntimeSandboxBackupName(metadata.name) === input.stagingId + ); } export async function createRuntimeSandboxBackup( bindings: ApiBindings, input: { readonly dir: string; + readonly incarnation: number; readonly sandboxId: string; readonly sessionId: string | null; + readonly stagingId: SandboxBackupId; readonly ttlSeconds: number; }, ): Promise { const { getRuntimeSubjectKeepAliveHandle } = await import("./runtime-subject-lifecycle/runtime-subject-lifecycle.service"); - return withDisposedRpcResource( - await getRuntimeSubjectKeepAliveHandle(bindings, input.sandboxId), + await getRuntimeSubjectKeepAliveHandle(bindings, input.sandboxId, input.incarnation), async (sandbox) => { - if (input.sessionId !== null) { - await prepareRuntimeSessionWorkspaceCheckpoint(sandbox, { - cwd: input.dir, - sessionId: input.sessionId, - }); - } else { - await sandbox.mkdir(input.dir, { recursive: true }); - } - + const forbiddenPaths = + input.sessionId === null + ? undefined + : [`${getSessionRuntimeStatePath(input.sessionId, "openai-runtime")}/auth.json`]; return withDisposedRpcResult( - sandbox.createBackup({ + toRuntimeSubjectIncarnationHandle(sandbox).createRuntimeSubjectBackup(input.incarnation, { dir: input.dir, + ...(forbiddenPaths === undefined ? {} : { forbiddenPaths }), + name: createRuntimeSandboxBackupName(input.stagingId), ttl: input.ttlSeconds, }), (result) => ({ @@ -118,15 +139,16 @@ export async function createRuntimeSandboxBackup( ); } -export async function deleteSandboxBackupObjects( - bindings: ApiBindings, - backupIds: readonly string[], +export async function deleteAuthorizedSandboxBackupObjects( + bindings: Pick, + backupIds: readonly SandboxBackupId[], ): Promise { - if (backupIds.length === 0) { - return; + for (const backupId of new Set(backupIds)) { + if (!(await isSandboxBackupDeletionAuthorized(bindings.DB, backupId))) { + throw new Error("Sandbox backup object deletion lacks a durable D1 intent."); + } + await beginSandboxBackupDeletionAttempt(bindings.DB, backupId); + await bindings.SANDBOX_STATE_BUCKET.delete([...getSandboxBackupObjectKeys(backupId)]); + await completeSandboxBackupDeletion(bindings.DB, backupId); } - - const objectKeys = backupIds.flatMap((backupId) => getSandboxBackupObjectKeys(backupId)); - - await bindings.SANDBOX_STATE_BUCKET.delete(objectKeys); } diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-backup-pruning.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-backup-pruning.ts index de37061a..dca2ec2f 100644 --- a/apps/api/src/modules/runtime/infrastructure/sandbox-backup-pruning.ts +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-backup-pruning.ts @@ -1,13 +1,15 @@ +import type { SandboxBackupId } from "@mosoo/id"; + import type { ReadySandboxBackupForPruning } from "./sandbox-backup-store"; export function selectSandboxBackupPruneIds( backups: readonly ReadySandboxBackupForPruning[], -): string[] { +): SandboxBackupId[] { const keepIds = new Set(); const readyCountsByDir = new Map(); for (const backup of backups) { - if (backup.keep) { + if (backup.keep || backup.protected) { keepIds.add(backup.id); continue; } diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-backup-reconciliation.service.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-backup-reconciliation.service.ts new file mode 100644 index 00000000..b4cb5b39 --- /dev/null +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-backup-reconciliation.service.ts @@ -0,0 +1,520 @@ +import type { ApiCommandId } from "@mosoo/db"; +import type { SandboxBackupId } from "@mosoo/id"; +import { parsePlatformId } from "@mosoo/id"; + +import { createErrorLogContext, logWarn } from "../../../platform/cloudflare/logger"; +import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; +import { + environmentPackageArtifactMetadataBackupId, + publishEnvironmentPackageArtifactBackup, + resolveEnvironmentPackageArtifactBackup, +} from "../../environments/application/environment-package-artifact-backup"; +import { + claimEnvironmentPackageArtifactBackupActual, + clearMissingEnvironmentPackageArtifactBackupActual, + completeEnvironmentPackageArtifactBackupStage, + getEnvironmentPackageArtifactBackupStage, + getEnvironmentPackageArtifactCommandIntent, + retireExpiredEnvironmentPackageArtifactBackups, + revokeTerminalEnvironmentPackageArtifactBackupStage, + revokeTerminalEnvironmentPackageArtifactBackupStages, +} from "../../environments/application/environment-package-artifact-backup-store"; +import { + environmentPackageArtifactDir, + parseEnvironmentPackageArtifactBackupName, +} from "../../environments/domain/environment-package-artifact"; +import type { + EnvironmentPackageArtifactBackupAuthorityRef, + EnvironmentPackageArtifactKey, + EnvironmentPackageArtifactPaths, +} from "../../environments/domain/environment-package-artifact"; +import { + decodeSandboxBackupIdForPlatform, + encodeSandboxBackupIdForStorage, +} from "./sandbox-backup-id"; +import { + deleteAuthorizedSandboxBackupObjects, + getSandboxBackupObjectKeys, + parseRuntimeSandboxBackupName, + parseSandboxBackupMetadata, +} from "./sandbox-backup-platform"; +import type { SandboxBackupMetadata } from "./sandbox-backup-platform"; +import { + authorizeSandboxBackupDeletion, + claimSandboxBackupStageActual, + finalizeSandboxBackupStage, + getSandboxBackupRecord, + getSandboxBackupRecordByStagingId, + getSandboxBackupStage, + listPendingSandboxBackupDeletions, +} from "./sandbox-backup-store"; +import type { SandboxBackupDeletionAuthority } from "./sandbox-backup-store"; + +const RECONCILIATION_PAGE_SIZE = 64; +const ORPHAN_GRACE_MS = 24 * 60 * 60_000; +const BACKUP_OBJECT_KEY = /^backups\/([^/]+)\/(data\.sqsh|meta\.json)$/u; + +export interface SandboxBackupReconciliationPageResult { + readonly hasMore: boolean; + readonly nextCursor: string | null; + readonly processed: number; +} + +function parseBackupObjectKey(key: string): { + readonly backupId: SandboxBackupId; + readonly kind: "data" | "meta"; +} | null { + const match = BACKUP_OBJECT_KEY.exec(key); + if (match === null) { + return null; + } + try { + return { + backupId: encodeSandboxBackupIdForStorage(match[1] ?? ""), + kind: match[2] === "meta.json" ? "meta" : "data", + }; + } catch { + return null; + } +} + +function isPastGrace(uploaded: Date, nowMs: number): boolean { + return uploaded.getTime() <= nowMs - ORPHAN_GRACE_MS; +} + +async function deleteSandboxBackup( + bindings: ApiBindings, + backupId: SandboxBackupId, + authority: SandboxBackupDeletionAuthority, +): Promise { + if (!(await authorizeSandboxBackupDeletion(bindings.DB, { authority, backupId }))) { + return false; + } + await deleteAuthorizedSandboxBackupObjects(bindings, [backupId]); + return true; +} + +function environmentArtifactPathsEqual( + left: EnvironmentPackageArtifactPaths, + right: EnvironmentPackageArtifactPaths, +): boolean { + return JSON.stringify(left) === JSON.stringify(right); +} + +async function readEnvironmentArtifactBackup( + bindings: ApiBindings, + key: EnvironmentPackageArtifactKey, +): Promise<{ + readonly backupId: SandboxBackupId; + readonly paths: EnvironmentPackageArtifactPaths; +} | null> { + const metadata = await resolveEnvironmentPackageArtifactBackup(bindings, key); + if (metadata === null) { + return null; + } + const backupId = environmentPackageArtifactMetadataBackupId(metadata); + if (backupId === null) { + throw new Error("Environment package artifact metadata has an invalid backup ID."); + } + return { backupId, paths: metadata.paths }; +} + +async function deleteEnvironmentArtifactBackupIfUnreferenced( + bindings: ApiBindings, + key: EnvironmentPackageArtifactKey, + backupId: SandboxBackupId, + authority: { + readonly attemptCount: number; + readonly commandId: ApiCommandId; + readonly deliveryGeneration: number; + readonly invalid?: boolean; + }, +): Promise { + if ((await readEnvironmentArtifactBackup(bindings, key))?.backupId !== backupId) { + await deleteSandboxBackup(bindings, backupId, { + attemptCount: authority.attemptCount, + commandId: authority.commandId, + deliveryGeneration: authority.deliveryGeneration, + kind: authority.invalid === true ? "environment_invalid" : "environment_candidate", + }); + } +} + +function parseEnvironmentArtifactCommandId( + authority: EnvironmentPackageArtifactBackupAuthorityRef, +): ApiCommandId | null { + try { + return parsePlatformId(authority.commandId, "environment artifact command ID"); + } catch { + return null; + } +} + +async function reconcileEnvironmentArtifactMetadataObject( + bindings: ApiBindings, + input: { + readonly authority: EnvironmentPackageArtifactBackupAuthorityRef; + readonly backupId: SandboxBackupId; + readonly dataExists: boolean; + readonly metadata: SandboxBackupMetadata; + readonly nowMs: number; + readonly uploaded: Date; + }, +): Promise { + const commandId = parseEnvironmentArtifactCommandId(input.authority); + if (commandId === null) { + return; + } + const [stage, intent] = await Promise.all([ + getEnvironmentPackageArtifactBackupStage(bindings.DB, commandId), + getEnvironmentPackageArtifactCommandIntent(bindings.DB, commandId), + ]); + const keySource = stage ?? intent; + if (keySource === null) { + return; + } + const key: EnvironmentPackageArtifactKey = { + appId: keySource.appId, + inputDigest: keySource.inputDigest, + }; + const current = await readEnvironmentArtifactBackup(bindings, key); + if (current !== null) { + if (stage !== null) { + if (!environmentArtifactPathsEqual(stage.paths, current.paths)) { + logWarn("runtime.environment_artifact.manifest_paths_mismatch", { + backupId: input.backupId, + commandId, + }); + return; + } + await completeEnvironmentPackageArtifactBackupStage(bindings.DB, { + actualBackupId: stage.actualBackupId, + attemptCount: stage.attemptCount, + claimOwner: stage.claimOwner, + commandId: stage.commandId, + deliveryGeneration: stage.deliveryGeneration, + }); + } + if (current.backupId !== input.backupId && isPastGrace(input.uploaded, input.nowMs)) { + await deleteEnvironmentArtifactBackupIfUnreferenced(bindings, key, input.backupId, { + attemptCount: input.authority.attemptCount, + commandId, + deliveryGeneration: input.authority.deliveryGeneration, + }); + } + return; + } + + const stageMatchesObjectAuthority = + stage !== null && + stage.attemptCount === input.authority.attemptCount && + stage.commandId === commandId && + stage.deliveryGeneration === input.authority.deliveryGeneration; + const objectMatchesStage = + stageMatchesObjectAuthority && + input.dataExists && + input.metadata.id === decodeSandboxBackupIdForPlatform(input.backupId) && + input.metadata.dir === stage.dir && + stage.dir === environmentPackageArtifactDir(key); + + if (!objectMatchesStage) { + if (stageMatchesObjectAuthority && stage.actualBackupId === input.backupId) { + await clearMissingEnvironmentPackageArtifactBackupActual(bindings.DB, { + actualBackupId: input.backupId, + attemptCount: stage.attemptCount, + claimOwner: stage.claimOwner, + commandId: stage.commandId, + deliveryGeneration: stage.deliveryGeneration, + }); + } + if (stageMatchesObjectAuthority) { + await revokeTerminalEnvironmentPackageArtifactBackupStage(bindings.DB, commandId); + } + if (isPastGrace(input.uploaded, input.nowMs)) { + await deleteEnvironmentArtifactBackupIfUnreferenced(bindings, key, input.backupId, { + attemptCount: input.authority.attemptCount, + commandId, + deliveryGeneration: input.authority.deliveryGeneration, + invalid: true, + }); + } + return; + } + + const claimed = await claimEnvironmentPackageArtifactBackupActual(bindings.DB, { + actualBackupId: input.backupId, + authority: { + attemptCount: stage.attemptCount, + claimOwner: stage.claimOwner, + deliveryGeneration: stage.deliveryGeneration, + }, + commandId, + dir: stage.dir, + }); + if (claimed?.actualBackupId === input.backupId) { + await publishEnvironmentPackageArtifactBackup(bindings, { + attemptCount: stage.attemptCount, + backupId: input.backupId, + claimOwner: stage.claimOwner, + commandId, + deliveryGeneration: stage.deliveryGeneration, + key, + paths: stage.paths, + }); + const committed = await readEnvironmentArtifactBackup(bindings, key); + if (committed !== null && environmentArtifactPathsEqual(stage.paths, committed.paths)) { + await completeEnvironmentPackageArtifactBackupStage(bindings.DB, { + actualBackupId: stage.actualBackupId ?? input.backupId, + attemptCount: stage.attemptCount, + claimOwner: stage.claimOwner, + commandId: stage.commandId, + deliveryGeneration: stage.deliveryGeneration, + }); + if (committed.backupId === input.backupId) { + return; + } + } + } else if (claimed === null) { + await revokeTerminalEnvironmentPackageArtifactBackupStage(bindings.DB, commandId); + } + if (isPastGrace(input.uploaded, input.nowMs)) { + await deleteEnvironmentArtifactBackupIfUnreferenced(bindings, key, input.backupId, { + attemptCount: input.authority.attemptCount, + commandId, + deliveryGeneration: input.authority.deliveryGeneration, + }); + } +} + +async function reconcileMetadataObject( + bindings: ApiBindings, + input: { readonly backupId: SandboxBackupId; readonly nowMs: number; readonly uploaded: Date }, +): Promise { + const publicRecord = await getSandboxBackupRecord(bindings.DB, input.backupId); + if (publicRecord?.status === "pruned") { + await deleteSandboxBackup(bindings, input.backupId, { kind: "pruned" }); + return; + } + + const [, metadataKey] = getSandboxBackupObjectKeys(input.backupId); + const stored = await bindings.SANDBOX_STATE_BUCKET.get(metadataKey); + if (stored === null) { + return; + } + let metadata = null; + try { + metadata = parseSandboxBackupMetadata(JSON.parse(await stored.text())); + } catch { + // Unknown or legacy metadata is never owned by this reconciler. + } + if (metadata === null) { + if (publicRecord?.status === "ready") { + logWarn("runtime.sandbox_backup.ready_metadata_invalid", { backupId: input.backupId }); + } + return; + } + const [dataKey] = getSandboxBackupObjectKeys(input.backupId); + const data = await bindings.SANDBOX_STATE_BUCKET.head(dataKey); + if (publicRecord?.status === "ready") { + if (data === null) { + logWarn("runtime.sandbox_backup.ready_data_missing", { backupId: input.backupId }); + } + return; + } + + const environmentAuthority = parseEnvironmentPackageArtifactBackupName(metadata.name); + if (environmentAuthority !== null) { + await reconcileEnvironmentArtifactMetadataObject(bindings, { + authority: environmentAuthority, + backupId: input.backupId, + dataExists: data !== null, + metadata, + nowMs: input.nowMs, + uploaded: input.uploaded, + }); + return; + } + + const stagingId = parseRuntimeSandboxBackupName(metadata.name); + if (stagingId === null) { + return; + } + const finalized = await getSandboxBackupRecordByStagingId(bindings.DB, stagingId); + if (finalized !== null) { + if (finalized.id !== input.backupId) { + await deleteSandboxBackup(bindings, input.backupId, { + kind: "runtime_candidate", + stagingId, + }); + } else if (finalized.status === "pruned") { + await deleteSandboxBackup(bindings, input.backupId, { kind: "pruned" }); + } + return; + } + + const stage = await getSandboxBackupStage(bindings.DB, stagingId); + if (data === null) { + if (isPastGrace(input.uploaded, input.nowMs)) { + await deleteSandboxBackup(bindings, input.backupId, { + kind: "runtime_invalid", + stagingId, + }); + } + return; + } + if (metadata.id !== decodeSandboxBackupIdForPlatform(input.backupId)) { + if (isPastGrace(input.uploaded, input.nowMs)) { + await deleteSandboxBackup(bindings, input.backupId, { + kind: "runtime_invalid", + stagingId, + }); + } + return; + } + if (stage === null) { + if (isPastGrace(input.uploaded, input.nowMs)) { + await deleteSandboxBackup(bindings, input.backupId, { + kind: "runtime_candidate", + stagingId, + }); + } + return; + } + if (stage.dir !== metadata.dir) { + if (isPastGrace(input.uploaded, input.nowMs)) { + await deleteSandboxBackup(bindings, input.backupId, { + kind: "runtime_invalid", + stagingId, + }); + } + return; + } + + const claimed = await claimSandboxBackupStageActual(bindings.DB, { + actualBackupId: input.backupId, + dir: metadata.dir, + sandboxIncarnation: stage.sandboxIncarnation, + stagingId, + }); + if (claimed === null) { + await deleteSandboxBackup(bindings, input.backupId, { + kind: "runtime_candidate", + stagingId, + }); + return; + } + if (claimed.actualBackupId !== input.backupId) { + await deleteSandboxBackup(bindings, input.backupId, { + kind: "runtime_candidate", + stagingId, + }); + return; + } + const result = await finalizeSandboxBackupStage(bindings.DB, { + actualBackupId: input.backupId, + stagingId, + }); + if (result === null || !result.candidateAccepted) { + await deleteSandboxBackup(bindings, input.backupId, { + kind: "runtime_candidate", + stagingId, + }); + } +} + +async function reconcileDataObject( + bindings: ApiBindings, + input: { readonly backupId: SandboxBackupId; readonly nowMs: number; readonly uploaded: Date }, +): Promise { + const publicRecord = await getSandboxBackupRecord(bindings.DB, input.backupId); + if (publicRecord?.status === "pruned") { + await deleteSandboxBackup(bindings, input.backupId, { kind: "pruned" }); + return; + } + const [, metadataKey] = getSandboxBackupObjectKeys(input.backupId); + if ((await bindings.SANDBOX_STATE_BUCKET.head(metadataKey)) !== null) { + return; + } + if (publicRecord?.status === "ready") { + logWarn("runtime.sandbox_backup.ready_metadata_missing", { backupId: input.backupId }); + return; + } + if (!isPastGrace(input.uploaded, input.nowMs)) { + return; + } + await deleteSandboxBackup(bindings, input.backupId, { kind: "unattributed" }); +} + +export async function reconcileSandboxBackupPage( + bindings: ApiBindings, + input: { readonly cursor: string | null }, +): Promise { + const clock = await bindings.DB.prepare( + "SELECT CAST(unixepoch('subsec') * 1000 AS INTEGER) AS now_ms", + ).first<{ now_ms: number }>(); + if (clock === null) { + throw new Error("Sandbox backup reconciliation could not read the D1 clock."); + } + const revokedStages = await revokeTerminalEnvironmentPackageArtifactBackupStages(bindings.DB); + const retired = await retireExpiredEnvironmentPackageArtifactBackups( + bindings.DB, + RECONCILIATION_PAGE_SIZE, + ); + const pending = await listPendingSandboxBackupDeletions(bindings.DB, RECONCILIATION_PAGE_SIZE); + for (const backupId of pending) { + try { + await deleteAuthorizedSandboxBackupObjects(bindings, [backupId]); + } catch (error) { + logWarn("runtime.sandbox_backup.pending_deletion_failed", { + ...createErrorLogContext(error), + backupId, + }); + } + } + const pendingAfter = + pending.length === RECONCILIATION_PAGE_SIZE + ? await listPendingSandboxBackupDeletions(bindings.DB, RECONCILIATION_PAGE_SIZE) + : []; + const databaseHasMore = + revokedStages === RECONCILIATION_PAGE_SIZE || + retired.length === RECONCILIATION_PAGE_SIZE || + (pendingAfter.length > 0 && + (pendingAfter.length !== pending.length || + pendingAfter.some((backupId, index) => backupId !== pending[index]))); + const page = await bindings.SANDBOX_STATE_BUCKET.list({ + ...(input.cursor === null ? {} : { cursor: input.cursor }), + limit: RECONCILIATION_PAGE_SIZE, + prefix: "backups/", + }); + for (const object of page.objects) { + const parsed = parseBackupObjectKey(object.key); + if (parsed === null) { + continue; + } + try { + if (parsed.kind === "meta") { + await reconcileMetadataObject(bindings, { + backupId: parsed.backupId, + nowMs: clock.now_ms, + uploaded: object.uploaded, + }); + } else { + await reconcileDataObject(bindings, { + backupId: parsed.backupId, + nowMs: clock.now_ms, + uploaded: object.uploaded, + }); + } + } catch (error) { + logWarn("runtime.sandbox_backup.reconciliation_object_failed", { + ...createErrorLogContext(error), + key: object.key, + }); + } + } + return { + hasMore: page.truncated || databaseHasMore, + nextCursor: page.truncated ? page.cursor : null, + processed: page.objects.length, + }; +} diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-backup-store.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-backup-store.ts index cbce1d7c..739271fc 100644 --- a/apps/api/src/modules/runtime/infrastructure/sandbox-backup-store.ts +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-backup-store.ts @@ -1,311 +1,1263 @@ import type { SessionStatus } from "@mosoo/contracts/session"; +import type { ApiCommandId } from "@mosoo/db"; import { sandboxBackupsTable, + sandboxBackupStagingTable, sandboxSessionsTable, - sandboxesTable, - nativeResumeRefsTable, sessionsTable, } from "@mosoo/db"; -import { parsePlatformId } from "@mosoo/id"; +import { createPlatformId, parsePlatformId } from "@mosoo/id"; import type { + DriverInstanceId, RuntimeOperationId, SandboxBackupId, SandboxId, SessionId, SessionRunId, } from "@mosoo/id"; -import { and, asc, desc, eq, inArray, isNull, sql } from "drizzle-orm"; +import { and, asc, eq, inArray, isNull } from "drizzle-orm"; -import type { AppDatabase } from "../../../platform/db/drizzle"; -import { - getAppDatabase, - getD1ChangeCount, - runAppDatabaseBatch, -} from "../../../platform/db/drizzle"; -import { currentTimestampMs } from "../../../time"; +import { getAppDatabase } from "../../../platform/db/drizzle"; +import type { RuntimeSubjectOperationLease } from "./runtime-subject-lifecycle/runtime-subject-store"; -// Each row binds ten values; nine rows stay below D1's per-statement variable limit. -const SANDBOX_BACKUP_INSERT_BATCH_SIZE = 9; -type AppDatabaseBatchItem = Parameters[0][number]; +const TERMINAL_BACKUP_DRIVER_STATUSES_SQL = + "('provisioning', 'connecting', 'ready', 'stopping', 'stopped', 'failed')"; +const DATABASE_NOW_MS_SQL = "CAST(unixepoch('subsec') * 1000 AS INTEGER)"; -export interface CreatedSandboxBackupRecord { +export interface SandboxBackupRecord { + readonly createdAt: number; readonly dir: string; - readonly id: string; + readonly id: SandboxBackupId; + readonly keep: boolean; + readonly operationId: RuntimeOperationId | null; + readonly sandboxId: SandboxId; + readonly sandboxIncarnation: number; + readonly sessionRunId: SessionRunId | null; + readonly stagingId: SandboxBackupId; + readonly status: "pruned" | "ready"; + readonly ttlSeconds: number; + readonly updatedAt: number; + readonly workspaceSessionId: SessionId | null; +} + +export interface SandboxBackupStage { + readonly actualBackupId: SandboxBackupId | null; + readonly claimOwner: string | null; + readonly createdAt: number; + readonly dir: string; + readonly driverGeneration: number | null; + readonly driverInstanceId: DriverInstanceId | null; + readonly id: SandboxBackupId; + readonly operationId: RuntimeOperationId | null; + readonly sandboxId: SandboxId; + readonly sandboxIncarnation: number; + readonly sessionRunId: SessionRunId | null; + readonly ttlSeconds: number; + readonly updatedAt: number; + readonly updatesSubjectBackup: boolean; + readonly workspaceSessionId: SessionId | null; } -export interface CreatedSandboxBackupWrite { - readonly backup: CreatedSandboxBackupRecord; +export interface SandboxBackupTarget { + readonly dir: string; readonly updateSandboxLastBackup: boolean; + readonly workspaceSessionId: string | null; } -export interface ReadySandboxBackupForPruning { +export type SandboxBackupAdmission = + | { readonly kind: "operation"; readonly lease: RuntimeSubjectOperationLease } + | { + readonly driverGeneration: number; + readonly driverInstanceId: DriverInstanceId; + readonly incarnation: number; + readonly kind: "terminal"; + readonly sessionId: SessionId; + readonly sessionRunId: SessionRunId; + }; + +export type SandboxBackupWrite = + | { readonly backup: SandboxBackupRecord; readonly kind: "finalized" } + | { readonly isNew: boolean; readonly kind: "staged"; readonly stage: SandboxBackupStage }; + +interface SandboxBackupScope { readonly dir: string; - readonly id: SandboxBackupId; - readonly keep: boolean; + readonly operationId: RuntimeOperationId | null; + readonly sandboxId: SandboxId; + readonly sandboxIncarnation: number; + readonly sessionRunId: SessionRunId | null; } -export async function listReadySandboxBackupsForSessionRun( +function assertPositiveSafeInteger(value: number, label: string): void { + if (!Number.isSafeInteger(value) || value <= 0) { + throw new TypeError(`${label} must be a positive safe integer.`); + } +} + +function scopeForTarget( + sandboxId: SandboxId, + admission: SandboxBackupAdmission, + dir: string, +): SandboxBackupScope { + return { + dir, + operationId: admission.kind === "operation" ? admission.lease.operationId : null, + sandboxId, + sandboxIncarnation: + admission.kind === "operation" ? admission.lease.incarnation : admission.incarnation, + sessionRunId: admission.kind === "terminal" ? admission.sessionRunId : null, + }; +} + +function mapBackup(row: typeof sandboxBackupsTable.$inferSelect): SandboxBackupRecord { + if (row.status !== "ready" && row.status !== "pruned") { + throw new Error("Sandbox backup storage returned an invalid public status."); + } + return { ...row, status: row.status }; +} + +export async function getSandboxBackupStage( + database: D1Database, + stagingId: SandboxBackupId, +): Promise { + return ( + (await getAppDatabase(database) + .select() + .from(sandboxBackupStagingTable) + .where(eq(sandboxBackupStagingTable.id, stagingId)) + .limit(1) + .get()) ?? null + ); +} + +export async function getSandboxBackupRecord( + database: D1Database, + backupId: SandboxBackupId, +): Promise { + const row = await getAppDatabase(database) + .select() + .from(sandboxBackupsTable) + .where(eq(sandboxBackupsTable.id, backupId)) + .limit(1) + .get(); + return row === undefined ? null : mapBackup(row); +} + +export async function getSandboxBackupRecordByStagingId( + database: D1Database, + stagingId: SandboxBackupId, +): Promise { + const row = await getAppDatabase(database) + .select() + .from(sandboxBackupsTable) + .where(eq(sandboxBackupsTable.stagingId, stagingId)) + .limit(1) + .get(); + return row === undefined ? null : mapBackup(row); +} + +export type SandboxBackupDeletionAuthority = + | { readonly kind: "pruned" } + | { readonly kind: "unattributed" } + | { + readonly kind: "runtime_candidate" | "runtime_invalid"; + readonly stagingId: SandboxBackupId; + } + | { + readonly attemptCount: number; + readonly commandId: ApiCommandId; + readonly deliveryGeneration: number; + readonly kind: "environment_candidate" | "environment_invalid"; + }; + +export async function authorizeSandboxBackupDeletion( database: D1Database, input: { - readonly sandboxId: string; - readonly sessionRunId: string; + readonly authority: SandboxBackupDeletionAuthority; + readonly backupId: SandboxBackupId; }, -): Promise { - const sandboxId = parsePlatformId(input.sandboxId, "sandbox id"); - const sessionRunId = parsePlatformId(input.sessionRunId, "session run id"); - const rows = await getAppDatabase(database) - .select({ - dir: sandboxBackupsTable.dir, - id: sandboxBackupsTable.id, - }) - .from(sandboxBackupsTable) - .where( - and( - eq(sandboxBackupsTable.sandboxId, sandboxId), - eq(sandboxBackupsTable.sessionRunId, sessionRunId), - eq(sandboxBackupsTable.status, "ready"), - ), +): Promise { + const { authority, backupId } = input; + const invalidRuntime = authority.kind === "runtime_invalid"; + const invalidEnvironment = authority.kind === "environment_invalid"; + const preparations: D1PreparedStatement[] = []; + if (invalidRuntime) { + preparations.push( + database + .prepare( + `UPDATE sandbox_backup_staging AS stage + SET actual_backup_id = NULL, updated_at = ${DATABASE_NOW_MS_SQL} + WHERE stage.id = ? AND stage.actual_backup_id = ? + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup AS backup + WHERE backup.id = ? AND backup.status = 'ready' + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup AS artifact + WHERE artifact.backup_id = ? + )`, + ) + .bind(authority.stagingId, backupId, backupId, backupId), + ); + } else if (invalidEnvironment) { + preparations.push( + database + .prepare( + `UPDATE environment_package_artifact_backup_staging AS stage + SET actual_backup_id = NULL, updated_at = ${DATABASE_NOW_MS_SQL} + WHERE stage.command_id = ? AND stage.delivery_generation = ? + AND stage.attempt_count = ? AND stage.actual_backup_id = ? + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup AS artifact + WHERE artifact.backup_id = ? + )`, + ) + .bind( + authority.commandId, + authority.deliveryGeneration, + authority.attemptCount, + backupId, + backupId, + ), + ); + } + + const requiresPruned = + authority.kind === "pruned" + ? `AND EXISTS ( + SELECT 1 FROM sandbox_backup AS pruned + WHERE pruned.id = ? AND pruned.status = 'pruned' + )` + : ""; + const blocksUnclaimedRuntime = + authority.kind === "runtime_candidate" + ? `AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_staging AS candidate_stage + WHERE candidate_stage.id = ? AND candidate_stage.actual_backup_id IS NULL + )` + : ""; + const blocksUnclaimedEnvironment = + authority.kind === "environment_candidate" + ? `AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup_staging AS candidate_stage + WHERE candidate_stage.command_id = ? AND candidate_stage.delivery_generation = ? + AND candidate_stage.attempt_count = ? + AND candidate_stage.actual_backup_id IS NULL + )` + : ""; + const authorityBindings = + authority.kind === "pruned" + ? [backupId] + : authority.kind === "runtime_candidate" + ? [authority.stagingId] + : authority.kind === "environment_candidate" + ? [authority.commandId, authority.deliveryGeneration, authority.attemptCount] + : []; + const authorize = database + .prepare( + `INSERT INTO sandbox_backup_delete_intent (backup_id, created_at, delete_after, deleted_at) + SELECT ?, ${DATABASE_NOW_MS_SQL}, ${DATABASE_NOW_MS_SQL}, NULL + WHERE NOT EXISTS ( + SELECT 1 FROM sandbox_backup_delete_intent AS deletion + WHERE deletion.backup_id = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup AS backup + WHERE backup.id = ? AND backup.status = 'ready' + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_staging AS stage + WHERE stage.actual_backup_id = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup_staging AS stage + WHERE stage.actual_backup_id = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup AS artifact + WHERE artifact.backup_id = ? + ) + ${requiresPruned} + ${blocksUnclaimedRuntime} + ${blocksUnclaimedEnvironment} + `, ) - .all(); + .bind(backupId, backupId, backupId, backupId, backupId, backupId, ...authorityBindings); + if (preparations.length === 0) { + await authorize.run(); + } else { + await database.batch([...preparations, authorize]); + } + return ( + (await database + .prepare( + `SELECT 1 FROM sandbox_backup_delete_intent + WHERE backup_id = ? AND delete_after <= ${DATABASE_NOW_MS_SQL}`, + ) + .bind(backupId) + .first()) !== null + ); +} - return rows; +export async function listPendingSandboxBackupDeletions( + database: D1Database, + limit: number, +): Promise { + assertPositiveSafeInteger(limit, "Sandbox backup deletion limit"); + const rows = await database + .prepare( + `SELECT backup_id + FROM sandbox_backup_delete_intent + WHERE deleted_at IS NULL AND delete_after <= ${DATABASE_NOW_MS_SQL} + ORDER BY coalesce(attempted_at, delete_after), attempted_at IS NOT NULL, + created_at, backup_id + LIMIT ?`, + ) + .bind(limit) + .all<{ backup_id: SandboxBackupId }>(); + return rows.results.map((row) => row.backup_id); } -export interface SandboxSessionBackupCandidate { - readonly cwd: string; - readonly lastMessageAt: number | null; - readonly sessionId: SessionId; - readonly sessionStatus: SessionStatus; +export async function beginSandboxBackupDeletionAttempt( + database: D1Database, + backupId: SandboxBackupId, +): Promise { + const result = await database + .prepare( + `UPDATE sandbox_backup_delete_intent + SET attempted_at = ${DATABASE_NOW_MS_SQL} + WHERE backup_id = ? AND deleted_at IS NULL AND delete_after <= ${DATABASE_NOW_MS_SQL}`, + ) + .bind(backupId) + .run(); + return (result.meta.changes ?? 0) === 1; } -function parseSandboxBackupIds(values: readonly string[], label: string): SandboxBackupId[] { - return values.map((value, index) => - parsePlatformId(value, `${label}[${index}]`), +export async function isSandboxBackupDeletionAuthorized( + database: D1Database, + backupId: SandboxBackupId, +): Promise { + return ( + (await database + .prepare( + `SELECT 1 FROM sandbox_backup_delete_intent + WHERE backup_id = ? AND delete_after <= ${DATABASE_NOW_MS_SQL}`, + ) + .bind(backupId) + .first()) !== null ); } -function sandboxStatusOperationCondition(operationId: RuntimeOperationId | null | undefined) { - if (operationId === undefined) { - return []; - } - - return operationId === null - ? [isNull(sandboxesTable.statusOperationId)] - : [eq(sandboxesTable.statusOperationId, operationId)]; +export async function completeSandboxBackupDeletion( + database: D1Database, + backupId: SandboxBackupId, +): Promise { + const result = await database + .prepare( + `UPDATE sandbox_backup_delete_intent + SET deleted_at = coalesce(deleted_at, ${DATABASE_NOW_MS_SQL}) + WHERE backup_id = ? AND attempted_at IS NOT NULL`, + ) + .bind(backupId) + .run(); + return (result.meta.changes ?? 0) === 1; } -export async function listReadySandboxBackupsForPruning( +async function getStageByScope( database: D1Database, - sandboxId: string, -): Promise { - const parsedSandboxId = parsePlatformId(sandboxId, "sandbox id"); + scope: SandboxBackupScope, +): Promise { + return ( + (await getAppDatabase(database) + .select() + .from(sandboxBackupStagingTable) + .where( + and( + eq(sandboxBackupStagingTable.sandboxId, scope.sandboxId), + eq(sandboxBackupStagingTable.sandboxIncarnation, scope.sandboxIncarnation), + eq(sandboxBackupStagingTable.dir, scope.dir), + scope.operationId === null + ? isNull(sandboxBackupStagingTable.operationId) + : eq(sandboxBackupStagingTable.operationId, scope.operationId), + scope.sessionRunId === null + ? isNull(sandboxBackupStagingTable.sessionRunId) + : eq(sandboxBackupStagingTable.sessionRunId, scope.sessionRunId), + ), + ) + .limit(1) + .get()) ?? null + ); +} - return getAppDatabase(database) - .select({ - dir: sandboxBackupsTable.dir, - id: sandboxBackupsTable.id, - keep: sandboxBackupsTable.keep, - }) +async function getFinalizedSandboxBackupByScope( + database: D1Database, + scope: SandboxBackupScope, +): Promise { + const row = await getAppDatabase(database) + .select() .from(sandboxBackupsTable) .where( and( - eq(sandboxBackupsTable.sandboxId, parsedSandboxId), - eq(sandboxBackupsTable.status, "ready"), + eq(sandboxBackupsTable.sandboxId, scope.sandboxId), + eq(sandboxBackupsTable.sandboxIncarnation, scope.sandboxIncarnation), + eq(sandboxBackupsTable.dir, scope.dir), + scope.operationId === null + ? isNull(sandboxBackupsTable.operationId) + : eq(sandboxBackupsTable.operationId, scope.operationId), + scope.sessionRunId === null + ? isNull(sandboxBackupsTable.sessionRunId) + : eq(sandboxBackupsTable.sessionRunId, scope.sessionRunId), ), ) - .orderBy(asc(sandboxBackupsTable.dir), desc(sandboxBackupsTable.createdAt)) - .all(); + .limit(1) + .get(); + return row === undefined ? null : mapBackup(row); } -export async function markSandboxBackupsPruned( - database: D1Database, - backupIds: readonly string[], -): Promise { - if (backupIds.length === 0) { - return; +function assertTargetMatches( + value: Pick, + target: SandboxBackupTarget, +): void { + if (value.workspaceSessionId !== target.workspaceSessionId) { + throw new Error("A sandbox backup retry changed its immutable workspace target."); } + if ( + "updatesSubjectBackup" in value && + value.updatesSubjectBackup !== target.updateSandboxLastBackup + ) { + throw new Error("A sandbox backup retry changed its subject checkpoint target."); + } +} - const parsedBackupIds = parseSandboxBackupIds(backupIds, "sandbox backup id"); +function workspaceAdmission( + target: SandboxBackupTarget, + scope: SandboxBackupScope, +): { + readonly bindings: readonly unknown[]; + readonly predicate: string; +} { + if (target.workspaceSessionId === null) { + return { bindings: [], predicate: "1 = 1" }; + } + return { + bindings: [target.workspaceSessionId, scope.sandboxId, scope.sandboxIncarnation, scope.dir], + predicate: `EXISTS ( + SELECT 1 + FROM sandbox_session AS workspace + JOIN session AS logical_session ON logical_session.id = workspace.session_id + WHERE workspace.session_id = ? + AND workspace.sandbox_id = ? + AND workspace.sandbox_incarnation = ? + AND workspace.cwd = ? + AND workspace.status IN ('active', 'closed') + AND workspace.cleanup_operation_id IS NULL + AND logical_session.archived_at IS NULL + AND logical_session.cleanup_operation_kind IS NULL + )`, + }; +} - await getAppDatabase(database) - .update(sandboxBackupsTable) - .set({ - status: "pruned", - updatedAt: currentTimestampMs(), - }) - .where(inArray(sandboxBackupsTable.id, [...new Set(parsedBackupIds)])) +async function insertStage( + database: D1Database, + input: { + readonly admission: SandboxBackupAdmission; + readonly sandboxId: SandboxId; + readonly target: SandboxBackupTarget; + readonly ttlSeconds: number; + }, +): Promise { + const id = createPlatformId(); + const scope = scopeForTarget(input.sandboxId, input.admission, input.target.dir); + const workspace = workspaceAdmission(input.target, scope); + const authority = + input.admission.kind === "operation" + ? { + bindings: [ + input.sandboxId, + input.admission.lease.incarnation, + input.admission.lease.operationId, + input.admission.lease.claimOwner, + ], + predicate: `EXISTS ( + SELECT 1 FROM sandbox + WHERE id = ? AND incarnation = ? AND status = 'backing_up' + AND status_operation_id = ? AND claim_owner = ? + AND claim_expires_at > ${DATABASE_NOW_MS_SQL} + )`, + } + : { + bindings: [ + input.admission.sessionRunId, + input.admission.sessionId, + input.admission.driverInstanceId, + input.admission.driverInstanceId, + input.admission.driverGeneration, + input.sandboxId, + input.admission.incarnation, + input.admission.sessionId, + input.admission.sessionRunId, + ], + predicate: `EXISTS ( + SELECT 1 + FROM session_run AS terminal_run + JOIN driver_instance AS terminal_driver + ON terminal_driver.id = terminal_run.driver_instance_id + WHERE terminal_run.id = ? AND terminal_run.session_id = ? + AND terminal_run.driver_instance_id = ? AND terminal_run.status = 'completed' + AND terminal_driver.id = ? AND terminal_driver.generation = ? + AND terminal_driver.sandbox_id = ? AND terminal_driver.sandbox_incarnation = ? + AND terminal_driver.sandbox_session_id = ? + AND terminal_driver.status IN ${TERMINAL_BACKUP_DRIVER_STATUSES_SQL} + AND terminal_driver.status_operation_id = ? + AND NOT EXISTS ( + SELECT 1 FROM session_run AS successor + WHERE successor.session_id = terminal_run.session_id + AND successor.id <> terminal_run.id + AND successor.status IN ('queued', 'booting', 'running', 'waiting_input') + ) + )`, + }; + const result = await database + .prepare( + `INSERT INTO sandbox_backup_staging ( + actual_backup_id, claim_owner, created_at, dir, driver_generation, driver_instance_id, + id, operation_id, sandbox_id, sandbox_incarnation, session_run_id, ttl_seconds, updated_at, + updates_subject_backup, workspace_session_id + ) + SELECT NULL, ?, ${DATABASE_NOW_MS_SQL}, ?, ?, ?, ?, ?, ?, ?, ?, ?, + ${DATABASE_NOW_MS_SQL}, ?, ? + WHERE ${authority.predicate} AND ${workspace.predicate} + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup AS finalized + WHERE finalized.staging_id = ? + OR (finalized.sandbox_id = ? AND finalized.sandbox_incarnation = ? + AND finalized.dir = ? AND finalized.operation_id IS ? + AND finalized.session_run_id IS ?) + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_staging AS existing + WHERE existing.id = ? + OR (existing.sandbox_id = ? AND existing.sandbox_incarnation = ? + AND existing.dir = ? AND existing.operation_id IS ? + AND existing.session_run_id IS ?) + ) + ON CONFLICT DO NOTHING`, + ) + .bind( + input.admission.kind === "operation" ? input.admission.lease.claimOwner : null, + scope.dir, + input.admission.kind === "terminal" ? input.admission.driverGeneration : null, + input.admission.kind === "terminal" ? input.admission.driverInstanceId : null, + id, + scope.operationId, + scope.sandboxId, + scope.sandboxIncarnation, + scope.sessionRunId, + input.ttlSeconds, + input.target.updateSandboxLastBackup ? 1 : 0, + input.target.workspaceSessionId, + ...authority.bindings, + ...workspace.bindings, + id, + scope.sandboxId, + scope.sandboxIncarnation, + scope.dir, + scope.operationId, + scope.sessionRunId, + id, + scope.sandboxId, + scope.sandboxIncarnation, + scope.dir, + scope.operationId, + scope.sessionRunId, + ) .run(); + return (result.meta.changes ?? 0) === 1; } -export async function recordCreatedSandboxBackups( +export async function stageSandboxBackupWrites( database: D1Database, input: { - readonly backups: readonly CreatedSandboxBackupWrite[]; - readonly checkpointSessionId?: string; - readonly operationId?: string | null; + readonly admission: SandboxBackupAdmission; readonly sandboxId: string; - readonly sessionRunId?: string; + readonly targets: readonly SandboxBackupTarget[]; readonly ttlSeconds: number; }, -): Promise { - if (input.backups.length === 0) { - return; +): Promise { + assertPositiveSafeInteger(input.ttlSeconds, "Sandbox backup TTL"); + assertPositiveSafeInteger( + input.admission.kind === "operation" + ? input.admission.lease.incarnation + : input.admission.incarnation, + "Sandbox backup incarnation", + ); + if ( + input.admission.kind === "terminal" && + (!Number.isSafeInteger(input.admission.driverGeneration) || + input.admission.driverGeneration < 0) + ) { + throw new TypeError("Terminal sandbox backup identity is invalid."); } - const sandboxId = parsePlatformId(input.sandboxId, "sandbox id"); - const checkpointSessionId = - input.checkpointSessionId === undefined - ? null - : parsePlatformId(input.checkpointSessionId, "checkpoint session id"); - const operationId = - input.operationId === undefined || input.operationId === null - ? input.operationId - : parsePlatformId(input.operationId, "runtime operation id"); - const sessionRunId = - input.sessionRunId === undefined - ? null - : parsePlatformId(input.sessionRunId, "session run id"); - const now = currentTimestampMs(); - const backupRows = input.backups.map((entry, index) => ({ - createdAt: now, - dir: entry.backup.dir, - errorMessage: null, - id: parsePlatformId(entry.backup.id, `sandbox backup id ${index}`), - keep: false, - sandboxId, - sessionRunId, - status: "ready" as const, - ttlSeconds: input.ttlSeconds, - updatedAt: now, - })); - let subjectCheckpointBackup: CreatedSandboxBackupRecord | null = null; + const writes: SandboxBackupWrite[] = []; - for (const entry of input.backups) { - if (entry.updateSandboxLastBackup) { - subjectCheckpointBackup = entry.backup; + for (const rawTarget of input.targets) { + const target: SandboxBackupTarget = { + ...rawTarget, + workspaceSessionId: + rawTarget.workspaceSessionId === null + ? null + : parsePlatformId(rawTarget.workspaceSessionId, "workspace session id"), + }; + if (input.admission.kind === "terminal" && target.workspaceSessionId === null) { + throw new Error("Terminal sandbox backups require an exact workspace session."); } - } - - const results = await runAppDatabaseBatch(database, (appDb) => { - const queries: [AppDatabaseBatchItem, ...AppDatabaseBatchItem[]] = [ - appDb - .insert(sandboxBackupsTable) - .values(backupRows.slice(0, SANDBOX_BACKUP_INSERT_BATCH_SIZE)), - ]; - - for ( - let index = SANDBOX_BACKUP_INSERT_BATCH_SIZE; - index < backupRows.length; - index += SANDBOX_BACKUP_INSERT_BATCH_SIZE + const scope = scopeForTarget(sandboxId, input.admission, target.dir); + const finalized = await getFinalizedSandboxBackupByScope(database, scope); + if (finalized !== null) { + assertTargetMatches(finalized, target); + writes.push({ backup: finalized, kind: "finalized" }); + continue; + } + let stage = await getStageByScope(database, scope); + let isNew = false; + if ( + stage !== null && + input.admission.kind === "operation" && + stage.claimOwner !== input.admission.lease.claimOwner ) { - queries.push( - appDb - .insert(sandboxBackupsTable) - .values(backupRows.slice(index, index + SANDBOX_BACKUP_INSERT_BATCH_SIZE)), - ); + if (await isSandboxBackupStageCurrent(database, stage.id)) { + throw new Error("Sandbox backup stage belongs to another live lease owner."); + } + const revoked = await revokeSandboxBackupStage(database, { + onlyIfStale: true, + stagingId: stage.id, + }); + if (revoked?.actualBackupId !== null && revoked?.actualBackupId !== undefined) { + await authorizeSandboxBackupDeletion(database, { + authority: { kind: "runtime_candidate", stagingId: stage.id }, + backupId: revoked.actualBackupId, + }); + } + stage = null; } - - if (checkpointSessionId !== null && sessionRunId !== null) { - queries.push( - appDb - .update(nativeResumeRefsTable) - .set({ - committedSessionRunId: sessionRunId, - committedValue: sql`${nativeResumeRefsTable.value}`, - }) - .where( - and( - eq(nativeResumeRefsTable.sessionId, checkpointSessionId), - eq(nativeResumeRefsTable.observedSessionRunId, sessionRunId), - ), - ), - ); + if (stage === null) { + isNew = await insertStage(database, { + admission: input.admission, + sandboxId, + target, + ttlSeconds: input.ttlSeconds, + }); + const racedFinalized = await getFinalizedSandboxBackupByScope(database, scope); + if (racedFinalized !== null) { + assertTargetMatches(racedFinalized, target); + writes.push({ backup: racedFinalized, kind: "finalized" }); + continue; + } + stage = await getStageByScope(database, scope); } - - if (subjectCheckpointBackup) { - queries.push( - appDb - .update(sandboxesTable) - .set({ - lastBackupId: parsePlatformId( - subjectCheckpointBackup.id, - "checkpoint sandbox backup id", - ), - updatedAt: now, - }) - .where( - and( - eq(sandboxesTable.id, sandboxId), - inArray(sandboxesTable.status, ["backing_up", "destroying"]), - ...sandboxStatusOperationCondition(operationId), - ), - ), - ); + if (stage === null) { + throw new Error("Sandbox backup admission lost its exact lifecycle authority."); } + if ( + input.admission.kind === "operation" + ? stage.claimOwner !== input.admission.lease.claimOwner + : stage.claimOwner !== null || + stage.driverInstanceId !== input.admission.driverInstanceId || + stage.driverGeneration !== input.admission.driverGeneration + ) { + throw new Error("Sandbox backup stage belongs to another admission authority."); + } + if (!(await isSandboxBackupStageCurrent(database, stage.id))) { + throw new Error("Sandbox backup stage lost its lifecycle authority before use."); + } + assertTargetMatches(stage, target); + writes.push({ isNew, kind: "staged", stage }); + } + return writes; +} + +function currentAuthority(alias: string): string { + return `( + (${alias}.operation_id IS NOT NULL AND EXISTS ( + SELECT 1 FROM sandbox AS subject + WHERE subject.id = ${alias}.sandbox_id + AND subject.incarnation = ${alias}.sandbox_incarnation + AND subject.status = 'backing_up' + AND subject.status_operation_id = ${alias}.operation_id + AND ${alias}.claim_owner IS NOT NULL + AND subject.claim_owner = ${alias}.claim_owner + AND subject.claim_expires_at > ${DATABASE_NOW_MS_SQL} + )) OR (${alias}.session_run_id IS NOT NULL AND EXISTS ( + SELECT 1 FROM sandbox AS subject + JOIN session_run AS terminal_run ON terminal_run.id = ${alias}.session_run_id + JOIN driver_instance AS terminal_driver + ON terminal_driver.id = ${alias}.driver_instance_id + WHERE subject.id = ${alias}.sandbox_id + AND subject.incarnation = ${alias}.sandbox_incarnation + AND terminal_run.session_id = ${alias}.workspace_session_id + AND terminal_run.status = 'completed' + AND terminal_run.driver_instance_id = ${alias}.driver_instance_id + AND terminal_driver.generation = ${alias}.driver_generation + AND terminal_driver.sandbox_id = ${alias}.sandbox_id + AND terminal_driver.sandbox_incarnation = ${alias}.sandbox_incarnation + AND terminal_driver.sandbox_session_id = ${alias}.workspace_session_id + AND terminal_driver.status IN ${TERMINAL_BACKUP_DRIVER_STATUSES_SQL} + AND terminal_driver.status_operation_id = ${alias}.session_run_id + AND NOT EXISTS ( + SELECT 1 FROM session_run AS successor + WHERE successor.session_id = terminal_run.session_id + AND successor.id <> terminal_run.id + AND successor.status IN ('queued', 'booting', 'running', 'waiting_input') + ) + )) + ) AND (${alias}.workspace_session_id IS NULL OR EXISTS ( + SELECT 1 + FROM sandbox_session AS workspace + JOIN session AS logical_session ON logical_session.id = workspace.session_id + WHERE workspace.session_id = ${alias}.workspace_session_id + AND workspace.sandbox_id = ${alias}.sandbox_id + AND workspace.sandbox_incarnation = ${alias}.sandbox_incarnation + AND workspace.cwd = ${alias}.dir + AND workspace.status IN ('active', 'closed') + AND workspace.cleanup_operation_id IS NULL + AND logical_session.archived_at IS NULL + AND logical_session.cleanup_operation_kind IS NULL + ))`; +} - return queries; - }); +export async function isSandboxBackupStageCurrent( + database: D1Database, + stagingId: SandboxBackupId, +): Promise { + return ( + (await database + .prepare( + `SELECT id FROM sandbox_backup_staging AS stage + WHERE id = ? AND ${currentAuthority("stage")}`, + ) + .bind(stagingId) + .first()) !== null + ); +} - if (!subjectCheckpointBackup) { - return; +export async function claimSandboxBackupStageActual( + database: D1Database, + input: { + readonly actualBackupId: SandboxBackupId; + readonly dir: string; + readonly sandboxIncarnation: number; + readonly stagingId: SandboxBackupId; + }, +): Promise<{ readonly actualBackupId: SandboxBackupId } | null> { + const claimed = await database + .prepare( + `UPDATE sandbox_backup_staging AS stage + SET actual_backup_id = ?, updated_at = ${DATABASE_NOW_MS_SQL} + WHERE id = ? AND sandbox_incarnation = ? AND dir = ? + AND actual_backup_id IS NULL AND ${currentAuthority("stage")} + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_delete_intent AS deletion + WHERE deletion.backup_id = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup + WHERE id = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup + WHERE backup_id = ? + UNION ALL + SELECT 1 FROM environment_package_artifact_backup_staging + WHERE actual_backup_id = ? + ) + RETURNING actual_backup_id`, + ) + .bind( + input.actualBackupId, + input.stagingId, + input.sandboxIncarnation, + input.dir, + input.actualBackupId, + input.actualBackupId, + input.actualBackupId, + input.actualBackupId, + ) + .first<{ actual_backup_id: SandboxBackupId }>(); + if (claimed !== null) { + return { actualBackupId: claimed.actual_backup_id }; } - const updated = results.at(-1); + const winner = await database + .prepare( + `SELECT actual_backup_id + FROM ( + SELECT actual_backup_id, 0 AS priority + FROM sandbox_backup_staging + WHERE id = ? AND actual_backup_id IS NOT NULL + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_delete_intent AS deletion + WHERE deletion.backup_id = sandbox_backup_staging.actual_backup_id + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup + WHERE id = sandbox_backup_staging.actual_backup_id + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup + WHERE backup_id = sandbox_backup_staging.actual_backup_id + UNION ALL + SELECT 1 FROM environment_package_artifact_backup_staging + WHERE actual_backup_id = sandbox_backup_staging.actual_backup_id + ) + UNION ALL + SELECT id AS actual_backup_id, 1 AS priority + FROM sandbox_backup + WHERE staging_id = ? AND status = 'ready' + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_delete_intent AS deletion + WHERE deletion.backup_id = sandbox_backup.id + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup + WHERE backup_id = sandbox_backup.id + UNION ALL + SELECT 1 FROM environment_package_artifact_backup_staging + WHERE actual_backup_id = sandbox_backup.id + ) + ) + ORDER BY priority + LIMIT 1`, + ) + .bind(input.stagingId, input.stagingId) + .first<{ actual_backup_id: SandboxBackupId }>(); + return winner === null ? null : { actualBackupId: winner.actual_backup_id }; +} + +export async function clearMissingSandboxBackupStageActual( + database: D1Database, + input: { readonly actualBackupId: SandboxBackupId; readonly stagingId: SandboxBackupId }, +): Promise { + const result = await database + .prepare( + `UPDATE sandbox_backup_staging + SET actual_backup_id = NULL, updated_at = ${DATABASE_NOW_MS_SQL} + WHERE id = ? AND actual_backup_id = ? + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup + WHERE staging_id = sandbox_backup_staging.id + )`, + ) + .bind(input.stagingId, input.actualBackupId) + .run(); + return (result.meta.changes ?? 0) === 1; +} + +export async function revokeSandboxBackupStage( + database: D1Database, + input: { readonly onlyIfStale?: boolean; readonly stagingId: SandboxBackupId }, +): Promise<{ readonly actualBackupId: SandboxBackupId | null } | null> { + const row = await database + .prepare( + `DELETE FROM sandbox_backup_staging AS stage + WHERE id = ?${input.onlyIfStale === true ? ` AND NOT (${currentAuthority("stage")})` : ""} + RETURNING actual_backup_id`, + ) + .bind(input.stagingId) + .first<{ actual_backup_id: SandboxBackupId | null }>(); + return row === null ? null : { actualBackupId: row.actual_backup_id }; +} + +export async function listSandboxBackupStages( + database: D1Database, + limit: number, +): Promise { + assertPositiveSafeInteger(limit, "Sandbox backup stage limit"); + return getAppDatabase(database) + .select() + .from(sandboxBackupStagingTable) + .orderBy(asc(sandboxBackupStagingTable.updatedAt), asc(sandboxBackupStagingTable.id)) + .limit(limit) + .all(); +} + +export async function deferSandboxBackupStageRepair( + database: D1Database, + stage: Pick, +): Promise { + return ( + (await database + .prepare( + `UPDATE sandbox_backup_staging + SET updated_at = max(${DATABASE_NOW_MS_SQL}, updated_at + 1) + WHERE id = ? AND updated_at = ? AND updated_at < 9007199254740991 + RETURNING id`, + ) + .bind(stage.id, stage.updatedAt) + .first()) !== null + ); +} + +export interface FinalizeSandboxBackupResult { + readonly backup: SandboxBackupRecord; + readonly candidateAccepted: boolean; + readonly complete: boolean; +} + +export async function finalizeSandboxBackupStage( + database: D1Database, + input: { readonly actualBackupId: SandboxBackupId; readonly stagingId: SandboxBackupId }, +): Promise { + const stage = await getSandboxBackupStage(database, input.stagingId); + if (stage === null) { + const finalized = await getSandboxBackupRecordByStagingId(database, input.stagingId); + return finalized === null + ? null + : { + backup: finalized, + candidateAccepted: finalized.id === input.actualBackupId, + complete: true, + }; + } + const scope: SandboxBackupScope = { + dir: stage.dir, + operationId: stage.operationId, + sandboxId: stage.sandboxId, + sandboxIncarnation: stage.sandboxIncarnation, + sessionRunId: stage.sessionRunId, + }; + const existing = await getFinalizedSandboxBackupByScope(database, scope); + if (existing !== null && existing.stagingId !== stage.id) { + await revokeSandboxBackupStage(database, { stagingId: stage.id }); + return { backup: existing, candidateAccepted: false, complete: true }; + } + if (stage.actualBackupId !== input.actualBackupId) { + return null; + } + + await database.batch([ + database + .prepare( + `INSERT INTO sandbox_backup ( + created_at, dir, id, keep, operation_id, sandbox_id, sandbox_incarnation, + session_run_id, staging_id, status, ttl_seconds, updated_at, workspace_session_id + ) + SELECT created_at, dir, actual_backup_id, 0, operation_id, sandbox_id, + sandbox_incarnation, session_run_id, id, 'ready', ttl_seconds, + ${DATABASE_NOW_MS_SQL}, workspace_session_id + FROM sandbox_backup_staging AS stage + WHERE id = ? AND actual_backup_id = ? AND ${currentAuthority("stage")} + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_delete_intent AS deletion + WHERE deletion.backup_id = stage.actual_backup_id + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup AS existing + WHERE existing.id = stage.actual_backup_id + OR existing.staging_id = stage.id + OR (existing.sandbox_id = stage.sandbox_id + AND existing.sandbox_incarnation = stage.sandbox_incarnation + AND existing.dir = stage.dir + AND ((stage.operation_id IS NOT NULL + AND existing.operation_id = stage.operation_id) + OR (stage.session_run_id IS NOT NULL + AND existing.session_run_id = stage.session_run_id))) + ) + AND NOT EXISTS ( + SELECT 1 FROM environment_package_artifact_backup + WHERE backup_id = stage.actual_backup_id + UNION ALL + SELECT 1 FROM environment_package_artifact_backup_staging + WHERE actual_backup_id = stage.actual_backup_id + ) + ON CONFLICT DO NOTHING`, + ) + .bind(stage.id, input.actualBackupId), + database + .prepare( + `UPDATE native_resume_ref + SET committed_session_run_id = observed_session_run_id, committed_value = value + WHERE EXISTS ( + SELECT 1 FROM sandbox_backup_staging AS stage + JOIN sandbox_backup AS finalized + ON finalized.staging_id = stage.id AND finalized.id = stage.actual_backup_id + WHERE stage.id = ? AND stage.session_run_id = observed_session_run_id + AND stage.workspace_session_id = native_resume_ref.session_id + AND ${currentAuthority("stage")} + )`, + ) + .bind(stage.id), + database + .prepare( + `UPDATE sandbox + SET last_backup_id = ?, updated_at = ${DATABASE_NOW_MS_SQL} + WHERE id = ? AND incarnation = ? + AND status = 'backing_up' AND status_operation_id IS ? + AND EXISTS ( + SELECT 1 FROM sandbox_backup_staging AS stage + JOIN sandbox_backup AS finalized + ON finalized.staging_id = stage.id AND finalized.id = stage.actual_backup_id + WHERE stage.id = ? AND stage.updates_subject_backup = 1 + AND ${currentAuthority("stage")} + )`, + ) + .bind( + input.actualBackupId, + stage.sandboxId, + stage.sandboxIncarnation, + stage.operationId, + stage.id, + ), + database + .prepare( + `DELETE FROM sandbox_backup_staging AS stage + WHERE id = ? AND actual_backup_id = ? + AND EXISTS ( + SELECT 1 FROM sandbox_backup AS finalized + WHERE finalized.staging_id = stage.id AND finalized.id = stage.actual_backup_id + AND finalized.status = 'ready' + ) + AND (updates_subject_backup = 0 OR EXISTS ( + SELECT 1 FROM sandbox AS subject + WHERE subject.id = stage.sandbox_id + AND subject.incarnation = stage.sandbox_incarnation + AND subject.last_backup_id = stage.actual_backup_id + )) + AND NOT EXISTS ( + SELECT 1 FROM native_resume_ref AS native_ref + WHERE native_ref.session_id = stage.workspace_session_id + AND native_ref.observed_session_run_id = stage.session_run_id + AND native_ref.committed_session_run_id IS NOT stage.session_run_id + )`, + ) + .bind(stage.id, input.actualBackupId), + ]); + const finalized = + (await getSandboxBackupRecordByStagingId(database, stage.id)) ?? + (await getFinalizedSandboxBackupByScope(database, scope)); + return finalized === null + ? null + : { + backup: finalized, + candidateAccepted: finalized.id === input.actualBackupId, + complete: (await getSandboxBackupStage(database, stage.id)) === null, + }; +} + +export async function listReadySandboxBackupsForSessionRun( + database: D1Database, + input: { readonly sandboxId: string; readonly sessionRunId: string }, +): Promise> { + const sandboxId = parsePlatformId(input.sandboxId, "sandbox id"); + const sessionRunId = parsePlatformId(input.sessionRunId, "session run id"); + return getAppDatabase(database) + .select({ dir: sandboxBackupsTable.dir, id: sandboxBackupsTable.id }) + .from(sandboxBackupsTable) + .where( + and( + eq(sandboxBackupsTable.sandboxId, sandboxId), + eq(sandboxBackupsTable.sessionRunId, sessionRunId), + eq(sandboxBackupsTable.status, "ready"), + ), + ) + .all(); +} + +export interface ReadySandboxBackupForPruning { + readonly createdAt: number; + readonly dir: string; + readonly id: SandboxBackupId; + readonly keep: boolean; + readonly protected: boolean; + readonly sandboxId: SandboxId; +} - if (getD1ChangeCount(updated) === 0) { - throw new Error("Runtime subject changed before checkpoint backup was recorded."); +export async function listReadySandboxBackupsForPruning( + database: D1Database, + sandboxIdInput: string, +): Promise { + const sandboxId = parsePlatformId(sandboxIdInput, "sandbox id"); + const rows = await database + .prepare( + `SELECT backup.created_at, backup.dir, backup.id, backup.keep, backup.sandbox_id, + CASE WHEN EXISTS (SELECT 1 FROM sandbox WHERE last_backup_id = backup.id) + OR EXISTS ( + SELECT 1 FROM sandbox + WHERE status = 'restoring' AND last_restore_backup_id = backup.id + ) + OR EXISTS ( + SELECT 1 FROM session + WHERE id = backup.workspace_session_id + AND runtime_provisioning_operation_id IS NOT NULL + ) + OR EXISTS ( + SELECT 1 FROM sandbox_backup_staging AS stage + WHERE stage.id = backup.staging_id OR stage.actual_backup_id = backup.id + ) + THEN 1 ELSE 0 END AS protected + FROM sandbox_backup AS backup + WHERE backup.sandbox_id = ? AND backup.status = 'ready' + ORDER BY backup.dir ASC, backup.created_at DESC, backup.id DESC`, + ) + .bind(sandboxId) + .all<{ + created_at: number; + dir: string; + id: SandboxBackupId; + keep: number; + protected: number; + sandbox_id: SandboxId; + }>(); + return rows.results.map((row) => ({ + createdAt: row.created_at, + dir: row.dir, + id: row.id, + keep: row.keep === 1, + protected: row.protected === 1, + sandboxId: row.sandbox_id, + })); +} + +export async function markSandboxBackupsPruned( + database: D1Database, + backupIds: readonly SandboxBackupId[], +): Promise { + const pruned: SandboxBackupId[] = []; + for (const id of backupIds) { + const row = await database + .prepare( + `UPDATE sandbox_backup AS backup + SET status = 'pruned', updated_at = ${DATABASE_NOW_MS_SQL} + WHERE id = ? AND status = 'ready' AND keep = 0 + AND NOT EXISTS (SELECT 1 FROM sandbox WHERE last_backup_id = backup.id) + AND NOT EXISTS ( + SELECT 1 FROM sandbox + WHERE status = 'restoring' AND last_restore_backup_id = backup.id + ) + AND NOT EXISTS ( + SELECT 1 FROM session + WHERE id = backup.workspace_session_id + AND runtime_provisioning_operation_id IS NOT NULL + ) + AND NOT EXISTS ( + SELECT 1 FROM sandbox_backup_staging AS stage + WHERE stage.id = backup.staging_id OR stage.actual_backup_id = backup.id + ) + RETURNING id`, + ) + .bind(id) + .first<{ id: SandboxBackupId }>(); + if (row !== null) { + pruned.push(row.id); + } } + return pruned; +} + +export interface SandboxSessionBackupCandidate { + readonly cwd: string; + readonly lastMessageAt: number | null; + readonly sessionId: SessionId; + readonly sessionStatus: SessionStatus; } export async function listSandboxSessionBackupCandidates( database: D1Database, - sandboxId: string, + sandboxIdInput: string, + sandboxIncarnation: number, ): Promise { - const parsedSandboxId = parsePlatformId(sandboxId, "sandbox id"); - const results = await getAppDatabase(database) + assertPositiveSafeInteger(sandboxIncarnation, "Sandbox backup candidate incarnation"); + const sandboxId = parsePlatformId(sandboxIdInput, "sandbox id"); + return getAppDatabase(database) .select({ cwd: sandboxSessionsTable.cwd, - last_message_at: sessionsTable.lastMessageAt, - session_id: sandboxSessionsTable.sessionId, - session_status: sessionsTable.status, + lastMessageAt: sessionsTable.lastMessageAt, + sessionId: sandboxSessionsTable.sessionId, + sessionStatus: sessionsTable.status, }) .from(sandboxSessionsTable) .innerJoin(sessionsTable, eq(sessionsTable.id, sandboxSessionsTable.sessionId)) .where( and( - eq(sandboxSessionsTable.sandboxId, parsedSandboxId), + eq(sandboxSessionsTable.sandboxId, sandboxId), + eq(sandboxSessionsTable.sandboxIncarnation, sandboxIncarnation), inArray(sandboxSessionsTable.status, ["active", "closed"]), + isNull(sandboxSessionsTable.cleanupOperationId), + isNull(sessionsTable.archivedAt), + isNull(sessionsTable.cleanupOperationKind), ), ) .all(); - - return results.map((row) => ({ - cwd: row.cwd, - lastMessageAt: row.last_message_at, - sessionId: row.session_id, - sessionStatus: row.session_status, - })); } -export async function listSandboxBackupIdsByDir( +export async function revokeSandboxBackupsForSessionDelete( database: D1Database, - dir: string, + input: { + readonly cwd: string; + readonly operationId: RuntimeOperationId; + readonly sandboxId: SandboxId; + readonly sessionId: SessionId; + }, ): Promise { - const results = await getAppDatabase(database) - .select({ id: sandboxBackupsTable.id }) - .from(sandboxBackupsTable) - .where(eq(sandboxBackupsTable.dir, dir)) - .all(); - - return results.map((backup) => backup.id); -} - -export async function deleteSandboxBackupRecordsForDir( - database: D1Database, - dir: string, -): Promise { - await getAppDatabase(database) - .delete(sandboxBackupsTable) - .where(eq(sandboxBackupsTable.dir, dir)) - .run(); + const owned = (alias: string) => `EXISTS ( + SELECT 1 FROM session AS deleting_session + JOIN sandbox_session AS workspace ON workspace.session_id = deleting_session.id + WHERE deleting_session.id = ? + AND deleting_session.cleanup_operation_kind = 'delete' + AND deleting_session.status_operation_id = ? + AND workspace.sandbox_id = ? AND workspace.cwd = ? + AND ${alias}.sandbox_id = workspace.sandbox_id + AND ${alias}.dir = workspace.cwd + )`; + const results = await database.batch([ + database + .prepare( + `DELETE FROM sandbox_backup_staging AS stage + WHERE stage.workspace_session_id = ? AND stage.sandbox_id = ? AND stage.dir = ? + AND ${owned("stage")} + RETURNING actual_backup_id`, + ) + .bind( + input.sessionId, + input.sandboxId, + input.cwd, + input.sessionId, + input.operationId, + input.sandboxId, + input.cwd, + ), + database + .prepare( + `UPDATE sandbox_backup AS backup + SET status = 'pruned', updated_at = ${DATABASE_NOW_MS_SQL} + WHERE status = 'ready' + AND workspace_session_id = ? AND sandbox_id = ? AND dir = ? + AND ${owned("backup")} + RETURNING id`, + ) + .bind( + input.sessionId, + input.sandboxId, + input.cwd, + input.sessionId, + input.operationId, + input.sandboxId, + input.cwd, + ), + ]); + const ids = new Set(); + const [stagedResult, finalizedResult] = results; + if (stagedResult === undefined || finalizedResult === undefined) { + throw new Error("Session backup cleanup lost its D1 batch results."); + } + for (const row of stagedResult.results as Array<{ + actual_backup_id: SandboxBackupId | null; + }>) { + if (row.actual_backup_id !== null) { + ids.add(row.actual_backup_id); + } + } + for (const row of finalizedResult.results as Array<{ id: SandboxBackupId }>) { + ids.add(row.id); + } + return [...ids]; } diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-backup.service.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-backup.service.ts index 7c3920e9..58f7d8c5 100644 --- a/apps/api/src/modules/runtime/infrastructure/sandbox-backup.service.ts +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-backup.service.ts @@ -1,220 +1,229 @@ +import type { + DriverInstanceId, + RuntimeOperationId, + SandboxBackupId, + SandboxId, + SessionId, + SessionRunId, +} from "@mosoo/id"; +import { parsePlatformId } from "@mosoo/id"; + import { logWarn } from "../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; import { shouldBackupSandboxSession } from "../../sessions/domain/session-lifecycle"; import type { RuntimeCheckpointRule } from "../domain/runtime-kind-policy"; import { RuntimeSubjectCheckpointFailedError } from "./runtime-subject-lifecycle/runtime-subject-errors"; -import { createRuntimeSandboxBackup, deleteSandboxBackupObjects } from "./sandbox-backup-platform"; +import type { RuntimeSubjectOperationLease } from "./runtime-subject-lifecycle/runtime-subject-store"; +import { + createRuntimeSandboxBackup, + deleteAuthorizedSandboxBackupObjects, + isRuntimeSandboxBackupObjectReady, +} from "./sandbox-backup-platform"; import { selectSandboxBackupPruneIds } from "./sandbox-backup-pruning"; -import type { CreatedSandboxBackupWrite } from "./sandbox-backup-store"; +import type { + SandboxBackupAdmission, + SandboxBackupDeletionAuthority, + SandboxBackupStage, + SandboxBackupTarget, +} from "./sandbox-backup-store"; import { - deleteSandboxBackupRecordsForDir, - listReadySandboxBackupsForSessionRun, + authorizeSandboxBackupDeletion, + claimSandboxBackupStageActual, + clearMissingSandboxBackupStageActual, + deferSandboxBackupStageRepair, + finalizeSandboxBackupStage, + getSandboxBackupStage, + isSandboxBackupStageCurrent, listReadySandboxBackupsForPruning, - listSandboxBackupIdsByDir, + listReadySandboxBackupsForSessionRun, + listSandboxBackupStages, listSandboxSessionBackupCandidates, markSandboxBackupsPruned, - recordCreatedSandboxBackups, + revokeSandboxBackupsForSessionDelete, + revokeSandboxBackupStage, + stageSandboxBackupWrites, } from "./sandbox-backup-store"; const BACKUP_TTL_SECONDS = 10 * 365 * 24 * 60 * 60; -export interface SandboxSessionBackupTarget { - cwd: string; - sessionId: string; +async function deleteSandboxBackup( + bindings: ApiBindings, + backupId: SandboxBackupId, + authority: SandboxBackupDeletionAuthority, +): Promise { + if (!(await authorizeSandboxBackupDeletion(bindings.DB, { authority, backupId }))) { + return false; + } + await deleteAuthorizedSandboxBackupObjects(bindings, [backupId]); + return true; +} + +export interface TerminalSandboxBackupAuthority { + readonly driverGeneration: number; + readonly driverInstanceId: DriverInstanceId; + readonly incarnation: number; + readonly sessionId: SessionId; + readonly sessionRunId: SessionRunId; } -interface SandboxCheckpointBackupTarget { - readonly dir: string; - readonly sessionId: string | null; - readonly updateSandboxLastBackup: boolean; +export interface SandboxSessionBackupTarget { + readonly cwd: string; + readonly sessionId: string; } async function pruneSandboxBackups(bindings: ApiBindings, sandboxId: string): Promise { const backups = await listReadySandboxBackupsForPruning(bindings.DB, sandboxId); - const pruneIds = selectSandboxBackupPruneIds(backups); - - await deleteSandboxBackupObjects(bindings, pruneIds); - await markSandboxBackupsPruned(bindings.DB, pruneIds); + const candidates = selectSandboxBackupPruneIds(backups); + const pruned = await markSandboxBackupsPruned(bindings.DB, candidates); + for (const backupId of pruned) { + await deleteSandboxBackup(bindings, backupId, { kind: "pruned" }); + } } async function listSandboxSessionBackupTargets( database: D1Database, sandboxId: string, + sandboxIncarnation: number, ): Promise { - const candidates = await listSandboxSessionBackupCandidates(database, sandboxId); - - return candidates + return (await listSandboxSessionBackupCandidates(database, sandboxId, sandboxIncarnation)) .filter((candidate) => shouldBackupSandboxSession({ lastMessageAt: candidate.lastMessageAt, sessionStatus: candidate.sessionStatus, }), ) - .map((candidate) => ({ - cwd: candidate.cwd, - sessionId: candidate.sessionId, - })); + .map((candidate) => ({ cwd: candidate.cwd, sessionId: candidate.sessionId })); } -async function listSandboxCheckpointBackupTargets( +async function listCheckpointTargets( database: D1Database, input: { readonly rules: readonly RuntimeCheckpointRule[]; readonly sandboxId: string; + readonly sandboxIncarnation: number; }, -): Promise { - const targets: SandboxCheckpointBackupTarget[] = []; +): Promise { + const targets: SandboxBackupTarget[] = []; let sessionTargets: SandboxSessionBackupTarget[] | null = null; - for (const rule of input.rules) { switch (rule.type) { case "subject_memory": { targets.push({ dir: rule.path, - sessionId: null, updateSandboxLastBackup: rule.updateSubjectCheckpoint, + workspaceSessionId: null, }); break; } case "session_workspaces": { - sessionTargets ??= await listSandboxSessionBackupTargets(database, input.sandboxId); + sessionTargets ??= await listSandboxSessionBackupTargets( + database, + input.sandboxId, + input.sandboxIncarnation, + ); targets.push( ...sessionTargets.map((target) => ({ dir: target.cwd, - sessionId: rule.sanitizeTransientState ? target.sessionId : null, updateSandboxLastBackup: false, + workspaceSessionId: target.sessionId, })), ); break; } } } - return targets; } -async function createSandboxBackupsForTargets( +async function resolveStageActual( bindings: ApiBindings, - input: { - readonly sandboxId: string; - readonly targets: readonly SandboxCheckpointBackupTarget[]; - }, -): Promise { - const results = await Promise.allSettled( - input.targets.map(async (target) => ({ - backup: await createRuntimeSandboxBackup(bindings, { - dir: target.dir, - sandboxId: input.sandboxId, - sessionId: target.sessionId, - ttlSeconds: BACKUP_TTL_SECONDS, - }).catch((error: unknown) => { - throw new RuntimeSubjectCheckpointFailedError({ - cause: error, - dir: target.dir, - runtimeSubjectId: input.sandboxId, - }); - }), - updateSandboxLastBackup: target.updateSandboxLastBackup, - })), - ); - const createdBackups = results.flatMap((result) => - result.status === "fulfilled" ? [result.value] : [], - ); - const failedBackup = results.find((result) => result.status === "rejected"); - - if (failedBackup?.status === "rejected") { - await deleteSandboxBackupObjects( - bindings, - createdBackups.map((entry) => entry.backup.id), - ); - throw failedBackup.reason; + initial: SandboxBackupStage, +): Promise<{ + readonly actualId: NonNullable; + readonly stage: SandboxBackupStage; +}> { + let stage = initial; + if (stage.actualBackupId !== null) { + if ( + await isRuntimeSandboxBackupObjectReady(bindings, { + backupId: stage.actualBackupId, + dir: stage.dir, + stagingId: stage.id, + }) + ) { + return { actualId: stage.actualBackupId, stage }; + } + await clearMissingSandboxBackupStageActual(bindings.DB, { + actualBackupId: stage.actualBackupId, + stagingId: stage.id, + }); + const cleared = await getSandboxBackupStage(bindings.DB, stage.id); + if (cleared === null) { + throw new Error("Sandbox backup stage was revoked while its object was verified."); + } + stage = cleared; } - return createdBackups; + const candidate = await createRuntimeSandboxBackup(bindings, { + dir: stage.dir, + incarnation: stage.sandboxIncarnation, + sandboxId: stage.sandboxId, + sessionId: stage.workspaceSessionId, + stagingId: stage.id, + ttlSeconds: stage.ttlSeconds, + }); + if (candidate.dir !== stage.dir) { + await deleteSandboxBackup(bindings, candidate.id, { + kind: "runtime_invalid", + stagingId: stage.id, + }); + throw new Error("Sandbox backup platform result changed its staged directory."); + } + const claimed = await claimSandboxBackupStageActual(bindings.DB, { + actualBackupId: candidate.id, + dir: stage.dir, + sandboxIncarnation: stage.sandboxIncarnation, + stagingId: stage.id, + }); + if (claimed === null || claimed.actualBackupId === null) { + await deleteSandboxBackup(bindings, candidate.id, { + kind: "runtime_candidate", + stagingId: stage.id, + }); + throw new Error("Sandbox backup creation lost its durable stage."); + } + if (claimed.actualBackupId !== candidate.id) { + await deleteSandboxBackup(bindings, candidate.id, { + kind: "runtime_candidate", + stagingId: stage.id, + }); + } + return { actualId: claimed.actualBackupId, stage }; } -async function recordCreatedCheckpointBackups( +async function createAndFinalizeStage( bindings: ApiBindings, - input: { - readonly backups: readonly CreatedSandboxBackupWrite[]; - readonly checkpointSessionId?: string; - readonly operationId?: string | null; - readonly sandboxId: string; - readonly sessionRunId?: string; - }, + stage: SandboxBackupStage, ): Promise { - try { - await recordCreatedSandboxBackups(bindings.DB, { - backups: input.backups, - ...(input.checkpointSessionId === undefined - ? {} - : { checkpointSessionId: input.checkpointSessionId }), - ...(input.operationId === undefined ? {} : { operationId: input.operationId }), - sandboxId: input.sandboxId, - ...(input.sessionRunId === undefined ? {} : { sessionRunId: input.sessionRunId }), - ttlSeconds: BACKUP_TTL_SECONDS, + const { actualId } = await resolveStageActual(bindings, stage); + const finalized = await finalizeSandboxBackupStage(bindings.DB, { + actualBackupId: actualId, + stagingId: stage.id, + }); + if (finalized === null) { + await deleteSandboxBackup(bindings, actualId, { + kind: "runtime_candidate", + stagingId: stage.id, }); - } catch (error) { - if (input.sessionRunId !== undefined) { - let recordedBackups: readonly CreatedSandboxBackupWrite["backup"][]; - - try { - recordedBackups = await listReadySandboxBackupsForSessionRun(bindings.DB, { - sandboxId: input.sandboxId, - sessionRunId: input.sessionRunId, - }); - } catch { - // An ambiguous database failure must not delete objects that may already - // back a committed marker. Unreferenced objects are safer than a ready - // row whose Cloudflare backup was deleted. - throw new RuntimeSubjectCheckpointFailedError({ - cause: error, - runtimeSubjectId: input.sandboxId, - }); - } - - const recordedIds = new Set(recordedBackups.map((backup) => backup.id)); - const recordedDirs = new Set(recordedBackups.map((backup) => backup.dir)); - const checkpointRecorded = input.backups.every((entry) => recordedDirs.has(entry.backup.dir)); - const orphanedBackupIds = input.backups - .map((entry) => entry.backup.id) - .filter((backupId) => !recordedIds.has(backupId)); - - try { - await deleteSandboxBackupObjects(bindings, orphanedBackupIds); - } catch (cleanupError) { - if (!checkpointRecorded) { - throw new RuntimeSubjectCheckpointFailedError({ - cause: cleanupError, - runtimeSubjectId: input.sandboxId, - }); - } - - logWarn("runtime.sandbox_checkpoint.orphan_cleanup_failed", { - backupCount: orphanedBackupIds.length, - error: cleanupError instanceof Error ? cleanupError.message : String(cleanupError), - sandboxId: input.sandboxId, - sessionRunId: input.sessionRunId, - }); - } - - if (checkpointRecorded) { - return; - } - - throw new RuntimeSubjectCheckpointFailedError({ - cause: error, - runtimeSubjectId: input.sandboxId, - }); - } - - await deleteSandboxBackupObjects( - bindings, - input.backups.map((entry) => entry.backup.id), - ); - throw new RuntimeSubjectCheckpointFailedError({ - cause: error, - runtimeSubjectId: input.sandboxId, + throw new Error("Sandbox backup finalization lost its durable stage."); + } + if (!finalized.complete) { + throw new Error("Sandbox backup finalization did not commit every required relation."); + } + if (!finalized.candidateAccepted) { + await deleteSandboxBackup(bindings, actualId, { + kind: "runtime_candidate", + stagingId: stage.id, }); } } @@ -222,73 +231,80 @@ async function recordCreatedCheckpointBackups( async function createSandboxCheckpointBackups( bindings: ApiBindings, input: { - readonly operationId?: string | null; + readonly admission: SandboxBackupAdmission; readonly requiredSessionId?: string; readonly rules: readonly RuntimeCheckpointRule[]; readonly sandboxId: string; - readonly sessionRunId?: string; }, ): Promise { - let targets = await listSandboxCheckpointBackupTargets(bindings.DB, input); - + let targets = await listCheckpointTargets(bindings.DB, { + ...input, + sandboxIncarnation: + input.admission.kind === "operation" + ? input.admission.lease.incarnation + : input.admission.incarnation, + }); if ( input.requiredSessionId !== undefined && - !targets.some((target) => target.sessionId === input.requiredSessionId) + !targets.some((target) => target.workspaceSessionId === input.requiredSessionId) ) { throw new RuntimeSubjectCheckpointFailedError({ - cause: new Error( - `Session ${input.requiredSessionId} has no eligible workspace checkpoint target.`, - ), + cause: new Error(`Session ${input.requiredSessionId} has no eligible checkpoint target.`), runtimeSubjectId: input.sandboxId, }); } - if (targets.length === 0) { return; } - - if (input.sessionRunId !== undefined) { + if (input.admission.kind === "terminal") { const readyDirs = new Set( ( await listReadySandboxBackupsForSessionRun(bindings.DB, { sandboxId: input.sandboxId, - sessionRunId: input.sessionRunId, + sessionRunId: input.admission.sessionRunId, }) ).map((backup) => backup.dir), ); targets = targets.filter((target) => !readyDirs.has(target.dir)); - if (targets.length === 0) { return; } } - const backups = await createSandboxBackupsForTargets(bindings, { + const writes = await stageSandboxBackupWrites(bindings.DB, { + admission: input.admission, sandboxId: input.sandboxId, targets, + ttlSeconds: BACKUP_TTL_SECONDS, }); - - await recordCreatedCheckpointBackups(bindings, { - backups, - ...(input.requiredSessionId === undefined - ? {} - : { checkpointSessionId: input.requiredSessionId }), - ...(input.operationId === undefined ? {} : { operationId: input.operationId }), - sandboxId: input.sandboxId, - ...(input.sessionRunId === undefined ? {} : { sessionRunId: input.sessionRunId }), - }); + for (const write of writes) { + if (write.kind === "finalized") { + if (write.backup.status !== "ready") { + throw new Error("A finalized sandbox checkpoint is no longer restorable."); + } + continue; + } + try { + await createAndFinalizeStage(bindings, write.stage); + } catch (cause) { + throw new RuntimeSubjectCheckpointFailedError({ + cause, + dir: write.stage.dir, + runtimeSubjectId: input.sandboxId, + }); + } + } try { await pruneSandboxBackups(bindings, input.sandboxId); } catch (error) { - if (input.sessionRunId === undefined) { + if (input.admission.kind === "operation") { throw error; } - logWarn("runtime.sandbox_checkpoint.prune_failed", { error: error instanceof Error ? error.message : String(error), sandboxId: input.sandboxId, - sessionRunId: input.sessionRunId, + sessionRunId: input.admission.sessionRunId, }); } } @@ -296,24 +312,78 @@ async function createSandboxCheckpointBackups( export async function createSandboxCheckpoints( bindings: ApiBindings, input: { - operationId?: string | null; - requiredSessionId?: string; - rules: readonly RuntimeCheckpointRule[]; - sandboxId: string; - sessionRunId?: string; + readonly operationLease?: RuntimeSubjectOperationLease; + readonly requiredSessionId?: string; + readonly rules: readonly RuntimeCheckpointRule[]; + readonly sandboxId: string; + readonly terminalAuthority?: TerminalSandboxBackupAuthority; }, ): Promise { - await createSandboxCheckpointBackups(bindings, input); + if ((input.operationLease === undefined) === (input.terminalAuthority === undefined)) { + throw new Error("Sandbox checkpoints require exactly one durable admission authority."); + } + const admission: SandboxBackupAdmission = + input.operationLease === undefined + ? { ...input.terminalAuthority!, kind: "terminal" } + : { kind: "operation", lease: input.operationLease }; + await createSandboxCheckpointBackups(bindings, { ...input, admission }); +} + +async function repairSandboxBackupStage( + bindings: ApiBindings, + stage: SandboxBackupStage, +): Promise { + if (!(await isSandboxBackupStageCurrent(bindings.DB, stage.id))) { + const revoked = await revokeSandboxBackupStage(bindings.DB, { + onlyIfStale: true, + stagingId: stage.id, + }); + if (revoked?.actualBackupId !== null && revoked?.actualBackupId !== undefined) { + await deleteSandboxBackup(bindings, revoked.actualBackupId, { + kind: "runtime_candidate", + stagingId: stage.id, + }); + } + return; + } + await createAndFinalizeStage(bindings, stage); } -export async function deleteSandboxBackupsForDir( +export async function repairStagedSandboxBackups( + bindings: ApiBindings, + limit: number, +): Promise { + const stages = await listSandboxBackupStages(bindings.DB, limit); + for (const stage of stages) { + try { + await repairSandboxBackupStage(bindings, stage); + } catch (error) { + await deferSandboxBackupStageRepair(bindings.DB, stage); + logWarn("runtime.sandbox_checkpoint.repair_failed", { + error: error instanceof Error ? error.message : String(error), + stagingId: stage.id, + }); + } + } + return stages.length; +} + +export async function deleteSandboxBackupsForSession( bindings: ApiBindings, input: { - dir: string; + readonly cwd: string; + readonly operationId: string; + readonly sandboxId: string; + readonly sessionId: string; }, ): Promise { - const backupIds = await listSandboxBackupIdsByDir(bindings.DB, input.dir); - - await deleteSandboxBackupObjects(bindings, backupIds); - await deleteSandboxBackupRecordsForDir(bindings.DB, input.dir); + const backupIds = await revokeSandboxBackupsForSessionDelete(bindings.DB, { + cwd: input.cwd, + operationId: parsePlatformId(input.operationId, "session delete operation"), + sandboxId: parsePlatformId(input.sandboxId, "sandbox id"), + sessionId: parsePlatformId(input.sessionId, "session id"), + }); + for (const backupId of backupIds) { + await deleteSandboxBackup(bindings, backupId, { kind: "unattributed" }); + } } diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-file-bytes.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-file-bytes.ts index 91ea5cc0..eb6b5c24 100644 --- a/apps/api/src/modules/runtime/infrastructure/sandbox-file-bytes.ts +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-file-bytes.ts @@ -1,49 +1,43 @@ +import { fromBase64, toBase64 } from "../../../shared/bytes"; +import { quoteShellArg } from "../../../shared/shell"; import type { ExecutionSessionHandle } from "./sandbox-handles"; -function decodeBase64(value: string): Uint8Array { - if (value.length === 0) { - return new Uint8Array(); - } - - const binary = atob(value); - const bytes = new Uint8Array(binary.length); - - for (let index = 0; index < binary.length; index += 1) { - bytes[index] = binary.codePointAt(index) ?? 0; - } - - return bytes; -} - export async function readSandboxFileBytes( handle: ExecutionSessionHandle, path: string, + maxBytes?: number, ): Promise { + if (maxBytes !== undefined) { + if (!Number.isSafeInteger(maxBytes) || maxBytes < 0) { + throw new Error("Sandbox file byte limit is invalid."); + } + const command = [ + "runtime_file_read=$(mktemp)", + "trap 'rm -f \"$runtime_file_read\"' EXIT", + `head -c ${maxBytes + 1} -- ${quoteShellArg(path)} > "$runtime_file_read"`, + 'base64 -w 0 "$runtime_file_read"', + ].join(" && "); + const result = await handle.exec(`sh -lc ${quoteShellArg(command)}`); + if (!result.success || result.exitCode !== 0) { + throw new Error( + result.stderr.trim() || result.stdout.trim() || `Failed to read sandbox file ${path}.`, + ); + } + return fromBase64(result.stdout); + } const file = await handle.readFile(path, { encoding: "base64" }); if (file.encoding === "base64") { - return decodeBase64(file.content); + return fromBase64(file.content); } return new TextEncoder().encode(file.content); } -function encodeBase64(bytes: Uint8Array): string { - let binary = ""; - // Chunked conversion keeps String.fromCharCode off argument-count limits. - const chunkSize = 0x8000; - - for (let index = 0; index < bytes.length; index += chunkSize) { - binary += String.fromCharCode(...bytes.subarray(index, index + chunkSize)); - } - - return btoa(binary); -} - export async function writeSandboxFileBytes( handle: ExecutionSessionHandle, path: string, bytes: Uint8Array, ): Promise { - await handle.writeFile(path, encodeBase64(bytes), { encoding: "base64" }); + await handle.writeFile(path, toBase64(bytes), { encoding: "base64" }); } diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-handles.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-handles.ts index 56ee685b..dc256eef 100644 --- a/apps/api/src/modules/runtime/infrastructure/sandbox-handles.ts +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-handles.ts @@ -71,7 +71,9 @@ export interface SandboxHandle extends ExecutionSessionHandle { configureNetworkConstraints(constraints: SandboxNetworkConstraints): Promise; createBackup(options: { dir: string; + excludes?: string[]; localBucket?: boolean; + name?: string; ttl?: number; }): Promise<{ dir: string; id: string }>; createSession(options?: { @@ -100,6 +102,33 @@ export interface SandboxHandle extends ExecutionSessionHandle { wsConnect(request: Request, port: number): Promise; } +export interface RuntimeSubjectIncarnationHandle { + activateRuntimeSubjectIncarnation( + incarnation: number, + networkConstraintsHash: string, + ): Promise; + createRuntimeSubjectBackup( + incarnation: number, + options: { + dir: string; + excludes?: string[]; + forbiddenPaths?: string[]; + localBucket?: boolean; + name: string; + ttl?: number; + }, + ): Promise<{ dir: string; id: string }>; + destroyRuntimeSubjectIncarnation(incarnation: number): Promise<{ kind: "destroyed" | "stale" }>; + inspectRuntimeSubjectIncarnation( + incarnation: number, + networkConstraintsHash: string, + ): Promise<{ kind: "healthy" | "missing" | "retired" | "stale" | "unknown" }>; + markRuntimeSubjectIncarnationReady( + incarnation: number, + networkConstraintsHash: string, + ): Promise; +} + const SANDBOX_HANDLE_METHODS = [ "configureNetworkConstraints", "createBackup", @@ -134,3 +163,23 @@ export function toSandboxHandle(value: unknown): SandboxHandle { return value as SandboxHandle; } + +export function toRuntimeSubjectIncarnationHandle(value: unknown): RuntimeSubjectIncarnationHandle { + if (typeof value !== "object" || value === null) { + throw new Error("Cloudflare Sandbox incarnation handle is not an object."); + } + + for (const method of [ + "activateRuntimeSubjectIncarnation", + "createRuntimeSubjectBackup", + "destroyRuntimeSubjectIncarnation", + "inspectRuntimeSubjectIncarnation", + "markRuntimeSubjectIncarnationReady", + ] as const) { + if (typeof Reflect.get(value, method) !== "function") { + throw new TypeError(`Cloudflare Sandbox incarnation handle is missing ${method}.`); + } + } + + return value as RuntimeSubjectIncarnationHandle; +} diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-session.service.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-session.service.ts index a65aa899..c5dd0552 100644 --- a/apps/api/src/modules/runtime/infrastructure/sandbox-session.service.ts +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-session.service.ts @@ -2,4 +2,5 @@ export { closeIdleCattleConversationSession, closeSandboxConversationSession, ensureSandboxConversationSession, + repairPendingSandboxConversationSessionCleanups, } from "./sandbox-session/sandbox-conversation-session.service"; diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session-delete.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session-delete.ts index c1b1aae1..0be8a4d7 100644 --- a/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session-delete.ts +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session-delete.ts @@ -5,28 +5,46 @@ import { withDisposedRpcResult, } from "../../../../platform/cloudflare/rpc-disposal"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; +import { withRuntimeProvisionTimeout } from "../runtime-provision-timeout"; import { getRuntimeSubjectKeepAliveHandle } from "../runtime-subject-lifecycle/runtime-subject-platform"; +function isMissingSandboxSession(error: unknown): boolean { + return ( + error instanceof Error && + (error.name === "SessionNotFoundError" || + error.message.includes("SessionNotFoundError") || + (error.message.includes("Session '") && error.message.includes("' not found"))) + ); +} + export async function deleteActiveSandboxConversationSession( bindings: ApiBindings, input: { readonly sandboxSessionId: SandboxSessionId; readonly sandboxId: SandboxId; + readonly sandboxIncarnation: number; }, ): Promise { - await withDisposedRpcResource( - await getRuntimeSubjectKeepAliveHandle(bindings, input.sandboxId), - async (sandbox) => { - const deleted = await withDisposedRpcResult( - sandbox.deleteSession(input.sandboxSessionId), - (deletedSession) => ({ - success: deletedSession.success, - }), - ); + try { + await withDisposedRpcResource( + await getRuntimeSubjectKeepAliveHandle(bindings, input.sandboxId, input.sandboxIncarnation), + async (sandbox) => { + const deleted = await withDisposedRpcResult( + withRuntimeProvisionTimeout( + sandbox.deleteSession(input.sandboxSessionId), + `Sandbox session deletion for ${input.sandboxSessionId}`, + ), + (deletedSession) => ({ success: deletedSession.success }), + ); - if (!deleted.success) { - throw new Error(`Sandbox session ${input.sandboxSessionId} could not be deleted.`); - } - }, - ); + if (!deleted.success) { + throw new Error(`Sandbox session ${input.sandboxSessionId} could not be deleted.`); + } + }, + ); + } catch (error) { + if (!isMissingSandboxSession(error)) { + throw error; + } + } } diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session-platform.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session-platform.ts index 0c030ae9..ff2477d6 100644 --- a/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session-platform.ts +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session-platform.ts @@ -2,6 +2,7 @@ import { discardPromiseResult } from "@mosoo/effects"; import type { SandboxBackupId, SandboxSessionId } from "@mosoo/id"; import { withDisposedRpcResult } from "../../../../platform/cloudflare/rpc-disposal"; +import { quoteShellArg } from "../../../../shared/shell"; import { getRuntimeSessionOutputDirectory } from "../driver-instance/runtime-session-outputs"; import { withRuntimeProvisionTimeout } from "../runtime-provision-timeout"; import { decodeSandboxBackupIdForPlatform } from "../sandbox-backup-id"; @@ -12,10 +13,6 @@ interface SandboxConversationDirectoryBackup { readonly id: SandboxBackupId; } -function quoteShellArg(value: string): string { - return `'${value.replaceAll("'", `'"'"'`)}'`; -} - function isSessionAlreadyExistsError(error: unknown): boolean { return ( error instanceof Error && @@ -63,8 +60,14 @@ export async function prepareSandboxConversationDirectories(input: { readonly cwd: string; readonly sandbox: SandboxHandle; }): Promise { - await input.sandbox.mkdir(input.cwd, { recursive: true }); - await input.sandbox.mkdir(getRuntimeSessionOutputDirectory(input.cwd), { recursive: true }); + await withRuntimeProvisionTimeout( + input.sandbox.mkdir(input.cwd, { recursive: true }), + `Sandbox session cwd creation for ${input.cwd}`, + ); + await withRuntimeProvisionTimeout( + input.sandbox.mkdir(getRuntimeSessionOutputDirectory(input.cwd), { recursive: true }), + `Sandbox session output directory creation for ${input.cwd}`, + ); } export async function deleteSandboxConversationSessionBestEffort(input: { @@ -72,7 +75,10 @@ export async function deleteSandboxConversationSessionBestEffort(input: { readonly sandbox: SandboxHandle; }): Promise { try { - await input.sandbox.deleteSession(input.sandboxSessionId); + await withRuntimeProvisionTimeout( + input.sandbox.deleteSession(input.sandboxSessionId), + `Sandbox session deletion for ${input.sandboxSessionId}`, + ); } catch { // Best-effort cleanup for a partially configured session. } @@ -88,10 +94,13 @@ export async function openSandboxConversationSession(input: { try { return { created: true, - session: await input.sandbox.createSession({ - cwd: input.cwd, - id: input.sandboxSessionId, - }), + session: await withRuntimeProvisionTimeout( + input.sandbox.createSession({ + cwd: input.cwd, + id: input.sandboxSessionId, + }), + `Sandbox session creation for ${input.sandboxSessionId}`, + ), }; } catch (error) { if (!isSessionAlreadyExistsError(error)) { @@ -100,13 +109,19 @@ export async function openSandboxConversationSession(input: { return { created: false, - session: await input.sandbox.getSession(input.sandboxSessionId), + session: await withRuntimeProvisionTimeout( + input.sandbox.getSession(input.sandboxSessionId), + `Sandbox session lookup for ${input.sandboxSessionId}`, + ), }; } } return { created: false, - session: await input.sandbox.getSession(input.sandboxSessionId), + session: await withRuntimeProvisionTimeout( + input.sandbox.getSession(input.sandboxSessionId), + `Sandbox session lookup for ${input.sandboxSessionId}`, + ), }; } diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session.service.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session.service.ts index 1fef7bc5..45bcc9f8 100644 --- a/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session.service.ts +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session.service.ts @@ -1,6 +1,6 @@ import { getSessionOrganizationPath } from "@mosoo/agent-driver/paths"; import { createPlatformId } from "@mosoo/id"; -import type { SandboxId, SandboxSessionId, SessionId } from "@mosoo/id"; +import type { RuntimeOperationId, SandboxId, SandboxSessionId, SessionId } from "@mosoo/id"; import { RUNTIME_DIAGNOSTIC_EVENT } from "@mosoo/runtime-events"; import { disposeRpcResource } from "../../../../platform/cloudflare/rpc-disposal"; @@ -15,12 +15,18 @@ import { getRuntimeSubjectInactiveDeadline, runtimeCheckpointRulesInclude, } from "../../domain/runtime-kind-policy"; +import { + heartbeatRuntimeRunProvisioningLease, + recordRuntimeProvisioningConversationTarget, +} from "../runtime-subject-lifecycle/runtime-provisioning-lease-store"; import type { RuntimeConversationSessionRecord } from "../runtime-subject-lifecycle/runtime-subject-store"; import { + claimRuntimeConversationSessionCleanup, claimIdleSessionScopedConversationForClose, ensureRuntimeConversationSessionRecord, getRuntimeConversationSession, getRuntimeConversationSessionState, + listPendingRuntimeConversationSessionCleanups, recordRuntimeConversationSessionActive, recordRuntimeConversationSessionClosed, recordRuntimeConversationSessionError, @@ -52,6 +58,7 @@ function measureOptional( function resolveConversationContinuationPlan(input: { existingSession: RuntimeConversationSessionRecord | null; kind: EnsureSandboxConversationSessionInput["kind"]; + replaceClosedExecutionSession: boolean; }): { sandboxSessionId?: SandboxSessionId; requireCwdCheckpoint: boolean; @@ -93,12 +100,17 @@ function resolveConversationContinuationPlan(input: { }; } + if (input.existingSession.status === "cleanup_pending") { + throw new Error("Sandbox conversation cleanup is still pending; retry after maintenance."); + } + const shouldRestoreCwd = runtimeCheckpointRulesInclude( policy.checkpoint.restoreOnActivate, "session_workspaces", ); const shouldUseNewCloudflareSession = - input.existingSession.status === "closed" && policy.subject.scope === "session"; + input.existingSession.status === "closed" && + (policy.subject.scope === "session" || input.replaceClosedExecutionSession); return { ...(shouldUseNewCloudflareSession @@ -160,12 +172,20 @@ export async function ensureSandboxConversationSession( const continuation = resolveConversationContinuationPlan({ existingSession, kind: input.kind, + replaceClosedExecutionSession: input.replaceClosedExecutionSession ?? false, }); const cwd = existingSession?.cwd ?? getSessionOrganizationPath(input.sessionId); if (existingSession && existingSession.sandboxId !== input.sandboxId) { throw new Error("Sandbox session is already bound to a different sandbox."); } + if ( + existingSession && + existingSession.status !== "closed" && + existingSession.sandboxIncarnation !== input.sandboxIncarnation + ) { + throw new Error("Sandbox session belongs to a retired sandbox incarnation."); + } const frozenOrigin = existingSession ? parseSandboxConversationOrigin(existingSession.originJson) @@ -181,13 +201,44 @@ export async function ensureSandboxConversationSession( now, originJson: JSON.stringify(frozenOrigin), runtimeSubjectId: input.sandboxId, + sandboxIncarnation: input.sandboxIncarnation, sessionId: input.sessionId, }), )); const sandboxSessionId = continuation.sandboxSessionId ?? sessionRecord.sandboxSessionId; + let provisioningLease = input.provisioningLease; + + if (provisioningLease !== undefined) { + const recordedLease = await recordRuntimeProvisioningConversationTarget(bindings.DB, { + lease: provisioningLease, + sandboxIncarnation: input.sandboxIncarnation, + sandboxSessionId, + }); + if (recordedLease === null) { + throw new Error("Sandbox conversation provisioning lost lifecycle ownership."); + } + provisioningLease = recordedLease; + } + + const measureProvisioning = async (name: string, task: () => Promise): Promise => { + if ( + provisioningLease !== undefined && + !(await heartbeatRuntimeRunProvisioningLease(bindings.DB, provisioningLease)) + ) { + throw new Error("Sandbox conversation provisioning lost lifecycle ownership."); + } + const result = await measureOptional(input.timing, name, task); + if ( + provisioningLease !== undefined && + !(await heartbeatRuntimeRunProvisioningLease(bindings.DB, provisioningLease)) + ) { + throw new Error("Sandbox conversation provisioning lost lifecycle ownership."); + } + return result; + }; if (continuation.shouldRestoreCwd && existingSession) { - await measureOptional(input.timing, "conversation.restoreCwd", () => + await measureProvisioning("conversation.restoreCwd", () => restoreSandboxSessionCwdIfMissing({ cwd, latestReadyBackup: existingSession.latestReadyBackup, @@ -199,7 +250,7 @@ export async function ensureSandboxConversationSession( } if (continuation.shouldCreateCloudflareSession) { - await measureOptional(input.timing, "conversation.prepareDirectories", () => + await measureProvisioning("conversation.prepareDirectories", () => prepareSandboxConversationDirectories({ cwd, sandbox: input.sandbox, @@ -207,7 +258,7 @@ export async function ensureSandboxConversationSession( ); if (continuation.shouldRestoreSessionArtifacts) { - await measureOptional(input.timing, "conversation.restoreSessionArtifacts", () => + await measureProvisioning("conversation.restoreSessionArtifacts", () => restoreSessionArtifactsToWorkspace(bindings, { agentId: input.agentId, cwd, @@ -220,7 +271,7 @@ export async function ensureSandboxConversationSession( } if (input.mountSessionResources) { - await measureOptional(input.timing, "conversation.mountResources", () => + await measureProvisioning("conversation.mountResources", () => ensureSessionResourcesMounted({ bindings, sandbox: input.sandbox, @@ -230,7 +281,7 @@ export async function ensureSandboxConversationSession( } if (continuation.shouldDeleteErrorSession) { - await measureOptional(input.timing, "conversation.deleteErrorSession", () => + await measureProvisioning("conversation.deleteErrorSession", () => deleteSandboxConversationSessionBestEffort({ sandboxSessionId: sessionRecord.sandboxSessionId, sandbox: input.sandbox, @@ -238,46 +289,61 @@ export async function ensureSandboxConversationSession( ); } - const openedCloudflareSession = await measureOptional( - input.timing, - "conversation.openSession", - () => - openSandboxConversationSession({ - sandboxSessionId, - cwd, - sandbox: input.sandbox, - shouldCreate: continuation.shouldCreateCloudflareSession, - }), + const openedCloudflareSession = await measureProvisioning("conversation.openSession", () => + openSandboxConversationSession({ + sandboxSessionId, + cwd, + sandbox: input.sandbox, + shouldCreate: continuation.shouldCreateCloudflareSession, + }), ); const cloudflareSession = openedCloudflareSession.session; + const activatedAt = currentTimestampMs(); try { - await measureOptional(input.timing, "conversation.activateRecord", () => + const activated = await measureProvisioning("conversation.activateRecord", () => recordRuntimeConversationSessionActive(bindings.DB, { sandboxSessionId, cwd, - now, + ...(provisioningLease === undefined + ? {} + : { expectedProvisioningOperationId: provisioningLease.operationId }), + now: activatedAt, originJson: JSON.stringify(frozenOrigin), runtimeSubjectId: input.sandboxId, + sandboxIncarnation: input.sandboxIncarnation, sessionId: input.sessionId, }), ); + if (!activated) { + throw new Error("Sandbox conversation activation lost lifecycle ownership."); + } } catch (error) { const message = error instanceof Error ? error.message : "Sandbox conversation session activation failed."; - await recordRuntimeConversationSessionError(bindings.DB, { + const recordedError = await recordRuntimeConversationSessionError(bindings.DB, { sandboxSessionId, + sandboxIncarnation: input.sandboxIncarnation, cwd, errorCode: "runtime.conversation_mount_failed", + ...(provisioningLease === undefined + ? {} + : { expectedProvisioningOperationId: provisioningLease.operationId }), message, - now, + now: activatedAt, originJson: JSON.stringify(frozenOrigin), runtimeSubjectId: input.sandboxId, sessionId: input.sessionId, }); disposeRpcResource(cloudflareSession); + if (!recordedError) { + await deleteSandboxConversationSessionBestEffort({ + sandbox: input.sandbox, + sandboxSessionId, + }); + } throw new Error(message, { cause: error }); } @@ -286,29 +352,68 @@ export async function ensureSandboxConversationSession( sandboxSessionId, cwd, origin: frozenOrigin, + ...(provisioningLease === undefined ? {} : { provisioningLease }), }; } export async function closeSandboxConversationSession( bindings: ApiBindings, input: { + expectedProvisioningOperationId?: RuntimeOperationId; + expectedSandboxSessionId?: SandboxSessionId; sandboxId: SandboxId; sessionId: SessionId; }, ): Promise { const state = await getRuntimeConversationSessionState(bindings.DB, { + ...(input.expectedProvisioningOperationId === undefined + ? {} + : { expectedProvisioningOperationId: input.expectedProvisioningOperationId }), + ...(input.expectedSandboxSessionId === undefined + ? {} + : { expectedSandboxSessionId: input.expectedSandboxSessionId }), runtimeSubjectId: input.sandboxId, sessionId: input.sessionId, }); - if (!state || state.status !== "active") { + if (!state) { return; } + let cleanupOperationId = + state.status === "cleanup_pending" + ? state.cleanupOperationId + : await claimRuntimeConversationSessionCleanup(bindings.DB, { + ...(input.expectedProvisioningOperationId === undefined + ? {} + : { expectedProvisioningOperationId: input.expectedProvisioningOperationId }), + now: currentTimestampMs(), + runtimeSubjectId: input.sandboxId, + sandboxIncarnation: state.sandboxIncarnation, + sandboxSessionId: state.sandboxSessionId, + sessionId: input.sessionId, + }); + if (cleanupOperationId === null) { + const adopted = await getRuntimeConversationSessionState(bindings.DB, { + expectedSandboxIncarnation: state.sandboxIncarnation, + expectedSandboxSessionId: state.sandboxSessionId, + runtimeSubjectId: input.sandboxId, + sessionId: input.sessionId, + }); + if (adopted?.status !== "cleanup_pending" || adopted.cleanupOperationId === null) { + throw new Error("Sandbox conversation cleanup lost lifecycle ownership."); + } + cleanupOperationId = adopted.cleanupOperationId; + } + // Force-close: session-end / cleanup callers must tear down regardless of // idleness. The idle sweep uses closeIdleCattleConversationSession instead. await finalizeSandboxConversationClose(bindings, { sandboxId: input.sandboxId, + cleanupOperationId, + ...(input.expectedProvisioningOperationId === undefined + ? {} + : { expectedProvisioningOperationId: input.expectedProvisioningOperationId }), sessionId: input.sessionId, state, }); @@ -337,20 +442,22 @@ export async function closeIdleCattleConversationSession( return false; } - const claimed = await claimIdleSessionScopedConversationForClose(bindings.DB, { + const cleanupOperationId = await claimIdleSessionScopedConversationForClose(bindings.DB, { idleSinceLte: input.idleSinceLte, now: currentTimestampMs(), runtimeSubjectId: input.sandboxId, + sandboxIncarnation: state.sandboxIncarnation, sandboxSessionId: state.sandboxSessionId, sessionId: input.sessionId, }); - if (!claimed) { + if (cleanupOperationId === null) { return false; } await finalizeSandboxConversationClose(bindings, { sandboxId: input.sandboxId, + cleanupOperationId, sessionId: input.sessionId, state, }); @@ -358,10 +465,32 @@ export async function closeIdleCattleConversationSession( return true; } +export async function repairPendingSandboxConversationSessionCleanups( + bindings: ApiBindings, + limit: number, +): Promise { + const pending = await listPendingRuntimeConversationSessionCleanups(bindings.DB, limit); + + await Promise.allSettled( + pending.map((cleanup) => + finalizeSandboxConversationClose(bindings, { + cleanupOperationId: cleanup.cleanupOperationId, + sandboxId: cleanup.sandboxId, + sessionId: cleanup.sessionId, + state: cleanup, + }), + ), + ); + + return pending.length; +} + async function finalizeSandboxConversationClose( bindings: ApiBindings, input: { sandboxId: SandboxId; + cleanupOperationId: RuntimeOperationId; + expectedProvisioningOperationId?: RuntimeOperationId; sessionId: SessionId; state: RuntimeConversationSessionState; }, @@ -370,36 +499,43 @@ async function finalizeSandboxConversationClose( const { deleteActiveSandboxConversationSession } = await import("./sandbox-conversation-session-delete"); - try { - await deleteActiveSandboxConversationSession(bindings, { - sandboxSessionId: input.state.sandboxSessionId, - sandboxId: input.sandboxId, - }); + await deleteActiveSandboxConversationSession(bindings, { + sandboxSessionId: input.state.sandboxSessionId, + sandboxId: input.sandboxId, + sandboxIncarnation: input.state.sandboxIncarnation, + }); - if (input.state.agentId) { - await appendRuntimeDiagnosticEvent(bindings, { - eventName: RUNTIME_DIAGNOSTIC_EVENT.sandboxSessionDestroyed.name, - sessionId: input.sessionId, - value: { - ...toRuntimeDiagnosticBaseValue({ - agentId: input.state.agentId, - sessionId: input.sessionId, - }), - reason: "runtime_subject_session_closed", - sandboxId: input.sandboxId, - }, - }); - } - } finally { - // Remote cleanup must not strand the local subject outside reclamation. - await recordRuntimeConversationSessionClosed(bindings.DB, { - inactiveDeadlineAt: getRuntimeSubjectInactiveDeadline( - getRuntimeKindPolicy(input.state.kind), - now, - ), + const recorded = await recordRuntimeConversationSessionClosed(bindings.DB, { + cleanupOperationId: input.cleanupOperationId, + ...(input.expectedProvisioningOperationId === undefined + ? {} + : { expectedProvisioningOperationId: input.expectedProvisioningOperationId }), + inactiveDeadlineAt: getRuntimeSubjectInactiveDeadline( + getRuntimeKindPolicy(input.state.kind), now, - runtimeSubjectId: input.sandboxId, + ), + now, + runtimeSubjectId: input.sandboxId, + sandboxIncarnation: input.state.sandboxIncarnation, + sandboxSessionId: input.state.sandboxSessionId, + sessionId: input.sessionId, + }); + if (!recorded) { + throw new Error("Sandbox conversation cleanup lost lifecycle ownership."); + } + + if (input.state.agentId) { + await appendRuntimeDiagnosticEvent(bindings, { + eventName: RUNTIME_DIAGNOSTIC_EVENT.sandboxSessionDestroyed.name, sessionId: input.sessionId, + value: { + ...toRuntimeDiagnosticBaseValue({ + agentId: input.state.agentId, + sessionId: input.sessionId, + }), + reason: "runtime_subject_session_closed", + sandboxId: input.sandboxId, + }, }); } } diff --git a/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-session.types.ts b/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-session.types.ts index 0c5fd2b3..9641a9dc 100644 --- a/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-session.types.ts +++ b/apps/api/src/modules/runtime/infrastructure/sandbox-session/sandbox-session.types.ts @@ -3,6 +3,7 @@ import type { AgentId, SandboxId, SandboxSessionId, SessionId } from "@mosoo/id" import type { RuntimeTimingRecorder } from "../../application/session-runs/session-runtime-timing"; import type { DriverOrigin as DriverOriginValue } from "../../domain/driver-snapshot"; +import type { RuntimeRunProvisioningLease } from "../runtime-subject-lifecycle/runtime-provisioning-lease-store"; import type { ExecutionSessionHandle, SandboxHandle } from "../sandbox-handles"; export interface EnsureSandboxConversationSessionInput { @@ -10,8 +11,11 @@ export interface EnsureSandboxConversationSessionInput { kind: AgentKind; mountSessionResources: boolean; origin: DriverOriginValue; + provisioningLease?: RuntimeRunProvisioningLease; + replaceClosedExecutionSession?: boolean; sandbox: SandboxHandle; sandboxId: SandboxId; + sandboxIncarnation: number; sessionId: SessionId; timing?: RuntimeTimingRecorder; } @@ -21,4 +25,5 @@ export interface SandboxConversationSessionResult { sandboxSessionId: SandboxSessionId; cwd: string; origin: DriverOriginValue; + provisioningLease?: RuntimeRunProvisioningLease; } diff --git a/apps/api/src/modules/runtime/infrastructure/session-resources/session-resource-mount.service.ts b/apps/api/src/modules/runtime/infrastructure/session-resources/session-resource-mount.service.ts index 37e5e886..8c7f7d7d 100644 --- a/apps/api/src/modules/runtime/infrastructure/session-resources/session-resource-mount.service.ts +++ b/apps/api/src/modules/runtime/infrastructure/session-resources/session-resource-mount.service.ts @@ -1,6 +1,11 @@ -import { getSessionResourceRootPath } from "@mosoo/agent-driver/paths"; +import { + SANDBOX_WORKSPACE_ROOT, + getSessionResourceBackingPath, + getSessionResourceRootPath, +} from "@mosoo/agent-driver/paths"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; +import { quoteShellArg } from "../../../../shared/shell"; import { createRuntimeSandboxBucketMountOptions, isRuntimeSandboxLocalBucketEnabled, @@ -10,12 +15,8 @@ import { toRuntimeBucketMountConflictError } from "../runtime-sandbox-mount-erro import { RuntimeBucketMountConflictError } from "../runtime-subject-lifecycle/runtime-subject-errors"; import type { SandboxHandle } from "../sandbox-handles"; -function quoteShellArg(value: string): string { - return `'${value.replaceAll("'", `'"'"'`)}'`; -} - function getSessionResourceMountPath(sessionId: string): string { - return getSessionResourceRootPath(sessionId); + return getSessionResourceBackingPath(sessionId); } function getSessionResourceBucketPrefix(sessionId: string): string { @@ -37,64 +38,82 @@ async function sandboxBucketMountIsReady(input: { return probe.success && probe.exitCode === 0; } +async function ensureSessionResourceAlias(input: { + readonly mountPath: string; + readonly publicPath: string; + readonly sandbox: SandboxHandle; +}): Promise { + const workspacePrefix = `${SANDBOX_WORKSPACE_ROOT}/`; + const publicParent = input.publicPath.slice(0, input.publicPath.lastIndexOf("/")); + if (!publicParent.startsWith(workspacePrefix) || !input.mountPath.startsWith(workspacePrefix)) { + throw new Error("Session resource paths must be inside the workspace root."); + } + const parentDepth = publicParent.slice(workspacePrefix.length).split("/").length; + const target = `${"../".repeat(parentDepth)}${input.mountPath.slice(workspacePrefix.length)}`; + const command = [ + "set -eu", + `link=${quoteShellArg(input.publicPath)}`, + `target=${quoteShellArg(target)}`, + 'if [ -L "$link" ]; then [ "$(readlink "$link")" = "$target" ];', + 'elif [ -e "$link" ]; then exit 42;', + 'else ln -s "$target" "$link" 2>/dev/null || { [ -L "$link" ] && [ "$(readlink "$link")" = "$target" ]; }; fi', + ].join("; "); + const result = await input.sandbox.exec(`sh -lc ${quoteShellArg(command)}`); + if (!result.success || result.exitCode !== 0) { + throw new Error("Session resource alias conflicts with its reserved workspace path."); + } +} + export async function ensureSessionResourcesMounted(input: { bindings: ApiBindings; sandbox: SandboxHandle; sessionId: string; }): Promise { const mountPath = getSessionResourceMountPath(input.sessionId); + const publicPath = getSessionResourceRootPath(input.sessionId); const bucket = resolveRuntimeSandboxBucketMountTarget(input.bindings); const prefix = getSessionResourceBucketPrefix(input.sessionId); const localBucket = isRuntimeSandboxLocalBucketEnabled(input.bindings); - if ( - await sandboxBucketMountIsReady({ - localBucket, - mountPath, - sandbox: input.sandbox, - }) - ) { - return; - } + const ready = await sandboxBucketMountIsReady({ + localBucket, + mountPath, + sandbox: input.sandbox, + }); - await input.sandbox.mkdir(mountPath, { recursive: true }); + if (!ready) { + await input.sandbox.mkdir(mountPath, { recursive: true }); - try { - await input.sandbox.mountBucket( - bucket, - mountPath, - createRuntimeSandboxBucketMountOptions(input.bindings, { - prefix, - readOnly: true, - }), - ); - } catch (cause) { - const error = - toRuntimeBucketMountConflictError(cause, { + try { + await input.sandbox.mountBucket( + bucket, mountPath, - }) ?? cause; - - if ( - error instanceof RuntimeBucketMountConflictError && - (localBucket || - (await sandboxBucketMountIsReady({ - localBucket, + createRuntimeSandboxBucketMountOptions(input.bindings, { + prefix, + readOnly: true, + }), + ); + } catch (cause) { + const error = + toRuntimeBucketMountConflictError(cause, { mountPath, - sandbox: input.sandbox, - }))) - ) { - return; - } + }) ?? cause; - if ( - !localBucket && - error instanceof RuntimeBucketMountConflictError && - error.bucket === bucket && - error.prefix === prefix - ) { - return; - } + const sameMount = + error instanceof RuntimeBucketMountConflictError && + (localBucket || + (await sandboxBucketMountIsReady({ + localBucket, + mountPath, + sandbox: input.sandbox, + })) || + (error.bucket === bucket && error.prefix === prefix)); - throw error; + if (!sameMount) { + throw error; + } + } } + + await ensureSessionResourceAlias({ mountPath, publicPath, sandbox: input.sandbox }); } diff --git a/apps/api/src/modules/runtime/infrastructure/session-runs/external-tool-effect-store.repository.ts b/apps/api/src/modules/runtime/infrastructure/session-runs/external-tool-effect-store.repository.ts index 17d66daa..f19bc1a1 100644 --- a/apps/api/src/modules/runtime/infrastructure/session-runs/external-tool-effect-store.repository.ts +++ b/apps/api/src/modules/runtime/infrastructure/session-runs/external-tool-effect-store.repository.ts @@ -1,10 +1,18 @@ +import { + ExternalToolEffectClaimToken, + ExternalToolEffectSettlement, +} from "@mosoo/contracts/external-tool-effect"; import type { ExternalToolEffectClaim, + ExternalToolEffectState, ExternalToolEffectStatus, } from "@mosoo/contracts/external-tool-effect"; -import { McpExecuteCommandResult } from "@mosoo/contracts/runtime-command"; +import { McpExecuteCommandResult, RuntimeCommand } from "@mosoo/contracts/runtime-command"; +import type { McpExecuteCommand } from "@mosoo/contracts/runtime-command"; import { parseSchemaValue } from "@mosoo/contracts/validation"; import { + driverCommandsTable, + driverInstancesTable, externalToolEffectAttemptsTable, externalToolEffectsTable, sessionRunsTable, @@ -17,7 +25,7 @@ import type { McpServerId, SessionRunId, } from "@mosoo/id"; -import { and, eq, inArray } from "drizzle-orm"; +import { and, eq, exists, inArray, sql } from "drizzle-orm"; import { getAppDatabase, @@ -25,37 +33,115 @@ import { runAppDatabaseBatch, } from "../../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../../time"; +import { LIVE_DRIVER_INSTANCE_STATUSES } from "../../domain/driver-instance-lifecycle.machine"; import { ACTIVE_SESSION_RUN_STATUSES } from "../../domain/session-run-lifecycle.machine"; interface ExternalToolEffectRow { attemptCount: number; + claimToken: string | null; + command: McpExecuteCommand; id: ExternalToolEffectId; idempotencyKey: string; + requestId: string; resultJson: string | null; + serverId: string; status: ExternalToolEffectStatus; + toolName: string; +} + +type EffectLookup = { + commandId: DriverCommandId; + driverGeneration: number; + driverInstanceId: DriverInstanceId; +}; + +function serializeRuntimeSettlement(value: unknown): string { + const serialized = JSON.stringify(value); + if (serialized === undefined) { + throw new TypeError("External tool effect settlement is not JSON serializable."); + } + return serialized; } function parseMcpResult(value: string): typeof McpExecuteCommandResult.infer { return parseSchemaValue(McpExecuteCommandResult, JSON.parse(value)); } +function selectedValue(value: Value, alias: string) { + return sql`${value}`.as(alias); +} + +function toExternalToolEffectState(effect: ExternalToolEffectRow): ExternalToolEffectState { + switch (effect.status) { + case "intent": + return { effectId: effect.id, kind: "intent" }; + case "claimed": + if (effect.claimToken === null || effect.attemptCount < 1) { + throw new Error("Claimed external tool effect is missing its claim audit data."); + } + return { + attempt: effect.attemptCount, + effectId: effect.id, + idempotencyKey: effect.idempotencyKey, + kind: "claimed", + }; + case "succeeded": + if (effect.resultJson === null) { + throw new Error("Succeeded external tool effect is missing its command result."); + } + return { + effectId: effect.id, + kind: "succeeded", + result: parseMcpResultForCommand(effect.resultJson, effect.command), + }; + case "unknown": + return { effectId: effect.id, kind: "unknown" }; + } +} + +function parseMcpResultForCommand( + resultJson: string, + command: McpExecuteCommand, +): typeof McpExecuteCommandResult.infer { + const result = parseMcpResult(resultJson); + if ( + result.requestId !== command.requestId || + result.serverId !== command.serverId || + result.toolName !== command.toolName + ) { + throw new Error("External tool effect result does not match its immutable command intent."); + } + + return result; +} + async function getExternalToolEffect( database: D1Database, - input: { - commandId: DriverCommandId; - driverInstanceId: DriverInstanceId; - }, + input: EffectLookup, ): Promise { const effect = (await getAppDatabase(database) .select({ attemptCount: externalToolEffectsTable.attemptCount, + claimToken: externalToolEffectsTable.claimToken, id: externalToolEffectsTable.id, idempotencyKey: externalToolEffectsTable.idempotencyKey, + payloadJson: driverCommandsTable.payloadJson, resultJson: externalToolEffectsTable.resultJson, + serverId: externalToolEffectsTable.serverId, + sessionRunId: externalToolEffectsTable.sessionRunId, status: externalToolEffectsTable.status, + toolName: externalToolEffectsTable.toolName, }) .from(externalToolEffectsTable) + .innerJoin( + driverCommandsTable, + and( + eq(driverCommandsTable.id, externalToolEffectsTable.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + eq(driverCommandsTable.driverInstanceId, externalToolEffectsTable.driverInstanceId), + ), + ) .where( and( eq(externalToolEffectsTable.commandId, input.commandId), @@ -68,8 +154,18 @@ async function getExternalToolEffect( if (effect === null) { throw new Error("External tool effect intent was not found for the runtime command."); } + const command = parseSchemaValue(RuntimeCommand, JSON.parse(effect.payloadJson)); + if ( + command.kind !== "mcp.execute" || + command.commandId !== input.commandId || + command.runId !== effect.sessionRunId || + command.serverId !== effect.serverId || + command.toolName !== effect.toolName + ) { + throw new Error("External tool effect does not match its immutable command intent."); + } - return effect; + return { ...effect, command, requestId: command.requestId }; } /** @@ -77,41 +173,24 @@ async function getExternalToolEffect( * caller commits it in the same D1 batch as the command record, so either both * become visible to a Driver or neither does. */ -export async function prepareExternalToolEffectIntent( - database: D1Database, - input: { - command: { - commandId: string; - serverId: string; - toolName: string; - }; - driverInstanceId: DriverInstanceId; - }, -): Promise { +export function prepareExternalToolEffectIntent(input: { + command: { + commandId: string; + runId: string; + serverId: string; + toolName: string; + }; + driverInstanceId: DriverInstanceId; +}) { const commandId = parsePlatformId(input.command.commandId, "driver command id"); + const sessionRunId = parsePlatformId(input.command.runId, "Session Run id"); const serverId = parsePlatformId(input.command.serverId, "MCP server id"); - const activeRun = - (await getAppDatabase(database) - .select({ id: sessionRunsTable.id }) - .from(sessionRunsTable) - .where( - and( - eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), - inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), - ), - ) - .limit(1) - .get()) ?? null; - - if (activeRun === null) { - throw new Error("MCP external tool effects require an active Session Run."); - } - const effectId = createPlatformId(); const nowMs = currentTimestampMs(); return { attemptCount: 0, + claimToken: null, commandId, createdAt: nowMs, driverInstanceId: input.driverInstanceId, @@ -120,248 +199,339 @@ export async function prepareExternalToolEffectIntent( providerReceiptJson: null, resultJson: null, serverId, - sessionRunId: activeRun.id, + // This scalar subquery is evaluated in the command+intent D1 batch. If the + // exact Run is no longer active it yields NULL, so the NOT NULL constraint + // rolls the whole batch back instead of leaving a command without a fence. + sessionRunId: sql`( + SELECT ${sessionRunsTable.id} + FROM ${sessionRunsTable} + WHERE ${sessionRunsTable.id} = ${sessionRunId} + AND ${sessionRunsTable.driverInstanceId} = ${input.driverInstanceId} + AND ${sessionRunsTable.status} IN (${sql.join( + ACTIVE_SESSION_RUN_STATUSES.map((status) => sql`${status}`), + sql`, `, + )}) + LIMIT 1 + )`, status: "intent" as const, toolName: input.command.toolName, updatedAt: nowMs, }; } +/** Reads the authoritative ledger without acquiring permission to execute. */ +export async function observeExternalToolEffect( + database: D1Database, + input: EffectLookup, +): Promise { + return toExternalToolEffectState(await getExternalToolEffect(database, input)); +} + /** - * Atomically fences the only permitted provider invocation. An interrupted - * execution becomes unknown instead of being replayed by a new Driver. + * Fences the only permitted provider invocation. A repeated token recovers a + * lost claim response; any different token closes the ambiguous claim instead + * of replaying the provider call. */ export async function claimExternalToolEffect( database: D1Database, - input: { - commandId: DriverCommandId; - driverInstanceId: DriverInstanceId; - }, -): Promise { - const effect = await getExternalToolEffect(database, input); + input: EffectLookup & { claimToken: string }, +): Promise { + parseSchemaValue(ExternalToolEffectClaimToken, input.claimToken); + + for (let retry = 0; retry < 3; retry += 1) { + const effect = await getExternalToolEffect(database, input); - if (effect.status === "succeeded") { - if (effect.resultJson === null) { - throw new Error("Succeeded external tool effect is missing its command result."); + if (effect.status === "succeeded" || effect.status === "unknown") { + const canonical = toExternalToolEffectState(effect); + if (canonical.kind === "intent" || canonical.kind === "claimed") { + throw new Error("A terminal external tool effect returned a non-terminal state."); + } + return canonical; } - return { - effectId: effect.id, - kind: "completed", - result: parseMcpResult(effect.resultJson), - }; - } + if (effect.status === "claimed") { + if (effect.claimToken === input.claimToken) { + const canonical = toExternalToolEffectState(effect); + if (canonical.kind !== "claimed") { + throw new Error("A claimed external tool effect returned a different state."); + } + return canonical; + } - if (effect.status === "unknown") { - return { effectId: effect.id, kind: "unknown" }; - } - - if (effect.status === "executing") { - await markExternalToolEffectUnknown(database, input); - return { effectId: effect.id, kind: "unknown" }; - } + await markClaimedExternalToolEffectUnknown(database, effect); + const canonical = await observeExternalToolEffect(database, input); + if (canonical.kind === "intent") { + throw new Error("A claimed external tool effect unexpectedly returned to intent."); + } + return canonical; + } - const attempt = effect.attemptCount + 1; - const nowMs = currentTimestampMs(); - const [update] = await runAppDatabaseBatch(database, (appDatabase) => [ - appDatabase - .update(externalToolEffectsTable) - .set({ - attemptCount: attempt, - status: "executing", - updatedAt: nowMs, - }) - .where( - and( - eq(externalToolEffectsTable.id, effect.id), - eq(externalToolEffectsTable.attemptCount, effect.attemptCount), - eq(externalToolEffectsTable.status, "intent"), + const attempt = effect.attemptCount + 1; + const nowMs = currentTimestampMs(); + const [update] = await runAppDatabaseBatch(database, (appDatabase) => [ + appDatabase + .update(externalToolEffectsTable) + .set({ + attemptCount: attempt, + claimToken: input.claimToken, + status: "claimed", + updatedAt: nowMs, + }) + .where( + and( + eq(externalToolEffectsTable.id, effect.id), + eq(externalToolEffectsTable.attemptCount, effect.attemptCount), + eq(externalToolEffectsTable.status, "intent"), + exists( + appDatabase + .select({ id: driverInstancesTable.id }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.driverGeneration), + inArray(driverInstancesTable.status, LIVE_DRIVER_INSTANCE_STATUSES), + ), + ), + ), + exists( + appDatabase + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, externalToolEffectsTable.sessionRunId), + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ), + ), + exists( + appDatabase + .select({ id: driverCommandsTable.id }) + .from(driverCommandsTable) + .where( + and( + eq(driverCommandsTable.id, input.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + eq(driverCommandsTable.status, "accepted"), + ), + ), + ), + ), ), - ), - appDatabase - .insert(externalToolEffectAttemptsTable) - .values({ + appDatabase + .insert(externalToolEffectAttemptsTable) + .select( + appDatabase + .select({ + attempt: selectedValue(attempt, "attempt"), + claimToken: selectedValue(input.claimToken, "claim_token"), + completedAt: selectedValue(null, "completed_at"), + createdAt: selectedValue(nowMs, "created_at"), + effectId: externalToolEffectsTable.id, + providerReceiptJson: selectedValue(null, "provider_receipt_json"), + resultJson: selectedValue(null, "result_json"), + status: selectedValue("claimed" as const, "status"), + }) + .from(externalToolEffectsTable) + .where( + and( + eq(externalToolEffectsTable.id, effect.id), + eq(externalToolEffectsTable.attemptCount, attempt), + eq(externalToolEffectsTable.claimToken, input.claimToken), + eq(externalToolEffectsTable.status, "claimed"), + ), + ), + ) + .onConflictDoNothing(), + ]); + + if (getD1ChangeCount(update) > 0) { + return { attempt, - completedAt: null, - createdAt: nowMs, effectId: effect.id, - providerReceiptJson: null, - resultJson: null, - status: "executing", - }) - .onConflictDoNothing(), - ]); - - if (getD1ChangeCount(update) === 0) { - return claimExternalToolEffect(database, input); + idempotencyKey: effect.idempotencyKey, + kind: "claimed", + }; + } } - return { - attempt, - effectId: effect.id, - idempotencyKey: effect.idempotencyKey, - kind: "execute", - }; + throw new Error("External tool effect claim did not reach a stable state."); } /** - * Stores the outcome before the Driver sends the command terminal receipt. - * The Driver preserves provider `_meta` as an opaque receipt when one exists. - * MCP does not standardize reconciliation, so absent receipts remain null and - * a later delivery uncertainty is fenced as unknown rather than replayed. + * Persists an owner's terminal observation and always returns the state that + * actually won the ledger CAS. This makes settlement-response loss retryable. */ -export async function completeExternalToolEffect( +export async function settleExternalToolEffect( database: D1Database, - input: { - commandId: DriverCommandId; - driverInstanceId: DriverInstanceId; - providerReceiptJson?: string | null; - result: typeof McpExecuteCommandResult.infer; + input: EffectLookup & { + claimToken: string; + effectId: ExternalToolEffectId; + settlement: ExternalToolEffectSettlement; }, -): Promise { +): Promise { + parseSchemaValue(ExternalToolEffectClaimToken, input.claimToken); + const settlement = parseSchemaValue( + ExternalToolEffectSettlement, + JSON.parse(serializeRuntimeSettlement(input.settlement)), + ); + const effect = await getExternalToolEffect(database, input); - if (effect.status === "succeeded") { - return; + if (effect.id !== input.effectId) { + throw new Error("External tool effect settlement does not match the command's effect id."); + } + + if (effect.status !== "claimed" || effect.claimToken !== input.claimToken) { + return toExternalToolEffectState(effect); } - if (effect.status !== "executing") { - throw new Error(`Cannot complete an external tool effect in ${effect.status} state.`); + + if ( + settlement.kind === "succeeded" && + (settlement.result.requestId !== effect.requestId || + settlement.result.serverId !== effect.serverId || + settlement.result.toolName !== effect.toolName) + ) { + throw new Error("External tool effect result does not match its immutable command intent."); } const nowMs = currentTimestampMs(); - const resultJson = JSON.stringify(input.result); - const [update] = await runAppDatabaseBatch(database, (appDatabase) => [ + const providerReceiptJson = + settlement.kind === "succeeded" ? (settlement.providerReceiptJson ?? null) : null; + const resultJson = settlement.kind === "succeeded" ? JSON.stringify(settlement.result) : null; + const status = settlement.kind === "succeeded" ? "succeeded" : "unknown"; + + await runAppDatabaseBatch(database, (appDatabase) => [ appDatabase .update(externalToolEffectsTable) .set({ - providerReceiptJson: input.providerReceiptJson ?? null, + providerReceiptJson, resultJson, - status: "succeeded", + status, updatedAt: nowMs, }) .where( and( eq(externalToolEffectsTable.id, effect.id), - eq(externalToolEffectsTable.status, "executing"), + eq(externalToolEffectsTable.claimToken, input.claimToken), + eq(externalToolEffectsTable.status, "claimed"), ), ), appDatabase .update(externalToolEffectAttemptsTable) - .set({ - completedAt: nowMs, - providerReceiptJson: input.providerReceiptJson ?? null, - resultJson, - status: "succeeded", - }) + .set({ completedAt: nowMs, providerReceiptJson, resultJson, status }) .where( and( eq(externalToolEffectAttemptsTable.effectId, effect.id), eq(externalToolEffectAttemptsTable.attempt, effect.attemptCount), - eq(externalToolEffectAttemptsTable.status, "executing"), + eq(externalToolEffectAttemptsTable.claimToken, input.claimToken), + eq(externalToolEffectAttemptsTable.status, "claimed"), ), ), ]); - if (getD1ChangeCount(update) === 0) { - throw new Error("External tool effect completion lost its execution claim."); - } + return observeExternalToolEffect(database, input); } -/** - * No generic MCP provider reconciliation or compensation exists. This is a - * deliberate terminal fence: callers must resolve an unknown effect explicitly - * before they create another external action. - */ -export async function markExternalToolEffectUnknown( +async function markClaimedExternalToolEffectUnknown( database: D1Database, - input: { - commandId: DriverCommandId; - driverInstanceId: DriverInstanceId; - }, + effect: ExternalToolEffectRow, ): Promise { - const effect = await getExternalToolEffect(database, input); - - if (effect.status !== "executing") { + if (effect.status !== "claimed" || effect.claimToken === null) { return; } + const claimToken = effect.claimToken; const nowMs = currentTimestampMs(); - const [update] = await runAppDatabaseBatch(database, (appDatabase) => [ + await runAppDatabaseBatch(database, (appDatabase) => [ appDatabase .update(externalToolEffectsTable) - .set({ - status: "unknown", - updatedAt: nowMs, - }) + .set({ status: "unknown", updatedAt: nowMs }) .where( and( eq(externalToolEffectsTable.id, effect.id), - eq(externalToolEffectsTable.status, "executing"), + eq(externalToolEffectsTable.claimToken, claimToken), + eq(externalToolEffectsTable.status, "claimed"), ), ), appDatabase .update(externalToolEffectAttemptsTable) - .set({ - completedAt: nowMs, - status: "unknown", - }) + .set({ completedAt: nowMs, status: "unknown" }) .where( and( eq(externalToolEffectAttemptsTable.effectId, effect.id), eq(externalToolEffectAttemptsTable.attempt, effect.attemptCount), - eq(externalToolEffectAttemptsTable.status, "executing"), + eq(externalToolEffectAttemptsTable.claimToken, claimToken), + eq(externalToolEffectAttemptsTable.status, "claimed"), ), ), ]); - - if (getD1ChangeCount(update) === 0) { - return; - } } -/** Marks all in-flight effects for a terminal Driver in two bounded writes. */ -export async function markExecutingExternalToolEffectsUnknownForDriver( +/** Fences every unresolved claim owned by a terminal Driver. */ +export async function markClaimedExternalToolEffectsUnknownForDriver( database: D1Database, - driverInstanceId: DriverInstanceId, + input: { + driverGeneration: number; + driverInstanceId: DriverInstanceId; + }, ): Promise { - const effects = await getAppDatabase(database) - .select({ - attemptCount: externalToolEffectsTable.attemptCount, - id: externalToolEffectsTable.id, - }) - .from(externalToolEffectsTable) - .where( - and( - eq(externalToolEffectsTable.driverInstanceId, driverInstanceId), - eq(externalToolEffectsTable.status, "executing"), - ), - ) - .all(); - - if (effects.length === 0) { - return; - } - - const effectIds = effects.map((effect) => effect.id); const nowMs = currentTimestampMs(); await runAppDatabaseBatch(database, (appDatabase) => [ appDatabase - .update(externalToolEffectsTable) - .set({ status: "unknown", updatedAt: nowMs }) + .update(externalToolEffectAttemptsTable) + .set({ completedAt: nowMs, status: "unknown" }) .where( and( - eq(externalToolEffectsTable.driverInstanceId, driverInstanceId), - eq(externalToolEffectsTable.status, "executing"), - inArray(externalToolEffectsTable.id, effectIds), + eq(externalToolEffectAttemptsTable.status, "claimed"), + exists( + appDatabase + .select({ id: externalToolEffectsTable.id }) + .from(externalToolEffectsTable) + .where( + and( + eq(externalToolEffectsTable.id, externalToolEffectAttemptsTable.effectId), + eq(externalToolEffectsTable.driverInstanceId, input.driverInstanceId), + eq(externalToolEffectsTable.status, "claimed"), + exists( + appDatabase + .select({ id: driverCommandsTable.id }) + .from(driverCommandsTable) + .where( + and( + eq(driverCommandsTable.id, externalToolEffectsTable.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + ), + ), + ), + ), + ), + ), ), ), appDatabase - .update(externalToolEffectAttemptsTable) - .set({ completedAt: nowMs, status: "unknown" }) + .update(externalToolEffectsTable) + .set({ status: "unknown", updatedAt: nowMs }) .where( and( - inArray(externalToolEffectAttemptsTable.effectId, effectIds), - eq(externalToolEffectAttemptsTable.status, "executing"), + eq(externalToolEffectsTable.driverInstanceId, input.driverInstanceId), + eq(externalToolEffectsTable.status, "claimed"), + exists( + appDatabase + .select({ id: driverCommandsTable.id }) + .from(driverCommandsTable) + .where( + and( + eq(driverCommandsTable.id, externalToolEffectsTable.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + ), + ), + ), ), ), ]); @@ -369,14 +539,14 @@ export async function markExecutingExternalToolEffectsUnknownForDriver( export async function getExternalToolEffectForCommand( database: D1Database, - input: { - commandId: DriverCommandId; - driverInstanceId: DriverInstanceId; - }, + input: EffectLookup, ): Promise<{ attemptCount: number; + claimToken: string | null; + id: ExternalToolEffectId; idempotencyKey: string; providerReceiptJson: string | null; + resultJson: string | null; sessionRunId: SessionRunId; status: ExternalToolEffectStatus; } | null> { @@ -384,12 +554,23 @@ export async function getExternalToolEffectForCommand( (await getAppDatabase(database) .select({ attemptCount: externalToolEffectsTable.attemptCount, + claimToken: externalToolEffectsTable.claimToken, + id: externalToolEffectsTable.id, idempotencyKey: externalToolEffectsTable.idempotencyKey, providerReceiptJson: externalToolEffectsTable.providerReceiptJson, + resultJson: externalToolEffectsTable.resultJson, sessionRunId: externalToolEffectsTable.sessionRunId, status: externalToolEffectsTable.status, }) .from(externalToolEffectsTable) + .innerJoin( + driverCommandsTable, + and( + eq(driverCommandsTable.id, externalToolEffectsTable.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + eq(driverCommandsTable.driverInstanceId, externalToolEffectsTable.driverInstanceId), + ), + ) .where( and( eq(externalToolEffectsTable.commandId, input.commandId), diff --git a/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-record.mapper.ts b/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-record.mapper.ts index 6021ee41..390fa885 100644 --- a/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-record.mapper.ts +++ b/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-record.mapper.ts @@ -1,7 +1,10 @@ import { RuntimeCommandRecord } from "@mosoo/contracts/runtime-command"; import type { RuntimeCommand, RuntimeCommandStatus } from "@mosoo/contracts/runtime-command"; +import { DurableRunError } from "@mosoo/contracts/session-run"; import { parseSchemaValue } from "@mosoo/contracts/validation"; +import { NonEmptyString } from "@mosoo/contracts/validation"; import type { DriverCommandId, DriverInstanceId } from "@mosoo/id"; +import { type } from "arktype"; import { toIsoString } from "../../../../time"; @@ -25,6 +28,7 @@ class RuntimeCommandStoreCorruptionError extends Error { export interface RuntimeCommandRecordRow { ackedAt: number | null; completedAt: number | null; + driverGeneration: number | null; driverInstanceId: DriverInstanceId; errorJson: string | null; expiresAt: number | null; @@ -37,6 +41,56 @@ export interface RuntimeCommandRecordRow { status: RuntimeCommandStatus; } +const legacyTerminalRuntimeCommandRecordBase = { + ackedAt: "string | null", + completedAt: "string | null", + driverInstanceId: NonEmptyString, + error: DurableRunError.or("null"), + expiresAt: "string | null", + id: NonEmptyString, + issuedAt: "string", + seq: "number >= 0", + status: '"completed" | "failed" | "expired" | "cancelled"', +} as const; + +const LegacyTerminalRuntimeCommandRecord = type({ + ...legacyTerminalRuntimeCommandRecordBase, + kind: '"turn.cancel"', + payload: type({ + commandId: NonEmptyString, + kind: '"turn.cancel"', + "reason?": "string", + }).onUndeclaredKey("reject"), + result: "null", +}) + .onUndeclaredKey("reject") + .or( + type({ + ...legacyTerminalRuntimeCommandRecordBase, + kind: '"permission.resolve"', + payload: type({ + commandId: NonEmptyString, + decision: '"allow_once" | "reject_once"', + kind: '"permission.resolve"', + requestId: NonEmptyString, + }).onUndeclaredKey("reject"), + result: "null", + }).onUndeclaredKey("reject"), + ); +export type LegacyTerminalRuntimeCommandRecord = typeof LegacyTerminalRuntimeCommandRecord.infer; + +export type RuntimeCommandStorageRecord = + | { + driverGeneration: number; + format: "v3"; + record: RuntimeCommandRecord; + } + | { + driverGeneration: null; + format: "legacy-v2-terminal"; + record: LegacyTerminalRuntimeCommandRecord | RuntimeCommandRecord; + }; + type RuntimeCommandJsonColumn = "errorJson" | "payloadJson" | "resultJson"; function readRuntimeCommandJsonRaw( @@ -92,9 +146,11 @@ function parseRuntimeCommandJsonColumn( } } -export function toRuntimeCommandRecordFromRow(row: RuntimeCommandRecordRow): RuntimeCommandRecord { +export function toRuntimeCommandStorageRecordFromRow( + row: RuntimeCommandRecordRow, +): RuntimeCommandStorageRecord { try { - return parseSchemaValue(RuntimeCommandRecord, { + const recordInput = { ackedAt: row.ackedAt === null ? null : toIsoString(row.ackedAt), completedAt: row.completedAt === null ? null : toIsoString(row.completedAt), driverInstanceId: row.driverInstanceId, @@ -107,7 +163,47 @@ export function toRuntimeCommandRecordFromRow(row: RuntimeCommandRecordRow): Run result: row.resultJson === null ? null : parseRuntimeCommandJsonColumn(row, "resultJson"), seq: row.seq, status: row.status, - }); + }; + + if ( + (recordInput.result !== null && recordInput.error !== null) || + (recordInput.result !== null && recordInput.status !== "completed") || + (recordInput.error !== null && recordInput.status === "completed") || + ((recordInput.result !== null || recordInput.error !== null) && + !["completed", "failed", "expired", "cancelled"].includes(recordInput.status)) + ) { + throw new TypeError("Runtime command terminal payload does not match its stored status."); + } + + if (row.driverGeneration === null) { + if (!["completed", "failed", "expired", "cancelled"].includes(row.status)) { + throw new TypeError("Only terminal legacy runtime commands may omit driver generation."); + } + + if (!LegacyTerminalRuntimeCommandRecord.allows(recordInput)) { + return { + driverGeneration: null, + format: "legacy-v2-terminal", + record: parseSchemaValue(RuntimeCommandRecord, recordInput), + }; + } + + return { + driverGeneration: null, + format: "legacy-v2-terminal", + record: parseSchemaValue(LegacyTerminalRuntimeCommandRecord, recordInput), + }; + } + + if (!Number.isSafeInteger(row.driverGeneration) || row.driverGeneration < 0) { + throw new TypeError("Runtime command driver generation is invalid."); + } + + return { + driverGeneration: row.driverGeneration, + format: "v3", + record: parseSchemaValue(RuntimeCommandRecord, recordInput), + }; } catch (error) { if (error instanceof RuntimeCommandStoreCorruptionError) { throw error; @@ -116,7 +212,7 @@ export function toRuntimeCommandRecordFromRow(row: RuntimeCommandRecordRow): Run throw new RuntimeCommandStoreCorruptionError({ cause: error, commandId: row.id, - message: `Runtime command ${row.id} does not match the runtime command contract.`, + message: `Runtime command ${row.id} does not match a supported storage contract.`, }); } } diff --git a/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-store.repository.ts b/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-store.repository.ts index b190c407..0ed32750 100644 --- a/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-store.repository.ts +++ b/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-store.repository.ts @@ -1,26 +1,46 @@ -import { RuntimeCommandRecord } from "@mosoo/contracts/runtime-command"; -import type { +import { createMcpExecuteFailedEventIdentity } from "@mosoo/agent-driver/events"; +import { + InputStartCommandResult, + McpExecuteCommandResult, + RUNTIME_COMMAND_MAX_UTF8_BYTES, + RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES, RuntimeCommand, RuntimeCommandResult, RuntimeCommandStatus, } from "@mosoo/contracts/runtime-command"; +import type { RuntimeCommandRecord } from "@mosoo/contracts/runtime-command"; +import { DurableRunError } from "@mosoo/contracts/session-run"; import type { RunError } from "@mosoo/contracts/session-run"; -import { parseSchemaValue } from "@mosoo/contracts/validation"; -import { driverCommandsTable, driverInstancesTable, externalToolEffectsTable } from "@mosoo/db"; +import { PrimitiveRecord, parseSchemaValue } from "@mosoo/contracts/validation"; +import { + driverCommandsTable, + driverInstancesTable, + externalToolEffectsTable, + sessionEventsTable, + sessionRunsTable, + sessionsTable, +} from "@mosoo/db"; import { parsePlatformId } from "@mosoo/id"; -import type { DriverCommandId, DriverInstanceId, SessionRunId } from "@mosoo/id"; +import type { DriverCommandId, DriverInstanceId, SessionId, SessionRunId } from "@mosoo/id"; +import { stringifyRuntimeEventSemanticValue } from "@mosoo/runtime-events"; import { and, asc, eq, exists, gt, inArray, isNull, lte, ne, or, sql } from "drizzle-orm"; +import type { SQL } from "drizzle-orm"; import { getAppDatabase, getD1ChangeCount, runAppDatabaseBatch, } from "../../../../platform/db/drizzle"; -import { currentTimestampMs, toIsoString } from "../../../../time"; +import { currentTimestampMs } from "../../../../time"; import { LIVE_DRIVER_INSTANCE_STATUSES } from "../../domain/driver-instance-lifecycle.machine"; +import { ACTIVE_SESSION_RUN_STATUSES } from "../../domain/session-run-lifecycle.machine"; +import { createSessionRunTerminalSourceId } from "../../domain/session-run-terminal-event-id"; import { prepareExternalToolEffectIntent } from "./external-tool-effect-store.repository"; -import { toRuntimeCommandRecordFromRow } from "./runtime-command-record.mapper"; -import type { RuntimeCommandRecordRow } from "./runtime-command-record.mapper"; +import { toRuntimeCommandStorageRecordFromRow } from "./runtime-command-record.mapper"; +import type { + RuntimeCommandRecordRow, + RuntimeCommandStorageRecord, +} from "./runtime-command-record.mapper"; import { createRuntimeCommandBatchTransitionOutcome, decideRuntimeCommandTransition, @@ -39,29 +59,80 @@ export interface RuntimeCommandMaintenanceOutcome { recovered: RuntimeCommandBatchTransitionOutcome; } -export interface RuntimeCommandGlobalMaintenanceOutcome extends RuntimeCommandMaintenanceOutcome { +export interface RuntimeCommandTerminalRepairOutcome { + completed: RuntimeCommandBatchTransitionOutcome; failed: RuntimeCommandBatchTransitionOutcome; } -async function getNextRuntimeCommandSeq( - database: D1Database, - driverInstanceId: DriverInstanceId, -): Promise { - const row = - (await getAppDatabase(database) - .update(driverInstancesTable) - .set({ - commandSeqCursor: sql`${driverInstancesTable.commandSeqCursor} + 1`, - }) - .where(eq(driverInstancesTable.id, driverInstanceId)) - .returning({ seq: driverInstancesTable.commandSeqCursor }) - .get()) ?? null; +export interface AcceptedMcpCommandRepair { + command: Extract; + commandId: DriverCommandId; + effectId: string; + effectStatus: "intent" | "succeeded" | "unknown"; + runtimeId: string; + sessionId: SessionId; + terminal: + | { + result: typeof McpExecuteCommandResult.infer; + status: "completed"; + } + | { + error: RunError; + status: "failed"; + }; +} + +export interface AcceptedInputStartCommandRepair { + command: Extract; + commandId: DriverCommandId; + runtimeId: string; + sessionId: SessionId; + terminal: + | { + result: typeof InputStartCommandResult.infer; + status: "completed"; + } + | { + status: "cancelled"; + } + | { + error: RunError; + status: "failed"; + }; +} - if (row === null) { - throw new Error("Driver instance not found while allocating a runtime command sequence."); +function parseSessionRunRepairError(input: { + errorCode: string | null; + errorDetailsJson: string | null; + errorMessage: string | null; + errorRetryable: boolean | null; +}): RunError | null { + if (input.errorCode === null && input.errorDetailsJson === null && input.errorMessage === null) { + return null; } + if ( + input.errorCode === null || + input.errorCode.length === 0 || + input.errorMessage === null || + input.errorMessage.length === 0 + ) { + throw new Error("Session Run has an incomplete authoritative durable error."); + } + + const details = + input.errorDetailsJson === null + ? {} + : parseSchemaValue(PrimitiveRecord, JSON.parse(input.errorDetailsJson)); + return parseSchemaValue(DurableRunError, { + code: input.errorCode, + details, + message: input.errorMessage, + retryable: input.errorRetryable ?? false, + }); +} - return row.seq; +function selectedValue(value: Value, alias: string) { + return sql`${value}`.as(alias); } function isDriverCommandSeqConflict(error: unknown): boolean { @@ -75,6 +146,10 @@ function isDriverCommandSeqConflict(error: unknown): boolean { ); } +function isExternalToolEffectRunAdmissionConflict(error: unknown): boolean { + return error instanceof Error && error.message.includes("external_tool_effect.session_run_id"); +} + type RuntimeCommandErrorDetails = Record; function createRuntimeCommandDeliveryExpiredError(details: RuntimeCommandErrorDetails): RunError { @@ -95,33 +170,522 @@ function createRuntimeCommandDriverTerminalError(details: RuntimeCommandErrorDet }; } +function createRuntimeCommandExternalToolEffectUnknownError(input: { + command: Extract; + effectId: string; +}): RunError { + const message = `External effect ${input.effectId} for MCP tool ${input.command.toolName} has an unknown outcome and will not be replayed.`; + return { + code: "driver.external_tool_effect_unknown", + details: { + commandId: input.command.commandId, + effectId: input.effectId, + requestId: input.command.requestId, + runId: input.command.runId, + serverId: input.command.serverId, + toolName: input.command.toolName, + }, + message, + retryable: false, + }; +} + +function createRuntimeCommandExternalToolEffectNotExecutedError( + details: RuntimeCommandErrorDetails, +): RunError { + return { + code: "driver.external_tool_effect_not_executed", + details, + message: "The Driver stopped before the external tool invocation was claimed.", + retryable: true, + }; +} + +const runtimeCommandJsonEncoder = new TextEncoder(); + +function serializeRuntimeCommandJson(value: unknown, label: string): string { + const serialized = JSON.stringify(value); + + if (serialized === undefined) { + throw new TypeError(`${label} is not JSON serializable.`); + } + + return serialized; +} + +function assertRuntimeCommandJsonLimit(serialized: string, limit: number, label: string): void { + const byteLength = runtimeCommandJsonEncoder.encode(serialized).byteLength; + + if (byteLength > limit) { + throw new RangeError(`${label} exceeds ${limit} UTF-8 bytes.`); + } +} + +function externalToolEffectCommandTerminalGuard( + database: ReturnType, + input: { + commandId: DriverCommandId; + resultJson: string | null; + status: RuntimeCommandStatus; + }, +): SQL { + if (input.status !== "completed") { + return exists( + database + .select({ id: externalToolEffectsTable.id }) + .from(externalToolEffectsTable) + .where( + and( + eq(externalToolEffectsTable.commandId, input.commandId), + inArray(externalToolEffectsTable.status, ["intent", "unknown"]), + ), + ), + ); + } + + if (input.resultJson === null) { + return sql`FALSE`; + } + + const proposed = parseSchemaValue(McpExecuteCommandResult, JSON.parse(input.resultJson)); + return exists( + database + .select({ id: externalToolEffectsTable.id }) + .from(externalToolEffectsTable) + .where( + and( + eq(externalToolEffectsTable.commandId, input.commandId), + eq(externalToolEffectsTable.status, "succeeded"), + sql`json_extract(${externalToolEffectsTable.resultJson}, '$.outputText') = ${proposed.outputText}`, + sql`json_extract(${externalToolEffectsTable.resultJson}, '$.requestId') = ${proposed.requestId}`, + sql`json_extract(${externalToolEffectsTable.resultJson}, '$.serverId') = ${proposed.serverId}`, + sql`json_extract(${externalToolEffectsTable.resultJson}, '$.toolName') = ${proposed.toolName}`, + sql`COALESCE(json_extract(${externalToolEffectsTable.resultJson}, '$.isError'), 0) = ${proposed.isError === true ? 1 : 0}`, + ), + ), + ); +} + +interface RuntimeCommandTerminalState { + readonly ackedAt: number | null; + readonly errorJson: string | null; + readonly kind: RuntimeCommand["kind"]; + readonly payloadJson: string; + readonly resultJson: string | null; + readonly status: RuntimeCommandStatus; +} + +function mcpTerminalSourceEventId( + commandId: DriverCommandId, + status: "cancelled" | "completed", +): string { + return `mcp.execute.${status}:${commandId}`; +} + +function primitiveRecordJsonEquals(column: SQL, value: Record): SQL { + const serialized = JSON.stringify(value); + + return sql`NOT EXISTS ( + SELECT 1 + FROM json_each(${column}) AS stored + WHERE NOT EXISTS ( + SELECT 1 + FROM json_each(${serialized}) AS proposed + WHERE proposed.key = stored.key + AND proposed.type = stored.type + AND proposed.value IS stored.value + ) + ) AND NOT EXISTS ( + SELECT 1 + FROM json_each(${serialized}) AS proposed + WHERE NOT EXISTS ( + SELECT 1 + FROM json_each(${column}) AS stored + WHERE stored.key = proposed.key + AND stored.type = proposed.type + AND stored.value IS proposed.value + ) + )`; +} + +function runtimeCommandTerminalEventGuard( + database: ReturnType, + input: { + command: RuntimeCommand; + driverInstanceId: DriverInstanceId; + error: RunError | null; + result: RuntimeCommandResult | null; + status: RuntimeCommandStatus; + }, +): SQL { + if (!isRuntimeCommandTerminalStatus(input.status)) { + return sql`TRUE`; + } + + if (input.command.kind === "input.start") { + const runId = parsePlatformId( + input.command.runId, + "input.start terminal Session Run ID", + ); + const eventType = + input.status === "completed" + ? "run.completed" + : input.status === "failed" + ? "run.failed" + : "run.cancelled"; + const runStatusGuard = + input.status === "cancelled" || input.status === "expired" + ? inArray(sessionRunsTable.status, ["cancelled", "expired"]) + : input.status === "completed" || input.status === "failed" + ? eq(sessionRunsTable.status, input.status) + : sql`FALSE`; + const runErrorGuard = (() => { + if (input.status === "failed") { + return input.error === null + ? sql`FALSE` + : and( + eq(sessionRunsTable.errorCode, input.error.code), + eq(sessionRunsTable.errorMessage, input.error.message), + eq(sessionRunsTable.errorRetryable, input.error.retryable), + primitiveRecordJsonEquals( + sql`${sessionRunsTable.errorDetailsJson}`, + input.error.details, + ), + ); + } + + return input.status === "completed" + ? and( + isNull(sessionRunsTable.errorCode), + isNull(sessionRunsTable.errorDetailsJson), + isNull(sessionRunsTable.errorMessage), + isNull(sessionRunsTable.errorRetryable), + ) + : sql`TRUE`; + })(); + const terminalRun = database + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, runId), + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + runStatusGuard, + runErrorGuard, + ), + ); + + return exists( + database + .select({ id: sessionEventsTable.id }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.runId, runId), + eq(sessionEventsTable.eventType, eventType), + eq( + sessionEventsTable.sourceEventId, + createSessionRunTerminalSourceId(runId, eventType), + ), + exists(terminalRun), + ), + ), + ); + } + + if (input.command.kind !== "mcp.execute") { + return sql`TRUE`; + } + + const commandId = parsePlatformId( + input.command.commandId, + "mcp.execute terminal command ID", + ); + const runId = parsePlatformId( + input.command.runId, + "mcp.execute terminal Session Run ID", + ); + + const toolStatus = + input.status === "completed" ? "completed" : input.status === "failed" ? "failed" : "cancelled"; + const outputGuard = + input.status === "completed" && input.result !== null + ? eq( + sessionEventsTable.toolOutputText, + parseSchemaValue(McpExecuteCommandResult, input.result).outputText, + ) + : input.status === "failed" && input.error !== null + ? eq(sessionEventsTable.toolOutputText, input.error.message) + : isNull(sessionEventsTable.toolOutputText); + const sourceEventId = + input.status === "failed" + ? input.error === null + ? null + : createMcpExecuteFailedEventIdentity({ + commandId, + rawInput: input.command.argumentsJson, + rawOutput: input.error.message, + title: input.command.toolName, + toolCallId: input.command.toolCallId, + }).sourceEventId + : mcpTerminalSourceEventId( + commandId, + input.status === "completed" ? "completed" : "cancelled", + ); + + return exists( + database + .select({ id: sessionEventsTable.id }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.mcpCommandId, commandId), + eq(sessionEventsTable.runId, runId), + eq(sessionEventsTable.eventType, "tool.call.updated"), + sourceEventId === null ? sql`FALSE` : eq(sessionEventsTable.sourceEventId, sourceEventId), + eq(sessionEventsTable.toolCallId, input.command.toolCallId), + eq(sessionEventsTable.toolInputJson, input.command.argumentsJson), + eq(sessionEventsTable.toolName, input.command.toolName), + eq(sessionEventsTable.toolStatus, toolStatus), + outputGuard, + ), + ), + ); +} + +async function assertRuntimeCommandTerminalEvent( + database: D1Database, + input: Parameters[1], +): Promise { + const row = await getAppDatabase(database) + .select({ + valid: sql`CASE WHEN ${runtimeCommandTerminalEventGuard( + getAppDatabase(database), + input, + )} THEN 1 ELSE 0 END`, + }) + .from(driverCommandsTable) + .where( + eq( + driverCommandsTable.id, + parsePlatformId(input.command.commandId, "Runtime command ID"), + ), + ) + .limit(1) + .get(); + + if (row?.valid !== 1) { + throw new Error( + `Runtime command ${input.command.commandId} terminal update is missing its durable event.`, + ); + } +} + +function parseRuntimeCommandResultJson( + kind: RuntimeCommand["kind"], + value: string | null, +): RuntimeCommandResult | null { + if (value === null) { + return null; + } + + const result = parseSchemaValue(RuntimeCommandResult, JSON.parse(value)); + + if (kind === "input.start") { + return parseSchemaValue(InputStartCommandResult, result); + } + + if (kind === "mcp.execute") { + const parsed = parseSchemaValue(McpExecuteCommandResult, result); + return { ...parsed, isError: parsed.isError ?? false }; + } + + throw new Error(`${kind} runtime commands cannot contain a durable result.`); +} + +function parseRuntimeCommandErrorJson(value: string | null): RunError | null { + return value === null ? null : parseSchemaValue(DurableRunError, JSON.parse(value)); +} + +function assertEqualRuntimeCommandPayload( + expected: unknown, + actual: unknown, + commandId: DriverCommandId, +): void { + if (stringifyRuntimeEventSemanticValue(expected) !== stringifyRuntimeEventSemanticValue(actual)) { + throw new Error( + `Runtime command ${commandId} duplicate conflicts with its durable terminal payload.`, + ); + } +} + +async function assertMcpTerminalEffectState( + database: D1Database, + input: { + commandId: DriverCommandId; + driverGeneration: number; + driverInstanceId: DriverInstanceId; + proposedResult: RuntimeCommandResult | null; + storedResult: RuntimeCommandResult | null; + status: RuntimeCommandStatus; + }, +): Promise { + const effect = + (await getAppDatabase(database) + .select({ + resultJson: externalToolEffectsTable.resultJson, + status: externalToolEffectsTable.status, + }) + .from(externalToolEffectsTable) + .innerJoin( + driverCommandsTable, + and( + eq(driverCommandsTable.id, externalToolEffectsTable.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + ), + ) + .where( + and( + eq(externalToolEffectsTable.commandId, input.commandId), + eq(externalToolEffectsTable.driverInstanceId, input.driverInstanceId), + ), + ) + .limit(1) + .get()) ?? null; + + if (effect === null) { + throw new Error(`MCP runtime command ${input.commandId} is missing its durable effect.`); + } + + if (input.status === "completed") { + if (effect.status !== "succeeded" || effect.resultJson === null) { + throw new Error(`MCP runtime command ${input.commandId} has no succeeded durable effect.`); + } + + const effectResult = parseRuntimeCommandResultJson("mcp.execute", effect.resultJson); + assertEqualRuntimeCommandPayload(effectResult, input.storedResult, input.commandId); + assertEqualRuntimeCommandPayload(effectResult, input.proposedResult, input.commandId); + return; + } + + if (!isRuntimeCommandTerminalStatus(input.status)) { + return; + } + + if ((effect.status !== "intent" && effect.status !== "unknown") || effect.resultJson !== null) { + throw new Error( + `MCP runtime command ${input.commandId} terminal state conflicts with its effect.`, + ); + } +} + +async function assertRuntimeCommandDuplicatePayload( + database: D1Database, + input: { + commandId: DriverCommandId; + driverGeneration: number; + driverInstanceId: DriverInstanceId; + errorJson: string | null; + resultJson: string | null; + stored: RuntimeCommandTerminalState; + }, +): Promise { + const proposedError = parseRuntimeCommandErrorJson(input.errorJson); + const storedError = parseRuntimeCommandErrorJson(input.stored.errorJson); + const proposedResult = parseRuntimeCommandResultJson(input.stored.kind, input.resultJson); + const storedResult = parseRuntimeCommandResultJson(input.stored.kind, input.stored.resultJson); + const command = parseSchemaValue(RuntimeCommand, JSON.parse(input.stored.payloadJson)); + + assertEqualRuntimeCommandPayload(storedError, proposedError, input.commandId); + assertEqualRuntimeCommandPayload(storedResult, proposedResult, input.commandId); + + if (isRuntimeCommandTerminalStatus(input.stored.status)) { + await assertRuntimeCommandTerminalEvent(database, { + command, + driverInstanceId: input.driverInstanceId, + error: storedError, + result: storedResult, + status: input.stored.status, + }); + } + + if (input.stored.kind === "mcp.execute") { + await assertMcpTerminalEffectState(database, { + commandId: input.commandId, + driverGeneration: input.driverGeneration, + driverInstanceId: input.driverInstanceId, + proposedResult, + status: input.stored.status, + storedResult, + }); + } +} + +function runtimeCommandActiveRunGuard(database: ReturnType): SQL { + return or( + eq(driverCommandsTable.kind, "session.stop"), + exists( + database + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq( + sessionRunsTable.id, + sql`json_extract(${driverCommandsTable.payloadJson}, '$.runId')`, + ), + eq(sessionRunsTable.driverInstanceId, driverCommandsTable.driverInstanceId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ), + ), + )!; +} + export async function createRuntimeCommandRecord( database: D1Database, input: { command: RuntimeCommand; + driverGeneration: number; driverInstanceId: DriverInstanceId; expiresAt?: number | null; status?: RuntimeCommandStatus; }, ): Promise { + const payloadJson = serializeRuntimeCommandJson(input.command, "Runtime command"); + assertRuntimeCommandJsonLimit(payloadJson, RUNTIME_COMMAND_MAX_UTF8_BYTES, "Runtime command"); + const command = parseSchemaValue(RuntimeCommand, JSON.parse(payloadJson)); + if (!Number.isSafeInteger(input.driverGeneration) || input.driverGeneration < 0) { + throw new TypeError("Driver generation must be a non-negative safe integer."); + } const issuedAt = currentTimestampMs(); - const commandId = parsePlatformId(input.command.commandId, "Runtime command ID"); - const status = input.status ?? "queued"; - const payloadJson = JSON.stringify(input.command); - - return createRuntimeCommandRecordAttempt(database, input, { - attempt: 0, - commandId, - issuedAt, - payloadJson, - status, - }); + const commandId = parsePlatformId(command.commandId, "Runtime command ID"); + const sessionRunId = + command.kind === "session.stop" + ? null + : parsePlatformId(command.runId, "Runtime command Session Run ID"); + const status = parseSchemaValue(RuntimeCommandStatus, input.status ?? "queued"); + + return createRuntimeCommandRecordAttempt( + database, + { ...input, command }, + { + attempt: 0, + commandId, + issuedAt, + payloadJson, + sessionRunId, + status, + }, + ); } async function createRuntimeCommandRecordAttempt( database: D1Database, input: { command: RuntimeCommand; + driverGeneration: number; driverInstanceId: DriverInstanceId; expiresAt?: number | null; status?: RuntimeCommandStatus; @@ -131,6 +695,7 @@ async function createRuntimeCommandRecordAttempt( commandId: DriverCommandId; issuedAt: number; payloadJson: string; + sessionRunId: SessionRunId | null; status: RuntimeCommandStatus; }, ): Promise { @@ -141,57 +706,87 @@ async function createRuntimeCommandRecordAttempt( try { const externalToolEffectIntent = input.command.kind === "mcp.execute" - ? await prepareExternalToolEffectIntent(database, { + ? prepareExternalToolEffectIntent({ command: input.command, driverInstanceId: input.driverInstanceId, }) : null; - const seq = await getNextRuntimeCommandSeq(database, input.driverInstanceId); const commandValues = { - ackedAt: null, - completedAt: null, - deliveryConnectionId: null, - driverInstanceId: input.driverInstanceId, - errorJson: null, - expiresAt: input.expiresAt ?? null, - id: state.commandId, - issuedAt: state.issuedAt, - kind: input.command.kind, - payloadJson: state.payloadJson, - resultJson: null, - seq, - status: state.status, + ackedAt: selectedValue(null, "acked_at"), + completedAt: selectedValue(null, "completed_at"), + deliveryConnectionId: selectedValue(null, "delivery_connection_id"), + driverGeneration: driverInstancesTable.generation, + driverInstanceId: driverInstancesTable.id, + errorJson: selectedValue(null, "error_json"), + expiresAt: selectedValue(input.expiresAt ?? null, "expires_at"), + id: selectedValue(state.commandId, "id"), + issuedAt: selectedValue(state.issuedAt, "issued_at"), + kind: selectedValue(input.command.kind, "kind"), + payloadJson: selectedValue(state.payloadJson, "payload_json"), + resultJson: selectedValue(null, "result_json"), + seq: driverInstancesTable.commandSeqCursor, + status: selectedValue(state.status, "status"), }; + const results = await runAppDatabaseBatch(database, (appDatabase) => { + const activeRun = + state.sessionRunId === null + ? null + : appDatabase + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, state.sessionRunId), + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ); + const admission = and( + eq(driverInstancesTable.id, input.driverInstanceId), + eq(driverInstancesTable.generation, input.driverGeneration), + inArray(driverInstancesTable.status, [...LIVE_DRIVER_INSTANCE_STATUSES]), + ...(activeRun === null ? [] : [exists(activeRun)]), + ); + const commandInsert = appDatabase + .insert(driverCommandsTable) + .select(appDatabase.select(commandValues).from(driverInstancesTable).where(admission)); + const cursorUpdate = appDatabase + .update(driverInstancesTable) + .set({ commandSeqCursor: sql`${driverInstancesTable.commandSeqCursor} + 1` }) + .where(admission); + + return externalToolEffectIntent === null + ? [cursorUpdate, commandInsert] + : [ + cursorUpdate, + commandInsert, + appDatabase.insert(externalToolEffectsTable).values(externalToolEffectIntent), + ]; + }); - if (externalToolEffectIntent === null) { - await getAppDatabase(database).insert(driverCommandsTable).values(commandValues).run(); - } else { - await runAppDatabaseBatch(database, (appDatabase) => [ - appDatabase.insert(driverCommandsTable).values(commandValues), - appDatabase.insert(externalToolEffectsTable).values(externalToolEffectIntent), - ]); + if (getD1ChangeCount((results as readonly unknown[])[1]) === 0) { + throw new Error("Driver generation is no longer current."); } - const record = parseSchemaValue(RuntimeCommandRecord, { - ackedAt: null, - completedAt: null, - driverInstanceId: input.driverInstanceId, - error: null, - expiresAt: - input.expiresAt === null || input.expiresAt === undefined - ? null - : toIsoString(input.expiresAt), - id: state.commandId, - issuedAt: toIsoString(state.issuedAt), - kind: input.command.kind, - payload: input.command, - result: null, - seq, - status: state.status, - }); + const record = await getRuntimeCommandRecord( + database, + input.driverInstanceId, + input.driverGeneration, + state.commandId, + ); + + if (record === null) { + throw new Error("Runtime command was not persisted."); + } return record; } catch (error) { + if (input.command.kind === "mcp.execute" && isExternalToolEffectRunAdmissionConflict(error)) { + throw new Error("MCP external tool effects require the command's active Session Run.", { + cause: error, + }); + } + if (state.attempt < 4 && isDriverCommandSeqConflict(error)) { return createRuntimeCommandRecordAttempt(database, input, { ...state, @@ -208,22 +803,67 @@ export async function updateRuntimeCommandRecord( input: { commandId: DriverCommandId; deliveryConnectionId?: string; + driverGeneration: number; driverInstanceId: DriverInstanceId; error?: RunError; result?: RuntimeCommandResult; status: RuntimeCommandStatus; }, ): Promise { + const status = parseSchemaValue(RuntimeCommandStatus, input.status); + if (input.error !== undefined && input.result !== undefined) { + throw new TypeError("Runtime command terminal update cannot contain both error and result."); + } + + const errorJson = + input.error === undefined + ? null + : serializeRuntimeCommandJson(input.error, "Runtime command terminal error"); + const resultJson = + input.result === undefined + ? null + : serializeRuntimeCommandJson(input.result, "Runtime command result"); + const terminalPayloadJson = errorJson ?? resultJson; + if (terminalPayloadJson !== null) { + assertRuntimeCommandJsonLimit( + terminalPayloadJson, + RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES, + "Runtime command terminal payload", + ); + } + if (errorJson !== null) { + parseSchemaValue(DurableRunError, JSON.parse(errorJson)); + } + const parsedResult = + resultJson === null + ? undefined + : parseSchemaValue(RuntimeCommandResult, JSON.parse(resultJson)); + const storedResultJson = parsedResult === null ? null : resultJson; + if ( + (!isRuntimeCommandTerminalStatus(status) && + (errorJson !== null || storedResultJson !== null)) || + (status === "completed" && errorJson !== null) || + (status !== "completed" && isRuntimeCommandTerminalStatus(status) && storedResultJson !== null) + ) { + throw new TypeError("Runtime command terminal payload does not match its target status."); + } + const current = (await getAppDatabase(database) .select({ + ackedAt: driverCommandsTable.ackedAt, deliveryConnectionId: driverCommandsTable.deliveryConnectionId, + errorJson: driverCommandsTable.errorJson, + kind: driverCommandsTable.kind, + payloadJson: driverCommandsTable.payloadJson, + resultJson: driverCommandsTable.resultJson, status: driverCommandsTable.status, }) .from(driverCommandsTable) .where( and( eq(driverCommandsTable.id, input.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), ), ) @@ -235,10 +875,42 @@ export async function updateRuntimeCommandRecord( currentStatus: null, kind: "rejected", reason: "command_not_found", - targetStatus: input.status, + targetStatus: status, }; } + const command = parseSchemaValue(RuntimeCommand, JSON.parse(current.payloadJson)); + if (command.commandId !== input.commandId || command.kind !== current.kind) { + throw new Error(`Runtime command ${input.commandId} has inconsistent immutable payload.`); + } + + if (parsedResult !== undefined && parsedResult !== null) { + if (current.kind === "input.start") { + parseSchemaValue(InputStartCommandResult, parsedResult); + if (Object.keys(parsedResult).some((key) => key !== "requestId")) { + throw new TypeError("input.start runtime command results may only contain requestId."); + } + } else if (current.kind === "mcp.execute") { + parseSchemaValue(McpExecuteCommandResult, parsedResult); + const allowedKeys = new Set(["isError", "outputText", "requestId", "serverId", "toolName"]); + if (Object.keys(parsedResult).some((key) => !allowedKeys.has(key))) { + throw new TypeError("mcp.execute runtime command result contains an unknown field."); + } + } else { + throw new TypeError(`${current.kind} runtime commands cannot store a result.`); + } + } + + if ( + command.kind === "input.start" && + status === "completed" && + (parsedResult === undefined || + parsedResult === null || + parseSchemaValue(InputStartCommandResult, parsedResult).requestId !== command.requestId) + ) { + throw new TypeError("input.start completion must carry its exact requestId."); + } + if ( input.deliveryConnectionId !== undefined && current.deliveryConnectionId !== input.deliveryConnectionId @@ -247,20 +919,42 @@ export async function updateRuntimeCommandRecord( currentStatus: current.status, kind: "rejected", reason: "stale_delivery_connection", - targetStatus: input.status, + targetStatus: status, }; } - const transition = decideRuntimeCommandTransition(current.status, input.status); + const transition = decideRuntimeCommandTransition(current.status, status); + + if (transition.kind === "duplicate") { + await assertRuntimeCommandDuplicatePayload(database, { + commandId: input.commandId, + driverGeneration: input.driverGeneration, + driverInstanceId: input.driverInstanceId, + errorJson, + resultJson: storedResultJson, + stored: { ...current, kind: command.kind }, + }); + return transition; + } if (transition.kind !== "applied") { return transition; } const timestampMs = currentTimestampMs(); - const ackedAt = isRuntimeCommandAcknowledgedStatus(input.status) ? timestampMs : null; - const completedAt = isRuntimeCommandTerminalStatus(input.status) ? timestampMs : null; - + const ackedAt = isRuntimeCommandAcknowledgedStatus(status) ? timestampMs : null; + const completedAt = isRuntimeCommandTerminalStatus(status) ? timestampMs : null; + + const authoritativeResultJson = + current.kind === "mcp.execute" && status === "completed" + ? sql`( + SELECT ${externalToolEffectsTable.resultJson} + FROM ${externalToolEffectsTable} + WHERE ${externalToolEffectsTable.commandId} = ${input.commandId} + AND ${externalToolEffectsTable.status} = 'succeeded' + LIMIT 1 + )` + : storedResultJson; const result = await getAppDatabase(database) .update(driverCommandsTable) .set({ @@ -270,30 +964,95 @@ export async function updateRuntimeCommandRecord( completedAt === null ? undefined : sql`COALESCE(${completedAt}, ${driverCommandsTable.completedAt})`, - errorJson: input.error === undefined ? null : JSON.stringify(input.error), - resultJson: input.result === undefined ? null : JSON.stringify(input.result), - status: input.status, + errorJson, + resultJson: authoritativeResultJson, + status, }) .where( and( eq(driverCommandsTable.id, input.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), eq(driverCommandsTable.status, current.status), ...(input.deliveryConnectionId === undefined ? [] : [eq(driverCommandsTable.deliveryConnectionId, input.deliveryConnectionId)]), + ...(status === "accepted" ? [runtimeCommandActiveRunGuard(getAppDatabase(database))] : []), + ...(current.kind === "mcp.execute" && isRuntimeCommandTerminalStatus(status) + ? [ + externalToolEffectCommandTerminalGuard(getAppDatabase(database), { + commandId: input.commandId, + resultJson: storedResultJson, + status, + }), + ] + : []), + ...(isRuntimeCommandTerminalStatus(status) && + (command.kind === "input.start" || command.kind === "mcp.execute") + ? [ + runtimeCommandTerminalEventGuard(getAppDatabase(database), { + command, + driverInstanceId: input.driverInstanceId, + error: parseRuntimeCommandErrorJson(errorJson), + result: + storedResultJson === null + ? null + : parseRuntimeCommandResultJson(command.kind, storedResultJson), + status, + }), + ] + : []), ), ) .run(); - return getD1ChangeCount(result) > 0 - ? transition - : { - currentStatus: current.status, - kind: "rejected", - reason: "illegal_transition", - targetStatus: input.status, - }; + if (getD1ChangeCount(result) > 0) { + return transition; + } + + const winner = + (await getAppDatabase(database) + .select({ + ackedAt: driverCommandsTable.ackedAt, + errorJson: driverCommandsTable.errorJson, + kind: driverCommandsTable.kind, + payloadJson: driverCommandsTable.payloadJson, + resultJson: driverCommandsTable.resultJson, + status: driverCommandsTable.status, + }) + .from(driverCommandsTable) + .where( + and( + eq(driverCommandsTable.id, input.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + ), + ) + .limit(1) + .get()) ?? null; + + if (winner?.status === status) { + const winnerCommand = parseSchemaValue(RuntimeCommand, JSON.parse(winner.payloadJson)); + if (winnerCommand.commandId !== input.commandId || winnerCommand.kind !== winner.kind) { + throw new Error(`Runtime command ${input.commandId} has inconsistent immutable payload.`); + } + await assertRuntimeCommandDuplicatePayload(database, { + commandId: input.commandId, + driverGeneration: input.driverGeneration, + driverInstanceId: input.driverInstanceId, + errorJson, + resultJson: storedResultJson, + stored: { ...winner, kind: winnerCommand.kind }, + }); + return { kind: "duplicate", status }; + } + + return { + currentStatus: winner?.status ?? current.status, + kind: "rejected", + reason: "illegal_transition", + targetStatus: status, + }; } export async function markRuntimeCommandRecordDelivered( @@ -301,6 +1060,7 @@ export async function markRuntimeCommandRecordDelivered( input: { commandId: DriverCommandId; connectionId: string; + driverGeneration: number; driverInstanceId: DriverInstanceId; expiresAfter?: number; }, @@ -311,12 +1071,15 @@ export async function markRuntimeCommandRecordDelivered( (await db .select({ deliveryConnectionId: driverCommandsTable.deliveryConnectionId, + kind: driverCommandsTable.kind, + payloadJson: driverCommandsTable.payloadJson, status: driverCommandsTable.status, }) .from(driverCommandsTable) .where( and( eq(driverCommandsTable.id, input.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), ), ) @@ -330,6 +1093,7 @@ export async function markRuntimeCommandRecordDelivered( and( eq(driverInstancesTable.id, input.driverInstanceId), eq(driverInstancesTable.connectionId, input.connectionId), + eq(driverInstancesTable.generation, input.driverGeneration), ), ) .limit(1) @@ -344,6 +1108,50 @@ export async function markRuntimeCommandRecordDelivered( }; } + if (current.kind !== "session.stop") { + const command = parseSchemaValue(RuntimeCommand, JSON.parse(current.payloadJson)); + + if (command.kind === "session.stop") { + throw new Error("Runtime command kind does not match its payload."); + } + + const runId = parsePlatformId(command.runId, "Runtime command Session Run ID"); + const activeRun = + (await db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, runId), + eq(sessionRunsTable.driverInstanceId, input.driverInstanceId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ) + .limit(1) + .get()) ?? null; + + if (activeRun === null) { + await db + .update(driverCommandsTable) + .set({ completedAt: currentTimestampMs(), status: "cancelled" }) + .where( + and( + eq(driverCommandsTable.id, input.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + eq(driverCommandsTable.status, current.status), + ), + ) + .run(); + return { + currentStatus: current.status, + kind: "rejected", + reason: "inactive_session_run", + targetStatus, + }; + } + } + if (activeConnection === null) { return { currentStatus: current.status, @@ -379,6 +1187,7 @@ export async function markRuntimeCommandRecordDelivered( and( eq(driverInstancesTable.id, input.driverInstanceId), eq(driverInstancesTable.connectionId, input.connectionId), + eq(driverInstancesTable.generation, input.driverGeneration), ), ); const result = await db @@ -390,6 +1199,7 @@ export async function markRuntimeCommandRecordDelivered( .where( and( eq(driverCommandsTable.id, input.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), eq(driverCommandsTable.status, current.status), ...(input.expiresAfter === undefined @@ -401,6 +1211,7 @@ export async function markRuntimeCommandRecordDelivered( )!, ]), exists(activeConnectionQuery), + runtimeCommandActiveRunGuard(db), ), ) .run(); @@ -418,13 +1229,27 @@ export async function markRuntimeCommandRecordDelivered( export async function getRuntimeCommandRecord( database: D1Database, driverInstanceId: DriverInstanceId, + driverGeneration: number, commandId: DriverCommandId, ): Promise { + const stored = await getRuntimeCommandStorageRecord(database, driverInstanceId, commandId); + return stored?.format === "v3" && stored.driverGeneration === driverGeneration + ? stored.record + : null; +} + +/** Reads immutable pre-v3 terminal history without making it executable wire data. */ +export async function getRuntimeCommandStorageRecord( + database: D1Database, + driverInstanceId: DriverInstanceId, + commandId: DriverCommandId, +): Promise { const row = (await getAppDatabase(database) .select({ ackedAt: driverCommandsTable.ackedAt, completedAt: driverCommandsTable.completedAt, + driverGeneration: driverCommandsTable.driverGeneration, driverInstanceId: driverCommandsTable.driverInstanceId, errorJson: driverCommandsTable.errorJson, expiresAt: driverCommandsTable.expiresAt, @@ -450,12 +1275,13 @@ export async function getRuntimeCommandRecord( return null; } - return toRuntimeCommandRecordFromRow(row); + return toRuntimeCommandStorageRecordFromRow(row); } async function expireRuntimeCommandDeliveryLeases( database: D1Database, driverInstanceId: DriverInstanceId, + driverGeneration: number, nowMs: number, ): Promise { const expirableStatuses = getRuntimeCommandDeliveryLeaseExpirableStatuses(); @@ -473,6 +1299,7 @@ async function expireRuntimeCommandDeliveryLeases( .where( and( eq(driverCommandsTable.driverInstanceId, driverInstanceId), + eq(driverCommandsTable.driverGeneration, driverGeneration), inArray(driverCommandsTable.status, [...expirableStatuses]), or(eq(driverCommandsTable.status, queuedStatus), isNull(driverCommandsTable.ackedAt)), lte(driverCommandsTable.expiresAt, nowMs), @@ -553,6 +1380,7 @@ async function recoverRuntimeCommandsDeliveredToStaleConnections( database: D1Database, input: { connectionId: string; + driverGeneration: number; driverInstanceId: DriverInstanceId; }, ): Promise { @@ -568,6 +1396,7 @@ async function recoverRuntimeCommandsDeliveredToStaleConnections( .where( and( eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), inArray(driverCommandsTable.status, [...recoverableStatuses]), isNull(driverCommandsTable.ackedAt), or( @@ -593,6 +1422,7 @@ async function recoverRuntimeCommandsDeliveredToStaleConnectionsGlobally( .where( and( eq(driverInstancesTable.id, driverCommandsTable.driverInstanceId), + eq(driverInstancesTable.generation, driverCommandsTable.driverGeneration), inArray(driverInstancesTable.status, [...LIVE_DRIVER_INSTANCE_STATUSES]), or( isNull(driverCommandsTable.deliveryConnectionId), @@ -620,71 +1450,382 @@ async function recoverRuntimeCommandsDeliveredToStaleConnectionsGlobally( return createRuntimeCommandBatchTransitionOutcome(targetStatus, getD1ChangeCount(result)); } -export async function failAcceptedRuntimeCommandsForTerminalDriver( +/** Returns accepted input commands whose authoritative Session Run is terminal. */ +export async function listAcceptedInputStartCommandRepairsForTerminalDriver( database: D1Database, input: { + driverGeneration: number; driverInstanceId: DriverInstanceId; - nowMs?: number; }, -): Promise { - const nowMs = input.nowMs ?? currentTimestampMs(); - const targetStatus = "failed" satisfies RuntimeCommandStatus; - const error = createRuntimeCommandDriverTerminalError({ - driverInstanceId: input.driverInstanceId, - }); +): Promise { + const rows = await getAppDatabase(database) + .select({ + commandId: driverCommandsTable.id, + payloadJson: driverCommandsTable.payloadJson, + run: { + errorCode: sessionRunsTable.errorCode, + errorDetailsJson: sessionRunsTable.errorDetailsJson, + errorMessage: sessionRunsTable.errorMessage, + errorRetryable: sessionRunsTable.errorRetryable, + sessionId: sessionRunsTable.sessionId, + status: sessionRunsTable.status, + }, + runtimeId: sql`COALESCE(${sessionRunsTable.runtimeId}, ${sessionsTable.runtimeId})`, + }) + .from(driverCommandsTable) + .innerJoin( + sessionRunsTable, + and( + sql`json_extract(${driverCommandsTable.payloadJson}, '$.runId') = ${sessionRunsTable.id}`, + eq(sessionRunsTable.driverInstanceId, driverCommandsTable.driverInstanceId), + ), + ) + .innerJoin(sessionsTable, eq(sessionsTable.id, sessionRunsTable.sessionId)) + .where( + and( + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + eq(driverCommandsTable.kind, "input.start"), + eq(driverCommandsTable.status, "accepted"), + ), + ) + .orderBy(asc(driverCommandsTable.seq)) + .all(); + + const repairs: AcceptedInputStartCommandRepair[] = []; + for (const row of rows) { + const command = parseSchemaValue(RuntimeCommand, JSON.parse(row.payloadJson)); + if (command.kind !== "input.start" || command.commandId !== row.commandId) { + throw new Error("Input command repair does not match its immutable command intent."); + } + if (row.runtimeId.length === 0) { + throw new Error("Input command repair is missing its durable completion identity."); + } - const result = await getAppDatabase(database) - .update(driverCommandsTable) - .set({ - completedAt: sql`COALESCE(${driverCommandsTable.completedAt}, ${nowMs})`, - errorJson: JSON.stringify(error), - status: targetStatus, + const runError = parseSessionRunRepairError(row.run); + const common = { + command, + commandId: row.commandId, + runtimeId: row.runtimeId, + sessionId: row.run.sessionId, + }; + + if (row.run.status === "completed") { + if (runError !== null) { + throw new Error("Completed Session Run unexpectedly contains a durable error."); + } + repairs.push({ + ...common, + terminal: { + result: { requestId: command.requestId }, + status: "completed" as const, + }, + }); + continue; + } + if (row.run.status === "cancelled" || row.run.status === "expired") { + repairs.push({ ...common, terminal: { status: "cancelled" as const } }); + continue; + } + if (row.run.status === "failed") { + if (runError === null) { + throw new Error("Failed Session Run is missing its authoritative durable error."); + } + repairs.push({ ...common, terminal: { error: runError, status: "failed" as const } }); + continue; + } + + throw new Error("Input command repair requires an authoritative terminal Session Run."); + } + return repairs; +} + +/** Returns canonical repairs that still need a durable terminal tool event. */ +export async function listAcceptedMcpCommandRepairsForTerminalDriver( + database: D1Database, + input: { + driverGeneration: number; + driverInstanceId: DriverInstanceId; + }, +): Promise { + const rows = await getAppDatabase(database) + .select({ + commandId: driverCommandsTable.id, + effectId: externalToolEffectsTable.id, + effectStatus: externalToolEffectsTable.status, + payloadJson: driverCommandsTable.payloadJson, + resultJson: externalToolEffectsTable.resultJson, + runtimeId: sql`COALESCE(${sessionRunsTable.runtimeId}, ${sessionsTable.runtimeId})`, + sessionId: sessionRunsTable.sessionId, }) + .from(driverCommandsTable) + .innerJoin( + externalToolEffectsTable, + and( + eq(externalToolEffectsTable.commandId, driverCommandsTable.id), + eq(externalToolEffectsTable.driverInstanceId, driverCommandsTable.driverInstanceId), + ), + ) + .innerJoin(sessionRunsTable, eq(sessionRunsTable.id, externalToolEffectsTable.sessionRunId)) + .innerJoin(sessionsTable, eq(sessionsTable.id, sessionRunsTable.sessionId)) .where( and( eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + eq(driverCommandsTable.kind, "mcp.execute"), eq(driverCommandsTable.status, "accepted"), + sql`json_extract(${driverCommandsTable.payloadJson}, '$.runId') = ${externalToolEffectsTable.sessionRunId}`, + sql`json_extract(${driverCommandsTable.payloadJson}, '$.serverId') = ${externalToolEffectsTable.serverId}`, + sql`json_extract(${driverCommandsTable.payloadJson}, '$.toolName') = ${externalToolEffectsTable.toolName}`, ), ) - .run(); + .orderBy(asc(driverCommandsTable.seq)) + .all(); - return createRuntimeCommandBatchTransitionOutcome(targetStatus, getD1ChangeCount(result)); + const repairs: AcceptedMcpCommandRepair[] = []; + for (const row of rows) { + if (row.runtimeId.length === 0) { + throw new Error("MCP command repair is missing its durable completion identity."); + } + + const command = parseSchemaValue(RuntimeCommand, JSON.parse(row.payloadJson)); + if (command.kind !== "mcp.execute" || command.commandId !== row.commandId) { + throw new Error("MCP command repair does not match its immutable command intent."); + } + + if (row.effectStatus === "claimed") { + throw new Error("Terminal Driver repair must fence claimed MCP effects before listing them."); + } + + const common = { + command, + commandId: row.commandId, + effectId: row.effectId, + effectStatus: row.effectStatus, + runtimeId: row.runtimeId, + sessionId: row.sessionId, + }; + + if (row.effectStatus === "succeeded") { + if (row.resultJson === null) { + throw new Error("Succeeded MCP command is missing its durable result."); + } + const result = parseSchemaValue(McpExecuteCommandResult, JSON.parse(row.resultJson)); + if ( + result.requestId !== command.requestId || + result.serverId !== command.serverId || + result.toolName !== command.toolName + ) { + throw new Error("Succeeded MCP result does not match its immutable command intent."); + } + repairs.push({ ...common, terminal: { result, status: "completed" as const } }); + continue; + } + + if (row.resultJson !== null) { + throw new Error("Unresolved MCP command unexpectedly contains a durable result."); + } + const error = + row.effectStatus === "unknown" + ? createRuntimeCommandExternalToolEffectUnknownError({ + command, + effectId: row.effectId, + }) + : createRuntimeCommandExternalToolEffectNotExecutedError({ + commandId: row.commandId, + driverInstanceId: input.driverInstanceId, + effectId: row.effectId, + }); + repairs.push({ ...common, terminal: { error, status: "failed" as const } }); + } + return repairs; } -async function failAcceptedRuntimeCommandsForTerminalDriversGlobally( +export async function repairAcceptedRuntimeCommandsForTerminalDriver( database: D1Database, - nowMs: number, -): Promise { - const targetStatus = "failed" satisfies RuntimeCommandStatus; - const error = createRuntimeCommandDriverTerminalError({}); + input: { + driverGeneration: number; + driverInstanceId: DriverInstanceId; + nowMs?: number; + }, +): Promise { + const nowMs = input.nowMs ?? currentTimestampMs(); const db = getAppDatabase(database); - const terminalDriverQuery = db - .select({ id: driverInstancesTable.id }) - .from(driverInstancesTable) + const unacceptedIntents = await db + .select({ + commandId: driverCommandsTable.id, + effectId: externalToolEffectsTable.id, + }) + .from(driverCommandsTable) + .innerJoin( + externalToolEffectsTable, + and( + eq(externalToolEffectsTable.commandId, driverCommandsTable.id), + eq(externalToolEffectsTable.driverInstanceId, driverCommandsTable.driverInstanceId), + ), + ) .where( and( - eq(driverInstancesTable.id, driverCommandsTable.driverInstanceId), - inArray(driverInstancesTable.status, ["failed", "stopped"]), + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + inArray(driverCommandsTable.status, ["queued", "delivered"]), + eq(externalToolEffectsTable.status, "intent"), ), - ); + ) + .all(); - const result = await db + let failedCount = 0; + for (const effect of unacceptedIntents) { + const result = await db + .update(driverCommandsTable) + .set({ + completedAt: sql`COALESCE(${driverCommandsTable.completedAt}, ${nowMs})`, + errorJson: JSON.stringify( + createRuntimeCommandExternalToolEffectNotExecutedError({ + commandId: effect.commandId, + driverInstanceId: input.driverInstanceId, + effectId: effect.effectId, + }), + ), + status: "failed", + }) + .where( + and( + eq(driverCommandsTable.id, effect.commandId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + inArray(driverCommandsTable.status, ["queued", "delivered"]), + exists( + db + .select({ id: externalToolEffectsTable.id }) + .from(externalToolEffectsTable) + .where( + and( + eq(externalToolEffectsTable.id, effect.effectId), + eq(externalToolEffectsTable.commandId, effect.commandId), + eq(externalToolEffectsTable.driverInstanceId, input.driverInstanceId), + eq(externalToolEffectsTable.status, "intent"), + ), + ), + ), + ), + ) + .run(); + failedCount += getD1ChangeCount(result); + } + + const pendingEventFirstCommand = + (await db + .select({ id: driverCommandsTable.id }) + .from(driverCommandsTable) + .where( + and( + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + or( + and( + eq(driverCommandsTable.kind, "mcp.execute"), + inArray(driverCommandsTable.status, ["queued", "delivered", "accepted"]), + ), + and( + eq(driverCommandsTable.kind, "input.start"), + eq(driverCommandsTable.status, "accepted"), + ), + ), + ), + ) + .limit(1) + .get()) ?? null; + if (pendingEventFirstCommand !== null) { + throw new Error("Accepted input and MCP commands require event-first terminal reconciliation."); + } + + const genericResult = await db .update(driverCommandsTable) .set({ completedAt: sql`COALESCE(${driverCommandsTable.completedAt}, ${nowMs})`, - errorJson: JSON.stringify(error), - status: targetStatus, + errorJson: JSON.stringify( + createRuntimeCommandDriverTerminalError({ + driverInstanceId: input.driverInstanceId, + }), + ), + status: "failed", }) - .where(and(eq(driverCommandsTable.status, "accepted"), exists(terminalDriverQuery))) + .where( + and( + eq(driverCommandsTable.driverInstanceId, input.driverInstanceId), + eq(driverCommandsTable.driverGeneration, input.driverGeneration), + ne(driverCommandsTable.kind, "mcp.execute"), + or( + inArray(driverCommandsTable.status, ["queued", "delivered"]), + and( + eq(driverCommandsTable.status, "accepted"), + ne(driverCommandsTable.kind, "input.start"), + ), + ), + ), + ) .run(); + failedCount += getD1ChangeCount(genericResult); - return createRuntimeCommandBatchTransitionOutcome(targetStatus, getD1ChangeCount(result)); + return { + completed: createRuntimeCommandBatchTransitionOutcome("completed", 0), + failed: createRuntimeCommandBatchTransitionOutcome("failed", failedCount), + }; +} + +export async function listTerminalDriversWithPendingRuntimeCommands( + database: D1Database, +): Promise { + const db = getAppDatabase(database); + const pendingCommand = db + .select({ id: driverCommandsTable.id }) + .from(driverCommandsTable) + .where( + and( + eq(driverCommandsTable.driverInstanceId, driverInstancesTable.id), + eq(driverCommandsTable.driverGeneration, driverInstancesTable.generation), + inArray(driverCommandsTable.status, ["queued", "delivered", "accepted"]), + ), + ); + const claimedEffect = db + .select({ id: externalToolEffectsTable.id }) + .from(externalToolEffectsTable) + .where( + and( + eq(externalToolEffectsTable.driverInstanceId, driverInstancesTable.id), + eq(externalToolEffectsTable.status, "claimed"), + exists( + db + .select({ id: driverCommandsTable.id }) + .from(driverCommandsTable) + .where( + and( + eq(driverCommandsTable.id, externalToolEffectsTable.commandId), + eq(driverCommandsTable.driverGeneration, driverInstancesTable.generation), + ), + ), + ), + ), + ); + + return db + .select({ generation: driverInstancesTable.generation, id: driverInstancesTable.id }) + .from(driverInstancesTable) + .where( + and( + inArray(driverInstancesTable.status, ["failed", "stopped"]), + or(exists(pendingCommand), exists(claimedEffect)), + ), + ) + .all(); } export async function maintainRuntimeCommandRecords( database: D1Database, input: { connectionId: string; + driverGeneration: number; driverInstanceId: DriverInstanceId; nowMs?: number; }, @@ -693,9 +1834,15 @@ export async function maintainRuntimeCommandRecords( const recovered = await recoverRuntimeCommandsDeliveredToStaleConnections(database, { connectionId: input.connectionId, + driverGeneration: input.driverGeneration, driverInstanceId: input.driverInstanceId, }); - const expired = await expireRuntimeCommandDeliveryLeases(database, input.driverInstanceId, nowMs); + const expired = await expireRuntimeCommandDeliveryLeases( + database, + input.driverInstanceId, + input.driverGeneration, + nowMs, + ); return { expired, @@ -708,16 +1855,14 @@ export async function repairRuntimeCommandRecords( input: { nowMs?: number; } = {}, -): Promise { +): Promise { const nowMs = input.nowMs ?? currentTimestampMs(); const recovered = await recoverRuntimeCommandsDeliveredToStaleConnectionsGlobally(database); const expired = await expireRuntimeCommandDeliveryLeasesGlobally(database, nowMs); - const failed = await failAcceptedRuntimeCommandsForTerminalDriversGlobally(database, nowMs); return { expired, - failed, recovered, }; } @@ -725,6 +1870,7 @@ export async function repairRuntimeCommandRecords( export async function claimNextQueuedRuntimeCommandRecord( database: D1Database, driverInstanceId: DriverInstanceId, + driverGeneration: number, connectionId: string, ): Promise { const nowMs = currentTimestampMs(); @@ -732,6 +1878,7 @@ export async function claimNextQueuedRuntimeCommandRecord( await maintainRuntimeCommandRecords(database, { connectionId, + driverGeneration, driverInstanceId, nowMs, }); @@ -744,7 +1891,9 @@ export async function claimNextQueuedRuntimeCommandRecord( .where( and( eq(driverCommandsTable.driverInstanceId, driverInstanceId), + eq(driverCommandsTable.driverGeneration, driverGeneration), eq(driverCommandsTable.status, "queued"), + runtimeCommandActiveRunGuard(db), or(isNull(driverCommandsTable.expiresAt), gt(driverCommandsTable.expiresAt, nowMs)), ), ) @@ -759,6 +1908,7 @@ export async function claimNextQueuedRuntimeCommandRecord( const deliveryOutcome = await markRuntimeCommandRecordDelivered(database, { commandId: nextQueued.id, connectionId, + driverGeneration, driverInstanceId, expiresAfter: nowMs, }); @@ -772,7 +1922,7 @@ export async function claimNextQueuedRuntimeCommandRecord( const claimed = deliveryOutcome.kind === "applied" - ? await getRuntimeCommandRecord(database, driverInstanceId, nextQueued.id) + ? await getRuntimeCommandRecord(database, driverInstanceId, driverGeneration, nextQueued.id) : null; if (claimed !== null) { diff --git a/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-transition.ts b/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-transition.ts index b31f3fb7..5796776b 100644 --- a/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-transition.ts +++ b/apps/api/src/modules/runtime/infrastructure/session-runs/runtime-command-transition.ts @@ -16,6 +16,7 @@ export type RuntimeCommandTransitionOutcome = | "command_not_found" | "illegal_transition" | "inactive_delivery_connection" + | "inactive_session_run" | "stale_delivery_connection"; targetStatus: RuntimeCommandStatus; }; @@ -32,7 +33,7 @@ const previousStatusesByTarget = { completed: ["delivered", "accepted"], delivered: ["queued"], expired: ["queued", "delivered", "accepted"], - failed: ["delivered", "accepted"], + failed: ["queued", "delivered", "accepted"], queued: ["delivered"], } as const satisfies Record; diff --git a/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-admission.repository.ts b/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-admission.repository.ts index d00162aa..1d06d123 100644 --- a/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-admission.repository.ts +++ b/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-admission.repository.ts @@ -18,6 +18,7 @@ import type { SessionMessageId, SessionRunId, } from "@mosoo/id"; +import { createRuntimeEventSemanticHash } from "@mosoo/runtime-events"; import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; import { and, eq, exists, inArray, isNull, ne, notExists, or, sql } from "drizzle-orm"; import type { SQL } from "drizzle-orm"; @@ -57,6 +58,10 @@ interface QueuedMessageAdmissionRecord { timestampMs: number; } +type SourcedRuntimeEventEnvelope = RuntimeEventEnvelope & { + readonly sourceEventId: string; +}; + export interface CommitQueuedSessionRunAdmissionInput { apiCommand: PreparedApiCommand; clientRequestId: string | null; @@ -82,9 +87,50 @@ function admissionSessionPredicate(input: CommitQueuedSessionRunAdmissionInput) eq(sessionsTable.appId, input.session.appId), eq(sessionsTable.lastRunId, input.run.id), eq(sessionsTable.status, "RUNNING"), + isNull(sessionsTable.runtimeProvisioningOperationId), ); } +export function cattleTerminalCheckpointReadyPredicate(db: AppDatabase): SQL { + return or( + ne(sessionsTable.kind, "cattle"), + eq(sessionsTable.workspaceCheckpointRequired, false), + isNull(sessionsTable.lastRunId), + notExists( + db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, sessionsTable.lastRunId), + eq(sessionRunsTable.status, "completed"), + ), + ), + ), + exists( + db + .select({ id: sandboxBackupsTable.id }) + .from(sandboxBackupsTable) + .innerJoin( + sandboxSessionsTable, + and( + eq(sandboxSessionsTable.sessionId, sessionsTable.id), + eq(sandboxSessionsTable.sandboxId, sandboxBackupsTable.sandboxId), + eq(sandboxSessionsTable.sandboxIncarnation, sandboxBackupsTable.sandboxIncarnation), + eq(sandboxSessionsTable.cwd, sandboxBackupsTable.dir), + eq(sandboxBackupsTable.workspaceSessionId, sandboxSessionsTable.sessionId), + ), + ) + .where( + and( + eq(sandboxBackupsTable.sessionRunId, sessionsTable.lastRunId), + eq(sandboxBackupsTable.status, "ready"), + ), + ), + ), + )!; +} + function claimableSessionPredicate(db: AppDatabase, input: CommitQueuedSessionRunAdmissionInput) { return and( eq(sessionsTable.id, input.session.id), @@ -93,41 +139,8 @@ function claimableSessionPredicate(db: AppDatabase, input: CommitQueuedSessionRu isNull(sessionsTable.archivedAt), eq(sessionsTable.status, "IDLE"), isNull(sessionsTable.statusOperationId), - or( - ne(sessionsTable.kind, "cattle"), - eq(sessionsTable.workspaceCheckpointRequired, false), - isNull(sessionsTable.lastRunId), - notExists( - db - .select({ id: sessionRunsTable.id }) - .from(sessionRunsTable) - .where( - and( - eq(sessionRunsTable.id, sessionsTable.lastRunId), - eq(sessionRunsTable.status, "completed"), - ), - ), - ), - exists( - db - .select({ id: sandboxBackupsTable.id }) - .from(sandboxBackupsTable) - .innerJoin( - sandboxSessionsTable, - and( - eq(sandboxSessionsTable.sessionId, sessionsTable.id), - eq(sandboxSessionsTable.sandboxId, sandboxBackupsTable.sandboxId), - eq(sandboxSessionsTable.cwd, sandboxBackupsTable.dir), - ), - ) - .where( - and( - eq(sandboxBackupsTable.sessionRunId, sessionsTable.lastRunId), - eq(sandboxBackupsTable.status, "ready"), - ), - ), - ), - ), + isNull(sessionsTable.runtimeProvisioningOperationId), + cattleTerminalCheckpointReadyPredicate(db), notExists( db .select({ id: sessionRunsTable.id }) @@ -161,52 +174,21 @@ export async function isCattleTerminalCheckpointReadyForNextRun( sessionId: SessionId, ): Promise { const appDb = getAppDatabase(database); - const session = - (await appDb - .select({ - kind: sessionsTable.kind, - lastRunId: sessionsTable.lastRunId, - lastRunStatus: sessionRunsTable.status, - workspaceCheckpointRequired: sessionsTable.workspaceCheckpointRequired, - }) + const [session, ready] = await Promise.all([ + appDb + .select({ id: sessionsTable.id }) .from(sessionsTable) - .leftJoin(sessionRunsTable, eq(sessionRunsTable.id, sessionsTable.lastRunId)) .where(eq(sessionsTable.id, sessionId)) .limit(1) - .get()) ?? null; - - if ( - session === null || - session.kind !== "cattle" || - !session.workspaceCheckpointRequired || - session.lastRunId === null || - session.lastRunStatus !== "completed" - ) { - return true; - } - - const checkpoint = - (await appDb - .select({ id: sandboxBackupsTable.id }) - .from(sandboxBackupsTable) - .innerJoin( - sandboxSessionsTable, - and( - eq(sandboxSessionsTable.sessionId, sessionId), - eq(sandboxSessionsTable.sandboxId, sandboxBackupsTable.sandboxId), - eq(sandboxSessionsTable.cwd, sandboxBackupsTable.dir), - ), - ) - .where( - and( - eq(sandboxBackupsTable.sessionRunId, session.lastRunId), - eq(sandboxBackupsTable.status, "ready"), - ), - ) + .get(), + appDb + .select({ id: sessionsTable.id }) + .from(sessionsTable) + .where(and(eq(sessionsTable.id, sessionId), cattleTerminalCheckpointReadyPredicate(appDb))) .limit(1) - .get()) ?? null; - - return checkpoint !== null; + .get(), + ]); + return session === undefined || ready !== undefined; } export async function hasSessionRunAdmissionClientRequestReceipt( @@ -274,6 +256,7 @@ function createRunInsertQuery(db: AppDatabase, input: CommitQueuedSessionRunAdmi errorCode: selectedValue(null, "error_code"), errorDetailsJson: selectedValue(null, "error_details_json"), errorMessage: selectedValue(null, "error_message"), + errorRetryable: selectedValue(null, "error_retryable"), id: selectedValue(input.run.id, "id"), model: selectedValue(input.run.model, "model"), provider: selectedValue(input.run.provider, "provider"), @@ -286,6 +269,10 @@ function createRunInsertQuery(db: AppDatabase, input: CommitQueuedSessionRunAdmi statusOperationId: selectedValue(null, "status_operation_id"), statusSeq: selectedValue(0, "status_seq"), statusSource: selectedValue("api", "status_source"), + terminalReconciliationAttemptedAt: selectedValue( + null, + "terminal_reconciliation_attempted_at", + ), traceId: selectedValue(input.run.traceId, "trace_id"), trigger: selectedValue(input.run.trigger, "trigger"), updatedAt: selectedValue(input.run.timestampMs, "updated_at"), @@ -307,6 +294,7 @@ function createMessageInsertQuery(db: AppDatabase, input: CommitQueuedSessionRun ), id: selectedValue(input.message.id, "id"), planJson: selectedValue(null, "plan_json"), + projectionFormat: selectedValue("materialized" as const, "projection_format"), role: selectedValue("user" as const, "role"), segmentsJson: selectedValue(null, "segments_json"), seq: sessionsTable.messageSeqCursor, @@ -333,41 +321,74 @@ function runtimeEventOccurredAt(event: RuntimeEventEnvelope, fallbackMs: number) return Number.isFinite(occurredAt) ? occurredAt : fallbackMs; } +function sourceAdmissionEvent( + event: RuntimeEventEnvelope, + index: number, + clientRequestId: string | null, +): SourcedRuntimeEventEnvelope { + return { + ...event, + sourceEventId: event.sourceEventId ?? (index === 0 ? clientRequestId : null) ?? event.id, + }; +} + function createEventInsertQuery( db: AppDatabase, input: CommitQueuedSessionRunAdmissionInput, - event: RuntimeEventEnvelope, + event: SourcedRuntimeEventEnvelope, index: number, + semanticHash: string, ) { const projection = createSessionRuntimeEventProjection(event); const timestampMs = input.run.timestampMs + index; const occurredAt = runtimeEventOccurredAt(event, timestampMs); - const sourceEventId = - event.sourceEventId ?? (index === 0 ? input.clientRequestId : null) ?? event.id; return db.insert(sessionEventsTable).select( db .select({ agentId: sessionsTable.agentId, + artifactAttemptId: selectedValue(null, "artifact_attempt_id"), + artifactManifestJson: selectedValue(null, "artifact_manifest_json"), + artifactManifestSha256: selectedValue(null, "artifact_manifest_sha256"), contentText: selectedValue(projection.contentText, "content_text"), createdAt: selectedValue(timestampMs, "created_at"), endedAt: selectedValue(Math.max(occurredAt, timestampMs), "ended_at"), eventType: selectedValue(projection.eventType, "event_type"), family: selectedValue(projection.family, "family"), id: selectedValue(event.id, "id"), + mcpCommandId: selectedValue(projection.mcpCommandId, "mcp_command_id"), occurredAt: selectedValue(occurredAt, "occurred_at"), processStatus: selectedValue(projection.processStatus, "process_status"), processType: selectedValue(projection.processType, "process_type"), runId: selectedValue(projection.runId, "run_id"), + runtimeOperationEventJson: selectedValue(null, "runtime_operation_event_json"), + semanticHash: selectedValue(semanticHash, "semantic_hash"), seq: sql`${sessionsTable.runtimeEventSeqCursor} - ${input.events.length - index - 1}`.as( "seq", ), sessionId: sessionsTable.id, - sourceEventId: selectedValue(sourceEventId, "source_event_id"), + sourceEventId: selectedValue(event.sourceEventId, "source_event_id"), source: selectedValue(projection.source, "source"), + streamId: selectedValue(projection.streamId, "stream_id"), + terminalEventJson: selectedValue(null, "terminal_event_json"), toolCallId: selectedValue(projection.toolCallId, "tool_call_id"), + toolInputDeltaJson: selectedValue(projection.toolInputDeltaJson, "tool_input_delta_json"), toolInputJson: selectedValue(projection.toolInputJson, "tool_input_json"), toolName: selectedValue(projection.toolName, "tool_name"), + toolOutputDeltaText: selectedValue( + projection.toolOutputDeltaText, + "tool_output_delta_text", + ), + toolOutputText: selectedValue(projection.toolOutputText, "tool_output_text"), + toolParentMessageId: selectedValue( + projection.toolParentMessageId, + "tool_parent_message_id", + ), + toolResultMessageId: selectedValue( + projection.toolResultMessageId, + "tool_result_message_id", + ), + toolStatus: selectedValue(projection.toolStatus, "tool_status"), tokens: selectedValue(projection.tokens, "tokens"), traceId: selectedValue(projection.traceId, "trace_id"), visibility: selectedValue(projection.visibility, "visibility"), @@ -399,6 +420,7 @@ function createApiCommandInsertQuery(db: AppDatabase, input: CommitQueuedSession completedAt: selectedValue(record.completedAt, "completed_at"), createdAt: selectedValue(record.createdAt, "created_at"), dedupeKey: selectedValue(record.dedupeKey, "dedupe_key"), + deliveryGeneration: selectedValue(record.deliveryGeneration, "delivery_generation"), id: selectedValue(record.id, "id"), kind: selectedValue(record.kind, "kind"), lastErrorCode: selectedValue(record.lastErrorCode, "last_error_code"), @@ -436,42 +458,58 @@ export async function commitQueuedSessionRunAdmission( } } - const results = await runAppDatabaseBatch(database, (db) => [ - createRunInsertQuery(db, input), - db - .update(sessionsTable) - .set({ - lastMessageAt: input.message.timestampMs, - lastRunId: input.run.id, - messageSeqCursor: sql`${sessionsTable.messageSeqCursor} + 1`, - model: sql`COALESCE(${input.run.model}, ${sessionsTable.model})`, - provider: sql`COALESCE(${input.run.provider}, ${sessionsTable.provider})`, - runtimeEventSeqCursor: sql`${sessionsTable.runtimeEventSeqCursor} + ${input.events.length}`, - ...createSessionStatusTransitionPatch({ - status: "RUNNING", - timestampMs: input.run.timestampMs, - }), - }) - .where( - and( - eq(sessionsTable.id, input.session.id), - eq(sessionsTable.agentId, input.session.agentId), - eq(sessionsTable.appId, input.session.appId), - isNull(sessionsTable.archivedAt), - eq(sessionsTable.status, "IDLE"), - isNull(sessionsTable.statusOperationId), - exists( - db - .select({ id: sessionRunsTable.id }) - .from(sessionRunsTable) - .where(eq(sessionRunsTable.id, input.run.id)), + const events = input.events.map((event, index) => + sourceAdmissionEvent(event, index, input.clientRequestId), + ); + const semanticHashes = await Promise.all(events.map(createRuntimeEventSemanticHash)); + + const results = await runAppDatabaseBatch(database, (db) => { + const eventInserts = events.map((event, index) => { + const semanticHash = semanticHashes[index]; + if (semanticHash === undefined) { + throw new Error("Queued Session Run admission event hash is missing."); + } + return createEventInsertQuery(db, input, event, index, semanticHash); + }); + + return [ + createRunInsertQuery(db, input), + db + .update(sessionsTable) + .set({ + lastMessageAt: input.message.timestampMs, + lastRunId: input.run.id, + messageSeqCursor: sql`${sessionsTable.messageSeqCursor} + 1`, + model: sql`COALESCE(${input.run.model}, ${sessionsTable.model})`, + provider: sql`COALESCE(${input.run.provider}, ${sessionsTable.provider})`, + runtimeEventSeqCursor: sql`${sessionsTable.runtimeEventSeqCursor} + ${input.events.length}`, + ...createSessionStatusTransitionPatch({ + status: "RUNNING", + timestampMs: input.run.timestampMs, + }), + }) + .where( + and( + eq(sessionsTable.id, input.session.id), + eq(sessionsTable.agentId, input.session.agentId), + eq(sessionsTable.appId, input.session.appId), + isNull(sessionsTable.archivedAt), + eq(sessionsTable.status, "IDLE"), + isNull(sessionsTable.statusOperationId), + isNull(sessionsTable.runtimeProvisioningOperationId), + exists( + db + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where(eq(sessionRunsTable.id, input.run.id)), + ), ), ), - ), - createMessageInsertQuery(db, input), - ...input.events.map((event, index) => createEventInsertQuery(db, input, event, index)), - createApiCommandInsertQuery(db, input), - ]); + createMessageInsertQuery(db, input), + ...eventInserts, + createApiCommandInsertQuery(db, input), + ]; + }); return getD1ChangeCount((results as readonly unknown[])[0]) > 0; } diff --git a/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-read.repository.ts b/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-read.repository.ts index ac084e40..e62ca9b8 100644 --- a/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-read.repository.ts +++ b/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-read.repository.ts @@ -23,6 +23,7 @@ function sessionRunSummaryColumns() { error_code: sessionRunsTable.errorCode, error_details_json: sessionRunsTable.errorDetailsJson, error_message: sessionRunsTable.errorMessage, + error_retryable: sessionRunsTable.errorRetryable, id: sessionRunsTable.id, model: sessionRunsTable.model, provider: sessionRunsTable.provider, diff --git a/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-row.mapper.ts b/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-row.mapper.ts index d8ac7393..ef534c10 100644 --- a/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-row.mapper.ts +++ b/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-row.mapper.ts @@ -1,5 +1,5 @@ +import { DurableRunError } from "@mosoo/contracts/session-run"; import type { - RunError, SessionRunStatus, SessionRunSummary, SessionRunTrigger, @@ -27,6 +27,7 @@ export interface SessionRunRow { error_code: string | null; error_details_json: string | null; error_message: string | null; + error_retryable?: boolean | null; id: SessionRunId; model: string | null; provider: string | null; @@ -86,7 +87,7 @@ function parseRunErrorDetailsJson(raw: string): unknown { } } -function toRunError(row: SessionRunRow): RunError | null { +function toRunError(row: SessionRunRow): DurableRunError | null { if ( row.error_code === null || row.error_code === "" || @@ -96,10 +97,10 @@ function toRunError(row: SessionRunRow): RunError | null { return null; } - return { + return parseSchemaValue(DurableRunError, { code: row.error_code, details: parseJsonRecord(row.error_details_json), message: row.error_message, - retryable: false, - }; + retryable: row.error_retryable ?? false, + }); } diff --git a/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-store.repository.ts b/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-store.repository.ts index 7f190df5..f1ceccd6 100644 --- a/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-store.repository.ts +++ b/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-store.repository.ts @@ -5,7 +5,6 @@ export { hasActiveSessionRun, } from "./session-run-read.repository"; export { - cancelActiveSessionRunsForRuntimeOperation, createSessionRunRecordIfSessionIdle, SessionRunCreationGuardRejectedError, setSessionRunStatus, diff --git a/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-write.repository.ts b/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-write.repository.ts index ce1e4c3f..b71d48f4 100644 --- a/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-write.repository.ts +++ b/apps/api/src/modules/runtime/infrastructure/session-runs/session-run-write.repository.ts @@ -1,7 +1,5 @@ -import type { AgentKind } from "@mosoo/contracts/agent"; -import type { SessionStatus, SessionType } from "@mosoo/contracts/session"; +import type { SessionStatus } from "@mosoo/contracts/session"; import type { - RunError, SessionRunStatus, SessionRunSummary, SessionRunTrigger, @@ -17,10 +15,9 @@ import type { SessionRunId, } from "@mosoo/id"; import { generateTraceId } from "@mosoo/observability"; -import { and, eq, exists, inArray, notInArray, sql } from "drizzle-orm"; +import { and, eq, exists, isNull, notInArray, sql } from "drizzle-orm"; import type { SQL } from "drizzle-orm"; -import { createErrorLogContext, logInfo, logWarn } from "../../../../platform/cloudflare/logger"; import { getAppDatabase, getD1ChangeCount, @@ -30,7 +27,6 @@ import { currentTimestampMs, toIsoString } from "../../../../time"; import { toSessionLifecycleStatusForRunStatus } from "../../../sessions/domain/session-lifecycle"; import type { BoundCapabilityRunProvenance } from "../../domain/bound-capability-run-provenance"; import { - ACTIVE_SESSION_RUN_STATUSES, decideSessionRunTransition, isTerminalSessionRunStatus, toSessionRunStatusLifecycleEventName, @@ -42,10 +38,8 @@ import type { ActiveSessionRunStatus } from "./session-run-row.mapper"; import { updateSessionLastRun } from "./session-run-session.repository"; type SessionRunStatusUpdateInput = { - error?: RunError | null; - operationId?: RuntimeOperationId | null; source?: SessionRunTransitionSource; - status: SessionRunStatus; + status: ActiveSessionRunStatus; }; type UpdateSessionRunStatusInput = SessionRunStatusUpdateInput & { @@ -54,7 +48,6 @@ type UpdateSessionRunStatusInput = SessionRunStatusUpdateInput & { * check is atomic with the write via the status_seq optimistic guard. */ expectedCurrentStatus?: SessionRunStatus; - preserveSessionLifecycle?: boolean; runId: SessionRunId; }; @@ -77,8 +70,6 @@ type SessionRunTransitionSource = | "system" | "viewer"; -const SESSION_RUN_STATUS_WRITE_BATCH_SIZE = 50; - /** * The caller supplied an atomic predicate for Run creation and it no longer * held when the INSERT statement executed. No Run record was created. @@ -98,15 +89,19 @@ interface LoadedSessionRunLifecycleRow { error_code: string | null; error_details_json: string | null; error_message: string | null; + error_retryable: boolean | null; id: SessionRunId; model: string | null; provider: string | null; runtime_id: string; + session_archived_at: number | null; + session_cleanup_operation_kind: "archive" | "delete" | null; session_id: SessionId; - session_kind: AgentKind; session_last_run_id: SessionRunId | null; session_status: SessionStatus; - session_type: SessionType; + session_status_operation_id: RuntimeOperationId | null; + session_status_seq: number; + session_updated_at: number; started_at: number | null; status: SessionRunStatus; status_seq: number; @@ -120,7 +115,7 @@ export type SessionRunTransitionOutcome = kind: "applied"; previousStatus: SessionRunStatus; run: SessionRunSummary; - sessionLifecycle: "current_last_run_updated" | "not_current_last_run" | "preserved"; + sessionLifecycle: "current_last_run_updated" | "not_current_last_run"; statusSeq: number; } | { @@ -154,17 +149,13 @@ export type SessionRunTransitionOutcome = statusSeq: number; }; -export interface NonTerminalSessionRunsStatusUpdateResult { - readonly runIds: readonly SessionRunId[]; - readonly timestampMs: number; -} - function createSessionRunStatusUpdate(input: SessionRunStatusUpdateInput, timestampMs: number) { return { - completedAt: isTerminalSessionRunStatus(input.status) ? timestampMs : undefined, - errorCode: input.error?.code ?? null, - errorDetailsJson: input.error ? JSON.stringify(input.error.details) : null, - errorMessage: input.error?.message ?? null, + completedAt: undefined, + errorCode: null, + errorDetailsJson: null, + errorMessage: null, + errorRetryable: null, startedAt: input.status === "queued" ? undefined @@ -172,7 +163,7 @@ function createSessionRunStatusUpdate(input: SessionRunStatusUpdateInput, timest status: input.status, statusChangedAt: timestampMs, statusEvent: toSessionRunStatusLifecycleEventName(input.status), - statusOperationId: input.operationId ?? null, + statusOperationId: null, statusSeq: sql`${sessionRunsTable.statusSeq} + 1`, statusSource: input.source ?? "system", updatedAt: timestampMs, @@ -180,51 +171,13 @@ function createSessionRunStatusUpdate(input: SessionRunStatusUpdateInput, timest } function createCurrentSessionRunProjectionPatch(input: { - readonly sessionKind: AgentKind; - readonly status: SessionRunStatus; + readonly status: ActiveSessionRunStatus; readonly timestampMs: number; }) { - return { - ...createSessionStatusTransitionPatch({ - status: toSessionLifecycleStatusForRunStatus(input.status), - timestampMs: input.timestampMs, - }), - ...(input.sessionKind === "cattle" && input.status === "completed" - ? { workspaceCheckpointRequired: true } - : {}), - }; -} - -function logTerminalSessionRun( - current: LoadedSessionRunLifecycleRow, - input: SessionRunStatusUpdateInput, - timestampMs: number, -): void { - if (!isTerminalSessionRunStatus(input.status)) return; - - try { - logInfo("session.run.terminal", { - durationMs: Math.max(0, timestampMs - (current.started_at ?? current.created_at)), - endToEndMs: Math.max(0, timestampMs - current.created_at), - errorCode: input.error?.code ?? null, - runId: current.id, - runtimeId: current.runtime_id, - sessionType: current.session_type, - source: input.source ?? "system", - status: input.status, - traceId: current.trace_id, - trigger: current.trigger, - }); - } catch (error) { - try { - logWarn("session.run.terminal_log.failed", { - ...createErrorLogContext(error), - runId: current.id, - }); - } catch { - // Observability must never turn a committed Run transition into an API failure. - } - } + return createSessionStatusTransitionPatch({ + status: toSessionLifecycleStatusForRunStatus(input.status), + timestampMs: input.timestampMs, + }); } function applySessionRunStatusUpdate( @@ -234,17 +187,8 @@ function applySessionRunStatusUpdate( ): SessionRunSummary { return { ...run, - completedAt: isTerminalSessionRunStatus(input.status) - ? toIsoString(timestampMs) - : run.completedAt, - error: input.error - ? { - code: input.error.code, - details: input.error.details, - message: input.error.message, - retryable: input.error.retryable, - } - : null, + completedAt: run.completedAt, + error: null, startedAt: input.status === "queued" ? run.startedAt : (run.startedAt ?? toIsoString(timestampMs)), status: input.status, @@ -261,15 +205,19 @@ function sessionRunLifecycleColumns() { error_code: sessionRunsTable.errorCode, error_details_json: sessionRunsTable.errorDetailsJson, error_message: sessionRunsTable.errorMessage, + error_retryable: sessionRunsTable.errorRetryable, id: sessionRunsTable.id, model: sessionRunsTable.model, provider: sessionRunsTable.provider, runtime_id: sessionsTable.runtimeId, + session_archived_at: sessionsTable.archivedAt, + session_cleanup_operation_kind: sessionsTable.cleanupOperationKind, session_id: sessionRunsTable.sessionId, - session_kind: sessionsTable.kind, session_last_run_id: sessionsTable.lastRunId, session_status: sessionsTable.status, - session_type: sessionsTable.type, + session_status_operation_id: sessionsTable.statusOperationId, + session_status_seq: sessionsTable.statusSeq, + session_updated_at: sessionsTable.updatedAt, started_at: sessionRunsTable.startedAt, status: sessionRunsTable.status, status_seq: sessionRunsTable.statusSeq, @@ -279,6 +227,22 @@ function sessionRunLifecycleColumns() { }; } +function writableObservedSessionCondition(current: LoadedSessionRunLifecycleRow) { + return and( + eq(sessionsTable.id, current.session_id), + current.session_last_run_id === null + ? isNull(sessionsTable.lastRunId) + : eq(sessionsTable.lastRunId, current.session_last_run_id), + eq(sessionsTable.status, current.session_status), + notInArray(sessionsTable.status, ["TERMINATED"]), + eq(sessionsTable.statusSeq, current.session_status_seq), + eq(sessionsTable.updatedAt, current.session_updated_at), + isNull(sessionsTable.archivedAt), + isNull(sessionsTable.cleanupOperationKind), + isNull(sessionsTable.statusOperationId), + ); +} + function toSessionRunSummaryFromLifecycleRow(row: LoadedSessionRunLifecycleRow): SessionRunSummary { return toSessionRunSummary({ completed_at: row.completed_at, @@ -288,6 +252,7 @@ function toSessionRunSummaryFromLifecycleRow(row: LoadedSessionRunLifecycleRow): error_code: row.error_code, error_details_json: row.error_details_json, error_message: row.error_message, + error_retryable: row.error_retryable, id: row.id, model: row.model, provider: row.provider, @@ -324,6 +289,7 @@ export function createInsertedSessionRunSummary( error_code: null, error_details_json: null, error_message: null, + error_retryable: null, id: identifiers.runId, model: input.model ?? null, provider: input.provider ?? null, @@ -393,6 +359,7 @@ export async function createSessionRunRecordIfSessionIdle( error_code, error_message, error_details_json, + error_retryable, started_at, completed_at, created_by_account_id, @@ -425,6 +392,7 @@ export async function createSessionRunRecordIfSessionIdle( NULL, NULL, NULL, + NULL, ${input.startedAt ?? null}, NULL, ${input.createdBy}, @@ -498,130 +466,10 @@ export async function setSessionRunStatus( database: D1Database, input: UpdateSessionRunStatusInput, ): Promise { - return transitionSessionRunStatusAt(database, input, currentTimestampMs()); -} - -export async function cancelActiveSessionRunsForRuntimeOperation( - database: D1Database, - input: { - readonly error: RunError; - readonly operationId: RuntimeOperationId; - readonly runIds: readonly SessionRunId[]; - }, -): Promise { - const runIds = [...new Set(input.runIds)].filter((runId) => runId !== ""); - - if (runIds.length === 0) { - return { - runIds: [], - timestampMs: currentTimestampMs(), - }; + if (isTerminalSessionRunStatus(input.status)) { + throw new Error("Terminal Session Run transitions require the atomic terminal projection."); } - - const cancelledRunIds = new Set(); - let timestampMs = currentTimestampMs(); - - for (let index = 0; index < runIds.length; index += SESSION_RUN_STATUS_WRITE_BATCH_SIZE) { - const runIdBatch = runIds.slice(index, index + SESSION_RUN_STATUS_WRITE_BATCH_SIZE); - const updated = await setNonTerminalSessionRunsStatus(database, { - error: input.error, - operationId: input.operationId, - preserveSessionLifecycle: true, - runIds: runIdBatch, - source: "runtime_operation", - status: "cancelled", - }); - timestampMs = updated.timestampMs; - - const rows = await getAppDatabase(database) - .select({ id: sessionRunsTable.id }) - .from(sessionRunsTable) - .where( - and( - inArray(sessionRunsTable.id, runIdBatch), - eq(sessionRunsTable.status, "cancelled"), - eq(sessionRunsTable.statusOperationId, input.operationId), - ), - ) - .all(); - - for (const row of rows) { - cancelledRunIds.add(row.id); - } - } - - return { - runIds: [...cancelledRunIds], - timestampMs, - }; -} - -async function setNonTerminalSessionRunsStatus( - database: D1Database, - input: SessionRunStatusUpdateInput & { - readonly preserveSessionLifecycle?: boolean; - readonly runIds: readonly SessionRunId[]; - }, -): Promise { - const runIds = [...new Set(input.runIds)].filter((runId) => runId !== ""); - const timestampMs = currentTimestampMs(); - - if (runIds.length === 0) { - return { - runIds: [], - timestampMs, - }; - } - - const updatedRuns = await getAppDatabase(database) - .update(sessionRunsTable) - .set(createSessionRunStatusUpdate(input, timestampMs)) - .where( - and( - inArray(sessionRunsTable.id, runIds), - inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), - ), - ) - .returning({ - id: sessionRunsTable.id, - sessionId: sessionRunsTable.sessionId, - }) - .all(); - - if (input.preserveSessionLifecycle === true || updatedRuns.length === 0) { - return { - runIds: updatedRuns.map((run) => run.id), - timestampMs, - }; - } - - await getAppDatabase(database) - .update(sessionsTable) - .set( - createSessionStatusTransitionPatch({ - status: toSessionLifecycleStatusForRunStatus(input.status), - timestampMs, - }), - ) - .where( - and( - inArray( - sessionsTable.id, - updatedRuns.map((run) => run.sessionId), - ), - inArray( - sessionsTable.lastRunId, - updatedRuns.map((run) => run.id), - ), - notInArray(sessionsTable.status, ["TERMINATED"]), - ), - ) - .run(); - - return { - runIds: updatedRuns.map((run) => run.id), - timestampMs, - }; + return transitionSessionRunStatusAt(database, input, currentTimestampMs()); } async function transitionSessionRunStatusAt( @@ -646,12 +494,22 @@ async function repairCurrentSessionRunProjection( const projectedStatus = toSessionLifecycleStatusForRunStatus(input.targetStatus); - if (input.current.session_status === projectedStatus) { + if ( + input.current.session_status === projectedStatus && + input.current.session_archived_at === null && + input.current.session_cleanup_operation_kind === null && + input.current.session_status_operation_id === null + ) { return "already_projected"; } - if (input.current.session_status === "TERMINATED") { - return "already_projected"; + if ( + input.current.session_archived_at !== null || + input.current.session_cleanup_operation_kind !== null || + input.current.session_status_operation_id !== null || + input.current.session_status === "TERMINATED" + ) { + return "repair_needed"; } const sessionUpdateResult = await getAppDatabase(database) @@ -662,13 +520,7 @@ async function repairCurrentSessionRunProjection( timestampMs: input.timestampMs, }), ) - .where( - and( - eq(sessionsTable.id, input.current.session_id), - eq(sessionsTable.lastRunId, input.current.id), - notInArray(sessionsTable.status, ["TERMINATED"]), - ), - ) + .where(writableObservedSessionCondition(input.current)) .run(); return getD1ChangeCount(sessionUpdateResult) > 0 ? "repaired" : "repair_needed"; @@ -715,22 +567,20 @@ async function transitionSessionRunStatus( break; } case "duplicate": { - if (input.preserveSessionLifecycle !== true) { - const projection = await repairCurrentSessionRunProjection(database, { - current, - targetStatus: input.status, - timestampMs, - }); - - if (projection === "repair_needed") { - return { - kind: "repair_needed", - previousStatus: current.status, - reason: "session_lifecycle_not_updated", - run: toSessionRunSummaryFromLifecycleRow(current), - statusSeq: current.status_seq, - }; - } + const projection = await repairCurrentSessionRunProjection(database, { + current, + targetStatus: input.status, + timestampMs, + }); + + if (projection === "repair_needed") { + return { + kind: "repair_needed", + previousStatus: current.status, + reason: "session_lifecycle_not_updated", + run: toSessionRunSummaryFromLifecycleRow(current), + statusSeq: current.status_seq, + }; } return { @@ -761,39 +611,6 @@ async function transitionSessionRunStatus( const run = toUpdatedSessionRunSummary(current, input, timestampMs); const statusSeq = current.status_seq + 1; - if (input.preserveSessionLifecycle === true) { - const runUpdateResult = await getAppDatabase(database) - .update(sessionRunsTable) - .set(createSessionRunStatusUpdate(input, timestampMs)) - .where( - and( - eq(sessionRunsTable.id, input.runId), - eq(sessionRunsTable.status, current.status), - eq(sessionRunsTable.statusSeq, current.status_seq), - ), - ) - .run(); - - if (getD1ChangeCount(runUpdateResult) === 0) { - return { - currentStatus: current.status, - kind: "stale", - reason: "concurrent_transition", - targetStatus: input.status, - }; - } - - logTerminalSessionRun(current, input, timestampMs); - - return { - kind: "applied", - previousStatus: current.status, - run, - sessionLifecycle: "preserved", - statusSeq, - }; - } - if (current.session_last_run_id !== input.runId) { const runUpdateResult = await getAppDatabase(database) .update(sessionRunsTable) @@ -803,6 +620,12 @@ async function transitionSessionRunStatus( eq(sessionRunsTable.id, input.runId), eq(sessionRunsTable.status, current.status), eq(sessionRunsTable.statusSeq, current.status_seq), + exists( + getAppDatabase(database) + .select({ id: sessionsTable.id }) + .from(sessionsTable) + .where(writableObservedSessionCondition(current)), + ), ), ) .run(); @@ -816,8 +639,6 @@ async function transitionSessionRunStatus( }; } - logTerminalSessionRun(current, input, timestampMs); - return { kind: "applied", previousStatus: current.status, @@ -836,22 +657,25 @@ async function transitionSessionRunStatus( eq(sessionRunsTable.id, input.runId), eq(sessionRunsTable.status, current.status), eq(sessionRunsTable.statusSeq, current.status_seq), + exists( + db + .select({ id: sessionsTable.id }) + .from(sessionsTable) + .where(writableObservedSessionCondition(current)), + ), ), ), db .update(sessionsTable) .set( createCurrentSessionRunProjectionPatch({ - sessionKind: current.session_kind, status: input.status, timestampMs, }), ) .where( and( - eq(sessionsTable.id, current.session_id), - eq(sessionsTable.lastRunId, input.runId), - notInArray(sessionsTable.status, ["TERMINATED"]), + writableObservedSessionCondition(current), exists( db .select({ id: sessionRunsTable.id }) @@ -877,8 +701,6 @@ async function transitionSessionRunStatus( }; } - logTerminalSessionRun(current, input, timestampMs); - if (getD1ChangeCount(sessionUpdateResult) === 0 && current.session_status !== "TERMINATED") { return { kind: "repair_needed", diff --git a/apps/api/src/modules/sessions/application/agent-session-retrieve.service.ts b/apps/api/src/modules/sessions/application/agent-session-retrieve.service.ts index 9fbcc12c..f991f428 100644 --- a/apps/api/src/modules/sessions/application/agent-session-retrieve.service.ts +++ b/apps/api/src/modules/sessions/application/agent-session-retrieve.service.ts @@ -4,6 +4,7 @@ import type { AgentSessionNativeRuntimeRefDiagnostics, AgentSessionRecoverability, AgentSessionRetrieveResult, + AgentTaskSnapshot, SessionExecutionSkillReference, SessionExecutionToolReference, SessionSummary, @@ -17,6 +18,7 @@ import { eq } from "drizzle-orm"; import { getAppDatabase } from "../../../platform/db/drizzle"; import type { AuthenticatedViewer } from "../../auth/application/viewer-auth.service"; import { findSessionExecutionPlan } from "../../runtime/application/session-definition/session-execution.repository"; +import { loadSessionAgentTaskSnapshot } from "../infrastructure/session-agent-task-snapshot.repository"; import { loadSessionViewerState } from "./session-live-state.service"; import { getSessionSummaryAccessById, @@ -68,8 +70,12 @@ export async function retrieveAgentSession( input: AgentSessionLookupInput, ): Promise { const access = await getSessionSummaryAccessById(database, viewer.id, input); + const taskSnapshot = selectCurrentAgentTaskSnapshot( + access.session, + await loadSessionAgentTaskSnapshot(database, input.sessionId), + ); - return toAgentSessionRetrieveResult(access); + return toAgentSessionRetrieveResult({ ...access, taskSnapshot }); } export async function retrieveThreadAgentSession( @@ -78,16 +84,42 @@ export async function retrieveThreadAgentSession( input: AgentSessionLookupInput, ): Promise { const session = await getSessionSummaryForCreator(database, viewer.id, input); + const taskSnapshot = selectCurrentAgentTaskSnapshot( + session, + await loadSessionAgentTaskSnapshot(database, input.sessionId), + ); return toAgentSessionRetrieveResult({ isSessionCreator: true, session, + taskSnapshot, }); } +function selectCurrentAgentTaskSnapshot( + session: SessionSummary, + snapshot: AgentTaskSnapshot | null, +): AgentTaskSnapshot | null { + if ( + snapshot === null || + session.archivedAt !== null || + session.status !== "RUNNING" || + session.lastRun?.id !== snapshot.runId || + session.lastRun.status === "cancelled" || + session.lastRun.status === "completed" || + session.lastRun.status === "expired" || + session.lastRun.status === "failed" + ) { + return null; + } + + return snapshot; +} + export function toAgentSessionRetrieveResult(input: { isSessionCreator: boolean; session: SessionSummary; + taskSnapshot?: AgentTaskSnapshot | null; }): AgentSessionRetrieveResult { return { capabilities: getAgentSessionActionCapabilities({ @@ -96,6 +128,7 @@ export function toAgentSessionRetrieveResult(input: { }), recoverability: getAgentSessionRecoverability(input.session), session: input.session, + taskSnapshot: input.taskSnapshot ?? null, }; } diff --git a/apps/api/src/modules/sessions/application/session-cleanup.service.ts b/apps/api/src/modules/sessions/application/session-cleanup.service.ts index 935f3c79..b420fcb7 100644 --- a/apps/api/src/modules/sessions/application/session-cleanup.service.ts +++ b/apps/api/src/modules/sessions/application/session-cleanup.service.ts @@ -8,24 +8,26 @@ import { import { createPlatformId } from "@mosoo/id"; import type { RuntimeOperationId, SessionId, SessionRunId } from "@mosoo/id"; import type { SQL } from "drizzle-orm"; -import { and, asc, eq, exists, inArray, isNotNull, lte, or, sql } from "drizzle-orm"; +import { and, asc, eq, exists, inArray, isNotNull, isNull, lte, or, sql } from "drizzle-orm"; import { createErrorLogContext, logWarn } from "../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; import { getAppDatabase } from "../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../time"; import { fileStore } from "../../files/application/file-store"; +import { recordCanonicalSessionRunTerminal } from "../../runtime/application/session-runs/session-run-terminal-failure.service"; +import { assertCanonicalTerminalSessionRunProjection } from "../../runtime/application/session-runs/terminal-run-reconciliation.service"; +import { + ACTIVE_SESSION_RUN_STATUSES, + isTerminalSessionRunStatus, +} from "../../runtime/domain/session-run-lifecycle.machine"; import { destroyDriverInstanceDurableObject } from "../../runtime/infrastructure/driver-instance/client"; -import { listLiveDriverInstanceIdsForSandboxSessions } from "../../runtime/infrastructure/driver-instance/live-driver-instance.repository"; +import { listLiveDriverInstanceRefsForSandboxSessions } from "../../runtime/infrastructure/driver-instance/live-driver-instance.repository"; import { stopDriverSession } from "../../runtime/infrastructure/driver-session-stop.service"; -import { deleteSandboxBackupsForDir } from "../../runtime/infrastructure/sandbox-backup.service"; +import { deleteSandboxBackupsForSession } from "../../runtime/infrastructure/sandbox-backup.service"; import { closeSandboxConversationSession } from "../../runtime/infrastructure/sandbox-session.service"; -import { - SESSION_DELETE_CLEANUP_STEPS, - completeSessionDeleteCleanupStep, - shouldSkipSessionDeleteCleanupStep, - skipSessionDeleteCleanupStep, -} from "../domain/session-cleanup-plan"; +import { getSessionRunSummary } from "../../runtime/infrastructure/session-runs/session-run-store.repository"; +import { SESSION_DELETE_CLEANUP_STEPS } from "../domain/session-cleanup-plan"; import type { SessionDeleteCleanupStep, SessionDeleteCleanupStepOutcome, @@ -36,14 +38,26 @@ import { destroySessionDurableObject } from "../infrastructure/session/client"; type AppDatabase = ReturnType; export interface SessionDeleteCleanupRepairCandidate { + readonly archivedAt: number; + readonly cleanupOperationKind: "delete" | null; readonly operationId: RuntimeOperationId; readonly sessionId: SessionId; + readonly status: "IDLE" | "RESCHEDULING" | "TERMINATED"; + readonly statusSeq: number; + readonly updatedAt: number; } export interface DeleteSessionCascadeOptions { readonly operationId?: RuntimeOperationId; } +const DELETED_RUN_ERROR = { + code: "session.deleted", + details: {}, + message: "Session was deleted before the run completed.", + retryable: false, +} as const; + function driverInstancesForSessionCondition( db: AppDatabase, sessionId: SessionId, @@ -76,6 +90,8 @@ async function resolveSessionDeleteCleanupOperationId( const existing = (await getAppDatabase(database) .select({ + archived_at: sessionsTable.archivedAt, + cleanup_operation_kind: sessionsTable.cleanupOperationKind, operation_id: sessionsTable.statusOperationId, status: sessionsTable.status, }) @@ -84,7 +100,21 @@ async function resolveSessionDeleteCleanupOperationId( .limit(1) .get()) ?? null; - if (existing?.status === "TERMINATED" && existing.operation_id !== null) { + if ( + existing?.operation_id !== null && + existing?.operation_id !== undefined && + existing.cleanup_operation_kind === "delete" && + (existing.status === "IDLE" || existing.status === "RESCHEDULING") + ) { + return existing.operation_id; + } + if ( + existing?.archived_at !== null && + existing?.archived_at !== undefined && + existing.cleanup_operation_kind === null && + existing.status === "TERMINATED" && + existing.operation_id !== null + ) { return existing.operation_id; } @@ -98,18 +128,59 @@ async function admitSessionDeleteCleanup( readonly sessionId: SessionId; readonly timestampMs: number; }, -): Promise { - await getAppDatabase(database) +): Promise { + const row = await getAppDatabase(database) .update(sessionsTable) .set({ archivedAt: sql`COALESCE(${sessionsTable.archivedAt}, ${input.timestampMs})`, - status: "TERMINATED", + cleanupOperationKind: "delete", + status: sql`CASE + WHEN ${sessionsTable.statusOperationId} = ${input.operationId} + AND ${sessionsTable.status} IN ('IDLE', 'RESCHEDULING') + THEN ${sessionsTable.status} + ELSE 'RESCHEDULING' + END`, statusOperationId: input.operationId, - statusSeq: sql`${sessionsTable.statusSeq} + 1`, - updatedAt: input.timestampMs, + statusSeq: sql`${sessionsTable.statusSeq} + CASE + WHEN ${sessionsTable.statusOperationId} = ${input.operationId} + AND ${sessionsTable.status} IN ('IDLE', 'RESCHEDULING') + THEN 0 + ELSE 1 + END`, + updatedAt: sql`MAX(${sessionsTable.updatedAt}, ${input.timestampMs})`, }) - .where(eq(sessionsTable.id, input.sessionId)) - .run(); + .where( + and( + eq(sessionsTable.id, input.sessionId), + isNull(sessionsTable.runtimeProvisioningOperationId), + or( + and( + eq(sessionsTable.cleanupOperationKind, "delete"), + eq(sessionsTable.statusOperationId, input.operationId), + inArray(sessionsTable.status, ["IDLE", "RESCHEDULING"]), + ), + and(isNull(sessionsTable.cleanupOperationKind), isNull(sessionsTable.statusOperationId)), + and( + eq(sessionsTable.cleanupOperationKind, "archive"), + eq(sessionsTable.status, "IDLE"), + isNull(sessionsTable.statusOperationId), + ), + and( + isNotNull(sessionsTable.archivedAt), + isNull(sessionsTable.cleanupOperationKind), + eq(sessionsTable.status, "TERMINATED"), + eq(sessionsTable.statusOperationId, input.operationId), + ), + ), + ), + ) + .returning({ archivedAt: sessionsTable.archivedAt }) + .get(); + + if (row?.archivedAt === null || row?.archivedAt === undefined) { + throw new Error("Session delete cleanup could not acquire lifecycle ownership."); + } + return row.archivedAt; } async function listSessionDeleteCleanupRepairCandidates( @@ -125,14 +196,25 @@ async function listSessionDeleteCleanupRepairCandidates( const rows = await getAppDatabase(database) .select({ + archivedAt: sessionsTable.archivedAt, + cleanupOperationKind: sessionsTable.cleanupOperationKind, operationId: sessionsTable.statusOperationId, sessionId: sessionsTable.id, + status: sessionsTable.status, + statusSeq: sessionsTable.statusSeq, + updatedAt: sessionsTable.updatedAt, }) .from(sessionsTable) .where( and( isNotNull(sessionsTable.archivedAt), - eq(sessionsTable.status, "TERMINATED"), + or( + and( + eq(sessionsTable.cleanupOperationKind, "delete"), + inArray(sessionsTable.status, ["IDLE", "RESCHEDULING"]), + ), + and(isNull(sessionsTable.cleanupOperationKind), eq(sessionsTable.status, "TERMINATED")), + ), isNotNull(sessionsTable.statusOperationId), lte(sessionsTable.updatedAt, input.staleUpdatedAtLte), ), @@ -141,9 +223,120 @@ async function listSessionDeleteCleanupRepairCandidates( .limit(input.limit) .all(); - return rows.flatMap((row) => - row.operationId === null ? [] : [{ operationId: row.operationId, sessionId: row.sessionId }], - ); + return rows.flatMap((row) => { + if ( + row.archivedAt === null || + row.operationId === null || + !["IDLE", "RESCHEDULING", "TERMINATED"].includes(row.status) || + (row.cleanupOperationKind !== null && row.cleanupOperationKind !== "delete") + ) { + return []; + } + return [ + { + archivedAt: row.archivedAt, + cleanupOperationKind: row.cleanupOperationKind, + operationId: row.operationId, + sessionId: row.sessionId, + status: row.status as "IDLE" | "RESCHEDULING" | "TERMINATED", + statusSeq: row.statusSeq, + updatedAt: row.updatedAt, + }, + ]; + }); +} + +async function normalizeSessionDeleteRuntimeLifecycle( + bindings: ApiBindings, + input: { + readonly operationId: RuntimeOperationId; + readonly sessionId: SessionId; + readonly timestampMs: number; + }, +): Promise { + const activeRuns = await getAppDatabase(bindings.DB) + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.sessionId, input.sessionId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ) + .all(); + + if (activeRuns.length > 1) { + throw new Error("Session delete found more than one active Session Run."); + } + const [activeRun] = activeRuns; + if (activeRun !== undefined) { + const outcome = await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + deliver: false, + error: DELETED_RUN_ERROR, + expectedSessionOperationId: input.operationId, + lifecycle: "IDLE", + runId: activeRun.id, + sessionId: input.sessionId, + source: "runtime_operation", + status: "cancelled", + timestampMs: input.timestampMs, + }); + if ( + outcome.kind === "stale" && + !["cancelled", "completed", "expired", "failed"].includes(outcome.run.status) + ) { + throw new Error("Session delete lost ownership of its active Session Run."); + } + } + + await getAppDatabase(bindings.DB) + .update(sessionsTable) + .set({ + status: "IDLE", + statusSeq: sql`${sessionsTable.statusSeq} + 1`, + updatedAt: sql`MAX(${sessionsTable.updatedAt}, ${input.timestampMs})`, + }) + .where( + and( + eq(sessionsTable.id, input.sessionId), + eq(sessionsTable.status, "RESCHEDULING"), + eq(sessionsTable.statusOperationId, input.operationId), + ), + ) + .run(); + + const projection = await getAppDatabase(bindings.DB) + .select({ + cleanupOperationKind: sessionsTable.cleanupOperationKind, + lastRunId: sessionsTable.lastRunId, + operationId: sessionsTable.statusOperationId, + status: sessionsTable.status, + }) + .from(sessionsTable) + .where(eq(sessionsTable.id, input.sessionId)) + .limit(1) + .get(); + if ( + projection === undefined || + projection.cleanupOperationKind !== "delete" || + projection.operationId !== input.operationId || + projection.status !== "IDLE" + ) { + throw new Error("Session delete cleanup lost its lifecycle ownership."); + } + if (projection.lastRunId === null) { + return; + } + const lastRun = await getSessionRunSummary(bindings.DB, projection.lastRunId); + if (lastRun === null || !isTerminalSessionRunStatus(lastRun.status)) { + throw new Error("Session delete cleanup did not reach a terminal current Run."); + } + await assertCanonicalTerminalSessionRunProjection(bindings, { + runId: lastRun.id, + sessionId: input.sessionId, + status: lastRun.status, + }); } export async function deleteSessionCascade( @@ -159,6 +352,7 @@ export async function deleteSessionCascade( sessionId, }); const outcomes: SessionDeleteCleanupStepOutcome[] = []; + let cleanupTimestampMs = timestampMs; let targets: SessionDeleteCleanupTargets | null = null; async function loadCleanupTargets(): Promise { @@ -170,7 +364,7 @@ export async function deleteSessionCascade( .limit(1) .get()) ?? null; - const liveDriverInstanceIds = await listLiveDriverInstanceIdsForSandboxSessions(bindings.DB, [ + const liveDriverInstances = await listLiveDriverInstanceRefsForSandboxSessions(bindings.DB, [ sessionId, ]); @@ -181,23 +375,41 @@ export async function deleteSessionCascade( .all(); const runIds = sessionRuns.map((row) => row.id); const associatedDriverInstanceRows = await db - .select({ id: driverInstancesTable.id }) + .select({ generation: driverInstancesTable.generation, id: driverInstancesTable.id }) .from(driverInstancesTable) .where(driverInstancesForSessionCondition(db, sessionId, runIds)) .all(); return { - associatedDriverInstanceIds: associatedDriverInstanceRows.map((row) => row.id), - liveDriverInstanceIds, + associatedDriverInstances: associatedDriverInstanceRows, + liveDriverInstances, sandboxId: sandboxSession?.sandbox_id ?? null, - sessionId, }; } async function executeStep(step: SessionDeleteCleanupStep): Promise { + if (step !== "archive_session_row") { + const owned = await db + .update(sessionsTable) + .set({ updatedAt: sql`MAX(${sessionsTable.updatedAt}, ${currentTimestampMs()})` }) + .where( + and( + eq(sessionsTable.id, sessionId), + eq(sessionsTable.cleanupOperationKind, "delete"), + eq(sessionsTable.statusOperationId, operationId), + inArray(sessionsTable.status, ["IDLE", "RESCHEDULING"]), + ), + ) + .returning({ id: sessionsTable.id }) + .get(); + if (owned === undefined) { + throw new Error("Session delete cleanup lost its operation ownership."); + } + } + switch (step) { case "archive_session_row": { - await admitSessionDeleteCleanup(bindings.DB, { + cleanupTimestampMs = await admitSessionDeleteCleanup(bindings.DB, { operationId, sessionId, timestampMs, @@ -209,23 +421,29 @@ export async function deleteSessionCascade( return; } case "stop_live_drivers": { - await Promise.all( - requireCleanupTargets(targets).liveDriverInstanceIds.map((driverInstanceId) => + const results = await Promise.allSettled( + requireCleanupTargets(targets).liveDriverInstances.map((driver) => stopDriverSession(bindings, { - driverInstanceId, + driverInstanceId: driver.id, + expectedDriverGeneration: driver.generation, reason: "session.deleted", - terminalRun: { - error: { - code: "session.deleted", - details: {}, - message: "Session was deleted before the run completed.", - retryable: false, - }, - status: "cancelled", - }, }), ), ); + const failure = results.find( + (result): result is PromiseRejectedResult => result.status === "rejected", + ); + if (failure !== undefined) { + throw failure.reason; + } + return; + } + case "normalize_runtime_lifecycle": { + await normalizeSessionDeleteRuntimeLifecycle(bindings, { + operationId, + sessionId, + timestampMs: cleanupTimestampMs, + }); return; } case "close_sandbox_session": { @@ -241,11 +459,22 @@ export async function deleteSessionCascade( return; } case "destroy_driver_objects": { - await Promise.all( - requireCleanupTargets(targets).liveDriverInstanceIds.map((driverInstanceId) => - destroyDriverInstanceDurableObject(bindings, driverInstanceId, "session.deleted"), + const results = await Promise.allSettled( + requireCleanupTargets(targets).associatedDriverInstances.map((driver) => + destroyDriverInstanceDurableObject( + bindings, + driver.id, + driver.generation, + "session.deleted", + ), ), ); + const failure = results.find( + (result): result is PromiseRejectedResult => result.status === "rejected", + ); + if (failure !== undefined) { + throw failure.reason; + } return; } case "destroy_session_object": { @@ -253,7 +482,16 @@ export async function deleteSessionCascade( return; } case "delete_session_backups": { - await deleteSandboxBackupsForDir(bindings, { dir: sessionCwd }); + const sandboxId = requireCleanupTargets(targets).sandboxId; + if (sandboxId === null) { + return; + } + await deleteSandboxBackupsForSession(bindings, { + cwd: sessionCwd, + operationId, + sandboxId, + sessionId, + }); return; } case "delete_session_files": { @@ -264,8 +502,9 @@ export async function deleteSessionCascade( return; } case "delete_driver_rows": { - const associatedDriverInstanceIds = - requireCleanupTargets(targets).associatedDriverInstanceIds; + const associatedDriverInstanceIds = requireCleanupTargets( + targets, + ).associatedDriverInstances.map((driver) => driver.id); if (associatedDriverInstanceIds.length === 0) { return; } @@ -286,20 +525,26 @@ export async function deleteSessionCascade( } } + function shouldSkipStep(step: SessionDeleteCleanupStep): boolean { + if (targets === null) { + return false; + } + return ( + (step === "close_sandbox_session" && targets.sandboxId === null) || + (step === "stop_live_drivers" && targets.liveDriverInstances.length === 0) || + ((step === "delete_driver_rows" || step === "destroy_driver_objects") && + targets.associatedDriverInstances.length === 0) + ); + } + for (const step of SESSION_DELETE_CLEANUP_STEPS) { - if ( - targets !== null && - shouldSkipSessionDeleteCleanupStep({ - step, - targets, - }) - ) { - outcomes.push(skipSessionDeleteCleanupStep(step)); + if (shouldSkipStep(step)) { + outcomes.push({ status: "skipped", step }); continue; } await executeStep(step); - outcomes.push(completeSessionDeleteCleanupStep(step)); + outcomes.push({ status: "completed", step }); } return outcomes; @@ -312,15 +557,59 @@ export async function repairStaleSessionDeleteCleanups( readonly staleUpdatedAtLte: number; }, ): Promise { - const candidates = await listSessionDeleteCleanupRepairCandidates(bindings.DB, input); + const observed = await listSessionDeleteCleanupRepairCandidates(bindings.DB, input); + const candidates: SessionDeleteCleanupRepairCandidate[] = []; + for (const candidate of observed) { + const attemptAt = Math.max(currentTimestampMs(), candidate.updatedAt + 1); + const claimed = await getAppDatabase(bindings.DB) + .update(sessionsTable) + .set({ + cleanupOperationKind: "delete", + status: candidate.status === "TERMINATED" ? "RESCHEDULING" : candidate.status, + statusSeq: + candidate.status === "TERMINATED" + ? sql`${sessionsTable.statusSeq} + 1` + : sessionsTable.statusSeq, + updatedAt: attemptAt, + }) + .where( + and( + eq(sessionsTable.id, candidate.sessionId), + eq(sessionsTable.archivedAt, candidate.archivedAt), + candidate.cleanupOperationKind === null + ? isNull(sessionsTable.cleanupOperationKind) + : eq(sessionsTable.cleanupOperationKind, candidate.cleanupOperationKind), + eq(sessionsTable.status, candidate.status), + eq(sessionsTable.statusOperationId, candidate.operationId), + eq(sessionsTable.statusSeq, candidate.statusSeq), + eq(sessionsTable.updatedAt, candidate.updatedAt), + ), + ) + .returning({ id: sessionsTable.id }) + .get(); + if (claimed !== undefined) { + candidates.push(candidate); + } + } - await Promise.all( + await Promise.allSettled( candidates.map(async (candidate) => { try { await deleteSessionCascade(bindings, candidate.sessionId, { operationId: candidate.operationId, }); } catch (error) { + await getAppDatabase(bindings.DB) + .update(sessionsTable) + .set({ updatedAt: sql`MAX(${sessionsTable.updatedAt}, ${currentTimestampMs()})` }) + .where( + and( + eq(sessionsTable.id, candidate.sessionId), + eq(sessionsTable.cleanupOperationKind, "delete"), + eq(sessionsTable.statusOperationId, candidate.operationId), + ), + ) + .run(); logWarn("session.delete_cleanup.repair_failed", { ...createErrorLogContext(error), operationId: candidate.operationId, diff --git a/apps/api/src/modules/sessions/application/session-event-write.service.ts b/apps/api/src/modules/sessions/application/session-event-write.service.ts index d84ff2c8..6bfee501 100644 --- a/apps/api/src/modules/sessions/application/session-event-write.service.ts +++ b/apps/api/src/modules/sessions/application/session-event-write.service.ts @@ -1,6 +1,5 @@ -import type { AgUiSessionEvent } from "@mosoo/ag-ui-session"; import { createPlatformId } from "@mosoo/id"; -import type { RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; +import type { DriverCommandId, RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; import { createRuntimeEvent } from "@mosoo/runtime-events"; import type { RuntimeEventActor, @@ -19,8 +18,7 @@ import { persistSessionRuntimeEvents, } from "../infrastructure/session-runtime-event-store.repository"; import type { PersistSessionRuntimeEventsResult } from "../infrastructure/session-runtime-event-store.repository"; -import { appRuntimeEventsToSessionDeliveryEvents } from "./session-live-state.service"; -import { publishSessionViewerEvents } from "./session-viewer-events.service"; +import { syncSessionViewerState } from "./session-viewer-events.service"; export interface AppendOneSessionEventPerSessionResult { readonly persistedCount: number; @@ -42,17 +40,15 @@ export interface CreateSessionRuntimeEventInput { visibility?: RuntimeEventVisibility; } -async function publishSessionViewerEventsSafely( +async function syncSessionViewerStateSafely( bindings: ApiBindings, sessionId: SessionId, - events: AgUiSessionEvent[], ): Promise { try { - await publishSessionViewerEvents(bindings, sessionId, events); + await syncSessionViewerState(bindings, sessionId); } catch (error) { - logWarn("session.runtime_event.live_delivery_failed", { + logWarn("session.runtime_event.live_sync_failed", { ...createErrorLogContext(error), - eventCount: events.length, sessionId, }); } @@ -63,10 +59,8 @@ export async function publishPersistedSessionRuntimeEvents(input: { events: readonly RuntimeEventEnvelope[]; sessionId: SessionId; }): Promise { - const deliveryEvents = appRuntimeEventsToSessionDeliveryEvents(input.events); - - if (deliveryEvents.length > 0) { - await publishSessionViewerEventsSafely(input.bindings, input.sessionId, deliveryEvents); + if (input.events.length > 0) { + await syncSessionViewerStateSafely(input.bindings, input.sessionId); } } @@ -74,6 +68,7 @@ export interface AppendSessionRuntimeEventsInput { bindings: ApiBindings; deliver?: boolean; events: RuntimeEventEnvelope[]; + provenMcpCommandId?: DriverCommandId | null; sessionId: SessionId; sourceEventId?: string | null; } @@ -122,11 +117,15 @@ export async function appendSessionRuntimeEvents( persistedSourceEventIds: [], }; } + if (input.provenMcpCommandId != null && input.events.length !== 1) { + throw new Error("Proven MCP command provenance requires exactly one runtime event."); + } const result = await persistSessionRuntimeEvents(input.bindings.DB, { records: input.events.map((event, index) => ({ event, occurredAt: toRuntimeEventOccurredAtMs(event), + provenMcpCommandId: index === 0 ? (input.provenMcpCommandId ?? null) : null, sourceEventId: event.sourceEventId ?? (index === 0 ? (input.sourceEventId ?? null) : null), })), sessionId: input.sessionId, @@ -135,7 +134,7 @@ export async function appendSessionRuntimeEvents( if (input.deliver !== false) { await publishPersistedSessionRuntimeEvents({ bindings: input.bindings, - events: result.persistedEvents, + events: input.events, sessionId: input.sessionId, }); } @@ -167,19 +166,10 @@ export async function appendOneSessionRuntimeEventPerSession(input: { return result; } - const skippedSessionIds = new Set(result.skippedSessionIds); await Promise.all( - input.records.flatMap((record) => { - if (skippedSessionIds.has(record.sessionId)) { - return []; - } - - const deliveryEvents = appRuntimeEventsToSessionDeliveryEvents([record.event]); - - return deliveryEvents.length === 0 - ? [] - : [publishSessionViewerEventsSafely(input.bindings, record.sessionId, deliveryEvents)]; - }), + [...new Set(input.records.map((record) => record.sessionId))].map((sessionId) => + syncSessionViewerStateSafely(input.bindings, sessionId), + ), ); return result; diff --git a/apps/api/src/modules/sessions/application/session-lifecycle-mutation.service.ts b/apps/api/src/modules/sessions/application/session-lifecycle-mutation.service.ts index 8fc681f1..5fe95f04 100644 --- a/apps/api/src/modules/sessions/application/session-lifecycle-mutation.service.ts +++ b/apps/api/src/modules/sessions/application/session-lifecycle-mutation.service.ts @@ -1,22 +1,30 @@ import type { AgentSessionActionCapabilityName } from "@mosoo/contracts/session"; import { sandboxSessionsTable, sessionRunsTable, sessionsTable } from "@mosoo/db"; -import type { AppId, SessionId, SessionRunId } from "@mosoo/id"; +import { createPlatformId } from "@mosoo/id"; +import type { AppId, RuntimeOperationId, SessionId, SessionRunId } from "@mosoo/id"; import { getAvailableAgentSessionActionCapability } from "@mosoo/session-policy"; -import { and, eq, inArray } from "drizzle-orm"; +import { and, asc, eq, inArray, isNotNull, isNull, lte, ne, or, sql } from "drizzle-orm"; +import { createErrorLogContext, logWarn } from "../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../platform/cloudflare/worker-types"; import { getAppDatabase } from "../../../platform/db/drizzle"; import { forbiddenError } from "../../../platform/errors"; import { currentTimestampMs } from "../../../time"; import type { AuthenticatedViewer } from "../../auth/application/viewer-auth.service"; +import { recordCanonicalSessionRunTerminal } from "../../runtime/application/session-runs/session-run-terminal-failure.service"; +import { assertCanonicalTerminalSessionRunProjection } from "../../runtime/application/session-runs/terminal-run-reconciliation.service"; import { - createSessionStatusTransitionPatch, - setSystemSessionRunStatus, -} from "../../runtime/application/session-lifecycle-transition.service"; -import { ACTIVE_SESSION_RUN_STATUSES } from "../../runtime/domain/session-run-lifecycle.machine"; -import { listLiveDriverInstanceIdsForSandboxSessions } from "../../runtime/infrastructure/driver-instance/live-driver-instance.repository"; + ACTIVE_SESSION_RUN_STATUSES, + isTerminalSessionRunStatus, +} from "../../runtime/domain/session-run-lifecycle.machine"; +import { listLiveDriverInstanceRefsForSandboxSessions } from "../../runtime/infrastructure/driver-instance/live-driver-instance.repository"; import { stopDriverSession } from "../../runtime/infrastructure/driver-session-stop.service"; import { closeSandboxConversationSession } from "../../runtime/infrastructure/sandbox-session.service"; +import { + cattleTerminalCheckpointReadyPredicate, + isCattleTerminalCheckpointReadyForNextRun, +} from "../../runtime/infrastructure/session-runs/session-run-admission.repository"; +import { getSessionRunSummary } from "../../runtime/infrastructure/session-runs/session-run-store.repository"; import type { SessionActionAuthorization, SessionParticipantCapabilityAccessRow, @@ -26,12 +34,7 @@ import { lookupAppSessionParticipantCapabilityAccess, resolveSessionActionCreatorFlag, } from "../domain/session-access.policy"; -import { - SESSION_ARCHIVE_CLEANUP_STEPS, - completeSessionArchiveCleanupStep, - shouldSkipSessionArchiveCleanupStep, - skipSessionArchiveCleanupStep, -} from "../domain/session-cleanup-plan"; +import { SESSION_ARCHIVE_CLEANUP_STEPS } from "../domain/session-cleanup-plan"; import type { SessionArchiveCleanupStep, SessionArchiveCleanupStepOutcome, @@ -71,6 +74,130 @@ const ARCHIVED_RUN_ERROR = { retryable: false, } as const; +interface SessionArchiveCleanupClaim { + readonly operationId: RuntimeOperationId; + readonly timestampMs: number; +} + +async function readCurrentSessionArchiveCleanupClaim( + database: D1Database, + sessionId: SessionId, +): Promise { + const row = await getAppDatabase(database) + .select({ + cleanupOperationKind: sessionsTable.cleanupOperationKind, + operationId: sessionsTable.statusOperationId, + timestampMs: sessionsTable.archivedAt, + }) + .from(sessionsTable) + .where(eq(sessionsTable.id, sessionId)) + .limit(1) + .get(); + + return row?.cleanupOperationKind !== "archive" || + row.operationId === null || + row.timestampMs === null + ? null + : { operationId: row.operationId, timestampMs: row.timestampMs }; +} + +async function claimSessionArchiveCleanup( + bindings: ApiBindings, + input: { + readonly appId: AppId; + readonly sessionId: SessionId; + }, +): Promise { + const db = getAppDatabase(bindings.DB); + const current = await db + .select({ + archivedAt: sessionsTable.archivedAt, + cleanupOperationKind: sessionsTable.cleanupOperationKind, + operationId: sessionsTable.statusOperationId, + provisioningOperationId: sessionsTable.runtimeProvisioningOperationId, + status: sessionsTable.status, + statusSeq: sessionsTable.statusSeq, + updatedAt: sessionsTable.updatedAt, + }) + .from(sessionsTable) + .where(and(eq(sessionsTable.id, input.sessionId), eq(sessionsTable.appId, input.appId))) + .limit(1) + .get(); + if (current === undefined) { + throw new Error("Session was not found while claiming archive cleanup."); + } + if (current.provisioningOperationId !== null) { + throw new Error("Session runtime provisioning is still in progress."); + } + if (current.cleanupOperationKind === "archive") { + if (current.archivedAt === null) { + throw new Error("Session archive cleanup has an incomplete durable claim."); + } + if (current.operationId === null) { + return null; + } + return { operationId: current.operationId, timestampMs: current.archivedAt }; + } + const isLegacyArchive = + current.archivedAt !== null && + current.cleanupOperationKind === null && + (current.operationId === null || current.status !== "TERMINATED"); + if (current.cleanupOperationKind !== null || (current.operationId !== null && !isLegacyArchive)) { + throw new Error("Session is owned by another lifecycle operation."); + } + if (!(await isCattleTerminalCheckpointReadyForNextRun(bindings.DB, input.sessionId))) { + throw new Error("Session archive is waiting for its completed Run checkpoint."); + } + + const operationId = createPlatformId(); + const timestampMs = current.archivedAt ?? currentTimestampMs(); + const updatedAt = Math.max(currentTimestampMs(), current.updatedAt + 1); + const claimed = await db + .update(sessionsTable) + .set({ + archivedAt: timestampMs, + cleanupOperationKind: "archive", + status: "RESCHEDULING", + statusOperationId: operationId, + statusSeq: sql`${sessionsTable.statusSeq} + 1`, + updatedAt, + }) + .where( + and( + eq(sessionsTable.id, input.sessionId), + eq(sessionsTable.appId, input.appId), + current.archivedAt === null + ? isNull(sessionsTable.archivedAt) + : eq(sessionsTable.archivedAt, current.archivedAt), + isNull(sessionsTable.cleanupOperationKind), + current.operationId === null + ? isNull(sessionsTable.statusOperationId) + : eq(sessionsTable.statusOperationId, current.operationId), + isNull(sessionsTable.runtimeProvisioningOperationId), + eq(sessionsTable.status, current.status), + eq(sessionsTable.statusSeq, current.statusSeq), + eq(sessionsTable.updatedAt, current.updatedAt), + cattleTerminalCheckpointReadyPredicate(db), + ), + ) + .returning({ id: sessionsTable.id }) + .get(); + if (claimed !== undefined) { + return { operationId, timestampMs }; + } + + const winner = await readCurrentSessionArchiveCleanupClaim(bindings.DB, input.sessionId); + if (winner !== null) { + return winner; + } + + if (!(await isCattleTerminalCheckpointReadyForNextRun(bindings.DB, input.sessionId))) { + throw new Error("Session archive is waiting for its completed Run checkpoint."); + } + + throw new Error("Session archive cleanup lost its admission race."); +} + function ensureLifecycleActionCapability(input: { action: AgentSessionActionCapabilityName; authorization?: SessionActionAuthorization | undefined; @@ -106,46 +233,94 @@ async function listActiveSessionRunIds( return rows.map((row) => row.id); } -async function cancelActiveSessionRunsForLifecycle( - database: D1Database, +async function terminalizeActiveSessionRunsForLifecycle( + bindings: ApiBindings, sessionId: SessionId, + timestampMs: number, + operationId: RuntimeOperationId, ): Promise { - const activeRunIds = await listActiveSessionRunIds(database, sessionId); + const activeRunIds = await listActiveSessionRunIds(bindings.DB, sessionId); + + if (activeRunIds.length > 1) { + throw new Error("Session archive found more than one active Session Run."); + } for (const runId of activeRunIds) { - const outcome = await setSystemSessionRunStatus(database, { + const outcome = await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + deliver: false, error: ARCHIVED_RUN_ERROR, + expectedSessionOperationId: operationId, + lifecycle: "IDLE", runId, + sessionId, + source: "runtime_operation", status: "cancelled", + timestampMs, }); - if (outcome.kind === "repair_needed") { - throw new Error("Session archive left the session lifecycle projection stale."); + if ( + outcome.kind === "stale" && + !["cancelled", "completed", "expired", "failed"].includes(outcome.run.status) + ) { + throw new Error("Session archive lost ownership of its active Session Run."); } } } async function normalizeSessionRuntimeLifecycle( - database: D1Database, + bindings: ApiBindings, sessionId: SessionId, + timestampMs: number, + operationId: RuntimeOperationId, ): Promise { - await cancelActiveSessionRunsForLifecycle(database, sessionId); + await terminalizeActiveSessionRunsForLifecycle(bindings, sessionId, timestampMs, operationId); - await getAppDatabase(database) + await getAppDatabase(bindings.DB) .update(sessionsTable) - .set( - createSessionStatusTransitionPatch({ - status: "IDLE", - timestampMs: currentTimestampMs(), - }), - ) + .set({ + status: "IDLE", + statusSeq: sql`${sessionsTable.statusSeq} + 1`, + updatedAt: sql`MAX(${sessionsTable.updatedAt}, ${timestampMs})`, + }) .where( and( eq(sessionsTable.id, sessionId), - inArray(sessionsTable.status, ["RUNNING", "RESCHEDULING"]), + eq(sessionsTable.status, "RESCHEDULING"), + eq(sessionsTable.statusOperationId, operationId), ), ) .run(); + + const projection = await getAppDatabase(bindings.DB) + .select({ + lastRunId: sessionsTable.lastRunId, + operationId: sessionsTable.statusOperationId, + status: sessionsTable.status, + }) + .from(sessionsTable) + .where(eq(sessionsTable.id, sessionId)) + .limit(1) + .get(); + if ( + projection === undefined || + projection.operationId !== operationId || + projection.status !== "IDLE" + ) { + throw new Error("Session archive cleanup lost its lifecycle ownership."); + } + if (projection.lastRunId === null) { + return; + } + const lastRun = await getSessionRunSummary(bindings.DB, projection.lastRunId); + if (lastRun === null || !isTerminalSessionRunStatus(lastRun.status)) { + throw new Error("Session archive cleanup did not reach a terminal current Run."); + } + await assertCanonicalTerminalSessionRunProjection(bindings, { + runId: lastRun.id, + sessionId, + status: lastRun.status, + }); } export async function archiveAgentSession({ @@ -159,12 +334,31 @@ export async function archiveAgentSession({ appId, sessionId, }); - ensureLifecycleActionCapability({ - action: "archive_session", - authorization, - session, - }); - const timestampMs = currentTimestampMs(); + const existingClaim = await readCurrentSessionArchiveCleanupClaim(bindings.DB, sessionId); + if (existingClaim === null && session.archived_at === null) { + ensureLifecycleActionCapability({ + action: "archive_session", + authorization, + session, + }); + } + const claim = existingClaim ?? (await claimSessionArchiveCleanup(bindings, { appId, sessionId })); + if (claim === null) { + return []; + } + return executeSessionArchiveCleanup(bindings, { appId, claim, sessionId }); +} + +async function executeSessionArchiveCleanup( + bindings: ApiBindings, + input: { + readonly appId: AppId; + readonly claim: SessionArchiveCleanupClaim; + readonly sessionId: SessionId; + }, +): Promise { + const { appId, sessionId } = input; + const { operationId, timestampMs } = input.claim; const outcomes: SessionArchiveCleanupStepOutcome[] = []; let targets: SessionArchiveCleanupTargets | null = null; @@ -177,28 +371,38 @@ export async function archiveAgentSession({ .limit(1) .get()) ?? null; - const liveDriverInstanceIds = await listLiveDriverInstanceIdsForSandboxSessions(bindings.DB, [ + const liveDriverInstances = await listLiveDriverInstanceRefsForSandboxSessions(bindings.DB, [ sessionId, ]); return { - liveDriverInstanceIds, + liveDriverInstances, sandboxId: sandboxSession?.sandbox_id ?? null, - sessionId, }; } async function executeStep(step: SessionArchiveCleanupStep): Promise { + const owned = await getAppDatabase(bindings.DB) + .update(sessionsTable) + .set({ updatedAt: sql`MAX(${sessionsTable.updatedAt}, ${currentTimestampMs()})` }) + .where( + and( + eq(sessionsTable.id, sessionId), + eq(sessionsTable.appId, appId), + eq(sessionsTable.archivedAt, timestampMs), + eq(sessionsTable.cleanupOperationKind, "archive"), + eq(sessionsTable.statusOperationId, operationId), + inArray(sessionsTable.status, ["IDLE", "RESCHEDULING"]), + ), + ) + .returning({ id: sessionsTable.id }) + .get(); + if (owned === undefined) { + throw new Error("Session archive cleanup lost its operation ownership."); + } + switch (step) { case "archive_session_row": { - await getAppDatabase(bindings.DB) - .update(sessionsTable) - .set({ - archivedAt: timestampMs, - updatedAt: timestampMs, - }) - .where(and(eq(sessionsTable.id, sessionId), eq(sessionsTable.appId, appId))) - .run(); return; } case "close_viewer_sockets": { @@ -210,22 +414,25 @@ export async function archiveAgentSession({ return; } case "stop_live_drivers": { - await Promise.all( - requireArchiveCleanupTargets(targets).liveDriverInstanceIds.map((driverInstanceId) => + const stopOutcomes = await Promise.allSettled( + requireArchiveCleanupTargets(targets).liveDriverInstances.map((driver) => stopDriverSession(bindings, { - driverInstanceId, + driverInstanceId: driver.id, + expectedDriverGeneration: driver.generation, reason: "session.archived", - terminalRun: { - error: ARCHIVED_RUN_ERROR, - status: "cancelled", - }, }), ), ); + const failure = stopOutcomes.find( + (outcome): outcome is PromiseRejectedResult => outcome.status === "rejected", + ); + if (failure !== undefined) { + throw failure.reason; + } return; } case "normalize_runtime_lifecycle": { - await normalizeSessionRuntimeLifecycle(bindings.DB, sessionId); + await normalizeSessionRuntimeLifecycle(bindings, sessionId, timestampMs, operationId); return; } case "close_sandbox_session": { @@ -240,26 +447,58 @@ export async function archiveAgentSession({ }); return; } + case "complete_archive_session": { + const completed = await getAppDatabase(bindings.DB) + .update(sessionsTable) + .set({ + archivedAt: timestampMs, + cleanupOperationKind: "archive", + status: "IDLE", + statusOperationId: null, + statusSeq: sql`${sessionsTable.statusSeq} + 1`, + updatedAt: sql`MAX(${sessionsTable.updatedAt}, ${timestampMs})`, + }) + .where( + and( + eq(sessionsTable.id, sessionId), + eq(sessionsTable.appId, appId), + eq(sessionsTable.archivedAt, timestampMs), + eq(sessionsTable.cleanupOperationKind, "archive"), + eq(sessionsTable.status, "IDLE"), + eq(sessionsTable.statusOperationId, operationId), + ), + ) + .returning({ id: sessionsTable.id }) + .get(); + if (completed === undefined) { + throw new Error("Session archive cleanup lost its operation ownership."); + } + return; + } default: { throw new Error("Unknown session archive cleanup step."); } } } + function shouldSkipStep(step: SessionArchiveCleanupStep): boolean { + if (targets === null) { + return false; + } + return ( + (step === "close_sandbox_session" && targets.sandboxId === null) || + (step === "stop_live_drivers" && targets.liveDriverInstances.length === 0) + ); + } + for (const step of SESSION_ARCHIVE_CLEANUP_STEPS) { - if ( - targets !== null && - shouldSkipSessionArchiveCleanupStep({ - step, - targets, - }) - ) { - outcomes.push(skipSessionArchiveCleanupStep(step)); + if (shouldSkipStep(step)) { + outcomes.push({ status: "skipped", step }); continue; } await executeStep(step); - outcomes.push(completeSessionArchiveCleanupStep(step)); + outcomes.push({ status: "completed", step }); } return outcomes; @@ -275,6 +514,126 @@ function requireArchiveCleanupTargets( return targets; } +export async function repairStaleSessionArchiveCleanups( + bindings: ApiBindings, + input: { + readonly limit: number; + readonly staleUpdatedAtLte: number; + }, +): Promise { + if (!Number.isSafeInteger(input.limit) || input.limit <= 0) { + throw new Error("Session archive cleanup repair limit must be a positive integer."); + } + + const rows = await getAppDatabase(bindings.DB) + .select({ + appId: sessionsTable.appId, + cleanupOperationKind: sessionsTable.cleanupOperationKind, + operationId: sessionsTable.statusOperationId, + sessionId: sessionsTable.id, + status: sessionsTable.status, + statusSeq: sessionsTable.statusSeq, + timestampMs: sessionsTable.archivedAt, + updatedAt: sessionsTable.updatedAt, + }) + .from(sessionsTable) + .where( + and( + isNotNull(sessionsTable.archivedAt), + or( + and( + eq(sessionsTable.cleanupOperationKind, "archive"), + inArray(sessionsTable.status, ["IDLE", "RESCHEDULING"]), + isNotNull(sessionsTable.statusOperationId), + ), + and( + isNull(sessionsTable.cleanupOperationKind), + or(isNull(sessionsTable.statusOperationId), ne(sessionsTable.status, "TERMINATED")), + ), + ), + lte(sessionsTable.updatedAt, input.staleUpdatedAtLte), + ), + ) + .orderBy(asc(sessionsTable.updatedAt), asc(sessionsTable.id)) + .limit(input.limit) + .all(); + const candidates: Array<{ + appId: AppId; + claim: SessionArchiveCleanupClaim; + sessionId: SessionId; + }> = []; + for (const row of rows) { + if (row.timestampMs === null) { + continue; + } + if (row.cleanupOperationKind === null) { + const claim = await claimSessionArchiveCleanup(bindings, { + appId: row.appId, + sessionId: row.sessionId, + }); + if (claim !== null) { + candidates.push({ appId: row.appId, claim, sessionId: row.sessionId }); + } + continue; + } + if (row.operationId === null) { + continue; + } + + const attemptAt = Math.max(currentTimestampMs(), row.updatedAt + 1); + const claimed = await getAppDatabase(bindings.DB) + .update(sessionsTable) + .set({ updatedAt: attemptAt }) + .where( + and( + eq(sessionsTable.id, row.sessionId), + eq(sessionsTable.archivedAt, row.timestampMs), + eq(sessionsTable.cleanupOperationKind, "archive"), + eq(sessionsTable.status, row.status), + eq(sessionsTable.statusOperationId, row.operationId), + eq(sessionsTable.statusSeq, row.statusSeq), + eq(sessionsTable.updatedAt, row.updatedAt), + ), + ) + .returning({ id: sessionsTable.id }) + .get(); + if (claimed !== undefined) { + candidates.push({ + appId: row.appId, + claim: { operationId: row.operationId, timestampMs: row.timestampMs }, + sessionId: row.sessionId, + }); + } + } + + await Promise.allSettled( + candidates.map(async (candidate) => { + try { + await executeSessionArchiveCleanup(bindings, candidate); + } catch (error) { + await getAppDatabase(bindings.DB) + .update(sessionsTable) + .set({ updatedAt: sql`MAX(${sessionsTable.updatedAt}, ${currentTimestampMs()})` }) + .where( + and( + eq(sessionsTable.id, candidate.sessionId), + eq(sessionsTable.cleanupOperationKind, "archive"), + eq(sessionsTable.statusOperationId, candidate.claim.operationId), + ), + ) + .run(); + logWarn("session.archive_cleanup.repair_failed", { + ...createErrorLogContext(error), + operationId: candidate.claim.operationId, + sessionId: candidate.sessionId, + }); + } + }), + ); + + return candidates.length; +} + export async function unarchiveAgentSession({ authorization, database, @@ -292,16 +651,31 @@ export async function unarchiveAgentSession({ session, }); - await normalizeSessionRuntimeLifecycle(database, sessionId); - - await getAppDatabase(database) + if (session.archived_at === null) { + throw new Error("Session is not archived."); + } + const unarchived = await getAppDatabase(database) .update(sessionsTable) .set({ archivedAt: null, + cleanupOperationKind: null, updatedAt: currentTimestampMs(), }) - .where(and(eq(sessionsTable.id, sessionId), eq(sessionsTable.appId, appId))) - .run(); + .where( + and( + eq(sessionsTable.id, sessionId), + eq(sessionsTable.appId, appId), + eq(sessionsTable.archivedAt, session.archived_at), + eq(sessionsTable.cleanupOperationKind, "archive"), + eq(sessionsTable.status, "IDLE"), + isNull(sessionsTable.statusOperationId), + ), + ) + .returning({ id: sessionsTable.id }) + .get(); + if (unarchived === undefined) { + throw new Error("Session cleanup is still in progress."); + } } export async function deleteAgentSession({ diff --git a/apps/api/src/modules/sessions/application/session-live-state.service.ts b/apps/api/src/modules/sessions/application/session-live-state.service.ts index fcfc69e6..ab951205 100644 --- a/apps/api/src/modules/sessions/application/session-live-state.service.ts +++ b/apps/api/src/modules/sessions/application/session-live-state.service.ts @@ -27,12 +27,6 @@ export function appRuntimeEventToSessionDeliveryEvents( return appRuntimeEventToAgUiSessionEvents(event); } -export function appRuntimeEventsToSessionDeliveryEvents( - events: readonly RuntimeEventEnvelope[], -): SessionDeliveryEvent[] { - return events.flatMap((event) => appRuntimeEventToSessionDeliveryEvents(event)); -} - export function applyRuntimeEventToSessionLiveState( state: SessionLiveState, event: RuntimeEventEnvelope, @@ -44,5 +38,11 @@ function applyRuntimeEventsToSessionLiveState( state: SessionLiveState, events: readonly RuntimeEventEnvelope[], ): SessionLiveState { - return applyAgUiEventsToSessionLiveState(state, appRuntimeEventsToSessionDeliveryEvents(events)); + let next = state; + + for (const event of events) { + next = applyAgUiEventsToSessionLiveState(next, appRuntimeEventToSessionDeliveryEvents(event)); + } + + return next; } diff --git a/apps/api/src/modules/sessions/application/session-message-query.service.ts b/apps/api/src/modules/sessions/application/session-message-query.service.ts index e9150aa8..ca18b719 100644 --- a/apps/api/src/modules/sessions/application/session-message-query.service.ts +++ b/apps/api/src/modules/sessions/application/session-message-query.service.ts @@ -6,6 +6,7 @@ import { asc, eq } from "drizzle-orm"; import { getAppDatabase } from "../../../platform/db/drizzle"; import type { AuthenticatedViewer } from "../../auth/application/viewer-auth.service"; import { ensureAppSessionParticipantAccess } from "../domain/session-access.policy"; +import { resolveStoredSessionMessageReferences } from "../infrastructure/session-message-reference.repository"; import { toSessionMessage } from "./session-message-mappers"; import { getSessionReadAccess } from "./session-read-access.service"; @@ -46,13 +47,17 @@ async function listSessionMessages( created_by_account_id: sessionMessagesTable.createdByAccountId, id: sessionMessagesTable.id, plan_json: sessionMessagesTable.planJson, + projection_format: sessionMessagesTable.projectionFormat, role: sessionMessagesTable.role, segments_json: sessionMessagesTable.segmentsJson, + session_run_id: sessionMessagesTable.sessionRunId, }) .from(sessionMessagesTable) .where(eq(sessionMessagesTable.sessionId, sessionId)) .orderBy(asc(sessionMessagesTable.seq)) .all(); - return results.map((row) => toSessionMessage(row)); + const resolved = await resolveStoredSessionMessageReferences(database, sessionId, results); + + return resolved.map((row) => toSessionMessage(row)); } diff --git a/apps/api/src/modules/sessions/application/session-message-write.service.ts b/apps/api/src/modules/sessions/application/session-message-write.service.ts deleted file mode 100644 index 30e01572..00000000 --- a/apps/api/src/modules/sessions/application/session-message-write.service.ts +++ /dev/null @@ -1,6 +0,0 @@ -export { - insertSessionMessageRecord, - insertSessionMessage, - type InsertedSessionMessage, - type InsertSessionMessageInput, -} from "../infrastructure/session-message-store.repository"; diff --git a/apps/api/src/modules/sessions/application/session-model-call.service.ts b/apps/api/src/modules/sessions/application/session-model-call.service.ts deleted file mode 100644 index 2dda2922..00000000 --- a/apps/api/src/modules/sessions/application/session-model-call.service.ts +++ /dev/null @@ -1,4 +0,0 @@ -export { - upsertSessionModelCallUsage, - type SessionModelCallStatus, -} from "../infrastructure/session-model-call.repository"; diff --git a/apps/api/src/modules/sessions/application/session-process-events.service.ts b/apps/api/src/modules/sessions/application/session-process-events.service.ts index 90544d14..0d2b4f43 100644 --- a/apps/api/src/modules/sessions/application/session-process-events.service.ts +++ b/apps/api/src/modules/sessions/application/session-process-events.service.ts @@ -5,14 +5,21 @@ import { import type { SessionProcessEvent } from "@mosoo/contracts/session"; import { sessionEventsTable } from "@mosoo/db"; import type { AccountId, AppId, RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; -import { and, desc, eq } from "drizzle-orm"; +import { and, desc, eq, lt } from "drizzle-orm"; import { getAppDatabase } from "../../../platform/db/drizzle"; import { validationError } from "../../../platform/errors"; import { toIsoString } from "../../../time"; import type { AuthenticatedViewer } from "../../auth/application/viewer-auth.service"; import { getAppSessionParticipantTimelineAccess } from "../domain/session-access.policy"; -import { foldStreamedSessionEventRows } from "../domain/session-event-stream-fold"; +import { + excludeSessionEventStreams, + findLeftIncompleteSessionEventStreamKeys, + foldStreamedSessionEventRows, + getSessionEventStreamKey, +} from "../domain/session-event-stream-fold"; +import { formatStoredSessionEventContent } from "../domain/session-event-tool-content"; +import type { StoredToolStatus } from "../domain/session-event-tool-content"; export interface SessionEventProcessRow { content_text: string; @@ -24,11 +31,16 @@ export interface SessionEventProcessRow { process_type: SessionProcessEvent["type"]; run_id: SessionRunId | null; seq: number; + stream_id: string | null; + tool_name?: string | null; + tool_status?: StoredToolStatus | null; tokens: number | null; } const DEFAULT_PROCESS_EVENT_LIMIT = 500; const MAX_PROCESS_EVENT_LIMIT = 1000; +const PROCESS_EVENT_ROW_PAGE_SIZE = MAX_PROCESS_EVENT_LIMIT + 1; +const PROCESS_EVENT_RAW_ROW_SCAN_LIMIT = MAX_PROCESS_EVENT_LIMIT * 20; // Longer gaps may be permission or user waits, not continuous execution. const MAX_INFERRED_PROCESS_EVENT_DURATION_MS = 5 * 60 * 1000; @@ -45,6 +57,12 @@ interface SessionProcessEventAccess { updatedAt: string; } +interface SessionProcessEventWindow { + events: SessionProcessEvent[]; + recorded: boolean; + truncated: boolean; +} + function normalizeProcessEventLimit(limit: number | null | undefined): number { if (limit === null || limit === undefined) { return DEFAULT_PROCESS_EVENT_LIMIT; @@ -60,9 +78,7 @@ function normalizeProcessEventLimit(limit: number | null | undefined): number { function finalizeProcessEventDurations( projections: ProcessEventProjection[], ): SessionProcessEvent[] { - const sortedProjections = projections.toSorted( - (a, b) => a.startMs - b.startMs || a.order - b.order, - ); + const sortedProjections = projections.toSorted((a, b) => a.order - b.order); return sortedProjections.map((projection, index) => { const next = sortedProjections[index + 1] ?? null; @@ -91,13 +107,19 @@ function finalizeProcessEventDurations( }); } -function toProcessEventProjectionFromSessionEventRow( - row: SessionEventProcessRow, -): ProcessEventProjection { - return { +export function createSessionProcessEventsFromSessionEventRows( + rows: SessionEventProcessRow[], +): SessionProcessEvent[] { + const foldedRows = foldStreamedSessionEventRows(rows); + const projections = foldedRows.map((row) => ({ endMs: row.ended_at, event: { - content: row.content_text, + content: formatStoredSessionEventContent({ + contentText: row.content_text, + eventType: row.event_type, + toolName: row.tool_name ?? null, + toolStatus: row.tool_status ?? null, + }), durationMs: 0, id: row.id, occurredAt: toIsoString(row.occurred_at), @@ -108,18 +130,7 @@ function toProcessEventProjectionFromSessionEventRow( order: row.seq, runId: row.run_id, startMs: row.occurred_at, - }; -} - -export function createSessionProcessEventsFromSessionEventRows( - rows: SessionEventProcessRow[], - options: { foldStreamedRows?: boolean } = {}, -): SessionProcessEvent[] { - const foldedRows = - options.foldStreamedRows === false - ? rows - : foldStreamedSessionEventRows(rows, { flushOpenStreams: true }).rows; - const projections = foldedRows.map(toProcessEventProjectionFromSessionEventRow); + })); return finalizeProcessEventDurations(projections); } @@ -170,6 +181,107 @@ async function getThreadSessionProcessEventAccess( }; } +async function readSessionProcessEventWindow(input: { + database: D1Database; + limit: number; + sessionId: SessionId; +}): Promise { + const scannedRows: SessionEventProcessRow[] = []; + let beforeSeq: number | null = null; + let reachedStart = false; + + while (scannedRows.length < PROCESS_EVENT_RAW_ROW_SCAN_LIMIT) { + const rowCapacity = Math.min( + PROCESS_EVENT_ROW_PAGE_SIZE, + PROCESS_EVENT_RAW_ROW_SCAN_LIMIT - scannedRows.length, + ); + const querySize = rowCapacity + 1; + const filters = [ + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.visibility, "all_consumers"), + ]; + + if (beforeSeq !== null) { + filters.push(lt(sessionEventsTable.seq, beforeSeq)); + } + + const page = await getAppDatabase(input.database) + .select({ + content_text: sessionEventsTable.contentText, + ended_at: sessionEventsTable.endedAt, + event_type: sessionEventsTable.eventType, + id: sessionEventsTable.id, + occurred_at: sessionEventsTable.occurredAt, + process_status: sessionEventsTable.processStatus, + process_type: sessionEventsTable.processType, + run_id: sessionEventsTable.runId, + seq: sessionEventsTable.seq, + stream_id: sessionEventsTable.streamId, + tool_name: sessionEventsTable.toolName, + tool_status: sessionEventsTable.toolStatus, + tokens: sessionEventsTable.tokens, + }) + .from(sessionEventsTable) + .where(and(...filters)) + .orderBy(desc(sessionEventsTable.seq)) + .limit(querySize) + .all(); + + if (page.length === 0) { + reachedStart = true; + break; + } + + const scannedPage = page.slice(0, rowCapacity); + scannedRows.push(...scannedPage); + beforeSeq = scannedPage[scannedPage.length - 1]?.seq ?? beforeSeq; + const chronologicalRows = scannedRows.toReversed(); + const foldedRows = foldStreamedSessionEventRows(chronologicalRows); + const events = createSessionProcessEventsFromSessionEventRows(chronologicalRows); + + if (events.length > input.limit) { + const reachedDatabaseStart = page.length <= rowCapacity; + const incompleteStreams = reachedDatabaseStart + ? new Set() + : findLeftIncompleteSessionEventStreamKeys(chronologicalRows); + const rowsByEventId = new Map(foldedRows.map((row) => [row.id, row])); + const retainedStreamsAreComplete = events.slice(-input.limit).every((event) => { + const row = rowsByEventId.get(event.id); + const key = row === undefined ? null : getSessionEventStreamKey(row); + return key === null || !incompleteStreams.has(key); + }); + + if (!retainedStreamsAreComplete) { + continue; + } + + return { + events: events.slice(-input.limit), + recorded: true, + truncated: true, + }; + } + + if (page.length <= rowCapacity) { + reachedStart = true; + break; + } + } + + const chronologicalRows = scannedRows.toReversed(); + const incompleteStreams = reachedStart + ? new Set() + : findLeftIncompleteSessionEventStreamKeys(chronologicalRows); + const completeRows = excludeSessionEventStreams(chronologicalRows, incompleteStreams); + const events = createSessionProcessEventsFromSessionEventRows(completeRows); + + return { + events: events.slice(-input.limit), + recorded: scannedRows.length > 0, + truncated: !reachedStart || events.length > input.limit, + }; +} + async function listSessionProcessEvents( database: D1Database, viewer: AuthenticatedViewer, @@ -184,56 +296,29 @@ async function listSessionProcessEvents( appId: input.appId, sessionId: input.sessionId, }); - const rows = await getAppDatabase(database) - .select({ - content_text: sessionEventsTable.contentText, - ended_at: sessionEventsTable.endedAt, - event_type: sessionEventsTable.eventType, - id: sessionEventsTable.id, - occurred_at: sessionEventsTable.occurredAt, - process_status: sessionEventsTable.processStatus, - process_type: sessionEventsTable.processType, - run_id: sessionEventsTable.runId, - seq: sessionEventsTable.seq, - tokens: sessionEventsTable.tokens, - }) - .from(sessionEventsTable) - .where( - and( - eq(sessionEventsTable.sessionId, input.sessionId), - eq(sessionEventsTable.visibility, "all_consumers"), - ), - ) - .orderBy(desc(sessionEventsTable.seq)) - .limit(limit + 1) - .all(); - - if (rows.length === 0) { + const window = await readSessionProcessEventWindow({ + database, + limit, + sessionId: input.sessionId, + }); + + if (!window.recorded) { return [createNoRuntimeEventsRecordedEvent(session)]; } - const hasMore = rows.length > limit; - const processEvents = createSessionProcessEventsFromSessionEventRows( - rows.slice(0, limit).toReversed(), - ); - - if (!hasMore) { - return processEvents; + if (!window.truncated) { + return window.events; } - const firstEvent = processEvents[0] ?? null; - - if (firstEvent === null) { - return processEvents; - } + const firstEvent = window.events[0] ?? null; return [ createProcessEventsTruncatedEvent({ limit, - occurredAt: firstEvent.occurredAt, + occurredAt: firstEvent?.occurredAt ?? session.updatedAt, sessionId: session.id, }), - ...processEvents, + ...window.events, ]; } diff --git a/apps/api/src/modules/sessions/application/session-runtime-recovery-query.service.ts b/apps/api/src/modules/sessions/application/session-runtime-recovery-query.service.ts index f33cebb5..76280aa0 100644 --- a/apps/api/src/modules/sessions/application/session-runtime-recovery-query.service.ts +++ b/apps/api/src/modules/sessions/application/session-runtime-recovery-query.service.ts @@ -5,6 +5,7 @@ import { and, desc, eq, isNull, ne, or } from "drizzle-orm"; import { getAppDatabase } from "../../../platform/db/drizzle"; import { sanitizeProviderPrivateMarkup } from "../domain/provider-private-markup"; +import { resolveStoredSessionMessageContentReferences } from "../infrastructure/session-message-reference.repository"; const MAX_RUNTIME_RECOVERY_MESSAGES = 100; // Replayed history rides the driver boot payload and, for prompt-replay @@ -22,8 +23,11 @@ export async function getSessionRuntimeRecoveryMessages( ): Promise { const rows = await getAppDatabase(database) .select({ - content: sessionMessagesTable.contentText, + content_text: sessionMessagesTable.contentText, + id: sessionMessagesTable.id, + projection_format: sessionMessagesTable.projectionFormat, role: sessionMessagesTable.role, + session_run_id: sessionMessagesTable.sessionRunId, }) .from(sessionMessagesTable) .where( @@ -40,31 +44,42 @@ export async function getSessionRuntimeRecoveryMessages( .orderBy(desc(sessionMessagesTable.seq)) .limit(MAX_RUNTIME_RECOVERY_MESSAGES) .all(); - const newestFirst: DriverRecoveryMessage[] = []; let remainingChars = MAX_RUNTIME_RECOVERY_CONTENT_CHARS; for (const row of rows) { + const resolvedRows = await resolveStoredSessionMessageContentReferences( + database, + input.sessionId, + [row], + remainingChars + 1, + ); + + const [resolved] = resolvedRows; + if (resolved === undefined) { + continue; + } const content = - row.role === "assistant" ? sanitizeProviderPrivateMarkup(row.content).text : row.content; + resolved.role === "assistant" + ? sanitizeProviderPrivateMarkup(resolved.content_text).text + : resolved.content_text; if (content.trim().length === 0) { continue; } if (content.length > remainingChars) { - // Keep the window contiguous: stop at the first message that no longer - // fits instead of skipping past it. The newest message alone is kept - // truncated so an oversized latest exchange cannot erase all context. if (newestFirst.length === 0) { - newestFirst.push({ content: content.slice(0, remainingChars), role: row.role }); + newestFirst.push({ content: content.slice(0, remainingChars), role: resolved.role }); } - - break; + return newestFirst.toReversed(); } - newestFirst.push({ content, role: row.role }); + newestFirst.push({ content, role: resolved.role }); remainingChars -= content.length; + if (remainingChars === 0) { + return newestFirst.toReversed(); + } } return newestFirst.toReversed(); diff --git a/apps/api/src/modules/sessions/application/session-title.service.ts b/apps/api/src/modules/sessions/application/session-title.service.ts index 163546fe..834ea743 100644 --- a/apps/api/src/modules/sessions/application/session-title.service.ts +++ b/apps/api/src/modules/sessions/application/session-title.service.ts @@ -1,6 +1,7 @@ import type { RenameSessionInput, SessionSummary } from "@mosoo/contracts/session"; import { sessionsTable } from "@mosoo/db"; -import { and, eq, isNull } from "drizzle-orm"; +import type { PlatformId, RuntimeEventId, SessionId } from "@mosoo/id"; +import { and, eq, isNotNull, isNull, lt, or } from "drizzle-orm"; import { getAppDatabase } from "../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../time"; @@ -20,6 +21,91 @@ export interface RenameSessionRequest { viewer: AuthenticatedViewer; } +export interface DurableSessionAutoTitleInput { + creatorAccountId: PlatformId; + eventSeq: number; + sessionId: SessionId; + title: string; +} + +export function prepareDurableSessionAutoTitleProjection( + database: D1Database, + input: { + createdAt: number; + eventId: RuntimeEventId; + eventSeq: number; + semanticHash: string; + sessionId: SessionId; + title: string; + }, +): D1PreparedStatement[] { + if (!Number.isSafeInteger(input.eventSeq) || input.eventSeq < 0) { + throw new Error("Session auto-title event seq must be a non-negative safe integer."); + } + + const title = normalizeSessionTitle(input.title); + const receiptFence = `EXISTS ( + SELECT 1 + FROM session_event AS receipt + WHERE receipt.id = ? + AND receipt.session_id = ? + AND receipt.event_type = 'session.info.updated' + AND receipt.semantic_hash = ? + AND receipt.seq = ? + )`; + const update = database + .prepare( + `UPDATE session + SET auto_title_event_seq = ?, title = ?, updated_at = ? + WHERE id = ? + AND renamed = 0 + AND (title IS NULL OR auto_title_event_seq IS NOT NULL) + AND (auto_title_event_seq IS NULL OR auto_title_event_seq < ?) + AND ${receiptFence}`, + ) + .bind( + input.eventSeq, + title, + input.createdAt, + input.sessionId, + input.eventSeq, + input.eventId, + input.sessionId, + input.semanticHash, + input.eventSeq, + ); + const guard = database + .prepare( + `INSERT INTO session_event (id) + SELECT ? + WHERE ${receiptFence} + AND NOT EXISTS ( + SELECT 1 + FROM session + WHERE id = ? + AND ( + renamed = 1 + OR (auto_title_event_seq IS NULL AND title IS NOT NULL) + OR auto_title_event_seq > ? + OR (auto_title_event_seq = ? AND title = ?) + ) + )`, + ) + .bind( + input.eventId, + input.eventId, + input.sessionId, + input.semanticHash, + input.eventSeq, + input.sessionId, + input.eventSeq, + input.eventSeq, + title, + ); + + return [update, guard]; +} + async function hydrateUpdatedSessionSummary( database: D1Database, row: SessionSummaryRow, @@ -102,3 +188,69 @@ export async function autoTitleSession( return hydrateUpdatedSessionSummary(database, updated); } + +export async function applyDurableSessionAutoTitle( + database: D1Database, + input: DurableSessionAutoTitleInput, +): Promise { + if (!Number.isSafeInteger(input.eventSeq) || input.eventSeq < 0) { + throw new Error("Session auto-title event seq must be a non-negative safe integer."); + } + + const title = normalizeSessionTitle(input.title); + const appDatabase = getAppDatabase(database); + + await appDatabase + .update(sessionsTable) + .set({ + autoTitleEventSeq: input.eventSeq, + title, + updatedAt: currentTimestampMs(), + }) + .where( + and( + eq(sessionsTable.id, input.sessionId), + eq(sessionsTable.creatorAccountId, input.creatorAccountId), + eq(sessionsTable.renamed, false), + or(isNull(sessionsTable.title), isNotNull(sessionsTable.autoTitleEventSeq)), + or( + isNull(sessionsTable.autoTitleEventSeq), + lt(sessionsTable.autoTitleEventSeq, input.eventSeq), + ), + ), + ) + .run(); + + const stored = + (await appDatabase + .select({ + autoTitleEventSeq: sessionsTable.autoTitleEventSeq, + renamed: sessionsTable.renamed, + title: sessionsTable.title, + }) + .from(sessionsTable) + .where( + and( + eq(sessionsTable.id, input.sessionId), + eq(sessionsTable.creatorAccountId, input.creatorAccountId), + ), + ) + .limit(1) + .get()) ?? null; + + if (stored === null) { + throw new Error("Session not found for durable auto-title."); + } + + if (stored.renamed || (stored.autoTitleEventSeq === null && stored.title !== null)) { + return; + } + + if (stored.autoTitleEventSeq === null || stored.autoTitleEventSeq < input.eventSeq) { + throw new Error("Session auto-title CAS did not persist the durable event."); + } + + if (stored.autoTitleEventSeq === input.eventSeq && stored.title !== title) { + throw new Error("Session auto-title event seq was replayed with conflicting content."); + } +} diff --git a/apps/api/src/modules/sessions/application/session-viewer-events.service.ts b/apps/api/src/modules/sessions/application/session-viewer-events.service.ts index 1a89b17f..447dd447 100644 --- a/apps/api/src/modules/sessions/application/session-viewer-events.service.ts +++ b/apps/api/src/modules/sessions/application/session-viewer-events.service.ts @@ -1,4 +1 @@ -export { - publishSessionViewerEvents, - syncSessionViewerState, -} from "../infrastructure/session/client"; +export { syncSessionViewerState } from "../infrastructure/session/client"; diff --git a/apps/api/src/modules/sessions/domain/provider-private-markup.ts b/apps/api/src/modules/sessions/domain/provider-private-markup.ts index 6a3aee07..33a4a6db 100644 --- a/apps/api/src/modules/sessions/domain/provider-private-markup.ts +++ b/apps/api/src/modules/sessions/domain/provider-private-markup.ts @@ -1,7 +1,11 @@ -import { filterOpenAiPrivateCitations } from "@mosoo/agent-driver/provider-output"; +import { + filterOpenAiPrivateCitations, + OpenAiPrivateCitationStreamFilter, +} from "@mosoo/agent-driver/provider-output"; import type { SessionMessageSegment } from "@mosoo/contracts/session"; export const sanitizeProviderPrivateMarkup = filterOpenAiPrivateCitations; +export { OpenAiPrivateCitationStreamFilter as ProviderPrivateMarkupStreamFilter }; export function sanitizeAssistantMessageSegments( segments: SessionMessageSegment[], diff --git a/apps/api/src/modules/sessions/domain/session-cleanup-plan.ts b/apps/api/src/modules/sessions/domain/session-cleanup-plan.ts index 78a42d63..3401fd3e 100644 --- a/apps/api/src/modules/sessions/domain/session-cleanup-plan.ts +++ b/apps/api/src/modules/sessions/domain/session-cleanup-plan.ts @@ -1,4 +1,4 @@ -import type { DriverInstanceId, SandboxId, SessionId } from "@mosoo/id"; +import type { DriverInstanceId, SandboxId } from "@mosoo/id"; export const SESSION_ARCHIVE_CLEANUP_STEPS = [ "archive_session_row", @@ -7,12 +7,14 @@ export const SESSION_ARCHIVE_CLEANUP_STEPS = [ "stop_live_drivers", "normalize_runtime_lifecycle", "close_sandbox_session", + "complete_archive_session", ] as const; export const SESSION_DELETE_CLEANUP_STEPS = [ "archive_session_row", "load_cleanup_targets", "stop_live_drivers", + "normalize_runtime_lifecycle", "close_sandbox_session", "destroy_driver_objects", "destroy_session_object", @@ -24,104 +26,24 @@ export const SESSION_DELETE_CLEANUP_STEPS = [ export type SessionArchiveCleanupStep = (typeof SESSION_ARCHIVE_CLEANUP_STEPS)[number]; export type SessionDeleteCleanupStep = (typeof SESSION_DELETE_CLEANUP_STEPS)[number]; -export type SessionDeleteCleanupStepStatus = "completed" | "skipped"; -export type SessionArchiveCleanupStepStatus = "completed" | "skipped"; -export interface SessionArchiveCleanupStepOutcome { +export type SessionArchiveCleanupStepOutcome = { step: SessionArchiveCleanupStep; - status: SessionArchiveCleanupStepStatus; -} + status: "completed" | "skipped"; +}; -export interface SessionDeleteCleanupStepOutcome { +export type SessionDeleteCleanupStepOutcome = { step: SessionDeleteCleanupStep; - status: SessionDeleteCleanupStepStatus; -} + status: "completed" | "skipped"; +}; export interface SessionArchiveCleanupTargets { - liveDriverInstanceIds: readonly DriverInstanceId[]; + liveDriverInstances: readonly { generation: number; id: DriverInstanceId }[]; sandboxId: SandboxId | null; - sessionId: SessionId; } export interface SessionDeleteCleanupTargets { - associatedDriverInstanceIds: readonly DriverInstanceId[]; - liveDriverInstanceIds: readonly DriverInstanceId[]; + associatedDriverInstances: readonly { generation: number; id: DriverInstanceId }[]; + liveDriverInstances: readonly { generation: number; id: DriverInstanceId }[]; sandboxId: SandboxId | null; - sessionId: SessionId; -} - -export function completeSessionArchiveCleanupStep( - step: SessionArchiveCleanupStep, -): SessionArchiveCleanupStepOutcome { - return { status: "completed", step }; -} - -export function completeSessionDeleteCleanupStep( - step: SessionDeleteCleanupStep, -): SessionDeleteCleanupStepOutcome { - return { status: "completed", step }; -} - -export function skipSessionArchiveCleanupStep( - step: SessionArchiveCleanupStep, -): SessionArchiveCleanupStepOutcome { - return { status: "skipped", step }; -} - -export function skipSessionDeleteCleanupStep( - step: SessionDeleteCleanupStep, -): SessionDeleteCleanupStepOutcome { - return { status: "skipped", step }; -} - -export function shouldSkipSessionArchiveCleanupStep(input: { - step: SessionArchiveCleanupStep; - targets: SessionArchiveCleanupTargets; -}): boolean { - switch (input.step) { - case "close_sandbox_session": { - return input.targets.sandboxId === null; - } - case "stop_live_drivers": { - return input.targets.liveDriverInstanceIds.length === 0; - } - case "archive_session_row": - case "close_viewer_sockets": - case "load_runtime_targets": - case "normalize_runtime_lifecycle": { - return false; - } - default: { - throw new Error("Unknown session archive cleanup step."); - } - } -} - -export function shouldSkipSessionDeleteCleanupStep(input: { - step: SessionDeleteCleanupStep; - targets: SessionDeleteCleanupTargets; -}): boolean { - switch (input.step) { - case "close_sandbox_session": { - return input.targets.sandboxId === null; - } - case "destroy_driver_objects": - case "stop_live_drivers": { - return input.targets.liveDriverInstanceIds.length === 0; - } - case "delete_driver_rows": { - return input.targets.associatedDriverInstanceIds.length === 0; - } - case "archive_session_row": - case "delete_session_backups": - case "delete_session_files": - case "delete_session_row": - case "destroy_session_object": - case "load_cleanup_targets": { - return false; - } - default: { - throw new Error("Unknown session delete cleanup step."); - } - } } diff --git a/apps/api/src/modules/sessions/domain/session-event-stream-fold.ts b/apps/api/src/modules/sessions/domain/session-event-stream-fold.ts index 1ef897df..6f467607 100644 --- a/apps/api/src/modules/sessions/domain/session-event-stream-fold.ts +++ b/apps/api/src/modules/sessions/domain/session-event-stream-fold.ts @@ -3,18 +3,9 @@ import type { RuntimeEventId, SessionRunId } from "@mosoo/id"; // Streamed text events (message.delta / thought.delta) are persisted one row // per fragment so every accepted source identity stays durable (#274). Reading // them back verbatim renders each fragment as its own timeline entry, so read -// paths fold a fragment stream into a single row before projecting process -// events. The merge rules mirror the pre-persistence compactor -// (runtime-event-compaction.ts): deltas append, snapshots prefer the longer -// prefix-matching text. -// -// Rows carry no message identity, so a stream whose driver-side identity -// fractured (a dropped message_start splits one reply across several -// started/delta groups, YEF-884) still folds into several fragment rows, and -// the trailing message.added snapshot lands after its stream already closed. -// To heal both shapes, closed fragment rows are kept as supersede candidates: -// a snapshot whose text prefix-matches the concatenated fragments replaces -// them with a single row instead of rendering a duplicate. +// paths fold each identified stream into one row before projecting process +// events. message.added starts an authoritative snapshot, and later deltas +// append to it. export interface StreamFoldableSessionEventRow { content_text: string; @@ -25,21 +16,33 @@ export interface StreamFoldableSessionEventRow { process_type: string; run_id: SessionRunId | null; seq: number; + stream_id: string | null; tokens: number | null; } -export interface FoldedStreamedSessionEventRows { - /** - * Raw rows of streams that have not seen their closing event yet, in seq - * order. Only populated when `flushOpenStreams` is false; callers carry them - * into the next fold so a stream spanning reads still emits exactly once. - */ - openStreamRows: R[]; - rows: R[]; -} - type StreamRowPhase = "added" | "completed" | "delta" | "started"; +const MESSAGE_STREAM_AUTHORITY_RESET_EVENT_TYPES = [ + "message.cancelled", + "message.failed", + "message.started", +] as const; + +export const MESSAGE_STREAM_AUTHORITY_BOUNDARY_EVENT_TYPES = [ + "message.added", + ...MESSAGE_STREAM_AUTHORITY_RESET_EVENT_TYPES, +] as const; + +export const MESSAGE_STREAM_EVENT_TYPES = [ + ...MESSAGE_STREAM_AUTHORITY_BOUNDARY_EVENT_TYPES, + "message.completed", + "message.delta", +] as const; + +const messageStreamAuthorityResetEventTypes: ReadonlySet = new Set( + MESSAGE_STREAM_AUTHORITY_RESET_EVENT_TYPES, +); + interface StreamRowClassification { phase: StreamRowPhase; placeholder: string; @@ -50,9 +53,12 @@ const THOUGHT_PLACEHOLDER = "Agent thinking updated."; const streamRowClassifications: Readonly> = { "message.added": { phase: "added", placeholder: MESSAGE_PLACEHOLDER }, + "message.cancelled": { phase: "completed", placeholder: MESSAGE_PLACEHOLDER }, "message.completed": { phase: "completed", placeholder: MESSAGE_PLACEHOLDER }, "message.delta": { phase: "delta", placeholder: MESSAGE_PLACEHOLDER }, + "message.failed": { phase: "completed", placeholder: MESSAGE_PLACEHOLDER }, "message.started": { phase: "started", placeholder: MESSAGE_PLACEHOLDER }, + "thought.cancelled": { phase: "completed", placeholder: THOUGHT_PLACEHOLDER }, "thought.completed": { phase: "completed", placeholder: THOUGHT_PLACEHOLDER }, "thought.delta": { phase: "delta", placeholder: THOUGHT_PLACEHOLDER }, "thought.started": { phase: "started", placeholder: THOUGHT_PLACEHOLDER }, @@ -60,163 +66,304 @@ const streamRowClassifications: Readonly const terminalRunEventTypes = new Set(["run.cancelled", "run.completed", "run.failed"]); -interface OpenStreamGroup { +export interface MessageStreamTextFragment { + kind: "append" | "reset"; + text: string; +} + +export function getMessageStreamTextFragment( + row: StreamFoldableSessionEventRow, +): MessageStreamTextFragment | null { + const classification = streamRowClassifications[row.event_type]; + + if ( + classification === undefined || + !row.event_type.startsWith("message.") || + (classification.phase !== "added" && classification.phase !== "delta") + ) { + return null; + } + + return { + kind: classification.phase === "added" ? "reset" : "append", + text: row.content_text, + }; +} + +interface StreamGroup { contentText: string; + firstRow: R; + latestRow: R; + outputIndex: number | null; placeholder: string; - rows: R[]; + representative: R | null; runId: SessionRunId | null; + terminal: boolean; } -interface ClosedFragmentSegment { - content: string; - outputIndex: number; +interface ScannedStreamState { + leftBoundarySeen: boolean; + runId: SessionRunId | null; } -function appendStreamText(current: string, next: string): string { - if (next.length === 0) { - return current; - } - - return current.length === 0 ? next : `${current}${next}`; +function createStreamGroup( + row: R, + classification: StreamRowClassification, +): StreamGroup { + return { + contentText: "", + firstRow: row, + latestRow: row, + outputIndex: null, + placeholder: classification.placeholder, + representative: null, + runId: row.run_id, + terminal: false, + }; } -function mergeSnapshotText(current: string, next: string): string { - if (next.length === 0) { - return current; +export function getSessionEventStreamKey(row: StreamFoldableSessionEventRow): string | null { + if (streamRowClassifications[row.event_type] === undefined) { + return null; } - if (current.length === 0) { - return next; + if (row.stream_id === null) { + throw new Error(`Streamed session event ${row.id} has no stream identity.`); } - if (next.length > current.length && next.startsWith(current)) { - return next; - } + return JSON.stringify([row.run_id, row.process_type, row.stream_id]); +} - if (current.length >= next.length && current.startsWith(next)) { - return current; - } +// Reverse window scans must cross a stream's left boundary before treating +// its folded content and first sequence as complete. Migration identities +// that equal their row ID are deliberately row-scoped and already complete. +export function findLeftIncompleteSessionEventStreamKeys( + rows: readonly StreamFoldableSessionEventRow[], +): Set { + const runStarts = new Set(); + const streams = new Map(); - return `${current}${next}`; -} + for (const row of rows) { + if (row.event_type === "run.started" && row.run_id !== null) { + runStarts.add(row.run_id); + } + + const key = getSessionEventStreamKey(row); -function isSnapshotOfFragments(fragments: string, snapshot: string): boolean { - if (fragments.length === 0 || snapshot.length === 0) { - return false; + if (key === null) { + continue; + } + + const stream = streams.get(key) ?? { + leftBoundarySeen: false, + runId: row.run_id, + }; + stream.leftBoundarySeen ||= + row.event_type === "message.added" || + row.event_type === "message.started" || + row.event_type === "thought.started" || + row.stream_id === row.id; + streams.set(key, stream); } - return snapshot.startsWith(fragments) || fragments.startsWith(snapshot); + return new Set( + [...streams] + .filter( + ([, stream]) => + !stream.leftBoundarySeen && (stream.runId === null || !runStarts.has(stream.runId)), + ) + .map(([key]) => key), + ); } -function createOpenStreamGroup( - row: R, - classification: StreamRowClassification, -): OpenStreamGroup { - return { - contentText: "", - placeholder: classification.placeholder, - rows: [], - runId: row.run_id, - }; +export function excludeSessionEventStreams( + rows: readonly R[], + excludedKeys: ReadonlySet, +): R[] { + return rows.filter((row) => { + const key = getSessionEventStreamKey(row); + return key === null || !excludedKeys.has(key); + }); } function mergeStreamRow( - group: OpenStreamGroup, + group: StreamGroup, row: R, classification: StreamRowClassification, ): void { - group.rows.push(row); + group.latestRow = row; + group.contentText = mergeStreamText( + group.contentText, + row.content_text, + classification, + group.placeholder, + ); + + if (classification.phase === "added") { + if (!group.terminal) { + group.representative = row; + } + } else if (classification.phase === "completed") { + if (group.outputIndex === null) { + group.representative = row; + } + group.terminal = true; + } +} - // Fragments that carried no text are persisted with the process-draft - // placeholder; they mark stream boundaries and must not leak into the text. - if (row.content_text === group.placeholder) { - return; +function mergeStreamText( + currentText: string, + rowText: string, + classification: StreamRowClassification, + placeholder: string, +): string { + const text = + classification.phase === "added" || classification.phase === "delta" + ? rowText + : rowText === placeholder + ? "" + : rowText; + + if (classification.phase === "added") { + return text; } + if (classification.phase === "delta") { + return currentText + text; + } + if (classification.phase === "completed") { + return currentText || text; + } + return currentText; +} - group.contentText = - classification.phase === "delta" - ? appendStreamText(group.contentText, row.content_text) - : mergeSnapshotText(group.contentText, row.content_text); +/** + * Resolves one already-ordered assistant message stream only when its latest + * lifecycle is sealed by message.completed. A later start, delta, or + * authoritative snapshot opens the stream again, so an old terminal receipt + * cannot make an incomplete replacement look final. + */ +interface AuthoritativeMessageStream { + sealed: boolean; + text: string; } -function createFoldedStreamRow( - group: OpenStreamGroup, -): R | null { - const firstRow = group.rows[0]; - const lastRow = group.rows[group.rows.length - 1]; +export interface MessageStreamLifecycle { + authoritative: boolean; + sealed: boolean; +} - if (firstRow === undefined || lastRow === undefined || group.contentText.length === 0) { - return null; +export interface MessageStreamReducerState + extends AuthoritativeMessageStream, MessageStreamLifecycle {} + +export function createMessageStreamLifecycle(): MessageStreamLifecycle { + return { authoritative: false, sealed: false }; +} + +export function reduceMessageStreamLifecycle( + state: MessageStreamLifecycle, + eventType: string, +): MessageStreamLifecycle { + if (eventType === "message.added") { + return { authoritative: true, sealed: false }; + } + if (messageStreamAuthorityResetEventTypes.has(eventType)) { + return { authoritative: false, sealed: false }; } + if (eventType === "message.completed") { + return { authoritative: state.authoritative, sealed: state.authoritative }; + } + if (eventType === "message.delta") { + return { authoritative: state.authoritative, sealed: false }; + } + return state; +} - return { - ...lastRow, - content_text: group.contentText, - occurred_at: firstRow.occurred_at, - seq: firstRow.seq, - }; +export function createMessageStreamReducerState(): MessageStreamReducerState { + return { ...createMessageStreamLifecycle(), text: "" }; } -export function foldStreamedSessionEventRows( - rows: readonly R[], - options: { flushOpenStreams: boolean }, -): FoldedStreamedSessionEventRows { - const output: (R | null)[] = []; - const openGroups = new Map>(); - const fragmentSegments = new Map(); +export function reduceMessageStreamRow( + state: MessageStreamReducerState, + row: StreamFoldableSessionEventRow, + options?: { maxTextLength: number }, +): void { + const classification = streamRowClassifications[row.event_type]; + + if (classification === undefined || !row.event_type.startsWith("message.")) { + return; + } - function closeGroupAsFragment(key: string, group: OpenStreamGroup): void { - const folded = createFoldedStreamRow(group); + state.text = mergeStreamText(state.text, row.content_text, classification, MESSAGE_PLACEHOLDER); + if (options !== undefined && state.text.length > options.maxTextLength) { + state.text = state.text.slice(0, options.maxTextLength); + } + const lifecycle = reduceMessageStreamLifecycle(state, row.event_type); + state.authoritative = lifecycle.authoritative; + state.sealed = lifecycle.sealed; +} - if (folded === null) { - return; - } +function resolveAuthoritativeMessageStream( + rows: readonly StreamFoldableSessionEventRow[], +): AuthoritativeMessageStream | null { + const state = createMessageStreamReducerState(); - output.push(folded); - const segments = fragmentSegments.get(key) ?? []; - segments.push({ content: folded.content_text, outputIndex: output.length - 1 }); - fragmentSegments.set(key, segments); + for (const row of rows) { + reduceMessageStreamRow(state, row); } - // A snapshot row is the authoritative text of the message it closes. When - // its text extends (or repeats) the fragment rows already emitted for the - // same stream key, those fragments were partial views of this snapshot: - // collapse them into one row at the first fragment's timeline position. - function closeGroupWithSnapshot(key: string, folded: R | null): void { - if (folded === null) { - return; - } + return state.authoritative ? { sealed: state.sealed, text: state.text } : null; +} - const segments = fragmentSegments.get(key) ?? []; - fragmentSegments.delete(key); - const fragments = segments.map((segment) => segment.content).join(""); +export function resolveSealedMessageStream( + rows: readonly StreamFoldableSessionEventRow[], +): { text: string } | null { + const stream = resolveAuthoritativeMessageStream(rows); + return stream?.sealed === true ? { text: stream.text } : null; +} - if (!isSnapshotOfFragments(fragments, folded.content_text)) { - output.push(folded); - return; - } +function createFoldedStreamRow( + group: StreamGroup, + keepEmpty: boolean, +): R | null { + if (!keepEmpty && group.contentText.length === 0) { + return null; + } - const firstSegment = segments[0]; - const anchorRow = firstSegment === undefined ? null : output[firstSegment.outputIndex]; + const lastRow = group.representative ?? group.latestRow; - if (firstSegment === undefined || anchorRow === null || anchorRow === undefined) { - output.push(folded); - return; - } + return { + ...lastRow, + content_text: group.contentText, + occurred_at: group.firstRow.occurred_at, + seq: group.firstRow.seq, + }; +} - for (const segment of segments) { - output[segment.outputIndex] = null; - } +function emitStreamGroup( + output: R[], + group: StreamGroup, + keepEmpty: boolean, +): void { + const folded = createFoldedStreamRow(group, keepEmpty); - output[firstSegment.outputIndex] = { - ...folded, - content_text: - folded.content_text.length >= fragments.length ? folded.content_text : fragments, - occurred_at: anchorRow.occurred_at, - seq: anchorRow.seq, - }; + if (folded === null) { + return; } + if (group.outputIndex === null) { + group.outputIndex = output.length; + output.push(folded); + } else { + output[group.outputIndex] = folded; + } +} + +export function foldStreamedSessionEventRows( + rows: readonly R[], +): R[] { + const output: R[] = []; + const groups = new Map>(); + for (const row of rows) { const classification = streamRowClassifications[row.event_type]; @@ -225,10 +372,9 @@ export function foldStreamedSessionEventRows row !== null) }; } - return { - openStreamRows: [...openGroups.values()] - .flatMap((group) => group.rows) - .toSorted((a, b) => a.seq - b.seq), - rows: output.filter((row): row is R => row !== null), - }; + return output.toSorted((a, b) => a.seq - b.seq); } diff --git a/apps/api/src/modules/sessions/domain/session-event-tool-content.ts b/apps/api/src/modules/sessions/domain/session-event-tool-content.ts new file mode 100644 index 00000000..f626ec2e --- /dev/null +++ b/apps/api/src/modules/sessions/domain/session-event-tool-content.ts @@ -0,0 +1,35 @@ +export type StoredToolStatus = "cancelled" | "completed" | "failed" | "running"; + +function normalizeToolResult(value: string): string { + const normalized = value.replaceAll(/\s+/g, " ").trim(); + return normalized.length > 0 ? normalized : value; +} + +export function formatStoredSessionEventContent(input: { + contentText: string; + eventType: string; + toolName?: string | null; + toolStatus?: StoredToolStatus | null; +}): string { + if (input.eventType === "session.commands.updated") { + return "Session commands updated."; + } + if (input.eventType === "usage.updated") { + return "Usage updated."; + } + if (input.eventType !== "tool.call.updated" || input.toolStatus == null) { + return input.contentText; + } + + const name = input.toolName ?? "Tool"; + + if (input.toolStatus === "running" || input.toolStatus === "cancelled") { + return name; + } + + if (input.contentText.length > 0) { + return `${name} result: ${normalizeToolResult(input.contentText)}`; + } + + return input.toolStatus === "failed" ? `${name} failed.` : `${name} completed.`; +} diff --git a/apps/api/src/modules/sessions/domain/session-message-projection-parser.ts b/apps/api/src/modules/sessions/domain/session-message-projection-parser.ts index 1281c6e5..36e82e03 100644 --- a/apps/api/src/modules/sessions/domain/session-message-projection-parser.ts +++ b/apps/api/src/modules/sessions/domain/session-message-projection-parser.ts @@ -1,4 +1,6 @@ import type { SessionMessagePlanEntry, SessionMessageSegment } from "@mosoo/contracts/session"; +import { parsePlatformId } from "@mosoo/id"; +import type { SessionRunId } from "@mosoo/id"; export interface StoredSessionMessageProjectionInput { planJson: string | null; @@ -50,6 +52,10 @@ function readNullableString(value: unknown, fieldName: string): string | null { return value; } +function readOptionalNullableString(value: unknown, fieldName: string): string | null | undefined { + return value === undefined ? undefined : readNullableString(value, fieldName); +} + function readPlanPriority(value: unknown): SessionMessagePlanEntry["priority"] { if (value === "high" || value === "medium" || value === "low") { return value; @@ -91,12 +97,18 @@ function parseMessageSegment(raw: unknown): SessionMessageSegment { const tool = readString(raw["tool"], "tool segment tool"); const toolCallId = readString(raw["toolCallId"], "tool segment toolCallId"); + const rawRunId = readOptionalNullableString(raw["runId"], "tool segment runId"); + const runId = + rawRunId === undefined || rawRunId === null + ? rawRunId + : parsePlatformId(rawRunId, "tool segment runId"); if (kind === "tool_use") { return { argsText: readString(raw["argsText"], "tool_use argsText"), kind: "tool_use", path: readNullableString(raw["path"], "tool_use path"), + ...(runId === undefined ? {} : { runId }), tool, toolCallId, }; @@ -106,6 +118,7 @@ function parseMessageSegment(raw: unknown): SessionMessageSegment { return { kind: "tool_result", output: readString(raw["output"], "tool_result output"), + ...(runId === undefined ? {} : { runId }), tool, toolCallId, }; diff --git a/apps/api/src/modules/sessions/domain/session-runtime-event-authority.ts b/apps/api/src/modules/sessions/domain/session-runtime-event-authority.ts new file mode 100644 index 00000000..25d25f79 --- /dev/null +++ b/apps/api/src/modules/sessions/domain/session-runtime-event-authority.ts @@ -0,0 +1,33 @@ +import { + createRuntimeEventSemanticHash, + parseRuntimeEventEnvelope, + stringifyRuntimeEventSemanticValue, +} from "@mosoo/runtime-events"; +import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; + +export async function readSessionRuntimeEventSemanticAuthority(input: { + readonly eventJson: string | null; + readonly invalidMessage: string; + readonly missingMessage: string; + readonly semanticHash: string | null; +}): Promise { + if (input.eventJson === null || input.semanticHash === null) { + throw new Error(input.missingMessage); + } + + let event: RuntimeEventEnvelope; + try { + event = parseRuntimeEventEnvelope(JSON.parse(input.eventJson)); + } catch { + throw new Error(input.invalidMessage); + } + + if ( + stringifyRuntimeEventSemanticValue(event) !== input.eventJson || + (await createRuntimeEventSemanticHash(event)) !== input.semanticHash + ) { + throw new Error(input.invalidMessage); + } + + return event; +} diff --git a/apps/api/src/modules/sessions/domain/session-runtime-event-projection.ts b/apps/api/src/modules/sessions/domain/session-runtime-event-projection.ts index 56daa299..bd0afb4e 100644 --- a/apps/api/src/modules/sessions/domain/session-runtime-event-projection.ts +++ b/apps/api/src/modules/sessions/domain/session-runtime-event-projection.ts @@ -1,3 +1,4 @@ +import { createMcpExecuteFailedEventIdentity } from "@mosoo/agent-driver/events"; import type { SessionProcessEventStatus, SessionProcessEventType, @@ -5,16 +6,20 @@ import type { SessionRuntimeEventSource, SessionRuntimeEventVisibility, } from "@mosoo/contracts/session"; -import { parseJsonObject } from "@mosoo/contracts/validation"; -import type { JsonValue } from "@mosoo/contracts/validation"; -import type { SessionRunId } from "@mosoo/id"; +import { parsePlatformId } from "@mosoo/id"; +import type { DriverCommandId, SessionRunId } from "@mosoo/id"; import { + createRuntimeToolResultMessageId, createProcessDraftFromRuntimeEvent, getRuntimeEventSessionFamily, getRuntimeEventParticipantVisibility, getRuntimeEventSource, + readRuntimeEventMessageKey, + readRuntimeEventPayload, readRuntimeEventPermissionRequest, + readRuntimeEventString, readRuntimeEventToolCallUpdate, + readRuntimeEventToolOutputSnapshot, } from "@mosoo/runtime-events"; import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; @@ -22,13 +27,21 @@ export interface SessionRuntimeEventProjection { contentText: string; eventType: string; family: SessionRuntimeEventFamily; + mcpCommandId: DriverCommandId | null; processStatus: SessionProcessEventStatus; processType: SessionProcessEventType; runId: SessionRunId | null; source: SessionRuntimeEventSource; + streamId: string | null; toolCallId: string | null; + toolInputDeltaJson: string | null; toolInputJson: string | null; toolName: string | null; + toolOutputDeltaText: string | null; + toolOutputText: string | null; + toolParentMessageId: string | null; + toolResultMessageId: string | null; + toolStatus: "cancelled" | "completed" | "failed" | "running" | null; traceId: string | null; tokens: number | null; visibility: SessionRuntimeEventVisibility; @@ -40,67 +53,97 @@ function isKnownRuntimeEventSource(value: unknown): value is SessionRuntimeEvent return typeof value === "string" && knownRuntimeEventSources.has(value); } -function normalizeContentText(value: string): string { - const normalized = value.replaceAll(/\s+/g, " ").trim(); - return normalized.length > 0 ? normalized : value; -} - -function sortJsonValue(value: JsonValue): JsonValue { - if (Array.isArray(value)) { - return value.map(sortJsonValue); - } - - if (value !== null && typeof value === "object") { - return Object.fromEntries( - Object.entries(value) - .toSorted(([left], [right]) => (left < right ? -1 : left > right ? 1 : 0)) - .map(([key, entry]) => [key, sortJsonValue(entry)]), - ); - } - - return value; -} - -function toCanonicalToolInputJson(rawInput: string | null): string | null { - if (rawInput === null || rawInput.length === 0) { - return null; - } - - try { - const input = parseJsonObject(JSON.parse(rawInput), "Runtime tool input"); - - return Object.keys(input).length === 0 ? null : JSON.stringify(sortJsonValue(input)); - } catch { - return null; - } -} - -function readProjectedToolCall(event: RuntimeEventEnvelope): { +function readProjectedToolCall( + event: RuntimeEventEnvelope, + provenMcpCommandId: DriverCommandId | null, +): { + mcpCommandId: DriverCommandId | null; toolCallId: string | null; + toolInputDeltaJson: string | null; toolInputJson: string | null; toolName: string | null; + toolOutputDeltaText: string | null; + toolOutputText: string | null; + toolParentMessageId: string | null; + toolResultMessageId: string | null; + toolStatus: "cancelled" | "completed" | "failed" | "running" | null; } { if (event.kind === "tool.call.updated") { const toolCall = readRuntimeEventToolCallUpdate(event); + if (provenMcpCommandId !== null) { + const commandId = parsePlatformId( + event.correlationId, + "MCP command correlation ID", + ); + const sourceEventId = + toolCall.status === "failed" && + toolCall.kind === "mcp" && + toolCall.rawInput !== null && + toolCall.rawOutput !== null && + toolCall.title !== null + ? createMcpExecuteFailedEventIdentity({ + commandId, + rawInput: toolCall.rawInput, + rawOutput: toolCall.rawOutput, + title: toolCall.title, + toolCallId: toolCall.toolCallId, + }).sourceEventId + : `mcp.execute.${toolCall.status}:${commandId}`; + if ( + commandId !== provenMcpCommandId || + toolCall.kind !== "mcp" || + toolCall.status === "running" || + (toolCall.status === "failed" && + (toolCall.rawInput === null || toolCall.rawOutput === null || toolCall.title === null)) || + event.sourceEventId !== sourceEventId + ) { + throw new Error("Proven MCP command does not match its terminal tool event."); + } + } return { + mcpCommandId: provenMcpCommandId, toolCallId: toolCall.toolCallId, - // Running input is streamed and may be valid JSON before it is complete. - toolInputJson: - toolCall.status === "running" ? null : toCanonicalToolInputJson(toolCall.rawInput), - // Terminal titles are provider display labels and may differ from the stable start name. - toolName: toolCall.status === "running" ? toolCall.title || toolCall.kind : null, + toolInputDeltaJson: toolCall.rawInputDelta, + toolInputJson: toolCall.rawInput, + toolName: toolCall.title ?? toolCall.kind, + toolOutputDeltaText: toolCall.rawOutputDelta, + toolOutputText: readRuntimeEventToolOutputSnapshot(toolCall), + toolParentMessageId: toolCall.parentMessageId ?? toolCall.messageId, + toolResultMessageId: createRuntimeToolResultMessageId({ + runId: event.runId ?? null, + toolCallId: toolCall.toolCallId, + }), + toolStatus: toolCall.status, }; } const permission = readRuntimeEventPermissionRequest(event); return permission === null - ? { toolCallId: null, toolInputJson: null, toolName: null } + ? { + toolCallId: null, + mcpCommandId: null, + toolInputDeltaJson: null, + toolInputJson: null, + toolName: null, + toolOutputDeltaText: null, + toolOutputText: null, + toolParentMessageId: null, + toolResultMessageId: null, + toolStatus: null, + } : { toolCallId: permission.toolCallId, + mcpCommandId: null, + toolInputDeltaJson: null, toolInputJson: null, toolName: null, + toolOutputDeltaText: null, + toolOutputText: null, + toolParentMessageId: null, + toolResultMessageId: null, + toolStatus: null, }; } @@ -109,23 +152,33 @@ function readProjectedContentText( draft: ReturnType, ): string { if (event.kind !== "tool.call.updated") { - return draft.content; - } + const payload = readRuntimeEventPayload(event); + if (event.kind === "plan.updated") { + return JSON.stringify(Array.isArray(payload["entries"]) ? payload["entries"] : []); + } + if (event.kind === "session.commands.updated") { + return JSON.stringify(Array.isArray(payload["commands"]) ? payload["commands"] : []); + } + if (event.kind === "session.config.updated") { + return JSON.stringify({ + configOptions: Array.isArray(payload["options"]) ? payload["options"] : [], + }); + } + if (event.kind === "session.mode.updated") { + return JSON.stringify({ + currentModeId: readRuntimeEventString(payload, "currentMode"), + visibleModes: Array.isArray(payload["availableModes"]) ? payload["availableModes"] : [], + }); + } + if (event.kind === "usage.updated") { + return JSON.stringify(payload); + } - const toolCall = readRuntimeEventToolCallUpdate(event); - - if (toolCall.status !== "completed" && toolCall.status !== "failed") { return draft.content; } - const name = toolCall.title ?? toolCall.kind ?? "Tool"; - const result = toolCall.rawOutput ?? toolCall.content; - - if (result === null) { - return toolCall.status === "failed" ? `${name} failed.` : `${name} completed.`; - } - - return `${name} result: ${normalizeContentText(result)}`; + const toolCall = readRuntimeEventToolCallUpdate(event); + return toolCall.rawOutput ?? toolCall.rawOutputDelta ?? toolCall.content ?? ""; } function readProjectedProcessStatus( @@ -142,12 +195,27 @@ function readProjectedProcessStatus( return draft.status ?? "available"; } +function readProjectedVisibility(event: RuntimeEventEnvelope): SessionRuntimeEventVisibility { + // The canonical participant event still drives live AG-UI delivery. Its + // durable row is a state receipt, not a historical Process timeline entry. + return event.kind === "agent.tasks.replaced" + ? "owner_debug" + : getRuntimeEventParticipantVisibility(event); +} + +function readProjectedStreamId(event: RuntimeEventEnvelope): string | null { + return event.kind === "run.completed" + ? readRuntimeEventString(readRuntimeEventPayload(event), "finalMessageId") + : readRuntimeEventMessageKey(event); +} + export function createSessionRuntimeEventProjection( event: RuntimeEventEnvelope, + options?: { provenMcpCommandId: DriverCommandId | null }, ): SessionRuntimeEventProjection { const draft = createProcessDraftFromRuntimeEvent(event); const source = getRuntimeEventSource(event); - const toolCall = readProjectedToolCall(event); + const toolCall = readProjectedToolCall(event, options?.provenMcpCommandId ?? null); return { contentText: readProjectedContentText(event, draft), @@ -157,9 +225,10 @@ export function createSessionRuntimeEventProjection( processType: draft.type, runId: event.runId ?? null, source: isKnownRuntimeEventSource(source) ? source : "system", + streamId: readProjectedStreamId(event), ...toolCall, traceId: event.traceId ?? null, tokens: draft.tokens ?? null, - visibility: getRuntimeEventParticipantVisibility(event), + visibility: readProjectedVisibility(event), }; } diff --git a/apps/api/src/modules/sessions/domain/session-terminal-event-authority.ts b/apps/api/src/modules/sessions/domain/session-terminal-event-authority.ts new file mode 100644 index 00000000..c0b18b0f --- /dev/null +++ b/apps/api/src/modules/sessions/domain/session-terminal-event-authority.ts @@ -0,0 +1,55 @@ +import { + readRuntimeEventPayload, + readRuntimeEventString, + readRuntimeRunPayload, +} from "@mosoo/runtime-events"; +import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; + +import { readSessionRuntimeEventSemanticAuthority } from "./session-runtime-event-authority"; +import { createSessionRuntimeEventProjection } from "./session-runtime-event-projection"; + +export interface TerminalEventSemanticAuthority { + readonly event: RuntimeEventEnvelope; + readonly finalMessageId: string | null; + readonly lifecycle: "IDLE" | "TERMINATED"; +} + +export async function readTerminalEventSemanticAuthority(input: { + eventJson: string | null; + eventType: string; + runId: string; + semanticHash: string | null; + sessionId: string; + sourceEventId: string; + streamId: string | null; +}): Promise { + const event: RuntimeEventEnvelope = await readSessionRuntimeEventSemanticAuthority({ + eventJson: input.eventJson, + invalidMessage: `Session run ${input.runId} has an invalid durable terminal semantic authority.`, + missingMessage: `Session run ${input.runId} has no durable terminal semantic authority.`, + semanticHash: input.semanticHash, + }); + + const projection = createSessionRuntimeEventProjection(event); + const runtimeRun = readRuntimeRunPayload(event); + const finalMessageId = + event.kind === "run.completed" + ? readRuntimeEventString(readRuntimeEventPayload(event), "finalMessageId") + : null; + if ( + event.kind !== input.eventType || + event.runId !== input.runId || + event.sessionId !== input.sessionId || + event.sourceEventId !== input.sourceEventId || + projection.eventType !== input.eventType || + projection.runId !== input.runId || + projection.streamId !== input.streamId || + runtimeRun.run?.id !== input.runId || + (runtimeRun.lifecycle !== "IDLE" && runtimeRun.lifecycle !== "TERMINATED") || + finalMessageId !== input.streamId + ) { + throw new Error(`Session run ${input.runId} terminal semantic authority is invalid.`); + } + + return { event, finalMessageId, lifecycle: runtimeRun.lifecycle }; +} diff --git a/apps/api/src/modules/sessions/infrastructure/session-agent-task-snapshot.repository.ts b/apps/api/src/modules/sessions/infrastructure/session-agent-task-snapshot.repository.ts new file mode 100644 index 00000000..0702f929 --- /dev/null +++ b/apps/api/src/modules/sessions/infrastructure/session-agent-task-snapshot.repository.ts @@ -0,0 +1,160 @@ +import { AgentTaskSnapshot } from "@mosoo/contracts/session"; +import type { SessionRunStatus } from "@mosoo/contracts/session-run"; +import { parseSchemaValue } from "@mosoo/contracts/validation"; +import { sessionAgentTaskSnapshotsTable, sessionRunsTable, sessionsTable } from "@mosoo/db"; +import type { RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; +import { and, eq, inArray, isNull } from "drizzle-orm"; + +import { createErrorLogContext, logWarn } from "../../../platform/cloudflare/logger"; +import { getAppDatabase } from "../../../platform/db/drizzle"; +import { ACTIVE_SESSION_RUN_STATUSES } from "../../runtime/domain/session-run-lifecycle.machine"; + +export function prepareSessionAgentTaskSnapshotUpsert( + database: D1Database, + input: { + eventId: RuntimeEventId; + snapshot: AgentTaskSnapshot; + }, +): D1PreparedStatement { + return database + .prepare( + `INSERT INTO session_agent_task_snapshot ( + session_id, + run_id, + driver_instance_id, + seq, + tasks_json + ) + SELECT + receipt.session_id, + receipt.run_id, + current_run.driver_instance_id, + receipt.seq, + ? + FROM session_event AS receipt + JOIN session AS current_session + ON current_session.id = receipt.session_id + JOIN session_run AS current_run + ON current_run.id = receipt.run_id + AND current_run.session_id = receipt.session_id + WHERE receipt.id = ? + AND receipt.event_type = 'agent.tasks.replaced' + AND receipt.run_id = ? + AND current_session.last_run_id = receipt.run_id + AND current_session.archived_at IS NULL + AND current_session.status = 'RUNNING' + AND current_run.driver_instance_id = ? + AND current_run.status IN ('queued', 'booting', 'running', 'waiting_input') + ON CONFLICT (session_id) DO UPDATE SET + run_id = excluded.run_id, + driver_instance_id = excluded.driver_instance_id, + seq = excluded.seq, + tasks_json = excluded.tasks_json + WHERE excluded.seq > session_agent_task_snapshot.seq`, + ) + .bind( + JSON.stringify({ tasks: input.snapshot.tasks }), + input.eventId, + input.snapshot.runId, + input.snapshot.driverInstanceId, + ); +} + +export function parseStoredAgentTaskSnapshot(input: { + driverInstanceId: string; + runId: string; + sessionId: string; + tasksJson: string; +}): AgentTaskSnapshot | null { + try { + return parseSchemaValue(AgentTaskSnapshot, { + ...JSON.parse(input.tasksJson), + driverInstanceId: input.driverInstanceId, + runId: input.runId, + }); + } catch (error) { + logWarn("session.agent_task_snapshot.invalid", { + ...createErrorLogContext(error), + runId: input.runId, + sessionId: input.sessionId, + }); + return null; + } +} + +export interface SessionAgentTaskState { + driverInstanceId: string | null; + runId: SessionRunId; + runStatus: SessionRunStatus; + snapshot: AgentTaskSnapshot | null; +} + +export async function loadSessionAgentTaskState( + database: D1Database, + sessionId: SessionId, +): Promise { + const row = + (await getAppDatabase(database) + .select({ + runDriverInstanceId: sessionRunsTable.driverInstanceId, + runId: sessionRunsTable.id, + runStatus: sessionRunsTable.status, + taskDriverInstanceId: sessionAgentTaskSnapshotsTable.driverInstanceId, + taskRunId: sessionAgentTaskSnapshotsTable.runId, + tasksJson: sessionAgentTaskSnapshotsTable.tasksJson, + }) + .from(sessionsTable) + .innerJoin( + sessionRunsTable, + and( + eq(sessionRunsTable.id, sessionsTable.lastRunId), + eq(sessionRunsTable.sessionId, sessionsTable.id), + ), + ) + .leftJoin( + sessionAgentTaskSnapshotsTable, + and( + eq(sessionAgentTaskSnapshotsTable.sessionId, sessionsTable.id), + eq(sessionAgentTaskSnapshotsTable.runId, sessionRunsTable.id), + eq(sessionAgentTaskSnapshotsTable.driverInstanceId, sessionRunsTable.driverInstanceId), + ), + ) + .where( + and( + eq(sessionsTable.id, sessionId), + isNull(sessionsTable.archivedAt), + eq(sessionsTable.status, "RUNNING"), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ) + .limit(1) + .get()) ?? null; + + if (row === null) { + return null; + } + + const snapshot = + row.taskDriverInstanceId === null || row.taskRunId === null || row.tasksJson === null + ? null + : parseStoredAgentTaskSnapshot({ + driverInstanceId: row.taskDriverInstanceId, + runId: row.taskRunId, + sessionId, + tasksJson: row.tasksJson, + }); + + return { + driverInstanceId: row.runDriverInstanceId, + runId: row.runId, + runStatus: row.runStatus, + snapshot, + }; +} + +export async function loadSessionAgentTaskSnapshot( + database: D1Database, + sessionId: SessionId, +): Promise { + return (await loadSessionAgentTaskState(database, sessionId))?.snapshot ?? null; +} diff --git a/apps/api/src/modules/sessions/infrastructure/session-message-event-stream.repository.ts b/apps/api/src/modules/sessions/infrastructure/session-message-event-stream.repository.ts new file mode 100644 index 00000000..afbcec01 --- /dev/null +++ b/apps/api/src/modules/sessions/infrastructure/session-message-event-stream.repository.ts @@ -0,0 +1,301 @@ +import type { SessionProcessEvent } from "@mosoo/contracts/session"; +import { sessionEventsTable } from "@mosoo/db"; +import type { SessionId, SessionRunId } from "@mosoo/id"; +import { + and, + asc, + desc, + eq, + gt, + gte, + inArray, + isNotNull, + isNull, + lte, + ne, + notInArray, + or, +} from "drizzle-orm"; + +import { getAppDatabase } from "../../../platform/db/drizzle"; +import { + createMessageStreamLifecycle, + createMessageStreamReducerState, + MESSAGE_STREAM_EVENT_TYPES, + reduceMessageStreamLifecycle, + reduceMessageStreamRow, +} from "../domain/session-event-stream-fold"; +import type { StreamFoldableSessionEventRow } from "../domain/session-event-stream-fold"; + +const SESSION_MESSAGE_EVENT_PAGE_SIZE = 1_000; +const SESSION_MESSAGE_EVENT_TYPE_SET: ReadonlySet = new Set(MESSAGE_STREAM_EVENT_TYPES); + +export interface SessionMessageEventStreamCursor { + endSeq: number; + finish: boolean; + outputOffset: number; + startSeq: number; +} + +export interface SessionMessageEventStreamIdentity { + processType: SessionProcessEvent["type"]; + runId: SessionRunId | null; + sessionId: SessionId; + streamId: string; +} + +export interface SessionMessageEventRow extends StreamFoldableSessionEventRow { + process_status: SessionProcessEvent["status"]; + process_type: SessionProcessEvent["type"]; +} + +async function assertExactMessageStreamIdentity( + database: D1Database, + input: SessionMessageEventStreamIdentity & { endSeq?: number }, +): Promise { + const conflictingRun = + input.runId === null + ? isNotNull(sessionEventsTable.runId) + : or(isNull(sessionEventsTable.runId), ne(sessionEventsTable.runId, input.runId)); + const collision = await getAppDatabase(database) + .select({ id: sessionEventsTable.id }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.streamId, input.streamId), + inArray(sessionEventsTable.eventType, [...MESSAGE_STREAM_EVENT_TYPES]), + input.endSeq === undefined ? undefined : lte(sessionEventsTable.seq, input.endSeq), + or(ne(sessionEventsTable.processType, input.processType), conflictingRun), + ), + ) + .limit(1) + .get(); + + if (collision !== undefined) { + throw new Error(`Session message stream ${input.streamId} has conflicting identity rows.`); + } +} + +export async function* iteratePublicSessionMessageEventRows( + database: D1Database, + input: SessionMessageEventStreamIdentity & { cursor: SessionMessageEventStreamCursor }, +): AsyncGenerator { + let afterSeq: number | null = null; + await assertExactMessageStreamIdentity(database, { ...input, endSeq: input.cursor.endSeq }); + const streamScope = [ + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.streamId, input.streamId), + eq(sessionEventsTable.processType, input.processType), + lte(sessionEventsTable.seq, input.cursor.endSeq), + input.runId === null + ? isNull(sessionEventsTable.runId) + : eq(sessionEventsTable.runId, input.runId), + ]; + const invalidRow = await getAppDatabase(database) + .select({ + eventType: sessionEventsTable.eventType, + visibility: sessionEventsTable.visibility, + }) + .from(sessionEventsTable) + .where( + and( + ...streamScope, + or( + ne(sessionEventsTable.visibility, "all_consumers"), + notInArray(sessionEventsTable.eventType, [...MESSAGE_STREAM_EVENT_TYPES]), + ), + ), + ) + .limit(1) + .get(); + if (invalidRow?.visibility !== undefined) { + throw new Error( + invalidRow.visibility === "all_consumers" + ? `Session message stream ${input.streamId} has unsupported event ${invalidRow.eventType}.` + : `Session message stream ${input.streamId} has mixed visibility.`, + ); + } + + for (;;) { + const filters = [ + ...streamScope, + eq(sessionEventsTable.visibility, "all_consumers"), + gte(sessionEventsTable.seq, input.cursor.startSeq), + ]; + + if (afterSeq !== null) { + filters.push(gt(sessionEventsTable.seq, afterSeq)); + } + + const page = await getAppDatabase(database) + .select({ + content_text: sessionEventsTable.contentText, + ended_at: sessionEventsTable.endedAt, + event_type: sessionEventsTable.eventType, + id: sessionEventsTable.id, + occurred_at: sessionEventsTable.occurredAt, + process_status: sessionEventsTable.processStatus, + process_type: sessionEventsTable.processType, + run_id: sessionEventsTable.runId, + seq: sessionEventsTable.seq, + stream_id: sessionEventsTable.streamId, + tokens: sessionEventsTable.tokens, + }) + .from(sessionEventsTable) + .where(and(...filters)) + .orderBy(asc(sessionEventsTable.seq)) + .limit(SESSION_MESSAGE_EVENT_PAGE_SIZE) + .all(); + + for (const row of page) { + if (!SESSION_MESSAGE_EVENT_TYPE_SET.has(row.event_type)) { + throw new Error( + `Session message stream ${input.streamId} has unsupported event ${row.event_type}.`, + ); + } + yield row; + } + + if (page.length < SESSION_MESSAGE_EVENT_PAGE_SIZE) { + return; + } + + afterSeq = page[page.length - 1]?.seq ?? afterSeq; + } +} + +export async function readSealedPublicSessionMessage( + database: D1Database, + input: SessionMessageEventStreamIdentity & { endSeq?: number }, +): Promise<{ text: string } | null> { + const scope = [ + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.streamId, input.streamId), + eq(sessionEventsTable.processType, input.processType), + input.runId === null + ? isNull(sessionEventsTable.runId) + : eq(sessionEventsTable.runId, input.runId), + ]; + const endSeq = + input.endSeq ?? + ( + await getAppDatabase(database) + .select({ seq: sessionEventsTable.seq }) + .from(sessionEventsTable) + .where(and(...scope)) + .orderBy(desc(sessionEventsTable.seq)) + .limit(1) + .get() + )?.seq; + if (endSeq === undefined) { + return null; + } + const state = createMessageStreamReducerState(); + + for await (const row of iteratePublicSessionMessageEventRows(database, { + ...input, + cursor: { + endSeq, + finish: true, + outputOffset: 0, + startSeq: 0, + }, + })) { + reduceMessageStreamRow(state, row); + } + + return state.authoritative && state.sealed ? { text: state.text } : null; +} + +export interface PublicSessionMessageStreamSealState { + authoritative: boolean; + sealed: boolean; +} + +export async function readPublicSessionMessageStreamSealState( + database: D1Database, + input: SessionMessageEventStreamIdentity, +): Promise { + await assertExactMessageStreamIdentity(database, input); + const scope = [ + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.streamId, input.streamId), + eq(sessionEventsTable.processType, input.processType), + input.runId === null + ? isNull(sessionEventsTable.runId) + : eq(sessionEventsTable.runId, input.runId), + ]; + const endSeq = ( + await getAppDatabase(database) + .select({ seq: sessionEventsTable.seq }) + .from(sessionEventsTable) + .where(and(...scope)) + .orderBy(desc(sessionEventsTable.seq)) + .limit(1) + .get() + )?.seq; + if (endSeq === undefined) { + return { authoritative: false, sealed: false }; + } + + const invalidRow = await getAppDatabase(database) + .select({ + eventType: sessionEventsTable.eventType, + visibility: sessionEventsTable.visibility, + }) + .from(sessionEventsTable) + .where( + and( + ...scope, + lte(sessionEventsTable.seq, endSeq), + or( + ne(sessionEventsTable.visibility, "all_consumers"), + notInArray(sessionEventsTable.eventType, [...MESSAGE_STREAM_EVENT_TYPES]), + ), + ), + ) + .limit(1) + .get(); + if (invalidRow?.visibility !== undefined) { + throw new Error( + invalidRow.visibility === "all_consumers" + ? `Session message stream ${input.streamId} has unsupported event ${invalidRow.eventType}.` + : `Session message stream ${input.streamId} has mixed visibility.`, + ); + } + + let afterSeq: number | null = null; + let state = createMessageStreamLifecycle(); + for (;;) { + const page = await getAppDatabase(database) + .select({ eventType: sessionEventsTable.eventType, seq: sessionEventsTable.seq }) + .from(sessionEventsTable) + .where( + and( + ...scope, + eq(sessionEventsTable.visibility, "all_consumers"), + lte(sessionEventsTable.seq, endSeq), + afterSeq === null ? undefined : gt(sessionEventsTable.seq, afterSeq), + ), + ) + .orderBy(asc(sessionEventsTable.seq)) + .limit(SESSION_MESSAGE_EVENT_PAGE_SIZE) + .all(); + + for (const row of page) { + state = reduceMessageStreamLifecycle(state, row.eventType); + } + if (page.length < SESSION_MESSAGE_EVENT_PAGE_SIZE) { + return state; + } + afterSeq = page[page.length - 1]?.seq ?? afterSeq; + } +} + +export async function isSealedPublicSessionMessageStream( + database: D1Database, + input: SessionMessageEventStreamIdentity, +): Promise { + return (await readPublicSessionMessageStreamSealState(database, input)).sealed; +} diff --git a/apps/api/src/modules/sessions/infrastructure/session-message-reference.repository.ts b/apps/api/src/modules/sessions/infrastructure/session-message-reference.repository.ts new file mode 100644 index 00000000..19e0a87c --- /dev/null +++ b/apps/api/src/modules/sessions/infrastructure/session-message-reference.repository.ts @@ -0,0 +1,1310 @@ +import { + applyAgUiEventToSessionLiveState, + applyToolCallUpdateToSessionLiveState, + createInitialSessionLiveState, + createServerCustomEvent, + EventType, + MOSOO_CUSTOM_EVENT, + parseAgUiSessionEvent, + parseNullableSessionUsageSummary, +} from "@mosoo/ag-ui-session"; +import type { + AgUiSessionEvent, + SessionLiveState, + SessionViewPlanEntry, +} from "@mosoo/ag-ui-session"; +import { sessionEventsTable } from "@mosoo/db"; +import type { SessionId, SessionMessageId, SessionRunId } from "@mosoo/id"; +import { createSessionRunTerminalSourceId } from "@mosoo/runtime-events"; +import { and, asc, desc, eq, gt, gte, inArray, lte, ne, notInArray, or, sql } from "drizzle-orm"; +import type { SQL } from "drizzle-orm"; + +import { getAppDatabase } from "../../../platform/db/drizzle"; +import { + ProviderPrivateMarkupStreamFilter, + sanitizeProviderPrivateMarkup, +} from "../domain/provider-private-markup"; +import { + createMessageStreamReducerState, + MESSAGE_STREAM_AUTHORITY_BOUNDARY_EVENT_TYPES, + MESSAGE_STREAM_EVENT_TYPES, + reduceMessageStreamRow, +} from "../domain/session-event-stream-fold"; +import type { + MessageStreamReducerState, + StreamFoldableSessionEventRow, +} from "../domain/session-event-stream-fold"; +import { readTerminalEventSemanticAuthority } from "../domain/session-terminal-event-authority"; + +const REFERENCE_EVENT_PAGE_SIZE = 1_000; +const BOUNDED_REFERENCE_EVENT_PAGE_SIZE = 1_000; +const BOUNDED_REFERENCE_EVENT_PAGE_CHARS = 512 * 1_024; +const RUN_TERMINAL_EVENT_TYPES = ["run.cancelled", "run.completed", "run.failed"]; +const MESSAGE_EVENT_TYPE_SET: ReadonlySet = new Set(MESSAGE_STREAM_EVENT_TYPES); +const THOUGHT_EVENT_TYPES = new Set([ + "thought.cancelled", + "thought.completed", + "thought.delta", + "thought.started", +]); +const THOUGHT_PROCESS_TYPE = "agent.thinking.delta"; +const USER_MESSAGE_PROCESS_TYPE = "user.message"; +const ACTIVE_SESSION_STATE_EVENT_TYPES = [ + "session.commands.updated", + "session.config.updated", + "session.mode.updated", + "usage.updated", +] as const; + +export interface StoredSessionMessageReferenceRow { + content_text: string; + id: SessionMessageId; + plan_json: string | null; + projection_format: "event_stream_v3" | "materialized"; + role: "assistant" | "user"; + segments_json: string | null; + session_run_id: SessionRunId | null; +} + +interface ReferenceEventRow extends StreamFoldableSessionEventRow { + tool_call_id: string | null; + tool_input_delta_json: string | null; + tool_input_json: string | null; + tool_name: string | null; + tool_output_delta_text: string | null; + tool_output_text: string | null; + tool_parent_message_id: string | null; + tool_result_message_id: string | null; + tool_status: "cancelled" | "completed" | "failed" | "running" | null; + visibility: string; +} + +interface ReferenceState { + isCarrier: boolean; + isFinal: boolean; + live: SessionLiveState; + message: MessageStreamReducerState; + planJson: string | null; +} + +interface ContentReferenceState { + filter: ProviderPrivateMarkupStreamFilter; + hasStreamText: boolean; + message: MessageStreamReducerState; +} + +interface SealedReferenceCursor { + endSeq: number; + startSeq: number; +} + +interface LightweightReference { + messageId: SessionMessageId; + runId: SessionRunId; +} + +interface CanonicalReference extends LightweightReference { + isCarrier: boolean; + isFinal: boolean; +} + +interface TerminalReference { + messageId: string | null; + seq: number; +} + +interface SessionStateEventRow { + content_text: string; + event_type: string; + seq: number; + visibility: string; +} + +interface ToolRoute { + authority: "carrier" | "final" | null; + firstOutputSeq: number | null; + firstParentSeq: number | null; + parentMessageId: string | null; + referenceKeys: Set; + resultMessageId: string | null; + runId: SessionRunId; + toolCallId: string; +} + +export interface StoredSessionMessageContentReferenceRow { + content_text: string; + id: SessionMessageId; + projection_format: "event_stream_v3" | "materialized"; + role: "assistant" | "user"; + session_run_id: SessionRunId | null; +} + +function referenceKey(runId: SessionRunId, messageId: string): string { + return JSON.stringify([runId, messageId]); +} + +function potentialReference( + row: StoredSessionMessageContentReferenceRow, +): LightweightReference | null { + return row.projection_format === "event_stream_v3" && + row.role === "assistant" && + row.session_run_id !== null + ? { messageId: row.id, runId: row.session_run_id } + : null; +} + +function requireTerminalReferences( + candidates: readonly LightweightReference[], + terminalReferences: ReadonlyMap, +): CanonicalReference[] { + const references: CanonicalReference[] = []; + for (const reference of candidates) { + const terminal = terminalReferences.get(reference.runId); + const isCarrier = String(reference.messageId) === String(reference.runId); + const isFinal = terminal?.messageId === reference.messageId; + if (terminal === undefined || (!isCarrier && !isFinal)) { + throw new Error( + `Stored event-stream assistant ${reference.messageId} has no exact terminal authority.`, + ); + } + references.push({ ...reference, isCarrier, isFinal }); + } + return references; +} + +function createReferenceState( + sessionId: SessionId, + reference: Pick = { + isCarrier: false, + isFinal: false, + }, +): ReferenceState { + return { + ...reference, + live: createInitialSessionLiveState({ sessionId, title: null, viewerId: "reference" }), + message: createMessageStreamReducerState(), + planJson: null, + }; +} + +function createContentReferenceState(): ContentReferenceState { + return { + filter: new ProviderPrivateMarkupStreamFilter(), + hasStreamText: false, + message: createMessageStreamReducerState(), + }; +} + +function reduceSanitizedMessageStreamRow( + state: ContentReferenceState, + row: StreamFoldableSessionEventRow, + maxTextLength: number, +): void { + if (row.event_type === "message.started" || row.event_type === "message.added") { + state.filter = new ProviderPrivateMarkupStreamFilter(); + state.hasStreamText = false; + } + + if (row.event_type === "message.added" || row.event_type === "message.delta") { + state.hasStreamText ||= row.content_text.length > 0; + reduceMessageStreamRow( + state.message, + { ...row, content_text: state.filter.push(row.content_text).text }, + { maxTextLength }, + ); + return; + } + + if (row.event_type === "message.completed") { + const trailingText = state.filter.finish().text; + if (trailingText.length > 0) { + reduceMessageStreamRow( + state.message, + { ...row, content_text: trailingText, event_type: "message.delta" }, + { maxTextLength }, + ); + } + reduceMessageStreamRow( + state.message, + { + ...row, + content_text: state.hasStreamText + ? "" + : sanitizeProviderPrivateMarkup(row.content_text).text, + }, + { maxTextLength }, + ); + return; + } + + reduceMessageStreamRow(state.message, row, { maxTextLength }); +} + +function toolKey(runId: SessionRunId, toolCallId: string): string { + return JSON.stringify([runId, toolCallId]); +} + +function mergeToolRouteIdentity( + route: ToolRoute, + row: { + tool_parent_message_id: string | null; + tool_result_message_id: string | null; + }, +): void { + if ( + (route.parentMessageId !== null && + row.tool_parent_message_id !== null && + route.parentMessageId !== row.tool_parent_message_id) || + (route.resultMessageId !== null && + row.tool_result_message_id !== null && + route.resultMessageId !== row.tool_result_message_id) + ) { + throw new Error(`Stored tool call ${route.toolCallId} changed its message identity.`); + } + route.parentMessageId ??= row.tool_parent_message_id; + route.resultMessageId ??= row.tool_result_message_id; +} + +function applyAgUiEvents(state: ReferenceState, events: readonly AgUiSessionEvent[]): void { + for (const event of events) { + state.live = applyAgUiEventToSessionLiveState(state.live, event); + } +} + +function messageAgUiEvent(row: ReferenceEventRow): AgUiSessionEvent { + const messageId = row.stream_id; + if (messageId === null) { + throw new Error(`Stored ${row.event_type} event is missing its stream identity.`); + } + + if (row.process_type === THOUGHT_PROCESS_TYPE) { + switch (row.event_type) { + case "thought.started": + return { messageId, role: "reasoning", type: EventType.REASONING_MESSAGE_START }; + case "thought.delta": + return { + delta: row.content_text, + messageId, + type: EventType.REASONING_MESSAGE_CONTENT, + }; + default: + return { messageId, type: EventType.REASONING_MESSAGE_END }; + } + } + + switch (row.event_type) { + case "message.added": + return { + delta: row.content_text, + messageId, + role: row.process_type === USER_MESSAGE_PROCESS_TYPE ? "user" : "assistant", + type: EventType.TEXT_MESSAGE_CHUNK, + }; + case "message.started": + return { + messageId, + role: row.process_type === USER_MESSAGE_PROCESS_TYPE ? "user" : "assistant", + type: EventType.TEXT_MESSAGE_START, + }; + case "message.delta": + return { delta: row.content_text, messageId, type: EventType.TEXT_MESSAGE_CONTENT }; + default: + return { messageId, type: EventType.TEXT_MESSAGE_END }; + } +} + +function applyReferenceEvent( + states: ReadonlyMap, + statesByRun: ReadonlyMap, + ceilings: ReadonlyMap, + toolRoutes: ReadonlyMap, + row: ReferenceEventRow, + options?: { activeArtifacts?: boolean }, +): void { + if (row.run_id === null || row.seq > (ceilings.get(row.run_id) ?? -1)) { + return; + } + + if (row.event_type === "plan.updated") { + const runStates = statesByRun.get(row.run_id) ?? []; + if (runStates.length > 0 && row.visibility !== "all_consumers") { + throw new Error(`Session run ${row.run_id} has a mixed-visibility plan stream.`); + } + const plan: unknown = JSON.parse(row.content_text); + if (!Array.isArray(plan)) { + throw new Error(`Session run ${row.run_id} has an invalid plan projection.`); + } + for (const state of runStates) { + state.planJson = row.content_text; + state.live = applyAgUiEventToSessionLiveState( + state.live, + createServerCustomEvent(MOSOO_CUSTOM_EVENT.sessionPlanUpdated.name, { + plan: plan as SessionViewPlanEntry[], + }), + ); + } + return; + } + + if ( + options?.activeArtifacts === true && + (row.process_type === THOUGHT_PROCESS_TYPE || row.process_type === USER_MESSAGE_PROCESS_TYPE) + ) { + if (row.visibility !== "all_consumers") { + throw new Error(`Session run ${row.run_id} has a mixed-visibility message stream.`); + } + const supported = + row.process_type === THOUGHT_PROCESS_TYPE + ? THOUGHT_EVENT_TYPES.has(row.event_type) + : MESSAGE_EVENT_TYPE_SET.has(row.event_type); + if (!supported) { + throw new Error(`Stored active message stream has unsupported event ${row.event_type}.`); + } + const event = messageAgUiEvent(row); + for (const state of new Set(statesByRun.get(row.run_id) ?? [])) { + state.live = applyAgUiEventToSessionLiveState(state.live, event); + } + return; + } + + if (row.process_type === "agent.message.delta") { + if (row.stream_id === null) { + return; + } + const state = states.get(referenceKey(row.run_id, row.stream_id)); + if (state === undefined || !state.isFinal) { + return; + } + if (row.visibility !== "all_consumers") { + throw new Error("Stored assistant reference has a mixed-visibility message stream."); + } + if (!MESSAGE_EVENT_TYPE_SET.has(row.event_type)) { + throw new Error(`Stored assistant reference has unsupported event ${row.event_type}.`); + } + reduceMessageStreamRow(state.message, row); + applyAgUiEvents(state, [messageAgUiEvent(row)]); + return; + } + + if (row.event_type !== "tool.call.updated") { + return; + } + if (row.visibility !== "all_consumers") { + throw new Error("Stored assistant reference has a mixed-visibility tool stream."); + } + if (row.tool_call_id === null) { + return; + } + const route = toolRoutes.get(toolKey(row.run_id, row.tool_call_id)); + if (route === undefined) { + return; + } + mergeToolRouteIdentity(route, row); + const resultMessageId = route.resultMessageId ?? route.parentMessageId; + if (resultMessageId === null) { + throw new Error(`Stored tool call ${row.tool_call_id} has no message identity.`); + } + const routedStates = new Set(); + for (const key of route.referenceKeys) { + const state = states.get(key); + if (state === undefined) { + continue; + } + routedStates.add(state); + } + for (const state of routedStates) { + const carrierAuthority = route.authority === "carrier"; + state.live = applyToolCallUpdateToSessionLiveState(state.live, { + inputDelta: row.tool_input_delta_json, + inputSnapshot: row.tool_input_json, + outputDelta: row.tool_output_delta_text, + outputSnapshot: row.tool_output_text, + parentMessageId: + row.tool_parent_message_id === null + ? null + : carrierAuthority + ? route.runId + : row.tool_parent_message_id, + resultMessageId: carrierAuthority ? route.runId : resultMessageId, + runId: row.run_id, + toolCallId: row.tool_call_id, + toolName: row.tool_name ?? "Tool", + }); + } +} + +function sessionStateAgUiEvent(row: SessionStateEventRow): AgUiSessionEvent { + if (row.visibility !== "all_consumers") { + throw new Error(`Session state event ${row.event_type} has mixed visibility.`); + } + const content: unknown = JSON.parse(row.content_text); + if (row.event_type === "plan.updated") { + if (!Array.isArray(content)) { + throw new Error("Session plan event has an invalid projection."); + } + return createServerCustomEvent(MOSOO_CUSTOM_EVENT.sessionPlanUpdated.name, { + plan: content as SessionViewPlanEntry[], + }); + } + const names = { + "session.commands.updated": MOSOO_CUSTOM_EVENT.sessionCommandsUpdated.name, + "session.config.updated": MOSOO_CUSTOM_EVENT.sessionConfigUpdated.name, + "session.mode.updated": MOSOO_CUSTOM_EVENT.sessionModeUpdated.name, + "usage.updated": MOSOO_CUSTOM_EVENT.sessionUsageUpdated.name, + } as const; + const name = names[row.event_type as keyof typeof names]; + if (name === undefined) { + throw new Error(`Unsupported Session state event ${row.event_type}.`); + } + return parseAgUiSessionEvent({ + name, + type: EventType.CUSTOM, + value: + row.event_type === "session.commands.updated" + ? { commands: content } + : row.event_type === "usage.updated" + ? { usage: parseNullableSessionUsageSummary(content) } + : content, + }); +} + +async function readTerminalReferences( + database: D1Database, + sessionId: SessionId, + runIdsInput: readonly SessionRunId[], +): Promise> { + const runIds = [...new Set(runIdsInput)]; + const terminalRows = await getAppDatabase(database) + .select({ + eventType: sessionEventsTable.eventType, + runId: sessionEventsTable.runId, + semanticHash: sessionEventsTable.semanticHash, + seq: sessionEventsTable.seq, + sourceEventId: sessionEventsTable.sourceEventId, + streamId: sessionEventsTable.streamId, + terminalEventJson: sessionEventsTable.terminalEventJson, + }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, sessionId), + sql`${sessionEventsTable.runId} IN (SELECT value FROM json_each(${JSON.stringify(runIds)}))`, + inArray(sessionEventsTable.eventType, RUN_TERMINAL_EVENT_TYPES), + ), + ) + .all(); + + const referencesByRun = new Map(); + const seenRuns = new Set(); + for (const row of terminalRows) { + if (row.runId === null) { + continue; + } + if (seenRuns.has(row.runId)) { + throw new Error( + `Stored assistant reference has multiple terminal events for run ${row.runId}.`, + ); + } + seenRuns.add(row.runId); + if ( + (row.eventType !== "run.cancelled" && + row.eventType !== "run.completed" && + row.eventType !== "run.failed") || + row.sourceEventId !== createSessionRunTerminalSourceId(row.runId, row.eventType) + ) { + throw new Error(`Stored assistant reference has a non-canonical terminal source.`); + } + if (row.semanticHash === null) { + continue; + } + const semanticAuthority = await readTerminalEventSemanticAuthority({ + eventJson: row.terminalEventJson, + eventType: row.eventType, + runId: row.runId, + semanticHash: row.semanticHash, + sessionId, + sourceEventId: row.sourceEventId, + streamId: row.streamId, + }); + referencesByRun.set(row.runId, { + messageId: semanticAuthority.finalMessageId, + seq: row.seq, + }); + } + return referencesByRun; +} + +async function scanReferenceEvents( + database: D1Database, + input: { + ceilings: ReadonlyMap; + includePlan?: boolean; + messageIds?: readonly string[]; + references: readonly LightweightReference[]; + runIds?: readonly SessionRunId[]; + sessionId: SessionId; + toolRoutes: ReadonlyMap; + }, + apply: (row: ReferenceEventRow) => void, +): Promise { + const runIds = [...new Set(input.runIds ?? input.references.map((reference) => reference.runId))]; + const messageIds = [ + ...new Set(input.messageIds ?? input.references.map((reference) => reference.messageId)), + ]; + const toolCallIds = [...new Set([...input.toolRoutes.values()].map((route) => route.toolCallId))]; + const endSeq = Math.max(...runIds.map((runId) => input.ceilings.get(runId) ?? -1)); + let afterSeq: number | null = null; + + for (;;) { + const page = await getAppDatabase(database) + .select({ + content_text: sessionEventsTable.contentText, + ended_at: sessionEventsTable.endedAt, + event_type: sessionEventsTable.eventType, + id: sessionEventsTable.id, + occurred_at: sessionEventsTable.occurredAt, + process_type: sessionEventsTable.processType, + run_id: sessionEventsTable.runId, + seq: sessionEventsTable.seq, + stream_id: sessionEventsTable.streamId, + tokens: sessionEventsTable.tokens, + tool_call_id: sessionEventsTable.toolCallId, + tool_input_delta_json: sessionEventsTable.toolInputDeltaJson, + tool_input_json: sessionEventsTable.toolInputJson, + tool_name: sessionEventsTable.toolName, + tool_output_delta_text: sessionEventsTable.toolOutputDeltaText, + tool_output_text: sessionEventsTable.toolOutputText, + tool_parent_message_id: sessionEventsTable.toolParentMessageId, + tool_result_message_id: sessionEventsTable.toolResultMessageId, + tool_status: sessionEventsTable.toolStatus, + visibility: sessionEventsTable.visibility, + }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, input.sessionId), + sql`${sessionEventsTable.runId} IN (SELECT value FROM json_each(${JSON.stringify(runIds)}))`, + lte(sessionEventsTable.seq, endSeq), + or( + and( + inArray(sessionEventsTable.processType, [ + "agent.message.delta", + THOUGHT_PROCESS_TYPE, + USER_MESSAGE_PROCESS_TYPE, + ]), + sql`${sessionEventsTable.streamId} IN (SELECT value FROM json_each(${JSON.stringify(messageIds)}))`, + ), + and( + eq(sessionEventsTable.eventType, "tool.call.updated"), + eq(sessionEventsTable.visibility, "all_consumers"), + toolCallIds.length === 0 + ? sql`0` + : sql`${sessionEventsTable.toolCallId} IN (SELECT value FROM json_each(${JSON.stringify(toolCallIds)}))`, + ), + input.includePlan === false ? sql`0` : eq(sessionEventsTable.eventType, "plan.updated"), + ), + afterSeq === null ? undefined : gt(sessionEventsTable.seq, afterSeq), + ), + ) + .orderBy(asc(sessionEventsTable.seq)) + .limit(REFERENCE_EVENT_PAGE_SIZE) + .all(); + + for (const row of page) { + apply(row); + } + if (page.length < REFERENCE_EVENT_PAGE_SIZE) { + return; + } + afterSeq = page[page.length - 1]?.seq ?? afterSeq; + } +} + +async function discoverToolRoutes( + database: D1Database, + input: { + ceilings: ReadonlyMap; + references: readonly CanonicalReference[]; + sessionId: SessionId; + }, +): Promise> { + const runIds = [...new Set(input.references.map((reference) => reference.runId))]; + const scopes = runIds.map((runId) => ({ + endSeq: input.ceilings.get(runId) ?? -1, + runId, + })); + const targetKeys = new Set( + input.references.map((reference) => referenceKey(reference.runId, reference.messageId)), + ); + const finalKeys = new Set( + input.references + .filter((reference) => reference.isFinal) + .map((reference) => referenceKey(reference.runId, reference.messageId)), + ); + const endSeq = Math.max(...runIds.map((runId) => input.ceilings.get(runId) ?? -1)); + const routes = new Map(); + let afterSeq: number | null = null; + + for (;;) { + const page = await getAppDatabase(database) + .select({ + run_id: sessionEventsTable.runId, + seq: sessionEventsTable.seq, + tool_call_id: sessionEventsTable.toolCallId, + tool_output_delta_text: sessionEventsTable.toolOutputDeltaText, + tool_output_text: sessionEventsTable.toolOutputText, + tool_parent_message_id: sessionEventsTable.toolParentMessageId, + tool_result_message_id: sessionEventsTable.toolResultMessageId, + }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, input.sessionId), + sql`${sessionEventsTable.runId} IN (SELECT value FROM json_each(${JSON.stringify(runIds)}))`, + eq(sessionEventsTable.eventType, "tool.call.updated"), + eq(sessionEventsTable.visibility, "all_consumers"), + lte(sessionEventsTable.seq, endSeq), + sql`EXISTS ( + SELECT 1 + FROM json_each(${JSON.stringify(scopes)}) AS scope + WHERE json_extract(scope.value, '$.runId') = ${sessionEventsTable.runId} + AND ${sessionEventsTable.seq} <= json_extract(scope.value, '$.endSeq') + )`, + afterSeq === null ? undefined : gt(sessionEventsTable.seq, afterSeq), + ), + ) + .orderBy(asc(sessionEventsTable.seq)) + .limit(REFERENCE_EVENT_PAGE_SIZE) + .all(); + + for (const row of page) { + if ( + row.run_id === null || + row.tool_call_id === null || + row.seq > (input.ceilings.get(row.run_id) ?? -1) + ) { + continue; + } + const key = toolKey(row.run_id, row.tool_call_id); + const route = routes.get(key) ?? { + authority: null, + firstOutputSeq: null, + firstParentSeq: null, + parentMessageId: null, + referenceKeys: new Set(), + resultMessageId: null, + runId: row.run_id, + toolCallId: row.tool_call_id, + }; + mergeToolRouteIdentity(route, row); + if (row.tool_parent_message_id !== null) { + route.firstParentSeq ??= row.seq; + } + if (row.tool_output_delta_text !== null || row.tool_output_text !== null) { + route.firstOutputSeq ??= row.seq; + } + routes.set(key, route); + } + if (page.length < REFERENCE_EVENT_PAGE_SIZE) { + break; + } + afterSeq = page[page.length - 1]?.seq ?? afterSeq; + } + + for (const [key, route] of routes) { + const parentOwnsRoute = + route.firstParentSeq !== null && + (route.firstOutputSeq === null || route.firstParentSeq <= route.firstOutputSeq); + const parentKey = + route.parentMessageId === null ? null : referenceKey(route.runId, route.parentMessageId); + const parentIsFinal = parentOwnsRoute && parentKey !== null && finalKeys.has(parentKey); + const authorityMessageId = parentIsFinal ? route.parentMessageId : route.runId; + const hasEffectiveSegment = parentOwnsRoute || route.firstOutputSeq !== null; + if (hasEffectiveSegment && authorityMessageId !== null) { + const matchedKey = referenceKey(route.runId, authorityMessageId); + route.authority = parentIsFinal ? "final" : "carrier"; + if (!targetKeys.has(matchedKey)) { + throw new Error(`Stored tool call ${route.toolCallId} has no canonical carrier.`); + } + route.referenceKeys.add(matchedKey); + } + if (route.referenceKeys.size === 0) { + routes.delete(key); + } + } + + return routes; +} + +async function discoverActiveRunArtifactIdentities( + database: D1Database, + input: { + endSeq: number; + runId: SessionRunId; + sessionId: SessionId; + }, +): Promise<{ messageIds: string[]; toolRoutes: Map }> { + const messageIds = new Set(); + const toolRoutes = new Map(); + let afterSeq: number | null = null; + + for (;;) { + const page = await getAppDatabase(database) + .select({ + event_type: sessionEventsTable.eventType, + process_type: sessionEventsTable.processType, + seq: sessionEventsTable.seq, + stream_id: sessionEventsTable.streamId, + tool_call_id: sessionEventsTable.toolCallId, + tool_parent_message_id: sessionEventsTable.toolParentMessageId, + tool_result_message_id: sessionEventsTable.toolResultMessageId, + }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.runId, input.runId), + lte(sessionEventsTable.seq, input.endSeq), + or( + and( + inArray(sessionEventsTable.processType, [ + "agent.message.delta", + USER_MESSAGE_PROCESS_TYPE, + ]), + inArray(sessionEventsTable.eventType, [...MESSAGE_STREAM_EVENT_TYPES]), + ), + and( + eq(sessionEventsTable.processType, THOUGHT_PROCESS_TYPE), + inArray(sessionEventsTable.eventType, [...THOUGHT_EVENT_TYPES]), + ), + eq(sessionEventsTable.eventType, "tool.call.updated"), + ), + afterSeq === null ? undefined : gt(sessionEventsTable.seq, afterSeq), + ), + ) + .orderBy(asc(sessionEventsTable.seq)) + .limit(REFERENCE_EVENT_PAGE_SIZE) + .all(); + + for (const row of page) { + if ( + row.process_type === "agent.message.delta" || + row.process_type === THOUGHT_PROCESS_TYPE || + row.process_type === USER_MESSAGE_PROCESS_TYPE + ) { + if (row.stream_id === null) { + throw new Error(`Stored ${row.event_type} event has no message stream identity.`); + } + messageIds.add(row.stream_id); + } + if (row.event_type !== "tool.call.updated" || row.tool_call_id === null) { + continue; + } + const key = toolKey(input.runId, row.tool_call_id); + const route = toolRoutes.get(key) ?? { + authority: null, + firstOutputSeq: null, + firstParentSeq: null, + parentMessageId: null, + referenceKeys: new Set(), + resultMessageId: null, + runId: input.runId, + toolCallId: row.tool_call_id, + }; + mergeToolRouteIdentity(route, row); + toolRoutes.set(key, route); + if (row.tool_parent_message_id !== null) { + messageIds.add(row.tool_parent_message_id); + } + if (row.tool_result_message_id !== null) { + messageIds.add(row.tool_result_message_id); + } + } + if (page.length < REFERENCE_EVENT_PAGE_SIZE) { + return { + messageIds: [...messageIds], + toolRoutes: new Map( + [...toolRoutes].filter( + ([, route]) => route.parentMessageId !== null || route.resultMessageId !== null, + ), + ), + }; + } + afterSeq = page[page.length - 1]?.seq ?? afterSeq; + } +} + +async function applyStoredSessionStateArtifacts( + database: D1Database, + input: { + endSeq: number; + sessionId: SessionId; + state: SessionLiveState; + }, +): Promise { + let state = input.state; + let afterSeq: number | null = null; + + for (;;) { + const page = await getAppDatabase(database) + .select({ + content_text: sessionEventsTable.contentText, + event_type: sessionEventsTable.eventType, + seq: sessionEventsTable.seq, + visibility: sessionEventsTable.visibility, + }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, input.sessionId), + inArray(sessionEventsTable.eventType, [ + ...ACTIVE_SESSION_STATE_EVENT_TYPES, + "plan.updated", + ]), + lte(sessionEventsTable.seq, input.endSeq), + afterSeq === null ? undefined : gt(sessionEventsTable.seq, afterSeq), + ), + ) + .orderBy(asc(sessionEventsTable.seq)) + .limit(REFERENCE_EVENT_PAGE_SIZE) + .all(); + for (const row of page) { + state = applyAgUiEventToSessionLiveState(state, sessionStateAgUiEvent(row)); + } + if (page.length < REFERENCE_EVENT_PAGE_SIZE) { + return state; + } + afterSeq = page[page.length - 1]?.seq ?? afterSeq; + } +} + +export async function applyStoredSessionArtifacts( + database: D1Database, + input: { + endSeq: number; + includeActiveRunArtifacts: boolean; + runId: SessionRunId | null; + sessionId: SessionId; + state: SessionLiveState; + }, +): Promise { + const metadataState = await applyStoredSessionStateArtifacts(database, input); + if (input.runId === null || !input.includeActiveRunArtifacts) { + return metadataState; + } + + const { messageIds, toolRoutes } = await discoverActiveRunArtifactIdentities(database, { + endSeq: input.endSeq, + runId: input.runId, + sessionId: input.sessionId, + }); + const state: ReferenceState = { + isCarrier: false, + isFinal: true, + live: metadataState, + message: createMessageStreamReducerState(), + planJson: null, + }; + const activeKey = referenceKey(input.runId, "$active"); + const states = new Map([[activeKey, state]]); + for (const messageId of messageIds) { + states.set(referenceKey(input.runId, messageId), state); + } + for (const route of toolRoutes.values()) { + route.referenceKeys.add(activeKey); + } + const ceilings = new Map([[input.runId, input.endSeq]]); + const statesByRun = new Map([[input.runId, [state]]]); + + await scanReferenceEvents( + database, + { + ceilings, + includePlan: false, + messageIds, + references: [], + runIds: [input.runId], + sessionId: input.sessionId, + toolRoutes, + }, + (row) => { + applyReferenceEvent(states, statesByRun, ceilings, toolRoutes, row, { + activeArtifacts: true, + }); + }, + ); + + return state.live; +} + +async function readSealedReferenceCursor( + database: D1Database, + input: { + endSeq: number; + reference: LightweightReference; + sessionId: SessionId; + }, +): Promise { + const databaseClient = getAppDatabase(database); + const scope = [ + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.runId, input.reference.runId), + eq(sessionEventsTable.streamId, input.reference.messageId), + eq(sessionEventsTable.processType, "agent.message.delta"), + lte(sessionEventsTable.seq, input.endSeq), + ]; + const [invalid, latest, boundary] = await Promise.all([ + databaseClient + .select({ id: sessionEventsTable.id }) + .from(sessionEventsTable) + .where( + and( + ...scope, + or( + ne(sessionEventsTable.visibility, "all_consumers"), + notInArray(sessionEventsTable.eventType, [...MESSAGE_STREAM_EVENT_TYPES]), + ), + ), + ) + .limit(1) + .get(), + databaseClient + .select({ eventType: sessionEventsTable.eventType, seq: sessionEventsTable.seq }) + .from(sessionEventsTable) + .where(and(...scope, inArray(sessionEventsTable.eventType, [...MESSAGE_STREAM_EVENT_TYPES]))) + .orderBy(desc(sessionEventsTable.seq)) + .limit(1) + .get(), + databaseClient + .select({ eventType: sessionEventsTable.eventType, seq: sessionEventsTable.seq }) + .from(sessionEventsTable) + .where( + and( + ...scope, + inArray(sessionEventsTable.eventType, [...MESSAGE_STREAM_AUTHORITY_BOUNDARY_EVENT_TYPES]), + ), + ) + .orderBy(desc(sessionEventsTable.seq)) + .limit(1) + .get(), + ]); + if (invalid !== undefined) { + throw new Error(`Stored assistant reference ${input.reference.messageId} is not public.`); + } + return latest?.eventType === "message.completed" && boundary?.eventType === "message.added" + ? { endSeq: latest.seq, startSeq: boundary.seq } + : null; +} + +async function readBoundedReferenceContent( + database: D1Database, + input: { + cursor: SealedReferenceCursor; + maxTextLength: number; + reference: LightweightReference; + sessionId: SessionId; + }, +): Promise { + const state = createContentReferenceState(); + if (input.maxTextLength === 0) { + return ""; + } + let afterSeq: number | null = null; + + for (;;) { + const filters: (SQL | undefined)[] = [ + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.runId, input.reference.runId), + eq(sessionEventsTable.processType, "agent.message.delta"), + eq(sessionEventsTable.streamId, input.reference.messageId), + gte(sessionEventsTable.seq, input.cursor.startSeq), + lte(sessionEventsTable.seq, input.cursor.endSeq), + afterSeq === null ? undefined : gt(sessionEventsTable.seq, afterSeq), + ]; + const metadata: { contentLength: number; seq: number }[] = await getAppDatabase(database) + .select({ + contentLength: sql`length(${sessionEventsTable.contentText})`, + seq: sessionEventsTable.seq, + }) + .from(sessionEventsTable) + .where(and(...filters)) + .orderBy(asc(sessionEventsTable.seq)) + .limit(BOUNDED_REFERENCE_EVENT_PAGE_SIZE) + .all(); + if (metadata.length === 0) { + throw new Error(`Stored assistant reference ${input.reference.messageId} has no content.`); + } + let pageChars = 0; + let pageRowCount = 0; + for (const row of metadata) { + if (pageRowCount > 0 && pageChars + row.contentLength > BOUNDED_REFERENCE_EVENT_PAGE_CHARS) { + break; + } + pageChars += row.contentLength; + pageRowCount += 1; + } + const pageEndSeq: number | undefined = metadata[pageRowCount - 1]?.seq; + if (pageEndSeq === undefined) { + throw new Error(`Stored assistant reference ${input.reference.messageId} has no content.`); + } + const page = await getAppDatabase(database) + .select({ + content_text: sessionEventsTable.contentText, + ended_at: sessionEventsTable.endedAt, + event_type: sessionEventsTable.eventType, + id: sessionEventsTable.id, + occurred_at: sessionEventsTable.occurredAt, + process_type: sessionEventsTable.processType, + run_id: sessionEventsTable.runId, + seq: sessionEventsTable.seq, + stream_id: sessionEventsTable.streamId, + tokens: sessionEventsTable.tokens, + visibility: sessionEventsTable.visibility, + }) + .from(sessionEventsTable) + .where(and(...filters, lte(sessionEventsTable.seq, pageEndSeq))) + .orderBy(asc(sessionEventsTable.seq)) + .all(); + + for (const row of page) { + reduceSanitizedMessageStreamRow(state, row, input.maxTextLength); + if (state.message.text.length >= input.maxTextLength) { + return state.message.text; + } + } + if ( + pageEndSeq === input.cursor.endSeq || + (pageRowCount === metadata.length && metadata.length < BOUNDED_REFERENCE_EVENT_PAGE_SIZE) + ) { + if (!state.message.authoritative || !state.message.sealed) { + throw new Error( + `Stored assistant reference ${input.reference.messageId} is not sealed and authoritative.`, + ); + } + return state.message.text; + } + afterSeq = pageEndSeq; + } +} + +async function assertReferenceMessageIdentities( + database: D1Database, + input: { + ceilings: ReadonlyMap; + references: readonly LightweightReference[]; + sessionId: SessionId; + }, +): Promise { + const scopes = input.references.map((reference) => ({ + endSeq: input.ceilings.get(reference.runId) ?? -1, + messageId: reference.messageId, + runId: reference.runId, + })); + const collision = await getAppDatabase(database) + .select({ id: sessionEventsTable.id }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.sessionId, input.sessionId), + inArray(sessionEventsTable.eventType, [...MESSAGE_STREAM_EVENT_TYPES]), + sql`EXISTS ( + SELECT 1 + FROM json_each(${JSON.stringify(scopes)}) AS scope + WHERE json_extract(scope.value, '$.messageId') = ${sessionEventsTable.streamId} + AND ${sessionEventsTable.seq} <= json_extract(scope.value, '$.endSeq') + AND ( + ${sessionEventsTable.runId} IS NOT json_extract(scope.value, '$.runId') + OR ${sessionEventsTable.processType} <> 'agent.message.delta' + ) + )`, + ), + ) + .limit(1) + .get(); + + if (collision !== undefined) { + throw new Error("Stored assistant reference has conflicting message-stream identity rows."); + } +} + +export async function resolveStoredSessionMessageContentReferences< + Row extends StoredSessionMessageContentReferenceRow, +>( + database: D1Database, + sessionId: SessionId, + rows: readonly Row[], + maxTextLength: number, +): Promise { + for (const row of rows) { + if ( + row.projection_format === "event_stream_v3" && + (row.role !== "assistant" || row.session_run_id === null || row.content_text !== "") + ) { + throw new Error(`Stored event-stream assistant ${row.id} is not lightweight.`); + } + } + const candidates = rows.flatMap((row): LightweightReference[] => { + const reference = potentialReference(row); + return reference === null ? [] : [reference]; + }); + if (candidates.length === 0) { + return [...rows]; + } + const terminalReferences = await readTerminalReferences( + database, + sessionId, + candidates.map((candidate) => candidate.runId), + ); + const references = requireTerminalReferences(candidates, terminalReferences); + const ceilings = new Map( + references.map((reference) => [reference.runId, terminalReferences.get(reference.runId)!.seq]), + ); + const finalReferences = references.filter((reference) => reference.isFinal); + await assertReferenceMessageIdentities(database, { + ceilings, + references: finalReferences, + sessionId, + }); + const resolvedText = new Map(); + for (const reference of finalReferences) { + const endSeq = ceilings.get(reference.runId); + if (endSeq === undefined) { + throw new Error( + `Stored event-stream assistant ${reference.messageId} has no terminal ceiling.`, + ); + } + const cursor = await readSealedReferenceCursor(database, { + endSeq, + reference, + sessionId, + }); + if (cursor === null) { + throw new Error( + `Stored assistant reference ${reference.messageId} is not sealed and authoritative.`, + ); + } + resolvedText.set( + referenceKey(reference.runId, reference.messageId), + await readBoundedReferenceContent(database, { + cursor, + maxTextLength, + reference, + sessionId, + }), + ); + } + + return rows.map((row) => { + if (row.session_run_id === null) { + return row; + } + const contentText = resolvedText.get(referenceKey(row.session_run_id, row.id)); + if (contentText === undefined) { + return row; + } + return { ...row, content_text: contentText }; + }); +} + +async function resolveStoredSessionMessageSnapshot( + database: D1Database, + sessionId: SessionId, + rows: readonly Row[], +): Promise { + for (const row of rows) { + if ( + row.projection_format === "event_stream_v3" && + (row.role !== "assistant" || + row.session_run_id === null || + row.content_text !== "" || + row.plan_json !== null || + row.segments_json !== null) + ) { + throw new Error(`Stored event-stream assistant ${row.id} is not lightweight.`); + } + } + const candidates = rows.flatMap((row): LightweightReference[] => { + const reference = potentialReference(row); + return reference === null ? [] : [reference]; + }); + if (candidates.length === 0) { + return [...rows]; + } + const terminalReferences = await readTerminalReferences( + database, + sessionId, + candidates.map((candidate) => candidate.runId), + ); + const references = requireTerminalReferences(candidates, terminalReferences); + const ceilings = new Map( + references.map((reference) => [reference.runId, terminalReferences.get(reference.runId)!.seq]), + ); + await assertReferenceMessageIdentities(database, { + ceilings, + references: references.filter((reference) => reference.isFinal), + sessionId, + }); + const states = new Map(); + const statesByRun = new Map(); + for (const reference of references) { + const state = createReferenceState(sessionId, reference); + states.set(referenceKey(reference.runId, reference.messageId), state); + if (reference.isFinal) { + const runStates = statesByRun.get(reference.runId) ?? []; + runStates.push(state); + statesByRun.set(reference.runId, runStates); + } + } + + const toolRoutes = await discoverToolRoutes(database, { ceilings, references, sessionId }); + await scanReferenceEvents(database, { ceilings, references, sessionId, toolRoutes }, (row) => + applyReferenceEvent(states, statesByRun, ceilings, toolRoutes, row), + ); + + return rows.map((row) => { + if (row.session_run_id === null) { + return row; + } + const state = states.get(referenceKey(row.session_run_id, row.id)); + if (state === undefined) { + return row; + } + if (state.isFinal && (!state.message.authoritative || !state.message.sealed)) { + throw new Error(`Stored assistant reference ${row.id} is not sealed and authoritative.`); + } + const message = state.live.messages.find((candidate) => candidate.id === row.id); + if ( + message === undefined || + (state.isFinal && message.content !== state.message.text) || + (!state.isFinal && message.content !== "") || + (!state.isFinal && + !message.segments.some( + (segment) => segment.kind === "tool_result" || segment.kind === "tool_use", + )) + ) { + throw new Error(`Stored assistant reference ${row.id} diverges from the live reducer.`); + } + if (state.live.messages.some((candidate) => candidate.id !== row.id)) { + throw new Error(`Stored assistant reference ${row.id} produced a detached message.`); + } + + return { + ...row, + content_text: state.isFinal ? state.message.text : "", + plan_json: state.isFinal ? state.planJson : null, + segments_json: JSON.stringify(message.segments), + }; + }); +} + +export async function resolveStoredSessionMessageReferences< + Row extends StoredSessionMessageReferenceRow, +>(database: D1Database, sessionId: SessionId, rows: readonly Row[]): Promise { + return resolveStoredSessionMessageSnapshot(database, sessionId, rows); +} + +export async function resolveStoredSessionMessageSnapshotReferences< + Row extends StoredSessionMessageReferenceRow, +>(database: D1Database, sessionId: SessionId, rows: readonly Row[]): Promise { + return resolveStoredSessionMessageSnapshot(database, sessionId, rows); +} diff --git a/apps/api/src/modules/sessions/infrastructure/session-message-snapshot.repository.ts b/apps/api/src/modules/sessions/infrastructure/session-message-snapshot.repository.ts index 564b9a3c..cd1e94ee 100644 --- a/apps/api/src/modules/sessions/infrastructure/session-message-snapshot.repository.ts +++ b/apps/api/src/modules/sessions/infrastructure/session-message-snapshot.repository.ts @@ -1,5 +1,5 @@ import { sessionMessagesTable } from "@mosoo/db"; -import type { SessionId, SessionMessageId } from "@mosoo/id"; +import type { SessionId, SessionMessageId, SessionRunId } from "@mosoo/id"; import { asc, eq } from "drizzle-orm"; import { getAppDatabase } from "../../../platform/db/drizzle"; @@ -10,6 +10,7 @@ import { } from "../domain/provider-private-markup"; import { parseStoredSessionMessageProjection } from "../domain/session-message-projection-parser"; import type { SessionLiveStateMessage } from "./session-live-state.types"; +import { resolveStoredSessionMessageSnapshotReferences } from "./session-message-reference.repository"; export interface StoredSessionMessageRow { content_text: string; @@ -19,6 +20,7 @@ export interface StoredSessionMessageRow { role: "assistant" | "user"; segments_json: string | null; seq: number; + session_run_id: SessionRunId | null; } function compareStoredSessionMessageRows( @@ -50,7 +52,7 @@ function toLiveStateMessage(row: StoredSessionMessageRow): SessionLiveStateMessa function storedSessionMessageRowsToLiveMessages( rows: StoredSessionMessageRow[], ): SessionLiveStateMessage[] { - return [...rows].toSorted(compareStoredSessionMessageRows).map((row) => toLiveStateMessage(row)); + return [...rows].toSorted(compareStoredSessionMessageRows).map(toLiveStateMessage); } export async function loadStoredSessionMessages( @@ -63,14 +65,22 @@ export async function loadStoredSessionMessages( created_at: sessionMessagesTable.createdAt, id: sessionMessagesTable.id, plan_json: sessionMessagesTable.planJson, + projection_format: sessionMessagesTable.projectionFormat, role: sessionMessagesTable.role, segments_json: sessionMessagesTable.segmentsJson, seq: sessionMessagesTable.seq, + session_run_id: sessionMessagesTable.sessionRunId, }) .from(sessionMessagesTable) .where(eq(sessionMessagesTable.sessionId, sessionId)) .orderBy(asc(sessionMessagesTable.seq)) .all(); - return storedSessionMessageRowsToLiveMessages(results); + const resolved = await resolveStoredSessionMessageSnapshotReferences( + database, + sessionId, + results, + ); + + return storedSessionMessageRowsToLiveMessages(resolved); } diff --git a/apps/api/src/modules/sessions/infrastructure/session-model-call.repository.ts b/apps/api/src/modules/sessions/infrastructure/session-model-call.repository.ts index 8de48d19..6c5cea09 100644 --- a/apps/api/src/modules/sessions/infrastructure/session-model-call.repository.ts +++ b/apps/api/src/modules/sessions/infrastructure/session-model-call.repository.ts @@ -1,6 +1,8 @@ +import type { SessionRunStatus } from "@mosoo/contracts/session-run"; import { agentsTable, appsTable, + sessionEventsTable, sessionModelCallsTable, sessionRunsTable, sessionsTable, @@ -13,16 +15,23 @@ import type { DriverInstanceId, OrganizationId, AppId, + RuntimeEventId, SessionId, SessionModelCallId, SessionRunId, } from "@mosoo/id"; -import { and, eq, sql } from "drizzle-orm"; +import { and, eq, exists, sql } from "drizzle-orm"; +import type { SQL } from "drizzle-orm"; import { getAppDatabase, runAppDatabaseBatch } from "../../../platform/db/drizzle"; import { isTruthy } from "../../../shared/truthiness"; import { currentTimestampMs } from "../../../time"; -import { createRuntimeUsageEventUpsert } from "../../cost/application/cost-usage-event.service"; +import { + createRuntimeUsageEventConvergencePredicate, + createRuntimeUsageEventUnrolledPredicate, + createRuntimeUsageEventUpsert, + hasRuntimeUsageEventRollupReceipt, +} from "../../cost/application/cost-usage-event.service"; import type { SessionUsageSummary } from "./session-live-state.types"; interface SessionModelCallRunRow { agent_id: AgentId; @@ -31,6 +40,7 @@ interface SessionModelCallRunRow { agent_status: "draft" | "published"; actor_user_id: AccountId; completed_at: number | null; + created_at: number; model: string | null; app_organization_id: OrganizationId; app_id: AppId; @@ -41,20 +51,52 @@ interface SessionModelCallRunRow { session_provider: string; session_runtime_id: string; started_at: number | null; + status: SessionRunStatus; trigger: "resume" | "retry" | "system" | "user_prompt"; } export type SessionModelCallStatus = "completed" | "failed" | "started"; export interface UpsertSessionModelCallUsageInput { + createdAtMs: number; driverInstanceId: DriverInstanceId; sessionId: SessionId; sessionRunId: SessionRunId; - status: SessionModelCallStatus; + sourceEventSeq: number; traceId: string; usage: SessionUsageSummary | null; } +export interface DurableSessionModelCallUsageProjectionInput extends UpsertSessionModelCallUsageInput { + eventId: RuntimeEventId; + semanticHash: string; +} + +type SessionModelCallInsert = typeof sessionModelCallsTable.$inferInsert; + +interface StoredSessionModelCall { + cacheCreationTokens: number | null; + cacheReadTokens: number | null; + costCurrency: string | null; + driverInstanceId: string | null; + inputTokens: number | null; + metadataJson: string | null; + model: string; + nativeCallId: string | null; + outputTokens: number | null; + provider: string; + sessionId: string; + sessionRunId: string; + sourceEventSeq: number; + startedAt: number | null; + totalCostUsdMicros: number | null; + traceId: string; +} + +function selectedValue(value: Value, alias: string) { + return sql`${value}`.as(alias); +} + function toTokenCount(value: number | null | undefined): number | null { if (typeof value !== "number" || !Number.isFinite(value) || value < 0) { return null; @@ -73,9 +115,14 @@ function toUsdMicros(value: number | null | undefined): number | null { function buildUsageMetadata(usage: SessionUsageSummary): string { return JSON.stringify({ + cachedReadTokens: usage.cachedReadTokens ?? null, cachedWriteTokens: usage.cachedWriteTokens ?? null, callId: usage.callId ?? null, + costAmount: usage.costAmount ?? null, + costCurrency: usage.costCurrency ?? null, + inputTokens: usage.inputTokens ?? null, model: usage.model ?? null, + outputTokens: usage.outputTokens ?? null, provider: usage.provider ?? null, size: usage.size ?? null, source: usage.source, @@ -99,6 +146,7 @@ async function getSessionModelCallRunRow( agent_revision_id: sessionRunsTable.deploymentVersionId, agent_status: sql<"draft" | "published">`${agentsTable.status}`, completed_at: sessionRunsTable.completedAt, + created_at: sessionRunsTable.createdAt, model: sql`${sessionRunsTable.model}`.mapWith(sessionRunsTable.model).as("model"), app_organization_id: appsTable.organizationId, app_id: sessionsTable.appId, @@ -117,6 +165,7 @@ async function getSessionModelCallRunRow( .mapWith(sessionsTable.runtimeId) .as("session_runtime_id"), started_at: sessionRunsTable.startedAt, + status: sessionRunsTable.status, trigger: sessionRunsTable.trigger, }) .from(sessionRunsTable) @@ -135,31 +184,217 @@ async function getSessionModelCallRunRow( ); } -export async function upsertSessionModelCallUsage( +function toSessionModelCallStatus(status: SessionRunStatus): SessionModelCallStatus { + if (status === "completed") { + return "completed"; + } + + if (status === "failed" || status === "cancelled" || status === "expired") { + return "failed"; + } + + return "started"; +} + +function createSessionModelCallInsertSelect( + database: ReturnType, + values: SessionModelCallInsert, + writeFence: SQL, +) { + return database + .select({ + cacheCreationTokens: selectedValue(values.cacheCreationTokens, "cache_creation_tokens"), + cacheReadTokens: selectedValue(values.cacheReadTokens, "cache_read_tokens"), + callKey: selectedValue(values.callKey, "call_key"), + completedAt: selectedValue(values.completedAt, "completed_at"), + costCurrency: selectedValue(values.costCurrency, "cost_currency"), + createdAt: selectedValue(values.createdAt, "created_at"), + driverInstanceId: selectedValue(values.driverInstanceId, "driver_instance_id"), + errorCode: selectedValue(values.errorCode, "error_code"), + errorMessage: selectedValue(values.errorMessage, "error_message"), + id: selectedValue(values.id, "id"), + inputTokens: selectedValue(values.inputTokens, "input_tokens"), + metadataJson: selectedValue(values.metadataJson, "metadata_json"), + model: selectedValue(values.model, "model"), + nativeCallId: selectedValue(values.nativeCallId, "native_call_id"), + outputTokens: selectedValue(values.outputTokens, "output_tokens"), + provider: selectedValue(values.provider, "provider"), + sourceEventSeq: selectedValue(values.sourceEventSeq, "source_event_seq"), + sessionId: selectedValue(values.sessionId, "session_id"), + sessionRunId: selectedValue(values.sessionRunId, "session_run_id"), + startedAt: selectedValue(values.startedAt, "started_at"), + status: selectedValue(values.status, "status"), + totalCostUsdMicros: selectedValue(values.totalCostUsdMicros, "total_cost_usd_micros"), + traceId: selectedValue(values.traceId, "trace_id"), + updatedAt: selectedValue(values.updatedAt, "updated_at"), + }) + .from(sql`(SELECT 1)`) + .where(writeFence); +} + +function createSessionModelCallUsageUpsert( + database: ReturnType, + values: SessionModelCallInsert, + writeFence: SQL, +) { + return database + .insert(sessionModelCallsTable) + .select(createSessionModelCallInsertSelect(database, values, writeFence)) + .onConflictDoUpdate({ + set: { + cacheCreationTokens: sql`COALESCE(excluded.cache_creation_tokens, ${sessionModelCallsTable.cacheCreationTokens})`, + cacheReadTokens: sql`COALESCE(excluded.cache_read_tokens, ${sessionModelCallsTable.cacheReadTokens})`, + completedAt: sql`COALESCE(excluded.completed_at, ${sessionModelCallsTable.completedAt})`, + costCurrency: sql`COALESCE(excluded.cost_currency, ${sessionModelCallsTable.costCurrency})`, + driverInstanceId: sql`excluded.driver_instance_id`, + inputTokens: sql`COALESCE(excluded.input_tokens, ${sessionModelCallsTable.inputTokens})`, + metadataJson: sql`excluded.metadata_json`, + model: sql`excluded.model`, + outputTokens: sql`COALESCE(excluded.output_tokens, ${sessionModelCallsTable.outputTokens})`, + provider: sql`excluded.provider`, + sourceEventSeq: sql`excluded.source_event_seq`, + startedAt: sql`COALESCE(${sessionModelCallsTable.startedAt}, excluded.started_at)`, + status: sql`excluded.status`, + totalCostUsdMicros: sql`COALESCE(excluded.total_cost_usd_micros, ${sessionModelCallsTable.totalCostUsdMicros})`, + traceId: sql`excluded.trace_id`, + updatedAt: sql`excluded.updated_at`, + }, + setWhere: sql`${sessionModelCallsTable.sourceEventSeq} < excluded.source_event_seq AND ${writeFence}`, + target: [sessionModelCallsTable.sessionRunId, sessionModelCallsTable.callKey], + }); +} + +function createSessionModelCallUsageConvergenceGuard( + database: ReturnType, + modelCallValues: SessionModelCallInsert, + usageEventInput: Parameters[1], +) { + return database.insert(sessionModelCallsTable).select( + createSessionModelCallInsertSelect( + database, + modelCallValues, + sql`${createRuntimeUsageEventUnrolledPredicate(database, usageEventInput)} + AND NOT (${createRuntimeUsageEventConvergencePredicate(database, usageEventInput)})`, + ), + ); +} + +async function getStoredSessionModelCall( + database: D1Database, + input: { callKey: string; sessionRunId: SessionRunId }, +): Promise { + return ( + (await getAppDatabase(database) + .select({ + cacheCreationTokens: sessionModelCallsTable.cacheCreationTokens, + cacheReadTokens: sessionModelCallsTable.cacheReadTokens, + costCurrency: sessionModelCallsTable.costCurrency, + driverInstanceId: sessionModelCallsTable.driverInstanceId, + inputTokens: sessionModelCallsTable.inputTokens, + metadataJson: sessionModelCallsTable.metadataJson, + model: sessionModelCallsTable.model, + nativeCallId: sessionModelCallsTable.nativeCallId, + outputTokens: sessionModelCallsTable.outputTokens, + provider: sessionModelCallsTable.provider, + sessionId: sessionModelCallsTable.sessionId, + sessionRunId: sessionModelCallsTable.sessionRunId, + sourceEventSeq: sessionModelCallsTable.sourceEventSeq, + startedAt: sessionModelCallsTable.startedAt, + totalCostUsdMicros: sessionModelCallsTable.totalCostUsdMicros, + traceId: sessionModelCallsTable.traceId, + }) + .from(sessionModelCallsTable) + .where( + and( + eq(sessionModelCallsTable.sessionRunId, input.sessionRunId), + eq(sessionModelCallsTable.callKey, input.callKey), + ), + ) + .limit(1) + .get()) ?? null + ); +} + +function assertSessionModelCallConverged( + stored: StoredSessionModelCall | null, + expected: SessionModelCallInsert, +): void { + const sourceEventSeq = expected.sourceEventSeq ?? 0; + + if (stored === null || stored.sourceEventSeq < sourceEventSeq) { + throw new Error("Session model call CAS did not persist the durable event."); + } + + if (stored.sourceEventSeq > sourceEventSeq) { + return; + } + + const requiredValuesMatch = + stored.driverInstanceId === (expected.driverInstanceId ?? null) && + stored.metadataJson === (expected.metadataJson ?? null) && + stored.model === expected.model && + stored.nativeCallId === (expected.nativeCallId ?? null) && + stored.provider === expected.provider && + stored.sessionId === expected.sessionId && + stored.sessionRunId === expected.sessionRunId && + stored.startedAt === (expected.startedAt ?? null) && + stored.traceId === expected.traceId; + const optionalValuesMatch = + (expected.cacheCreationTokens == null || + stored.cacheCreationTokens === expected.cacheCreationTokens) && + (expected.cacheReadTokens == null || stored.cacheReadTokens === expected.cacheReadTokens) && + (expected.costCurrency == null || stored.costCurrency === expected.costCurrency) && + (expected.inputTokens == null || stored.inputTokens === expected.inputTokens) && + (expected.outputTokens == null || stored.outputTokens === expected.outputTokens) && + (expected.totalCostUsdMicros == null || + stored.totalCostUsdMicros === expected.totalCostUsdMicros); + + if (!requiredValuesMatch || !optionalValuesMatch) { + throw new Error("Session model call event seq was replayed with conflicting content."); + } +} + +async function prepareSessionModelCallUsageValues( database: D1Database, input: UpsertSessionModelCallUsageInput, -): Promise { +): Promise<{ + modelCallValues: SessionModelCallInsert; + usageEventInput: Parameters[1]; +} | null> { if (!input.usage) { - return; + return null; + } + + if (!Number.isSafeInteger(input.sourceEventSeq) || input.sourceEventSeq < 0) { + throw new Error("Session model call source event seq must be a non-negative safe integer."); + } + if (!Number.isSafeInteger(input.createdAtMs) || input.createdAtMs < 0) { + throw new Error("Session model call createdAtMs must be a non-negative safe integer."); } - const usage = input.usage; + const usage = input.usage; const run = await getSessionModelCallRunRow(database, input.sessionRunId); if (!run) { throw new Error("Session run not found for model call usage."); } + if (run.session_id !== input.sessionId) { + throw new Error("Session model call usage does not belong to the durable Session Run."); + } + const timestampMs = currentTimestampMs(); - const completedAt = - input.status === "completed" || input.status === "failed" - ? (run.completed_at ?? timestampMs) - : null; + const status = toSessionModelCallStatus(run.status); + const completedAt = status === "started" ? null : run.completed_at; + + if (status !== "started" && completedAt === null) { + throw new Error("Terminal Session Run is missing completed_at for model call usage."); + } + const nativeCallId = normalizeUsageCallId(usage.callId); const callKey = isTruthy(nativeCallId) ? `model_call:${nativeCallId}` : "run_usage"; const provider = run.provider ?? run.session_provider; const model = run.model ?? run.session_model; - const usageEventInput = { callKey, driverInstanceId: input.driverInstanceId, @@ -170,7 +405,7 @@ export async function upsertSessionModelCallUsage( agentOwnerUserId: run.agent_owner_user_id, agentRevisionId: run.agent_revision_id, agentStatus: run.agent_status, - createdAtMs: completedAt ?? timestampMs, + createdAtMs: input.createdAtMs, model, organizationId: run.app_organization_id, appId: run.app_id, @@ -180,66 +415,235 @@ export async function upsertSessionModelCallUsage( sessionRunId: input.sessionRunId, trigger: run.trigger, }, + sourceEventSeq: input.sourceEventSeq, usage, } satisfies Parameters[1]; + const modelCallValues = { + cacheCreationTokens: toTokenCount(usage.cachedWriteTokens), + cacheReadTokens: toTokenCount(usage.cachedReadTokens), + callKey, + completedAt, + costCurrency: usage.costCurrency ?? null, + createdAt: input.createdAtMs, + driverInstanceId: input.driverInstanceId, + errorCode: null, + errorMessage: null, + id: createPlatformId(), + inputTokens: toTokenCount(usage.inputTokens), + metadataJson: buildUsageMetadata(usage), + model, + nativeCallId, + outputTokens: toTokenCount(usage.outputTokens), + provider, + sessionId: input.sessionId, + sessionRunId: input.sessionRunId, + sourceEventSeq: input.sourceEventSeq, + startedAt: run.started_at ?? run.created_at, + status, + totalCostUsdMicros: toUsdMicros(usage.costAmount), + traceId: input.traceId, + updatedAt: timestampMs, + } satisfies SessionModelCallInsert; + + return { modelCallValues, usageEventInput }; +} + +export async function prepareDurableSessionModelCallUsageProjection( + database: D1Database, + input: DurableSessionModelCallUsageProjectionInput, +): Promise { + const prepared = await prepareSessionModelCallUsageValues(database, input); + + if (prepared === null) { + return []; + } + + const { modelCallValues, usageEventInput } = prepared; + const appDatabase = getAppDatabase(database); + const receiptFence = exists( + appDatabase + .select({ id: sessionEventsTable.id }) + .from(sessionEventsTable) + .where( + and( + eq(sessionEventsTable.id, input.eventId), + eq(sessionEventsTable.sessionId, input.sessionId), + eq(sessionEventsTable.eventType, "usage.updated"), + eq(sessionEventsTable.runId, input.sessionRunId), + eq(sessionEventsTable.semanticHash, input.semanticHash), + eq(sessionEventsTable.seq, input.sourceEventSeq), + ), + ), + ); + const writeFence = sql`${createRuntimeUsageEventUnrolledPredicate( + appDatabase, + usageEventInput, + )} AND ${receiptFence}`; + const modelCallQuery = createSessionModelCallUsageUpsert( + appDatabase, + modelCallValues, + writeFence, + ).toSQL(); + const statements = [database.prepare(modelCallQuery.sql).bind(...modelCallQuery.params)]; + const usageEventUpsert = createRuntimeUsageEventUpsert( + appDatabase, + usageEventInput, + receiptFence, + ); + + if (usageEventUpsert !== null) { + const usageEventQuery = usageEventUpsert.toSQL(); + statements.push(database.prepare(usageEventQuery.sql).bind(...usageEventQuery.params)); + } + + const usageEventConvergenceQuery = + usageEventUpsert === null + ? null + : appDatabase + .select({ converged: sql`1` }) + .from(sql`(SELECT 1)`) + .where(createRuntimeUsageEventConvergencePredicate(appDatabase, usageEventInput)) + .toSQL(); + const usageEventFailureSql = + usageEventConvergenceQuery === null ? "" : ` OR NOT EXISTS (${usageEventConvergenceQuery.sql})`; + + const sourceEventSeq = modelCallValues.sourceEventSeq ?? 0; + statements.push( + database + .prepare( + `INSERT INTO session_event (id) + SELECT ? + WHERE EXISTS ( + SELECT 1 + FROM session_event AS receipt + WHERE receipt.id = ? + AND receipt.session_id = ? + AND receipt.event_type = 'usage.updated' + AND receipt.run_id = ? + AND receipt.semantic_hash = ? + AND receipt.seq = ? + ) + AND ( + NOT EXISTS ( + SELECT 1 + FROM session_model_call AS stored + WHERE stored.session_run_id = ? + AND stored.call_key = ? + AND ( + stored.source_event_seq > ? + OR ( + stored.source_event_seq = ? + AND stored.driver_instance_id IS ? + AND stored.metadata_json IS ? + AND stored.model = ? + AND stored.native_call_id IS ? + AND stored.provider = ? + AND stored.session_id = ? + AND stored.started_at IS ? + AND stored.trace_id = ? + AND (? IS NULL OR stored.cache_creation_tokens = ?) + AND (? IS NULL OR stored.cache_read_tokens = ?) + AND (? IS NULL OR stored.cost_currency = ?) + AND (? IS NULL OR stored.input_tokens = ?) + AND (? IS NULL OR stored.output_tokens = ?) + AND (? IS NULL OR stored.total_cost_usd_micros = ?) + ) + ) + )${usageEventFailureSql} + )`, + ) + .bind( + input.eventId, + input.eventId, + input.sessionId, + input.sessionRunId, + input.semanticHash, + input.sourceEventSeq, + input.sessionRunId, + modelCallValues.callKey, + sourceEventSeq, + sourceEventSeq, + modelCallValues.driverInstanceId ?? null, + modelCallValues.metadataJson ?? null, + modelCallValues.model, + modelCallValues.nativeCallId ?? null, + modelCallValues.provider, + modelCallValues.sessionId, + modelCallValues.startedAt ?? null, + modelCallValues.traceId, + modelCallValues.cacheCreationTokens ?? null, + modelCallValues.cacheCreationTokens ?? null, + modelCallValues.cacheReadTokens ?? null, + modelCallValues.cacheReadTokens ?? null, + modelCallValues.costCurrency ?? null, + modelCallValues.costCurrency ?? null, + modelCallValues.inputTokens ?? null, + modelCallValues.inputTokens ?? null, + modelCallValues.outputTokens ?? null, + modelCallValues.outputTokens ?? null, + modelCallValues.totalCostUsdMicros ?? null, + modelCallValues.totalCostUsdMicros ?? null, + ...(usageEventConvergenceQuery?.params ?? []), + ), + ); + + return statements; +} + +export async function upsertSessionModelCallUsage( + database: D1Database, + input: UpsertSessionModelCallUsageInput, +): Promise { + const prepared = await prepareSessionModelCallUsageValues(database, input); + + if (prepared === null) { + return; + } + const { modelCallValues, usageEventInput } = prepared; + const callKey = modelCallValues.callKey; + + if (await hasRuntimeUsageEventRollupReceipt(database, usageEventInput)) { + const stored = await getStoredSessionModelCall(database, { + callKey, + sessionRunId: input.sessionRunId, + }); + + if (stored !== null && stored.sourceEventSeq >= input.sourceEventSeq) { + assertSessionModelCallConverged(stored, modelCallValues); + return; + } + + throw new Error( + "Session model call usage was already rolled up and cannot be replaced safely.", + ); + } await runAppDatabaseBatch(database, (appDatabase) => { - const modelCallUpsert = appDatabase - .insert(sessionModelCallsTable) - .values({ - cacheCreationTokens: toTokenCount(usage.cachedWriteTokens), - cacheReadTokens: toTokenCount(usage.cachedReadTokens), - callKey, - completedAt, - costCurrency: usage.costCurrency ?? null, - createdAt: timestampMs, - driverInstanceId: input.driverInstanceId, - errorCode: null, - errorMessage: null, - id: createPlatformId(), - inputTokens: toTokenCount(usage.inputTokens), - metadataJson: buildUsageMetadata(usage), - model, - nativeCallId, - outputTokens: toTokenCount(usage.outputTokens), - provider, - sessionId: input.sessionId, - sessionRunId: input.sessionRunId, - startedAt: run.started_at ?? timestampMs, - status: input.status, - totalCostUsdMicros: toUsdMicros(usage.costAmount), - traceId: input.traceId, - updatedAt: timestampMs, - }) - .onConflictDoUpdate({ - set: { - cacheCreationTokens: sql`COALESCE(excluded.cache_creation_tokens, ${sessionModelCallsTable.cacheCreationTokens})`, - cacheReadTokens: sql`COALESCE(excluded.cache_read_tokens, ${sessionModelCallsTable.cacheReadTokens})`, - completedAt: sql`COALESCE(excluded.completed_at, ${sessionModelCallsTable.completedAt})`, - costCurrency: sql`COALESCE(excluded.cost_currency, ${sessionModelCallsTable.costCurrency})`, - driverInstanceId: sql`excluded.driver_instance_id`, - inputTokens: sql`COALESCE(excluded.input_tokens, ${sessionModelCallsTable.inputTokens})`, - metadataJson: sql`excluded.metadata_json`, - model: sql`excluded.model`, - outputTokens: sql`COALESCE(excluded.output_tokens, ${sessionModelCallsTable.outputTokens})`, - provider: sql`excluded.provider`, - startedAt: sql`COALESCE(${sessionModelCallsTable.startedAt}, excluded.started_at)`, - status: sql`CASE - WHEN ${sessionModelCallsTable.status} IN ('completed', 'failed') - AND excluded.status = 'started' - THEN ${sessionModelCallsTable.status} - ELSE excluded.status - END`, - totalCostUsdMicros: sql`COALESCE(excluded.total_cost_usd_micros, ${sessionModelCallsTable.totalCostUsdMicros})`, - traceId: sql`excluded.trace_id`, - updatedAt: sql`excluded.updated_at`, - }, - target: [sessionModelCallsTable.sessionRunId, sessionModelCallsTable.callKey], - }); + const writeFence = createRuntimeUsageEventUnrolledPredicate(appDatabase, usageEventInput); + const modelCallUpsert = createSessionModelCallUsageUpsert( + appDatabase, + modelCallValues, + writeFence, + ); const usageEventUpsert = createRuntimeUsageEventUpsert(appDatabase, usageEventInput); - return usageEventUpsert === null ? [modelCallUpsert] : [modelCallUpsert, usageEventUpsert]; + return usageEventUpsert === null + ? [modelCallUpsert] + : [ + modelCallUpsert, + usageEventUpsert, + createSessionModelCallUsageConvergenceGuard( + appDatabase, + modelCallValues, + usageEventInput, + ), + ]; }); + + assertSessionModelCallConverged( + await getStoredSessionModelCall(database, { callKey, sessionRunId: input.sessionRunId }), + modelCallValues, + ); } function normalizeUsageCallId(value: string | null | undefined): string | null { diff --git a/apps/api/src/modules/sessions/infrastructure/session-runtime-event-store.repository.ts b/apps/api/src/modules/sessions/infrastructure/session-runtime-event-store.repository.ts index 13232410..4fffb02c 100644 --- a/apps/api/src/modules/sessions/infrastructure/session-runtime-event-store.repository.ts +++ b/apps/api/src/modules/sessions/infrastructure/session-runtime-event-store.repository.ts @@ -1,12 +1,34 @@ -import { sessionEventsTable, sessionRunsTable, sessionsTable } from "@mosoo/db"; +import { parseNullableSessionUsageSummary } from "@mosoo/ag-ui-session"; +import { + driverInstancesTable, + sessionEventsTable, + sessionRunsTable, + sessionsTable, +} from "@mosoo/db"; import { createPlatformId } from "@mosoo/id"; import type { RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; -import { and, eq, inArray, isNull, sql } from "drizzle-orm"; +import { + createRuntimeEventSemanticHash, + readRuntimeAgentTaskSnapshot, + readRuntimeEventPayload, + readRuntimeEventString, +} from "@mosoo/runtime-events"; +import { and, eq, exists, inArray, isNull, sql } from "drizzle-orm"; import { getAppDatabase } from "../../../platform/db/drizzle"; +import type { AppDatabase } from "../../../platform/db/drizzle"; import { currentTimestampMs } from "../../../time"; +import { ACTIVE_SESSION_RUN_STATUSES } from "../../runtime/domain/session-run-lifecycle.machine"; +import { readNativeResumeRef } from "../../runtime/infrastructure/driver-instance/native-resume-ref-event"; +import { prepareRuntimeArtifactPromotion } from "../../runtime/infrastructure/driver-instance/runtime-artifact-attempt.repository"; +import { prepareNativeResumeRefProjection } from "../../runtime/infrastructure/native-resume-ref.repository"; +import { prepareDurableSessionAutoTitleProjection } from "../application/session-title.service"; import { createSessionRuntimeEventProjection } from "../domain/session-runtime-event-projection"; +import { normalizeSessionTitle } from "../domain/session-title"; +import { prepareSessionAgentTaskSnapshotUpsert } from "./session-agent-task-snapshot.repository"; +import { prepareDurableSessionModelCallUsageProjection } from "./session-model-call.repository"; import type { + DriverRuntimeEventFence, InsertSessionEventResult, OneRuntimeEventPerSessionAllocation, OneRuntimeEventPerSessionInput, @@ -23,7 +45,7 @@ import type { SessionRuntimeEventRecord, SessionRuntimeEventSourceReceipt, } from "./session-runtime-event-store.types"; -import { appSessionViewerRuntimeEvents } from "./session-viewer-event-projection.repository"; +import { prepareSessionViewerRuntimeEventProjection } from "./session-viewer-event-projection.repository"; export type { OneRuntimeEventPerSessionInput, @@ -35,13 +57,18 @@ export type { } from "./session-runtime-event-store.types"; const MAX_SESSION_RUNTIME_EVENT_INSERT_ATTEMPTS = 5; -// D1 accepts at most 100 bound parameters; each session_event row binds 21. -const MAX_SESSION_EVENT_ROWS_PER_INSERT = 4; +// D1 accepts at most 100 bound parameters; the run-active fence adds its own binds. +const MAX_SESSION_EVENT_ROWS_PER_INSERT = 2; const WRITABLE_SESSION_STATUSES = ["IDLE", "RUNNING", "RESCHEDULING"] as const; const TERMINAL_LIFECYCLE_WRITABLE_SESSION_STATUSES = [ ...WRITABLE_SESSION_STATUSES, "TERMINATED", ] as const; +const RUN_TERMINAL_EVENT_TYPES: ReadonlySet = new Set([ + "run.cancelled", + "run.completed", + "run.failed", +]); function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); @@ -55,6 +82,12 @@ function isTerminalSessionLifecycleEvent(event: SessionRuntimeEventRecord): bool ); } +function assertNoRunTerminalEvents(records: readonly { event: SessionRuntimeEventRecord }[]): void { + if (records.some(({ event }) => RUN_TERMINAL_EVENT_TYPES.has(event.kind))) { + throw new Error("Run terminal events require the atomic terminal-run projection."); + } +} + function canWriteAfterTerminatedSession( records: readonly { event: SessionRuntimeEventRecord }[], ): boolean { @@ -104,25 +137,37 @@ function isSessionRuntimeEventSeqConflict(error: unknown): boolean { ); } +function isSessionRuntimeEventBatchFenceConflict(error: unknown): boolean { + const errorText = readErrorMessageTree(error); + return ( + errorText.includes("NOT NULL constraint failed: session_event.agent_id") || + errorText.includes("NOT NULL constraint failed: session_event.session_id") + ); +} + function readRuntimeEventEndedAt(event: SessionRuntimeEventRecord, fallbackMs: number): number { const endedAt = Date.parse(event.occurredAt); return Number.isFinite(endedAt) && endedAt >= fallbackMs ? endedAt : fallbackMs; } -async function allocateSessionRuntimeEventBatch( +function selectedValue(value: T, alias: string) { + return sql`${value}`.as(alias); +} + +async function readSessionRuntimeEventBatchAllocation( database: D1Database, input: { allowTerminatedSession: boolean; - count: number; sessionId: SessionId; }, ): Promise { const session = (await getAppDatabase(database) - .update(sessionsTable) - .set({ - runtimeEventSeqCursor: sql`${sessionsTable.runtimeEventSeqCursor} + ${input.count}`, + .select({ + agentId: sessionsTable.agentId, + seqCursor: sessionsTable.runtimeEventSeqCursor, }) + .from(sessionsTable) .where( and( eq(sessionsTable.id, input.sessionId), @@ -130,10 +175,6 @@ async function allocateSessionRuntimeEventBatch( inArray(sessionsTable.status, sessionWritableStatusValues(input.allowTerminatedSession)), ), ) - .returning({ - agentId: sessionsTable.agentId, - seqCursor: sessionsTable.runtimeEventSeqCursor, - }) .get()) ?? null; if (session === null) { @@ -142,7 +183,8 @@ async function allocateSessionRuntimeEventBatch( return { agentId: session.agentId, - firstSeq: session.seqCursor - input.count + 1, + firstSeq: session.seqCursor + 1, + previousCursor: session.seqCursor, }; } @@ -150,6 +192,7 @@ async function persistSessionRuntimeEventRows( database: D1Database, input: { allowTerminatedSession: boolean; + driverFence?: DriverRuntimeEventFence; rows: SerializedSessionRuntimeEventInput[]; sessionId: SessionId; }, @@ -163,38 +206,72 @@ async function persistSessionRuntimeEventRows( }; } - const projectedRows = input.rows.map((row, sourceIndex) => ({ - row: { - ...row, - projection: createSessionRuntimeEventProjection(row.event), - }, - sourceIndex, + let pendingRows = input.rows.map((row) => ({ + ...row, + projection: createSessionRuntimeEventProjection(row.event, { + provenMcpCommandId: row.provenMcpCommandId, + }), })); const timestampMs = currentTimestampMs(); for (let attempt = 0; attempt < MAX_SESSION_RUNTIME_EVENT_INSERT_ATTEMPTS; attempt += 1) { - const allocation = await allocateSessionRuntimeEventBatch(database, { + if (pendingRows.length === 0) { + return { + insertedCount: 0, + insertedRows: [], + insertedSessionIds: [], + insertedSourceEventIds: [], + }; + } + const allocation = await readSessionRuntimeEventBatchAllocation(database, { allowTerminatedSession: input.allowTerminatedSession, - count: projectedRows.length, sessionId: input.sessionId, }); + const projectedRows = pendingRows.map((row, sourceIndex) => ({ row, sourceIndex })); try { return await insertSessionRuntimeEventRows(database, { allocation, + allowTerminatedSession: input.allowTerminatedSession, + ...(input.driverFence === undefined ? {} : { driverFence: input.driverFence }), rows: projectedRows, sessionId: input.sessionId, timestampMs, }); } catch (error) { if ( - attempt < MAX_SESSION_RUNTIME_EVENT_INSERT_ATTEMPTS - 1 && - isSessionRuntimeEventSeqConflict(error) + !isSessionRuntimeEventSeqConflict(error) && + !isSessionRuntimeEventBatchFenceConflict(error) ) { + throw error; + } + const receipts = await getSessionRuntimeEventSourceReceipts(database, { + sessionId: input.sessionId, + sourceEventIds: pendingRows.map((row) => row.sourceEventId), + }); + pendingRows = pendingRows.filter((row) => { + const receipt = receipts.get(row.sourceEventId); + if (receipt === undefined) { + return true; + } + if ( + receipt.semanticHash !== row.semanticHash || + receipt.type !== row.projection.eventType + ) { + throw new Error( + `Runtime event source ${row.sourceEventId} conflicts with its durable receipt.`, + { cause: error }, + ); + } + return false; + }); + if (attempt < MAX_SESSION_RUNTIME_EVENT_INSERT_ATTEMPTS - 1) { continue; } - throw error; + throw new Error("Runtime event batch lost its atomic session or active-run fence.", { + cause: error, + }); } } @@ -329,10 +406,12 @@ async function allocateOneRuntimeEventPerSession( record === undefined ? false : canWriteAfterTerminatedSession([record]); const session = (await appDb - .update(sessionsTable) - .set({ - runtimeEventSeqCursor: sql`${sessionsTable.runtimeEventSeqCursor} + 1`, + .select({ + agentId: sessionsTable.agentId, + previousCursor: sessionsTable.runtimeEventSeqCursor, + sessionId: sessionsTable.id, }) + .from(sessionsTable) .where( and( eq(sessionsTable.id, sessionId), @@ -340,17 +419,13 @@ async function allocateOneRuntimeEventPerSession( inArray(sessionsTable.status, sessionWritableStatusValues(allowTerminatedSession)), ), ) - .returning({ - agentId: sessionsTable.agentId, - seq: sessionsTable.runtimeEventSeqCursor, - sessionId: sessionsTable.id, - }) .get()) ?? null; if (session !== null) { allocations.set(session.sessionId, { agentId: session.agentId, - seq: session.seq, + previousCursor: session.previousCursor, + seq: session.previousCursor + 1, sessionId: session.sessionId, }); } @@ -359,19 +434,26 @@ async function allocateOneRuntimeEventPerSession( return allocations; } -function toOneRuntimeEventPerSessionRows( +async function toOneRuntimeEventPerSessionRows( records: readonly OneRuntimeEventPerSessionInput[], -): OneRuntimeEventPerSessionRowInput[] { - return records.map((record) => ({ - event: record.event, - occurredAt: record.occurredAt, - projection: createSessionRuntimeEventProjection(record.event), - sessionId: record.sessionId, - sourceEventId: readSessionRuntimeEventSourceEventId({ +): Promise { + return Promise.all( + records.map(async (record) => ({ + artifactAttemptId: null, + artifactManifestJson: null, + artifactManifestSha256: null, event: record.event, - sourceEventId: null, - }), - })); + occurredAt: record.occurredAt, + projection: createSessionRuntimeEventProjection(record.event), + provenMcpCommandId: null, + semanticHash: await createRuntimeEventSemanticHash(record.event), + sessionId: record.sessionId, + sourceEventId: readSessionRuntimeEventSourceEventId({ + event: record.event, + sourceEventId: null, + }), + })), + ); } function readSessionRuntimeEventSourceEventId(input: { @@ -408,9 +490,124 @@ function toOneRuntimeEventPerSessionInsertValues(input: { }); } +async function filterNewOneRuntimeEventPerSessionRows( + database: D1Database, + rows: readonly OneRuntimeEventPerSessionRowInput[], +): Promise { + const results = await Promise.all( + rows.map(async (row) => { + const receipt = ( + await getSessionRuntimeEventSourceReceipts(database, { + sessionId: row.sessionId, + sourceEventIds: [row.sourceEventId], + }) + ).get(row.sourceEventId); + if (receipt === undefined) { + return row; + } + if (receipt.semanticHash !== row.semanticHash || receipt.type !== row.projection.eventType) { + throw new Error( + `Runtime event source ${row.sourceEventId} conflicts with its durable receipt.`, + ); + } + return null; + }), + ); + + return results.filter((row): row is OneRuntimeEventPerSessionRowInput => row !== null); +} + +async function prepareDurableRuntimeEventSideEffectProjection( + database: D1Database, + value: SessionEventInsertValue, + driverFence: DriverRuntimeEventFence | undefined, +): Promise { + if (driverFence === undefined) { + return []; + } + + if (value.event.kind === "session.info.updated") { + const title = readRuntimeEventString(readRuntimeEventPayload(value.event), "title"); + + if (title === null || title.trim().length === 0) { + return []; + } + + return prepareDurableSessionAutoTitleProjection(database, { + createdAt: value.createdAt, + eventId: value.id, + eventSeq: value.seq, + semanticHash: value.semanticHash, + sessionId: value.sessionId, + title: normalizeSessionTitle(title), + }); + } + + if (value.event.kind === "usage.updated") { + const usage = parseNullableSessionUsageSummary(value.event.payload); + + if (usage === null || driverFence.sessionRunId === null) { + return []; + } + if (value.event.driverInstanceId !== driverFence.driverInstanceId) { + throw new Error("Durable usage event is missing its exact Driver identity."); + } + if (value.runId !== driverFence.sessionRunId) { + throw new Error("Durable usage event is missing its exact Session Run identity."); + } + + return prepareDurableSessionModelCallUsageProjection(database, { + createdAtMs: value.createdAt, + driverInstanceId: driverFence.driverInstanceId, + eventId: value.id, + semanticHash: value.semanticHash, + sessionId: value.sessionId, + sessionRunId: driverFence.sessionRunId, + sourceEventSeq: value.seq, + traceId: value.traceId ?? driverFence.sessionRunId, + usage, + }); + } + + if (value.event.kind !== "runtime.resume.updated") { + return []; + } + + const nativeResumeRef = readNativeResumeRef(value.event); + + if (nativeResumeRef === null || driverFence.sessionRunId === null) { + return []; + } + if (value.event.driverInstanceId !== driverFence.driverInstanceId) { + throw new Error("Durable native resume event is missing its exact Driver identity."); + } + if (value.runId !== driverFence.sessionRunId) { + throw new Error("Durable native resume event is missing its exact Session Run identity."); + } + + return prepareNativeResumeRefProjection(database, { + createdAt: value.createdAt, + driverInstanceId: driverFence.driverInstanceId, + eventId: value.id, + nativeResumeRef, + observedEventSeq: value.seq, + semanticHash: value.semanticHash, + sessionId: value.sessionId, + sessionRunId: driverFence.sessionRunId, + }); +} + async function insertSessionEventRows( database: D1Database, values: readonly SessionEventInsertValue[], + atomicAllocations?: ReadonlyMap< + SessionId, + { + allowTerminatedSession: boolean; + driverFence?: DriverRuntimeEventFence; + previousCursor: number; + } + >, ): Promise { if (values.length === 0) { return { @@ -423,11 +620,62 @@ async function insertSessionEventRows( const appDatabase = getAppDatabase(database); const statements: D1PreparedStatement[] = []; + const receiptStatementIndexes: number[] = []; + const valuesBySession = Map.groupBy(values, (value) => value.sessionId); + const nextCursorBySession = new Map(); + + for (const [sessionId, allocation] of atomicAllocations ?? []) { + const sessionValues = valuesBySession.get(sessionId) ?? []; + if (sessionValues.length === 0) { + continue; + } + const nextCursor = allocation.previousCursor + sessionValues.length; + nextCursorBySession.set(sessionId, nextCursor); + const writableStatuses = sessionWritableStatusValues(allocation.allowTerminatedSession); + statements.push( + database + .prepare( + `UPDATE session + SET runtime_event_seq_cursor = ? + WHERE id = ? + AND runtime_event_seq_cursor = ? + AND archived_at IS NULL + AND status IN (${writableStatuses.map(() => "?").join(", ")})`, + ) + .bind(nextCursor, sessionId, allocation.previousCursor, ...writableStatuses), + ); + } + + const insertSize = atomicAllocations === undefined ? MAX_SESSION_EVENT_ROWS_PER_INSERT : 1; + for (let index = 0; index < values.length; index += insertSize) { + const chunk = values.slice(index, index + insertSize); + const firstValue = chunk[0]; + + if (firstValue === undefined) { + continue; + } + + const selection = createSessionEventInsertSelect( + appDatabase, + firstValue, + nextCursorBySession.get(firstValue.sessionId) ?? null, + atomicAllocations?.get(firstValue.sessionId)?.driverFence, + ); + + for (const value of chunk.slice(1)) { + selection.unionAll( + createSessionEventInsertSelect( + appDatabase, + value, + nextCursorBySession.get(value.sessionId) ?? null, + atomicAllocations?.get(value.sessionId)?.driverFence, + ), + ); + } - for (let index = 0; index < values.length; index += MAX_SESSION_EVENT_ROWS_PER_INSERT) { const query = appDatabase .insert(sessionEventsTable) - .values(values.slice(index, index + MAX_SESSION_EVENT_ROWS_PER_INSERT)) + .select(selection) .onConflictDoNothing({ target: [sessionEventsTable.sessionId, sessionEventsTable.sourceEventId], }) @@ -437,18 +685,128 @@ async function insertSessionEventRows( }) .toSQL(); + receiptStatementIndexes.push(statements.length); statements.push(database.prepare(query.sql).bind(...query.params)); + + for (const value of chunk) { + if (value.agentTaskSnapshot !== null) { + statements.push( + prepareSessionAgentTaskSnapshotUpsert(database, { + eventId: value.id, + snapshot: value.agentTaskSnapshot, + }), + ); + } + statements.push( + ...prepareSessionViewerRuntimeEventProjection(database, { + createdAt: value.createdAt, + event: value.event, + eventId: value.id, + sessionId: value.sessionId, + }), + ); + statements.push( + ...(await prepareDurableRuntimeEventSideEffectProjection( + database, + value, + atomicAllocations?.get(value.sessionId)?.driverFence, + )), + ); + if ( + value.artifactAttemptId !== null && + value.artifactManifestJson !== null && + value.artifactManifestSha256 !== null + ) { + statements.push( + ...prepareRuntimeArtifactPromotion(database, { + attemptId: value.artifactAttemptId, + eventId: value.id, + manifestJson: value.artifactManifestJson, + manifestSha256: value.artifactManifestSha256, + timestampMs: value.createdAt, + }), + ); + } + } + } + + for (const [sessionId, nextCursor] of nextCursorBySession) { + const sessionValues = valuesBySession.get(sessionId) ?? []; + statements.push( + database + .prepare( + `INSERT INTO session_event (id) + SELECT ? + WHERE NOT EXISTS ( + SELECT 1 + FROM session AS s + WHERE s.id = ? + AND s.runtime_event_seq_cursor = ? + AND ( + SELECT COUNT(*) + FROM session_event AS e + WHERE e.session_id = s.id + AND e.id IN (SELECT value FROM json_each(?)) + ) = ? + )`, + ) + .bind( + createPlatformId(), + sessionId, + nextCursor, + JSON.stringify(sessionValues.map((value) => value.id)), + sessionValues.length, + ), + ); } const results = await database.batch<{ session_id: SessionId; source_event_id: string }>( statements, ); - const insertedRows = results.flatMap((result) => - result.results.map((row) => ({ + const insertedRows = receiptStatementIndexes.flatMap((resultIndex) => { + const result = results[resultIndex]; + + return (result?.results ?? []).map((row) => ({ sessionId: row.session_id, sourceEventId: row.source_event_id, - })), + })); + }); + const insertedKeys = new Set(insertedRows.map(sessionSourceEventKey)); + const replayCandidates = values.filter( + (value) => + !insertedKeys.has( + sessionSourceEventKey({ + sessionId: value.sessionId, + sourceEventId: value.sourceEventId, + }), + ), ); + const replayCandidatesBySession = new Map(); + for (const candidate of replayCandidates) { + replayCandidatesBySession.set(candidate.sessionId, [ + ...(replayCandidatesBySession.get(candidate.sessionId) ?? []), + candidate, + ]); + } + + for (const [sessionId, candidates] of replayCandidatesBySession) { + const receipts = await getSessionRuntimeEventSourceReceipts(database, { + sessionId, + sourceEventIds: candidates.map((candidate) => candidate.sourceEventId), + }); + for (const candidate of candidates) { + const receipt = receipts.get(candidate.sourceEventId); + if ( + receipt === undefined || + receipt.semanticHash !== candidate.semanticHash || + receipt.type !== candidate.eventType + ) { + throw new Error( + `Runtime event source ${candidate.sourceEventId} conflicts with its durable receipt.`, + ); + } + } + } return { insertedCount: insertedRows.length, @@ -473,43 +831,37 @@ export async function persistOneRuntimeEventPerSession( assertUniqueRuntimeEventSessions(input.records); assertOneRuntimeEventPerSessionMatches(input.records); + assertNoRunTerminalEvents(input.records); await ensureRuntimeEventRunsMatchSessions(database, readRuntimeEventRunScopes(input.records)); - const rows = toOneRuntimeEventPerSessionRows(input.records); + const rows = await toOneRuntimeEventPerSessionRows(input.records); + let pendingRows = await filterNewOneRuntimeEventPerSessionRows(database, rows); const timestampMs = currentTimestampMs(); for (let attempt = 0; attempt < MAX_SESSION_RUNTIME_EVENT_INSERT_ATTEMPTS; attempt += 1) { - const allocations = await allocateOneRuntimeEventPerSession(database, input.records); + const allocations = await allocateOneRuntimeEventPerSession(database, pendingRows); const values = toOneRuntimeEventPerSessionInsertValues({ allocations, - rows, + rows: pendingRows, timestampMs, }); + const atomicAllocations = new Map( + [...allocations].map(([sessionId, allocation]) => [ + sessionId, + { + allowTerminatedSession: + pendingRows.find((row) => row.sessionId === sessionId) !== undefined && + canWriteAfterTerminatedSession( + pendingRows.filter((row) => row.sessionId === sessionId), + ), + previousCursor: allocation.previousCursor, + }, + ]), + ); try { - const insertResult = await insertSessionEventRows(database, values); + const insertResult = await insertSessionEventRows(database, values, atomicAllocations); const insertedSessionIds = new Set(insertResult.insertedSessionIds); - const insertedKeys = new Set(insertResult.insertedRows.map(sessionSourceEventKey)); - - await appSessionViewerRuntimeEvents( - database, - rows.flatMap((row) => - insertedKeys.has( - sessionSourceEventKey({ - sessionId: row.sessionId, - sourceEventId: row.sourceEventId, - }), - ) - ? [ - { - event: row.event, - occurredAt: row.occurredAt, - sessionId: row.sessionId, - }, - ] - : [], - ), - ); return { persistedCount: insertResult.insertedCount, @@ -521,13 +873,19 @@ export async function persistOneRuntimeEventPerSession( }; } catch (error) { if ( - attempt < MAX_SESSION_RUNTIME_EVENT_INSERT_ATTEMPTS - 1 && - isSessionRuntimeEventSeqConflict(error) + !isSessionRuntimeEventSeqConflict(error) && + !isSessionRuntimeEventBatchFenceConflict(error) ) { + throw error; + } + pendingRows = await filterNewOneRuntimeEventPerSessionRows(database, pendingRows); + if (attempt < MAX_SESSION_RUNTIME_EVENT_INSERT_ATTEMPTS - 1) { continue; } - throw error; + throw new Error("Runtime event batch lost its atomic session or active-run fence.", { + cause: error, + }); } } @@ -538,7 +896,7 @@ export async function persistOneRuntimeEventPerSession( } function toSessionRuntimeEventInsertValue(input: { - allocation: SessionRuntimeEventBatchAllocation; + allocation: Pick; row: ProjectedSessionRuntimeEventInput; sessionId: SessionId; sourceIndex: number; @@ -548,30 +906,163 @@ function toSessionRuntimeEventInsertValue(input: { const occurredAt = input.row.occurredAt ?? input.timestampMs + input.sourceIndex; return { + agentTaskSnapshot: + input.row.event.kind === "agent.tasks.replaced" + ? readRuntimeAgentTaskSnapshot(input.row.event) + : null, agentId: input.allocation.agentId, + artifactAttemptId: input.row.artifactAttemptId, + artifactManifestJson: input.row.artifactManifestJson, + artifactManifestSha256: input.row.artifactManifestSha256, contentText: input.row.projection.contentText, createdAt: input.timestampMs + input.sourceIndex, endedAt: readRuntimeEventEndedAt(input.row.event, occurredAt), + event: input.row.event, eventType: input.row.projection.eventType, family: input.row.projection.family, id, + mcpCommandId: input.row.projection.mcpCommandId, occurredAt, processStatus: input.row.projection.processStatus, processType: input.row.projection.processType, runId: input.row.projection.runId, + runtimeOperationEventJson: null, + semanticHash: input.row.semanticHash, + terminalEventJson: null, seq: input.allocation.firstSeq + input.sourceIndex, sessionId: input.sessionId, sourceEventId: input.row.sourceEventId, source: input.row.projection.source, + streamId: input.row.projection.streamId, toolCallId: input.row.projection.toolCallId, + toolInputDeltaJson: input.row.projection.toolInputDeltaJson, toolInputJson: input.row.projection.toolInputJson, toolName: input.row.projection.toolName, + toolOutputDeltaText: input.row.projection.toolOutputDeltaText, + toolOutputText: input.row.projection.toolOutputText, + toolParentMessageId: input.row.projection.toolParentMessageId, + toolResultMessageId: input.row.projection.toolResultMessageId, + toolStatus: input.row.projection.toolStatus, tokens: input.row.projection.tokens, traceId: input.row.projection.traceId, visibility: input.row.projection.visibility, }; } +function createSessionEventInsertSelect( + database: AppDatabase, + value: SessionEventInsertValue, + requiredCursor: number | null = null, + driverFence?: DriverRuntimeEventFence, +) { + return database + .select({ + agentId: selectedValue(value.agentId, "agent_id"), + artifactAttemptId: selectedValue(value.artifactAttemptId, "artifact_attempt_id"), + artifactManifestJson: selectedValue(value.artifactManifestJson, "artifact_manifest_json"), + artifactManifestSha256: selectedValue( + value.artifactManifestSha256, + "artifact_manifest_sha256", + ), + contentText: selectedValue(value.contentText, "content_text"), + createdAt: selectedValue(value.createdAt, "created_at"), + endedAt: selectedValue(value.endedAt, "ended_at"), + eventType: selectedValue(value.eventType, "event_type"), + family: selectedValue(value.family, "family"), + id: selectedValue(value.id, "id"), + mcpCommandId: selectedValue(value.mcpCommandId, "mcp_command_id"), + occurredAt: selectedValue(value.occurredAt, "occurred_at"), + processStatus: selectedValue(value.processStatus, "process_status"), + processType: selectedValue(value.processType, "process_type"), + runId: selectedValue(value.runId, "run_id"), + runtimeOperationEventJson: selectedValue( + value.runtimeOperationEventJson, + "runtime_operation_event_json", + ), + semanticHash: selectedValue(value.semanticHash, "semantic_hash"), + seq: selectedValue(value.seq, "seq"), + sessionId: selectedValue(value.sessionId, "session_id"), + sourceEventId: selectedValue(value.sourceEventId, "source_event_id"), + source: selectedValue(value.source, "source"), + streamId: selectedValue(value.streamId, "stream_id"), + terminalEventJson: selectedValue(value.terminalEventJson, "terminal_event_json"), + toolCallId: selectedValue(value.toolCallId, "tool_call_id"), + toolInputDeltaJson: selectedValue(value.toolInputDeltaJson, "tool_input_delta_json"), + toolInputJson: selectedValue(value.toolInputJson, "tool_input_json"), + toolName: selectedValue(value.toolName, "tool_name"), + toolOutputDeltaText: selectedValue(value.toolOutputDeltaText, "tool_output_delta_text"), + toolOutputText: selectedValue(value.toolOutputText, "tool_output_text"), + toolParentMessageId: selectedValue(value.toolParentMessageId, "tool_parent_message_id"), + toolResultMessageId: selectedValue(value.toolResultMessageId, "tool_result_message_id"), + toolStatus: selectedValue(value.toolStatus, "tool_status"), + tokens: selectedValue(value.tokens, "tokens"), + traceId: selectedValue(value.traceId, "trace_id"), + visibility: selectedValue(value.visibility, "visibility"), + }) + .from(sessionsTable) + .where( + and( + eq(sessionsTable.id, value.sessionId), + isNull(sessionsTable.archivedAt), + requiredCursor === null + ? undefined + : eq(sessionsTable.runtimeEventSeqCursor, requiredCursor), + driverFence === undefined + ? undefined + : and( + exists( + database + .select({ id: driverInstancesTable.id }) + .from(driverInstancesTable) + .where( + and( + eq(driverInstancesTable.id, driverFence.driverInstanceId), + eq(driverInstancesTable.connectionId, driverFence.connectionId), + eq(driverInstancesTable.generation, driverFence.generation), + eq(driverInstancesTable.sandboxSessionId, value.sessionId), + ), + ), + ), + driverFence.sessionRunId === null + ? undefined + : and( + eq(sessionsTable.lastRunId, driverFence.sessionRunId), + eq(sessionsTable.status, "RUNNING"), + isNull(sessionsTable.statusOperationId), + exists( + database + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, driverFence.sessionRunId), + eq(sessionRunsTable.sessionId, value.sessionId), + eq(sessionRunsTable.driverInstanceId, driverFence.driverInstanceId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ), + ), + ), + ), + value.runId === null + ? undefined + : exists( + database + .select({ id: sessionRunsTable.id }) + .from(sessionRunsTable) + .where( + and( + eq(sessionRunsTable.id, value.runId), + eq(sessionRunsTable.sessionId, value.sessionId), + inArray(sessionRunsTable.status, ACTIVE_SESSION_RUN_STATUSES), + ), + ), + ), + ), + ) + .$dynamic(); +} + function toSessionRuntimeEventInsertValues(input: { allocation: SessionRuntimeEventBatchAllocation; rows: ProjectedSessionRuntimeEventRowInput[]; @@ -593,18 +1084,33 @@ async function insertSessionRuntimeEventRows( database: D1Database, input: { allocation: SessionRuntimeEventBatchAllocation; + allowTerminatedSession: boolean; + driverFence?: DriverRuntimeEventFence; rows: ProjectedSessionRuntimeEventRowInput[]; sessionId: SessionId; timestampMs: number; }, ): Promise { - return insertSessionEventRows(database, toSessionRuntimeEventInsertValues(input)); + return insertSessionEventRows( + database, + toSessionRuntimeEventInsertValues(input), + new Map([ + [ + input.sessionId, + { + allowTerminatedSession: input.allowTerminatedSession, + ...(input.driverFence === undefined ? {} : { driverFence: input.driverFence }), + previousCursor: input.allocation.previousCursor, + }, + ], + ]), + ); } async function filterNewSessionRuntimeEventInputs( database: D1Database, input: PersistSessionRuntimeEventsInput, -): Promise { +): Promise<(SessionRuntimeEventInput & { semanticHash: string })[]> { const sourceEventIds = input.records.map((record) => readSessionRuntimeEventSourceEventId({ event: record.event, @@ -616,19 +1122,34 @@ async function filterNewSessionRuntimeEventInputs( sessionId: input.sessionId, sourceEventIds, }); - const acceptedSourceIds = new Set(); + const acceptedHashes = new Map(); + const records = await Promise.all( + input.records.map(async (record) => ({ + ...record, + semanticHash: await createRuntimeEventSemanticHash(record.event), + })), + ); - return input.records.filter((record) => { + return records.filter((record) => { const sourceEventId = readSessionRuntimeEventSourceEventId({ event: record.event, sourceEventId: record.sourceEventId, }); + const receipt = persistedReceipts.get(sourceEventId); + const acceptedHash = acceptedHashes.get(sourceEventId); + + if ( + (receipt !== undefined && receipt.semanticHash !== record.semanticHash) || + (acceptedHash !== undefined && acceptedHash !== record.semanticHash) + ) { + throw new Error(`Runtime event source ${sourceEventId} conflicts with its durable receipt.`); + } - if (persistedReceipts.has(sourceEventId) || acceptedSourceIds.has(sourceEventId)) { + if (receipt !== undefined || acceptedHash !== undefined) { return false; } - acceptedSourceIds.add(sourceEventId); + acceptedHashes.set(sourceEventId, record.semanticHash); return true; }); } @@ -638,6 +1159,16 @@ export async function persistSessionRuntimeEvents( input: PersistSessionRuntimeEventsInput, ): Promise { assertRuntimeEventBatchSessionMatches(input); + assertNoRunTerminalEvents(input.records); + if ( + input.driverFence !== undefined && + input.records.some( + (record) => + record.event.runId !== undefined && record.event.runId !== input.driverFence?.sessionRunId, + ) + ) { + throw new Error("Driver runtime event does not match its fenced Session Run."); + } await ensureRuntimeEventRunsMatchSessions( database, readRuntimeEventRunScopes( @@ -658,36 +1189,29 @@ export async function persistSessionRuntimeEvents( }; } - const rows = records.map((record) => ({ - event: record.event, - occurredAt: record.occurredAt, - sourceEventId: readSessionRuntimeEventSourceEventId({ + const rows = await Promise.all( + records.map(async (record) => ({ + artifactAttemptId: record.artifactAttemptId ?? null, + artifactManifestJson: record.artifactManifestJson ?? null, + artifactManifestSha256: record.artifactManifestSha256 ?? null, event: record.event, - sourceEventId: record.sourceEventId, - }), - })); + occurredAt: record.occurredAt, + provenMcpCommandId: record.provenMcpCommandId ?? null, + semanticHash: record.semanticHash, + sourceEventId: readSessionRuntimeEventSourceEventId({ + event: record.event, + sourceEventId: record.sourceEventId, + }), + })), + ); const result = await persistSessionRuntimeEventRows(database, { allowTerminatedSession: canWriteAfterTerminatedSession(records), + ...(input.driverFence === undefined ? {} : { driverFence: input.driverFence }), rows, sessionId: input.sessionId, }); const insertedSourceEventIds = new Set(result.insertedSourceEventIds); - await appSessionViewerRuntimeEvents( - database, - rows.flatMap((row) => - insertedSourceEventIds.has(row.sourceEventId) - ? [ - { - event: row.event, - occurredAt: row.occurredAt, - sessionId: input.sessionId, - }, - ] - : [], - ), - ); - return { persistedCount: result.insertedCount, persistedEvents: rows @@ -713,6 +1237,7 @@ export async function getSessionRuntimeEventSourceReceipts( const rows = await getAppDatabase(database) .select({ event_id: sessionEventsTable.sourceEventId, + semantic_hash: sessionEventsTable.semanticHash, seq: sessionEventsTable.seq, type: sessionEventsTable.eventType, }) @@ -730,6 +1255,7 @@ export async function getSessionRuntimeEventSourceReceipts( for (const row of rows) { receipts.set(row.event_id, { eventId: row.event_id, + semanticHash: row.semantic_hash, seq: row.seq, type: row.type, }); diff --git a/apps/api/src/modules/sessions/infrastructure/session-runtime-event-store.types.ts b/apps/api/src/modules/sessions/infrastructure/session-runtime-event-store.types.ts index d6f8ee2b..f65f78a1 100644 --- a/apps/api/src/modules/sessions/infrastructure/session-runtime-event-store.types.ts +++ b/apps/api/src/modules/sessions/infrastructure/session-runtime-event-store.types.ts @@ -1,4 +1,12 @@ -import type { AgentId, RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; +import type { AgentTaskSnapshot } from "@mosoo/contracts/session"; +import type { + AgentId, + DriverCommandId, + DriverInstanceId, + RuntimeEventId, + SessionId, + SessionRunId, +} from "@mosoo/id"; import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; import type { createSessionRuntimeEventProjection } from "../domain/session-runtime-event-projection"; @@ -6,10 +14,18 @@ import type { createSessionRuntimeEventProjection } from "../domain/session-runt type SessionRuntimeEventProjection = ReturnType; export interface PersistSessionRuntimeEventsInput { + readonly driverFence?: DriverRuntimeEventFence; records: readonly SessionRuntimeEventInput[]; readonly sessionId: SessionId; } +export interface DriverRuntimeEventFence { + readonly connectionId: string; + readonly driverInstanceId: DriverInstanceId; + readonly generation: number; + readonly sessionRunId: SessionRunId | null; +} + export interface PersistSessionRuntimeEventsResult { readonly persistedCount: number; readonly persistedEvents: readonly SessionRuntimeEventRecord[]; @@ -17,16 +33,25 @@ export interface PersistSessionRuntimeEventsResult { } export interface SessionRuntimeEventInput { + readonly artifactAttemptId?: string | null; + readonly artifactManifestJson?: string | null; + readonly artifactManifestSha256?: string | null; readonly event: SessionRuntimeEventRecord; readonly occurredAt: number | null; + readonly provenMcpCommandId?: DriverCommandId | null; readonly sourceEventId: string | null; } export type SessionRuntimeEventRecord = RuntimeEventEnvelope; export interface SerializedSessionRuntimeEventInput { + readonly artifactAttemptId: string | null; + readonly artifactManifestJson: string | null; + readonly artifactManifestSha256: string | null; readonly event: SessionRuntimeEventRecord; readonly occurredAt: number | null; + readonly provenMcpCommandId: DriverCommandId | null; + readonly semanticHash: string; readonly sourceEventId: string; } @@ -42,10 +67,12 @@ export interface ProjectedSessionRuntimeEventRowInput { export interface SessionRuntimeEventBatchAllocation { readonly agentId: AgentId; readonly firstSeq: number; + readonly previousCursor: number; } export interface OneRuntimeEventPerSessionAllocation { readonly agentId: AgentId; + readonly previousCursor: number; readonly seq: number; readonly sessionId: SessionId; } @@ -68,29 +95,46 @@ export interface PersistOneRuntimeEventPerSessionResult { export interface SessionRuntimeEventSourceReceipt { readonly eventId: string; + readonly semanticHash: string | null; readonly seq: number; readonly type: string; } export interface SessionEventInsertValue { + readonly agentTaskSnapshot: AgentTaskSnapshot | null; readonly agentId: AgentId; + readonly artifactAttemptId: string | null; + readonly artifactManifestJson: string | null; + readonly artifactManifestSha256: string | null; readonly contentText: string; readonly createdAt: number; readonly endedAt: number; + readonly event: SessionRuntimeEventRecord; readonly eventType: string; readonly family: SessionRuntimeEventProjection["family"]; readonly id: RuntimeEventId; + readonly mcpCommandId: DriverCommandId | null; readonly occurredAt: number; readonly processStatus: SessionRuntimeEventProjection["processStatus"]; readonly processType: SessionRuntimeEventProjection["processType"]; readonly runId: SessionRunId | null; + readonly runtimeOperationEventJson: string | null; + readonly semanticHash: string; + readonly terminalEventJson: string | null; readonly seq: number; readonly sessionId: SessionId; readonly source: SessionRuntimeEventProjection["source"]; readonly sourceEventId: string; + readonly streamId: string | null; readonly toolCallId: string | null; + readonly toolInputDeltaJson: string | null; readonly toolInputJson: string | null; readonly toolName: string | null; + readonly toolOutputDeltaText: string | null; + readonly toolOutputText: string | null; + readonly toolParentMessageId: string | null; + readonly toolResultMessageId: string | null; + readonly toolStatus: SessionRuntimeEventProjection["toolStatus"]; readonly tokens: number | null; readonly traceId: string | null; readonly visibility: SessionRuntimeEventProjection["visibility"]; diff --git a/apps/api/src/modules/sessions/infrastructure/session-viewer-event-projection.repository.ts b/apps/api/src/modules/sessions/infrastructure/session-viewer-event-projection.repository.ts index c18521cb..1750b3dd 100644 --- a/apps/api/src/modules/sessions/infrastructure/session-viewer-event-projection.repository.ts +++ b/apps/api/src/modules/sessions/infrastructure/session-viewer-event-projection.repository.ts @@ -2,246 +2,331 @@ import type { SessionPermissionRequestView, SessionReadinessSnapshotView, } from "@mosoo/ag-ui-session"; -import { - SessionPermissionRequestViewSchema, - SessionReadinessSnapshotViewSchema, -} from "@mosoo/ag-ui-session"; +import { SessionReadinessSnapshotViewSchema } from "@mosoo/ag-ui-session"; import { parseSchemaValue } from "@mosoo/contracts/validation"; -import { sessionPermissionRequestsTable, sessionReadinessSnapshotsTable } from "@mosoo/db"; -import type { DriverInstanceId, SessionId, SessionRunId } from "@mosoo/id"; +import { createPlatformId } from "@mosoo/id"; +import type { RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; import { readRuntimeEventPayload, readRuntimeEventPermissionRequest, + readRuntimeRunPayload, readRuntimeEventString, } from "@mosoo/runtime-events"; import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; -import { and, eq } from "drizzle-orm"; - -import { getAppDatabase } from "../../../platform/db/drizzle"; -import { currentTimestampMs } from "../../../time"; export interface SessionViewerProjectionRuntimeEvent { + readonly createdAt: number; readonly event: RuntimeEventEnvelope; - readonly occurredAt: number | null; + readonly eventId: RuntimeEventId; readonly sessionId: SessionId; } -function toProjectionTimestamp(record: SessionViewerProjectionRuntimeEvent): number { - if (record.occurredAt !== null) { - return record.occurredAt; - } +const INSERTED_EVENT_FENCE = `EXISTS ( + SELECT 1 + FROM session_event AS receipt + WHERE receipt.id = ? + AND receipt.session_id = ? + AND receipt.event_type = ? +)`; - const occurredAt = Date.parse(record.event.occurredAt); - return Number.isFinite(occurredAt) ? occurredAt : currentTimestampMs(); -} - -function parsePermissionRequestViews(value: unknown): SessionPermissionRequestView[] | null { - if (!Array.isArray(value)) { - return null; - } - - const requests: SessionPermissionRequestView[] = []; - - for (const entry of value) { - requests.push(parseSchemaValue(SessionPermissionRequestViewSchema, entry)); - } - - return requests; +function preparePermissionRequestUpsert( + database: D1Database, + record: SessionViewerProjectionRuntimeEvent, + request: SessionPermissionRequestView, +): D1PreparedStatement { + return database + .prepare( + `INSERT INTO session_permission_request ( + created_at, driver_instance_id, raw_input, request_id, run_id, + session_id, title, tool_call_id, tool_kind, updated_at + ) + SELECT ?, ?, ?, ?, ?, ?, ?, ?, ?, ? + WHERE ${INSERTED_EVENT_FENCE} + ON CONFLICT (session_id, request_id) DO UPDATE SET + driver_instance_id = excluded.driver_instance_id, + raw_input = excluded.raw_input, + run_id = excluded.run_id, + title = excluded.title, + tool_call_id = excluded.tool_call_id, + tool_kind = excluded.tool_kind, + updated_at = excluded.updated_at`, + ) + .bind( + record.createdAt, + request.driverInstanceId, + request.rawInput, + request.requestId, + request.runId, + record.sessionId, + request.title, + request.toolCallId, + request.toolKind, + record.createdAt, + record.eventId, + record.sessionId, + record.event.kind, + ); } -async function upsertPermissionRequest( +function preparePermissionRequestDelete( database: D1Database, record: SessionViewerProjectionRuntimeEvent, -): Promise { - const request = readRuntimeEventPermissionRequest(record.event); + runId: SessionRunId, + requestId: string, +): D1PreparedStatement { + return database + .prepare( + `DELETE FROM session_permission_request + WHERE session_id = ? + AND run_id = ? + AND request_id = ? + AND ${INSERTED_EVENT_FENCE}`, + ) + .bind(record.sessionId, runId, requestId, record.eventId, record.sessionId, record.event.kind); +} - if (request === null) { - return; +function requirePermissionRunId(record: SessionViewerProjectionRuntimeEvent): SessionRunId { + if (record.event.runId === undefined) { + throw new Error(`Runtime event ${record.event.kind} requires an exact Session Run identity.`); } - - const timestamp = toProjectionTimestamp(record); - - await getAppDatabase(database) - .insert(sessionPermissionRequestsTable) - .values({ - createdAt: timestamp, - driverInstanceId: request.driverInstanceId, - rawInput: request.rawInput, - requestId: request.requestId, - runId: request.runId, - sessionId: record.sessionId, - title: request.title, - toolCallId: request.toolCallId, - toolKind: request.toolKind, - updatedAt: timestamp, - }) - .onConflictDoUpdate({ - set: { - driverInstanceId: request.driverInstanceId, - rawInput: request.rawInput, - runId: request.runId, - title: request.title, - toolCallId: request.toolCallId, - toolKind: request.toolKind, - updatedAt: timestamp, - }, - target: [sessionPermissionRequestsTable.sessionId, sessionPermissionRequestsTable.requestId], - }) - .run(); + return record.event.runId; } -async function replacePermissionRequests( +function preparePermissionRunStatusUpdate( database: D1Database, record: SessionViewerProjectionRuntimeEvent, - permissionRequests: readonly SessionPermissionRequestView[], -): Promise { - const timestamp = toProjectionTimestamp(record); - const db = getAppDatabase(database); - - await db - .delete(sessionPermissionRequestsTable) - .where(eq(sessionPermissionRequestsTable.sessionId, record.sessionId)) - .run(); - - if (permissionRequests.length === 0) { - return; - } - - await db - .insert(sessionPermissionRequestsTable) - .values( - permissionRequests.map((request) => ({ - createdAt: timestamp, - driverInstanceId: request.driverInstanceId as DriverInstanceId, - rawInput: request.rawInput, - requestId: request.requestId, - runId: request.runId as SessionRunId, - sessionId: record.sessionId, - title: request.title, - toolCallId: request.toolCallId, - toolKind: request.toolKind, - updatedAt: timestamp, - })), + runId: SessionRunId, +): D1PreparedStatement { + return database + .prepare( + `WITH desired AS ( + SELECT CASE WHEN EXISTS ( + SELECT 1 FROM session_permission_request AS permission + WHERE permission.session_id = ? AND permission.run_id = ? + ) THEN 'waiting_input' ELSE 'running' END AS next_status + ) + UPDATE session_run + SET status = (SELECT next_status FROM desired), + status_changed_at = CASE + WHEN status <> (SELECT next_status FROM desired) THEN ? ELSE status_changed_at END, + status_event = CASE + WHEN status <> (SELECT next_status FROM desired) + THEN CASE (SELECT next_status FROM desired) + WHEN 'waiting_input' THEN 'run.wait_for_input' ELSE 'run.start' END + ELSE status_event END, + status_operation_id = CASE + WHEN status <> (SELECT next_status FROM desired) THEN NULL ELSE status_operation_id END, + status_seq = status_seq + CASE + WHEN status <> (SELECT next_status FROM desired) THEN 1 ELSE 0 END, + status_source = CASE + WHEN status <> (SELECT next_status FROM desired) THEN ? ELSE status_source END, + updated_at = CASE + WHEN status <> (SELECT next_status FROM desired) THEN ? ELSE updated_at END + WHERE id = ? + AND session_id = ? + AND status IN ('running', 'waiting_input') + AND ${INSERTED_EVENT_FENCE}`, ) - .run(); + .bind( + record.sessionId, + runId, + record.createdAt, + record.event.origin, + record.createdAt, + runId, + record.sessionId, + record.eventId, + record.sessionId, + record.event.kind, + ); } -async function removePermissionRequestById( +function preparePermissionProjectionGuard( database: D1Database, - input: { - readonly requestId: string; - readonly sessionId: SessionId; - }, -): Promise { - await getAppDatabase(database) - .delete(sessionPermissionRequestsTable) - .where( - and( - eq(sessionPermissionRequestsTable.sessionId, input.sessionId), - eq(sessionPermissionRequestsTable.requestId, input.requestId), - ), + record: SessionViewerProjectionRuntimeEvent, + runId: SessionRunId, +): D1PreparedStatement { + return database + .prepare( + `INSERT INTO session_event (id) + SELECT ? + WHERE NOT EXISTS ( + SELECT 1 + FROM session_event AS receipt + INNER JOIN session_run AS run + ON run.id = ? AND run.session_id = receipt.session_id + WHERE receipt.id = ? + AND receipt.session_id = ? + AND receipt.event_type = ? + AND run.status = CASE WHEN EXISTS ( + SELECT 1 FROM session_permission_request AS permission + WHERE permission.session_id = receipt.session_id AND permission.run_id = run.id + ) THEN 'waiting_input' ELSE 'running' END + )`, ) - .run(); + .bind( + createPlatformId(), + runId, + record.eventId, + record.sessionId, + record.event.kind, + ); } -async function clearRunPermissionRequests( +function preparePermissionRequested( database: D1Database, record: SessionViewerProjectionRuntimeEvent, -): Promise { - const runId = record.event.runId; - - if (runId === undefined) { - return; +): D1PreparedStatement[] { + const request = readRuntimeEventPermissionRequest(record.event); + if (request === null) { + return []; } - - await getAppDatabase(database) - .delete(sessionPermissionRequestsTable) - .where( - and( - eq(sessionPermissionRequestsTable.sessionId, record.sessionId), - eq(sessionPermissionRequestsTable.runId, runId), - ), - ) - .run(); + const runId = requirePermissionRunId(record); + if (request.runId !== runId) { + throw new Error("Permission request payload conflicts with its event Run identity."); + } + return [ + preparePermissionRequestUpsert(database, record, request), + preparePermissionRunStatusUpdate(database, record, runId), + preparePermissionProjectionGuard(database, record, runId), + ]; } -async function appPermissionResolution( +function preparePermissionResolved( database: D1Database, record: SessionViewerProjectionRuntimeEvent, -): Promise { +): D1PreparedStatement[] { const payload = readRuntimeEventPayload(record.event); - const permissionRequests = parsePermissionRequestViews(payload["permissionRequests"]); - - if (permissionRequests !== null) { - await replacePermissionRequests(database, record, permissionRequests); - return; - } - const requestId = readRuntimeEventString(payload, "requestId"); - - if (requestId !== null) { - await removePermissionRequestById(database, { - requestId, - sessionId: record.sessionId, - }); + if (requestId === null) { + return []; } + const runId = requirePermissionRunId(record); + return [ + preparePermissionRequestDelete(database, record, runId, requestId), + preparePermissionRunStatusUpdate(database, record, runId), + preparePermissionProjectionGuard(database, record, runId), + ]; } -async function upsertReadinessSnapshot( +function prepareReadinessUpsert( database: D1Database, record: SessionViewerProjectionRuntimeEvent, -): Promise { +): D1PreparedStatement { const readiness = parseSchemaValue( SessionReadinessSnapshotViewSchema, record.event.payload, ) satisfies SessionReadinessSnapshotView; - const timestamp = toProjectionTimestamp(record); - const readinessJson = JSON.stringify(readiness); - await getAppDatabase(database) - .insert(sessionReadinessSnapshotsTable) - .values({ - readinessJson, - sessionId: record.sessionId, - updatedAt: timestamp, - }) - .onConflictDoUpdate({ - set: { - readinessJson, - updatedAt: timestamp, - }, - target: sessionReadinessSnapshotsTable.sessionId, - }) - .run(); + return database + .prepare( + `INSERT INTO session_readiness_snapshot (readiness_json, session_id, updated_at) + SELECT ?, ?, ? + WHERE ${INSERTED_EVENT_FENCE} + ON CONFLICT (session_id) DO UPDATE SET + readiness_json = excluded.readiness_json, + updated_at = excluded.updated_at`, + ) + .bind( + JSON.stringify(readiness), + record.sessionId, + record.createdAt, + record.eventId, + record.sessionId, + record.event.kind, + ); } -export async function appSessionViewerRuntimeEvents( +function prepareRunStarted( database: D1Database, - records: readonly SessionViewerProjectionRuntimeEvent[], -): Promise { - for (const record of records) { - switch (record.event.kind) { - case "permission.requested": { - await upsertPermissionRequest(database, record); - break; - } - case "permission.resolved": { - await appPermissionResolution(database, record); - break; - } - case "run.cancelled": - case "run.completed": - case "run.failed": { - await clearRunPermissionRequests(database, record); - break; - } - case "session.readiness.updated": { - await upsertReadinessSnapshot(database, record); - break; - } - default: { - break; - } + record: SessionViewerProjectionRuntimeEvent, +): D1PreparedStatement[] { + const runId = requirePermissionRunId(record); + const run = readRuntimeRunPayload(record.event).run; + const startedAt = + run?.startedAt === null || run?.startedAt === undefined + ? Number.NaN + : Date.parse(run.startedAt); + if (run === null || run.id !== runId || run.status !== "running" || !Number.isFinite(startedAt)) { + throw new Error("Runtime run.started projection requires one exact running Run view."); + } + + const update = database + .prepare( + `UPDATE session_run + SET error_code = NULL, + error_details_json = NULL, + error_message = NULL, + error_retryable = NULL, + started_at = COALESCE(started_at, ?), + status = CASE WHEN status IN ('queued', 'booting') THEN 'running' ELSE status END, + status_changed_at = CASE + WHEN status IN ('queued', 'booting') THEN ? ELSE status_changed_at END, + status_event = CASE + WHEN status IN ('queued', 'booting') THEN 'run.start' ELSE status_event END, + status_operation_id = CASE + WHEN status IN ('queued', 'booting') THEN NULL ELSE status_operation_id END, + status_seq = status_seq + CASE WHEN status IN ('queued', 'booting') THEN 1 ELSE 0 END, + status_source = CASE + WHEN status IN ('queued', 'booting') THEN ? ELSE status_source END, + updated_at = CASE + WHEN status IN ('queued', 'booting') THEN ? ELSE updated_at END + WHERE id = ? + AND session_id = ? + AND status IN ('queued', 'booting', 'running', 'waiting_input') + AND ${INSERTED_EVENT_FENCE}`, + ) + .bind( + startedAt, + record.createdAt, + record.event.origin, + record.createdAt, + runId, + record.sessionId, + record.eventId, + record.sessionId, + record.event.kind, + ); + const guard = database + .prepare( + `INSERT INTO session_event (id) + SELECT ? + WHERE NOT EXISTS ( + SELECT 1 + FROM session_event AS receipt + INNER JOIN session_run AS run + ON run.id = ? AND run.session_id = receipt.session_id + WHERE receipt.id = ? + AND receipt.session_id = ? + AND receipt.event_type = 'run.started' + AND run.status IN ('running', 'waiting_input') + AND run.started_at IS NOT NULL + )`, + ) + .bind(createPlatformId(), runId, record.eventId, record.sessionId); + + return [update, guard]; +} + +export function prepareSessionViewerRuntimeEventProjection( + database: D1Database, + record: SessionViewerProjectionRuntimeEvent, +): D1PreparedStatement[] { + switch (record.event.kind) { + case "permission.requested": { + return preparePermissionRequested(database, record); + } + case "permission.resolved": { + return preparePermissionResolved(database, record); + } + case "run.started": { + return prepareRunStarted(database, record); + } + case "session.readiness.updated": { + return [prepareReadinessUpsert(database, record)]; + } + default: { + return []; } } } diff --git a/apps/api/src/modules/sessions/infrastructure/session-viewer-live-snapshot.repository.ts b/apps/api/src/modules/sessions/infrastructure/session-viewer-live-snapshot.repository.ts index 9b35e3a7..fbb721ce 100644 --- a/apps/api/src/modules/sessions/infrastructure/session-viewer-live-snapshot.repository.ts +++ b/apps/api/src/modules/sessions/infrastructure/session-viewer-live-snapshot.repository.ts @@ -18,23 +18,27 @@ import { parseSchemaValue, } from "@mosoo/contracts/validation"; import { + sessionAgentTaskSnapshotsTable, sessionPermissionRequestsTable, sessionReadinessSnapshotsTable, sessionRunsTable, sessionsTable, } from "@mosoo/db"; import type { AgentDeploymentVersionId, PlatformId, SessionId, SessionRunId } from "@mosoo/id"; -import { asc, eq } from "drizzle-orm"; +import { and, asc, eq, isNull } from "drizzle-orm"; import { getAppDatabase } from "../../../platform/db/drizzle"; import { isTruthy } from "../../../shared/truthiness"; import { toIsoString } from "../../../time"; import { fileStore } from "../../files/application/file-store"; +import { parseStoredAgentTaskSnapshot } from "./session-agent-task-snapshot.repository"; import { createInitialSessionLiveState } from "./session-live-state.reducer"; +import { applyStoredSessionArtifacts } from "./session-message-reference.repository"; import { loadStoredSessionMessages } from "./session-message-snapshot.repository"; interface SessionViewerStateSessionRow { id: SessionId; + runtime_event_seq_cursor: number; status: SessionStatus; title: string | null; updated_at: number; @@ -45,9 +49,11 @@ interface SessionViewerStateJoinedRow extends SessionViewerStateSessionRow { run_created_at: number | null; run_deployment_version_id: AgentDeploymentVersionId | null; run_deployment_version_number: number | null; + run_driver_instance_id: string | null; run_error_code: string | null; run_error_details_json: string | null; run_error_message: string | null; + run_error_retryable: boolean | null; run_id: SessionRunId | null; run_model: string | null; run_provider: string | null; @@ -56,6 +62,9 @@ interface SessionViewerStateJoinedRow extends SessionViewerStateSessionRow { run_trace_id: string | null; run_trigger: SessionRunTrigger | null; run_updated_at: number | null; + task_driver_instance_id: string | null; + task_run_id: string | null; + task_tasks_json: string | null; } interface SessionViewerStateSnapshotRow { @@ -70,6 +79,7 @@ interface SessionViewerStateRunRow { error_code: string | null; error_details_json: string | null; error_message: string | null; + error_retryable: boolean | null; id: SessionRunId; model: string | null; provider: string | null; @@ -96,6 +106,13 @@ export interface LoadSessionViewerStateInput { viewerId: PlatformId; } +export interface LoadedSessionViewerState { + runtimeEventSeqCursor: number; + state: SessionLiveState; +} + +const MAX_CONSISTENT_SNAPSHOT_ATTEMPTS = 5; + async function listSessionViewerStateSnapshotRows( database: D1Database, sessionId: SessionId, @@ -108,9 +125,11 @@ async function listSessionViewerStateSnapshotRows( run_created_at: sessionRunsTable.createdAt, run_deployment_version_id: sessionRunsTable.deploymentVersionId, run_deployment_version_number: sessionRunsTable.deploymentVersionNumber, + run_driver_instance_id: sessionRunsTable.driverInstanceId, run_error_code: sessionRunsTable.errorCode, run_error_details_json: sessionRunsTable.errorDetailsJson, run_error_message: sessionRunsTable.errorMessage, + run_error_retryable: sessionRunsTable.errorRetryable, run_id: sessionRunsTable.id, run_model: sessionRunsTable.model, run_provider: sessionRunsTable.provider, @@ -119,13 +138,26 @@ async function listSessionViewerStateSnapshotRows( run_trace_id: sessionRunsTable.traceId, run_trigger: sessionRunsTable.trigger, run_updated_at: sessionRunsTable.updatedAt, + runtime_event_seq_cursor: sessionsTable.runtimeEventSeqCursor, status: sessionsTable.status, + task_driver_instance_id: sessionAgentTaskSnapshotsTable.driverInstanceId, + task_run_id: sessionAgentTaskSnapshotsTable.runId, + task_tasks_json: sessionAgentTaskSnapshotsTable.tasksJson, title: sessionsTable.title, updated_at: sessionsTable.updatedAt, }, }) .from(sessionsTable) .leftJoin(sessionRunsTable, eq(sessionRunsTable.id, sessionsTable.lastRunId)) + .leftJoin( + sessionAgentTaskSnapshotsTable, + and( + eq(sessionAgentTaskSnapshotsTable.sessionId, sessionsTable.id), + eq(sessionAgentTaskSnapshotsTable.runId, sessionRunsTable.id), + eq(sessionAgentTaskSnapshotsTable.driverInstanceId, sessionRunsTable.driverInstanceId), + isNull(sessionsTable.archivedAt), + ), + ) .where(eq(sessionsTable.id, sessionId)) .limit(1) .all(); @@ -226,7 +258,7 @@ function toRunError(row: SessionViewerStateRunRow): RunError | null { code: row.error_code, details: parseJsonRecord(row.error_details_json), message: row.error_message, - retryable: false, + retryable: row.error_retryable ?? false, }; } @@ -269,6 +301,7 @@ function toJoinedSessionRunSummary(row: SessionViewerStateJoinedRow): SessionRun error_code: row.run_error_code, error_details_json: row.run_error_details_json, error_message: row.run_error_message, + error_retryable: row.run_error_retryable, id: row.run_id, model: row.run_model, provider: row.run_provider, @@ -313,6 +346,28 @@ function isTerminalRunStatus(status: SessionRunView["status"]): boolean { ); } +function toJoinedAgentTaskSnapshot( + row: SessionViewerStateJoinedRow, +): SessionLiveState["taskSnapshot"] { + if ( + row.status !== "RUNNING" || + row.run_status === null || + isTerminalRunStatus(row.run_status) || + row.task_driver_instance_id === null || + row.task_run_id === null || + row.task_tasks_json === null + ) { + return null; + } + + return parseStoredAgentTaskSnapshot({ + driverInstanceId: row.task_driver_instance_id, + runId: row.task_run_id, + sessionId: row.id, + tasksJson: row.task_tasks_json, + }); +} + function toCanonicalLifecycleStatus( sessionStatus: SessionStatus, runStatus: SessionRunView["status"], @@ -334,6 +389,7 @@ function applyCanonicalSessionState( input: { files: SessionViewFile[]; latestRun: SessionRunSummary | null; + runDriverInstanceId: string | null; session: SessionViewerStateSessionRow; viewerId: PlatformId; }, @@ -347,6 +403,13 @@ function applyCanonicalSessionState( return { ...state, files: input.files, + infra: { + ...state.infra, + driverInstanceId: + input.session.status === "RUNNING" && !isTerminalRunStatus(run.status) + ? input.runDriverInstanceId + : null, + }, lifecycle: toCanonicalLifecycleStatus(input.session.status, run.status), permissionRequests, run, @@ -357,24 +420,24 @@ function applyCanonicalSessionState( }; } -export async function loadSessionViewerState( +async function loadSessionViewerStateSnapshotOnce( database: D1Database, input: LoadSessionViewerStateInput, -): Promise { - const snapshotRowsPromise = listSessionViewerStateSnapshotRows(database, input.sessionId); +): Promise { + const snapshotRows = await listSessionViewerStateSnapshotRows(database, input.sessionId); + const session = getFirstSnapshotRow(snapshotRows).session; const messagesPromise = loadStoredSessionMessages(database, input.sessionId); const permissionRequestsPromise = listActivePermissionRequests(database, input.sessionId); const readinessPromise = getLatestReadinessSnapshot(database, input.sessionId); const sessionFilesPromise = fileStore.listReadySessionFiles(database, input.sessionId); - const [snapshotRows, messages, permissionRequests, readiness, sessionFiles] = await Promise.all([ - snapshotRowsPromise, + const [messages, permissionRequests, readiness, sessionFiles] = await Promise.all([ messagesPromise, permissionRequestsPromise, readinessPromise, sessionFilesPromise, ]); - const session = getFirstSnapshotRow(snapshotRows).session; const latestRun = toJoinedSessionRunSummary(session); + const taskSnapshot = toJoinedAgentTaskSnapshot(session); const baseState = createInitialSessionLiveState({ sessionId: input.sessionId, title: session.title, @@ -385,13 +448,60 @@ export async function loadSessionViewerState( messages, permissionRequests, readiness, + taskSnapshot, }; const state = applyCanonicalSessionState(stateWithMessages, { files: sessionFiles, latestRun, + runDriverInstanceId: session.run_driver_instance_id, session, viewerId: input.viewerId, }); - return state; + const resolvedState = + session.runtime_event_seq_cursor === 0 + ? state + : await applyStoredSessionArtifacts(database, { + endSeq: session.runtime_event_seq_cursor, + includeActiveRunArtifacts: latestRun !== null && !isTerminalRunStatus(latestRun.status), + runId: latestRun?.id ?? null, + sessionId: input.sessionId, + state, + }); + + return { + runtimeEventSeqCursor: session.runtime_event_seq_cursor, + state: resolvedState, + }; +} + +export async function loadSessionViewerStateSnapshot( + database: D1Database, + input: LoadSessionViewerStateInput, +): Promise { + for (let attempt = 0; attempt < MAX_CONSISTENT_SNAPSHOT_ATTEMPTS; attempt += 1) { + const snapshot = await loadSessionViewerStateSnapshotOnce(database, input); + const current = await getAppDatabase(database) + .select({ runtimeEventSeqCursor: sessionsTable.runtimeEventSeqCursor }) + .from(sessionsTable) + .where(eq(sessionsTable.id, input.sessionId)) + .limit(1) + .get(); + + if (current === undefined) { + throw new Error("Session not found."); + } + if (current.runtimeEventSeqCursor === snapshot.runtimeEventSeqCursor) { + return snapshot; + } + } + + throw new Error("Session changed while its viewer state snapshot was loading."); +} + +export async function loadSessionViewerState( + database: D1Database, + input: LoadSessionViewerStateInput, +): Promise { + return (await loadSessionViewerStateSnapshot(database, input)).state; } diff --git a/apps/api/src/modules/sessions/infrastructure/session/client.ts b/apps/api/src/modules/sessions/infrastructure/session/client.ts index b246e09b..bf369e24 100644 --- a/apps/api/src/modules/sessions/infrastructure/session/client.ts +++ b/apps/api/src/modules/sessions/infrastructure/session/client.ts @@ -15,6 +15,12 @@ function getSessionStub(env: ApiBindings, sessionId: SessionId): DurableObjectSt return binding.get(binding.idFromName(sessionId)); } +export interface SessionViewerEventBatch { + events: AgUiSessionEvent[]; + previousRuntimeEventSeqCursor: number | null; + runtimeEventSeqCursor: number | null; +} + function createSessionDoRequest(sessionId: SessionId, path: string, init?: RequestInit): Request { const headers = new Headers(init?.headers); headers.set(SESSION_ID_HEADER, sessionId); @@ -52,16 +58,19 @@ export async function connectSessionViewerWebSocket( ); } -export async function publishSessionViewerEvents( +export async function publishSessionViewerEventBatches( env: ApiBindings, sessionId: SessionId | null, - events: AgUiSessionEvent[], + batches: SessionViewerEventBatch[], ): Promise { - if (sessionId === null || sessionId === "" || events.length === 0) { + if ( + sessionId === null || + sessionId === "" || + batches.every((batch) => batch.events.length === 0) + ) { return; } - - await getSessionStub(env, sessionId).publishEvents(sessionId, events); + await getSessionStub(env, sessionId).publishEventBatches(sessionId, batches); } export async function connectSessionPublicEventWebSocket( diff --git a/apps/api/src/modules/sessions/infrastructure/session/do.ts b/apps/api/src/modules/sessions/infrastructure/session/do.ts index 96146dc7..800e0484 100644 --- a/apps/api/src/modules/sessions/infrastructure/session/do.ts +++ b/apps/api/src/modules/sessions/infrastructure/session/do.ts @@ -139,16 +139,41 @@ export class Session extends DurableObject { return normalizedSessionId; } - async publishEvents(sessionId: string, events: AgUiSessionEvent[]): Promise { + async publishEvents( + sessionId: string, + events: AgUiSessionEvent[], + runtimeEventSeqCursor: number | null = null, + previousRuntimeEventSeqCursor: number | null = null, + ): Promise { this.#ensureActiveRpcSession(sessionId); if (events.length > 0) { this.#publicEventSockets.notifyEventsAvailable(); } - await this.#viewerSockets.broadcastEvents(events); + await this.#viewerSockets.broadcastEvents( + events, + runtimeEventSeqCursor, + previousRuntimeEventSeqCursor, + ); + } + + async publishEventBatches( + sessionId: string, + batches: Array<{ + events: AgUiSessionEvent[]; + previousRuntimeEventSeqCursor: number | null; + runtimeEventSeqCursor: number | null; + }>, + ): Promise { + this.#ensureActiveRpcSession(sessionId); + if (batches.some((batch) => batch.events.length > 0)) { + this.#publicEventSockets.notifyEventsAvailable(); + } + await this.#viewerSockets.broadcastEventBatches(batches); } async syncViewers(sessionId: string): Promise { this.#ensureActiveRpcSession(sessionId); + this.#publicEventSockets.notifyEventsAvailable(); await this.#viewerSockets.broadcastStateSync(); } diff --git a/apps/api/src/modules/sessions/infrastructure/session/viewer-live-state.ts b/apps/api/src/modules/sessions/infrastructure/session/viewer-live-state.ts index b9a7d54a..3dede25c 100644 --- a/apps/api/src/modules/sessions/infrastructure/session/viewer-live-state.ts +++ b/apps/api/src/modules/sessions/infrastructure/session/viewer-live-state.ts @@ -2,6 +2,7 @@ import type { SessionId } from "@mosoo/id"; import type { AuthenticatedViewer } from "../../../auth/application/viewer-auth.service"; import { reconcileStaleActiveSessionRun } from "../../../runtime/application/session-runs/stale-run-reconciliation.service"; +import { loadSessionAgentTaskState } from "../session-agent-task-snapshot.repository"; import type { SessionLiveState } from "../session-live-state.types"; import { loadSessionViewerState } from "../session-viewer-live-snapshot.repository"; @@ -21,7 +22,20 @@ export async function loadViewerLiveState( (await reconcileStaleActiveSessionRun(input.database, input.sessionId)); if (input.cachedState && !reconciledStaleRun) { - return normalizeViewerLiveState(input.cachedState, input); + const taskState = await loadSessionAgentTaskState(input.database, input.sessionId); + const cacheGenerationIsCurrent = + taskState !== null && + input.cachedState.lifecycle === "RUNNING" && + taskState.runId === input.cachedState.run.id && + taskState.runStatus === input.cachedState.run.status && + taskState.driverInstanceId === input.cachedState.infra.driverInstanceId; + + if (cacheGenerationIsCurrent) { + return normalizeViewerLiveState( + { ...input.cachedState, taskSnapshot: taskState.snapshot }, + input, + ); + } } return loadSessionViewerState(input.database, { diff --git a/apps/api/src/modules/sessions/infrastructure/session/viewer-permissions.ts b/apps/api/src/modules/sessions/infrastructure/session/viewer-permissions.ts index adbc18ba..a753d1b3 100644 --- a/apps/api/src/modules/sessions/infrastructure/session/viewer-permissions.ts +++ b/apps/api/src/modules/sessions/infrastructure/session/viewer-permissions.ts @@ -34,7 +34,7 @@ export async function rejectDisconnectedViewerPermissionRequests( await appendSessionRuntimeEvents({ bindings: input.env, deliver: false, - events: [result.event], + events: result.events, sessionId: input.attachment.sessionId, }); diff --git a/apps/api/src/modules/sessions/infrastructure/session/viewer-socket-hub.ts b/apps/api/src/modules/sessions/infrastructure/session/viewer-socket-hub.ts index 1ca33a7f..e4bddbea 100644 --- a/apps/api/src/modules/sessions/infrastructure/session/viewer-socket-hub.ts +++ b/apps/api/src/modules/sessions/infrastructure/session/viewer-socket-hub.ts @@ -7,6 +7,7 @@ import { } from "../../../../platform/cloudflare/durable-object-support"; import { createErrorLogContext, logError, logInfo } from "../../../../platform/cloudflare/logger"; import type { ApiBindings } from "../../../../platform/cloudflare/worker-types"; +import { currentTimestampMs } from "../../../../time"; import type { SessionLiveState } from "../session-live-state.types"; import { json } from "./requests"; import { readSessionViewerSocketHeaders } from "./socket-headers"; @@ -28,6 +29,8 @@ import { declare const WebSocketPair: new () => [WebSocket, WebSocket]; +const VIEWER_CURSOR_RECONCILIATION_DELAY_MS = 10_000; + interface SessionViewerSocketHubOptions { ctx: DurableObjectState; env: ApiBindings; @@ -36,6 +39,12 @@ interface SessionViewerSocketHubOptions { withSessionLogContext: (fn: () => T) => T; } +export interface SessionViewerEventBatch { + events: AgUiSessionEvent[]; + previousRuntimeEventSeqCursor: number | null; + runtimeEventSeqCursor: number | null; +} + function getSocketAttachment(ws: WebSocket): SessionSocketAttachment | null { const attachment: unknown = ws.deserializeAttachment(); return isViewerSocketAttachment(attachment) ? attachment : null; @@ -47,6 +56,7 @@ export class SessionViewerSocketHub { readonly #getSessionId: () => string | null; #liveStateCache: SessionLiveState | null = null; readonly #rememberSessionId: (sessionId: string) => void; + #stateOperationTail: Promise = Promise.resolve(); readonly #withSessionLogContext: (fn: () => T) => T; constructor(options: SessionViewerSocketHubOptions) { @@ -57,55 +67,126 @@ export class SessionViewerSocketHub { this.#withSessionLogContext = options.withSessionLogContext; } - async broadcastEvents(events: AgUiSessionEvent[]): Promise { - if (events.length === 0) { - return; - } - - const broadcast = buildViewerBroadcastFrames({ - cachedState: this.#liveStateCache, - events, - }); + async broadcastEvents( + events: AgUiSessionEvent[], + runtimeEventSeqCursor: number | null = null, + previousRuntimeEventSeqCursor: number | null = null, + ): Promise { + await this.broadcastEventBatches([ + { events, previousRuntimeEventSeqCursor, runtimeEventSeqCursor }, + ]); + } - if (!broadcast) { + async broadcastEventBatches(batches: SessionViewerEventBatch[]): Promise { + if (batches.every((batch) => batch.events.length === 0)) { return; } - if (broadcast.state) { - this.#liveStateCache = broadcast.state; - } - - for (const socket of this.#getViewerSockets()) { - const attachment = getSocketAttachment(socket); + await this.#runStateOperation(async () => { + const sockets = this.#getViewerSockets().flatMap((socket) => { + const attachment = getSocketAttachment(socket); + return attachment === null || socket.readyState !== WebSocket.OPEN + ? [] + : [{ attachment, frames: [] as string[], requiresStateSync: false, socket }]; + }); - if (!attachment || socket.readyState !== WebSocket.OPEN) { - continue; + for (const batch of batches) { + const broadcast = buildViewerBroadcastFrames({ + cachedState: batch.runtimeEventSeqCursor === null ? this.#liveStateCache : null, + events: batch.events, + }); + if (broadcast === null) { + continue; + } + if (batch.runtimeEventSeqCursor !== null) { + this.#liveStateCache = null; + } else if (broadcast.state !== null) { + this.#liveStateCache = broadcast.state; + } + for (const target of sockets) { + if (target.requiresStateSync) { + continue; + } + if (target.attachment.runtimeEventSeqCursor === undefined) { + target.frames = []; + target.requiresStateSync = true; + continue; + } + if ( + batch.runtimeEventSeqCursor !== null && + target.attachment.runtimeEventSeqCursor >= batch.runtimeEventSeqCursor + ) { + continue; + } + if ( + batch.runtimeEventSeqCursor !== null && + (batch.previousRuntimeEventSeqCursor === null || + target.attachment.runtimeEventSeqCursor !== batch.previousRuntimeEventSeqCursor) + ) { + target.frames = []; + target.requiresStateSync = true; + continue; + } + target.frames.push(...broadcast.frames); + if (batch.runtimeEventSeqCursor !== null) { + target.attachment = { + ...target.attachment, + runtimeEventSeqCursor: batch.runtimeEventSeqCursor, + }; + } + } } - sendFrames(socket, broadcast.frames); - } + for (const target of sockets) { + if (target.requiresStateSync) { + continue; + } + if (target.frames.length === 0) { + continue; + } + try { + sendFrames(target.socket, target.frames); + target.socket.serializeAttachment(target.attachment); + } catch { + closeOpenSocket(target.socket, 1011, "session.viewer.event-delivery-failed"); + } + } + const stateSyncTargets = sockets + .filter((target) => target.requiresStateSync) + .map(({ attachment, socket }) => ({ attachment, socket })); + if (stateSyncTargets.length > 0) { + await sendViewerSocketStateSyncBatch({ + database: this.#env.DB, + sockets: stateSyncTargets, + updateLiveStateCache: (state) => { + this.#rememberLoadedLiveState(state); + }, + }); + } + }); } async broadcastStateSync(): Promise { - const sockets = this.#getViewerSockets() - .map((socket) => ({ attachment: getSocketAttachment(socket), socket })) - .filter( - ( - candidate, - ): candidate is { - attachment: ViewerSocketAttachment; - socket: WebSocket; - } => candidate.attachment !== null, - ); - - await sendViewerSocketStateSyncBatch({ - cachedState: this.#liveStateCache, - database: this.#env.DB, - getLatestCachedState: () => this.#liveStateCache, - sockets, - updateLiveStateCache: (state) => { - this.#rememberLoadedLiveState(state); - }, + await this.#runStateOperation(async () => { + this.#liveStateCache = null; + const sockets = this.#getViewerSockets() + .map((socket) => ({ attachment: getSocketAttachment(socket), socket })) + .filter( + ( + candidate, + ): candidate is { + attachment: ViewerSocketAttachment; + socket: WebSocket; + } => candidate.attachment !== null, + ); + + await sendViewerSocketStateSyncBatch({ + database: this.#env.DB, + sockets, + updateLiveStateCache: (state) => { + this.#rememberLoadedLiveState(state); + }, + }); }); } @@ -132,7 +213,11 @@ export class SessionViewerSocketHub { this.#rememberSessionId(attachment.sessionId); this.#ctx.acceptWebSocket(server, ["viewer"]); server.serializeAttachment(attachment); - this.#ctx.waitUntil(clearViewerPermissionCleanupAlarm({ storage: this.#ctx.storage })); + this.#ctx.waitUntil( + clearViewerPermissionCleanupAlarm({ storage: this.#ctx.storage }).then(() => + this.#scheduleViewerCursorReconciliation(), + ), + ); this.#ctx.waitUntil(this.#sendViewerStateSync(server, attachment)); this.#withSessionLogContext(() => { @@ -213,30 +298,49 @@ export class SessionViewerSocketHub { } async handleAlarm(): Promise { - await runViewerPermissionCleanupAlarm({ - cachedState: this.#liveStateCache, - env: this.#env, - hasOpenViewer: (sessionId) => this.#hasOpenViewer(sessionId), - storage: this.#ctx.storage, - updateLiveStateCache: (state) => { - this.#rememberLoadedLiveState(state); - }, + await this.#runStateOperation(async () => { + try { + await runViewerPermissionCleanupAlarm({ + cachedState: this.#liveStateCache, + env: this.#env, + hasOpenViewer: (sessionId) => this.#hasOpenViewer(sessionId), + storage: this.#ctx.storage, + updateLiveStateCache: (state) => { + this.#rememberLoadedLiveState(state); + }, + }); + } finally { + try { + await this.#reconcileViewerCursors(); + } finally { + await this.#scheduleViewerCursorReconciliation(); + } + } }); } async #sendViewerStateSync(ws: WebSocket, attachment: ViewerSocketAttachment): Promise { - await sendViewerSocketStateSync({ - attachment, - cachedState: this.#liveStateCache, - database: this.#env.DB, - getLatestCachedState: () => this.#liveStateCache, - updateLiveStateCache: (state) => { - this.#rememberLoadedLiveState(state); - }, - ws, + await this.#runStateOperation(async () => { + await sendViewerSocketStateSync({ + attachment, + database: this.#env.DB, + updateLiveStateCache: (state) => { + this.#rememberLoadedLiveState(state); + }, + ws, + }); }); } + #runStateOperation(operation: () => Promise): Promise { + const result = this.#stateOperationTail.then(operation); + this.#stateOperationTail = result.then( + () => undefined, + () => undefined, + ); + return result; + } + #rememberLoadedLiveState(state: SessionLiveState | null): void { if (!state) { return; @@ -256,10 +360,87 @@ export class SessionViewerSocketHub { }); } + async #reconcileViewerCursors(): Promise { + const sessionId = this.#resolveSessionId(); + if (sessionId === null) { + return; + } + const durable = await this.#env.DB.prepare( + "SELECT runtime_event_seq_cursor FROM session WHERE id = ?", + ) + .bind(sessionId) + .first<{ runtime_event_seq_cursor: number }>(); + const sockets = this.#getViewerSockets().flatMap((socket) => { + const attachment = getSocketAttachment(socket); + return attachment !== null && + durable !== null && + (attachment.runtimeEventSeqCursor === undefined || + attachment.runtimeEventSeqCursor !== durable.runtime_event_seq_cursor) + ? [{ attachment, socket }] + : []; + }); + if (sockets.length === 0) { + return; + } + + await sendViewerSocketStateSyncBatch({ + database: this.#env.DB, + sockets, + updateLiveStateCache: (state) => { + this.#rememberLoadedLiveState(state); + }, + }); + } + + async #scheduleViewerCursorReconciliation(): Promise { + const sessionId = this.#resolveSessionId(); + if (sessionId === null || !this.#hasOpenViewer(sessionId)) { + return; + } + await this.#ctx.storage.setAlarm(currentTimestampMs() + VIEWER_CURSOR_RECONCILIATION_DELAY_MS); + } + + #resolveSessionId(): string | null { + const rememberedSessionId = this.#getSessionId(); + const attachedSessionIds = new Set(); + + for (const socket of this.#getViewerSockets()) { + const attachment = getSocketAttachment(socket); + if (attachment !== null) { + attachedSessionIds.add(attachment.sessionId); + } + } + + if (rememberedSessionId !== null) { + for (const socket of this.#getViewerSockets()) { + const attachment = getSocketAttachment(socket); + if (attachment !== null && attachment.sessionId !== rememberedSessionId) { + closeOpenSocket(socket, 1008, "session.viewer.session-mismatch"); + } + } + return rememberedSessionId; + } + + if (attachedSessionIds.size !== 1) { + if (attachedSessionIds.size > 1) { + this.closeSockets("session.viewer.session-mismatch"); + } + return null; + } + + const [sessionId] = attachedSessionIds; + if (sessionId === undefined) { + return null; + } + this.#rememberSessionId(sessionId); + return sessionId; + } + async #scheduleViewerPermissionsOnLastDisconnect( attachment: ViewerSocketAttachment, ): Promise { if (this.#hasOpenViewer(attachment.sessionId)) { + await this.#scheduleViewerCursorReconciliation(); return; } diff --git a/apps/api/src/modules/sessions/infrastructure/session/viewer-socket-state-sync.ts b/apps/api/src/modules/sessions/infrastructure/session/viewer-socket-state-sync.ts index 89edba6e..a70ba7d5 100644 --- a/apps/api/src/modules/sessions/infrastructure/session/viewer-socket-state-sync.ts +++ b/apps/api/src/modules/sessions/infrastructure/session/viewer-socket-state-sync.ts @@ -8,14 +8,12 @@ import { import { reconcileStaleActiveSessionRun } from "../../../runtime/application/session-runs/stale-run-reconciliation.service"; import { getActiveAppSessionParticipantAccess } from "../../domain/session-access.policy"; import type { SessionLiveState } from "../session-live-state.types"; -import { loadViewerLiveState } from "./viewer-live-state"; +import { loadSessionViewerStateSnapshot } from "../session-viewer-live-snapshot.repository"; import type { ViewerSocketAttachment } from "./viewer-socket"; interface SendViewerSocketStateSyncOptions { attachment: ViewerSocketAttachment; - cachedState: SessionLiveState | null; database: D1Database; - getLatestCachedState(): SessionLiveState | null; updateLiveStateCache(state: SessionLiveState | null): void; ws: WebSocket; } @@ -26,9 +24,7 @@ interface ViewerSocketStateSyncTarget { } interface SendViewerSocketStateSyncBatchOptions { - cachedState: SessionLiveState | null; database: D1Database; - getLatestCachedState(): SessionLiveState | null; sockets: ViewerSocketStateSyncTarget[]; updateLiveStateCache(state: SessionLiveState | null): void; } @@ -56,9 +52,7 @@ export async function sendViewerSocketStateSync( options: SendViewerSocketStateSyncOptions, ): Promise { await sendViewerSocketStateSyncBatch({ - cachedState: options.cachedState, database: options.database, - getLatestCachedState: options.getLatestCachedState, sockets: [ { attachment: options.attachment, @@ -73,7 +67,6 @@ export async function sendViewerSocketStateSyncBatch( options: SendViewerSocketStateSyncBatchOptions, ): Promise { const groupedTargets = groupOpenViewerSockets(options.sockets); - let cachedState = options.cachedState; const reconciledStaleRunsBySessionId = new Map(); for (const targets of groupedTargets.values()) { @@ -97,28 +90,30 @@ export async function sendViewerSocketStateSyncBatch( throw error; } - cachedState ??= options.getLatestCachedState(); - const reconciledStaleRun = await getReconciledStaleRun( + await getReconciledStaleRun( options.database, firstTarget.attachment.sessionId, reconciledStaleRunsBySessionId, ); - const state = await loadViewerLiveState({ - cachedState, - database: options.database, - reconciledStaleRun, + const snapshot = await loadSessionViewerStateSnapshot(options.database, { sessionId: firstTarget.attachment.sessionId, - viewer: firstTarget.attachment.viewer, + viewerId: firstTarget.attachment.viewer.id, }); - const latestState = options.getLatestCachedState(); - const stateToSend = latestState !== null && latestState !== cachedState ? latestState : state; + const stateToSend = snapshot.state; const frames = createStateSyncFrames(stateToSend); options.updateLiveStateCache(stateToSend); - cachedState = stateToSend; for (const target of targets) { - sendFrames(target.socket, frames); + try { + sendFrames(target.socket, frames); + target.socket.serializeAttachment({ + ...target.attachment, + runtimeEventSeqCursor: snapshot.runtimeEventSeqCursor, + } satisfies ViewerSocketAttachment); + } catch { + closeOpenSocket(target.socket, 1011, "session.viewer.state-sync-failed"); + } } } } diff --git a/apps/api/src/modules/sessions/infrastructure/session/viewer-socket.ts b/apps/api/src/modules/sessions/infrastructure/session/viewer-socket.ts index 6646a904..e1a73ba4 100644 --- a/apps/api/src/modules/sessions/infrastructure/session/viewer-socket.ts +++ b/apps/api/src/modules/sessions/infrastructure/session/viewer-socket.ts @@ -2,6 +2,7 @@ import type { SessionViewerSocketContext } from "./socket-headers"; export interface ViewerSocketAttachment extends SessionViewerSocketContext { role: "viewer"; + runtimeEventSeqCursor?: number; } export type SessionSocketAttachment = ViewerSocketAttachment; @@ -21,6 +22,10 @@ export function isViewerSocketAttachment(value: unknown): value is ViewerSocketA typeof value["publicOrigin"] === "string" && typeof value["appId"] === "string" && typeof value["sessionId"] === "string" && + (value["runtimeEventSeqCursor"] === undefined || + (typeof value["runtimeEventSeqCursor"] === "number" && + Number.isSafeInteger(value["runtimeEventSeqCursor"]) && + value["runtimeEventSeqCursor"] >= 0)) && isRecord(viewer) && typeof viewer["email"] === "string" && typeof viewer["emailVerified"] === "boolean" && diff --git a/apps/api/src/modules/skills/application/skill-package.shared.ts b/apps/api/src/modules/skills/application/skill-package.shared.ts index 6ef518ca..24edd355 100644 --- a/apps/api/src/modules/skills/application/skill-package.shared.ts +++ b/apps/api/src/modules/skills/application/skill-package.shared.ts @@ -35,7 +35,7 @@ export async function sha256Hex(bytes: Uint8Array): Promise { ? bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength) : Uint8Array.from(bytes).buffer; const digest = await crypto.subtle.digest("SHA-256", buffer); - return [...new Uint8Array(digest)].map((byte) => byte.toString(16).padStart(2, "0")).join(""); + return new Uint8Array(digest).toHex(); } export function inferMimeType(path: string): string | null { diff --git a/apps/api/src/modules/users/application/viewer-context.service.ts b/apps/api/src/modules/users/application/viewer-context.service.ts index 85f04c43..adb40c1c 100644 --- a/apps/api/src/modules/users/application/viewer-context.service.ts +++ b/apps/api/src/modules/users/application/viewer-context.service.ts @@ -111,19 +111,8 @@ async function getViewerAccountState( function getViewerAuth(bindings: ApiBindings, viewer: AuthenticatedViewer | null): ViewerAuth { const methods: AuthMethod[] = ["email_otp"]; - const authBindings = bindings as ApiBindings & { - GOOGLE_OAUTH_CLIENT_ID?: string; - GOOGLE_OAUTH_CLIENT_SECRET?: string; - }; - if ( - authBindings.GOOGLE_OAUTH_CLIENT_ID?.trim() !== null && - authBindings.GOOGLE_OAUTH_CLIENT_ID?.trim() !== undefined && - authBindings.GOOGLE_OAUTH_CLIENT_ID?.trim() !== "" && - authBindings.GOOGLE_OAUTH_CLIENT_SECRET?.trim() !== null && - authBindings.GOOGLE_OAUTH_CLIENT_SECRET?.trim() !== undefined && - authBindings.GOOGLE_OAUTH_CLIENT_SECRET?.trim() !== "" - ) { + if (bindings.GOOGLE_OAUTH_CLIENT_ID?.trim() && bindings.GOOGLE_OAUTH_CLIENT_SECRET?.trim()) { methods.push("google_oauth"); } diff --git a/apps/api/src/platform/cloudflare/app-deployment-workflow.ts b/apps/api/src/platform/cloudflare/app-deployment-workflow.ts new file mode 100644 index 00000000..9634a17c --- /dev/null +++ b/apps/api/src/platform/cloudflare/app-deployment-workflow.ts @@ -0,0 +1,15 @@ +import { WorkflowEntrypoint } from "cloudflare:workers"; +import type { WorkflowEvent, WorkflowStep } from "cloudflare:workers"; + +import type { ApiCommandMessage } from "../../modules/api-command/application/api-command-message"; +import { runAppDeploymentWorkflow } from "../../modules/api-command/application/api-command-workflow"; +import type { ApiBindings } from "./worker-types"; + +export class AppDeploymentWorkflow extends WorkflowEntrypoint { + override async run( + event: Readonly>, + step: WorkflowStep, + ): Promise { + await runAppDeploymentWorkflow(this.env, event.payload, step); + } +} diff --git a/apps/api/src/platform/cloudflare/create-api-worker.ts b/apps/api/src/platform/cloudflare/create-api-worker.ts index 0d4ad3ed..8ad5721a 100644 --- a/apps/api/src/platform/cloudflare/create-api-worker.ts +++ b/apps/api/src/platform/cloudflare/create-api-worker.ts @@ -1,12 +1,20 @@ import { MOSOO_CONSOLE_HOST, MOSOO_LEGACY_CONSOLE_HOST } from "@mosoo/contracts/origin"; import { enqueueScheduledMaintenanceCommand } from "../../modules/api-command/application/api-command-enqueue"; -import { redriveFailedApiCommandEnqueues } from "../../modules/api-command/application/api-command-ledger"; +import { + reconcileAppDeploymentWorkflowDeliveries, + redriveFailedApiCommandEnqueues, +} from "../../modules/api-command/application/api-command-ledger"; import type { ApiCommandMessage } from "../../modules/api-command/application/api-command-message"; import { processApiCommandDeadLetterMessage, processApiCommandMessage, } from "../../modules/api-command/application/api-command-processor"; +import { + dispatchAppDeploymentGatewayRequest, + resolveActiveAppDeploymentScriptName, +} from "../../modules/apps/application/app-deployment-gateway"; +import { createErrorLogContext, logError } from "./logger"; import type { ApiBindings } from "./worker-types"; interface ApiHttpApp { @@ -36,6 +44,17 @@ export function createApiWorker(): ExportedHandler { return Response.redirect(url.toString(), 308); } + const gatewayResponse = await dispatchAppDeploymentGatewayRequest(request, { + appDeploymentDomain: env.MOSOO_APP_DEPLOYMENT_DOMAIN, + dispatcher: env.APP_DEPLOYMENT_DISPATCHER, + resolveActiveScriptName: (mosooSubdomain) => + resolveActiveAppDeploymentScriptName(env.DB, mosooSubdomain), + }); + + if (gatewayResponse !== null) { + return gatewayResponse; + } + const app = await getHttpApp(); const response = await app.fetch(request, env, ctx); @@ -43,6 +62,14 @@ export function createApiWorker(): ExportedHandler { }, async scheduled(controller: ScheduledController, env: ApiBindings): Promise { await redriveFailedApiCommandEnqueues(env); + try { + await reconcileAppDeploymentWorkflowDeliveries(env, controller.scheduledTime); + } catch (error) { + logError("api-command.app_deployment_workflow_reconciliation_failed", { + ...createErrorLogContext(error), + scheduledTime: controller.scheduledTime, + }); + } await enqueueScheduledMaintenanceCommand(env, { scheduledTime: controller.scheduledTime, }); diff --git a/apps/api/src/platform/cloudflare/worker-types.ts b/apps/api/src/platform/cloudflare/worker-types.ts index e429ce75..3bd5b5c8 100644 --- a/apps/api/src/platform/cloudflare/worker-types.ts +++ b/apps/api/src/platform/cloudflare/worker-types.ts @@ -19,6 +19,10 @@ interface ApiCommandQueueBinding { API_COMMAND_QUEUE: Queue; } +interface AppDeploymentWorkflowBinding { + APP_DEPLOYMENT_WORKFLOW: Workflow; +} + interface OptionalLocalProviderFetchProxyBindings { MOSOO_PROVIDER_FETCH_PROXY_TOKEN?: string; MOSOO_PROVIDER_FETCH_PROXY_URL?: string; @@ -58,6 +62,7 @@ interface OptionalRuntimeSubjectPlatformBindings { } export type ApiBindings = Env & + AppDeploymentWorkflowBinding & ApiCommandQueueBinding & OptionalSandboxBinding & OptionalDriverConnectionBinding & diff --git a/apps/api/src/platform/db/drizzle.ts b/apps/api/src/platform/db/drizzle.ts index e435c746..11530724 100644 --- a/apps/api/src/platform/db/drizzle.ts +++ b/apps/api/src/platform/db/drizzle.ts @@ -33,6 +33,7 @@ import { sandboxesTable, sandboxSessionsTable, sessionExecutionSnapshotsTable, + sessionAgentTaskSnapshotsTable, sessionEventsTable, sessionMessagesTable, sessionModelCallsTable, @@ -88,6 +89,7 @@ const schema = { sandboxesTable, sandboxSessionsTable, sessionExecutionSnapshotsTable, + sessionAgentTaskSnapshotsTable, sessionEventsTable, sessionMessagesTable, sessionModelCallsTable, @@ -142,10 +144,10 @@ async function readAllRows(statement: D1PreparedStatement): Promise const row = await statement.first(); return row === null ? result - : ({ + : { ...result, results: [row], - } as D1Result); + }; } catch (error) { if (!isCompatMethodNotImplementedError(error) || typeof statement.first !== "function") { throw error; diff --git a/apps/api/src/shared/bytes.ts b/apps/api/src/shared/bytes.ts index c1c3f331..16c257f0 100644 --- a/apps/api/src/shared/bytes.ts +++ b/apps/api/src/shared/bytes.ts @@ -15,39 +15,18 @@ export function toArrayBuffer(bytes: Uint8Array): ArrayBuffer { } export function toBase64(bytes: Uint8Array): string { - let binary = ""; - - for (const byte of bytes) { - binary += String.fromCodePoint(byte); - } - - return btoa(binary); + return bytes.toBase64(); } -export function fromBase64(value: string): Uint8Array { - const binary = atob(value); - const bytes = new Uint8Array(binary.length); - - for (let index = 0; index < binary.length; index += 1) { - const byte = binary.codePointAt(index); - - if (byte === undefined) { - throw new Error("Base64 payload is invalid."); - } - - bytes[index] = byte; - } - - return bytes; +export function fromBase64(value: string): Uint8Array { + return Uint8Array.fromBase64(value); } export function toBase64Url(value: Uint8Array | string): string { const bytes = typeof value === "string" ? new TextEncoder().encode(value) : value; - return toBase64(bytes).replaceAll("+", "-").replaceAll("/", "_").replace(/=+$/u, ""); + return bytes.toBase64({ alphabet: "base64url", omitPadding: true }); } -export function fromBase64Url(value: string): Uint8Array { - const normalized = value.replaceAll("-", "+").replaceAll("_", "/"); - const padded = normalized.padEnd(normalized.length + ((4 - (normalized.length % 4)) % 4), "="); - return fromBase64(padded); +export function fromBase64Url(value: string): Uint8Array { + return Uint8Array.fromBase64(value, { alphabet: "base64url" }); } diff --git a/apps/api/src/shared/shell.ts b/apps/api/src/shared/shell.ts new file mode 100644 index 00000000..9a9df617 --- /dev/null +++ b/apps/api/src/shared/shell.ts @@ -0,0 +1,3 @@ +export function quoteShellArg(value: string): string { + return `'${value.replaceAll("'", `'"'"'`)}'`; +} diff --git a/apps/api/tests/agent-package-file-import.test.ts b/apps/api/tests/agent-package-file-import.test.ts index 7ed70387..fb809e65 100644 --- a/apps/api/tests/agent-package-file-import.test.ts +++ b/apps/api/tests/agent-package-file-import.test.ts @@ -264,7 +264,7 @@ async function createFixture(input: { archiveBytes?: Uint8Array } = {}) { `); const bucket = new MemoryByteBucket(); const bindings = createPublicHttpTestBindings(database, { - fileBucket: bucket as unknown as R2Bucket, + fileBucket: bucket, }) as ApiBindings; const archiveBytes = input.archiveBytes ?? createAgentPackageArchiveBytes(createPackageFixture(input)); diff --git a/apps/api/tests/agent-runtime-events-projection.test.ts b/apps/api/tests/agent-runtime-events-projection.test.ts index d775313d..e351572d 100644 --- a/apps/api/tests/agent-runtime-events-projection.test.ts +++ b/apps/api/tests/agent-runtime-events-projection.test.ts @@ -1,7 +1,8 @@ import { describe, expect, test } from "bun:test"; +import { createMcpExecuteFailedEventIdentity } from "@mosoo/agent-driver/events"; +import type { DriverCommandId } from "@mosoo/id"; import { createRuntimeEvent } from "@mosoo/runtime-events"; -import type { RuntimeEventKind } from "@mosoo/runtime-events"; import { readPermissionRequestViews, @@ -109,13 +110,56 @@ const PROJECTION_CASES = [ ] as const; describe("agent runtime event projection", () => { + test("accepts only the content-addressed identity of a proven MCP failure", () => { + const commandId = "01J0000000000000000000000X" as DriverCommandId; + const failed = createMcpExecuteFailedEventIdentity({ + commandId, + rawInput: '{"issue":"A-1"}', + rawOutput: "provider rejected the request", + title: "createIssue", + toolCallId: "tool-A-1", + }); + const event = createRuntimeEvent({ + correlationId: commandId, + id: "failed-mcp-tool", + kind: "tool.call.updated", + occurredAt: "2026-08-31T00:00:00.000Z", + payload: failed.payload, + sessionId: "session-1", + sourceEventId: failed.sourceEventId, + }); + + expect( + createSessionRuntimeEventProjection(event, { provenMcpCommandId: commandId }), + ).toMatchObject({ + mcpCommandId: commandId, + toolOutputText: "provider rejected the request", + toolStatus: "failed", + }); + expect(() => + createSessionRuntimeEventProjection( + { ...event, sourceEventId: `${failed.sourceEventId}0` }, + { provenMcpCommandId: commandId }, + ), + ).toThrow("Proven MCP command does not match its terminal tool event"); + expect(() => + createSessionRuntimeEventProjection( + { + ...event, + payload: { ...failed.payload, rawOutput: "tampered provider failure" }, + }, + { provenMcpCommandId: commandId }, + ), + ).toThrow("Proven MCP command does not match its terminal tool event"); + }); + test("derives v2 runtime families from the runtime event contract", () => { for (const event of PROJECTION_CASES) { const projection = createSessionRuntimeEventProjection( createRuntimeEvent({ actor: "system", id: `${event.kind}-${event.family}`, - kind: event.kind as RuntimeEventKind, + kind: event.kind, occurredAt: "2026-05-26T00:00:00.000Z", origin: "system", payload: event.value, @@ -136,7 +180,7 @@ describe("agent runtime event projection", () => { } }); - test("apps completed and failed tool output as process-ready content", () => { + test("projects completed, failed, and cancelled tool snapshots", () => { const completed = createSessionRuntimeEventProjection( createRuntimeEvent({ actor: "driver", @@ -170,27 +214,88 @@ describe("agent runtime event projection", () => { sessionId: "session-1", }), ); + const cancelled = createSessionRuntimeEventProjection( + createRuntimeEvent({ + actor: "driver", + id: "tool-3", + kind: "tool.call.updated", + occurredAt: "2026-05-26T00:00:02.000Z", + origin: "driver", + payload: { + status: "cancelled", + title: "Shell", + toolCallId: "tool-3", + }, + sessionId: "session-1", + }), + ); expect(completed).toMatchObject({ - contentText: "Shell result: hello world", + contentText: "hello\nworld", processStatus: "available", processType: "tool.use.completed", toolCallId: "tool-1", - toolInputJson: '{"command":"echo hello","timeout":30}', - toolName: null, + toolInputDeltaJson: null, + toolInputJson: '{"timeout":30,"command":"echo hello"}', + toolName: "Shell", + toolOutputDeltaText: null, + toolOutputText: "hello\nworld", + toolStatus: "completed", }); expect(failed).toMatchObject({ - contentText: "Shell result: permission denied", + contentText: "permission denied", processStatus: "error", processType: "tool.use.completed", toolCallId: "tool-2", + toolInputDeltaJson: null, toolInputJson: null, - toolName: null, + toolName: "Shell", + toolOutputDeltaText: null, + toolOutputText: "permission denied", + toolStatus: "failed", + }); + expect(cancelled).toMatchObject({ + contentText: "", + processStatus: "available", + processType: "tool.use.completed", + toolCallId: "tool-3", + toolInputDeltaJson: null, + toolInputJson: null, + toolName: "Shell", + toolOutputDeltaText: null, + toolOutputText: null, + toolStatus: "cancelled", + }); + }); + + test("projects message failure as an error boundary without changing streamed text", () => { + const projection = createSessionRuntimeEventProjection( + createRuntimeEvent({ + actor: "driver", + id: "message-failed", + kind: "message.failed", + occurredAt: "2026-05-26T00:00:00.000Z", + origin: "driver", + payload: { + error: { code: "runtime.failed", message: "Provider failed." }, + messageId: "message-1", + role: "agent", + }, + sessionId: "session-1", + }), + ); + + expect(projection).toMatchObject({ + contentText: "Message updated.", + eventType: "message.failed", + family: "message", + processStatus: "error", + processType: "agent.message.delta", }); }); test.each(["{}", '{"command":', '{"cwd":"/workspace"}'])( - "keeps running streamed tool input %s out of the canonical projection", + "preserves running tool input snapshot %s without guessing delta semantics", (rawInput) => { const projection = createSessionRuntimeEventProjection( createRuntimeEvent({ @@ -208,10 +313,36 @@ describe("agent runtime event projection", () => { }), ); - expect(projection.toolInputJson).toBeNull(); + expect(projection).toMatchObject({ + toolInputDeltaJson: null, + toolInputJson: rawInput, + }); }, ); + test("projects running tool input deltas separately from snapshots", () => { + const projection = createSessionRuntimeEventProjection( + createRuntimeEvent({ + actor: "driver", + id: "tool-stream", + kind: "tool.call.updated", + occurredAt: "2026-05-26T00:00:00.000Z", + origin: "driver", + payload: { + rawInputDelta: '{"command":', + status: "running", + toolCallId: "tool-stream", + }, + sessionId: "session-1", + }), + ); + + expect(projection).toMatchObject({ + toolInputDeltaJson: '{"command":', + toolInputJson: null, + }); + }); + test("rejects malformed completed tool output before process projection", () => { const event = createRuntimeEvent({ actor: "driver", @@ -256,8 +387,9 @@ describe("agent runtime event projection", () => { exitCode: 1, }, message: "Runtime failed.", - recoverable: true, + retryable: true, }, + recoverable: true, }, runId: "run-1", sessionId: "session-1", diff --git a/apps/api/tests/agent-session-retrieve.test.ts b/apps/api/tests/agent-session-retrieve.test.ts index d272228a..7df37cae 100644 --- a/apps/api/tests/agent-session-retrieve.test.ts +++ b/apps/api/tests/agent-session-retrieve.test.ts @@ -59,12 +59,14 @@ function createAgentSessionRetrieveDatabase(): SqliteD1Database { created_at integer, deployment_version_id text, deployment_version_number integer, + driver_instance_id text, error_code text, error_details_json text, error_message text, id text PRIMARY KEY NOT NULL, model text, provider text, + session_id text, started_at integer, status text, trace_id text, @@ -72,6 +74,14 @@ function createAgentSessionRetrieveDatabase(): SqliteD1Database { updated_at integer ); + CREATE TABLE session_agent_task_snapshot ( + driver_instance_id text NOT NULL, + run_id text NOT NULL, + seq integer NOT NULL, + session_id text PRIMARY KEY NOT NULL, + tasks_json text NOT NULL + ); + INSERT INTO session ( agent_id, created_at, @@ -151,6 +161,108 @@ describe("agent session retrieve", () => { expect(result.session.id).toBe("session-1"); }); + test("returns the current schema-validated task snapshot", async () => { + const database = createAgentSessionRetrieveDatabase(); + database.execute(` + UPDATE session + SET last_run_id = 'run-1', status = 'RUNNING' + WHERE id = 'session-1'; + + INSERT INTO session_run ( + created_at, + driver_instance_id, + id, + session_id, + status, + trace_id, + trigger, + updated_at + ) + VALUES ( + 2, + 'driver-1', + 'run-1', + 'session-1', + 'running', + 'trace-1', + 'user_message', + 2 + ); + + INSERT INTO session_agent_task_snapshot ( + driver_instance_id, + run_id, + seq, + session_id, + tasks_json + ) + VALUES ( + 'driver-1', + 'run-1', + 3, + 'session-1', + '{"tasks":[{"taskId":"task-1","taskType":"review"}]}' + ); + `); + + const result = await retrieveAgentSession(database, VIEWER, { + appId: APP_ID, + sessionId: "session-1", + }); + + expect(result.taskSnapshot).toEqual({ + driverInstanceId: "driver-1", + runId: "run-1", + tasks: [{ taskId: "task-1", taskType: "review" }], + }); + }); + + test.each([ + ["terminal run", null, "completed", '{"tasks":[{"taskId":"stale"}]}'], + ["malformed state", null, "running", '{"tasks":"invalid"}'], + ["archived running session", 3, "running", '{"tasks":[{"taskId":"stale"}]}'], + ])("fails closed for %s task snapshots", async (_label, archivedAt, runStatus, tasksJson) => { + const database = createAgentSessionRetrieveDatabase(); + await database + .prepare("UPDATE session SET archived_at = ?, last_run_id = ?, status = ? WHERE id = ?") + .bind(archivedAt, "run-1", "RUNNING", "session-1") + .run(); + await database + .prepare( + `INSERT INTO session_run ( + created_at, + driver_instance_id, + id, + session_id, + status, + trace_id, + trigger, + updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind(2, "driver-1", "run-1", "session-1", runStatus, "trace-1", "user_message", 2) + .run(); + await database + .prepare( + `INSERT INTO session_agent_task_snapshot ( + driver_instance_id, + run_id, + seq, + session_id, + tasks_json + ) VALUES (?, ?, ?, ?, ?)`, + ) + .bind("driver-1", "run-1", 3, "session-1", tasksJson) + .run(); + + const result = await retrieveAgentSession(database, VIEWER, { + appId: APP_ID, + sessionId: "session-1", + }); + + expect(result.taskSnapshot).toBeNull(); + }); + test("apps terminal cleanup rows as not recoverable even with archive marker", async () => { const database = createAgentSessionRetrieveDatabase(); diff --git a/apps/api/tests/api-command-queue.test.ts b/apps/api/tests/api-command-queue.test.ts index 0ab181b1..0d3e1aab 100644 --- a/apps/api/tests/api-command-queue.test.ts +++ b/apps/api/tests/api-command-queue.test.ts @@ -1,21 +1,59 @@ import { describe, expect, test } from "bun:test"; -import { apiCommandsTable } from "@mosoo/db"; +import { apiCommandsTable, appDeploymentRunsTable, appDeploymentsTable } from "@mosoo/db"; +import { createPlatformId } from "@mosoo/id"; +import type { SandboxBackupId } from "@mosoo/id"; +import type { WorkflowStep } from "cloudflare:workers"; import { eq } from "drizzle-orm"; +import { + enqueueAppDeploymentScriptReconciliationCommand, + enqueueSandboxBackupReconciliationCommand, +} from "../src/modules/api-command/application/api-command-enqueue"; import { API_COMMAND_LEASE_MS, API_COMMAND_QUEUE_DELIVERY_PENDING_CODE, API_COMMAND_QUEUE_SEND_FAILED_CODE, admitApiCommand, claimApiCommand, + completeApiCommand, deliverApiCommand, enqueueApiCommand, + getAppDeploymentWorkflowInstanceId, + markApiCommandDeadLettered, + markApiCommandFailed, + prepareApiCommand, + reconcileAppDeploymentWorkflowDeliveries, redriveFailedApiCommandEnqueues, + releaseApiCommandForRetry, renewApiCommandClaim, } from "../src/modules/api-command/application/api-command-ledger"; +import type { ApiCommandClaim } from "../src/modules/api-command/application/api-command-ledger"; import type { ApiCommandMessage } from "../src/modules/api-command/application/api-command-message"; -import { processApiCommandMessage } from "../src/modules/api-command/application/api-command-processor"; +import { parseApiCommandMessage } from "../src/modules/api-command/application/api-command-message"; +import { parseApiCommandPayload } from "../src/modules/api-command/application/api-command-payload"; +import { + APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS, + APP_DEPLOYMENT_RUN_DISPATCH_RETRY_EXHAUSTED_CODE, +} from "../src/modules/api-command/application/api-command-policy"; +import { + cleanupAppDeploymentScript, + deleteAppDeploymentArtifactPrefix, + processApiCommandDeadLetterMessage, + processApiCommandDelivery, + processApiCommandMessage, +} from "../src/modules/api-command/application/api-command-processor"; +import { + APP_DEPLOYMENT_WORKFLOW_EXECUTION_BUDGET_MS, + ApiCommandWorkflowRetryError, + runAppDeploymentWorkflow, +} from "../src/modules/api-command/application/api-command-workflow"; +import { appDeploymentCandidateScriptName } from "../src/modules/apps/application/app-deployment-executor.service"; +import type { AppDeploymentScriptCandidateAuthority } from "../src/modules/apps/application/app-deployment-script-reconciliation.service"; +import { + markAppDeploymentScriptUploadStarted, + registerAppDeploymentScriptCandidate, +} from "../src/modules/apps/application/app-deployment-script-reconciliation.service"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { createApiCommandQueueStub, @@ -23,10 +61,374 @@ import { createPublicHttpTestBindings, createRecordedQueueMessage, nowMsForTest, + PUBLIC_API_TEST_IDS, } from "./helpers/public-api-http-test-fixture"; +async function claimCommand(input: { + commandId: ApiCommandMessage["commandId"]; + database: D1Database; + deliveryGeneration?: number; + nowMs: number; + claimOwner: string; +}): Promise { + const result = await claimApiCommand({ + ...input, + deliveryGeneration: input.deliveryGeneration ?? 1, + }); + if (result.kind !== "claimed") { + throw new Error(`Expected claimed API command, received ${result.kind}.`); + } + return result.claim; +} + +async function processAppWorkflowDeliveryForTest( + bindings: ApiBindings, + message: Message, + nowMs: () => number, + options: Parameters[3] = {}, +): Promise { + const disposition = await processApiCommandDelivery(bindings, message.body, nowMs, options); + if (disposition.kind === "finished") { + message.ack(); + } else { + message.retry({ delaySeconds: disposition.delaySeconds }); + } +} + +function createInvokingWorkflowStep() { + const calls: { config: unknown; name: string }[] = []; + const step = { + async do(name: string, config: unknown, callback: () => Promise) { + calls.push({ config, name }); + return callback(); + }, + } as unknown as WorkflowStep; + return { calls, step }; +} + +type TestWorkflowStatus = InstanceStatus["status"]; + +function createRecordingAppDeploymentWorkflow( + input: { + createMode?: "created" | "empty" | "throw_after_create" | "throw_before_create"; + initialStatus?: TestWorkflowStatus; + } = {}, +) { + const created: { id: string; params: ApiCommandMessage }[] = []; + const getCalls: string[] = []; + const instances = new Map(); + const restarted: { id: string; options?: WorkflowInstanceRestartOptions }[] = []; + const resumed: string[] = []; + let statusHook: ((id: string) => Promise | void) | undefined; + const workflow = { + async createBatch(batch: { id?: string; params?: ApiCommandMessage }[]) { + if (input.createMode === "throw_before_create") { + throw new Error("Workflow create response failed before acceptance."); + } + const createdInstances = batch.map(({ id, params }) => { + if (id === undefined || params === undefined) { + throw new Error("Expected an explicit Workflow instance ID and params."); + } + created.push({ id, params }); + if (!instances.has(id)) { + instances.set(id, input.initialStatus ?? "queued"); + } + return { id }; + }); + if (input.createMode === "throw_after_create") { + throw new Error("Workflow create response was lost after acceptance."); + } + if (input.createMode === "empty") { + return []; + } + return createdInstances; + }, + async get(id: string) { + getCalls.push(id); + if (!instances.has(id)) { + throw new Error("Workflow instance does not exist."); + } + return { + id, + async restart(options?: WorkflowInstanceRestartOptions) { + restarted.push({ id, options }); + }, + async resume() { + resumed.push(id); + }, + async status() { + await statusHook?.(id); + return { status: instances.get(id) ?? "unknown" }; + }, + }; + }, + } as unknown as Workflow; + return { + created, + deleteInstance: (id: string) => instances.delete(id), + getCalls, + instances, + restarted, + resumed, + setStatus: (id: string, status: TestWorkflowStatus) => instances.set(id, status), + setStatusHook: (hook: (id: string) => Promise | void) => { + statusHook = hook; + }, + workflow, + }; +} + +async function createAppDeploymentCommandFixture(input: { + errorCode?: string; + errorMessage?: string; + runStatus: + | "activating" + | "building" + | "failed" + | "preparing" + | "queued" + | "submitted" + | "submitting"; + workflow?: ReturnType; +}) { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const appDeploymentWorkflow = input.workflow ?? createRecordingAppDeploymentWorkflow(); + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + bindings.APP_DEPLOYMENT_WORKFLOW = appDeploymentWorkflow.workflow; + const deploymentId = PUBLIC_API_TEST_IDS.deployment; + const runId = PUBLIC_API_TEST_IDS.run; + + await database + .app() + .insert(appDeploymentsTable) + .values({ + appId: PUBLIC_API_TEST_IDS.app, + createdAt: nowMsForTest(), + defaultBranch: "main", + id: deploymentId, + latestRunId: runId, + mosooSubdomain: "api-command-test", + ownerAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + repoName: "repo", + repoOwner: "owner", + repoUrl: "https://github.com/owner/repo", + sourceKind: "github_public", + updatedAt: nowMsForTest(), + }) + .run(); + await database + .app() + .insert(appDeploymentRunsTable) + .values({ + appId: PUBLIC_API_TEST_IDS.app, + createdAt: nowMsForTest(), + deploymentId, + errorCode: input.errorCode ?? null, + errorMessage: input.errorMessage ?? null, + id: runId, + sourceBranch: "main", + sourceCommitSha: "a".repeat(40), + status: input.runStatus, + updatedAt: nowMsForTest(), + }) + .run(); + const commandId = await enqueueApiCommand(bindings, { + dedupeKey: `app_deployment_run_dispatch:${runId}`, + kind: "app_deployment_run_dispatch", + payload: { appDeploymentRunId: runId }, + }); + const body = appDeploymentWorkflow.created[0]?.params ?? { + commandId, + deliveryGeneration: 1, + }; + + return { + appDeploymentWorkflow, + bindings, + body, + commandId, + database, + deploymentId, + queue, + runId, + }; +} + +async function registerAppDeploymentCandidateForTest( + fixture: Awaited>, + authority: AppDeploymentScriptCandidateAuthority, + nowMs: number, + readyForOutcome = false, +): Promise { + const scriptName = appDeploymentCandidateScriptName(authority); + await registerAppDeploymentScriptCandidate(fixture.database, { + authority, + deploymentId: fixture.deploymentId, + nowMs, + runId: fixture.runId, + scriptName, + }); + if (readyForOutcome) { + await markAppDeploymentScriptUploadStarted(fixture.database, { + authority, + deploymentId: fixture.deploymentId, + nowMs, + runId: fixture.runId, + scriptName, + }); + await fixture.database + .app() + .update(appDeploymentRunsTable) + .set({ status: "activating" }) + .where(eq(appDeploymentRunsTable.id, fixture.runId)) + .run(); + } + return scriptName; +} + describe("API command queue", () => { - test("dedupes producer-side and sends only the command id", async () => { + test("requires a positive safe delivery generation in every queue message", () => { + const commandId = "01J0000000000000000000000C"; + + expect(parseApiCommandMessage({ commandId, deliveryGeneration: 1 })).toEqual({ + commandId, + deliveryGeneration: 1, + }); + expect( + parseApiCommandMessage({ commandId, deliveryGeneration: Number.MAX_SAFE_INTEGER }), + ).toEqual({ commandId, deliveryGeneration: Number.MAX_SAFE_INTEGER }); + + for (const deliveryGeneration of [undefined, 0, 1.5, Number.MAX_SAFE_INTEGER + 1]) { + expect(() => parseApiCommandMessage({ commandId, deliveryGeneration })).toThrow(); + } + }); + + test("keeps the D1 backup page independent from the opaque R2 cursor", () => { + const payload = { cursor: null, databasePage: 0, scheduledTime: nowMsForTest() }; + expect( + parseApiCommandPayload("sandbox_backup_reconciliation", JSON.stringify(payload)), + ).toEqual(payload); + for (const databasePage of [-1, 0.5, Number.MAX_SAFE_INTEGER + 1]) { + expect(() => + parseApiCommandPayload( + "sandbox_backup_reconciliation", + JSON.stringify({ ...payload, databasePage }), + ), + ).toThrow(); + } + }); + + test("continues a D1 backup backlog without an R2 cursor", async () => { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const bucket = { + async delete(): Promise {}, + async list(): Promise { + return { delimitedPrefixes: [], objects: [], truncated: false }; + }, + } as unknown as R2Bucket; + const bindings = { + ...createPublicHttpTestBindings(database, { apiCommandQueue: queue }), + SANDBOX_STATE_BUCKET: bucket, + } as ApiBindings; + database.execute("DROP TRIGGER sandbox_backup_delete_intent_authority"); + await database.batch( + Array.from({ length: 65 }, () => + database + .prepare( + `INSERT INTO sandbox_backup_delete_intent ( + attempted_at, backup_id, created_at, delete_after, deleted_at + ) VALUES (NULL, ?, 0, 0, NULL)`, + ) + .bind(createPlatformId()), + ), + ); + const scheduledTime = nowMsForTest(); + await enqueueSandboxBackupReconciliationCommand(bindings, { + cursor: null, + databasePage: 0, + scheduledTime, + }); + + const first = createRecordedQueueMessage({ body: queue.sent[0].body }); + await processApiCommandMessage(bindings, first.message, nowMsForTest); + + expect(first.recorded).toEqual([{ type: "ack" }]); + expect(queue.sent).toHaveLength(2); + const rows = await database + .app() + .select({ dedupeKey: apiCommandsTable.dedupeKey, payloadJson: apiCommandsTable.payloadJson }) + .from(apiCommandsTable) + .orderBy(apiCommandsTable.dedupeKey) + .all(); + expect(rows.map(({ dedupeKey }) => dedupeKey)).toEqual([ + `sandbox_backup_reconciliation:${scheduledTime}:0:start`, + `sandbox_backup_reconciliation:${scheduledTime}:1:start`, + ]); + expect(parseApiCommandPayload("sandbox_backup_reconciliation", rows[1].payloadJson)).toEqual({ + cursor: null, + databasePage: 1, + scheduledTime, + }); + + const second = createRecordedQueueMessage({ body: queue.sent[1].body }); + await processApiCommandMessage(bindings, second.message, nowMsForTest); + + expect(second.recorded).toEqual([{ type: "ack" }]); + expect(queue.sent).toHaveLength(2); + await expect( + database + .prepare( + "SELECT count(*) AS count FROM sandbox_backup_delete_intent WHERE deleted_at IS NOT NULL", + ) + .first(), + ).resolves.toEqual({ count: 65 }); + }); + + test("enqueues and parses an App deployment script reconciliation page", async () => { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + const scheduledTime = Date.UTC(2026, 7, 30, 4); + + await enqueueAppDeploymentScriptReconciliationCommand(bindings, { + cursor: null, + scheduledTime, + }); + + const [row] = await database + .app() + .select({ + dedupeKey: apiCommandsTable.dedupeKey, + id: apiCommandsTable.id, + kind: apiCommandsTable.kind, + payloadJson: apiCommandsTable.payloadJson, + }) + .from(apiCommandsTable) + .all(); + expect(row).toMatchObject({ + dedupeKey: `app-deployment-script-reconciliation:${scheduledTime}:root`, + kind: "app_deployment_script_reconciliation", + }); + expect(queue.sent[0]?.body).toEqual({ commandId: row?.id, deliveryGeneration: 1 }); + expect(parseApiCommandPayload(row.kind, row.payloadJson)).toEqual({ + cursor: null, + scheduledTime, + }); + expect(() => + parseApiCommandPayload( + "app_deployment_script_reconciliation", + JSON.stringify({ cursor: "", scheduledTime }), + ), + ).toThrow("cursor must not be empty"); + }); + + test("dedupes producer-side and sends the durable delivery generation", async () => { const database = await createPublicHttpContractDatabase(); const queue = createApiCommandQueueStub(); const bindings = createPublicHttpTestBindings(database, { @@ -47,7 +449,7 @@ describe("API command queue", () => { expect(duplicateId).toBe(firstId); expect(queue.sent).toEqual([ { - body: { commandId: firstId }, + body: { commandId: firstId, deliveryGeneration: 1 }, contentType: "json", delaySeconds: null, id: "queued-1", @@ -91,7 +493,649 @@ describe("API command queue", () => { }); await deliverApiCommand(bindings, admission); - expect(queue.sent[0]?.body).toEqual({ commandId: admission.commandId }); + expect(queue.sent[0]?.body).toEqual({ + commandId: admission.commandId, + deliveryGeneration: 1, + }); + }); + + test("routes App deployment commands to one fixed Workflow instance", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + + expect(fixture.queue.sent).toEqual([]); + expect(fixture.appDeploymentWorkflow.created).toEqual([ + { + id: getAppDeploymentWorkflowInstanceId(fixture.commandId, 1), + params: { commandId: fixture.commandId, deliveryGeneration: 1 }, + }, + ]); + }); + + test("runs App deployment in one bounded Workflow step with one absolute deadline", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + const workflowStep = createInvokingWorkflowStep(); + const stepStartedAt = nowMsForTest() + 100; + let nowCallCount = 0; + let receivedDeadlineMs: number | undefined; + + await runAppDeploymentWorkflow( + fixture.bindings, + fixture.body, + workflowStep.step, + () => { + nowCallCount += 1; + return nowCallCount === 1 ? stepStartedAt : stepStartedAt + 5_000; + }, + { + dispatchAppDeploymentRun: async (_bindings, _payload, authority, options) => { + receivedDeadlineMs = options?.deadlineMs; + await authority.requireOwnership(); + return { + errorCode: "workflow_test_terminal", + errorMessage: "The Workflow test completed without an external deployment.", + kind: "terminal_failure", + runId: fixture.runId, + }; + }, + }, + ); + + expect(workflowStep.calls).toEqual([ + { + config: { + retries: { backoff: "constant", delay: "30 seconds", limit: 1_000 }, + timeout: "30 minutes", + }, + name: "dispatch", + }, + ]); + expect(receivedDeadlineMs).toBe(stepStartedAt + APP_DEPLOYMENT_WORKFLOW_EXECUTION_BUDGET_MS); + }); + + test("throws the dedicated Workflow retry error while the D1 claim is busy", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + const claimedAt = nowMsForTest() + 100; + await claimCommand({ + claimOwner: "workflow-busy-owner", + commandId: fixture.commandId, + database: fixture.database, + nowMs: claimedAt, + }); + + await expect( + runAppDeploymentWorkflow( + fixture.bindings, + fixture.body, + createInvokingWorkflowStep().step, + () => claimedAt + 1, + ), + ).rejects.toBeInstanceOf(ApiCommandWorkflowRetryError); + }); + + test("hands residual main and DLQ App messages to Workflow without dispatching", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + const main = createRecordedQueueMessage({ body: fixture.body }); + const deadLetter = createRecordedQueueMessage({ body: fixture.body }); + let dispatchCount = 0; + + await processApiCommandMessage(fixture.bindings, main.message, nowMsForTest, { + dispatchAppDeploymentRun: async () => { + dispatchCount += 1; + return { kind: "skipped" }; + }, + }); + await processApiCommandDeadLetterMessage(fixture.bindings, deadLetter.message, nowMsForTest); + + expect(dispatchCount).toBe(0); + expect(main.recorded).toEqual([{ type: "ack" }]); + expect(deadLetter.recorded).toEqual([{ type: "ack" }]); + expect(fixture.appDeploymentWorkflow.created.map(({ id }) => id)).toEqual([ + getAppDeploymentWorkflowInstanceId(fixture.commandId, 1), + getAppDeploymentWorkflowInstanceId(fixture.commandId, 1), + getAppDeploymentWorkflowInstanceId(fixture.commandId, 1), + ]); + }); + + test("retries a residual App message when Workflow handoff is ambiguous", async () => { + const workflow = createRecordingAppDeploymentWorkflow({ + createMode: "throw_after_create", + }); + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued", workflow }); + const recorded = createRecordedQueueMessage({ body: fixture.body }); + let dispatchCount = 0; + + await processApiCommandMessage(fixture.bindings, recorded.message, nowMsForTest, { + dispatchAppDeploymentRun: async () => { + dispatchCount += 1; + return { kind: "skipped" }; + }, + }); + + expect(dispatchCount).toBe(0); + expect(recorded.recorded).toEqual([{ delaySeconds: 30, type: "retry" }]); + }); + + test.each(["missing", "stale", "terminal"] as const)( + "acks a %s residual App message without Workflow handoff", + async (state) => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + if (state === "missing") { + await fixture.database + .app() + .delete(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .run(); + } else { + await fixture.database + .app() + .update(apiCommandsTable) + .set(state === "stale" ? { deliveryGeneration: 2 } : { status: "failed" }) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .run(); + } + const recorded = createRecordedQueueMessage({ body: fixture.body }); + + await processApiCommandMessage(fixture.bindings, recorded.message, nowMsForTest); + + expect(recorded.recorded).toEqual([{ type: "ack" }]); + expect(fixture.appDeploymentWorkflow.created).toHaveLength(1); + }, + ); + + test("collects every paginated artifact key before deleting the prefix", async () => { + const remaining = new Set(["apps/a/one", "apps/a/two", "apps/a/three", "other"]); + const listCursors: Array = []; + const deleted: string[][] = []; + const bucket = { + async delete(keys: string | string[]) { + const page = typeof keys === "string" ? [keys] : keys; + deleted.push(page); + for (const key of page) { + remaining.delete(key); + } + }, + async list(options: R2ListOptions) { + listCursors.push(options.cursor); + const keys = options.cursor === undefined ? ["apps/a/one", "apps/a/two"] : ["apps/a/three"]; + return { + cursor: options.cursor === undefined ? "next" : undefined, + delimitedPrefixes: [], + objects: keys.map((key) => ({ key })), + truncated: options.cursor === undefined, + } as unknown as R2Objects; + }, + } as unknown as R2Bucket; + + await deleteAppDeploymentArtifactPrefix(bucket, "apps/a/"); + + expect(listCursors).toEqual([undefined, "next"]); + expect(deleted).toEqual([["apps/a/one", "apps/a/two", "apps/a/three"]]); + expect([...remaining]).toEqual(["other"]); + }); + + test("fails artifact cleanup before deleting when a truncated page has no cursor", async () => { + let deleteCount = 0; + const bucket = { + async delete() { + deleteCount += 1; + }, + async list() { + return { + delimitedPrefixes: [], + objects: [{ key: "apps/a/one" }], + truncated: true, + } as unknown as R2Objects; + }, + } as unknown as R2Bucket; + + await expect(deleteAppDeploymentArtifactPrefix(bucket, "apps/a/")).rejects.toThrow( + "no continuation cursor", + ); + expect(deleteCount).toBe(0); + }); + + test("stops both artifact producers before deleting their outputs", async () => { + const database = await createPublicHttpContractDatabase(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const buildDestroy = Promise.withResolvers(); + const deployDestroy = Promise.withResolvers(); + const bothDestroying = Promise.withResolvers(); + let destroying = 0; + let listCount = 0; + bindings.runtimeSubjectHandleFactory = (sandboxId) => + new Proxy( + {}, + { + get(_target, property) { + if (typeof property === "symbol" || property === "then") { + return undefined; + } + if (property === "setKeepAlive") { + return async () => {}; + } + if (property === "destroy") { + return async () => { + destroying += 1; + if (destroying === 2) { + bothDestroying.resolve(); + } + await (sandboxId === "build-sandbox" + ? buildDestroy.promise + : deployDestroy.promise); + }; + } + return async () => { + throw new Error(`Unexpected sandbox call: ${property}`); + }; + }, + }, + ); + bindings.FILE_BUCKET = { + async delete() {}, + async list() { + listCount += 1; + return { delimitedPrefixes: [], objects: [], truncated: false }; + }, + }; + const commandId = prepareApiCommand({ + dedupeKey: "app-script-cleanup:fence", + kind: "app_deployment_run_dispatch", + payload: { appDeploymentRunId: PUBLIC_API_TEST_IDS.run }, + }).commandId; + + const cleanup = cleanupAppDeploymentScript(bindings, { + artifactPrefix: "apps/script/", + attemptCount: 1, + buildSandboxId: "build-sandbox", + commandId, + deliveryGeneration: 1, + deploySandboxId: "deploy-sandbox", + deploymentId: PUBLIC_API_TEST_IDS.deployment, + reconcileCount: 0, + runId: PUBLIC_API_TEST_IDS.run, + scriptName: "script", + uploadStartedAt: null, + }); + await bothDestroying.promise; + expect(listCount).toBe(0); + + buildDestroy.resolve(); + await Promise.resolve(); + expect(listCount).toBe(0); + + deployDestroy.resolve(); + await cleanup; + expect(listCount).toBe(1); + }); + + test("does not delete outputs when either artifact producer cannot stop", async () => { + const database = await createPublicHttpContractDatabase(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + let listCount = 0; + bindings.runtimeSubjectHandleFactory = (sandboxId) => + new Proxy( + {}, + { + get(_target, property) { + if (typeof property === "symbol" || property === "then") { + return undefined; + } + if (property === "setKeepAlive") { + return async () => {}; + } + if (property === "destroy") { + return async () => { + if (sandboxId === "build-sandbox") { + throw new Error("build sandbox is still producing artifacts"); + } + }; + } + return async () => { + throw new Error(`Unexpected sandbox call: ${property}`); + }; + }, + }, + ); + bindings.FILE_BUCKET = { + async delete() {}, + async list() { + listCount += 1; + return { delimitedPrefixes: [], objects: [], truncated: false }; + }, + }; + const commandId = prepareApiCommand({ + dedupeKey: "app-script-cleanup:producer-failure", + kind: "app_deployment_run_dispatch", + payload: { appDeploymentRunId: PUBLIC_API_TEST_IDS.run }, + }).commandId; + + await expect( + cleanupAppDeploymentScript(bindings, { + artifactPrefix: "apps/script/", + attemptCount: 1, + buildSandboxId: "build-sandbox", + commandId, + deliveryGeneration: 1, + deploySandboxId: "deploy-sandbox", + deploymentId: PUBLIC_API_TEST_IDS.deployment, + reconcileCount: 0, + runId: PUBLIC_API_TEST_IDS.run, + scriptName: "script", + uploadStartedAt: null, + }), + ).rejects.toBeInstanceOf(AggregateError); + expect(listCount).toBe(0); + }); + + test.each([ + "complete", + "errored", + "paused", + "queued", + "running", + "terminated", + "waiting", + "waitingForPause", + ] as const)("redrives an ambiguous Workflow create after an ACK loss in %s", async (status) => { + const workflow = createRecordingAppDeploymentWorkflow({ + createMode: "throw_after_create", + initialStatus: status, + }); + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued", workflow }); + + await expect( + fixture.database + .app() + .select({ lastErrorCode: apiCommandsTable.lastErrorCode }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ lastErrorCode: API_COMMAND_QUEUE_SEND_FAILED_CODE }); + expect(workflow.getCalls).toEqual([]); + }); + + test("confirms a duplicate Workflow create that resolves with an empty batch", async () => { + const workflow = createRecordingAppDeploymentWorkflow({ createMode: "empty" }); + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued", workflow }); + + await expect( + fixture.database + .app() + .select({ lastErrorCode: apiCommandsTable.lastErrorCode }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ lastErrorCode: null }); + expect(workflow.getCalls).toEqual([]); + }); + + test.each([ + { createMode: "throw_before_create", initialStatus: "queued" }, + { createMode: "throw_after_create", initialStatus: "unknown" }, + ] as const)("keeps an unconfirmed Workflow delivery pending: %o", async (workflowInput) => { + const workflow = createRecordingAppDeploymentWorkflow(workflowInput); + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued", workflow }); + + await expect( + fixture.database + .app() + .select({ lastErrorCode: apiCommandsTable.lastErrorCode }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ lastErrorCode: API_COMMAND_QUEUE_SEND_FAILED_CODE }); + }); + + test.each(["queued", "running", "waiting"] as const)( + "leaves an active %s Workflow instance alone", + async (status) => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + const instanceId = getAppDeploymentWorkflowInstanceId(fixture.commandId, 1); + fixture.appDeploymentWorkflow.setStatus(instanceId, status); + + await reconcileAppDeploymentWorkflowDeliveries(fixture.bindings, nowMsForTest() + 1); + + expect(fixture.appDeploymentWorkflow.resumed).toEqual([]); + expect(fixture.appDeploymentWorkflow.restarted).toEqual([]); + await expect( + fixture.database + .app() + .select({ deliveryGeneration: apiCommandsTable.deliveryGeneration }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ deliveryGeneration: 1 }); + }, + ); + + test.each(["paused", "waitingForPause"] as const)( + "resumes a %s Workflow because D1 remains active", + async (status) => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + const instanceId = getAppDeploymentWorkflowInstanceId(fixture.commandId, 1); + fixture.appDeploymentWorkflow.setStatus(instanceId, status); + + await reconcileAppDeploymentWorkflowDeliveries(fixture.bindings, nowMsForTest() + 1); + + expect(fixture.appDeploymentWorkflow.resumed).toEqual([instanceId]); + }, + ); + + test.each(["complete", "errored", "terminated"] as const)( + "restarts an unexpected %s Workflow from the beginning", + async (status) => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + const instanceId = getAppDeploymentWorkflowInstanceId(fixture.commandId, 1); + fixture.appDeploymentWorkflow.setStatus(instanceId, status); + + await reconcileAppDeploymentWorkflowDeliveries(fixture.bindings, nowMsForTest() + 1); + + expect(fixture.appDeploymentWorkflow.restarted).toEqual([ + { id: instanceId, options: undefined }, + ]); + }, + ); + + test("re-ensures the same Workflow ID when status is unknown", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + const instanceId = getAppDeploymentWorkflowInstanceId(fixture.commandId, 1); + fixture.appDeploymentWorkflow.setStatus(instanceId, "unknown"); + + await reconcileAppDeploymentWorkflowDeliveries(fixture.bindings, nowMsForTest() + 1); + + expect(fixture.appDeploymentWorkflow.created.map(({ id }) => id)).toEqual([ + instanceId, + instanceId, + ]); + }); + + test("re-ensures the same Workflow ID when status cannot be observed", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + const instanceId = getAppDeploymentWorkflowInstanceId(fixture.commandId, 1); + fixture.appDeploymentWorkflow.setStatusHook(() => { + throw new Error("Workflow status is unavailable."); + }); + + await reconcileAppDeploymentWorkflowDeliveries(fixture.bindings, nowMsForTest() + 1); + + expect(fixture.appDeploymentWorkflow.created.map(({ id }) => id)).toEqual([ + instanceId, + instanceId, + ]); + await expect( + fixture.database + .app() + .select({ deliveryGeneration: apiCommandsTable.deliveryGeneration }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ deliveryGeneration: 1 }); + }); + + test("recreates a missing Workflow with the same generation and attempt budget", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + const instanceId = getAppDeploymentWorkflowInstanceId(fixture.commandId, 1); + await fixture.database + .app() + .update(apiCommandsTable) + .set({ attemptCount: 2 }) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .run(); + + for (const nowMs of [nowMsForTest() + 1, nowMsForTest() + 2]) { + fixture.appDeploymentWorkflow.deleteInstance(instanceId); + await reconcileAppDeploymentWorkflowDeliveries(fixture.bindings, nowMs); + } + + expect(fixture.appDeploymentWorkflow.created.map(({ id }) => id)).toEqual([ + instanceId, + instanceId, + instanceId, + ]); + await expect( + fixture.database + .app() + .select({ + attemptCount: apiCommandsTable.attemptCount, + deliveryGeneration: apiCommandsTable.deliveryGeneration, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ attemptCount: 2, deliveryGeneration: 1 }); + }); + + test("terminalizes attempt four after restarting a terminated Workflow", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "activating" }); + await fixture.database + .app() + .update(apiCommandsTable) + .set({ attemptCount: APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS }) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .run(); + fixture.appDeploymentWorkflow.setStatus( + getAppDeploymentWorkflowInstanceId(fixture.commandId, 1), + "terminated", + ); + + await reconcileAppDeploymentWorkflowDeliveries(fixture.bindings, nowMsForTest() + 1); + expect(fixture.appDeploymentWorkflow.restarted).toEqual([ + { + id: getAppDeploymentWorkflowInstanceId(fixture.commandId, 1), + options: undefined, + }, + ]); + const recorded = createRecordedQueueMessage({ body: fixture.body }); + let dispatchCount = 0; + await processAppWorkflowDeliveryForTest( + fixture.bindings, + recorded.message, + () => nowMsForTest() + 2, + { + dispatchAppDeploymentRun: async () => { + dispatchCount += 1; + return { kind: "skipped" }; + }, + }, + ); + + expect(dispatchCount).toBe(0); + expect(recorded.recorded).toEqual([{ type: "ack" }]); + await expect( + fixture.database + .app() + .select({ + attemptCount: apiCommandsTable.attemptCount, + deliveryGeneration: apiCommandsTable.deliveryGeneration, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ + attemptCount: APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS + 1, + deliveryGeneration: 1, + status: "failed", + }); + }); + + test("does not inspect a Workflow while its D1 claim lease is live", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + await claimCommand({ + claimOwner: "live-workflow-owner", + commandId: fixture.commandId, + database: fixture.database, + nowMs: nowMsForTest(), + }); + + await reconcileAppDeploymentWorkflowDeliveries(fixture.bindings, nowMsForTest() + 1); + + expect(fixture.appDeploymentWorkflow.getCalls).toEqual([]); + }); + + test("round-robins a bounded Workflow reconciliation page without starvation", async () => { + const database = await createPublicHttpContractDatabase(); + const workflow = createRecordingAppDeploymentWorkflow(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + bindings.APP_DEPLOYMENT_WORKFLOW = workflow.workflow; + const records = Array.from({ length: 201 }, (_, index) => + prepareApiCommand( + { + dedupeKey: `app_deployment_run_dispatch:fairness:${index}`, + kind: "app_deployment_run_dispatch", + payload: { appDeploymentRunId: PUBLIC_API_TEST_IDS.run }, + }, + { timestampMs: 100 }, + ), + ); + for (let offset = 0; offset < records.length; offset += 40) { + await database + .app() + .insert(apiCommandsTable) + .values(records.slice(offset, offset + 40).map(({ record }) => record)) + .run(); + } + for (const { commandId } of records) { + workflow.instances.set(getAppDeploymentWorkflowInstanceId(commandId, 1), "queued"); + } + + for (const nowMs of [2, 3, 4]) { + await reconcileAppDeploymentWorkflowDeliveries(bindings, nowMs); + } + + expect(workflow.getCalls).toHaveLength(300); + expect(new Set(workflow.getCalls).size).toBe(201); + }); + + test("does not move the reconciliation cursor after a concurrent generation change", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "queued" }); + const concurrentUpdatedAt = nowMsForTest() + 7; + fixture.appDeploymentWorkflow.setStatusHook(async () => { + await fixture.database + .app() + .update(apiCommandsTable) + .set({ + deliveryGeneration: 2, + status: "succeeded", + updatedAt: concurrentUpdatedAt, + }) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .run(); + }); + + await reconcileAppDeploymentWorkflowDeliveries(fixture.bindings, nowMsForTest() + 100); + + await expect( + fixture.database + .app() + .select({ + deliveryGeneration: apiCommandsTable.deliveryGeneration, + updatedAt: apiCommandsTable.updatedAt, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ deliveryGeneration: 2, updatedAt: concurrentUpdatedAt }); }); test("marks malformed payload commands failed and acks the message", async () => { @@ -132,6 +1176,40 @@ describe("API command queue", () => { expect(recorded.recorded).toEqual([{ type: "ack" }]); }); + test("terminalizes a malformed Environment main-lane command before ACK", async () => { + const database = await createPublicHttpContractDatabase(); + const apiQueue = createApiCommandQueueStub(); + const environmentQueue = createApiCommandQueueStub(); + const bindings = { + ...createPublicHttpTestBindings(database, { apiCommandQueue: apiQueue }), + ENVIRONMENT_ARTIFACT_BUILD_QUEUE: environmentQueue, + } as ApiBindings; + const commandId = await enqueueApiCommand(bindings, { + dedupeKey: "environment-package-artifact:malformed", + kind: "environment_package_artifact_build", + payload: {}, + }); + const queued = environmentQueue.sent[0]?.body; + if (!queued) throw new Error("Expected an Environment artifact Queue message."); + const recorded = createRecordedQueueMessage({ body: queued }); + + await processApiCommandMessage(bindings, recorded.message, nowMsForTest); + + expect(apiQueue.sent).toHaveLength(0); + expect(recorded.recorded).toEqual([{ type: "ack" }]); + await expect( + database + .app() + .select({ + lastErrorCode: apiCommandsTable.lastErrorCode, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, commandId)) + .get(), + ).resolves.toEqual({ lastErrorCode: "invalid_payload", status: "failed" }); + }); + test("keeps a command claimable when Queue accepts it but send reports a timeout", async () => { const database = await createPublicHttpContractDatabase(); const retainedMessages: ApiCommandMessage[] = []; @@ -173,12 +1251,16 @@ describe("API command queue", () => { expect(commandId).toBe(retained.commandId); await expect( claimApiCommand({ + claimOwner: "consumer-after-timeout", commandId: retained.commandId, database, + deliveryGeneration: retained.deliveryGeneration, nowMs: nowMsForTest(), - ownerId: "consumer-after-timeout", }), - ).resolves.toMatchObject({ commandId: retained.commandId }); + ).resolves.toMatchObject({ + claim: { commandId: retained.commandId }, + kind: "claimed", + }); }); test("redrives a durable command after a definite Queue send failure", async () => { @@ -266,7 +1348,7 @@ describe("API command queue", () => { .get(); expect(queue.sent).toHaveLength(1); - expect(queue.sent[0]?.body).toEqual({ commandId }); + expect(queue.sent[0]?.body).toEqual({ commandId, deliveryGeneration: 1 }); expect(row).toEqual({ lastErrorCode: null, lastErrorMessage: null }); }); @@ -282,19 +1364,18 @@ describe("API command queue", () => { payload: { scheduledTime: nowMsForTest() }, }); - await claimApiCommand({ + const claim = await claimCommand({ + claimOwner: "owner-1", commandId, database, nowMs: 1_000, - ownerId: "owner-1", }); await expect( renewApiCommandClaim({ - commandId, + claim, database, nowMs: 2_000, - ownerId: "owner-1", }), ).resolves.toBe(true); @@ -309,4 +1390,942 @@ describe("API command queue", () => { expect(row?.claimExpiresAt).toBe(2_000 + API_COMMAND_LEASE_MS); }); + + test("classifies claimed, busy, stale, terminal, and missing deliveries", async () => { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + const commandId = await enqueueApiCommand(bindings, { + dedupeKey: "scheduled:claim-dispositions", + kind: "scheduled_maintenance", + payload: { scheduledTime: nowMsForTest() }, + }); + const claimed = await claimApiCommand({ + claimOwner: "owner-a", + commandId, + database, + deliveryGeneration: 1, + nowMs: 1_000, + }); + + expect(claimed).toMatchObject({ + claim: { attemptCount: 1, claimOwner: "owner-a", commandId, deliveryGeneration: 1 }, + kind: "claimed", + }); + await expect( + claimApiCommand({ + claimOwner: "owner-b", + commandId, + database, + deliveryGeneration: 1, + nowMs: 2_000, + }), + ).resolves.toEqual({ + claimExpiresAt: 1_000 + API_COMMAND_LEASE_MS, + kind: "busy", + }); + await expect( + claimApiCommand({ + claimOwner: "owner-b", + commandId, + database, + deliveryGeneration: 2, + nowMs: 2_000, + }), + ).resolves.toEqual({ kind: "stale" }); + + if (claimed.kind !== "claimed") { + throw new Error("Expected claimed API command."); + } + await completeApiCommand({ claim: claimed.claim, database, nowMs: 2_000 }); + await expect( + claimApiCommand({ + claimOwner: "owner-b", + commandId, + database, + deliveryGeneration: 1, + nowMs: 3_000, + }), + ).resolves.toEqual({ kind: "terminal" }); + await expect( + claimApiCommand({ + claimOwner: "owner-b", + commandId: "01J0000000000000000000000Z", + database, + deliveryGeneration: 1, + nowMs: 3_000, + }), + ).resolves.toEqual({ kind: "missing" }); + }); + + test("fences completion at expiry and lets the next attempt take over", async () => { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + const commandId = await enqueueApiCommand(bindings, { + dedupeKey: "scheduled:expiry-boundary", + kind: "scheduled_maintenance", + payload: { scheduledTime: nowMsForTest() }, + }); + const first = await claimCommand({ + claimOwner: "owner-a", + commandId, + database, + nowMs: 1_000, + }); + const expiresAt = 1_000 + API_COMMAND_LEASE_MS; + + await expect(completeApiCommand({ claim: first, database, nowMs: expiresAt })).resolves.toBe( + false, + ); + await expect( + claimApiCommand({ + claimOwner: "owner-b", + commandId, + database, + deliveryGeneration: 1, + nowMs: expiresAt, + }), + ).resolves.toMatchObject({ + claim: { attemptCount: 2, claimOwner: "owner-b" }, + kind: "claimed", + }); + }); + + test("rejects every stale attempt mutation after a successor takes over", async () => { + const mutations: readonly { + apply(claim: ApiCommandClaim, database: D1Database, nowMs: number): Promise; + name: string; + }[] = [ + { + apply: (claim, database, nowMs) => renewApiCommandClaim({ claim, database, nowMs }), + name: "renew", + }, + { + apply: (claim, database, nowMs) => completeApiCommand({ claim, database, nowMs }), + name: "complete", + }, + { + apply: (claim, database, nowMs) => + releaseApiCommandForRetry({ + claim, + database, + errorCode: "stale", + errorMessage: "stale", + nowMs, + }), + name: "release", + }, + { + apply: (claim, database, nowMs) => + markApiCommandFailed({ + claim, + database, + errorCode: "stale", + errorMessage: "stale", + nowMs, + }), + name: "fail", + }, + { + apply: (claim, database, nowMs) => + markApiCommandDeadLettered({ + claim, + database, + errorCode: "stale", + errorMessage: "stale", + nowMs, + }), + name: "dead-letter", + }, + ]; + + for (const mutation of mutations) { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + const commandId = await enqueueApiCommand(bindings, { + dedupeKey: `scheduled:stale-${mutation.name}`, + kind: "scheduled_maintenance", + payload: { scheduledTime: nowMsForTest() }, + }); + const first = await claimCommand({ + claimOwner: "owner-a", + commandId, + database, + nowMs: 1_000, + }); + const takeoverAt = 1_000 + API_COMMAND_LEASE_MS; + await claimCommand({ + claimOwner: "owner-b", + commandId, + database, + nowMs: takeoverAt, + }); + const before = await database + .app() + .select() + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, commandId)) + .get(); + + await expect(mutation.apply(first, database, takeoverAt + 1)).resolves.toBe(false); + await expect( + database + .app() + .select() + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, commandId)) + .get(), + ).resolves.toEqual(before); + } + }); + + test("increments the delivery generation for an explicit terminal retry", async () => { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + const commandId = await enqueueApiCommand(bindings, { + dedupeKey: "scheduled:terminal-retry", + kind: "scheduled_maintenance", + payload: { scheduledTime: 1 }, + }); + const first = await claimCommand({ + claimOwner: "owner-a", + commandId, + database, + nowMs: 1_000, + }); + await markApiCommandFailed({ + claim: first, + database, + errorCode: "failed", + errorMessage: "failed", + nowMs: 2_000, + }); + + const admission = await admitApiCommand(bindings, { + dedupeKey: "scheduled:terminal-retry", + kind: "scheduled_maintenance", + payload: { scheduledTime: 2 }, + retryTerminal: true, + }); + await deliverApiCommand(bindings, admission); + + await expect( + database + .app() + .select({ + attemptCount: apiCommandsTable.attemptCount, + deliveryGeneration: apiCommandsTable.deliveryGeneration, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, commandId)) + .get(), + ).resolves.toEqual({ attemptCount: 0, deliveryGeneration: 2, status: "queued" }); + expect(queue.sent.at(-1)?.body).toEqual({ commandId, deliveryGeneration: 2 }); + await expect( + claimApiCommand({ + claimOwner: "old-message", + commandId, + database, + deliveryGeneration: 1, + nowMs: 3_000, + }), + ).resolves.toEqual({ kind: "stale" }); + }); + + test("fails closed when delivery generation or attempt count is exhausted", async () => { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + const commandId = await enqueueApiCommand(bindings, { + dedupeKey: "scheduled:exhausted", + kind: "scheduled_maintenance", + payload: { scheduledTime: 1 }, + }); + + database.execute( + `UPDATE api_command SET delivery_generation = ${Number.MAX_SAFE_INTEGER}, status = 'failed' WHERE id = '${commandId}'`, + ); + await expect( + admitApiCommand(bindings, { + dedupeKey: "scheduled:exhausted", + kind: "scheduled_maintenance", + payload: { scheduledTime: 2 }, + retryTerminal: true, + }), + ).rejects.toThrow("delivery generation is exhausted"); + + database.execute( + `UPDATE api_command SET attempt_count = ${Number.MAX_SAFE_INTEGER}, delivery_generation = 1, status = 'queued' WHERE id = '${commandId}'`, + ); + await expect( + claimApiCommand({ + claimOwner: "owner-a", + commandId, + database, + deliveryGeneration: 1, + nowMs: 1_000, + }), + ).rejects.toThrow("attempt count is exhausted"); + database.execute(`UPDATE api_command SET status = 'failed' WHERE id = '${commandId}'`); + await expect( + claimApiCommand({ + claimOwner: "owner-a", + commandId, + database, + deliveryGeneration: 1, + nowMs: 1_000, + }), + ).resolves.toEqual({ kind: "terminal" }); + }); + + test("rejects a dedupe key reused by a different command kind", async () => { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + await enqueueApiCommand(bindings, { + dedupeKey: "shared:kind-collision", + kind: "scheduled_maintenance", + payload: { scheduledTime: 1 }, + }); + + await expect( + enqueueApiCommand(bindings, { + dedupeKey: "shared:kind-collision", + kind: "cost_ledger_reconciliation", + payload: { cursor: null, mode: "audit", scheduledTime: 1 }, + }), + ).rejects.toThrow("different command kind"); + expect(queue.sent).toHaveLength(1); + }); + + test("retries busy main and DLQ deliveries until the live lease expires", async () => { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + const commandId = await enqueueApiCommand(bindings, { + dedupeKey: "scheduled:busy-redelivery", + kind: "scheduled_maintenance", + payload: { scheduledTime: 1 }, + }); + await claimCommand({ + claimOwner: "live-owner", + commandId, + database, + nowMs: 1_000, + }); + const before = await database + .app() + .select() + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, commandId)) + .get(); + const recorded = createRecordedQueueMessage({ + body: { commandId, deliveryGeneration: 1 }, + }); + const deadLetter = createRecordedQueueMessage({ + body: { commandId, deliveryGeneration: 1 }, + }); + + await processApiCommandMessage(bindings, recorded.message, () => 2_000); + await processApiCommandDeadLetterMessage(bindings, deadLetter.message, () => 2_000); + + expect(recorded.recorded).toEqual([{ delaySeconds: 299, type: "retry" }]); + expect(deadLetter.recorded).toEqual([{ delaySeconds: 299, type: "retry" }]); + await expect( + database + .app() + .select() + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, commandId)) + .get(), + ).resolves.toEqual(before); + }); + + test("acks old-generation main and DLQ deliveries without touching the live successor", async () => { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + const commandId = await enqueueApiCommand(bindings, { + dedupeKey: "scheduled:old-deliveries", + kind: "scheduled_maintenance", + payload: { scheduledTime: 1 }, + }); + const first = await claimCommand({ + claimOwner: "owner-a", + commandId, + database, + nowMs: 1_000, + }); + await markApiCommandFailed({ + claim: first, + database, + errorCode: "failed", + errorMessage: "failed", + nowMs: 2_000, + }); + const admission = await admitApiCommand(bindings, { + dedupeKey: "scheduled:old-deliveries", + kind: "scheduled_maintenance", + payload: { scheduledTime: 2 }, + retryTerminal: true, + }); + const successor = await claimCommand({ + claimOwner: "owner-b", + commandId, + database, + deliveryGeneration: 2, + nowMs: 3_000, + }); + expect(admission.shouldDeliver).toBe(true); + const before = await database + .app() + .select() + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, commandId)) + .get(); + const main = createRecordedQueueMessage({ + body: { commandId, deliveryGeneration: 1 }, + }); + const deadLetter = createRecordedQueueMessage({ + body: { commandId, deliveryGeneration: 1 }, + }); + + await processApiCommandMessage(bindings, main.message, () => 4_000); + await processApiCommandDeadLetterMessage(bindings, deadLetter.message, () => 4_000); + + expect(successor.deliveryGeneration).toBe(2); + expect(main.recorded).toEqual([{ type: "ack" }]); + expect(deadLetter.recorded).toEqual([{ type: "ack" }]); + await expect( + database + .app() + .select() + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, commandId)) + .get(), + ).resolves.toEqual(before); + }); + + test("uses a new owner for takeover and stops the stale attempt before its effect", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "activating" }); + const firstEntered = Promise.withResolvers(); + const releaseFirst = Promise.withResolvers(); + const owners: string[] = []; + let effectCount = 0; + let currentTimeMs = Date.now(); + const first = createRecordedQueueMessage({ + body: fixture.body, + id: "same-queue-delivery", + }); + const second = createRecordedQueueMessage({ + body: fixture.body, + id: "same-queue-delivery", + }); + + const firstTask = processAppWorkflowDeliveryForTest( + fixture.bindings, + first.message, + () => currentTimeMs, + { + dispatchAppDeploymentRun: async (_bindings, _payload, authority) => { + owners.push(authority.claimOwner); + await authority.requireOwnership(); + await registerAppDeploymentCandidateForTest(fixture, authority, currentTimeMs); + await authority.requireOwnership(); + firstEntered.resolve(); + await releaseFirst.promise; + await authority.requireOwnership(); + effectCount += 1; + return { + errorCode: "stale_effect", + errorMessage: "The stale attempt must not reach this effect.", + kind: "terminal_failure", + runId: fixture.runId, + }; + }, + }, + ); + await firstEntered.promise; + currentTimeMs += API_COMMAND_LEASE_MS; + + try { + await processAppWorkflowDeliveryForTest( + fixture.bindings, + second.message, + () => currentTimeMs, + { + dispatchAppDeploymentRun: async (_bindings, _payload, authority) => { + owners.push(authority.claimOwner); + await authority.requireOwnership(); + effectCount += 1; + return { + errorCode: "takeover_failure", + errorMessage: "The takeover owns the terminal result.", + kind: "terminal_failure", + runId: fixture.runId, + }; + }, + }, + ); + } finally { + releaseFirst.resolve(); + await firstTask; + } + + expect(owners).toHaveLength(2); + expect(owners[0]).not.toBe(owners[1]); + expect(effectCount).toBe(1); + expect(second.recorded).toEqual([{ type: "ack" }]); + expect(first.recorded).toEqual([{ delaySeconds: 30, type: "retry" }]); + await expect( + fixture.database + .app() + .select({ + activeScriptName: appDeploymentsTable.activeScriptName, + lastSuccessfulUrl: appDeploymentsTable.lastSuccessfulUrl, + }) + .from(appDeploymentsTable) + .where(eq(appDeploymentsTable.id, fixture.deploymentId)) + .get(), + ).resolves.toEqual({ activeScriptName: null, lastSuccessfulUrl: null }); + await expect( + fixture.database + .app() + .select({ + errorCode: appDeploymentRunsTable.errorCode, + status: appDeploymentRunsTable.status, + }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, fixture.runId)) + .get(), + ).resolves.toEqual({ errorCode: "takeover_failure", status: "failed" }); + }); + + test("latches an uncertain renewal error before the handler can continue", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "activating" }); + let failNextRenewal = false; + let failedRenewalCount = 0; + const failingDatabase = new Proxy(fixture.database, { + get(target, property, receiver) { + if (property === "prepare") { + return (query: string) => { + if ( + failNextRenewal && + query.startsWith('update "api_command" set "claim_expires_at" = ?') + ) { + failNextRenewal = false; + failedRenewalCount += 1; + throw new Error("renewal response was lost"); + } + return target.prepare(query); + }; + } + const value = Reflect.get(target, property, receiver) as unknown; + return typeof value === "function" ? value.bind(target) : value; + }, + }) as D1Database; + const bindings = { ...fixture.bindings, DB: failingDatabase }; + const recorded = createRecordedQueueMessage({ body: fixture.body }); + let effectCount = 0; + + await processAppWorkflowDeliveryForTest(bindings, recorded.message, () => 1_000, { + dispatchAppDeploymentRun: async (_bindings, _payload, authority) => { + failNextRenewal = true; + await authority.requireOwnership().catch(() => undefined); + await authority.requireOwnership(); + effectCount += 1; + return { + errorCode: "unreachable", + errorMessage: "The handler must stop after an uncertain renewal.", + kind: "terminal_failure", + runId: fixture.runId, + }; + }, + }); + + expect(failedRenewalCount).toBe(1); + expect(effectCount).toBe(0); + expect(recorded.recorded).toEqual([{ delaySeconds: 30, type: "retry" }]); + await expect( + fixture.database + .app() + .select({ + claimOwner: apiCommandsTable.claimOwner, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ claimOwner: null, status: "queued" }); + }); + + test("terminalizes a timed-out App command before a fourth dispatch effect", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "activating" }); + fixture.database.execute(` + UPDATE api_command + SET attempt_count = ${APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS}, + claim_expires_at = 1, + claim_owner = 'timed-out-attempt', + status = 'running' + WHERE id = '${fixture.commandId}' + `); + const recorded = createRecordedQueueMessage({ body: fixture.body }); + let dispatchCount = 0; + + await processAppWorkflowDeliveryForTest(fixture.bindings, recorded.message, () => 2, { + dispatchAppDeploymentRun: async () => { + dispatchCount += 1; + return { kind: "skipped" }; + }, + }); + + expect(dispatchCount).toBe(0); + expect(recorded.recorded).toEqual([{ type: "ack" }]); + await expect( + fixture.database + .app() + .select({ + errorCode: appDeploymentRunsTable.errorCode, + status: appDeploymentRunsTable.status, + }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, fixture.runId)) + .get(), + ).resolves.toEqual({ + errorCode: APP_DEPLOYMENT_RUN_DISPATCH_RETRY_EXHAUSTED_CODE, + status: "failed", + }); + await expect( + fixture.database + .app() + .select({ + attemptCount: apiCommandsTable.attemptCount, + lastErrorCode: apiCommandsTable.lastErrorCode, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ + attemptCount: APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS + 1, + lastErrorCode: APP_DEPLOYMENT_RUN_DISPATCH_RETRY_EXHAUSTED_CODE, + status: "failed", + }); + }); + + test("commits App deployment success and command success in one owned batch", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "activating" }); + const currentTimeMs = Date.now(); + let activeScriptName = ""; + const targetUrl = "https://api-command-test.example.com"; + const recorded = createRecordedQueueMessage({ body: fixture.body }); + + await processAppWorkflowDeliveryForTest( + fixture.bindings, + recorded.message, + () => currentTimeMs, + { + dispatchAppDeploymentRun: async (_bindings, _payload, authority) => { + await authority.requireOwnership(); + activeScriptName = await registerAppDeploymentCandidateForTest( + fixture, + authority, + currentTimeMs, + true, + ); + await authority.requireOwnership(); + return { + activeScriptName, + deploymentId: fixture.deploymentId, + externalDeploymentId: "external-deployment", + externalProjectId: "external-project", + externalVersionId: "external-version", + kind: "succeeded", + runId: fixture.runId, + url: targetUrl, + }; + }, + }, + ); + + await expect( + fixture.database + .app() + .select({ + activeScriptName: appDeploymentsTable.activeScriptName, + lastSuccessfulUrl: appDeploymentsTable.lastSuccessfulUrl, + }) + .from(appDeploymentsTable) + .where(eq(appDeploymentsTable.id, fixture.deploymentId)) + .get(), + ).resolves.toEqual({ activeScriptName, lastSuccessfulUrl: targetUrl }); + await expect( + fixture.database + .app() + .select({ + externalDeploymentId: appDeploymentRunsTable.externalDeploymentId, + externalProjectId: appDeploymentRunsTable.externalProjectId, + externalVersionId: appDeploymentRunsTable.externalVersionId, + status: appDeploymentRunsTable.status, + targetScriptName: appDeploymentRunsTable.targetScriptName, + url: appDeploymentRunsTable.url, + }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, fixture.runId)) + .get(), + ).resolves.toEqual({ + externalDeploymentId: "external-deployment", + externalProjectId: "external-project", + externalVersionId: "external-version", + status: "success", + targetScriptName: activeScriptName, + url: targetUrl, + }); + await expect( + fixture.database + .app() + .select({ + claimOwner: apiCommandsTable.claimOwner, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ claimOwner: null, status: "succeeded" }); + expect(recorded.recorded).toEqual([{ type: "ack" }]); + }); + + test("does not publish an App deployment outcome for a different candidate", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "activating" }); + const currentTimeMs = Date.now(); + let persistedCandidate = ""; + let reportedCandidate = ""; + const recorded = createRecordedQueueMessage({ body: fixture.body }); + + await processAppWorkflowDeliveryForTest( + fixture.bindings, + recorded.message, + () => currentTimeMs, + { + dispatchAppDeploymentRun: async (_bindings, _payload, authority) => { + await authority.requireOwnership(); + persistedCandidate = await registerAppDeploymentCandidateForTest( + fixture, + authority, + currentTimeMs, + true, + ); + reportedCandidate = appDeploymentCandidateScriptName({ + ...authority, + attemptCount: authority.attemptCount + 1, + }); + await authority.requireOwnership(); + return { + activeScriptName: reportedCandidate, + deploymentId: fixture.deploymentId, + externalDeploymentId: "external-deployment", + externalProjectId: "external-project", + externalVersionId: "external-version", + kind: "succeeded", + runId: fixture.runId, + url: "https://api-command-test.example.com", + }; + }, + }, + ); + + await expect( + fixture.database + .app() + .select({ + activeScriptName: appDeploymentsTable.activeScriptName, + lastSuccessfulUrl: appDeploymentsTable.lastSuccessfulUrl, + }) + .from(appDeploymentsTable) + .where(eq(appDeploymentsTable.id, fixture.deploymentId)) + .get(), + ).resolves.toEqual({ activeScriptName: null, lastSuccessfulUrl: null }); + await expect( + fixture.database + .app() + .select({ + status: appDeploymentRunsTable.status, + targetScriptName: appDeploymentRunsTable.targetScriptName, + url: appDeploymentRunsTable.url, + }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, fixture.runId)) + .get(), + ).resolves.toEqual({ + status: "activating", + targetScriptName: persistedCandidate, + url: null, + }); + await expect( + fixture.database + .app() + .select({ status: apiCommandsTable.status }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ status: "running" }); + expect(recorded.recorded).toEqual([{ delaySeconds: 30, type: "retry" }]); + }); + + test("preserves a pre-existing App failure while terminalizing its command", async () => { + const fixture = await createAppDeploymentCommandFixture({ + errorCode: "existing_failure", + errorMessage: "The deployment was already failed elsewhere.", + runStatus: "failed", + }); + const recorded = createRecordedQueueMessage({ body: fixture.body }); + + await processAppWorkflowDeliveryForTest( + fixture.bindings, + recorded.message, + () => nowMsForTest() + 1, + { + dispatchAppDeploymentRun: async (_bindings, _payload, authority) => { + await authority.requireOwnership(); + return { + errorCode: "later_failure", + errorMessage: "This error must not replace the business terminal state.", + kind: "terminal_failure", + runId: fixture.runId, + }; + }, + }, + ); + + await expect( + fixture.database + .app() + .select({ + errorCode: appDeploymentRunsTable.errorCode, + errorMessage: appDeploymentRunsTable.errorMessage, + status: appDeploymentRunsTable.status, + }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, fixture.runId)) + .get(), + ).resolves.toEqual({ + errorCode: "existing_failure", + errorMessage: "The deployment was already failed elsewhere.", + status: "failed", + }); + await expect( + fixture.database + .app() + .select({ + lastErrorCode: apiCommandsTable.lastErrorCode, + lastErrorMessage: apiCommandsTable.lastErrorMessage, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ + lastErrorCode: "existing_failure", + lastErrorMessage: "The deployment was already failed elsewhere.", + status: "failed", + }); + expect(recorded.recorded).toEqual([{ type: "ack" }]); + }); + + test("fails an App command closed when its terminal run is missing", async () => { + const fixture = await createAppDeploymentCommandFixture({ runStatus: "activating" }); + await fixture.database + .app() + .delete(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, fixture.runId)) + .run(); + const recorded = createRecordedQueueMessage({ body: fixture.body }); + + await processAppWorkflowDeliveryForTest( + fixture.bindings, + recorded.message, + () => nowMsForTest() + 1, + { + dispatchAppDeploymentRun: async (_bindings, _payload, authority) => { + await authority.requireOwnership(); + return { + errorCode: "missing_run_failure", + errorMessage: "The missing run cannot remain queued forever.", + kind: "terminal_failure", + runId: fixture.runId, + }; + }, + }, + ); + + await expect( + fixture.database + .app() + .select({ + lastErrorCode: apiCommandsTable.lastErrorCode, + lastErrorMessage: apiCommandsTable.lastErrorMessage, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, fixture.commandId)) + .get(), + ).resolves.toEqual({ + lastErrorCode: "missing_run_failure", + lastErrorMessage: "The missing run cannot remain queued forever.", + status: "failed", + }); + expect(recorded.recorded).toEqual([{ type: "ack" }]); + }); + + test("retries a DLQ message when its ledger claim cannot be persisted", async () => { + const database = await createPublicHttpContractDatabase(); + const queue = createApiCommandQueueStub(); + const baseBindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + const commandId = await enqueueApiCommand(baseBindings, { + dedupeKey: "scheduled:dlq-database-error", + kind: "scheduled_maintenance", + payload: { scheduledTime: 1 }, + }); + const failingDatabase = new Proxy(database, { + get(target, property, receiver) { + if (property === "prepare") { + return () => { + throw new Error("database unavailable"); + }; + } + const value = Reflect.get(target, property, receiver) as unknown; + return typeof value === "function" ? value.bind(target) : value; + }, + }) as D1Database; + const bindings = { ...baseBindings, DB: failingDatabase }; + const recorded = createRecordedQueueMessage({ + body: { commandId, deliveryGeneration: 1 }, + }); + + await processApiCommandDeadLetterMessage(bindings, recorded.message, nowMsForTest); + + expect(recorded.recorded).toEqual([{ delaySeconds: 30, type: "retry" }]); + }); }); diff --git a/apps/api/tests/api-driver-boundary-fixtures.ts b/apps/api/tests/api-driver-boundary-fixtures.ts index d0d74336..cd440dd3 100644 --- a/apps/api/tests/api-driver-boundary-fixtures.ts +++ b/apps/api/tests/api-driver-boundary-fixtures.ts @@ -173,7 +173,9 @@ export function createDriverEvent(value: object): DriverEvent { createId: () => API_DRIVER_BOUNDARY_IDS.runtimeEvent, driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, occurredAt: "1970-01-01T00:00:00.010Z", + runtimeId: "openai-runtime", sessionId: API_DRIVER_BOUNDARY_IDS.session, + traceId: "trace-1", }, value, ); @@ -195,6 +197,7 @@ export function createRuntimeSessionLink(): RuntimeSessionLink { sandboxId: API_DRIVER_BOUNDARY_IDS.sandbox, sandboxKind: "cattle", sandboxSubjectKind: "session", + runtimeId: "openai-runtime", sessionId: API_DRIVER_BOUNDARY_IDS.session, sessionRunId: API_DRIVER_BOUNDARY_IDS.sessionRun, sessionRunStatus: "running", diff --git a/apps/api/tests/api-driver-boundary.test.ts b/apps/api/tests/api-driver-boundary.test.ts index 10eb2e48..3cc1e759 100644 --- a/apps/api/tests/api-driver-boundary.test.ts +++ b/apps/api/tests/api-driver-boundary.test.ts @@ -1,22 +1,32 @@ import { describe, expect, test } from "bun:test"; -import { readFileSync } from "node:fs"; -import { EventType, MOSOO_CUSTOM_EVENT } from "@mosoo/ag-ui-session"; +import { + EventType, + MOSOO_CUSTOM_EVENT, + applyAgUiEventToSessionLiveState, +} from "@mosoo/ag-ui-session"; import { DRIVER_CONTROL_PORT_MAX, DRIVER_CONTROL_PORT_MIN, DRIVER_PROTOCOL_VERSION, parseDriverBootPayloadJson, } from "@mosoo/agent-driver/boot"; +import { + RUNTIME_EVENT_KINDS as DRIVER_RUNTIME_EVENT_KINDS, + RUNTIME_EVENT_SCHEMA_VERSION as DRIVER_RUNTIME_EVENT_SCHEMA_VERSION, + toRuntimeEventInput as toDriverRuntimeEventInput, +} from "@mosoo/agent-driver/events"; import { createDefaultAgentBuiltInTools } from "@mosoo/contracts/agent"; import { PLATFORM_ID_FIXTURES } from "@mosoo/id/testing"; -import { RUNTIME_EVENT_SCHEMA_VERSION, createRuntimeEvent } from "@mosoo/runtime-events"; +import { + RUNTIME_EVENT_KINDS, + RUNTIME_EVENT_SCHEMA_VERSION, + createRuntimeEvent, +} from "@mosoo/runtime-events"; import { getDriverControlPort } from "../src/modules/runtime/domain/sandbox-layout"; -import { - assertRuntimeEventMatchesDriverEnvelope, - assertRuntimeEventMatchesDriverLink, -} from "../src/modules/runtime/infrastructure/driver-instance/event-link-assertion"; +import { canonicalizeDriverEventEnvelope } from "../src/modules/runtime/infrastructure/driver-instance/driver-event-canonicalization"; +import { assertRuntimeEventMatchesDriverLink } from "../src/modules/runtime/infrastructure/driver-instance/event-link-assertion"; import { createBaseLiveState, readPermissionRequestViews, @@ -24,7 +34,10 @@ import { } from "../src/modules/runtime/infrastructure/driver-instance/event-projection"; import { appRuntimeDriverEvents } from "../src/modules/runtime/infrastructure/driver-instance/events"; import { readNativeResumeRef } from "../src/modules/runtime/infrastructure/driver-instance/native-resume-ref-event"; -import { parseDriverEventBatchInput } from "../src/modules/runtime/infrastructure/driver-instance/rpc-wire"; +import { + parseDriverEventBatchInput, + parseDriverEventBatchOutput, +} from "../src/modules/runtime/infrastructure/driver-instance/rpc-wire"; import { createDriverBootPayload, verifyRuntimeActionToken, @@ -56,10 +69,6 @@ const artifactPaths = { python: ["/workspace/.mosoo/environment-artifacts/artifact/python/site-packages"], }; -function readText(path: string): string { - return readFileSync(new URL(path, import.meta.url), "utf8"); -} - describe("API to driver boundary", () => { test("assigns driver control ports inside the sandbox image contract", () => { const port = getDriverControlPort("driver-01KRZRFGXAA788FW1GDBT7F0EZ"); @@ -72,195 +81,6 @@ describe("API to driver boundary", () => { expect(AGENT_DRIVER_PROCESS_COMMAND).toBe("agent-driver"); }); - test("uses the agent-driver runtime contract for runtime selection", () => { - const runtimeConfig = readText("../src/modules/runtime/domain/runtime-config.ts"); - const agentConfig = readText( - "../src/modules/agents/application/agent-versioned-config.service.ts", - ); - const nativeResumeRef = readText( - "../src/modules/runtime/infrastructure/native-resume-ref.repository.ts", - ); - const nativeResumeRefEvent = readText( - "../src/modules/runtime/infrastructure/driver-instance/native-resume-ref-event.ts", - ); - - expect(runtimeConfig).toContain('from "@mosoo/agent-driver/runtime"'); - expect(runtimeConfig).not.toContain('from "@mosoo/driver-protocol"'); - expect(agentConfig).toContain('from "@mosoo/agent-driver/runtime"'); - expect(agentConfig).not.toContain('from "@mosoo/driver-protocol"'); - expect(nativeResumeRef).toContain('from "@mosoo/agent-driver/runtime"'); - expect(nativeResumeRef).not.toContain('from "@mosoo/driver-protocol"'); - expect(nativeResumeRefEvent).toContain('from "@mosoo/agent-driver/runtime"'); - expect(nativeResumeRefEvent).not.toContain('from "@mosoo/driver-protocol"'); - }); - - test("uses agent-driver boot constants for process startup", () => { - const bootToken = readText("../src/modules/runtime/infrastructure/runtime-boot-token.ts"); - const provisioning = readText( - "../src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-provisioning.service.ts", - ); - const driverRecord = readText( - "../src/modules/runtime/infrastructure/driver-instance/driver-instance-record.repository.ts", - ); - const sandboxLayout = readText("../src/modules/runtime/domain/sandbox-layout.ts"); - - expect(bootToken).toContain('from "@mosoo/agent-driver/boot"'); - expect(bootToken).toContain("DRIVER_PROTOCOL_VERSION"); - expect(bootToken).not.toContain('from "@mosoo/driver-protocol"'); - expect(provisioning).toContain('from "@mosoo/agent-driver/boot"'); - expect(provisioning).toContain("DRIVER_BOOT_PAYLOAD_FILE_ENV_NAME"); - expect(provisioning).toContain("const bootPayload = createDriverBootPayload"); - expect(provisioning).toContain("JSON.stringify(bootPayload)"); - expect(driverRecord).toContain('from "@mosoo/agent-driver/boot"'); - expect(driverRecord).toContain("DRIVER_PROTOCOL_VERSION"); - expect(sandboxLayout).toContain('from "@mosoo/agent-driver/boot"'); - expect(sandboxLayout).toContain("DRIVER_CONTROL_PORT_COUNT"); - expect(sandboxLayout).toContain("DRIVER_CONTROL_PORT_MIN"); - }); - - test("uses agent-driver sandbox path contracts", () => { - const runtimeProfile = readText("../src/modules/runtime/application/agent-runtime-profile.ts"); - const subjectPlatform = readText( - "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-platform.ts", - ); - - expect(runtimeProfile).toContain('from "@mosoo/agent-driver/paths"'); - expect(subjectPlatform).toContain('from "@mosoo/agent-driver/paths"'); - }); - - test("publishes agent-driver event envelope contracts", () => { - const driverPublicEvents = readText("../../driver/src/events.ts"); - const rpc = readText("../src/modules/runtime/infrastructure/driver-instance/rpc.ts"); - const rpcWire = readText("../src/modules/runtime/infrastructure/driver-instance/rpc-wire.ts"); - const ingestion = readText( - "../src/modules/runtime/infrastructure/driver-instance/rpc-event-ingestion-controller.ts", - ); - const projection = readText("../src/modules/runtime/infrastructure/driver-instance/events.ts"); - const receipts = readText( - "../src/modules/runtime/infrastructure/driver-instance/driver-event-receipts.ts", - ); - const replayFilter = readText( - "../src/modules/runtime/infrastructure/driver-instance/runtime-event-replay-filter.ts", - ); - const fixtures = readText("./api-driver-boundary-fixtures.ts"); - - expect(driverPublicEvents).toContain("./protocol/events"); - expect(rpcWire).toContain('from "@mosoo/agent-driver/events"'); - expect(rpcWire).toContain("parseDriverEventEnvelope"); - expect(rpcWire).not.toContain('from "@mosoo/driver-protocol"'); - expect(rpc).not.toContain("toAgentDriverEventEnvelopes"); - expect(ingestion).toContain("state.readProcessedDriverEventReceipts(input.events)"); - expect(ingestion).toContain("state.filterUnprocessedDriverEvents(input.events)"); - expect(projection).toContain('from "@mosoo/agent-driver/events"'); - expect(projection).not.toContain('from "@mosoo/driver-protocol"'); - expect(receipts).toContain('from "@mosoo/agent-driver/events"'); - expect(replayFilter).toContain('from "@mosoo/agent-driver/events"'); - expect(fixtures).toContain('from "@mosoo/agent-driver/events"'); - }); - - test("uses agent-driver ORPC contracts behind the API wire parser", () => { - const rpc = readText("../src/modules/runtime/infrastructure/driver-instance/rpc.ts"); - const controller = readText( - "../src/modules/runtime/infrastructure/driver-instance/rpc-controller.ts", - ); - const command = readText( - "../src/modules/runtime/infrastructure/driver-instance/rpc-command-controller.ts", - ); - const eventIngestion = readText( - "../src/modules/runtime/infrastructure/driver-instance/rpc-event-ingestion-controller.ts", - ); - const handshake = readText( - "../src/modules/runtime/infrastructure/driver-instance/rpc-handshake-controller.ts", - ); - const state = readText("../src/modules/runtime/infrastructure/driver-instance/state.ts"); - const runtimeState = readText( - "../src/modules/runtime/infrastructure/driver-instance/runtime-state.ts", - ); - const runtimeStateStore = readText( - "../src/modules/runtime/infrastructure/driver-instance/runtime-state-store.ts", - ); - const lifecycle = readText( - "../src/modules/runtime/infrastructure/driver-instance/lifecycle.ts", - ); - const handler = readText( - "../src/modules/runtime/infrastructure/driver-instance/rpc-handler.ts", - ); - const rpcWire = readText("../src/modules/runtime/infrastructure/driver-instance/rpc-wire.ts"); - - expect(rpc).toContain('from "@mosoo/agent-driver/orpc"'); - expect(rpc).toContain('from "./rpc-wire"'); - expect(controller).toContain('from "@mosoo/agent-driver/orpc"'); - expect(command).toContain('from "@mosoo/agent-driver/orpc"'); - expect(eventIngestion).toContain('from "@mosoo/agent-driver/orpc"'); - expect(handshake).toContain('from "@mosoo/agent-driver/orpc"'); - expect(state).toContain('from "@mosoo/agent-driver/orpc"'); - expect(runtimeState).toContain('from "@mosoo/agent-driver/orpc"'); - expect(runtimeStateStore).toContain('from "@mosoo/agent-driver/orpc"'); - expect(runtimeStateStore).toContain("parseDriverHelloInput"); - expect(runtimeStateStore).not.toContain('from "@mosoo/driver-protocol"'); - expect(lifecycle).toContain('from "@mosoo/agent-driver/orpc"'); - expect(controller).not.toContain('from "@mosoo/driver-protocol"'); - expect(eventIngestion).not.toContain('from "@mosoo/driver-protocol"'); - expect(rpc).not.toContain('from "@mosoo/driver-protocol"'); - expect(rpcWire).toContain("runtimeOrpcRouter"); - expect(rpcWire).toContain('from "@mosoo/agent-driver/orpc"'); - expect(rpcWire).toContain('from "@mosoo/agent-driver/events"'); - expect(rpcWire).not.toContain('from "@mosoo/driver-protocol"'); - expect(handler).not.toContain('from "@mosoo/driver-protocol"'); - expect(handler).toContain("runtimeOrpcRouter"); - }); - - test("builds session config traces from the agent-driver boot payload", () => { - const dispatchRun = readText( - "../src/modules/runtime/application/session-runs/dispatch-run.service.ts", - ); - const callbackContract = readText( - "../src/modules/runtime/application/execution-plane/driver-boot-payload-prepared.ts", - ); - const configTrace = readText( - "../src/modules/runtime/application/session-definition/session-config-trace-event.ts", - ); - - expect(callbackContract).toContain('from "@mosoo/agent-driver/boot"'); - expect(dispatchRun).toContain("onBootPayloadPrepared: async ({ bootPayload })"); - expect(dispatchRun).not.toContain("toAgentDriverBootPayload(bootPayload)"); - expect(dispatchRun).toContain("buildSessionConfigTraceValue(bootPayload)"); - expect(dispatchRun).toContain("bootPayload.execution.session.mcpServers.length"); - expect(dispatchRun).toContain("bootPayload.execution.provider"); - expect(configTrace).toContain('from "@mosoo/agent-driver/boot"'); - expect(configTrace).not.toContain('from "@mosoo/driver-protocol"'); - }); - - test("owns platform driver snapshots inside the API runtime domain", () => { - const driverSnapshot = readText("../src/modules/runtime/domain/driver-snapshot.ts"); - const runtimeProfile = readText("../src/modules/runtime/application/agent-runtime-profile.ts"); - const executionTypes = readText( - "../src/modules/runtime/application/session-definition/session-execution.types.ts", - ); - const executionSpec = readText( - "../src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-execution-spec.builder.ts", - ); - const sandboxSessionTypes = readText( - "../src/modules/runtime/infrastructure/sandbox-session/sandbox-session.types.ts", - ); - const sandboxConversationCodec = readText( - "../src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session-codec.ts", - ); - const mcpRuntime = readText("../src/modules/mcp/application/mcp-runtime.service.ts"); - const fixtures = readText("./api-driver-boundary-fixtures.ts"); - - expect(driverSnapshot).toContain('from "@mosoo/agent-driver/runtime"'); - expect(driverSnapshot).not.toContain('from "@mosoo/driver-protocol"'); - expect(runtimeProfile).toContain('from "../domain/driver-snapshot"'); - expect(executionTypes).toContain('from "../../domain/driver-snapshot"'); - expect(executionSpec).toContain('from "../../domain/driver-snapshot"'); - expect(sandboxSessionTypes).toContain('from "../../domain/driver-snapshot"'); - expect(sandboxConversationCodec).toContain('from "../../domain/driver-snapshot"'); - expect(sandboxConversationCodec).toContain("readSandboxConversationOriginRecord"); - expect(mcpRuntime).toContain('from "../../runtime/domain/driver-snapshot"'); - expect(fixtures).toContain('from "../src/modules/runtime/domain/driver-snapshot"'); - }); - test("builds a driver execution spec with scoped grants and profile env", async () => { const execution = await buildExecutionSpec(bindings, { builtInTools: [ @@ -479,8 +299,10 @@ describe("API to driver boundary", () => { messageId: "message-1", role: "agent", }, + runtimeId: "openai-runtime", schemaVersion: RUNTIME_EVENT_SCHEMA_VERSION, sessionId: API_DRIVER_BOUNDARY_IDS.session, + traceId: "trace-1", visibility: "participant", }); }); @@ -510,8 +332,264 @@ describe("API to driver boundary", () => { expect(envelope.occurredAt).toBe("1970-01-01T00:00:00.010Z"); }); + test("requires source identity on every accepted driver event receipt", () => { + const receipt = { + eventId: "source-1", + seq: 1, + type: "message.delta", + }; + + expect(parseDriverEventBatchOutput({ accepted: [receipt] })).toEqual({ + accepted: [receipt], + }); + expect(() => + parseDriverEventBatchOutput({ + accepted: [{ seq: 1, type: "message.delta" }], + }), + ).toThrow(); + expect(() => + parseDriverEventBatchOutput({ + accepted: [{ ...receipt, eventId: "" }], + }), + ).toThrow(); + }); + + test("keeps the Driver and API canonical event vocabularies exactly aligned", () => { + expect(RUNTIME_EVENT_SCHEMA_VERSION).toBe(DRIVER_RUNTIME_EVENT_SCHEMA_VERSION); + expect(RUNTIME_EVENT_KINDS).toEqual(DRIVER_RUNTIME_EVENT_KINDS); + }); + + test("admits an envelope built by the production Driver event builder", () => { + const occurredAt = "2026-08-29T00:00:00.000Z"; + const [driverEvent] = toDriverRuntimeEventInput( + { + createId: () => API_DRIVER_BOUNDARY_IDS.runtimeEvent, + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + occurredAt, + runId: API_DRIVER_BOUNDARY_IDS.sessionRun, + sessionId: API_DRIVER_BOUNDARY_IDS.session, + sourceEventId: "driver-builder-message-delta", + } as Parameters[0], + { + kind: "message.delta", + payload: { + contentDelta: "Driver-built payload", + messageId: "driver-builder-message", + role: "agent", + }, + }, + ); + + expect(driverEvent).toBeDefined(); + expect( + parseDriverEventBatchInput({ + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + events: [ + { + event: driverEvent, + eventId: "driver-builder-message-delta", + occurredAt, + }, + ], + }).events, + ).toEqual([ + { + event: driverEvent, + eventId: "driver-builder-message-delta", + occurredAt, + }, + ]); + }); + + test("admits Driver terminal wire fixtures through the independent API boundary", async () => { + const link = createRuntimeSessionLink(); + const occurredAt = "1970-01-01T00:00:01.000Z"; + const fixtures = [ + { + id: "01J0000000000000000000001C", + kind: "message.cancelled", + payload: { messageId: "message-cancelled", role: "agent" }, + }, + { + id: "01J0000000000000000000001D", + kind: "message.failed", + payload: { + error: { + code: "runtime.failed", + details: {}, + message: "Runtime failed.", + retryable: false, + }, + messageId: "message-failed", + role: "agent", + }, + }, + { + id: "01J0000000000000000000001E", + kind: "thought.cancelled", + payload: { thoughtId: "thought-cancelled" }, + }, + { + id: "01J0000000000000000000001F", + kind: "tool.call.updated", + payload: { status: "cancelled", toolCallId: "tool-cancelled" }, + }, + ] as const; + const batch = parseDriverEventBatchInput({ + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + events: fixtures.map(({ id, kind, payload }, index) => ({ + event: { + actor: "driver", + delivery: "lossless", + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + id, + kind, + occurredAt, + origin: "driver", + payload, + runId: API_DRIVER_BOUNDARY_IDS.sessionRun, + runtimeId: "openai-runtime", + schemaVersion: DRIVER_RUNTIME_EVENT_SCHEMA_VERSION, + sessionId: API_DRIVER_BOUNDARY_IDS.session, + traceId: "trace-1", + visibility: "participant", + }, + eventId: `source-terminal-${index}`, + occurredAt, + })), + }); + const projection = await appRuntimeDriverEvents({ DB: new SqliteD1Database() } as ApiBindings, { + currentLiveState: createBaseLiveState({ + callerId: link.callerId, + creatorId: link.creatorId, + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + sessionId: link.sessionId, + }), + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + events: batch.events, + link, + }); + + expect( + projection.runtimeEvents.map(({ event }) => ({ + id: event.id, + kind: event.kind, + payload: event.payload, + schemaVersion: event.schemaVersion, + })), + ).toEqual( + fixtures.map(({ id, kind, payload }) => ({ + id, + kind, + payload, + schemaVersion: RUNTIME_EVENT_SCHEMA_VERSION, + })), + ); + expect(projection.sessionDeliveryEvents.map(({ event }) => event.type)).toEqual([ + EventType.TEXT_MESSAGE_END, + EventType.TEXT_MESSAGE_END, + EventType.REASONING_MESSAGE_END, + EventType.CUSTOM, + ]); + }); + + test("replays tool snapshots and deltas through the same live-state reducer", async () => { + const link = createRuntimeSessionLink(); + const currentLiveState = createBaseLiveState({ + callerId: link.callerId, + creatorId: link.creatorId, + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + sessionId: link.sessionId, + }); + const updates = [ + { + parentMessageId: "assistant-1", + rawInputDelta: '{"cmd":', + rawOutputDelta: "partial", + status: "running", + title: "Shell", + toolCallId: "tool-1", + }, + { + rawInput: '{"cmd":"ls"', + rawOutput: "snapshot", + status: "running", + toolCallId: "tool-1", + }, + { + rawInputDelta: ',"tail":true}', + rawOutputDelta: " tail", + status: "completed", + toolCallId: "tool-1", + }, + ] as const; + const batch = parseDriverEventBatchInput({ + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + events: updates.map((payload, index) => ({ + event: createDriverEvent({ + kind: "tool.call.updated", + payload, + runId: API_DRIVER_BOUNDARY_IDS.sessionRun, + }), + eventId: `source-tool-${index}`, + occurredAt: `1970-01-01T00:00:0${index + 1}.000Z`, + })), + }); + const projection = await appRuntimeDriverEvents({ DB: new SqliteD1Database() } as ApiBindings, { + currentLiveState, + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + events: batch.events, + link, + }); + const replayed = projection.sessionDeliveryEvents.reduce( + (state, record) => applyAgUiEventToSessionLiveState(state, record.event), + currentLiveState, + ); + + expect({ ...projection.nextLiveState, updatedAt: null }).toEqual({ + ...replayed, + updatedAt: null, + }); + expect(replayed.messages).toEqual([ + expect.objectContaining({ + id: "assistant-1", + segments: [ + { + argsText: '{"cmd":"ls","tail":true}', + kind: "tool_use", + path: null, + runId: API_DRIVER_BOUNDARY_IDS.sessionRun, + tool: "Shell", + toolCallId: "tool-1", + }, + { + kind: "tool_result", + output: "snapshot tail", + runId: API_DRIVER_BOUNDARY_IDS.sessionRun, + tool: "Shell", + toolCallId: "tool-1", + }, + ], + }), + ]); + expect( + projection.sessionDeliveryEvents.map(({ event }) => + event.type === EventType.CUSTOM ? event.name : event.type, + ), + ).toEqual(Array(3).fill(MOSOO_CUSTOM_EVENT.sessionToolUpdated.name)); + }); + test("apps admitted driver wire events into API runtime and viewer events", async () => { const link = createRuntimeSessionLink(); + const permissionRequest = { + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + rawInput: "pwd", + requestId: "permission-1", + runId: API_DRIVER_BOUNDARY_IDS.sessionRun, + title: "Allow shell command?", + toolCallId: "tool-1", + toolKind: "shell", + }; const permissionRequested = createDriverEvent({ kind: "permission.requested", payload: { @@ -537,14 +615,23 @@ describe("API to driver boundary", () => { }, ], }); + const baseLiveState = createBaseLiveState({ + callerId: link.callerId, + creatorId: link.creatorId, + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + sessionId: link.sessionId, + }); const projection = await appRuntimeDriverEvents({ DB: new SqliteD1Database() } as ApiBindings, { - currentLiveState: createBaseLiveState({ - callerId: link.callerId, - creatorId: link.creatorId, - driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, - sessionId: link.sessionId, - }), + currentLiveState: { + ...baseLiveState, + lifecycle: "RUNNING", + run: { + ...baseLiveState.run, + id: API_DRIVER_BOUNDARY_IDS.sessionRun, + status: "running", + }, + }, driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, events: batch.events, link, @@ -564,23 +651,22 @@ describe("API to driver boundary", () => { expect(projection.liveStateChanged).toBe(true); expect(projection.sessionDeliveryEvents[0]?.event).toMatchObject({ name: MOSOO_CUSTOM_EVENT.sessionPermissionsUpdated.name, + type: EventType.CUSTOM, value: { - permissionRequests: [ - { - driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, - rawInput: "pwd", - requestId: "permission-1", - runId: API_DRIVER_BOUNDARY_IDS.sessionRun, - title: "Allow shell command?", - toolCallId: "tool-1", - toolKind: "shell", - }, - ], + permissionRequest, + permissionRequests: [], + }, + }); + expect(projection.nextLiveState).toMatchObject({ + permissionRequests: [permissionRequest], + run: { + id: API_DRIVER_BOUNDARY_IDS.sessionRun, + status: "waiting_input", }, }); }); - test("adds failed tool result delivery before terminal run update", async () => { + test("accepts a production Driver terminal envelope without inventing failed tool output", async () => { const link = createRuntimeSessionLink(); const baseLiveState = createBaseLiveState({ callerId: link.callerId, @@ -588,26 +674,51 @@ describe("API to driver boundary", () => { driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, sessionId: link.sessionId, }); - const runFailed = createDriverEvent({ - kind: "run.failed", - payload: { - error: { - code: "runtime.failed", - message: "Runtime driver control socket is not connected.", + const pendingToolUse = { + argsText: '{"cmd":"pwd"}', + kind: "tool_use", + path: null, + tool: "Shell", + toolCallId: "tool-1", + } as const; + const runError = { + code: "runtime.failed", + details: {}, + message: "Runtime driver control socket is not connected.", + retryable: false, + }; + const transportSourceId = "driver-builder-run-failed"; + const [driverRunFailed] = toDriverRuntimeEventInput( + { + createId: () => API_DRIVER_BOUNDARY_IDS.runtimeEvent, + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + occurredAt: "1970-01-01T00:00:01.000Z", + runId: API_DRIVER_BOUNDARY_IDS.sessionRun, + runtimeId: "openai-runtime", + sessionId: API_DRIVER_BOUNDARY_IDS.session, + sourceEventId: transportSourceId, + } as Parameters[0], + { + kind: "run.failed", + payload: { + error: runError, + recoverable: false, }, }, - runId: API_DRIVER_BOUNDARY_IDS.sessionRun, - }); + ); const batch = parseDriverEventBatchInput({ driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, events: [ { - event: runFailed, - eventId: "source-run-failed", + event: driverRunFailed, + eventId: transportSourceId, occurredAt: "1970-01-01T00:00:01.000Z", }, ], }); + const canonicalEnvelope = canonicalizeDriverEventEnvelope(batch.events[0], { + traceId: link.traceId, + }); const projection = await appRuntimeDriverEvents({ DB: new SqliteD1Database() } as ApiBindings, { currentLiveState: { @@ -620,15 +731,7 @@ describe("API to driver boundary", () => { id: "assistant-1", plan: [], role: "assistant", - segments: [ - { - argsText: '{"cmd":"pwd"}', - kind: "tool_use", - path: null, - tool: "Shell", - toolCallId: "tool-1", - }, - ], + segments: [pendingToolUse], }, ], run: { @@ -638,35 +741,45 @@ describe("API to driver boundary", () => { }, }, driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, - events: batch.events, + events: [canonicalEnvelope], link, }); const canonicalFailureSourceId = `session-run-terminal:${API_DRIVER_BOUNDARY_IDS.sessionRun}:run.failed`; + expect(canonicalEnvelope).toMatchObject({ + event: { sourceEventId: canonicalFailureSourceId, traceId: link.traceId }, + eventId: transportSourceId, + }); expect(projection.runtimeEvents).toMatchObject([{ sourceEventId: canonicalFailureSourceId }]); expect(projection.sessionDeliveryEvents.map((record) => record.sourceEventId)).toEqual([ canonicalFailureSourceId, - canonicalFailureSourceId, - canonicalFailureSourceId, - ]); - expect(projection.sessionDeliveryEvents.map((record) => record.event.type)).toEqual([ - EventType.TOOL_CALL_RESULT, - EventType.TOOL_CALL_END, - EventType.CUSTOM, ]); expect(projection.sessionDeliveryEvents[0]?.event).toMatchObject({ - content: - "Tool failed before returning a result: Runtime driver control socket is not connected.", - toolCallId: "tool-1", - type: EventType.TOOL_CALL_RESULT, + name: MOSOO_CUSTOM_EVENT.sessionRunUpdated.name, + type: EventType.CUSTOM, + value: { + lifecycle: "IDLE", + run: { + error: runError, + id: API_DRIVER_BOUNDARY_IDS.sessionRun, + status: "failed", + }, + }, }); - expect(projection.nextLiveState?.messages[0]?.segments.at(-1)).toEqual({ - kind: "tool_result", - output: - "Tool failed before returning a result: Runtime driver control socket is not connected.", - tool: "Shell", - toolCallId: "tool-1", + expect(projection.nextLiveState).toMatchObject({ + infra: { + driverInstanceId: null, + lastFailureMessage: runError.message, + lastFailureReason: runError.code, + }, + lifecycle: "IDLE", + run: { + error: runError, + id: API_DRIVER_BOUNDARY_IDS.sessionRun, + status: "failed", + }, }); + expect(projection.nextLiveState?.messages[0]?.segments).toEqual([pendingToolUse]); }); test("rejects legacy event shapes from the driver channel", () => { @@ -693,7 +806,9 @@ describe("API to driver boundary", () => { contentDelta: "wrong session", messageId: "message-1", }, + runtimeId: "openai-runtime", sessionId: "01J0000000000000000000000M", + traceId: "trace-1", }), { driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, @@ -713,7 +828,9 @@ describe("API to driver boundary", () => { startedAt: "1970-01-01T00:00:00.010Z", }, runId: "01J0000000000000000000000P", + runtimeId: "openai-runtime", sessionId: API_DRIVER_BOUNDARY_IDS.session, + traceId: "trace-1", }), { driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, @@ -733,7 +850,9 @@ describe("API to driver boundary", () => { contentDelta: "wrong driver", messageId: "message-1", }, + runtimeId: "openai-runtime", sessionId: API_DRIVER_BOUNDARY_IDS.session, + traceId: "trace-1", }), { driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, @@ -753,7 +872,9 @@ describe("API to driver boundary", () => { contentDelta: "missing run", messageId: "message-1", }, + runtimeId: "openai-runtime", sessionId: API_DRIVER_BOUNDARY_IDS.session, + traceId: "trace-1", }), { driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, @@ -773,7 +894,9 @@ describe("API to driver boundary", () => { messageId: "message-1", }, runId: API_DRIVER_BOUNDARY_IDS.sessionRun, + runtimeId: "openai-runtime", sessionId: API_DRIVER_BOUNDARY_IDS.session, + traceId: "trace-1", }), { driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, @@ -794,7 +917,9 @@ describe("API to driver boundary", () => { messageId: "message-1", }, runId: API_DRIVER_BOUNDARY_IDS.sessionRun, + runtimeId: "openai-runtime", sessionId: API_DRIVER_BOUNDARY_IDS.session, + traceId: "trace-1", }), { driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, @@ -804,42 +929,76 @@ describe("API to driver boundary", () => { ).not.toThrow(); }); - test("rejects canonical driver events whose source id disagrees with the envelope", () => { + test("rejects non-empty task snapshots after run terminal but permits an explicit clear", () => { + const event = (tasks: { taskId: string }[]) => + createRuntimeEvent({ + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + id: API_DRIVER_BOUNDARY_IDS.runtimeEvent, + kind: "agent.tasks.replaced", + occurredAt: "1970-01-01T00:00:00.010Z", + payload: { tasks }, + runId: API_DRIVER_BOUNDARY_IDS.sessionRun, + runtimeId: "openai-runtime", + sessionId: API_DRIVER_BOUNDARY_IDS.session, + traceId: "trace-1", + }); + const terminalLink = { + ...createRuntimeSessionLink(), + sessionRunStatus: "completed" as const, + }; + expect(() => - assertRuntimeEventMatchesDriverEnvelope( - createRuntimeEvent({ - id: API_DRIVER_BOUNDARY_IDS.runtimeEvent, - kind: "message.delta", - occurredAt: "1970-01-01T00:00:00.010Z", - payload: { - contentDelta: "wrong source", - messageId: "message-1", - }, - sessionId: API_DRIVER_BOUNDARY_IDS.session, - sourceEventId: "source-inner", - }), + assertRuntimeEventMatchesDriverLink(event([{ taskId: "stale" }]), { + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + link: terminalLink, + }), + ).toThrow("requires an active session run"); + expect(() => + assertRuntimeEventMatchesDriverLink(event([]), { + driverInstanceId: API_DRIVER_BOUNDARY_IDS.driverInstance, + link: terminalLink, + }), + ).not.toThrow(); + }); + + test("rejects Driver transport source ids that disagree with their envelopes", () => { + expect(() => + canonicalizeDriverEventEnvelope( { + event: createRuntimeEvent({ + id: API_DRIVER_BOUNDARY_IDS.runtimeEvent, + kind: "message.delta", + occurredAt: "1970-01-01T00:00:00.010Z", + payload: { + contentDelta: "wrong source", + messageId: "message-1", + }, + sessionId: API_DRIVER_BOUNDARY_IDS.session, + sourceEventId: "source-inner", + }), eventId: "source-outer", }, + { traceId: null }, ), ).toThrow("Runtime driver event source id does not match the driver envelope."); expect(() => - assertRuntimeEventMatchesDriverEnvelope( - createRuntimeEvent({ - id: "01J0000000000000000000000H", - kind: "message.delta", - occurredAt: "1970-01-01T00:00:00.010Z", - payload: { - contentDelta: "ok", - messageId: "message-1", - }, - sessionId: API_DRIVER_BOUNDARY_IDS.session, - sourceEventId: "source-1", - }), + canonicalizeDriverEventEnvelope( { + event: createRuntimeEvent({ + id: "01J0000000000000000000000H", + kind: "message.delta", + occurredAt: "1970-01-01T00:00:00.010Z", + payload: { + contentDelta: "ok", + messageId: "message-1", + }, + sessionId: API_DRIVER_BOUNDARY_IDS.session, + sourceEventId: "source-1", + }), eventId: "source-1", }, + { traceId: null }, ), ).not.toThrow(); }); diff --git a/apps/api/tests/app-agent-bound-run-revocation.test.ts b/apps/api/tests/app-agent-bound-run-revocation.test.ts index ee480edb..ac8ed7ab 100644 --- a/apps/api/tests/app-agent-bound-run-revocation.test.ts +++ b/apps/api/tests/app-agent-bound-run-revocation.test.ts @@ -15,6 +15,12 @@ import { SessionRunCreationGuardRejectedError, } from "../src/modules/runtime/application/session-run.service"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; +import { + BOUND_DEPLOYMENT_ID, + BOUND_DEPLOYMENT_RUN_ID, + deleteBoundDeployment, + insertBoundDeployment, +} from "./bound-capability-fixtures"; import { PUBLIC_API_TEST_IDS, createPublicHttpContractDatabase, @@ -24,9 +30,6 @@ import { } from "./helpers/public-api-http-test-fixture"; import type { SqliteD1Database } from "./helpers/public-api-http-test-fixture"; -const DEPLOYMENT_ID = "01J0000000000000000000000D"; -const DEPLOYMENT_RUN_ID = "01J0000000000000000000000R"; - const CLAIMS: AppAgentCapabilityClaims = { agentId: PUBLIC_API_TEST_IDS.agent, appId: PUBLIC_API_TEST_IDS.app, @@ -35,48 +38,11 @@ const CLAIMS: AppAgentCapabilityClaims = { expose: "public_thread", name: "Public API Agent", }, - deploymentId: DEPLOYMENT_ID, - deploymentRunId: DEPLOYMENT_RUN_ID, + deploymentId: BOUND_DEPLOYMENT_ID, + deploymentRunId: BOUND_DEPLOYMENT_RUN_ID, exp: Date.now() + 60_000, }; -async function insertDeploymentAuthority(database: SqliteD1Database): Promise { - database.execute(` - CREATE TABLE app_deployment ( - app_id text NOT NULL, - deleted_at integer, - id text PRIMARY KEY NOT NULL - ); - - CREATE TABLE app_deployment_run ( - app_id text NOT NULL, - deployment_id text NOT NULL, - id text PRIMARY KEY NOT NULL, - plan_json text, - status text NOT NULL - ); - - CREATE INDEX app_deployment_run_deployment_id_idx - ON app_deployment_run (deployment_id, id); - `); - - await database - .prepare("INSERT INTO app_deployment (app_id, deleted_at, id) VALUES (?, NULL, ?)") - .bind(PUBLIC_API_TEST_IDS.app, DEPLOYMENT_ID) - .run(); - await database - .prepare( - "INSERT INTO app_deployment_run (app_id, deployment_id, id, plan_json, status) VALUES (?, ?, ?, ?, 'success')", - ) - .bind( - PUBLIC_API_TEST_IDS.app, - DEPLOYMENT_ID, - DEPLOYMENT_RUN_ID, - JSON.stringify({ agentBindings: [CLAIMS.binding] }), - ) - .run(); -} - function revokeDeploymentWhenRunInsertStarts(database: SqliteD1Database): D1Database { let revoked = false; @@ -99,10 +65,7 @@ function revokeDeploymentWhenRunInsertStarts(database: SqliteD1Database): D1Data if (typeof method === "function") { return async (...args: unknown[]) => { revoked = true; - await database - .prepare("UPDATE app_deployment SET deleted_at = ? WHERE id = ?") - .bind(Date.now(), DEPLOYMENT_ID) - .run(); + await deleteBoundDeployment(database); return method.apply(target, args); }; } @@ -178,7 +141,7 @@ describe("bound Agent Run revocation boundary", () => { test("creates a Run while the claimed deployment authority remains current", async () => { const database = await createPublicHttpContractDatabase(); await insertOwnerSession(database); - await insertDeploymentAuthority(database); + await insertBoundDeployment(database, { agentBindings: [CLAIMS.binding] }); const viewer = await getAccountViewer(database, PUBLIC_API_TEST_IDS.ownerAccount); if (viewer === null) { @@ -258,7 +221,7 @@ describe("bound Agent Run revocation boundary", () => { test("does not expose accepted Run provenance to a non-owner", async () => { const database = await createPublicHttpContractDatabase(); await insertOwnerSession(database); - await insertDeploymentAuthority(database); + await insertBoundDeployment(database, { agentBindings: [CLAIMS.binding] }); const owner = await getAccountViewer(database, PUBLIC_API_TEST_IDS.ownerAccount); const nonOwner = await getAccountViewer(database, PUBLIC_API_TEST_IDS.nonOwnerAccount); @@ -358,7 +321,7 @@ describe("bound Agent Run revocation boundary", () => { test("does not insert a Run when deletion commits after preflight authorization", async () => { const database = await createPublicHttpContractDatabase(); await insertOwnerSession(database); - await insertDeploymentAuthority(database); + await insertBoundDeployment(database, { agentBindings: [CLAIMS.binding] }); const viewer = await getAccountViewer(database, PUBLIC_API_TEST_IDS.ownerAccount); if (viewer === null) { diff --git a/apps/api/tests/app-agent-capability-revocation-http.test.ts b/apps/api/tests/app-agent-capability-revocation-http.test.ts index 5e969a43..7b409734 100644 --- a/apps/api/tests/app-agent-capability-revocation-http.test.ts +++ b/apps/api/tests/app-agent-capability-revocation-http.test.ts @@ -6,6 +6,12 @@ import { registerPublicApiRoute } from "../src/adapters/http/routes/public-api-r import { mintAppAgentCapabilityToken } from "../src/modules/public-api/app-agent-capability"; import type { AppAgentCapabilityClaims } from "../src/modules/public-api/app-agent-capability"; import type { ApiBindings, ApiGatewayEnvironment } from "../src/platform/cloudflare/worker-types"; +import { + BOUND_DEPLOYMENT_ID, + BOUND_DEPLOYMENT_RUN_ID, + deleteBoundDeployment, + insertBoundDeployment, +} from "./bound-capability-fixtures"; import { PUBLIC_API_TEST_IDS, createPublicHttpContractDatabase, @@ -14,9 +20,6 @@ import { } from "./helpers/public-api-http-test-fixture"; import type { SqliteD1Database } from "./helpers/public-api-http-test-fixture"; -const DEPLOYMENT_ID = "01J0000000000000000000000D"; -const DEPLOYMENT_RUN_ID = "01J0000000000000000000000R"; - function createBoundAgentRouteTestApp(): Hono { const app = new Hono(); const publicApi = new Hono(); @@ -37,54 +40,13 @@ function capabilityClaims( expose: "public_thread", name: "Public API Agent", }, - deploymentId: DEPLOYMENT_ID, - deploymentRunId: DEPLOYMENT_RUN_ID, + deploymentId: BOUND_DEPLOYMENT_ID, + deploymentRunId: BOUND_DEPLOYMENT_RUN_ID, exp: Date.now() + 60_000, ...overrides, }; } -async function insertDeploymentAuthority( - database: SqliteD1Database, - input: { agentBindings: unknown[]; deletedAt: number | null }, -): Promise { - database.execute(` - CREATE TABLE app_deployment ( - app_id text NOT NULL, - deleted_at integer, - id text PRIMARY KEY NOT NULL - ); - - CREATE TABLE app_deployment_run ( - app_id text NOT NULL, - deployment_id text NOT NULL, - id text PRIMARY KEY NOT NULL, - plan_json text, - status text NOT NULL - ); - - CREATE INDEX app_deployment_run_deployment_id_idx - ON app_deployment_run (deployment_id, id); - `); - - await database - .prepare("INSERT INTO app_deployment (app_id, deleted_at, id) VALUES (?, ?, ?)") - .bind(PUBLIC_API_TEST_IDS.app, input.deletedAt, DEPLOYMENT_ID) - .run(); - await database - .prepare( - "INSERT INTO app_deployment_run (app_id, deployment_id, id, plan_json, status) VALUES (?, ?, ?, ?, ?)", - ) - .bind( - PUBLIC_API_TEST_IDS.app, - DEPLOYMENT_ID, - DEPLOYMENT_RUN_ID, - JSON.stringify({ agentBindings: input.agentBindings }), - "success", - ) - .run(); -} - async function requestBoundAgent( database: D1Database, claims: AppAgentCapabilityClaims, @@ -139,10 +101,7 @@ function revokeDeploymentWhenRunInsertStarts(database: SqliteD1Database): D1Data if (typeof method === "function") { return async (...args: unknown[]) => { revoked = true; - await database - .prepare("UPDATE app_deployment SET deleted_at = ? WHERE id = ?") - .bind(Date.now(), DEPLOYMENT_ID) - .run(); + await deleteBoundDeployment(database); return method.apply(target, args); }; } @@ -168,7 +127,7 @@ async function expectNoSessions(database: SqliteD1Database): Promise { describe("bound Agent capability revocation HTTP boundary", () => { test("rejects a deleted deployment capability before it can create a Session", async () => { const database = await createPublicHttpContractDatabase(); - await insertDeploymentAuthority(database, { + await insertBoundDeployment(database, { agentBindings: [capabilityClaims().binding], deletedAt: Date.now(), }); @@ -219,7 +178,7 @@ describe("bound Agent capability revocation HTTP boundary", () => { test("rejects a capability whose current successful revision removed its binding", async () => { const database = await createPublicHttpContractDatabase(); - await insertDeploymentAuthority(database, { + await insertBoundDeployment(database, { agentBindings: [], deletedAt: null, }); @@ -238,7 +197,7 @@ describe("bound Agent capability revocation HTTP boundary", () => { test("cleans up the new Session when deletion wins the final Run creation race", async () => { const database = await createPublicHttpContractDatabase(); - await insertDeploymentAuthority(database, { + await insertBoundDeployment(database, { agentBindings: [capabilityClaims().binding], deletedAt: null, }); diff --git a/apps/api/tests/app-deployment-cloudflare-client.test.ts b/apps/api/tests/app-deployment-cloudflare-client.test.ts index 9721ee3b..58a5b291 100644 --- a/apps/api/tests/app-deployment-cloudflare-client.test.ts +++ b/apps/api/tests/app-deployment-cloudflare-client.test.ts @@ -1,38 +1,81 @@ import { describe, expect, test } from "bun:test"; -import { createWorkerModuleUpload } from "../src/modules/apps/application/app-deployment-cloudflare-client"; +import { + createStaticAssetsUploadForm, + createWorkerModuleUpload, + validateStaticAssetsManifest, +} from "../src/modules/apps/application/app-deployment-cloudflare-client"; describe("app deployment Cloudflare client", () => { - test("uses the module name and metadata as multipart part names", async () => { + test("builds an atomic WfP module upload", async () => { const scriptContent = "export default { fetch() {} };"; + const tags = ["mosoo-managed", "d-01j0000000000000000000000d"]; const upload = createWorkerModuleUpload({ compatibilityDate: "2026-07-14", mainModuleName: "worker.js", scriptContent, scriptName: "example", + tags, vars: { MOSOO_AGENT_URL: "https://example.com/bound/token" }, }); - const modulePart = upload.get("worker.js"); + const modulePart = upload.files[0]; - expect(upload).toBeInstanceOf(FormData); expect(modulePart).toBeInstanceOf(File); - expect((modulePart as File).name).toBe("worker.js"); - expect((modulePart as File).type).toBe("application/javascript+module"); - expect(await (modulePart as File).text()).toBe(scriptContent); - expect(upload.get("files")).toBeNull(); - expect(upload.get("metadata")).toBe( - JSON.stringify({ - bindings: [ - { - name: "MOSOO_AGENT_URL", - text: "https://example.com/bound/token", - type: "plain_text", - }, - ], - compatibility_date: "2026-07-14", - main_module: "worker.js", - }), - ); + expect(modulePart.name).toBe("worker.js"); + expect(modulePart.type).toBe("application/javascript+module"); + expect(await modulePart.text()).toBe(scriptContent); + expect(upload.metadata).toEqual({ + bindings: [ + { + name: "MOSOO_AGENT_URL", + text: "https://example.com/bound/token", + type: "secret_text", + }, + ], + compatibility_date: "2026-07-14", + main_module: "worker.js", + tags, + }); + }); + + test("builds the official assets-only multipart upload without a wrapper module", () => { + const form = createStaticAssetsUploadForm({ + compatibilityDate: "2026-07-14", + completionToken: "completion-token", + headers: "/assets/*\n Cache-Control: public\n", + redirects: "/old /new 301\n", + tags: ["mosoo-managed"], + }); + + expect([...form.keys()]).toEqual(["metadata"]); + const metadata = form.get("metadata"); + + if (typeof metadata !== "string") { + throw new Error("Expected Static Assets upload metadata."); + } + + expect(JSON.parse(metadata)).toEqual({ + assets: { + config: { + _headers: "/assets/*\n Cache-Control: public\n", + _redirects: "/old /new 301\n", + }, + jwt: "completion-token", + }, + compatibility_date: "2026-07-14", + tags: ["mosoo-managed"], + }); }); + + test.each(["/../secret", "/assets//app.js", "/assets/./app.js", "/assets\\app.js"])( + "rejects an untrusted Static Assets manifest path %s", + (path) => { + expect(() => + validateStaticAssetsManifest({ + [path]: { hash: "0123456789abcdef0123456789abcdef", size: 1 }, + }), + ).toThrow("invalid path"); + }, + ); }); diff --git a/apps/api/tests/app-deployment-detector.test.ts b/apps/api/tests/app-deployment-detector.test.ts index e5b4c784..6a37bee8 100644 --- a/apps/api/tests/app-deployment-detector.test.ts +++ b/apps/api/tests/app-deployment-detector.test.ts @@ -5,10 +5,8 @@ import { detectAppDeploymentPlan, } from "../src/modules/apps/application/app-deployment-detector"; -const RESOURCE_NAME = "app-01j00000000000000000000054"; - function detect(files: Record) { - return detectAppDeploymentPlan({ files }, { resourceName: RESOURCE_NAME }); + return detectAppDeploymentPlan({ files }); } describe("app deployment detector", () => { @@ -19,17 +17,31 @@ describe("app deployment detector", () => { outputDir: ".", packageManager: "none", rootDir: ".", - targetKind: "cloudflare_pages", + targetKind: "cloudflare_static_assets", targetMode: "static_assets", }); }); - test("uses the caller-provided Cloudflare resource name", () => { - expect(detect({ "index.html": "
Hello
" }).generatedWranglerConfig).toContain( - `name = "${RESOURCE_NAME}"`, - ); + test("keeps the static plan physical-name independent", () => { + const config = detect({ "index.html": "
Hello
" }).generatedWranglerConfig; + + expect(config).toContain("[assets]"); + expect(config).toContain('directory = "artifact"'); + expect(config).not.toContain("name ="); }); + test.each(["_worker.js", "_routes.json", "functions/api.js"])( + "rejects unsupported static asset path %s", + (path) => { + expect(() => + detect({ + "index.html": "
Hello
", + [path]: "export default {};", + }), + ).toThrow(AppDeploymentDetectionError); + }, + ); + test("detects Vite static output", () => { expect( detect({ @@ -44,7 +56,7 @@ describe("app deployment detector", () => { installCommand: "pnpm install --frozen-lockfile", outputDir: "dist", packageManager: "pnpm", - targetKind: "cloudflare_pages", + targetKind: "cloudflare_static_assets", }); }); @@ -101,7 +113,7 @@ describe("app deployment detector", () => { installCommand: "npm ci", outputDir: "out", packageManager: "npm", - targetKind: "cloudflare_pages", + targetKind: "cloudflare_static_assets", }); }); @@ -129,7 +141,7 @@ fallback = "index.html" outputDir: "public", routesFallback: "index.html", rootDir: "site", - targetKind: "cloudflare_pages", + targetKind: "cloudflare_static_assets", }); }); diff --git a/apps/api/tests/app-deployment-gateway.test.ts b/apps/api/tests/app-deployment-gateway.test.ts new file mode 100644 index 00000000..482a6fd8 --- /dev/null +++ b/apps/api/tests/app-deployment-gateway.test.ts @@ -0,0 +1,402 @@ +import { describe, expect, test } from "bun:test"; + +import { + APP_DEPLOYMENT_PROBE_HOST_LABEL, + APP_DEPLOYMENT_PROBE_PATH_PREFIX, + APP_DEPLOYMENT_PROBE_SCRIPT_NAME, + dispatchAppDeploymentGatewayRequest, + parseAppDeploymentSubdomain, +} from "../src/modules/apps/application/app-deployment-gateway"; +import type { + AppDeploymentDispatchNamespace, + AppDeploymentGatewayOptions, +} from "../src/modules/apps/application/app-deployment-gateway"; +import { createApiWorker } from "../src/platform/cloudflare/create-api-worker"; +import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const APP_ID = "01J0000000000000000000000Q"; +const OTHER_APP_ID = "01J0000000000000000000000R"; +const MOSOO_SUBDOMAIN = `app-${APP_ID.toLowerCase()}`; +const OTHER_MOSOO_SUBDOMAIN = `app-${OTHER_APP_ID.toLowerCase()}`; +const DOMAIN = "apps.mosoo.ai"; + +function gatewayOptions( + overrides: Partial = {}, +): AppDeploymentGatewayOptions { + return { + appDeploymentDomain: DOMAIN, + dispatcher: { + get() { + throw new Error("Unexpected dispatch"); + }, + }, + resolveActiveScriptName: async () => null, + ...overrides, + }; +} + +function createGatewayDatabase(): SqliteD1Database { + const database = new SqliteD1Database(); + database.execute(` + CREATE TABLE app_deployment ( + active_script_name text, + deleted_at integer, + mosoo_subdomain text NOT NULL + ); + `); + return database; +} + +async function insertGatewayDeployment( + database: D1Database, + input: { activeScriptName: string | null; deletedAt: number | null; mosooSubdomain: string }, +): Promise { + await database + .prepare( + `INSERT INTO app_deployment (active_script_name, deleted_at, mosoo_subdomain) + VALUES (?, ?, ?)`, + ) + .bind(input.activeScriptName, input.deletedAt, input.mosooSubdomain) + .run(); +} + +function getApiWorkerFetch(): NonNullable["fetch"]> { + const fetch = createApiWorker().fetch; + + if (fetch === undefined || typeof fetch !== "function") { + throw new Error("API Worker fetch handler is unavailable."); + } + + return fetch; +} + +function createGatewayBindings( + database: D1Database, + dispatcher: AppDeploymentDispatchNamespace, +): ApiBindings { + return { + APP_DEPLOYMENT_DISPATCHER: dispatcher, + DB: database, + MOSOO_APP_DEPLOYMENT_DOMAIN: DOMAIN, + } as ApiBindings; +} + +describe("app deployment gateway", () => { + test("parses only the exact one-label app hostname", () => { + expect(parseAppDeploymentSubdomain(`${MOSOO_SUBDOMAIN}.${DOMAIN}`, DOMAIN)).toBe( + MOSOO_SUBDOMAIN, + ); + expect(parseAppDeploymentSubdomain(`${MOSOO_SUBDOMAIN.toUpperCase()}.${DOMAIN}.`, DOMAIN)).toBe( + MOSOO_SUBDOMAIN, + ); + + for (const hostname of [ + DOMAIN, + `foo.${DOMAIN}`, + `x.${MOSOO_SUBDOMAIN}.${DOMAIN}`, + `app-81j0000000000000000000000q.${DOMAIN}`, + `${MOSOO_SUBDOMAIN}.evil-${DOMAIN}`, + ]) { + expect(parseAppDeploymentSubdomain(hostname, DOMAIN)).toBeNull(); + } + }); + + test("leaves unrelated hosts for the API router", async () => { + let resolved = false; + const response = await dispatchAppDeploymentGatewayRequest( + new Request("https://cloud.mosoo.ai/api/graphql"), + gatewayOptions({ + resolveActiveScriptName: async () => { + resolved = true; + return null; + }, + }), + ); + + expect(response).toBeNull(); + expect(resolved).toBe(false); + }); + + test("fails closed for an invalid app-domain host", async () => { + let resolved = false; + const response = await dispatchAppDeploymentGatewayRequest( + new Request(`https://x.${MOSOO_SUBDOMAIN}.${DOMAIN}/private`), + gatewayOptions({ + resolveActiveScriptName: async () => { + resolved = true; + return null; + }, + }), + ); + + expect(response?.status).toBe(404); + expect(response?.headers.get("Cache-Control")).toBe("no-store"); + expect(resolved).toBe(false); + }); + + test("exposes the operational probe only on its exact hostname and path", async () => { + let dispatched = false; + let resolved = false; + const options = gatewayOptions({ + dispatcher: { + get() { + dispatched = true; + throw new Error("Unexpected dispatch"); + }, + }, + resolveActiveScriptName: async () => { + resolved = true; + return null; + }, + }); + + for (const url of [ + `https://${APP_DEPLOYMENT_PROBE_HOST_LABEL}.${DOMAIN}/private`, + `https://other.${DOMAIN}${APP_DEPLOYMENT_PROBE_PATH_PREFIX}http`, + ]) { + const response = await dispatchAppDeploymentGatewayRequest(new Request(url), options); + expect(response?.status).toBe(404); + expect(response?.headers.get("Cache-Control")).toBe("no-store"); + } + + expect(dispatched).toBe(false); + expect(resolved).toBe(false); + }); + + test("returns an uncached 404 when no script is active", async () => { + const resolvedSubdomains: string[] = []; + const response = await dispatchAppDeploymentGatewayRequest( + new Request(`https://${MOSOO_SUBDOMAIN}.${DOMAIN}/`), + gatewayOptions({ + resolveActiveScriptName: async (subdomain) => { + resolvedSubdomains.push(subdomain); + return null; + }, + }), + ); + + expect(resolvedSubdomains).toEqual([MOSOO_SUBDOMAIN]); + expect(response?.status).toBe(404); + expect(response?.headers.get("Cache-Control")).toBe("no-store"); + }); + + test("returns 503 for the app domain when the dispatch binding is missing", async () => { + let resolved = false; + const response = await dispatchAppDeploymentGatewayRequest( + new Request(`https://${MOSOO_SUBDOMAIN}.${DOMAIN}/`), + gatewayOptions({ + dispatcher: undefined, + resolveActiveScriptName: async () => { + resolved = true; + return "must-not-dispatch"; + }, + }), + ); + + expect(response?.status).toBe(503); + expect(response?.headers.get("Cache-Control")).toBe("no-store"); + expect(resolved).toBe(false); + }); + + test("dispatches the original request and returns the user Worker response unchanged", async () => { + const request = new Request(`https://${MOSOO_SUBDOMAIN}.${DOMAIN}/events?after=42`, { + body: "stream me", + headers: { "Content-Type": "text/plain" }, + method: "POST", + }); + const workerResponse = new Response("worker stream", { + headers: { "X-Worker": "attempt-7" }, + }); + const calls: Array<{ request: Request; scriptName: string }> = []; + const dispatcher: AppDeploymentDispatchNamespace = { + get(scriptName) { + return { + fetch(receivedRequest) { + calls.push({ request: receivedRequest, scriptName }); + return workerResponse; + }, + }; + }, + }; + + const response = await dispatchAppDeploymentGatewayRequest( + request, + gatewayOptions({ + dispatcher, + resolveActiveScriptName: async () => "mosoo-run-g2-a7", + }), + ); + + expect(calls).toEqual([{ request, scriptName: "mosoo-run-g2-a7" }]); + expect(request.bodyUsed).toBe(false); + expect(response).toBe(workerResponse); + }); + + test("returns 503 without dispatching when active-script resolution fails", async () => { + let dispatched = false; + const lookupError = new Error("D1 unavailable"); + + const response = await dispatchAppDeploymentGatewayRequest( + new Request(`https://${MOSOO_SUBDOMAIN}.${DOMAIN}/`), + gatewayOptions({ + dispatcher: { + get() { + dispatched = true; + throw new Error("Unexpected dispatch"); + }, + }, + resolveActiveScriptName: async () => { + throw lookupError; + }, + }), + ); + + expect(response?.status).toBe(503); + expect(response?.headers.get("Cache-Control")).toBe("no-store"); + expect(await response?.text()).toBe(""); + expect(dispatched).toBe(false); + }); + + test("returns 503 without leaking missing or failed user Worker details", async () => { + const dispatchers: AppDeploymentDispatchNamespace[] = [ + { + get() { + throw new Error("Worker not found"); + }, + }, + { + get() { + return { + fetch() { + throw new Error("User Worker failed"); + }, + }; + }, + }, + ]; + + for (const dispatcher of dispatchers) { + const response = await dispatchAppDeploymentGatewayRequest( + new Request(`https://${MOSOO_SUBDOMAIN}.${DOMAIN}/`), + gatewayOptions({ + dispatcher, + resolveActiveScriptName: async () => "missing-attempt-script", + }), + ); + + expect(response?.status).toBe(503); + expect(response?.headers.get("Cache-Control")).toBe("no-store"); + expect(await response?.text()).toBe(""); + } + }); + + test("routes the API Worker through the active D1 script before Hono", async () => { + const database = createGatewayDatabase(); + await insertGatewayDeployment(database, { + activeScriptName: "mosoo-run-g2-a7", + deletedAt: null, + mosooSubdomain: MOSOO_SUBDOMAIN, + }); + + const request = new Request(`https://${MOSOO_SUBDOMAIN}.${DOMAIN}/events?after=42`, { + body: "stream me", + method: "POST", + }); + const workerResponse = new Response("worker stream"); + let dispatchedRequest: Request | null = null; + const bindings = createGatewayBindings(database, { + get(scriptName) { + expect(scriptName).toBe("mosoo-run-g2-a7"); + return { + fetch(receivedRequest) { + dispatchedRequest = receivedRequest; + return workerResponse; + }, + }; + }, + }); + + const response = await getApiWorkerFetch()(request, bindings, {} as ExecutionContext); + + expect(dispatchedRequest).toBe(request); + expect(request.bodyUsed).toBe(false); + expect(response).toBe(workerResponse); + }); + + test("routes the pre-provisioned bootstrap probe through the API Worker without D1", async () => { + const request = new Request( + `https://${APP_DEPLOYMENT_PROBE_HOST_LABEL}.${DOMAIN}${APP_DEPLOYMENT_PROBE_PATH_PREFIX}http`, + { body: "probe", method: "POST" }, + ); + const workerResponse = new Response("probe-ok", { + headers: { "Cache-Control": "no-store" }, + }); + let dispatchedRequest: Request | null = null; + const bindings = createGatewayBindings(new SqliteD1Database(), { + get(scriptName) { + expect(scriptName).toBe(APP_DEPLOYMENT_PROBE_SCRIPT_NAME); + return { + fetch(receivedRequest) { + dispatchedRequest = receivedRequest; + return workerResponse; + }, + }; + }, + }); + + const response = await getApiWorkerFetch()(request, bindings, {} as ExecutionContext); + + expect(dispatchedRequest).toBe(request); + expect(request.bodyUsed).toBe(false); + expect(response).toBe(workerResponse); + }); + + test("the API Worker fails closed on an app host without the dispatch binding", async () => { + const response = await getApiWorkerFetch()( + new Request(`https://${MOSOO_SUBDOMAIN}.${DOMAIN}/`), + { + DB: createGatewayDatabase(), + MOSOO_APP_DEPLOYMENT_DOMAIN: DOMAIN, + } as ApiBindings, + {} as ExecutionContext, + ); + + expect(response.status).toBe(503); + expect(response.headers.get("Cache-Control")).toBe("no-store"); + }); + + test("does not dispatch unpublished or deleted D1 deployments", async () => { + const database = createGatewayDatabase(); + await insertGatewayDeployment(database, { + activeScriptName: null, + deletedAt: null, + mosooSubdomain: MOSOO_SUBDOMAIN, + }); + await insertGatewayDeployment(database, { + activeScriptName: "stale-deleted-script", + deletedAt: 1, + mosooSubdomain: OTHER_MOSOO_SUBDOMAIN, + }); + + let dispatchCount = 0; + const bindings = createGatewayBindings(database, { + get() { + dispatchCount += 1; + throw new Error("Unexpected dispatch"); + }, + }); + + for (const subdomain of [MOSOO_SUBDOMAIN, OTHER_MOSOO_SUBDOMAIN]) { + const response = await getApiWorkerFetch()( + new Request(`https://${subdomain}.${DOMAIN}/`), + bindings, + {} as ExecutionContext, + ); + + expect(response.status).toBe(404); + expect(response.headers.get("Cache-Control")).toBe("no-store"); + } + expect(dispatchCount).toBe(0); + }); +}); diff --git a/apps/api/tests/app-deployment-script-reconciliation.test.ts b/apps/api/tests/app-deployment-script-reconciliation.test.ts new file mode 100644 index 00000000..43b85ff5 --- /dev/null +++ b/apps/api/tests/app-deployment-script-reconciliation.test.ts @@ -0,0 +1,809 @@ +import { describe, expect, test } from "bun:test"; + +import type { ApiCommandId } from "@mosoo/db"; +import { createPlatformId } from "@mosoo/id"; +import type { AppDeploymentId, AppDeploymentRunId, AppId } from "@mosoo/id"; + +import { MANAGED_PROD_SCHEMA_TRIGGERS } from "../bin/prod-schema-guard"; +import { + APP_DEPLOYMENT_SCRIPT_GRACE_MS, + markAppDeploymentScriptUploadStarted, + reconcileAppDeploymentScriptPage, + registerAppDeploymentScriptCandidate, +} from "../src/modules/apps/application/app-deployment-script-reconciliation.service"; +import type { + AppDeploymentScriptCandidateAuthority, + AppDeploymentScriptCleanupAction, +} from "../src/modules/apps/application/app-deployment-script-reconciliation.service"; +import { applyDrizzleMigration, drizzleMigrations } from "./helpers/drizzle-migrations"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const ACTIVE_RUN_STATUSES = [ + "queued", + "preparing", + "building", + "submitting", + "submitted", + "activating", +] as const; + +interface AttemptFixture { + readonly appId: AppId; + readonly authority: AppDeploymentScriptCandidateAuthority; + readonly commandId: ApiCommandId; + readonly database: SqliteD1Database; + readonly deploymentId: AppDeploymentId; + readonly registeredAt: number; + readonly runId: AppDeploymentRunId; + readonly scriptName: string; +} + +function createDatabase(): SqliteD1Database { + const database = new SqliteD1Database(); + for (const migration of drizzleMigrations) { + applyDrizzleMigration(database, migration.tag); + } + return database; +} + +function candidateScriptName( + commandId: ApiCommandId, + deliveryGeneration: number, + attemptCount: number, +): string { + return `app-${commandId.toLowerCase()}-g${deliveryGeneration.toString(36)}-a${attemptCount.toString(36)}`; +} + +function authority( + commandId: ApiCommandId, + input: { attemptCount: number; claimOwner: string; deliveryGeneration?: number }, +): AppDeploymentScriptCandidateAuthority { + return { + attemptCount: input.attemptCount, + claimOwner: input.claimOwner, + commandId, + deliveryGeneration: input.deliveryGeneration ?? 1, + requireOwnership: async () => undefined, + }; +} + +async function insertAttemptFoundation( + database: SqliteD1Database, + input: { + readonly registeredAt: number; + readonly status?: (typeof ACTIVE_RUN_STATUSES)[number]; + }, +): Promise> { + const appId = createPlatformId(); + const commandId = createPlatformId(); + const deploymentId = createPlatformId(); + const runId = createPlatformId(); + const claimOwner = crypto.randomUUID(); + + await database + .prepare( + `INSERT INTO app_deployment ( + active_script_name, app_id, created_at, default_branch, deleted_at, id, + last_successful_url, latest_run_id, mosoo_subdomain, owner_account_id, + repo_name, repo_owner, repo_url, source_kind, updated_at + ) VALUES (NULL, ?, ?, 'main', NULL, ?, NULL, ?, ?, ?, 'repo', 'owner', + 'https://example.test/repo', 'github_public', ?)`, + ) + .bind( + appId, + input.registeredAt, + deploymentId, + runId, + `app-${deploymentId.toLowerCase()}`, + createPlatformId(), + input.registeredAt, + ) + .run(); + await database + .prepare( + `INSERT INTO app_deployment_run ( + app_id, created_at, deployment_id, id, source_branch, source_commit_sha, + status, updated_at + ) VALUES (?, ?, ?, ?, 'main', '0123456789abcdef', ?, ?)`, + ) + .bind( + appId, + input.registeredAt, + deploymentId, + runId, + input.status ?? "queued", + input.registeredAt, + ) + .run(); + await database + .prepare( + `INSERT INTO api_command ( + attempt_count, claim_expires_at, claim_owner, completed_at, created_at, + dedupe_key, delivery_generation, id, kind, last_error_code, + last_error_message, payload_json, status, updated_at + ) VALUES (1, ?, ?, NULL, ?, ?, 1, ?, 'app_deployment_run_dispatch', NULL, + NULL, ?, 'running', ?)`, + ) + .bind( + Date.now() + 7 * APP_DEPLOYMENT_SCRIPT_GRACE_MS, + claimOwner, + input.registeredAt, + `dispatch:${commandId}`, + commandId, + JSON.stringify({ appDeploymentRunId: runId }), + input.registeredAt, + ) + .run(); + + return { appId, commandId, database, deploymentId, registeredAt: input.registeredAt, runId }; +} + +async function createAttempt( + database: SqliteD1Database, + input: { + readonly registeredAt?: number; + readonly status?: (typeof ACTIVE_RUN_STATUSES)[number]; + } = {}, +): Promise { + const foundation = await insertAttemptFoundation(database, { + registeredAt: input.registeredAt ?? Date.now(), + status: input.status, + }); + const claimOwner = ( + await database + .prepare("SELECT claim_owner FROM api_command WHERE id = ?") + .bind(foundation.commandId) + .first<{ claim_owner: string }>() + )?.claim_owner; + if (claimOwner === undefined) { + throw new Error("App deployment command fixture is missing its claim owner."); + } + const attemptAuthority = authority(foundation.commandId, { attemptCount: 1, claimOwner }); + const scriptName = candidateScriptName(foundation.commandId, 1, 1); + await registerAppDeploymentScriptCandidate(database, { + authority: attemptAuthority, + deploymentId: foundation.deploymentId, + nowMs: foundation.registeredAt, + runId: foundation.runId, + scriptName, + }); + return { ...foundation, authority: attemptAuthority, scriptName }; +} + +async function createReconciliationAuthority( + database: SqliteD1Database, + nowMs: number, +): Promise { + const commandId = createPlatformId(); + const claimOwner = crypto.randomUUID(); + await database + .prepare( + `INSERT INTO api_command ( + attempt_count, claim_expires_at, claim_owner, completed_at, created_at, + dedupe_key, delivery_generation, id, kind, last_error_code, + last_error_message, payload_json, status, updated_at + ) VALUES (1, ?, ?, NULL, ?, ?, 1, ?, 'app_deployment_script_reconciliation', + NULL, NULL, '{}', 'running', ?)`, + ) + .bind( + nowMs + 7 * APP_DEPLOYMENT_SCRIPT_GRACE_MS, + claimOwner, + nowMs, + `script-reconciliation:${commandId}`, + commandId, + nowMs, + ) + .run(); + return authority(commandId, { attemptCount: 1, claimOwner }); +} + +async function expireAttemptAuthority(fixture: AttemptFixture): Promise { + await fixture.database + .prepare( + `UPDATE api_command + SET claim_expires_at = NULL, claim_owner = NULL, status = 'failed' + WHERE id = ?`, + ) + .bind(fixture.commandId) + .run(); +} + +async function retireScriptThroughTerminalRun(fixture: AttemptFixture): Promise { + await fixture.database + .prepare("UPDATE app_deployment_run SET status = 'failed' WHERE id = ?") + .bind(fixture.runId) + .run(); + const retireAfter = (await readScript(fixture)).retire_after; + if (typeof retireAfter !== "number") { + throw new Error("Terminal App deployment run did not retire its script candidate."); + } + return retireAfter; +} + +async function seedExpiredHistoricalScript(fixture: AttemptFixture, dueAt: number): Promise { + const retireAuthority = MANAGED_PROD_SCHEMA_TRIGGERS.find( + ({ name }) => name === "app_deployment_script_retire_authority", + ); + if (retireAuthority === undefined) { + throw new Error("App deployment script retirement trigger fixture is missing."); + } + fixture.database.execute("DROP TRIGGER app_deployment_script_retire_authority"); + await fixture.database + .prepare( + `UPDATE app_deployment_script + SET next_reconcile_at = ?, retire_after = ? + WHERE script_name = ?`, + ) + .bind(dueAt, fixture.registeredAt, fixture.scriptName) + .run(); + fixture.database.execute(retireAuthority.sql); +} + +async function readScript( + fixture: Pick, +): Promise> { + const row = await fixture.database + .prepare("SELECT * FROM app_deployment_script WHERE script_name = ?") + .bind(fixture.scriptName) + .first>(); + if (row === null) { + throw new Error("App deployment script fixture disappeared."); + } + return row; +} + +describe("App deployment script ledger", () => { + test("plain INSERT atomically resets any active run and retires the previous attempt", async () => { + const database = createDatabase(); + const first = await createAttempt(database, { status: "building" }); + + expect( + await database + .prepare("SELECT status, target_script_name FROM app_deployment_run WHERE id = ?") + .bind(first.runId) + .first(), + ).toEqual({ status: "preparing", target_script_name: first.scriptName }); + await expect( + registerAppDeploymentScriptCandidate(database, { + authority: first.authority, + deploymentId: first.deploymentId, + nowMs: first.registeredAt, + runId: first.runId, + scriptName: first.scriptName, + }), + ).rejects.toThrow(); + + const secondClaimOwner = crypto.randomUUID(); + await database + .prepare( + `UPDATE api_command + SET attempt_count = 2, claim_expires_at = ?, claim_owner = ?, status = 'running' + WHERE id = ?`, + ) + .bind(Date.now() + APP_DEPLOYMENT_SCRIPT_GRACE_MS, secondClaimOwner, first.commandId) + .run(); + await database + .prepare("UPDATE app_deployment_run SET status = 'submitting' WHERE id = ?") + .bind(first.runId) + .run(); + + const secondAuthority = authority(first.commandId, { + attemptCount: 2, + claimOwner: secondClaimOwner, + }); + const secondScriptName = candidateScriptName(first.commandId, 1, 2); + const beforeRegister = Date.now(); + await registerAppDeploymentScriptCandidate(database, { + authority: secondAuthority, + deploymentId: first.deploymentId, + nowMs: beforeRegister, + runId: first.runId, + scriptName: secondScriptName, + }); + const afterRegister = Date.now(); + + expect( + await database + .prepare("SELECT status, target_script_name FROM app_deployment_run WHERE id = ?") + .bind(first.runId) + .first(), + ).toEqual({ status: "preparing", target_script_name: secondScriptName }); + const old = await readScript(first); + expect(old.retire_after).toBeGreaterThanOrEqual( + beforeRegister + APP_DEPLOYMENT_SCRIPT_GRACE_MS, + ); + expect(old.retire_after).toBeLessThanOrEqual(afterRegister + APP_DEPLOYMENT_SCRIPT_GRACE_MS); + await database + .prepare("UPDATE app_deployment_run SET status = 'failed' WHERE id = ?") + .bind(first.runId) + .run(); + expect( + (await readScript({ database, scriptName: secondScriptName })).retire_after, + ).toBeNumber(); + expect( + await database.prepare("SELECT count(*) AS count FROM app_deployment_script").first(), + ).toEqual({ count: 2 }); + }); + + test("requires an upload fence before promotion and drains replaced or deleted active scripts", async () => { + const database = createDatabase(); + const first = await createAttempt(database); + await database + .prepare("UPDATE app_deployment_run SET status = 'activating' WHERE id = ?") + .bind(first.runId) + .run(); + await expect( + database + .prepare( + `UPDATE app_deployment + SET active_script_name = ?, last_successful_url = 'https://app.test', updated_at = ? + WHERE id = ?`, + ) + .bind(first.scriptName, first.registeredAt + 1, first.deploymentId) + .run(), + ).rejects.toThrow(); + + await expect( + database + .prepare( + `UPDATE app_deployment_script + SET retire_after = CAST(unixepoch('subsec') * 1000 AS INTEGER) + 86400000, + next_reconcile_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) + 86400000 + WHERE script_name = ?`, + ) + .bind(first.scriptName) + .run(), + ).rejects.toThrow(); + expect((await readScript(first)).retire_after).toBeNull(); + + await database + .prepare("UPDATE api_command SET payload_json = '{}' WHERE id = ?") + .bind(first.commandId) + .run(); + await expect( + markAppDeploymentScriptUploadStarted(database, { + authority: first.authority, + deploymentId: first.deploymentId, + nowMs: first.registeredAt + 1, + runId: first.runId, + scriptName: first.scriptName, + }), + ).rejects.toThrow(); + await database + .prepare("UPDATE api_command SET payload_json = ? WHERE id = ?") + .bind(JSON.stringify({ appDeploymentRunId: first.runId }), first.commandId) + .run(); + await markAppDeploymentScriptUploadStarted(database, { + authority: first.authority, + deploymentId: first.deploymentId, + nowMs: first.registeredAt + 1, + runId: first.runId, + scriptName: first.scriptName, + }); + await database + .prepare( + `UPDATE app_deployment + SET active_script_name = ?, last_successful_url = 'https://app.test', updated_at = ? + WHERE id = ?`, + ) + .bind(first.scriptName, first.registeredAt + 2, first.deploymentId) + .run(); + expect(await readScript(first)).toMatchObject({ next_reconcile_at: null, retire_after: null }); + + const secondClaimOwner = crypto.randomUUID(); + await database + .prepare( + `UPDATE api_command + SET attempt_count = 2, claim_expires_at = ?, claim_owner = ?, status = 'running' + WHERE id = ?`, + ) + .bind(Date.now() + APP_DEPLOYMENT_SCRIPT_GRACE_MS, secondClaimOwner, first.commandId) + .run(); + const secondAuthority = authority(first.commandId, { + attemptCount: 2, + claimOwner: secondClaimOwner, + }); + const secondScriptName = candidateScriptName(first.commandId, 1, 2); + await registerAppDeploymentScriptCandidate(database, { + authority: secondAuthority, + deploymentId: first.deploymentId, + nowMs: Date.now(), + runId: first.runId, + scriptName: secondScriptName, + }); + await markAppDeploymentScriptUploadStarted(database, { + authority: secondAuthority, + deploymentId: first.deploymentId, + nowMs: Date.now(), + runId: first.runId, + scriptName: secondScriptName, + }); + await database + .prepare("UPDATE app_deployment_run SET status = 'activating' WHERE id = ?") + .bind(first.runId) + .run(); + await database + .prepare( + `UPDATE app_deployment + SET active_script_name = ?, last_successful_url = 'https://app.test', updated_at = ? + WHERE id = ?`, + ) + .bind(secondScriptName, Date.now(), first.deploymentId) + .run(); + + expect((await readScript(first)).retire_after).toBeNumber(); + await database + .prepare( + `UPDATE app_deployment + SET active_script_name = NULL, deleted_at = ?, last_successful_url = NULL, updated_at = ? + WHERE id = ?`, + ) + .bind(Date.now(), Date.now(), first.deploymentId) + .run(); + const second = { database, scriptName: secondScriptName }; + const retired = await readScript(second); + expect(retired.retire_after).toBeNumber(); + await expect( + database + .prepare( + `INSERT OR REPLACE INTO app_deployment_script + SELECT * FROM app_deployment_script WHERE script_name = ?`, + ) + .bind(secondScriptName) + .run(), + ).rejects.toThrow(); + expect(await readScript(second)).toEqual(retired); + await expect( + database + .prepare("UPDATE app_deployment_script SET retire_after = NULL WHERE script_name = ?") + .bind(secondScriptName) + .run(), + ).rejects.toThrow(); + await expect( + database + .prepare("DELETE FROM app_deployment_script WHERE script_name = ?") + .bind(secondScriptName) + .run(), + ).rejects.toThrow(); + }); + + test("linearizes promotion and orphan retirement in either commit order", async () => { + const nowMs = Date.now(); + const database = createDatabase(); + const gcAuthority = await createReconciliationAuthority(database, nowMs); + const gcWinner = await createAttempt(database, { registeredAt: nowMs - 1_000 }); + await markAppDeploymentScriptUploadStarted(database, { + authority: gcWinner.authority, + deploymentId: gcWinner.deploymentId, + nowMs, + runId: gcWinner.runId, + scriptName: gcWinner.scriptName, + }); + await expireAttemptAuthority(gcWinner); + await database + .prepare("UPDATE app_deployment_script SET next_reconcile_at = ? WHERE script_name = ?") + .bind(nowMs, gcWinner.scriptName) + .run(); + + let cleanupCalls = 0; + const armed = await reconcileAppDeploymentScriptPage( + database, + { authority: gcAuthority, cursor: null }, + async () => { + cleanupCalls += 1; + }, + ); + expect(armed.armed).toBe(1); + expect(cleanupCalls).toBe(0); + + await database + .prepare( + `UPDATE api_command + SET claim_expires_at = ?, claim_owner = ?, status = 'running' + WHERE id = ?`, + ) + .bind( + nowMs + APP_DEPLOYMENT_SCRIPT_GRACE_MS, + gcWinner.authority.claimOwner, + gcWinner.commandId, + ) + .run(); + await database + .prepare("UPDATE app_deployment_run SET status = 'activating' WHERE id = ?") + .bind(gcWinner.runId) + .run(); + await expect( + database + .prepare( + `UPDATE app_deployment + SET active_script_name = ?, last_successful_url = 'https://app.test', updated_at = ? + WHERE id = ?`, + ) + .bind(gcWinner.scriptName, nowMs, gcWinner.deploymentId) + .run(), + ).rejects.toThrow(); + + const promotionWinner = await createAttempt(database, { registeredAt: nowMs }); + await markAppDeploymentScriptUploadStarted(database, { + authority: promotionWinner.authority, + deploymentId: promotionWinner.deploymentId, + nowMs, + runId: promotionWinner.runId, + scriptName: promotionWinner.scriptName, + }); + await database + .prepare("UPDATE app_deployment_run SET status = 'activating' WHERE id = ?") + .bind(promotionWinner.runId) + .run(); + await database + .prepare( + `UPDATE app_deployment + SET active_script_name = ?, last_successful_url = 'https://app.test', updated_at = ? + WHERE id = ?`, + ) + .bind(promotionWinner.scriptName, nowMs, promotionWinner.deploymentId) + .run(); + await database + .prepare("UPDATE app_deployment_script SET next_reconcile_at = ? WHERE script_name = ?") + .bind(nowMs, promotionWinner.scriptName) + .run(); + const protectedResult = await reconcileAppDeploymentScriptPage( + database, + { authority: gcAuthority, cursor: null }, + async () => { + cleanupCalls += 1; + }, + ); + expect(protectedResult.deferred).toBeGreaterThanOrEqual(1); + expect(await readScript(promotionWinner)).toMatchObject({ + next_reconcile_at: null, + retire_after: null, + }); + expect(cleanupCalls).toBe(0); + }); + + test("protects every active candidate status only under the exact current command claim", async () => { + const nowMs = Date.now(); + const database = createDatabase(); + const fixture = await createAttempt(database, { registeredAt: nowMs }); + const gcAuthority = await createReconciliationAuthority(database, nowMs); + + for (const status of ACTIVE_RUN_STATUSES) { + await database + .prepare("UPDATE app_deployment_run SET status = ? WHERE id = ?") + .bind(status, fixture.runId) + .run(); + await database + .prepare("UPDATE app_deployment_script SET next_reconcile_at = ? WHERE script_name = ?") + .bind(nowMs, fixture.scriptName) + .run(); + const result = await reconcileAppDeploymentScriptPage( + database, + { authority: gcAuthority, cursor: null }, + async () => { + throw new Error("A live candidate must not be cleaned."); + }, + ); + expect(result.deferred).toBe(1); + expect((await readScript(fixture)).retire_after).toBeNull(); + } + + await expireAttemptAuthority(fixture); + await database + .prepare("UPDATE app_deployment_script SET next_reconcile_at = ? WHERE script_name = ?") + .bind(nowMs, fixture.scriptName) + .run(); + expect( + ( + await reconcileAppDeploymentScriptPage( + database, + { authority: gcAuthority, cursor: null }, + async () => undefined, + ) + ).armed, + ).toBe(1); + }); + + test("reconciles expired historical rows with permanent tombstones and recoverable leases", async () => { + const nowMs = Date.now(); + const database = createDatabase(); + const fixture = await createAttempt(database, { + registeredAt: nowMs - 2 * APP_DEPLOYMENT_SCRIPT_GRACE_MS, + }); + await markAppDeploymentScriptUploadStarted(database, { + authority: fixture.authority, + deploymentId: fixture.deploymentId, + nowMs: fixture.registeredAt + 1, + runId: fixture.runId, + scriptName: fixture.scriptName, + }); + await expireAttemptAuthority(fixture); + await seedExpiredHistoricalScript(fixture, nowMs); + const gcAuthority = await createReconciliationAuthority(database, nowMs); + const callerClockAhead = { + authority: gcAuthority, + cursor: null, + nowMs: Number.MAX_SAFE_INTEGER, + }; + const expectedBackoffDays = [1, 2, 4, 8, 16, 30, 30]; + let firstTombstone: number | null = null; + + for (const [index, days] of expectedBackoffDays.entries()) { + await database + .prepare("UPDATE app_deployment_script SET next_reconcile_at = ? WHERE script_name = ?") + .bind(fixture.registeredAt, fixture.scriptName) + .run(); + const actions: AppDeploymentScriptCleanupAction[] = []; + const result = await reconcileAppDeploymentScriptPage( + database, + callerClockAhead, + async (action) => { + actions.push(action); + }, + ); + expect(result.cleaned).toBe(1); + expect(actions[0]).toMatchObject({ + artifactPrefix: `app-deployments/${fixture.scriptName}/`, + buildSandboxId: `${fixture.runId}-g1-a1-build`, + deploySandboxId: `${fixture.runId}-g1-a1-deploy`, + scriptName: fixture.scriptName, + uploadStartedAt: fixture.registeredAt + 1, + }); + const row = await readScript(fixture); + firstTombstone ??= row.external_deleted_at as number; + expect(row.external_deleted_at).toBe(firstTombstone); + expect(row.next_reconcile_at).toBe( + (row.last_reconciled_at as number) + days * APP_DEPLOYMENT_SCRIPT_GRACE_MS, + ); + expect(row.reconcile_count).toBe(index + 1); + } + + await database + .prepare( + `UPDATE app_deployment_script + SET next_reconcile_at = ?, reconcile_count = ? + WHERE script_name = ?`, + ) + .bind(fixture.registeredAt, Number.MAX_SAFE_INTEGER, fixture.scriptName) + .run(); + expect( + (await reconcileAppDeploymentScriptPage(database, callerClockAhead, async () => undefined)) + .cleaned, + ).toBe(1); + const saturated = await readScript(fixture); + expect(saturated.reconcile_count).toBe(Number.MAX_SAFE_INTEGER); + expect(Number.isSafeInteger(saturated.next_reconcile_at)).toBeTrue(); + expect( + await database.prepare("SELECT count(*) AS count FROM app_deployment_script").first(), + ).toEqual({ count: 1 }); + + const noUpload = await createAttempt(database, { + registeredAt: nowMs - 2 * APP_DEPLOYMENT_SCRIPT_GRACE_MS, + }); + await expireAttemptAuthority(noUpload); + await seedExpiredHistoricalScript(noUpload, nowMs); + const failed = await reconcileAppDeploymentScriptPage(database, callerClockAhead, async () => { + throw new Error("transient cleanup failure"); + }); + expect(failed.failures).toHaveLength(1); + const failedRow = await readScript(noUpload); + expect(failedRow).toMatchObject({ + external_deleted_at: null, + reconcile_count: 0, + reconcile_expires_at: null, + reconcile_owner: null, + }); + expect(failedRow.next_reconcile_at).toBe( + (failedRow.last_reconciled_at as number) + APP_DEPLOYMENT_SCRIPT_GRACE_MS, + ); + + await database + .prepare( + `UPDATE app_deployment_script + SET next_reconcile_at = ?, reconcile_expires_at = ?, reconcile_owner = 'busy' + WHERE script_name = ?`, + ) + .bind(noUpload.registeredAt, Date.now() + 60_000, noUpload.scriptName) + .run(); + expect( + ( + await reconcileAppDeploymentScriptPage(database, callerClockAhead, async () => { + throw new Error("A live row lease must not be stolen."); + }) + ).deferred, + ).toBe(1); + await database + .prepare("UPDATE app_deployment_script SET reconcile_expires_at = ? WHERE script_name = ?") + .bind(0, noUpload.scriptName) + .run(); + let recoveredAction: AppDeploymentScriptCleanupAction | null = null; + await expect( + reconcileAppDeploymentScriptPage( + database, + { authority: gcAuthority, cursor: null }, + async (action) => { + recoveredAction = action; + await database + .prepare( + "UPDATE app_deployment_script SET reconcile_expires_at = 0 WHERE script_name = ?", + ) + .bind(noUpload.scriptName) + .run(); + }, + ), + ).rejects.toThrow("lost durable authority"); + expect(await readScript(noUpload)).toMatchObject({ + reconcile_count: 0, + reconcile_expires_at: 0, + }); + + const recovered = await reconcileAppDeploymentScriptPage( + database, + { authority: gcAuthority, cursor: null }, + async (action) => { + recoveredAction = action; + }, + ); + expect(recovered.cleaned).toBe(1); + expect(recoveredAction?.uploadStartedAt).toBeNull(); + expect((await readScript(noUpload)).external_deleted_at).toBeNull(); + }); + + test("uses D1 time rather than a caller clock to enforce the retirement grace", async () => { + const nowMs = Date.now(); + const database = createDatabase(); + const fixture = await createAttempt(database, { registeredAt: nowMs }); + await expireAttemptAuthority(fixture); + const retireAfter = await retireScriptThroughTerminalRun(fixture); + const gcAuthority = await createReconciliationAuthority(database, nowMs); + const callerClockAhead = { + authority: gcAuthority, + cursor: null, + nowMs: Number.MAX_SAFE_INTEGER, + }; + + let cleanupCalls = 0; + const result = await reconcileAppDeploymentScriptPage(database, callerClockAhead, async () => { + cleanupCalls += 1; + }); + + expect(result.processed).toBe(0); + expect(cleanupCalls).toBe(0); + expect((await readScript(fixture)).retire_after).toBe(retireAfter); + }); + + test("pages expired historical rows without skipping a permanent ledger entry", async () => { + const nowMs = Date.now(); + const database = createDatabase(); + for (let index = 0; index < 2; index += 1) { + const fixture = await createAttempt(database, { + registeredAt: nowMs - 2 * APP_DEPLOYMENT_SCRIPT_GRACE_MS - index, + }); + await expireAttemptAuthority(fixture); + await seedExpiredHistoricalScript(fixture, nowMs); + } + const gcAuthority = await createReconciliationAuthority(database, nowMs); + let cleaned = 0; + const first = await reconcileAppDeploymentScriptPage( + database, + { authority: gcAuthority, cursor: null }, + async () => { + cleaned += 1; + }, + ); + expect(first).toMatchObject({ hasMore: true, processed: 1 }); + expect(first.nextCursor).not.toBeNull(); + const second = await reconcileAppDeploymentScriptPage( + database, + { authority: gcAuthority, cursor: first.nextCursor }, + async () => { + cleaned += 1; + }, + ); + expect(second).toMatchObject({ hasMore: false, nextCursor: null, processed: 1 }); + expect(cleaned).toBe(2); + expect( + await database.prepare("SELECT count(*) AS count FROM app_deployment_script").first(), + ).toEqual({ count: 2 }); + }); +}); diff --git a/apps/api/tests/app-deployment-service.test.ts b/apps/api/tests/app-deployment-service.test.ts index 57c8e284..2ea679a1 100644 --- a/apps/api/tests/app-deployment-service.test.ts +++ b/apps/api/tests/app-deployment-service.test.ts @@ -1,19 +1,40 @@ import { describe, expect, test } from "bun:test"; +import { createHash } from "node:crypto"; +import { mkdtemp, mkdir, readFile, rm, symlink, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; -import { apiCommandsTable, appDeploymentRunsTable, appDeploymentsTable } from "@mosoo/db"; +import { + apiCommandsTable, + appDeploymentRunsTable, + appDeploymentScriptsTable, + appDeploymentsTable, +} from "@mosoo/db"; +import { MANAGED_PROD_SCHEMA_TRIGGERS } from "@mosoo/db/deploy-schema-guard"; +import type { AppDeploymentRunId } from "@mosoo/id"; import { eq } from "drizzle-orm"; import { createAppDeploymentRunDispatchDedupeKey } from "../src/modules/api-command/application/api-command-enqueue"; import { API_COMMAND_QUEUE_SEND_FAILED_CODE } from "../src/modules/api-command/application/api-command-ledger"; -import { - APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS, - APP_DEPLOYMENT_RUN_DISPATCH_RETRY_EXHAUSTED_CODE, -} from "../src/modules/api-command/application/api-command-policy"; -import { processApiCommandDeadLetterMessage } from "../src/modules/api-command/application/api-command-processor"; +import type { ApiCommandMessage } from "../src/modules/api-command/application/api-command-message"; +import { APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS } from "../src/modules/api-command/application/api-command-policy"; import { getDeploymentAgentCapabilityAuthority } from "../src/modules/apps/application/app-deployment-capability-authority.service"; import type { CloudflareDeploymentClient } from "../src/modules/apps/application/app-deployment-cloudflare-client"; -import type { AppDeploymentBuildRunner } from "../src/modules/apps/application/app-deployment-executor.service"; -import { dispatchAppDeploymentRun } from "../src/modules/apps/application/app-deployment-executor.service"; +import type { + AppDeploymentAttemptAuthority, + AppDeploymentBuildRunner, + AppDeploymentDispatchOutcome, + DispatchAppDeploymentRunOptions, +} from "../src/modules/apps/application/app-deployment-executor.service"; +import { + AppDeploymentNonRetryableError, + APP_DEPLOYMENT_STATIC_ASSET_TOOL_SOURCE, + appDeploymentBuildSandboxId, + appDeploymentCandidateScriptName, + appDeploymentDeploySandboxId, + dispatchAppDeploymentRun, +} from "../src/modules/apps/application/app-deployment-executor.service"; +import { registerAppDeploymentScriptCandidate } from "../src/modules/apps/application/app-deployment-script-reconciliation.service"; import { deleteAppDeployment, deployApp, @@ -25,13 +46,8 @@ import type { AuthenticatedViewer } from "../src/modules/auth/application/viewer import type { SandboxHandle } from "../src/modules/runtime/infrastructure/sandbox-handles"; import { createApiWorker } from "../src/platform/cloudflare/create-api-worker"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; -import { API_ERROR_CODE } from "../src/platform/errors"; -import type { ApiError } from "../src/platform/errors"; import { currentTimestampMs } from "../src/time"; -import { - createApiCommandQueueStub, - createRecordedQueueMessage, -} from "./helpers/api-command-queue-fixture"; +import { createApiCommandQueueStub } from "./helpers/api-command-queue-fixture"; import { SqliteD1Database } from "./helpers/sqlite-d1"; const OWNER_ID = "01J00000000000000000000001"; @@ -64,6 +80,7 @@ function createDatabase(): SqliteD1Database { ); CREATE TABLE app_deployment ( + active_script_name text, app_id text NOT NULL, created_at integer NOT NULL, default_branch text NOT NULL, @@ -77,7 +94,17 @@ function createDatabase(): SqliteD1Database { repo_owner text NOT NULL, repo_url text NOT NULL, source_kind text NOT NULL, - updated_at integer NOT NULL + updated_at integer NOT NULL, + CONSTRAINT app_deployment_traffic_authority_check CHECK ( + (active_script_name IS NULL AND last_successful_url IS NULL) + OR ( + deleted_at IS NULL + AND typeof(active_script_name) = 'text' + AND length(active_script_name) > 0 + AND typeof(last_successful_url) = 'text' + AND length(last_successful_url) > 0 + ) + ) ); CREATE UNIQUE INDEX app_deployment_active_app_idx @@ -111,6 +138,25 @@ function createDatabase(): SqliteD1Database { ON app_deployment_run (app_id) WHERE status IN ('queued', 'preparing', 'building', 'submitting', 'submitted', 'activating'); + CREATE TABLE app_deployment_script ( + attempt_count integer NOT NULL, + command_id text NOT NULL, + delivery_generation integer NOT NULL, + deployment_id text NOT NULL, + external_deleted_at integer, + last_reconciled_at integer, + next_reconcile_at integer, + reconcile_count integer DEFAULT 0 NOT NULL, + reconcile_expires_at integer, + reconcile_owner text, + registered_at integer NOT NULL, + registered_claim_owner text NOT NULL, + retire_after integer, + run_id text NOT NULL, + script_name text PRIMARY KEY NOT NULL, + upload_started_at integer + ); + CREATE TABLE api_command ( attempt_count integer DEFAULT 0 NOT NULL, claim_expires_at integer, @@ -118,6 +164,7 @@ function createDatabase(): SqliteD1Database { completed_at integer, created_at integer NOT NULL, dedupe_key text NOT NULL, + delivery_generation integer DEFAULT 1 NOT NULL, id text PRIMARY KEY NOT NULL, kind text NOT NULL, last_error_code text, @@ -140,30 +187,74 @@ function createDatabase(): SqliteD1Database { VALUES ('${APP_ID}', '01J00000000000000000000006', '${OWNER_ID}', 'App', 1, 1); `); + for (const trigger of MANAGED_PROD_SCHEMA_TRIGGERS.filter( + ({ name }) => + name.startsWith("app_deployment_script_") || + name.startsWith("app_deployment_run_target_script_") || + name === "app_deployment_run_terminal_script_retire", + )) { + database.execute(trigger.sql); + } + return database; } +function createAppDeploymentWorkflowStub( + options: { + failCreate?: () => boolean; + onCreate?: (message: ApiCommandMessage) => void; + } = {}, +) { + const created: ApiCommandMessage[] = []; + const instanceIds = new Set(); + const binding = { + async createBatch(batch: { id?: string; params?: ApiCommandMessage }[]) { + if (options.failCreate?.() === true) { + throw new Error("Workflow response timed out."); + } + return batch.map(({ id, params }) => { + if (id === undefined || params === undefined) { + throw new Error("Expected an exact Workflow instance ID and payload."); + } + created.push(params); + instanceIds.add(id); + options.onCreate?.(params); + return { id }; + }); + }, + async get(id: string) { + return { + id, + async status() { + return { status: instanceIds.has(id) ? "queued" : "unknown" }; + }, + }; + }, + } as unknown as Workflow; + return { binding, created }; +} + function createBindings(database: SqliteD1Database) { const queue = createApiCommandQueueStub(); + const workflow = createAppDeploymentWorkflowStub(); return { bindings: { + APP_DEPLOYMENT_DISPATCHER: { + get: () => ({ fetch: async () => new Response("candidate ready") }), + }, + APP_DEPLOYMENT_WORKFLOW: workflow.binding, API_COMMAND_QUEUE: queue, CLOUDFLARE_ACCOUNT_ID: "test-account", CLOUDFLARE_API_TOKEN: "test-token", CLOUDFLARE_ZONE_ID: "test-zone", DB: database, MOSOO_APP_DEPLOYMENT_DOMAIN: "apps.localhost", - } as Pick< - ApiBindings, - | "API_COMMAND_QUEUE" - | "CLOUDFLARE_ACCOUNT_ID" - | "CLOUDFLARE_API_TOKEN" - | "CLOUDFLARE_ZONE_ID" - | "DB" - | "MOSOO_APP_DEPLOYMENT_DOMAIN" - >, + MOSOO_APP_DISPATCH_NAMESPACE: "mosoo-app-deployments-stage", + SANDBOX_FILE_BUCKET_LOCAL: "true", + } as ApiBindings, queue, + workflow, }; } @@ -213,9 +304,10 @@ async function seedExpiredRunningDispatch( await setDeploymentRunUpdatedAt(database, runId, 1); } -async function seedExhaustedRunningDispatch( +async function seedThirdRunningDispatch( database: SqliteD1Database, runId: string, + claimExpiresAt: number, ): Promise { await database .prepare( @@ -229,7 +321,7 @@ async function seedExhaustedRunningDispatch( WHERE dedupe_key = ?`, ) .bind( - NOW_MS + 60_000, + claimExpiresAt, APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS, createAppDeploymentRunDispatchDedupeKey(runId), ) @@ -273,6 +365,172 @@ async function seedQueuedDispatch(database: SqliteD1Database, runId: string): Pr .run(); } +async function claimDeploymentAttemptForTest( + database: SqliteD1Database, + runId: string, + input: { attemptCount?: number; claimOwner?: string; deliveryGeneration?: number } = {}, +): Promise { + const command = await database + .app() + .select({ + deliveryGeneration: apiCommandsTable.deliveryGeneration, + id: apiCommandsTable.id, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.dedupeKey, createAppDeploymentRunDispatchDedupeKey(runId))) + .limit(1) + .get(); + + if (command === undefined) { + throw new Error("Expected an App deployment dispatch command."); + } + + const attemptCount = input.attemptCount ?? 1; + const claimOwner = input.claimOwner ?? `test-claim-${attemptCount}`; + const deliveryGeneration = input.deliveryGeneration ?? command.deliveryGeneration; + await database + .prepare( + `UPDATE api_command + SET attempt_count = ?, + claim_expires_at = ?, + claim_owner = ?, + delivery_generation = ?, + status = 'running' + WHERE id = ?`, + ) + .bind(attemptCount, currentTimestampMs() + 60_000, claimOwner, deliveryGeneration, command.id) + .run(); + + const requireOwnership = async (): Promise => { + const owned = await database + .prepare( + `SELECT 1 + FROM api_command + WHERE id = ? + AND delivery_generation = ? + AND attempt_count = ? + AND claim_owner = ? + AND status = 'running' + AND claim_expires_at > ?`, + ) + .bind(command.id, deliveryGeneration, attemptCount, claimOwner, currentTimestampMs()) + .first(); + + if (owned === null) { + throw new Error("App deployment test attempt lost ownership."); + } + }; + + return { + attemptCount, + claimOwner, + commandId: command.id, + deliveryGeneration, + requireOwnership, + }; +} + +async function commitDeploymentOutcomeForTest( + database: SqliteD1Database, + authority: AppDeploymentAttemptAuthority, + outcome: AppDeploymentDispatchOutcome, +): Promise { + const nowMs = currentTimestampMs(); + const commandTerminal = database + .prepare( + `UPDATE api_command + SET claim_expires_at = NULL, + claim_owner = NULL, + completed_at = ?, + status = ?, + updated_at = ? + WHERE id = ? + AND delivery_generation = ? + AND attempt_count = ? + AND claim_owner = ? + AND status = 'running'`, + ) + .bind( + nowMs, + outcome.kind === "terminal_failure" ? "failed" : "succeeded", + nowMs, + authority.commandId, + authority.deliveryGeneration, + authority.attemptCount, + authority.claimOwner, + ); + + if (outcome.kind === "skipped") { + await commandTerminal.run(); + return; + } + + if (outcome.kind === "terminal_failure") { + await database.batch([ + database + .prepare( + `UPDATE app_deployment_run + SET error_code = ?, error_message = ?, status = 'failed', updated_at = ? + WHERE id = ? + AND status IN ('queued', 'preparing', 'building', 'submitting', 'submitted', 'activating')`, + ) + .bind(outcome.errorCode, outcome.errorMessage, nowMs, outcome.runId), + commandTerminal, + ]); + return; + } + + await database.batch([ + database + .prepare( + `UPDATE app_deployment + SET active_script_name = ?, last_successful_url = ?, updated_at = ? + WHERE id = ? AND latest_run_id = ? AND deleted_at IS NULL`, + ) + .bind(outcome.activeScriptName, outcome.url, nowMs, outcome.deploymentId, outcome.runId), + database + .prepare( + `UPDATE app_deployment_run + SET error_code = NULL, + error_message = NULL, + external_deployment_id = ?, + external_project_id = ?, + external_version_id = ?, + status = 'success', + updated_at = ?, + url = ? + WHERE id = ? AND status = 'activating' AND target_script_name = ?`, + ) + .bind( + outcome.externalDeploymentId, + outcome.externalProjectId, + outcome.externalVersionId, + nowMs, + outcome.url, + outcome.runId, + outcome.activeScriptName, + ), + commandTerminal, + ]); +} + +async function dispatchDeploymentForTest( + database: SqliteD1Database, + bindings: ApiBindings, + runId: string, + options: DispatchAppDeploymentRunOptions = {}, +): Promise { + const authority = await claimDeploymentAttemptForTest(database, runId); + const outcome = await dispatchAppDeploymentRun( + bindings, + { appDeploymentRunId: runId as AppDeploymentRunId }, + authority, + options, + ); + await commitDeploymentOutcomeForTest(database, authority, outcome); + return outcome; +} + async function seedDeployment( database: SqliteD1Database, input: { appId: string; deploymentId: string }, @@ -348,39 +606,18 @@ function createCloudflareDeleteRecorder( overrides: Partial = {}, ): CloudflareDeploymentClient { return { - async deletePagesDomain(input) { - deleted.push(`pages-domain:${input.hostname}`); - }, - async deletePagesProject(input) { - deleted.push(`pages:${input.projectName}`); - }, - async deleteWorkerDomain(input) { - deleted.push(`worker-domain:${input.hostname}`); - }, - async deleteWorkerRoute(input) { - deleted.push(`worker-route:${input.hostname}`); + async createStaticAssetsUploadSession() { + throw new Error("Unexpected Static Assets upload session."); }, async deleteWorkerScript(input) { deleted.push(`worker:${input.scriptName}`); }, + async deployStaticAssets() { + throw new Error("Unexpected Static Assets deploy."); + }, async deployWorkerModule() { throw new Error("Unexpected Worker deploy."); }, - async ensurePagesProject() { - throw new Error("Unexpected Pages project creation."); - }, - async ensurePagesDomain() { - throw new Error("Unexpected Pages domain creation."); - }, - async ensureWorkerDomain() { - throw new Error("Unexpected Worker domain creation."); - }, - async ensureWorkerRoute() { - throw new Error("Unexpected Worker route creation."); - }, - async getLatestPagesDeployment() { - throw new Error("Unexpected Pages deployment read."); - }, ...overrides, }; } @@ -416,16 +653,19 @@ function createTestSandboxHandle( readFile: unexpectedSandboxCall, setKeepAlive: async () => {}, writeFile: unexpectedSandboxCall, - } as SandboxHandle; + }; } const session = { - exec: async (command: string) => { - events.push(`${id}:session:${command}`); + exec: async (command: string, options?: { timeout?: number }) => { + events.push(`${id}:session:timeout=${String(options?.timeout)}:${command}`); return successfulCommandResult(); }, mkdir: unexpectedSandboxCall, - readFile: unexpectedSandboxCall, + readFile: async (path: string) => ({ + content: path.endsWith("completion-token") ? "completion-jwt" : "", + encoding: "utf8" as const, + }), startProcess: unexpectedSandboxCall, watch: unexpectedSandboxCall, writeFile: unexpectedSandboxCall, @@ -433,28 +673,163 @@ function createTestSandboxHandle( return { ...base, - createSession: async () => session, + createSession: async (options) => { + events.push( + `${id}:session-env=${Object.keys(options?.env ?? {}) + .toSorted() + .join(",")}`, + ); + return session; + }, destroy: async () => { events.push(`${id}:destroy`); }, - exec: async (command) => { - events.push(`${id}:${command}`); - return successfulCommandResult(command.includes("find . -type f") ? "./index.html\n" : ""); + exec: async (command, options) => { + events.push(`${id}:timeout=${String(options?.timeout)}:${command}`); + return successfulCommandResult( + command.includes("find . -type f") + ? "./.mosoo.toml\n./_headers\n./_redirects\n./index.html\n" + : "", + ); }, - readFile: async (_path, options) => ({ - content: options?.encoding === "base64" ? "YXJjaGl2ZQ==" : "
Hello
", + readFile: async (path, options) => ({ + content: path.endsWith("inventory.json") + ? JSON.stringify({ + assetsIgnore: "private/**\n", + headers: "/assets/*\n Cache-Control: public\n", + paths: [".assetsignore", "_headers", "_redirects", "index.html", "private/secret.txt"], + redirects: "/old /new 301\n\n/* /index.html 200\n", + }) + : path.endsWith("manifest.json") + ? JSON.stringify({ + "/index.html": { + hash: "0123456789abcdef0123456789abcdef", + size: 18, + }, + }) + : options?.encoding === "base64" + ? "YXJjaGl2ZQ==" + : path.endsWith(".mosoo.toml") + ? 'type = "static"\n\n[build]\noutput = "."\n\n[routes]\nfallback = "index.html"\n' + : "
Hello
", encoding: options?.encoding ?? "utf8", }), + mountBucket: async (bucket, mountPath, options) => { + events.push( + `${id}:mount:${bucket}:${mountPath}:${options.prefix}:${String(options.readOnly)}`, + ); + }, setKeepAlive: async (keepAlive) => { events.push(`${id}:keep-alive:${String(keepAlive)}`); }, - writeFile: async (path) => { - events.push(`${id}:write:${path}`); + unmountBucket: async (mountPath) => { + events.push(`${id}:unmount:${mountPath}`); + }, + writeFile: async (path, content) => { + events.push( + path.endsWith("selection.json") || path.endsWith("buckets.json") + ? `${id}:write:${path}:${content}` + : `${id}:write:${path}`, + ); + }, + }; +} + +function createLocalStaticDeploymentSandboxHandle( + id: string, + root: string, + populateRepository: (repoDir: string) => Promise, +): SandboxHandle { + const sandboxWorkDir = `/tmp/mosoo-app-deployment-${id}`; + const workDir = join(root, id); + const repoDir = join(workDir, "repo"); + const bucketDir = join(root, "bucket"); + const mapPath = (path: string) => + path + .replaceAll(sandboxWorkDir, workDir) + .replaceAll("/mnt/mosoo-app-deployment-artifact", bucketDir); + const execute = async (command: string) => { + const process = Bun.spawn(["sh", "-lc", mapPath(command)], { + stderr: "pipe", + stdout: "pipe", + }); + const [exitCode, stderr, stdout] = await Promise.all([ + process.exited, + new Response(process.stderr).text(), + new Response(process.stdout).text(), + ]); + + return { exitCode, stderr, stdout, success: exitCode === 0 }; + }; + const base = { + configureNetworkConstraints: async () => {}, + createBackup: unexpectedSandboxCall, + deleteSession: unexpectedSandboxCall, + getSession: unexpectedSandboxCall, + mkdir: unexpectedSandboxCall, + restoreBackup: unexpectedSandboxCall, + startProcess: unexpectedSandboxCall, + terminal: unexpectedSandboxCall, + watch: unexpectedSandboxCall, + wsConnect: unexpectedSandboxCall, + }; + + return { + ...base, + createSession: async () => ({ + exec: async (command) => + command.includes(" upload ") ? successfulCommandResult() : execute(command), + mkdir: unexpectedSandboxCall, + readFile: async (path, options) => ({ + content: path.endsWith("completion-token") + ? "completion-jwt" + : await readFile(mapPath(path), options?.encoding === "base64" ? "base64" : "utf8"), + encoding: options?.encoding ?? "utf8", + }), + startProcess: unexpectedSandboxCall, + watch: unexpectedSandboxCall, + writeFile: unexpectedSandboxCall, + }), + destroy: async () => {}, + exec: async (command) => { + if (command.includes("git clone --no-tags")) { + await rm(workDir, { force: true, recursive: true }); + await mkdir(repoDir, { recursive: true }); + await populateRepository(repoDir); + return successfulCommandResult(); + } + + return execute(command); + }, + mountBucket: async () => { + await mkdir(bucketDir, { recursive: true }); }, - } as SandboxHandle; + readFile: async (path, options) => ({ + content: path.endsWith("completion-token") + ? "completion-jwt" + : await readFile(mapPath(path), options?.encoding === "base64" ? "base64" : "utf8"), + encoding: options?.encoding ?? "utf8", + }), + setKeepAlive: async () => {}, + unmountBucket: async () => {}, + writeFile: async (path, content) => { + await writeFile(mapPath(path), content); + }, + }; } -function createWorkerDeploymentSandboxHandle(id: string, events: string[]): SandboxHandle { +function createWorkerDeploymentSandboxHandle( + id: string, + events: string[], + fixture: { + bundleContent?: string; + bundleError?: string; + bundleProjectRoot?: string; + rootDir?: string; + } = {}, +): SandboxHandle { + const rootDir = fixture.rootDir ?? "."; + let generatedBundleContent: string | undefined; const base = { configureNetworkConstraints: async () => {}, createBackup: unexpectedSandboxCall, @@ -476,44 +851,232 @@ function createWorkerDeploymentSandboxHandle(id: string, events: string[]): Sand destroy: async () => { events.push(`${id}:destroy`); }, - exec: async (command) => { - events.push(`${id}:${command}`); + exec: async (command, execOptions) => { + events.push(`${id}:timeout=${String(execOptions?.timeout)}:${command}`); + + if (command.includes("bun build") && fixture.bundleError !== undefined) { + return { + exitCode: 1, + stderr: fixture.bundleError, + stdout: "", + success: false as const, + }; + } + + if (command.includes("bun build") && fixture.bundleProjectRoot !== undefined) { + let result: Awaited>; + + try { + result = await Bun.build({ + entrypoints: [join(fixture.bundleProjectRoot, "src/index.js")], + external: ["cloudflare:*"], + format: "esm", + target: "browser", + }); + } catch (error) { + return { + exitCode: 1, + stderr: + error instanceof AggregateError + ? error.errors + .map((cause) => (cause instanceof Error ? cause.message : String(cause))) + .join("\n") + : error instanceof Error + ? error.message + : String(error), + stdout: "", + success: false as const, + }; + } + + if (!result.success) { + return { + exitCode: 1, + stderr: result.logs.map((log) => log.message).join("\n"), + stdout: "", + success: false as const, + }; + } + + generatedBundleContent = await result.outputs[0].text(); + } + return successfulCommandResult( command.includes("find . -type f -print") - ? "./.mosoo.toml\n./wrangler.toml\n./src/index.js\n" + ? `./.mosoo.toml\n./${rootDir === "." ? "" : `${rootDir}/`}src/index.js\n` : "", ); }, - readFile: async (path, options) => { - let content = "export default { fetch() { return new Response('ok'); } };\n"; + readFile: async (path, readOptions) => { + let content = + generatedBundleContent ?? + fixture.bundleContent ?? + "export default { fetch() { return new Response('ok'); } };\n"; if (path.endsWith(".mosoo.toml")) { content = [ "schema = 1", 'name = "worker-app"', + ...(rootDir === "." ? [] : [`root = "${rootDir}"`]), "", "[deploy]", 'adapter = "cloudflare-workers"', - 'wrangler = "wrangler.toml"', + "", + "[worker]", + 'entry = "src/index.js"', "", ].join("\n"); - } else if (path.endsWith("wrangler.toml")) { - content = 'name = "worker-app"\nmain = "src/index.js"\n'; } - return { content, encoding: options?.encoding ?? "utf8" }; + return { content, encoding: readOptions?.encoding ?? "utf8" }; }, setKeepAlive: async (keepAlive) => { events.push(`${id}:keep-alive:${String(keepAlive)}`); }, writeFile: unexpectedSandboxCall, - } as SandboxHandle; + }; } describe("app deployment service", () => { - test("creates a deployment run and queues dispatch", async () => { + test("salts Static Assets hashes by script and extension and remains deterministic", async () => { + const root = await mkdtemp(join(tmpdir(), "mosoo-static-assets-hash-")); + + try { + const artifactDir = join(root, "artifact"); + const selectionPath = join(root, "selection.json"); + const toolPath = join(root, "static-assets.mjs"); + const firstManifestPath = join(root, "manifest-a.json"); + const repeatedManifestPath = join(root, "manifest-a-repeated.json"); + const secondManifestPath = join(root, "manifest-b.json"); + const content = "same bytes"; + await mkdir(artifactDir); + await writeFile(join(artifactDir, "index.html"), content); + await writeFile(join(artifactDir, "same.css"), content); + await writeFile(join(artifactDir, "same.js"), content); + await writeFile(selectionPath, JSON.stringify(["index.html", "same.css", "same.js"])); + await writeFile(toolPath, APP_DEPLOYMENT_STATIC_ASSET_TOOL_SOURCE); + + for (const [scriptName, outputPath] of [ + ["app-command-g1-a1", firstManifestPath], + ["app-command-g1-a1", repeatedManifestPath], + ["app-command-g1-a2", secondManifestPath], + ] as const) { + const process = Bun.spawn([ + "node", + toolPath, + "manifest", + artifactDir, + selectionPath, + outputPath, + scriptName, + ]); + expect(await process.exited).toBe(0); + } + + const first = JSON.parse(await readFile(firstManifestPath, "utf8")); + const repeated = JSON.parse(await readFile(repeatedManifestPath, "utf8")); + const second = JSON.parse(await readFile(secondManifestPath, "utf8")); + const expectedHash = createHash("sha256") + .update("app-command-g1-a1") + .update(Buffer.from([0])) + .update("html") + .update(Buffer.from([0])) + .update(content) + .digest("hex") + .slice(0, 32); + + expect(first["/index.html"]).toEqual({ hash: expectedHash, size: content.length }); + expect(repeated).toEqual(first); + expect(second["/index.html"].hash).not.toBe(expectedHash); + expect(first["/same.css"].hash).not.toBe(first["/same.js"].hash); + } finally { + await rm(root, { force: true, recursive: true }); + } + }); + + test("prunes Git metadata from a root Static Assets artifact before manifest creation", async () => { + const root = await mkdtemp(join(tmpdir(), "mosoo-static-assets-root-")); + + try { + const database = createDatabase(); + const { bindings } = createBindings(database); + let manifest: Record = {}; + bindings.runtimeSubjectHandleFactory = (runtimeSubjectId) => + createLocalStaticDeploymentSandboxHandle(runtimeSubjectId, root, async (repoDir) => { + await mkdir(join(repoDir, ".git"), { recursive: true }); + await writeFile(join(repoDir, ".assetsignore"), "!/.git/**\n"); + await writeFile(join(repoDir, ".git", "config"), "credential = secret\n"); + await writeFile(join(repoDir, "index.html"), "
Root static
"); + }); + const cloudflareClient = createCloudflareDeleteRecorder([], { + async createStaticAssetsUploadSession(input) { + manifest = input.manifest; + return { + buckets: [Object.values(input.manifest).map(({ hash }) => hash)], + uploadToken: "eyJhbGciOiJub25lIn0.e30.signature", + }; + }, + async deployStaticAssets() { + return { deploymentId: null, versionId: null }; + }, + }); + const run = await deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS }, + ); + + const outcome = await dispatchDeploymentForTest(database, bindings, run.id, { + cloudflareClient, + }); + + expect(outcome).toMatchObject({ kind: "succeeded" }); + expect(Object.keys(manifest)).toEqual(["/index.html"]); + } finally { + await rm(root, { force: true, recursive: true }); + } + }); + + test("rejects a Static Assets output symlink outside the repository without retrying", async () => { + const root = await mkdtemp(join(tmpdir(), "mosoo-static-assets-escape-")); + const escapedDir = join(root, "escaped"); + + try { + await mkdir(escapedDir); + await writeFile(join(escapedDir, "index.html"), "
Escaped
"); + await writeFile(join(escapedDir, "_redirects"), "sentinel\n"); + const database = createDatabase(); + const { bindings } = createBindings(database); + bindings.runtimeSubjectHandleFactory = (runtimeSubjectId) => + createLocalStaticDeploymentSandboxHandle(runtimeSubjectId, root, async (repoDir) => { + await writeFile( + join(repoDir, ".mosoo.toml"), + 'type = "static"\n\n[build]\noutput = "public"\n\n[routes]\nfallback = "index.html"\n', + ); + await symlink(escapedDir, join(repoDir, "public"), "dir"); + }); + const run = await deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS }, + ); + + await expect(dispatchDeploymentForTest(database, bindings, run.id)).resolves.toMatchObject({ + errorCode: "AppDeploymentNonRetryableError", + errorMessage: "Static Assets output escapes repository root.", + kind: "terminal_failure", + }); + expect(await readFile(join(escapedDir, "_redirects"), "utf8")).toBe("sentinel\n"); + } finally { + await rm(root, { force: true, recursive: true }); + } + }); + + test("creates a deployment run and starts its durable Workflow", async () => { const database = createDatabase(); - const { bindings, queue } = createBindings(database); + const { bindings, workflow } = createBindings(database); const run = await deployApp( bindings, @@ -533,7 +1096,7 @@ describe("app deployment service", () => { sourceCommitSha: "abc123", status: "queued", }); - expect(queue.sent).toHaveLength(1); + expect(workflow.created).toHaveLength(1); const command = await database.app().select().from(apiCommandsTable).limit(1).get(); @@ -544,6 +1107,10 @@ describe("app deployment service", () => { expect(JSON.parse(command?.payloadJson ?? "{}")).toEqual({ appDeploymentRunId: run.id, }); + expect(workflow.created[0]).toEqual({ + commandId: command?.id, + deliveryGeneration: 1, + }); const deployment = await getAppDeployment(bindings, VIEWER, APP_ID); expect(deployment?.latestRun?.id).toBe(run.id); @@ -558,18 +1125,14 @@ describe("app deployment service", () => { const database = createDatabase(); const { bindings } = createBindings(database); const deliveredCommandIds: string[] = []; - let queueUnavailable = true; + let workflowUnavailable = true; + const workflow = createAppDeploymentWorkflowStub({ + failCreate: () => workflowUnavailable, + onCreate: (message) => deliveredCommandIds.push(message.commandId), + }); const deferredBindings = { ...bindings, - API_COMMAND_QUEUE: { - async send(input: { commandId: string }): Promise { - if (queueUnavailable) { - throw new Error("Queue response timed out."); - } - - deliveredCommandIds.push(input.commandId); - }, - }, + APP_DEPLOYMENT_WORKFLOW: workflow.binding, }; const run = await deployApp( @@ -603,7 +1166,7 @@ describe("app deployment service", () => { }); expect(runRow).toEqual({ errorCode: null, status: "queued" }); - queueUnavailable = false; + workflowUnavailable = false; await createApiWorker().scheduled( { scheduledTime: NOW_MS } as ScheduledController, deferredBindings as ApiBindings, @@ -623,8 +1186,8 @@ describe("app deployment service", () => { async build() {}, async deploy() { return { - externalDeploymentId: "pages-deployment-after-redrive", - externalProjectId: "pages-project-after-redrive", + externalDeploymentId: null, + externalProjectId: null, externalVersionId: null, url: `https://app-${APP_ID.toLowerCase()}.apps.localhost`, }; @@ -640,11 +1203,7 @@ describe("app deployment service", () => { expect(deliveredCommandIds).toContain(redrivenCommand?.id); expect(redrivenCommand).toMatchObject({ lastErrorCode: null, status: "queued" }); - await dispatchAppDeploymentRun( - deferredBindings as ApiBindings, - { appDeploymentRunId: run.id }, - { runner }, - ); + await dispatchDeploymentForTest(database, deferredBindings as ApiBindings, run.id, { runner }); const status = await getAppDeploymentStatus(deferredBindings, VIEWER, APP_ID); expect(status).toMatchObject({ status: "success" }); @@ -671,45 +1230,103 @@ describe("app deployment service", () => { ).rejects.toThrow("An App deployment run is already active."); }); - test("recovers an active deployment run without a dispatch command", async () => { + test("a losing concurrent deploy cannot overwrite the winner repository", async () => { const database = createDatabase(); const { bindings } = createBindings(database); - const firstRun = await deployApp( bindings, VIEWER, { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, { fetch: githubFetch, nowMs: () => NOW_MS }, ); + await database + .prepare("UPDATE app_deployment_run SET status = 'failed' WHERE id = ?") + .bind(firstRun.id) + .run(); + await database.prepare("UPDATE api_command SET status = 'failed'").run(); + + const originalPrepare = database.prepare.bind(database); + let injected = false; + const wrapRunInsert = (statement: D1PreparedStatement): D1PreparedStatement => ({ + ...statement, + bind: (...values) => wrapRunInsert(statement.bind(...values)), + run: async () => { + if (!injected) { + injected = true; + database.execute(` + INSERT INTO app_deployment_run ( + app_id, created_at, deployment_id, id, source_branch, source_commit_sha, + status, updated_at + ) VALUES ( + '${APP_ID}', ${NOW_MS + 1}, '${firstRun.deploymentId}', + '${runListRunId(999)}', 'main', 'winner-sha', 'queued', ${NOW_MS + 1} + ) + `); + } + + return statement.run(); + }, + }); + database.prepare = ((query: string) => { + const statement = originalPrepare(query); + return /insert\s+into\s+"app_deployment_run"/iu.test(query) + ? wrapRunInsert(statement) + : statement; + }) as D1Database["prepare"]; + + const losingFetch: typeof fetch = async (input) => { + const url = input instanceof URL ? input.href : typeof input === "string" ? input : input.url; + + if (url === "https://api.github.com/repos/loser/repository") { + return Response.json({ + clone_url: "https://github.com/loser/repository.git", + default_branch: "trunk", + name: "repository", + owner: { login: "loser" }, + private: false, + }); + } + if (url === "https://api.github.com/repos/loser/repository/branches/trunk") { + return Response.json({ commit: { sha: "loser-sha" } }); + } - await database.prepare("DELETE FROM api_command").run(); - await setDeploymentRunUpdatedAt(database, firstRun.id, 1); - - const secondRun = await deployApp( - bindings, - VIEWER, - { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, - { fetch: githubFetch, nowMs: () => NOW_MS + 1 }, - ); + throw new Error(`Unexpected GitHub request: ${url}`); + }; - expect(secondRun.id).not.toBe(firstRun.id); - expect(secondRun.status).toBe("queued"); + try { + await expect( + deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/loser/repository" }, + { fetch: losingFetch, nowMs: () => NOW_MS + 1 }, + ), + ).rejects.toThrow(); + } finally { + database.prepare = originalPrepare; + } - const firstRunRow = await database + const deployment = await database .app() - .select() - .from(appDeploymentRunsTable) - .where(eq(appDeploymentRunsTable.id, firstRun.id)) - .limit(1) + .select({ + latestRunId: appDeploymentsTable.latestRunId, + repoName: appDeploymentsTable.repoName, + repoOwner: appDeploymentsTable.repoOwner, + repoUrl: appDeploymentsTable.repoUrl, + }) + .from(appDeploymentsTable) + .where(eq(appDeploymentsTable.id, firstRun.deploymentId)) .get(); - expect(firstRunRow).toMatchObject({ - errorCode: "deployment_dispatch_missing", - status: "failed", + expect(deployment).toEqual({ + latestRunId: firstRun.id, + repoName: "awire", + repoOwner: "samzong", + repoUrl: "https://github.com/samzong/awire.git", }); }); - test("recovers an active deployment run with an expired running dispatch command", async () => { + test("recovers an active deployment run without a dispatch command", async () => { const database = createDatabase(); const { bindings } = createBindings(database); @@ -720,7 +1337,8 @@ describe("app deployment service", () => { { fetch: githubFetch, nowMs: () => NOW_MS }, ); - await seedExpiredRunningDispatch(database, firstRun.id); + await database.prepare("DELETE FROM api_command").run(); + await setDeploymentRunUpdatedAt(database, firstRun.id, 1); const secondRun = await deployApp( bindings, @@ -741,12 +1359,12 @@ describe("app deployment service", () => { .get(); expect(firstRunRow).toMatchObject({ - errorCode: "deployment_dispatch_expired", + errorCode: "deployment_dispatch_missing", status: "failed", }); }); - test("recovers an active deployment run with an expired running dispatch from reads", async () => { + test("keeps a Workflow deployment active after its command lease expires", async () => { const database = createDatabase(); const { bindings } = createBindings(database); @@ -760,27 +1378,36 @@ describe("app deployment service", () => { await seedExpiredRunningDispatch(database, run.id); await expect(getAppDeploymentStatus(bindings, VIEWER, APP_ID)).resolves.toMatchObject({ - errorCode: "deployment_dispatch_expired", + errorCode: null, id: run.id, - status: "failed", + status: "queued", }); await expect(getAppDeployment(bindings, VIEWER, APP_ID)).resolves.toMatchObject({ latestRun: { - errorCode: "deployment_dispatch_expired", + errorCode: null, id: run.id, - status: "failed", + status: "queued", }, }); await expect(listAppDeploymentRuns(bindings, VIEWER, APP_ID, 10)).resolves.toEqual([ expect.objectContaining({ - errorCode: "deployment_dispatch_expired", + errorCode: null, id: run.id, - status: "failed", + status: "queued", }), ]); + await expect( + deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS + 1 }, + ), + ).rejects.toThrow("An App deployment run is already active."); + const runRow = await database .app() .select() @@ -790,12 +1417,58 @@ describe("app deployment service", () => { .get(); expect(runRow).toMatchObject({ - errorCode: "deployment_dispatch_expired", - status: "failed", + errorCode: null, + status: "queued", }); }); - test("fails an active deployment run when dispatch retries are exhausted", async () => { + test("never terminalizes a third running Workflow attempt from a read path", async () => { + for (const claimExpiresAt of [currentTimestampMs() + 60_000, 1]) { + const database = createDatabase(); + const { bindings } = createBindings(database); + const run = await deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS }, + ); + + await seedThirdRunningDispatch(database, run.id, claimExpiresAt); + + await expect(getAppDeploymentStatus(bindings, VIEWER, APP_ID)).resolves.toMatchObject({ + errorCode: null, + id: run.id, + status: "queued", + }); + await expect( + deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS + 1 }, + ), + ).rejects.toThrow("An App deployment run is already active."); + + const command = await database + .app() + .select({ + attemptCount: apiCommandsTable.attemptCount, + lastErrorCode: apiCommandsTable.lastErrorCode, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.dedupeKey, createAppDeploymentRunDispatchDedupeKey(run.id))) + .limit(1) + .get(); + expect(command).toMatchObject({ + attemptCount: APP_DEPLOYMENT_RUN_DISPATCH_MAX_ATTEMPTS, + lastErrorCode: "SandboxError", + status: "running", + }); + } + }); + + test("keeps a fresh active deployment run without a dispatch command active", async () => { const database = createDatabase(); const { bindings } = createBindings(database); @@ -803,58 +1476,11 @@ describe("app deployment service", () => { bindings, VIEWER, { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, - { fetch: githubFetch, nowMs: () => NOW_MS }, + { fetch: githubFetch, nowMs: currentTimestampMs }, ); - await seedExhaustedRunningDispatch(database, firstRun.id); - - await expect(getAppDeploymentStatus(bindings, VIEWER, APP_ID)).resolves.toMatchObject({ - errorCode: APP_DEPLOYMENT_RUN_DISPATCH_RETRY_EXHAUSTED_CODE, - errorMessage: expect.stringContaining("Container is starting"), - id: firstRun.id, - status: "failed", - }); - - const dispatchCommand = await database - .app() - .select({ - lastErrorCode: apiCommandsTable.lastErrorCode, - status: apiCommandsTable.status, - }) - .from(apiCommandsTable) - .where(eq(apiCommandsTable.dedupeKey, createAppDeploymentRunDispatchDedupeKey(firstRun.id))) - .limit(1) - .get(); - - expect(dispatchCommand).toMatchObject({ - lastErrorCode: APP_DEPLOYMENT_RUN_DISPATCH_RETRY_EXHAUSTED_CODE, - status: "failed", - }); - - const secondRun = await deployApp( - bindings, - VIEWER, - { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, - { fetch: githubFetch, nowMs: () => NOW_MS + 1 }, - ); - - expect(secondRun.id).not.toBe(firstRun.id); - expect(secondRun.status).toBe("queued"); - }); - - test("keeps a fresh active deployment run without a dispatch command active", async () => { - const database = createDatabase(); - const { bindings } = createBindings(database); - - const firstRun = await deployApp( - bindings, - VIEWER, - { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, - { fetch: githubFetch, nowMs: currentTimestampMs }, - ); - - await database.prepare("DELETE FROM api_command").run(); - await setDeploymentRunUpdatedAt(database, firstRun.id, currentTimestampMs()); + await database.prepare("DELETE FROM api_command").run(); + await setDeploymentRunUpdatedAt(database, firstRun.id, currentTimestampMs()); await expect( deployApp( @@ -884,16 +1510,24 @@ describe("app deployment service", () => { const { bindings } = createBindings(database); const targetUrl = `https://app-${APP_ID.toLowerCase()}.apps.localhost`; let buildPlanName: string | null = null; + let uploadWasFenced = false; const runner: AppDeploymentBuildRunner = { async build({ plan }) { buildPlanName = plan.generatedWranglerConfig; }, - async deploy() { + async deploy({ activeScriptName }) { + const script = await database + .app() + .select({ uploadStartedAt: appDeploymentScriptsTable.uploadStartedAt }) + .from(appDeploymentScriptsTable) + .where(eq(appDeploymentScriptsTable.scriptName, activeScriptName)) + .limit(1) + .get(); + uploadWasFenced = typeof script?.uploadStartedAt === "number"; return { - externalDeploymentId: "pages-deployment-1", - externalProjectId: "pages-project-1", + externalDeploymentId: null, + externalProjectId: null, externalVersionId: null, - url: targetUrl, }; }, async prepare() { @@ -911,34 +1545,384 @@ describe("app deployment service", () => { { fetch: githubFetch, nowMs: () => NOW_MS }, ); - await dispatchAppDeploymentRun( - bindings as ApiBindings, - { appDeploymentRunId: run.id }, - { - runner, - }, - ); + await dispatchDeploymentForTest(database, bindings, run.id, { runner }); const status = await getAppDeploymentStatus(bindings, VIEWER, APP_ID); const deployment = await getAppDeployment(bindings, VIEWER, APP_ID); const runRow = await database.app().select().from(appDeploymentRunsTable).limit(1).get(); - expect(buildPlanName).toContain(`name = "app-${APP_ID.toLowerCase()}"`); + expect(buildPlanName).toContain("[assets]"); + expect(uploadWasFenced).toBe(true); + expect(buildPlanName).not.toContain("name ="); expect(status).toMatchObject({ liveUrl: targetUrl, status: "success", }); expect(deployment?.liveUrl).toBe(targetUrl); expect(runRow).toMatchObject({ - externalDeploymentId: "pages-deployment-1", - externalProjectId: "pages-project-1", + externalDeploymentId: null, + externalProjectId: null, status: "success", - targetKind: "cloudflare_pages", - targetProjectName: `app-${APP_ID.toLowerCase()}`, + targetKind: "cloudflare_static_assets", + targetProjectName: null, + targetScriptName: expect.stringMatching(/^app-[0-9a-z]{26}-g1-a1$/u), url: targetUrl, }); }); + test("isolates deployment sandboxes by delivery generation and attempt", () => { + const runId = runListRunId(1) as AppDeploymentRunId; + const first = { attemptCount: 1, deliveryGeneration: 1 }; + const takeover = { attemptCount: 2, deliveryGeneration: 1 }; + const explicitRetry = { attemptCount: 1, deliveryGeneration: 2 }; + + expect( + new Set([ + appDeploymentBuildSandboxId(runId, first), + appDeploymentBuildSandboxId(runId, takeover), + appDeploymentBuildSandboxId(runId, explicitRetry), + appDeploymentDeploySandboxId(runId, first), + appDeploymentDeploySandboxId(runId, takeover), + appDeploymentDeploySandboxId(runId, explicitRetry), + ]).size, + ).toBe(6); + }); + + test.each(["queued", "preparing", "building", "submitting", "submitted", "activating"] as const)( + "restarts an interrupted %s run from preparing under a new attempt", + async (interruptedStatus) => { + const database = createDatabase(); + const { bindings } = createBindings(database); + const run = await deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS }, + ); + const firstAuthority = await claimDeploymentAttemptForTest(database, run.id); + const firstScriptName = appDeploymentCandidateScriptName(firstAuthority); + await registerAppDeploymentScriptCandidate(database, { + authority: firstAuthority, + deploymentId: run.deploymentId, + nowMs: currentTimestampMs(), + runId: run.id, + scriptName: firstScriptName, + }); + await database + .prepare("UPDATE app_deployment_run SET status = ? WHERE id = ?") + .bind(interruptedStatus, run.id) + .run(); + const authority = await claimDeploymentAttemptForTest(database, run.id, { + attemptCount: 2, + claimOwner: `resume-${interruptedStatus}`, + }); + const secondScriptName = appDeploymentCandidateScriptName(authority); + let statusAtPrepare: string | null = null; + const outcome = await dispatchAppDeploymentRun( + bindings, + { appDeploymentRunId: run.id }, + authority, + { + runner: { + async build() {}, + async deploy() { + return { + externalDeploymentId: null, + externalProjectId: null, + externalVersionId: null, + }; + }, + async prepare() { + statusAtPrepare = + ( + await database + .app() + .select({ status: appDeploymentRunsTable.status }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, run.id)) + .limit(1) + .get() + )?.status ?? null; + return { + repoDir: "/repo", + snapshot: { files: { "index.html": "
Recovered
" } }, + }; + }, + }, + }, + ); + await commitDeploymentOutcomeForTest(database, authority, outcome); + const scripts = await database + .app() + .select() + .from(appDeploymentScriptsTable) + .orderBy(appDeploymentScriptsTable.attemptCount); + const runRow = await database + .app() + .select({ + status: appDeploymentRunsTable.status, + targetScriptName: appDeploymentRunsTable.targetScriptName, + }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, run.id)) + .limit(1) + .get(); + + expect(statusAtPrepare).toBe("preparing"); + expect(secondScriptName).not.toBe(firstScriptName); + expect(outcome).toMatchObject({ kind: "succeeded", runId: run.id }); + expect(runRow).toEqual({ status: "success", targetScriptName: secondScriptName }); + expect(scripts).toHaveLength(2); + expect(scripts[0]).toMatchObject({ + attemptCount: 1, + retireAfter: expect.any(Number), + scriptName: firstScriptName, + }); + expect(scripts[1]).toMatchObject({ + attemptCount: 2, + retireAfter: null, + scriptName: secondScriptName, + uploadStartedAt: expect.any(Number), + }); + }, + ); + + test("bounds the whole dispatch by one deadline and lets the next attempt recover", async () => { + const database = createDatabase(); + const { bindings } = createBindings(database); + const run = await deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS }, + ); + const firstAuthority = await claimDeploymentAttemptForTest(database, run.id); + + await expect( + dispatchAppDeploymentRun(bindings, { appDeploymentRunId: run.id }, firstAuthority, { + deadlineMs: Date.now() + 5, + runner: { + async build() {}, + async deploy() { + throw new Error("Expired dispatch must not deploy."); + }, + async prepare() { + await new Promise((resolve) => setTimeout(resolve, 25)); + return { + repoDir: "/repo", + snapshot: { files: { "index.html": "
Late
" } }, + }; + }, + }, + }), + ).rejects.toThrow(/execution (?:budget|deadline)/u); + + const secondAuthority = await claimDeploymentAttemptForTest(database, run.id, { + attemptCount: 2, + claimOwner: "deadline-recovery", + }); + const recovered = await dispatchAppDeploymentRun( + bindings, + { appDeploymentRunId: run.id }, + secondAuthority, + { + runner: { + async build() {}, + async deploy() { + return { + externalDeploymentId: null, + externalProjectId: null, + externalVersionId: null, + }; + }, + async prepare() { + return { + repoDir: "/repo", + snapshot: { files: { "index.html": "
Recovered
" } }, + }; + }, + }, + }, + ); + + expect(recovered).toMatchObject({ kind: "succeeded", runId: run.id }); + }); + + test("returns deployment success without writing either terminal row", async () => { + const database = createDatabase(); + const { bindings } = createBindings(database); + const run = await deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS }, + ); + const authority = await claimDeploymentAttemptForTest(database, run.id); + const outcome = await dispatchAppDeploymentRun( + bindings, + { appDeploymentRunId: run.id }, + authority, + { + runner: { + async build() {}, + async deploy() { + return { + externalDeploymentId: null, + externalProjectId: null, + externalVersionId: null, + url: "https://owned.apps.localhost", + }; + }, + async prepare() { + return { + repoDir: "/repo", + snapshot: { files: { "index.html": "
Owned
" } }, + }; + }, + }, + }, + ); + const deployment = await database.app().select().from(appDeploymentsTable).limit(1).get(); + const runRow = await database.app().select().from(appDeploymentRunsTable).limit(1).get(); + const command = await database.app().select().from(apiCommandsTable).limit(1).get(); + + expect(outcome).toMatchObject({ kind: "succeeded", runId: run.id }); + expect(deployment?.lastSuccessfulUrl).toBeNull(); + expect(runRow).toMatchObject({ status: "activating", url: null }); + expect(command).toMatchObject({ claimOwner: authority.claimOwner, status: "running" }); + }); + + test("returns non-retryable failure without letting the executor fail the run", async () => { + const database = createDatabase(); + const { bindings } = createBindings(database); + const run = await deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS }, + ); + const authority = await claimDeploymentAttemptForTest(database, run.id); + const outcome = await dispatchAppDeploymentRun( + bindings, + { appDeploymentRunId: run.id }, + authority, + { + runner: { + async build() { + throw new AppDeploymentNonRetryableError("Unsupported build."); + }, + async deploy() { + throw new Error("Deploy must not run."); + }, + async prepare() { + return { + repoDir: "/repo", + snapshot: { files: { "index.html": "
Owned
" } }, + }; + }, + }, + }, + ); + const runRow = await database.app().select().from(appDeploymentRunsTable).limit(1).get(); + + expect(outcome).toEqual({ + errorCode: "AppDeploymentNonRetryableError", + errorMessage: "Unsupported build.", + kind: "terminal_failure", + runId: run.id, + }); + expect(runRow?.status).toBe("building"); + }); + + test("stops a stale attempt after takeover without moving state backward or deploying", async () => { + const database = createDatabase(); + const { bindings } = createBindings(database); + const run = await deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS }, + ); + const firstAuthority = await claimDeploymentAttemptForTest(database, run.id, { + attemptCount: 1, + claimOwner: "first-attempt", + }); + let unblockFirstBuild!: () => void; + let markFirstBuildStarted!: () => void; + const firstBuildStarted = new Promise((resolve) => { + markFirstBuildStarted = resolve; + }); + const firstBuildGate = new Promise((resolve) => { + unblockFirstBuild = resolve; + }); + let firstDeployCount = 0; + let takeoverDeployCount = 0; + const prepared = { + repoDir: "/repo", + snapshot: { files: { "index.html": "
Owned
" } }, + }; + const firstDispatch = dispatchAppDeploymentRun( + bindings, + { appDeploymentRunId: run.id }, + firstAuthority, + { + runner: { + async build() { + markFirstBuildStarted(); + await firstBuildGate; + }, + async deploy() { + firstDeployCount += 1; + throw new Error("Stale deployment must not run."); + }, + async prepare() { + return prepared; + }, + }, + }, + ); + + await firstBuildStarted; + const takeoverAuthority = await claimDeploymentAttemptForTest(database, run.id, { + attemptCount: 2, + claimOwner: "takeover-attempt", + }); + const takeoverOutcome = await dispatchAppDeploymentRun( + bindings, + { appDeploymentRunId: run.id }, + takeoverAuthority, + { + runner: { + async build() {}, + async deploy() { + takeoverDeployCount += 1; + return { + externalDeploymentId: "takeover-deployment", + externalProjectId: "takeover-project", + externalVersionId: null, + url: "https://takeover.apps.localhost", + }; + }, + async prepare() { + return prepared; + }, + }, + }, + ); + await commitDeploymentOutcomeForTest(database, takeoverAuthority, takeoverOutcome); + unblockFirstBuild(); + + await expect(firstDispatch).rejects.toThrow("lost ownership"); + const runRow = await database.app().select().from(appDeploymentRunsTable).limit(1).get(); + + expect(firstDeployCount).toBe(0); + expect(takeoverDeployCount).toBe(1); + expect(runRow).toMatchObject({ + externalDeploymentId: "takeover-deployment", + status: "success", + }); + }); + test("does not delete stable Cloudflare resources when an inactive run finishes", async () => { const database = createDatabase(); const { bindings } = createBindings(database); @@ -958,8 +1942,8 @@ describe("app deployment service", () => { .bind(run.id) .run(); return { - externalDeploymentId: "pages-deployment-1", - externalProjectId: "pages-project-1", + externalDeploymentId: null, + externalProjectId: null, externalVersionId: null, url: targetUrl, }; @@ -972,14 +1956,10 @@ describe("app deployment service", () => { }, }; - await dispatchAppDeploymentRun( - bindings as ApiBindings, - { appDeploymentRunId: run.id }, - { - cloudflareClient: createCloudflareDeleteRecorder(deleted), - runner, - }, - ); + await dispatchDeploymentForTest(database, bindings, run.id, { + cloudflareClient: createCloudflareDeleteRecorder(deleted), + runner, + }); const runRow = await database .app() @@ -993,7 +1973,7 @@ describe("app deployment service", () => { expect(runRow?.status).toBe("failed"); }); - test("compensates resources created after deployment deletion", async () => { + test("leaves a late WfP candidate retired for ledger GC after deployment deletion", async () => { const database = createDatabase(); const { bindings } = createBindings(database); const deleted: string[] = []; @@ -1008,12 +1988,12 @@ describe("app deployment service", () => { const runner: AppDeploymentBuildRunner = { async build() {}, async deploy() { - await deleteAppDeployment(bindings, VIEWER, { appId: APP_ID }, { cloudflareClient }); - externallyCreated.push(`pages:${APP_ID}`); + await deleteAppDeployment(bindings, VIEWER, { appId: APP_ID }); + externallyCreated.push(`wfp:${APP_ID}`); return { - externalDeploymentId: "pages-deployment-after-delete", - externalProjectId: "pages-project-after-delete", + externalDeploymentId: null, + externalProjectId: null, externalVersionId: null, url: `https://app-${APP_ID.toLowerCase()}.apps.localhost`, }; @@ -1026,11 +2006,12 @@ describe("app deployment service", () => { }, }; - await dispatchAppDeploymentRun( - bindings as ApiBindings, - { appDeploymentRunId: run.id }, - { cloudflareClient, runner }, - ); + await expect( + dispatchDeploymentForTest(database, bindings, run.id, { + cloudflareClient, + runner, + }), + ).rejects.toThrow("lost ownership"); const deployment = await database .app() @@ -1048,15 +2029,13 @@ describe("app deployment service", () => { .where(eq(appDeploymentRunsTable.id, run.id)) .get(); - expect(externallyCreated).toEqual([`pages:${APP_ID}`]); + expect(externallyCreated).toEqual([`wfp:${APP_ID}`]); expect(deployment?.deletedAt).toBeNumber(); expect(runRow).toEqual({ errorCode: "deployment_deleted", status: "failed" }); - expect(deleted).toHaveLength(10); - expect(deleted.filter((entry) => entry.startsWith("pages:"))).toHaveLength(2); - expect(deleted.filter((entry) => entry.startsWith("worker:"))).toHaveLength(2); + expect(deleted).toEqual([]); }); - test("does not compensate a deleted deployment after a replacement is active", async () => { + test("cannot let a late retired candidate disturb a replacement deployment", async () => { const database = createDatabase(); const { bindings } = createBindings(database); const deleted: string[] = []; @@ -1070,7 +2049,7 @@ describe("app deployment service", () => { const runner: AppDeploymentBuildRunner = { async build() {}, async deploy() { - await deleteAppDeployment(bindings, VIEWER, { appId: APP_ID }, { cloudflareClient }); + await deleteAppDeployment(bindings, VIEWER, { appId: APP_ID }); await deployApp( bindings, VIEWER, @@ -1079,8 +2058,8 @@ describe("app deployment service", () => { ); return { - externalDeploymentId: "pages-deployment-after-replacement", - externalProjectId: "pages-project-after-replacement", + externalDeploymentId: null, + externalProjectId: null, externalVersionId: null, url: `https://app-${APP_ID.toLowerCase()}.apps.localhost`, }; @@ -1093,24 +2072,132 @@ describe("app deployment service", () => { }, }; - await dispatchAppDeploymentRun( - bindings as ApiBindings, - { appDeploymentRunId: run.id }, - { cloudflareClient, runner }, - ); + await expect( + dispatchDeploymentForTest(database, bindings, run.id, { + cloudflareClient, + runner, + }), + ).rejects.toThrow("lost ownership"); const activeDeployment = await getAppDeployment(bindings, VIEWER, APP_ID); expect(activeDeployment).not.toBeNull(); - expect(deleted).toHaveLength(5); + expect(deleted).toEqual([]); + }); + + test("an old delete cannot terminate or destroy a replacement deployment", async () => { + const database = createDatabase(); + const { bindings } = createBindings(database); + const destroyed: string[] = []; + bindings.runtimeSubjectHandleFactory = (runtimeSubjectId) => + createTestSandboxHandle(runtimeSubjectId, destroyed, "destroy-only"); + await deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS }, + ); + + let replacement: Awaited> | null = null; + let raced = false; + const racingDatabase = new Proxy(database as D1Database, { + get(target, property, receiver) { + if (property !== "prepare") { + const value = Reflect.get(target, property, receiver) as unknown; + return typeof value === "function" ? value.bind(target) : value; + } + return (query: string) => { + const statement = target.prepare(query); + if (!query.includes("FROM app_deployment_run AS run")) { + return statement; + } + const interceptAll = (prepared: D1PreparedStatement) => + new Proxy(prepared, { + get(statementTarget, statementProperty, statementReceiver) { + if (statementProperty === "bind") { + return (...values: unknown[]) => interceptAll(statementTarget.bind(...values)); + } + if (statementProperty !== "all") { + const value = Reflect.get( + statementTarget, + statementProperty, + statementReceiver, + ) as unknown; + return typeof value === "function" ? value.bind(statementTarget) : value; + } + return async () => { + if (!raced) { + raced = true; + await deleteAppDeployment( + bindings, + VIEWER, + { appId: APP_ID }, + { nowMs: () => NOW_MS + 1 }, + ); + replacement = await deployApp( + bindings, + VIEWER, + { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS + 2 }, + ); + await database + .app() + .update(apiCommandsTable) + .set({ attemptCount: 1 }) + .where( + eq( + apiCommandsTable.dedupeKey, + createAppDeploymentRunDispatchDedupeKey(replacement.id), + ), + ) + .run(); + } + return statementTarget.all(); + }; + }, + }); + return interceptAll(statement); + }; + }, + }); + + await deleteAppDeployment( + { ...bindings, DB: racingDatabase }, + VIEWER, + { appId: APP_ID }, + { nowMs: () => NOW_MS + 3 }, + ); + + if (replacement === null) { + throw new Error("Expected a replacement deployment."); + } + const replacementRun = await database + .app() + .select({ status: appDeploymentRunsTable.status }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, replacement.id)) + .get(); + const replacementCommand = await database + .app() + .select({ status: apiCommandsTable.status }) + .from(apiCommandsTable) + .where( + eq(apiCommandsTable.dedupeKey, createAppDeploymentRunDispatchDedupeKey(replacement.id)), + ) + .get(); + + await expect(getAppDeployment(bindings, VIEWER, APP_ID)).resolves.not.toBeNull(); + expect(replacementRun).toEqual({ status: "queued" }); + expect(replacementCommand).toEqual({ status: "queued" }); + expect(destroyed).not.toContain(`${replacement.id}-g1-a1-build`); + expect(destroyed).not.toContain(`${replacement.id}-g1-a1-deploy`); }); test("deletes the active deployment and fails the active run", async () => { const database = createDatabase(); const { bindings } = createBindings(database); - const deleted: string[] = []; const destroyed: string[] = []; - (bindings as ApiBindings).runtimeSubjectHandleFactory = (runtimeSubjectId) => + bindings.runtimeSubjectHandleFactory = (runtimeSubjectId) => createTestSandboxHandle(runtimeSubjectId, destroyed, "destroy-only"); const run = await deployApp( @@ -1119,14 +2206,34 @@ describe("app deployment service", () => { { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, { fetch: githubFetch, nowMs: () => NOW_MS }, ); - + const activeScriptName = "app-active-g1-a1"; + const liveUrl = `https://app-${APP_ID.toLowerCase()}.apps.localhost`; + await database + .prepare( + `UPDATE app_deployment + SET active_script_name = ?, last_successful_url = ? + WHERE id = ?`, + ) + .bind(activeScriptName, liveUrl, run.deploymentId) + .run(); + await database + .prepare( + `UPDATE api_command + SET attempt_count = 1, + claim_expires_at = ?, + claim_owner = 'active-delete-worker', + status = 'running' + WHERE dedupe_key = ?`, + ) + .bind(NOW_MS + 60_000, createAppDeploymentRunDispatchDedupeKey(run.id)) + .run(); await expect( deleteAppDeployment( bindings, VIEWER, { appId: APP_ID }, { - cloudflareClient: createCloudflareDeleteRecorder(deleted), + nowMs: () => NOW_MS + 1, }, ), ).resolves.toEqual({ ok: true }); @@ -1142,42 +2249,38 @@ describe("app deployment service", () => { id: run.id, status: "failed", }); - expect(deploymentRow?.deletedAt).toBeNumber(); + expect(deploymentRow).toMatchObject({ + activeScriptName: null, + deletedAt: NOW_MS + 1, + lastSuccessfulUrl: null, + }); expect(runRow?.status).toBe("failed"); - expect(deleted).toContain(`pages-domain:app-${APP_ID.toLowerCase()}.apps.localhost`); - expect(deleted).toContain(`pages:app-${APP_ID.toLowerCase()}`); - expect(deleted).toContain(`worker-domain:app-${APP_ID.toLowerCase()}.apps.localhost`); - expect(deleted).toContain(`worker-route:app-${APP_ID.toLowerCase()}.apps.localhost`); - expect(deleted).toContain(`worker:app-${APP_ID.toLowerCase()}`); - expect(destroyed).toContain(`${run.id}-build`); - expect(destroyed).toContain(`${run.id}-deploy`); + expect(destroyed).toContain(`${run.id}-g1-a1-build`); + expect(destroyed).toContain(`${run.id}-g1-a1-deploy`); }); - test("keeps a deployment retryable when Cloudflare cleanup fails", async () => { + test("revokes traffic without deleting the retired WfP script inline", async () => { const database = createDatabase(); const { bindings } = createBindings(database); - const failedDeletes: string[] = []; - const successfulDeletes: string[] = []; const run = await deployApp( bindings, VIEWER, { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, { fetch: githubFetch, nowMs: () => NOW_MS }, ); - const failingClient = createCloudflareDeleteRecorder(failedDeletes, { - async deletePagesProject(input) { - failedDeletes.push(`pages:${input.projectName}`); - throw new Error("Cloudflare API unavailable."); - }, + const activeScriptName = "app-deferred-g1-a1"; + await database + .prepare( + `UPDATE app_deployment + SET active_script_name = ?, last_successful_url = ? + WHERE id = ?`, + ) + .bind(activeScriptName, "https://app.example", run.deploymentId) + .run(); + await expect(deleteAppDeployment(bindings, VIEWER, { appId: APP_ID })).resolves.toEqual({ + ok: true, }); - await expect( - deleteAppDeployment(bindings, VIEWER, { appId: APP_ID }, { cloudflareClient: failingClient }), - ).rejects.toMatchObject({ - code: API_ERROR_CODE.appDeploymentCleanupFailed, - name: "ApiError", - } satisfies Partial); - const deploymentAfterFailure = await database .app() .select() @@ -1191,31 +2294,18 @@ describe("app deployment service", () => { .where(eq(appDeploymentRunsTable.id, run.id)) .get(); - expect(deploymentAfterFailure).toMatchObject({ deletedAt: null }); + expect(deploymentAfterFailure).toMatchObject({ + activeScriptName: null, + lastSuccessfulUrl: null, + }); + expect(deploymentAfterFailure?.deletedAt).toBeNumber(); expect(runAfterFailure).toMatchObject({ errorCode: "deployment_deleted", status: "failed", }); - expect(failedDeletes).toHaveLength(5); - - await expect( - deleteAppDeployment( - bindings, - VIEWER, - { appId: APP_ID }, - { cloudflareClient: createCloudflareDeleteRecorder(successfulDeletes) }, - ), - ).resolves.toEqual({ ok: true }); - - const deploymentAfterRetry = await database - .app() - .select({ deletedAt: appDeploymentsTable.deletedAt }) - .from(appDeploymentsTable) - .where(eq(appDeploymentsTable.id, run.deploymentId)) - .get(); - - expect(deploymentAfterRetry?.deletedAt).toBeNumber(); - expect(successfulDeletes).toHaveLength(5); + await expect(deleteAppDeployment(bindings, VIEWER, { appId: APP_ID })).resolves.toEqual({ + ok: true, + }); }); test("does not expose a live URL after deleting a successful deployment", async () => { @@ -1226,8 +2316,8 @@ describe("app deployment service", () => { async build() {}, async deploy() { return { - externalDeploymentId: "pages-deployment-1", - externalProjectId: "pages-project-1", + externalDeploymentId: null, + externalProjectId: null, externalVersionId: null, url: targetUrl, }; @@ -1246,21 +2336,8 @@ describe("app deployment service", () => { { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, { fetch: githubFetch, nowMs: () => NOW_MS }, ); - await dispatchAppDeploymentRun( - bindings as ApiBindings, - { appDeploymentRunId: run.id }, - { - runner, - }, - ); - await deleteAppDeployment( - bindings, - VIEWER, - { appId: APP_ID }, - { - cloudflareClient: createCloudflareDeleteRecorder([]), - }, - ); + await dispatchDeploymentForTest(database, bindings, run.id, { runner }); + await deleteAppDeployment(bindings, VIEWER, { appId: APP_ID }); await expect(getAppDeploymentStatus(bindings, VIEWER, APP_ID)).resolves.toMatchObject({ liveUrl: null, @@ -1276,8 +2353,8 @@ describe("app deployment service", () => { async build() {}, async deploy() { return { - externalDeploymentId: "pages-deployment-1", - externalProjectId: "pages-project-1", + externalDeploymentId: null, + externalProjectId: null, externalVersionId: null, url: targetUrl, }; @@ -1296,11 +2373,7 @@ describe("app deployment service", () => { { fetch: githubFetch, nowMs: () => NOW_MS }, ); - await dispatchAppDeploymentRun( - bindings as ApiBindings, - { appDeploymentRunId: run.id }, - { runner }, - ); + await dispatchDeploymentForTest(database, bindings, run.id, { runner }); await database .prepare("UPDATE app_deployment_run SET plan_json = ? WHERE id = ?") .bind( @@ -1322,14 +2395,7 @@ describe("app deployment service", () => { authorized: true, }); - await deleteAppDeployment( - bindings, - VIEWER, - { appId: APP_ID }, - { - cloudflareClient: createCloudflareDeleteRecorder([]), - }, - ); + await deleteAppDeployment(bindings, VIEWER, { appId: APP_ID }); await expect(getDeploymentAgentCapabilityAuthority(database, authority)).resolves.toEqual({ authorized: false, @@ -1337,20 +2403,29 @@ describe("app deployment service", () => { }); }); - test("uses the Pages deployment URL while the custom domain is pending", async () => { + test("deploys native Static Assets to an attempt-specific WfP script", async () => { const database = createDatabase(); const { bindings } = createBindings(database); const calls: string[] = []; - const pagesUrl = "https://app-example.pages.dev"; + const cloudflareCalls: string[] = []; const cloudflareClient = createCloudflareDeleteRecorder([], { - ensurePagesDomain: async () => ({ status: "initializing" }), - ensurePagesProject: async () => ({ projectId: "pages-project-1" }), - getLatestPagesDeployment: async () => ({ - deploymentId: "pages-deployment-1", - url: pagesUrl, - }), + async createStaticAssetsUploadSession(input) { + cloudflareCalls.push( + `session:${input.scriptName}:${Object.keys(input.manifest).join(",")}`, + ); + return { + buckets: [["0123456789abcdef0123456789abcdef"]], + uploadToken: "eyJhbGciOiJub25lIn0.e30.signature", + }; + }, + async deployStaticAssets(input) { + cloudflareCalls.push( + `deploy:${input.scriptName}:${input.completionToken}:${input.tags.join(",")}:${input.headers}:${input.redirects}`, + ); + return { deploymentId: null, versionId: null }; + }, }); - (bindings as ApiBindings).runtimeSubjectHandleFactory = (runtimeSubjectId) => + bindings.runtimeSubjectHandleFactory = (runtimeSubjectId) => createTestSandboxHandle(runtimeSubjectId, calls, "deployment"); const run = await deployApp( @@ -1360,49 +2435,83 @@ describe("app deployment service", () => { { fetch: githubFetch, nowMs: () => NOW_MS }, ); - await dispatchAppDeploymentRun( - bindings as ApiBindings, - { appDeploymentRunId: run.id }, - { cloudflareClient }, - ); + await dispatchDeploymentForTest(database, bindings, run.id, { + cloudflareClient, + }); const status = await getAppDeploymentStatus(bindings, VIEWER, APP_ID); const runRow = await database.app().select().from(appDeploymentRunsTable).limit(1).get(); + const scriptRow = await database.app().select().from(appDeploymentScriptsTable).limit(1).get(); + const targetUrl = `https://app-${APP_ID.toLowerCase()}.apps.localhost`; + const scriptName = runRow?.targetScriptName; expect(status).toMatchObject({ - liveUrl: pagesUrl, + liveUrl: targetUrl, status: "success", }); expect(runRow).toMatchObject({ - externalDeploymentId: "pages-deployment-1", - externalProjectId: "pages-project-1", + externalDeploymentId: null, + externalProjectId: null, status: "success", - url: pagesUrl, + targetKind: "cloudflare_static_assets", + targetProjectName: null, + targetScriptName: expect.stringMatching(/^app-[0-9a-z]{26}-g1-a1$/u), + url: targetUrl, }); - expect(calls.some((call) => call.includes("wrangler pages deploy"))).toBe(true); + expect(calls.every((call) => !call.includes("wrangler"))).toBe(true); + const fallbackArchive = calls.find((call) => call.includes(">>")); + const commandCalls = calls.filter((call) => call.includes("sh -lc")); + + expect(fallbackArchive).toContain(">>"); + expect(fallbackArchive).toContain("find . -name .git -prune -o -type f -print0"); + expect(fallbackArchive).toContain("output_root=$(pwd -P)"); + expect(fallbackArchive).not.toContain("/_headers"); + expect(commandCalls.length).toBeGreaterThan(0); + expect( + commandCalls.every((call) => { + const timeout = Number(/timeout=(\d+)/u.exec(call)?.[1]); + return timeout > 0 && timeout <= 600_000; + }), + ).toBe(true); + expect( + calls + .filter((call) => call.includes(":mount:FILE_BUCKET:")) + .map((call) => call.slice(call.lastIndexOf(":"))), + ).toEqual([":false", ":true"]); + expect( + calls + .filter((call) => call.includes(":mount:FILE_BUCKET:")) + .every((call) => call.includes(`/app-deployments/${scriptName}/`)), + ).toBe(true); + expect(calls.some((call) => call.endsWith('/selection.json:["index.html"]'))).toBe(true); + expect( + calls.some((call) => + call.includes("session-env=CLOUDFLARE_ACCOUNT_ID,CLOUDFLARE_ASSET_UPLOAD_JWT"), + ), + ).toBe(true); + expect(calls.every((call) => !call.includes("CLOUDFLARE_API_TOKEN="))).toBe(true); + expect(scriptRow?.uploadStartedAt).toBeNumber(); + expect(cloudflareCalls).toEqual([ + `session:${scriptName}:/index.html`, + `deploy:${scriptName}:completion-jwt:mosoo-managed,d-${run.deploymentId.toLowerCase()},r-${run.id.toLowerCase()}:/assets/*\n Cache-Control: public\n:/old /new 301\n\n/* /index.html 200\n`, + ]); }); - test("deploys worker modules with a Worker route and custom domain", async () => { + test("deploys a Worker module as an attempt-specific WfP script", async () => { const database = createDatabase(); const { bindings } = createBindings(database); const calls: string[] = []; const cloudflareCalls: string[] = []; - (bindings as ApiBindings).runtimeSubjectHandleFactory = (runtimeSubjectId) => + bindings.runtimeSubjectHandleFactory = (runtimeSubjectId) => createWorkerDeploymentSandboxHandle(runtimeSubjectId, calls); const cloudflareClient = createCloudflareDeleteRecorder([], { async deployWorkerModule(input) { cloudflareCalls.push( - `worker:${input.scriptName}:${input.mainModuleName}:${input.scriptContent.trim()}`, + `worker:${input.scriptName}:${input.mainModuleName}:${input.tags.join(",")}:${input.scriptContent.trim()}`, ); return { deploymentId: "worker-deployment-1", versionId: "worker-version-1" }; }, - async ensureWorkerDomain(input) { - cloudflareCalls.push(`worker-domain:${input.hostname}:${input.scriptName}`); - }, - async ensureWorkerRoute(input) { - cloudflareCalls.push(`worker-route:${input.hostname}:${input.scriptName}`); - }, }); const run = await deployApp( @@ -1412,17 +2521,14 @@ describe("app deployment service", () => { { fetch: githubFetch, nowMs: () => NOW_MS }, ); - await dispatchAppDeploymentRun( - bindings as ApiBindings, - { appDeploymentRunId: run.id }, - { cloudflareClient }, - ); + await dispatchDeploymentForTest(database, bindings, run.id, { + cloudflareClient, + }); - const targetName = `app-${APP_ID.toLowerCase()}`; - const hostname = `${targetName}.apps.localhost`; - const targetUrl = `https://${hostname}`; + const targetUrl = `https://app-${APP_ID.toLowerCase()}.apps.localhost`; const status = await getAppDeploymentStatus(bindings, VIEWER, APP_ID); const runRow = await database.app().select().from(appDeploymentRunsTable).limit(1).get(); + const scriptName = runRow?.targetScriptName; expect(status).toMatchObject({ liveUrl: targetUrl, @@ -1434,65 +2540,147 @@ describe("app deployment service", () => { externalVersionId: "worker-version-1", status: "success", targetKind: "cloudflare_worker", - targetScriptName: targetName, + targetProjectName: null, + targetScriptName: expect.stringMatching(/^app-[0-9a-z]{26}-g1-a1$/u), url: targetUrl, }); expect(cloudflareCalls).toEqual([ - `worker:${targetName}:index.js:export default { fetch() { return new Response('ok'); } };`, - `worker-route:${hostname}:${targetName}`, - `worker-domain:${hostname}:${targetName}`, + `worker:${scriptName}:worker.mjs:mosoo-managed,d-${run.deploymentId.toLowerCase()},r-${run.id.toLowerCase()}:export default { fetch() { return new Response('ok'); } };`, ]); + expect( + calls.some( + (call) => + call.includes("/repo/.") && call.includes("bun build") && call.includes("src/index.js"), + ), + ).toBe(true); }); - test("dead letters active deployment runs without overwriting terminal runs", async () => { - const database = createDatabase(); - const { bindings, queue } = createBindings(database); + test("bundles relative and package Worker imports from the configured root", async () => { + const projectRoot = await mkdtemp(join(tmpdir(), "mosoo-worker-bundle-")); - const run = await deployApp( - bindings, - VIEWER, - { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, - { fetch: githubFetch, nowMs: () => NOW_MS }, - ); - const queued = queue.sent[0]; + try { + await mkdir(join(projectRoot, "node_modules", "fixture-package"), { recursive: true }); + await mkdir(join(projectRoot, "src"), { recursive: true }); + await writeFile( + join(projectRoot, "node_modules", "fixture-package", "package.json"), + JSON.stringify({ exports: "./index.js", name: "fixture-package", type: "module" }), + ); + await writeFile( + join(projectRoot, "node_modules", "fixture-package", "index.js"), + 'export const packageValue = "package-value";\n', + ); + await writeFile( + join(projectRoot, "src", "relative.js"), + 'export const relativeValue = "relative-value";\n', + ); + await writeFile( + join(projectRoot, "src", "index.js"), + [ + 'import { relativeValue } from "./relative.js";', + 'import { packageValue } from "fixture-package";', + "export default {", + " fetch() { return new Response(`${relativeValue}:${packageValue}`); },", + "};", + "", + ].join("\n"), + ); + + const database = createDatabase(); + const { bindings } = createBindings(database); + const calls: string[] = []; + let uploaded = ""; + bindings.runtimeSubjectHandleFactory = (runtimeSubjectId) => + createWorkerDeploymentSandboxHandle(runtimeSubjectId, calls, { + bundleProjectRoot: projectRoot, + rootDir: "packages/worker", + }); + const cloudflareClient = createCloudflareDeleteRecorder([], { + async deployWorkerModule(input) { + uploaded = input.scriptContent; + return { deploymentId: "worker-deployment-1", versionId: "worker-version-1" }; + }, + }); + const run = await deployApp( + bindings, + VIEWER, + { appId: APP_ID, configPath: ".mosoo.toml", repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS }, + ); - if (queued === undefined) { - throw new Error("Expected deployment dispatch queue message."); + await expect( + dispatchDeploymentForTest(database, bindings, run.id, { + cloudflareClient, + }), + ).resolves.toMatchObject({ kind: "succeeded" }); + + expect(uploaded).toContain("relative-value"); + expect(uploaded).toContain("package-value"); + expect(uploaded).not.toContain('from "./relative.js"'); + expect(uploaded).not.toContain('from "fixture-package"'); + const bundleCall = calls.find((call) => call.includes("bun build")); + + expect(bundleCall).toContain("/repo/packages/worker"); + expect(bundleCall?.indexOf("/repo/packages/worker")).toBeLessThan( + bundleCall?.indexOf("bun build") ?? -1, + ); + } finally { + await rm(projectRoot, { force: true, recursive: true }); } + }); - await database.app().update(apiCommandsTable).set({ payloadJson: "{}" }).run(); + test("fails closed when a Worker dependency cannot be bundled", async () => { + const projectRoot = await mkdtemp(join(tmpdir(), "mosoo-worker-missing-dependency-")); - await processApiCommandDeadLetterMessage( - bindings as ApiBindings, - createRecordedQueueMessage({ body: queued.body }).message, - () => NOW_MS + 1, - ); + try { + await mkdir(join(projectRoot, "src"), { recursive: true }); + await writeFile( + join(projectRoot, "src", "index.js"), + 'import "missing-worker-dependency";\nexport default { fetch() {} };\n', + ); - await expect(getAppDeploymentStatus(bindings, VIEWER, APP_ID)).resolves.toMatchObject({ - errorCode: "queue_dead_lettered", - id: run.id, - status: "failed", - }); + const database = createDatabase(); + const { bindings } = createBindings(database); + const calls: string[] = []; + let deployCalls = 0; + bindings.runtimeSubjectHandleFactory = (runtimeSubjectId) => + createWorkerDeploymentSandboxHandle(runtimeSubjectId, calls, { + bundleProjectRoot: projectRoot, + rootDir: "packages/worker", + }); + const cloudflareClient = createCloudflareDeleteRecorder([], { + async deployWorkerModule() { + deployCalls += 1; + return { deploymentId: "unexpected", versionId: "unexpected" }; + }, + }); + const run = await deployApp( + bindings, + VIEWER, + { appId: APP_ID, configPath: ".mosoo.toml", repoUrl: "https://github.com/samzong/awire" }, + { fetch: githubFetch, nowMs: () => NOW_MS }, + ); - await deleteAppDeployment( - bindings, - VIEWER, - { appId: APP_ID }, - { - cloudflareClient: createCloudflareDeleteRecorder([]), - }, - ); - await processApiCommandDeadLetterMessage( - bindings as ApiBindings, - createRecordedQueueMessage({ body: queued.body }).message, - () => NOW_MS + 2, - ); + const outcome = await dispatchDeploymentForTest(database, bindings, run.id, { + cloudflareClient, + }); - await expect(getAppDeploymentStatus(bindings, VIEWER, APP_ID)).resolves.toMatchObject({ - errorCode: "queue_dead_lettered", - id: run.id, - status: "failed", - }); + expect(outcome).toMatchObject({ + errorCode: "AppDeploymentNonRetryableError", + errorMessage: expect.stringContaining('Could not resolve: "missing-worker-dependency"'), + kind: "terminal_failure", + }); + expect(deployCalls).toBe(0); + await expect( + database + .app() + .select({ status: appDeploymentRunsTable.status }) + .from(appDeploymentRunsTable) + .where(eq(appDeploymentRunsTable.id, run.id)) + .get(), + ).resolves.toEqual({ status: "failed" }); + } finally { + await rm(projectRoot, { force: true, recursive: true }); + } }); test("lists deployment runs newest-first", async () => { @@ -1613,8 +2801,8 @@ describe("app deployment service", () => { async build() {}, async deploy() { return { - externalDeploymentId: "pages-deployment-1", - externalProjectId: "pages-project-1", + externalDeploymentId: null, + externalProjectId: null, externalVersionId: null, url: targetUrl, }; @@ -1633,21 +2821,8 @@ describe("app deployment service", () => { { appId: APP_ID, repoUrl: "https://github.com/samzong/awire" }, { fetch: githubFetch, nowMs: () => NOW_MS }, ); - await dispatchAppDeploymentRun( - bindings as ApiBindings, - { appDeploymentRunId: run.id }, - { - runner, - }, - ); - await deleteAppDeployment( - bindings, - VIEWER, - { appId: APP_ID }, - { - cloudflareClient: createCloudflareDeleteRecorder([]), - }, - ); + await dispatchDeploymentForTest(database, bindings, run.id, { runner }); + await deleteAppDeployment(bindings, VIEWER, { appId: APP_ID }); // The deployment is soft-deleted: run history stays listed, but liveUrl is // suppressed because the deployment row carries deletedAt. diff --git a/apps/api/tests/app-overview.test.ts b/apps/api/tests/app-overview.test.ts index b1e83887..f9fde426 100644 --- a/apps/api/tests/app-overview.test.ts +++ b/apps/api/tests/app-overview.test.ts @@ -1,9 +1,16 @@ import { describe, expect, test } from "bun:test"; -import { isInputObjectType, isObjectType } from "graphql"; +import { isEnumType, isInputObjectType, isObjectType } from "graphql"; import { createGraphQLSchema } from "../src/adapters/graphql/create-graphql-schema"; import { createAppDeploymentRunDispatchDedupeKey } from "../src/modules/api-command/application/api-command-enqueue"; +import { admitApiCommand } from "../src/modules/api-command/application/api-command-ledger"; +import { processApiCommandDelivery } from "../src/modules/api-command/application/api-command-processor"; +import { appDeploymentCandidateScriptName } from "../src/modules/apps/application/app-deployment-executor.service"; +import { + markAppDeploymentScriptUploadStarted, + registerAppDeploymentScriptCandidate, +} from "../src/modules/apps/application/app-deployment-script-reconciliation.service"; import { getAppOverview, getControlPlaneOverview, @@ -13,6 +20,7 @@ import { createApiTestFixture } from "./helpers/api-test-fixture"; const OVERVIEW_DEPLOYMENT_ID = "01J000000000000000000000D1"; const OVERVIEW_DEPLOYMENT_RUN_ID = "01J000000000000000000000D2"; +const OVERVIEW_AGENT_DEPLOYMENT_VERSION_ID = "01J000000000000000000000D3"; function createOverviewDeploymentUrl(appId: string, domain: string): string { return `https://app-${appId.toLowerCase()}.${domain}`; @@ -36,6 +44,12 @@ async function insertOverviewAgent( updatedAt: number; }, ): Promise { + const configJson = JSON.stringify({ + packageMcpServers: [], + packageResolution: null, + packageSkills: [], + providerOptions: {}, + }); await fixture.database .prepare( `INSERT INTO agent ( @@ -57,12 +71,7 @@ async function insertOverviewAgent( ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, ) .bind( - JSON.stringify({ - packageMcpServers: [], - packageResolution: null, - packageSkills: [], - providerOptions: {}, - }), + configJson, 1, "Extra overview fixture.", input.id, @@ -74,11 +83,46 @@ async function insertOverviewAgent( "Help with overview tests.", "openai", "openai-runtime", - "published", + "draft", input.updatedAt, "private", ) .run(); + + await fixture.database.batch([ + fixture.database + .prepare( + `INSERT INTO agent_deployment_version ( + agent_id, config_json, created_at, created_by_account_id, environment_id, + id, kind, mcp_bindings_json, model, prompt, provider, runtime_id, + skills_json, summary, version_number + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + input.id, + configJson, + input.updatedAt, + fixture.viewer.id, + null, + OVERVIEW_AGENT_DEPLOYMENT_VERSION_ID, + "cattle", + "[]", + "gpt-5.4", + "Help with overview tests.", + "openai", + "openai-runtime", + "[]", + "Initial publish", + 1, + ), + fixture.database + .prepare( + `UPDATE agent + SET live_deployment_version_id = ?, status = 'published', updated_at = ? + WHERE id = ? AND status = 'draft'`, + ) + .bind(OVERVIEW_AGENT_DEPLOYMENT_VERSION_ID, input.updatedAt, input.id), + ]); } async function insertOverviewCredentialMetadata( @@ -131,6 +175,7 @@ async function insertOverviewDeploymentMetadata( deleted_at, id, last_successful_url, + latest_run_id, mosoo_subdomain, owner_account_id, repo_name, @@ -138,7 +183,7 @@ async function insertOverviewDeploymentMetadata( repo_url, source_kind, updated_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, ) .bind( fixture.ids.appId, @@ -146,7 +191,8 @@ async function insertOverviewDeploymentMetadata( "main", null, OVERVIEW_DEPLOYMENT_ID, - liveUrl, + null, + OVERVIEW_DEPLOYMENT_RUN_ID, `app-${fixture.ids.appId.toLowerCase()}`, fixture.viewer.id, "awire", @@ -183,13 +229,56 @@ async function insertOverviewDeploymentMetadata( OVERVIEW_DEPLOYMENT_RUN_ID, "main", "abc123", - "success", - "cloudflare_pages", + "preparing", + "cloudflare_static_assets", 2, - liveUrl, + null, ) .run(); + const nowMs = Date.now(); + const admission = await admitApiCommand(fixture.bindings, { + dedupeKey: createAppDeploymentRunDispatchDedupeKey(OVERVIEW_DEPLOYMENT_RUN_ID), + kind: "app_deployment_run_dispatch", + payload: { appDeploymentRunId: OVERVIEW_DEPLOYMENT_RUN_ID }, + }); + const disposition = await processApiCommandDelivery( + fixture.bindings, + { commandId: admission.commandId, deliveryGeneration: 1 }, + () => nowMs, + { + dispatchAppDeploymentRun: async (_bindings, _input, authority) => { + const activeScriptName = appDeploymentCandidateScriptName(authority); + const candidate = { + authority, + deploymentId: OVERVIEW_DEPLOYMENT_ID, + nowMs, + runId: OVERVIEW_DEPLOYMENT_RUN_ID, + scriptName: activeScriptName, + } as const; + await registerAppDeploymentScriptCandidate(fixture.database, candidate); + await markAppDeploymentScriptUploadStarted(fixture.database, candidate); + await fixture.database + .prepare("UPDATE app_deployment_run SET status = 'activating' WHERE id = ?") + .bind(OVERVIEW_DEPLOYMENT_RUN_ID) + .run(); + return { + activeScriptName, + deploymentId: OVERVIEW_DEPLOYMENT_ID, + externalDeploymentId: null, + externalProjectId: null, + externalVersionId: null, + kind: "succeeded" as const, + runId: OVERVIEW_DEPLOYMENT_RUN_ID, + url: liveUrl, + }; + }, + }, + ); + if (disposition.kind !== "finished") { + throw new Error("Overview deployment fixture did not reach its durable success state."); + } + return { liveUrl }; } @@ -202,6 +291,7 @@ describe("App overview", () => { const credential = schema.getType("AppOverviewProviderCredential"); const deployment = schema.getType("AppDeployment"); const deploymentRun = schema.getType("AppDeploymentRun"); + const deploymentTargetKind = schema.getType("AppDeploymentTargetKind"); const deployInput = schema.getType("DeployAppInput"); if ( @@ -211,6 +301,7 @@ describe("App overview", () => { !isObjectType(credential) || !isObjectType(deployment) || !isObjectType(deploymentRun) || + !isEnumType(deploymentTargetKind) || !isInputObjectType(deployInput) ) { throw new Error("Expected App overview GraphQL types."); @@ -229,6 +320,10 @@ describe("App overview", () => { expect(String(appOverview.getFields().deployment.type)).toBe("AppDeployment"); expect(String(deployment.getFields().latestRun.type)).toBe("AppDeploymentRun"); expect(String(deploymentRun.getFields().status.type)).toBe("AppDeploymentRunStatus!"); + expect(deploymentTargetKind.getValues().map(({ name }) => name)).toEqual([ + "cloudflare_static_assets", + "cloudflare_worker", + ]); expect(String(deploymentStatus.type)).toBe("AppDeploymentRun"); expect(String(deploy.type)).toBe("AppDeploymentRun!"); expect(String(deleteDeployment.type)).toBe("OperationResult!"); @@ -266,7 +361,7 @@ describe("App overview", () => { latestRun: { liveUrl: deploymentFixture.liveUrl, status: "success", - targetKind: "cloudflare_pages", + targetKind: "cloudflare_static_assets", }, liveUrl: deploymentFixture.liveUrl, plannedUrl: deploymentFixture.liveUrl, diff --git a/apps/api/tests/bound-agent-idempotency.e2e.test.ts b/apps/api/tests/bound-agent-idempotency.e2e.test.ts index 10ca1712..a1c5964a 100644 --- a/apps/api/tests/bound-agent-idempotency.e2e.test.ts +++ b/apps/api/tests/bound-agent-idempotency.e2e.test.ts @@ -1,7 +1,12 @@ import { describe, expect, test } from "bun:test"; import { createPlatformId, parsePlatformId } from "@mosoo/id"; -import type { AgentDeploymentVersionId, SessionId, SessionMessageId } from "@mosoo/id"; +import type { + AgentDeploymentVersionId, + SessionId, + SessionMessageId, + SessionRunId, +} from "@mosoo/id"; import { Hono } from "hono"; import { registerPublicApiRoute } from "../src/adapters/http/routes/public-api-route"; @@ -16,20 +21,28 @@ import { import { mintAppAgentCapabilityToken } from "../src/modules/public-api/app-agent-capability"; import type { AppAgentCapabilityClaims } from "../src/modules/public-api/app-agent-capability"; import { queueSessionRun } from "../src/modules/runtime/application/session-run.service"; +import { recordCanonicalSessionRunTerminal } from "../src/modules/runtime/application/session-runs/session-run-terminal-failure.service"; +import { prepareAssistantMessageProjection } from "../src/modules/runtime/infrastructure/driver-instance/assistant-message-projection"; +import { + appendSessionRuntimeEvents, + createSessionRuntimeEvent, +} from "../src/modules/sessions/application/session-event-write.service"; import type { ApiBindings, ApiGatewayEnvironment } from "../src/platform/cloudflare/worker-types"; +import { + BOUND_DEPLOYMENT_ID, + BOUND_DEPLOYMENT_RUN_ID, + deleteBoundDeployment, + insertBoundDeployment, +} from "./bound-capability-fixtures"; import { PUBLIC_API_TEST_IDS, createApiCommandQueueStub, createPublicHttpContractDatabase, createPublicHttpTestBindings, createTestExecutionContext, - nowMsForTest, } from "./helpers/public-api-http-test-fixture"; import type { ApiCommandQueueStub, SqliteD1Database } from "./helpers/public-api-http-test-fixture"; -const DEPLOYMENT_ID = "01J0000000000000000000000D"; -const DEPLOYMENT_RUN_ID = "01J0000000000000000000000R"; - interface DispatchCommandPayload { session: { id: string }; sessionRunId: string; @@ -62,57 +75,17 @@ function capabilityClaims( expose: "public_thread", name: "Public API Agent", }, - deploymentId: DEPLOYMENT_ID, - deploymentRunId: DEPLOYMENT_RUN_ID, + deploymentId: BOUND_DEPLOYMENT_ID, + deploymentRunId: BOUND_DEPLOYMENT_RUN_ID, exp: Date.now() + 60_000, ...overrides, }; } -async function insertDeploymentAuthority( - database: SqliteD1Database, - bindings: AppAgentCapabilityClaims["binding"][], -): Promise { - database.execute(` - CREATE TABLE app_deployment ( - app_id text NOT NULL, - deleted_at integer, - id text PRIMARY KEY NOT NULL - ); - - CREATE TABLE app_deployment_run ( - app_id text NOT NULL, - deployment_id text NOT NULL, - id text PRIMARY KEY NOT NULL, - plan_json text, - status text NOT NULL - ); - - CREATE INDEX app_deployment_run_deployment_id_idx - ON app_deployment_run (deployment_id, id); - `); - - await database - .prepare("INSERT INTO app_deployment (app_id, deleted_at, id) VALUES (?, NULL, ?)") - .bind(PUBLIC_API_TEST_IDS.app, DEPLOYMENT_ID) - .run(); - await database - .prepare( - "INSERT INTO app_deployment_run (app_id, deployment_id, id, plan_json, status) VALUES (?, ?, ?, ?, 'success')", - ) - .bind( - PUBLIC_API_TEST_IDS.app, - DEPLOYMENT_ID, - DEPLOYMENT_RUN_ID, - JSON.stringify({ agentBindings: bindings }), - ) - .run(); -} - function createCompletingApiCommandQueue(database: SqliteD1Database): ApiCommandQueueStub { const sent: ApiCommandQueueStub["sent"] = []; - return { + const queue: ApiCommandQueueStub = { sent, async send(body: ApiCommandMessage, options): Promise { sent.push({ @@ -132,69 +105,65 @@ function createCompletingApiCommandQueue(database: SqliteD1Database): ApiCommand } const payload = JSON.parse(command.payloadJson) as DispatchCommandPayload; - const timestampMs = nowMsForTest() + sent.length; - - await database - .prepare( - `UPDATE session_run - SET status = 'completed', - completed_at = ?, - status_changed_at = ?, - status_event = 'run.complete', - status_seq = status_seq + 1, - status_source = 'driver', - updated_at = ? - WHERE id = ?`, - ) - .bind(timestampMs, timestampMs, timestampMs, payload.sessionRunId) - .run(); - await database - .prepare( - `UPDATE session - SET status = 'IDLE', - status_seq = status_seq + 1, - message_seq_cursor = message_seq_cursor + 1, - last_message_at = ?, - updated_at = ? - WHERE id = ?`, - ) - .bind(timestampMs, timestampMs, payload.session.id) - .run(); - - const session = await database - .prepare("SELECT message_seq_cursor AS seq FROM session WHERE id = ?") - .bind(payload.session.id) - .first<{ seq: number }>(); - - if (session === null) { - throw new Error("Queued Session is missing."); - } - - await database - .prepare( - `INSERT INTO session_message ( - content_text, - created_at, - created_by_account_id, - id, - plan_json, - role, - segments_json, - seq, - session_id, - session_run_id - ) VALUES (?, ?, ?, ?, NULL, 'assistant', NULL, ?, ?, ?)`, - ) - .bind( - "The original bound request completed.", - timestampMs, - PUBLIC_API_TEST_IDS.ownerAccount, - createPlatformId(), - session.seq, - payload.session.id, - payload.sessionRunId, - ) - .run(); + const sessionId = parsePlatformId(payload.session.id, "queued Session ID"); + const sessionRunId = parsePlatformId( + payload.sessionRunId, + "queued Session Run ID", + ); + const messageId = createPlatformId(); + const timestampMs = Date.now(); + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: queue, + }) as ApiBindings; + + await appendSessionRuntimeEvents({ + bindings, + deliver: false, + events: [ + createSessionRuntimeEvent({ + actor: "driver", + kind: "message.added", + occurredAtMs: timestampMs, + origin: "driver", + payload: { + content: "The original bound request completed.", + messageId, + role: "agent", + }, + runId: sessionRunId, + sessionId, + sourceEventId: createPlatformId(), + visibility: "participant", + }), + createSessionRuntimeEvent({ + actor: "driver", + kind: "message.completed", + occurredAtMs: timestampMs, + origin: "driver", + payload: { messageId, role: "agent" }, + runId: sessionRunId, + sessionId, + sourceEventId: createPlatformId(), + visibility: "participant", + }), + ], + sessionId, + }); + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: prepareAssistantMessageProjection({ + createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + messageId, + sessionId, + sessionRunId, + }), + deliver: false, + error: null, + runId: sessionRunId, + sessionId, + source: "api", + status: "completed", + timestampMs, + }); await database .prepare( `UPDATE api_command @@ -209,6 +178,8 @@ function createCompletingApiCommandQueue(database: SqliteD1Database): ApiCommand .run(); }, }; + + return queue; } function failFirstMatchingStatement(database: D1Database, pattern: RegExp): D1Database { @@ -324,7 +295,7 @@ async function withAcceleratedClock(operation: () => Promise): Promise async function createFixture(bindings = [capabilityClaims().binding]) { const database = await createPublicHttpContractDatabase(); - await insertDeploymentAuthority(database, bindings); + await insertBoundDeployment(database, { agentBindings: bindings }); const queue = createCompletingApiCommandQueue(database); return { database, queue }; @@ -496,10 +467,10 @@ describe("bound Agent HTTP idempotency", () => { const reservation = await database .prepare( - `SELECT run_id AS runId, session_id AS sessionId + `SELECT id, run_id AS runId, session_id AS sessionId FROM bound_agent_call_idempotency_key`, ) - .first<{ runId: string; sessionId: string }>(); + .first<{ id: string; runId: string; sessionId: string }>(); const viewer = await getAccountViewer(database, PUBLIC_API_TEST_IDS.ownerAccount); if (reservation === null || viewer === null) { @@ -510,7 +481,10 @@ describe("bound Agent HTTP idempotency", () => { // reserved Session itself must still prevent a later Run from becoming // the replay target for the original key. await database.prepare("UPDATE bound_agent_call_idempotency_key SET run_id = NULL").run(); - await database.prepare("UPDATE session_event SET source_event_id = id").run(); + await database + .prepare("UPDATE session_event SET source_event_id = id WHERE source_event_id = ?") + .bind(reservation.id) + .run(); const later = await queueSessionRun({ bindings: createPublicHttpTestBindings(database, { @@ -797,10 +771,7 @@ describe("bound Agent HTTP idempotency", () => { }); expect(first.status).toBe(200); - await database - .prepare("UPDATE app_deployment SET deleted_at = ? WHERE id = ?") - .bind(Date.now(), DEPLOYMENT_ID) - .run(); + await deleteBoundDeployment(database); const revoked = await requestBoundAgent({ claims: capabilityClaims(), diff --git a/apps/api/tests/bound-capability-fixtures.ts b/apps/api/tests/bound-capability-fixtures.ts index b4bf29e4..dba5d7e0 100644 --- a/apps/api/tests/bound-capability-fixtures.ts +++ b/apps/api/tests/bound-capability-fixtures.ts @@ -1,3 +1,7 @@ +import { appDeploymentRunsTable, appDeploymentsTable } from "@mosoo/db"; +import { parsePlatformId } from "@mosoo/id"; +import type { AppDeploymentId, AppDeploymentRunId } from "@mosoo/id"; +import { eq } from "drizzle-orm"; import type { Hono } from "hono"; import { mintAppAgentCapabilityToken } from "../src/modules/public-api/app-agent-capability"; @@ -6,15 +10,24 @@ import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { PUBLIC_API_TEST_IDS, createPublicHttpTestBindings, + nowMsForTest, } from "./helpers/public-api-http-test-fixture"; import type { SqliteD1Database } from "./helpers/public-api-http-test-fixture"; import { requestPublicApiWithBindings } from "./public-thread-api-fixtures"; -export const BOUND_DEPLOYMENT_ID = "01J0000000000000000000000D"; -export const BOUND_DEPLOYMENT_RUN_ID = "01J0000000000000000000000R"; -export const BOUND_REPLACEMENT_DEPLOYMENT_RUN_ID = "01J0000000000000000000000S"; -export const BOUND_OTHER_DEPLOYMENT_ID = "01J0000000000000000000000E"; -export const BOUND_OTHER_DEPLOYMENT_RUN_ID = "01J0000000000000000000000T"; +export const BOUND_DEPLOYMENT_ID = parsePlatformId("01J0000000000000000000000D"); +export const BOUND_DEPLOYMENT_RUN_ID = parsePlatformId( + "01J0000000000000000000000R", +); +export const BOUND_REPLACEMENT_DEPLOYMENT_RUN_ID = parsePlatformId( + "01J0000000000000000000000S", +); +export const BOUND_OTHER_DEPLOYMENT_ID = parsePlatformId( + "01J0000000000000000000000E", +); +export const BOUND_OTHER_DEPLOYMENT_RUN_ID = parsePlatformId( + "01J0000000000000000000000T", +); export const BOUND_BINDING = { env: "MOSOO_AGENT_URL", @@ -36,51 +49,42 @@ export function boundCapabilityClaims( }; } -/** - * The minimal Deployment authority tables the capability checks read. Mirrors - * the production columns the authority service touches; the public HTTP core - * schema does not include them. - */ -export function createBoundDeploymentAuthoritySchema(database: SqliteD1Database): void { - database.execute(` - CREATE TABLE app_deployment ( - app_id text NOT NULL, - deleted_at integer, - id text PRIMARY KEY NOT NULL - ); - - CREATE TABLE app_deployment_run ( - app_id text NOT NULL, - deployment_id text NOT NULL, - id text PRIMARY KEY NOT NULL, - plan_json text, - status text NOT NULL - ); - - CREATE INDEX app_deployment_run_deployment_id_idx - ON app_deployment_run (deployment_id, id); - `); -} - export async function insertBoundDeployment( database: SqliteD1Database, input: { agentBindings?: unknown[]; deletedAt?: number | null; - deploymentId?: string; - deploymentRunId?: string; + deploymentId?: AppDeploymentId; + deploymentRunId?: AppDeploymentRunId; } = {}, ): Promise { const deploymentId = input.deploymentId ?? BOUND_DEPLOYMENT_ID; + const deploymentRunId = input.deploymentRunId ?? BOUND_DEPLOYMENT_RUN_ID; + const nowMs = nowMsForTest(); await database - .prepare("INSERT INTO app_deployment (app_id, deleted_at, id) VALUES (?, ?, ?)") - .bind(PUBLIC_API_TEST_IDS.app, input.deletedAt ?? null, deploymentId) + .app() + .insert(appDeploymentsTable) + .values({ + appId: PUBLIC_API_TEST_IDS.app, + createdAt: nowMs, + defaultBranch: "main", + deletedAt: input.deletedAt ?? null, + id: deploymentId, + latestRunId: deploymentRunId, + mosooSubdomain: `bound-${deploymentId.toLowerCase()}`, + ownerAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + repoName: "bound-capability", + repoOwner: "mosoo", + repoUrl: "https://github.com/mosoo/bound-capability", + sourceKind: "github_public", + updatedAt: nowMs, + }) .run(); await insertBoundDeploymentRun(database, { agentBindings: input.agentBindings ?? [BOUND_BINDING], deploymentId, - deploymentRunId: input.deploymentRunId ?? BOUND_DEPLOYMENT_RUN_ID, + deploymentRunId, }); } @@ -88,22 +92,33 @@ export async function insertBoundDeploymentRun( database: SqliteD1Database, input: { agentBindings: unknown[]; - deploymentId: string; - deploymentRunId: string; - status?: string; + deploymentId: AppDeploymentId; + deploymentRunId: AppDeploymentRunId; + status?: "failed" | "success"; }, ): Promise { + const nowMs = nowMsForTest(); + await database + .app() + .insert(appDeploymentRunsTable) + .values({ + appId: PUBLIC_API_TEST_IDS.app, + createdAt: nowMs, + deploymentId: input.deploymentId, + id: input.deploymentRunId, + planJson: JSON.stringify({ agentBindings: input.agentBindings }), + sourceBranch: "main", + sourceCommitSha: "a".repeat(40), + status: input.status ?? "success", + updatedAt: nowMs, + }) + .run(); + await database - .prepare( - "INSERT INTO app_deployment_run (app_id, deployment_id, id, plan_json, status) VALUES (?, ?, ?, ?, ?)", - ) - .bind( - PUBLIC_API_TEST_IDS.app, - input.deploymentId, - input.deploymentRunId, - JSON.stringify({ agentBindings: input.agentBindings }), - input.status ?? "success", - ) + .app() + .update(appDeploymentsTable) + .set({ latestRunId: input.deploymentRunId, updatedAt: nowMs }) + .where(eq(appDeploymentsTable.id, input.deploymentId)) .run(); } @@ -112,19 +127,21 @@ export async function deleteBoundDeployment( deploymentId = BOUND_DEPLOYMENT_ID, ): Promise { await database - .prepare("UPDATE app_deployment SET deleted_at = ? WHERE id = ?") - .bind(Date.now(), deploymentId) + .app() + .update(appDeploymentsTable) + .set({ deletedAt: Date.now() }) + .where(eq(appDeploymentsTable.id, deploymentId)) .run(); } -export async function mintBoundCapabilityToken( +async function mintBoundCapabilityToken( bindings: ApiBindings, claims: AppAgentCapabilityClaims = boundCapabilityClaims(), ): Promise { return mintAppAgentCapabilityToken(bindings.RUNTIME_ACTION_TOKEN_SECRET, claims); } -export function boundCapabilityUrl(token: string, path = ""): string { +function boundCapabilityUrl(token: string, path = ""): string { return `https://api.example.com/api/v1/bound/${token}${path}`; } diff --git a/apps/api/tests/bound-capability-public-thread-api.e2e.test.ts b/apps/api/tests/bound-capability-public-thread-api.e2e.test.ts index 22b9c7cb..7ca60d88 100644 --- a/apps/api/tests/bound-capability-public-thread-api.e2e.test.ts +++ b/apps/api/tests/bound-capability-public-thread-api.e2e.test.ts @@ -1,9 +1,13 @@ import { describe, expect, test } from "bun:test"; -import { sessionRunsTable, sessionsTable } from "@mosoo/db"; -import { eq } from "drizzle-orm"; +import { createPlatformId } from "@mosoo/id"; -import { insertSessionMessage } from "../src/modules/sessions/infrastructure/session-message-store.repository"; +import { recordCanonicalSessionRunTerminal } from "../src/modules/runtime/application/session-runs/session-run-terminal-failure.service"; +import { prepareAssistantMessageProjection } from "../src/modules/runtime/infrastructure/driver-instance/assistant-message-projection"; +import { + appendSessionRuntimeEvents, + createSessionRuntimeEvent, +} from "../src/modules/sessions/application/session-event-write.service"; import { BOUND_DEPLOYMENT_ID, BOUND_DEPLOYMENT_RUN_ID, @@ -13,7 +17,6 @@ import { BOUND_BINDING, boundCapabilityClaims, createBoundCapabilityClient, - createBoundDeploymentAuthoritySchema, createBoundTestBindings, deleteBoundDeployment, insertBoundDeployment, @@ -52,13 +55,12 @@ interface BoundSurface { async function createBoundSurface(claims = boundCapabilityClaims()): Promise { const database = await createPublicHttpContractDatabase(); - createBoundDeploymentAuthoritySchema(database); await insertBoundDeployment(database); const app = createPublicThreadApiTestApp(); const bucket = new PublicApiMemoryFileBucket(); const bindings = createBoundTestBindings(database, { - fileBucket: bucket as unknown as R2Bucket, + fileBucket: bucket, }); const client = await createBoundCapabilityClient({ app, bindings, claims }); @@ -190,33 +192,55 @@ async function simulateArtifactAndCompletion( httpMetadata: { contentType: "application/zip" }, }); - await insertSessionMessage(surface.database, { - content: FINAL_OUTPUT_TEXT, - createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, - role: "assistant", - segments: [{ kind: "text", text: FINAL_OUTPUT_TEXT }], + const finalMessageId = createPlatformId(); + await appendSessionRuntimeEvents({ + bindings: surface.bindings, + deliver: false, + events: [ + createSessionRuntimeEvent({ + actor: "driver", + kind: "message.added", + occurredAtMs: 1_100, + origin: "driver", + payload: { + content: FINAL_OUTPUT_TEXT, + messageId: finalMessageId, + role: "agent", + }, + runId: input.runId, + sessionId: input.threadId, + sourceEventId: createPlatformId(), + visibility: "participant", + }), + createSessionRuntimeEvent({ + actor: "driver", + kind: "message.completed", + occurredAtMs: 1_125, + origin: "driver", + payload: { messageId: finalMessageId, role: "agent" }, + runId: input.runId, + sessionId: input.threadId, + sourceEventId: createPlatformId(), + visibility: "participant", + }), + ], sessionId: input.threadId, - sessionRunId: input.runId, }); - await surface.database - .app() - .update(sessionRunsTable) - .set({ - completedAt: 1_150, - errorCode: null, - errorDetailsJson: null, - errorMessage: null, - status: "completed", - updatedAt: 1_150, - }) - .where(eq(sessionRunsTable.id, input.runId)) - .run(); - await surface.database - .app() - .update(sessionsTable) - .set({ lastRunId: input.runId, status: "IDLE", updatedAt: 1_150 }) - .where(eq(sessionsTable.id, input.threadId)) - .run(); + await recordCanonicalSessionRunTerminal(surface.bindings, { + assistantMessage: prepareAssistantMessageProjection({ + createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + messageId: finalMessageId, + sessionId: input.threadId, + sessionRunId: input.runId, + }), + deliver: false, + error: null, + runId: input.runId, + sessionId: input.threadId, + source: "api", + status: "completed", + timestampMs: 1_150, + }); return artifactId; } @@ -376,18 +400,6 @@ describe("bound capability Public Thread API e2e", () => { test("keeps the capability inside its App, declared Agent, and Deployment", async () => { const surface = await createBoundSurface(); - await insertBoundDeployment(surface.database, { - deploymentId: BOUND_OTHER_DEPLOYMENT_ID, - deploymentRunId: BOUND_OTHER_DEPLOYMENT_RUN_ID, - }); - const otherDeployment = await createBoundCapabilityClient({ - app: createPublicThreadApiTestApp(), - bindings: surface.bindings, - claims: boundCapabilityClaims({ - deploymentId: BOUND_OTHER_DEPLOYMENT_ID, - deploymentRunId: BOUND_OTHER_DEPLOYMENT_RUN_ID, - }), - }); await withProviderProbeMock(async () => { const fileId = await uploadAttachment(surface); @@ -423,7 +435,22 @@ describe("bound capability Public Thread API e2e", () => { }); expect(ownerThreadContinue.status).toBe(404); - // Another Deployment of the same App and Agent cannot see this Thread. + // Rotate the App to a new Deployment. The production schema permits one + // active Deployment per App, and the successor cannot see its + // predecessor's Threads. + await deleteBoundDeployment(surface.database); + await insertBoundDeployment(surface.database, { + deploymentId: BOUND_OTHER_DEPLOYMENT_ID, + deploymentRunId: BOUND_OTHER_DEPLOYMENT_RUN_ID, + }); + const otherDeployment = await createBoundCapabilityClient({ + app: createPublicThreadApiTestApp(), + bindings: surface.bindings, + claims: boundCapabilityClaims({ + deploymentId: BOUND_OTHER_DEPLOYMENT_ID, + deploymentRunId: BOUND_OTHER_DEPLOYMENT_RUN_ID, + }), + }); const crossDeployment = await otherDeployment.request(`/threads/${threadId}`); expect(crossDeployment.status).toBe(404); const crossDeploymentList = await otherDeployment.request("/threads"); @@ -561,8 +588,8 @@ describe("bound capability Public Thread API e2e", () => { test("cleans up the Thread when deletion wins the guarded Run insert race", async () => { const surface = await createBoundSurface(); const revoking = revokeDeploymentWhenRunInsertStarts(surface.database); - const bindings = createBoundTestBindings(revoking as unknown as SqliteD1Database, { - fileBucket: surface.bucket as unknown as R2Bucket, + const bindings = createBoundTestBindings(revoking, { + fileBucket: surface.bucket, }); const client = await createBoundCapabilityClient({ app: createPublicThreadApiTestApp(), diff --git a/apps/api/tests/cattle-continuation-restore.test.ts b/apps/api/tests/cattle-continuation-restore.test.ts index 920cb710..7af652ab 100644 --- a/apps/api/tests/cattle-continuation-restore.test.ts +++ b/apps/api/tests/cattle-continuation-restore.test.ts @@ -105,9 +105,14 @@ async function createContinuationFixture(): Promise<{ .app() .insert(sandboxesTable) .values({ + agentId: PUBLIC_API_TEST_IDS.agent, + appId: PUBLIC_API_TEST_IDS.app, createdAt: now, id: PUBLIC_API_TEST_IDS.sandbox, + incarnation: 1, kind: "cattle", + networkConstraintsHash: "0".repeat(64), + ownerAccountId: PUBLIC_API_TEST_IDS.ownerAccount, status: "active", subjectId: PUBLIC_API_TEST_IDS.ownerSession, subjectKind: "session", @@ -122,6 +127,7 @@ async function createContinuationFixture(): Promise<{ cwd: SESSION_CWD, originJson: JSON.stringify(ORIGIN), sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxIncarnation: 1, sandboxSessionId: PRIOR_SANDBOX_SESSION_ID, sessionId: PUBLIC_API_TEST_IDS.ownerSession, status: "closed", @@ -136,10 +142,14 @@ async function createContinuationFixture(): Promise<{ dir: SESSION_CWD, id: STORED_BACKUP_ID, keep: false, + operationId: PUBLIC_API_TEST_IDS.operation, sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxIncarnation: 1, + stagingId: STORED_BACKUP_ID, status: "ready", ttlSeconds: 10 * 365 * 24 * 60 * 60, updatedAt: now, + workspaceSessionId: PUBLIC_API_TEST_IDS.ownerSession, }) .run(); @@ -154,6 +164,7 @@ function createInput(sandbox: SandboxHandle) { origin: ORIGIN, sandbox, sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxIncarnation: 1, sessionId: PUBLIC_API_TEST_IDS.ownerSession, }; } diff --git a/apps/api/tests/cattle-terminal-checkpoint.test.ts b/apps/api/tests/cattle-terminal-checkpoint.test.ts index bbcd895e..5518977a 100644 --- a/apps/api/tests/cattle-terminal-checkpoint.test.ts +++ b/apps/api/tests/cattle-terminal-checkpoint.test.ts @@ -1,8 +1,24 @@ import { describe, expect, test } from "bun:test"; -import { releaseTerminalDriverInstanceSessionRun } from "../src/modules/runtime/infrastructure/driver-instance/terminal-run-release"; +import { getSessionRuntimeStatePath } from "@mosoo/agent-driver/paths"; + +import { cleanupDriverInstances } from "../src/modules/runtime/infrastructure/driver-instance/maintenance"; +import { + releaseTerminalDriverInstanceSessionRun, + repairTerminalDriverRuntimeCommandsGlobally, +} from "../src/modules/runtime/infrastructure/driver-instance/terminal-run-release"; +import { repairClaimedDriverStopsGlobally } from "../src/modules/runtime/infrastructure/driver-session-stop.service"; import { encodeSandboxBackupIdForStorage } from "../src/modules/runtime/infrastructure/sandbox-backup-id"; -import type { SandboxHandle } from "../src/modules/runtime/infrastructure/sandbox-handles"; +import { + claimSandboxBackupStageActual, + finalizeSandboxBackupStage, + getSandboxBackupStage, + stageSandboxBackupWrites, +} from "../src/modules/runtime/infrastructure/sandbox-backup-store"; +import type { + RuntimeSubjectIncarnationHandle, + SandboxHandle, +} from "../src/modules/runtime/infrastructure/sandbox-handles"; import { isCattleTerminalCheckpointReadyForNextRun } from "../src/modules/runtime/infrastructure/session-runs/session-run-admission.repository"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { @@ -19,13 +35,19 @@ const CREATED_BACKUP_ID = "550e8400-e29b-41d4-a716-446655440002"; interface CheckpointSandboxState { backupAvailable: boolean; - backupOptions: Array<{ dir: string; ttl: number | undefined }>; + backupOptions: Array<{ + dir: string; + excludes: string[] | undefined; + forbiddenPaths: string[] | undefined; + name: string; + ttl: number | undefined; + }>; createBackupCalls: number; } function createCheckpointSandbox(state: CheckpointSandboxState): { commands: string[]; - sandbox: SandboxHandle; + sandbox: RuntimeSubjectIncarnationHandle & SandboxHandle; } { const commands: string[] = []; const unavailable = async (): Promise => { @@ -35,10 +57,18 @@ function createCheckpointSandbox(state: CheckpointSandboxState): { return { commands, sandbox: { + activateRuntimeSubjectIncarnation: unavailable, configureNetworkConstraints: unavailable, - async createBackup(options) { + createBackup: unavailable, + async createRuntimeSubjectBackup(_incarnation, options) { state.createBackupCalls += 1; - state.backupOptions.push({ dir: options.dir, ttl: options.ttl }); + state.backupOptions.push({ + dir: options.dir, + excludes: options.excludes, + forbiddenPaths: options.forbiddenPaths, + name: options.name, + ttl: options.ttl, + }); if (!state.backupAvailable) { throw new Error("backup service unavailable"); @@ -49,11 +79,14 @@ function createCheckpointSandbox(state: CheckpointSandboxState): { createSession: unavailable, deleteSession: unavailable, destroy: unavailable, + destroyRuntimeSubjectIncarnation: unavailable, async exec(command) { commands.push(command); return { exitCode: 0, stderr: "", stdout: "", success: true }; }, getSession: unavailable, + inspectRuntimeSubjectIncarnation: unavailable, + markRuntimeSubjectIncarnationReady: unavailable, mkdir: async () => {}, mountBucket: unavailable, readFile: unavailable, @@ -84,33 +117,36 @@ async function createTerminalCheckpointFixture(): Promise<{ WHERE id = '${PUBLIC_API_TEST_IDS.ownerSession}'; INSERT INTO sandbox ( - id, kind, subject_kind, subject_id, status, bind_mount_ready, + agent_id, app_id, id, incarnation, kind, network_constraints_hash, + owner_account_id, subject_kind, subject_id, status, bind_mount_ready, global_mounts_json, created_at, updated_at ) VALUES ( - '${PUBLIC_API_TEST_IDS.sandbox}', 'cattle', 'session', '${PUBLIC_API_TEST_IDS.ownerSession}', + '${PUBLIC_API_TEST_IDS.agent}', '${PUBLIC_API_TEST_IDS.app}', + '${PUBLIC_API_TEST_IDS.sandbox}', 1, 'cattle', '${"0".repeat(64)}', + '${PUBLIC_API_TEST_IDS.ownerAccount}', 'session', '${PUBLIC_API_TEST_IDS.ownerSession}', 'active', 1, '[]', 1, 1 ); INSERT INTO sandbox_session ( cloudflare_session_id, created_at, cwd, origin_json, sandbox_id, - session_id, status, updated_at + sandbox_incarnation, session_id, status, updated_at ) VALUES ( '01J0000000000000000000000Z', 1, '${SESSION_CWD}', '{"callerUserId":"${PUBLIC_API_TEST_IDS.ownerAccount}","entrypoint":"api","executionOwnerUserId":"${PUBLIC_API_TEST_IDS.ownerAccount}","type":"agent"}', - '${PUBLIC_API_TEST_IDS.sandbox}', '${PUBLIC_API_TEST_IDS.ownerSession}', 'active', 1 + '${PUBLIC_API_TEST_IDS.sandbox}', 1, '${PUBLIC_API_TEST_IDS.ownerSession}', 'active', 1 ); INSERT INTO driver_instance ( id, boot_token_expires_at, boot_token_hash, connection_id, created_at, expires_at, heartbeat_count, protocol, protocol_version, runtime, - sandbox_id, sandbox_session_id, status, updated_at + sandbox_id, sandbox_incarnation, sandbox_session_id, status, updated_at ) VALUES ( '${PUBLIC_API_TEST_IDS.driverOwner}', 1, X'01', 'checkpoint-connection', 1, 1, 0, 'orpc-ws', 1, 'openai-runtime', '${PUBLIC_API_TEST_IDS.sandbox}', - '${PUBLIC_API_TEST_IDS.ownerSession}', 'ready', 1 + 1, '${PUBLIC_API_TEST_IDS.ownerSession}', 'ready', 1 ); INSERT INTO session_run ( @@ -125,19 +161,6 @@ async function createTerminalCheckpointFixture(): Promise<{ 'openai-runtime', 'trace-checkpoint', 1, 2, 1, 2 ); - CREATE TABLE native_resume_ref ( - committed_session_run_id text, - committed_value text, - created_at integer NOT NULL, - kind text NOT NULL, - observed_driver_instance_id text, - observed_session_run_id text, - runtime_id text NOT NULL, - session_id text PRIMARY KEY NOT NULL, - updated_at integer NOT NULL, - value text NOT NULL - ); - INSERT INTO native_resume_ref ( created_at, kind, observed_driver_instance_id, observed_session_run_id, runtime_id, session_id, updated_at, value @@ -149,11 +172,13 @@ async function createTerminalCheckpointFixture(): Promise<{ ); INSERT INTO sandbox_backup ( - created_at, dir, id, keep, sandbox_id, session_run_id, status, ttl_seconds, updated_at + created_at, dir, id, keep, sandbox_id, sandbox_incarnation, session_run_id, + staging_id, status, ttl_seconds, updated_at, workspace_session_id ) VALUES ( - 1, '${SESSION_CWD}', '${PRIOR_BACKUP_ID}', 0, '${PUBLIC_API_TEST_IDS.sandbox}', - '${PUBLIC_API_TEST_IDS.runAlt}', 'ready', 315360000, 1 + 1, '${SESSION_CWD}', '${PRIOR_BACKUP_ID}', 0, '${PUBLIC_API_TEST_IDS.sandbox}', 0, + '${PUBLIC_API_TEST_IDS.runAlt}', '${PRIOR_BACKUP_ID}', 'ready', 315360000, 1, + '${PUBLIC_API_TEST_IDS.ownerSession}' ); `); @@ -175,11 +200,293 @@ async function createTerminalCheckpointFixture(): Promise<{ } describe("cattle terminal checkpoint", () => { + test("an old terminal Run cannot commit its resume ref after a successor starts", async () => { + const { database } = await createTerminalCheckpointFixture(); + database.execute(` + UPDATE driver_instance + SET status_operation_id = '${PUBLIC_API_TEST_IDS.run}' + WHERE id = '${PUBLIC_API_TEST_IDS.driverOwner}' + `); + const [write] = await stageSandboxBackupWrites(database, { + admission: { + driverGeneration: 0, + driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner, + incarnation: 1, + kind: "terminal", + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + sessionRunId: PUBLIC_API_TEST_IDS.run, + }, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + targets: [ + { + dir: SESSION_CWD, + updateSandboxLastBackup: false, + workspaceSessionId: PUBLIC_API_TEST_IDS.ownerSession, + }, + ], + ttlSeconds: 100, + }); + if (write?.kind !== "staged") { + throw new Error("Terminal backup stage was not created."); + } + const actualBackupId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440004"); + await claimSandboxBackupStageActual(database, { + actualBackupId, + dir: SESSION_CWD, + sandboxIncarnation: 1, + stagingId: write.stage.id, + }); + await database + .prepare( + `INSERT INTO sandbox_backup ( + created_at, dir, id, keep, sandbox_id, sandbox_incarnation, session_run_id, + staging_id, status, ttl_seconds, updated_at, workspace_session_id + ) SELECT created_at, dir, actual_backup_id, 0, sandbox_id, sandbox_incarnation, + session_run_id, id, 'ready', ttl_seconds, + CAST(unixepoch('subsec') * 1000 AS INTEGER), workspace_session_id + FROM sandbox_backup_staging WHERE id = ?`, + ) + .bind(write.stage.id) + .run(); + database.execute(` + INSERT INTO session_run ( + id, session_id, agent_id, created_by_account_id, deployment_version_id, + deployment_version_number, driver_instance_id, trigger, status, provider, + model, runtime_id, trace_id, started_at, created_at, updated_at + ) VALUES ( + '${PUBLIC_API_TEST_IDS.runAlt}', '${PUBLIC_API_TEST_IDS.ownerSession}', + '${PUBLIC_API_TEST_IDS.agent}', '${PUBLIC_API_TEST_IDS.ownerAccount}', + '${PUBLIC_API_TEST_IDS.deployment}', 1, '${PUBLIC_API_TEST_IDS.driverOwner}', + 'resume', 'running', 'openai', 'gpt-5.4', 'openai-runtime', + 'trace-successor', 3, 3, 3 + ) + `); + + await expect( + finalizeSandboxBackupStage(database, { + actualBackupId, + stagingId: write.stage.id, + }), + ).resolves.toMatchObject({ candidateAccepted: true, complete: false }); + await expect( + database + .prepare("SELECT committed_session_run_id, committed_value FROM native_resume_ref") + .first(), + ).resolves.toEqual({ committed_session_run_id: null, committed_value: null }); + await expect(getSandboxBackupStage(database, write.stage.id)).resolves.not.toBeNull(); + }); + + test("does not admit a Thread from another workspace's same-path checkpoint", async () => { + const foreignBackupId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440003"); + const insertCheckpoint = async ( + database: D1Database, + workspaceSessionId: string, + ): Promise => { + await database + .prepare( + `INSERT INTO sandbox_backup ( + created_at, dir, id, keep, sandbox_id, sandbox_incarnation, session_run_id, + staging_id, status, ttl_seconds, updated_at, workspace_session_id + ) VALUES (2, ?, ?, 0, ?, 1, ?, ?, 'ready', 315360000, 2, ?)`, + ) + .bind( + SESSION_CWD, + foreignBackupId, + PUBLIC_API_TEST_IDS.sandbox, + PUBLIC_API_TEST_IDS.run, + foreignBackupId, + workspaceSessionId, + ) + .run(); + }; + + const { database } = await createTerminalCheckpointFixture(); + await insertCheckpoint(database, PUBLIC_API_TEST_IDS.nonOwnerSession); + + await expect( + isCattleTerminalCheckpointReadyForNextRun(database, PUBLIC_API_TEST_IDS.ownerSession), + ).resolves.toBe(false); + + const { database: ownerDatabase } = await createTerminalCheckpointFixture(); + await insertCheckpoint(ownerDatabase, PUBLIC_API_TEST_IDS.ownerSession); + await expect( + isCattleTerminalCheckpointReadyForNextRun(ownerDatabase, PUBLIC_API_TEST_IDS.ownerSession), + ).resolves.toBe(true); + }); + + test("maintenance hands a stopped foreign owner to the terminal Run without a command ledger", async () => { + const fixture = await createTerminalCheckpointFixture(); + fixture.sandboxState.backupAvailable = true; + fixture.database.execute(` + UPDATE driver_instance + SET expires_at = 9999999999999, status = 'stopped', + status_operation_id = '${PUBLIC_API_TEST_IDS.operation}' + WHERE id = '${PUBLIC_API_TEST_IDS.driverOwner}' + `); + const controlRequests: string[] = []; + const bindings = { + ...fixture.bindings, + DriverConnection: { + get: () => ({ + fetch: async (request: Request) => { + controlRequests.push(new URL(request.url).pathname); + return Response.json({ ok: true }); + }, + }), + idFromName: () => "driver-do-id", + }, + } as unknown as ApiBindings; + + await expect( + fixture.database.prepare("SELECT COUNT(*) AS count FROM driver_command").first(), + ).resolves.toEqual({ count: 0 }); + + await cleanupDriverInstances(bindings); + + expect(controlRequests).toEqual(["/control/destroy"]); + expect(fixture.sandboxState.createBackupCalls).toBe(1); + await expect( + fixture.database + .prepare("SELECT status_operation_id FROM driver_instance WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.driverOwner) + .first(), + ).resolves.toEqual({ status_operation_id: null }); + await expect( + fixture.database + .prepare("SELECT driver_instance_id FROM session_run WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.run) + .first(), + ).resolves.toEqual({ driver_instance_id: PUBLIC_API_TEST_IDS.driverOwner }); + }); + + test("does not run foreign-owner terminal side effects after the Driver generation rotates", async () => { + const fixture = await createTerminalCheckpointFixture(); + fixture.sandboxState.backupAvailable = true; + fixture.database.execute(` + UPDATE driver_instance + SET status = 'stopped', status_operation_id = '${PUBLIC_API_TEST_IDS.operation}' + WHERE id = '${PUBLIC_API_TEST_IDS.driverOwner}' + `); + let rotateAfterCandidateRead = true; + const interceptCandidateRead = (statement: D1PreparedStatement): D1PreparedStatement => + new Proxy(statement, { + get(target, property, receiver) { + if (property === "bind") { + return (...values: unknown[]) => + interceptCandidateRead(Reflect.apply(target.bind, target, values)); + } + if (property === "all" || property === "raw") { + return async (...args: unknown[]) => { + const read = property === "all" ? target.all : target.raw; + const result = await Reflect.apply(read, target, args); + rotateAfterCandidateRead = false; + fixture.database.execute(` + UPDATE driver_instance + SET generation = generation + 1 + WHERE id = '${PUBLIC_API_TEST_IDS.driverOwner}' + `); + return result; + }; + } + return Reflect.get(target, property, receiver); + }, + }); + const racingDatabase = new Proxy(fixture.database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => { + const statement = target.prepare(query); + return rotateAfterCandidateRead && query.includes("status_operation_id") + ? interceptCandidateRead(statement) + : statement; + }; + } + const value = Reflect.get(target, property, target) as unknown; + return typeof value === "function" ? value.bind(target) : value; + }, + }); + let controlRequests = 0; + const bindings = { + ...fixture.bindings, + DB: racingDatabase, + DriverConnection: { + get: () => ({ + fetch: async () => { + controlRequests += 1; + return Response.json({ ok: true }); + }, + }), + idFromName: () => "driver-do-id", + }, + } as unknown as ApiBindings; + + await repairClaimedDriverStopsGlobally(bindings); + + expect(rotateAfterCandidateRead).toBe(false); + expect(controlRequests).toBe(0); + expect(fixture.sandboxState.createBackupCalls).toBe(0); + await expect( + fixture.database + .prepare("SELECT generation, status_operation_id FROM driver_instance WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.driverOwner) + .first(), + ).resolves.toEqual({ generation: 1, status_operation_id: PUBLIC_API_TEST_IDS.operation }); + await expect( + fixture.database + .prepare("SELECT driver_instance_id FROM session_run WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.run) + .first(), + ).resolves.toEqual({ driver_instance_id: PUBLIC_API_TEST_IDS.driverOwner }); + }); + + test("rejects an old Driver generation before checkpoint or conversation side effects", async () => { + const { bindings, commands, database, sandboxState } = await createTerminalCheckpointFixture(); + sandboxState.backupAvailable = true; + database.execute(` + UPDATE driver_instance + SET generation = 1 + WHERE id = '${PUBLIC_API_TEST_IDS.driverOwner}'; + + INSERT INTO session_run ( + id, session_id, agent_id, created_by_account_id, deployment_version_id, + deployment_version_number, driver_instance_id, trigger, status, provider, + model, runtime_id, trace_id, started_at, created_at, updated_at + ) + VALUES ( + '${PUBLIC_API_TEST_IDS.runAlt}', '${PUBLIC_API_TEST_IDS.ownerSession}', + '${PUBLIC_API_TEST_IDS.agent}', '${PUBLIC_API_TEST_IDS.ownerAccount}', + '${PUBLIC_API_TEST_IDS.deployment}', 1, '${PUBLIC_API_TEST_IDS.driverOwner}', + 'resume', 'running', 'openai', 'gpt-5.4', 'openai-runtime', 'trace-r2', 3, 3, 3 + ); + + UPDATE session + SET last_run_id = '${PUBLIC_API_TEST_IDS.runAlt}', status = 'RUNNING', updated_at = 3 + WHERE id = '${PUBLIC_API_TEST_IDS.ownerSession}'; + `); + + await expect( + releaseTerminalDriverInstanceSessionRun(bindings, { + driverGeneration: 0, + driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner, + sessionRunId: PUBLIC_API_TEST_IDS.run, + }), + ).rejects.toThrow("exact Driver ownership"); + expect(sandboxState.createBackupCalls).toBe(0); + expect(commands).toEqual([]); + await expect( + database + .prepare("SELECT status FROM sandbox_session WHERE session_id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first("status"), + ).resolves.toBe("active"); + }); + test("keeps the last good checkpoint and blocks continuation until a retry commits the Run", async () => { const { bindings, commands, database, sandboxState } = await createTerminalCheckpointFixture(); await expect( releaseTerminalDriverInstanceSessionRun(bindings, { + driverGeneration: 0, driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner, sessionRunId: PUBLIC_API_TEST_IDS.run, }), @@ -201,10 +508,7 @@ describe("cattle terminal checkpoint", () => { ).resolves.toEqual({ committed_session_run_id: null, committed_value: null }); sandboxState.backupAvailable = true; - await releaseTerminalDriverInstanceSessionRun(bindings, { - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner, - sessionRunId: PUBLIC_API_TEST_IDS.run, - }); + await repairTerminalDriverRuntimeCommandsGlobally(bindings); await expect( isCattleTerminalCheckpointReadyForNextRun(database, PUBLIC_API_TEST_IDS.ownerSession), @@ -235,17 +539,37 @@ describe("cattle terminal checkpoint", () => { committed_value: "thread-checkpointed", }); expect(sandboxState.backupOptions).toEqual([ - { dir: SESSION_CWD, ttl: 10 * 365 * 24 * 60 * 60 }, - { dir: SESSION_CWD, ttl: 10 * 365 * 24 * 60 * 60 }, + expect.objectContaining({ + dir: SESSION_CWD, + excludes: undefined, + forbiddenPaths: [ + `${getSessionRuntimeStatePath(PUBLIC_API_TEST_IDS.ownerSession, "openai-runtime")}/auth.json`, + ], + ttl: 10 * 365 * 24 * 60 * 60, + }), + expect.objectContaining({ + dir: SESSION_CWD, + excludes: undefined, + forbiddenPaths: [ + `${getSessionRuntimeStatePath(PUBLIC_API_TEST_IDS.ownerSession, "openai-runtime")}/auth.json`, + ], + ttl: 10 * 365 * 24 * 60 * 60, + }), ]); - expect(commands.join("\n")).toContain("session-files"); - expect(commands.join("\n")).toContain("driver-boot-payload-*.json"); - expect(commands.join("\n")).toContain("openai-runtime/auth.json"); + expect( + sandboxState.backupOptions.every(({ name }) => name.startsWith("mosoo:runtime-backup:v1:")), + ).toBe(true); + expect(commands).toEqual([]); + await repairTerminalDriverRuntimeCommandsGlobally(bindings); + expect(sandboxState.createBackupCalls).toBe(2); + const commandCount = commands.length; await releaseTerminalDriverInstanceSessionRun(bindings, { + driverGeneration: 0, driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner, sessionRunId: PUBLIC_API_TEST_IDS.run, }); expect(sandboxState.createBackupCalls).toBe(2); + expect(commands).toHaveLength(commandCount); }); }); diff --git a/apps/api/tests/cloudflare-sandbox-network-contract.test.ts b/apps/api/tests/cloudflare-sandbox-network-contract.test.ts index c3d1f301..a95f3978 100644 --- a/apps/api/tests/cloudflare-sandbox-network-contract.test.ts +++ b/apps/api/tests/cloudflare-sandbox-network-contract.test.ts @@ -8,12 +8,12 @@ interface PackageMetadata { } describe("pinned Cloudflare Sandbox network contract", () => { - test("keeps the 0.12.6 start-time and persisted interception semantics we enforce", async () => { + test("keeps the 0.12.9 start-time and persisted interception semantics we enforce", async () => { const workspaceRequire = createRequire(import.meta.url); const sandboxPackagePath = workspaceRequire.resolve("@cloudflare/sandbox/package.json"); const sandboxPackage = (await Bun.file(sandboxPackagePath).json()) as PackageMetadata; - expect(sandboxPackage.version).toBe("0.12.6"); + expect(sandboxPackage.version).toBe("0.12.9"); const sandboxRequire = createRequire(sandboxPackagePath); const containersPackagePath = sandboxRequire.resolve("@cloudflare/containers/package.json"); diff --git a/apps/api/tests/cost-ledger-reconciliation.test.ts b/apps/api/tests/cost-ledger-reconciliation.test.ts index 44b35eb0..cf43a7c5 100644 --- a/apps/api/tests/cost-ledger-reconciliation.test.ts +++ b/apps/api/tests/cost-ledger-reconciliation.test.ts @@ -147,6 +147,7 @@ async function createReconciliationDatabase(): Promise { native_call_id text, output_tokens integer, provider text NOT NULL, + source_event_seq integer DEFAULT 0 NOT NULL, session_id text NOT NULL, session_run_id text NOT NULL, started_at integer, @@ -157,6 +158,14 @@ async function createReconciliationDatabase(): Promise { UNIQUE (session_run_id, call_key) ); + CREATE TABLE session_event ( + content_text text NOT NULL, + event_type text NOT NULL, + run_id text, + seq integer NOT NULL, + session_id text NOT NULL + ); + CREATE TABLE usage_event ( actor_user_id text NOT NULL, agent_id text NOT NULL, @@ -181,6 +190,7 @@ async function createReconciliationDatabase(): Promise { session_run_id text, source text NOT NULL, source_event_id text NOT NULL, + source_event_seq integer DEFAULT 0 NOT NULL, total_cost_usd_micros integer NOT NULL, usage_contract text NOT NULL, UNIQUE (source, source_event_id) @@ -232,13 +242,17 @@ async function createReconciliationDatabase(): Promise { payload_json text NOT NULL, status text NOT NULL, attempt_count integer DEFAULT 0 NOT NULL, + delivery_generation integer DEFAULT 1 NOT NULL, claim_owner text, claim_expires_at integer, last_error_code text, last_error_message text, completed_at integer, created_at integer NOT NULL, - updated_at integer NOT NULL + updated_at integer NOT NULL, + CONSTRAINT api_command_delivery_generation_check + CHECK (typeof(delivery_generation) = 'integer' + AND delivery_generation BETWEEN 1 AND 9007199254740991) ); CREATE UNIQUE INDEX api_command_dedupe_idx ON api_command (dedupe_key); @@ -663,6 +677,13 @@ describe("cost ledger reconciliation", () => { id: "01J00000000000000000000203", sourceEventId, }); + await database + .prepare( + `INSERT INTO session_event (content_text, event_type, run_id, seq, session_id) + VALUES (?, 'usage.updated', ?, 0, ?)`, + ) + .bind(JSON.stringify({ callId: "rolled-history" }), RUN_ID, SESSION_ID) + .run(); const bindings = createPublicHttpTestBindings(database) as ApiBindings; await runUsageDailyRollup(bindings, new Date(NOW_MS)); diff --git a/apps/api/tests/cost-usage-event.test.ts b/apps/api/tests/cost-usage-event.test.ts index 4b8a4952..b9b2c5ab 100644 --- a/apps/api/tests/cost-usage-event.test.ts +++ b/apps/api/tests/cost-usage-event.test.ts @@ -80,6 +80,7 @@ function createUsageEventDatabase(): SqliteD1Database { session_run_id text, source text NOT NULL, source_event_id text NOT NULL, + source_event_seq integer DEFAULT 0 NOT NULL, total_cost_usd_micros integer NOT NULL, usage_contract text NOT NULL, UNIQUE (source, source_event_id) @@ -394,7 +395,7 @@ describe("cost usage event", () => { }); }); - test("uses reported USD cost for a known model when token counters are unavailable", async () => { + test("preserves reported USD cost across explicit zero token corrections", async () => { const database = createUsageEventDatabase(); const usage = { costAmount: 0.42, @@ -403,7 +404,7 @@ describe("cost usage event", () => { usageContract: "openai_total_with_cached_breakdown", } satisfies SessionUsageSummary; - await recordRuntimeUsageEvent(database, { + const input = { callKey: "known-cost-only-call", driverInstanceId: DRIVER_INSTANCE_ID, nativeCallId: "known-cost-only-native-call", @@ -412,24 +413,39 @@ describe("cost usage event", () => { model: "gpt-5.4", provider: "openai", }, + sourceEventSeq: 1, usage, + }; + + await recordRuntimeUsageEvent(database, input); + await recordRuntimeUsageEvent(database, { + ...input, + sourceEventSeq: 2, + usage: { + inputTokens: 0, + outputTokens: 0, + source: "session_update", + usageContract: "openai_total_with_cached_breakdown", + }, }); const row = await database .prepare( ` - SELECT price_snapshot_json, pricing_status, total_cost_usd_micros + SELECT price_snapshot_json, pricing_status, source_event_seq, total_cost_usd_micros FROM usage_event `, ) .first<{ price_snapshot_json: string | null; pricing_status: string; + source_event_seq: number; total_cost_usd_micros: number; }>(); expect(row).toMatchObject({ pricing_status: "priced", + source_event_seq: 2, total_cost_usd_micros: 420_000, }); expect(JSON.parse(row?.price_snapshot_json ?? "{}")).toEqual({ @@ -440,4 +456,152 @@ describe("cost usage event", () => { tokenCountersUnavailable: true, }); }); + + test("updates durable usage only from a higher source event seq", async () => { + const database = createUsageEventDatabase(); + const input = { + callKey: "sequenced-usage", + driverInstanceId: DRIVER_INSTANCE_ID, + nativeCallId: "sequenced-native-call", + run: RUN_CONTEXT, + sourceEventSeq: 5, + usage: { + inputTokens: 10, + outputTokens: 5, + source: "prompt_response" as const, + usageContract: "openai_total_with_cached_breakdown" as const, + }, + }; + + await recordRuntimeUsageEvent(database, input); + await recordRuntimeUsageEvent(database, { + ...input, + sourceEventSeq: 4, + usage: { ...input.usage, inputTokens: 1 }, + }); + await recordRuntimeUsageEvent(database, input); + await expect( + recordRuntimeUsageEvent(database, { + ...input, + usage: { ...input.usage, inputTokens: 20 }, + }), + ).rejects.toThrow("replayed with conflicting content"); + await recordRuntimeUsageEvent(database, { + ...input, + sourceEventSeq: 6, + usage: { ...input.usage, inputTokens: 30 }, + }); + + expect( + await database + .prepare("SELECT input_tokens, source_event_seq FROM usage_event") + .first<{ input_tokens: number; source_event_seq: number }>(), + ).toEqual({ input_tokens: 30, source_event_seq: 6 }); + }); + + test("merges partial Anthropic cache buckets without losing prior counters", async () => { + const database = createUsageEventDatabase(); + const input = { + callKey: "partial-anthropic", + driverInstanceId: DRIVER_INSTANCE_ID, + nativeCallId: "partial-anthropic", + run: { + ...RUN_CONTEXT, + createdAtMs: Date.UTC(2026, 7, 31), + model: "claude-sonnet-5", + provider: "anthropic", + }, + sourceEventSeq: 5, + usage: { + cachedReadTokens: 2, + inputTokens: 10, + outputTokens: 5, + source: "prompt_response" as const, + usageContract: "anthropic_bucketed" as const, + }, + }; + const partialInput = { + ...input, + run: { ...input.run, createdAtMs: Date.UTC(2026, 8, 1) }, + sourceEventSeq: 6, + usage: { + cachedReadTokens: 3, + source: "prompt_response" as const, + usageContract: "anthropic_bucketed" as const, + }, + }; + + await recordRuntimeUsageEvent(database, input); + await recordRuntimeUsageEvent(database, partialInput); + await recordRuntimeUsageEvent(database, partialInput); + + expect( + await database + .prepare( + `SELECT cache_read_tokens, input_tokens, output_tokens, price_snapshot_json, + source_event_seq, total_cost_usd_micros + FROM usage_event`, + ) + .first(), + ).toEqual({ + cache_read_tokens: 3, + input_tokens: 13, + output_tokens: 5, + price_snapshot_json: JSON.stringify({ + billableInputTokens: 10, + cacheReadUsdPerMillion: 0.2, + cacheWriteUsdPerMillion: 2.5, + inputUsdPerMillion: 2, + longContextApplied: false, + model: "claude-sonnet-5", + outputUsdPerMillion: 10, + provider: "anthropic", + source: "mosoo_seed_2026_07_10", + }), + source_event_seq: 6, + total_cost_usd_micros: 71, + }); + }); + + test("applies explicit zero corrections without creating empty ledger rows", async () => { + const database = createUsageEventDatabase(); + const input = { + callKey: "zero-correction", + driverInstanceId: DRIVER_INSTANCE_ID, + nativeCallId: "zero-correction", + run: { ...RUN_CONTEXT, model: "gpt-5.4", provider: "openai" }, + sourceEventSeq: 1, + usage: { + inputTokens: 0, + outputTokens: 0, + source: "prompt_response" as const, + usageContract: "openai_total_with_cached_breakdown" as const, + }, + }; + + await recordRuntimeUsageEvent(database, input); + expect(await database.prepare("SELECT COUNT(*) AS count FROM usage_event").first()).toEqual({ + count: 0, + }); + + await recordRuntimeUsageEvent(database, { + ...input, + usage: { ...input.usage, inputTokens: 10, outputTokens: 5 }, + }); + await recordRuntimeUsageEvent(database, { ...input, sourceEventSeq: 2 }); + + expect( + await database + .prepare( + `SELECT input_tokens, output_tokens, source_event_seq, total_cost_usd_micros + FROM usage_event`, + ) + .first(), + ).toEqual({ + input_tokens: 0, + output_tokens: 0, + source_event_seq: 2, + total_cost_usd_micros: 0, + }); + }); }); diff --git a/apps/api/tests/cost-usage-idempotency.test.ts b/apps/api/tests/cost-usage-idempotency.test.ts index f797e157..ea92c77f 100644 --- a/apps/api/tests/cost-usage-idempotency.test.ts +++ b/apps/api/tests/cost-usage-idempotency.test.ts @@ -47,12 +47,30 @@ function createUsageDatabase(): SqliteD1Database { session_run_id text, source text NOT NULL, source_event_id text NOT NULL, + source_event_seq integer DEFAULT 0 NOT NULL, total_cost_usd_micros integer NOT NULL, usage_contract text NOT NULL ); CREATE UNIQUE INDEX usage_event_source_event_idx ON usage_event (source, source_event_id); + CREATE TABLE session_model_call ( + call_key text NOT NULL, + driver_instance_id text, + native_call_id text, + session_id text NOT NULL, + session_run_id text NOT NULL, + source_event_seq integer DEFAULT 0 NOT NULL + ); + + CREATE TABLE session_event ( + content_text text NOT NULL, + event_type text NOT NULL, + run_id text, + seq integer NOT NULL, + session_id text NOT NULL + ); + CREATE TABLE usage_daily_rollup ( organization_id text NOT NULL, app_id text NOT NULL, @@ -131,6 +149,40 @@ function createUsageEventInput(): RecordRuntimeUsageEventInput { }; } +async function seedDurableUsageAuthority( + database: SqliteD1Database, + input: { eventSeq: number; nativeCallId?: string | null }, +): Promise { + const nativeCallId = input.nativeCallId === undefined ? "native-call-1" : input.nativeCallId; + const callKey = nativeCallId === null ? "run_usage" : `model_call:${nativeCallId}`; + + await database + .prepare( + `INSERT INTO session_model_call ( + call_key, + driver_instance_id, + native_call_id, + session_id, + session_run_id, + source_event_seq + ) VALUES (?, ?, ?, ?, ?, ?)`, + ) + .bind(callKey, DRIVER_INSTANCE_ID, nativeCallId, SESSION_ID, SESSION_RUN_ID, input.eventSeq) + .run(); + await database + .prepare( + `INSERT INTO session_event (content_text, event_type, run_id, seq, session_id) + VALUES (?, 'usage.updated', ?, ?, ?)`, + ) + .bind( + JSON.stringify(nativeCallId === null ? {} : { callId: nativeCallId }), + SESSION_RUN_ID, + input.eventSeq, + SESSION_ID, + ) + .run(); +} + async function readRollupTotals( database: SqliteD1Database, ): Promise<{ requestCount: number; totalCostUsdMicros: number }> { @@ -151,6 +203,7 @@ describe("runtime usage idempotency across rollup", () => { const database = createUsageDatabase(); const env = { DB: database } as unknown as ApiBindings; + await seedDurableUsageAuthority(database, { eventSeq: 0 }); await recordRuntimeUsageEvent(database, createUsageEventInput()); await runUsageDailyRollup(env, ROLLUP_TIME); @@ -182,6 +235,211 @@ describe("runtime usage idempotency across rollup", () => { expect(afterSecondRollup).toEqual({ requestCount: 1, totalCostUsdMicros: 5_000_000 }); }); + test("does not strand billable usage behind a newer non-recordable receipt", async () => { + const database = createUsageDatabase(); + const env = { DB: database } as unknown as ApiBindings; + const input = { ...createUsageEventInput(), sourceEventSeq: 1 }; + + await seedDurableUsageAuthority(database, { eventSeq: 1 }); + await recordRuntimeUsageEvent(database, input); + await database + .prepare( + `INSERT INTO session_event (content_text, event_type, run_id, seq, session_id) + VALUES (?, 'usage.updated', ?, 2, ?)`, + ) + .bind( + JSON.stringify({ callId: "native-call-1", source: "session_update", totalTokens: 150 }), + SESSION_RUN_ID, + SESSION_ID, + ) + .run(); + await database.prepare("UPDATE session_model_call SET source_event_seq = 2").run(); + await recordRuntimeUsageEvent(database, { + ...input, + sourceEventSeq: 2, + usage: { + callId: "native-call-1", + source: "session_update", + totalTokens: 150, + }, + }); + + await runUsageDailyRollup(env, ROLLUP_TIME); + + expect(await readRollupTotals(database)).toEqual({ + requestCount: 1, + totalCostUsdMicros: 5_000_000, + }); + expect( + await database + .prepare("SELECT COUNT(*) AS count FROM usage_event") + .first<{ count: number }>(), + ).toEqual({ count: 0 }); + expect( + await database + .prepare("SELECT COUNT(*) AS count FROM usage_event_rollup_receipt") + .first<{ count: number }>(), + ).toEqual({ count: 1 }); + + await runUsageDailyRollup(env, ROLLUP_TIME); + expect(await readRollupTotals(database)).toEqual({ + requestCount: 1, + totalCostUsdMicros: 5_000_000, + }); + }); + + test("rolls up only the latest recordable usage once", async () => { + const database = createUsageDatabase(); + const env = { DB: database } as unknown as ApiBindings; + const input = { ...createUsageEventInput(), sourceEventSeq: 1 }; + + await seedDurableUsageAuthority(database, { eventSeq: 1 }); + await recordRuntimeUsageEvent(database, input); + await database + .prepare( + `INSERT INTO session_event (content_text, event_type, run_id, seq, session_id) + VALUES (?, 'usage.updated', ?, 2, ?)`, + ) + .bind(JSON.stringify({ callId: "native-call-1" }), SESSION_RUN_ID, SESSION_ID) + .run(); + await database.prepare("UPDATE session_model_call SET source_event_seq = 2").run(); + + await recordRuntimeUsageEvent(database, { + ...input, + sourceEventSeq: 2, + usage: { ...input.usage, costAmount: 7 }, + }); + + expect( + await database + .prepare("SELECT COUNT(*) AS count, source_event_seq FROM usage_event") + .first<{ count: number; source_event_seq: number }>(), + ).toEqual({ count: 1, source_event_seq: 2 }); + + await runUsageDailyRollup(env, ROLLUP_TIME); + + expect(await readRollupTotals(database)).toEqual({ + requestCount: 1, + totalCostUsdMicros: 7_000_000, + }); + expect( + await database + .prepare("SELECT COUNT(*) AS count FROM usage_event") + .first<{ count: number }>(), + ).toEqual({ count: 0 }); + + await runUsageDailyRollup(env, ROLLUP_TIME); + expect(await readRollupTotals(database)).toEqual({ + requestCount: 1, + totalCostUsdMicros: 7_000_000, + }); + expect( + await database + .prepare("SELECT COUNT(*) AS count FROM usage_event_rollup_receipt") + .first<{ count: number }>(), + ).toEqual({ count: 1 }); + }); + + test("reprices the merged token snapshot before rolling up a partial update", async () => { + const database = createUsageDatabase(); + const env = { DB: database } as unknown as ApiBindings; + const input = { + ...createUsageEventInput(), + run: { + ...createUsageEventInput().run, + model: "gpt-5.4", + provider: "openai", + }, + sourceEventSeq: 1, + usage: { + callId: "native-call-1", + inputTokens: 300_000, + outputTokens: 5, + source: "prompt_response" as const, + usageContract: "openai_total_with_cached_breakdown" as const, + }, + }; + + await seedDurableUsageAuthority(database, { eventSeq: 1 }); + await recordRuntimeUsageEvent(database, input); + await database.prepare("UPDATE session_model_call SET source_event_seq = 2").run(); + const partialInput = { + ...input, + sourceEventSeq: 2, + usage: { + callId: "native-call-1", + outputTokens: 7, + source: "prompt_response", + usageContract: "openai_total_with_cached_breakdown", + }, + }; + + await recordRuntimeUsageEvent(database, partialInput); + await recordRuntimeUsageEvent(database, partialInput); + + expect( + await database + .prepare( + `SELECT input_tokens, output_tokens, price_snapshot_json, + source_event_seq, total_cost_usd_micros + FROM usage_event`, + ) + .first(), + ).toEqual({ + input_tokens: 300_000, + output_tokens: 7, + price_snapshot_json: JSON.stringify({ + billableInputTokens: 300_000, + cacheReadUsdPerMillion: 0.5, + cacheWriteUsdPerMillion: 6.25, + inputUsdPerMillion: 5, + longContextApplied: true, + model: "gpt-5.4", + outputUsdPerMillion: 22.5, + provider: "openai", + source: "mosoo_seed_2026_07_10", + }), + source_event_seq: 2, + total_cost_usd_micros: 1_500_158, + }); + + await runUsageDailyRollup(env, ROLLUP_TIME); + + expect( + await database + .prepare( + `SELECT input_tokens, output_tokens, request_count, total_cost_usd_micros + FROM usage_daily_rollup`, + ) + .first(), + ).toEqual({ + input_tokens: 300_000, + output_tokens: 7, + request_count: 1, + total_cost_usd_micros: 1_500_158, + }); + }); + + test("matches call-less usage receipts through the durable run_usage identity", async () => { + const database = createUsageDatabase(); + const env = { DB: database } as unknown as ApiBindings; + + await seedDurableUsageAuthority(database, { eventSeq: 3, nativeCallId: null }); + await recordRuntimeUsageEvent(database, { + ...createUsageEventInput(), + callKey: "run_usage", + nativeCallId: null, + sourceEventSeq: 3, + usage: { ...createUsageEventInput().usage, callId: null }, + }); + await runUsageDailyRollup(env, ROLLUP_TIME); + + expect(await readRollupTotals(database)).toEqual({ + requestCount: 1, + totalCostUsdMicros: 5_000_000, + }); + }); + test("prunes rollup receipts past the daily rollup retention window", async () => { const database = createUsageDatabase(); const env = { DB: database } as unknown as ApiBindings; diff --git a/apps/api/tests/driver-command-terminal-ack.test.ts b/apps/api/tests/driver-command-terminal-ack.test.ts index 6bcf29cb..8c91050d 100644 --- a/apps/api/tests/driver-command-terminal-ack.test.ts +++ b/apps/api/tests/driver-command-terminal-ack.test.ts @@ -83,6 +83,7 @@ function createDatabase(): GatedTerminalLookupDatabase { acked_at integer, completed_at integer, delivery_connection_id text, + driver_generation integer, driver_instance_id text NOT NULL, error_json text, expires_at integer, @@ -96,13 +97,75 @@ function createDatabase(): GatedTerminalLookupDatabase { ); CREATE TABLE session_run ( + created_at integer DEFAULT 0 NOT NULL, + created_by_account_id text, driver_instance_id text, + error_code text, + error_details_json text, + error_message text, + error_retryable integer, id text PRIMARY KEY NOT NULL, - status text NOT NULL + runtime_id text, + session_id text, + status text NOT NULL, + trace_id text, + updated_at integer DEFAULT 0 NOT NULL + ); + + CREATE TABLE driver_instance ( + connection_id text, + generation integer DEFAULT 0 NOT NULL, + id text PRIMARY KEY NOT NULL, + sandbox_id text NOT NULL, + sandbox_session_id text NOT NULL, + status text NOT NULL, + status_changed_at integer DEFAULT 0 NOT NULL, + status_event text DEFAULT 'driver.provision' NOT NULL, + status_operation_id text, + status_seq integer DEFAULT 0 NOT NULL, + status_source text DEFAULT 'api' NOT NULL, + updated_at integer DEFAULT 0 NOT NULL + ); + + CREATE TABLE sandbox_session ( + origin_json text, + sandbox_id text, + status text, + session_id text PRIMARY KEY NOT NULL + ); + + CREATE TABLE session ( + agent_id text, + app_id text, + creator_account_id text, + id text PRIMARY KEY NOT NULL, + runtime_id text, + type text + ); + + CREATE TABLE agent ( + id text PRIMARY KEY NOT NULL, + owner_account_id text + ); + + CREATE TABLE sandbox ( + id text PRIMARY KEY NOT NULL, + inactive_deadline_at integer, + kind text, + subject_kind text, + updated_at integer DEFAULT 0 NOT NULL + ); + + CREATE TABLE session_event ( + event_type text NOT NULL, + id text PRIMARY KEY NOT NULL, + run_id text, + source_event_id text ); INSERT INTO driver_command ( delivery_connection_id, + driver_generation, driver_instance_id, id, issued_at, @@ -112,6 +175,7 @@ function createDatabase(): GatedTerminalLookupDatabase { status ) VALUES ( 'connection-1', + 0, '${DRIVER_INSTANCE_ID}', '${COMMAND_ID}', 0, @@ -120,30 +184,46 @@ function createDatabase(): GatedTerminalLookupDatabase { 1, 'delivered' ); + + INSERT INTO driver_instance (id, sandbox_id, sandbox_session_id, status) + VALUES ( + '${DRIVER_INSTANCE_ID}', + '01J00000000000000000000004', + '01J00000000000000000000005', + 'ready' + ); + + INSERT INTO session_run (driver_instance_id, id, status) + VALUES ('${DRIVER_INSTANCE_ID}', '${SESSION_RUN_ID}', 'completed'); + + INSERT INTO session_event (event_type, id, run_id, source_event_id) + VALUES ( + 'run.completed', + '01J00000000000000000000006', + '${SESSION_RUN_ID}', + 'session-run-terminal:${SESSION_RUN_ID}:run.completed' + ); `); return database; } describe("terminal runtime command acknowledgement", () => { - test("returns before linked run cleanup finishes", async () => { + test("returns after linked run cleanup finishes", async () => { const database = createDatabase(); - const backgroundTasks: Promise[] = []; const controller = new DriverInstanceRpcCommandController({ env: { DB: database } as ApiBindings, state: { + requireDriverGeneration: () => 0, requireDriverInstanceId: () => DRIVER_INSTANCE_ID, }, - waitUntil: (task: Promise) => { - backgroundTasks.push(task); - }, - withRuntimeLogContext: (fn: () => unknown) => fn(), } as never); const update = controller.handleCommandUpdate( { commandId: COMMAND_ID, driverInstanceId: DRIVER_INSTANCE_ID, + result: { requestId: "request-1" }, status: "completed", }, { @@ -158,15 +238,13 @@ describe("terminal runtime command acknowledgement", () => { new Promise((resolve) => setTimeout(() => resolve(false), 20)), ]); - expect(backgroundTasks).toHaveLength(1); + expect(acknowledgedBeforeCleanup).toBeFalse(); database.cleanupGate.resolve(); const result = await update; - await Promise.all(backgroundTasks); - expect(acknowledgedBeforeCleanup).toBeTrue(); expect(result).toEqual({ ok: true }); await expect( - getRuntimeCommandRecord(database, DRIVER_INSTANCE_ID, COMMAND_ID), + getRuntimeCommandRecord(database, DRIVER_INSTANCE_ID, 0, COMMAND_ID), ).resolves.toMatchObject({ status: "completed" }); }); }); diff --git a/apps/api/tests/driver-event-batch-admission.test.ts b/apps/api/tests/driver-event-batch-admission.test.ts new file mode 100644 index 00000000..5077aa70 --- /dev/null +++ b/apps/api/tests/driver-event-batch-admission.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, test } from "bun:test"; + +import type { DriverEventEnvelope } from "@mosoo/agent-driver/events"; +import { createPlatformId } from "@mosoo/id"; +import type { RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id"; +import { createRuntimeEvent } from "@mosoo/runtime-events"; +import type { RuntimeEventKind } from "@mosoo/runtime-events"; + +import { assertDriverEventBatchTerminalOrder } from "../src/modules/runtime/infrastructure/driver-instance/rpc-event-ingestion-controller"; + +const SESSION_ID = "01J00000000000000000000001" as SessionId; +const RUN_ID = "01J00000000000000000000002" as SessionRunId; + +function driverEvent(kind: RuntimeEventKind, sourceEventId: string): DriverEventEnvelope { + const event = createRuntimeEvent({ + id: createPlatformId(), + kind, + occurredAt: "2026-08-29T00:00:00.000Z", + payload: + kind === "run.failed" + ? { + error: { code: "driver.failed", message: "Driver failed.", retryable: false }, + recoverable: false, + } + : kind === "message.completed" + ? { messageId: "message-1", role: "agent" } + : {}, + runId: RUN_ID, + sessionId: SESSION_ID, + sourceEventId, + }); + + return { event, eventId: sourceEventId, occurredAt: event.occurredAt }; +} + +describe("Driver event batch terminal admission", () => { + test("accepts one terminal event only when it is last", () => { + expect(() => + assertDriverEventBatchTerminalOrder([ + driverEvent("message.completed", "message-completed"), + driverEvent("run.completed", "run-completed"), + ]), + ).not.toThrow(); + }); + + test("rejects multiple terminal events before persistence", () => { + expect(() => + assertDriverEventBatchTerminalOrder([ + driverEvent("run.completed", "run-completed"), + driverEvent("run.failed", "run-failed"), + ]), + ).toThrow("multiple run terminal events"); + }); + + test("rejects any event after a run terminal", () => { + expect(() => + assertDriverEventBatchTerminalOrder([ + driverEvent("run.completed", "run-completed"), + driverEvent("message.completed", "late-message"), + ]), + ).toThrow("must be last"); + }); +}); diff --git a/apps/api/tests/driver-finalization-repair.test.ts b/apps/api/tests/driver-finalization-repair.test.ts index 9a65fafc..24ab57a6 100644 --- a/apps/api/tests/driver-finalization-repair.test.ts +++ b/apps/api/tests/driver-finalization-repair.test.ts @@ -1,32 +1,58 @@ import { describe, expect, test } from "bun:test"; +import { createMcpExecuteFailedEventIdentity } from "@mosoo/agent-driver/events"; +import { + ExternalToolEffectSettlement, + MCP_EXTERNAL_TOOL_EFFECT_SETTLEMENT_MAX_UTF8_BYTES, + measureMcpExternalToolEffectSettlement, +} from "@mosoo/contracts/external-tool-effect"; +import { + RUNTIME_COMMAND_MAX_UTF8_BYTES, + RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES, + measureRuntimeCommandJson, +} from "@mosoo/contracts/runtime-command"; import type { RuntimeCommand } from "@mosoo/contracts/runtime-command"; -import type { DriverCommandId, DriverInstanceId, SessionRunId } from "@mosoo/id"; +import { DURABLE_RUN_ERROR_MAX_UTF8_BYTES } from "@mosoo/contracts/session-run"; +import type { RunError } from "@mosoo/contracts/session-run"; +import type { + DriverCommandId, + DriverInstanceId, + ExternalToolEffectId, + SessionRunId, +} from "@mosoo/id"; -import type { AgentDriverBackend } from "../../driver/src/core/agent-driver-backend"; -import { createAgentDriverContext } from "../../driver/src/core/agent-driver-backend"; -import { DriverCommandDispatcher } from "../../driver/src/core/driver-command-dispatcher"; -import { DriverPermissionBroker } from "../../driver/src/core/driver-permission-broker"; import type { DriverRuntimeIo } from "../../driver/src/core/driver-runtime-io"; import { DriverRuntimeStateMachine } from "../../driver/src/core/driver-runtime-state"; -import type { AgentDriverMcpPort } from "../../driver/src/host-ports"; -import { createBufferedSinkLogger } from "../../driver/src/observability"; -import { createDriverStartInputFromBootPayload } from "../../driver/src/protocol/start"; +import type { AgentDriverMcpExecution } from "../../driver/src/host-ports"; +import { parseRunId } from "../../driver/src/protocol/id"; import type { RuntimeCommand as DriverRuntimeCommand } from "../../driver/src/runtime-command"; -import { driverBootPayload } from "../../driver/tests/driver-boot-payload-fixture"; +import { promiseWithTimeout } from "../../driver/src/utils/async"; +import { + createBackend, + createDispatcher, + FakeDriverRuntimeIo, +} from "../../driver/tests/driver-runtime-boundary-fixtures"; import { recordCanonicalSessionRunFailure } from "../src/modules/runtime/application/session-runs/session-run-terminal-failure.service"; +import { + createFailedSessionRunRuntimeEvent, + createSessionRunUpdatedEvent, +} from "../src/modules/runtime/application/session-runs/session-run-view-events.service"; +import { createSessionRunTerminalSourceId } from "../src/modules/runtime/domain/session-run-terminal-event-id"; +import { commitTerminalRunProjection } from "../src/modules/runtime/infrastructure/driver-instance/completed-run-commit.repository"; import { cleanupDriverInstances } from "../src/modules/runtime/infrastructure/driver-instance/maintenance"; -import { repairFinalizedTerminalDriverRunState } from "../src/modules/runtime/infrastructure/driver-instance/terminal-run-release"; +import { repairFinalizedTerminalDriverRunState as repairFinalizedTerminalDriverRunStateForGeneration } from "../src/modules/runtime/infrastructure/driver-instance/terminal-run-release"; import { - claimExternalToolEffect, - completeExternalToolEffect, - getExternalToolEffectForCommand, - markExternalToolEffectUnknown, + claimExternalToolEffect as claimExternalToolEffectRecord, + getExternalToolEffectForCommand as getExternalToolEffectForCommandRecord, + observeExternalToolEffect as observeExternalToolEffectRecord, + settleExternalToolEffect as settleExternalToolEffectRecord, } from "../src/modules/runtime/infrastructure/session-runs/external-tool-effect-store.repository"; import { - createRuntimeCommandRecord, - getRuntimeCommandRecord, + createRuntimeCommandRecord as createRuntimeCommandRecordForGeneration, + getRuntimeCommandRecord as getRuntimeCommandRecordForGeneration, + updateRuntimeCommandRecord as updateRuntimeCommandRecordForGeneration, } from "../src/modules/runtime/infrastructure/session-runs/runtime-command-store.repository"; +import { getSessionRunSummary } from "../src/modules/runtime/infrastructure/session-runs/session-run-store.repository"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { createPublicHttpContractDatabase, @@ -37,8 +63,12 @@ import { import type { SqliteD1Database } from "./helpers/public-api-http-test-fixture"; const FINALIZE_RUN_ID = "01J0000000000000000000000T" as SessionRunId; +const DRIVER_GENERATION = 0; const FINALIZE_COMMAND_ID = "01J0000000000000000000000V" as DriverCommandId; const MCP_COMMAND_ID = "01J0000000000000000000000X" as DriverCommandId; +const MCP_CLAIM_TOKEN = "00000000-0000-4000-8000-000000000001"; +const MCP_OTHER_CLAIM_TOKEN = "00000000-0000-4000-8000-000000000002"; +const DISPATCHER_TEST_PHASE_TIMEOUT_MS = 1_500; const FINALIZE_CLOUDFLARE_SESSION_ID = "01J0000000000000000000000W"; const TURN_INTERRUPTED_MESSAGE = "This turn was interrupted before it completed. Please resend your last request."; @@ -48,6 +78,168 @@ const PROVISION_ERROR = { message: "Driver command dispatch failed.", retryable: false, } as const; +const MCP_SUCCESS_RESULT = { + outputText: "created issue A-1", + requestId: "request-01J0000000000000000000000X", + serverId: "01J0000000000000000000000Y", + toolName: "createIssue", +} as const; + +const repairFinalizedTerminalDriverRunState = ( + bindings: ApiBindings, + input: Omit< + Parameters[1], + "driverGeneration" | "sessionRunId" + >, +) => + bindings.DB.prepare("UPDATE driver_instance SET status = ? WHERE id = ?") + .bind(input.status, PUBLIC_API_TEST_IDS.driverOwner) + .run() + .then(() => + repairFinalizedTerminalDriverRunStateForGeneration(bindings, { + ...input, + driverGeneration: DRIVER_GENERATION, + sessionRunId: FINALIZE_RUN_ID, + }), + ); + +const createRuntimeCommandRecord = ( + database: D1Database, + input: Omit[1], "driverGeneration">, +) => + createRuntimeCommandRecordForGeneration(database, { + ...input, + driverGeneration: DRIVER_GENERATION, + }); + +const getRuntimeCommandRecord = ( + database: D1Database, + driverInstanceId: DriverInstanceId, + commandId: DriverCommandId, +) => getRuntimeCommandRecordForGeneration(database, driverInstanceId, DRIVER_GENERATION, commandId); + +const updateRuntimeCommandRecord = ( + database: D1Database, + input: Omit[1], "driverGeneration">, +) => + updateRuntimeCommandRecordForGeneration(database, { + ...input, + driverGeneration: DRIVER_GENERATION, + }); + +type EffectLookupInput unknown> = Omit< + Parameters[1], + "driverGeneration" +>; + +const claimExternalToolEffect = ( + database: D1Database, + input: EffectLookupInput, +) => claimExternalToolEffectRecord(database, { ...input, driverGeneration: DRIVER_GENERATION }); + +const getExternalToolEffectForCommand = ( + database: D1Database, + input: EffectLookupInput, +) => + getExternalToolEffectForCommandRecord(database, { + ...input, + driverGeneration: DRIVER_GENERATION, + }); + +const observeExternalToolEffect = ( + database: D1Database, + input: EffectLookupInput, +) => observeExternalToolEffectRecord(database, { ...input, driverGeneration: DRIVER_GENERATION }); + +const settleExternalToolEffect = ( + database: D1Database, + input: EffectLookupInput, +) => settleExternalToolEffectRecord(database, { ...input, driverGeneration: DRIVER_GENERATION }); + +function succeededSettlementAtSize(byteLength: number) { + const empty = { + kind: "succeeded", + result: { ...MCP_SUCCESS_RESULT, outputText: "" }, + } as const; + const outputBytes = byteLength - measureMcpExternalToolEffectSettlement(empty); + + if (outputBytes < 0) { + throw new Error("Requested settlement size is smaller than its fixed fields."); + } + + return { + kind: "succeeded" as const, + result: { ...MCP_SUCCESS_RESULT, outputText: "x".repeat(outputBytes) }, + }; +} + +function textFieldAtJsonSize(byteLength: number, create: (text: string) => Value): Value { + const empty = create(""); + const value = create("x".repeat(byteLength - measureRuntimeCommandJson(empty))); + + expect(measureRuntimeCommandJson(value)).toBe(byteLength); + return value; +} + +function mcpCommandAtSize(byteLength: number) { + return textFieldAtJsonSize(byteLength, (argumentsJson) => ({ + ...mcpExecuteCommand(MCP_COMMAND_ID), + argumentsJson, + })); +} + +function inputStartCommandAtSize(byteLength: number) { + return textFieldAtJsonSize(byteLength, (text) => ({ + ...inputStartCommand(MCP_COMMAND_ID), + input: { text }, + })); +} + +async function commitInputCommandTerminalAuthority( + database: D1Database, + input: { error: RunError; status: "failed" } | { error: null; status: "completed" }, +): Promise { + const current = await getSessionRunSummary(database, FINALIZE_RUN_ID); + if (current === null) { + throw new Error("Missing Session Run fixture."); + } + + const timestampMs = Date.now(); + const timestamp = new Date(timestampMs).toISOString(); + const run = { + ...current, + completedAt: timestamp, + error: input.error, + startedAt: current.startedAt ?? timestamp, + status: input.status, + updatedAt: timestamp, + }; + const kind = input.status === "completed" ? "run.completed" : "run.failed"; + const sourceEventId = createSessionRunTerminalSourceId(FINALIZE_RUN_ID, kind); + const event = + input.status === "completed" + ? createSessionRunUpdatedEvent(run, PUBLIC_API_TEST_IDS.ownerSession, "IDLE", sourceEventId) + : createFailedSessionRunRuntimeEvent({ + run, + runError: input.error, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + sourceEventId, + }); + + const outcome = await commitTerminalRunProjection(database, { + assistantMessage: null, + error: input.error, + runId: FINALIZE_RUN_ID, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + source: "api", + targetStatus: input.status, + terminalEvent: { event, occurredAt: timestampMs, sourceEventId }, + timestampMs, + }); + if (outcome.kind === "stale") { + throw new Error("Session Run terminal authority lost its test setup race."); + } +} interface TerminalEventRow { content_text: string; @@ -72,15 +264,14 @@ interface DriverInterruptionBenchmarkMetrics { viewerTerminalRate: string; } -class PersistentEffectDriverIo implements DriverRuntimeIo { +class PersistentEffectDriverIo extends FakeDriverRuntimeIo { readonly completedReceipt = Promise.withResolvers(); readonly effectPersisted = Promise.withResolvers(); - readonly updates: Parameters[0][] = []; - readonly #commands: readonly DriverRuntimeCommand[]; + readonly lostCompletedReceipt = Promise.withResolvers(); readonly #database: SqliteD1Database; readonly #driverInstanceId: DriverInstanceId; readonly #loseCompletedReceipt: boolean; - #commandIndex = 0; + #terminalReceiptObserved = false; constructor(input: { commands: readonly DriverRuntimeCommand[]; @@ -88,141 +279,87 @@ class PersistentEffectDriverIo implements DriverRuntimeIo { driverInstanceId: DriverInstanceId; loseCompletedReceipt?: boolean; }) { - this.#commands = input.commands; + super(input.commands, parseRunId(FINALIZE_RUN_ID)); this.#database = input.database; this.#driverInstanceId = input.driverInstanceId; this.#loseCompletedReceipt = input.loseCompletedReceipt ?? false; } - beginRun(): void {} - - async claimExternalToolEffect( + override async claimExternalToolEffect( input: Parameters[0], _signal: AbortSignal, ): ReturnType { return claimExternalToolEffect(this.#database, { + claimToken: input.claimToken, commandId: input.commandId as DriverCommandId, driverInstanceId: this.#driverInstanceId, }); } - async commandUpdate( + override async commandUpdate( input: Parameters[0], _signal: AbortSignal, ): Promise { this.updates.push(input); - if (this.#loseCompletedReceipt && input.status === "completed") { - throw new Error("injected terminal receipt loss after durable effect completion"); + if (input.status !== "completed") { + return; } - if (input.status === "completed") { - this.completedReceipt.resolve(); + this.#terminalReceiptObserved = true; + if (this.#loseCompletedReceipt) { + this.lostCompletedReceipt.resolve(); + throw new Error("injected terminal receipt loss after durable effect completion"); } + + this.completedReceipt.resolve(); } - async completeExternalToolEffect( - input: Parameters[0], + terminalReceiptObserved(): boolean { + return this.#terminalReceiptObserved; + } + + override async observeExternalToolEffect( + input: Parameters[0], _signal: AbortSignal, - ): Promise { - await completeExternalToolEffect(this.#database, { + ): ReturnType { + return observeExternalToolEffect(this.#database, { commandId: input.commandId as DriverCommandId, driverInstanceId: this.#driverInstanceId, - ...(input.providerReceiptJson === undefined - ? {} - : { providerReceiptJson: input.providerReceiptJson }), - result: input.result, }); - this.effectPersisted.resolve(); } - async completeRun(): Promise {} - - endRun(): void {} - - async failRun(): Promise {} - - async heartbeat(): ReturnType { - return { heartbeatCount: 1, ok: true }; - } - - isDrained(): boolean { - return this.#commandIndex >= this.#commands.length; - } - - async markExternalToolEffectUnknown( - input: Parameters[0], + override async settleExternalToolEffect( + input: Parameters[0], _signal: AbortSignal, - ): Promise { - await markExternalToolEffectUnknown(this.#database, { + ): ReturnType { + const state = await settleExternalToolEffect(this.#database, { + claimToken: input.claimToken, commandId: input.commandId as DriverCommandId, driverInstanceId: this.#driverInstanceId, + effectId: input.effectId as ExternalToolEffectId, + settlement: input.settlement, }); - } - - async nextCommand(_signal: AbortSignal): Promise { - const command = this.#commands[this.#commandIndex] ?? null; - - if (command !== null) { - this.#commandIndex += 1; + if (state.kind === "succeeded") { + this.effectPersisted.resolve(); } - - return command; - } - - async pushEvents( - input: Parameters[0], - ): ReturnType { - return { - accepted: input.events.map((event, index) => ({ seq: index + 1, type: event.kind })), - }; + return state; } } function createPersistentEffectDispatcher(input: { io: PersistentEffectDriverIo; - mcpExecute: AgentDriverMcpPort["execute"]; -}): { dispatcher: DriverCommandDispatcher; logger: ReturnType } { - const logger = createBufferedSinkLogger({ - level: "debug", - service: "persistent-effect-dispatcher-test", - sink: async () => {}, - }); - const backend: AgentDriverBackend = { - cancelActiveTurn: async () => {}, - handleInput: async () => {}, - runtime: "openai-runtime", - start: async () => {}, - stop: async () => {}, - }; - const payload = createDriverStartInputFromBootPayload(driverBootPayload); - const dispatcher = new DriverCommandDispatcher({ - backend, - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner, - isShuttingDown: () => input.io.isDrained(), - permissionRequests: new DriverPermissionBroker(() => logger), - rememberRunFailure: () => {}, - runtimeContextFactory: (socket, runtimeLogger) => - createAgentDriverContext({ - eventSink: socket, - logger: runtimeLogger, - payload, - permission: { request: async () => "reject_once" }, - ports: { - commandSource: { nextCommand: (signal) => socket.nextCommand(signal) }, - mcp: { execute: input.mcpExecute }, - }, - }), + mcpExecute: AgentDriverMcpExecution["execute"]; +}) { + return createDispatcher({ + backend: createBackend(), + isShuttingDown: () => input.io.terminalReceiptObserved(), + mcpExecute: (_command, effect) => input.mcpExecute(effect), runtimeState: new DriverRuntimeStateMachine("ready"), - sandboxId: payload.sandboxId, - shutdown: async () => {}, - shutdownSignal: new AbortController().signal, }); - - return { dispatcher, logger }; } -function inputStartCommand(id: DriverCommandId): RuntimeCommand { +function inputStartCommand(id: DriverCommandId): Extract { return { commandId: id, input: { @@ -234,12 +371,13 @@ function inputStartCommand(id: DriverCommandId): RuntimeCommand { }; } -function mcpExecuteCommand(id: DriverCommandId): RuntimeCommand { +function mcpExecuteCommand(id: DriverCommandId): Extract { return { argumentsJson: '{"title":"do not duplicate"}', commandId: id, kind: "mcp.execute", requestId: `request-${id}`, + runId: FINALIZE_RUN_ID, serverId: "01J0000000000000000000000Y", toolCallId: `tool-${id}`, toolName: "createIssue", @@ -248,56 +386,17 @@ function mcpExecuteCommand(id: DriverCommandId): RuntimeCommand { async function insertFinalizedDriverLeaseFixture(database: SqliteD1Database): Promise { await insertOwnerSession(database); - database.execute(` - CREATE TABLE IF NOT EXISTS driver_command ( - acked_at integer, - completed_at integer, - delivery_connection_id text, - driver_instance_id text NOT NULL, - error_json text, - expires_at integer, - id text PRIMARY KEY NOT NULL, - issued_at integer NOT NULL, - kind text NOT NULL, - payload_json text NOT NULL, - result_json text, - seq integer NOT NULL, - status text NOT NULL - ); - - CREATE TABLE IF NOT EXISTS external_tool_effect ( - attempt_count integer NOT NULL, - command_id text NOT NULL UNIQUE, - created_at integer NOT NULL, - driver_instance_id text NOT NULL, - id text PRIMARY KEY NOT NULL, - idempotency_key text NOT NULL UNIQUE, - provider_receipt_json text, - result_json text, - server_id text NOT NULL, - session_run_id text NOT NULL, - status text NOT NULL, - tool_name text NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE IF NOT EXISTS external_tool_effect_attempt ( - attempt integer NOT NULL, - completed_at integer, - created_at integer NOT NULL, - effect_id text NOT NULL, - provider_receipt_json text, - result_json text, - status text NOT NULL, - PRIMARY KEY (effect_id, attempt) - ); - `); await database .prepare( ` INSERT INTO sandbox ( id, + agent_id, + app_id, + owner_account_id, + incarnation, kind, + network_constraints_hash, subject_kind, subject_id, status, @@ -306,12 +405,17 @@ async function insertFinalizedDriverLeaseFixture(database: SqliteD1Database): Pr created_at, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( PUBLIC_API_TEST_IDS.sandbox, + PUBLIC_API_TEST_IDS.agent, + PUBLIC_API_TEST_IDS.app, + PUBLIC_API_TEST_IDS.ownerAccount, + 1, "pet", + "0".repeat(64), "agent", PUBLIC_API_TEST_IDS.agent, "active", @@ -330,11 +434,12 @@ async function insertFinalizedDriverLeaseFixture(database: SqliteD1Database): Pr cwd, origin_json, sandbox_id, + sandbox_incarnation, session_id, status, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( @@ -348,6 +453,7 @@ async function insertFinalizedDriverLeaseFixture(database: SqliteD1Database): Pr type: "agent", }), PUBLIC_API_TEST_IDS.sandbox, + 1, PUBLIC_API_TEST_IDS.ownerSession, "active", 1, @@ -368,11 +474,12 @@ async function insertFinalizedDriverLeaseFixture(database: SqliteD1Database): Pr protocol_version, runtime, sandbox_id, + sandbox_incarnation, sandbox_session_id, status, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( @@ -387,8 +494,9 @@ async function insertFinalizedDriverLeaseFixture(database: SqliteD1Database): Pr 1, "openai-runtime", PUBLIC_API_TEST_IDS.sandbox, + 1, PUBLIC_API_TEST_IDS.ownerSession, - "stopped", + "ready", 1, ) .run(); @@ -441,6 +549,58 @@ async function insertFinalizedDriverLeaseFixture(database: SqliteD1Database): Pr .run(); } +async function setDriverStatus( + database: SqliteD1Database, + status: "ready" | "failed" | "stopped", +): Promise { + await database + .prepare("UPDATE driver_instance SET status = ?, updated_at = ? WHERE id = ?") + .bind(status, Date.now(), PUBLIC_API_TEST_IDS.driverOwner) + .run(); +} + +async function createTerminalDriverFixture() { + const database = await createPublicHttpContractDatabase(); + await insertFinalizedDriverLeaseFixture(database); + + return { + bindings: createPublicHttpTestBindings(database) as ApiBindings, + database, + driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + }; +} + +async function createReadyMcpEffectFixture( + command: RuntimeCommand = mcpExecuteCommand(MCP_COMMAND_ID), +) { + const fixture = await createTerminalDriverFixture(); + await setDriverStatus(fixture.database, "ready"); + await createRuntimeCommandRecord(fixture.database, { + command, + driverInstanceId: fixture.driverInstanceId, + status: "accepted", + }); + + return fixture; +} + +async function createClaimedMcpEffectFixture( + command: Extract = mcpExecuteCommand(MCP_COMMAND_ID), +) { + const fixture = await createReadyMcpEffectFixture(command); + const claim = await claimExternalToolEffect(fixture.database, { + claimToken: MCP_CLAIM_TOKEN, + commandId: command.commandId as DriverCommandId, + driverInstanceId: fixture.driverInstanceId, + }); + + if (claim.kind !== "claimed") { + throw new Error("Expected the test effect to be claimed."); + } + + return { ...fixture, claim }; +} + async function readTerminalEvents(database: SqliteD1Database): Promise { return database .prepare( @@ -516,13 +676,11 @@ function createBeforeBenchmarkSample() { } async function createAfterBenchmarkSample() { - const database = await createPublicHttpContractDatabase(); - await insertFinalizedDriverLeaseFixture(database); - const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const { bindings, database, driverInstanceId } = await createTerminalDriverFixture(); const startedAt = performance.now(); await repairFinalizedTerminalDriverRunState(bindings, { - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, status: "stopped", }); const visibleInterruptedMs = performance.now() - startedAt; @@ -544,22 +702,20 @@ async function createAfterBenchmarkSample() { describe("driver finalization repair", () => { test("fails active run lease, accepted commands, and publishes a replayable terminal event", async () => { - const database = await createPublicHttpContractDatabase(); - await insertFinalizedDriverLeaseFixture(database); - const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const { bindings, database, driverInstanceId } = await createTerminalDriverFixture(); await createRuntimeCommandRecord(database, { command: inputStartCommand(FINALIZE_COMMAND_ID), - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, status: "accepted", }); await repairFinalizedTerminalDriverRunState(bindings, { - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, status: "stopped", }); await repairFinalizedTerminalDriverRunState(bindings, { - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, status: "stopped", }); @@ -573,11 +729,7 @@ describe("driver finalization repair", () => { ) .bind(PUBLIC_API_TEST_IDS.driverOwner) .first<{ id: string }>(); - const command = await getRuntimeCommandRecord( - database, - PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, - FINALIZE_COMMAND_ID, - ); + const command = await getRuntimeCommandRecord(database, driverInstanceId, FINALIZE_COMMAND_ID); const terminalEvents = await readTerminalEvents(database); expect(run).toEqual({ @@ -586,7 +738,7 @@ describe("driver finalization repair", () => { }); expect(activeLease).toBeNull(); expect(command?.status).toBe("failed"); - expect(command?.error?.code).toBe("driver.command_driver_terminal"); + expect(command?.error?.code).toBe("runtime.turn_interrupted"); expect(terminalEvents).toEqual([ { content_text: TURN_INTERRUPTED_MESSAGE, @@ -596,7 +748,7 @@ describe("driver finalization repair", () => { process_type: "run.failed", run_id: FINALIZE_RUN_ID, seq: 1, - source: "api", + source: "driver", source_event_id: `session-run-terminal:${FINALIZE_RUN_ID}:run.failed`, trace_id: "trace-finalize", visibility: "all_consumers", @@ -605,12 +757,10 @@ describe("driver finalization repair", () => { }); test("deduplicates dispatch repair after driver finalization", async () => { - const database = await createPublicHttpContractDatabase(); - await insertFinalizedDriverLeaseFixture(database); - const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const { bindings, database, driverInstanceId } = await createTerminalDriverFixture(); await repairFinalizedTerminalDriverRunState(bindings, { - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, status: "failed", }); await recordCanonicalSessionRunFailure(bindings, { @@ -630,8 +780,7 @@ describe("driver finalization repair", () => { }); test("does not queue an MCP command when its durable effect intent cannot be prepared", async () => { - const database = await createPublicHttpContractDatabase(); - await insertFinalizedDriverLeaseFixture(database); + const { database, driverInstanceId } = await createTerminalDriverFixture(); await database .prepare("UPDATE session_run SET status = 'failed' WHERE id = ?") .bind(FINALIZE_RUN_ID) @@ -640,34 +789,62 @@ describe("driver finalization repair", () => { await expect( createRuntimeCommandRecord(database, { command: mcpExecuteCommand(MCP_COMMAND_ID), - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, status: "accepted", }), - ).rejects.toThrow("MCP external tool effects require an active Session Run."); + ).rejects.toThrow("MCP external tool effects require the command's active Session Run."); - expect( - await getRuntimeCommandRecord( - database, - PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, - MCP_COMMAND_ID, - ), - ).toBeNull(); + expect(await getRuntimeCommandRecord(database, driverInstanceId, MCP_COMMAND_ID)).toBeNull(); }); - test("persists the MCP effect intent and fences it as unknown after Driver loss", async () => { - const database = await createPublicHttpContractDatabase(); - await insertFinalizedDriverLeaseFixture(database); - const bindings = createPublicHttpTestBindings(database) as ApiBindings; - - await createRuntimeCommandRecord(database, { - command: mcpExecuteCommand(MCP_COMMAND_ID), - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, - status: "accepted", + for (const repair of ["finalizer", "maintenance"] as const) { + test(`${repair} reports an unclaimed MCP intent as safely retryable`, async () => { + const { bindings, database, driverInstanceId } = await createTerminalDriverFixture(); + await createRuntimeCommandRecord(database, { + command: mcpExecuteCommand(MCP_COMMAND_ID), + driverInstanceId, + status: "accepted", + }); + const effect = await getExternalToolEffectForCommand(database, { + commandId: MCP_COMMAND_ID, + driverInstanceId, + }); + + if (repair === "finalizer") { + await repairFinalizedTerminalDriverRunState(bindings, { + driverInstanceId, + status: "stopped", + }); + } else { + await database + .prepare("UPDATE driver_instance SET expires_at = ? WHERE id = ?") + .bind(Date.now() + 60_000, PUBLIC_API_TEST_IDS.driverOwner) + .run(); + await cleanupDriverInstances(bindings); + } + + expect( + await getRuntimeCommandRecord(database, driverInstanceId, MCP_COMMAND_ID), + ).toMatchObject({ + error: { + code: "driver.external_tool_effect_not_executed", + details: { + commandId: MCP_COMMAND_ID, + effectId: effect?.id, + }, + retryable: true, + }, + status: "failed", + }); }); + } + + test("persists the MCP effect intent and fences it as unknown after Driver loss", async () => { + const { bindings, database, driverInstanceId } = await createReadyMcpEffectFixture(); const intent = await getExternalToolEffectForCommand(database, { commandId: MCP_COMMAND_ID, - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, }); expect(intent).toMatchObject({ attemptCount: 0, @@ -677,26 +854,30 @@ describe("driver finalization repair", () => { expect(intent?.idempotencyKey).toHaveLength(26); const claim = await claimExternalToolEffect(database, { + claimToken: MCP_CLAIM_TOKEN, commandId: MCP_COMMAND_ID, - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, }); - expect(claim).toMatchObject({ attempt: 1, kind: "execute" }); + expect(claim).toMatchObject({ attempt: 1, kind: "claimed" }); + + await setDriverStatus(database, "failed"); await repairFinalizedTerminalDriverRunState(bindings, { - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, status: "failed", }); await expect( claimExternalToolEffect(database, { + claimToken: MCP_OTHER_CLAIM_TOKEN, commandId: MCP_COMMAND_ID, - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, }), ).resolves.toMatchObject({ kind: "unknown" }); expect( await getExternalToolEffectForCommand(database, { commandId: MCP_COMMAND_ID, - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, }), ).toMatchObject({ attemptCount: 1, status: "unknown" }); expect( @@ -704,27 +885,496 @@ describe("driver finalization repair", () => { .prepare("SELECT completed_at, status FROM external_tool_effect_attempt") .first<{ completed_at: number; status: string }>(), ).toMatchObject({ status: "unknown" }); + const record = await getRuntimeCommandRecord(database, driverInstanceId, MCP_COMMAND_ID); + expect(record).toMatchObject({ + error: { + code: "driver.external_tool_effect_unknown", + details: { effectId: claim.effectId }, + }, + status: "failed", + }); + if (record?.status !== "failed") { + throw new Error("Expected repaired MCP failure."); + } + + const command = mcpExecuteCommand(MCP_COMMAND_ID); + const failedEvent = createMcpExecuteFailedEventIdentity({ + commandId: MCP_COMMAND_ID, + rawInput: command.argumentsJson, + rawOutput: record.error.message, + title: command.toolName, + toolCallId: command.toolCallId, + }); + expect( + await database + .prepare( + "SELECT source_event_id, tool_input_json, tool_output_text, tool_status FROM session_event WHERE mcp_command_id = ?", + ) + .bind(MCP_COMMAND_ID) + .first(), + ).toEqual({ + source_event_id: failedEvent.sourceEventId, + tool_input_json: failedEvent.payload.rawInput, + tool_output_text: failedEvent.payload.rawOutput, + tool_status: "failed", + }); }); - test("retains an unknown MCP effect after terminal Driver maintenance", async () => { - const database = await createPublicHttpContractDatabase(); - await insertFinalizedDriverLeaseFixture(database); - const bindings = createPublicHttpTestBindings(database) as ApiBindings; + test("recovers a lost claim response only for the same claim token", async () => { + const { database, driverInstanceId } = await createReadyMcpEffectFixture(); - await createRuntimeCommandRecord(database, { - command: mcpExecuteCommand(MCP_COMMAND_ID), - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, - status: "accepted", + await expect( + observeExternalToolEffect(database, { + commandId: MCP_COMMAND_ID, + driverInstanceId, + }), + ).resolves.toMatchObject({ kind: "intent" }); + + const first = await claimExternalToolEffect(database, { + claimToken: MCP_CLAIM_TOKEN, + commandId: MCP_COMMAND_ID, + driverInstanceId, + }); + await expect( + claimExternalToolEffect(database, { + claimToken: MCP_CLAIM_TOKEN, + commandId: MCP_COMMAND_ID, + driverInstanceId, + }), + ).resolves.toEqual(first); + await expect( + observeExternalToolEffect(database, { + commandId: MCP_COMMAND_ID, + driverInstanceId, + }), + ).resolves.toEqual(first); + expect( + await database + .prepare("SELECT COUNT(*) AS count FROM external_tool_effect_attempt") + .first<{ count: number }>(), + ).toEqual({ count: 1 }); + + await expect( + claimExternalToolEffect(database, { + claimToken: MCP_OTHER_CLAIM_TOKEN, + commandId: MCP_COMMAND_ID, + driverInstanceId, + }), + ).resolves.toMatchObject({ effectId: first.effectId, kind: "unknown" }); + await expect( + updateRuntimeCommandRecord(database, { + commandId: MCP_COMMAND_ID, + driverInstanceId, + result: MCP_SUCCESS_RESULT, + status: "completed", + }), + ).resolves.toMatchObject({ kind: "rejected", reason: "illegal_transition" }); + expect( + await database + .prepare("SELECT claim_token, status FROM external_tool_effect_attempt") + .first<{ claim_token: string; status: string }>(), + ).toEqual({ claim_token: MCP_CLAIM_TOKEN, status: "unknown" }); + }); + + test("atomically fences both command-terminal and claim-first interleavings", async () => { + for (const first of ["command", "claim"] as const) { + const { bindings, database, driverInstanceId } = await createReadyMcpEffectFixture(); + + if (first === "claim") { + await expect( + claimExternalToolEffect(database, { + claimToken: MCP_CLAIM_TOKEN, + commandId: MCP_COMMAND_ID, + driverInstanceId, + }), + ).resolves.toMatchObject({ kind: "claimed" }); + } + + await setDriverStatus(database, "failed"); + await repairFinalizedTerminalDriverRunState(bindings, { + driverInstanceId, + status: "failed", + }); + + await expect( + observeExternalToolEffect(database, { + commandId: MCP_COMMAND_ID, + driverInstanceId, + }), + ).resolves.toMatchObject({ kind: first === "claim" ? "unknown" : "intent" }); + const replayClaim = claimExternalToolEffect(database, { + claimToken: MCP_OTHER_CLAIM_TOKEN, + commandId: MCP_COMMAND_ID, + driverInstanceId, + }); + if (first === "claim") { + await expect(replayClaim).resolves.toMatchObject({ kind: "unknown" }); + } else { + await expect(replayClaim).rejects.toThrow("did not reach a stable state"); + } + expect( + await getRuntimeCommandRecord(database, driverInstanceId, MCP_COMMAND_ID), + ).toMatchObject({ + error: { + code: + first === "claim" + ? "driver.external_tool_effect_unknown" + : "driver.external_tool_effect_not_executed", + }, + status: "failed", + }); + + const attemptCount = await database + .prepare("SELECT COUNT(*) AS count FROM external_tool_effect_attempt") + .first<{ count: number }>(); + expect(attemptCount).toEqual({ count: first === "claim" ? 1 : 0 }); + } + }); + + test("returns stored success when settlement acknowledgement is lost", async () => { + const { claim, database, driverInstanceId } = await createClaimedMcpEffectFixture(); + const settlement = { + claimToken: MCP_CLAIM_TOKEN, + commandId: MCP_COMMAND_ID, + driverInstanceId, + effectId: claim.effectId as ExternalToolEffectId, + settlement: { + kind: "succeeded" as const, + providerReceiptJson: '{"orderId":"A-1"}', + result: MCP_SUCCESS_RESULT, + }, + }; + + await expect( + settleExternalToolEffect(database, { + ...settlement, + claimToken: MCP_OTHER_CLAIM_TOKEN, + settlement: { kind: "unknown" }, + }), + ).resolves.toEqual(claim); + await settleExternalToolEffect(database, settlement); + await expect( + settleExternalToolEffect(database, { + ...settlement, + settlement: { + kind: "succeeded", + result: { ...MCP_SUCCESS_RESULT, outputText: "conflicting retry" }, + }, + }), + ).resolves.toEqual({ + effectId: claim.effectId, + kind: "succeeded", + result: MCP_SUCCESS_RESULT, }); - await claimExternalToolEffect(database, { + }); + + test("rejects an oversized succeeded settlement before its ledger batch", async () => { + const command = mcpCommandAtSize(RUNTIME_COMMAND_MAX_UTF8_BYTES); + const exact = succeededSettlementAtSize(MCP_EXTERNAL_TOOL_EFFECT_SETTLEMENT_MAX_UTF8_BYTES); + const oversized = succeededSettlementAtSize( + MCP_EXTERNAL_TOOL_EFFECT_SETTLEMENT_MAX_UTF8_BYTES + 1, + ); + expect(measureMcpExternalToolEffectSettlement(exact)).toBe( + MCP_EXTERNAL_TOOL_EFFECT_SETTLEMENT_MAX_UTF8_BYTES, + ); + expect(ExternalToolEffectSettlement.allows(exact)).toBeTrue(); + expect(ExternalToolEffectSettlement.allows(oversized)).toBeFalse(); + + expect(measureRuntimeCommandJson(command) + measureMcpExternalToolEffectSettlement(exact)).toBe( + 2_000_000 - 128 * 1_024, + ); + + const { bindings, claim, database, driverInstanceId } = + await createClaimedMcpEffectFixture(command); + const originalBatch = database.batch.bind(database); + let batchCalls = 0; + database.batch = ((statements: D1PreparedStatement[]) => { + batchCalls += 1; + return originalBatch(statements); + }) as typeof database.batch; + + await expect( + settleExternalToolEffect(database, { + claimToken: MCP_CLAIM_TOKEN, + commandId: MCP_COMMAND_ID, + driverInstanceId, + effectId: claim.effectId as ExternalToolEffectId, + settlement: oversized, + }), + ).rejects.toThrow(`${MCP_EXTERNAL_TOOL_EFFECT_SETTLEMENT_MAX_UTF8_BYTES} UTF-8 bytes`); + expect(batchCalls).toBe(0); + await expect( + observeExternalToolEffect(database, { commandId: MCP_COMMAND_ID, driverInstanceId }), + ).resolves.toEqual(claim); + expect( + await database + .prepare( + "SELECT completed_at, result_json, status FROM external_tool_effect_attempt WHERE effect_id = ?", + ) + .bind(claim.effectId) + .first(), + ).toEqual({ completed_at: null, result_json: null, status: "claimed" }); + + await expect( + settleExternalToolEffect(database, { + claimToken: MCP_CLAIM_TOKEN, + commandId: MCP_COMMAND_ID, + driverInstanceId, + effectId: claim.effectId as ExternalToolEffectId, + settlement: exact, + }), + ).resolves.toMatchObject({ kind: "succeeded", result: exact.result }); + expect(batchCalls).toBe(1); + await setDriverStatus(database, "failed"); + await repairFinalizedTerminalDriverRunState(bindings, { + driverInstanceId, + status: "failed", + }); + await expect( + updateRuntimeCommandRecord(database, { + commandId: MCP_COMMAND_ID, + driverInstanceId, + result: exact.result, + status: "completed", + }), + ).resolves.toMatchObject({ kind: "duplicate" }); + + const row = await database + .prepare("SELECT payload_json, result_json, status FROM driver_command WHERE id = ?") + .bind(MCP_COMMAND_ID) + .first<{ payload_json: string; result_json: string; status: string }>(); + expect(row?.status).toBe("completed"); + expect(measureRuntimeCommandJson(JSON.parse(row!.payload_json))).toBe( + RUNTIME_COMMAND_MAX_UTF8_BYTES, + ); + expect(measureRuntimeCommandJson(JSON.parse(row!.result_json))).toBe( + measureRuntimeCommandJson(exact.result), + ); + }); + + test("rejects an oversized command before allocating sequence or effect state", async () => { + const { database, driverInstanceId } = await createTerminalDriverFixture(); + await setDriverStatus(database, "ready"); + const command = mcpCommandAtSize(RUNTIME_COMMAND_MAX_UTF8_BYTES + 1); + const originalPrepare = database.prepare; + let prepareCalls = 0; + database.prepare = ((query: string) => { + prepareCalls += 1; + return originalPrepare.call(database, query); + }) as typeof database.prepare; + + try { + await expect( + createRuntimeCommandRecord(database, { + command, + driverInstanceId, + status: "accepted", + }), + ).rejects.toThrow(`Runtime command exceeds ${RUNTIME_COMMAND_MAX_UTF8_BYTES} UTF-8 bytes.`); + } finally { + database.prepare = originalPrepare; + } + + expect(prepareCalls).toBe(0); + expect( + await database + .prepare("SELECT command_seq_cursor FROM driver_instance WHERE id = ?") + .bind(driverInstanceId) + .first(), + ).toEqual({ command_seq_cursor: 0 }); + expect(await database.prepare("SELECT COUNT(*) AS count FROM driver_command").first()).toEqual({ + count: 0, + }); + expect( + await database.prepare("SELECT COUNT(*) AS count FROM external_tool_effect").first(), + ).toEqual({ count: 0 }); + }); + + test("stores the maximum command beside the maximum error payload", async () => { + expect(DURABLE_RUN_ERROR_MAX_UTF8_BYTES).toBe(RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES); + const createError = (message: string): RunError => ({ + code: "driver.failed", + details: {}, + message, + retryable: false, + }); + const exact = textFieldAtJsonSize(RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES, createError); + const oversized = textFieldAtJsonSize( + RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES + 1, + createError, + ); + const command = inputStartCommandAtSize(RUNTIME_COMMAND_MAX_UTF8_BYTES); + const { database, driverInstanceId } = await createReadyMcpEffectFixture(command); + const update = (error: RunError) => + updateRuntimeCommandRecord(database, { + commandId: MCP_COMMAND_ID, + driverInstanceId, + error, + status: "failed", + }); + const originalPrepare = database.prepare; + let prepareCalls = 0; + database.prepare = ((query: string) => { + prepareCalls += 1; + return originalPrepare.call(database, query); + }) as typeof database.prepare; + + try { + await expect(update(oversized)).rejects.toThrow( + `Runtime command terminal payload exceeds ${RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES} UTF-8 bytes.`, + ); + } finally { + database.prepare = originalPrepare; + } + + expect(prepareCalls).toBe(0); + await commitInputCommandTerminalAuthority(database, { error: exact, status: "failed" }); + await expect(update(exact)).resolves.toMatchObject({ kind: "applied" }); + + const row = await database + .prepare( + "SELECT error_json, payload_json, result_json, status FROM driver_command WHERE id = ?", + ) + .bind(MCP_COMMAND_ID) + .first<{ + error_json: string | null; + payload_json: string; + result_json: string | null; + status: string; + }>(); + expect(row?.status).toBe("failed"); + expect(measureRuntimeCommandJson(JSON.parse(row!.payload_json))).toBe( + RUNTIME_COMMAND_MAX_UTF8_BYTES, + ); + expect(measureRuntimeCommandJson(JSON.parse(row!.error_json))).toBe( + RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES, + ); + expect(row!.result_json).toBeNull(); + }); + + test.each(["requestId", "serverId", "toolName"] as const)( + "rejects a succeeded settlement whose %s does not match the immutable intent", + async (field) => { + const { bindings, claim, database, driverInstanceId } = await createClaimedMcpEffectFixture(); + + await expect( + settleExternalToolEffect(database, { + claimToken: MCP_CLAIM_TOKEN, + commandId: MCP_COMMAND_ID, + driverInstanceId, + effectId: claim.effectId as ExternalToolEffectId, + settlement: { + kind: "succeeded", + result: { ...MCP_SUCCESS_RESULT, [field]: `wrong-${field}` }, + }, + }), + ).rejects.toThrow("External tool effect result does not match its immutable command intent."); + await expect( + observeExternalToolEffect(database, { commandId: MCP_COMMAND_ID, driverInstanceId }), + ).resolves.toEqual(claim); + + await setDriverStatus(database, "failed"); + await repairFinalizedTerminalDriverRunState(bindings, { + driverInstanceId, + status: "failed", + }); + + await expect( + observeExternalToolEffect(database, { commandId: MCP_COMMAND_ID, driverInstanceId }), + ).resolves.toMatchObject({ effectId: claim.effectId, kind: "unknown" }); + expect( + await getRuntimeCommandRecord(database, driverInstanceId, MCP_COMMAND_ID), + ).toMatchObject({ + error: { code: "driver.external_tool_effect_unknown" }, + status: "failed", + }); + }, + ); + + test("does not collapse an MCP error result into a successful command result", async () => { + const { bindings, claim, database, driverInstanceId } = await createClaimedMcpEffectFixture(); + const errorResult = { ...MCP_SUCCESS_RESULT, isError: true } as const; + await settleExternalToolEffect(database, { + claimToken: MCP_CLAIM_TOKEN, commandId: MCP_COMMAND_ID, - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, + effectId: claim.effectId as ExternalToolEffectId, + settlement: { kind: "succeeded", result: errorResult }, }); + await setDriverStatus(database, "failed"); await repairFinalizedTerminalDriverRunState(bindings, { - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, - status: "stopped", + driverInstanceId, + status: "failed", + }); + + await expect( + updateRuntimeCommandRecord(database, { + commandId: MCP_COMMAND_ID, + driverInstanceId, + result: MCP_SUCCESS_RESULT, + status: "completed", + }), + ).rejects.toThrow("duplicate conflicts with its durable terminal payload"); + await expect( + updateRuntimeCommandRecord(database, { + commandId: MCP_COMMAND_ID, + driverInstanceId, + result: errorResult, + status: "completed", + }), + ).resolves.toMatchObject({ kind: "duplicate" }); + expect(await getRuntimeCommandRecord(database, driverInstanceId, MCP_COMMAND_ID)).toMatchObject( + { + result: errorResult, + status: "completed", + }, + ); + }); + + test("projects the canonical winner of settlement and terminal repair", async () => { + const { bindings, claim, database, driverInstanceId } = await createClaimedMcpEffectFixture(); + await setDriverStatus(database, "failed"); + + await Promise.all([ + settleExternalToolEffect(database, { + claimToken: MCP_CLAIM_TOKEN, + commandId: MCP_COMMAND_ID, + driverInstanceId, + effectId: claim.effectId as ExternalToolEffectId, + settlement: { kind: "succeeded", result: MCP_SUCCESS_RESULT }, + }), + repairFinalizedTerminalDriverRunState(bindings, { driverInstanceId, status: "failed" }), + ]); + + const state = await observeExternalToolEffect(database, { + commandId: MCP_COMMAND_ID, + driverInstanceId, }); + await repairFinalizedTerminalDriverRunState(bindings, { driverInstanceId, status: "failed" }); + const command = await getRuntimeCommandRecord(database, driverInstanceId, MCP_COMMAND_ID); + + if (state.kind === "succeeded") { + expect(command).toMatchObject({ result: MCP_SUCCESS_RESULT, status: "completed" }); + } else { + expect(state.kind).toBe("unknown"); + expect(command).toMatchObject({ + error: { + code: "driver.external_tool_effect_unknown", + details: { effectId: claim.effectId }, + }, + status: "failed", + }); + } + expect( + await database + .prepare("SELECT completed_at, status FROM external_tool_effect_attempt") + .first<{ completed_at: number | null; status: string }>(), + ).toMatchObject({ completed_at: expect.any(Number), status: state.kind }); + }); + test("retains an unknown MCP effect after terminal Driver maintenance", async () => { + const { bindings, database, driverInstanceId } = await createClaimedMcpEffectFixture(); + await setDriverStatus(database, "stopped"); await cleanupDriverInstances(bindings); expect( @@ -736,43 +1386,39 @@ describe("driver finalization repair", () => { expect( await getExternalToolEffectForCommand(database, { commandId: MCP_COMMAND_ID, - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, }), ).toMatchObject({ status: "unknown" }); + expect(await getRuntimeCommandRecord(database, driverInstanceId, MCP_COMMAND_ID)).toMatchObject( + { + error: { code: "driver.external_tool_effect_unknown" }, + status: "failed", + }, + ); }); test("redelivers a persisted successful MCP result without a second provider invocation", async () => { - const database = await createPublicHttpContractDatabase(); - await insertFinalizedDriverLeaseFixture(database); - - await createRuntimeCommandRecord(database, { - command: mcpExecuteCommand(MCP_COMMAND_ID), - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, - status: "accepted", - }); - await claimExternalToolEffect(database, { - commandId: MCP_COMMAND_ID, - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, - }); - await completeExternalToolEffect(database, { + const { bindings, claim, database, driverInstanceId } = await createClaimedMcpEffectFixture(); + await settleExternalToolEffect(database, { + claimToken: MCP_CLAIM_TOKEN, commandId: MCP_COMMAND_ID, - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, - providerReceiptJson: '{"orderId":"A-1"}', - result: { - outputText: "created issue A-1", - requestId: "request-01J0000000000000000000000X", - serverId: "01J0000000000000000000000Y", - toolName: "createIssue", + driverInstanceId, + effectId: claim.effectId as ExternalToolEffectId, + settlement: { + kind: "succeeded", + providerReceiptJson: '{"orderId":"A-1"}', + result: MCP_SUCCESS_RESULT, }, }); await expect( claimExternalToolEffect(database, { + claimToken: MCP_OTHER_CLAIM_TOKEN, commandId: MCP_COMMAND_ID, - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, }), ).resolves.toMatchObject({ - kind: "completed", + kind: "succeeded", result: { outputText: "created issue A-1", requestId: "request-01J0000000000000000000000X", @@ -780,10 +1426,26 @@ describe("driver finalization repair", () => { toolName: "createIssue", }, }); + await expect( + updateRuntimeCommandRecord(database, { + commandId: MCP_COMMAND_ID, + driverInstanceId, + result: { ...MCP_SUCCESS_RESULT, outputText: "not the stored result" }, + status: "completed", + }), + ).resolves.toMatchObject({ kind: "rejected", reason: "illegal_transition" }); + await setDriverStatus(database, "failed"); + await repairFinalizedTerminalDriverRunState(bindings, { + driverInstanceId, + status: "failed", + }); + expect(await getRuntimeCommandRecord(database, driverInstanceId, MCP_COMMAND_ID)).toMatchObject( + { result: MCP_SUCCESS_RESULT, status: "completed" }, + ); expect( await getExternalToolEffectForCommand(database, { commandId: MCP_COMMAND_ID, - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId, + driverInstanceId, }), ).toMatchObject({ providerReceiptJson: '{"orderId":"A-1"}', status: "succeeded" }); expect( @@ -794,15 +1456,8 @@ describe("driver finalization repair", () => { }); test("restarts a real Driver dispatcher without replaying a persisted MCP effect", async () => { - const database = await createPublicHttpContractDatabase(); - await insertFinalizedDriverLeaseFixture(database); const command = mcpExecuteCommand(MCP_COMMAND_ID) as DriverRuntimeCommand; - const driverInstanceId = PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId; - await createRuntimeCommandRecord(database, { - command, - driverInstanceId, - status: "accepted", - }); + const { database, driverInstanceId } = await createReadyMcpEffectFixture(command); let providerCalls = 0; const firstIo = new PersistentEffectDriverIo({ @@ -813,20 +1468,30 @@ describe("driver finalization repair", () => { }); const first = createPersistentEffectDispatcher({ io: firstIo, - mcpExecute: async (request) => { + mcpExecute: async () => { providerCalls += 1; return { outputText: "created issue A-1", providerReceiptJson: '{"orderId":"A-1"}', - requestId: request.requestId, - serverId: request.serverId, - toolName: request.toolName, + requestId: command.requestId, + serverId: command.serverId, + toolName: command.toolName, }; }, }); - await first.dispatcher.run(firstIo, first.logger); - await firstIo.effectPersisted.promise; + const firstRun = first.dispatcher.run(firstIo, first.logger); + await Promise.all([ + expect(firstRun).rejects.toThrow("terminal status could not be delivered"), + promiseWithTimeout(firstIo.effectPersisted.promise, { + label: "durable MCP effect persistence", + timeoutMs: DISPATCHER_TEST_PHASE_TIMEOUT_MS, + }), + promiseWithTimeout(firstIo.lostCompletedReceipt.promise, { + label: "lost MCP completion receipt", + timeoutMs: DISPATCHER_TEST_PHASE_TIMEOUT_MS, + }), + ]); await first.logger.destroy(); const secondIo = new PersistentEffectDriverIo({ @@ -843,7 +1508,10 @@ describe("driver finalization repair", () => { }); await second.dispatcher.run(secondIo, second.logger); - await secondIo.completedReceipt.promise; + await promiseWithTimeout(secondIo.completedReceipt.promise, { + label: "replayed MCP completion receipt", + timeoutMs: DISPATCHER_TEST_PHASE_TIMEOUT_MS, + }); await second.logger.destroy(); expect(providerCalls).toBe(1); diff --git a/apps/api/tests/driver-instance-http.test.ts b/apps/api/tests/driver-instance-http.test.ts index fae1840c..63603f75 100644 --- a/apps/api/tests/driver-instance-http.test.ts +++ b/apps/api/tests/driver-instance-http.test.ts @@ -6,28 +6,35 @@ import type { DriverInstanceHttpHandler } from "../src/modules/runtime/infrastru import { handleDriverInstanceRequest } from "../src/modules/runtime/infrastructure/driver-instance/http"; interface CapturingDriverInstanceHttpHandler extends DriverInstanceHttpHandler { + readonly destroyGenerations: number[]; readonly destroyReasons: string[]; readonly driverSocketRequests: Request[]; + readonly readyGenerations: number[]; } function createDriverInstanceHttpHandler(): CapturingDriverInstanceHttpHandler { + const destroyGenerations: number[] = []; const destroyReasons: string[] = []; const driverSocketRequests: Request[] = []; + const readyGenerations: number[] = []; return { + destroyGenerations, destroyReasons, driverSocketRequests, + readyGenerations, async acceptDriverSocket(request: Request): Promise { driverSocketRequests.push(request); return Response.json({ ok: true }); }, - async destroy(reason: string): Promise { + async destroy(generation: number, reason: string): Promise { + destroyGenerations.push(generation); destroyReasons.push(reason); }, - async fail(_message: string): Promise { + async fail(_generation: number, _message: string): Promise { throw new Error("Unexpected fail call."); }, - async sendControlCommand(_command: RuntimeCommand): Promise { + async sendControlCommand(_generation: number, _command: RuntimeCommand): Promise { throw new Error("Unexpected command call."); }, snapshot() { @@ -36,14 +43,17 @@ function createDriverInstanceHttpHandler(): CapturingDriverInstanceHttpHandler { async waitForClose() { throw new Error("Unexpected close wait call."); }, - async waitForHeartbeat() { - throw new Error("Unexpected heartbeat wait call."); - }, - async waitForHello() { - throw new Error("Unexpected hello wait call."); - }, - async waitForReady() { - throw new Error("Unexpected ready wait call."); + async waitForReady(generation: number) { + readyGenerations.push(generation); + return { + heartbeatCount: 0, + lastHeartbeatAt: null, + ready: { + at: "2026-01-01T00:00:00.000Z", + driverInstanceId: "driver-instance", + pid: 1, + }, + }; }, }; } @@ -94,14 +104,51 @@ describe("driver instance HTTP boundary", () => { expect(handler.driverSocketRequests).toHaveLength(0); }); - test("uses the default destroy reason for an empty destroy body", async () => { - const { handler, payload, response } = await postDestroyRequest(); + test("requires and preserves the exact ready generation", async () => { + const handler = createDriverInstanceHttpHandler(); + const missing = await handleDriverInstanceRequest( + handler, + new Request("https://driver.local/wait/ready?timeoutMs=1000"), + ).catch((error: unknown) => error); + + expect(missing).toBeInstanceOf(TypeError); + const response = await handleDriverInstanceRequest( + handler, + new Request("https://driver.local/wait/ready?generation=7&timeoutMs=1000"), + ); + expect(response.status).toBe(200); + expect(handler.readyGenerations).toEqual([7]); + }); + + test("removes unused hello and heartbeat wait routes", async () => { + const handler = createDriverInstanceHttpHandler(); + + for (const path of ["/wait/hello", "/wait/heartbeat"]) { + const response = await handleDriverInstanceRequest( + handler, + new Request(`https://driver.local${path}?timeoutMs=1000`), + ); + expect(response.status).toBe(404); + } + }); + + test("uses the default destroy reason while preserving the exact generation", async () => { + const { handler, payload, response } = await postDestroyRequest('{"generation":3}'); expect(response.status).toBe(200); expect(payload).toEqual({ ok: true }); + expect(handler.destroyGenerations).toEqual([3]); expect(handler.destroyReasons).toEqual(["runtime.driver_instance.destroyed"]); }); + test("rejects a destroy without a generation", async () => { + const { handler, response } = await postDestroyRequest(); + + expect(response.status).toBe(400); + expect(handler.destroyGenerations).toEqual([]); + expect(handler.destroyReasons).toEqual([]); + }); + test("rejects malformed destroy JSON instead of silently using the default reason", async () => { const { handler, payload, response } = await postDestroyRequest("{"); @@ -111,4 +158,18 @@ describe("driver instance HTTP boundary", () => { }); expect(handler.destroyReasons).toEqual([]); }); + + test("rejects an invalid runtime command at the HTTP boundary", async () => { + const handler = createDriverInstanceHttpHandler(); + const response = await handleDriverInstanceRequest( + handler, + new Request("https://driver.local/control/send", { + body: JSON.stringify({ commandId: "command-1", kind: "input.start" }), + method: "POST", + }), + ); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toEqual({ error: expect.any(String) }); + }); }); diff --git a/apps/api/tests/driver-instance-record.test.ts b/apps/api/tests/driver-instance-record.test.ts index c6528322..55413068 100644 --- a/apps/api/tests/driver-instance-record.test.ts +++ b/apps/api/tests/driver-instance-record.test.ts @@ -1,7 +1,8 @@ import { describe, expect, test } from "bun:test"; +import { DRIVER_PROTOCOL_VERSION } from "@mosoo/agent-driver/boot"; import { parsePlatformId } from "@mosoo/id"; -import type { DriverInstanceId, SandboxId, SessionId, SessionRunId } from "@mosoo/id"; +import type { DriverInstanceId, McpServerId, SandboxId, SessionId, SessionRunId } from "@mosoo/id"; import { PLATFORM_ID_FIXTURES } from "@mosoo/id/testing"; import { @@ -15,7 +16,10 @@ import { markDriverInstanceFailedIfBootTokenMatches, recordRuntimeProcessStarted, } from "../src/modules/runtime/infrastructure/driver-instance/driver-instance-record.repository"; -import { claimDriverInstanceByBootTokenHash } from "../src/modules/runtime/infrastructure/driver-instance/driver-instance-token.repository"; +import { + claimDriverInstanceByBootTokenHash, + validateDriverInstanceBootTokenHash, +} from "../src/modules/runtime/infrastructure/driver-instance/driver-instance-token.repository"; import { finalizeDriverInstance, markDriverInstanceConnected, @@ -32,37 +36,55 @@ const REPLACEMENT_DRIVER_INSTANCE_ID = parsePlatformId( "replacement driver instance id", ); const SANDBOX_ID = PLATFORM_ID_FIXTURES.sandbox; +const SANDBOX_INCARNATION = 1; const SESSION_ID = PLATFORM_ID_FIXTURES.session; const SESSION_RUN_ID = PLATFORM_ID_FIXTURES.sessionRun; const NEXT_SESSION_RUN_ID = parsePlatformId( "01J0000000000000000000000Q", "next session run id", ); +const EFFECT_COMMAND_ID = "01J0000000000000000000000T"; +const EFFECT_ID = "01J0000000000000000000000V"; +const MCP_SERVER_ID = parsePlatformId("01J0000000000000000000000W", "MCP server id"); +const REPLACEMENT_MCP_SERVER_ID = parsePlatformId( + "01J0000000000000000000000X", + "replacement MCP server id", +); function createDriverInstanceRecordDatabase(): SqliteD1Database { - const database = new SqliteD1Database({ foreignKeys: false }); + const database = new SqliteD1Database(); database.execute(` CREATE TABLE driver_command ( - driver_instance_id text NOT NULL + driver_generation integer DEFAULT 0 NOT NULL, + driver_instance_id text NOT NULL, + id text PRIMARY KEY NOT NULL, + status text NOT NULL, + FOREIGN KEY (driver_instance_id) REFERENCES driver_instance(id) ON DELETE CASCADE ); CREATE TABLE driver_instance_mcp_grant ( + app_id text NOT NULL, auth_type text NOT NULL, - authorization_state text, + authorization_state text NOT NULL, can_invalidate integer NOT NULL, can_refresh integer NOT NULL, created_at integer NOT NULL, credential_id text, driver_instance_id text NOT NULL, server_id text NOT NULL, - updated_at integer NOT NULL + updated_at integer NOT NULL, + FOREIGN KEY (driver_instance_id) REFERENCES driver_instance(id) ON DELETE CASCADE, + UNIQUE (driver_instance_id, server_id) ); CREATE TABLE external_tool_effect ( + command_id text NOT NULL, driver_instance_id text NOT NULL, id text PRIMARY KEY NOT NULL, - status text NOT NULL + status text NOT NULL, + FOREIGN KEY (command_id) REFERENCES driver_command(id) ON DELETE CASCADE, + FOREIGN KEY (driver_instance_id) REFERENCES driver_instance(id) ON DELETE CASCADE ); CREATE TABLE driver_instance ( @@ -89,6 +111,7 @@ function createDriverInstanceRecordDatabase(): SqliteD1Database { restart_count integer NOT NULL, runtime text NOT NULL, sandbox_id text NOT NULL, + sandbox_incarnation integer DEFAULT 0 NOT NULL, sandbox_session_id text NOT NULL, status text NOT NULL, status_changed_at integer DEFAULT 0 NOT NULL, @@ -102,8 +125,15 @@ function createDriverInstanceRecordDatabase(): SqliteD1Database { CREATE TABLE session_run ( driver_instance_id text, id text PRIMARY KEY NOT NULL, + status text DEFAULT 'running' NOT NULL, updated_at integer NOT NULL ); + + CREATE TABLE session ( + id text PRIMARY KEY NOT NULL, + runtime_provisioning_operation_id text, + runtime_provisioning_sandbox_id text + ); `); return database; @@ -192,6 +222,7 @@ function insertDriverRecord( restart_count, runtime, sandbox_id, + sandbox_incarnation, sandbox_session_id, status, status_changed_at, @@ -224,6 +255,7 @@ function insertDriverRecord( 0, 'openai-runtime', '${SANDBOX_ID}', + ${SANDBOX_INCARNATION}, '${SESSION_ID}', '${input.status ?? "provisioning"}', 1, @@ -239,6 +271,28 @@ function insertDriverRecord( .run(); } +function insertExternalEffectFixture( + database: SqliteD1Database, + input: { + commandStatus: "accepted" | "completed" | "failed" | "queued"; + effectStatus: "claimed" | "intent" | "succeeded" | "unknown"; + }, +): void { + database.execute(` + INSERT INTO driver_command (driver_instance_id, id, status) + VALUES ('${DRIVER_INSTANCE_ID}', '${EFFECT_COMMAND_ID}', '${input.commandStatus}'); + INSERT INTO external_tool_effect (command_id, driver_instance_id, id, status) + VALUES ('${EFFECT_COMMAND_ID}', '${DRIVER_INSTANCE_ID}', '${EFFECT_ID}', '${input.effectStatus}'); + INSERT INTO driver_instance_mcp_grant ( + app_id, auth_type, authorization_state, can_invalidate, can_refresh, + created_at, credential_id, driver_instance_id, server_id, updated_at + ) VALUES ( + '${PLATFORM_ID_FIXTURES.app}', 'oauth', 'expired', 0, 0, + 1, NULL, '${DRIVER_INSTANCE_ID}', '${MCP_SERVER_ID}', 1 + ); + `); +} + describe("driver instance records", () => { test("insert-only creation does not overwrite an existing record", async () => { const database = createDriverInstanceRecordDatabase(); @@ -250,6 +304,7 @@ describe("driver instance records", () => { driverInstanceId: DRIVER_INSTANCE_ID, runtime: "openai-runtime", sandboxId: SANDBOX_ID, + sandboxIncarnation: SANDBOX_INCARNATION, sandboxSessionId: SESSION_ID, }); @@ -269,6 +324,10 @@ describe("driver instance records", () => { test("replace creation rotates generation on an existing record", async () => { const database = createDriverInstanceRecordDatabase(); insertDriverRecord(database, {}); + insertExternalEffectFixture(database, { + commandStatus: "queued", + effectStatus: "intent", + }); database.execute(` INSERT INTO session_run (driver_instance_id, id, updated_at) VALUES (NULL, '${NEXT_SESSION_RUN_ID}', 1) @@ -279,7 +338,19 @@ describe("driver instance records", () => { driverInstanceId: DRIVER_INSTANCE_ID, runtime: "openai-runtime", sandboxId: SANDBOX_ID, + sandboxIncarnation: SANDBOX_INCARNATION, sandboxSessionId: SESSION_ID, + mcpGrants: [ + { + appId: PLATFORM_ID_FIXTURES.app, + authType: "bearer", + authorizationState: "active", + canInvalidate: false, + canRefresh: true, + credentialId: null, + serverId: MCP_SERVER_ID, + }, + ], }); await expect(readDriverRecord(database)).resolves.toMatchObject({ @@ -289,6 +360,132 @@ describe("driver instance records", () => { }); expect(result.status).toBe("created"); expect(result.generation).toBe(1); + await expect( + database + .prepare("SELECT id FROM driver_command WHERE id = ?") + .bind(EFFECT_COMMAND_ID) + .first(), + ).resolves.toBeNull(); + await expect( + database.prepare("SELECT id FROM external_tool_effect WHERE id = ?").bind(EFFECT_ID).first(), + ).resolves.toBeNull(); + await expect( + database + .prepare( + "SELECT auth_type, authorization_state, can_refresh FROM driver_instance_mcp_grant WHERE driver_instance_id = ?", + ) + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ + auth_type: "bearer", + authorization_state: "active", + can_refresh: 1, + }); + }); + + test("does not rotate a Driver generation owned by terminal cleanup", async () => { + const database = createDriverInstanceRecordDatabase(); + insertDriverRecord(database, { status: "failed" }); + database.execute(` + UPDATE driver_instance + SET status_operation_id = '${SESSION_RUN_ID}' + WHERE id = '${DRIVER_INSTANCE_ID}' + `); + + await expect( + createDriverInstanceRecord(createBindings(database), { + bootTokenHash: token(2), + driverInstanceId: DRIVER_INSTANCE_ID, + runtime: "openai-runtime", + sandboxId: SANDBOX_ID, + sandboxIncarnation: SANDBOX_INCARNATION, + sandboxSessionId: SESSION_ID, + }), + ).rejects.toThrow("replacement is blocked"); + await expect(readDriverRecord(database)).resolves.toMatchObject({ + bootTokenHex: "01", + generation: 0, + status: "failed", + }); + }); + + for (const effectStatus of ["intent", "claimed", "unknown", "succeeded"] as const) { + test(`replace preserves a protected ${effectStatus} external effect`, async () => { + const database = createDriverInstanceRecordDatabase(); + insertDriverRecord(database, { status: "failed" }); + insertExternalEffectFixture(database, { commandStatus: "accepted", effectStatus }); + + await expect( + createDriverInstanceRecord(createBindings(database), { + bootTokenHash: token(2), + driverInstanceId: DRIVER_INSTANCE_ID, + runtime: "openai-runtime", + sandboxId: SANDBOX_ID, + sandboxIncarnation: SANDBOX_INCARNATION, + sandboxSessionId: SESSION_ID, + mcpGrants: [ + { + appId: PLATFORM_ID_FIXTURES.app, + authType: "bearer", + authorizationState: "active", + canInvalidate: false, + canRefresh: false, + credentialId: null, + serverId: REPLACEMENT_MCP_SERVER_ID, + }, + ], + }), + ).rejects.toThrow("blocked by a protected external effect"); + + await expect(readDriverRecord(database)).resolves.toMatchObject({ + bootTokenHex: "01", + generation: 0, + status: "failed", + }); + await expect( + database + .prepare("SELECT status FROM driver_command WHERE id = ?") + .bind(EFFECT_COMMAND_ID) + .first(), + ).resolves.toEqual({ status: "accepted" }); + await expect( + database + .prepare("SELECT status FROM external_tool_effect WHERE id = ?") + .bind(EFFECT_ID) + .first(), + ).resolves.toEqual({ status: effectStatus }); + await expect( + database + .prepare( + "SELECT COUNT(*) AS count, MIN(auth_type) AS auth_type FROM driver_instance_mcp_grant WHERE driver_instance_id = ?", + ) + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ auth_type: "oauth", count: 1 }); + }); + } + + test("replace discards a succeeded effect only after its command is terminal", async () => { + const database = createDriverInstanceRecordDatabase(); + insertDriverRecord(database, { status: "failed" }); + insertExternalEffectFixture(database, { + commandStatus: "completed", + effectStatus: "succeeded", + }); + + await expect( + createDriverInstanceRecord(createBindings(database), { + bootTokenHash: token(2), + driverInstanceId: DRIVER_INSTANCE_ID, + runtime: "openai-runtime", + sandboxId: SANDBOX_ID, + sandboxIncarnation: SANDBOX_INCARNATION, + sandboxSessionId: SESSION_ID, + }), + ).resolves.toMatchObject({ generation: 1, status: "created" }); + await expect( + database.prepare("SELECT id FROM external_tool_effect WHERE id = ?").bind(EFFECT_ID).first(), + ).resolves.toBeNull(); }); test("does not reuse stopping driver records", async () => { @@ -307,6 +504,7 @@ describe("driver instance records", () => { await expect( getReusableDriverInstanceRecord(database, { sandboxId: SANDBOX_ID, + sandboxIncarnation: SANDBOX_INCARNATION, sandboxSessionId: SESSION_ID, }), ).resolves.toMatchObject({ @@ -334,8 +532,10 @@ describe("driver instance records", () => { const database = createDriverInstanceRecordDatabase(); insertDriverRecord(database, { status: "failed" }); database.execute(` - INSERT INTO external_tool_effect (driver_instance_id, id, status) - VALUES ('${DRIVER_INSTANCE_ID}', '01J0000000000000000000000Z', 'unknown') + INSERT INTO driver_command (driver_instance_id, id, status) + VALUES ('${DRIVER_INSTANCE_ID}', '${EFFECT_COMMAND_ID}', 'failed'); + INSERT INTO external_tool_effect (command_id, driver_instance_id, id, status) + VALUES ('${EFFECT_COMMAND_ID}', '${DRIVER_INSTANCE_ID}', '${EFFECT_ID}', 'unknown') `); await cleanupDriverInstances(createBindings(database)); @@ -343,6 +543,35 @@ describe("driver instance records", () => { await expect(readDriverRecord(database)).resolves.toMatchObject({ status: "failed" }); }); + test("maintenance retains a provisioning cleanup target until its lease is released", async () => { + const database = createDriverInstanceRecordDatabase(); + insertDriverRecord(database, { status: "failed" }); + database.execute(` + INSERT INTO session ( + id, runtime_provisioning_operation_id, runtime_provisioning_sandbox_id + ) VALUES ( + '${SESSION_ID}', '01J0000000000000000000000R', '${SANDBOX_ID}' + ) + `); + + await cleanupDriverInstances(createBindings(database)); + await expect(readDriverRecord(database)).resolves.toMatchObject({ status: "failed" }); + + database.execute(` + UPDATE session + SET runtime_provisioning_operation_id = NULL, + runtime_provisioning_sandbox_id = NULL + WHERE id = '${SESSION_ID}' + `); + await cleanupDriverInstances(createBindings(database)); + await expect( + database + .prepare("SELECT id FROM driver_instance WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toBeNull(); + }); + test("maintenance gives connecting drivers the cold ready budget", async () => { const database = createDriverInstanceRecordDatabase(); insertDriverRecord(database, { @@ -493,26 +722,39 @@ describe("driver instance records", () => { generation: 0, }), ).resolves.toBe(true); - await recordDriverInstanceHello(bindings, { - connectionId: "connection-1", - driverInstanceId: DRIVER_INSTANCE_ID, - generation: 0, - hello: { - capabilities: [], - driverVersion: "driver-test", - pid: 11, - protocolVersion: 2, - runtime: "openai-runtime", - startedAt: "2026-05-08T00:00:00.000Z", - }, - }); - await markDriverInstanceReady(bindings, { + const hello = { + capabilities: [], + driverVersion: "driver-test", + pid: 11, + protocolVersion: DRIVER_PROTOCOL_VERSION, + runtime: "openai-runtime" as const, + startedAt: "2026-05-08T00:00:00.000Z", + }; + await expect( + recordDriverInstanceHello(bindings, { + connectionId: "connection-1", + driverInstanceId: DRIVER_INSTANCE_ID, + generation: 0, + hello, + }), + ).resolves.toBe("applied"); + await expect( + recordDriverInstanceHello(bindings, { + connectionId: "connection-1", + driverInstanceId: DRIVER_INSTANCE_ID, + generation: 0, + hello, + }), + ).resolves.toBe("replay"); + const ready = { at: "2026-05-08T00:00:01.000Z", connectionId: "connection-1", driverInstanceId: DRIVER_INSTANCE_ID, generation: 0, pid: 22, - }); + }; + await expect(markDriverInstanceReady(bindings, ready)).resolves.toBe("applied"); + await expect(markDriverInstanceReady(bindings, ready)).resolves.toBe("replay"); await expect( recordDriverInstanceHello(bindings, { connectionId: "connection-1", @@ -522,12 +764,12 @@ describe("driver instance records", () => { capabilities: [], driverVersion: "late-driver-test", pid: 99, - protocolVersion: 2, + protocolVersion: DRIVER_PROTOCOL_VERSION, runtime: "openai-runtime", startedAt: "2026-05-08T00:00:02.000Z", }, }), - ).resolves.toBe(false); + ).resolves.toBe("conflict"); const row = await database .prepare( @@ -550,6 +792,47 @@ describe("driver instance records", () => { }); }); + test("replays the same boot claim after D1 connected but before DO state persisted", async () => { + const database = createDriverInstanceRecordDatabase(); + insertDriverRecord(database, {}); + const bindings = createBindings(database); + + await claimDriverInstanceByBootTokenHash(bindings, token(1)); + await expect( + markDriverInstanceConnected(bindings, { + bootTokenHash: token(1), + connectedAt: 1, + connectionId: "lost-connection", + driverInstanceId: DRIVER_INSTANCE_ID, + generation: 0, + }), + ).resolves.toBe(true); + + await expect(validateDriverInstanceBootTokenHash(bindings, token(1))).resolves.toEqual({ + driverInstanceId: DRIVER_INSTANCE_ID, + error: null, + generation: 0, + }); + await expect(claimDriverInstanceByBootTokenHash(bindings, token(1))).resolves.toEqual({ + driverInstanceId: DRIVER_INSTANCE_ID, + error: null, + generation: 0, + }); + await expect( + markDriverInstanceConnected(bindings, { + bootTokenHash: token(1), + connectedAt: 2, + connectionId: "recovered-connection", + driverInstanceId: DRIVER_INSTANCE_ID, + generation: 0, + }), + ).resolves.toBe(true); + await expect(readDriverRecord(database)).resolves.toMatchObject({ + connectionId: "recovered-connection", + status: "connecting", + }); + }); + test("does not let terminal driver finalization overwrite a failed driver", async () => { const database = createDriverInstanceRecordDatabase(); insertDriverRecord(database, {}); @@ -571,7 +854,7 @@ describe("driver instance records", () => { heartbeatCount: 3, status: "stopped", }), - ).resolves.toBe(false); + ).resolves.toBeNull(); await expect(readDriverRecord(database)).resolves.toMatchObject({ errorMessage: "failed first", @@ -595,6 +878,7 @@ describe("driver instance records", () => { driverInstanceId: DRIVER_INSTANCE_ID, runtime: "openai-runtime", sandboxId: SANDBOX_ID, + sandboxIncarnation: SANDBOX_INCARNATION, sandboxSessionId: SESSION_ID, }); @@ -628,7 +912,7 @@ describe("driver instance records", () => { heartbeatCount: 1, status: "failed", }), - ).resolves.toBe(false); + ).resolves.toBeNull(); await expect(readDriverRecord(database)).resolves.toMatchObject({ bootTokenHex: "02", diff --git a/apps/api/tests/driver-instance-runtime-state.test.ts b/apps/api/tests/driver-instance-runtime-state.test.ts new file mode 100644 index 00000000..3c50daa2 --- /dev/null +++ b/apps/api/tests/driver-instance-runtime-state.test.ts @@ -0,0 +1,199 @@ +import { describe, expect, test } from "bun:test"; + +import { DRIVER_PROTOCOL_VERSION } from "@mosoo/agent-driver/boot"; +import type { + DriverHelloInput, + DriverHelloOutput, + DriverReadyInput, +} from "@mosoo/agent-driver/orpc"; +import { PLATFORM_ID_FIXTURES } from "@mosoo/id/testing"; + +import { DriverInstanceRuntimeState } from "../src/modules/runtime/infrastructure/driver-instance/runtime-state"; + +class MemoryStorage { + failHelloCommitOnce = false; + failReadyCommitOnce = false; + readonly values = new Map(); + + async deleteAll(): Promise { + this.values.clear(); + } + + async get(key: string): Promise { + return structuredClone(this.values.get(key)) as T | undefined; + } + + async put(key: string, value: unknown): Promise { + if ( + this.failHelloCommitOnce && + typeof value === "object" && + value !== null && + "hello" in value && + value.hello !== null && + "pendingHello" in value && + value.pendingHello === null + ) { + this.failHelloCommitOnce = false; + throw new Error("DO hello commit failed"); + } + if ( + this.failReadyCommitOnce && + typeof value === "object" && + value !== null && + "ready" in value && + value.ready !== null && + "pendingReady" in value && + value.pendingReady === null + ) { + this.failReadyCommitOnce = false; + throw new Error("DO ready commit failed"); + } + + this.values.set(key, structuredClone(value)); + } +} + +const HELLO: DriverHelloInput = { + capabilities: [], + driverVersion: "test-driver", + pid: 42, + protocolVersion: DRIVER_PROTOCOL_VERSION, + runtime: "openai-runtime", + startedAt: "2026-08-30T00:00:00.000Z", +}; + +const HELLO_OUTPUT: DriverHelloOutput = { + acceptedCapabilities: [], + connectionId: "connection-0", + driverInstanceId: PLATFORM_ID_FIXTURES.driverInstance, + heartbeatIntervalMs: 1_000, + runConfig: { + commandLeaseMs: 1_000, + envPolicy: "strict", + eventBatchMaxSize: 64, + organizationPath: "/workspace", + }, + runId: PLATFORM_ID_FIXTURES.sessionRun, +}; + +const READY: DriverReadyInput = { + at: "2026-08-30T00:00:01.000Z", + driverInstanceId: PLATFORM_ID_FIXTURES.driverInstance, + pid: 42, +}; + +async function createState(storage = new MemoryStorage(), generation = 0) { + const state = new DriverInstanceRuntimeState({ storage }); + await state.load(); + await state.initializeDriverInstance(PLATFORM_ID_FIXTURES.driverInstance, generation); + await state.recordAcceptedConnection({ + connectedAt: 1, + connectionId: `connection-${generation}`, + driverGeneration: generation, + traceId: null, + }); + return { state, storage }; +} + +describe("driver instance durable handshake receipts", () => { + test("resumes an exact hello after D1 committed but the DO commit failed", async () => { + const { state, storage } = await createState(); + const epoch = state.requireConnectionEpoch(); + + await expect(state.stageHello(epoch, HELLO, HELLO_OUTPUT)).resolves.toBe("applied"); + storage.failHelloCommitOnce = true; + await expect(state.commitHello(epoch)).rejects.toThrow("DO hello commit failed"); + + const restarted = new DriverInstanceRuntimeState({ storage }); + await restarted.load(); + await expect(restarted.stageHello(epoch, HELLO, HELLO_OUTPUT)).resolves.toBe("resume"); + await expect(restarted.commitHello(epoch)).resolves.toEqual(HELLO_OUTPUT); + await expect(restarted.stageHello(epoch, HELLO, HELLO_OUTPUT)).resolves.toBe("replay"); + await expect(restarted.stageHello(epoch, { ...HELLO, pid: 43 }, HELLO_OUTPUT)).rejects.toThrow( + "conflicts with the canonical receipt", + ); + }); + + test("a same-generation successor requires a fresh hello receipt", async () => { + const { state } = await createState(); + const firstEpoch = state.requireConnectionEpoch(); + await state.stageHello(firstEpoch, HELLO, HELLO_OUTPUT); + await state.commitHello(firstEpoch); + + await state.recordAcceptedConnection({ + connectedAt: 2, + connectionId: "connection-successor", + driverGeneration: 0, + traceId: null, + }); + + expect(state.hello).toBeNull(); + expect(state.pendingHello).toBeNull(); + expect(state.ready).toBeNull(); + await expect(state.stageHello(firstEpoch, HELLO, HELLO_OUTPUT)).rejects.toThrow( + "no longer current", + ); + }); + + test("resumes exact ready projection after its DO commit fails", async () => { + const { state, storage } = await createState(); + const epoch = state.requireConnectionEpoch(); + await state.stageHello(epoch, HELLO, HELLO_OUTPUT); + await state.commitHello(epoch); + await expect(state.stageReady(epoch, READY)).resolves.toBe("applied"); + storage.failReadyCommitOnce = true; + await expect(state.commitReady(epoch)).rejects.toThrow("DO ready commit failed"); + + const restarted = new DriverInstanceRuntimeState({ storage }); + await restarted.load(); + await expect(restarted.stageReady(epoch, READY)).resolves.toBe("resume"); + await expect(restarted.commitReady(epoch)).resolves.toMatchObject({ ready: READY }); + await expect(restarted.stageReady(epoch, READY)).resolves.toBe("replay"); + await expect(restarted.stageReady(epoch, { ...READY, pid: 99 })).rejects.toThrow( + "conflicts with the canonical receipt", + ); + }); +}); + +describe("driver instance generation waiters", () => { + test("a generation-zero waiter cannot be resolved by generation one", async () => { + const { state } = await createState(); + const oldWait = state.waitForReady(0, 10_000); + const oldRejection = oldWait.catch((error: unknown) => error); + expect(state.readyWaiters).toHaveLength(1); + + await state.resetForReuse({ beforeReset: async () => {}, driverGeneration: 1 }); + expect(state.readyWaiters).toHaveLength(0); + expect(await oldRejection).toEqual( + expect.objectContaining({ + message: expect.stringContaining("generation is no longer current"), + }), + ); + await state.recordAcceptedConnection({ + connectedAt: 2, + connectionId: "connection-1", + driverGeneration: 1, + traceId: null, + }); + const epoch = state.requireConnectionEpoch(); + const output = { ...HELLO_OUTPUT, connectionId: epoch.connectionId }; + await state.stageHello(epoch, HELLO, output); + await state.commitHello(epoch); + + const currentWait = state.waitForReady(1, 10_000); + await state.stageReady(epoch, READY); + const result = await state.commitReady(epoch); + state.resolveReadyWaiters(result, 1); + await expect(currentWait).resolves.toEqual(result); + }); + + test("removes a timed-out waiter", async () => { + const { state } = await createState(); + + await expect(state.waitForReady(0, 1)).rejects.toThrow("timed out"); + expect(state.readyWaiters).toHaveLength(0); + expect(state.closeWaiters).toHaveLength(0); + expect(await state.waitForClose(0, 1).catch(() => null)).toBeNull(); + expect(state.closeWaiters).toHaveLength(0); + }); +}); diff --git a/apps/api/tests/driver-instance-sockets.test.ts b/apps/api/tests/driver-instance-sockets.test.ts index f198ede3..8b880bad 100644 --- a/apps/api/tests/driver-instance-sockets.test.ts +++ b/apps/api/tests/driver-instance-sockets.test.ts @@ -4,22 +4,33 @@ import { DriverInstanceSocketRegistry } from "../src/modules/runtime/infrastruct interface FakeSocket { acceptedTags: string[][]; + attachment: unknown; closes: { code: number; reason: string }[]; readyState: number; close(code?: number, reason?: string): void; + deserializeAttachment(): unknown; + serializeAttachment(value: unknown): void; } const SOCKET_OPEN = 1; +const SOCKET_CLOSING = 2; const SOCKET_CLOSED = 3; +const EPOCH_A = { connectionId: "connection-a", generation: 1 } as const; +const EPOCH_B = { connectionId: "connection-b", generation: 1 } as const; function createFakeSocket(): FakeSocket { const socket: FakeSocket = { acceptedTags: [], + attachment: null, closes: [], readyState: SOCKET_OPEN, close(code = 1000, reason = "") { socket.closes.push({ code, reason }); - socket.readyState = SOCKET_CLOSED; + socket.readyState = SOCKET_CLOSING; + }, + deserializeAttachment: () => socket.attachment, + serializeAttachment(value: unknown) { + socket.attachment = structuredClone(value); }, }; @@ -58,52 +69,98 @@ describe("driver instance socket registry", () => { const registry = new DriverInstanceSocketRegistry(ctx); const socket = createFakeSocket(); - registry.acceptDriverSocket(socket as unknown as WebSocket); + registry.acceptDriverSocket(socket as unknown as WebSocket, EPOCH_A); expect(accepted).toHaveLength(1); expect(accepted[0]?.tags).toEqual(["driver"]); - expect(registry.getDriverSocket()).toBe(socket as unknown as WebSocket); + expect(socket.attachment).toEqual(EPOCH_A); + expect(registry.getDriverSocket(EPOCH_A)).toBe(socket); }); test("finds the driver socket via tags after a hibernation wake", () => { const { ctx } = createFakeContext(); const bootRegistry = new DriverInstanceSocketRegistry(ctx); const socket = createFakeSocket(); - bootRegistry.acceptDriverSocket(socket as unknown as WebSocket); + bootRegistry.acceptDriverSocket(socket as unknown as WebSocket, EPOCH_A); // A wake after eviction constructs a fresh registry with no in-memory // active socket; the tagged socket must still be discoverable. const wokenRegistry = new DriverInstanceSocketRegistry(ctx); - expect(wokenRegistry.getDriverSocket()).toBe(socket as unknown as WebSocket); - expect(wokenRegistry.isActiveDriverSocket(socket as unknown as WebSocket)).toBe(true); - expect(wokenRegistry.isSupersededDriverSocket(socket as unknown as WebSocket)).toBe(false); + expect(wokenRegistry.getDriverSocket(EPOCH_A)).toBe(socket); + expect(wokenRegistry.socketMatchesEpoch(socket as unknown as WebSocket, EPOCH_A)).toBe(true); }); test("marks replaced sockets as superseded without affecting the successor", () => { const { ctx } = createFakeContext(); const registry = new DriverInstanceSocketRegistry(ctx); const first = createFakeSocket(); - registry.acceptDriverSocket(first as unknown as WebSocket); + registry.acceptDriverSocket(first as unknown as WebSocket, EPOCH_A); registry.replaceDriverSockets(); const second = createFakeSocket(); - registry.acceptDriverSocket(second as unknown as WebSocket); + registry.acceptDriverSocket(second as unknown as WebSocket, EPOCH_B); expect(first.closes).toEqual([{ code: 1012, reason: "runtime.socket.replaced" }]); - expect(registry.isSupersededDriverSocket(first as unknown as WebSocket)).toBe(true); - expect(registry.isSupersededDriverSocket(second as unknown as WebSocket)).toBe(false); - expect(registry.getDriverSocket()).toBe(second as unknown as WebSocket); + expect(first.readyState).toBe(SOCKET_CLOSING); + expect(registry.getDriverSocket(EPOCH_A)).toBeNull(); + expect(registry.getDriverSocket(EPOCH_B)).toBe(second); }); - test("does not treat the last closing socket as superseded", () => { + test("hibernation ignores a closing predecessor even when it is listed before the successor", () => { const { ctx } = createFakeContext(); const registry = new DriverInstanceSocketRegistry(ctx); + const first = createFakeSocket(); + registry.acceptDriverSocket(first as unknown as WebSocket, EPOCH_A); + registry.replaceDriverSockets(); + const second = createFakeSocket(); + registry.acceptDriverSocket(second as unknown as WebSocket, EPOCH_B); + + const wokenRegistry = new DriverInstanceSocketRegistry(ctx); + + expect(wokenRegistry.getDriverSocket(EPOCH_B)).toBe(second); + expect(wokenRegistry.getDriverSocket(EPOCH_A)).toBeNull(); + }); + + test("ignores missing and malformed persistent attachments", () => { + const { ctx } = createFakeContext(); const socket = createFakeSocket(); - registry.acceptDriverSocket(socket as unknown as WebSocket); + const registry = new DriverInstanceSocketRegistry(ctx); + ctx.acceptWebSocket(socket, ["driver"]); - socket.readyState = SOCKET_CLOSED; + expect(registry.getDriverSocket(EPOCH_A)).toBeNull(); + socket.attachment = { connectionId: "connection-a", generation: -1 }; + expect(registry.getDriverSocket(EPOCH_A)).toBeNull(); + }); + + test("an old callback resumed after a successor can only close its own socket", async () => { + const { ctx } = createFakeContext(); + const registry = new DriverInstanceSocketRegistry(ctx); + const first = createFakeSocket(); + registry.acceptDriverSocket(first as unknown as WebSocket, EPOCH_A); + const capturedEpoch = registry.getSocketEpoch(first as unknown as WebSocket); + let resume: () => void = () => {}; + const gate = new Promise((resolve) => { + resume = resolve; + }); + const staleCallback = (async () => { + await gate; + + if ( + capturedEpoch !== null && + !registry.isCurrentDriverSocket(first as unknown as WebSocket, capturedEpoch, EPOCH_B) + ) { + first.close(1000, "runtime.socket.superseded"); + } + })(); + + const second = createFakeSocket(); + registry.acceptDriverSocket(second as unknown as WebSocket, EPOCH_B); + resume(); + await staleCallback; - expect(registry.isSupersededDriverSocket(socket as unknown as WebSocket)).toBe(false); + expect(first.closes).toEqual([{ code: 1000, reason: "runtime.socket.superseded" }]); + expect(second.closes).toEqual([]); + expect(registry.getDriverSocket(EPOCH_B)).toBe(second); }); }); diff --git a/apps/api/tests/driver-llm-proxy-route.test.ts b/apps/api/tests/driver-llm-proxy-route.test.ts index edefdb30..913f9401 100644 --- a/apps/api/tests/driver-llm-proxy-route.test.ts +++ b/apps/api/tests/driver-llm-proxy-route.test.ts @@ -1,5 +1,6 @@ import { afterEach, describe, expect, test } from "bun:test"; +import { DRIVER_PROTOCOL_VERSION } from "@mosoo/agent-driver/boot"; import { driverInstancesTable, vendorCredentialsTable } from "@mosoo/db"; import { parsePlatformId } from "@mosoo/id"; import type { DriverInstanceId, AppId, VendorCredentialId } from "@mosoo/id"; @@ -61,7 +62,7 @@ function createDriverRouteTestApp(): Hono { function captureUpstreamFetch(response?: () => Response): CapturedUpstreamRequest[] { const captured: CapturedUpstreamRequest[] = []; - globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + globalThis.fetch = async (input: RequestInfo | URL, init?: RequestInit) => { const request = new Request(input, init); captured.push({ body: request.method === "GET" || request.method === "HEAD" ? null : await request.text(), @@ -83,7 +84,7 @@ function captureUpstreamFetch(response?: () => Response): CapturedUpstreamReques status: 200, }) ); - }) as typeof fetch; + }; return captured; } @@ -93,9 +94,11 @@ async function insertDriverInstance( status: "provisioning" | "connecting" | "ready" | "failed", input: { bootTokenExpiresAt?: number; + bootTokenHash?: Uint8Array; driverInstanceId?: DriverInstanceId; generation?: number; lastHeartbeatAt?: number | null; + sandboxIncarnation?: number; updatedAt?: number; } = {}, ) { @@ -105,7 +108,7 @@ async function insertDriverInstance( .insert(driverInstancesTable) .values({ bootTokenExpiresAt: input.bootTokenExpiresAt ?? nowMs + 60_000, - bootTokenHash: new Uint8Array([1, 2, 3]), + bootTokenHash: input.bootTokenHash ?? new Uint8Array([1, 2, 3]), bootTokenUsedAt: null, closeCode: null, closeReason: null, @@ -122,9 +125,10 @@ async function insertDriverInstance( lastHeartbeatAt: input.lastHeartbeatAt ?? null, processId: null, protocol: "orpc-ws", - protocolVersion: 2, + protocolVersion: DRIVER_PROTOCOL_VERSION, runtime: "claude-agent-sdk", sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxIncarnation: input.sandboxIncarnation ?? 1, sandboxSessionId: PUBLIC_API_TEST_IDS.ownerSession, status, statusChangedAt: nowMs, @@ -580,7 +584,9 @@ describe("driver LLM proxy route", () => { const { bindings, database } = await setupFixture(); await insertDriverInstance(database, "provisioning", { bootTokenExpiresAt: Date.now() - 1, + bootTokenHash: new Uint8Array([4, 5, 6]), driverInstanceId: OTHER_DRIVER_INSTANCE_ID, + sandboxIncarnation: 2, }); const response = await dispatch(bindings, llmProxyRequest("/v1/messages", { method: "POST" })); @@ -856,9 +862,9 @@ describe("driver LLM proxy route", () => { test("maps upstream failures to 502", async () => { const { bindings } = await setupFixture(); - globalThis.fetch = (async () => { + globalThis.fetch = async () => { throw new Error("boom"); - }) as typeof fetch; + }; const grant = await createLlmProxyGrant(bindings); const response = await dispatch( diff --git a/apps/api/tests/driver-log-pre-hello-buffering.test.ts b/apps/api/tests/driver-log-pre-hello-buffering.test.ts index 48e76257..9e597475 100644 --- a/apps/api/tests/driver-log-pre-hello-buffering.test.ts +++ b/apps/api/tests/driver-log-pre-hello-buffering.test.ts @@ -41,7 +41,7 @@ function createBatch(message: string): DriverLogBatchInput { timestamp: new Date(0).toISOString(), }, ], - } as DriverLogBatchInput; + }; } const activeContext = { diff --git a/apps/api/tests/driver-session-link.test.ts b/apps/api/tests/driver-session-link.test.ts index f179ae5a..c34ac01d 100644 --- a/apps/api/tests/driver-session-link.test.ts +++ b/apps/api/tests/driver-session-link.test.ts @@ -27,12 +27,15 @@ function createDriverSessionLinkDatabase(): SqliteD1Database { ); CREATE TABLE session_run ( + created_at integer DEFAULT 0 NOT NULL, created_by_account_id text, driver_instance_id text, id text PRIMARY KEY NOT NULL, + runtime_id text, session_id text NOT NULL, status text NOT NULL, - trace_id text + trace_id text, + updated_at integer DEFAULT 0 NOT NULL ); CREATE TABLE session ( @@ -40,6 +43,7 @@ function createDriverSessionLinkDatabase(): SqliteD1Database { app_id text, creator_account_id text NOT NULL, id text PRIMARY KEY NOT NULL, + runtime_id text, type text DEFAULT 'preview' NOT NULL ); @@ -101,12 +105,14 @@ describe("driver runtime session link", () => { expect( runtimeSessionLinkNeedsRefresh({ agentId: AGENT_ID, + appId: null, callerId: CALLER_FROM_ORIGIN_ID, creatorId: CREATOR_ID, executionOwnerId: EXECUTION_OWNER_FROM_ORIGIN_ID, sandboxId: SANDBOX_ID, sandboxKind: "cattle", sandboxSubjectKind: "session", + runtimeId: null, sessionId: SESSION_ID, sessionRunId: null, sessionRunStatus: null, diff --git a/apps/api/tests/driver-session-ready.test.ts b/apps/api/tests/driver-session-ready.test.ts index 37bb8ef6..22fb02c9 100644 --- a/apps/api/tests/driver-session-ready.test.ts +++ b/apps/api/tests/driver-session-ready.test.ts @@ -25,6 +25,7 @@ type ProvisionSessionDriverInput = { }; type ProvisionSessionDriverResult = { + driverGeneration: number; driverInstanceId: string; process: RuntimeProcessHandle; sandboxId: string; @@ -60,6 +61,17 @@ const SANDBOX_ID = PLATFORM_ID_FIXTURES.sandbox; const SESSION_ID = PLATFORM_ID_FIXTURES.session; const SESSION_RUN_ID = PLATFORM_ID_FIXTURES.sessionRun; const SANDBOX_SESSION_ID = SESSION_ID as unknown as SandboxSessionId; +const SANDBOX_INCARNATION = 1; + +const RUNTIME_PROVISIONING_LEASE = { + heartbeatAt: 1, + operationId: PLATFORM_ID_FIXTURES.runtimeOperation, + runId: SESSION_RUN_ID, + sandboxId: SANDBOX_ID, + sandboxIncarnation: SANDBOX_INCARNATION, + sandboxSessionId: SANDBOX_SESSION_ID, + sessionId: SESSION_ID, +} as const; const PROFILE: DriverProfileConfig = { agentId: AGENT_ID, @@ -112,10 +124,13 @@ function createDriverSessionDatabase(): SqliteD1Database { database.execute(` CREATE TABLE driver_instance ( command_seq_cursor integer DEFAULT 0 NOT NULL, + generation integer DEFAULT 0 NOT NULL, id text PRIMARY KEY NOT NULL, sandbox_id text NOT NULL, + sandbox_incarnation integer NOT NULL, sandbox_session_id text NOT NULL, status text NOT NULL, + status_operation_id text, updated_at integer NOT NULL ); @@ -123,6 +138,7 @@ function createDriverSessionDatabase(): SqliteD1Database { acked_at integer, completed_at integer, delivery_connection_id text, + driver_generation integer, driver_instance_id text NOT NULL, error_json text, expires_at integer, @@ -136,14 +152,20 @@ function createDriverSessionDatabase(): SqliteD1Database { ); CREATE TABLE sandbox ( + claim_owner text, id text PRIMARY KEY NOT NULL, inactive_deadline_at integer, + incarnation integer NOT NULL, kind text NOT NULL, + operation_kind text, + status text NOT NULL, + status_operation_id text, updated_at integer NOT NULL ); CREATE TABLE sandbox_session ( sandbox_id text NOT NULL, + sandbox_incarnation integer NOT NULL, session_id text PRIMARY KEY NOT NULL, status text NOT NULL ); @@ -157,20 +179,28 @@ function createDriverSessionDatabase(): SqliteD1Database { updated_at integer NOT NULL ); - INSERT INTO sandbox (id, inactive_deadline_at, kind, updated_at) - VALUES ('${SANDBOX_ID}', 1, 'pet', 1); + INSERT INTO sandbox (id, inactive_deadline_at, incarnation, kind, status, updated_at) + VALUES ('${SANDBOX_ID}', 1, ${SANDBOX_INCARNATION}, 'pet', 'active', 1); - INSERT INTO sandbox_session (sandbox_id, session_id, status) - VALUES ('${SANDBOX_ID}', '${SESSION_ID}', 'active'); + INSERT INTO sandbox_session (sandbox_id, sandbox_incarnation, session_id, status) + VALUES ('${SANDBOX_ID}', ${SANDBOX_INCARNATION}, '${SESSION_ID}', 'active'); INSERT INTO driver_instance ( id, sandbox_id, + sandbox_incarnation, sandbox_session_id, status, updated_at ) - VALUES ('${DRIVER_INSTANCE_ID}', '${SANDBOX_ID}', '${SESSION_ID}', 'provisioning', 1); + VALUES ( + '${DRIVER_INSTANCE_ID}', + '${SANDBOX_ID}', + ${SANDBOX_INCARNATION}, + '${SESSION_ID}', + 'provisioning', + 1 + ); INSERT INTO session_run (id, session_id, status, status_seq, updated_at) VALUES ('${SESSION_RUN_ID}', '${SESSION_ID}', 'running', 0, 1); @@ -228,12 +258,15 @@ describe("driver session readiness", () => { const bindings = createBindings(database, requests); const driver = await ensureDriverSessionReady(bindings, "https://api.test/runtime", { + builtInTools: [], cloudflareSession: {} as ExecutionSessionHandle, profile: PROFILE, resolvedMcpServers: [], resolvedSkillCatalog: [], resolvedSkills: [], + runtimeProvisioningLease: RUNTIME_PROVISIONING_LEASE, sandbox: {} as SandboxHandle, + sandboxIncarnation: SANDBOX_INCARNATION, sandboxSessionId: SESSION_ID, sessionId: SESSION_ID, sessionRunId: SESSION_RUN_ID, @@ -271,17 +304,24 @@ describe("driver session readiness", () => { INSERT INTO driver_instance ( id, sandbox_id, + sandbox_incarnation, sandbox_session_id, status, updated_at ) - VALUES (?, ?, ?, 'connecting', 2) + VALUES (?, ?, ?, ?, 'connecting', 2) `, ) - .bind(input.driverInstanceId, input.profile.sandbox.id, input.sandboxSessionId) + .bind( + input.driverInstanceId, + input.profile.sandbox.id, + SANDBOX_INCARNATION, + input.sandboxSessionId, + ) .run(); return { + driverGeneration: 0, driverInstanceId: input.driverInstanceId, process: createHangingProcess(), sandboxId: input.profile.sandbox.id, @@ -299,12 +339,15 @@ describe("driver session readiness", () => { }); const driver = await ensureDriverSessionReady(bindings, "https://api.test/runtime", { + builtInTools: [], cloudflareSession: {} as ExecutionSessionHandle, profile: PROFILE, resolvedMcpServers: [], resolvedSkillCatalog: [], resolvedSkills: [], + runtimeProvisioningLease: RUNTIME_PROVISIONING_LEASE, sandbox: {} as SandboxHandle, + sandboxIncarnation: SANDBOX_INCARNATION, sandboxSessionId: SESSION_ID, sessionId: SESSION_ID, sessionRunId: SESSION_RUN_ID, @@ -335,12 +378,15 @@ describe("driver session readiness", () => { const requests = { count: 0 }; const bindings = createBindings(database, requests); const driver = await ensureDriverSessionReady(bindings, "https://api.test/runtime", { + builtInTools: [], cloudflareSession: {} as ExecutionSessionHandle, profile: PROFILE, resolvedMcpServers: [], resolvedSkillCatalog: [], resolvedSkills: [], + runtimeProvisioningLease: RUNTIME_PROVISIONING_LEASE, sandbox: {} as SandboxHandle, + sandboxIncarnation: SANDBOX_INCARNATION, sandboxSessionId: SESSION_ID, sessionId: SESSION_ID, sessionRunId: SESSION_RUN_ID, @@ -349,6 +395,7 @@ describe("driver session readiness", () => { await dispatchDriverTurn(bindings, { attachmentIds: [], + driverGeneration: driver.driverGeneration, driverInstanceId: driver.driverInstanceId, prompt: "hello", sessionRunId: SESSION_RUN_ID, diff --git a/apps/api/tests/driver-session-stop.test.ts b/apps/api/tests/driver-session-stop.test.ts index 3071c2aa..c0964d30 100644 --- a/apps/api/tests/driver-session-stop.test.ts +++ b/apps/api/tests/driver-session-stop.test.ts @@ -1,6 +1,11 @@ import { describe, expect, test } from "bun:test"; -import { stopDriverSession } from "../src/modules/runtime/infrastructure/driver-session-stop.service"; +import type { DriverInstanceId, SessionRunId } from "@mosoo/id"; + +import { + repairClaimedDriverStopsGlobally, + stopDriverSession, +} from "../src/modules/runtime/infrastructure/driver-session-stop.service"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { SqliteD1Database } from "./helpers/sqlite-d1"; @@ -10,9 +15,15 @@ function createDriverStopDatabase(): SqliteD1Database { database.execute(` CREATE TABLE driver_instance ( id text PRIMARY KEY NOT NULL, + generation integer NOT NULL, sandbox_id text NOT NULL, sandbox_session_id text NOT NULL, status text NOT NULL, + status_changed_at integer DEFAULT 0 NOT NULL, + status_event text DEFAULT 'driver.provision' NOT NULL, + status_operation_id text, + status_seq integer DEFAULT 0 NOT NULL, + status_source text DEFAULT 'system' NOT NULL, updated_at integer NOT NULL ); @@ -120,33 +131,25 @@ function createDriverStopDatabase(): SqliteD1Database { INSERT INTO driver_instance ( id, + generation, sandbox_id, sandbox_session_id, status, updated_at ) - VALUES ('driver-1', '01J0000000000000000000000D', 'sandbox-session-1', 'failed', 1); + VALUES ('driver-1', 0, '01J0000000000000000000000D', 'sandbox-session-1', 'failed', 1); `); return database; } describe("driver session stop", () => { - test("releases linked runs and records the terminal status", async () => { + test("releases the linked lease without competing with Driver terminal events", async () => { const database = createDriverStopDatabase(); await stopDriverSession({ DB: database } as ApiBindings, { driverInstanceId: "driver-1", reason: "test.stop", - terminalRun: { - error: { - code: "agent.runtime_state_operation", - details: {}, - message: "Stopped by runtime operation.", - retryable: false, - }, - status: "cancelled", - }, }); const run = await database @@ -154,20 +157,205 @@ describe("driver session stop", () => { .bind("run-1") .first<{ error_code: string | null; status: string }>(); expect(run).toEqual({ - error_code: "agent.runtime_state_operation", - status: "cancelled", + error_code: null, + status: "running", }); const session = await database .prepare("SELECT status FROM session WHERE id = ?") .bind("session-1") .first<{ status: string }>(); - expect(session).toEqual({ status: "IDLE" }); + expect(session).toEqual({ status: "RUNNING" }); const runLink = await database .prepare("SELECT driver_instance_id FROM session_run WHERE id = ?") .bind("run-1") .first<{ driver_instance_id: string | null }>(); - expect(runLink).toEqual({ driver_instance_id: "driver-1" }); + expect(runLink).toEqual({ driver_instance_id: null }); + }); + + for (const preclaimed of [false, true]) { + test(`claims the exact Driver and Run before sending the unscoped stop command${ + preclaimed ? " from its terminal owner" : "" + }`, async () => { + const database = createDriverStopDatabase(); + database.execute(` + UPDATE driver_instance + SET status = 'ready', status_operation_id = ${preclaimed ? "'run-1'" : "NULL"} + WHERE id = 'driver-1' + `); + const observedClaims: unknown[] = []; + const bindings = { + DB: database, + DriverConnection: { + get: () => ({ + fetch: async (request: Request) => { + const path = new URL(request.url).pathname; + if (path === "/control/send") { + observedClaims.push( + await database + .prepare( + "SELECT status, status_operation_id FROM driver_instance WHERE id = 'driver-1'", + ) + .first(), + ); + return Response.json({ ok: true }); + } + if (path === "/wait/close") { + database.execute( + "UPDATE driver_instance SET status = 'stopped' WHERE id = 'driver-1'", + ); + return Response.json({ close: null, terminalized: true }); + } + throw new Error(`Unexpected Driver control request: ${path}`); + }, + }), + idFromName: () => "driver-do-id", + }, + } as unknown as ApiBindings; + + await stopDriverSession(bindings, { + driverInstanceId: "driver-1" as DriverInstanceId, + expectedDriverGeneration: 0, + expectedSessionRunId: "run-1" as SessionRunId, + reason: "test.claimed-stop", + }); + + expect(observedClaims).toEqual([ + { + status: "stopping", + status_operation_id: expect.any(String), + }, + ]); + await expect( + database + .prepare("SELECT status, status_operation_id FROM driver_instance WHERE id = 'driver-1'") + .first(), + ).resolves.toEqual({ status: "stopped", status_operation_id: null }); + await expect( + database.prepare("SELECT driver_instance_id FROM session_run WHERE id = 'run-1'").first(), + ).resolves.toEqual({ driver_instance_id: null }); + }); + } + + test("does not send a stale stop after a successor Run wins on the same Driver generation", async () => { + const database = createDriverStopDatabase(); + database.execute("UPDATE driver_instance SET status = 'ready' WHERE id = 'driver-1'"); + const originalPrepare = database.prepare.bind(database); + let injectSuccessor = true; + database.prepare = ((query: string) => { + const wrap = (statement: D1PreparedStatement): D1PreparedStatement => + new Proxy(statement, { + get(target, property, receiver) { + if (property === "bind") { + return (...values: unknown[]) => wrap(target.bind(...values)); + } + if ( + property === "all" || + property === "first" || + property === "raw" || + property === "run" + ) { + return async (...args: unknown[]) => { + if (injectSuccessor) { + injectSuccessor = false; + database.execute(` + UPDATE session_run SET status = 'completed' WHERE id = 'run-1'; + INSERT INTO session_run ( + id, driver_instance_id, session_id, agent_id, created_at, + created_by_account_id, status, trace_id, trigger, updated_at + ) VALUES ( + 'run-2', 'driver-1', 'session-1', '01J00000000000000000000009', 2, + 'account-1', 'running', 'trace-2', 'resume', 2 + ); + UPDATE session SET last_run_id = 'run-2', updated_at = 2 WHERE id = 'session-1'; + `); + } + return Reflect.apply(target[property], target, args); + }; + } + return Reflect.get(target, property, receiver); + }, + }); + + return /update\s+(?:"|`)?driver_instance(?:"|`)?/i.test(query) + ? wrap(originalPrepare(query)) + : originalPrepare(query); + }) as typeof database.prepare; + let controlRequests = 0; + const bindings = { + DB: database, + DriverConnection: { + get: () => ({ + fetch: async () => { + controlRequests += 1; + return Response.json({ ok: true }); + }, + }), + idFromName: () => "driver-do-id", + }, + } as unknown as ApiBindings; + + await expect( + stopDriverSession(bindings, { + driverInstanceId: "driver-1" as DriverInstanceId, + expectedDriverGeneration: 0, + expectedSessionRunId: "run-1" as SessionRunId, + reason: "test.stale-stop", + }), + ).rejects.toThrow("exact Driver and Session Run ownership"); + + expect(controlRequests).toBe(0); + await expect( + database + .prepare("SELECT status, status_operation_id FROM driver_instance WHERE id = 'driver-1'") + .first(), + ).resolves.toEqual({ status: "ready", status_operation_id: null }); + await expect( + database + .prepare("SELECT driver_instance_id, status FROM session_run WHERE id = 'run-2'") + .first(), + ).resolves.toEqual({ driver_instance_id: "driver-1", status: "running" }); + }); + + test("maintenance resumes a stop that crashed after its durable claim", async () => { + const database = createDriverStopDatabase(); + const operationId = "01J0000000000000000000000X"; + database.execute(` + UPDATE driver_instance + SET status = 'stopping', status_operation_id = '${operationId}' + WHERE id = 'driver-1' + `); + const requests: string[] = []; + const bindings = { + DB: database, + DriverConnection: { + get: () => ({ + fetch: async (request: Request) => { + const path = new URL(request.url).pathname; + requests.push(path); + if (path === "/control/fail") { + database.execute( + "UPDATE driver_instance SET status = 'failed' WHERE id = 'driver-1'", + ); + } + return Response.json({ ok: true }); + }, + }), + idFromName: () => "driver-do-id", + }, + } as unknown as ApiBindings; + + await repairClaimedDriverStopsGlobally(bindings); + + expect(requests).toEqual(["/control/fail", "/wait/close"]); + await expect( + database + .prepare("SELECT status_operation_id FROM driver_instance WHERE id = 'driver-1'") + .first(), + ).resolves.toEqual({ status_operation_id: null }); + await expect( + database.prepare("SELECT driver_instance_id FROM session_run WHERE id = 'run-1'").first(), + ).resolves.toEqual({ driver_instance_id: null }); }); }); diff --git a/apps/api/tests/driver-skill-package-route.test.ts b/apps/api/tests/driver-skill-package-route.test.ts index 0925503f..e7b6f9ff 100644 --- a/apps/api/tests/driver-skill-package-route.test.ts +++ b/apps/api/tests/driver-skill-package-route.test.ts @@ -1,5 +1,6 @@ import { describe, expect, test } from "bun:test"; +import { DRIVER_PROTOCOL_VERSION } from "@mosoo/agent-driver/boot"; import { driverInstancesTable, skillSnapshotsTable } from "@mosoo/db"; import { Hono } from "hono"; @@ -111,9 +112,10 @@ async function insertDriverInstance( lastHeartbeatAt: null, processId: null, protocol: "orpc-ws", - protocolVersion: 2, + protocolVersion: DRIVER_PROTOCOL_VERSION, runtime: "openai-runtime", sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxIncarnation: 1, sandboxSessionId: PUBLIC_API_TEST_IDS.ownerSession, status, statusChangedAt: nowMs, @@ -144,7 +146,7 @@ describe("driver skill package route", () => { const database = await createPublicHttpContractDatabase(); const bucket = new PublicApiMemoryFileBucket(); const bindings = createPublicHttpTestBindings(database, { - fileBucket: bucket as unknown as R2Bucket, + fileBucket: bucket, }) as ApiBindings; ensureSkillRouteTables(database); @@ -168,7 +170,7 @@ describe("driver skill package route", () => { const database = await createPublicHttpContractDatabase(); const bucket = new PublicApiMemoryFileBucket(); const bindings = createPublicHttpTestBindings(database, { - fileBucket: bucket as unknown as R2Bucket, + fileBucket: bucket, }) as ApiBindings; ensureSkillRouteTables(database); diff --git a/apps/api/tests/driver-terminal-state-coordinator.test.ts b/apps/api/tests/driver-terminal-state-coordinator.test.ts new file mode 100644 index 00000000..4c728190 --- /dev/null +++ b/apps/api/tests/driver-terminal-state-coordinator.test.ts @@ -0,0 +1,479 @@ +import { describe, expect, test } from "bun:test"; + +import { PLATFORM_ID_FIXTURES } from "@mosoo/id/testing"; + +import { finalizeDriverInstance } from "../src/modules/runtime/infrastructure/driver-instance/lifecycle"; +import type { RuntimeSessionViewCache } from "../src/modules/runtime/infrastructure/driver-instance/runtime-session-view-cache"; +import { DriverInstanceRuntimeState } from "../src/modules/runtime/infrastructure/driver-instance/runtime-state"; +import { DRIVER_INSTANCE_STATE_STORAGE_KEY } from "../src/modules/runtime/infrastructure/driver-instance/runtime-state-store"; +import type { DriverInstanceStoredState } from "../src/modules/runtime/infrastructure/driver-instance/runtime-state-store"; +import type { SessionViewerEventDeliveryBuffer } from "../src/modules/runtime/infrastructure/driver-instance/session-viewer-event-delivery-buffer"; +import type { repairFinalizedTerminalDriverRunState } from "../src/modules/runtime/infrastructure/driver-instance/terminal-run-release"; +import { DriverInstanceTerminalStateCoordinator } from "../src/modules/runtime/infrastructure/driver-instance/terminal-state-coordinator"; +import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const DRIVER_INSTANCE_ID = PLATFORM_ID_FIXTURES.driverInstance; +const TERMINAL_REPAIR_RESULT = { link: null, released: false } as const; + +class MemoryDriverInstanceStorage { + cleanupWriteFailures = 0; + readonly values = new Map(); + + async deleteAll(): Promise { + this.values.clear(); + } + + async get(key: string): Promise { + return structuredClone(this.values.get(key)) as T | undefined; + } + + async put(key: string, value: unknown): Promise { + if ( + this.cleanupWriteFailures > 0 && + typeof value === "object" && + value !== null && + "terminalCleanupComplete" in value && + value.terminalCleanupComplete === true + ) { + this.cleanupWriteFailures -= 1; + throw new Error("terminal snapshot write failed"); + } + + this.values.set(key, structuredClone(value)); + } + + storedState(): DriverInstanceStoredState { + const value = this.values.get(DRIVER_INSTANCE_STATE_STORAGE_KEY); + + if (value === undefined) { + throw new Error("Driver instance state was not persisted."); + } + + return value as DriverInstanceStoredState; + } +} + +type FinalizeDriver = typeof finalizeDriverInstance; +type RepairFinalizedRunState = typeof repairFinalizedTerminalDriverRunState; + +async function createConnectedState( + storage: MemoryDriverInstanceStorage, +): Promise { + const state = new DriverInstanceRuntimeState({ storage }); + await state.load(); + + if (state.driverInstanceId === null) { + await state.initializeDriverInstance(DRIVER_INSTANCE_ID, 1); + await state.recordAcceptedConnection({ + connectedAt: 1, + connectionId: "connection-1", + driverGeneration: 1, + traceId: null, + }); + } + + state.setRuntimeSessionLink({ + agentId: null, + appId: null, + callerId: null, + creatorId: null, + executionOwnerId: null, + sandboxId: null, + sandboxKind: null, + sandboxSubjectKind: null, + runtimeId: null, + sessionId: PLATFORM_ID_FIXTURES.session, + sessionRunId: PLATFORM_ID_FIXTURES.sessionRun, + sessionRunStatus: "running", + sessionType: null, + traceId: null, + }); + + return state; +} + +function createCoordinator( + storage: MemoryDriverInstanceStorage, + state: DriverInstanceRuntimeState, + input: { + finalizeDriver: FinalizeDriver; + flush?: () => Promise; + repairFinalizedRunState: RepairFinalizedRunState; + requestStateSync?: (sessionId: string | null) => void; + }, +): DriverInstanceTerminalStateCoordinator { + return new DriverInstanceTerminalStateCoordinator({ + clearStorage: async () => storage.deleteAll(), + env: {} as ApiBindings, + finalizeDriver: input.finalizeDriver, + repairFinalizedRunState: input.repairFinalizedRunState, + state, + viewCache: { reset: () => {} } as unknown as RuntimeSessionViewCache, + viewerEventDelivery: { + flush: input.flush ?? (async () => {}), + requestStateSync: input.requestStateSync ?? (() => {}), + resetAfterFlush: () => {}, + } as unknown as SessionViewerEventDeliveryBuffer, + withRuntimeLogContext: (fn) => fn(), + }); +} + +function createDriverInstanceDatabase(status: "failed" | "ready"): SqliteD1Database { + const database = new SqliteD1Database(); + database.execute(` + CREATE TABLE driver_instance ( + close_code integer, + close_reason text, + connection_id text, + driver_pid integer, + driver_started_at integer, + error_message text, + expires_at integer NOT NULL, + generation integer NOT NULL, + heartbeat_count integer NOT NULL, + id text PRIMARY KEY NOT NULL, + last_heartbeat_at integer, + status text NOT NULL, + status_changed_at integer NOT NULL, + status_event text NOT NULL, + status_seq integer NOT NULL, + status_source text NOT NULL, + updated_at integer NOT NULL + ); + INSERT INTO driver_instance ( + connection_id, expires_at, generation, heartbeat_count, id, status, + status_changed_at, status_event, status_seq, status_source, updated_at + ) VALUES ( + 'connection-1', 1, 1, 0, '${DRIVER_INSTANCE_ID}', '${status}', + 1, 'driver.${status}', 1, 'driver', 1 + ); + `); + return database; +} + +describe("driver terminal state coordinator", () => { + test("joins concurrent finalization callers onto one in-flight transition", async () => { + const storage = new MemoryDriverInstanceStorage(); + const state = await createConnectedState(storage); + let finalizeCalls = 0; + let markStarted: () => void = () => {}; + let releaseFinalization: () => void = () => {}; + const started = new Promise((resolve) => { + markStarted = resolve; + }); + const finalizationGate = new Promise((resolve) => { + releaseFinalization = resolve; + }); + const coordinator = createCoordinator(storage, state, { + finalizeDriver: async () => { + finalizeCalls += 1; + markStarted(); + await finalizationGate; + return "stopped"; + }, + repairFinalizedRunState: async () => TERMINAL_REPAIR_RESULT, + }); + + const epoch = state.requireConnectionEpoch(); + const first = coordinator.finalize(epoch); + const second = coordinator.finalize(epoch); + await started; + + expect(finalizeCalls).toBe(1); + releaseFinalization(); + await expect(Promise.all([first, second])).resolves.toEqual([undefined, undefined]); + expect(state.terminalCleanupComplete).toBe(true); + }); + + test("an old finalizer cannot settle a successor connection after its await", async () => { + const storage = new MemoryDriverInstanceStorage(); + const state = await createConnectedState(storage); + const oldEpoch = state.requireConnectionEpoch(); + let markStarted: () => void = () => {}; + let resume: () => void = () => {}; + const started = new Promise((resolve) => { + markStarted = resolve; + }); + const gate = new Promise((resolve) => { + resume = resolve; + }); + let repairCalls = 0; + const coordinator = createCoordinator(storage, state, { + finalizeDriver: async () => { + markStarted(); + await gate; + return "stopped"; + }, + repairFinalizedRunState: async () => { + repairCalls += 1; + return TERMINAL_REPAIR_RESULT; + }, + }); + + const oldFinalizer = coordinator.finalize(oldEpoch); + await started; + state.connectionId = "connection-2"; + state.close = null; + state.terminalized = false; + const successorWait = state.waitForReady(1, 10_000).catch((error: unknown) => error); + resume(); + await oldFinalizer; + + expect(repairCalls).toBe(0); + expect(state.connectionId).toBe("connection-2"); + expect(state.close).toBeNull(); + expect(state.terminalCleanupComplete).toBe(false); + expect(state.readyWaiters).toHaveLength(1); + state.resetAfterDestroy("test cleanup"); + await successorWait; + }); + + test("retries driver finalization in the same object after close intent is durable", async () => { + const storage = new MemoryDriverInstanceStorage(); + const state = await createConnectedState(storage); + let finalizeCalls = 0; + let repairCalls = 0; + const coordinator = createCoordinator(storage, state, { + finalizeDriver: async () => { + finalizeCalls += 1; + + if (finalizeCalls === 1) { + throw new Error("driver finalization failed"); + } + + return "stopped"; + }, + repairFinalizedRunState: async () => { + repairCalls += 1; + return TERMINAL_REPAIR_RESULT; + }, + }); + + const epoch = state.requireConnectionEpoch(); + await expect(coordinator.finalize(epoch)).rejects.toThrow("driver finalization failed"); + expect(state.close).not.toBeNull(); + expect(state.terminalized).toBe(true); + expect(state.terminalCleanupComplete).toBe(false); + + await expect(coordinator.finalize(epoch)).resolves.toBeUndefined(); + expect(finalizeCalls).toBe(2); + expect(repairCalls).toBe(1); + expect(state.terminalCleanupComplete).toBe(true); + }); + + test("finalizes without waiting for derived viewer event delivery", async () => { + const storage = new MemoryDriverInstanceStorage(); + const state = await createConnectedState(storage); + let flushCalls = 0; + let finalizeCalls = 0; + const syncSessionIds: Array = []; + const coordinator = createCoordinator(storage, state, { + finalizeDriver: async () => { + finalizeCalls += 1; + return "stopped"; + }, + flush: async () => { + flushCalls += 1; + }, + repairFinalizedRunState: async () => TERMINAL_REPAIR_RESULT, + requestStateSync: (sessionId) => syncSessionIds.push(sessionId), + }); + + await expect(coordinator.finalize(state.requireConnectionEpoch())).resolves.toBeUndefined(); + expect(flushCalls).toBe(0); + expect(finalizeCalls).toBe(1); + expect(syncSessionIds).toEqual([PLATFORM_ID_FIXTURES.session]); + expect(state.terminalCleanupComplete).toBe(true); + }); + + test("restarts repair after the terminal driver CAS was already committed", async () => { + const storage = new MemoryDriverInstanceStorage(); + const firstState = await createConnectedState(storage); + let casCommitted = false; + let casWrites = 0; + let finalizeCalls = 0; + let repairCalls = 0; + const finalizeDriver: FinalizeDriver = async () => { + finalizeCalls += 1; + + if (!casCommitted) { + casCommitted = true; + casWrites += 1; + } + + return "stopped"; + }; + const repairFinalizedRunState: RepairFinalizedRunState = async () => { + repairCalls += 1; + + if (repairCalls === 1) { + throw new Error("run repair failed"); + } + + return TERMINAL_REPAIR_RESULT; + }; + const firstCoordinator = createCoordinator(storage, firstState, { + finalizeDriver, + repairFinalizedRunState, + }); + + await expect(firstCoordinator.finalize(firstState.requireConnectionEpoch())).rejects.toThrow( + "run repair failed", + ); + expect(firstState.terminalCleanupComplete).toBe(false); + expect(casWrites).toBe(1); + + const restartedState = await createConnectedState(storage); + expect(restartedState.terminalized).toBe(true); + expect(restartedState.terminalCleanupComplete).toBe(false); + await expect( + createCoordinator(storage, restartedState, { + finalizeDriver, + repairFinalizedRunState, + }).finalize(restartedState.requireConnectionEpoch()), + ).resolves.toBeUndefined(); + + expect(casWrites).toBe(1); + expect(finalizeCalls).toBe(2); + expect(repairCalls).toBe(2); + expect(storage.storedState().terminalCleanupComplete).toBe(true); + + const completedState = await createConnectedState(storage); + await expect( + createCoordinator(storage, completedState, { + finalizeDriver: async () => { + throw new Error("completed finalization must not run again"); + }, + repairFinalizedRunState, + }).finalize(completedState.requireConnectionEpoch()), + ).resolves.toBeUndefined(); + }); + + test("restarts after terminal snapshot persistence fails without duplicating side effects", async () => { + const storage = new MemoryDriverInstanceStorage(); + const firstState = await createConnectedState(storage); + storage.cleanupWriteFailures = 1; + let finalizeCalls = 0; + let finalizeWrites = 0; + let finalized = false; + let repairCalls = 0; + let repairWrites = 0; + let repaired = false; + const finalizeDriver: FinalizeDriver = async () => { + finalizeCalls += 1; + + if (!finalized) { + finalized = true; + finalizeWrites += 1; + } + + return "stopped"; + }; + const repairFinalizedRunState: RepairFinalizedRunState = async () => { + repairCalls += 1; + + if (!repaired) { + repaired = true; + repairWrites += 1; + } + + return TERMINAL_REPAIR_RESULT; + }; + const firstCoordinator = createCoordinator(storage, firstState, { + finalizeDriver, + repairFinalizedRunState, + }); + + await expect(firstCoordinator.finalize(firstState.requireConnectionEpoch())).rejects.toThrow( + "terminal snapshot write failed", + ); + expect(firstState.terminalCleanupComplete).toBe(false); + expect(storage.storedState().terminalCleanupComplete).toBe(false); + + const restartedState = await createConnectedState(storage); + await expect( + createCoordinator(storage, restartedState, { + finalizeDriver, + repairFinalizedRunState, + }).finalize(restartedState.requireConnectionEpoch()), + ).resolves.toBeUndefined(); + + expect(finalizeCalls).toBe(2); + expect(finalizeWrites).toBe(1); + expect(repairCalls).toBe(2); + expect(repairWrites).toBe(1); + expect(restartedState.terminalCleanupComplete).toBe(true); + }); + + test("repairs the canonical failed state when maintenance wins a clean-close race", async () => { + const storage = new MemoryDriverInstanceStorage(); + const state = await createConnectedState(storage); + let repairedStatus: "failed" | "stopped" | null = null; + const coordinator = createCoordinator(storage, state, { + finalizeDriver: async () => "failed", + repairFinalizedRunState: async (_bindings, input) => { + repairedStatus = input.status; + return TERMINAL_REPAIR_RESULT; + }, + }); + + await coordinator.finalize(state.requireConnectionEpoch()); + + expect(state.close?.code).toBe(1000); + expect(repairedStatus).toBe("failed"); + expect(state.terminalCleanupComplete).toBe(true); + }); +}); + +describe("driver terminal finalization CAS", () => { + test("recognizes an exact replay without advancing the terminal transition twice", async () => { + const database = createDriverInstanceDatabase("ready"); + const bindings = { DB: database } as ApiBindings; + const input = { + connectionId: "connection-1", + generation: 1, + heartbeatCount: 2, + status: "stopped" as const, + }; + + await expect(finalizeDriverInstance(bindings, DRIVER_INSTANCE_ID, input)).resolves.toBe( + "stopped", + ); + await expect(finalizeDriverInstance(bindings, DRIVER_INSTANCE_ID, input)).resolves.toBe( + "stopped", + ); + await expect( + finalizeDriverInstance(bindings, DRIVER_INSTANCE_ID, { + ...input, + connectionId: "stale-connection", + }), + ).resolves.toBeNull(); + await expect( + database + .prepare("SELECT status, status_seq AS statusSeq FROM driver_instance WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ status: "stopped", statusSeq: 2 }); + }); + + test("returns a prior canonical failed status for the same connection generation", async () => { + const database = createDriverInstanceDatabase("failed"); + const bindings = { DB: database } as ApiBindings; + + await expect( + finalizeDriverInstance(bindings, DRIVER_INSTANCE_ID, { + connectionId: "connection-1", + generation: 1, + heartbeatCount: 2, + status: "stopped", + }), + ).resolves.toBe("failed"); + await expect( + database + .prepare("SELECT status, status_seq AS statusSeq FROM driver_instance WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ status: "failed", statusSeq: 1 }); + }); +}); diff --git a/apps/api/tests/environment-package-artifact.test.ts b/apps/api/tests/environment-package-artifact.test.ts index 7461081a..42d36530 100644 --- a/apps/api/tests/environment-package-artifact.test.ts +++ b/apps/api/tests/environment-package-artifact.test.ts @@ -17,9 +17,15 @@ import { resolveReadyEnvironmentPackageArtifact, } from "../src/modules/environments/application/environment-package-artifact.service"; import { resolveEnvironmentSetupScriptForExecution } from "../src/modules/environments/application/environment-runtime-snapshot"; -import { createEnvironmentPackageArtifactKey } from "../src/modules/environments/domain/environment-package-artifact"; -import { environmentPackageArtifactSandboxId } from "../src/modules/environments/domain/environment-package-artifact"; +import { + createEnvironmentPackageArtifactKey, + environmentPackageArtifactBuildSandboxId, +} from "../src/modules/environments/domain/environment-package-artifact"; import { exposeEnvironmentNodeModules } from "../src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-environment-artifact"; +import { + createEphemeralSandboxOptions, + getEphemeralUnversionedSandboxHandle, +} from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-platform"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { createApiCommandQueueStub, @@ -57,7 +63,9 @@ describe("Environment package artifacts", () => { { manager: "pip", packages: ["jsonschema==4.25.1", "requests==2.32.4"] }, ]); expect(key.inputDigest).toMatch(/^[0-9a-f]{64}$/u); - expect(environmentPackageArtifactSandboxId(key).length).toBeLessThanOrEqual(63); + expect(environmentPackageArtifactBuildSandboxId("01J0000000000000000000000B", 2, 3)).toBe( + "envpkg-01j0000000000000000000000b-2-3", + ); expect( parseApiCommandPayload( "environment_package_artifact_build", @@ -104,10 +112,26 @@ describe("Environment package artifacts", () => { if (!queued) { throw new Error("Expected queued artifact command."); } - await processApiCommandDeadLetterMessage( - bindings, - createRecordedQueueMessage({ body: queued.body }).message, - ); + const deadLetter = createRecordedQueueMessage({ body: queued.body }); + await processApiCommandDeadLetterMessage(bindings, deadLetter.message); + + expect(deadLetter.recorded).toEqual([{ type: "ack" }]); + await expect( + database + .app() + .select({ + lastErrorCode: apiCommandsTable.lastErrorCode, + lastErrorMessage: apiCommandsTable.lastErrorMessage, + status: apiCommandsTable.status, + }) + .from(apiCommandsTable) + .where(eq(apiCommandsTable.id, command.id)) + .get(), + ).resolves.toEqual({ + lastErrorCode: "package_install_failed", + lastErrorMessage: "Package installation failed.", + status: "dead_lettered", + }); await expect( resolveReadyEnvironmentPackageArtifact(bindings, APP_ID, JSON.stringify(packages)), @@ -179,4 +203,43 @@ describe("Environment package artifacts", () => { rmSync(root, { force: true, recursive: true }); } }); + + test("isolates attempt cleanup and lets abandoned build sandboxes expire", async () => { + const created: string[] = []; + const destroyed: string[] = []; + const bindings = { + runtimeSubjectHandleFactory: (id: string) => { + created.push(id); + return new Proxy( + {}, + { + get: (_target, property) => { + if (property === "then") { + return undefined; + } + return property === "destroy" + ? async () => { + destroyed.push(id); + } + : async () => {}; + }, + }, + ); + }, + } as ApiBindings; + const attemptAId = environmentPackageArtifactBuildSandboxId("01J0000000000000000000000B", 1, 1); + const attemptBId = environmentPackageArtifactBuildSandboxId("01J0000000000000000000000B", 1, 2); + + const attemptA = await getEphemeralUnversionedSandboxHandle(bindings, attemptAId, 15 * 60); + await getEphemeralUnversionedSandboxHandle(bindings, attemptBId, 15 * 60); + await attemptA.destroy(); + + expect(createEphemeralSandboxOptions(15 * 60)).toEqual({ + keepAlive: false, + normalizeId: true, + sleepAfter: 15 * 60, + }); + expect(created).toEqual([attemptAId, attemptBId]); + expect(destroyed).toEqual([attemptAId]); + }); }); diff --git a/apps/api/tests/external-tool-effect-migration.test.ts b/apps/api/tests/external-tool-effect-migration.test.ts new file mode 100644 index 00000000..1e0e53aa --- /dev/null +++ b/apps/api/tests/external-tool-effect-migration.test.ts @@ -0,0 +1,2327 @@ +import { describe, expect, test } from "bun:test"; + +import { + RUNTIME_COMMAND_MAX_UTF8_BYTES, + RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES, + measureRuntimeCommandJson, +} from "@mosoo/contracts/runtime-command"; +import type { DriverCommandId, DriverInstanceId, SessionRunId } from "@mosoo/id"; + +import { + acquireProdDeployLeaseStatements, + assertProdDeployLeaseOwned, + PROD_DEPLOY_LEASE_TABLE, +} from "../bin/prod-deploy-lease"; +import { + assertProtocolV3LossyMigrationInventory, + authorizeProtocolV3LegacyRewriteSql, + ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL, + installProtocolV3CutoverSql, + parseProtocolV3LossyMigrationInventory, + parseProtocolV3LegacyTerminalIntegrity, + PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + PROTOCOL_V3_CUTOVER_OBJECTS_SQL, + PROTOCOL_V3_CUTOVER_TABLE, + PROTOCOL_V3_LEGACY_TERMINAL_INTEGRITY_SQL, + PROTOCOL_V3_LOSSY_MIGRATION_INVENTORY_SQL, + REMOVE_PROTOCOL_V3_CUTOVER_SQL, + storeProtocolV3CutoverBookmarkSql, +} from "../bin/protocol-v3-cutover"; +import { + createRuntimeCommandRecord, + getRuntimeCommandRecord, + getRuntimeCommandStorageRecord, + updateRuntimeCommandRecord, +} from "../src/modules/runtime/infrastructure/session-runs/runtime-command-store.repository"; +import { + applyDrizzleMigrationAsync as applyMigration, + applyDrizzleMigrationsThrough, +} from "./helpers/drizzle-migrations"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const OLD_V2_PAYLOAD_MAX_UTF8_BYTES = 1024 * 1024; +const RELEASE_TREE_OID = "0123456789abcdef0123456789abcdef01234567"; +const INSTALL_PROTOCOL_V3_CUTOVER_SQL = installProtocolV3CutoverSql(RELEASE_TREE_OID); + +const EFFECT_UNKNOWN_ID = "01J0000000000000000000001A"; +const EFFECT_SUCCEEDED_ID = "01J0000000000000000000001B"; +const COMMAND_UNKNOWN_ID = "01J0000000000000000000001C" as DriverCommandId; +const DRIVER_ID = "01J0000000000000000000001D" as DriverInstanceId; +const SERVER_ID = "01J0000000000000000000001E"; +const RUN_ID = "01J0000000000000000000001F" as SessionRunId; +const COMMAND_SUCCEEDED_ID = "01J0000000000000000000001G" as DriverCommandId; +const SESSION_ID = "01J0000000000000000000001H"; +const AGENT_ID = "01J0000000000000000000001J"; +const ACCOUNT_ID = "01J0000000000000000000001K"; +const APP_ID = "01J0000000000000000000001M"; +const SANDBOX_ID = "01J0000000000000000000001N"; +const SANDBOX_SESSION_ID = "01J0000000000000000000001P"; +const COMMAND_ERROR_ID = "01J0000000000000000000001Q" as DriverCommandId; +const COMMAND_CANONICAL_BOUNDARY_ID = "01J0000000000000000000001R" as DriverCommandId; +const EFFECT_CANONICAL_BOUNDARY_ID = "01J0000000000000000000001S"; +const EFFECT_PROVIDER_RECEIPT_ID = "01J0000000000000000000001T"; +const COMMAND_PROVIDER_RECEIPT_ID = "01J0000000000000000000001V" as DriverCommandId; +const COMMAND_CONTROL_ID = "01J0000000000000000000001W" as DriverCommandId; +const COMMAND_PERMISSION_ID = "01J0000000000000000000001X" as DriverCommandId; +const COMMAND_GENERATION_SEVEN_ID = "01J0000000000000000000001Y" as DriverCommandId; +const COMMAND_GENERATION_EIGHT_ID = "01J0000000000000000000001Z" as DriverCommandId; +const LEGACY_MESSAGE_ID = "01J00000000000000000000020"; +const LEGACY_TERMINAL_EVENT_ID = "01J00000000000000000000021"; +const V3_REFERENCE_MESSAGE_ID = "01J00000000000000000000022"; +const DEPLOY_OWNER = "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa"; +const OTHER_DEPLOY_OWNER = "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb"; + +const utf8Encoder = new TextEncoder(); + +function jsonAtUtf8Size( + targetBytes: number, + createValue: (padding: string) => Value, +): string { + const empty = JSON.stringify(createValue("")); + const paddingBytes = targetBytes - utf8Encoder.encode(empty).byteLength; + if (paddingBytes < 0) { + throw new Error("JSON fixture base exceeds its requested byte size."); + } + + const value = JSON.stringify(createValue("x".repeat(paddingBytes))); + if (utf8Encoder.encode(value).byteLength !== targetBytes) { + throw new Error("JSON fixture did not reach its requested UTF-8 byte size."); + } + return value; +} + +function padJsonObjectToUtf8Size(value: string, targetBytes: number): string { + if (!value.endsWith("}")) throw new Error("JSON fixture must be an object."); + const paddingBytes = targetBytes - utf8Encoder.encode(value).byteLength; + if (paddingBytes < 0) throw new Error("JSON fixture base exceeds its requested byte size."); + const padded = `${value.slice(0, -1)}${" ".repeat(paddingBytes)}}`; + if (utf8Encoder.encode(padded).byteLength !== targetBytes) { + throw new Error("JSON fixture did not reach its requested UTF-8 byte size."); + } + return padded; +} + +async function readLegacyTerminalIntegrity(database: SqliteD1Database) { + const integrityRow = await database + .prepare(PROTOCOL_V3_LEGACY_TERMINAL_INTEGRITY_SQL) + .first>(); + return parseProtocolV3LegacyTerminalIntegrity( + JSON.stringify([{ results: [integrityRow], success: true }]), + ); +} + +async function readLossyMigrationInventory(database: SqliteD1Database) { + const row = await database + .prepare(PROTOCOL_V3_LOSSY_MIGRATION_INVENTORY_SQL) + .first>(); + if (row === null) throw new Error("Lossy migration inventory returned no row."); + return parseProtocolV3LossyMigrationInventory( + JSON.stringify([{ results: [row], success: true }]), + ); +} + +function installLegacyRewriteAuthorizationInfrastructure(database: SqliteD1Database): void { + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + database.execute(storeProtocolV3CutoverBookmarkSql("test-bookmark")); + for (const statement of acquireProdDeployLeaseStatements(DEPLOY_OWNER)) { + database.execute(statement); + } +} + +async function authorizeLegacyTerminalRewrite(database: SqliteD1Database): Promise { + installLegacyRewriteAuthorizationInfrastructure(database); + const integrity = await readLegacyTerminalIntegrity(database); + database.execute( + authorizeProtocolV3LegacyRewriteSql( + DEPLOY_OWNER, + integrity.noncanonicalTerminalSources, + integrity.rewriteCandidateManifestJson, + ), + ); +} + +async function createV2Database(): Promise { + const database = new SqliteD1Database(); + applyDrizzleMigrationsThrough(database, "0012_agent-task-snapshot-state"); + + return database; +} + +async function insertV2Owners(database: SqliteD1Database): Promise { + database.execute(` + INSERT INTO driver_instance ( + boot_token_expires_at, boot_token_hash, created_at, expires_at, + heartbeat_count, id, protocol, protocol_version, runtime, sandbox_id, + sandbox_session_id, status, updated_at + ) VALUES ( + 1000, x'01', 1, 1000, 0, '${DRIVER_ID}', 'rpc', 2, 'codex', + '${SANDBOX_ID}', '${SANDBOX_SESSION_ID}', 'failed', 10 + ); + + INSERT INTO session ( + agent_id, created_at, creator_account_id, id, kind, model, app_id, + provider, renamed, runtime_id, status, updated_at + ) VALUES ( + '${AGENT_ID}', 1, '${ACCOUNT_ID}', '${SESSION_ID}', 'agent', 'model', + '${APP_ID}', 'provider', 0, 'codex', 'TERMINATED', 10 + ); + + INSERT INTO session_run ( + agent_id, completed_at, created_at, created_by_account_id, + driver_instance_id, error_code, error_details_json, error_message, id, + session_id, status, trace_id, trigger, updated_at + ) VALUES ( + '${AGENT_ID}', 10, 1, '${ACCOUNT_ID}', '${DRIVER_ID}', NULL, NULL, NULL, + '${RUN_ID}', '${SESSION_ID}', 'failed', 'trace-1', 'user_prompt', 10 + ); + `); +} + +async function insertLegacySessionTerminal( + database: SqliteD1Database, + input: { + eventId: string; + eventType: "run.cancelled" | "run.completed" | "run.failed"; + runId?: string; + seq?: number; + sessionId?: string; + sourceEventId: string; + }, +): Promise { + await database + .prepare( + `INSERT INTO session_event ( + agent_id, content_text, created_at, ended_at, event_type, family, id, + occurred_at, process_status, process_type, run_id, seq, session_id, + source_event_id, source, visibility + ) VALUES (?, ?, 3, 3, ?, 'run', ?, 3, 'available', ?, ?, ?, ?, ?, 'api', 'all_consumers')`, + ) + .bind( + AGENT_ID, + input.eventType, + input.eventType, + input.eventId, + input.eventType, + input.runId ?? RUN_ID, + input.seq ?? 1, + input.sessionId ?? SESSION_ID, + input.sourceEventId, + ) + .run(); +} + +async function insertV2Command( + database: SqliteD1Database, + input: { + errorJson?: string | null; + id: DriverCommandId; + kind: string; + payloadJson: string; + resultJson?: string | null; + seq: number; + status: string; + }, +): Promise { + await database + .prepare( + `INSERT INTO driver_command ( + acked_at, completed_at, driver_instance_id, error_json, id, issued_at, + kind, payload_json, result_json, seq, status + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + 2, + 3, + DRIVER_ID, + input.errorJson ?? null, + input.id, + 1, + input.kind, + input.payloadJson, + input.resultJson ?? null, + input.seq, + input.status, + ) + .run(); +} + +async function insertV2Effect( + database: SqliteD1Database, + input: { + commandId: DriverCommandId; + effectId: string; + providerReceiptJson?: string | null; + resultJson: string | null; + status: "executing" | "succeeded"; + toolName: string; + }, +): Promise { + await database + .prepare( + `INSERT INTO external_tool_effect ( + attempt_count, command_id, created_at, driver_instance_id, id, + idempotency_key, provider_receipt_json, result_json, server_id, + session_run_id, status, tool_name, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + 1, + input.commandId, + 1, + DRIVER_ID, + input.effectId, + input.effectId, + input.providerReceiptJson === undefined + ? input.status === "succeeded" + ? '{"receipt":"legacy"}' + : null + : input.providerReceiptJson, + input.resultJson, + SERVER_ID, + RUN_ID, + input.status, + input.toolName, + 2, + ) + .run(); +} + +async function readMigratedEffectState( + database: SqliteD1Database, + commandId: DriverCommandId, +): Promise> { + const row = await database + .prepare("SELECT id, result_json, status FROM external_tool_effect WHERE command_id = ?") + .bind(commandId) + .first<{ id: string; result_json: string | null; status: string }>(); + if (row === null) { + throw new Error("Migrated external tool effect was not found."); + } + + return row.status === "succeeded" + ? { effectId: row.id, kind: row.status, result: JSON.parse(row.result_json ?? "null") } + : { effectId: row.id, kind: row.status }; +} + +async function insertV2Attempt( + database: SqliteD1Database, + input: { + effectId: string; + providerReceiptJson?: string | null; + resultJson: string | null; + status: "executing" | "succeeded"; + }, +): Promise { + await database + .prepare( + `INSERT INTO external_tool_effect_attempt ( + attempt, completed_at, created_at, effect_id, provider_receipt_json, + result_json, status + ) VALUES (?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + 1, + input.status === "succeeded" ? 3 : null, + 2, + input.effectId, + input.providerReceiptJson === undefined + ? input.status === "succeeded" + ? '{"receipt":"legacy"}' + : null + : input.providerReceiptJson, + input.resultJson, + input.status, + ) + .run(); +} + +async function expectV2MigrationRejection( + seed: (database: SqliteD1Database) => Promise, +): Promise { + const database = await createV2Database(); + await insertV2Owners(database); + await seed(database); + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + return database; +} + +describe("durable MCP effect v3 migration", () => { + test("rejects every lossy v2 rewrite before changing schema or history", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + + const unknownPayloadJson = jsonAtUtf8Size(OLD_V2_PAYLOAD_MAX_UTF8_BYTES, (argumentsJson) => ({ + argumentsJson, + commandId: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + requestId: "legacy-unknown-request", + serverId: SERVER_ID, + toolCallId: "legacy-unknown-tool-call", + toolName: "createIssue", + })); + await insertV2Command(database, { + id: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + payloadJson: unknownPayloadJson, + seq: 1, + status: "failed", + }); + await insertV2Effect(database, { + commandId: COMMAND_UNKNOWN_ID, + effectId: EFFECT_UNKNOWN_ID, + resultJson: null, + status: "executing", + toolName: "createIssue", + }); + await insertV2Attempt(database, { + effectId: EFFECT_UNKNOWN_ID, + resultJson: null, + status: "executing", + }); + + const succeededPayloadJson = jsonAtUtf8Size( + RUNTIME_COMMAND_MAX_UTF8_BYTES + 1, + (argumentsJson) => ({ + argumentsJson, + commandId: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + requestId: "legacy-succeeded-request", + serverId: SERVER_ID, + toolCallId: "legacy-succeeded-tool-call", + toolName: "createIssue", + }), + ); + const historicalToolName = "createIssue"; + const createSucceededResult = (targetBytes: number) => + jsonAtUtf8Size(targetBytes, (outputText) => ({ + outputText, + requestId: "legacy-succeeded-request", + serverId: SERVER_ID, + toolName: historicalToolName, + })); + const effectResultJson = createSucceededResult(OLD_V2_PAYLOAD_MAX_UTF8_BYTES); + const attemptResultJson = createSucceededResult(OLD_V2_PAYLOAD_MAX_UTF8_BYTES - 1); + const commandResultJson = createSucceededResult(OLD_V2_PAYLOAD_MAX_UTF8_BYTES - 2); + await insertV2Command(database, { + id: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + payloadJson: succeededPayloadJson, + resultJson: commandResultJson, + seq: 2, + status: "completed", + }); + await insertV2Effect(database, { + commandId: COMMAND_SUCCEEDED_ID, + effectId: EFFECT_SUCCEEDED_ID, + resultJson: effectResultJson, + status: "succeeded", + toolName: historicalToolName, + }); + await insertV2Attempt(database, { + effectId: EFFECT_SUCCEEDED_ID, + resultJson: attemptResultJson, + status: "succeeded", + }); + + const canonicalBoundaryPayloadJson = jsonAtUtf8Size( + RUNTIME_COMMAND_MAX_UTF8_BYTES + 1, + (argumentsJson) => ({ + argumentsJson, + commandId: COMMAND_CANONICAL_BOUNDARY_ID, + kind: "mcp.execute", + requestId: "legacy-boundary-request", + serverId: SERVER_ID, + toolCallId: "legacy-boundary-tool-call", + toolName: "createIssue", + }), + ); + const canonicalBoundaryResultJson = jsonAtUtf8Size( + RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES, + (outputText) => ({ + outputText, + requestId: "legacy-boundary-request", + serverId: SERVER_ID, + toolName: "createIssue", + }), + ); + await insertV2Command(database, { + errorJson: JSON.stringify({ + code: "legacy.finalization_race", + details: {}, + message: "The old finalizer raced the durable effect.", + retryable: false, + }), + id: COMMAND_CANONICAL_BOUNDARY_ID, + kind: "mcp.execute", + payloadJson: canonicalBoundaryPayloadJson, + resultJson: canonicalBoundaryResultJson, + seq: 3, + status: "failed", + }); + await insertV2Effect(database, { + commandId: COMMAND_CANONICAL_BOUNDARY_ID, + effectId: EFFECT_CANONICAL_BOUNDARY_ID, + resultJson: canonicalBoundaryResultJson, + status: "succeeded", + toolName: "createIssue", + }); + await insertV2Attempt(database, { + effectId: EFFECT_CANONICAL_BOUNDARY_ID, + resultJson: canonicalBoundaryResultJson, + status: "succeeded", + }); + + const providerReceiptResult = JSON.stringify({ + outputText: "provider response", + requestId: "legacy-provider-request", + serverId: SERVER_ID, + toolName: "createIssue", + }); + const oversizedProviderReceipt = "x".repeat(RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES); + await insertV2Command(database, { + id: COMMAND_PROVIDER_RECEIPT_ID, + kind: "mcp.execute", + payloadJson: JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_PROVIDER_RECEIPT_ID, + kind: "mcp.execute", + requestId: "legacy-provider-request", + serverId: SERVER_ID, + toolCallId: "legacy-provider-tool-call", + toolName: "createIssue", + }), + resultJson: providerReceiptResult, + seq: 4, + status: "completed", + }); + await insertV2Effect(database, { + commandId: COMMAND_PROVIDER_RECEIPT_ID, + effectId: EFFECT_PROVIDER_RECEIPT_ID, + providerReceiptJson: oversizedProviderReceipt, + resultJson: providerReceiptResult, + status: "succeeded", + toolName: "createIssue", + }); + await insertV2Attempt(database, { + effectId: EFFECT_PROVIDER_RECEIPT_ID, + providerReceiptJson: oversizedProviderReceipt, + resultJson: providerReceiptResult, + status: "succeeded", + }); + + const commandErrorJson = jsonAtUtf8Size(OLD_V2_PAYLOAD_MAX_UTF8_BYTES, (context) => ({ + code: "legacy.command_error", + details: { context }, + message: "Legacy command failed.", + retryable: false, + })); + await insertV2Command(database, { + errorJson: commandErrorJson, + id: COMMAND_ERROR_ID, + kind: "input.start", + payloadJson: JSON.stringify({ + commandId: COMMAND_ERROR_ID, + input: { text: "hello" }, + kind: "input.start", + requestId: "legacy-input-request", + runId: RUN_ID, + }), + seq: 5, + status: "failed", + }); + await insertV2Command(database, { + id: COMMAND_CONTROL_ID, + kind: "turn.cancel", + payloadJson: JSON.stringify({ commandId: COMMAND_CONTROL_ID, kind: "turn.cancel" }), + resultJson: "null", + seq: 6, + status: "completed", + }); + await insertV2Command(database, { + id: COMMAND_PERMISSION_ID, + kind: "permission.resolve", + payloadJson: JSON.stringify({ + commandId: COMMAND_PERMISSION_ID, + decision: "allow_once", + kind: "permission.resolve", + requestId: "legacy-permission-request", + }), + seq: 7, + status: "completed", + }); + + const sessionErrorJson = jsonAtUtf8Size(OLD_V2_PAYLOAD_MAX_UTF8_BYTES, (context) => ({ + code: "legacy.session_error", + details: { context }, + message: "Legacy Session Run failed.", + retryable: false, + })); + const sessionError = JSON.parse(sessionErrorJson) as { + code: string; + details: Record; + message: string; + }; + await database + .prepare( + "UPDATE session_run SET error_code = ?, error_details_json = ?, error_message = ? WHERE id = ?", + ) + .bind(sessionError.code, JSON.stringify(sessionError.details), sessionError.message, RUN_ID) + .run(); + + // The same Driver instance may have been reused after these terminal commands. + // Its current generation is not evidence of the historical command generation. + await database + .prepare("UPDATE driver_instance SET generation = ? WHERE id = ?") + .bind(7, DRIVER_ID) + .run(); + + const inventory = await readLossyMigrationInventory(database); + expect(inventory).toEqual({ + attemptCompletionTimeFabrications: 0, + candidateIds: [ + { category: "command_error_omission", id: COMMAND_ERROR_ID }, + { category: "mcp_argument_omission", id: COMMAND_UNKNOWN_ID }, + { category: "mcp_argument_omission", id: COMMAND_SUCCEEDED_ID }, + { category: "mcp_argument_omission", id: COMMAND_CANONICAL_BOUNDARY_ID }, + { category: "mcp_command_terminal_conflict", id: EFFECT_CANONICAL_BOUNDARY_ID }, + { category: "mcp_result_omission", id: EFFECT_SUCCEEDED_ID }, + { category: "mcp_result_omission", id: EFFECT_CANONICAL_BOUNDARY_ID }, + { category: "provider_receipt_loss", id: EFFECT_PROVIDER_RECEIPT_ID }, + { category: "session_run_error_omission", id: RUN_ID }, + ], + commandErrorOmissions: 1, + commandPayloadConflicts: 0, + controlReasonOmissions: 0, + inputStartResultOmissions: 0, + inputTextOmissions: 0, + mcpArgumentOmissions: 3, + mcpCommandTerminalConflicts: 1, + mcpResultConflicts: 0, + mcpResultOmissions: 2, + orphanEffects: 0, + providerReceiptLosses: 1, + permissionPayloadRewrites: 0, + sessionRunErrorOmissions: 1, + totalCandidates: 9, + }); + expect(() => assertProtocolV3LossyMigrationInventory(inventory)).toThrow( + "No lossy migration candidates are authorized", + ); + + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + + await expect( + database + .prepare("SELECT payload_json FROM driver_command WHERE id = ?") + .bind(COMMAND_UNKNOWN_ID) + .first(), + ).resolves.toEqual({ payload_json: unknownPayloadJson }); + await expect( + database + .prepare("SELECT error_json FROM driver_command WHERE id = ?") + .bind(COMMAND_ERROR_ID) + .first(), + ).resolves.toEqual({ error_json: commandErrorJson }); + await expect( + database + .prepare("SELECT provider_receipt_json FROM external_tool_effect WHERE id = ?") + .bind(EFFECT_PROVIDER_RECEIPT_ID) + .first(), + ).resolves.toEqual({ provider_receipt_json: oversizedProviderReceipt }); + await expect( + database + .prepare( + "SELECT error_code, error_details_json, error_message FROM session_run WHERE id = ?", + ) + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ + error_code: sessionError.code, + error_details_json: JSON.stringify(sessionError.details), + error_message: sessionError.message, + }); + await expect( + database + .prepare( + "SELECT count(*) AS count FROM pragma_table_info('driver_command') WHERE name = 'driver_generation'", + ) + .first(), + ).resolves.toEqual({ count: 0 }); + }); + + test("rejects a real v2 database with an old-limit nonterminal command", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + const payloadJson = jsonAtUtf8Size(OLD_V2_PAYLOAD_MAX_UTF8_BYTES, (text) => ({ + commandId: COMMAND_ERROR_ID, + input: { text }, + kind: "input.start", + requestId: "legacy-input-request", + runId: RUN_ID, + })); + await insertV2Command(database, { + id: COMMAND_ERROR_ID, + kind: "input.start", + payloadJson, + seq: 1, + status: "accepted", + }); + await expect( + database + .prepare("SELECT COUNT(*) AS count FROM driver_command WHERE status = 'accepted'") + .first(), + ).resolves.toEqual({ count: 1 }); + await expect(database.prepare("PRAGMA ignore_check_constraints").first()).resolves.toEqual({ + ignore_check_constraints: 0, + }); + + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + await expect( + database + .prepare("SELECT payload_json, status FROM driver_command WHERE id = ?") + .bind(COMMAND_ERROR_ID) + .first(), + ).resolves.toEqual({ payload_json: payloadJson, status: "accepted" }); + }); + + test("rejects an orphan effect before its inner-join rebuild can delete it", async () => { + const database = new SqliteD1Database({ foreignKeys: false }); + applyDrizzleMigrationsThrough(database, "0012_agent-task-snapshot-state"); + await database + .prepare( + `INSERT INTO external_tool_effect ( + attempt_count, command_id, created_at, driver_instance_id, id, + idempotency_key, provider_receipt_json, result_json, server_id, + session_run_id, status, tool_name, updated_at + ) VALUES (0, ?, 1, ?, ?, ?, NULL, NULL, ?, ?, 'intent', 'createIssue', 2)`, + ) + .bind(COMMAND_UNKNOWN_ID, DRIVER_ID, EFFECT_UNKNOWN_ID, EFFECT_UNKNOWN_ID, SERVER_ID, RUN_ID) + .run(); + database.execute("PRAGMA foreign_keys = ON"); + await expect(database.prepare("PRAGMA foreign_keys").first()).resolves.toEqual({ + foreign_keys: 1, + }); + + const inventory = await readLossyMigrationInventory(database); + expect(inventory.orphanEffects).toBe(1); + expect(inventory.candidateIds).toEqual([{ category: "orphan_effect", id: EFFECT_UNKNOWN_ID }]); + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + await expect( + database + .prepare("SELECT command_id, status FROM external_tool_effect WHERE id = ?") + .bind(EFFECT_UNKNOWN_ID) + .first(), + ).resolves.toEqual({ command_id: COMMAND_UNKNOWN_ID, status: "intent" }); + }); + + test("rejects missing terminal attempt times instead of fabricating migration time", async () => { + for (const status of ["succeeded", "unknown"] as const) { + const database = await createV2Database(); + await insertV2Owners(database); + const payloadJson = JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + }); + const resultJson = + status === "succeeded" + ? JSON.stringify({ + outputText: "done", + requestId: "legacy-request", + serverId: SERVER_ID, + toolName: "createIssue", + }) + : null; + await insertV2Command(database, { + id: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + payloadJson, + resultJson, + seq: 1, + status: status === "succeeded" ? "completed" : "failed", + }); + await insertV2Effect(database, { + commandId: COMMAND_SUCCEEDED_ID, + effectId: EFFECT_SUCCEEDED_ID, + providerReceiptJson: status === "succeeded" ? '{"receipt":"legacy"}' : null, + resultJson, + status: status === "succeeded" ? "succeeded" : "executing", + toolName: "createIssue", + }); + await insertV2Attempt(database, { + effectId: EFFECT_SUCCEEDED_ID, + providerReceiptJson: status === "succeeded" ? '{"receipt":"legacy"}' : null, + resultJson, + status: status === "succeeded" ? "succeeded" : "executing", + }); + await database + .prepare("UPDATE external_tool_effect SET status = ? WHERE id = ?") + .bind(status, EFFECT_SUCCEEDED_ID) + .run(); + await database + .prepare( + "UPDATE external_tool_effect_attempt SET completed_at = NULL, status = ? WHERE effect_id = ?", + ) + .bind(status, EFFECT_SUCCEEDED_ID) + .run(); + + const inventory = await readLossyMigrationInventory(database); + expect(inventory.attemptCompletionTimeFabrications).toBe(1); + expect(inventory.candidateIds).toEqual([ + { category: "attempt_completion_time_fabrication", id: EFFECT_SUCCEEDED_ID }, + ]); + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + await expect( + database + .prepare( + "SELECT completed_at, status FROM external_tool_effect_attempt WHERE effect_id = ?", + ) + .bind(EFFECT_SUCCEEDED_ID) + .first(), + ).resolves.toEqual({ completed_at: null, status }); + } + }); + + test("keeps the executing-to-unknown transition time monotonic", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + await insertV2Command(database, { + id: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + payloadJson: JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + }), + seq: 1, + status: "failed", + }); + await insertV2Effect(database, { + commandId: COMMAND_UNKNOWN_ID, + effectId: EFFECT_UNKNOWN_ID, + providerReceiptJson: null, + resultJson: null, + status: "executing", + toolName: "createIssue", + }); + await insertV2Attempt(database, { + effectId: EFFECT_UNKNOWN_ID, + providerReceiptJson: null, + resultJson: null, + status: "executing", + }); + const futureTimestamp = 9_007_199_254_740_000; + await database + .prepare("UPDATE external_tool_effect SET updated_at = ? WHERE id = ?") + .bind(futureTimestamp, EFFECT_UNKNOWN_ID) + .run(); + await database + .prepare("UPDATE external_tool_effect_attempt SET created_at = ? WHERE effect_id = ?") + .bind(futureTimestamp, EFFECT_UNKNOWN_ID) + .run(); + const inventory = await readLossyMigrationInventory(database); + expect(inventory.totalCandidates).toBe(0); + + await applyMigration(database, "0013_durable-mcp-effect-v3"); + const transition = await database + .prepare( + `SELECT + "attempt"."completed_at", + "attempt"."status" AS "attempt_status", + "effect"."status" AS "effect_status", + "effect"."updated_at" + FROM "external_tool_effect_attempt" AS "attempt" + INNER JOIN "external_tool_effect" AS "effect" ON "effect"."id" = "attempt"."effect_id" + WHERE "attempt"."effect_id" = ?`, + ) + .bind(EFFECT_UNKNOWN_ID) + .first(); + expect(transition).toEqual({ + attempt_status: "unknown", + completed_at: futureTimestamp, + effect_status: "unknown", + updated_at: futureTimestamp, + }); + }); + + test("keeps legacy terminal history separate from two new Driver generations", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + await insertV2Command(database, { + id: COMMAND_CONTROL_ID, + kind: "turn.cancel", + payloadJson: JSON.stringify({ commandId: COMMAND_CONTROL_ID, kind: "turn.cancel" }), + seq: 1, + status: "completed", + }); + database.execute(` + UPDATE session + SET runtime_event_seq_cursor = 1 + WHERE id = '${SESSION_ID}'; + UPDATE session_run + SET error_code = 'legacy.failed', + error_details_json = '{}', + error_message = 'Legacy failure.', + status_event = 'run.fail' + WHERE id = '${RUN_ID}'; + `); + await insertLegacySessionTerminal(database, { + eventId: LEGACY_TERMINAL_EVENT_ID, + eventType: "run.failed", + sourceEventId: `session-run-terminal:${RUN_ID}:run.failed`, + }); + const inventory = await readLossyMigrationInventory(database); + expect(inventory.totalCandidates).toBe(0); + expect(() => assertProtocolV3LossyMigrationInventory(inventory)).not.toThrow(); + await applyMigration(database, "0013_durable-mcp-effect-v3"); + await applyMigration(database, "0014_session-event-stream-identity"); + await database + .prepare( + "UPDATE driver_instance SET command_seq_cursor = 10, connection_id = ?, generation = 7, status = 'ready' WHERE id = ?", + ) + .bind("generation-7", DRIVER_ID) + .run(); + await database + .prepare("UPDATE session_run SET completed_at = NULL, status = 'running' WHERE id = ?") + .bind(RUN_ID) + .run(); + + await createRuntimeCommandRecord(database, { + command: { + commandId: COMMAND_GENERATION_SEVEN_ID, + kind: "turn.cancel", + reason: "generation seven", + runId: RUN_ID, + }, + driverGeneration: 7, + driverInstanceId: DRIVER_ID, + status: "accepted", + }); + await updateRuntimeCommandRecord(database, { + commandId: COMMAND_GENERATION_SEVEN_ID, + driverGeneration: 7, + driverInstanceId: DRIVER_ID, + status: "cancelled", + }); + await database + .prepare("UPDATE driver_instance SET connection_id = ?, generation = 8 WHERE id = ?") + .bind("generation-8", DRIVER_ID) + .run(); + await createRuntimeCommandRecord(database, { + command: { + commandId: COMMAND_GENERATION_EIGHT_ID, + input: { text: "generation eight" }, + kind: "input.start", + requestId: "generation-8-request", + runId: RUN_ID, + }, + driverGeneration: 8, + driverInstanceId: DRIVER_ID, + }); + + await expect( + getRuntimeCommandStorageRecord(database, DRIVER_ID, COMMAND_CONTROL_ID), + ).resolves.toMatchObject({ driverGeneration: null, format: "legacy-v2-terminal" }); + await expect( + getRuntimeCommandRecord(database, DRIVER_ID, 7, COMMAND_GENERATION_SEVEN_ID), + ).resolves.toMatchObject({ status: "cancelled" }); + await expect( + getRuntimeCommandRecord(database, DRIVER_ID, 8, COMMAND_GENERATION_SEVEN_ID), + ).resolves.toBeNull(); + await expect( + getRuntimeCommandRecord(database, DRIVER_ID, 8, COMMAND_GENERATION_EIGHT_ID), + ).resolves.toMatchObject({ status: "queued" }); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_GENERATION_SEVEN_ID, + driverGeneration: 8, + driverInstanceId: DRIVER_ID, + status: "failed", + }), + ).resolves.toMatchObject({ kind: "rejected", reason: "command_not_found" }); + await expect( + database + .prepare( + "SELECT driver_generation, status FROM driver_command ORDER BY driver_generation, id", + ) + .all(), + ).resolves.toMatchObject({ + results: [ + { driver_generation: null, status: "completed" }, + { driver_generation: 7, status: "cancelled" }, + { driver_generation: 8, status: "queued" }, + ], + }); + }); + + test("preserves bounded MCP identities without inventing a 256-byte contract", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + const toolName = "t".repeat(300); + const payloadJson = JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName, + }); + const result = { + outputText: "done", + requestId: "legacy-request", + serverId: SERVER_ID, + toolName, + }; + const resultJson = JSON.stringify(result); + await insertV2Command(database, { + id: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + payloadJson, + resultJson, + seq: 1, + status: "completed", + }); + await insertV2Effect(database, { + commandId: COMMAND_SUCCEEDED_ID, + effectId: EFFECT_SUCCEEDED_ID, + resultJson, + status: "succeeded", + toolName, + }); + await insertV2Attempt(database, { + effectId: EFFECT_SUCCEEDED_ID, + resultJson, + status: "succeeded", + }); + + await applyMigration(database, "0013_durable-mcp-effect-v3"); + + await expect(readMigratedEffectState(database, COMMAND_SUCCEEDED_ID)).resolves.toEqual({ + effectId: EFFECT_SUCCEEDED_ID, + kind: "succeeded", + result, + }); + const command = await getRuntimeCommandStorageRecord(database, DRIVER_ID, COMMAND_SUCCEEDED_ID); + expect(command?.record).toMatchObject({ payload: { toolName }, result: { toolName } }); + expect(measureRuntimeCommandJson(command?.record.payload)).toBeLessThanOrEqual( + RUNTIME_COMMAND_MAX_UTF8_BYTES, + ); + }); + + test("rejects an identity-only oversized MCP record instead of replacing its identity", async () => { + const payloadJson = jsonAtUtf8Size(RUNTIME_COMMAND_MAX_UTF8_BYTES + 1, (requestId) => ({ + argumentsJson: "{}", + commandId: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + requestId, + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + })); + const database = await expectV2MigrationRejection(async (fixture) => { + await insertV2Command(fixture, { + id: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + payloadJson, + seq: 1, + status: "failed", + }); + await insertV2Effect(fixture, { + commandId: COMMAND_UNKNOWN_ID, + effectId: EFFECT_UNKNOWN_ID, + resultJson: null, + status: "executing", + toolName: "createIssue", + }); + await insertV2Attempt(fixture, { + effectId: EFFECT_UNKNOWN_ID, + resultJson: null, + status: "executing", + }); + }); + + await expect( + database + .prepare("SELECT payload_json FROM driver_command WHERE id = ?") + .bind(COMMAND_UNKNOWN_ID) + .first(), + ).resolves.toEqual({ payload_json: payloadJson }); + }); + + test("rejects oversized input text without changing attachments or Run identity", async () => { + const payloadJson = jsonAtUtf8Size(RUNTIME_COMMAND_MAX_UTF8_BYTES + 1, (text) => ({ + commandId: COMMAND_ERROR_ID, + input: { attachmentIds: ["file-1", "file-2"], text }, + kind: "input.start", + requestId: "legacy-input-request", + runId: RUN_ID, + })); + const database = await createV2Database(); + await insertV2Owners(database); + await insertV2Command(database, { + id: COMMAND_ERROR_ID, + kind: "input.start", + payloadJson, + seq: 1, + status: "failed", + }); + + const inventory = await readLossyMigrationInventory(database); + expect(inventory.inputTextOmissions).toBe(1); + expect(inventory.candidateIds).toEqual([ + { category: "input_text_omission", id: COMMAND_ERROR_ID }, + ]); + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + + await expect( + database + .prepare("SELECT payload_json FROM driver_command WHERE id = ?") + .bind(COMMAND_ERROR_ID) + .first(), + ).resolves.toEqual({ payload_json: payloadJson }); + }); + + test("rejects an oversized input result even when its command payload is bounded", async () => { + const resultJson = jsonAtUtf8Size( + RUNTIME_COMMAND_TERMINAL_PAYLOAD_MAX_UTF8_BYTES + 1, + (requestId) => ({ requestId }), + ); + const database = await createV2Database(); + await insertV2Owners(database); + await insertV2Command(database, { + id: COMMAND_ERROR_ID, + kind: "input.start", + payloadJson: JSON.stringify({ + commandId: COMMAND_ERROR_ID, + input: { text: "hello" }, + kind: "input.start", + requestId: "legacy-input-request", + runId: RUN_ID, + }), + resultJson, + seq: 1, + status: "completed", + }); + + const inventory = await readLossyMigrationInventory(database); + expect(inventory.inputStartResultOmissions).toBe(1); + expect(inventory.candidateIds).toEqual([ + { category: "input_start_result_omission", id: COMMAND_ERROR_ID }, + ]); + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + await expect( + database + .prepare("SELECT result_json FROM driver_command WHERE id = ?") + .bind(COMMAND_ERROR_ID) + .first(), + ).resolves.toEqual({ result_json: resultJson }); + }); + + test("rejects oversized control reasons instead of replacing them", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + const fixtures = [ + { id: COMMAND_CONTROL_ID, kind: "turn.cancel" }, + { id: COMMAND_PERMISSION_ID, kind: "session.stop" }, + ] as const; + for (const [index, fixture] of fixtures.entries()) { + await insertV2Command(database, { + id: fixture.id, + kind: fixture.kind, + payloadJson: jsonAtUtf8Size(RUNTIME_COMMAND_MAX_UTF8_BYTES + 1, (reason) => ({ + commandId: fixture.id, + kind: fixture.kind, + reason, + })), + seq: index + 1, + status: "failed", + }); + } + + const inventory = await readLossyMigrationInventory(database); + expect(inventory.controlReasonOmissions).toBe(2); + expect(inventory.candidateIds).toEqual([ + { category: "control_reason_omission", id: COMMAND_CONTROL_ID }, + { category: "control_reason_omission", id: COMMAND_PERMISSION_ID }, + ]); + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + }); + + test("rejects duplicate command payload keys before SQLite can choose a different copy", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + const inputPayloadJson = `{"commandId":"${COMMAND_ERROR_ID}","input":{"text":"first","text":"last"},"kind":"input.start","requestId":"legacy-input-request","runId":"${RUN_ID}"}`; + const permissionPayloadJson = padJsonObjectToUtf8Size( + `{"commandId":"${COMMAND_PERMISSION_ID}","decision":"allow_once","decision":"reject_once","kind":"permission.resolve","requestId":"legacy-permission-request"}`, + RUNTIME_COMMAND_MAX_UTF8_BYTES + 1, + ); + expect(JSON.parse(inputPayloadJson)).toMatchObject({ input: { text: "last" } }); + expect(JSON.parse(permissionPayloadJson)).toMatchObject({ decision: "reject_once" }); + + await insertV2Command(database, { + id: COMMAND_ERROR_ID, + kind: "input.start", + payloadJson: inputPayloadJson, + seq: 1, + status: "failed", + }); + await insertV2Command(database, { + id: COMMAND_PERMISSION_ID, + kind: "permission.resolve", + payloadJson: permissionPayloadJson, + seq: 2, + status: "completed", + }); + + const inventory = await readLossyMigrationInventory(database); + expect(inventory.commandPayloadConflicts).toBe(2); + expect(inventory.permissionPayloadRewrites).toBe(1); + expect(inventory.candidateIds).toEqual([ + { category: "command_payload_conflict", id: COMMAND_ERROR_ID }, + { category: "command_payload_conflict", id: COMMAND_PERMISSION_ID }, + { category: "permission_payload_rewrite", id: COMMAND_PERMISSION_ID }, + ]); + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + await expect( + database + .prepare("SELECT payload_json FROM driver_command WHERE id = ?") + .bind(COMMAND_ERROR_ID) + .first(), + ).resolves.toEqual({ payload_json: inputPayloadJson }); + await expect( + database + .prepare("SELECT payload_json FROM driver_command WHERE id = ?") + .bind(COMMAND_PERMISSION_ID) + .first(), + ).resolves.toEqual({ payload_json: permissionPayloadJson }); + }); + + test("rejects duplicate MCP result keys before SQLite can erase JS-visible conflicts", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + const payloadJson = JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + }); + const result = (lastOutput: string) => + `{"outputText":"shared-first","outputText":"${lastOutput}","requestId":"legacy-request","serverId":"${SERVER_ID}","toolName":"createIssue"}`; + const commandResultJson = result("command-last"); + const effectResultJson = result("effect-last"); + const attemptResultJson = result("attempt-last"); + + await insertV2Command(database, { + id: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + payloadJson, + resultJson: commandResultJson, + seq: 1, + status: "completed", + }); + await insertV2Effect(database, { + commandId: COMMAND_SUCCEEDED_ID, + effectId: EFFECT_SUCCEEDED_ID, + resultJson: effectResultJson, + status: "succeeded", + toolName: "createIssue", + }); + await insertV2Attempt(database, { + effectId: EFFECT_SUCCEEDED_ID, + resultJson: attemptResultJson, + status: "succeeded", + }); + + const inventory = await readLossyMigrationInventory(database); + expect(inventory.mcpResultConflicts).toBe(1); + expect(inventory.candidateIds).toEqual([ + { category: "mcp_result_conflict", id: EFFECT_SUCCEEDED_ID }, + ]); + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + expect(JSON.parse(commandResultJson)).toMatchObject({ outputText: "command-last" }); + expect(JSON.parse(effectResultJson)).toMatchObject({ outputText: "effect-last" }); + expect(JSON.parse(attemptResultJson)).toMatchObject({ outputText: "attempt-last" }); + await expect( + database + .prepare("SELECT result_json FROM driver_command WHERE id = ?") + .bind(COMMAND_SUCCEEDED_ID) + .first(), + ).resolves.toEqual({ result_json: commandResultJson }); + await expect( + database + .prepare("SELECT result_json FROM external_tool_effect WHERE id = ?") + .bind(EFFECT_SUCCEEDED_ID) + .first(), + ).resolves.toEqual({ result_json: effectResultJson }); + await expect( + database + .prepare("SELECT result_json FROM external_tool_effect_attempt WHERE effect_id = ?") + .bind(EFFECT_SUCCEEDED_ID) + .first(), + ).resolves.toEqual({ result_json: attemptResultJson }); + }); + + test("preserves the authoritative MCP result without a SQLite JSON round trip", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + const payloadJson = JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + }); + const resultJson = JSON.stringify({ + outputText: "\ud800", + requestId: "legacy-request", + serverId: SERVER_ID, + toolName: "createIssue", + }); + expect((JSON.parse(resultJson) as { outputText: string }).outputText.charCodeAt(0)).toBe( + 0xd800, + ); + await insertV2Command(database, { + id: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + payloadJson, + resultJson, + seq: 1, + status: "completed", + }); + await insertV2Effect(database, { + commandId: COMMAND_SUCCEEDED_ID, + effectId: EFFECT_SUCCEEDED_ID, + resultJson, + status: "succeeded", + toolName: "createIssue", + }); + await insertV2Attempt(database, { + effectId: EFFECT_SUCCEEDED_ID, + resultJson, + status: "succeeded", + }); + expect((await readLossyMigrationInventory(database)).totalCandidates).toBe(0); + + await applyMigration(database, "0013_durable-mcp-effect-v3"); + for (const [table, predicate] of [ + ["driver_command", "id"], + ["external_tool_effect", "id"], + ["external_tool_effect_attempt", "effect_id"], + ] as const) { + const row = await database + .prepare(`SELECT result_json FROM ${table} WHERE ${predicate} = ?`) + .bind(table === "driver_command" ? COMMAND_SUCCEEDED_ID : EFFECT_SUCCEEDED_ID) + .first<{ result_json: string }>(); + expect(row?.result_json).toBe(resultJson); + expect( + (JSON.parse(row?.result_json ?? "null") as { outputText: string }).outputText.charCodeAt(0), + ).toBe(0xd800); + } + }); + + test("rejects conflicting MCP result and provider-receipt histories", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + const payloadJson = JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + }); + const result = (outputText: string) => + JSON.stringify({ + outputText, + requestId: "legacy-request", + serverId: SERVER_ID, + toolName: "createIssue", + }); + await insertV2Command(database, { + id: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + payloadJson, + resultJson: result("command"), + seq: 1, + status: "completed", + }); + await insertV2Effect(database, { + commandId: COMMAND_SUCCEEDED_ID, + effectId: EFFECT_SUCCEEDED_ID, + providerReceiptJson: "effect-receipt", + resultJson: result("effect"), + status: "succeeded", + toolName: "createIssue", + }); + await insertV2Attempt(database, { + effectId: EFFECT_SUCCEEDED_ID, + providerReceiptJson: "attempt-receipt", + resultJson: result("attempt"), + status: "succeeded", + }); + + const inventory = await readLossyMigrationInventory(database); + expect(inventory.mcpResultConflicts).toBe(1); + expect(inventory.providerReceiptLosses).toBe(1); + expect(inventory.candidateIds).toEqual([ + { category: "mcp_result_conflict", id: EFFECT_SUCCEEDED_ID }, + { category: "provider_receipt_loss", id: EFFECT_SUCCEEDED_ID }, + ]); + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + }); + + test("rejects result and receipt data attached to an unsettled MCP effect", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + const resultJson = JSON.stringify({ + outputText: "unsettled output", + requestId: "legacy-request", + serverId: SERVER_ID, + toolName: "createIssue", + }); + await insertV2Command(database, { + id: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + payloadJson: JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + }), + seq: 1, + status: "failed", + }); + await insertV2Effect(database, { + commandId: COMMAND_UNKNOWN_ID, + effectId: EFFECT_UNKNOWN_ID, + providerReceiptJson: "unsettled-effect-receipt", + resultJson, + status: "executing", + toolName: "createIssue", + }); + await insertV2Attempt(database, { + effectId: EFFECT_UNKNOWN_ID, + providerReceiptJson: "unsettled-attempt-receipt", + resultJson, + status: "executing", + }); + + const inventory = await readLossyMigrationInventory(database); + expect(inventory.mcpResultConflicts).toBe(1); + expect(inventory.providerReceiptLosses).toBe(1); + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + }); + + test("rejects a real v2 terminal MCP command without its effect fence", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + const payloadJson = JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + }); + await insertV2Command(database, { + id: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + payloadJson, + seq: 1, + status: "failed", + }); + await expect( + database + .prepare( + "SELECT COUNT(*) AS count FROM driver_command AS command WHERE command.kind = 'mcp.execute' AND NOT EXISTS (SELECT 1 FROM external_tool_effect AS effect WHERE effect.command_id = command.id)", + ) + .first(), + ).resolves.toEqual({ count: 1 }); + + await expect(applyMigration(database, "0013_durable-mcp-effect-v3")).rejects.toThrow(); + await expect( + database + .prepare("SELECT payload_json FROM driver_command WHERE id = ?") + .bind(COMMAND_UNKNOWN_ID) + .first(), + ).resolves.toEqual({ payload_json: payloadJson }); + }); + + test("rejects an oversized source error with non-primitive details before omission", async () => { + const errorJson = jsonAtUtf8Size(OLD_V2_PAYLOAD_MAX_UTF8_BYTES, (padding) => ({ + code: "legacy.invalid_error", + details: { nested: { padding } }, + message: "Invalid legacy error.", + retryable: false, + })); + const database = await expectV2MigrationRejection(async (fixture) => { + await insertV2Command(fixture, { + errorJson, + id: COMMAND_ERROR_ID, + kind: "input.start", + payloadJson: JSON.stringify({ + commandId: COMMAND_ERROR_ID, + input: { text: "hello" }, + kind: "input.start", + requestId: "legacy-input-request", + runId: RUN_ID, + }), + seq: 1, + status: "failed", + }); + }); + + await expect( + database + .prepare("SELECT error_json FROM driver_command WHERE id = ?") + .bind(COMMAND_ERROR_ID) + .first(), + ).resolves.toEqual({ error_json: errorJson }); + }); + + test("rejects one malformed MCP result copy before authoritative synchronization", async () => { + const payloadJson = JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + }); + const validResultJson = JSON.stringify({ + outputText: "done", + requestId: "legacy-request", + serverId: SERVER_ID, + toolName: "createIssue", + }); + const invalidAttemptResultJson = JSON.stringify({ + requestId: "legacy-request", + serverId: SERVER_ID, + toolName: "createIssue", + }); + const database = await expectV2MigrationRejection(async (fixture) => { + await insertV2Command(fixture, { + id: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + payloadJson, + resultJson: validResultJson, + seq: 1, + status: "completed", + }); + await insertV2Effect(fixture, { + commandId: COMMAND_SUCCEEDED_ID, + effectId: EFFECT_SUCCEEDED_ID, + resultJson: validResultJson, + status: "succeeded", + toolName: "createIssue", + }); + await insertV2Attempt(fixture, { + effectId: EFFECT_SUCCEEDED_ID, + resultJson: invalidAttemptResultJson, + status: "succeeded", + }); + }); + + await expect( + database + .prepare("SELECT result_json FROM external_tool_effect_attempt WHERE effect_id = ?") + .bind(EFFECT_SUCCEEDED_ID) + .first(), + ).resolves.toEqual({ result_json: invalidAttemptResultJson }); + }); + + test("rejects undeclared MCP result fields before authoritative synchronization", async () => { + const payloadJson = JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + }); + const resultJson = JSON.stringify({ + debug: "provider-specific data must not enter the public result", + outputText: "done", + requestId: "legacy-request", + serverId: SERVER_ID, + toolName: "createIssue", + }); + await expectV2MigrationRejection(async (fixture) => { + await insertV2Command(fixture, { + id: COMMAND_SUCCEEDED_ID, + kind: "mcp.execute", + payloadJson, + resultJson, + seq: 1, + status: "completed", + }); + await insertV2Effect(fixture, { + commandId: COMMAND_SUCCEEDED_ID, + effectId: EFFECT_SUCCEEDED_ID, + resultJson, + status: "succeeded", + toolName: "createIssue", + }); + await insertV2Attempt(fixture, { + effectId: EFFECT_SUCCEEDED_ID, + resultJson, + status: "succeeded", + }); + }); + }); + + test("rejects a missing input Run identity instead of fabricating one", async () => { + const payloadJson = jsonAtUtf8Size(OLD_V2_PAYLOAD_MAX_UTF8_BYTES, (text) => ({ + commandId: COMMAND_ERROR_ID, + input: { text }, + kind: "input.start", + requestId: "legacy-input-request", + })); + await expectV2MigrationRejection(async (fixture) => { + await insertV2Command(fixture, { + id: COMMAND_ERROR_ID, + kind: "input.start", + payloadJson, + seq: 1, + status: "failed", + }); + }); + }); + + test("rejects empty input attachment identities that Driver cannot execute", async () => { + await expectV2MigrationRejection(async (fixture) => { + await insertV2Command(fixture, { + id: COMMAND_ERROR_ID, + kind: "input.start", + payloadJson: JSON.stringify({ + commandId: COMMAND_ERROR_ID, + input: { attachmentIds: [""], text: "hello" }, + kind: "input.start", + requestId: "legacy-input-request", + runId: RUN_ID, + }), + seq: 1, + status: "failed", + }); + }); + }); + + test("rejects an invalid effect attempt audit instead of inventing a current claim", async () => { + const payloadJson = JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + }); + await expectV2MigrationRejection(async (fixture) => { + await insertV2Command(fixture, { + id: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + payloadJson, + seq: 1, + status: "failed", + }); + await insertV2Effect(fixture, { + commandId: COMMAND_UNKNOWN_ID, + effectId: EFFECT_UNKNOWN_ID, + resultJson: null, + status: "executing", + toolName: "createIssue", + }); + await insertV2Attempt(fixture, { + effectId: EFFECT_UNKNOWN_ID, + resultJson: null, + status: "executing", + }); + await fixture + .prepare("UPDATE external_tool_effect SET attempt_count = 2 WHERE id = ?") + .bind(EFFECT_UNKNOWN_ID) + .run(); + }); + }); + + test("rejects a completed command whose effect remains unknown", async () => { + const payloadJson = JSON.stringify({ + argumentsJson: "{}", + commandId: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + requestId: "legacy-request", + serverId: SERVER_ID, + toolCallId: "legacy-tool-call", + toolName: "createIssue", + }); + const resultJson = JSON.stringify({ + outputText: "unproven", + requestId: "legacy-request", + serverId: SERVER_ID, + toolName: "createIssue", + }); + await expectV2MigrationRejection(async (fixture) => { + await insertV2Command(fixture, { + id: COMMAND_UNKNOWN_ID, + kind: "mcp.execute", + payloadJson, + resultJson, + seq: 1, + status: "completed", + }); + await insertV2Effect(fixture, { + commandId: COMMAND_UNKNOWN_ID, + effectId: EFFECT_UNKNOWN_ID, + resultJson: null, + status: "executing", + toolName: "createIssue", + }); + await insertV2Attempt(fixture, { + effectId: EFFECT_UNKNOWN_ID, + resultJson: null, + status: "executing", + }); + }); + }); + + test("rejects nested Session Run error details before oversized repair", async () => { + await expectV2MigrationRejection(async (fixture) => { + await fixture + .prepare( + "UPDATE session_run SET error_code = ?, error_details_json = ?, error_message = ? WHERE id = ?", + ) + .bind( + "legacy.invalid_session_error", + JSON.stringify({ nested: { invalid: true } }), + "Invalid Session Run error.", + RUN_ID, + ) + .run(); + }); + }); + + test("rejects missing permission decisions and RunError retryability", async () => { + await expectV2MigrationRejection(async (fixture) => { + await insertV2Command(fixture, { + errorJson: JSON.stringify({ + code: "legacy.missing_retryable", + details: {}, + message: "Missing retryability.", + }), + id: COMMAND_ERROR_ID, + kind: "permission.resolve", + payloadJson: JSON.stringify({ + commandId: COMMAND_ERROR_ID, + kind: "permission.resolve", + requestId: "legacy-permission-request", + }), + seq: 1, + status: "failed", + }); + }); + }); +}); + +describe("session event v3 migration", () => { + test("leaves the canonical deploy lease table on a direct no-candidate migration", async () => { + const database = await createV2Database(); + await applyMigration(database, "0013_durable-mcp-effect-v3"); + await applyMigration(database, "0014_session-event-stream-identity"); + + const results = await database.batch( + acquireProdDeployLeaseStatements(DEPLOY_OWNER).map((sql) => database.prepare(sql)), + ); + expect(() => assertProdDeployLeaseOwned(JSON.stringify(results), DEPLOY_OWNER)).not.toThrow(); + }); + + for (const fixture of [ + { eventType: "run.completed", status: "completed", statusEvent: "run.complete" }, + { eventType: "run.cancelled", status: "cancelled", statusEvent: "run.cancel" }, + { eventType: "run.failed", status: "failed", statusEvent: "run.fail" }, + ] as const) { + test(`normalizes a provider-source ${fixture.eventType} through the exact cutover gate`, async () => { + const database = await createV2Database(); + await insertV2Owners(database); + await database + .prepare( + `UPDATE session_run + SET error_code = ?, + error_details_json = ?, + error_message = ?, + status = ?, + status_event = ? + WHERE id = ?`, + ) + .bind( + fixture.status === "failed" ? "legacy.failed" : null, + fixture.status === "failed" ? "{}" : null, + fixture.status === "failed" ? "Legacy failure." : null, + fixture.status, + fixture.statusEvent, + RUN_ID, + ) + .run(); + database.execute( + `UPDATE session SET runtime_event_seq_cursor = 1 WHERE id = '${SESSION_ID}'`, + ); + await insertLegacySessionTerminal(database, { + eventId: LEGACY_TERMINAL_EVENT_ID, + eventType: fixture.eventType, + sourceEventId: `provider-${fixture.status}-event`, + }); + + await applyMigration(database, "0013_durable-mcp-effect-v3"); + await authorizeLegacyTerminalRewrite(database); + await applyMigration(database, "0014_session-event-stream-identity"); + + await expect(database.prepare(PROTOCOL_V3_CUTOVER_OBJECTS_SQL).first()).resolves.toEqual({ + exact_object_count: PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + object_count: PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + }); + + await expect( + database + .prepare( + "SELECT count(*) AS count FROM sqlite_master WHERE name = '__protocol_v3_legacy_rewrite_authorization'", + ) + .first(), + ).resolves.toEqual({ count: 0 }); + + await expect( + database + .prepare("SELECT id, semantic_hash, seq, source_event_id FROM session_event WHERE id = ?") + .bind(LEGACY_TERMINAL_EVENT_ID) + .first(), + ).resolves.toEqual({ + id: LEGACY_TERMINAL_EVENT_ID, + semantic_hash: null, + seq: 1, + source_event_id: `session-run-terminal:${RUN_ID}:${fixture.eventType}`, + }); + }); + } + + test("rejects an otherwise valid provider-source rewrite without deploy authorization", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + database.execute(` + UPDATE session SET runtime_event_seq_cursor = 1 WHERE id = '${SESSION_ID}'; + UPDATE session_run + SET error_code = NULL, error_details_json = NULL, error_message = NULL, + status = 'completed', status_event = 'run.complete' + WHERE id = '${RUN_ID}'; + `); + await insertLegacySessionTerminal(database, { + eventId: LEGACY_TERMINAL_EVENT_ID, + eventType: "run.completed", + sourceEventId: "provider-completed-event", + }); + await applyMigration(database, "0013_durable-mcp-effect-v3"); + + await expect(applyMigration(database, "0014_session-event-stream-identity")).rejects.toThrow(); + await expect( + database + .prepare("SELECT source_event_id FROM session_event WHERE id = ?") + .bind(LEGACY_TERMINAL_EVENT_ID) + .first(), + ).resolves.toEqual({ source_event_id: "provider-completed-event" }); + }); + + test("does not authorize a candidate set that changed after the integrity preflight", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + database.execute(` + UPDATE session SET runtime_event_seq_cursor = 1 WHERE id = '${SESSION_ID}'; + UPDATE session_run + SET error_code = NULL, error_details_json = NULL, error_message = NULL, + status = 'completed', status_event = 'run.complete' + WHERE id = '${RUN_ID}'; + `); + await insertLegacySessionTerminal(database, { + eventId: LEGACY_TERMINAL_EVENT_ID, + eventType: "run.completed", + sourceEventId: "provider-completed-event", + }); + await applyMigration(database, "0013_durable-mcp-effect-v3"); + const integrity = await readLegacyTerminalIntegrity(database); + installLegacyRewriteAuthorizationInfrastructure(database); + + database.execute( + `UPDATE session_event SET source_event_id = 'provider-changed-event' WHERE id = '${LEGACY_TERMINAL_EVENT_ID}'`, + ); + database.execute( + authorizeProtocolV3LegacyRewriteSql( + DEPLOY_OWNER, + integrity.noncanonicalTerminalSources, + integrity.rewriteCandidateManifestJson, + ), + ); + + await expect( + database + .prepare( + "SELECT count(*) AS count FROM __protocol_v3_legacy_rewrite_authorization WHERE id = 1", + ) + .first(), + ).resolves.toEqual({ count: 0 }); + await expect(applyMigration(database, "0014_session-event-stream-identity")).rejects.toThrow(); + }); + + test("revokes stale authorization when a retry no longer proves the frozen gate", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + database.execute(` + UPDATE session SET runtime_event_seq_cursor = 1 WHERE id = '${SESSION_ID}'; + UPDATE session_run + SET error_code = NULL, error_details_json = NULL, error_message = NULL, + status = 'completed', status_event = 'run.complete' + WHERE id = '${RUN_ID}'; + `); + await insertLegacySessionTerminal(database, { + eventId: LEGACY_TERMINAL_EVENT_ID, + eventType: "run.completed", + sourceEventId: "provider-completed-event", + }); + await applyMigration(database, "0013_durable-mcp-effect-v3"); + installLegacyRewriteAuthorizationInfrastructure(database); + const integrity = await readLegacyTerminalIntegrity(database); + const authorizationSql = authorizeProtocolV3LegacyRewriteSql( + DEPLOY_OWNER, + integrity.noncanonicalTerminalSources, + integrity.rewriteCandidateManifestJson, + ); + database.execute(authorizationSql); + database.execute(`UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" SET "command_freeze" = 0`); + database.execute(authorizationSql); + + await expect( + database + .prepare( + "SELECT count(*) AS count FROM __protocol_v3_legacy_rewrite_authorization WHERE id = 1", + ) + .first(), + ).resolves.toEqual({ count: 0 }); + await expect(applyMigration(database, "0014_session-event-stream-identity")).rejects.toThrow(); + }); + + for (const fixture of [ + { + authorizationRemains: true, + label: "the authorized candidate identity changes", + mutate: (database: SqliteD1Database) => + database.execute( + `UPDATE session_event SET source_event_id = 'provider-changed-event' WHERE id = '${LEGACY_TERMINAL_EVENT_ID}'`, + ), + }, + { + authorizationRemains: true, + label: "the durable deploy mutex owner changes", + mutate: (database: SqliteD1Database) => + database.execute( + `UPDATE "${PROD_DEPLOY_LEASE_TABLE}" SET "owner" = '${OTHER_DEPLOY_OWNER}'`, + ), + }, + { + authorizationRemains: true, + label: "the rewrite authorization expires", + mutate: (database: SqliteD1Database) => + database.execute( + "UPDATE __protocol_v3_legacy_rewrite_authorization SET expires_at = 0 WHERE id = 1", + ), + }, + { + authorizationRemains: true, + label: "the terminal projection falls behind its Session cursor proof", + mutate: (database: SqliteD1Database) => + database.execute( + `UPDATE session SET runtime_event_seq_cursor = 0 WHERE id = '${SESSION_ID}'`, + ), + }, + { + authorizationRemains: false, + label: "the admission freeze is disabled", + mutate: (database: SqliteD1Database) => + database.execute(`UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" SET command_freeze = 0`), + }, + { + authorizationRemains: false, + label: "the bound bookmark changes", + mutate: (database: SqliteD1Database) => + database.execute( + `UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" SET pre_migration_bookmark = 'changed-bookmark'`, + ), + }, + { + authorizationRemains: false, + label: "the bound release tree changes", + mutate: (database: SqliteD1Database) => + database.execute( + `UPDATE "${PROTOCOL_V3_CUTOVER_TABLE}" SET release_tree_oid = '89abcdef0123456789abcdef0123456789abcdef'`, + ), + }, + { + authorizationRemains: false, + label: "the admission gate is removed", + mutate: (database: SqliteD1Database) => database.execute(REMOVE_PROTOCOL_V3_CUTOVER_SQL), + }, + { + authorizationRemains: true, + label: "an extra trigger is attached to an admission table", + mutate: (database: SqliteD1Database) => + database.execute(` + CREATE TRIGGER disable_protocol_v3_gate + BEFORE INSERT ON session_run + WHEN 0 + BEGIN + SELECT 1; + END + `), + }, + { + authorizationRemains: true, + label: "an extra trigger is attached to App deployment admission", + mutate: (database: SqliteD1Database) => + database.execute(` + CREATE TRIGGER disable_app_deployment_gate + BEFORE INSERT ON app_deployment_run + WHEN 0 + BEGIN + SELECT 1; + END + `), + }, + { + authorizationRemains: true, + label: "an extra trigger is attached to the rewrite authorization table", + mutate: (database: SqliteD1Database) => + database.execute(` + CREATE TRIGGER retain_protocol_v3_authorization + AFTER DELETE ON __protocol_v3_legacy_rewrite_authorization + WHEN 0 + BEGIN + SELECT 1; + END + `), + }, + { + authorizationRemains: true, + label: "an extra trigger is attached to the deploy lease table", + mutate: (database: SqliteD1Database) => + database.execute(` + CREATE TRIGGER spoof_protocol_v3_deploy_lease + AFTER UPDATE ON "${PROD_DEPLOY_LEASE_TABLE}" + WHEN 0 + BEGIN + SELECT 1; + END + `), + }, + { + authorizationRemains: true, + label: "a canonical admission trigger is replaced by WHEN 0", + mutate: (database: SqliteD1Database) => + database.execute(` + DROP TRIGGER "__protocol_v3_cutover_session_run_insert"; + CREATE TRIGGER "__protocol_v3_cutover_session_run_insert" + BEFORE INSERT ON session_run + WHEN 0 + BEGIN + SELECT 1; + END + `), + }, + ]) { + test(`rejects the legacy rewrite after ${fixture.label}`, async () => { + const database = await createV2Database(); + await insertV2Owners(database); + database.execute(` + UPDATE session SET runtime_event_seq_cursor = 1 WHERE id = '${SESSION_ID}'; + UPDATE session_run + SET error_code = NULL, error_details_json = NULL, error_message = NULL, + status = 'completed', status_event = 'run.complete' + WHERE id = '${RUN_ID}'; + `); + await insertLegacySessionTerminal(database, { + eventId: LEGACY_TERMINAL_EVENT_ID, + eventType: "run.completed", + sourceEventId: "provider-completed-event", + }); + await applyMigration(database, "0013_durable-mcp-effect-v3"); + await authorizeLegacyTerminalRewrite(database); + fixture.mutate(database); + + await expect( + applyMigration(database, "0014_session-event-stream-identity"), + ).rejects.toThrow(); + await expect( + database + .prepare( + "SELECT candidate_count FROM __protocol_v3_legacy_rewrite_authorization WHERE id = 1", + ) + .first(), + ).resolves.toEqual(fixture.authorizationRemains ? { candidate_count: 1 } : null); + }); + } + + test("leaves an already-canonical legacy terminal identity unchanged", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + await database + .prepare( + `UPDATE session_run + SET error_code = 'legacy.failed', error_details_json = '{}', + error_message = 'Legacy failure.', status_event = 'run.fail' + WHERE id = ?`, + ) + .bind(RUN_ID) + .run(); + database.execute(`UPDATE session SET runtime_event_seq_cursor = 1 WHERE id = '${SESSION_ID}'`); + const canonicalSource = `session-run-terminal:${RUN_ID}:run.failed`; + await insertLegacySessionTerminal(database, { + eventId: LEGACY_TERMINAL_EVENT_ID, + eventType: "run.failed", + sourceEventId: canonicalSource, + }); + + await applyMigration(database, "0013_durable-mcp-effect-v3"); + await applyMigration(database, "0014_session-event-stream-identity"); + + await expect( + database + .prepare("SELECT id, semantic_hash, seq, source_event_id FROM session_event WHERE id = ?") + .bind(LEGACY_TERMINAL_EVENT_ID) + .first(), + ).resolves.toEqual({ + id: LEGACY_TERMINAL_EVENT_ID, + semantic_hash: null, + seq: 1, + source_event_id: canonicalSource, + }); + }); + + test("rejects a legacy terminal source rewrite collision atomically", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + await database + .prepare( + "UPDATE session_run SET status = 'completed', status_event = 'run.complete' WHERE id = ?", + ) + .bind(RUN_ID) + .run(); + await insertLegacySessionTerminal(database, { + eventId: LEGACY_TERMINAL_EVENT_ID, + eventType: "run.completed", + sourceEventId: "provider-completed-event", + }); + await database + .prepare( + `INSERT INTO session_event ( + agent_id, content_text, created_at, ended_at, event_type, family, id, + occurred_at, process_status, process_type, run_id, seq, session_id, + source_event_id, source, visibility + ) VALUES (?, '', 4, 4, 'message.added', 'message', ?, 4, 'available', + 'agent.message.delta', ?, 2, ?, ?, 'api', 'all_consumers')`, + ) + .bind( + AGENT_ID, + "01J00000000000000000000022", + RUN_ID, + SESSION_ID, + `session-run-terminal:${RUN_ID}:run.completed`, + ) + .run(); + + await applyMigration(database, "0013_durable-mcp-effect-v3"); + await expect(applyMigration(database, "0014_session-event-stream-identity")).rejects.toThrow(); + await expect( + database + .prepare("SELECT source_event_id FROM session_event WHERE id = ?") + .bind(LEGACY_TERMINAL_EVENT_ID) + .first(), + ).resolves.toEqual({ source_event_id: "provider-completed-event" }); + }); + + test("rejects multiple legacy terminal winners before rewriting either source", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + await database + .prepare( + "UPDATE session_run SET status = 'completed', status_event = 'run.complete' WHERE id = ?", + ) + .bind(RUN_ID) + .run(); + await insertLegacySessionTerminal(database, { + eventId: LEGACY_TERMINAL_EVENT_ID, + eventType: "run.completed", + sourceEventId: "provider-completed-event-1", + }); + await insertLegacySessionTerminal(database, { + eventId: "01J00000000000000000000022", + eventType: "run.completed", + seq: 2, + sourceEventId: "provider-completed-event-2", + }); + + await applyMigration(database, "0013_durable-mcp-effect-v3"); + await expect(applyMigration(database, "0014_session-event-stream-identity")).rejects.toThrow(); + await expect( + database + .prepare("SELECT source_event_id FROM session_event WHERE run_id = ? ORDER BY seq") + .bind(RUN_ID) + .all(), + ).resolves.toMatchObject({ + results: [ + { source_event_id: "provider-completed-event-1" }, + { source_event_id: "provider-completed-event-2" }, + ], + }); + }); + + for (const fixture of [ + { + eventType: "run.completed", + label: "a terminal kind that conflicts with the Run status", + runId: RUN_ID, + }, + { + eventType: "run.failed", + label: "a terminal event whose Run link does not exist", + runId: "01J00000000000000000000023", + }, + ] as const) { + test(`rejects ${fixture.label} before source normalization`, async () => { + const database = await createV2Database(); + await insertV2Owners(database); + await database + .prepare("UPDATE session_run SET status_event = 'run.fail' WHERE id = ?") + .bind(RUN_ID) + .run(); + await insertLegacySessionTerminal(database, { + eventId: LEGACY_TERMINAL_EVENT_ID, + eventType: fixture.eventType, + runId: fixture.runId, + sourceEventId: "provider-terminal-event", + }); + + await applyMigration(database, "0013_durable-mcp-effect-v3"); + await expect( + applyMigration(database, "0014_session-event-stream-identity"), + ).rejects.toThrow(); + await expect( + database + .prepare("SELECT source_event_id FROM session_event WHERE id = ?") + .bind(LEGACY_TERMINAL_EVENT_ID) + .first(), + ).resolves.toEqual({ source_event_id: "provider-terminal-event" }); + }); + } + + test("labels legacy projections and backfills retryable Run errors without inventing receipts", async () => { + const database = await createV2Database(); + await insertV2Owners(database); + database.execute(` + INSERT INTO session_message ( + content_text, created_at, created_by_account_id, id, plan_json, role, + segments_json, seq, session_id, session_run_id + ) VALUES ( + 'legacy materialized message', 2, '${ACCOUNT_ID}', '${LEGACY_MESSAGE_ID}', + NULL, 'user', NULL, 1, '${SESSION_ID}', NULL + ); + + UPDATE session_run + SET error_code = 'legacy.failed', + error_details_json = NULL, + error_message = 'Legacy failure.', + status_event = 'run.fail' + WHERE id = '${RUN_ID}'; + + INSERT INTO session_event ( + agent_id, content_text, created_at, ended_at, event_type, family, id, + occurred_at, process_status, process_type, run_id, seq, session_id, + source_event_id, source, visibility + ) VALUES ( + '${AGENT_ID}', 'Legacy failure.', 3, 3, 'run.failed', 'run', + '${LEGACY_TERMINAL_EVENT_ID}', 3, 'error', 'run.failed', '${RUN_ID}', 1, + '${SESSION_ID}', 'legacy-terminal-source', 'api', 'all_consumers' + ); + UPDATE session SET runtime_event_seq_cursor = 1 WHERE id = '${SESSION_ID}'; + `); + + await applyMigration(database, "0013_durable-mcp-effect-v3"); + await authorizeLegacyTerminalRewrite(database); + await applyMigration(database, "0014_session-event-stream-identity"); + + await expect( + database + .prepare("SELECT projection_format FROM session_message WHERE id = ?") + .bind(LEGACY_MESSAGE_ID) + .first(), + ).resolves.toEqual({ projection_format: "materialized" }); + await expect( + database + .prepare("SELECT error_details_json, error_retryable FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ error_details_json: "{}", error_retryable: 0 }); + await expect( + database + .prepare("SELECT semantic_hash, source_event_id FROM session_event WHERE id = ?") + .bind(LEGACY_TERMINAL_EVENT_ID) + .first(), + ).resolves.toEqual({ + semantic_hash: null, + source_event_id: `session-run-terminal:${RUN_ID}:run.failed`, + }); + await expect( + database + .prepare("UPDATE session_message SET projection_format = 'invalid' WHERE id = ?") + .bind(LEGACY_MESSAGE_ID) + .run(), + ).rejects.toThrow(); + await database + .prepare( + `INSERT INTO session_message ( + content_text, created_at, created_by_account_id, id, plan_json, + projection_format, role, segments_json, seq, session_id, session_run_id + ) VALUES ('', 4, ?, ?, NULL, 'event_stream_v3', 'assistant', NULL, 2, ?, ?)`, + ) + .bind(ACCOUNT_ID, V3_REFERENCE_MESSAGE_ID, SESSION_ID, RUN_ID) + .run(); + for (const statement of [ + "UPDATE session_message SET role = 'user' WHERE id = ?", + "UPDATE session_message SET session_run_id = NULL WHERE id = ?", + "UPDATE session_message SET content_text = 'materialized' WHERE id = ?", + "UPDATE session_message SET plan_json = '{}' WHERE id = ?", + "UPDATE session_message SET segments_json = '[]' WHERE id = ?", + ]) { + await expect( + database.prepare(statement).bind(V3_REFERENCE_MESSAGE_ID).run(), + ).rejects.toThrow(); + } + await expect(database.prepare("PRAGMA foreign_key_check").all()).resolves.toMatchObject({ + results: [], + success: true, + }); + }); +}); diff --git a/apps/api/tests/external-tool-effect-store.test.ts b/apps/api/tests/external-tool-effect-store.test.ts new file mode 100644 index 00000000..e6b5ea97 --- /dev/null +++ b/apps/api/tests/external-tool-effect-store.test.ts @@ -0,0 +1,711 @@ +import { describe, expect, test } from "bun:test"; + +import { + MCP_EXTERNAL_TOOL_EFFECT_SETTLEMENT_MAX_UTF8_BYTES, + measureMcpExternalToolEffectSettlement, +} from "@mosoo/contracts/external-tool-effect"; +import { + RUNTIME_COMMAND_MAX_UTF8_BYTES, + measureRuntimeCommandJson, +} from "@mosoo/contracts/runtime-command"; +import type { RuntimeCommand } from "@mosoo/contracts/runtime-command"; +import type { + DriverCommandId, + DriverInstanceId, + ExternalToolEffectId, + SessionRunId, +} from "@mosoo/id"; + +import { + claimExternalToolEffect as claimExternalToolEffectRecord, + getExternalToolEffectForCommand as getExternalToolEffectForCommandRecord, + markClaimedExternalToolEffectsUnknownForDriver, + settleExternalToolEffect as settleExternalToolEffectRecord, +} from "../src/modules/runtime/infrastructure/session-runs/external-tool-effect-store.repository"; +import { + createRuntimeCommandRecord as persistRuntimeCommandRecord, + getRuntimeCommandRecord as readRuntimeCommandRecord, + listAcceptedInputStartCommandRepairsForTerminalDriver, + listAcceptedMcpCommandRepairsForTerminalDriver, + repairAcceptedRuntimeCommandsForTerminalDriver, +} from "../src/modules/runtime/infrastructure/session-runs/runtime-command-store.repository"; +import { + createPublicHttpContractDatabase, + insertActiveSandboxSessionFixture, + insertOwnerSession, + PUBLIC_API_TEST_IDS, +} from "./helpers/public-api-http-test-fixture"; +import type { SqliteD1Database } from "./helpers/public-api-http-test-fixture"; + +const RUN_ID = "01J0000000000000000000000T" as SessionRunId; +const NEXT_RUN_ID = "01J0000000000000000000000V" as SessionRunId; +const COMMAND_ID = "01J0000000000000000000000X" as DriverCommandId; +const BOUNDARY_COMMAND_ID = "01J0000000000000000000000Z" as DriverCommandId; +const OVERSIZED_COMMAND_ID = "01J00000000000000000000010" as DriverCommandId; +const DRIVER_INSTANCE_ID = PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId; +const DRIVER_GENERATION = 0; +const CLAIM_TOKEN = "123e4567-e89b-42d3-a456-426614174000"; + +function createRuntimeCommandRecord( + database: D1Database, + input: Omit[1], "driverGeneration">, +) { + return persistRuntimeCommandRecord(database, { ...input, driverGeneration: DRIVER_GENERATION }); +} + +function getRuntimeCommandRecord( + database: D1Database, + driverInstanceId: DriverInstanceId, + commandId: DriverCommandId, +) { + return readRuntimeCommandRecord(database, driverInstanceId, DRIVER_GENERATION, commandId); +} + +function claimExternalToolEffect( + database: D1Database, + input: Omit[1], "driverGeneration">, +) { + return claimExternalToolEffectRecord(database, { ...input, driverGeneration: DRIVER_GENERATION }); +} + +function settleExternalToolEffect( + database: D1Database, + input: Omit[1], "driverGeneration">, +) { + return settleExternalToolEffectRecord(database, { + ...input, + driverGeneration: DRIVER_GENERATION, + }); +} + +function getExternalToolEffectForCommand( + database: D1Database, + input: Omit[1], "driverGeneration">, +) { + return getExternalToolEffectForCommandRecord(database, { + ...input, + driverGeneration: DRIVER_GENERATION, + }); +} + +function mcpExecuteCommand( + runId: SessionRunId, + commandId: DriverCommandId = COMMAND_ID, +): Extract { + return { + argumentsJson: '{"title":"durable"}', + commandId, + kind: "mcp.execute", + requestId: "request-1", + runId, + serverId: "01J0000000000000000000000Y", + toolCallId: "tool-1", + toolName: "createIssue", + }; +} + +function mcpExecuteCommandAtSize( + targetBytes: number, + commandId: DriverCommandId = BOUNDARY_COMMAND_ID, +) { + const command = { + ...mcpExecuteCommand(RUN_ID, commandId), + argumentsJson: "", + }; + return { + ...command, + argumentsJson: "x".repeat(targetBytes - measureRuntimeCommandJson(command)), + }; +} + +function succeededSettlementAtSize(targetBytes: number) { + const settlement = { + kind: "succeeded" as const, + result: { + outputText: "", + requestId: "request-1", + serverId: "01J0000000000000000000000Y", + toolName: "createIssue", + }, + }; + return { + ...settlement, + result: { + ...settlement.result, + outputText: "x".repeat(targetBytes - measureMcpExternalToolEffectSettlement(settlement)), + }, + }; +} + +async function insertSessionRun(database: SqliteD1Database, runId: SessionRunId): Promise { + await database + .prepare( + "INSERT INTO session_run (id, session_id, agent_id, created_by_account_id, driver_instance_id, trigger, status, trace_id, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind( + runId, + PUBLIC_API_TEST_IDS.ownerSession, + PUBLIC_API_TEST_IDS.agent, + PUBLIC_API_TEST_IDS.ownerAccount, + DRIVER_INSTANCE_ID, + "user_prompt", + "running", + `trace-${runId}`, + 1, + 1, + ) + .run(); +} + +async function createEffectStoreFixture(): Promise { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await insertActiveSandboxSessionFixture(database, { + ownerAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxSessionId: "01J0000000000000000000000W", + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + }); + await database + .prepare( + "INSERT INTO driver_instance (id, boot_token_expires_at, boot_token_hash, connection_id, created_at, expires_at, heartbeat_count, protocol, protocol_version, runtime, sandbox_id, sandbox_incarnation, sandbox_session_id, status, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind( + DRIVER_INSTANCE_ID, + 1, + new Uint8Array([1]), + "connection-current", + 1, + Date.now() + 60_000, + 0, + "orpc-ws", + 3, + "openai-runtime", + PUBLIC_API_TEST_IDS.sandbox, + 1, + PUBLIC_API_TEST_IDS.ownerSession, + "ready", + 1, + ) + .run(); + await insertSessionRun(database, RUN_ID); + return database; +} + +describe("external tool effect store", () => { + test("atomically rejects commands after the exact Driver generation becomes terminal", async () => { + const database = await createEffectStoreFixture(); + await database + .prepare("UPDATE driver_instance SET status = 'failed' WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .run(); + + await expect( + createRuntimeCommandRecord(database, { + command: { + commandId: COMMAND_ID, + kind: "session.stop", + reason: "already terminal", + }, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).rejects.toThrow("Driver generation is no longer current."); + await expect( + database + .prepare( + "SELECT command_seq_cursor, (SELECT count(*) FROM driver_command) AS command_count FROM driver_instance WHERE id = ?", + ) + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ command_count: 0, command_seq_cursor: 0 }); + }); + + test("atomically rejects a command for an inactive exact Run", async () => { + const database = await createEffectStoreFixture(); + await database + .prepare("UPDATE session_run SET status = 'failed' WHERE id = ?") + .bind(RUN_ID) + .run(); + await insertSessionRun(database, NEXT_RUN_ID); + + await expect( + createRuntimeCommandRecord(database, { + command: mcpExecuteCommand(RUN_ID), + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }), + ).rejects.toThrow("MCP external tool effects require the command's active Session Run."); + await expect( + getRuntimeCommandRecord(database, DRIVER_INSTANCE_ID, COMMAND_ID), + ).resolves.toBeNull(); + await expect( + database.prepare("SELECT COUNT(*) AS count FROM external_tool_effect").first(), + ).resolves.toEqual({ count: 0 }); + }); + + test("does not claim an intent after its exact Run becomes terminal", async () => { + const database = await createEffectStoreFixture(); + await createRuntimeCommandRecord(database, { + command: mcpExecuteCommand(RUN_ID), + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + await database + .prepare("UPDATE session_run SET status = 'failed' WHERE id = ?") + .bind(RUN_ID) + .run(); + + await expect( + claimExternalToolEffect(database, { + claimToken: CLAIM_TOKEN, + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).rejects.toThrow("External tool effect claim did not reach a stable state."); + await expect( + getExternalToolEffectForCommand(database, { + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).resolves.toMatchObject({ attemptCount: 0, claimToken: null, status: "intent" }); + await expect( + database.prepare("SELECT COUNT(*) AS count FROM external_tool_effect_attempt").first(), + ).resolves.toEqual({ count: 0 }); + }); + + test("rejects malformed claim tokens before claim or settlement mutation", async () => { + const database = await createEffectStoreFixture(); + await createRuntimeCommandRecord(database, { + command: mcpExecuteCommand(RUN_ID), + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + + await expect( + claimExternalToolEffect(database, { + claimToken: "x".repeat(1_000_000), + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).rejects.toThrow(); + const claim = await claimExternalToolEffect(database, { + claimToken: CLAIM_TOKEN, + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + }); + if (claim.kind !== "claimed") { + throw new Error("Expected the external effect to be claimed."); + } + + await expect( + settleExternalToolEffect(database, { + claimToken: "not-a-uuid", + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + effectId: claim.effectId as ExternalToolEffectId, + settlement: { kind: "unknown" }, + }), + ).rejects.toThrow(); + await expect( + getExternalToolEffectForCommand(database, { + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).resolves.toMatchObject({ claimToken: CLAIM_TOKEN, status: "claimed" }); + }); + + test("rejects an invalid settlement before mutating the effect or attempt", async () => { + const database = await createEffectStoreFixture(); + await createRuntimeCommandRecord(database, { + command: mcpExecuteCommand(RUN_ID), + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + const claim = await claimExternalToolEffect(database, { + claimToken: CLAIM_TOKEN, + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + }); + if (claim.kind !== "claimed") { + throw new Error("Expected the external effect to be claimed."); + } + + await expect( + settleExternalToolEffect(database, { + claimToken: CLAIM_TOKEN, + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + effectId: claim.effectId as ExternalToolEffectId, + settlement: { + kind: "succeeded", + result: { + requestId: "request-1", + serverId: "01J0000000000000000000000Y", + toolName: "createIssue", + }, + } as never, + }), + ).rejects.toThrow(); + await expect( + settleExternalToolEffect(database, { + claimToken: CLAIM_TOKEN, + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + effectId: claim.effectId as ExternalToolEffectId, + settlement: { + kind: "succeeded", + result: { + debug: "must not enter durable state", + outputText: "done", + requestId: "request-1", + serverId: "01J0000000000000000000000Y", + toolName: "createIssue", + }, + } as never, + }), + ).rejects.toThrow(); + await expect( + database + .prepare( + `SELECT + command.error_json AS command_error_json, + command.result_json AS command_result_json, + command.status AS command_status, + effect.status, + effect.result_json, + attempt.status AS attempt_status, + attempt.completed_at, + attempt.result_json AS attempt_result_json + FROM external_tool_effect AS effect + JOIN external_tool_effect_attempt AS attempt ON attempt.effect_id = effect.id + JOIN driver_command AS command ON command.id = effect.command_id + WHERE effect.command_id = ?`, + ) + .bind(COMMAND_ID) + .first(), + ).resolves.toEqual({ + attempt_result_json: null, + attempt_status: "claimed", + command_error_json: null, + command_result_json: null, + command_status: "accepted", + completed_at: null, + result_json: null, + status: "claimed", + }); + }); + + test("classifies an intent as unexecuted and fences a claimed effect before repair", async () => { + const database = await createEffectStoreFixture(); + const command = mcpExecuteCommand(RUN_ID); + await createRuntimeCommandRecord(database, { + command, + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + + await expect( + listAcceptedMcpCommandRepairsForTerminalDriver(database, { + driverGeneration: DRIVER_GENERATION, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).resolves.toMatchObject([ + { + commandId: COMMAND_ID, + terminal: { + error: { code: "driver.external_tool_effect_not_executed", retryable: true }, + status: "failed", + }, + }, + ]); + await claimExternalToolEffect(database, { + claimToken: CLAIM_TOKEN, + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + }); + await expect( + listAcceptedMcpCommandRepairsForTerminalDriver(database, { + driverGeneration: DRIVER_GENERATION, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).rejects.toThrow("must fence claimed MCP effects"); + + await markClaimedExternalToolEffectsUnknownForDriver(database, { + driverGeneration: DRIVER_GENERATION + 1, + driverInstanceId: DRIVER_INSTANCE_ID, + }); + await expect( + getExternalToolEffectForCommand(database, { + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).resolves.toMatchObject({ status: "claimed" }); + + await markClaimedExternalToolEffectsUnknownForDriver(database, { + driverGeneration: DRIVER_GENERATION, + driverInstanceId: DRIVER_INSTANCE_ID, + }); + const effect = await getExternalToolEffectForCommand(database, { + commandId: COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + }); + if (effect === null) { + throw new Error("Expected the exact-generation external effect."); + } + const message = `External effect ${effect.id} for MCP tool createIssue has an unknown outcome and will not be replayed.`; + await expect( + listAcceptedMcpCommandRepairsForTerminalDriver(database, { + driverGeneration: DRIVER_GENERATION, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).resolves.toMatchObject([ + { + command, + commandId: COMMAND_ID, + terminal: { + error: { + code: "driver.external_tool_effect_unknown", + details: { + commandId: COMMAND_ID, + effectId: effect.id, + requestId: "request-1", + runId: RUN_ID, + serverId: "01J0000000000000000000000Y", + toolName: "createIssue", + }, + message, + retryable: false, + }, + status: "failed", + }, + }, + ]); + }); + + test("stores one maximum command with one maximum settlement", async () => { + const database = await createEffectStoreFixture(); + const command = mcpExecuteCommandAtSize(RUNTIME_COMMAND_MAX_UTF8_BYTES); + const settlement = succeededSettlementAtSize( + MCP_EXTERNAL_TOOL_EFFECT_SETTLEMENT_MAX_UTF8_BYTES, + ); + expect(measureRuntimeCommandJson(command)).toBe(RUNTIME_COMMAND_MAX_UTF8_BYTES); + expect(measureMcpExternalToolEffectSettlement(settlement)).toBe( + MCP_EXTERNAL_TOOL_EFFECT_SETTLEMENT_MAX_UTF8_BYTES, + ); + + await expect( + createRuntimeCommandRecord(database, { + command: mcpExecuteCommandAtSize(RUNTIME_COMMAND_MAX_UTF8_BYTES + 1, OVERSIZED_COMMAND_ID), + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }), + ).rejects.toThrow(`${RUNTIME_COMMAND_MAX_UTF8_BYTES} UTF-8 bytes`); + await expect( + database + .prepare( + "SELECT (SELECT count(*) FROM driver_command WHERE id = ?) AS command_count, (SELECT count(*) FROM external_tool_effect WHERE command_id = ?) AS effect_count", + ) + .bind(OVERSIZED_COMMAND_ID, OVERSIZED_COMMAND_ID) + .first(), + ).resolves.toEqual({ command_count: 0, effect_count: 0 }); + + await createRuntimeCommandRecord(database, { + command, + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + const claim = await claimExternalToolEffect(database, { + claimToken: CLAIM_TOKEN, + commandId: BOUNDARY_COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + }); + if (claim.kind !== "claimed") { + throw new Error("Expected the boundary external effect to be claimed."); + } + + await expect( + settleExternalToolEffect(database, { + claimToken: CLAIM_TOKEN, + commandId: BOUNDARY_COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + effectId: claim.effectId as ExternalToolEffectId, + settlement: { + ...settlement, + result: { ...settlement.result, outputText: `${settlement.result.outputText}x` }, + }, + }), + ).rejects.toThrow(`${MCP_EXTERNAL_TOOL_EFFECT_SETTLEMENT_MAX_UTF8_BYTES} UTF-8 bytes`); + await expect( + getExternalToolEffectForCommand(database, { + commandId: BOUNDARY_COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).resolves.toMatchObject({ status: "claimed" }); + await expect( + getRuntimeCommandRecord(database, DRIVER_INSTANCE_ID, BOUNDARY_COMMAND_ID), + ).resolves.toMatchObject({ result: null, status: "accepted" }); + + await expect( + settleExternalToolEffect(database, { + claimToken: CLAIM_TOKEN, + commandId: BOUNDARY_COMMAND_ID, + driverInstanceId: DRIVER_INSTANCE_ID, + effectId: claim.effectId as ExternalToolEffectId, + settlement, + }), + ).resolves.toMatchObject({ kind: "succeeded", result: settlement.result }); + await expect( + listAcceptedMcpCommandRepairsForTerminalDriver(database, { + driverGeneration: DRIVER_GENERATION + 1, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).resolves.toEqual([]); + await expect( + listAcceptedMcpCommandRepairsForTerminalDriver(database, { + driverGeneration: DRIVER_GENERATION, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).resolves.toEqual([ + expect.objectContaining({ + command, + commandId: BOUNDARY_COMMAND_ID, + runtimeId: "openai-runtime", + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + terminal: { result: settlement.result, status: "completed" }, + }), + ]); + await expect( + database + .prepare( + `SELECT + length(CAST(command.payload_json AS BLOB)) AS payload_bytes, + effect.result_json = attempt.result_json AS effect_matches_attempt, + command.status AS command_status, + effect.status AS effect_status, + attempt.status AS attempt_status + FROM driver_command AS command + JOIN external_tool_effect AS effect ON effect.command_id = command.id + JOIN external_tool_effect_attempt AS attempt ON attempt.effect_id = effect.id + WHERE command.id = ?`, + ) + .bind(BOUNDARY_COMMAND_ID) + .first(), + ).resolves.toEqual({ + attempt_status: "succeeded", + command_status: "accepted", + effect_matches_attempt: 1, + effect_status: "succeeded", + payload_bytes: RUNTIME_COMMAND_MAX_UTF8_BYTES, + }); + }); + + test("derives accepted input completion only from its exact terminal Session Run", async () => { + const database = await createEffectStoreFixture(); + const command = { + commandId: COMMAND_ID, + input: { text: "continue" }, + kind: "input.start" as const, + requestId: "request-1", + runId: RUN_ID, + }; + await createRuntimeCommandRecord(database, { + command, + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + await database + .prepare("UPDATE session_run SET completed_at = 2, status = 'completed' WHERE id = ?") + .bind(RUN_ID) + .run(); + + await expect( + listAcceptedInputStartCommandRepairsForTerminalDriver(database, { + driverGeneration: DRIVER_GENERATION + 1, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).resolves.toEqual([]); + await expect( + listAcceptedInputStartCommandRepairsForTerminalDriver(database, { + driverGeneration: DRIVER_GENERATION, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).resolves.toEqual([ + { + command, + commandId: COMMAND_ID, + runtimeId: "openai-runtime", + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + terminal: { result: { requestId: "request-1" }, status: "completed" }, + }, + ]); + + await database + .prepare("UPDATE driver_instance SET status = 'failed' WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .run(); + await expect( + repairAcceptedRuntimeCommandsForTerminalDriver(database, { + driverGeneration: DRIVER_GENERATION, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).rejects.toThrow("event-first terminal reconciliation"); + await expect( + getRuntimeCommandRecord(database, DRIVER_INSTANCE_ID, COMMAND_ID), + ).resolves.toMatchObject({ result: null, status: "accepted" }); + }); + + test("requires an authoritative error before repairing an accepted failed input", async () => { + const database = await createEffectStoreFixture(); + const command = { + commandId: COMMAND_ID, + input: { text: "continue" }, + kind: "input.start" as const, + requestId: "request-1", + runId: RUN_ID, + }; + await createRuntimeCommandRecord(database, { + command, + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + await database + .prepare( + "UPDATE session_run SET completed_at = 2, error_code = 'driver.failed', error_details_json = ?, error_message = 'Driver failed', status = 'failed' WHERE id = ?", + ) + .bind(JSON.stringify({ driverInstanceId: DRIVER_INSTANCE_ID }), RUN_ID) + .run(); + + await expect( + listAcceptedInputStartCommandRepairsForTerminalDriver(database, { + driverGeneration: DRIVER_GENERATION, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).resolves.toMatchObject([ + { + command, + terminal: { + error: { + code: "driver.failed", + details: { driverInstanceId: DRIVER_INSTANCE_ID }, + message: "Driver failed", + retryable: false, + }, + status: "failed", + }, + }, + ]); + + await database + .prepare( + "UPDATE session_run SET error_code = NULL, error_details_json = NULL, error_message = NULL WHERE id = ?", + ) + .bind(RUN_ID) + .run(); + await expect( + listAcceptedInputStartCommandRepairsForTerminalDriver(database, { + driverGeneration: DRIVER_GENERATION, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).rejects.toThrow("missing its authoritative durable error"); + }); +}); diff --git a/apps/api/tests/file-upload-access.test.ts b/apps/api/tests/file-upload-access.test.ts index 3268ee5e..1088e844 100644 --- a/apps/api/tests/file-upload-access.test.ts +++ b/apps/api/tests/file-upload-access.test.ts @@ -74,6 +74,7 @@ function createFileUploadAccessDatabase(): SqliteD1Database { parent_path text NOT NULL, path text NOT NULL, purpose text NOT NULL, + runtime_event_seq integer, scope_id text NOT NULL, scope_kind text NOT NULL, session_kind text, diff --git a/apps/api/tests/file-upload-recovery.test.ts b/apps/api/tests/file-upload-recovery.test.ts index 4dd7a74c..ad3a8350 100644 --- a/apps/api/tests/file-upload-recovery.test.ts +++ b/apps/api/tests/file-upload-recovery.test.ts @@ -207,6 +207,7 @@ function createUploadRecoveryDatabase(): SqliteD1Database { parent_path text NOT NULL, path text NOT NULL, purpose text NOT NULL, + runtime_event_seq integer, size integer NOT NULL, updated_at integer NOT NULL, version integer NOT NULL diff --git a/apps/api/tests/helpers/api-test-fixture.ts b/apps/api/tests/helpers/api-test-fixture.ts index b07dec7e..7e259f6b 100644 --- a/apps/api/tests/helpers/api-test-fixture.ts +++ b/apps/api/tests/helpers/api-test-fixture.ts @@ -7,6 +7,7 @@ import type { AuthenticatedViewer } from "../../src/modules/auth/application/vie import { getViewer } from "../../src/modules/users/application/viewer-context.service"; import { storeVendorCredentialSecret } from "../../src/modules/vendor-credentials/application/vendor-credential.secret-resolution"; import type { ApiBindings } from "../../src/platform/cloudflare/worker-types"; +import { applyDrizzleMigrations } from "./drizzle-migrations"; import { createPublicHttpTestBindings } from "./public-api-http-test-fixture"; import { SqliteD1Database } from "./sqlite-d1"; @@ -187,9 +188,9 @@ export class ApiTestClient { } export async function createApiTestFixture(): Promise { - const database = new SqliteD1Database({ foreignKeys: false }); + const database = new SqliteD1Database(); - createApiTestSchema(database); + applyDrizzleMigrations(database); await seedApiTestFixture(database); const bindings = { @@ -255,463 +256,6 @@ export async function insertTestVendorCredential( .run(); } -function createApiTestSchema(database: SqliteD1Database): void { - database.execute(` - CREATE TABLE account ( - created_at integer NOT NULL, - email text NOT NULL, - email_verified integer NOT NULL, - id text PRIMARY KEY NOT NULL, - image_url text, - last_active_organization_id text, - name text NOT NULL, - system_agent_model text, - updated_at integer NOT NULL - ); - CREATE UNIQUE INDEX account_email_idx ON account (email); - - CREATE TABLE auth_account ( - access_token text, - access_token_expires_at integer, - provider_account_id text NOT NULL, - created_at integer NOT NULL, - id text PRIMARY KEY NOT NULL, - id_token text, - password text, - provider_id text NOT NULL, - refresh_token text, - refresh_token_expires_at integer, - scope text, - updated_at integer NOT NULL, - account_id text NOT NULL - ); - - CREATE TABLE auth_session ( - created_at integer NOT NULL, - expires_at integer NOT NULL, - id text PRIMARY KEY NOT NULL, - ip_address text, - token text NOT NULL, - updated_at integer NOT NULL, - user_agent text, - account_id text NOT NULL - ); - CREATE UNIQUE INDEX auth_session_token_idx ON auth_session (token); - - CREATE TABLE auth_verification ( - created_at integer NOT NULL, - expires_at integer NOT NULL, - id text PRIMARY KEY NOT NULL, - identifier text NOT NULL, - updated_at integer NOT NULL, - value text NOT NULL - ); - - CREATE TABLE organization ( - avatar_url text, - created_at integer NOT NULL, - creator_account_id text, - id text PRIMARY KEY NOT NULL, - name text NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE app ( - created_at integer NOT NULL, - default_environment_id text, - id text PRIMARY KEY NOT NULL, - name text NOT NULL, - organization_id text NOT NULL, - owner_account_id text NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE app_deployment ( - app_id text NOT NULL, - created_at integer NOT NULL, - default_branch text NOT NULL, - deleted_at integer, - id text PRIMARY KEY NOT NULL, - last_successful_url text, - latest_run_id text, - mosoo_subdomain text NOT NULL, - owner_account_id text NOT NULL, - repo_name text NOT NULL, - repo_owner text NOT NULL, - repo_url text NOT NULL, - source_kind text NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE app_deployment_run ( - app_id text NOT NULL, - created_at integer NOT NULL, - deployment_id text NOT NULL, - error_code text, - error_message text, - external_deployment_id text, - external_project_id text, - external_version_id text, - generated_wrangler_config_json text, - id text PRIMARY KEY NOT NULL, - mosoo_config_json text, - plan_json text, - source_branch text NOT NULL, - source_commit_sha text NOT NULL, - status text NOT NULL, - target_kind text, - target_project_name text, - target_script_name text, - updated_at integer NOT NULL, - url text - ); - - CREATE TABLE api_command ( - attempt_count integer DEFAULT 0 NOT NULL, - claim_expires_at integer, - claim_owner text, - completed_at integer, - created_at integer NOT NULL, - dedupe_key text NOT NULL, - id text PRIMARY KEY NOT NULL, - kind text NOT NULL, - last_error_code text, - last_error_message text, - payload_json text NOT NULL, - status text NOT NULL, - updated_at integer NOT NULL - ); - CREATE UNIQUE INDEX api_command_dedupe_idx ON api_command (dedupe_key); - - CREATE TABLE agent ( - config_json text NOT NULL, - created_at integer NOT NULL, - description text, - environment_id text, - id text PRIMARY KEY NOT NULL, - kind text DEFAULT 'pet' NOT NULL, - live_deployment_version_id text, - model text NOT NULL, - name text NOT NULL, - owner_account_id text NOT NULL, - app_id text NOT NULL, - prompt text NOT NULL, - provider text NOT NULL, - runtime_id text NOT NULL, - status text DEFAULT 'draft' NOT NULL, - updated_at integer NOT NULL, - visibility text DEFAULT 'private' NOT NULL - ); - - CREATE TABLE environment ( - created_at integer NOT NULL, - current_revision_id text NOT NULL, - description text NOT NULL, - forked_from_environment_id text, - forked_from_environment_name text, - forked_from_owner_name text, - id text PRIMARY KEY NOT NULL, - name text NOT NULL, - owner_account_id text, - app_id text NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE environment_revision ( - allow_mcp_servers integer NOT NULL, - allow_package_managers integer NOT NULL, - allowed_hosts_json text NOT NULL, - created_at integer NOT NULL, - created_by_account_id text, - env_vars_json text NOT NULL, - environment_id text NOT NULL, - id text PRIMARY KEY NOT NULL, - network_policy text NOT NULL, - packages_json text NOT NULL, - app_id text NOT NULL, - setup_script text NOT NULL - ); - - CREATE TABLE mcp_server ( - auth_type text NOT NULL, - byo_client_id text, - byo_client_secret_secret_id text, - created_at integer NOT NULL, - credential_scope text NOT NULL, - description text, - enabled integer DEFAULT true NOT NULL, - icon_url text, - id text PRIMARY KEY NOT NULL, - name text NOT NULL, - oauth_metadata_json text, - owner_account_id text NOT NULL, - app_id text NOT NULL, - source text NOT NULL, - updated_at integer NOT NULL, - url text NOT NULL - ); - - CREATE TABLE skill ( - author text NOT NULL, - created_at integer NOT NULL, - current_snapshot_id text NOT NULL, - description text NOT NULL, - forked_from_owner_name text, - forked_from_skill_id text, - forked_from_skill_name text, - id text PRIMARY KEY NOT NULL, - name text NOT NULL, - owner_account_id text NOT NULL, - app_id text NOT NULL, - source_kind text NOT NULL, - updated_at integer NOT NULL, - version text - ); - - CREATE TABLE file_record ( - committed integer NOT NULL, - created_at integer NOT NULL, - created_by_account_id text NOT NULL, - etag text, - expires_at integer, - id text PRIMARY KEY NOT NULL, - mime_type text, - name text NOT NULL, - object_key text NOT NULL, - owner_id text NOT NULL, - owner_kind text NOT NULL, - parent_path text NOT NULL, - path text NOT NULL, - purpose text NOT NULL, - scope_id text, - scope_kind text NOT NULL, - session_kind text, - size integer NOT NULL, - status text NOT NULL, - updated_at integer NOT NULL, - version integer NOT NULL - ); - - CREATE TABLE file_upload ( - content_type text NOT NULL, - created_at integer NOT NULL, - created_by_account_id text NOT NULL, - expected_size integer NOT NULL, - expires_at integer NOT NULL, - file_id text NOT NULL, - id text PRIMARY KEY NOT NULL, - if_match_etag text, - multipart_upload_id text, - overwrite integer NOT NULL, - part_size integer, - scope_id text, - scope_kind text NOT NULL, - status text NOT NULL, - strategy text NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE agent_deployment_version ( - agent_id text NOT NULL, - config_json text NOT NULL, - created_at integer NOT NULL, - created_by_account_id text NOT NULL, - environment_id text, - id text PRIMARY KEY NOT NULL, - kind text NOT NULL, - mcp_bindings_json text NOT NULL, - model text NOT NULL, - prompt text NOT NULL, - provider text NOT NULL, - runtime_id text NOT NULL, - skills_json text NOT NULL, - summary text NOT NULL, - version_number integer NOT NULL - ); - - CREATE TABLE agent_skill ( - agent_id text NOT NULL, - created_at integer NOT NULL, - skill_id text NOT NULL, - sort_order integer NOT NULL, - PRIMARY KEY (agent_id, skill_id) - ); - - CREATE TABLE session ( - agent_id text NOT NULL, - archived_at integer, - attributed_user_id text, - created_at integer NOT NULL, - creator_account_id text NOT NULL, - deployment_version_id text, - deployment_version_number integer, - id text PRIMARY KEY NOT NULL, - kind text NOT NULL, - last_message_at integer, - last_run_id text, - message_seq_cursor integer DEFAULT 0 NOT NULL, - metadata_json text DEFAULT '{}' NOT NULL, - model text NOT NULL, - app_id text NOT NULL, - provider text NOT NULL, - renamed integer NOT NULL, - runtime_id text NOT NULL, - status text NOT NULL, - status_operation_id text, - status_seq integer DEFAULT 0 NOT NULL, - runtime_event_seq_cursor integer DEFAULT 0 NOT NULL, - workspace_checkpoint_required integer DEFAULT 0 NOT NULL, - title text, - type text DEFAULT 'preview' NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE session_run ( - agent_id text NOT NULL, - completed_at integer, - created_at integer NOT NULL, - created_by_account_id text NOT NULL, - deployment_version_id text, - deployment_version_number integer, - driver_instance_id text, - error_code text, - error_details_json text, - error_message text, - id text PRIMARY KEY NOT NULL, - model text, - provider text, - runtime_id text, - session_id text NOT NULL, - started_at integer, - status text NOT NULL, - status_changed_at integer DEFAULT 0 NOT NULL, - status_event text DEFAULT 'run.queue' NOT NULL, - status_operation_id text, - status_seq integer DEFAULT 0 NOT NULL, - status_source text DEFAULT 'system' NOT NULL, - trace_id text NOT NULL, - trigger text NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE sandbox_session ( - cloudflare_session_id text NOT NULL, - created_at integer NOT NULL, - cwd text NOT NULL, - origin_json text NOT NULL, - sandbox_id text NOT NULL, - session_id text PRIMARY KEY NOT NULL, - status text NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE sandbox_backup ( - created_at integer NOT NULL, - dir text NOT NULL, - error_message text, - id text PRIMARY KEY NOT NULL, - keep integer DEFAULT 0 NOT NULL, - sandbox_id text NOT NULL, - session_run_id text, - status text NOT NULL, - ttl_seconds integer NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE driver_instance ( - boot_token_expires_at integer NOT NULL, - boot_token_hash blob NOT NULL, - boot_token_used_at integer, - close_code integer, - close_reason text, - connection_id text, - created_at integer NOT NULL, - command_seq_cursor integer DEFAULT 0 NOT NULL, - driver_pid integer, - driver_started_at integer, - driver_version text, - error_message text, - expires_at integer NOT NULL, - heartbeat_count integer NOT NULL, - generation integer DEFAULT 0 NOT NULL, - id text PRIMARY KEY NOT NULL, - last_heartbeat_at integer, - process_id text, - protocol text NOT NULL, - protocol_version integer NOT NULL, - restart_count integer DEFAULT 0 NOT NULL, - runtime text NOT NULL, - sandbox_id text NOT NULL, - sandbox_session_id text NOT NULL, - status text NOT NULL, - status_changed_at integer DEFAULT 0 NOT NULL, - status_event text DEFAULT 'driver.provision' NOT NULL, - status_operation_id text, - status_seq integer DEFAULT 0 NOT NULL, - status_source text DEFAULT 'system' NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE agent_mcp_binding ( - agent_credential_id text, - agent_id text NOT NULL, - created_at integer DEFAULT 1 NOT NULL, - credential_mode text DEFAULT 'runtime_resolved' NOT NULL, - enabled integer DEFAULT 1 NOT NULL, - id text PRIMARY KEY DEFAULT '01J00000000000000000000999' NOT NULL, - server_id text NOT NULL, - sort_order integer DEFAULT 0 NOT NULL, - updated_at integer DEFAULT 1 NOT NULL - ); - - CREATE TABLE mcp_credential ( - account_id text, - agent_id text, - auth_type text NOT NULL, - created_at integer NOT NULL, - expires_at integer, - id text PRIMARY KEY NOT NULL, - last_refreshed_at integer, - oauth_client_id text, - oauth_client_secret_secret_id text, - app_id text NOT NULL, - refresh_secret_id text, - scope text NOT NULL, - scope_values_json text, - secret_id text NOT NULL, - server_id text NOT NULL, - status text NOT NULL, - subject_label text, - updated_at integer NOT NULL - ); - - CREATE TABLE vendor_credential ( - api_base text, - api_key_secret_id text NOT NULL, - created_at integer NOT NULL, - id text PRIMARY KEY NOT NULL, - is_default integer DEFAULT false NOT NULL, - models text, - name text NOT NULL, - app_id text NOT NULL, - updated_at integer NOT NULL, - vendor_id text NOT NULL - ); - - CREATE TABLE vault_secret ( - algorithm text NOT NULL DEFAULT 'AES-GCM', - ciphertext text NOT NULL, - ciphertext_iv text NOT NULL, - created_at integer NOT NULL, - id text PRIMARY KEY NOT NULL, - kind text NOT NULL, - updated_at integer NOT NULL, - wrapped_dek text NOT NULL, - wrapped_dek_iv text NOT NULL - ); - `); -} - async function seedApiTestFixture(database: D1Database): Promise { await database .prepare( @@ -824,7 +368,7 @@ async function seedApiTestFixture(database: D1Database): Promise { "[]", API_TEST_IDS.environmentId, API_TEST_IDS.environmentRevisionId, - "sandbox", + "full", "[]", API_TEST_IDS.appId, "", diff --git a/apps/api/tests/helpers/drizzle-migrations.ts b/apps/api/tests/helpers/drizzle-migrations.ts new file mode 100644 index 00000000..f398fcf4 --- /dev/null +++ b/apps/api/tests/helpers/drizzle-migrations.ts @@ -0,0 +1 @@ +export * from "../../../../pkgs/db/scripts/drizzle-migrations"; diff --git a/apps/api/tests/helpers/public-api-http-core-schema.sql b/apps/api/tests/helpers/public-api-http-core-schema.sql deleted file mode 100644 index 5eebffcc..00000000 --- a/apps/api/tests/helpers/public-api-http-core-schema.sql +++ /dev/null @@ -1,241 +0,0 @@ -CREATE TABLE account ( - id text PRIMARY KEY NOT NULL, - email text NOT NULL, - email_verified integer NOT NULL, - image_url text, - last_active_organization_id text, - name text NOT NULL, - system_agent_model text, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE personal_access_token ( - id text PRIMARY KEY NOT NULL, - account_id text NOT NULL, - label text NOT NULL, - token_hash text NOT NULL, - created_at integer NOT NULL, - updated_at integer NOT NULL, - last_used_at integer, - revoked_at integer -); - -CREATE TABLE organization ( - id text PRIMARY KEY NOT NULL, - name text NOT NULL, - avatar_url text, - creator_account_id text, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE app ( - id text PRIMARY KEY NOT NULL, - organization_id text NOT NULL, - owner_account_id text NOT NULL, - name text NOT NULL, - default_environment_id text, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE agent ( - id text PRIMARY KEY NOT NULL, - app_id text NOT NULL, - owner_account_id text NOT NULL, - name text NOT NULL, - description text, - environment_id text, - live_deployment_version_id text, - kind text NOT NULL, - runtime_id text NOT NULL, - provider text NOT NULL, - model text NOT NULL, - prompt text NOT NULL, - created_at integer NOT NULL, - updated_at integer NOT NULL, - config_json text NOT NULL, - status text NOT NULL, - visibility text NOT NULL -); - -CREATE TABLE agent_deployment_version ( - id text PRIMARY KEY NOT NULL, - agent_id text NOT NULL, - version_number integer NOT NULL, - summary text NOT NULL, - kind text NOT NULL, - runtime_id text NOT NULL, - provider text NOT NULL, - model text NOT NULL, - prompt text NOT NULL, - config_json text NOT NULL, - environment_id text, - skills_json text NOT NULL, - mcp_bindings_json text NOT NULL, - created_by_account_id text NOT NULL, - created_at integer NOT NULL -); - -CREATE TABLE skill ( - author text NOT NULL, - created_at integer NOT NULL, - current_snapshot_id text NOT NULL, - description text NOT NULL, - forked_from_owner_name text, - forked_from_skill_id text, - forked_from_skill_name text, - id text PRIMARY KEY NOT NULL, - name text NOT NULL, - owner_account_id text NOT NULL, - app_id text NOT NULL, - source_kind text NOT NULL, - updated_at integer NOT NULL, - version text -); - -CREATE TABLE agent_skill ( - agent_id text NOT NULL, - skill_id text NOT NULL, - sort_order integer NOT NULL, - created_at integer NOT NULL, - PRIMARY KEY (agent_id, skill_id) -); - -CREATE TABLE mcp_server ( - id text PRIMARY KEY NOT NULL, - owner_account_id text NOT NULL, - name text NOT NULL, - description text, - source text DEFAULT 'app' NOT NULL, - auth_type text DEFAULT 'bearer' NOT NULL, - credential_scope text DEFAULT 'app' NOT NULL, - icon_url text, - byo_client_id text, - byo_client_secret_secret_id text, - oauth_metadata_json text, - app_id text NOT NULL, - url text NOT NULL, - enabled integer DEFAULT 1 NOT NULL, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE agent_mcp_binding ( - id text PRIMARY KEY NOT NULL, - agent_id text NOT NULL, - server_id text NOT NULL, - agent_credential_id text, - credential_mode text DEFAULT 'runtime_resolved' NOT NULL, - enabled integer DEFAULT 1 NOT NULL, - sort_order integer DEFAULT 0 NOT NULL, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE environment ( - id text PRIMARY KEY NOT NULL, - app_id text NOT NULL, - owner_account_id text, - current_revision_id text NOT NULL, - name text NOT NULL, - description text NOT NULL, - forked_from_environment_id text, - forked_from_environment_name text, - forked_from_owner_name text, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE environment_revision ( - id text PRIMARY KEY NOT NULL, - environment_id text NOT NULL, - app_id text NOT NULL, - created_by_account_id text, - setup_script text NOT NULL, - packages_json text NOT NULL, - env_vars_json text NOT NULL, - network_policy text NOT NULL, - allowed_hosts_json text NOT NULL, - allow_mcp_servers integer NOT NULL, - allow_package_managers integer NOT NULL, - created_at integer NOT NULL -); - -CREATE TABLE vendor_credential ( - id text PRIMARY KEY NOT NULL, - app_id text NOT NULL, - vendor_id text NOT NULL, - name text NOT NULL, - api_key_secret_id text NOT NULL, - api_base text, - is_default integer DEFAULT false NOT NULL, - models text, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE sandbox ( - id text PRIMARY KEY NOT NULL, - agent_id text, - app_id text, - kind text NOT NULL, - subject_kind text NOT NULL, - subject_id text NOT NULL, - status text NOT NULL, - status_event text DEFAULT 'runtime_subject.cold' NOT NULL, - status_source text DEFAULT 'system' NOT NULL, - status_seq integer DEFAULT 0 NOT NULL, - status_operation_id text, - status_changed_at integer DEFAULT 0 NOT NULL, - last_error text, - last_error_code text, - last_backup_id text, - last_restore_backup_id text, - bind_mount_ready integer DEFAULT 0 NOT NULL, - global_mounts_json text DEFAULT '[]' NOT NULL, - claim_owner text, - claim_expires_at integer, - inactive_deadline_at integer, - owner_account_id text, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE vault_secret ( - id text PRIMARY KEY NOT NULL, - kind text NOT NULL, - algorithm text DEFAULT 'AES-GCM' NOT NULL, - ciphertext text NOT NULL, - ciphertext_iv text NOT NULL, - wrapped_dek text NOT NULL, - wrapped_dek_iv text NOT NULL, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE api_command ( - id text PRIMARY KEY NOT NULL, - kind text NOT NULL, - dedupe_key text NOT NULL, - payload_json text NOT NULL, - status text NOT NULL, - attempt_count integer DEFAULT 0 NOT NULL, - claim_owner text, - claim_expires_at integer, - last_error_code text, - last_error_message text, - completed_at integer, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE UNIQUE INDEX api_command_dedupe_idx - ON api_command (dedupe_key); - -CREATE INDEX api_command_status_updated_idx - ON api_command (status, updated_at); - -CREATE INDEX api_command_claim_idx - ON api_command (status, claim_expires_at); diff --git a/apps/api/tests/helpers/public-api-http-runtime-schema.sql b/apps/api/tests/helpers/public-api-http-runtime-schema.sql deleted file mode 100644 index b95294ec..00000000 --- a/apps/api/tests/helpers/public-api-http-runtime-schema.sql +++ /dev/null @@ -1,342 +0,0 @@ -CREATE TABLE public_api_rate_limit_window ( - bucket_key text NOT NULL, - request_count integer DEFAULT 0 NOT NULL, - shard integer NOT NULL, - updated_at integer NOT NULL, - window_start integer NOT NULL, - PRIMARY KEY (bucket_key, window_start, shard) -); - -CREATE TABLE public_api_idempotency_key ( - id text PRIMARY KEY NOT NULL, - token_id text NOT NULL, - idempotency_key text NOT NULL, - method text NOT NULL, - route text NOT NULL, - body_hash text, - response_status integer, - response_json text, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE bound_agent_call_idempotency_key ( - id text PRIMARY KEY NOT NULL, - subject_hash text NOT NULL, - idempotency_key text NOT NULL, - body_hash text NOT NULL, - session_id text NOT NULL, - run_id text, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE UNIQUE INDEX bound_agent_call_idempotency_subject_key_idx - ON bound_agent_call_idempotency_key (subject_hash, idempotency_key); - -CREATE INDEX bound_agent_call_idempotency_updated_idx - ON bound_agent_call_idempotency_key (updated_at); - -CREATE TABLE session ( - id text PRIMARY KEY NOT NULL, - app_id text NOT NULL, - creator_account_id text NOT NULL, - attributed_user_id text, - end_user_id text, - agent_id text NOT NULL, - deployment_version_id text, - deployment_version_number integer, - kind text NOT NULL, - title text, - provider text NOT NULL, - model text NOT NULL, - runtime_id text NOT NULL, - status text NOT NULL, - type text DEFAULT 'ui' NOT NULL, - metadata_json text DEFAULT '{}' NOT NULL, - last_run_id text, - last_message_at integer, - message_seq_cursor integer DEFAULT 0 NOT NULL, - runtime_event_seq_cursor integer DEFAULT 0 NOT NULL, - workspace_checkpoint_required integer DEFAULT 0 NOT NULL, - archived_at integer, - renamed integer DEFAULT 0 NOT NULL, - status_operation_id text, - status_seq integer DEFAULT 0 NOT NULL, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE session_run ( - id text PRIMARY KEY NOT NULL, - session_id text NOT NULL, - agent_id text NOT NULL, - bound_capability_agent_id text, - bound_capability_app_id text, - bound_capability_binding_env text, - bound_capability_binding_name text, - bound_capability_deployment_id text, - bound_capability_deployment_run_id text, - created_by_account_id text NOT NULL, - deployment_version_id text, - deployment_version_number integer, - driver_instance_id text, - trigger text NOT NULL, - status text NOT NULL, - provider text, - model text, - runtime_id text, - trace_id text, - error_code text, - error_message text, - error_details_json text, - started_at integer, - completed_at integer, - created_at integer, - status_changed_at integer DEFAULT 0 NOT NULL, - status_event text DEFAULT 'run.queue' NOT NULL, - status_operation_id text, - status_seq integer DEFAULT 0 NOT NULL, - status_source text DEFAULT 'system' NOT NULL, - updated_at integer -); - -CREATE TABLE session_run_skill ( - session_run_id text NOT NULL, - skill_id text NOT NULL, - skill_name text NOT NULL, - snapshot_id text, - blob_sha256 text, - mount_path text NOT NULL, - resolution_mode text NOT NULL, - materialization_status text NOT NULL, - warning_code text, - created_at integer NOT NULL, - updated_at integer NOT NULL, - PRIMARY KEY (session_run_id, skill_id) -); - -CREATE TABLE session_execution_snapshot ( - session_id text PRIMARY KEY NOT NULL, - plan_json text NOT NULL, - created_at integer NOT NULL -); - -CREATE TABLE session_message ( - id text PRIMARY KEY NOT NULL, - session_id text NOT NULL, - session_run_id text, - seq integer NOT NULL, - role text NOT NULL, - content_text text NOT NULL, - segments_json text, - plan_json text, - created_by_account_id text NOT NULL, - created_at integer NOT NULL -); - -CREATE TABLE session_event ( - id text PRIMARY KEY NOT NULL, - session_id text NOT NULL, - run_id text, - agent_id text NOT NULL, - seq integer NOT NULL, - content_text text NOT NULL, - ended_at integer NOT NULL, - event_type text NOT NULL, - family text NOT NULL, - process_status text NOT NULL, - process_type text NOT NULL, - source text NOT NULL, - source_event_id text NOT NULL, - tool_call_id text, - tool_input_json text, - tool_name text, - tokens integer, - trace_id text, - visibility text NOT NULL, - occurred_at integer NOT NULL, - created_at integer NOT NULL -); - -CREATE TABLE session_permission_request ( - created_at integer NOT NULL, - driver_instance_id text NOT NULL, - raw_input text, - request_id text NOT NULL, - run_id text NOT NULL, - session_id text NOT NULL, - title text NOT NULL, - tool_call_id text, - tool_kind text, - updated_at integer NOT NULL, - PRIMARY KEY (session_id, request_id) -); - -CREATE INDEX session_permission_request_run_idx - ON session_permission_request (session_id, run_id); - -CREATE TABLE session_readiness_snapshot ( - readiness_json text NOT NULL, - session_id text PRIMARY KEY NOT NULL, - updated_at integer NOT NULL -); - -CREATE UNIQUE INDEX session_message_session_seq_idx - ON session_message (session_id, seq); -CREATE UNIQUE INDEX session_event_session_seq_idx - ON session_event (session_id, seq); -CREATE UNIQUE INDEX session_event_session_source_idx - ON session_event (session_id, source_event_id); - -CREATE TABLE sandbox_session ( - cloudflare_session_id text NOT NULL, - created_at integer NOT NULL, - cwd text NOT NULL, - origin_json text NOT NULL, - sandbox_id text NOT NULL, - session_id text PRIMARY KEY NOT NULL, - status text NOT NULL, - updated_at integer NOT NULL -); - -CREATE INDEX sandbox_session_sandbox_status_idx - ON sandbox_session (sandbox_id, status, updated_at); -CREATE UNIQUE INDEX sandbox_session_cloudflare_session_idx - ON sandbox_session (cloudflare_session_id); - -CREATE TABLE sandbox_backup ( - created_at integer NOT NULL, - dir text NOT NULL, - error_message text, - id text PRIMARY KEY NOT NULL, - keep integer DEFAULT 0 NOT NULL, - sandbox_id text NOT NULL, - session_run_id text, - status text NOT NULL, - ttl_seconds integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE INDEX sandbox_backup_sandbox_status_created_idx - ON sandbox_backup (sandbox_id, status, created_at); - -CREATE UNIQUE INDEX sandbox_backup_terminal_checkpoint_idx - ON sandbox_backup (sandbox_id, dir, session_run_id) - WHERE session_run_id IS NOT NULL AND status = 'ready'; - -CREATE TABLE driver_instance ( - id text PRIMARY KEY NOT NULL, - sandbox_id text NOT NULL, - sandbox_session_id text NOT NULL, - runtime text NOT NULL, - protocol text NOT NULL, - protocol_version integer NOT NULL, - status text NOT NULL, - status_changed_at integer DEFAULT 0 NOT NULL, - status_event text DEFAULT 'driver.provision' NOT NULL, - status_operation_id text, - status_seq integer DEFAULT 0 NOT NULL, - status_source text DEFAULT 'system' NOT NULL, - process_id text, - connection_id text, - command_seq_cursor integer DEFAULT 0 NOT NULL, - boot_token_hash blob NOT NULL, - boot_token_expires_at integer NOT NULL, - boot_token_used_at integer, - driver_pid integer, - driver_started_at integer, - driver_version text, - close_code integer, - close_reason text, - error_message text, - generation integer DEFAULT 0 NOT NULL, - heartbeat_count integer NOT NULL, - last_heartbeat_at integer, - restart_count integer DEFAULT 0 NOT NULL, - expires_at integer NOT NULL, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE external_tool_effect ( - attempt_count integer DEFAULT 0 NOT NULL, - command_id text NOT NULL, - created_at integer NOT NULL, - driver_instance_id text NOT NULL, - id text PRIMARY KEY NOT NULL, - idempotency_key text NOT NULL, - provider_receipt_json text, - result_json text, - server_id text NOT NULL, - session_run_id text NOT NULL, - status text NOT NULL, - tool_name text NOT NULL, - updated_at integer NOT NULL -); - -CREATE UNIQUE INDEX external_tool_effect_command_idx - ON external_tool_effect (command_id); -CREATE UNIQUE INDEX external_tool_effect_idempotency_key_idx - ON external_tool_effect (idempotency_key); -CREATE INDEX external_tool_effect_run_status_idx - ON external_tool_effect (session_run_id, status, id); -CREATE INDEX external_tool_effect_driver_status_idx - ON external_tool_effect (driver_instance_id, status); - -CREATE TABLE external_tool_effect_attempt ( - attempt integer NOT NULL, - completed_at integer, - created_at integer NOT NULL, - effect_id text NOT NULL, - provider_receipt_json text, - result_json text, - status text NOT NULL, - PRIMARY KEY (effect_id, attempt) -); - -CREATE INDEX external_tool_effect_attempt_status_idx - ON external_tool_effect_attempt (status, created_at); - -CREATE TABLE file_record ( - id text PRIMARY KEY NOT NULL, - scope_kind text NOT NULL, - scope_id text, - session_kind text, - status text NOT NULL, - name text NOT NULL, - path text NOT NULL, - parent_path text NOT NULL, - object_key text NOT NULL, - owner_id text NOT NULL, - owner_kind text NOT NULL, - purpose text NOT NULL, - expires_at integer, - mime_type text, - size integer NOT NULL, - etag text, - committed integer NOT NULL, - version integer NOT NULL, - created_by_account_id text NOT NULL, - created_at integer NOT NULL, - updated_at integer NOT NULL -); - -CREATE TABLE file_upload ( - id text PRIMARY KEY NOT NULL, - file_id text NOT NULL, - scope_kind text NOT NULL, - scope_id text NOT NULL, - strategy text NOT NULL, - status text NOT NULL, - content_type text NOT NULL, - expected_size integer NOT NULL, - overwrite integer NOT NULL, - if_match_etag text, - multipart_upload_id text, - part_size integer, - created_by_account_id text NOT NULL, - expires_at integer NOT NULL, - created_at integer NOT NULL, - updated_at integer NOT NULL -); diff --git a/apps/api/tests/helpers/public-api-http-test-fixture.ts b/apps/api/tests/helpers/public-api-http-test-fixture.ts index 18fc2bd2..7222a4b9 100644 --- a/apps/api/tests/helpers/public-api-http-test-fixture.ts +++ b/apps/api/tests/helpers/public-api-http-test-fixture.ts @@ -1,5 +1,3 @@ -import { readFileSync } from "node:fs"; - import type { SessionSummary } from "@mosoo/contracts/session"; import { accountsTable, @@ -9,18 +7,34 @@ import { environmentsTable, organizationsTable, personalAccessTokensTable, + sandboxSessionsTable, + sandboxesTable, appsTable, sessionExecutionSnapshotsTable, sessionsTable, vendorCredentialsTable, } from "@mosoo/db"; -import type { VendorCredentialId } from "@mosoo/id"; +import { parsePlatformId } from "@mosoo/id"; +import type { + AccountId, + AgentId, + AppId, + SandboxId, + SandboxSessionId, + SessionId, + VendorCredentialId, +} from "@mosoo/id"; import { hashTokenValue } from "../../src/modules/auth/application/personal-access-token.service"; import { storeVendorCredentialSecret } from "../../src/modules/vendor-credentials/application/vendor-credential.secret-resolution"; import type { ApiBindings } from "../../src/platform/cloudflare/worker-types"; import type { ApiCommandQueueStub } from "./api-command-queue-fixture"; import { createApiCommandQueueStub } from "./api-command-queue-fixture"; +import { + applyDrizzleMigrations, + applyDrizzleMigrationsFrom, + applyDrizzleMigrationsThrough, +} from "./drizzle-migrations"; import { SqliteD1Database } from "./sqlite-d1"; export { SqliteD1Database } from "./sqlite-d1"; export { @@ -32,13 +46,6 @@ export { type RecordedQueueMessageAction, } from "./api-command-queue-fixture"; -const CONTRACT_SCHEMA_SQL = readFileSync( - new URL("./public-api-http-core-schema.sql", import.meta.url), - "utf8", -) - .concat("\n") - .concat(readFileSync(new URL("./public-api-http-runtime-schema.sql", import.meta.url), "utf8")); - const INITIAL_AGENT_CONFIG_JSON = JSON.stringify({ packageMcpServers: [], packageResolution: null, @@ -75,7 +82,7 @@ export const PUBLIC_API_TEST_IDS = { driverOwner: "01J0000000000000000000000F", } as const; -const PUBLIC_API_VENDOR_CREDENTIAL_ID = "vendor-openai-app" as VendorCredentialId; +const PUBLIC_API_VENDOR_CREDENTIAL_ID = "01J0000000000000000000000S" as VendorCredentialId; export function createTestExecutionContext(): ExecutionContext { return { @@ -100,8 +107,9 @@ export function nowMsForTest(): number { } interface StoredObject { - body: string; + body: Uint8Array; contentType: string; + customMetadata: Record; etag: string; key: string; } @@ -129,6 +137,7 @@ export class PublicApiMemoryFileBucket { body: ReadableStream | ArrayBuffer | ArrayBufferView | string | Blob | null, options?: R2PutOptions, ): Promise { + const bytes = await this.#readBody(body); const existing = this.objects.get(key); const ifNoneMatch = this.#readOnlyIfHeader(options?.onlyIf, "If-None-Match"); const ifMatch = this.#readOnlyIfHeader(options?.onlyIf, "If-Match"); @@ -142,8 +151,9 @@ export class PublicApiMemoryFileBucket { } const stored: StoredObject = { - body: await this.#readBody(body), + body: bytes, contentType: options?.httpMetadata?.contentType ?? "application/octet-stream", + customMetadata: { ...options?.customMetadata }, etag: this.#createEtag(), key, }; @@ -160,27 +170,25 @@ export class PublicApiMemoryFileBucket { async #readBody( body: ReadableStream | ArrayBuffer | ArrayBufferView | string | Blob | null, - ): Promise { + ): Promise { if (body === null) { - return ""; + return new Uint8Array(); } if (typeof body === "string") { - return body; + return new TextEncoder().encode(body); } if (body instanceof Blob) { - return body.text(); + return new Uint8Array(await body.arrayBuffer()); } if (body instanceof ArrayBuffer) { - return new TextDecoder().decode(body); + return new Uint8Array(body).slice(); } if (ArrayBuffer.isView(body)) { - return new TextDecoder().decode( - new Uint8Array(body.buffer, body.byteOffset, body.byteLength), - ); + return new Uint8Array(body.buffer, body.byteOffset, body.byteLength).slice(); } const chunks: Uint8Array[] = []; @@ -205,7 +213,7 @@ export class PublicApiMemoryFileBucket { offset += chunk.byteLength; } - return new TextDecoder().decode(bytes); + return bytes; } #readOnlyIfHeader(onlyIf: R2PutOptions["onlyIf"] | undefined, name: string): string | null { @@ -214,14 +222,14 @@ export class PublicApiMemoryFileBucket { #toObject(stored: StoredObject): R2Object { return { - customMetadata: {}, + customMetadata: { ...stored.customMetadata }, etag: stored.etag, httpEtag: `"${stored.etag}"`, httpMetadata: { contentType: stored.contentType, }, key: stored.key, - size: new TextEncoder().encode(stored.body).byteLength, + size: stored.body.byteLength, uploaded: new Date(0), version: "", writeHttpMetadata(headers: Headers) { @@ -231,7 +239,7 @@ export class PublicApiMemoryFileBucket { } #toObjectBody(stored: StoredObject): R2ObjectBody { - const bytes = new TextEncoder().encode(stored.body); + const bytes = stored.body.slice(); return { ...this.#toObject(stored), @@ -249,10 +257,10 @@ export class PublicApiMemoryFileBucket { }), bodyUsed: false, async json() { - return JSON.parse(stored.body) as T; + return JSON.parse(new TextDecoder().decode(bytes)) as T; }, async text() { - return stored.body; + return new TextDecoder().decode(bytes); }, } as R2ObjectBody; } @@ -289,11 +297,14 @@ export function createPublicHttpTestBindings( }; } -export async function createPublicHttpContractDatabase(): Promise { - const database = new SqliteD1Database(); - const nowMs = nowMsForTest(); +export function migratePre0014PublicHttpContractDatabase(database: SqliteD1Database): void { + applyDrizzleMigrationsFrom(database, "0014_session-event-stream-identity"); +} - database.execute(CONTRACT_SCHEMA_SQL); +async function seedPublicHttpContractDatabase( + database: SqliteD1Database, +): Promise { + const nowMs = nowMsForTest(); const db = database.app(); await db @@ -490,6 +501,94 @@ export async function createPublicHttpContractDatabase(): Promise { + const database = new SqliteD1Database(); + applyDrizzleMigrations(database); + await seedPublicHttpContractDatabase(database); + return database.serialize(); +})(); + +export async function createPublicHttpContractDatabase(): Promise { + return new SqliteD1Database({ serialized: await publicHttpContractDatabaseTemplate }); +} + +export async function createPre0014PublicHttpContractDatabase(): Promise { + const database = new SqliteD1Database(); + applyDrizzleMigrationsThrough(database, "0013_durable-mcp-effect-v3"); + return seedPublicHttpContractDatabase(database); +} + +export async function insertActiveSandboxSessionFixture( + database: SqliteD1Database, + input: { + agentId?: string; + appId?: string; + cwd?: string; + inactiveDeadlineAt?: number | null; + kind?: "cattle" | "pet"; + ownerAccountId: string; + sandboxId: string; + sandboxSessionId?: string; + sessionId: string; + timestampMs?: number; + }, +): Promise { + const agentId = parsePlatformId( + input.agentId ?? PUBLIC_API_TEST_IDS.agent, + "fixture agent id", + ); + const appId = parsePlatformId(input.appId ?? PUBLIC_API_TEST_IDS.app, "fixture app id"); + const ownerAccountId = parsePlatformId(input.ownerAccountId, "fixture account id"); + const sandboxId = parsePlatformId(input.sandboxId, "fixture sandbox id"); + const sandboxSessionId = parsePlatformId( + input.sandboxSessionId ?? input.sandboxId, + "fixture sandbox session id", + ); + const sessionId = parsePlatformId(input.sessionId, "fixture session id"); + const kind = input.kind ?? "pet"; + const timestampMs = input.timestampMs ?? nowMsForTest(); + + await database + .app() + .insert(sandboxesTable) + .values({ + agentId, + appId, + createdAt: timestampMs, + id: sandboxId, + inactiveDeadlineAt: input.inactiveDeadlineAt ?? null, + incarnation: 1, + kind, + networkConstraintsHash: "0".repeat(64), + ownerAccountId, + status: "active", + subjectId: kind === "pet" ? agentId : sessionId, + subjectKind: kind === "pet" ? "agent" : "session", + updatedAt: timestampMs, + }) + .run(); + await database + .app() + .insert(sandboxSessionsTable) + .values({ + createdAt: timestampMs, + cwd: input.cwd ?? `/workspace/se/${sessionId}`, + originJson: JSON.stringify({ + callerUserId: ownerAccountId, + entrypoint: "api", + executionOwnerUserId: ownerAccountId, + type: "agent", + }), + sandboxId, + sandboxIncarnation: 1, + sandboxSessionId, + sessionId, + status: "active", + updatedAt: timestampMs, + }) + .run(); +} + export async function insertNonOwnerSession(database: SqliteD1Database): Promise { await insertSession(database, { creatorAccountId: PUBLIC_API_TEST_IDS.nonOwnerAccount, @@ -622,6 +721,7 @@ function createOkDurableObjectNamespace() { destroy: async () => {}, fetch: async () => new Response(null, { status: 204 }), publishEvents: async () => {}, + syncViewers: async () => {}, }), idFromName: (name: string) => name, }; diff --git a/apps/api/tests/helpers/runtime-output-sandbox.ts b/apps/api/tests/helpers/runtime-output-sandbox.ts new file mode 100644 index 00000000..617759ef --- /dev/null +++ b/apps/api/tests/helpers/runtime-output-sandbox.ts @@ -0,0 +1,104 @@ +import type { SandboxHandle } from "../../src/modules/runtime/infrastructure/sandbox-handles"; + +const encoder = new TextEncoder(); + +export interface RuntimeOutputSandboxOptions { + readonly files?: ReadonlyMap; + readonly fileSizes?: ReadonlyMap; + readonly listError?: Error; + readonly onExec?: (command: string) => void; + readonly onRead?: (path: string) => void; + readonly readError?: Error; + readonly root?: string; +} + +export function createRuntimeOutputSandbox( + options: RuntimeOutputSandboxOptions = {}, +): SandboxHandle { + const root = (options.root ?? "/workspace/outputs").replace(/\/+$/, ""); + const entries = () => + [...(options.files ?? [])] + .map(([path, content]) => { + const absolutePath = path.startsWith("/") ? path : `${root}/${path}`; + return { + absolutePath, + bytes: typeof content === "string" ? encoder.encode(content) : content, + path, + relativePath: absolutePath.slice(root.length + 1), + }; + }) + .filter(({ absolutePath }) => absolutePath.startsWith(`${root}/`)); + const sizeOf = (entry: ReturnType[number]) => + options.fileSizes?.get(entry.path) ?? + options.fileSizes?.get(entry.absolutePath) ?? + entry.bytes.byteLength; + const findCommandEntry = (command: string) => + entries().find(({ absolutePath }) => command.includes(absolutePath)); + const failed = (stderr: string, exitCode = 1) => ({ + exitCode, + stderr, + stdout: "", + success: false, + }); + const succeeded = (stdout = "") => ({ exitCode: 0, stderr: "", stdout, success: true }); + const unavailable = (): Promise => + Promise.reject(new Error("Unexpected sandbox test method call.")); + const exec: SandboxHandle["exec"] = async (command) => { + options.onExec?.(command); + if (command.includes("find . -type f")) { + if (options.listError !== undefined) { + return failed(options.listError.message); + } + return succeeded( + entries() + .toSorted((left, right) => left.relativePath.localeCompare(right.relativePath)) + .map((entry) => `./${entry.relativePath}\0${sizeOf(entry)}\0`) + .join(""), + ); + } + if (command.includes("stat --printf=")) { + const entry = findCommandEntry(command); + return entry === undefined + ? failed("missing test file", 44) + : succeeded(String(sizeOf(entry))); + } + if (command.includes("head -c")) { + const entry = findCommandEntry(command); + if (entry === undefined) { + return failed("missing test file"); + } + options.onRead?.(entry.absolutePath); + if (options.readError !== undefined) { + return failed(options.readError.message); + } + return succeeded( + entry.bytes.subarray(0, Number(command.match(/head -c (\d+)/)?.[1])).toBase64(), + ); + } + return succeeded(); + }; + const session = { + exec, + mkdir: unavailable, + readFile: unavailable, + startProcess: unavailable, + watch: unavailable, + writeFile: unavailable, + }; + + return { + ...session, + configureNetworkConstraints: unavailable, + createBackup: unavailable, + createSession: unavailable, + deleteSession: unavailable, + destroy: unavailable, + getSession: async () => session, + mountBucket: unavailable, + restoreBackup: unavailable, + setKeepAlive: unavailable, + terminal: unavailable, + unmountBucket: unavailable, + wsConnect: unavailable, + }; +} diff --git a/apps/api/tests/helpers/sqlite-d1.ts b/apps/api/tests/helpers/sqlite-d1.ts index be3340bb..632595d7 100644 --- a/apps/api/tests/helpers/sqlite-d1.ts +++ b/apps/api/tests/helpers/sqlite-d1.ts @@ -8,11 +8,17 @@ type RawOptions = { columnNames?: boolean } | undefined; const statementQueries = new WeakMap(); export class SqliteD1Database implements D1Database { - readonly #database = new Database(":memory:"); + readonly #database: Database; readonly #maxBoundParams: number | undefined; #batchTail: Promise = Promise.resolve(); - constructor(input: { foreignKeys?: boolean; maxBoundParams?: number } = {}) { + constructor( + input: { foreignKeys?: boolean; maxBoundParams?: number; serialized?: Uint8Array } = {}, + ) { + this.#database = + input.serialized === undefined + ? new Database(":memory:") + : Database.deserialize(input.serialized); this.#maxBoundParams = input.maxBoundParams; this.#database.run(`PRAGMA foreign_keys = ${input.foreignKeys === false ? "OFF" : "ON"}`); } @@ -25,6 +31,10 @@ export class SqliteD1Database implements D1Database { return createSqliteD1Statement(this.#database, query, [], this.#maxBoundParams); } + serialize(): Uint8Array { + return this.#database.serialize(); + } + app(): AppDatabase { return getAppDatabase(this); } diff --git a/apps/api/tests/lease-ownership-renewal.test.ts b/apps/api/tests/lease-ownership-renewal.test.ts new file mode 100644 index 00000000..58b1fed3 --- /dev/null +++ b/apps/api/tests/lease-ownership-renewal.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, test } from "bun:test"; + +import { createLeaseOwnershipRenewal } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/lease-ownership-renewal"; + +describe("lease ownership renewal", () => { + test("serializes concurrent renewals and latches ownership loss", async () => { + const result = Promise.withResolvers(); + let renewalCalls = 0; + let destructiveCalls = 0; + const requireOwnership = createLeaseOwnershipRenewal(() => { + renewalCalls += 1; + return result.promise; + }, "lost ownership"); + + const heartbeat = requireOwnership(); + const destructive = requireOwnership().then(() => { + destructiveCalls += 1; + }); + expect(renewalCalls).toBe(1); + + const settled = Promise.allSettled([heartbeat, destructive]); + result.resolve(false); + expect(await settled).toEqual([ + { reason: expect.objectContaining({ message: "lost ownership" }), status: "rejected" }, + { reason: expect.objectContaining({ message: "lost ownership" }), status: "rejected" }, + ]); + await expect(requireOwnership()).rejects.toThrow("lost ownership"); + expect(renewalCalls).toBe(1); + expect(destructiveCalls).toBe(0); + }); + + test("latches an uncertain renewal failure", async () => { + const cause = new Error("renewal failed"); + let renewalCalls = 0; + const requireOwnership = createLeaseOwnershipRenewal(async () => { + renewalCalls += 1; + throw cause; + }, "lost ownership"); + + await expect(requireOwnership()).rejects.toMatchObject({ cause, message: "lost ownership" }); + await expect(requireOwnership()).rejects.toMatchObject({ cause, message: "lost ownership" }); + expect(renewalCalls).toBe(1); + }); +}); diff --git a/apps/api/tests/native-resume-ref.test.ts b/apps/api/tests/native-resume-ref.test.ts index d8917675..11a8f41d 100644 --- a/apps/api/tests/native-resume-ref.test.ts +++ b/apps/api/tests/native-resume-ref.test.ts @@ -1,13 +1,25 @@ import { describe, expect, test } from "bun:test"; +import { parsePlatformId } from "@mosoo/id"; +import type { DriverInstanceId, SessionId, SessionRunId } from "@mosoo/id"; + import { deleteNativeResumeRefsForSessions, getNativeResumeRefForRuntime, + upsertNativeResumeRef, } from "../src/modules/runtime/infrastructure/native-resume-ref.repository"; import { SqliteD1Database } from "./helpers/sqlite-d1"; const SESSION_ID_1 = "01J000000000000000000000G1"; const SESSION_ID_2 = "01J000000000000000000000G2"; +const DRIVER_INSTANCE_ID = parsePlatformId( + "01J000000000000000000000G4", + "driver instance ID", +); +const SESSION_RUN_ID = parsePlatformId( + "01J000000000000000000000G3", + "session run ID", +); describe("native resume refs", () => { test("uses only the cursor committed with a Cattle checkpoint", async () => { @@ -23,6 +35,7 @@ describe("native resume refs", () => { created_at integer NOT NULL, kind text NOT NULL, observed_driver_instance_id text, + observed_event_seq integer DEFAULT 0 NOT NULL, observed_session_run_id text, runtime_id text NOT NULL, session_id text PRIMARY KEY NOT NULL, @@ -72,6 +85,7 @@ describe("native resume refs", () => { created_at integer NOT NULL, kind text NOT NULL, observed_driver_instance_id text, + observed_event_seq integer DEFAULT 0 NOT NULL, observed_session_run_id text, runtime_id text NOT NULL, session_id text PRIMARY KEY NOT NULL, @@ -92,4 +106,62 @@ describe("native resume refs", () => { .all<{ session_id: string; value: string }>(); expect(rows.results).toEqual([{ session_id: SESSION_ID_2, value: "thread-2" }]); }); + + test("only a higher durable event seq can replace an observed ref", async () => { + const database = new SqliteD1Database({ foreignKeys: false }); + database.execute(` + CREATE TABLE native_resume_ref ( + committed_session_run_id text, + committed_value text, + created_at integer NOT NULL, + kind text NOT NULL, + observed_driver_instance_id text, + observed_event_seq integer DEFAULT 0 NOT NULL, + observed_session_run_id text, + runtime_id text NOT NULL, + session_id text PRIMARY KEY NOT NULL, + updated_at integer NOT NULL, + value text NOT NULL + ); + `); + + const observation = { + driverInstanceId: DRIVER_INSTANCE_ID, + nativeResumeRef: { + kind: "openai_thread_id" as const, + runtimeId: "openai-runtime" as const, + value: "thread-5", + }, + observedEventSeq: 5, + sessionId: SESSION_ID_1 as SessionId, + sessionRunId: SESSION_RUN_ID, + }; + + await upsertNativeResumeRef(database, observation); + await upsertNativeResumeRef(database, { + ...observation, + nativeResumeRef: { ...observation.nativeResumeRef, value: "stale-thread" }, + observedEventSeq: 4, + }); + await upsertNativeResumeRef(database, observation); + await expect( + upsertNativeResumeRef(database, { + ...observation, + nativeResumeRef: { ...observation.nativeResumeRef, value: "conflicting-thread" }, + }), + ).rejects.toThrow("replayed with conflicting content"); + + await upsertNativeResumeRef(database, { + ...observation, + nativeResumeRef: { ...observation.nativeResumeRef, value: "thread-6" }, + observedEventSeq: 6, + }); + + expect( + await database + .prepare("SELECT observed_event_seq, value FROM native_resume_ref WHERE session_id = ?") + .bind(SESSION_ID_1) + .first<{ observed_event_seq: number; value: string }>(), + ).toEqual({ observed_event_seq: 6, value: "thread-6" }); + }); }); diff --git a/apps/api/tests/owner-debug-terminal.test.ts b/apps/api/tests/owner-debug-terminal.test.ts index f3f733c0..9859f076 100644 --- a/apps/api/tests/owner-debug-terminal.test.ts +++ b/apps/api/tests/owner-debug-terminal.test.ts @@ -37,24 +37,33 @@ interface TerminalSpy { createSessionCalls: { cwd?: string; id?: string }[]; handle: SandboxHandle; mkdirCalls: string[]; + sessionDisposeCalls: number; setKeepAliveCalls: boolean[]; + subjectDisposeCalls: number; } function createTerminalSandboxHandleSpy(): TerminalSpy { const mkdirCalls: string[] = []; const setKeepAliveCalls: boolean[] = []; const createSessionCalls: { cwd?: string; id?: string }[] = []; + let sessionDisposeCalls = 0; + let subjectDisposeCalls = 0; const unavailable = async () => { throw new Error("Unexpected sandbox test method call."); }; const sessionResponse = new Response("ok", { status: 200 }); const handle = { + activateRuntimeSubjectIncarnation: async () => {}, configureNetworkConstraints: async () => {}, createBackup: unavailable, + createRuntimeSubjectBackup: unavailable, createSession: async (options) => { createSessionCalls.push({ cwd: options?.cwd, id: options?.id }); return { + [Symbol.dispose]: () => { + sessionDisposeCalls++; + }, exec: unavailable, mkdir: unavailable, readFile: unavailable, @@ -66,12 +75,15 @@ function createTerminalSandboxHandleSpy(): TerminalSpy { }, deleteSession: unavailable, destroy: unavailable, + destroyRuntimeSubjectIncarnation: unavailable, exec: unavailable, getSession: unavailable, + inspectRuntimeSubjectIncarnation: unavailable, mkdir: async (path) => { mkdirCalls.push(path); }, mountBucket: unavailable, + markRuntimeSubjectIncarnationReady: async () => {}, readFile: unavailable, restoreBackup: unavailable, setKeepAlive: async (value) => { @@ -83,9 +95,23 @@ function createTerminalSandboxHandleSpy(): TerminalSpy { watch: unavailable, writeFile: unavailable, wsConnect: unavailable, + [Symbol.dispose]: () => { + subjectDisposeCalls++; + }, } as unknown as SandboxHandle; - return { createSessionCalls, handle, mkdirCalls, setKeepAliveCalls }; + return { + createSessionCalls, + handle, + mkdirCalls, + get sessionDisposeCalls() { + return sessionDisposeCalls; + }, + setKeepAliveCalls, + get subjectDisposeCalls() { + return subjectDisposeCalls; + }, + }; } describe("owner debug terminal", () => { @@ -137,5 +163,7 @@ describe("owner debug terminal", () => { ); expect(spy.createSessionCalls).toHaveLength(1); expect(spy.createSessionCalls[0]?.cwd).toBe("/workspace"); + expect(spy.sessionDisposeCalls).toBe(1); + expect(spy.subjectDisposeCalls).toBe(1); }); }); diff --git a/apps/api/tests/pet-stranded-recycle.test.ts b/apps/api/tests/pet-stranded-recycle.test.ts index 93e01ebc..22a372f5 100644 --- a/apps/api/tests/pet-stranded-recycle.test.ts +++ b/apps/api/tests/pet-stranded-recycle.test.ts @@ -8,7 +8,9 @@ import { repairStrandedRuntimeSubjectDeadlines } from "../src/modules/runtime/in import { listInactiveRuntimeSubjects } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { + PUBLIC_API_TEST_IDS, createPublicHttpContractDatabase, + insertActiveSandboxSessionFixture, insertNonOwnerSession, } from "./helpers/public-api-http-test-fixture"; import type { SqliteD1Database } from "./helpers/sqlite-d1"; @@ -26,36 +28,21 @@ function createBindings(database: D1Database): ApiBindings { async function insertSandbox( database: SqliteD1Database, - input: { - readonly id?: string; - readonly kind?: "cattle" | "pet"; - readonly subjectId?: string; - readonly subjectKind?: "agent" | "session"; - } = {}, + kind: "cattle" | "pet" = "pet", ): Promise { - await database - .prepare( - ` - INSERT INTO sandbox (id, kind, subject_kind, subject_id, status, inactive_deadline_at, created_at, updated_at) - VALUES (?, ?, ?, ?, 'active', NULL, 1, 1) - `, - ) - .bind( - input.id ?? SANDBOX_ID, - input.kind ?? "pet", - input.subjectKind ?? "agent", - input.subjectId ?? AGENT_ID, - ) - .run(); + await insertActiveSandboxSessionFixture(database, { + kind, + ownerAccountId: PUBLIC_API_TEST_IDS.nonOwnerAccount, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + timestampMs: 1, + }); } async function insertDriver( database: SqliteD1Database, input: { - readonly id?: string; readonly lastHeartbeatAt: number; - readonly sandboxId?: string; - readonly status?: string; }, ): Promise { await database @@ -64,6 +51,7 @@ async function insertDriver( INSERT INTO driver_instance ( id, sandbox_id, + sandbox_incarnation, sandbox_session_id, runtime, protocol, @@ -79,17 +67,18 @@ async function insertDriver( created_at, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( - input.id ?? DRIVER_ID, - input.sandboxId ?? SANDBOX_ID, + DRIVER_ID, + SANDBOX_ID, + 1, SESSION_ID, "cloudflare-container", "driver-ws", 1, - input.status ?? "ready", + "ready", new Uint8Array([1]), Date.now() + 10_000, Date.now(), @@ -107,8 +96,6 @@ async function insertRun( database: SqliteD1Database, input: { readonly driverInstanceId?: string | null; - readonly id?: string; - readonly sessionId?: string; readonly status?: string; } = {}, ): Promise { @@ -134,8 +121,8 @@ async function insertRun( `, ) .bind( - input.id ?? RUN_ID, - input.sessionId ?? SESSION_ID, + RUN_ID, + SESSION_ID, AGENT_ID, "01J00000000000000000000002", "user_prompt", @@ -237,43 +224,46 @@ describe("pet stranded recycle", () => { ).resolves.toEqual([{ id: SANDBOX_ID, kind: "pet" }]); }); - test("leaves live, busy, and resident subjects alone", async () => { + test("leaves a pet with a live Driver alone", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); const nowMs = Date.now(); - // A pet with a live driver is not reclaimable. - await insertSandbox(database, { id: "01J0000000000000000000000V" }); + await insertSandbox(database); await insertDriver(database, { - id: "01J0000000000000000000000W", lastHeartbeatAt: nowMs, - sandboxId: "01J0000000000000000000000V", }); - // A pet whose subject still has an active run is not reclaimable, even - // when no driver row links the run to the sandbox yet. - await insertSandbox(database, { - id: "01J0000000000000000000000X", - subjectId: SESSION_ID, - subjectKind: "session", + + await expect(repairStrandedRuntimeSubjectDeadlines(database, { now: nowMs })).resolves.toEqual({ + cattle: 0, + pet: 0, }); + await expect(readSandboxDeadline(database)).resolves.toBeNull(); + }); + + test("leaves a pet with an active subject Run alone", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + const nowMs = Date.now(); + await insertSandbox(database); await insertRun(database, { driverInstanceId: null, status: "queued" }); - // A cattle subject with an active conversation stays resident. - await insertSandbox(database, { id: "01J0000000000000000000000Y", kind: "cattle" }); - await database - .prepare( - ` - INSERT INTO sandbox_session (cloudflare_session_id, created_at, cwd, origin_json, sandbox_id, session_id, status, updated_at) - VALUES ('cf-session-1', 1, '/workspace', '{}', '01J0000000000000000000000Y', '01J0000000000000000000000C', 'active', 1) - `, - ) - .run(); await expect(repairStrandedRuntimeSubjectDeadlines(database, { now: nowMs })).resolves.toEqual({ cattle: 0, pet: 0, }); + await expect(readSandboxDeadline(database)).resolves.toBeNull(); + }); + + test("leaves cattle with an active conversation resident", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + const nowMs = Date.now(); + await insertSandbox(database, "cattle"); - await expect(readSandboxDeadline(database, "01J0000000000000000000000V")).resolves.toBeNull(); - await expect(readSandboxDeadline(database, "01J0000000000000000000000X")).resolves.toBeNull(); - await expect(readSandboxDeadline(database, "01J0000000000000000000000Y")).resolves.toBeNull(); + await expect(repairStrandedRuntimeSubjectDeadlines(database, { now: nowMs })).resolves.toEqual({ + cattle: 0, + pet: 0, + }); + await expect(readSandboxDeadline(database)).resolves.toBeNull(); }); }); diff --git a/apps/api/tests/prod-deploy-lease.test.ts b/apps/api/tests/prod-deploy-lease.test.ts new file mode 100644 index 00000000..45c1c139 --- /dev/null +++ b/apps/api/tests/prod-deploy-lease.test.ts @@ -0,0 +1,209 @@ +import { describe, expect, test } from "bun:test"; + +import { acquireProdDeployLease } from "../bin/deploy-prod"; +import { + acquireProdDeployLeaseStatements, + assertProdDeployLeaseOwned, + assertProdDeployLeaseReleased, + PROD_DEPLOY_LEASE_TABLE, + PROD_DEPLOY_LEASE_TABLE_SQL, + releaseProdDeployLeaseStatements, + verifyProdDeployLeaseStatements, +} from "../bin/prod-deploy-lease"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const OWNER_A = "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa"; +const OWNER_B = "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb"; +const OWNED_ROW = { + owner: OWNER_A, + table_sql: PROD_DEPLOY_LEASE_TABLE_SQL, + trigger_count: 0, +}; + +async function executeBatch(database: SqliteD1Database, statements: readonly string[]) { + return JSON.stringify(await database.batch(statements.map((sql) => database.prepare(sql)))); +} + +async function readOwner(database: SqliteD1Database): Promise { + const row = await database + .prepare(`SELECT "owner" FROM "${PROD_DEPLOY_LEASE_TABLE}" WHERE "id" = 1`) + .first<{ owner: string }>(); + return row?.owner ?? null; +} + +describe("production deploy lease", () => { + test("keeps a duplicate same-owner batch idempotent while the mutex remains held", async () => { + const database = new SqliteD1Database(); + await executeBatch(database, acquireProdDeployLeaseStatements(OWNER_A)); + + const retried = await executeBatch(database, acquireProdDeployLeaseStatements(OWNER_A)); + expect(() => + assertProdDeployLeaseOwned(`[wrangler warning]\n${retried}\n[trailing warning]`, OWNER_A), + ).not.toThrow(); + expect(await readOwner(database)).toBe(OWNER_A); + }); + + test("verifies ownership without mutating the durable mutex", async () => { + const database = new SqliteD1Database(); + assertProdDeployLeaseOwned( + await executeBatch(database, acquireProdDeployLeaseStatements(OWNER_A)), + OWNER_A, + ); + + assertProdDeployLeaseOwned( + await executeBatch(database, verifyProdDeployLeaseStatements()), + OWNER_A, + ); + expect(await readOwner(database)).toBe(OWNER_A); + }); + + test.each([ + ["no statement", JSON.stringify([])], + [ + "wrong statement count", + JSON.stringify([ + { results: [], success: true }, + { results: [OWNED_ROW], success: true }, + ]), + ], + ["no row", JSON.stringify([{ results: [], success: true }])], + ["multiple rows", JSON.stringify([{ results: [OWNED_ROW, OWNED_ROW], success: true }])], + [ + "malformed owner", + JSON.stringify([{ results: [{ ...OWNED_ROW, owner: null }], success: true }]), + ], + ])("rejects an owned-mutex readback with %s", (_label, raw) => { + expect(() => assertProdDeployLeaseOwned(raw, OWNER_A)).toThrow(); + }); + + test("retries release after the owner delete response is lost", async () => { + const database = new SqliteD1Database(); + assertProdDeployLeaseOwned( + await executeBatch(database, acquireProdDeployLeaseStatements(OWNER_A)), + OWNER_A, + ); + await executeBatch(database, releaseProdDeployLeaseStatements(OWNER_A)); + + const retried = await executeBatch(database, releaseProdDeployLeaseStatements(OWNER_A)); + expect(() => assertProdDeployLeaseReleased(retried)).not.toThrow(); + expect(await readOwner(database)).toBeNull(); + }); + + test("never transfers ownership until the current owner explicitly releases", async () => { + const database = new SqliteD1Database(); + assertProdDeployLeaseOwned( + await executeBatch(database, acquireProdDeployLeaseStatements(OWNER_A)), + OWNER_A, + ); + + const contended = await executeBatch(database, acquireProdDeployLeaseStatements(OWNER_B)); + expect(() => assertProdDeployLeaseOwned(contended, OWNER_B)).toThrow("another owner"); + expect(await readOwner(database)).toBe(OWNER_A); + + assertProdDeployLeaseOwned( + await executeBatch(database, verifyProdDeployLeaseStatements()), + OWNER_A, + ); + assertProdDeployLeaseReleased( + await executeBatch(database, releaseProdDeployLeaseStatements(OWNER_A)), + ); + + assertProdDeployLeaseOwned( + await executeBatch(database, acquireProdDeployLeaseStatements(OWNER_B)), + OWNER_B, + ); + const staleVerification = await executeBatch(database, verifyProdDeployLeaseStatements()); + expect(() => assertProdDeployLeaseOwned(staleVerification, OWNER_A)).toThrow("another owner"); + + const lateRelease = await executeBatch(database, releaseProdDeployLeaseStatements(OWNER_A)); + expect(() => assertProdDeployLeaseReleased(lateRelease)).toThrow(); + expect(await readOwner(database)).toBe(OWNER_B); + + const release = await executeBatch(database, releaseProdDeployLeaseStatements(OWNER_B)); + expect(() => assertProdDeployLeaseReleased(release)).not.toThrow(); + expect(await readOwner(database)).toBeNull(); + }); + + test("sends one acquisition request and retains a late commit for exact-owner recovery", async () => { + const database = new SqliteD1Database(); + const lostResponse = new Error("acquisition response was lost"); + let calls = 0; + let commitLate: (() => Promise) | undefined; + + expect(() => + acquireProdDeployLease(OWNER_A, (statements) => { + calls += 1; + commitLate = () => executeBatch(database, statements); + throw lostResponse; + }), + ).toThrow(lostResponse); + expect(calls).toBe(1); + + await commitLate?.(); + expect(await readOwner(database)).toBe(OWNER_A); + assertProdDeployLeaseReleased( + await executeBatch(database, releaseProdDeployLeaseStatements(OWNER_A)), + ); + }); + + test("retains the mutex after deployment failure and releases it only after success", async () => { + const source = await Bun.file(new URL("../bin/deploy-prod.ts", import.meta.url)).text(); + const retained = source.indexOf("✗ Retaining production deploy lease"); + const rethrown = source.indexOf("throw error;", retained); + const released = source.indexOf("releaseProdDeployLease();", rethrown); + + expect(retained).toBeGreaterThanOrEqual(0); + expect(rethrown).toBeGreaterThan(retained); + expect(released).toBeGreaterThan(rethrown); + }); + + test("rejects incomplete or non-empty release readbacks", () => { + const incomplete = JSON.stringify([{ results: [], success: true }]); + const multipleOwners = JSON.stringify([ + { results: [], success: true }, + { results: [{ owner: OWNER_B }, { owner: OWNER_B }], success: true }, + ]); + + expect(() => assertProdDeployLeaseReleased(incomplete)).toThrow(); + expect(() => assertProdDeployLeaseReleased(multipleOwners)).toThrow(); + }); + + test("fails closed when the reserved lease table has a spoofed schema", async () => { + const database = new SqliteD1Database(); + database.execute(` + CREATE TABLE "${PROD_DEPLOY_LEASE_TABLE}" ( + "id" integer PRIMARY KEY, + "owner" text NOT NULL, + "expires_at" integer NOT NULL + ) + `); + + const acquired = await executeBatch(database, acquireProdDeployLeaseStatements(OWNER_A)); + expect(() => assertProdDeployLeaseOwned(acquired, OWNER_A)).toThrow("schema is invalid"); + }); + + test("rejects lease-table triggers before they can change or release ownership", async () => { + const database = new SqliteD1Database(); + assertProdDeployLeaseOwned( + await executeBatch(database, acquireProdDeployLeaseStatements(OWNER_A)), + OWNER_A, + ); + database.execute(` + CREATE TRIGGER mutate_deploy_lease + AFTER DELETE ON "${PROD_DEPLOY_LEASE_TABLE}" + BEGIN + SELECT 1; + END + `); + + const verification = await executeBatch(database, verifyProdDeployLeaseStatements()); + expect(() => assertProdDeployLeaseOwned(verification, OWNER_A)).toThrow( + "must not have triggers", + ); + expect(await readOwner(database)).toBe(OWNER_A); + + const release = await executeBatch(database, releaseProdDeployLeaseStatements(OWNER_A)); + expect(() => assertProdDeployLeaseReleased(release)).toThrow(); + expect(await readOwner(database)).toBe(OWNER_A); + }); +}); diff --git a/apps/api/tests/prod-schema-guard.test.ts b/apps/api/tests/prod-schema-guard.test.ts index 8e37e7b3..6daaa9d0 100644 --- a/apps/api/tests/prod-schema-guard.test.ts +++ b/apps/api/tests/prod-schema-guard.test.ts @@ -1,106 +1,397 @@ import { describe, expect, test } from "bun:test"; +import { parseD1JsonResults } from "../bin/d1-json"; import { - extractTableNames, - findMissingProdTables, - getLatestSnapshotFilename, - parseExpectedTableNames, + assertProdSchemaMatches, + createProdSchemaCatalogFromIntrospectionRows, + createProdSchemaCatalogFromDrizzleSnapshot, + createProdSchemaIntrospectionStatements, + DRIZZLE_MIGRATION_PROD_SCHEMA_CATALOG, + findProdSchemaDifferences, + MANAGED_PROD_SCHEMA_TRIGGERS, + parseGeneratedProdSchemaCatalog, + PROTOCOL_V3_MIGRATION_INTENT_TABLE, + PROTOCOL_V3_MIGRATION_INTENT_TABLE_SQL, } from "../bin/prod-schema-guard"; +import { + applyDrizzleMigration, + assertDrizzleMigrationFiles, + drizzleMigrations, + latestDrizzleSnapshotFilename, +} from "./helpers/drizzle-migrations"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const VALID_SCHEMA_SQL = ` + CREATE TABLE parent ( + id text PRIMARY KEY NOT NULL, + state text DEFAULT 'ready' NOT NULL, + CONSTRAINT parent_state_check CHECK (state IN ('ready', 'retired')) + ); + CREATE TABLE child ( + tenant text NOT NULL, + id text NOT NULL, + parent_id text, + CONSTRAINT child_pk PRIMARY KEY (tenant, id), + CONSTRAINT child_parent_fk FOREIGN KEY (parent_id) + REFERENCES parent (id) ON UPDATE NO ACTION ON DELETE CASCADE, + CONSTRAINT child_id_check CHECK (length(id) > 0) + ); + CREATE UNIQUE INDEX child_parent_idx + ON child (parent_id) WHERE parent_id IS NOT NULL; + CREATE INDEX child_parent_expression_idx + ON child (coalesce(parent_id, id), tenant) WHERE parent_id IS NULL; + CREATE TABLE sequence_owner ( + id integer PRIMARY KEY AUTOINCREMENT + ); +`; + +const DRIZZLE_DIRECTORY = new URL("../../../pkgs/db/drizzle/", import.meta.url); +const PROTOCOL_V3_CUTOVER_TRIGGER_NAMES = [ + "__protocol_v3_cutover_api_command_insert", + "__protocol_v3_cutover_api_command_update", + "__protocol_v3_cutover_app_deployment_run_insert", + "__protocol_v3_cutover_app_deployment_run_update", + "__protocol_v3_cutover_command_insert", + "__protocol_v3_cutover_driver_insert", + "__protocol_v3_cutover_driver_update", + "__protocol_v3_cutover_environment_artifact_backup_staging_insert", + "__protocol_v3_cutover_sandbox_backup_insert", + "__protocol_v3_cutover_sandbox_backup_staging_insert", + "__protocol_v3_cutover_sandbox_backup_update", + "__protocol_v3_cutover_sandbox_insert", + "__protocol_v3_cutover_sandbox_session_insert", + "__protocol_v3_cutover_sandbox_session_update", + "__protocol_v3_cutover_sandbox_update", + "__protocol_v3_cutover_session_insert", + "__protocol_v3_cutover_session_run_insert", + "__protocol_v3_cutover_session_run_update", + "__protocol_v3_cutover_session_update", +] as const; + +async function readCatalog(database: SqliteD1Database) { + const tableNames = await readTableNames(database); + const rows = []; + for (const statement of createProdSchemaIntrospectionStatements(tableNames)) { + rows.push((await database.prepare(statement).all>()).results); + } + return createProdSchemaCatalogFromIntrospectionRows(rows, tableNames); +} + +async function readTableNames(database: SqliteD1Database): Promise { + const { results } = await database + .prepare( + "SELECT name FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%' ORDER BY name", + ) + .all<{ name: string }>(); + return results.map(({ name }) => name); +} + +async function catalogFor(sql: string) { + const database = new SqliteD1Database(); + database.execute(sql); + return readCatalog(database); +} + +describe("prod schema catalog", () => { + test("requires the journal to exactly match every ordered migration SQL file", () => { + const entries = [ + { idx: 0, tag: "0000_first" }, + { idx: 1, tag: "0001_second" }, + ]; + expect(() => + assertDrizzleMigrationFiles(entries, ["0000_first.sql", "0001_second.sql"], 2), + ).not.toThrow(); + for (const [candidateEntries, filenames, count] of [ + [entries, ["0000_first.sql", "0001_second.sql", "0002_untracked.sql"], 2], + [entries, ["0000_first.sql"], 2], + [[entries[0], { idx: 2, tag: "0002_second" }], ["0000_first.sql", "0002_second.sql"], 2], + [[entries[0], { idx: 1, tag: "0001/second" }], ["0000_first.sql", "0001/second.sql"], 2], + [entries, ["0000_first.sql", "0001_second.sql"], 1], + ] as const) { + expect(() => assertDrizzleMigrationFiles(candidateEntries, filenames, count)).toThrow(); + } + }); + + test("round-trips the multi-statement D1 JSON shape", async () => { + const database = new SqliteD1Database(); + database.execute(VALID_SCHEMA_SQL); + const tableNames = await readTableNames(database); + const results = []; + for (const statement of createProdSchemaIntrospectionStatements(tableNames)) { + results.push(await database.prepare(statement).all>()); + } -describe("getLatestSnapshotFilename", () => { - test("resolves the latest Drizzle snapshot from the journal", () => { expect( - getLatestSnapshotFilename( - JSON.stringify({ - entries: [ - { idx: 0, tag: "0000_baseline" }, - { idx: 1 }, - { idx: 2 }, - { idx: 3 }, - { idx: 4 }, - ], - }), + createProdSchemaCatalogFromIntrospectionRows( + parseD1JsonResults(`wrangler log\n${JSON.stringify(results)}`), + tableNames, ), - ).toBe("0004_snapshot.json"); + ).toEqual(await readCatalog(database)); + const missingSuccess: unknown[] = results.slice(); + missingSuccess[0] = { results: results[0]?.results }; + expect(() => parseD1JsonResults(JSON.stringify(missingSuccess))).toThrow(); }); - test("fails closed when the journal has no valid latest entry", () => { - expect(() => getLatestSnapshotFilename('{"entries":[]}')).toThrow(); + test("compares columns, composite PK, defaults, nullability, index predicate, FK, and checks", async () => { + const expected = await catalogFor(VALID_SCHEMA_SQL); + expect(() => assertProdSchemaMatches(expected, expected)).not.toThrow(); + + expect(expected.tables.find(({ name }) => name === "child")?.indexes).toContainEqual({ + columns: ["coalesce ( parent_id , id )", "tenant"], + name: "child_parent_expression_idx", + predicate: "parent_id is null", + unique: false, + }); + + const mutations = [ + VALID_SCHEMA_SQL.replace("state text", "state blob"), + VALID_SCHEMA_SQL.replace("DEFAULT 'ready'", "DEFAULT 'retired'"), + VALID_SCHEMA_SQL.replace("DEFAULT 'ready' NOT NULL", "DEFAULT 'ready'"), + VALID_SCHEMA_SQL.replace("parent_id text,", "parent_id text, extra text,"), + VALID_SCHEMA_SQL.replace("state text DEFAULT 'ready' NOT NULL", "state text"), + VALID_SCHEMA_SQL.replace("PRIMARY KEY (tenant, id)", "PRIMARY KEY (id, tenant)"), + VALID_SCHEMA_SQL.replace("child_parent_idx", "child_parent_other_idx"), + VALID_SCHEMA_SQL.replace( + "CREATE UNIQUE INDEX child_parent_idx", + "CREATE INDEX child_parent_idx", + ), + VALID_SCHEMA_SQL.replace("coalesce(parent_id, id)", "coalesce(parent_id, tenant)"), + VALID_SCHEMA_SQL.replace("WHERE parent_id IS NOT NULL", "WHERE parent_id IS NULL"), + VALID_SCHEMA_SQL.replace("ON UPDATE NO ACTION", "ON UPDATE CASCADE"), + VALID_SCHEMA_SQL.replace("ON DELETE CASCADE", "ON DELETE RESTRICT"), + VALID_SCHEMA_SQL.replace("length(id) > 0", "length(id) > 1"), + VALID_SCHEMA_SQL.replace("PRIMARY KEY AUTOINCREMENT", "PRIMARY KEY"), + ]; + for (const mutated of mutations) { + expect(findProdSchemaDifferences(expected, await catalogFor(mutated))).not.toEqual([]); + } + await expect( + catalogFor(VALID_SCHEMA_SQL.replace("parent_id text,", "parent_id text UNIQUE,")), + ).rejects.toThrow("unsupported inline UNIQUE constraints"); }); - test("fails closed when journal indexes are out of order, duplicated, or non-contiguous", () => { - for (const entries of [ - [{ idx: 1 }, { idx: 0 }], - [{ idx: 0 }, { idx: 0 }], - [{ idx: 0 }, { idx: 2 }], + test("compares table definitions for SQLite semantics missing from pragma metadata", async () => { + const plain = await catalogFor(`CREATE TABLE semantic ( + id text PRIMARY KEY NOT NULL, + source text NOT NULL, + derived text + )`); + const formatted = await catalogFor(`create table "semantic"( + "id" TEXT primary key not null, + "source" text not null, + "derived" text + )`); + expect(findProdSchemaDifferences(plain, formatted)).toEqual([]); + + for (const sql of [ + `CREATE TABLE semantic ( + id text COLLATE NOCASE PRIMARY KEY NOT NULL, + source text NOT NULL, + derived text + )`, + `CREATE TABLE semantic ( + id text PRIMARY KEY NOT NULL, + source text NOT NULL, + derived text + ) STRICT`, + `CREATE TABLE semantic ( + id text PRIMARY KEY NOT NULL, + source text NOT NULL, + derived text + ) WITHOUT ROWID`, ]) { - expect(() => getLatestSnapshotFilename(JSON.stringify({ entries }))).toThrow(); + expect(findProdSchemaDifferences(plain, await catalogFor(sql))).toContain( + "semantic.definition differs", + ); } - }); -}); -describe("parseExpectedTableNames", () => { - test("extracts sorted table names from a Drizzle snapshot", () => { - const snapshot = JSON.stringify({ - tables: { - session_run: { name: "session_run" }, - agent: { name: "agent" }, - }, - }); + const generatedVirtual = await catalogFor(`CREATE TABLE semantic ( + id text PRIMARY KEY NOT NULL, + source text NOT NULL, + derived text GENERATED ALWAYS AS (lower(source)) VIRTUAL + )`); + expect(findProdSchemaDifferences(plain, generatedVirtual)).toEqual([ + "semantic.columns differs", + "semantic.definition differs", + ]); - expect(parseExpectedTableNames(snapshot)).toEqual(["agent", "session_run"]); - }); + const generatedStored = await catalogFor(`CREATE TABLE semantic ( + id text PRIMARY KEY NOT NULL, + source text NOT NULL, + derived text GENERATED ALWAYS AS (lower(source)) STORED + )`); + expect(findProdSchemaDifferences(generatedVirtual, generatedStored)).toEqual([ + "semantic.columns differs", + "semantic.definition differs", + ]); - test("fails closed when the snapshot has no tables object", () => { - expect(() => parseExpectedTableNames("{}")).toThrow(); + const changedExpression = await catalogFor(`CREATE TABLE semantic ( + id text PRIMARY KEY NOT NULL, + source text NOT NULL, + derived text GENERATED ALWAYS AS (upper(source)) VIRTUAL + )`); + expect(findProdSchemaDifferences(generatedVirtual, changedExpression)).toEqual([ + "semantic.definition differs", + ]); }); - test("fails closed when the snapshot tables object is empty", () => { - expect(() => parseExpectedTableNames('{"tables":{}}')).toThrow(); - }); + test("ignores only known live internal tables and rejects application extras", async () => { + const expected = await catalogFor(VALID_SCHEMA_SQL); + const withInternalTables = await catalogFor(`${VALID_SCHEMA_SQL} + CREATE TABLE d1_migrations (id integer PRIMARY KEY AUTOINCREMENT); + CREATE TABLE _cf_KV (key text); + CREATE TABLE __production_deploy_lease (id integer); + CREATE TABLE __protocol_v3_cutover (id integer); + ${PROTOCOL_V3_MIGRATION_INTENT_TABLE_SQL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS")}; + `); + expect(findProdSchemaDifferences(expected, withInternalTables)).toEqual([]); - test("resolves the checked-in migration chain to a non-empty latest snapshot", async () => { - const metaDir = new URL("../../../pkgs/db/drizzle/meta/", import.meta.url); - const journal = await Bun.file(new URL("_journal.json", metaDir)).text(); - const snapshotFilename = getLatestSnapshotFilename(journal); - const snapshot = await Bun.file(new URL(snapshotFilename, metaDir)).text(); - const tableNames = parseExpectedTableNames(snapshot); + const malformedMigrationIntent = await catalogFor(`${VALID_SCHEMA_SQL} + CREATE TABLE __protocol_v3_cutover (id integer PRIMARY KEY); + ${PROTOCOL_V3_MIGRATION_INTENT_TABLE_SQL.replace("ON DELETE CASCADE", "ON DELETE RESTRICT")}; + `); + expect(findProdSchemaDifferences(expected, malformedMigrationIntent)).toContain( + `${PROTOCOL_V3_MIGRATION_INTENT_TABLE}.definition differs`, + ); - expect(tableNames).toContain("bound_agent_call_idempotency_key"); - expect(tableNames).toContain("usage_event_rollup_receipt"); - }); -}); + const triggeredMigrationIntent = await catalogFor(`${VALID_SCHEMA_SQL} + CREATE TABLE __protocol_v3_cutover (id integer PRIMARY KEY); + ${PROTOCOL_V3_MIGRATION_INTENT_TABLE_SQL}; + CREATE TRIGGER migration_intent_trigger + AFTER INSERT ON ${PROTOCOL_V3_MIGRATION_INTENT_TABLE} + BEGIN + SELECT 1; + END; + `); + expect(findProdSchemaDifferences(expected, triggeredMigrationIntent)).toContain( + "unexpected trigger migration_intent_trigger", + ); -describe("findMissingProdTables", () => { - test("returns expected tables absent from the live database", () => { - expect(findMissingProdTables(["agent", "session_run", "app"], ["agent", "app"])).toEqual([ - "session_run", - ]); + const withRetainedLegacyTable = await catalogFor( + `${VALID_SCHEMA_SQL} CREATE TABLE wechat_context_token (id integer);`, + ); + expect(findProdSchemaDifferences(expected, withRetainedLegacyTable)).toEqual([]); + + for (const name of [ + "WECHAT_CONTEXT_TOKEN", + "stale_application_table", + "__protocol_v3_legacy_rewrite_authorization", + ]) { + const withUnexpectedTable = await catalogFor( + `${VALID_SCHEMA_SQL} CREATE TABLE ${name} (id integer);`, + ); + expect(findProdSchemaDifferences(expected, withUnexpectedTable)).toEqual([ + `unexpected table ${name}`, + ]); + } }); - test("returns empty when every expected table is present (extra live tables are ignored)", () => { - expect(findMissingProdTables(["agent"], ["agent", "d1_migrations", "_cf_KV"])).toEqual([]); + test("requires the exact migration-owned trigger and rejects trigger extras", async () => { + const database = new SqliteD1Database(); + for (const migration of drizzleMigrations.filter(({ index }) => index <= 8)) { + applyDrizzleMigration(database, migration.tag); + } + const expected = await readCatalog(database); + const managed = MANAGED_PROD_SCHEMA_TRIGGERS.find( + ({ name }) => name === "session_event_tool_identity_consistency", + ); + if (!managed) throw new Error("Managed trigger fixture is missing."); + + database.execute(`DROP TRIGGER ${managed.name}`); + expect(findProdSchemaDifferences(expected, await readCatalog(database))).toContain( + `missing trigger ${managed.name}`, + ); + + database.execute(` + CREATE TRIGGER ${managed.name} + BEFORE INSERT ON session_event + WHEN 0 + BEGIN + SELECT 1; + END + `); + expect(findProdSchemaDifferences(expected, await readCatalog(database))).toContain( + `trigger ${managed.name} differs`, + ); + + database.execute(`DROP TRIGGER ${managed.name}`); + database.execute(managed.sql); + database.execute(` + CREATE TRIGGER extra_session_event_trigger + BEFORE UPDATE ON session_event + WHEN 0 + BEGIN + SELECT 1; + END + `); + expect(findProdSchemaDifferences(expected, await readCatalog(database))).toContain( + "unexpected trigger extra_session_event_trigger", + ); + + database.execute(` + CREATE TRIGGER __protocol_v3_cutover_api_command_insert + BEFORE UPDATE ON session_event + BEGIN + SELECT 1; + END + `); + expect(findProdSchemaDifferences(expected, await readCatalog(database))).toContain( + "trigger __protocol_v3_cutover_api_command_insert differs", + ); }); -}); -describe("extractTableNames", () => { - test("parses the wrangler d1 execute --json result shape", () => { - const stdout = JSON.stringify([ - { results: [{ name: "agent" }, { name: "session_run" }], success: true, meta: {} }, - ]); + test("pins the transient cutover trigger set through authority migrations", async () => { + const database = new SqliteD1Database(); + const stages = []; + for (const migration of drizzleMigrations) { + applyDrizzleMigration(database, migration.tag); + if ( + migration.tag === "0019_runtime-subject-operation-authority" || + migration.tag === "0020_sandbox-backup-object-authority" + ) { + stages.push({ + tag: migration.tag, + triggers: (await readCatalog(database)).triggers.filter(({ name }) => + name.startsWith("__protocol_v3_cutover_"), + ), + }); + } + } - expect(extractTableNames(stdout)).toEqual(["agent", "session_run"]); + const canonical = DRIZZLE_MIGRATION_PROD_SCHEMA_CATALOG.triggers.filter(({ name }) => + name.startsWith("__protocol_v3_cutover_"), + ); + expect(canonical.map(({ name }) => name)).toEqual(PROTOCOL_V3_CUTOVER_TRIGGER_NAMES); + expect(stages).toEqual([ + { tag: "0019_runtime-subject-operation-authority", triggers: canonical }, + { tag: "0020_sandbox-backup-object-authority", triggers: canonical }, + ]); }); - test("tolerates leading log lines before the JSON array", () => { - const stdout = `🌀 Executing on remote database DB\n${JSON.stringify([ - { results: [{ name: "agent" }] }, - ])}`; + test("matches the complete migration chain against the latest snapshot", async () => { + const snapshot = await Bun.file( + new URL(`meta/${latestDrizzleSnapshotFilename}`, DRIZZLE_DIRECTORY), + ).json(); + const expected = createProdSchemaCatalogFromDrizzleSnapshot(snapshot); + expect(findProdSchemaDifferences(expected, DRIZZLE_MIGRATION_PROD_SCHEMA_CATALOG)).toEqual([]); + const deployExpected = parseGeneratedProdSchemaCatalog(JSON.stringify(snapshot)); + expect(deployExpected.tables.every(({ definition }) => definition !== null)).toBeTrue(); + expect( + findProdSchemaDifferences(deployExpected, DRIZZLE_MIGRATION_PROD_SCHEMA_CATALOG), + ).toEqual([]); - expect(extractTableNames(stdout)).toEqual(["agent"]); + const firstTable = Object.values( + (snapshot as { tables: Record> }).tables, + )[0]; + if (firstTable === undefined) throw new Error("Drizzle snapshot table fixture is missing."); + firstTable["uniqueConstraints"] = { unsupported: {} }; + expect(() => createProdSchemaCatalogFromDrizzleSnapshot(snapshot)).toThrow( + "unsupported inline UNIQUE constraints", + ); }); - test("throws (fails closed) when no JSON array is present", () => { - expect(() => extractTableNames("error: could not connect")).toThrow(); + test("fails closed on malformed D1 output", () => { + expect(() => parseD1JsonResults("not json")).toThrow(); + expect(() => parseD1JsonResults("[]")).toThrow(); }); }); diff --git a/apps/api/tests/prod-wfp-preflight.test.ts b/apps/api/tests/prod-wfp-preflight.test.ts new file mode 100644 index 00000000..60c478fc --- /dev/null +++ b/apps/api/tests/prod-wfp-preflight.test.ts @@ -0,0 +1,511 @@ +import { describe, expect, test } from "bun:test"; + +import { + apiWorkerDeployArgs, + apiWorkerDryRunArgs, + assertProdWfpReadOnlyInfrastructure, + assertProdWfpWorkflowDeployGate, + PROD_API_WORKER_NAME, + PROD_APP_DEPLOYMENT_DOMAIN, + PROD_APP_DEPLOYMENT_WORKFLOW, + PROD_APP_DEPLOYMENT_WORKFLOW_BINDING, + PROD_APP_DEPLOYMENT_WORKFLOW_CLASS, + PROD_APP_DISPATCH_NAMESPACE, + PROD_APP_WILDCARD_DNS, + PROD_APP_WILDCARD_ROUTE, + PROD_WFP_WRITE_CANARY_SCRIPT, + shouldProbeProdWfpBeforeMutation, + verifyProdWfpGatewayProbe, + verifyProdWfpWriteCanary, + WRANGLER_DISABLE_AUTO_PROVISION, +} from "../bin/prod-wfp-preflight"; +import type { ProdWfpReadOnlyInventory, WfpWriteCanaryClient } from "../bin/prod-wfp-preflight"; +import { + APP_DEPLOYMENT_PROBE_HOST_LABEL, + APP_DEPLOYMENT_PROBE_SCRIPT_NAME, +} from "../src/modules/apps/application/app-deployment-gateway"; + +const ACCOUNT_ID = "a".repeat(32); +const ZONE_ID = "b".repeat(32); +const NOW_MS = Date.parse("2026-08-30T00:00:00.000Z"); +const APP_ID = "01J0000000000000000000000Q"; +const CONFIG = { + accountId: ACCOUNT_ID, + apiWorkerName: PROD_API_WORKER_NAME, + dispatchBinding: "APP_DEPLOYMENT_DISPATCHER", + dispatchNamespace: PROD_APP_DISPATCH_NAMESPACE, + requireProbeOnlyNamespace: false, + wildcardDns: PROD_APP_WILDCARD_DNS, + wildcardRoute: PROD_APP_WILDCARD_ROUTE, + workflowBinding: PROD_APP_DEPLOYMENT_WORKFLOW_BINDING, + workflowClass: PROD_APP_DEPLOYMENT_WORKFLOW_CLASS, + workflowName: PROD_APP_DEPLOYMENT_WORKFLOW, + workflowRepairAllowed: false, + zoneId: ZONE_ID, + zoneName: "mosoo.ai", +} as const; +function validInventory(): ProdWfpReadOnlyInventory { + return { + apiWorkerBindings: [ + { + className: null, + name: "APP_DEPLOYMENT_DISPATCHER", + namespace: PROD_APP_DISPATCH_NAMESPACE, + scriptName: null, + type: "dispatch_namespace", + workflowName: null, + }, + { + className: PROD_APP_DEPLOYMENT_WORKFLOW_CLASS, + name: PROD_APP_DEPLOYMENT_WORKFLOW_BINDING, + namespace: null, + scriptName: PROD_API_WORKER_NAME, + type: "workflow", + workflowName: PROD_APP_DEPLOYMENT_WORKFLOW, + }, + ], + certificates: [ + { + certificates: [ + { + expiresOn: "2026-12-31T00:00:00.000Z", + hosts: ["mosoo.ai", PROD_APP_WILDCARD_DNS], + status: "active", + }, + ], + hosts: ["mosoo.ai", PROD_APP_WILDCARD_DNS], + status: "active", + }, + ], + dnsRecords: [ + { name: PROD_APP_WILDCARD_DNS, proxied: true, type: "A" }, + { name: "api.mosoo.ai", proxied: true, type: "A" }, + { name: "mosoo.ai", proxied: null, type: "TXT" }, + ], + legacyResources: [], + namespace: { + id: "namespace-id", + name: PROD_APP_DISPATCH_NAMESPACE, + scriptCount: 2, + trustedWorkers: false, + }, + probeScript: { + id: APP_DEPLOYMENT_PROBE_SCRIPT_NAME, + namespace: PROD_APP_DISPATCH_NAMESPACE, + }, + routes: [ + { pattern: PROD_APP_WILDCARD_ROUTE, script: PROD_API_WORKER_NAME }, + { pattern: "api.mosoo.ai/*", script: "unrelated-worker" }, + ], + workflow: { + className: PROD_APP_DEPLOYMENT_WORKFLOW_CLASS, + id: "workflow-id", + name: PROD_APP_DEPLOYMENT_WORKFLOW, + scriptName: PROD_API_WORKER_NAME, + }, + zone: { accountId: ACCOUNT_ID, id: ZONE_ID, name: "mosoo.ai", status: "active" }, + }; +} + +describe("production WfP preflight", () => { + test("disables Wrangler auto-provision for API dry-run and deployment", () => { + expect(apiWorkerDryRunArgs()).toEqual([ + "deploy", + "--env", + "prod", + "--minify", + "--strict", + "--dry-run", + WRANGLER_DISABLE_AUTO_PROVISION, + ]); + expect(apiWorkerDeployArgs("release-tag")).toEqual([ + "deploy", + "--env", + "prod", + "--minify", + "--strict", + WRANGLER_DISABLE_AUTO_PROVISION, + "--containers-rollout", + "immediate", + "--tag", + "release-tag", + ]); + }); + + test("keeps every Workflow and the stage/prod dispatch topology exact", async () => { + const wrangler = Bun.TOML.parse( + await Bun.file(new URL("../wrangler.toml", import.meta.url)).text(), + ) as { + workflows: Array<{ binding: string; class_name: string; name: string }>; + env: Record< + string, + { + dispatch_namespaces: Array<{ binding: string; namespace: string }>; + routes: Array<{ pattern: string; zone_name: string }>; + vars: Record; + workflows: Array<{ binding: string; class_name: string; name: string }>; + } + >; + }; + + expect(wrangler.workflows).toEqual([ + { + binding: PROD_APP_DEPLOYMENT_WORKFLOW_BINDING, + class_name: PROD_APP_DEPLOYMENT_WORKFLOW_CLASS, + name: "mosoo-app-deployment-dev", + }, + ]); + + for (const [environment, domain, namespace, workflow] of [ + ["stage", "apps-stage.mosoo.ai", "mosoo-app-deployments-stage", "mosoo-app-deployment-stage"], + [ + "prod", + PROD_APP_DEPLOYMENT_DOMAIN, + PROD_APP_DISPATCH_NAMESPACE, + PROD_APP_DEPLOYMENT_WORKFLOW, + ], + ] as const) { + const config = wrangler.env[environment]; + expect(config?.vars["MOSOO_APP_DEPLOYMENT_DOMAIN"]).toBe(domain); + expect(config?.vars["MOSOO_APP_DISPATCH_NAMESPACE"]).toBe(namespace); + expect(config?.dispatch_namespaces).toEqual([ + { binding: "APP_DEPLOYMENT_DISPATCHER", namespace }, + ]); + expect(config?.workflows).toEqual([ + { + binding: PROD_APP_DEPLOYMENT_WORKFLOW_BINDING, + class_name: PROD_APP_DEPLOYMENT_WORKFLOW_CLASS, + name: workflow, + }, + ]); + expect(config?.routes).toContainEqual({ pattern: `*.${domain}/*`, zone_name: "mosoo.ai" }); + } + + expect(PROD_APP_WILDCARD_ROUTE).toBe("*.apps.mosoo.ai/*"); + }); + + test("accepts only the exact active account, zone, namespace, DNS, and certificate", () => { + expect(assertProdWfpReadOnlyInfrastructure(validInventory(), CONFIG, NOW_MS)).toBe("exact"); + + const trusted = validInventory(); + trusted.namespace.trustedWorkers = true; + expect(() => assertProdWfpReadOnlyInfrastructure(trusted, CONFIG, NOW_MS)).toThrow( + "missing, mismatched, or trusted", + ); + + const inconsistentCount = validInventory(); + inconsistentCount.namespace.scriptCount = 0; + expect(() => assertProdWfpReadOnlyInfrastructure(inconsistentCount, CONFIG, NOW_MS)).toThrow( + "missing, mismatched, or trusted", + ); + + const wrongAccount = validInventory(); + wrongAccount.zone.accountId = "c".repeat(32); + expect(() => assertProdWfpReadOnlyInfrastructure(wrongAccount, CONFIG, NOW_MS)).toThrow( + "zone identity", + ); + + const dnsOnly = validInventory(); + dnsOnly.dnsRecords = [{ name: PROD_APP_WILDCARD_DNS, proxied: false, type: "CNAME" }]; + expect(() => assertProdWfpReadOnlyInfrastructure(dnsOnly, CONFIG, NOW_MS)).toThrow( + "proxied traffic records", + ); + }); + + test("requires exact remote route, binding, probe script, and probe-only bootstrap namespace", () => { + const missingRoute = validInventory(); + missingRoute.routes = []; + expect(() => assertProdWfpReadOnlyInfrastructure(missingRoute, CONFIG, NOW_MS)).toThrow( + "route *.apps.mosoo.ai/* exactly once", + ); + + const wrongBinding = validInventory(); + wrongBinding.apiWorkerBindings[0].namespace = "wrong-namespace"; + expect(() => assertProdWfpReadOnlyInfrastructure(wrongBinding, CONFIG, NOW_MS)).toThrow( + "must bind APP_DEPLOYMENT_DISPATCHER", + ); + + for (const probeScript of [ + { id: null, namespace: null }, + { id: APP_DEPLOYMENT_PROBE_SCRIPT_NAME, namespace: "wrong-namespace" }, + { id: "wrong-probe", namespace: PROD_APP_DISPATCH_NAMESPACE }, + ]) { + const inventory = validInventory(); + inventory.probeScript = probeScript; + expect(() => assertProdWfpReadOnlyInfrastructure(inventory, CONFIG, NOW_MS)).toThrow( + `pre-provisioned ${APP_DEPLOYMENT_PROBE_SCRIPT_NAME}`, + ); + } + + const missingWorkflow = validInventory(); + missingWorkflow.workflow.id = null; + expect(() => assertProdWfpReadOnlyInfrastructure(missingWorkflow, CONFIG, NOW_MS)).toThrow( + "Bootstrap the exact AppDeploymentWorkflow", + ); + + const workflowBootstrap = validInventory(); + workflowBootstrap.workflow = { + className: null, + id: null, + name: null, + scriptName: null, + }; + workflowBootstrap.apiWorkerBindings = workflowBootstrap.apiWorkerBindings.filter( + (binding) => binding.type !== "workflow", + ); + expect(assertProdWfpReadOnlyInfrastructure(workflowBootstrap, CONFIG, NOW_MS)).toBe( + "bootstrap", + ); + expect(() => assertProdWfpWorkflowDeployGate("bootstrap", false)).toThrow( + "requires the closed production gate", + ); + expect(() => assertProdWfpWorkflowDeployGate("bootstrap", true)).not.toThrow(); + expect(shouldProbeProdWfpBeforeMutation("bootstrap")).toBe(false); + expect(shouldProbeProdWfpBeforeMutation("exact")).toBe(true); + + const interruptedBootstrap = validInventory(); + interruptedBootstrap.apiWorkerBindings = interruptedBootstrap.apiWorkerBindings.filter( + (binding) => binding.type !== "workflow", + ); + expect(() => assertProdWfpReadOnlyInfrastructure(interruptedBootstrap, CONFIG, NOW_MS)).toThrow( + "must bind APP_DEPLOYMENT_WORKFLOW", + ); + expect( + assertProdWfpReadOnlyInfrastructure( + interruptedBootstrap, + { ...CONFIG, workflowRepairAllowed: true }, + NOW_MS, + ), + ).toBe("repair"); + expect(() => assertProdWfpWorkflowDeployGate("repair", false)).toThrow( + "requires the closed production gate", + ); + expect(shouldProbeProdWfpBeforeMutation("repair")).toBe(false); + + const wrongWorkflowBinding = validInventory(); + wrongWorkflowBinding.apiWorkerBindings[1].workflowName = "wrong-workflow"; + expect(() => assertProdWfpReadOnlyInfrastructure(wrongWorkflowBinding, CONFIG, NOW_MS)).toThrow( + "must bind APP_DEPLOYMENT_WORKFLOW", + ); + + const removedWorkflowWithOldBinding = validInventory(); + removedWorkflowWithOldBinding.workflow = workflowBootstrap.workflow; + expect(() => + assertProdWfpReadOnlyInfrastructure(removedWorkflowWithOldBinding, CONFIG, NOW_MS), + ).toThrow("Bootstrap the exact AppDeploymentWorkflow"); + + const bootstrap = validInventory(); + expect(() => + assertProdWfpReadOnlyInfrastructure( + bootstrap, + { ...CONFIG, requireProbeOnlyNamespace: true }, + NOW_MS, + ), + ).toThrow("contain only the pre-provisioned"); + bootstrap.namespace.scriptCount = 1; + expect(() => + assertProdWfpReadOnlyInfrastructure( + bootstrap, + { ...CONFIG, requireProbeOnlyNamespace: true }, + NOW_MS, + ), + ).not.toThrow(); + }); + + test("rejects every Worker route that can shadow the deployment wildcard", () => { + for (const route of [ + { pattern: `${PROD_APP_WILDCARD_DNS}/admin*`, script: "other-worker" }, + { + pattern: `https://${APP_DEPLOYMENT_PROBE_HOST_LABEL}.${PROD_APP_DEPLOYMENT_DOMAIN}/*`, + script: "other-worker", + }, + { + pattern: `app-${APP_ID.toLowerCase()}.${PROD_APP_DEPLOYMENT_DOMAIN}/*`, + script: null, + }, + { pattern: "*.mosoo.ai/private*", script: null }, + { pattern: null, script: null }, + ]) { + const inventory = validInventory(); + inventory.routes = [...inventory.routes, route]; + expect(() => assertProdWfpReadOnlyInfrastructure(inventory, CONFIG, NOW_MS)).toThrow( + "must not shadow *.apps.mosoo.ai/*", + ); + } + }); + + test("rejects every specific DNS record below the deployment wildcard", () => { + for (const record of [ + { name: `probe.${PROD_APP_DEPLOYMENT_DOMAIN}`, proxied: true, type: "A" }, + { + name: `app-${APP_ID.toLowerCase()}.${PROD_APP_DEPLOYMENT_DOMAIN}`, + proxied: true, + type: "CNAME", + }, + { name: `disabled.${PROD_APP_DEPLOYMENT_DOMAIN}`, proxied: null, type: "TXT" }, + { name: `delegated.${PROD_APP_DEPLOYMENT_DOMAIN}`, proxied: null, type: "NS" }, + { + name: `_validation.probe.${PROD_APP_DEPLOYMENT_DOMAIN}`, + proxied: null, + type: "TXT", + }, + ]) { + const inventory = validInventory(); + inventory.dnsRecords = [...inventory.dnsRecords, record]; + expect(() => assertProdWfpReadOnlyInfrastructure(inventory, CONFIG, NOW_MS)).toThrow( + "DNS records must not shadow *.apps.mosoo.ai", + ); + } + }); + + test("reports legacy Pages/classic Worker resources without deleting them", () => { + const inventory = validInventory(); + inventory.legacyResources = [ + "Pages project app-legacy (delete project app-legacy)", + "classic Worker route route-id: app-legacy.apps.mosoo.ai/*", + ]; + + expect(() => assertProdWfpReadOnlyInfrastructure(inventory, CONFIG, NOW_MS)).toThrow( + "Pages project app-legacy", + ); + }); + + test("rejects shallow, single-host, inactive, and expiring certificates", () => { + for (const mutate of [ + (inventory: ProdWfpReadOnlyInventory) => { + inventory.certificates[0].hosts = ["*.mosoo.ai"]; + inventory.certificates[0].certificates[0].hosts = ["*.mosoo.ai"]; + }, + (inventory: ProdWfpReadOnlyInventory) => { + const exactHost = `app-${APP_ID.toLowerCase()}.${PROD_APP_DEPLOYMENT_DOMAIN}`; + inventory.certificates[0].hosts = [exactHost]; + inventory.certificates[0].certificates[0].hosts = [exactHost]; + }, + (inventory: ProdWfpReadOnlyInventory) => { + inventory.certificates[0].status = "pending_deployment"; + }, + (inventory: ProdWfpReadOnlyInventory) => { + inventory.certificates[0].certificates[0].expiresOn = "2026-09-20T00:00:00.000Z"; + }, + ]) { + const inventory = validInventory(); + mutate(inventory); + expect(() => assertProdWfpReadOnlyInfrastructure(inventory, CONFIG, NOW_MS)).toThrow( + "valid for at least 30 more days", + ); + } + }); + + test("proves Workers Scripts Write and always verifies exact cleanup", async () => { + const calls: string[] = []; + let exists = false; + const client: WfpWriteCanaryClient = { + async deleteScript(name) { + calls.push(`delete:${name}`); + exists = false; + }, + async readScript(name) { + calls.push(`read:${name}`); + return exists ? { id: name, namespace: PROD_APP_DISPATCH_NAMESPACE } : null; + }, + async uploadScript(name) { + calls.push(`upload:${name}`); + exists = true; + return { id: name }; + }, + }; + + await expect( + verifyProdWfpWriteCanary(client, PROD_APP_DISPATCH_NAMESPACE, PROD_WFP_WRITE_CANARY_SCRIPT), + ).resolves.toBeUndefined(); + expect(calls).toEqual([ + `delete:${PROD_WFP_WRITE_CANARY_SCRIPT}`, + `read:${PROD_WFP_WRITE_CANARY_SCRIPT}`, + `upload:${PROD_WFP_WRITE_CANARY_SCRIPT}`, + `read:${PROD_WFP_WRITE_CANARY_SCRIPT}`, + `delete:${PROD_WFP_WRITE_CANARY_SCRIPT}`, + `read:${PROD_WFP_WRITE_CANARY_SCRIPT}`, + ]); + }); + + test("fails closed when upload or deletion readback is not exact", async () => { + const remainingClient: WfpWriteCanaryClient = { + async deleteScript() {}, + async readScript(name) { + return { id: name, namespace: PROD_APP_DISPATCH_NAMESPACE }; + }, + async uploadScript() { + return { id: "wrong-script" }; + }, + }; + + await expect( + verifyProdWfpWriteCanary( + remainingClient, + PROD_APP_DISPATCH_NAMESPACE, + PROD_WFP_WRITE_CANARY_SCRIPT, + ), + ).rejects.toThrow("curl -X DELETE"); + }); + + test("probes request/body, stream, cancellation recovery, and WebSocket", async () => { + const calls: string[] = []; + const encoder = new TextEncoder(); + const fetcher: typeof fetch = async (input, init) => { + const url = new URL( + typeof input === "string" ? input : input instanceof URL ? input.href : input.url, + ); + calls.push(url.pathname); + const nonce = url.searchParams.get("nonce") ?? ""; + + if (url.pathname.endsWith("/http")) { + return Response.json( + { + body: typeof init?.body === "string" ? init.body : null, + hostname: url.hostname, + method: init?.method, + nonce, + pathname: url.pathname, + search: url.search, + }, + { headers: { "Cache-Control": "no-store" } }, + ); + } + + const chunks = url.pathname.endsWith("/stream") + ? [`start:${nonce}\n`, `end:${nonce}\n`] + : [`start:${nonce}\n`]; + return new Response( + new ReadableStream({ + start(controller) { + for (const chunk of chunks) controller.enqueue(encoder.encode(chunk)); + if (url.pathname.endsWith("/stream")) controller.close(); + }, + }), + ); + }; + let websocket: { nonce: string; url: string } | null = null; + + await verifyProdWfpGatewayProbe( + PROD_APP_DEPLOYMENT_DOMAIN, + { + fetch: fetcher, + probeWebSocket: async (url, nonce) => { + websocket = { nonce, url }; + }, + }, + () => "probe-nonce", + ); + + expect(calls).toEqual([ + "/.well-known/mosoo-wfp-probe/http", + "/.well-known/mosoo-wfp-probe/stream", + "/.well-known/mosoo-wfp-probe/cancel", + "/.well-known/mosoo-wfp-probe/http", + ]); + expect(websocket).toEqual({ + nonce: "probe-nonce", + url: `wss://${APP_DEPLOYMENT_PROBE_HOST_LABEL}.apps.mosoo.ai/.well-known/mosoo-wfp-probe/websocket?nonce=probe-nonce`, + }); + }); +}); diff --git a/apps/api/tests/production-deploy-safety.test.ts b/apps/api/tests/production-deploy-safety.test.ts index 6aaae4a6..41569b2c 100644 --- a/apps/api/tests/production-deploy-safety.test.ts +++ b/apps/api/tests/production-deploy-safety.test.ts @@ -1,9 +1,2740 @@ -import { expect, test } from "bun:test"; +import { describe, expect, test } from "bun:test"; -test("switches the production Worker and Driver image without a gradual protocol split", async () => { - const source = await Bun.file(new URL("../bin/deploy-prod.ts", import.meta.url)).text(); +import { MANAGED_PROD_SCHEMA_TRIGGERS } from "../bin/prod-schema-guard"; +import { + ACCEPT_PROTOCOL_V3_QUEUE_RESUME_SQL, + assertCutoverMigrationJournalAudited, + assertProtocolV3SmokeAgent, + assertProtocolV3Release, + assertProtocolV3WorkerVersion, + assertProtocolV3LegacyTerminalIntegrity, + assertProtocolV3LegacyTerminalSourceInventory, + assertProtocolV3LossyMigrationInventory, + assertProtocolV3RuntimeAuthorityPreflight, + AUDITED_MIGRATION_NAMES, + beginProtocolV3MigrationSql, + CLOSE_PROTOCOL_V3_SMOKE_WINDOW_SQL, + collectProtocolV3ContainerApplications, + collectProtocolV3ContainerInstances, + completeProtocolV3QueueResume, + DROP_PROTOCOL_V3_CUTOVER_TRIGGERS_SQL, + ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL, + ENTER_PROTOCOL_V3_DRAIN_SQL, + ENTER_PROTOCOL_V3_QUEUES_RESUMING_SQL, + findPendingProdMigrations, + installProtocolV3CutoverSql, + installProtocolV3PostMigrationCutoverSql, + INSTALL_PROTOCOL_V3_POST_MIGRATION_CUTOVER_SQL, + isProtocolV3ContainerRolloutConverged, + isProtocolV3CutoverDrained, + isProtocolV3RuntimeDrained, + isProtocolV3SmokeReady, + openProtocolV3SmokeWindowSql, + parseProtocolV3CommandFreeze, + parseCleanGitTreeOid, + parseProtocolV3ContainerManifestDigest, + parseProtocolV3CutoverDrain, + parseProtocolV3CutoverObjects, + parseProtocolV3CutoverProbe, + parseProtocolV3CutoverState, + parseProtocolV3LegacyTerminalIntegrity, + parseProtocolV3LegacyTerminalSourceInventory, + parseProtocolV3LossyMigrationInventory, + parseProtocolV3RuntimeAuthorityPreflight, + parseProtocolV3SmokeStatus, + parseProtocolV3WorkerDeployment, + parseStoredProtocolV3SmokeRequestKey, + parseStoredProtocolV3SmokeSession, + parseStoredProtocolV3CutoverBookmark, + parseTimeTravelBookmark, + PROTOCOL_V3_COMMAND_FREEZE_SQL, + PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + PROTOCOL_V3_CUTOVER_OBJECTS_SQL, + PROTOCOL_V3_CUTOVER_PROBE_SQL, + PROTOCOL_V3_CUTOVER_QUEUE_NAMES, + PROTOCOL_V3_LEGACY_TERMINAL_INTEGRITY_SQL, + PROTOCOL_V3_LEGACY_TERMINAL_SOURCE_INVENTORY_SQL, + PROTOCOL_V3_LOSSY_MIGRATION_INVENTORY_SQL, + PROTOCOL_V3_MIGRATION_INTENT_TABLE, + PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + PROTOCOL_V3_POST_MIGRATION_CUTOVER_DRAIN_SQL, + PROTOCOL_V3_POST_MIGRATION_UNSAFE_SANDBOXES_SQL, + PROTOCOL_V3_RUNTIME_AUTHORITY_PREFLIGHT_SQL, + PROTOCOL_V3_SMOKE_REQUEST_KEY_SQL, + PROTOCOL_V3_SMOKE_SESSION_SQL, + protocolV3ContainerImageTag, + protocolV3RuntimeAuthorityPreflightSql, + protocolV3SmokeAgentSql, + protocolV3SmokeStatusSql, + recoverProtocolV3CutoverFailure, + REMOVE_PROTOCOL_V3_CUTOVER_SQL, + storeProtocolV3CutoverBookmarkSql, + storeProtocolV3RolloutSql, + storeProtocolV3SmokeRequestKeySql, + storeProtocolV3SmokeSessionSql, + updateAndVerifyProtocolV3QueueDelivery, +} from "../bin/protocol-v3-cutover"; +import { applyDrizzleMigration, applyDrizzleMigrationsThrough } from "./helpers/drizzle-migrations"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; - expect(source).toContain( - 'run(["deploy", "--env", ENV, "--minify", "--containers-rollout", "immediate"]);', +const RELEASE_TREE_OID = "0123456789abcdef0123456789abcdef01234567"; +const CONTAINER_IMAGE_DIGEST = "a".repeat(64); +const INSTALL_PROTOCOL_V3_CUTOVER_SQL = installProtocolV3CutoverSql(RELEASE_TREE_OID); +const EMPTY_ROLLOUT = { + containerApplicationVersion: null, + containerImageDigest: null, + migrationStarted: false, + releaseTreeOid: RELEASE_TREE_OID, + workerVersionId: null, +} as const; +const DRAINED_CUTOVER_ROW = { + active_app_deployment_runs: 0, + active_runs: 0, + live_drivers: 0, + nonterminal_commands: 0, + nonterminal_api_commands: 0, + unsafe_environment_artifact_backup_staging: 0, + unsafe_sandbox_backups: 0, + unsafe_sandbox_backup_staging: 0, + unsafe_sandbox_sessions: 0, + unsafe_sandboxes: 0, + unsafe_sessions: 0, + unsettled_effects: 0, +} as const; +const INSERT_ENVIRONMENT_ARTIFACT_STAGING_SQL = ` + INSERT INTO environment_package_artifact_backup_staging ( + actual_backup_id, app_id, attempt_count, claim_owner, command_id, created_at, + delivery_generation, dir, input_digest, paths_json, updated_at + ) VALUES ( + NULL, 'app-1', 1, 'worker-1', 'artifact-command', 1, + 1, '/tmp/artifact', '${"0".repeat(64)}', + '{"executable":[],"node":[],"python":[]}', 1 + ) +`; + +function d1Json(row: Record): string { + return JSON.stringify([{ meta: {}, results: [row], success: true }]); +} + +function createCutoverGateDatabase(): SqliteD1Database { + const database = new SqliteD1Database(); + database.execute(` + CREATE TABLE d1_migrations (name text PRIMARY KEY); + CREATE TABLE session ( + id text PRIMARY KEY, + creator_account_id text NOT NULL, + end_user_id text, + status text NOT NULL, + status_operation_id text + ); + CREATE TABLE sandbox ( + id text PRIMARY KEY, + subject_kind text NOT NULL, + subject_id text NOT NULL, + status text NOT NULL, + status_operation_id text, + claim_owner text, + claim_expires_at integer + ); + CREATE TABLE sandbox_session ( + session_id text PRIMARY KEY, + sandbox_id text NOT NULL, + status text NOT NULL + ); + CREATE TABLE sandbox_backup ( + id text PRIMARY KEY, + status text NOT NULL, + error_message text + ); + CREATE TABLE session_run ( + id text PRIMARY KEY, + status text NOT NULL, + created_by_account_id text NOT NULL, + session_id text NOT NULL + ); + CREATE TABLE driver_instance ( + id text PRIMARY KEY, + status text NOT NULL, + sandbox_id text NOT NULL, + sandbox_session_id text NOT NULL + ); + CREATE TABLE driver_command ( + id text PRIMARY KEY, + driver_instance_id text NOT NULL, + kind text NOT NULL, + status text NOT NULL + ); + CREATE TABLE api_command ( + id text PRIMARY KEY, + status text NOT NULL, + claim_owner text, + claim_expires_at integer, + kind text NOT NULL, + created_at integer NOT NULL, + attempt_count integer NOT NULL DEFAULT 0, + payload_json text NOT NULL DEFAULT '{}' + ); + CREATE TABLE app_deployment_run ( + id text PRIMARY KEY, + status text NOT NULL + ); + CREATE TABLE external_tool_effect (id text PRIMARY KEY, status text NOT NULL); + `); + return database; +} + +function addPostRuntimeAuthoritySchema(database: SqliteD1Database): void { + database.execute(` + ALTER TABLE sandbox ADD operation_kind text; + ALTER TABLE session ADD archived_at integer; + ALTER TABLE session ADD cleanup_operation_kind text; + ALTER TABLE session ADD runtime_provisioning_operation_id text; + ALTER TABLE session ADD runtime_provisioning_run_id text; + ALTER TABLE session ADD runtime_provisioning_sandbox_id text; + ALTER TABLE session ADD runtime_provisioning_sandbox_session_id text; + ALTER TABLE session ADD runtime_provisioning_sandbox_incarnation integer; + ALTER TABLE session ADD runtime_provisioning_heartbeat_at integer; + ALTER TABLE api_command ADD delivery_generation integer NOT NULL DEFAULT 1; + CREATE TABLE sandbox_backup_staging ( + id text PRIMARY KEY, + sandbox_id text NOT NULL, + workspace_session_id text + ); + CREATE TABLE environment_package_artifact_backup_staging ( + actual_backup_id text, + app_id text NOT NULL, + attempt_count integer NOT NULL, + claim_owner text NOT NULL, + command_id text PRIMARY KEY, + created_at integer NOT NULL, + delivery_generation integer NOT NULL, + dir text NOT NULL, + input_digest text NOT NULL, + paths_json text NOT NULL, + updated_at integer NOT NULL + ); + `); +} + +async function readCutoverObjects(database: SqliteD1Database) { + const row = await database + .prepare(PROTOCOL_V3_CUTOVER_OBJECTS_SQL) + .first>(); + if (row === null) throw new Error("Cutover object inventory returned no row."); + return parseProtocolV3CutoverObjects(d1Json(row)); +} + +async function readCutoverState(database: SqliteD1Database) { + const row = await database + .prepare(PROTOCOL_V3_COMMAND_FREEZE_SQL) + .first>(); + if (row === null) throw new Error("Cutover state returned no row."); + return parseProtocolV3CutoverState(d1Json(row)); +} + +async function readLegacyTerminalIntegrity(database: SqliteD1Database) { + const row = await database + .prepare(PROTOCOL_V3_LEGACY_TERMINAL_INTEGRITY_SQL) + .first>(); + if (row === null) throw new Error("Legacy terminal integrity query returned no row."); + return parseProtocolV3LegacyTerminalIntegrity(d1Json(row)); +} + +async function readLegacyTerminalSourceInventory(database: SqliteD1Database) { + const row = await database + .prepare(PROTOCOL_V3_LEGACY_TERMINAL_SOURCE_INVENTORY_SQL) + .first>(); + if (row === null) throw new Error("Legacy terminal source inventory returned no row."); + return parseProtocolV3LegacyTerminalSourceInventory(d1Json(row)); +} + +function createLegacyTerminalDatabase(): SqliteD1Database { + const database = new SqliteD1Database(); + database.execute(` + CREATE TABLE session ( + id text PRIMARY KEY, + last_run_id text, + message_seq_cursor integer NOT NULL, + runtime_event_seq_cursor integer NOT NULL, + status text NOT NULL, + status_operation_id text + ); + CREATE TABLE session_run ( + completed_at integer, + error_code text, + error_details_json text, + error_message text, + id text PRIMARY KEY, + session_id text NOT NULL, + status text NOT NULL, + status_event text NOT NULL + ); + CREATE TABLE session_event ( + id text PRIMARY KEY, + event_type text NOT NULL, + run_id text, + seq integer NOT NULL, + session_id text NOT NULL, + source_event_id text NOT NULL + ); + CREATE TABLE session_message ( + id text PRIMARY KEY, + plan_json text, + role text NOT NULL, + seq integer NOT NULL, + segments_json text, + session_id text NOT NULL, + session_run_id text + ); + CREATE TABLE session_permission_request ( + id text PRIMARY KEY, + run_id text NOT NULL, + session_id text NOT NULL + ); + CREATE TABLE driver_command (id text PRIMARY KEY, status text NOT NULL); + CREATE TABLE external_tool_effect (id text PRIMARY KEY, status text NOT NULL); + `); + return database; +} + +function createValidLegacyCompletionDatabase(): SqliteD1Database { + const database = createLegacyTerminalDatabase(); + database.execute(` + INSERT INTO session VALUES ('session-1', 'completed-run', 1, 1, 'IDLE', NULL); + INSERT INTO session_run VALUES + (100, NULL, NULL, NULL, 'completed-run', 'session-1', 'completed', 'run.complete'); + INSERT INTO session_event VALUES + ('completed-event', 'run.completed', 'completed-run', 1, 'session-1', 'session-run-terminal:completed-run:run.completed'); + INSERT INTO session_message VALUES + ('assistant-1', '[]', 'assistant', 1, '[]', 'session-1', 'completed-run'); + `); + return database; +} + +function containerPage( + instances: readonly Record[], + pageToken: string | null, + nextPageToken: string | null, +): string { + return JSON.stringify({ + instances, + result_info: { next_page_token: nextPageToken, page_token: pageToken, per_page: 100 }, + }); +} + +function containerApplication(name: string, version: number): Record { + return { + configuration: { + image: `registry.cloudflare.com/account/${name}@sha256:${CONTAINER_IMAGE_DIGEST}`, + }, + health: { + instances: { active: 1, failed: 0, healthy: 1, scheduling: 0, starting: 0 }, + }, + id: crypto.randomUUID(), + name, + version, + }; +} + +function containerApplicationPage( + applications: readonly Record[], + pageToken: string | null, + nextPageToken: string | null, +): Record { + return { + result: applications, + result_info: { next_page_token: nextPageToken, page_token: pageToken, per_page: 100 }, + success: true, + }; +} + +interface PaginationFailureFixture { + readonly collect: (readPage: (pageToken: string | null) => T) => Promise; + readonly invalidMessage: string; + readonly invalidPage: T; + readonly label: string; + readonly page: (pageToken: string | null, nextPageToken: string | null) => T; +} + +function registerPaginationFailureTests(fixture: PaginationFailureFixture): void { + test(`fails closed on repeated ${fixture.label} page tokens`, async () => { + await expect(fixture.collect((pageToken) => fixture.page(pageToken, "page-2"))).rejects.toThrow( + "repeated page token", + ); + }); + + test(`propagates a later ${fixture.label} page failure`, async () => { + await expect( + fixture.collect((pageToken) => { + if (pageToken === null) return fixture.page(null, "page-2"); + throw new Error(`second ${fixture.label} page failed`); + }), + ).rejects.toThrow(`second ${fixture.label} page failed`); + }); + + test(`rejects an unexpected ${fixture.label} pagination shape`, async () => { + await expect(fixture.collect(() => fixture.invalidPage)).rejects.toThrow( + fixture.invalidMessage, + ); + }); + + test(`bounds ${fixture.label} pagination`, async () => { + let pagesRead = 0; + await expect( + fixture.collect((pageToken) => { + pagesRead += 1; + return fixture.page(pageToken, `page-${pagesRead}`); + }), + ).rejects.toThrow("exceeded its safety limit"); + expect(pagesRead).toBe(100); + }); +} + +describe("protocol v3 production cutover", () => { + test("closes new runtime admission while existing work drains", async () => { + const database = createCutoverGateDatabase(); + const smokeAccountId = "01J00000000000000000000001"; + const smokeSessionId = "01J0000000000000000000000A"; + const otherSessionId = "01J0000000000000000000000B"; + const smokeRequestKey = "protocol-v3-cutover-2de67417-f144-47e1-8b2e-1e71621a0d92"; + database.execute(` + INSERT INTO session VALUES + ('${smokeSessionId}', '${smokeAccountId}', '${smokeRequestKey}', 'IDLE', NULL), + ('${otherSessionId}', '${smokeAccountId}', 'ordinary-user', 'IDLE', NULL); + INSERT INTO sandbox VALUES + ('existing-sandbox', 'session', '${smokeSessionId}', 'active', NULL, 'worker-1', 100), + ('other-sandbox', 'session', '${otherSessionId}', 'cold', NULL, NULL, NULL); + INSERT INTO session_run VALUES + ('existing-run', 'queued', '${smokeAccountId}', '${smokeSessionId}'); + INSERT INTO app_deployment_run VALUES + ('existing-app-run', 'building'), + ('freeze-app-run', 'preparing'); + INSERT INTO driver_instance VALUES + ('existing-driver', 'ready', 'existing-sandbox', '${smokeSessionId}'), + ('other-existing-driver', 'stopped', 'other-sandbox', '${otherSessionId}'); + INSERT INTO api_command ( + id, status, claim_owner, claim_expires_at, kind, created_at + ) VALUES ( + 'existing-api-command', 'running', 'worker-1', 100, + 'environment_package_artifact_build', 0 + ); + `); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + + for (const kind of [ + "session_run_dispatch", + "app_deployment_run_dispatch", + "environment_package_artifact_build", + ]) { + for (let request = 0; request < 3; request += 1) { + await expect( + database + .prepare( + `INSERT INTO api_command ( + id, status, claim_owner, claim_expires_at, kind, created_at + ) VALUES (?, 'queued', NULL, NULL, ?, 1)`, + ) + .bind(`blocked-${kind}-${request}`, kind) + .run(), + ).rejects.toThrow("blocks new nonterminal API commands"); + } + } + for (const kind of [ + "app_deployment_script_reconciliation", + "cost_ledger_reconciliation", + "sandbox_backup_reconciliation", + "scheduled_maintenance", + ]) { + await expect( + database + .prepare( + `INSERT INTO api_command ( + id, status, claim_owner, claim_expires_at, kind, created_at + ) VALUES (?, 'queued', NULL, NULL, ?, 1)`, + ) + .bind(`drain-${kind}`, kind) + .run(), + ).resolves.toBeDefined(); + } + database.execute(` + INSERT INTO api_command ( + id, status, claim_owner, claim_expires_at, kind, created_at + ) VALUES ( + 'terminal-business-command', 'failed', NULL, NULL, + 'environment_package_artifact_build', 0 + ) + `); + await expect( + database + .prepare("UPDATE api_command SET status = 'queued' WHERE id = 'terminal-business-command'") + .run(), + ).rejects.toThrow("blocks API command admission"); + await expect( + database + .prepare( + `UPDATE api_command + SET kind = 'session_run_dispatch' + WHERE id = 'drain-scheduled_maintenance'`, + ) + .run(), + ).rejects.toThrow("blocks API command admission"); + + await expect( + database.prepare("UPDATE session_run SET status = 'running' WHERE id = 'existing-run'").run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare("UPDATE session_run SET status = 'completed' WHERE id = 'existing-run'") + .run(), + ).resolves.toBeDefined(); + await expect( + database.prepare("UPDATE session_run SET status = 'queued' WHERE id = 'existing-run'").run(), + ).rejects.toThrow("blocks Session Run reactivation"); + await expect( + database + .prepare( + `INSERT INTO session_run VALUES ('new-run', 'queued', '${smokeAccountId}', '${smokeSessionId}')`, + ) + .run(), + ).rejects.toThrow("blocks new active Session Runs"); + await expect( + database + .prepare( + "UPDATE app_deployment_run SET status = 'submitting' WHERE id = 'existing-app-run'", + ) + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare("UPDATE app_deployment_run SET status = 'success' WHERE id = 'existing-app-run'") + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare( + "UPDATE app_deployment_run SET status = 'activating' WHERE id = 'existing-app-run'", + ) + .run(), + ).rejects.toThrow("blocks App deployment Run reactivation"); + await expect( + database.prepare("INSERT INTO app_deployment_run VALUES ('new-app-run', 'queued')").run(), + ).rejects.toThrow("blocks new active App deployment Runs"); + + await expect( + database + .prepare("UPDATE driver_instance SET status = 'stopping' WHERE id = 'existing-driver'") + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare("UPDATE driver_instance SET status = 'stopped' WHERE id = 'existing-driver'") + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare("UPDATE driver_instance SET status = 'ready' WHERE id = 'existing-driver'") + .run(), + ).rejects.toThrow("blocks Driver reactivation"); + await expect( + database + .prepare( + `INSERT INTO driver_instance VALUES ('new-driver', 'provisioning', 'existing-sandbox', '${smokeSessionId}')`, + ) + .run(), + ).rejects.toThrow("blocks new live Driver instances"); + + for (const kind of ["input.start", "mcp.execute"]) { + await expect( + database + .prepare("INSERT INTO driver_command VALUES (?, 'existing-driver', ?, 'queued')") + .bind(kind, kind) + .run(), + ).rejects.toThrow("blocks new Driver commands"); + } + for (const kind of ["permission.resolve", "session.stop", "turn.cancel"]) { + await expect( + database + .prepare("INSERT INTO driver_command VALUES (?, 'existing-driver', ?, 'queued')") + .bind(kind, kind) + .run(), + ).resolves.toBeDefined(); + } + + await expect( + database + .prepare( + "UPDATE sandbox SET status = 'destroying', claim_owner = NULL, claim_expires_at = NULL WHERE id = 'existing-sandbox'", + ) + .run(), + ).resolves.toBeDefined(); + await expect( + database.prepare("UPDATE sandbox SET status = 'cold' WHERE id = 'existing-sandbox'").run(), + ).resolves.toBeDefined(); + await expect( + database.prepare("UPDATE sandbox SET status = 'active' WHERE id = 'existing-sandbox'").run(), + ).rejects.toThrow("blocks sandbox activation"); + await expect( + database + .prepare("INSERT INTO sandbox_backup VALUES ('staging-backup', 'creating', NULL)") + .run(), + ).rejects.toThrow("blocks new sandbox backup work"); + + database.execute(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + const stillDraining = await database + .prepare(PROTOCOL_V3_COMMAND_FREEZE_SQL) + .first>(); + expect(parseProtocolV3CommandFreeze(d1Json(stillDraining))).toBeFalse(); + await expect( + database + .prepare( + "UPDATE api_command SET status = 'succeeded', claim_owner = NULL, claim_expires_at = NULL WHERE id = 'existing-api-command'", + ) + .run(), + ).resolves.toBeDefined(); + database.execute(` + UPDATE api_command + SET status = 'succeeded' + WHERE id LIKE 'drain-%' + `); + + database.execute(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + const appRunStillActive = await database + .prepare(PROTOCOL_V3_COMMAND_FREEZE_SQL) + .first>(); + expect(parseProtocolV3CommandFreeze(d1Json(appRunStillActive))).toBeFalse(); + database.execute("UPDATE app_deployment_run SET status = 'failed' WHERE id = 'freeze-app-run'"); + + database.execute(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + const frozen = await database + .prepare(PROTOCOL_V3_COMMAND_FREEZE_SQL) + .first>(); + expect(parseProtocolV3CommandFreeze(d1Json(frozen))).toBeTrue(); + await expect( + database + .prepare( + "INSERT INTO driver_command VALUES ('frozen-control', 'existing-driver', 'session.stop', 'queued')", + ) + .run(), + ).rejects.toThrow("blocks new Driver commands"); + + await expect( + database + .prepare("UPDATE api_command SET status = 'queued' WHERE id = 'existing-api-command'") + .run(), + ).rejects.toThrow("blocks API command admission"); + await expect( + database + .prepare( + "UPDATE api_command SET status = 'queued' WHERE id = 'drain-scheduled_maintenance'", + ) + .run(), + ).rejects.toThrow("blocks API command admission"); + for (const status of ["queued", "running"]) { + await expect( + database + .prepare( + `INSERT INTO api_command ( + id, status, claim_owner, claim_expires_at, kind, created_at + ) VALUES (?, ?, NULL, NULL, 'scheduled_maintenance', 1)`, + ) + .bind(`frozen-${status}`, status) + .run(), + ).rejects.toThrow("blocks new nonterminal API commands"); + } + + database.execute(openProtocolV3SmokeWindowSql(smokeAccountId)); + await expect( + database + .prepare( + `INSERT INTO session_run VALUES ('smoke-run-before-request', 'queued', '${smokeAccountId}', '${smokeSessionId}')`, + ) + .run(), + ).rejects.toThrow("blocks new active Session Runs"); + database.execute(storeProtocolV3SmokeRequestKeySql(smokeRequestKey)); + await expect( + database + .prepare( + `INSERT INTO session_run VALUES ('smoke-run', 'queued', '${smokeAccountId}', '${smokeSessionId}')`, + ) + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare( + `INSERT INTO session_run VALUES ('same-account-run', 'queued', '${smokeAccountId}', '${otherSessionId}')`, + ) + .run(), + ).rejects.toThrow("blocks new active Session Runs"); + await expect( + database + .prepare( + `INSERT INTO driver_instance VALUES ('smoke-driver', 'provisioning', 'existing-sandbox', '${smokeSessionId}')`, + ) + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare( + `INSERT INTO driver_instance VALUES ('other-driver', 'provisioning', 'other-sandbox', '${otherSessionId}')`, + ) + .run(), + ).rejects.toThrow("blocks new live Driver instances"); + await expect( + database.prepare("UPDATE sandbox SET status = 'active' WHERE id = 'existing-sandbox'").run(), + ).resolves.toBeDefined(); + await expect( + database.prepare("UPDATE sandbox SET status = 'active' WHERE id = 'other-sandbox'").run(), + ).rejects.toThrow("blocks sandbox activation"); + await expect( + database + .prepare( + `INSERT INTO sandbox_session VALUES ('${smokeSessionId}', 'existing-sandbox', 'active')`, + ) + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare( + `INSERT INTO sandbox_session VALUES ('${otherSessionId}', 'other-sandbox', 'active')`, + ) + .run(), + ).rejects.toThrow("blocks new active sandbox Sessions"); + await expect( + database + .prepare( + `UPDATE session SET status_operation_id = 'smoke-operation' WHERE id = '${smokeSessionId}'`, + ) + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare( + `UPDATE session SET status_operation_id = 'ordinary-operation' WHERE id = '${otherSessionId}'`, + ) + .run(), + ).rejects.toThrow("blocks Session operation acquisition"); + await expect( + database + .prepare( + "INSERT INTO driver_command VALUES ('smoke-input', 'smoke-driver', 'input.start', 'queued')", + ) + .run(), + ).rejects.toThrow("blocks new Driver commands"); + await expect( + database + .prepare( + "INSERT INTO driver_command VALUES ('smoke-stop', 'smoke-driver', 'session.stop', 'queued')", + ) + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare( + "INSERT INTO driver_command VALUES ('other-stop', 'other-existing-driver', 'session.stop', 'queued')", + ) + .run(), + ).rejects.toThrow("blocks new Driver commands"); + + database.execute(CLOSE_PROTOCOL_V3_SMOKE_WINDOW_SQL); + await expect( + database + .prepare( + "INSERT INTO driver_command VALUES ('closed-again', 'existing-driver', 'session.stop', 'queued')", + ) + .run(), + ).rejects.toThrow("blocks new Driver commands"); + + database.execute(REMOVE_PROTOCOL_V3_CUTOVER_SQL); + await expect( + database + .prepare( + `INSERT INTO session_run VALUES ('after-cutover', 'queued', '${smokeAccountId}', '${smokeSessionId}')`, + ) + .run(), + ).resolves.toBeDefined(); + }); + + test("parses the fail-closed migration and drain probes", () => { + expect( + parseProtocolV3CutoverProbe( + `[wrangler warning]\n${d1Json({ + gate_present: 1, + })}`, + ), + ).toEqual({ gatePresent: true }); + expect(() => parseProtocolV3CutoverProbe(d1Json({ gate_present: 2 }))).toThrow("zero or one"); + + const localMigrations = [ + "0013_durable-mcp-effect-v3.sql", + "0014_session-event-stream-identity.sql", + "0019_runtime-subject-operation-authority.sql", + ]; + expect(() => assertCutoverMigrationJournalAudited(AUDITED_MIGRATION_NAMES)).not.toThrow(); + expect(() => + assertCutoverMigrationJournalAudited([...AUDITED_MIGRATION_NAMES, "0021_unreviewed.sql"]), + ).toThrow("audited only through 0020_sandbox-backup-object-authority.sql"); + expect(() => + assertCutoverMigrationJournalAudited([ + ...AUDITED_MIGRATION_NAMES.slice(0, -1), + "0018_z_unreviewed.sql", + AUDITED_MIGRATION_NAMES.at(-1) ?? "", + ]), + ).toThrow("audited only through 0020_sandbox-backup-object-authority.sql"); + expect(() => + assertCutoverMigrationJournalAudited(AUDITED_MIGRATION_NAMES.with(18, "0018_replaced.sql")), + ).toThrow("audited only through 0020_sandbox-backup-object-authority.sql"); + const remoteLedger = JSON.stringify([ + { + results: localMigrations.slice(0, 2).map((name) => ({ name })), + success: true, + }, + ]); + const pending = findPendingProdMigrations(remoteLedger, localMigrations); + expect(pending).toEqual(["0019_runtime-subject-operation-authority.sql"]); + expect( + findPendingProdMigrations( + JSON.stringify([ + { + results: [{ name: localMigrations[1] }, { name: localMigrations[0] }], + success: true, + }, + ]), + localMigrations, + ), + ).toEqual(["0019_runtime-subject-operation-authority.sql"]); + expect(() => + findPendingProdMigrations( + JSON.stringify([ + { + results: [{ name: localMigrations[0] }, { name: localMigrations[2] }], + success: true, + }, + ]), + localMigrations, + ), + ).toThrow("not an exact prefix"); + const blocked = parseProtocolV3CutoverDrain( + d1Json({ + ...DRAINED_CUTOVER_ROW, + nonterminal_commands: 1, + }), + ); + expect(isProtocolV3RuntimeDrained(blocked)).toBeTrue(); + expect(isProtocolV3CutoverDrained(blocked)).toBeFalse(); + const claimedEffect = parseProtocolV3CutoverDrain( + d1Json({ + ...DRAINED_CUTOVER_ROW, + unsettled_effects: 1, + }), + ); + expect(isProtocolV3RuntimeDrained(claimedEffect)).toBeFalse(); + expect(isProtocolV3CutoverDrained(claimedEffect)).toBeFalse(); + for (const blockedAuthority of [ + { active_app_deployment_runs: 1 }, + { nonterminal_api_commands: 1 }, + { unsafe_environment_artifact_backup_staging: 1 }, + ]) { + const state = parseProtocolV3CutoverDrain( + d1Json({ ...DRAINED_CUTOVER_ROW, ...blockedAuthority }), + ); + expect(isProtocolV3RuntimeDrained(state)).toBeFalse(); + expect(isProtocolV3CutoverDrained(state)).toBeFalse(); + } + for (const activeRuntime of [ + { active_runs: 1, live_drivers: 0 }, + { active_runs: 0, live_drivers: 1 }, + ]) { + const state = parseProtocolV3CutoverDrain( + d1Json({ + ...DRAINED_CUTOVER_ROW, + ...activeRuntime, + }), + ); + expect(isProtocolV3RuntimeDrained(state)).toBeFalse(); + expect(isProtocolV3CutoverDrained(state)).toBeFalse(); + } + expect( + isProtocolV3CutoverDrained( + parseProtocolV3CutoverDrain( + d1Json({ + ...DRAINED_CUTOVER_ROW, + }), + ), + ), + ).toBeTrue(); + }); + + test("bounds the read-only migration 0013 loss inventory and fails on any candidate", async () => { + const database = new SqliteD1Database(); + applyDrizzleMigrationsThrough(database, "0012_agent-task-snapshot-state"); + const row = await database + .prepare(PROTOCOL_V3_LOSSY_MIGRATION_INVENTORY_SQL) + .first>(); + const empty = parseProtocolV3LossyMigrationInventory(d1Json(row ?? {})); + expect(empty.totalCandidates).toBe(0); + expect(empty.candidateIds).toEqual([]); + expect(() => assertProtocolV3LossyMigrationInventory(empty)).not.toThrow(); + + const candidateIds = Array.from({ length: 50 }, (_, index) => [ + "mcp_argument_omission", + `01${"0".repeat(22)}${String(index).padStart(2, "0")}`, + ]); + const blockedRow = { + attempt_completion_time_fabrications: 0, + candidate_ids_json: JSON.stringify(candidateIds), + command_error_omissions: 0, + command_payload_conflicts: 0, + control_reason_omissions: 0, + input_start_result_omissions: 0, + input_text_omissions: 0, + mcp_argument_omissions: 51, + mcp_command_terminal_conflicts: 0, + mcp_result_conflicts: 0, + mcp_result_omissions: 0, + orphan_effects: 0, + provider_receipt_losses: 0, + permission_payload_rewrites: 0, + session_run_error_omissions: 0, + total_candidates: 51, + }; + const blocked = parseProtocolV3LossyMigrationInventory(d1Json(blockedRow)); + expect(blocked.candidateIds).toHaveLength(50); + expect(() => assertProtocolV3LossyMigrationInventory(blocked)).toThrow( + "No lossy migration candidates are authorized", + ); + expect(() => + parseProtocolV3LossyMigrationInventory( + d1Json({ + ...blockedRow, + candidate_ids_json: JSON.stringify([ + ...candidateIds, + ["mcp_argument_omission", "01000000000000000000000050"], + ]), + }), + ), + ).toThrow("candidate ID count is inconsistent"); + }); + + test("runs the migration 0013 loss preflight before migrations and keeps a direct guard", async () => { + const deploySource = await Bun.file(new URL("../bin/deploy-prod.ts", import.meta.url)).text(); + const cutoverSource = deploySource.slice( + deploySource.indexOf("async function runProtocolV3Cutover"), + ); + expect(cutoverSource.indexOf("verifyProdLossyMigrationInventory();")).toBeGreaterThanOrEqual(0); + expect(cutoverSource.indexOf("verifyProdLossyMigrationInventory();")).toBeLessThan( + cutoverSource.indexOf("applyD1Migrations();"), + ); + + const migrationSource = await Bun.file( + new URL("../../../pkgs/db/drizzle/0013_durable-mcp-effect-v3.sql", import.meta.url), + ).text(); + const lossGuard = migrationSource.indexOf("__durable_mcp_v3_loss_guard"); + expect(lossGuard).toBeGreaterThanOrEqual(0); + expect(lossGuard).toBeLessThan(migrationSource.indexOf("ALTER TABLE `driver_command`")); + expect(lossGuard).toBeLessThan(migrationSource.indexOf("UPDATE `driver_command`")); + }); + + test("rebuilds the exact post-migration gate before old Workers can reacquire work", async () => { + const database = createCutoverGateDatabase(); + const accountId = "01J00000000000000000000001"; + const smokeSessionId = "01J0000000000000000000000A"; + const otherSessionId = "01J0000000000000000000000B"; + const requestKey = "protocol-v3-cutover-2de67417-f144-47e1-8b2e-1e71621a0d92"; + database.execute(` + INSERT INTO session VALUES + ('${smokeSessionId}', '${accountId}', '${requestKey}', 'IDLE', NULL), + ('${otherSessionId}', '${accountId}', 'ordinary-user', 'IDLE', NULL); + INSERT INTO sandbox VALUES + ('smoke-sandbox', 'session', '${smokeSessionId}', 'cold', NULL, NULL, NULL), + ('other-sandbox', 'session', '${otherSessionId}', 'cold', NULL, NULL, NULL); + INSERT INTO api_command ( + id, status, claim_owner, claim_expires_at, kind, created_at, + attempt_count, payload_json + ) VALUES ( + 'artifact-command', 'running', 'worker-1', 9000000000000000, + 'environment_package_artifact_build', 0, 1, + '{"appId":"app-1","inputDigest":"${"0".repeat(64)}"}' + ); + `); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + addPostRuntimeAuthoritySchema(database); + database.execute(INSTALL_PROTOCOL_V3_POST_MIGRATION_CUTOVER_SQL); + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + objectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + }); + database.execute( + `INSERT INTO api_command ( + id, status, claim_owner, claim_expires_at, kind, created_at, + attempt_count, payload_json, delivery_generation + ) VALUES ( + 'terminal-command', 'succeeded', NULL, NULL, 'scheduled_maintenance', 0, + 0, '{}', 1 + )`, + ); + + const gate = await database + .prepare('SELECT started_at AS startedAt FROM "__protocol_v3_cutover" WHERE id = 1') + .first<{ startedAt: number }>(); + if (gate === null) throw new Error("Expected the protocol v3 cutover gate."); + database.execute( + `UPDATE api_command SET created_at = ${gate.startedAt + 1} + WHERE id = 'artifact-command'`, + ); + await expect(database.prepare(INSERT_ENVIRONMENT_ARTIFACT_STAGING_SQL).run()).rejects.toThrow( + "blocks new environment artifact backup staging", + ); + database.execute("UPDATE api_command SET created_at = 0 WHERE id = 'artifact-command'"); + await expect( + database + .prepare(INSERT_ENVIRONMENT_ARTIFACT_STAGING_SQL.replace("'worker-1'", "'wrong-owner'")) + .run(), + ).rejects.toThrow("blocks new environment artifact backup staging"); + await expect( + database.prepare(INSERT_ENVIRONMENT_ARTIFACT_STAGING_SQL).run(), + ).resolves.toBeDefined(); + database.execute(` + DELETE FROM environment_package_artifact_backup_staging + WHERE command_id = 'artifact-command'; + UPDATE api_command + SET status = 'succeeded', claim_owner = NULL, claim_expires_at = NULL + WHERE id = 'artifact-command'; + `); + database.execute(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + await expect( + database + .prepare("UPDATE api_command SET delivery_generation = 2 WHERE id = 'terminal-command'") + .run(), + ).rejects.toThrow("blocks API command admission"); + await expect(database.prepare(INSERT_ENVIRONMENT_ARTIFACT_STAGING_SQL).run()).rejects.toThrow( + "blocks new environment artifact backup staging", + ); + + await expect( + database + .prepare( + `UPDATE session SET runtime_provisioning_operation_id = 'ordinary-operation' WHERE id = '${otherSessionId}'`, + ) + .run(), + ).rejects.toThrow("blocks Session operation acquisition"); + await expect( + database + .prepare( + `UPDATE session SET archived_at = 1, cleanup_operation_kind = 'archive' WHERE id = '${otherSessionId}'`, + ) + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare( + `UPDATE session SET cleanup_operation_kind = 'delete' WHERE id = '${otherSessionId}'`, + ) + .run(), + ).rejects.toThrow("blocks Session operation acquisition"); + await expect( + database + .prepare("UPDATE sandbox SET operation_kind = 'activate' WHERE id = 'other-sandbox'") + .run(), + ).rejects.toThrow("blocks sandbox activation"); + await expect( + database + .prepare( + `INSERT INTO sandbox_backup_staging VALUES ('ordinary-staging', 'other-sandbox', '${otherSessionId}')`, + ) + .run(), + ).rejects.toThrow("blocks new sandbox backup staging"); + + database.execute(openProtocolV3SmokeWindowSql(accountId)); + database.execute(storeProtocolV3SmokeRequestKeySql(requestKey)); + await expect( + database + .prepare( + `UPDATE session SET runtime_provisioning_operation_id = 'smoke-operation' WHERE id = '${smokeSessionId}'`, + ) + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare("UPDATE sandbox SET operation_kind = 'activate' WHERE id = 'smoke-sandbox'") + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare( + `INSERT INTO sandbox_backup_staging VALUES ('smoke-staging', 'smoke-sandbox', '${smokeSessionId}')`, + ) + .run(), + ).resolves.toBeDefined(); + + database.execute(REMOVE_PROTOCOL_V3_CUTOVER_SQL); + database.execute(INSTALL_PROTOCOL_V3_POST_MIGRATION_CUTOVER_SQL); + expect( + await database + .prepare('SELECT count(*) AS count FROM "__protocol_v3_cutover"') + .first<{ count: number }>(), + ).toEqual({ count: 0 }); + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + objectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + }); + database.execute(installProtocolV3PostMigrationCutoverSql(RELEASE_TREE_OID)); + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + objectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + }); + }); + + test("migration 0020 preserves the exact preinstalled cutover gate", async () => { + const database = new SqliteD1Database(); + applyDrizzleMigrationsThrough(database, "0019_runtime-subject-operation-authority"); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + + applyDrizzleMigration(database, "0020_sandbox-backup-object-authority"); + + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + objectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + }); + }); + + test("reports legacy identity and backup rows that migration 0019 will reject atomically", async () => { + const database = new SqliteD1Database(); + database.execute(` + CREATE TABLE app (id text PRIMARY KEY); + CREATE TABLE agent ( + id text PRIMARY KEY, + app_id text NOT NULL, + kind text NOT NULL, + owner_account_id text NOT NULL + ); + CREATE TABLE session ( + id text PRIMARY KEY, + agent_id text NOT NULL, + app_id text NOT NULL, + kind text NOT NULL, + archived_at integer, + cleanup_operation_kind text, + runtime_provisioning_heartbeat_at integer, + runtime_provisioning_operation_id text, + runtime_provisioning_run_id text, + runtime_provisioning_sandbox_id text, + status text NOT NULL, + status_operation_id text + ); + CREATE TABLE session_run ( + id text PRIMARY KEY, + session_id text NOT NULL, + agent_id text NOT NULL, + status text NOT NULL + ); + CREATE TABLE sandbox_session ( + session_id text NOT NULL, + sandbox_id text NOT NULL, + cwd text NOT NULL, + status text NOT NULL + ); + CREATE TABLE sandbox ( + id text PRIMARY KEY, + kind text NOT NULL, + subject_kind text NOT NULL, + subject_id text NOT NULL, + agent_id text, + app_id text, + owner_account_id text, + last_backup_id text, + last_restore_backup_id text + ); + CREATE TABLE driver_instance (id text PRIMARY KEY, generation); + CREATE TABLE app_deployment (id text PRIMARY KEY, last_successful_url text); + CREATE TABLE preflight_parent (id text PRIMARY KEY); + CREATE TABLE preflight_child ( + id text PRIMARY KEY, + parent_id text REFERENCES preflight_parent(id) + ); + CREATE TABLE sandbox_backup ( + id text PRIMARY KEY, + sandbox_id text NOT NULL, + dir text NOT NULL, + session_run_id text, + status text NOT NULL, + error_message text, + keep integer NOT NULL, + ttl_seconds integer NOT NULL, + created_at integer NOT NULL, + updated_at integer NOT NULL + ); + INSERT INTO app VALUES ('app-1'); + INSERT INTO agent VALUES ('agent-1', 'app-1', 'cattle', 'account-1'); + INSERT INTO session ( + id, agent_id, app_id, kind, status + ) VALUES ('session-1', 'agent-1', 'app-1', 'cattle', 'IDLE'); + INSERT INTO session_run VALUES ('run-1', 'session-1', 'agent-1', 'completed'); + INSERT INTO sandbox_session VALUES ('session-1', 'sandbox-1', '/workspace', 'closed'); + INSERT INTO sandbox ( + id, kind, subject_kind, subject_id, agent_id, app_id, owner_account_id + ) VALUES ('sandbox-1', 'cattle', 'session', 'session-1', NULL, NULL, NULL); + INSERT INTO sandbox_backup VALUES ( + 'backup-1', 'sandbox-1', '/workspace', 'run-1', 'ready', NULL, 0, 3600, 1, 2 + ); + `); + const readPreflight = async () => + parseProtocolV3RuntimeAuthorityPreflight( + d1Json( + await database + .prepare(PROTOCOL_V3_RUNTIME_AUTHORITY_PREFLIGHT_SQL) + .first>(), + ), + ); + const valid = await readPreflight(); + expect(() => assertProtocolV3RuntimeAuthorityPreflight(valid)).not.toThrow(); + + database.execute("INSERT INTO app_deployment VALUES ('deployment-1', 'https://legacy.test')"); + expect((await readPreflight()).legacyAppDeploymentTraffic).toBe(1); + database.execute("DELETE FROM app_deployment"); + + database.execute("INSERT INTO driver_instance VALUES ('driver-1', 1.5)"); + expect((await readPreflight()).invalidDriverGenerations).toBe(1); + database.execute("DELETE FROM driver_instance"); + + database.execute("UPDATE sandbox SET last_backup_id = 'missing' WHERE id = 'sandbox-1'"); + expect((await readPreflight()).invalidSandboxBackupPointers).toBe(1); + database.execute("UPDATE sandbox SET last_backup_id = NULL WHERE id = 'sandbox-1'"); + + database.execute( + "UPDATE session SET runtime_provisioning_operation_id = 'operation-1' WHERE id = 'session-1'", + ); + expect((await readPreflight()).nonstaticSessions).toBe(1); + database.execute( + "UPDATE session SET runtime_provisioning_operation_id = NULL WHERE id = 'session-1'", + ); + + database.execute("UPDATE sandbox_session SET sandbox_id = 'other-sandbox'"); + expect((await readPreflight()).invalidSandboxSessionAuthorities).toBe(1); + database.execute("UPDATE sandbox_session SET sandbox_id = 'sandbox-1'"); + + database.execute(` + PRAGMA foreign_keys = OFF; + INSERT INTO preflight_child VALUES ('child-1', 'missing-parent'); + PRAGMA foreign_keys = ON; + `); + expect((await readPreflight()).foreignKeyViolations).toBe(1); + database.execute(` + PRAGMA foreign_keys = OFF; + DELETE FROM preflight_child; + PRAGMA foreign_keys = ON; + `); + + for (const [breakAuthority, restoreAuthority] of [ + [ + "UPDATE session_run SET status = 'failed' WHERE id = 'run-1'", + "UPDATE session_run SET status = 'completed' WHERE id = 'run-1'", + ], + [ + "UPDATE session_run SET agent_id = 'other-agent' WHERE id = 'run-1'", + "UPDATE session_run SET agent_id = 'agent-1' WHERE id = 'run-1'", + ], + [ + "UPDATE sandbox_session SET sandbox_id = 'other-sandbox'", + "UPDATE sandbox_session SET sandbox_id = 'sandbox-1'", + ], + [ + "UPDATE sandbox_session SET cwd = '/other'", + "UPDATE sandbox_session SET cwd = '/workspace'", + ], + [ + "UPDATE sandbox_backup SET dir = '' WHERE id = 'backup-1'", + "UPDATE sandbox_backup SET dir = '/workspace' WHERE id = 'backup-1'", + ], + ] as const) { + database.execute(breakAuthority); + expect((await readPreflight()).invalidSandboxBackups).toBe(1); + database.execute(restoreAuthority); + } + + database.execute("UPDATE sandbox SET agent_id = 'agent-1' WHERE id = 'sandbox-1'"); + expect((await readPreflight()).invalidSandboxIdentities).toBe(1); + database.execute( + "UPDATE sandbox SET agent_id = NULL, app_id = NULL, owner_account_id = NULL WHERE id = 'sandbox-1'", + ); + database.execute(` + INSERT INTO sandbox_backup VALUES ( + 'backup-2', 'sandbox-1', '/workspace', 'run-1', 'pruned', NULL, 1, 3600, 2, 3 + ); + `); + const invalid = await readPreflight(); + expect(invalid.duplicateSandboxBackups).toBe(1); + expect(() => assertProtocolV3RuntimeAuthorityPreflight(invalid)).toThrow( + "Migration 0019 remains the atomic authority", + ); + }); + + test("preflights the real 0012 schema without reading 0015 Session columns", async () => { + const database = new SqliteD1Database(); + applyDrizzleMigrationsThrough(database, "0012_agent-task-snapshot-state"); + const readPreflight = async () => + parseProtocolV3RuntimeAuthorityPreflight( + d1Json( + await database + .prepare(protocolV3RuntimeAuthorityPreflightSql(true)) + .first>(), + ), + ); + + const clean = await readPreflight(); + expect(() => assertProtocolV3RuntimeAuthorityPreflight(clean)).not.toThrow(); + database.execute(` + INSERT INTO session ( + agent_id, created_at, creator_account_id, id, kind, model, app_id, + provider, renamed, runtime_id, status, updated_at + ) VALUES ( + '01J0000000000000000000001A', 1, '01J0000000000000000000001B', + '01J0000000000000000000001C', 'cattle', 'model', + '01J0000000000000000000001D', 'provider', 0, 'codex', 'RUNNING', 1 + ) + `); + expect((await readPreflight()).nonstaticSessions).toBe(1); + }); + + test("checks 0015 Session authority columns once migration 0015 is applied", async () => { + const database = new SqliteD1Database(); + applyDrizzleMigrationsThrough(database, "0015_session-cleanup-operation"); + database.execute(` + INSERT INTO session ( + agent_id, created_at, creator_account_id, id, kind, model, app_id, + provider, renamed, runtime_id, runtime_provisioning_heartbeat_at, + runtime_provisioning_operation_id, runtime_provisioning_sandbox_id, + status, updated_at + ) VALUES ( + '01J0000000000000000000001A', 1, '01J0000000000000000000001B', + '01J0000000000000000000001C', 'cattle', 'model', + '01J0000000000000000000001D', 'provider', 0, 'codex', 1, + '01J0000000000000000000001E', '01J0000000000000000000001F', + 'IDLE', 1 + ) + `); + const preflight = parseProtocolV3RuntimeAuthorityPreflight( + d1Json( + await database + .prepare(protocolV3RuntimeAuthorityPreflightSql(false)) + .first>(), + ), + ); + + expect(preflight.nonstaticSessions).toBe(1); + expect(() => assertProtocolV3RuntimeAuthorityPreflight(preflight)).toThrow( + "nonstaticSessions=1", + ); + }); + + test("counts every post-migration authority lease before declaring the boundary drained", async () => { + const database = createCutoverGateDatabase(); + addPostRuntimeAuthoritySchema(database); + database.execute(` + INSERT INTO session VALUES ( + 'session-1', 'account-1', 'user-1', 'IDLE', NULL, + NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL + ); + INSERT INTO sandbox VALUES ( + 'sandbox-1', 'session', 'session-1', 'cold', NULL, NULL, NULL, NULL + ); + `); + const readDrain = async () => + parseProtocolV3CutoverDrain( + d1Json( + await database + .prepare(PROTOCOL_V3_POST_MIGRATION_CUTOVER_DRAIN_SQL) + .first>(), + ), + ); + expect(isProtocolV3CutoverDrained(await readDrain())).toBeTrue(); + + for (const [column, value] of [ + ["cleanup_operation_kind", "'archive'"], + ["runtime_provisioning_operation_id", "'operation-1'"], + ["runtime_provisioning_run_id", "'run-1'"], + ["runtime_provisioning_sandbox_id", "'sandbox-1'"], + ["runtime_provisioning_sandbox_session_id", "'sandbox-session-1'"], + ["runtime_provisioning_sandbox_incarnation", "1"], + ["runtime_provisioning_heartbeat_at", "1"], + ] as const) { + database.execute(`UPDATE session SET "${column}" = ${value} WHERE id = 'session-1'`); + expect((await readDrain()).unsafeSessions).toBe(1); + database.execute(`UPDATE session SET "${column}" = NULL WHERE id = 'session-1'`); + } + + database.execute( + "UPDATE session SET archived_at = 1, cleanup_operation_kind = 'archive' WHERE id = 'session-1'", + ); + expect(isProtocolV3CutoverDrained(await readDrain())).toBeTrue(); + database.execute("UPDATE session SET cleanup_operation_kind = 'delete' WHERE id = 'session-1'"); + expect((await readDrain()).unsafeSessions).toBe(1); + database.execute( + "UPDATE session SET archived_at = NULL, cleanup_operation_kind = NULL WHERE id = 'session-1'", + ); + + database.execute("UPDATE sandbox SET operation_kind = 'activate' WHERE id = 'sandbox-1'"); + expect((await readDrain()).unsafeSandboxes).toBe(1); + expect( + await database.prepare(PROTOCOL_V3_POST_MIGRATION_UNSAFE_SANDBOXES_SQL).first("id"), + ).toBe("sandbox-1"); + database.execute("UPDATE sandbox SET operation_kind = NULL WHERE id = 'sandbox-1'"); + database.execute( + "INSERT INTO sandbox_backup_staging VALUES ('staging-1', 'sandbox-1', 'session-1')", + ); + expect((await readDrain()).unsafeSandboxBackupStaging).toBe(1); + database.execute(INSERT_ENVIRONMENT_ARTIFACT_STAGING_SQL); + expect((await readDrain()).unsafeEnvironmentArtifactBackupStaging).toBe(1); + database.execute("INSERT INTO app_deployment_run VALUES ('active-app-run', 'activating')"); + expect((await readDrain()).activeAppDeploymentRuns).toBe(1); + expect(isProtocolV3CutoverDrained(await readDrain())).toBeFalse(); + }); + + test("allows only deterministic legacy terminal normalization before migration 0014", async () => { + const database = createLegacyTerminalDatabase(); + database.execute(` + INSERT INTO session VALUES ('session-1', 'failed-run', 1, 3, 'IDLE', NULL); + INSERT INTO session_run VALUES + (100, NULL, NULL, NULL, 'completed-run', 'session-1', 'completed', 'run.complete'), + (100, 'runtime.failed', NULL, 'Run failed', 'failed-run', 'session-1', 'failed', 'run.fail'); + INSERT INTO session_event VALUES + ('completed-event', 'run.completed', 'completed-run', 1, 'session-1', 'session-run-terminal:completed-run:run.completed'), + ('failed-event', 'run.failed', 'failed-run', 2, 'session-1', 'session-run-terminal:failed-run:run.failed'), + ('message-event', 'message.added', 'completed-run', 3, 'session-1', 'message-source'); + INSERT INTO session_message VALUES + ('assistant-1', '[]', 'assistant', 1, '[]', 'session-1', 'completed-run'); + `); + + const integrity = await readLegacyTerminalIntegrity(database); + expect(integrity).toEqual({ + ambiguousAssistantRuns: 0, + duplicateTerminalRuns: 0, + invalidFailedRuns: 0, + invalidNonfailedRunErrors: 0, + invalidTerminalLinks: 0, + legacyMaterializedMessages: 1, + legacyStreamRows: 1, + legacyTerminalEvents: 2, + mismatchedTerminalEvents: 0, + missingTerminalEvents: 0, + noncanonicalTerminalSources: 0, + nonterminalCommands: 0, + partialAssistantProjections: 0, + partialTerminalProjections: 0, + repairableFailedRuns: 1, + rewriteCandidateManifestJson: "[]", + unsettledEffects: 0, + }); + expect(() => assertProtocolV3LegacyTerminalIntegrity(integrity)).not.toThrow(); + const sourceInventory = await readLegacyTerminalSourceInventory(database); + expect(sourceInventory).toEqual({ + cancelled: { canonical: 0, canonicalTargetCollisions: 0, noncanonical: 0, total: 0 }, + completed: { canonical: 1, canonicalTargetCollisions: 0, noncanonical: 0, total: 1 }, + failed: { canonical: 1, canonicalTargetCollisions: 0, noncanonical: 0, total: 1 }, + invalidTerminalLinks: 0, + mismatchedTerminalEvents: 0, + multipleTerminalRuns: 0, + }); + expect(() => assertProtocolV3LegacyTerminalSourceInventory(sourceInventory)).not.toThrow(); + }); + + test("inventories deterministic source rewrites but rejects collisions and multiple terminals", async () => { + const database = createValidLegacyCompletionDatabase(); + database.execute(` + UPDATE session_event + SET source_event_id = 'provider-completion-event' + WHERE id = 'completed-event'; + INSERT INTO session_event VALUES + ('canonical-source-owner', 'message.added', 'completed-run', 2, 'session-1', 'session-run-terminal:completed-run:run.completed'), + ('second-terminal', 'run.failed', 'completed-run', 3, 'session-1', 'session-run-terminal:completed-run:run.failed'); + `); + + const inventory = await readLegacyTerminalSourceInventory(database); + expect(inventory).toEqual({ + cancelled: { canonical: 0, canonicalTargetCollisions: 0, noncanonical: 0, total: 0 }, + completed: { canonical: 0, canonicalTargetCollisions: 1, noncanonical: 1, total: 1 }, + failed: { canonical: 1, canonicalTargetCollisions: 0, noncanonical: 0, total: 1 }, + invalidTerminalLinks: 0, + mismatchedTerminalEvents: 1, + multipleTerminalRuns: 1, + }); + expect(() => assertProtocolV3LegacyTerminalSourceInventory(inventory)).toThrow( + "not ready for the protocol v3 production cutover", + ); + }); + + test("allows collision-free provider sources as deterministic migration rewrite candidates", async () => { + const database = createValidLegacyCompletionDatabase(); + database.execute( + "UPDATE session_event SET source_event_id = 'provider-completion-event' WHERE id = 'completed-event'", + ); + + const inventory = await readLegacyTerminalSourceInventory(database); + expect(inventory.completed).toEqual({ + canonical: 0, + canonicalTargetCollisions: 0, + noncanonical: 1, + total: 1, + }); + expect(() => assertProtocolV3LegacyTerminalSourceInventory(inventory)).not.toThrow(); + const integrity = await readLegacyTerminalIntegrity(database); + expect(JSON.parse(integrity.rewriteCandidateManifestJson)).toEqual([ + [ + "completed-event", + "session-1", + "completed-run", + "run.completed", + "provider-completion-event", + 1, + ], + ]); + }); + + for (const fixture of [ + { + blocker: "mismatchedTerminalEvents", + label: "a source rewrite whose event kind conflicts with its Run", + mutate: "UPDATE session_run SET status = 'failed' WHERE id = 'completed-run'", + }, + { + blocker: "invalidTerminalLinks", + label: "a source rewrite without an exact Run link", + mutate: "UPDATE session_event SET run_id = 'unknown-run' WHERE id = 'completed-event'", + }, + ] as const) { + test(`rejects ${fixture.label} in the read-only inventory`, async () => { + const database = createValidLegacyCompletionDatabase(); + database.execute(fixture.mutate); + + const inventory = await readLegacyTerminalSourceInventory(database); + expect(inventory[fixture.blocker]).toBe(1); + expect(() => assertProtocolV3LegacyTerminalSourceInventory(inventory)).toThrow( + "not ready for the protocol v3 production cutover", + ); + }); + } + + test("blocks every ambiguous legacy terminal and assistant projection", async () => { + const database = createLegacyTerminalDatabase(); + database.execute(` + INSERT INTO session VALUES ('session-1', 'partial-run', 3, 20, 'RUNNING', 'operation-1'); + INSERT INTO session_run VALUES + (100, NULL, NULL, NULL, 'duplicate-run', 'session-1', 'completed', 'run.complete'), + (100, NULL, NULL, NULL, 'mismatch-run', 'session-1', 'completed', 'run.complete'), + (100, 'runtime.failed', '{}', 'Run failed', 'missing-run', 'session-1', 'failed', 'run.fail'), + (NULL, 'runtime.failed', '{}', 'Run failed', 'partial-run', 'session-1', 'failed', 'run.complete'), + (100, NULL, NULL, NULL, 'ambiguous-run', 'session-1', 'completed', 'run.complete'), + (100, NULL, NULL, 'Run failed', 'invalid-error-run', 'session-1', 'failed', 'run.fail'), + (100, 'stale.error', '{}', 'Stale error', 'invalid-nonfailed-error-run', 'session-1', 'completed', 'run.complete'); + INSERT INTO session_event VALUES + ('duplicate-event-1', 'run.completed', 'duplicate-run', 1, 'session-1', 'session-run-terminal:duplicate-run:run.completed'), + ('duplicate-event-2', 'run.completed', 'duplicate-run', 2, 'session-1', 'legacy-duplicate-source'), + ('mismatch-event', 'run.failed', 'mismatch-run', 3, 'session-1', 'session-run-terminal:mismatch-run:run.failed'), + ('partial-event', 'run.failed', 'partial-run', 4, 'session-1', 'session-run-terminal:partial-run:run.failed'), + ('ambiguous-event', 'run.completed', 'ambiguous-run', 5, 'session-1', 'session-run-terminal:ambiguous-run:run.completed'), + ('invalid-error-event', 'run.failed', 'invalid-error-run', 6, 'session-1', 'session-run-terminal:invalid-error-run:run.failed'), + ('invalid-nonfailed-error-event', 'run.completed', 'invalid-nonfailed-error-run', 7, 'session-1', 'session-run-terminal:invalid-nonfailed-error-run:run.completed'), + ('orphan-event', 'run.completed', 'unknown-run', 8, 'session-1', 'session-run-terminal:unknown-run:run.completed'); + INSERT INTO session_message VALUES + ('partial-assistant', '[]', 'assistant', 1, '[]', 'session-1', 'partial-run'), + ('ambiguous-assistant-1', '[]', 'assistant', 2, '[]', 'session-1', 'ambiguous-run'), + ('ambiguous-assistant-2', '[]', 'assistant', 3, '[]', 'session-1', 'ambiguous-run'); + INSERT INTO session_permission_request VALUES ('partial-permission', 'partial-run', 'session-1'); + INSERT INTO driver_command VALUES ('accepted-command', 'accepted'); + INSERT INTO external_tool_effect VALUES ('claimed-effect', 'claimed'); + `); + + const integrity = await readLegacyTerminalIntegrity(database); + expect(integrity).toMatchObject({ + ambiguousAssistantRuns: 1, + duplicateTerminalRuns: 1, + invalidFailedRuns: 1, + invalidNonfailedRunErrors: 1, + invalidTerminalLinks: 1, + mismatchedTerminalEvents: 1, + missingTerminalEvents: 1, + noncanonicalTerminalSources: 1, + nonterminalCommands: 1, + partialAssistantProjections: 1, + partialTerminalProjections: 1, + unsettledEffects: 1, + }); + expect(() => assertProtocolV3LegacyTerminalIntegrity(integrity)).toThrow( + "Legacy terminal integrity is ambiguous", + ); + }); + + for (const fixture of [ + { + blocker: "partialTerminalProjections", + label: "a terminal event beyond the Session cursor", + mutate: "UPDATE session SET runtime_event_seq_cursor = 0 WHERE id = 'session-1'", + }, + { + blocker: "partialAssistantProjections", + label: "a materialized assistant beyond the Session cursor", + mutate: "UPDATE session SET message_seq_cursor = 0 WHERE id = 'session-1'", + }, + { + blocker: "partialTerminalProjections", + label: "a terminal Run with a permission request", + mutate: + "INSERT INTO session_permission_request VALUES ('permission-1', 'completed-run', 'session-1')", + }, + { + blocker: "partialTerminalProjections", + label: "a current terminal Run whose Session remains live", + mutate: "UPDATE session SET status = 'RUNNING' WHERE id = 'session-1'", + }, + { + blocker: "partialTerminalProjections", + label: "a current terminal Run whose Session keeps an operation fence", + mutate: "UPDATE session SET status_operation_id = 'operation-1' WHERE id = 'session-1'", + }, + { + blocker: "partialTerminalProjections", + label: "a terminal Run without a completion timestamp", + mutate: "UPDATE session_run SET completed_at = NULL WHERE id = 'completed-run'", + }, + { + blocker: "partialTerminalProjections", + label: "a terminal Run with the wrong lifecycle event", + mutate: "UPDATE session_run SET status_event = 'run.fail' WHERE id = 'completed-run'", + }, + ] as const) { + test(`blocks ${fixture.label} before migration 0014`, async () => { + const database = createValidLegacyCompletionDatabase(); + database.execute(fixture.mutate); + + const integrity = await readLegacyTerminalIntegrity(database); + expect(integrity[fixture.blocker]).toBe(1); + expect(() => assertProtocolV3LegacyTerminalIntegrity(integrity)).toThrow( + "Legacy terminal integrity is ambiguous", + ); + }); + } + + test("keeps the recovery bookmark in the one-shot D1 gate", async () => { + const database = createCutoverGateDatabase(); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(ENTER_PROTOCOL_V3_DRAIN_SQL); + + const bookmark = "00000085-0000024c-00004c6d-8e61117bf38d7adb71b934ebbf891683"; + expect(parseTimeTravelBookmark(JSON.stringify({ bookmark }))).toBe(bookmark); + database.execute(storeProtocolV3CutoverBookmarkSql(bookmark)); + + const stored = await database + .prepare( + 'SELECT pre_migration_bookmark AS bookmark FROM "__protocol_v3_cutover" WHERE id = 1', + ) + .first<{ bookmark: string }>(); + expect(parseStoredProtocolV3CutoverBookmark(d1Json(stored))).toBe(bookmark); + + const sessionId = "01J0000000000000000000000A"; + const requestKey = "protocol-v3-cutover-2de67417-f144-47e1-8b2e-1e71621a0d92"; + database.execute(storeProtocolV3SmokeRequestKeySql(requestKey)); + const storedRequestKey = await database + .prepare(PROTOCOL_V3_SMOKE_REQUEST_KEY_SQL) + .first>(); + expect(parseStoredProtocolV3SmokeRequestKey(d1Json(storedRequestKey))).toBe(requestKey); + database.execute(storeProtocolV3SmokeSessionSql(sessionId)); + const storedSession = await database + .prepare(PROTOCOL_V3_SMOKE_SESSION_SQL) + .first>(); + expect(parseStoredProtocolV3SmokeSession(d1Json(storedSession))).toBe(sessionId); + }); + + test("persists migration intent before the remote apply can outlive this process", async () => { + const database = createCutoverGateDatabase(); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + + database.execute(beginProtocolV3MigrationSql(RELEASE_TREE_OID)); + expect((await readCutoverState(database)).migrationStarted).toBe(false); + + database.execute(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + database.execute(beginProtocolV3MigrationSql(RELEASE_TREE_OID)); + expect(await readCutoverState(database)).toMatchObject({ + enabled: true, + migrationStarted: true, + }); + + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + expect((await readCutoverState(database)).migrationStarted).toBe(true); + }); + + test("requires exact trigger-free migration intent authority", async () => { + for (const mutate of [ + (database: SqliteD1Database) => + database.execute(` + DROP TABLE "${PROTOCOL_V3_MIGRATION_INTENT_TABLE}"; + CREATE TABLE "${PROTOCOL_V3_MIGRATION_INTENT_TABLE}" ( + id integer PRIMARY KEY, + started_at integer NOT NULL + ); + `), + (database: SqliteD1Database) => + database.execute(` + CREATE TRIGGER migration_intent_spoof + AFTER INSERT ON "${PROTOCOL_V3_MIGRATION_INTENT_TABLE}" + BEGIN + SELECT 1; + END; + `), + ]) { + const database = createCutoverGateDatabase(); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + mutate(database); + await expect(readCutoverState(database)).rejects.toThrow( + "migration intent authority is invalid", + ); + } + }); + + test("treats an isolated migration intent table as incomplete gate cleanup", async () => { + const database = new SqliteD1Database(); + database.execute(`CREATE TABLE "${PROTOCOL_V3_MIGRATION_INTENT_TABLE}" (id integer)`); + const probe = await database + .prepare(PROTOCOL_V3_CUTOVER_PROBE_SQL) + .first>(); + expect(parseProtocolV3CutoverProbe(d1Json(probe))).toEqual({ gatePresent: true }); + }); + + test("keeps a durable marker until queue resume survives a restart", async () => { + const database = createCutoverGateDatabase(); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + database.execute(ENTER_PROTOCOL_V3_QUEUES_RESUMING_SQL); + + const readState = async () => + parseProtocolV3CutoverState( + d1Json( + await database.prepare(PROTOCOL_V3_COMMAND_FREEZE_SQL).first>(), + ), + ); + expect(await readState()).toEqual({ + commandFreeze: true, + ...EMPTY_ROLLOUT, + enabled: true, + phase: "queues_resuming", + }); + + await expect( + database + .prepare( + "INSERT INTO session_run VALUES ('v3-run', 'queued', '01J00000000000000000000001', 'session-1')", + ) + .run(), + ).rejects.toThrow("protocol v3 cutover blocks new active Session Runs"); + + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + expect(await readState()).toEqual({ + commandFreeze: true, + ...EMPTY_ROLLOUT, + enabled: true, + phase: "queues_resuming", + }); + + database.execute(ACCEPT_PROTOCOL_V3_QUEUE_RESUME_SQL); + expect(await readState()).toEqual({ + commandFreeze: true, + ...EMPTY_ROLLOUT, + enabled: false, + phase: "queues_resuming", + }); + await expect( + database + .prepare("INSERT INTO session_run VALUES ('v3-run', 'queued', 'account-1', 'session-1')") + .run(), + ).resolves.toBeDefined(); + + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + objectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + }); + + database.execute(DROP_PROTOCOL_V3_CUTOVER_TRIGGERS_SQL); + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: 0, + objectCount: 1, + }); + await expect( + database + .prepare( + "INSERT INTO session_run VALUES ('cleanup-interrupted', 'queued', '01J00000000000000000000001', 'session-1')", + ) + .run(), + ).resolves.toBeDefined(); + + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(REMOVE_PROTOCOL_V3_CUTOVER_SQL); + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: 0, + objectCount: 0, + }); + }); + + test("binds crash recovery to one clean Git tree and one rollout", async () => { + const database = createCutoverGateDatabase(); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + const initial = await readCutoverState(database); + expect(() => assertProtocolV3Release(initial, RELEASE_TREE_OID)).not.toThrow(); + + const otherRelease = "89abcdef0123456789abcdef0123456789abcdef"; + database.execute(installProtocolV3CutoverSql(otherRelease)); + expect(() => assertProtocolV3Release(initial, otherRelease)).toThrow("belongs to release tree"); + + const workerVersionId = "2de67417-f144-47e1-8b2e-1e71621a0d92"; + database.execute( + storeProtocolV3RolloutSql(RELEASE_TREE_OID, workerVersionId, 17, CONTAINER_IMAGE_DIGEST), + ); + database.execute( + storeProtocolV3RolloutSql(RELEASE_TREE_OID, workerVersionId, 17, CONTAINER_IMAGE_DIGEST), + ); + expect(await readCutoverState(database)).toEqual({ + commandFreeze: false, + containerApplicationVersion: 17, + containerImageDigest: CONTAINER_IMAGE_DIGEST, + enabled: true, + migrationStarted: false, + phase: "draining", + releaseTreeOid: RELEASE_TREE_OID, + workerVersionId, + }); + + database.execute( + storeProtocolV3RolloutSql( + RELEASE_TREE_OID, + "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa", + 18, + "b".repeat(64), + ), + ); + expect((await readCutoverState(database)).workerVersionId).toBe(workerVersionId); + }); + + test("rejects dirty tracked, untracked, and submodule release trees", () => { + expect(parseCleanGitTreeOid(`${RELEASE_TREE_OID}\n`, "")).toBe(RELEASE_TREE_OID); + for (const status of [" M apps/api/src/index.ts", "?? untracked.ts", " m apps/driver"]) { + expect(() => parseCleanGitTreeOid(RELEASE_TREE_OID, status)).toThrow("clean Git worktree"); + } + for (const index of ["S apps/api/src/index.ts", "h apps/api/src/index.ts"]) { + expect(() => parseCleanGitTreeOid(RELEASE_TREE_OID, "", index)).toThrow("clean Git worktree"); + } + expect(() => parseCleanGitTreeOid(RELEASE_TREE_OID, "", "", "apps/web/.env.prod")).toThrow( + "Vite build inputs", + ); + expect(() => + parseCleanGitTreeOid(RELEASE_TREE_OID, "", "", "", ["PATH", "VITE_MOSOO_ENVIRONMENT"]), + ).toThrow("Vite build inputs"); + }); + + test("verifies the exact tagged Worker version before accepting rollout", () => { + const versionId = "2de67417-f144-47e1-8b2e-1e71621a0d92"; + expect( + parseProtocolV3WorkerDeployment( + JSON.stringify({ versions: [{ percentage: 100, version_id: versionId }] }), + ), + ).toEqual({ versionId }); + expect(() => + assertProtocolV3WorkerVersion( + JSON.stringify({ + annotations: { "workers/tag": `protocol-v3-${RELEASE_TREE_OID}` }, + id: versionId, + }), + versionId, + RELEASE_TREE_OID, + ), + ).not.toThrow(); + expect(() => + assertProtocolV3WorkerVersion( + JSON.stringify({ annotations: { "workers/tag": "protocol-v3-wrong" }, id: versionId }), + versionId, + RELEASE_TREE_OID, + ), + ).toThrow("not the exact protocol v3 release"); + + const repository = "registry.cloudflare.com/account/mosoo-api-prod-sandbox-prod"; + expect(protocolV3ContainerImageTag(repository, versionId)).toBe(`${repository}:2de67417`); + expect( + parseProtocolV3ContainerManifestDigest( + JSON.stringify({ Descriptor: { digest: `sha256:${CONTAINER_IMAGE_DIGEST}` } }), + ), + ).toBe(CONTAINER_IMAGE_DIGEST); + }); + + test("keeps admission closed when queue resume crashes partway and converges on retry", async () => { + const database = createCutoverGateDatabase(); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + database.execute(ENTER_PROTOCOL_V3_QUEUES_RESUMING_SQL); + const queues = [true, true, true]; + const actions: string[] = []; + + await expect( + completeProtocolV3QueueResume(await readCutoverState(database), { + commitAcceptance: () => { + actions.push("commit"); + database.execute(ACCEPT_PROTOCOL_V3_QUEUE_RESUME_SQL); + }, + removeMarker: () => { + actions.push("remove"); + database.execute(REMOVE_PROTOCOL_V3_CUTOVER_SQL); + }, + resumeAndVerifyQueues: async () => { + actions.push("resume-1"); + queues[0] = false; + queues[1] = false; + throw new Error("third queue readback failed"); + }, + }), + ).rejects.toThrow("third queue readback failed"); + expect(actions).toEqual(["resume-1"]); + expect(await readCutoverState(database)).toEqual({ + commandFreeze: true, + ...EMPTY_ROLLOUT, + enabled: true, + phase: "queues_resuming", + }); + + await completeProtocolV3QueueResume(await readCutoverState(database), { + commitAcceptance: () => { + actions.push("commit"); + database.execute(ACCEPT_PROTOCOL_V3_QUEUE_RESUME_SQL); + }, + removeMarker: () => { + actions.push("remove"); + database.execute(REMOVE_PROTOCOL_V3_CUTOVER_SQL); + }, + resumeAndVerifyQueues: async () => { + actions.push("resume-2"); + queues.fill(false); + }, + }); + expect(actions).toEqual(["resume-1", "resume-2", "commit", "remove"]); + expect(queues).toEqual([false, false, false]); + expect(await readCutoverObjects(database)).toEqual({ exactObjectCount: 0, objectCount: 0 }); + }); + + test("mutates and reads back every cutover Queue lane, then re-pauses all after failure", async () => { + const paused = new Map(PROTOCOL_V3_CUTOVER_QUEUE_NAMES.map((name) => [name, true])); + const calls: string[] = []; + let staleReadback: string | null = null; + const control = { + list: async () => PROTOCOL_V3_CUTOVER_QUEUE_NAMES.map((name) => ({ id: `id:${name}`, name })), + mutate: (queueName: string, action: "pause" | "resume") => { + calls.push(`${action}:${queueName}`); + paused.set(queueName, action === "pause"); + }, + read: async (queueId: string) => { + const name = queueId.slice("id:".length); + calls.push(`read:${name}`); + const deliveryPaused = paused.get(name); + if (deliveryPaused === undefined) throw new Error(`Unknown Queue ${name}.`); + return { + deliveryPaused: name === staleReadback ? !deliveryPaused : deliveryPaused, + name, + }; + }, + }; + + await updateAndVerifyProtocolV3QueueDelivery(control, "resume"); + expect([...paused.entries()]).toEqual( + PROTOCOL_V3_CUTOVER_QUEUE_NAMES.map((name) => [name, false]), + ); + await updateAndVerifyProtocolV3QueueDelivery(control, "pause"); + expect([...paused.entries()]).toEqual( + PROTOCOL_V3_CUTOVER_QUEUE_NAMES.map((name) => [name, true]), + ); + + staleReadback = "api-command-dlq"; + await expect(updateAndVerifyProtocolV3QueueDelivery(control, "resume")).rejects.toThrow( + "Production queue resume and readback failed", + ); + staleReadback = null; + await updateAndVerifyProtocolV3QueueDelivery(control, "pause"); + expect([...paused.entries()]).toEqual( + PROTOCOL_V3_CUTOVER_QUEUE_NAMES.map((name) => [name, true]), + ); + for (const action of ["pause", "resume"] as const) { + for (const name of PROTOCOL_V3_CUTOVER_QUEUE_NAMES) { + expect(calls).toContain(`${action}:${name}`); + expect(calls).toContain(`read:${name}`); + } + } + }); + + test("drains and re-pauses all three API command lanes before migration", async () => { + const deploySource = await Bun.file(new URL("../bin/deploy-prod.ts", import.meta.url)).text(); + const drainStart = deploySource.indexOf("if (initialPendingMigrations.length > 0)"); + const migrationStart = deploySource.indexOf("if (durableMcpMigrationPending)", drainStart); + const drainSource = deploySource.slice(drainStart, migrationStart); + + expect(drainSource).toContain("await resumeAndVerifyProdQueues(queueApiConfig);"); + expect(drainSource).toContain("await pauseAndVerifyProdQueues(queueApiConfig);"); + expect(drainSource).not.toContain('["api-command"]'); + expect(deploySource).toContain("let migrationStarted = true;"); + expect( + deploySource.indexOf("cutoverState = beginProtocolV3Migration(", migrationStart), + ).toBeLessThan(deploySource.indexOf("applyD1Migrations();", migrationStart)); + expect( + deploySource.indexOf("authorizeProdLegacyTerminalRewrite(", migrationStart), + ).toBeGreaterThan( + deploySource.indexOf("cutoverState = beginProtocolV3Migration(", migrationStart), + ); + expect( + deploySource.indexOf("authorizeProdLegacyTerminalRewrite(", migrationStart), + ).toBeLessThan(deploySource.indexOf("applyD1Migrations();", migrationStart)); + }); + + test("recovers a lost acceptance acknowledgement and an externally re-paused queue", async () => { + const database = createCutoverGateDatabase(); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + database.execute(ENTER_PROTOCOL_V3_QUEUES_RESUMING_SQL); + const queues = [true, true, true]; + const actions: string[] = []; + + await expect( + completeProtocolV3QueueResume(await readCutoverState(database), { + commitAcceptance: () => { + actions.push("commit-lost-ack"); + database.execute(ACCEPT_PROTOCOL_V3_QUEUE_RESUME_SQL); + throw new Error("acceptance readback unavailable"); + }, + removeMarker: () => { + actions.push("remove"); + database.execute(REMOVE_PROTOCOL_V3_CUTOVER_SQL); + }, + resumeAndVerifyQueues: async () => { + actions.push("resume-1"); + queues.fill(false); + }, + }), + ).rejects.toThrow("acceptance readback unavailable"); + expect(await readCutoverState(database)).toEqual({ + commandFreeze: true, + ...EMPTY_ROLLOUT, + enabled: false, + phase: "queues_resuming", + }); + + queues[1] = true; + await completeProtocolV3QueueResume(await readCutoverState(database), { + commitAcceptance: () => actions.push("unexpected-recommit"), + removeMarker: () => { + actions.push("remove"); + database.execute(REMOVE_PROTOCOL_V3_CUTOVER_SQL); + }, + resumeAndVerifyQueues: async () => { + actions.push("resume-2"); + queues.fill(false); + }, + }); + expect(actions).toEqual(["resume-1", "commit-lost-ack", "resume-2", "remove"]); + expect(queues).toEqual([false, false, false]); + expect(await readCutoverObjects(database)).toEqual({ exactObjectCount: 0, objectCount: 0 }); + }); + + for (const spoof of [ + { + label: "a WHEN 0 trigger", + name: "__protocol_v3_cutover_session_run_insert", + }, + { + label: "a differently-cased reserved trigger", + name: "__PROTOCOL_V3_CUTOVER_SESSION_RUN_INSERT", + }, + ]) { + test(`fails exact gate verification for ${spoof.label}`, async () => { + const database = createCutoverGateDatabase(); + database.execute(` + CREATE TRIGGER "${spoof.name}" + BEFORE INSERT ON "session_run" + WHEN 0 + BEGIN + SELECT 1; + END; + `); + + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT - 1, + objectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + }); + await expect( + database + .prepare( + "INSERT INTO session_run VALUES ('spoof-run', 'queued', 'account-1', 'session-1')", + ) + .run(), + ).resolves.toBeDefined(); + + const probe = await database + .prepare(PROTOCOL_V3_CUTOVER_PROBE_SQL) + .first>(); + expect(parseProtocolV3CutoverProbe(d1Json(probe)).gatePresent).toBe(true); + }); + } + + test("rejects an extra trigger attached to a protected admission table", async () => { + const database = createCutoverGateDatabase(); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(` + CREATE TRIGGER disable_protocol_v3_gate + BEFORE INSERT ON "session_run" + BEGIN + UPDATE "__protocol_v3_cutover" + SET "command_freeze" = 1, "enabled" = 0, "phase" = 'queues_resuming' + WHERE "id" = 1; + END + `); + + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: 0, + objectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT + 1, + }); + await expect( + database + .prepare("INSERT INTO session_run VALUES ('spoof-run', 'queued', 'account-1', 'session-1')") + .run(), + ).resolves.toBeDefined(); + const probe = await database + .prepare(PROTOCOL_V3_CUTOVER_PROBE_SQL) + .first>(); + expect(parseProtocolV3CutoverProbe(d1Json(probe)).gatePresent).toBe(true); + }); + + test("excludes only the exact permanent sandbox identity trigger", async () => { + const database = createCutoverGateDatabase(); + database.execute(` + ALTER TABLE sandbox ADD kind text; + ALTER TABLE sandbox ADD agent_id text; + ALTER TABLE sandbox ADD app_id text; + ALTER TABLE sandbox ADD owner_account_id text; + `); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + const managed = MANAGED_PROD_SCHEMA_TRIGGERS.find( + (trigger) => trigger.name === "sandbox_identity_immutable", + ); + if (managed === undefined) throw new Error("Sandbox identity trigger fixture is missing."); + database.execute(managed.sql); + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + objectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + }); + + database.execute(` + DROP TRIGGER sandbox_identity_immutable; + CREATE TRIGGER sandbox_identity_immutable + BEFORE INSERT ON sandbox WHEN 0 + BEGIN + SELECT 1; + END; + `); + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: 0, + objectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT + 1, + }); + }); + + test("excludes only the exact permanent environment staging triggers", async () => { + const database = createCutoverGateDatabase(); + addPostRuntimeAuthoritySchema(database); + database.execute(INSTALL_PROTOCOL_V3_POST_MIGRATION_CUTOVER_SQL); + for (const name of [ + "environment_package_artifact_backup_staging_authority", + "environment_package_artifact_backup_staging_immutable", + ]) { + const managed = MANAGED_PROD_SCHEMA_TRIGGERS.find((trigger) => trigger.name === name); + if (managed === undefined) throw new Error(`Managed trigger ${name} fixture is missing.`); + database.execute(managed.sql); + } + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + objectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + }); + + database.execute(` + DROP TRIGGER environment_package_artifact_backup_staging_authority; + CREATE TRIGGER environment_package_artifact_backup_staging_authority + BEFORE INSERT ON environment_package_artifact_backup_staging WHEN 0 + BEGIN + SELECT 1; + END; + `); + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: 0, + objectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT + 1, + }); + }); + + test("does not exempt a same-name spoof of the rewrite revocation trigger", async () => { + const database = createCutoverGateDatabase(); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(` + CREATE TRIGGER "__protocol_v3_legacy_rewrite_gate_update" + BEFORE INSERT ON "session_run" + BEGIN + UPDATE "__protocol_v3_cutover" + SET "command_freeze" = 1, "enabled" = 0, "phase" = 'queues_resuming' + WHERE "id" = 1; + END + `); + + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: 0, + objectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT + 1, + }); + await expect( + database + .prepare( + "INSERT INTO session_run VALUES ('auth-spoof-run', 'queued', 'account-1', 'session-1')", + ) + .run(), + ).resolves.toBeDefined(); + }); + + test("detects a same-name object from another sqlite namespace through cleanup", async () => { + const database = createCutoverGateDatabase(); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(` + DROP TRIGGER "__protocol_v3_cutover_session_run_insert"; + CREATE VIEW "__protocol_v3_cutover_session_run_insert" AS SELECT 1 AS value; + `); + + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: 0, + objectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT + 1, + }); + + database.execute(REMOVE_PROTOCOL_V3_CUTOVER_SQL); + expect(await readCutoverObjects(database)).toEqual({ + exactObjectCount: 0, + objectCount: 1, + }); + }); + + test("retries accepted queue verification and idempotent marker cleanup without re-pausing", async () => { + const database = createCutoverGateDatabase(); + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + database.execute(ENTER_PROTOCOL_V3_QUEUES_RESUMING_SQL); + database.execute(ACCEPT_PROTOCOL_V3_QUEUE_RESUME_SQL); + database.execute(REMOVE_PROTOCOL_V3_CUTOVER_SQL); + const calls: string[] = []; + const originalError = new Error("marker deletion committed before its readback failed"); + const state = { + bookmark: null, + initialPendingMigrations: [], + migrationStarted: false, + originalError, + queuesVerified: true, + }; + await expect( + recoverProtocolV3CutoverFailure(state, { + commitQueueAcceptance: () => calls.push("unexpected-commit"), + pauseAndVerifyQueues: async () => calls.push("unexpected-pause"), + printBookmark: () => calls.push("unexpected-bookmark"), + probe: () => { + calls.push("probe"); + return parseProtocolV3CutoverProbe( + d1Json({ + gate_present: 0, + }), + ); + }, + readBookmark: () => null, + readPendingMigrations: () => [], + removeMarker: () => calls.push("unexpected-remove"), + resumeAndVerifyQueues: async () => calls.push("resume"), + write: () => {}, + }), + ).resolves.toBeUndefined(); + expect(calls).toEqual(["resume", "probe"]); + expect(await readCutoverObjects(database)).toEqual({ exactObjectCount: 0, objectCount: 0 }); + }); + + test("keeps production closed when an unacknowledged migration commits after recovery starts", async () => { + const database = createCutoverGateDatabase(); + const migration = "0019_runtime-subject-operation-authority.sql"; + database.execute(INSTALL_PROTOCOL_V3_CUTOVER_SQL); + database.execute(ENABLE_PROTOCOL_V3_COMMAND_FREEZE_SQL); + database.execute(beginProtocolV3MigrationSql(RELEASE_TREE_OID)); + const calls: string[] = []; + const messages: string[] = []; + const originalError = new Error("migration apply acknowledgement timed out"); + + await expect( + recoverProtocolV3CutoverFailure( + { + bookmark: "emergency-bookmark", + initialPendingMigrations: [migration], + migrationStarted: (await readCutoverState(database)).migrationStarted, + originalError, + queuesVerified: false, + }, + { + commitQueueAcceptance: () => calls.push("unexpected-commit"), + pauseAndVerifyQueues: async () => { + calls.push("pause"); + database.execute(`INSERT INTO d1_migrations VALUES ('${migration}')`); + }, + printBookmark: () => calls.push("bookmark"), + probe: () => ({ gatePresent: true }), + readBookmark: () => null, + readPendingMigrations: () => { + calls.push("unexpected-read-pending"); + return [migration]; + }, + removeMarker: () => calls.push("unexpected-remove"), + resumeAndVerifyQueues: async () => calls.push("unexpected-resume"), + write: (message) => messages.push(message), + }, + ), + ).rejects.toThrow(originalError.message); + + expect(calls).toEqual(["pause", "bookmark"]); + expect(messages.join("\n")).toContain("migration request may have committed"); + expect(await readCutoverState(database)).toMatchObject({ + enabled: true, + migrationStarted: true, + }); + expect( + await database.prepare("SELECT name FROM d1_migrations").first<{ name: string }>(), + ).toEqual({ name: migration }); + }); + + test("keeps a gate-missing pre-acceptance failure closed after probing", async () => { + const infrastructureCalls: string[] = []; + const messages: string[] = []; + const originalError = new Error("cutover failed before queue acceptance"); + + await expect( + recoverProtocolV3CutoverFailure( + { + bookmark: "emergency-bookmark", + initialPendingMigrations: ["0019_runtime-subject-operation-authority.sql"], + migrationStarted: false, + originalError, + queuesVerified: false, + }, + { + commitQueueAcceptance: () => { + infrastructureCalls.push("commit"); + }, + pauseAndVerifyQueues: async () => { + infrastructureCalls.push("pause"); + }, + printBookmark: () => {}, + probe: () => { + infrastructureCalls.push("probe"); + return { + gatePresent: false, + }; + }, + readBookmark: () => { + infrastructureCalls.push("read-bookmark"); + return null; + }, + readPendingMigrations: () => { + infrastructureCalls.push("read-pending"); + return []; + }, + removeMarker: () => { + infrastructureCalls.push("remove"); + }, + resumeAndVerifyQueues: async () => { + infrastructureCalls.push("resume-and-verify"); + }, + write: (message) => messages.push(message), + }, + ), + ).rejects.toThrow(originalError.message); + + expect(infrastructureCalls).toEqual(["read-pending", "pause"]); + expect(messages.join("\n")).toContain("rolling forward this exact v3 release"); + expect(messages.join("\n")).not.toContain("manually resume"); + }); + + test("keeps queues closed when first Workflow publication fails its exact readback", async () => { + const calls: string[] = []; + const originalError = new Error("published Workflow binding is not exact"); + + await expect( + recoverProtocolV3CutoverFailure( + { + bookmark: null, + initialPendingMigrations: [], + migrationStarted: false, + originalError, + queuesVerified: false, + }, + { + commitQueueAcceptance: () => calls.push("unexpected-commit"), + pauseAndVerifyQueues: async () => calls.push("pause-and-verify"), + printBookmark: () => calls.push("unexpected-bookmark"), + probe: () => { + calls.push("unexpected-probe"); + return { gatePresent: true }; + }, + readBookmark: () => null, + readPendingMigrations: () => { + calls.push("unexpected-read-pending"); + return []; + }, + removeMarker: () => calls.push("unexpected-remove"), + resumeAndVerifyQueues: async () => calls.push("unexpected-resume"), + write: () => {}, + }, + ), + ).rejects.toThrow(originalError.message); + expect(calls).toEqual(["pause-and-verify"]); + }); + + test("restores the old service when legacy preflight fails before protocol migration", async () => { + const infrastructureCalls: string[] = []; + const originalError = new Error("Legacy terminal integrity is ambiguous"); + + await expect( + recoverProtocolV3CutoverFailure( + { + bookmark: null, + initialPendingMigrations: [ + "0013_durable-mcp-effect-v3.sql", + "0014_session-event-stream-identity.sql", + ], + migrationStarted: false, + originalError, + queuesVerified: false, + }, + { + commitQueueAcceptance: () => { + infrastructureCalls.push("commit"); + }, + pauseAndVerifyQueues: async () => { + infrastructureCalls.push("pause"); + }, + printBookmark: () => { + infrastructureCalls.push("print-bookmark"); + }, + probe: () => { + infrastructureCalls.push("probe"); + return { + gatePresent: true, + }; + }, + readBookmark: () => { + infrastructureCalls.push("read-bookmark"); + return null; + }, + readPendingMigrations: () => { + infrastructureCalls.push("read-pending"); + return ["0013_durable-mcp-effect-v3.sql", "0014_session-event-stream-identity.sql"]; + }, + removeMarker: () => { + infrastructureCalls.push("remove"); + }, + resumeAndVerifyQueues: async () => { + infrastructureCalls.push("resume-and-verify"); + }, + write: () => {}, + }, + ), + ).rejects.toThrow(originalError.message); + + expect(infrastructureCalls).toEqual(["read-pending", "resume-and-verify", "remove"]); + }); + + test("keeps old admission closed until every queue is verifiably resumed", async () => { + const calls: string[] = []; + const queues = [true, true, true]; + const originalError = new Error("legacy preflight failed"); + + await expect( + recoverProtocolV3CutoverFailure( + { + bookmark: null, + initialPendingMigrations: ["0019_runtime-subject-operation-authority.sql"], + migrationStarted: false, + originalError, + queuesVerified: false, + }, + { + commitQueueAcceptance: () => calls.push("commit"), + pauseAndVerifyQueues: async () => { + calls.push("pause-and-verify"); + queues.fill(true); + expect(queues).toEqual([true, true, true]); + }, + printBookmark: () => {}, + probe: () => ({ + gatePresent: true, + }), + readBookmark: () => null, + readPendingMigrations: () => { + calls.push("read-pending"); + return ["0019_runtime-subject-operation-authority.sql"]; + }, + removeMarker: () => calls.push("remove"), + resumeAndVerifyQueues: async () => { + calls.push("verify"); + queues[0] = false; + throw new Error("queue 2 remains paused"); + }, + write: () => {}, + }, + ), + ).rejects.toThrow(originalError.message); + expect(calls).toEqual(["read-pending", "verify", "pause-and-verify"]); + expect(queues).toEqual([true, true, true]); + }); + + test("re-pauses every queue when marker cleanup and its readback both fail", async () => { + const calls: string[] = []; + const originalError = new Error("pre-migration deploy failed"); + await expect( + recoverProtocolV3CutoverFailure( + { + bookmark: null, + initialPendingMigrations: ["0019_runtime-subject-operation-authority.sql"], + migrationStarted: false, + originalError, + queuesVerified: false, + }, + { + commitQueueAcceptance: () => calls.push("unexpected-commit"), + pauseAndVerifyQueues: async () => calls.push("pause-and-verify"), + printBookmark: () => calls.push("unexpected-bookmark"), + probe: () => { + calls.push("probe"); + throw new Error("D1 marker readback failed"); + }, + readBookmark: () => null, + readPendingMigrations: () => { + calls.push("read-pending"); + return ["0019_runtime-subject-operation-authority.sql"]; + }, + removeMarker: () => { + calls.push("remove"); + throw new Error("marker delete acknowledgement lost"); + }, + resumeAndVerifyQueues: async () => calls.push("resume-and-verify"), + write: () => {}, + }, + ), + ).rejects.toThrow(originalError.message); + expect(calls).toEqual([ + "read-pending", + "resume-and-verify", + "remove", + "pause-and-verify", + "probe", + ]); + }); + + test("re-pauses every queue when gate-absent old-service resume is only partial", async () => { + const calls: string[] = []; + const queues = [true, true, true]; + let resumeAttempts = 0; + const originalError = new Error("pre-migration deploy failed"); + + await expect( + recoverProtocolV3CutoverFailure( + { + bookmark: null, + initialPendingMigrations: ["0019_runtime-subject-operation-authority.sql"], + migrationStarted: false, + originalError, + queuesVerified: false, + }, + { + commitQueueAcceptance: () => calls.push("unexpected-commit"), + pauseAndVerifyQueues: async () => { + calls.push("pause-and-verify"); + queues.fill(true); + }, + printBookmark: () => calls.push("unexpected-bookmark"), + probe: () => { + calls.push("probe"); + return { gatePresent: false }; + }, + readBookmark: () => null, + readPendingMigrations: () => { + calls.push("read-pending"); + return ["0019_runtime-subject-operation-authority.sql"]; + }, + removeMarker: () => { + calls.push("remove"); + throw new Error("marker delete acknowledgement lost"); + }, + resumeAndVerifyQueues: async () => { + calls.push("resume-and-verify"); + resumeAttempts += 1; + if (resumeAttempts === 1) { + queues.fill(false); + return; + } + queues[0] = false; + throw new Error("queue resume was partial"); + }, + write: () => {}, + }, + ), + ).rejects.toThrow(originalError.message); + + expect(calls).toEqual([ + "read-pending", + "resume-and-verify", + "remove", + "pause-and-verify", + "probe", + "resume-and-verify", + "pause-and-verify", + ]); + expect(queues).toEqual([true, true, true]); + }); + + test("requires the live smoke Driver to complete protocol v3 hello and ready", () => { + const ready = parseProtocolV3SmokeStatus( + d1Json({ + boot_token_used_at: 1_787_942_399_000, + connection_id: "connection-1", + driver_pid: 42, + driver_started_at: 1_787_942_400_000, + driver_status: "ready", + driver_version: "3.0.0", + protocol_version: 3, + status_event: "driver.ready", + }), + ); + expect(isProtocolV3SmokeReady(ready)).toBeTrue(); + expect(isProtocolV3SmokeReady({ ...ready, driverPid: null })).toBeFalse(); + expect(isProtocolV3SmokeReady({ ...ready, connectionId: null })).toBeFalse(); + expect(isProtocolV3SmokeReady({ ...ready, driverStatus: "connecting" })).toBeFalse(); + expect(isProtocolV3SmokeReady({ ...ready, protocolVersion: 2 })).toBeFalse(); + expect(protocolV3SmokeStatusSql("01J0000000000000000000000A")).toContain( + '"latest"."driver_started_at"', + ); + }); + + test("requires a published cattle Agent for the production smoke", async () => { + const database = new SqliteD1Database(); + const agentId = "01J0000000000000000000000A"; + database.execute(` + CREATE TABLE agent (id text PRIMARY KEY, kind text NOT NULL, status text NOT NULL); + INSERT INTO agent VALUES ('${agentId}', 'pet', 'published'); + `); + const readAgent = async () => { + const row = await database + .prepare(protocolV3SmokeAgentSql(agentId)) + .first>(); + if (row === null) throw new Error("Smoke Agent query returned no row."); + return d1Json(row); + }; + + let raw = await readAgent(); + expect(() => assertProtocolV3SmokeAgent(raw)).toThrow("published cattle Agent"); + database.execute(`UPDATE agent SET kind = 'cattle', status = 'draft'`); + raw = await readAgent(); + expect(() => assertProtocolV3SmokeAgent(raw)).toThrow("published cattle Agent"); + database.execute(`UPDATE agent SET status = 'published'`); + raw = await readAgent(); + expect(() => assertProtocolV3SmokeAgent(raw)).not.toThrow(); + }); + + test("rejects an old Container version found only on a later page", async () => { + const requestedPages: Array = []; + const instances = await collectProtocolV3ContainerInstances((pageToken) => { + requestedPages.push(pageToken); + return pageToken === null + ? containerPage([{ state: "running", version: 3 }], null, "page-2") + : containerPage([{ state: "running", version: 2 }], "page-2", null); + }); + + expect(requestedPages).toEqual([null, "page-2"]); + expect(instances).toHaveLength(2); + expect( + isProtocolV3ContainerRolloutConverged({ state: "ready", version: 3 }, instances), + ).toBeFalse(); + expect( + isProtocolV3ContainerRolloutConverged({ state: "unknown", version: 3 }, [ + { state: "running", version: 3 }, + ]), + ).toBeFalse(); + }); + + test("accepts an exact scale-to-zero Container application with no old instances", async () => { + const [application] = await collectProtocolV3ContainerApplications(() => + containerApplicationPage( + [ + { + ...containerApplication("mosoo-api-prod-sandbox-prod", 3), + health: { + instances: { active: 0, failed: 0, healthy: 0, scheduling: 0, starting: 0 }, + }, + }, + ], + null, + null, + ), + ); + const instances = await collectProtocolV3ContainerInstances(() => + containerPage([], null, null), + ); + if (application === undefined) throw new Error("Container application fixture is missing."); + + expect(application.state).toBe("ready"); + expect(isProtocolV3ContainerRolloutConverged(application, instances)).toBeTrue(); + }); + + test.each(["degraded", "provisioning", "unknown"])( + "does not accept an empty %s Container application as converged", + (state) => { + expect(isProtocolV3ContainerRolloutConverged({ state, version: 3 }, [])).toBeFalse(); + }, + ); + + test.each(["active", "ready"])( + "accepts an empty exact %s Container application as converged", + (state) => { + expect(isProtocolV3ContainerRolloutConverged({ state, version: 3 }, [])).toBeTrue(); + }, ); + + test("finds the production Container application on a later page", async () => { + const requestedPages: Array = []; + const applications = await collectProtocolV3ContainerApplications((pageToken) => { + requestedPages.push(pageToken); + return pageToken === null + ? containerApplicationPage([containerApplication("other-app", 3)], null, "page-2") + : containerApplicationPage( + [containerApplication("mosoo-api-prod-sandbox-prod", 3)], + "page-2", + null, + ); + }); + + expect(requestedPages).toEqual([null, "page-2"]); + expect(applications.map(({ name }) => name)).toEqual([ + "other-app", + "mosoo-api-prod-sandbox-prod", + ]); + }); + + registerPaginationFailureTests({ + collect: collectProtocolV3ContainerApplications, + invalidMessage: "valid opaque token", + invalidPage: { result: [], result_info: { page_token: null }, success: true }, + label: "Container application", + page: (pageToken, nextPageToken) => containerApplicationPage([], pageToken, nextPageToken), + }); + + registerPaginationFailureTests({ + collect: collectProtocolV3ContainerInstances, + invalidMessage: "paginated object", + invalidPage: "[]", + label: "Container instance", + page: (pageToken, nextPageToken) => containerPage([], pageToken, nextPageToken), + }); }); diff --git a/apps/api/tests/public-thread-api-fixtures.ts b/apps/api/tests/public-thread-api-fixtures.ts index 62cebbb3..56e27594 100644 --- a/apps/api/tests/public-thread-api-fixtures.ts +++ b/apps/api/tests/public-thread-api-fixtures.ts @@ -3,7 +3,12 @@ import { expect } from "bun:test"; import { PUBLIC_API_PREFIX } from "@mosoo/contracts/public-api"; import type { SessionRuntimeEventVisibility } from "@mosoo/contracts/session"; import { sessionEventsTable } from "@mosoo/db"; -import { createRuntimeEvent } from "@mosoo/runtime-events"; +import { + createRuntimeEvent, + createRuntimeEventSemanticHash, + createSessionRunTerminalSourceId, + stringifyRuntimeEventSemanticValue, +} from "@mosoo/runtime-events"; import type { RuntimeEventKind, RuntimeEventVisibility } from "@mosoo/runtime-events"; import { Hono } from "hono"; @@ -119,6 +124,18 @@ const RUNTIME_EVENT_IDS_BY_SEQ = [ "01J00000000000000000000015", "01J00000000000000000000016", "01J00000000000000000000017", + "01J00000000000000000000018", + "01J00000000000000000000019", + "01J0000000000000000000001A", + "01J0000000000000000000001B", + "01J0000000000000000000001C", + "01J0000000000000000000001D", + "01J0000000000000000000001E", + "01J0000000000000000000001F", + "01J0000000000000000000001G", + "01J0000000000000000000001H", + "01J0000000000000000000001J", + "01J0000000000000000000001K", ]; function runtimeEventIdForSeq(seq: number): string { @@ -134,6 +151,7 @@ function runtimeEventIdForSeq(seq: number): string { export async function insertRuntimeEvent( database: SqliteD1Database, input: { + eventId?: string; kind: RuntimeEventKind; occurredAt: number; payload: unknown; @@ -147,19 +165,34 @@ export async function insertRuntimeEvent( const visibility = input.visibility ?? "participant"; const databaseVisibility: SessionRuntimeEventVisibility = visibility === "public" || visibility === "participant" ? "all_consumers" : "owner_debug"; - const eventId = runtimeEventIdForSeq(input.seq); + const eventId = input.eventId ?? runtimeEventIdForSeq(input.seq); + const sourceEventId = + runId !== null && + (input.kind === "run.cancelled" || + input.kind === "run.completed" || + input.kind === "run.failed") + ? createSessionRunTerminalSourceId(runId, input.kind) + : eventId; + const terminal = + input.kind === "run.cancelled" || input.kind === "run.completed" || input.kind === "run.failed"; + if (terminal && !isRecord(input.payload)) { + throw new Error("Terminal runtime event fixture payload must be an object."); + } const event = createRuntimeEvent({ actor: "driver", id: eventId, kind: input.kind, occurredAt: new Date(input.occurredAt).toISOString(), origin: "driver", - payload: input.payload, + payload: terminal ? { ...input.payload, lifecycle: "IDLE" } : input.payload, ...(runId === null ? {} : { runId }), sessionId: input.sessionId, + sourceEventId, visibility, }); const projection = createSessionRuntimeEventProjection(event); + const semanticHash = await createRuntimeEventSemanticHash(event); + const terminalEventJson = terminal ? stringifyRuntimeEventSemanticValue(event) : null; await database .app() @@ -172,17 +205,28 @@ export async function insertRuntimeEvent( eventType: input.kind, family: projection.family, id: eventId, + mcpCommandId: projection.mcpCommandId, occurredAt: input.occurredAt, processStatus: projection.processStatus, processType: projection.processType, runId, + runtimeOperationEventJson: null, + semanticHash, seq: input.seq, sessionId: input.sessionId, source: "driver", - sourceEventId: eventId, + sourceEventId, + streamId: projection.streamId, + terminalEventJson, toolCallId: projection.toolCallId, + toolInputDeltaJson: projection.toolInputDeltaJson, toolInputJson: projection.toolInputJson, toolName: projection.toolName, + toolOutputDeltaText: projection.toolOutputDeltaText, + toolOutputText: projection.toolOutputText, + toolParentMessageId: projection.toolParentMessageId, + toolResultMessageId: projection.toolResultMessageId, + toolStatus: projection.toolStatus, tokens: projection.tokens, traceId: null, visibility: databaseVisibility, @@ -223,7 +267,10 @@ export function createPublicEventSessionNamespace(): { fetch: async () => { const socket = new PublicEventTestSocket(); sockets.add(socket); - return { status: 101, webSocket: socket as unknown as WebSocket } as Response; + return { + status: 101, + webSocket: socket as unknown as WebSocket, + } as Response; }, publishEvents: async () => { for (const socket of sockets) { @@ -236,7 +283,7 @@ export function createPublicEventSessionNamespace(): { binding: { get: () => stub, idFromName: (name: string) => name, - } as unknown as ApiBindings["Session"], + }, close: () => { for (const socket of sockets) { socket.close(); diff --git a/apps/api/tests/public-thread-api.e2e.test.ts b/apps/api/tests/public-thread-api.e2e.test.ts index 000319f2..56718597 100644 --- a/apps/api/tests/public-thread-api.e2e.test.ts +++ b/apps/api/tests/public-thread-api.e2e.test.ts @@ -1,7 +1,12 @@ import { describe, expect, test } from "bun:test"; import { PUBLIC_THREAD_API_THREADS_MAX_LIMIT } from "@mosoo/contracts/public-api"; -import { sessionExecutionSnapshotsTable, sessionRunsTable, sessionsTable } from "@mosoo/db"; +import { + sessionExecutionSnapshotsTable, + sessionMessagesTable, + sessionRunsTable, + sessionsTable, +} from "@mosoo/db"; import { eq } from "drizzle-orm"; import { fileStore } from "../src/modules/files/application/file-store"; @@ -9,14 +14,18 @@ import { PUBLIC_API_RATE_LIMIT_REQUESTS_PER_MINUTE, enforcePublicApiRateLimit, } from "../src/modules/public-api/public-api-rate-limit.service"; +import { createSessionProcessEventsFromSessionEventRows } from "../src/modules/sessions/application/session-process-events.service"; +import type { SessionEventProcessRow } from "../src/modules/sessions/application/session-process-events.service"; import { insertSessionMessage } from "../src/modules/sessions/infrastructure/session-message-store.repository"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { PublicApiMemoryFileBucket, PUBLIC_API_TEST_IDS, TOKENS, + createPre0014PublicHttpContractDatabase, createPublicHttpContractDatabase, createPublicHttpTestBindings, + migratePre0014PublicHttpContractDatabase, } from "./helpers/public-api-http-test-fixture"; import { OWNER_VIEWER, @@ -63,6 +72,7 @@ const PROGRESS_OUTPUT_TEXTS = [ "进度 2：已调用工具校验表格，不能进入最终回答。", "进度 3：artifact 已创建，不能进入最终回答。", ] as const; +const FINAL_MESSAGE_ID = "01J0000000000000000000001M"; type PublicHttpTestDatabase = Awaited>; @@ -97,7 +107,7 @@ async function createReadyAppDraftFile(input: { name: string; }): Promise { const bindings = createPublicHttpTestBindings(input.database, { - fileBucket: input.bucket as unknown as R2Bucket, + fileBucket: input.bucket, }) as ApiBindings; const fileBytes = new TextEncoder().encode(input.body); const upload = await fileStore.createUpload(bindings, OWNER_VIEWER, { @@ -136,7 +146,7 @@ async function createPendingAppDraftFile(input: { name: string; }): Promise { const bindings = createPublicHttpTestBindings(input.database, { - fileBucket: input.bucket as unknown as R2Bucket, + fileBucket: input.bucket, }) as ApiBindings; const fileBytes = new TextEncoder().encode(input.body); const upload = await fileStore.createUpload(bindings, OWNER_VIEWER, { @@ -537,6 +547,18 @@ describe("Public Thread API e2e", () => { seq: 2, sessionId: threadId, }); + await insertRuntimeEvent(database, { + kind: "tool.call.updated", + occurredAt: 975, + payload: { + status: "cancelled", + title: "record_meal", + toolCallId: "tool-call-2", + }, + runId, + seq: 3, + sessionId: threadId, + }); const toolEventsResponse = await requestPublicApi( app, @@ -552,16 +574,23 @@ describe("Public Thread API e2e", () => { expect(toolEvents.map((event) => event["toolCallId"])).toEqual([ "tool-call-1", "tool-call-1", + "tool-call-2", ]); expect(toolEvents[0]).toMatchObject({ + toolInput: { calories: 420, mealId: "meal-1" }, toolName: "record_meal", type: "tool.use.started", }); - expectNoProperties(toolEvents[0], ["toolInput"]); expect(toolEvents[1]).toMatchObject({ toolInput: { calories: 420, mealId: "meal-1" }, type: "tool.use.completed", }); + expect(toolEvents[2]).toMatchObject({ + content: "record_meal", + toolName: "record_meal", + type: "tool.use.completed", + }); + expectNoProperties(toolEvents[2], ["toolInput"]); await database.prepare("DELETE FROM session_event WHERE session_id = ?").bind(threadId).run(); @@ -594,7 +623,7 @@ describe("Public Thread API e2e", () => { occurredAt: 1_085, payload: { content: FINAL_OUTPUT_TEXT, - messageId: "assistant-final", + messageId: FINAL_MESSAGE_ID, role: "agent", }, runId, @@ -602,20 +631,28 @@ describe("Public Thread API e2e", () => { sessionId: threadId, }); await insertRuntimeEvent(database, { - kind: "run.completed", - occurredAt: 1_125, - payload: { stopReason: "debug" }, + kind: "message.completed", + occurredAt: 1_100, + payload: { messageId: FINAL_MESSAGE_ID, role: "agent" }, runId, seq: 6, sessionId: threadId, + }); + await insertRuntimeEvent(database, { + kind: "run.started", + occurredAt: 1_125, + payload: { startedAt: "1970-01-01T00:00:01.125Z" }, + runId, + seq: 7, + sessionId: threadId, visibility: "owner_debug", }); await insertRuntimeEvent(database, { kind: "run.completed", occurredAt: 1_150, - payload: { stopReason: "end_turn" }, + payload: { finalMessageId: FINAL_MESSAGE_ID, stopReason: "end_turn" }, runId, - seq: 7, + seq: 8, sessionId: threadId, }); @@ -652,13 +689,19 @@ describe("Public Thread API e2e", () => { }); } await insertSessionMessage(database, { - content: FINAL_OUTPUT_TEXT, + content: "", createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + id: FINAL_MESSAGE_ID, role: "assistant", - segments: [{ kind: "text", text: FINAL_OUTPUT_TEXT }], sessionId: threadId, sessionRunId: runId, }); + await database + .app() + .update(sessionMessagesTable) + .set({ projectionFormat: "event_stream_v3" }) + .where(eq(sessionMessagesTable.id, FINAL_MESSAGE_ID)) + .run(); await database .app() @@ -673,6 +716,7 @@ describe("Public Thread API e2e", () => { errorCode: null, errorDetailsJson: null, errorMessage: null, + errorRetryable: null, status: "completed", updatedAt: 1_150, }) @@ -1165,39 +1209,138 @@ describe("Public Thread API e2e", () => { liveEvents.close(); await insertRuntimeEvent(database, { - kind: "message.completed", + kind: "message.added", occurredAt: 3_250, - payload: { messageId: "assistant-stream-live-1", role: "agent" }, + payload: { + content: "Live stream \uE200cite\uE202hidden\uE201delta A", + messageId: "assistant-stream-live-1", + role: "agent", + }, runId, seq: 6, sessionId: threadId, }); await insertRuntimeEvent(database, { - kind: "message.added", + kind: "message.completed", occurredAt: 3_300, + payload: { messageId: "assistant-stream-live-1", role: "agent" }, + runId, + seq: 7, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + kind: "message.delta", + occurredAt: 3_350, payload: { - content: "Live stream \uE200cite\uE202hidden\uE201delta A", - messageId: "assistant-stream-live-1", + contentDelta: "Same poll ", + messageId: "assistant-stream-same-poll", role: "agent", }, runId, - seq: 7, + seq: 8, sessionId: threadId, }); await insertRuntimeEvent(database, { kind: "message.added", - occurredAt: 3_350, + occurredAt: 3_400, + payload: { + content: "Same poll suffix", + messageId: "assistant-stream-same-poll", + role: "agent", + }, + runId, + seq: 9, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + kind: "message.completed", + occurredAt: 3_450, + payload: { messageId: "assistant-stream-same-poll", role: "agent" }, + runId, + seq: 10, + sessionId: threadId, + }); + const noDeltaTerminals = [ + ["message.completed", { messageId: "no-delta-completed", role: "agent" }], + ["message.cancelled", { messageId: "no-delta-cancelled", role: "agent" }], + [ + "message.failed", + { + error: { code: "runtime.failed", message: "Runtime failed." }, + messageId: "no-delta-failed", + role: "agent", + }, + ], + ] as const; + + for (const [index, [kind, payload]] of noDeltaTerminals.entries()) { + await insertRuntimeEvent(database, { + kind, + occurredAt: 3_550 + index * 50, + payload, + runId, + seq: 12 + index, + sessionId: threadId, + }); + } + const streamedTerminals = [ + [ + "message.cancelled", + "Cancelled stream \uE200cite\uE202hidden-cancelled\uE201", + { messageId: "assistant-stream-cancelled", role: "agent" }, + ], + [ + "message.failed", + "Failed stream \uE200cite\uE202hidden-failed\uE201", + { + error: { code: "runtime.failed", message: "Runtime failed." }, + messageId: "assistant-stream-failed", + role: "agent", + }, + ], + ] as const; + + for (const [index, [kind, contentDelta, payload]] of streamedTerminals.entries()) { + const seq = 15 + index * 2; + await insertRuntimeEvent(database, { + kind: "message.delta", + occurredAt: 3_700 + index * 100, + payload: { contentDelta, messageId: payload.messageId, role: "agent" }, + runId, + seq, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + kind, + occurredAt: 3_750 + index * 100, + payload, + runId, + seq: seq + 1, + sessionId: threadId, + }); + } + await insertRuntimeEvent(database, { + kind: "message.added", + occurredAt: 3_900, payload: { content: "Tail marker", messageId: "tail-message", role: "agent", }, runId, - seq: 8, + seq: 19, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + kind: "message.completed", + occurredAt: 3_950, + payload: { messageId: "tail-message", role: "agent" }, + runId, + seq: 20, sessionId: threadId, }); - const text = `${openDeltaText}${await readUntil("id: 01J00000000000000000000017")}`; + const text = `${openDeltaText}${await readUntil("id: 01J0000000000000000000001K")}`; await reader.cancel(); expect(text).toContain("event: thread.event"); expect(text).toContain("id: 01J00000000000000000000011"); @@ -1205,13 +1348,35 @@ describe("Public Thread API e2e", () => { expect(text).not.toContain("id: 01J00000000000000000000013"); expect(text).toContain("id: 01J00000000000000000000014"); expect(text).not.toContain("id: 01J00000000000000000000015"); - expect(text).not.toContain("id: 01J00000000000000000000016"); + expect(text).toContain("id: 01J00000000000000000000016"); expect(text).toContain("id: 01J00000000000000000000017"); + expect(text).not.toContain("id: 01J00000000000000000000018"); + expect(text).toContain("id: 01J00000000000000000000019"); + expect(text).not.toContain("id: 01J0000000000000000000001A"); + expect(text).toContain("id: 01J0000000000000000000001B"); + expect(text).toContain("id: 01J0000000000000000000001C"); + expect(text).toContain("id: 01J0000000000000000000001D"); + expect(text).not.toContain("id: 01J0000000000000000000001E"); + expect(text).toContain("id: 01J0000000000000000000001F"); + expect(text).not.toContain("id: 01J0000000000000000000001G"); + expect(text).toContain("id: 01J0000000000000000000001H"); + expect(text).not.toContain("id: 01J0000000000000000000001J"); + expect(text).toContain("id: 01J0000000000000000000001K"); expect(text.match(/Live stream /gu)).toHaveLength(1); expect(text.match(/delta A/gu)).toHaveLength(1); + expect(text.match(/Same poll /gu)).toHaveLength(1); + expect(text.match(/suffix/gu)).toHaveLength(1); + expect(text.match(/Cancelled stream /gu)).toHaveLength(1); + expect(text.match(/Failed stream /gu)).toHaveLength(1); + expect(text).not.toContain("hidden-cancelled"); + expect(text).not.toContain("hidden-failed"); + expect(text).not.toContain("Message updated."); + expect(text).not.toContain("Runtime failed."); expect(text).toContain('"type":"agent.message.delta"'); + expect(text).toContain('"status":"error"'); + expect(text).toMatch(/id: 01J0000000000000000000001H\ndata: [^\n]*"status":"error"/u); expect(text).toContain('"toolCallId":"tool-call-stream-1"'); - expect(text).not.toContain('"toolInput"'); + expect(text).toContain('"toolInput":{"calories":420,"mealId":"meal-1"}'); expect(text).toContain('"toolName":"record_meal"'); expect(text).toContain('"runId":null'); expect(text).toContain('"content":"'); @@ -1220,9 +1385,2204 @@ describe("Public Thread API e2e", () => { expect(text).not.toContain("private-diagnostic"); expect(text).not.toContain("traceId"); expect(text).not.toContain("event: thread.error"); + + const replayResponse = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events?limit=100`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + const replayEvents = expectArray((await readJson(replayResponse))["events"]).map( + expectRecord, + ); + + for (const id of [ + "01J00000000000000000000016", + "01J00000000000000000000019", + "01J0000000000000000000001B", + "01J0000000000000000000001C", + "01J0000000000000000000001D", + "01J0000000000000000000001F", + "01J0000000000000000000001H", + "01J0000000000000000000001K", + ]) { + expect(replayEvents.some((event) => event["id"] === id)).toBe(true); + } + for (const id of ["01J0000000000000000000001D", "01J0000000000000000000001H"]) { + expect(replayEvents.find((event) => event["id"] === id)).toMatchObject({ + status: "error", + }); + } }); }, 10_000); + test.each([ + { + canonicalText: "Hello world", + liveText: "Hello ", + provider: "OpenAI", + snapshotContent: "Hello ", + threadIndex: 250, + }, + { + canonicalText: "Hello world", + liveText: "Hello ", + provider: "ACP", + snapshotContent: "Hello ", + threadIndex: 251, + }, + { + canonicalText: "Hello world", + liveText: "Hello ", + provider: "Claude", + snapshotContent: [{ text: "Hello ", type: "text" }], + threadIndex: 252, + }, + { + canonicalText: "Final world", + liveText: "Draft ", + provider: "divergent", + snapshotContent: "Final ", + threadIndex: 253, + }, + ] as const)( + "projects $provider snapshot-before-terminal order without replaying live text", + async ({ canonicalText, liveText, provider, snapshotContent, threadIndex }) => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(threadIndex); + const messageId = `snapshot-before-terminal-${provider.toLowerCase()}`; + + await insertPublicThread(database, { + id: threadId, + title: `${provider} snapshot order`, + updatedAt: 1_000, + }); + await insertRuntimeEvent(database, { + kind: "message.started", + occurredAt: 1_100, + payload: { messageId, role: "agent" }, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + kind: "message.delta", + occurredAt: 1_200, + payload: { contentDelta: liveText, messageId, role: "agent" }, + seq: 2, + sessionId: threadId, + }); + + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + + const decoder = new TextDecoder(); + let sseText = ""; + const readUntil = async (marker: string): Promise => { + while (!sseText.includes(marker)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(3_000).then(() => { + throw new Error(`Timed out waiting for ${marker}.`); + }), + ]); + + if (chunk.done) { + throw new Error(`SSE closed before ${marker}.`); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + }; + + await readUntil("id: 01J00000000000000000000011"); + await insertRuntimeEvent(database, { + kind: "message.added", + occurredAt: 1_300, + payload: { content: snapshotContent, messageId, role: "agent" }, + seq: 3, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + kind: "message.delta", + occurredAt: 1_400, + payload: { contentDelta: "world", messageId, role: "agent" }, + seq: 4, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + kind: "message.completed", + occurredAt: 1_500, + payload: { messageId, role: "agent" }, + seq: 5, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil("id: 01J00000000000000000000014"); + await reader.cancel(); + liveEvents.close(); + + const liveContent = sseText + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => expectRecord(JSON.parse(line.slice("data: ".length)))["content"]) + .filter((content): content is string => typeof content === "string") + .join(""); + const replayResponse = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + const replayContent = expectArray((await readJson(replayResponse))["events"]) + .map((event) => expectRecord(event)["content"]) + .filter((content): content is string => typeof content === "string") + .join(""); + + expect(replayContent).toBe(canonicalText); + expect(liveContent).toBe(provider === "divergent" ? liveText : canonicalText); + }, + 10_000, + ); + + test("reconciles a paginated authoritative snapshot above the former payload limit", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(270); + const messageId = "paginated-authoritative-snapshot"; + const eventId = (seq: number) => generatedPublicThreadId(18_000 + seq); + const fragments = Array.from({ length: 1_001 }, (_, index) => + index === 1_000 ? "🙂终" : "x".repeat(400), + ); + + await insertPublicThread(database, { + id: threadId, + title: "Paginated authoritative snapshot", + updatedAt: 1_000, + }); + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + await reader.read(); + + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "message.started", + occurredAt: 1_000, + payload: { messageId, role: "agent" }, + seq: 1, + sessionId: threadId, + }); + for (const [index, content] of fragments.entries()) { + const seq = index + 2; + await insertRuntimeEvent(database, { + eventId: eventId(seq), + kind: index === 0 ? "message.added" : "message.delta", + occurredAt: seq * 1_000, + payload: + index === 0 + ? { content, messageId, role: "agent" } + : { contentDelta: content, messageId, role: "agent" }, + seq, + sessionId: threadId, + }); + } + const terminalSeq = fragments.length + 2; + const terminalEventId = eventId(terminalSeq); + await insertRuntimeEvent(database, { + eventId: terminalEventId, + kind: "message.completed", + occurredAt: terminalSeq * 1_000, + payload: { messageId, role: "agent" }, + seq: terminalSeq, + sessionId: threadId, + }); + liveEvents.emit(); + + const decoder = new TextDecoder(); + let sseText = ""; + while (!sseText.includes(`id: ${terminalEventId}`)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(8_000).then(() => { + throw new Error("Timed out waiting for the paginated authoritative snapshot."); + }), + ]); + if (chunk.done) { + throw new Error("SSE closed before the paginated authoritative snapshot."); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + await reader.cancel(); + liveEvents.close(); + + const liveMessages = sseText + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => expectRecord(JSON.parse(line.slice("data: ".length)))) + .filter((event) => event["type"] === "agent.message.delta"); + const replayResponse = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + const replayMessages = expectArray((await readJson(replayResponse))["events"]) + .map(expectRecord) + .filter((event) => event["type"] === "agent.message.delta"); + + expect(fragments.join("").length).toBeGreaterThan(384 * 1024); + expect(liveMessages).toEqual(replayMessages); + expect(liveMessages).toEqual([ + expect.objectContaining({ content: fragments.join(""), id: terminalEventId }), + ]); + }, 15_000); + + test("retains exact prefix metadata for more than one public event page of messages", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(271); + const runId = PUBLIC_API_TEST_IDS.run; + const messageId = "message-before-full-state-page"; + const eventId = (seq: number) => generatedPublicThreadId(20_000 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Unbounded active message metadata", + updatedAt: 1_000, + }); + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + const decoder = new TextDecoder(); + let sseText = ""; + const readUntil = async (marker: string): Promise => { + while (!sseText.includes(marker)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(5_000).then(() => { + throw new Error(`Timed out waiting for ${marker}.`); + }), + ]); + if (chunk.done) { + throw new Error(`SSE closed before ${marker}.`); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + }; + + await readUntil(": connected"); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "run.started", + occurredAt: 1_000, + payload: { startedAt: "1970-01-01T00:00:01.000Z" }, + runId, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(2), + kind: "message.started", + occurredAt: 2_000, + payload: { messageId, role: "agent" }, + runId, + seq: 2, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(3), + kind: "message.delta", + occurredAt: 3_000, + payload: { contentDelta: "prefix", messageId, role: "agent" }, + runId, + seq: 3, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(3)}`); + + for (let seq = 4; seq <= 1_004; seq += 1) { + await insertRuntimeEvent(database, { + eventId: eventId(seq), + kind: "message.started", + occurredAt: seq * 1_000, + payload: { messageId: `parallel-message-${seq}`, role: "agent" }, + runId, + seq, + sessionId: threadId, + }); + } + await insertRuntimeEvent(database, { + eventId: eventId(1_005), + kind: "tool.call.updated", + occurredAt: 1_005_000, + payload: { + rawInput: "{}", + status: "running", + title: "state_page_marker", + toolCallId: "state-page-marker", + }, + runId, + seq: 1_005, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(1_005)}`); + + await insertRuntimeEvent(database, { + eventId: eventId(1_006), + kind: "message.added", + occurredAt: 1_006_000, + payload: { content: "prefix", messageId, role: "agent" }, + runId, + seq: 1_006, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(1_007), + kind: "message.delta", + occurredAt: 1_007_000, + payload: { contentDelta: " suffix", messageId, role: "agent" }, + runId, + seq: 1_007, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(1_008), + kind: "message.completed", + occurredAt: 1_008_000, + payload: { messageId, role: "agent" }, + runId, + seq: 1_008, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(1_008)}`); + await reader.cancel(); + liveEvents.close(); + + const liveContent = sseText + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => expectRecord(JSON.parse(line.slice("data: ".length)))) + .filter((event) => event["type"] === "agent.message.delta") + .map((event) => event["content"]) + .filter((content): content is string => typeof content === "string") + .join(""); + + expect(liveContent).toBe("prefix suffix"); + }, 15_000); + + test("replaces standalone assistant snapshots before publishing the terminal text", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(258); + const messageId = "standalone-authoritative-snapshot"; + const eventId = (seq: number) => generatedPublicThreadId(9_000 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Standalone authoritative snapshot", + updatedAt: 1_000, + }); + + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "message.added", + occurredAt: 1_100, + payload: { content: "Obsolete snapshot", messageId, role: "agent" }, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(2), + kind: "message.added", + occurredAt: 1_200, + payload: { content: "Final \uE200ci", messageId, role: "agent" }, + seq: 2, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(3), + kind: "message.delta", + occurredAt: 1_300, + payload: { contentDelta: "te\uE202private\uE201world", messageId, role: "agent" }, + seq: 3, + sessionId: threadId, + }); + const terminalEventId = eventId(4); + await insertRuntimeEvent(database, { + eventId: terminalEventId, + kind: "message.completed", + occurredAt: 1_400, + payload: { messageId, role: "agent" }, + seq: 4, + sessionId: threadId, + }); + liveEvents.emit(); + + const decoder = new TextDecoder(); + let sseText = ""; + + while (!sseText.includes(`id: ${terminalEventId}`)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(3_000).then(() => { + throw new Error("Timed out waiting for the terminal assistant snapshot."); + }), + ]); + + if (chunk.done) { + throw new Error("SSE closed before the terminal assistant snapshot."); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + + await reader.cancel(); + liveEvents.close(); + const liveContent = sseText + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => expectRecord(JSON.parse(line.slice("data: ".length)))["content"]) + .filter((content): content is string => typeof content === "string") + .join(""); + const replayResponse = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + const replayContent = expectArray((await readJson(replayResponse))["events"]) + .map((event) => expectRecord(event)["content"]) + .filter((content): content is string => typeof content === "string") + .join(""); + + expect(liveContent).toBe("Final world"); + expect(replayContent).toBe(liveContent); + expect(sseText).not.toContain("Obsolete snapshot"); + expect(sseText).not.toContain("private"); + }, 10_000); + + test.each([ + { + deltaContent: "world", + snapshotContent: "Final ", + terminalKind: "message.completed", + terminalBeforeConnect: true, + threadIndex: 254, + timing: "after the terminal", + }, + { + deltaContent: "world", + snapshotContent: "Final ", + terminalKind: "message.completed", + terminalBeforeConnect: false, + threadIndex: 255, + timing: "between the snapshot and terminal", + }, + { + deltaContent: "world", + snapshotContent: "Final ", + terminalKind: "run.failed", + terminalBeforeConnect: false, + threadIndex: 259, + timing: "before a repaired run terminal", + }, + { + deltaContent: "te\uE202private\uE201world", + snapshotContent: "Final \uE200ci", + terminalKind: "message.completed", + terminalBeforeConnect: false, + threadIndex: 260, + timing: "across a private citation chunk boundary", + }, + ] as const)( + "hydrates canonical snapshot text when connecting $timing", + async ({ deltaContent, snapshotContent, terminalBeforeConnect, terminalKind, threadIndex }) => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(threadIndex); + const messageId = `initial-canonical-${String(threadIndex)}`; + const runId = terminalKind === "run.failed" ? PUBLIC_API_TEST_IDS.run : null; + + await insertPublicThread(database, { + id: threadId, + title: "Initial canonical snapshot", + updatedAt: 1_000, + }); + await insertRuntimeEvent(database, { + kind: "message.started", + occurredAt: 1_100, + payload: { messageId, role: "agent" }, + runId, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + kind: "message.delta", + occurredAt: 1_200, + payload: { contentDelta: "Draft ", messageId, role: "agent" }, + runId, + seq: 2, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + kind: "message.added", + occurredAt: 1_300, + payload: { content: snapshotContent, messageId, role: "agent" }, + runId, + seq: 3, + sessionId: threadId, + }); + + const completeSnapshot = async () => { + await insertRuntimeEvent(database, { + kind: "message.delta", + occurredAt: 1_400, + payload: { contentDelta: deltaContent, messageId, role: "agent" }, + runId, + seq: 4, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + kind: terminalKind, + occurredAt: 1_500, + payload: + terminalKind === "message.completed" + ? { messageId, role: "agent" } + : { + error: { + code: "runtime.driver_terminal", + details: {}, + message: "Driver disconnected.", + retryable: true, + }, + recoverable: true, + }, + runId, + seq: 5, + sessionId: threadId, + }); + }; + + if (terminalBeforeConnect) { + await completeSnapshot(); + } + + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + + const decoder = new TextDecoder(); + let sseText = ""; + const readUntil = async (marker: string): Promise => { + while (!sseText.includes(marker)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(3_000).then(() => { + throw new Error(`Timed out waiting for ${marker}.`); + }), + ]); + + if (chunk.done) { + throw new Error(`SSE closed before ${marker}.`); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + }; + + await readUntil( + terminalBeforeConnect ? "id: 01J00000000000000000000014" : "id: 01J00000000000000000000012", + ); + if (!terminalBeforeConnect) { + await completeSnapshot(); + liveEvents.emit(); + await readUntil("id: 01J00000000000000000000014"); + } + await reader.cancel(); + liveEvents.close(); + + const liveContent = sseText + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => expectRecord(JSON.parse(line.slice("data: ".length)))) + .filter((event) => event["type"] === "agent.message.delta") + .map((event) => event["content"]) + .filter((content): content is string => typeof content === "string") + .join(""); + const replayResponse = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + const replayContent = expectArray((await readJson(replayResponse))["events"]) + .map(expectRecord) + .filter((event) => event["type"] === "agent.message.delta") + .map((event) => event["content"]) + .filter((content): content is string => typeof content === "string") + .join(""); + + expect(liveContent).toBe("Final world"); + expect(replayContent).toBe(liveContent); + expect(liveContent).not.toContain("Draft"); + expect(sseText).not.toContain("private"); + expect(sseText).not.toContain("\uE200"); + }, + 10_000, + ); + + test("keeps a citation parser alive across a repaired run terminal", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(262); + const messageId = "terminal-citation-recovery"; + const runId = PUBLIC_API_TEST_IDS.run; + const eventId = (seq: number) => generatedPublicThreadId(12_000 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Terminal citation recovery", + updatedAt: 1_000, + }); + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + + const decoder = new TextDecoder(); + let sseText = ""; + const readUntil = async (marker: string): Promise => { + while (!sseText.includes(marker)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(3_000).then(() => { + throw new Error(`Timed out waiting for ${marker}.`); + }), + ]); + if (chunk.done) { + throw new Error(`SSE closed before ${marker}.`); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + }; + + await readUntil(": connected"); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "message.started", + occurredAt: 1_100, + payload: { messageId, role: "agent" }, + runId, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(2), + kind: "message.delta", + occurredAt: 1_200, + payload: { contentDelta: "before\uE200ci", messageId, role: "agent" }, + runId, + seq: 2, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(2)}`); + await insertRuntimeEvent(database, { + eventId: eventId(3), + kind: "run.failed", + occurredAt: 1_300, + payload: { + error: { + code: "runtime.driver_terminal", + details: {}, + message: "Driver disconnected.", + retryable: true, + }, + recoverable: true, + }, + runId, + seq: 3, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(3)}`); + await insertRuntimeEvent(database, { + eventId: eventId(4), + kind: "message.delta", + occurredAt: 1_250, + payload: { + contentDelta: "te\uE202SECRET\uE201after", + messageId, + role: "agent", + }, + runId, + seq: 4, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(4)}`); + await reader.cancel(); + liveEvents.close(); + + const liveMessageContent = sseText + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => expectRecord(JSON.parse(line.slice("data: ".length)))) + .filter((event) => event["type"] === "agent.message.delta") + .map((event) => event["content"]) + .filter((content): content is string => typeof content === "string") + .join(""); + const replayResponse = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + const replayEvents = expectArray((await readJson(replayResponse))["events"]).map(expectRecord); + const replayMessages = replayEvents.filter((event) => event["type"] === "agent.message.delta"); + + expect(liveMessageContent).toBe("beforeafter"); + expect(replayMessages).toHaveLength(1); + expect(replayMessages[0]?.["content"]).toBe(liveMessageContent); + expect(sseText).not.toContain("SECRET"); + expect(sseText).not.toContain("\uE200"); + }, 10_000); + + test("keeps one canonical stream when an authoritative snapshot arrives after repair", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(263); + const messageId = "terminal-snapshot-recovery"; + const runId = PUBLIC_API_TEST_IDS.run; + const eventId = (seq: number) => generatedPublicThreadId(12_100 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Terminal snapshot recovery", + updatedAt: 1_000, + }); + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + + const decoder = new TextDecoder(); + let sseText = ""; + const readUntil = async (marker: string): Promise => { + while (!sseText.includes(marker)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(3_000).then(() => { + throw new Error(`Timed out waiting for ${marker}.`); + }), + ]); + if (chunk.done) { + throw new Error(`SSE closed before ${marker}.`); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + }; + + await readUntil(": connected"); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "message.started", + occurredAt: 1_100, + payload: { messageId, role: "agent" }, + runId, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(2), + kind: "message.delta", + occurredAt: 1_200, + payload: { contentDelta: "Draft", messageId, role: "agent" }, + runId, + seq: 2, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(2)}`); + await insertRuntimeEvent(database, { + eventId: eventId(3), + kind: "run.failed", + occurredAt: 1_500, + payload: { + error: { + code: "runtime.driver_terminal", + details: {}, + message: "Driver disconnected.", + retryable: true, + }, + recoverable: true, + }, + runId, + seq: 3, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(3)}`); + await insertRuntimeEvent(database, { + eventId: eventId(4), + kind: "message.started", + occurredAt: 1_250, + payload: { messageId, role: "agent" }, + runId, + seq: 4, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(5), + kind: "message.added", + occurredAt: 1_300, + payload: { content: "Obsolete", messageId, role: "agent" }, + runId, + seq: 5, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(6), + kind: "message.added", + occurredAt: 1_350, + payload: { content: "Final ", messageId, role: "agent" }, + runId, + seq: 6, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(7), + kind: "message.delta", + occurredAt: 1_400, + payload: { contentDelta: "world", messageId, role: "agent" }, + runId, + seq: 7, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(8), + kind: "tool.call.updated", + occurredAt: 1_600, + payload: { + rawInput: "{}", + status: "running", + title: "recovery_marker", + toolCallId: "terminal-snapshot-marker", + }, + runId, + seq: 8, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(8)}`); + await reader.cancel(); + liveEvents.close(); + + const liveMessages = sseText + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => expectRecord(JSON.parse(line.slice("data: ".length)))) + .filter((event) => event["type"] === "agent.message.delta"); + const replayResponse = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + const replayMessages = expectArray((await readJson(replayResponse))["events"]) + .map(expectRecord) + .filter((event) => event["type"] === "agent.message.delta"); + + expect(liveMessages.map((event) => event["content"]).join("")).toBe("Draft"); + expect(replayMessages).toHaveLength(1); + expect(replayMessages[0]?.["content"]).toBe("Final world"); + expect(sseText).not.toContain("Obsolete"); + }, 10_000); + + test("keeps final stream identity stable when its message terminal arrives after repair", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(268); + const messageId = "late-message-terminal"; + const runId = PUBLIC_API_TEST_IDS.run; + const eventId = (seq: number) => generatedPublicThreadId(16_000 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Late message terminal identity", + updatedAt: 1_000, + }); + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + const decoder = new TextDecoder(); + let sseText = ""; + const readUntil = async (marker: string): Promise => { + while (!sseText.includes(marker)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(3_000).then(() => { + throw new Error(`Timed out waiting for ${marker}.`); + }), + ]); + if (chunk.done) { + throw new Error(`SSE closed before ${marker}.`); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + }; + + await readUntil(": connected"); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "message.started", + occurredAt: 1_100, + payload: { messageId, role: "agent" }, + runId, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(2), + kind: "message.added", + occurredAt: 1_200, + payload: { content: "Final ", messageId, role: "agent" }, + runId, + seq: 2, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(3), + kind: "message.delta", + occurredAt: 1_300, + payload: { contentDelta: "world", messageId, role: "agent" }, + runId, + seq: 3, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(4), + kind: "run.failed", + occurredAt: 1_400, + payload: { + error: { + code: "runtime.driver_terminal", + details: {}, + message: "Driver disconnected.", + retryable: true, + }, + recoverable: true, + }, + runId, + seq: 4, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(4)}`); + await insertRuntimeEvent(database, { + eventId: eventId(5), + kind: "message.completed", + occurredAt: 1_350, + payload: { messageId, role: "agent" }, + runId, + seq: 5, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(6), + kind: "tool.call.updated", + occurredAt: 1_500, + payload: { + rawInput: "{}", + status: "running", + title: "late_terminal_marker", + toolCallId: "late-terminal-marker", + }, + runId, + seq: 6, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(6)}`); + await reader.cancel(); + liveEvents.close(); + + const liveEventsBody = sseText + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => expectRecord(JSON.parse(line.slice("data: ".length)))); + const replayResponse = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + const replayEvents = expectArray((await readJson(replayResponse))["events"]).map(expectRecord); + const liveMessage = liveEventsBody.filter((event) => event["type"] === "agent.message.delta"); + const replayMessage = replayEvents.filter((event) => event["type"] === "agent.message.delta"); + + expect(liveMessage).toEqual(replayMessage); + expect(liveMessage).toEqual([ + expect.objectContaining({ content: "Final world", id: eventId(2) }), + ]); + expect(sseText).not.toContain(`id: ${eventId(5)}`); + }, 10_000); + + test("retains an undisplayed stream parser across the initial event limit", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(264); + const messageId = "limited-citation-stream"; + const eventId = (seq: number) => generatedPublicThreadId(12_200 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Limited citation stream", + updatedAt: 1_000, + }); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "message.started", + occurredAt: 1_100, + payload: { messageId, role: "agent" }, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(2), + kind: "message.delta", + occurredAt: 1_200, + payload: { contentDelta: "before\uE200ci", messageId, role: "agent" }, + seq: 2, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(3), + kind: "tool.call.updated", + occurredAt: 1_300, + payload: { + rawInput: "{}", + status: "running", + title: "limit_marker", + toolCallId: "limited-citation-marker", + }, + seq: 3, + sessionId: threadId, + }); + + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream?limit=1`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + + const decoder = new TextDecoder(); + let sseText = ""; + const readUntil = async (marker: string): Promise => { + while (!sseText.includes(marker)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(3_000).then(() => { + throw new Error(`Timed out waiting for ${marker}.`); + }), + ]); + if (chunk.done) { + throw new Error(`SSE closed before ${marker}.`); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + }; + + await readUntil(`id: ${eventId(3)}`); + await insertRuntimeEvent(database, { + eventId: eventId(4), + kind: "message.delta", + occurredAt: 1_400, + payload: { + contentDelta: "te\uE202SECRET\uE201after", + messageId, + role: "agent", + }, + seq: 4, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(4)}`); + await reader.cancel(); + liveEvents.close(); + + expect(sseText).toContain('"content":"after"'); + expect(sseText).not.toContain("SECRET"); + expect(sseText).not.toContain("\uE200"); + }, 10_000); + + test("fails closed for an old stream outside the initial scan", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(265); + const messageId = "unscanned-citation-stream"; + const eventId = (seq: number) => generatedPublicThreadId(14_000 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Unscanned citation stream", + updatedAt: 1_000, + }); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "message.started", + occurredAt: 1_000, + payload: { messageId, role: "agent" }, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(2), + kind: "message.delta", + occurredAt: 2_000, + payload: { contentDelta: "before\uE200ci", messageId, role: "agent" }, + seq: 2, + sessionId: threadId, + }); + + for (let seq = 3; seq <= 1_003; seq += 1) { + await insertRuntimeEvent(database, { + eventId: eventId(seq), + kind: "tool.call.updated", + occurredAt: seq * 1_000, + payload: { + rawInput: "{}", + status: "running", + title: "scan_filler", + toolCallId: `scan-filler-${seq}`, + }, + seq, + sessionId: threadId, + }); + } + + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream?limit=1`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + const decoder = new TextDecoder(); + let sseText = ""; + const readUntil = async (marker: string): Promise => { + while (!sseText.includes(marker)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(3_000).then(() => { + throw new Error(`Timed out waiting for ${marker}.`); + }), + ]); + if (chunk.done) { + throw new Error(`SSE closed before ${marker}.`); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + }; + + await readUntil(`id: ${eventId(1_003)}`); + await insertRuntimeEvent(database, { + eventId: eventId(1_004), + kind: "message.delta", + occurredAt: 1_004_000, + payload: { + contentDelta: "te\uE202SECRET\uE201after", + messageId, + role: "agent", + }, + seq: 1_004, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(1_005), + kind: "tool.call.updated", + occurredAt: 1_005_000, + payload: { + rawInput: "{}", + status: "running", + title: "post_scan_marker", + toolCallId: "post-scan-marker", + }, + seq: 1_005, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(1_005)}`); + await reader.cancel(); + liveEvents.close(); + + expect(sseText).not.toContain("SECRET"); + expect(sseText).not.toContain("\uE200"); + expect(sseText).not.toContain(`id: ${eventId(1_004)}`); + }, 10_000); + + test("trusts a delta parser when an exact full page reaches the database start", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(269); + const messageId = "exact-page-citation-stream"; + const eventId = (seq: number) => generatedPublicThreadId(17_000 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Exact page citation stream", + updatedAt: 1_000, + }); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "message.delta", + occurredAt: 1_000, + payload: { contentDelta: "before\uE200ci", messageId, role: "agent" }, + seq: 1, + sessionId: threadId, + }); + + for (let seq = 2; seq <= 1_000; seq += 1) { + await insertRuntimeEvent(database, { + eventId: eventId(seq), + kind: "tool.call.updated", + occurredAt: seq * 1_000, + payload: { + rawInput: "{}", + status: "running", + title: "page_filler", + toolCallId: `page-filler-${seq}`, + }, + seq, + sessionId: threadId, + }); + } + + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream?limit=1`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + const decoder = new TextDecoder(); + let sseText = ""; + const readUntil = async (marker: string): Promise => { + while (!sseText.includes(marker)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(3_000).then(() => { + throw new Error(`Timed out waiting for ${marker}.`); + }), + ]); + if (chunk.done) { + throw new Error(`SSE closed before ${marker}.`); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + }; + + await readUntil(`id: ${eventId(1_000)}`); + await insertRuntimeEvent(database, { + eventId: eventId(1_001), + kind: "message.delta", + occurredAt: 1_001_000, + payload: { + contentDelta: "te\uE202SECRET\uE201after", + messageId, + role: "agent", + }, + seq: 1_001, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(1_001)}`); + await reader.cancel(); + liveEvents.close(); + + expect(sseText).toContain('"content":"after"'); + expect(sseText).not.toContain("SECRET"); + expect(sseText).not.toContain("\uE200"); + }, 10_000); + + test("does not revive an omitted terminal snapshot on a later event", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(266); + const messageId = "omitted-terminal-snapshot"; + const runId = PUBLIC_API_TEST_IDS.run; + const eventId = (seq: number) => generatedPublicThreadId(15_100 + seq); + const failure = { + error: { + code: "runtime.driver_terminal", + details: {}, + message: "Driver disconnected.", + retryable: true, + }, + recoverable: true, + }; + + await insertPublicThread(database, { + id: threadId, + title: "Omitted terminal snapshot", + updatedAt: 1_000, + }); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "message.started", + occurredAt: 1_100, + payload: { messageId, role: "agent" }, + runId, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(2), + kind: "message.added", + occurredAt: 1_200, + payload: { content: "OLD OMITTED", messageId, role: "agent" }, + runId, + seq: 2, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(3), + kind: "message.completed", + occurredAt: 1_300, + payload: { messageId, role: "agent" }, + runId, + seq: 3, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(4), + kind: "run.failed", + occurredAt: 1_400, + payload: failure, + runId, + seq: 4, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(5), + kind: "tool.call.updated", + occurredAt: 1_500, + payload: { + rawInput: "{}", + status: "running", + title: "initial_marker", + toolCallId: "initial-terminal-marker", + }, + runId, + seq: 5, + sessionId: threadId, + }); + + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream?limit=1`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + const decoder = new TextDecoder(); + let sseText = ""; + const readUntil = async (marker: string): Promise => { + while (!sseText.includes(marker)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(3_000).then(() => { + throw new Error(`Timed out waiting for ${marker}.`); + }), + ]); + if (chunk.done) { + throw new Error(`SSE closed before ${marker}.`); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + }; + + await readUntil(`id: ${eventId(5)}`); + await insertRuntimeEvent(database, { + eventId: eventId(6), + kind: "tool.call.updated", + occurredAt: 1_600, + payload: { + rawInput: "{}", + status: "running", + title: "later_marker", + toolCallId: "later-terminal-marker", + }, + runId, + seq: 6, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(6)}`); + await reader.cancel(); + liveEvents.close(); + + expect(sseText).not.toContain("OLD OMITTED"); + expect(sseText).not.toContain(`id: ${eventId(2)}`); + }, 10_000); + + test("releases completed run state and rejects late fragments after the next run", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(267); + const firstRunId = PUBLIC_API_TEST_IDS.run; + const secondRunId = generatedPublicThreadId(15_250); + const messageId = "released-run-citation"; + const eventId = (seq: number) => generatedPublicThreadId(15_300 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Released run state", + updatedAt: 1_000, + }); + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + const decoder = new TextDecoder(); + let sseText = ""; + const readUntil = async (marker: string): Promise => { + while (!sseText.includes(marker)) { + const chunk = await Promise.race([ + reader.read(), + Bun.sleep(3_000).then(() => { + throw new Error(`Timed out waiting for ${marker}.`); + }), + ]); + if (chunk.done) { + throw new Error(`SSE closed before ${marker}.`); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + }; + + await readUntil(": connected"); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "message.started", + occurredAt: 1_100, + payload: { messageId, role: "agent" }, + runId: firstRunId, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(2), + kind: "message.delta", + occurredAt: 1_200, + payload: { contentDelta: "before\uE200ci", messageId, role: "agent" }, + runId: firstRunId, + seq: 2, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(3), + kind: "run.failed", + occurredAt: 1_300, + payload: { + error: { + code: "runtime.driver_terminal", + details: {}, + message: "Driver disconnected.", + retryable: true, + }, + recoverable: true, + }, + runId: firstRunId, + seq: 3, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(3)}`); + await insertRuntimeEvent(database, { + eventId: eventId(4), + kind: "run.started", + occurredAt: 1_400, + payload: { startedAt: "1970-01-01T00:00:01.400Z" }, + runId: secondRunId, + seq: 4, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(5), + kind: "message.delta", + occurredAt: 1_250, + payload: { + contentDelta: "te\uE202SECRET\uE201after", + messageId, + role: "agent", + }, + runId: firstRunId, + seq: 5, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(6), + kind: "tool.call.updated", + occurredAt: 1_500, + payload: { + rawInput: "{}", + status: "running", + title: "next_run_marker", + toolCallId: "next-run-marker", + }, + runId: secondRunId, + seq: 6, + sessionId: threadId, + }); + liveEvents.emit(); + await readUntil(`id: ${eventId(6)}`); + await reader.cancel(); + liveEvents.close(); + + expect(sseText).toContain('"content":"before"'); + expect(sseText).not.toContain("SECRET"); + expect(sseText).not.toContain(`id: ${eventId(5)}`); + }, 10_000); + + test("reads a complete retained stream across public event pages", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const threadId = generatedPublicThreadId(256); + const runId = PUBLIC_API_TEST_IDS.run; + const fragmentCount = 520; + const eventId = (seq: number) => generatedPublicThreadId(4_000 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Paged public streams", + updatedAt: 1_000, + }); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "run.started", + occurredAt: 1_000, + payload: { startedAt: "1970-01-01T00:00:01.000Z" }, + runId, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(2), + kind: "thought.started", + occurredAt: 2_000, + payload: { thoughtId: "thought-a" }, + runId, + seq: 2, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(3), + kind: "thought.started", + occurredAt: 3_000, + payload: { thoughtId: "thought-b" }, + runId, + seq: 3, + sessionId: threadId, + }); + + for (let index = 0; index < fragmentCount; index += 1) { + const seq = index * 2 + 4; + await insertRuntimeEvent(database, { + eventId: eventId(seq), + kind: "thought.delta", + occurredAt: seq * 1_000, + payload: { contentDelta: "a", thoughtId: "thought-a" }, + runId, + seq, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(seq + 1), + kind: "thought.delta", + occurredAt: (seq + 1) * 1_000, + payload: { contentDelta: "b", thoughtId: "thought-b" }, + runId, + seq: seq + 1, + sessionId: threadId, + }); + } + + const firstTerminalSeq = fragmentCount * 2 + 4; + for (const [offset, thoughtId] of ["thought-a", "thought-b"].entries()) { + const seq = firstTerminalSeq + offset; + await insertRuntimeEvent(database, { + eventId: eventId(seq), + kind: "thought.completed", + occurredAt: seq * 1_000, + payload: { thoughtId }, + runId, + seq, + sessionId: threadId, + }); + } + await insertRuntimeEvent(database, { + eventId: eventId(firstTerminalSeq + 2), + kind: "run.completed", + occurredAt: (firstTerminalSeq + 2) * 1_000, + payload: { stopReason: "end_turn" }, + runId, + seq: firstTerminalSeq + 2, + sessionId: threadId, + }); + + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events?limit=2`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + const body = await readJson(response); + const events = expectArray(body["events"]).map(expectRecord); + + expect(body["truncated"]).toBe(true); + expect(events.map((event) => event["type"])).toEqual(["agent.thinking.delta", "run.completed"]); + expect(events[0]?.["content"]).toBe("b".repeat(fragmentCount)); + }); + + test("orders public pages by durable sequence rather than driver time", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const threadId = generatedPublicThreadId(261); + const runId = PUBLIC_API_TEST_IDS.run; + const fragmentCount = 1_002; + const eventId = (seq: number) => generatedPublicThreadId(10_000 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Occurrence-ordered public stream", + updatedAt: 1_000, + }); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "run.started", + occurredAt: 1_000, + payload: { startedAt: "1970-01-01T00:00:01.000Z" }, + runId, + seq: 1, + sessionId: threadId, + }); + await insertRuntimeEvent(database, { + eventId: eventId(2), + kind: "message.started", + occurredAt: 100_000, + payload: { messageId: "out-of-order-public-message", role: "agent" }, + runId, + seq: 2, + sessionId: threadId, + }); + + for (let index = 0; index < fragmentCount; index += 1) { + const seq = index + 3; + await insertRuntimeEvent(database, { + eventId: eventId(seq), + kind: "message.delta", + occurredAt: 101_000 + index, + payload: { + contentDelta: "x", + messageId: "out-of-order-public-message", + role: "agent", + }, + runId, + seq, + sessionId: threadId, + }); + } + + const toolSeq = fragmentCount + 3; + await insertRuntimeEvent(database, { + eventId: eventId(toolSeq), + kind: "tool.call.updated", + occurredAt: 2_000, + payload: { + rawInput: "{}", + status: "running", + title: "read_file", + toolCallId: "out-of-order-tool", + }, + runId, + seq: toolSeq, + sessionId: threadId, + }); + const runTerminalSeq = toolSeq + 1; + await insertRuntimeEvent(database, { + eventId: eventId(runTerminalSeq), + kind: "run.completed", + occurredAt: 3_000, + payload: { stopReason: "end_turn" }, + runId, + seq: runTerminalSeq, + sessionId: threadId, + }); + + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events?limit=2`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + const body = await readJson(response); + const events = expectArray(body["events"]).map(expectRecord); + + expect(body["truncated"]).toBe(true); + expect(events.map((event) => event["type"])).toEqual(["tool.use.started", "run.completed"]); + expect(events[0]?.["toolCallId"]).toBe("out-of-order-tool"); + }); + + test("fails closed for a public message cut by the raw scan ceiling", async () => { + const database = await createPublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const liveEvents = createPublicEventSessionNamespace(); + const threadId = generatedPublicThreadId(257); + const runId = PUBLIC_API_TEST_IDS.run; + const fragmentCount = 20_001; + const eventId = (seq: number) => generatedPublicThreadId(7_000 + seq); + + await insertPublicThread(database, { + id: threadId, + title: "Ceiling-bounded public stream", + updatedAt: 1_000, + }); + await insertRuntimeEvent(database, { + eventId: eventId(1), + kind: "message.started", + occurredAt: 1_000, + payload: { messageId: "ceiling-message", role: "agent" }, + runId, + seq: 1, + sessionId: threadId, + }); + + for (let index = 0; index < fragmentCount; index += 1) { + const seq = index + 2; + await insertRuntimeEvent(database, { + eventId: eventId(seq), + kind: "message.delta", + occurredAt: seq * 1_000, + payload: { + contentDelta: index === 0 ? "before\uE200ci" : index === 1 ? "te\uE202SECRET" : "SECRET", + messageId: "ceiling-message", + role: "agent", + }, + runId, + seq, + sessionId: threadId, + }); + } + + const messageTerminalSeq = fragmentCount + 2; + await insertRuntimeEvent(database, { + eventId: eventId(messageTerminalSeq), + kind: "message.completed", + occurredAt: messageTerminalSeq * 1_000, + payload: { messageId: "ceiling-message", role: "agent" }, + runId, + seq: messageTerminalSeq, + sessionId: threadId, + }); + const finalMessageSeq = messageTerminalSeq + 1; + await insertRuntimeEvent(database, { + eventId: eventId(finalMessageSeq), + kind: "message.added", + occurredAt: finalMessageSeq * 1_000, + payload: { + content: "Complete final answer", + messageId: "complete-final-message", + role: "agent", + }, + runId, + seq: finalMessageSeq, + sessionId: threadId, + }); + const runTerminalSeq = finalMessageSeq + 1; + const runTerminalEventId = eventId(runTerminalSeq); + await insertRuntimeEvent(database, { + eventId: runTerminalEventId, + kind: "run.completed", + occurredAt: runTerminalSeq * 1_000, + payload: { stopReason: "end_turn" }, + runId, + seq: runTerminalSeq, + sessionId: threadId, + }); + + const listResponse = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events?limit=2`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + const listBody = await readJson(listResponse); + const listEvents = expectArray(listBody["events"]).map(expectRecord); + expect(listBody["truncated"]).toBe(true); + expect(listEvents.map((event) => event["type"])).toEqual([ + "agent.message.delta", + "run.completed", + ]); + expect(listEvents[0]?.["content"]).toBe("Complete final answer"); + + const streamResponse = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events/stream?limit=2`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + { sessionNamespace: liveEvents.binding }, + ); + const reader = streamResponse.body?.getReader(); + if (!reader) { + throw new Error("Expected stream response body."); + } + const decoder = new TextDecoder(); + let sseText = ""; + + while (!sseText.includes(`id: ${runTerminalEventId}`)) { + const chunk = await reader.read(); + if (chunk.done) { + throw new Error("SSE closed before the initial canonical event."); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + + const lateDeltaSeq = runTerminalSeq + 1; + await insertRuntimeEvent(database, { + eventId: eventId(lateDeltaSeq), + kind: "message.delta", + occurredAt: lateDeltaSeq * 1_000, + payload: { + contentDelta: "SECRET\uE201after", + messageId: "ceiling-message", + role: "agent", + }, + runId, + seq: lateDeltaSeq, + sessionId: threadId, + }); + const markerSeq = lateDeltaSeq + 1; + const markerEventId = eventId(markerSeq); + await insertRuntimeEvent(database, { + eventId: markerEventId, + kind: "tool.call.updated", + occurredAt: markerSeq * 1_000, + payload: { + rawInput: "{}", + status: "running", + title: "ceiling_marker", + toolCallId: "ceiling-marker", + }, + runId, + seq: markerSeq, + sessionId: threadId, + }); + liveEvents.emit(); + + while (!sseText.includes(`id: ${markerEventId}`)) { + const chunk = await reader.read(); + if (chunk.done) { + throw new Error("SSE closed before the post-ceiling marker."); + } + sseText += decoder.decode(chunk.value, { stream: true }); + } + + await reader.cancel(); + liveEvents.close(); + const liveEventsBody = sseText + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => expectRecord(JSON.parse(line.slice("data: ".length)))); + + expect(liveEventsBody.slice(0, listEvents.length)).toEqual(listEvents); + expect(liveEventsBody.at(-1)?.["id"]).toBe(markerEventId); + expect(sseText).not.toContain("SECRET"); + expect(sseText).not.toContain("\uE200"); + }, 30_000); + + test("reads row-scoped pre-0014 stream identities after migration", async () => { + const database = await createPre0014PublicHttpContractDatabase(); + const app = createPublicThreadApiTestApp(); + const threadId = generatedPublicThreadId(999); + const eventIds = ["01J0000000000000000000001F", "01J0000000000000000000001G"]; + + await database + .prepare( + `INSERT INTO session ( + agent_id, archived_at, end_user_id, created_at, creator_account_id, + deployment_version_id, deployment_version_number, id, kind, + metadata_json, model, app_id, provider, renamed, runtime_id, status, + title, type, updated_at + ) VALUES (?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + PUBLIC_API_TEST_IDS.agent, + "customer-123", + 1_000, + PUBLIC_API_TEST_IDS.ownerAccount, + PUBLIC_API_TEST_IDS.deployment, + 1, + threadId, + "pet", + JSON.stringify({ + public_api: { + created_by: { + token_id: PUBLIC_API_TEST_IDS.patOwner, + token_label: PUBLIC_API_TEST_IDS.patOwner, + }, + idempotency_key: null, + source: "public_api", + }, + }), + "gpt-5.4", + PUBLIC_API_TEST_IDS.app, + "openai", + false, + "openai-runtime", + "IDLE", + "Pre-0014 stream fixture", + "ui", + 1_000, + ) + .run(); + for (const [index, eventId] of eventIds.entries()) { + await database + .prepare( + `INSERT INTO session_event ( + id, session_id, agent_id, seq, content_text, ended_at, event_type, + family, process_status, process_type, source, source_event_id, + visibility, occurred_at, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + eventId, + threadId, + PUBLIC_API_TEST_IDS.agent, + index + 1, + `Legacy row ${index + 1}`, + 1_100 + index, + "message.delta", + "process", + "available", + "agent.message.delta", + "driver", + eventId, + "all_consumers", + 1_100 + index, + 1_100 + index, + ) + .run(); + } + + migratePre0014PublicHttpContractDatabase(database); + + const rows = await database + .prepare( + `SELECT content_text, ended_at, event_type, id, occurred_at, + process_status, process_type, run_id, seq, stream_id, tokens + FROM session_event + WHERE session_id = ? + ORDER BY seq`, + ) + .bind(threadId) + .all(); + expect(rows.results.map((row) => row.stream_id)).toEqual(eventIds); + await expect( + database + .prepare( + "SELECT semantic_hash, tool_parent_message_id, tool_result_message_id, tool_status FROM session_event WHERE session_id = ? ORDER BY seq", + ) + .bind(threadId) + .all(), + ).resolves.toMatchObject({ + results: eventIds.map(() => ({ + semantic_hash: null, + tool_parent_message_id: null, + tool_result_message_id: null, + tool_status: null, + })), + success: true, + }); + await expect( + database + .prepare("UPDATE session_event SET semantic_hash = 'not-a-sha256' WHERE id = ?") + .bind(eventIds[0]) + .run(), + ).rejects.toThrow(); + expect( + createSessionProcessEventsFromSessionEventRows(rows.results).map((event) => event.content), + ).toEqual(["Legacy row 1", "Legacy row 2"]); + + const response = await requestPublicApi( + app, + database, + new Request(`https://api.example.com/api/v1/threads/${threadId}/events`, { + headers: { Authorization: bearer(TOKENS.owner) }, + }), + ); + expect(response.status).toBe(200); + expect( + expectArray((await readJson(response))["events"]).map( + (event) => expectRecord(event)["content"], + ), + ).toEqual(["Legacy row 1", "Legacy row 2"]); + + await database + .prepare( + "UPDATE session_event SET event_type = 'tool.call.updated', run_id = ?, tool_call_id = 'tool-1', tool_status = 'completed', mcp_command_id = '01J0000000000000000000001J' WHERE id = ?", + ) + .bind(PUBLIC_API_TEST_IDS.run, eventIds[0]) + .run(); + await expect( + database + .prepare( + "UPDATE session_event SET event_type = 'tool.call.updated', run_id = ?, tool_call_id = 'tool-1', tool_status = 'running' WHERE id = ?", + ) + .bind(PUBLIC_API_TEST_IDS.run, eventIds[1]) + .run(), + ).resolves.toBeDefined(); + await expect( + database + .prepare( + "UPDATE session_event SET tool_status = 'failed', mcp_command_id = '01J0000000000000000000001J' WHERE id = ?", + ) + .bind(eventIds[1]) + .run(), + ).rejects.toThrow("UNIQUE constraint failed"); + }); + test("bounds public Thread lists on stable latest ordering", async () => { const database = await createPublicHttpContractDatabase(); const app = createPublicThreadApiTestApp(); @@ -1261,7 +3621,7 @@ describe("Public Thread API e2e", () => { const app = createPublicThreadApiTestApp(); const bucket = new PublicApiMemoryFileBucket(); const requestThreadApi = (request: Request) => - requestPublicApi(app, database, request, { fileBucket: bucket as unknown as R2Bucket }); + requestPublicApi(app, database, request, { fileBucket: bucket }); await withProviderProbeMock(async () => { const createThreadResponse = await requestThreadApi( @@ -1585,7 +3945,7 @@ describe("Public Thread API e2e", () => { const app = createPublicThreadApiTestApp(); const bucket = new PublicApiMemoryFileBucket(); const requestThreadApi = (request: Request) => - requestPublicApi(app, database, request, { fileBucket: bucket as unknown as R2Bucket }); + requestPublicApi(app, database, request, { fileBucket: bucket }); await withProviderProbeMock(async () => { const fileBody = "Launch note.\n"; @@ -1805,7 +4165,7 @@ describe("Public Thread API e2e", () => { const app = createPublicThreadApiTestApp(); const bucket = new PublicApiMemoryFileBucket(); const requestThreadApi = (request: Request) => - requestPublicApi(app, database, request, { fileBucket: bucket as unknown as R2Bucket }); + requestPublicApi(app, database, request, { fileBucket: bucket }); const threadId = generatedPublicThreadId(130); await insertPublicThread(database, { diff --git a/apps/api/tests/rpc-wire-v3.test.ts b/apps/api/tests/rpc-wire-v3.test.ts new file mode 100644 index 00000000..3f6bcfe1 --- /dev/null +++ b/apps/api/tests/rpc-wire-v3.test.ts @@ -0,0 +1,585 @@ +import { describe, expect, test } from "bun:test"; + +import { ExternalToolEffectClaim } from "@mosoo/contracts/external-tool-effect"; +import { + DURABLE_RUN_ERROR_MAX_UTF8_BYTES, + measureDurableRunErrorJson, +} from "@mosoo/contracts/session-run"; +import { parseSchemaValue } from "@mosoo/contracts/validation"; +import { PLATFORM_ID_FIXTURES } from "@mosoo/id/testing"; +import { createRuntimeEvent } from "@mosoo/runtime-events"; + +import { + parseDriverCommandUpdateInput, + parseDriverCompletionInput, + parseDriverEventBatchInput, + parseDriverEventBatchOutput, + parseDriverExternalToolEffectClaimInput, + parseDriverExternalToolEffectObserveInput, + parseDriverExternalToolEffectSettleInput, + parseDriverFailureInput, + parseDriverHeartbeatInput, + parseDriverHeartbeatOutput, + parseDriverHelloInput, + parseDriverHelloOutput, + parseDriverLogBatchInput, + parseDriverLogBatchOutput, + parseDriverNextCommandInput, + parseDriverNextCommandOutput, + parseDriverReadyInput, +} from "../src/modules/runtime/infrastructure/driver-instance/rpc-wire"; + +const EXTERNAL_TOOL_EFFECT_CLAIM_TOKEN = "123e4567-e89b-42d3-a456-426614174000"; + +const invalidPositiveSafeIntegers = [ + 0, + -1, + 0.5, + Number.NaN, + Number.POSITIVE_INFINITY, + Number.NEGATIVE_INFINITY, + Number.MAX_SAFE_INTEGER + 1, +] as const; +const invalidNonNegativeSafeIntegers = invalidPositiveSafeIntegers.slice(1); + +function helloInput(overrides: Record = {}) { + return { + capabilities: [], + driverVersion: "test", + pid: 1, + protocolVersion: 3, + runtime: "acp-fallback", + startedAt: "2026-08-29T00:00:00.000Z", + ...overrides, + }; +} + +function helloOutput(runConfig: Record = {}, runId: string | null = null) { + return { + acceptedCapabilities: [], + connectionId: "connection-1", + driverInstanceId: "driver-1", + heartbeatIntervalMs: 250, + runConfig: { + commandLeaseMs: 0, + envPolicy: "strict", + eventBatchMaxSize: 1, + organizationPath: "/workspace", + ...runConfig, + }, + runId, + }; +} + +function heartbeatInput(pid: number, at = "2026-08-29T00:00:00.000Z") { + return { at, pid, reason: "interval" }; +} + +function readyInput(pid: number) { + return { + at: "2026-08-29T00:00:00.000Z", + driverInstanceId: "driver-1", + pid, + }; +} + +function logBatch(seq: number) { + return { + driverInstanceId: "driver-1", + logs: [{ level: "info", message: "message", seq, timestamp: "now" }], + }; +} + +function eventBatchOutput(seq: number, type = "message.delta") { + return { accepted: [{ eventId: "source-1", seq, type }] }; +} + +describe("API Driver RPC wire v3", () => { + test("rejects undeclared keys throughout fixed RPC shapes", () => { + const event = createRuntimeEvent({ + id: PLATFORM_ID_FIXTURES.runtimeEvent, + kind: "diagnostic.reported", + occurredAt: "2026-08-29T00:00:00.000Z", + payload: { message: "ok" }, + sessionId: PLATFORM_ID_FIXTURES.session, + }); + const cases: ReadonlyArray unknown]> = [ + ["hello", () => parseDriverHelloInput(helloInput({ extra: true }))], + ["hello output", () => parseDriverHelloOutput({ ...helloOutput(), extra: true })], + ["run config", () => parseDriverHelloOutput(helloOutput({ extra: true }))], + ["heartbeat", () => parseDriverHeartbeatInput({ ...heartbeatInput(1), extra: true })], + [ + "heartbeat output", + () => parseDriverHeartbeatOutput({ heartbeatCount: 0, extra: true, ok: true }), + ], + ["ready", () => parseDriverReadyInput({ ...readyInput(1), extra: true })], + [ + "command update", + () => + parseDriverCommandUpdateInput({ + commandId: "command-1", + driverInstanceId: "driver-1", + extra: true, + status: "accepted", + }), + ], + [ + "effect claim", + () => + parseDriverExternalToolEffectClaimInput({ + claimToken: EXTERNAL_TOOL_EFFECT_CLAIM_TOKEN, + commandId: "command-1", + driverInstanceId: "driver-1", + extra: true, + }), + ], + [ + "effect observe", + () => + parseDriverExternalToolEffectObserveInput({ + commandId: "command-1", + driverInstanceId: "driver-1", + extra: true, + }), + ], + [ + "effect settle", + () => + parseDriverExternalToolEffectSettleInput({ + claimToken: EXTERNAL_TOOL_EFFECT_CLAIM_TOKEN, + commandId: "command-1", + driverInstanceId: "driver-1", + effectId: "effect-1", + extra: true, + settlement: { kind: "unknown" }, + }), + ], + [ + "next command", + () => parseDriverNextCommandInput({ driverInstanceId: "driver-1", extra: true }), + ], + ["next command output", () => parseDriverNextCommandOutput({ command: null, extra: true })], + [ + "completion", + () => + parseDriverCompletionInput({ driverInstanceId: "driver-1", extra: true, runId: "run-1" }), + ], + [ + "failure", + () => + parseDriverFailureInput({ + driverInstanceId: "driver-1", + error: { code: "failed", details: {}, message: "failed", retryable: false }, + extra: true, + runId: "run-1", + }), + ], + [ + "event batch", + () => + parseDriverEventBatchInput({ + driverInstanceId: "driver-1", + events: [{ event, eventId: "event-1" }], + extra: true, + }), + ], + [ + "event envelope", + () => + parseDriverEventBatchInput({ + driverInstanceId: "driver-1", + events: [{ event, eventId: "event-1", extra: true }], + }), + ], + [ + "event body", + () => + parseDriverEventBatchInput({ + driverInstanceId: "driver-1", + events: [{ event: { ...event, extra: true }, eventId: "event-1" }], + }), + ], + [ + "event context", + () => + parseDriverEventBatchInput({ + driverInstanceId: "driver-1", + events: [{ event: { ...event, context: {} }, eventId: "event-1" }], + }), + ], + [ + "event surface", + () => + parseDriverEventBatchInput({ + driverInstanceId: "driver-1", + events: [{ event: { ...event, surface: {} }, eventId: "event-1" }], + }), + ], + ["event output", () => parseDriverEventBatchOutput({ ...eventBatchOutput(0), extra: true })], + [ + "event receipt", + () => + parseDriverEventBatchOutput({ + accepted: [{ eventId: "event-1", extra: true, seq: 0, type: "message.delta" }], + }), + ], + ["log batch", () => parseDriverLogBatchInput({ ...logBatch(0), extra: true })], + [ + "log entry", + () => + parseDriverLogBatchInput({ + driverInstanceId: "driver-1", + logs: [{ ...logBatch(0).logs[0], extra: true }], + }), + ], + [ + "log context", + () => + parseDriverLogBatchInput({ + driverInstanceId: "driver-1", + logs: [{ ...logBatch(0).logs[0], context: { extra: true } }], + }), + ], + [ + "log error", + () => + parseDriverLogBatchInput({ + driverInstanceId: "driver-1", + logs: [ + { ...logBatch(0).logs[0], error: { extra: true, message: "failed", name: "Error" } }, + ], + }), + ], + ["log output", () => parseDriverLogBatchOutput({ extra: true, ok: true })], + ]; + + for (const [label, parse] of cases) { + expect(parse, label).toThrow(); + } + }); + + test("keeps the intentional primitive log fields extension point", () => { + expect( + parseDriverLogBatchInput({ + driverInstanceId: "driver-1", + logs: [{ ...logBatch(0).logs[0], fields: { futureField: true } }], + }).logs[0]?.fields, + ).toEqual({ futureField: true }); + }); + + test("bounds event and log batches at 64 entries", () => { + const event = createRuntimeEvent({ + id: PLATFORM_ID_FIXTURES.runtimeEvent, + kind: "diagnostic.reported", + occurredAt: "2026-08-29T00:00:00.000Z", + payload: { message: "ok" }, + sessionId: PLATFORM_ID_FIXTURES.session, + }); + const parsers = [ + (length: number) => + parseDriverEventBatchInput({ + driverInstanceId: "driver-1", + events: Array.from({ length }, () => ({ event, eventId: "event-1" })), + }), + (length: number) => + parseDriverLogBatchInput({ + driverInstanceId: "driver-1", + logs: Array.from({ length }, () => logBatch(0).logs[0]), + }), + ] as const; + + for (const parse of parsers) { + expect(() => parse(64)).not.toThrow(); + expect(() => parse(65)).toThrow(); + } + }); + + test.each(["accepted", "cancelled", "completed", "failed"] as const)( + "accepts Driver-owned command status %s", + (status) => { + const terminal = + status === "failed" + ? { + error: { + code: "driver.command_failed", + details: {}, + message: "failed", + retryable: false, + }, + } + : status === "completed" + ? { result: { requestId: "request-1" } } + : {}; + + expect( + parseDriverCommandUpdateInput({ + commandId: "command-1", + driverInstanceId: "driver-1", + ...terminal, + status, + }).status, + ).toBe(status); + }, + ); + + test.each([ + ["accepted with result", { result: { requestId: "request-1" }, status: "accepted" }], + [ + "accepted with error", + { + error: { code: "failed", details: {}, message: "failed", retryable: false }, + status: "accepted", + }, + ], + ["cancelled with result", { result: { requestId: "request-1" }, status: "cancelled" }], + [ + "cancelled with error", + { + error: { code: "failed", details: {}, message: "failed", retryable: false }, + status: "cancelled", + }, + ], + ["completed with null result", { result: null, status: "completed" }], + [ + "completed with error", + { + error: { code: "failed", details: {}, message: "failed", retryable: false }, + status: "completed", + }, + ], + ["failed without error", { status: "failed" }], + [ + "failed with result", + { + error: { code: "failed", details: {}, message: "failed", retryable: false }, + result: { requestId: "request-1" }, + status: "failed", + }, + ], + ] as const)("rejects command update payload mismatch: %s", (_label, update) => { + expect(() => + parseDriverCommandUpdateInput({ + commandId: "command-1", + driverInstanceId: "driver-1", + ...update, + }), + ).toThrow(); + }); + + test("rejects duplicate capability ids in hello input and output", () => { + const capability = { id: "text_stream", status: "supported", version: 1 } as const; + + expect(() => + parseDriverHelloInput(helloInput({ capabilities: [capability, capability] })), + ).toThrow(); + expect(() => + parseDriverHelloOutput({ + ...helloOutput(), + acceptedCapabilities: [capability, capability], + }), + ).toThrow(); + }); + + test.each(["queued", "delivered", "expired"] as const)( + "rejects host-owned command status %s at the RPC boundary", + (status) => { + expect(() => + parseDriverCommandUpdateInput({ + commandId: "command-1", + driverInstanceId: "driver-1", + status, + }), + ).toThrow(); + }, + ); + + test("accepts a canonical UUID external-effect claim token", () => { + expect( + parseDriverExternalToolEffectClaimInput({ + claimToken: EXTERNAL_TOOL_EFFECT_CLAIM_TOKEN, + commandId: "command-1", + driverInstanceId: "driver-1", + }), + ).toMatchObject({ claimToken: EXTERNAL_TOOL_EFFECT_CLAIM_TOKEN }); + expect( + parseDriverExternalToolEffectSettleInput({ + claimToken: EXTERNAL_TOOL_EFFECT_CLAIM_TOKEN, + commandId: "command-1", + driverInstanceId: "driver-1", + effectId: "effect-1", + settlement: { kind: "unknown" }, + }), + ).toMatchObject({ claimToken: EXTERNAL_TOOL_EFFECT_CLAIM_TOKEN }); + }); + + test.each([ + ["empty", ""], + ["uppercase UUID", EXTERNAL_TOOL_EFFECT_CLAIM_TOKEN.toUpperCase()], + ["extra-hyphen UUID", "123e4567-e89b-42d3-a456-42661417-00"], + ["non-v4 UUID", "123e4567-e89b-f2d3-a456-426614174000"], + ["invalid-variant UUID", "123e4567-e89b-42d3-c456-426614174000"], + ["non-UUID", "claim-token"], + ["37-character", "x".repeat(37)], + ["one-megabyte", "x".repeat(1_000_000)], + ] as const)("rejects a %s external-effect claim token", (_label, claimToken) => { + expect(() => + parseDriverExternalToolEffectClaimInput({ + claimToken, + commandId: "command-1", + driverInstanceId: "driver-1", + }), + ).toThrow(); + expect(() => + parseDriverExternalToolEffectSettleInput({ + claimToken, + commandId: "command-1", + driverInstanceId: "driver-1", + effectId: "effect-1", + settlement: { kind: "unknown" }, + }), + ).toThrow(); + }); + + test.each([2, 4] as const)("rejects protocol version %d", (protocolVersion) => { + expect(() => parseDriverHelloInput(helloInput({ protocolVersion }))).toThrow(); + }); + + test.each(invalidPositiveSafeIntegers)( + "rejects %p at every positive safe-integer field", + (value) => { + expect(() => parseDriverHelloInput(helloInput({ pid: value }))).toThrow(); + expect(() => parseDriverHeartbeatInput(heartbeatInput(value))).toThrow(); + expect(() => parseDriverReadyInput(readyInput(value))).toThrow(); + expect(() => parseDriverHelloOutput(helloOutput({ eventBatchMaxSize: value }))).toThrow(); + expect(() => + parseSchemaValue(ExternalToolEffectClaim, { + attempt: value, + effectId: "effect-1", + idempotencyKey: "idempotency-1", + kind: "claimed", + }), + ).toThrow(); + }, + ); + + test("bounds the negotiated event batch size to the wire limit", () => { + expect(parseDriverHelloOutput(helloOutput({ eventBatchMaxSize: 64 })).runConfig).toMatchObject({ + eventBatchMaxSize: 64, + }); + expect(() => parseDriverHelloOutput(helloOutput({ eventBatchMaxSize: 65 }))).toThrow(); + }); + + test.each([ + 0, + 249, + 250.5, + Number.NaN, + Number.POSITIVE_INFINITY, + Number.NEGATIVE_INFINITY, + Number.MAX_SAFE_INTEGER + 1, + ] as const)("rejects invalid heartbeat interval %p", (heartbeatIntervalMs) => { + expect(() => parseDriverHelloOutput({ ...helloOutput(), heartbeatIntervalMs })).toThrow(); + }); + + test.each(invalidNonNegativeSafeIntegers)( + "rejects %p at every non-negative safe-integer field", + (value) => { + expect(() => parseDriverHelloOutput(helloOutput({ commandLeaseMs: value }))).toThrow(); + expect(() => parseDriverHeartbeatOutput({ heartbeatCount: value, ok: true })).toThrow(); + expect(() => parseDriverLogBatchInput(logBatch(value))).toThrow(); + expect(() => parseDriverEventBatchOutput(eventBatchOutput(value))).toThrow(); + }, + ); + + test("accepts zero at every non-negative safe-integer field", () => { + expect(parseDriverHelloOutput(helloOutput()).runConfig.commandLeaseMs).toBe(0); + expect(parseDriverHeartbeatOutput({ heartbeatCount: 0, ok: true }).heartbeatCount).toBe(0); + expect(parseDriverLogBatchInput(logBatch(0)).logs[0]?.seq).toBe(0); + expect(parseDriverEventBatchOutput(eventBatchOutput(0)).accepted[0]?.seq).toBe(0); + }); + + test.each([ + [ + "failure message", + () => + parseDriverFailureInput({ + driverInstanceId: "driver-1", + error: { code: "failed", details: {}, message: "", retryable: false }, + runId: "run-1", + }), + ], + ["receipt type", () => parseDriverEventBatchOutput(eventBatchOutput(0, ""))], + ] as const)("rejects an empty %s", (_label, parse) => { + expect(parse).toThrow(); + }); + + test("bounds Driver failure errors with the durable JSON budget", () => { + const base = { code: "failed", details: {}, message: "", retryable: false }; + const exact = { + ...base, + message: "x".repeat(DURABLE_RUN_ERROR_MAX_UTF8_BYTES - measureDurableRunErrorJson(base)), + }; + + expect(measureDurableRunErrorJson(exact)).toBe(DURABLE_RUN_ERROR_MAX_UTF8_BYTES); + expect( + parseDriverFailureInput({ driverInstanceId: "driver-1", error: exact, runId: "run-1" }).error, + ).toEqual(exact); + expect(() => + parseDriverFailureInput({ + driverInstanceId: "driver-1", + error: { ...exact, message: `${exact.message}x` }, + runId: "run-1", + }), + ).toThrow(); + }); + + test("requires an exact run id on terminal RPCs", () => { + expect(() => parseDriverCompletionInput({ driverInstanceId: "driver-1" })).toThrow("runId"); + expect(() => + parseDriverFailureInput({ + driverInstanceId: "driver-1", + error: { code: "failed", details: {}, message: "failed", retryable: false }, + }), + ).toThrow("runId"); + }); + + test.each([ + ["missing", { accepted: [{ seq: 0, type: "message.delta" }] }], + ["empty", { accepted: [{ eventId: "", seq: 0, type: "message.delta" }] }], + ["non-string", { accepted: [{ eventId: 1, seq: 0, type: "message.delta" }] }], + ] as const)("rejects %s receipt eventId", (_label, input) => { + expect(() => parseDriverEventBatchOutput(input)).toThrow(); + }); + + test("rejects empty handshake identity fields", () => { + expect(() => parseDriverHelloInput(helloInput({ startedAt: "" }))).toThrow(); + expect(() => parseDriverHeartbeatInput(heartbeatInput(1, ""))).toThrow(); + expect(() => parseDriverHelloOutput(helloOutput({}, ""))).toThrow(); + }); + + test("rejects an unknown event receipt kind", () => { + expect(() => parseDriverEventBatchOutput(eventBatchOutput(0, "future.event"))).toThrow(); + }); + + test("preserves empty log and tracing strings", () => { + expect( + parseDriverLogBatchInput({ + driverInstanceId: "driver-1", + logs: [ + { + context: { spanId: "", traceId: "" }, + error: { message: "", name: "" }, + level: "error", + message: "", + seq: 0, + timestamp: "now", + }, + ], + }).logs[0], + ).toMatchObject({ + context: { spanId: "", traceId: "" }, + error: { message: "", name: "" }, + message: "", + }); + }); +}); diff --git a/apps/api/tests/runtime-artifact-migration.test.ts b/apps/api/tests/runtime-artifact-migration.test.ts new file mode 100644 index 00000000..e7a1f0d6 --- /dev/null +++ b/apps/api/tests/runtime-artifact-migration.test.ts @@ -0,0 +1,286 @@ +import { describe, expect, test } from "bun:test"; + +import type { FileId, SessionId } from "@mosoo/id"; + +import { fileStore } from "../src/modules/files/application/file-store"; +import { applyDrizzleMigration, applyDrizzleMigrationsBefore } from "./helpers/drizzle-migrations"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const MIGRATION_TAG = "0016_durable-event-side-effects"; + +const ACCOUNT_ID = "01J0000000000000000000001K"; +const AGENT_ID = "01J0000000000000000000001J"; +const APP_ID = "01J0000000000000000000001M"; +const SESSION_ID = "01J0000000000000000000001H" as SessionId; +const ORPHAN_SESSION_ID = "01J0000000000000000000001P" as SessionId; + +const OLDER_FILE_ID = "01J00000000000000000000039" as FileId; +const TIE_LOSER_FILE_ID = "01J00000000000000000000031" as FileId; +const WINNER_FILE_ID = "01J00000000000000000000032" as FileId; +const ORPHAN_FILE_ID = "01J00000000000000000000033" as FileId; +const INVALID_PARENT_FILE_ID = "01J00000000000000000000034" as FileId; +const INVALID_PATH_FILE_ID = "01J00000000000000000000035" as FileId; +const DELETING_FILE_ID = "01J00000000000000000000036" as FileId; +const LATE_REPORT_FILE_ID = "01J00000000000000000000037" as FileId; +const LATE_PATH_NEWER_FILE_ID = "01J0000000000000000000003A" as FileId; +const V3_FILE_ID = "01J0000000000000000000003D" as FileId; +const V3_LATE_LEGACY_FILE_ID = "01J0000000000000000000003E" as FileId; +const MISSING_HEAD_FILE_ID = "01J0000000000000000000003F" as FileId; +const MISSING_HEAD_LATE_FILE_ID = "01J0000000000000000000003G" as FileId; +const SAME_TIME_LOSER_FILE_ID = "01J0000000000000000000003H" as FileId; +const SAME_TIME_WINNER_FILE_ID = "01J0000000000000000000003J" as FileId; + +interface ArtifactHeadRow { + file_id: string | null; + runtime_event_seq: number; + session_id: string; + source_event_id: string; + source_path: string; + updated_at: number; +} + +function legacyRuntimeOutputParentPath(sourcePath: string, contentSha256: string): string { + return `runtime-output/${sourcePath}/${contentSha256}`; +} + +async function createPre0016Database(): Promise { + const database = new SqliteD1Database(); + applyDrizzleMigrationsBefore(database, MIGRATION_TAG); + + return database; +} + +async function insertSession(database: SqliteD1Database): Promise { + await database + .prepare( + `INSERT INTO session ( + agent_id, created_at, creator_account_id, id, kind, model, app_id, + provider, renamed, runtime_id, status, updated_at + ) VALUES (?, 1, ?, ?, 'agent', 'gpt-5.4', ?, 'openai', 0, 'codex', 'IDLE', 1)`, + ) + .bind(AGENT_ID, ACCOUNT_ID, SESSION_ID, APP_ID) + .run(); +} + +async function insertLegacyArtifact( + database: SqliteD1Database, + input: { + createdAt: number; + fileId: FileId; + parentPath: string; + sessionId?: SessionId; + status?: "deleting" | "ready"; + }, +): Promise { + const name = `artifact-${input.fileId}.txt`; + + await database + .prepare( + `INSERT INTO file_record ( + committed, created_at, created_by_account_id, etag, expires_at, id, + mime_type, name, object_key, owner_id, owner_kind, parent_path, path, + purpose, scope_id, scope_kind, session_kind, size, status, updated_at, + version + ) VALUES ( + 1, ?, ?, ?, NULL, ?, 'text/plain', ?, ?, ?, 'session', ?, ?, + 'session_artifact', ?, 'session', 'artifact', ?, ?, ?, 1 + )`, + ) + .bind( + input.createdAt, + ACCOUNT_ID, + `etag-${input.fileId}`, + input.fileId, + name, + `objects/${input.fileId}`, + input.sessionId ?? SESSION_ID, + input.parentPath, + `session-artifacts/${input.fileId}/${name}`, + input.sessionId ?? SESSION_ID, + input.fileId === WINNER_FILE_ID ? 32 : 1, + input.status ?? "ready", + input.createdAt, + ) + .run(); +} + +describe("runtime artifact migration", () => { + test("backfills only the deterministic visible and restorable legacy artifact head", async () => { + const database = await createPre0016Database(); + const preMigrationColumns = await database + .prepare("PRAGMA table_info(file_record)") + .all<{ name: string }>(); + + expect(preMigrationColumns.results.map(({ name }) => name)).not.toContain("runtime_event_seq"); + + await insertSession(database); + await insertLegacyArtifact(database, { + createdAt: 10, + fileId: OLDER_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/report.txt", "a".repeat(64)), + }); + await insertLegacyArtifact(database, { + createdAt: 20, + fileId: TIE_LOSER_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/report.txt", "b".repeat(64)), + }); + await insertLegacyArtifact(database, { + createdAt: 20, + fileId: WINNER_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/report.txt", "c".repeat(64)), + }); + await insertLegacyArtifact(database, { + createdAt: 30, + fileId: ORPHAN_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/orphan.txt", "d".repeat(64)), + sessionId: ORPHAN_SESSION_ID, + }); + await insertLegacyArtifact(database, { + createdAt: 30, + fileId: INVALID_PARENT_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/invalid-parent.txt", "g".repeat(64)), + }); + await insertLegacyArtifact(database, { + createdAt: 30, + fileId: INVALID_PATH_FILE_ID, + parentPath: `runtime-output/outputs/../invalid-path.txt/${"e".repeat(64)}`, + }); + await insertLegacyArtifact(database, { + createdAt: 30, + fileId: DELETING_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/deleting.txt", "f".repeat(64)), + status: "deleting", + }); + + applyDrizzleMigration(database, MIGRATION_TAG); + + const heads = await database + .prepare( + `SELECT file_id, runtime_event_seq, session_id, source_event_id, source_path, updated_at + FROM session_artifact_head + ORDER BY session_id, source_path`, + ) + .all(); + + expect(heads.results).toEqual([ + { + file_id: WINNER_FILE_ID, + runtime_event_seq: 0, + session_id: SESSION_ID, + source_event_id: `legacy-file:${WINNER_FILE_ID}`, + source_path: "outputs/report.txt", + updated_at: 20, + }, + ]); + + const viewerFiles = await fileStore.listReadySessionFiles(database, SESSION_ID); + expect(viewerFiles.map(({ id }) => id)).toEqual([WINNER_FILE_ID]); + + const restoreSources = await fileStore.listLatestReadySessionArtifactSources( + database, + SESSION_ID, + ); + expect(restoreSources).toEqual([ + { + objectKey: `objects/${WINNER_FILE_ID}`, + size: 32, + sourcePath: "outputs/report.txt", + }, + ]); + + // Once the migration creates the head ledger, later headless writes from + // an old Worker are not durable artifact authority. + await insertLegacyArtifact(database, { + createdAt: 40, + fileId: LATE_REPORT_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/report.txt", "1".repeat(64)), + }); + await insertLegacyArtifact(database, { + createdAt: 35, + fileId: LATE_PATH_NEWER_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/late.txt", "3".repeat(64)), + }); + + expect( + (await fileStore.listReadySessionFiles(database, SESSION_ID)).map(({ id }) => id), + ).toEqual([WINNER_FILE_ID]); + expect(await fileStore.listLatestReadySessionArtifactSources(database, SESSION_ID)).toEqual([ + { + objectKey: `objects/${WINNER_FILE_ID}`, + size: 32, + sourcePath: "outputs/report.txt", + }, + ]); + }); + + test("keeps v3 and invalid legacy heads authoritative over late legacy rows", async () => { + const database = await createPre0016Database(); + await insertSession(database); + applyDrizzleMigration(database, MIGRATION_TAG); + + await insertLegacyArtifact(database, { + createdAt: 10, + fileId: V3_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/v3.txt", "6".repeat(64)), + }); + await database + .prepare("UPDATE file_record SET runtime_event_seq = 2 WHERE id = ?") + .bind(V3_FILE_ID) + .run(); + await database + .prepare( + `INSERT INTO session_artifact_head ( + file_id, runtime_event_seq, session_id, source_event_id, source_path, updated_at + ) VALUES (?, 2, ?, 'v3:upsert', 'outputs/v3.txt', 10)`, + ) + .bind(V3_FILE_ID, SESSION_ID) + .run(); + await insertLegacyArtifact(database, { + createdAt: 20, + fileId: V3_LATE_LEGACY_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/v3.txt", "7".repeat(64)), + }); + + await database + .prepare( + `INSERT INTO session_artifact_head ( + file_id, runtime_event_seq, session_id, source_event_id, source_path, updated_at + ) VALUES (?, 0, ?, ?, 'outputs/missing.txt', 10)`, + ) + .bind(MISSING_HEAD_FILE_ID, SESSION_ID, `legacy-file:${MISSING_HEAD_FILE_ID}`) + .run(); + await insertLegacyArtifact(database, { + createdAt: 20, + fileId: MISSING_HEAD_LATE_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/missing.txt", "8".repeat(64)), + }); + + await insertLegacyArtifact(database, { + createdAt: 30, + fileId: SAME_TIME_LOSER_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/tie.txt", "9".repeat(64)), + }); + await insertLegacyArtifact(database, { + createdAt: 30, + fileId: SAME_TIME_WINNER_FILE_ID, + parentPath: legacyRuntimeOutputParentPath("outputs/tie.txt", "a".repeat(64)), + }); + + const viewerFileIds = (await fileStore.listReadySessionFiles(database, SESSION_ID)).map( + ({ id }) => id, + ); + expect(viewerFileIds).toContain(V3_FILE_ID); + expect(viewerFileIds).not.toContain(V3_LATE_LEGACY_FILE_ID); + expect(viewerFileIds).not.toContain(MISSING_HEAD_LATE_FILE_ID); + expect(viewerFileIds).not.toContain(SAME_TIME_WINNER_FILE_ID); + expect(viewerFileIds).not.toContain(SAME_TIME_LOSER_FILE_ID); + + expect(await fileStore.listLatestReadySessionArtifactSources(database, SESSION_ID)).toEqual([ + { + objectKey: `objects/${V3_FILE_ID}`, + size: 1, + sourcePath: "outputs/v3.txt", + }, + ]); + }); +}); diff --git a/apps/api/tests/runtime-command-store.test.ts b/apps/api/tests/runtime-command-store.test.ts index 36be482e..2373b851 100644 --- a/apps/api/tests/runtime-command-store.test.ts +++ b/apps/api/tests/runtime-command-store.test.ts @@ -1,18 +1,20 @@ import { describe, expect, test } from "bun:test"; +import { RUNTIME_COMMAND_MAX_UTF8_BYTES } from "@mosoo/contracts/runtime-command"; import type { RuntimeCommand, RuntimeCommandStatus } from "@mosoo/contracts/runtime-command"; +import type { RunError } from "@mosoo/contracts/session-run"; import { parsePlatformId } from "@mosoo/id"; import type { DriverCommandId, DriverInstanceId, SessionRunId } from "@mosoo/id"; import { - claimNextQueuedRuntimeCommandRecord, - createRuntimeCommandRecord, + claimNextQueuedRuntimeCommandRecord as claimNextQueuedRuntimeCommand, + createRuntimeCommandRecord as persistRuntimeCommandRecord, expireUndeliveredInputStartCommandsForRun, + getRuntimeCommandRecord as readRuntimeCommandRecord, + maintainRuntimeCommandRecords as maintainRuntimeCommands, + markRuntimeCommandRecordDelivered as markRuntimeCommandDelivered, repairRuntimeCommandRecords, - getRuntimeCommandRecord, - markRuntimeCommandRecordDelivered, - maintainRuntimeCommandRecords, - updateRuntimeCommandRecord, + updateRuntimeCommandRecord as updateStoredRuntimeCommandRecord, } from "../src/modules/runtime/infrastructure/session-runs/runtime-command-store.repository"; import { decideRuntimeCommandTransition, @@ -26,6 +28,7 @@ import { SqliteD1Database } from "./helpers/sqlite-d1"; const DRIVER_INSTANCE_ID = parsePlatformId("01J00000000000000000000009"); const TERMINAL_DRIVER_INSTANCE_ID = parsePlatformId("01J0000000000000000000000G"); const SESSION_RUN_ID = parsePlatformId("01J0000000000000000000000N"); +const DRIVER_GENERATION = 0; const COMMAND_IDS = { accepted: parsePlatformId("01J00000000000000000000015"), current: parsePlatformId("01J00000000000000000000017"), @@ -60,10 +63,57 @@ const EXPECTED_PREVIOUS_STATUSES = { completed: ["delivered", "accepted"], delivered: ["queued"], expired: ["queued", "delivered", "accepted"], - failed: ["delivered", "accepted"], + failed: ["queued", "delivered", "accepted"], queued: ["delivered"], } as const satisfies Record; +function createRuntimeCommandRecord( + database: D1Database, + input: Omit[1], "driverGeneration">, +) { + return persistRuntimeCommandRecord(database, { ...input, driverGeneration: DRIVER_GENERATION }); +} + +function getRuntimeCommandRecord( + database: D1Database, + driverInstanceId: DriverInstanceId, + commandId: DriverCommandId, +) { + return readRuntimeCommandRecord(database, driverInstanceId, DRIVER_GENERATION, commandId); +} + +function updateRuntimeCommandRecord( + database: D1Database, + input: Omit[1], "driverGeneration">, +) { + return updateStoredRuntimeCommandRecord(database, { + ...input, + driverGeneration: DRIVER_GENERATION, + }); +} + +function markRuntimeCommandRecordDelivered( + database: D1Database, + input: Omit[1], "driverGeneration">, +) { + return markRuntimeCommandDelivered(database, { ...input, driverGeneration: DRIVER_GENERATION }); +} + +function maintainRuntimeCommandRecords( + database: D1Database, + input: Omit[1], "driverGeneration">, +) { + return maintainRuntimeCommands(database, { ...input, driverGeneration: DRIVER_GENERATION }); +} + +function claimNextQueuedRuntimeCommandRecord( + database: D1Database, + driverInstanceId: DriverInstanceId, + connectionId: string, +) { + return claimNextQueuedRuntimeCommand(database, driverInstanceId, DRIVER_GENERATION, connectionId); +} + function createRuntimeCommandDatabase(): SqliteD1Database { const database = new SqliteD1Database({ foreignKeys: false }); @@ -71,6 +121,7 @@ function createRuntimeCommandDatabase(): SqliteD1Database { CREATE TABLE driver_instance ( command_seq_cursor integer DEFAULT 0 NOT NULL, connection_id text, + generation integer DEFAULT 0 NOT NULL, id text PRIMARY KEY NOT NULL, status text DEFAULT 'ready' NOT NULL ); @@ -79,6 +130,7 @@ function createRuntimeCommandDatabase(): SqliteD1Database { acked_at integer, completed_at integer, delivery_connection_id text, + driver_generation integer, driver_instance_id text NOT NULL, error_json text, expires_at integer, @@ -91,8 +143,60 @@ function createRuntimeCommandDatabase(): SqliteD1Database { status text NOT NULL ); + CREATE TABLE session_run ( + driver_instance_id text, + error_code text, + error_details_json text, + error_message text, + error_retryable integer, + id text PRIMARY KEY NOT NULL, + session_id text NOT NULL DEFAULT 'session-1', + status text NOT NULL + ); + + CREATE TABLE session_event ( + event_type text NOT NULL, + id text PRIMARY KEY NOT NULL, + mcp_command_id text, + run_id text, + source_event_id text NOT NULL, + tool_call_id text, + tool_input_json text, + tool_name text, + tool_output_text text, + tool_status text + ); + + CREATE TABLE external_tool_effect ( + attempt_count integer DEFAULT 0 NOT NULL, + claim_token text, + command_id text NOT NULL, + driver_instance_id text NOT NULL, + id text PRIMARY KEY NOT NULL, + idempotency_key text NOT NULL, + provider_receipt_json text, + result_json text, + status text NOT NULL, + updated_at integer NOT NULL + ); + + CREATE TABLE external_tool_effect_attempt ( + attempt integer NOT NULL, + claim_token text NOT NULL, + completed_at integer, + created_at integer NOT NULL, + effect_id text NOT NULL, + provider_receipt_json text, + result_json text, + status text NOT NULL, + PRIMARY KEY (effect_id, attempt) + ); + INSERT INTO driver_instance (id, connection_id, status) VALUES ('${DRIVER_INSTANCE_ID}', 'connection-1', 'ready'); + + INSERT INTO session_run (driver_instance_id, id, status) + VALUES ('${DRIVER_INSTANCE_ID}', '${SESSION_RUN_ID}', 'running'); `); return database; @@ -110,7 +214,510 @@ function inputStartCommand(id: DriverCommandId): RuntimeCommand { }; } +function sessionStopCommand(id: DriverCommandId): RuntimeCommand { + return { + commandId: id, + kind: "session.stop", + reason: "terminal Driver repair", + }; +} + +function mcpExecuteCommand(id: DriverCommandId): Extract { + return { + argumentsJson: '{"title":"durable"}', + commandId: id, + kind: "mcp.execute", + requestId: "request-1", + runId: SESSION_RUN_ID, + serverId: "01J0000000000000000000000Y", + toolCallId: "tool-1", + toolName: "createIssue", + }; +} + +function insertInputTerminalFact( + database: SqliteD1Database, + input: + | { error: RunError; status: "failed" } + | { error?: RunError; status: "cancelled" | "expired" } + | { status: "completed" }, +): void { + const error = "error" in input ? (input.error ?? null) : null; + const eventType = + input.status === "failed" + ? "run.failed" + : input.status === "completed" + ? "run.completed" + : "run.cancelled"; + database + .prepare( + `UPDATE session_run + SET error_code = ?, error_details_json = ?, error_message = ?, error_retryable = ?, status = ? + WHERE id = ?`, + ) + .bind( + error?.code ?? null, + error === null ? null : JSON.stringify(error.details), + error?.message ?? null, + error === null ? null : Number(error.retryable), + input.status, + SESSION_RUN_ID, + ) + .run(); + database + .prepare( + `INSERT INTO session_event (event_type, id, run_id, source_event_id) + VALUES (?, ?, ?, ?)`, + ) + .bind( + eventType, + `event-input-${input.status}`, + SESSION_RUN_ID, + `session-run-terminal:${SESSION_RUN_ID}:${eventType}`, + ) + .run(); +} + +function insertMcpTerminalFact( + database: SqliteD1Database, + command: Extract, + input: + | { outputText: string; status: "completed" } + | { error: RunError; status: "failed" } + | { status: "cancelled" }, +): void { + const outputText = + input.status === "completed" + ? input.outputText + : input.status === "failed" + ? input.error.message + : null; + database + .prepare( + `INSERT INTO session_event ( + event_type, id, mcp_command_id, run_id, source_event_id, tool_call_id, + tool_input_json, tool_name, tool_output_text, tool_status + ) VALUES ('tool.call.updated', ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + `event-mcp-${input.status}`, + command.commandId, + command.runId, + `mcp.execute.${input.status}:${command.commandId}`, + command.toolCallId, + command.argumentsJson, + command.toolName, + outputText, + input.status, + ) + .run(); +} + describe("runtime command store", () => { + test("measures and writes the same serialized command and terminal payload", async () => { + const database = createRuntimeCommandDatabase(); + let commandSerializations = 0; + const command = { + ...inputStartCommand(COMMAND_IDS.first), + toJSON() { + commandSerializations += 1; + return { + ...inputStartCommand(COMMAND_IDS.first), + input: { text: `serialized-${commandSerializations}` }, + }; + }, + }; + + await createRuntimeCommandRecord(database, { + command, + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + expect(commandSerializations).toBe(1); + await expect( + database + .prepare("SELECT payload_json FROM driver_command WHERE id = ?") + .bind(COMMAND_IDS.first) + .first(), + ).resolves.toEqual({ + payload_json: JSON.stringify({ + ...inputStartCommand(COMMAND_IDS.first), + input: { text: "serialized-1" }, + }), + }); + + let errorSerializations = 0; + const error = { + code: "test.serialized", + details: {}, + message: "original", + retryable: false, + toJSON() { + errorSerializations += 1; + return { + code: "test.serialized", + details: {}, + message: `serialized-${errorSerializations}`, + retryable: false, + }; + }, + }; + insertInputTerminalFact(database, { + error: { + code: "test.serialized", + details: {}, + message: "serialized-1", + retryable: false, + }, + status: "failed", + }); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + error, + status: "failed", + }), + ).resolves.toMatchObject({ kind: "applied" }); + expect(errorSerializations).toBe(1); + await expect( + database + .prepare("SELECT error_json FROM driver_command WHERE id = ?") + .bind(COMMAND_IDS.first) + .first(), + ).resolves.toEqual({ + error_json: JSON.stringify({ + code: "test.serialized", + details: {}, + message: "serialized-1", + retryable: false, + }), + }); + + const resultDatabase = createRuntimeCommandDatabase(); + await createRuntimeCommandRecord(resultDatabase, { + command: { ...inputStartCommand(COMMAND_IDS.second), requestId: "serialized-1" }, + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + insertInputTerminalFact(resultDatabase, { status: "completed" }); + let resultSerializations = 0; + const result = { + requestId: "original", + toJSON() { + resultSerializations += 1; + return { requestId: `serialized-${resultSerializations}` }; + }, + }; + await expect( + updateRuntimeCommandRecord(resultDatabase, { + commandId: COMMAND_IDS.second, + driverInstanceId: DRIVER_INSTANCE_ID, + result, + status: "completed", + }), + ).resolves.toMatchObject({ kind: "applied" }); + expect(resultSerializations).toBe(1); + await expect( + resultDatabase + .prepare("SELECT result_json FROM driver_command WHERE id = ?") + .bind(COMMAND_IDS.second) + .first(), + ).resolves.toEqual({ result_json: '{"requestId":"serialized-1"}' }); + }); + + test("rejects invalid command JSON before allocating a sequence or writing a row", async () => { + const database = createRuntimeCommandDatabase(); + const emptyCommand = { + ...inputStartCommand(COMMAND_IDS.first), + input: { text: "" }, + }; + const emptyPayloadBytes = new TextEncoder().encode(JSON.stringify(emptyCommand)).byteLength; + const oversizedCommand = { + ...emptyCommand, + input: { + text: "x".repeat(RUNTIME_COMMAND_MAX_UTF8_BYTES - emptyPayloadBytes + 1), + }, + }; + + await expect( + createRuntimeCommandRecord(database, { + command: { ...emptyCommand, input: { text: "" } }, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).rejects.toThrow(); + await expect( + createRuntimeCommandRecord(database, { + command: oversizedCommand, + driverInstanceId: DRIVER_INSTANCE_ID, + }), + ).rejects.toThrow(`${RUNTIME_COMMAND_MAX_UTF8_BYTES} UTF-8 bytes`); + await expect( + createRuntimeCommandRecord(database, { + command: inputStartCommand(COMMAND_IDS.second), + driverInstanceId: DRIVER_INSTANCE_ID, + status: "invalid" as never, + }), + ).rejects.toThrow(); + await expect( + database + .prepare( + "SELECT command_seq_cursor, (SELECT count(*) FROM driver_command) AS command_count FROM driver_instance WHERE id = ?", + ) + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ command_count: 0, command_seq_cursor: 0 }); + }); + + test("rejects terminal payloads that do not match the command kind without mutating it", async () => { + const database = createRuntimeCommandDatabase(); + await createRuntimeCommandRecord(database, { + command: inputStartCommand(COMMAND_IDS.first), + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + result: { + outputText: "wrong result kind", + requestId: "request-1", + serverId: "server-1", + toolName: "tool-1", + }, + status: "completed", + }), + ).rejects.toThrow(); + const validError = { + code: "test.invalid_status_payload", + details: {}, + message: "must not be stored", + retryable: false, + }; + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + error: validError, + status: "completed", + }), + ).rejects.toThrow(); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + result: { requestId: "request-1" }, + status: "failed", + }), + ).rejects.toThrow(); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + error: validError, + status: "delivered", + }), + ).rejects.toThrow(); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + status: "invalid" as never, + }), + ).rejects.toThrow(); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + error: { + code: "test.invalid", + details: { nested: { invalid: true } } as never, + message: "invalid nested details", + retryable: false, + }, + status: "failed", + }), + ).rejects.toThrow(); + await expect( + database + .prepare("SELECT error_json, result_json, status FROM driver_command WHERE id = ?") + .bind(COMMAND_IDS.first) + .first(), + ).resolves.toEqual({ error_json: null, result_json: null, status: "accepted" }); + }); + + test("mirrors the succeeded effect result bytes when completing an MCP command", async () => { + const database = createRuntimeCommandDatabase(); + const command = mcpExecuteCommand(COMMAND_IDS.first); + const authoritativeResult = + '{"toolName":"tool-1","outputText":"done","serverId":"server-1","requestId":"request-1"}'; + database.execute(` + INSERT INTO driver_command ( + driver_generation, driver_instance_id, id, issued_at, kind, payload_json, seq, status + ) VALUES ( + ${DRIVER_GENERATION}, '${DRIVER_INSTANCE_ID}', '${COMMAND_IDS.first}', 1, 'mcp.execute', '${JSON.stringify(command)}', 1, 'accepted' + ); + INSERT INTO external_tool_effect ( + attempt_count, command_id, driver_instance_id, id, idempotency_key, + result_json, status, updated_at + ) VALUES ( + 1, '${COMMAND_IDS.first}', '${DRIVER_INSTANCE_ID}', 'effect-1', + 'idempotency-1', '${authoritativeResult}', 'succeeded', 1 + ); + `); + insertMcpTerminalFact(database, command, { outputText: "done", status: "completed" }); + + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + result: { + isError: false, + outputText: "done", + requestId: "request-1", + serverId: "server-1", + toolName: "tool-1", + }, + status: "completed", + }), + ).resolves.toMatchObject({ kind: "applied" }); + await expect( + database + .prepare("SELECT result_json FROM driver_command WHERE id = ?") + .bind(COMMAND_IDS.first) + .first(), + ).resolves.toEqual({ result_json: authoritativeResult }); + + const exactReplay = { + isError: false, + outputText: "done", + requestId: "request-1", + serverId: "server-1", + toolName: "tool-1", + }; + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + result: exactReplay, + status: "completed", + }), + ).resolves.toEqual({ kind: "duplicate", status: "completed" }); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + result: { ...exactReplay, outputText: "changed" }, + status: "completed", + }), + ).rejects.toThrow("conflicts with its durable terminal payload"); + + database.execute(` + UPDATE external_tool_effect + SET result_json = NULL, status = 'unknown' + WHERE command_id = '${COMMAND_IDS.first}' + `); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + result: exactReplay, + status: "completed", + }), + ).rejects.toThrow("has no succeeded durable effect"); + }); + + test("accepts only exact failed command replays", async () => { + const database = createRuntimeCommandDatabase(); + const error = { + code: "test.command_failed", + details: { attempt: 1, source: "driver" }, + message: "Command failed.", + retryable: true, + }; + await createRuntimeCommandRecord(database, { + command: inputStartCommand(COMMAND_IDS.first), + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + insertInputTerminalFact(database, { error, status: "failed" }); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + error, + status: "failed", + }), + ).resolves.toMatchObject({ kind: "applied" }); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + error: { + message: error.message, + retryable: error.retryable, + details: { source: "driver", attempt: 1 }, + code: error.code, + }, + status: "failed", + }), + ).resolves.toEqual({ kind: "duplicate", status: "failed" }); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + error: { ...error, retryable: false }, + status: "failed", + }), + ).rejects.toThrow("conflicts with its durable terminal payload"); + }); + + test("settles cancelled input commands while preserving the Run cancellation error", async () => { + const database = createRuntimeCommandDatabase(); + const error = { + code: "runtime.operation_cancelled", + details: { operationId: "operation-1" }, + message: "The runtime operation cancelled this run.", + retryable: false, + }; + await createRuntimeCommandRecord(database, { + command: inputStartCommand(COMMAND_IDS.first), + driverInstanceId: DRIVER_INSTANCE_ID, + status: "accepted", + }); + insertInputTerminalFact(database, { error, status: "cancelled" }); + + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + status: "cancelled", + }), + ).resolves.toEqual({ kind: "applied", status: "cancelled" }); + await expect( + updateRuntimeCommandRecord(database, { + commandId: COMMAND_IDS.first, + driverInstanceId: DRIVER_INSTANCE_ID, + status: "cancelled", + }), + ).resolves.toEqual({ kind: "duplicate", status: "cancelled" }); + await expect( + database + .prepare( + "SELECT error_code, error_message, error_retryable, status FROM session_run WHERE id = ?", + ) + .bind(SESSION_RUN_ID) + .first(), + ).resolves.toEqual({ + error_code: error.code, + error_message: error.message, + error_retryable: 0, + status: "cancelled", + }); + }); + test("keeps runtime command transitions on the owner matrix", () => { for (const targetStatus of RUNTIME_COMMAND_STATUSES) { expect(getRuntimeCommandPreviousStatuses(targetStatus)).toEqual( @@ -410,13 +1017,13 @@ describe("runtime command store", () => { expect(recovered?.status).toBe("queued"); }); - test("repairs queued, delivered, and accepted commands globally", async () => { + test("recovers and expires commands globally without bypassing terminal orchestration", async () => { const database = createRuntimeCommandDatabase(); const nowMs = Date.now(); database.execute(` INSERT INTO driver_instance (id, connection_id, status) - VALUES ('${TERMINAL_DRIVER_INSTANCE_ID}', 'terminal-connection', 'stopped') + VALUES ('${TERMINAL_DRIVER_INSTANCE_ID}', 'terminal-connection', 'ready') `); await createRuntimeCommandRecord(database, { command: inputStartCommand(COMMAND_IDS.globalExpired), @@ -429,7 +1036,7 @@ describe("runtime command store", () => { expiresAt: nowMs + 60_000, }); await createRuntimeCommandRecord(database, { - command: inputStartCommand(COMMAND_IDS.globalAccepted), + command: sessionStopCommand(COMMAND_IDS.globalAccepted), driverInstanceId: TERMINAL_DRIVER_INSTANCE_ID, expiresAt: nowMs + 60_000, }); @@ -453,19 +1060,19 @@ describe("runtime command store", () => { UPDATE driver_instance SET connection_id = 'connection-2' WHERE id = '${DRIVER_INSTANCE_ID}' + ; + UPDATE driver_instance + SET status = 'stopped' + WHERE id = '${TERMINAL_DRIVER_INSTANCE_ID}' `); - await expect(repairRuntimeCommandRecords(database, { nowMs })).resolves.toEqual({ + const repair = await repairRuntimeCommandRecords(database, { nowMs }); + expect(repair).toEqual({ expired: { appliedCount: 1, kind: "batch_applied", status: "expired", }, - failed: { - appliedCount: 1, - kind: "batch_applied", - status: "failed", - }, recovered: { appliedCount: 1, kind: "batch_applied", @@ -473,6 +1080,10 @@ describe("runtime command store", () => { }, }); + await expect( + getRuntimeCommandRecord(database, TERMINAL_DRIVER_INSTANCE_ID, COMMAND_IDS.globalAccepted), + ).resolves.toMatchObject({ status: "accepted" }); + const expired = await getRuntimeCommandRecord( database, DRIVER_INSTANCE_ID, @@ -483,7 +1094,7 @@ describe("runtime command store", () => { DRIVER_INSTANCE_ID, COMMAND_IDS.globalStale, ); - const failed = await getRuntimeCommandRecord( + const pendingTerminalRepair = await getRuntimeCommandRecord( database, TERMINAL_DRIVER_INSTANCE_ID, COMMAND_IDS.globalAccepted, @@ -491,13 +1102,17 @@ describe("runtime command store", () => { expect(expired?.status).toBe("expired"); expect(recovered?.status).toBe("queued"); - expect(failed?.status).toBe("failed"); - expect(failed?.error?.code).toBe("driver.command_driver_terminal"); + expect(pendingTerminalRepair?.status).toBe("accepted"); + expect(pendingTerminalRepair?.error).toBeNull(); }); test("expires undelivered input commands for one run", async () => { const database = createRuntimeCommandDatabase(); const otherRunId = parsePlatformId("01J0000000000000000000000P"); + database.execute(` + INSERT INTO session_run (driver_instance_id, id, status) + VALUES ('${DRIVER_INSTANCE_ID}', '${otherRunId}', 'running') + `); await createRuntimeCommandRecord(database, { command: inputStartCommand(COMMAND_IDS.currentRunQueued), @@ -635,12 +1250,14 @@ describe("runtime command store", () => { expiresAt: Date.now() + 60_000, }); await claimNextQueuedRuntimeCommandRecord(database, DRIVER_INSTANCE_ID, "connection-1"); + insertInputTerminalFact(database, { status: "completed" }); await expect( updateRuntimeCommandRecord(database, { commandId: COMMAND_IDS.illegal, deliveryConnectionId: "connection-1", driverInstanceId: DRIVER_INSTANCE_ID, + result: { requestId: inputStartCommand(COMMAND_IDS.illegal).requestId }, status: "completed", }), ).resolves.toEqual({ diff --git a/apps/api/tests/runtime-conversation-idle-sweep.test.ts b/apps/api/tests/runtime-conversation-idle-sweep.test.ts index 50ca5a29..5a4638e6 100644 --- a/apps/api/tests/runtime-conversation-idle-sweep.test.ts +++ b/apps/api/tests/runtime-conversation-idle-sweep.test.ts @@ -15,10 +15,12 @@ function createDatabase(): SqliteD1Database { database.execute(` CREATE TABLE sandbox_session ( cloudflare_session_id text NOT NULL, + cleanup_operation_id text, created_at integer NOT NULL, cwd text NOT NULL, origin_json text NOT NULL, sandbox_id text NOT NULL, + sandbox_incarnation integer DEFAULT 0 NOT NULL, session_id text PRIMARY KEY NOT NULL, status text NOT NULL, updated_at integer NOT NULL @@ -37,6 +39,7 @@ function createDatabase(): SqliteD1Database { CREATE TABLE session_run ( id text PRIMARY KEY NOT NULL, driver_instance_id text, + session_id text NOT NULL, status text NOT NULL ); @@ -44,6 +47,7 @@ function createDatabase(): SqliteD1Database { id text PRIMARY KEY NOT NULL, kind text NOT NULL, last_run_id text, + runtime_provisioning_operation_id text, workspace_checkpoint_required integer DEFAULT 0 NOT NULL ); @@ -51,8 +55,10 @@ function createDatabase(): SqliteD1Database { dir text NOT NULL, id text PRIMARY KEY NOT NULL, sandbox_id text NOT NULL, + sandbox_incarnation integer NOT NULL, session_run_id text, - status text NOT NULL + status text NOT NULL, + workspace_session_id text ); `); @@ -86,8 +92,9 @@ async function insertConversation( await database .prepare( `INSERT INTO sandbox_session ( - cloudflare_session_id, created_at, cwd, origin_json, sandbox_id, session_id, status, updated_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + cloudflare_session_id, created_at, cwd, origin_json, sandbox_id, sandbox_incarnation, + session_id, status, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`, ) .bind( `cf-${input.sessionId}`, @@ -95,6 +102,7 @@ async function insertConversation( "cwd", "{}", input.sandboxId, + 1, input.sessionId, input.status, input.updatedAt, @@ -104,15 +112,17 @@ async function insertConversation( async function insertActiveRunLease( database: D1Database, - input: { readonly runId: string; readonly sandboxId: string }, + input: { readonly runId: string; readonly sandboxId: string; readonly sessionId: string }, ): Promise { await database .prepare("INSERT INTO driver_instance (id, sandbox_id) VALUES (?, ?)") .bind(`driver-${input.runId}`, input.sandboxId) .run(); await database - .prepare("INSERT INTO session_run (id, driver_instance_id, status) VALUES (?, ?, ?)") - .bind(input.runId, `driver-${input.runId}`, "running") + .prepare( + "INSERT INTO session_run (id, driver_instance_id, session_id, status) VALUES (?, ?, ?, ?)", + ) + .bind(input.runId, `driver-${input.runId}`, input.sessionId, "running") .run(); } @@ -127,8 +137,10 @@ describe("idle session-scoped conversation sweep", () => { updatedAt: NOW - GRACE_MS - 1, }); await database - .prepare("INSERT INTO session_run (id, driver_instance_id, status) VALUES (?, NULL, ?)") - .bind("run-checkpoint", "completed") + .prepare( + "INSERT INTO session_run (id, driver_instance_id, session_id, status) VALUES (?, NULL, ?, ?)", + ) + .bind("run-checkpoint", "session-checkpoint", "completed") .run(); await database .prepare("UPDATE session SET last_run_id = ?, workspace_checkpoint_required = 1 WHERE id = ?") @@ -146,16 +158,27 @@ describe("idle session-scoped conversation sweep", () => { idleSinceLte: NOW - GRACE_MS, now: NOW, runtimeSubjectId: "sb-checkpoint" as never, + sandboxIncarnation: 1, sandboxSessionId: "cf-session-checkpoint" as never, sessionId: "session-checkpoint" as never, }), - ).resolves.toBe(false); + ).resolves.toBeNull(); await database .prepare( - "INSERT INTO sandbox_backup (dir, id, sandbox_id, session_run_id, status) VALUES (?, ?, ?, ?, ?)", + `INSERT INTO sandbox_backup ( + dir, id, sandbox_id, sandbox_incarnation, session_run_id, status, workspace_session_id + ) VALUES (?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + "cwd", + "backup-checkpoint", + "sb-checkpoint", + 1, + "run-checkpoint", + "ready", + "session-checkpoint", ) - .bind("cwd", "backup-checkpoint", "sb-checkpoint", "run-checkpoint", "ready") .run(); await expect( @@ -176,8 +199,10 @@ describe("idle session-scoped conversation sweep", () => { updatedAt: NOW - GRACE_MS - 1, }); await database - .prepare("INSERT INTO session_run (id, driver_instance_id, status) VALUES (?, NULL, ?)") - .bind("run-legacy", "completed") + .prepare( + "INSERT INTO session_run (id, driver_instance_id, session_id, status) VALUES (?, NULL, ?, ?)", + ) + .bind("run-legacy", "session-legacy", "completed") .run(); await database .prepare("UPDATE session SET last_run_id = ? WHERE id = ?") @@ -195,10 +220,11 @@ describe("idle session-scoped conversation sweep", () => { idleSinceLte: NOW - GRACE_MS, now: NOW, runtimeSubjectId: "sb-legacy" as never, + sandboxIncarnation: 1, sandboxSessionId: "cf-session-legacy" as never, sessionId: "session-legacy" as never, }), - ).resolves.toBe(true); + ).resolves.not.toBeNull(); }); test("lists only idle active cattle conversations without a run lease", async () => { @@ -238,7 +264,11 @@ describe("idle session-scoped conversation sweep", () => { status: "active", updatedAt: NOW - GRACE_MS - 1, }); - await insertActiveRunLease(database, { runId: "run-busy", sandboxId: "sb-busy" }); + await insertActiveRunLease(database, { + runId: "run-busy", + sandboxId: "sb-busy", + sessionId: "session-busy", + }); const idle = await listIdleSessionScopedConversationSessions(database, { idleSinceLte: NOW - GRACE_MS, @@ -248,7 +278,7 @@ describe("idle session-scoped conversation sweep", () => { expect(idle).toEqual([{ sandboxId: "sb-idle", sessionId: "session-idle" }]); }); - test("atomic claim closes an idle conversation but loses to any re-activation", async () => { + test("atomic claim schedules an idle conversation cleanup but loses to any re-activation", async () => { const database = createDatabase(); const idleSinceLte = NOW - GRACE_MS; const claim = (sandboxId: string, sessionId: string, sandboxSessionId: string) => @@ -256,6 +286,7 @@ describe("idle session-scoped conversation sweep", () => { idleSinceLte, now: NOW, runtimeSubjectId: sandboxId as never, + sandboxIncarnation: 1, sandboxSessionId: sandboxSessionId as never, sessionId: sessionId as never, }); @@ -267,7 +298,7 @@ describe("idle session-scoped conversation sweep", () => { .first<{ status: string }>() )?.status; - // (1) still-idle, same session instance, no lease → claim wins, row closed. + // (1) still-idle, same session instance, no lease → claim wins. await insertConversation(database, { kind: "cattle", sandboxId: "sb-a", @@ -275,8 +306,8 @@ describe("idle session-scoped conversation sweep", () => { status: "active", updatedAt: NOW - GRACE_MS - 1, }); - expect(await claim("sb-a", "sess-a", "cf-sess-a")).toBe(true); - expect(await statusOf("sess-a")).toBe("closed"); + expect(await claim("sb-a", "sess-a", "cf-sess-a")).not.toBeNull(); + expect(await statusOf("sess-a")).toBe("cleanup_pending"); // (2) a follow-up refreshed updated_at past the grace → claim loses, untouched. await insertConversation(database, { @@ -286,7 +317,7 @@ describe("idle session-scoped conversation sweep", () => { status: "active", updatedAt: NOW, // refreshed by ensureSandboxConversationSession }); - expect(await claim("sb-b", "sess-b", "cf-sess-b")).toBe(false); + expect(await claim("sb-b", "sess-b", "cf-sess-b")).toBeNull(); expect(await statusOf("sess-b")).toBe("active"); // (3) the session was rebuilt (new cloudflare_session_id) → claim loses. @@ -297,7 +328,7 @@ describe("idle session-scoped conversation sweep", () => { status: "active", updatedAt: NOW - GRACE_MS - 1, }); - expect(await claim("sb-c", "sess-c", "cf-STALE")).toBe(false); + expect(await claim("sb-c", "sess-c", "cf-STALE")).toBeNull(); expect(await statusOf("sess-c")).toBe("active"); // (4) an active run lease appeared → claim loses. @@ -308,8 +339,12 @@ describe("idle session-scoped conversation sweep", () => { status: "active", updatedAt: NOW - GRACE_MS - 1, }); - await insertActiveRunLease(database, { runId: "run-d", sandboxId: "sb-d" }); - expect(await claim("sb-d", "sess-d", "cf-sess-d")).toBe(false); + await insertActiveRunLease(database, { + runId: "run-d", + sandboxId: "sb-d", + sessionId: "sess-d", + }); + expect(await claim("sb-d", "sess-d", "cf-sess-d")).toBeNull(); expect(await statusOf("sess-d")).toBe("active"); }); }); diff --git a/apps/api/tests/runtime-conversation-session-record.test.ts b/apps/api/tests/runtime-conversation-session-record.test.ts index 505e278e..2ae3516a 100644 --- a/apps/api/tests/runtime-conversation-session-record.test.ts +++ b/apps/api/tests/runtime-conversation-session-record.test.ts @@ -15,6 +15,7 @@ function createConversationSessionDatabase(): SqliteD1Database { cwd text NOT NULL, origin_json text NOT NULL, sandbox_id text NOT NULL, + sandbox_incarnation integer DEFAULT 0 NOT NULL, session_id text PRIMARY KEY NOT NULL, status text NOT NULL ); @@ -25,7 +26,8 @@ function createConversationSessionDatabase(): SqliteD1Database { id text PRIMARY KEY NOT NULL, sandbox_id text NOT NULL, session_run_id text, - status text NOT NULL + status text NOT NULL, + workspace_session_id text ); CREATE TABLE session ( @@ -71,13 +73,21 @@ async function insertBackup( readonly dir: string; readonly id: string; readonly status: string; + readonly workspaceSessionId?: string | null; }, ): Promise { await database .prepare( - "INSERT INTO sandbox_backup (created_at, dir, id, sandbox_id, status) VALUES (?, ?, ?, ?, ?)", + "INSERT INTO sandbox_backup (created_at, dir, id, sandbox_id, status, workspace_session_id) VALUES (?, ?, ?, ?, ?, ?)", + ) + .bind( + input.createdAt, + input.dir, + input.id, + "01J0000000000000000000000D", + input.status, + input.workspaceSessionId === undefined ? "session-1" : input.workspaceSessionId, ) - .bind(input.createdAt, input.dir, input.id, "01J0000000000000000000000D", input.status) .run(); } @@ -105,6 +115,20 @@ describe("runtime conversation session record", () => { }); await insertBackup(database, { createdAt: 4, + dir: SESSION_CWD, + id: "backup-other-owner", + status: "ready", + workspaceSessionId: "session-2", + }); + await insertBackup(database, { + createdAt: 5, + dir: SESSION_CWD, + id: "backup-legacy-unowned", + status: "ready", + workspaceSessionId: null, + }); + await insertBackup(database, { + createdAt: 6, dir: OTHER_CWD, id: "backup-other", status: "ready", diff --git a/apps/api/tests/runtime-event-persistence-compactor.test.ts b/apps/api/tests/runtime-event-persistence-compactor.test.ts deleted file mode 100644 index 9c84340e..00000000 --- a/apps/api/tests/runtime-event-persistence-compactor.test.ts +++ /dev/null @@ -1,551 +0,0 @@ -import { describe, expect, test } from "bun:test"; - -import type { DriverEventEnvelope } from "@mosoo/agent-driver/events"; -import { createRuntimeEvent } from "@mosoo/runtime-events"; -import type { RuntimeEventEnvelope, RuntimeEventKind } from "@mosoo/runtime-events"; - -import type { ProjectedRuntimeEventRecord } from "../src/modules/runtime/infrastructure/driver-instance/event-types"; -import { RuntimeEventPersistenceCompactor } from "../src/modules/runtime/infrastructure/driver-instance/runtime-event-persistence-compactor"; -import { filterDurablyAcceptedRuntimeStreamReplays } from "../src/modules/runtime/infrastructure/driver-instance/runtime-event-replay-filter"; - -function runtimeEvent(input: { - delivery?: RuntimeEventEnvelope["delivery"]; - id: string; - kind: RuntimeEventKind; - occurredAtMs: number; - payload: unknown; - runId?: string; -}): RuntimeEventEnvelope { - return createRuntimeEvent({ - id: input.id, - kind: input.kind, - occurredAt: new Date(input.occurredAtMs).toISOString(), - payload: input.payload, - ...(input.delivery === undefined ? {} : { delivery: input.delivery }), - ...(input.runId === undefined ? {} : { runId: input.runId }), - sessionId: "session-1", - }); -} - -function record(input: { - delivery?: RuntimeEventEnvelope["delivery"]; - id: string; - kind: RuntimeEventKind; - occurredAtMs: number; - payload: unknown; - runId?: string; -}): ProjectedRuntimeEventRecord { - const event = runtimeEvent(input); - - return { - event, - occurredAt: input.occurredAtMs, - sourceEventId: `source-${input.id}`, - }; -} - -function envelope(projectedRecord: ProjectedRuntimeEventRecord): DriverEventEnvelope { - return { - event: projectedRecord.event, - eventId: projectedRecord.sourceEventId ?? projectedRecord.event.id, - occurredAt: projectedRecord.occurredAt, - }; -} - -describe("runtime event persistence compactor", () => { - test("stores message and thinking streams as complete content events", () => { - const compactor = new RuntimeEventPersistenceCompactor(); - - expect( - compactor.compact([ - record({ - id: "message-start", - kind: "message.started", - occurredAtMs: 1_000, - payload: { messageId: "message-1", role: "agent" }, - runId: "run-1", - }), - record({ - id: "message-delta-1", - kind: "message.delta", - occurredAtMs: 1_010, - payload: { contentDelta: "Hello ", messageId: "message-1", role: "agent" }, - runId: "run-1", - }), - ]), - ).toEqual([]); - - const compacted = compactor.compact([ - record({ - id: "message-delta-2", - kind: "message.delta", - occurredAtMs: 1_020, - payload: { contentDelta: "world", messageId: "message-1", role: "agent" }, - runId: "run-1", - }), - record({ - id: "message-end", - kind: "message.completed", - occurredAtMs: 1_030, - payload: { messageId: "message-1", role: "agent" }, - runId: "run-1", - }), - record({ - id: "thought-start", - kind: "thought.started", - occurredAtMs: 1_040, - payload: { channel: "summary", thoughtId: "thought-1" }, - runId: "run-1", - }), - record({ - id: "thought-delta", - kind: "thought.delta", - occurredAtMs: 1_050, - payload: { channel: "summary", contentDelta: "Check inputs", thoughtId: "thought-1" }, - runId: "run-1", - }), - record({ - id: "thought-end", - kind: "thought.completed", - occurredAtMs: 1_060, - payload: { channel: "summary", thoughtId: "thought-1" }, - runId: "run-1", - }), - ]); - - expect(compacted.map((entry) => entry.event.kind)).toEqual([ - "message.added", - "thought.completed", - ]); - expect(compacted[0]?.occurredAt).toBe(1_000); - expect(compacted[0]?.sourceEventId).toBe("source-message-end"); - expect(compacted[0]?.event.payload).toMatchObject({ - content: "Hello world", - messageId: "message-1", - role: "agent", - }); - expect(compacted[1]?.event.payload).toMatchObject({ - channel: "summary", - content: "Check inputs", - thoughtId: "thought-1", - }); - }); - - test("appends repeated text deltas instead of treating them as snapshots", () => { - const compactor = new RuntimeEventPersistenceCompactor(); - - const compacted = compactor.compact([ - record({ - id: "message-delta-1", - kind: "message.delta", - occurredAtMs: 1_000, - payload: { contentDelta: "a", messageId: "message-1", role: "agent" }, - runId: "run-1", - }), - record({ - id: "message-delta-2", - kind: "message.delta", - occurredAtMs: 1_010, - payload: { contentDelta: "a", messageId: "message-1", role: "agent" }, - runId: "run-1", - }), - record({ - id: "message-end", - kind: "message.completed", - occurredAtMs: 1_020, - payload: { messageId: "message-1", role: "agent" }, - runId: "run-1", - }), - ]); - - expect(compacted[0]?.event.payload).toMatchObject({ - content: "aa", - messageId: "message-1", - }); - }); - - test("preserves structured message.added text blocks", () => { - const compactor = new RuntimeEventPersistenceCompactor(); - - const compacted = compactor.compact([ - record({ - id: "user-message", - kind: "message.added", - occurredAtMs: 1_000, - payload: { - content: [ - { text: "hello ", type: "text" }, - { text: "world", type: "text" }, - ], - messageId: "message-1", - role: "user", - }, - runId: "run-1", - }), - ]); - - expect(compacted.map((entry) => entry.event.kind)).toEqual(["message.added"]); - expect(compacted[0]?.event.payload).toMatchObject({ - content: "hello world", - messageId: "message-1", - role: "user", - }); - }); - - test("preserves an existing message role when terminal events omit it", () => { - const compactor = new RuntimeEventPersistenceCompactor(); - - const compacted = compactor.compact([ - record({ - id: "message-start", - kind: "message.started", - occurredAtMs: 1_000, - payload: { messageId: "message-1", role: "user" }, - runId: "run-1", - }), - record({ - id: "message-delta", - kind: "message.delta", - occurredAtMs: 1_010, - payload: { contentDelta: "hello", messageId: "message-1" }, - runId: "run-1", - }), - record({ - id: "message-end", - kind: "message.completed", - occurredAtMs: 1_020, - payload: { messageId: "message-1" }, - runId: "run-1", - }), - ]); - - expect(compacted[0]?.event.payload).toMatchObject({ - content: "hello", - messageId: "message-1", - role: "user", - }); - }); - - test("stores tool call snapshots as one completed call", () => { - const compactor = new RuntimeEventPersistenceCompactor(); - - expect( - compactor.compact([ - record({ - id: "tool-start", - kind: "item.started", - occurredAtMs: 2_000, - payload: { - itemId: "tool-1", - itemType: "tool_call", - parentMessageId: "message-1", - title: "Search", - }, - runId: "run-1", - }), - record({ - id: "tool-running", - kind: "tool.call.updated", - occurredAtMs: 2_010, - payload: { - rawInput: '{"q":', - status: "running", - toolCallId: "tool-1", - }, - runId: "run-1", - }), - ]), - ).toEqual([]); - - const compacted = compactor.compact([ - record({ - id: "tool-input", - kind: "tool.call.updated", - occurredAtMs: 2_020, - payload: { - rawInput: '"cats"}', - status: "running", - toolCallId: "tool-1", - }, - runId: "run-1", - }), - record({ - id: "tool-output", - kind: "tool.call.updated", - occurredAtMs: 2_050, - payload: { - messageId: "message-1", - rawOutput: "ok", - status: "completed", - toolCallId: "tool-1", - }, - runId: "run-1", - }), - record({ - id: "tool-end", - kind: "item.completed", - occurredAtMs: 2_060, - payload: { - itemId: "tool-1", - itemType: "tool_call", - status: "completed", - }, - runId: "run-1", - }), - ]); - - expect(compacted.map((entry) => entry.event.kind)).toEqual(["tool.call.updated"]); - expect(compacted[0]?.occurredAt).toBe(2_000); - expect(compacted[0]?.sourceEventId).toBe("source-tool-end"); - expect(compacted[0]?.event.payload).toMatchObject({ - messageId: "message-1", - parentMessageId: "message-1", - rawInput: '{"q":"cats"}', - rawOutput: "ok", - status: "completed", - title: "Search", - toolCallId: "tool-1", - }); - }); - - test("waits until run end before storing repeated completed tool snapshots", () => { - const compactor = new RuntimeEventPersistenceCompactor(); - - expect( - compactor.compact([ - record({ - id: "tool-start", - kind: "item.started", - occurredAtMs: 3_000, - payload: { - itemId: "tool-1", - itemType: "tool_call", - title: "Shell", - }, - runId: "run-1", - }), - ]), - ).toEqual([]); - expect( - compactor.compact([ - record({ - id: "tool-snapshot-1", - kind: "tool.call.updated", - occurredAtMs: 3_010, - payload: { - rawOutput: "one", - status: "completed", - toolCallId: "tool-1", - }, - runId: "run-1", - }), - ]), - ).toEqual([]); - expect( - compactor.compact([ - record({ - id: "tool-snapshot-2", - kind: "tool.call.updated", - occurredAtMs: 3_020, - payload: { - rawOutput: "one two", - status: "completed", - toolCallId: "tool-1", - }, - runId: "run-1", - }), - ]), - ).toEqual([]); - - const compacted = compactor.compact([ - record({ - id: "run-end", - kind: "run.completed", - occurredAtMs: 3_100, - payload: { stopReason: "end_turn" }, - runId: "run-1", - }), - ]); - - expect(compacted.map((entry) => entry.event.kind)).toEqual([ - "tool.call.updated", - "run.completed", - ]); - expect(compacted[0]?.event.payload).toMatchObject({ - rawOutput: "one two", - status: "completed", - title: "Shell", - toolCallId: "tool-1", - }); - }); - - test("stores terminal run failure as the pending tool result", () => { - const compactor = new RuntimeEventPersistenceCompactor(); - - expect( - compactor.compact([ - record({ - id: "tool-start", - kind: "item.started", - occurredAtMs: 3_000, - payload: { - itemId: "tool-1", - itemType: "tool_call", - title: "Shell", - }, - runId: "run-1", - }), - record({ - id: "tool-running", - kind: "tool.call.updated", - occurredAtMs: 3_010, - payload: { - rawInput: '{"cmd":"pwd"}', - status: "running", - toolCallId: "tool-1", - }, - runId: "run-1", - }), - ]), - ).toEqual([]); - - const compacted = compactor.compact([ - record({ - id: "run-failed", - kind: "run.failed", - occurredAtMs: 3_100, - payload: { - error: { - code: "runtime.failed", - message: "Runtime driver control socket is not connected.", - }, - }, - runId: "run-1", - }), - ]); - - expect(compacted.map((entry) => entry.event.kind)).toEqual(["tool.call.updated", "run.failed"]); - expect(compacted[0]?.event.payload).toMatchObject({ - rawInput: '{"cmd":"pwd"}', - rawOutput: - "Tool failed before returning a result: Runtime driver control socket is not connected.", - status: "failed", - title: "Shell", - toolCallId: "tool-1", - }); - }); - - test("drops stream fragment replays shadowed by durable terminal receipts", () => { - const messageStart = record({ - id: "message-start", - kind: "message.started", - occurredAtMs: 4_000, - payload: { messageId: "message-1", role: "agent" }, - runId: "run-1", - }); - const messageDelta = record({ - id: "message-delta", - kind: "message.delta", - occurredAtMs: 4_010, - payload: { contentDelta: "Hello", messageId: "message-1", role: "agent" }, - runId: "run-1", - }); - const messageEnd = record({ - id: "message-end", - kind: "message.completed", - occurredAtMs: 4_020, - payload: { messageId: "message-1", role: "agent" }, - runId: "run-1", - }); - - expect( - filterDurablyAcceptedRuntimeStreamReplays( - [messageStart, messageDelta, messageEnd].map(envelope), - new Set(["source-message-end"]), - ), - ).toEqual([]); - }); - - test("keeps a later final assistant message after a completed progress message replays", () => { - const progressStart = record({ - id: "progress-start", - kind: "message.started", - occurredAtMs: 4_000, - payload: { messageId: "message-progress", role: "agent" }, - runId: "run-1", - }); - const progressDelta = record({ - id: "progress-delta", - kind: "message.delta", - occurredAtMs: 4_010, - payload: { contentDelta: "进度：工具已完成。", messageId: "message-progress", role: "agent" }, - runId: "run-1", - }); - const progressEnd = record({ - id: "progress-end", - kind: "message.completed", - occurredAtMs: 4_020, - payload: { messageId: "message-progress", role: "agent" }, - runId: "run-1", - }); - const finalStart = record({ - id: "final-start", - kind: "message.started", - occurredAtMs: 4_030, - payload: { messageId: "message-final", role: "agent" }, - runId: "run-1", - }); - const finalDelta = record({ - id: "final-delta", - kind: "message.delta", - occurredAtMs: 4_040, - payload: { contentDelta: "最终：中文表格完整。", messageId: "message-final", role: "agent" }, - runId: "run-1", - }); - const finalEnd = record({ - id: "final-end", - kind: "message.completed", - occurredAtMs: 4_050, - payload: { messageId: "message-final", role: "agent" }, - runId: "run-1", - }); - - expect( - filterDurablyAcceptedRuntimeStreamReplays( - [progressStart, progressDelta, progressEnd, finalStart, finalDelta, finalEnd].map(envelope), - new Set(["source-progress-end"]), - ).map((entry) => entry.event.id), - ).toEqual(["final-start", "final-delta", "final-end"]); - }); - - test("drops open stream fragment replays shadowed by a durable run terminal receipt", () => { - const messageDelta = record({ - id: "message-delta", - kind: "message.delta", - occurredAtMs: 5_000, - payload: { contentDelta: "partial", messageId: "message-1", role: "agent" }, - runId: "run-1", - }); - const toolDelta = record({ - id: "tool-delta", - kind: "tool.call.updated", - occurredAtMs: 5_010, - payload: { rawInput: "{}", status: "running", toolCallId: "tool-1" }, - runId: "run-1", - }); - const runEnd = record({ - id: "run-end", - kind: "run.completed", - occurredAtMs: 5_100, - payload: { stopReason: "end_turn" }, - runId: "run-1", - }); - - expect( - filterDurablyAcceptedRuntimeStreamReplays( - [messageDelta, toolDelta, runEnd].map(envelope), - new Set(["source-run-end"]), - ), - ).toEqual([]); - }); -}); diff --git a/apps/api/tests/runtime-final-output-ingestion.test.ts b/apps/api/tests/runtime-final-output-ingestion.test.ts index ab3ee3a0..56ff94c0 100644 --- a/apps/api/tests/runtime-final-output-ingestion.test.ts +++ b/apps/api/tests/runtime-final-output-ingestion.test.ts @@ -1,46 +1,75 @@ import { afterEach, describe, expect, test } from "bun:test"; +import { createMcpExecuteFailedEventIdentity } from "@mosoo/agent-driver/events"; import type { DriverEventEnvelope } from "@mosoo/agent-driver/events"; -import type { DriverEventReceipt } from "@mosoo/agent-driver/orpc"; import { createPlatformId } from "@mosoo/id"; import type { DriverInstanceId, + DriverCommandId, RuntimeEventId, + RuntimeOperationId, SessionId, SessionMessageId, SessionRunId, } from "@mosoo/id"; -import { createRuntimeEvent } from "@mosoo/runtime-events"; +import { createRuntimeEvent, createRuntimeToolResultMessageId } from "@mosoo/runtime-events"; import type { RuntimeEventKind } from "@mosoo/runtime-events"; import { readPublicThreadRunFinalOutput } from "../src/modules/public-api/public-thread-events"; import { - createReceiptsForDriverEvents, - filterNewDriverEvents, - readReceiptsForProcessedDriverEvents, - rememberDriverEventReceipts, -} from "../src/modules/runtime/infrastructure/driver-instance/driver-event-receipts"; + recordCanonicalSessionRunFailure, + recordCanonicalSessionRunTerminal, +} from "../src/modules/runtime/application/session-runs/session-run-terminal-failure.service"; +import { + createFailedSessionRunRuntimeEvent, + createSessionRunUpdatedEvent, +} from "../src/modules/runtime/application/session-runs/session-run-view-events.service"; +import { prepareAssistantMessageProjection } from "../src/modules/runtime/infrastructure/driver-instance/assistant-message-projection"; +import { commitTerminalRunProjection } from "../src/modules/runtime/infrastructure/driver-instance/completed-run-commit.repository"; +import { canonicalizeDriverEventEnvelope } from "../src/modules/runtime/infrastructure/driver-instance/driver-event-canonicalization"; import type { RuntimeSessionLink } from "../src/modules/runtime/infrastructure/driver-instance/event-types"; import { DriverInstanceRpcEventIngestionController } from "../src/modules/runtime/infrastructure/driver-instance/rpc-event-ingestion-controller"; +import { + cleanupRuntimeArtifactAttempts, + parseRuntimeArtifactManifest, +} from "../src/modules/runtime/infrastructure/driver-instance/runtime-artifact-attempt.repository"; +import type { RuntimeArtifactManifest } from "../src/modules/runtime/infrastructure/driver-instance/runtime-artifact-attempt.repository"; +import { + RUNTIME_SESSION_OUTPUT_MAX_FILE_BYTES, + RUNTIME_SESSION_OUTPUT_MAX_TOTAL_BYTES, +} from "../src/modules/runtime/infrastructure/driver-instance/runtime-session-outputs"; import { RuntimeSessionViewCache } from "../src/modules/runtime/infrastructure/driver-instance/runtime-session-view-cache"; -import { recordDriverInstanceCompletion } from "../src/modules/runtime/infrastructure/driver-instance/terminal-driver-events"; +import { + recordDriverInstanceCompletion, + recordDriverInstanceFailure, +} from "../src/modules/runtime/infrastructure/driver-instance/terminal-driver-events"; +import { createRuntimeCommandRecord } from "../src/modules/runtime/infrastructure/session-runs/runtime-command-store.repository"; +import { getSessionRunSummary } from "../src/modules/runtime/infrastructure/session-runs/session-run-store.repository"; import { loadSessionViewerState } from "../src/modules/sessions/application/session-live-state.service"; +import type { SessionDeliveryEvent } from "../src/modules/sessions/application/session-live-state.service"; import { createSessionProcessEventsFromSessionEventRows } from "../src/modules/sessions/application/session-process-events.service"; import type { SessionEventProcessRow } from "../src/modules/sessions/application/session-process-events.service"; +import { getSessionRuntimeRecoveryMessages } from "../src/modules/sessions/application/session-runtime-recovery-query.service"; +import { syncSessionViewerState } from "../src/modules/sessions/infrastructure/session/client"; import { setServerProductAnalyticsTransportForTests } from "../src/platform/analytics/product-analytics"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { createPublicHttpContractDatabase, - createPublicHttpTestBindings, + createPublicHttpTestBindings as createBasePublicHttpTestBindings, + insertNonOwnerSession, insertOwnerSession, PUBLIC_API_TEST_IDS, + PublicApiMemoryFileBucket, } from "./helpers/public-api-http-test-fixture"; import type { SqliteD1Database } from "./helpers/public-api-http-test-fixture"; +import { createRuntimeOutputSandbox } from "./helpers/runtime-output-sandbox"; +import type { RuntimeOutputSandboxOptions } from "./helpers/runtime-output-sandbox"; +import { insertRuntimeEvent } from "./public-thread-api-fixtures"; const DRIVER_ID = PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId; const RUN_ID = PUBLIC_API_TEST_IDS.run as SessionRunId; const SESSION_ID = PUBLIC_API_TEST_IDS.ownerSession as SessionId; -const TERMINAL_SOURCE_EVENT_ID = "canary:run-completed"; +const TERMINAL_SOURCE_EVENT_ID = `session-run-terminal:${RUN_ID}:run.completed`; const CANARY_LINES = Array.from({ length: 160 }, (_, index) => { const lineNumber = String(index + 1).padStart(3, "0"); return `${lineNumber}|中文长文本校验-Aa${index % 10}-表格字符|END${lineNumber}`; @@ -61,6 +90,8 @@ const FINAL_TEXT_LINES = [ "CANARY-FINAL-END", ]; const FINAL_TEXT = FINAL_TEXT_LINES.join("\n"); +const LARGE_FINAL_TEXT = `${FINAL_TEXT}\n${"0123456789abcdef中文\n".repeat(90_000)}`; +const LARGE_FINAL_TEXT_CHUNK_CHARACTERS = 300; const PROGRESS_TEXTS = [ "进度 1：正在读取上游报告。", "进度 2：工具调用已经完成。", @@ -72,35 +103,18 @@ afterEach(() => { }); interface TestDriverState { - createDriverEventReceipts(events: readonly DriverEventEnvelope[]): DriverEventReceipt[]; - filterUnprocessedDriverEvents(events: readonly DriverEventEnvelope[]): DriverEventEnvelope[]; hello: { pid: number }; - readProcessedDriverEventReceipts(events: readonly DriverEventEnvelope[]): DriverEventReceipt[]; - rememberProcessedDriverEventReceipts(receipts: DriverEventReceipt[]): void; + requireDriverGeneration(): number; requireDriverInstanceId(): DriverInstanceId; runtimeSessionLink: RuntimeSessionLink | null; setRuntimeSessionLink(link: RuntimeSessionLink): void; } function createDriverState(): TestDriverState { - const processedReceipts = new Map(); - let nextSeq = 0; - return { - createDriverEventReceipts(events) { - const result = createReceiptsForDriverEvents({ events, nextSeq }); - nextSeq = result.nextSeq; - return result.receipts; - }, - filterUnprocessedDriverEvents(events) { - return filterNewDriverEvents({ events, processedReceipts }); - }, hello: { pid: 1 }, - readProcessedDriverEventReceipts(events) { - return readReceiptsForProcessedDriverEvents({ events, processedReceipts }); - }, - rememberProcessedDriverEventReceipts(receipts) { - rememberDriverEventReceipts({ processedReceipts, receipts }); + requireDriverGeneration() { + return 0; }, requireDriverInstanceId() { return DRIVER_ID; @@ -112,38 +126,105 @@ function createDriverState(): TestDriverState { }; } -function createController(bindings: ApiBindings): DriverInstanceRpcEventIngestionController { +function createController( + bindings: ApiBindings, + enqueue: (sessionId: SessionId | null, events: SessionDeliveryEvent[]) => void = () => undefined, +): DriverInstanceRpcEventIngestionController { return new DriverInstanceRpcEventIngestionController({ env: bindings, state: createDriverState(), viewCache: new RuntimeSessionViewCache(), viewerEventDelivery: { - enqueue: () => undefined, + enqueue, + flush: async () => {}, + flushSafely: async () => {}, + requestStateSync: (sessionId) => { + void syncSessionViewerState(bindings, sessionId); + }, + resetAfterFlush: () => {}, }, } as never); } +function createSessionSyncNamespace(syncedSessionIds: string[]): ApiBindings["Session"] { + return { + get: () => ({ + syncViewers: async (sessionId: string) => { + syncedSessionIds.push(sessionId); + }, + }), + idFromName: (name: string) => name, + } as never; +} + +interface ArtifactSandboxOptions extends RuntimeOutputSandboxOptions { + readonly resolveError?: Error; +} + +type PublicHttpTestBindingOptions = NonNullable< + Parameters[1] +> & { readonly artifactSandbox?: ArtifactSandboxOptions }; + +function createPublicHttpTestBindings( + database: D1Database, + options: PublicHttpTestBindingOptions = {}, +): Record { + const { artifactSandbox, ...baseOptions } = options; + + return { + ...createBasePublicHttpTestBindings(database, baseOptions), + runtimeSubjectHandleFactory: () => { + if (artifactSandbox?.resolveError !== undefined) { + throw artifactSandbox.resolveError; + } + return createRuntimeOutputSandbox(artifactSandbox); + }, + }; +} + +async function readRuntimeArtifactManifest( + database: D1Database, + sourceEventId: string, +): Promise { + const row = await database + .prepare("SELECT artifact_manifest_json FROM session_event WHERE source_event_id = ?") + .bind(sourceEventId) + .first<{ artifact_manifest_json: string | null }>(); + if (row?.artifact_manifest_json === null || row?.artifact_manifest_json === undefined) { + throw new Error(`Runtime artifact manifest is missing for ${sourceEventId}.`); + } + return parseRuntimeArtifactManifest(row.artifact_manifest_json); +} + function runtimeEvent(input: { + correlationId?: string; kind: RuntimeEventKind; + occurredAt?: number; payload: unknown; + runId?: SessionRunId | null; sourceEventId: string; + visibility?: "owner_debug"; }): DriverEventEnvelope { - const occurredAt = Date.now(); + const occurredAt = input.occurredAt ?? Date.now(); const event = createRuntimeEvent({ + ...(input.correlationId === undefined ? {} : { correlationId: input.correlationId }), driverInstanceId: DRIVER_ID, id: createPlatformId(), kind: input.kind, occurredAt: new Date(occurredAt).toISOString(), payload: input.payload, - runId: RUN_ID, + ...(input.runId === null ? {} : { runId: input.runId ?? RUN_ID }), + runtimeId: "openai-runtime", sessionId: SESSION_ID, sourceEventId: input.sourceEventId, + traceId: "trace-canary", + ...(input.visibility === undefined ? {} : { visibility: input.visibility }), }); return { event, eventId: input.sourceEventId, - occurredAt, + occurredAt: new Date(occurredAt).toISOString(), }; } @@ -154,14 +235,9 @@ function messageEvents(input: { }): DriverEventEnvelope[] { return [ runtimeEvent({ - kind: "message.started", - payload: { messageId: input.messageId, role: "agent" }, - sourceEventId: `${input.sourcePrefix}:started`, - }), - runtimeEvent({ - kind: "message.delta", - payload: { contentDelta: input.text, messageId: input.messageId, role: "agent" }, - sourceEventId: `${input.sourcePrefix}:delta`, + kind: "message.added", + payload: { content: input.text, messageId: input.messageId, role: "agent" }, + sourceEventId: `${input.sourcePrefix}:snapshot`, }), runtimeEvent({ kind: "message.completed", @@ -181,37 +257,60 @@ function splitIntoBatches(values: readonly T[], size: number): T[][] { return batches; } +function measurePreparedQueries(database: D1Database): { + database: D1Database; + readCount: () => number; +} { + let count = 0; + const measured = new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => { + count += 1; + return target.prepare(query); + }; + } + const value = Reflect.get(target, property, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + return { database: measured, readCount: () => count }; +} + async function insertRuntimeFixture(database: SqliteD1Database): Promise { await insertOwnerSession(database); database.execute(` INSERT INTO sandbox ( - id, kind, subject_kind, subject_id, status, bind_mount_ready, + agent_id, app_id, id, incarnation, kind, network_constraints_hash, + owner_account_id, subject_kind, subject_id, status, bind_mount_ready, global_mounts_json, created_at, updated_at ) VALUES ( - '${PUBLIC_API_TEST_IDS.sandbox}', 'pet', 'agent', '${PUBLIC_API_TEST_IDS.agent}', + '${PUBLIC_API_TEST_IDS.agent}', '${PUBLIC_API_TEST_IDS.app}', + '${PUBLIC_API_TEST_IDS.sandbox}', 1, 'pet', '${"0".repeat(64)}', + '${PUBLIC_API_TEST_IDS.ownerAccount}', 'agent', '${PUBLIC_API_TEST_IDS.agent}', 'active', 1, '[]', 1, 1 ); INSERT INTO sandbox_session ( cloudflare_session_id, created_at, cwd, origin_json, sandbox_id, - session_id, status, updated_at + sandbox_incarnation, session_id, status, updated_at ) VALUES ( - 'canary-cloudflare-session', 1, '/workspace', + '${PUBLIC_API_TEST_IDS.operation}', 1, '/workspace', '{"callerUserId":"${PUBLIC_API_TEST_IDS.ownerAccount}","entrypoint":"api","executionOwnerUserId":"${PUBLIC_API_TEST_IDS.ownerAccount}","type":"agent"}', - '${PUBLIC_API_TEST_IDS.sandbox}', '${SESSION_ID}', 'active', 1 + '${PUBLIC_API_TEST_IDS.sandbox}', 1, '${SESSION_ID}', 'active', 1 ); INSERT INTO driver_instance ( id, boot_token_expires_at, boot_token_hash, connection_id, created_at, expires_at, heartbeat_count, protocol, protocol_version, runtime, - sandbox_id, sandbox_session_id, status, updated_at + sandbox_id, sandbox_incarnation, sandbox_session_id, status, updated_at ) VALUES ( '${DRIVER_ID}', 1, X'01', 'canary-connection', 1, 1, 0, 'orpc-ws', 1, 'openai-runtime', '${PUBLIC_API_TEST_IDS.sandbox}', - '${SESSION_ID}', 'ready', 1 + 1, '${SESSION_ID}', 'ready', 1 ); INSERT INTO session_run ( @@ -275,7 +374,7 @@ function failTerminalSessionEventInsert(database: D1Database): D1Database { return (query: string) => wrapStatement( target.prepare(query), - /insert\s+into\s+["`]session_event["`]/iu.test(query), + /insert\s+into\s+(?:["`]session_event["`]|session_event)/iu.test(query), false, ); } @@ -286,6 +385,198 @@ function failTerminalSessionEventInsert(database: D1Database): D1Database { }); } +function rejectSessionEventContentReads(database: D1Database): D1Database { + return new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => { + if ( + /^\s*select\b/iu.test(query) && + /\bcontent_text\b/iu.test(query) && + /\bfrom\s+["`]?session_event["`]?/iu.test(query) + ) { + throw new Error("Terminal ingestion attempted to materialize session event content."); + } + return target.prepare(query); + }; + } + + const value = Reflect.get(target, property, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }); +} + +function failPreparedProjection(database: D1Database, pattern: RegExp): D1Database { + function wrap(statement: D1PreparedStatement, shouldFail: boolean): D1PreparedStatement { + return new Proxy(statement, { + get(target, property) { + if (property === "bind") { + return (...values: unknown[]) => wrap(target.bind(...values), shouldFail); + } + if (property === "run" && shouldFail) { + return async () => { + throw new Error("injected durable side-effect projection failure"); + }; + } + + const value = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + } + + return new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => wrap(target.prepare(query), pattern.test(query)); + } + + const value = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }); +} + +type RuntimeArtifactAckLossPhase = "claim" | "commit" | "create" | "put" | "seal"; + +function failAfterRuntimeArtifactDatabaseWrite( + database: D1Database, + phase: Exclude, +): { readonly database: D1Database; readonly wasInjected: () => boolean } { + const pattern = { + claim: "SET owned_object_keys_json = json_insert", + create: "INSERT INTO runtime_artifact_attempt", + seal: "SET manifest_json = ?, manifest_sha256 = ?, status = 'staged'", + }[phase]; + let injected = false; + + function wrap(statement: D1PreparedStatement): D1PreparedStatement { + return new Proxy(statement, { + get(target, property) { + if (property === "bind") { + return (...values: unknown[]) => wrap(target.bind(...values)); + } + if (property === "first") { + return async (...arguments_: unknown[]) => { + const result = await Reflect.apply(target.first, target, arguments_); + if (!injected) { + injected = true; + throw new Error(`injected artifact ${phase} acknowledgement loss`); + } + return result; + }; + } + + const value = Reflect.get(target, property, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + } + + return { + database: new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => { + const statement = target.prepare(query); + return query.includes(pattern) ? wrap(statement) : statement; + }; + } + + const value = Reflect.get(target, property, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }), + wasInjected: () => injected, + }; +} + +function failAfterRuntimeArtifactCommit( + database: D1Database, + sourceEventId: string, +): { readonly database: D1Database; readonly wasInjected: () => boolean } { + let injected = false; + return { + database: new Proxy(database, { + get(target, property) { + if (property === "batch") { + return async (statements: D1PreparedStatement[]) => { + const result = await target.batch(statements); + const receipt = await target + .prepare("SELECT 1 AS found FROM session_event WHERE source_event_id = ?") + .bind(sourceEventId) + .first(); + if (!injected && receipt !== null) { + injected = true; + throw new Error("injected artifact commit acknowledgement loss"); + } + return result; + }; + } + + const value = Reflect.get(target, property, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }), + wasInjected: () => injected, + }; +} + +function mutateBeforeFirstDatabaseBatch( + database: D1Database, + mutation: (database: D1Database) => Promise, +): { readonly database: D1Database; readonly wasInjected: () => boolean } { + let injected = false; + + return { + database: new Proxy(database, { + get(target, property) { + if (property === "batch") { + return async (statements: D1PreparedStatement[]) => { + if (!injected) { + injected = true; + await mutation(target); + } + return target.batch(statements); + }; + } + + const value = Reflect.get(target, property, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }), + wasInjected: () => injected, + }; +} + +function failAfterRuntimeArtifactPut(bucket: PublicApiMemoryFileBucket): { + readonly bucket: R2Bucket; + readonly wasInjected: () => boolean; +} { + let injected = false; + return { + bucket: new Proxy(bucket, { + get(target, property) { + if (property === "put") { + return async (...arguments_: Parameters) => { + const result = await target.put(...arguments_); + if (!injected) { + injected = true; + throw new Error("injected artifact put acknowledgement loss"); + } + return result; + }; + } + + const value = Reflect.get(target, property, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }), + wasInjected: () => injected, + }; +} + const activeContext = { assertActiveConnection: () => undefined, connectionId: "canary-connection", @@ -302,490 +593,4646 @@ async function pushFreshController( return result.accepted; } -describe("runtime final output ingestion", () => { - test.each([ - ["omits", false], - ["provides", true], - ] as const)( - "persists one final assistant snapshot when the driver %s it", - async (_driverBehavior, driverProvidesSnapshot) => { - const database = await createPublicHttpContractDatabase(); - await insertRuntimeFixture(database); - const capturedEvents: unknown[] = []; - setServerProductAnalyticsTransportForTests(async (_input, init) => { - capturedEvents.push(JSON.parse(init.body as string) as unknown); - return new Response(null, { status: 200 }); - }); - const bindings = { - ...createPublicHttpTestBindings(database), - POSTHOG_PROJECT_KEY: "phc_test", - } as ApiBindings; - const finalText = "The final answer."; - const fragmentTexts = ["The ", "final ", "answer."]; - const finalMessageId = createPlatformId(); - const events = [ - ...fragmentTexts.flatMap((text, index) => - messageEvents({ - messageId: createPlatformId(), - sourcePrefix: `fractured:${index + 1}`, - text, - }), - ), - ...(driverProvidesSnapshot - ? [ - runtimeEvent({ - kind: "message.added", - payload: { content: finalText, messageId: finalMessageId, role: "agent" }, - sourceEventId: "fractured:final-snapshot", - }), - ] - : []), - runtimeEvent({ - kind: "run.completed", - payload: { - finalMessageId, - finalMessageText: finalText, - stopReason: "end_turn", - }, - sourceEventId: "fractured:run-completed", - }), - ]; - - await pushFreshController(bindings, events); +async function insertRepairableCompletedAuthority( + database: D1Database, + sourcePrefix: string, +): Promise { + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + await pushFreshController(bindings, [ + ...messageEvents({ messageId: finalMessageId, sourcePrefix, text: FINAL_TEXT }), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: `${sourcePrefix}:terminal`, + }), + ]); + await database.prepare("DELETE FROM session_message WHERE id = ?").bind(finalMessageId).run(); + await database + .prepare("UPDATE session SET message_seq_cursor = 0 WHERE id = ?") + .bind(SESSION_ID) + .run(); + return finalMessageId; +} - const rows = await database - .prepare( - `SELECT content_text, ended_at, event_type, id, occurred_at, process_status, - process_type, run_id, seq, tokens - FROM session_event - WHERE session_id = ? AND run_id = ? - ORDER BY seq`, - ) - .bind(SESSION_ID, RUN_ID) - .all(); - const assistantMessages = createSessionProcessEventsFromSessionEventRows(rows.results).filter( - (event) => event.type === "agent.message.delta", - ); +const VALID_RUNTIME_ARTIFACT_MANIFEST = { + captureStatus: "complete", + files: [], + mode: "delta", + semanticHash: "0".repeat(64), + sourceEventId: "artifact-manifest:valid", + version: 1, +}; - expect( - rows.results - .filter((row) => row.event_type === "message.added") - .map((row) => row.content_text), - ).toEqual([finalText]); - expect(assistantMessages.map((event) => event.content)).toEqual([finalText]); - expect(capturedEvents).toEqual([ - expect.objectContaining({ - event: "task_succeeded", - properties: expect.objectContaining({ - run_duration_ms: expect.any(Number), - sandbox_id: PUBLIC_API_TEST_IDS.sandbox, - sandbox_kind: "pet", - sandbox_subject_kind: "agent", - session_type: "ui", - }), - }), - ]); - }, - ); +describe("runtime artifact manifest validation", () => { + test.each([ + ["malformed JSON", "{"], + ["non-object root", null], + ["unsupported version", { ...VALID_RUNTIME_ARTIFACT_MANIFEST, version: 2 }], + ["unsupported capture status", { ...VALID_RUNTIME_ARTIFACT_MANIFEST, captureStatus: "lost" }], + ["unsupported mode", { ...VALID_RUNTIME_ARTIFACT_MANIFEST, mode: "append" }], + ["non-array files", { ...VALID_RUNTIME_ARTIFACT_MANIFEST, files: {} }], + ["empty source event ID", { ...VALID_RUNTIME_ARTIFACT_MANIFEST, sourceEventId: "" }], + ["invalid semantic hash", { ...VALID_RUNTIME_ARTIFACT_MANIFEST, semanticHash: "invalid" }], + [ + "invalid file operation", + { + ...VALID_RUNTIME_ARTIFACT_MANIFEST, + files: [{ operation: "rename", sourcePath: "outputs/report.txt" }], + }, + ], + [ + "non-empty omitted capture", + { + ...VALID_RUNTIME_ARTIFACT_MANIFEST, + captureStatus: "omitted_size_limit", + files: [{ operation: "delete", sourcePath: "outputs/report.txt" }], + }, + ], + [ + "duplicate file identity", + { + ...VALID_RUNTIME_ARTIFACT_MANIFEST, + files: [ + { operation: "delete", sourcePath: "outputs/report.txt" }, + { operation: "delete", sourcePath: "outputs/report.txt" }, + ], + }, + ], + ] as const)("rejects the %s boundary", (_name, value) => { + const manifestJson = typeof value === "string" ? value : JSON.stringify(value); + expect(() => parseRuntimeArtifactManifest(manifestJson)).toThrow(); + }); +}); - test("preserves a long final snapshot across hibernation, terminal failure, and replay", async () => { +describe("runtime final output ingestion", () => { + test("admits only canonical content-addressed MCP failures", async () => { const database = await createPublicHttpContractDatabase(); await insertRuntimeFixture(database); const bindings = createPublicHttpTestBindings(database) as ApiBindings; - const progressMessageIds = PROGRESS_TEXTS.map(() => createPlatformId()); - const finalMessageId = createPlatformId(); - const progressEvents = [ - runtimeEvent({ - kind: "run.started", - payload: { startedAt: new Date(1).toISOString() }, - sourceEventId: "canary:run-started", - }), - ...PROGRESS_TEXTS.flatMap((text, index) => - messageEvents({ - messageId: progressMessageIds[index], - sourcePrefix: `canary:progress:${index + 1}`, - text, + const commandId = "01J0000000000000000000000X" as DriverCommandId; + const command = { + argumentsJson: '{"issue":"A-1"}', + commandId, + kind: "mcp.execute", + requestId: "request-A-1", + runId: RUN_ID, + serverId: "01J0000000000000000000000Y", + toolCallId: "tool-A-1", + toolName: "createIssue", + } as const; + await createRuntimeCommandRecord(database, { + command, + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + status: "accepted", + }); + const failed = createMcpExecuteFailedEventIdentity({ + commandId, + rawInput: command.argumentsJson, + rawOutput: "provider rejected the request", + title: command.toolName, + toolCallId: command.toolCallId, + }); + const event = runtimeEvent({ + correlationId: commandId, + kind: "tool.call.updated", + payload: failed.payload, + sourceEventId: failed.sourceEventId, + }); + + await expect(pushFreshController(bindings, [event])).resolves.toMatchObject([ + { eventId: failed.sourceEventId }, + ]); + await expect( + pushFreshController(bindings, [ + { + ...event, + event: { + ...event.event, + payload: { ...failed.payload, rawOutput: "tampered failure" }, + }, + }, + ]), + ).rejects.toThrow("canonical content identity"); + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + correlationId: commandId, + kind: "tool.call.updated", + payload: failed.payload, + sourceEventId: `mcp.execute.failed:${"0".repeat(64)}`, }), - ), - ]; + ]), + ).rejects.toThrow("canonical content identity"); + }); - expect(await pushFreshController(bindings, progressEvents)).toHaveLength(progressEvents.length); + test("commits canonical durable state before enqueue and replays without projection writes", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + await database.prepare("UPDATE session SET title = NULL WHERE id = ?").bind(SESSION_ID).run(); + const syncedSessionIds: string[] = []; + const bindings = createPublicHttpTestBindings(database, { + sessionNamespace: createSessionSyncNamespace(syncedSessionIds), + }) as ApiBindings; + let durableCommitCompleted = false; + const disconnectingDatabase = new Proxy(database, { + get(target, property) { + if (property === "batch") { + return async (statements: D1PreparedStatement[]) => { + const result = await target.batch(statements); + durableCommitCompleted = true; + return result; + }; + } - const toolEvents = [ + const value = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + let failureInjected = false; + const disconnectAfterCommitContext = { + ...activeContext, + assertActiveConnection() { + if (durableCommitCompleted && !failureInjected) { + failureInjected = true; + throw new Error("injected disconnect after durable event commit"); + } + }, + } as never; + const taskSourceEventId = "task-replay:tasks"; + const snapshotSourceEventId = "task-replay:message-snapshot"; + const deltaSourceEventId = "task-replay:message-delta"; + const messageId = createPlatformId(); + const thoughtId = createPlatformId(); + const userMessageId = createPlatformId(); + const events = [ runtimeEvent({ - kind: "item.started", - payload: { - itemId: "tool-canary", - itemType: "tool_call", - parentMessageId: finalMessageId, - title: "Create artifact", + kind: "agent.tasks.replaced", + payload: { tasks: [{ taskId: "canonical-task", title: "Inspect" }] }, + sourceEventId: taskSourceEventId, + }), + runtimeEvent({ + kind: "message.added", + payload: { content: "canonical ", messageId, role: "agent" }, + sourceEventId: snapshotSourceEventId, + }), + runtimeEvent({ + kind: "message.delta", + payload: { + contentDelta: "state", + messageId, + role: "agent", }, - sourceEventId: "canary:tool:started", + sourceEventId: deltaSourceEventId, + }), + runtimeEvent({ + kind: "thought.started", + payload: { thoughtId }, + sourceEventId: "task-replay:thought-started", + }), + runtimeEvent({ + kind: "thought.delta", + payload: { contentDelta: "canonical reasoning", thoughtId }, + sourceEventId: "task-replay:thought-delta", + }), + runtimeEvent({ + kind: "thought.completed", + payload: { thoughtId }, + sourceEventId: "task-replay:thought-completed", + }), + runtimeEvent({ + kind: "message.added", + payload: { content: "canonical question", messageId: userMessageId, role: "user" }, + sourceEventId: "task-replay:user-message", }), runtimeEvent({ kind: "tool.call.updated", payload: { - rawOutput: "artifact created", + parentMessageId: messageId, + rawInputDelta: '{"command":', + status: "running", + title: "Shell", + toolCallId: "stable-tool", + }, + sourceEventId: "task-replay:tool-started", + }), + runtimeEvent({ + kind: "tool.call.updated", + payload: { + messageId, + rawInput: '{"command":"pwd"}', + status: "running", + title: "Shell", + toolCallId: "stable-tool", + }, + sourceEventId: "task-replay:tool-snapshot", + }), + runtimeEvent({ + kind: "tool.call.updated", + payload: { + rawOutput: "canonical output", status: "completed", - toolCallId: "tool-canary", + title: "Shell", + toolCallId: "stable-tool", }, - sourceEventId: "canary:tool:updated", + sourceEventId: "task-replay:tool-completed", }), runtimeEvent({ - kind: "item.completed", - payload: { itemId: "tool-canary", itemType: "tool_call", status: "completed" }, - sourceEventId: "canary:tool:completed", + kind: "tool.call.updated", + payload: { + status: "running", + title: "Orphan tool", + toolCallId: "orphan-tool", + }, + sourceEventId: "task-replay:orphan-tool-started", + }), + runtimeEvent({ + kind: "tool.call.updated", + payload: { + rawOutput: "orphan output", + status: "completed", + title: "Orphan tool", + toolCallId: "orphan-tool", + }, + sourceEventId: "task-replay:orphan-tool-completed", + }), + runtimeEvent({ + kind: "plan.updated", + payload: { + entries: [{ content: "Verify durable state", priority: "high", status: "completed" }], + }, + sourceEventId: "task-replay:plan", + }), + runtimeEvent({ + kind: "session.commands.updated", + payload: { + commands: [{ description: "Inspect the workspace", name: "inspect" }], + }, + sourceEventId: "task-replay:commands", + }), + runtimeEvent({ + kind: "session.config.updated", + payload: { + options: [ + { + currentValue: "brief", + id: "tone", + name: "Tone", + type: "select", + values: [{ name: "Brief", value: "brief" }], + }, + ], + }, + sourceEventId: "task-replay:config", + }), + runtimeEvent({ + kind: "session.mode.updated", + payload: { + availableModes: [{ id: "plan", name: "Plan" }], + currentMode: "plan", + }, + sourceEventId: "task-replay:mode", + }), + runtimeEvent({ + kind: "session.info.updated", + payload: { title: "Atomic durable title" }, + sourceEventId: "task-replay:title", + }), + runtimeEvent({ + kind: "usage.updated", + payload: { + callId: "atomic-call", + inputTokens: 21, + outputTokens: 4, + source: "prompt_response", + usageContract: "openai_total_with_cached_breakdown", + }, + sourceEventId: "task-replay:usage", + }), + runtimeEvent({ + kind: "runtime.resume.updated", + payload: { resumePointer: "atomic-thread" }, + sourceEventId: "task-replay:native-resume", }), ]; - expect(await pushFreshController(bindings, toolEvents)).toHaveLength(toolEvents.length); + const firstDelivery: SessionDeliveryEvent[] = []; + + await expect( + createController( + { ...bindings, DB: disconnectingDatabase } as ApiBindings, + (_sessionId, deliveryEvents) => firstDelivery.push(...deliveryEvents), + ).handlePushEvents({ driverInstanceId: DRIVER_ID, events }, disconnectAfterCommitContext), + ).rejects.toThrow("injected disconnect after durable event commit"); + + expect(failureInjected).toBe(true); + expect(firstDelivery).toEqual([]); + const durableRows = await database + .prepare( + `SELECT event_type, source_event_id FROM session_event WHERE source_event_id IN (${events + .map(() => "?") + .join(", ")}) ORDER BY seq`, + ) + .bind(...events.map((event) => event.eventId)) + .all<{ event_type: string; source_event_id: string }>(); - const finalTextChunks = FINAL_TEXT_LINES.map((line, index) => - index === FINAL_TEXT_LINES.length - 1 ? line : `${line}\n`, + expect(durableRows.results).toEqual( + events.map((event) => ({ + event_type: event.event.kind, + source_event_id: event.eventId, + })), ); - const finalStreamEvents = [ - runtimeEvent({ - kind: "message.started", - payload: { messageId: finalMessageId, role: "agent" }, - sourceEventId: "canary:final:started", - }), - ...finalTextChunks.map((contentDelta, index) => - runtimeEvent({ - kind: "message.delta", - payload: { contentDelta, messageId: finalMessageId, role: "agent" }, - sourceEventId: `canary:final:delta:${index + 1}`, - }), + const titleSeq = events.findIndex((event) => event.event.kind === "session.info.updated") + 1; + const usageSeq = events.findIndex((event) => event.event.kind === "usage.updated") + 1; + const nativeResumeSeq = + events.findIndex((event) => event.event.kind === "runtime.resume.updated") + 1; + const usageReceipt = await database + .prepare("SELECT created_at FROM session_event WHERE source_event_id = ?") + .bind("task-replay:usage") + .first<{ created_at: number }>(); + + expect(usageReceipt).not.toBeNull(); + + expect( + await database + .prepare("SELECT auto_title_event_seq, title FROM session WHERE id = ?") + .bind(SESSION_ID) + .first(), + ).toEqual({ auto_title_event_seq: titleSeq, title: "Atomic durable title" }); + expect( + await database + .prepare( + `SELECT call_key, created_at, source_event_seq, status + FROM session_model_call + WHERE session_id = ? AND session_run_id = ?`, + ) + .bind(SESSION_ID, RUN_ID) + .first(), + ).toEqual({ + call_key: "model_call:atomic-call", + created_at: usageReceipt?.created_at, + source_event_seq: usageSeq, + status: "started", + }); + expect( + await database + .prepare( + `SELECT created_at, input_tokens, output_tokens, source_event_seq + FROM usage_event + WHERE session_id = ? AND session_run_id = ?`, + ) + .bind(SESSION_ID, RUN_ID) + .first(), + ).toEqual({ + created_at: usageReceipt?.created_at, + input_tokens: 21, + output_tokens: 4, + source_event_seq: usageSeq, + }); + expect( + await database + .prepare( + `SELECT observed_event_seq, observed_session_run_id, value + FROM native_resume_ref + WHERE session_id = ?`, + ) + .bind(SESSION_ID) + .first(), + ).toEqual({ + observed_event_seq: nativeResumeSeq, + observed_session_run_id: RUN_ID, + value: "atomic-thread", + }); + + const viewerUsageSeq = events.length + 1; + await insertRuntimeEvent(database, { + kind: "usage.updated", + occurredAt: Date.now(), + payload: { source: "session_update", totalTokens: 42 }, + runId: RUN_ID, + seq: viewerUsageSeq, + sessionId: SESSION_ID, + }); + await database + .prepare("UPDATE session SET runtime_event_seq_cursor = ? WHERE id = ?") + .bind(viewerUsageSeq, SESSION_ID) + .run(); + + const replayDelivery: SessionDeliveryEvent[] = []; + const replay = await createController( + { + ...bindings, + DB: failPreparedProjection( + database, + /SET auto_title_event_seq =|INSERT INTO ["`]?session_model_call|INSERT INTO native_resume_ref/iu, + ), + } as ApiBindings, + (_sessionId, deliveryEvents) => replayDelivery.push(...deliveryEvents), + ).handlePushEvents({ driverInstanceId: DRIVER_ID, events }, activeContext); + + expect(replay.accepted.map((receipt) => receipt.eventId)).toEqual( + events.map((event) => event.eventId), + ); + expect(replayDelivery).toEqual([]); + expect(syncedSessionIds).toEqual([SESSION_ID]); + expect( + await database + .prepare( + `SELECT + (SELECT COUNT(*) FROM session_model_call WHERE session_id = ?) AS model_call_count, + (SELECT created_at FROM session_model_call WHERE session_id = ?) AS model_call_created_at, + (SELECT COUNT(*) FROM usage_event WHERE session_id = ?) AS usage_event_count, + (SELECT created_at FROM usage_event WHERE session_id = ?) AS usage_event_created_at, + (SELECT observed_event_seq FROM native_resume_ref WHERE session_id = ?) AS native_seq, + (SELECT auto_title_event_seq FROM session WHERE id = ?) AS title_seq`, + ) + .bind(SESSION_ID, SESSION_ID, SESSION_ID, SESSION_ID, SESSION_ID, SESSION_ID) + .first(), + ).toEqual({ + model_call_count: 1, + model_call_created_at: usageReceipt?.created_at, + native_seq: nativeResumeSeq, + title_seq: titleSeq, + usage_event_count: 1, + usage_event_created_at: usageReceipt?.created_at, + }); + const hydrated = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }); + expect(hydrated.taskSnapshot).toEqual({ + driverInstanceId: DRIVER_ID, + runId: RUN_ID, + tasks: [{ taskId: "canonical-task", title: "Inspect" }], + }); + expect(hydrated.commands).toEqual([{ description: "Inspect the workspace", name: "inspect" }]); + expect(hydrated.configOptions).toEqual([ + { + currentValue: "brief", + id: "tone", + name: "Tone", + type: "select", + values: [{ name: "Brief", value: "brief" }], + }, + ]); + expect(hydrated.currentModeId).toBe("plan"); + expect(hydrated.visibleModes).toEqual([{ id: "plan", name: "Plan" }]); + expect(hydrated.usage).toMatchObject({ source: "session_update", totalTokens: 42 }); + expect(hydrated.plan).toEqual([ + { content: "Verify durable state", priority: "high", status: "completed" }, + ]); + expect(hydrated.messages.find((message) => message.id === userMessageId)).toMatchObject({ + content: "canonical question", + role: "user", + }); + expect(hydrated.messages.find((message) => message.id === thoughtId)?.segments).toEqual([ + { kind: "reasoning", text: "canonical reasoning" }, + ]); + expect(hydrated.messages.find((message) => message.id === messageId)).toMatchObject({ + content: "canonical state", + role: "assistant", + segments: expect.arrayContaining([ + { + argsText: '{"command":"pwd"}', + kind: "tool_use", + path: null, + runId: RUN_ID, + tool: "Shell", + toolCallId: "stable-tool", + }, + { + kind: "tool_result", + output: "canonical output", + runId: RUN_ID, + tool: "Shell", + toolCallId: "stable-tool", + }, + ]), + }); + const orphanResultMessageId = createRuntimeToolResultMessageId({ + runId: RUN_ID, + toolCallId: "orphan-tool", + }); + expect( + hydrated.messages.find((message) => message.id === orphanResultMessageId)?.segments, + ).toEqual([ + { + kind: "tool_result", + output: "orphan output", + runId: RUN_ID, + tool: "Orphan tool", + toolCallId: "orphan-tool", + }, + ]); + + await expect( + createController(bindings).handlePushEvents( + { + driverInstanceId: DRIVER_ID, + events: [ + runtimeEvent({ + kind: "agent.tasks.replaced", + payload: { tasks: [{ taskId: "non-canonical-retry-payload" }] }, + sourceEventId: taskSourceEventId, + }), + ], + }, + activeContext, ), + ).rejects.toThrow("conflicts with its durable receipt"); + + await database + .prepare("UPDATE session_run SET completed_at = ?, status = 'completed' WHERE id = ?") + .bind(Date.now(), RUN_ID) + .run(); + await database + .prepare("UPDATE session SET status = 'IDLE' WHERE id = ?") + .bind(SESSION_ID) + .run(); + const terminalHydrated = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }); + expect({ + commands: terminalHydrated.commands, + configOptions: terminalHydrated.configOptions, + currentModeId: terminalHydrated.currentModeId, + plan: terminalHydrated.plan, + usage: terminalHydrated.usage, + visibleModes: terminalHydrated.visibleModes, + }).toEqual({ + commands: hydrated.commands, + configOptions: hydrated.configOptions, + currentModeId: hydrated.currentModeId, + plan: hydrated.plan, + usage: hydrated.usage, + visibleModes: hydrated.visibleModes, + }); + }); + + test("rolls back receipts and every side effect when a durable projection fails", async () => { + const projectionPatterns = [ + /SET auto_title_event_seq =/iu, + /INSERT INTO ["`]?session_model_call/iu, + /INSERT INTO native_resume_ref/iu, ]; - const finalBatches = splitIntoBatches(finalStreamEvents, 50); - for (const batch of finalBatches.slice(0, -1)) { - expect(await pushFreshController(bindings, batch)).toHaveLength(batch.length); + for (const pattern of projectionPatterns) { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + await database.prepare("UPDATE session SET title = NULL WHERE id = ?").bind(SESSION_ID).run(); + const events = [ + runtimeEvent({ + kind: "session.info.updated", + payload: { title: "Must roll back" }, + sourceEventId: "atomic-rollback:title", + }), + runtimeEvent({ + kind: "usage.updated", + payload: { + inputTokens: 8, + outputTokens: 2, + source: "prompt_response", + usageContract: "openai_total_with_cached_breakdown", + }, + sourceEventId: "atomic-rollback:usage", + }), + runtimeEvent({ + kind: "runtime.resume.updated", + payload: { resumePointer: "must-roll-back" }, + sourceEventId: "atomic-rollback:native-resume", + }), + ]; + + await expect( + pushFreshController( + { + ...createPublicHttpTestBindings(database), + DB: failPreparedProjection(database, pattern), + } as ApiBindings, + events, + ), + ).rejects.toThrow("injected durable side-effect projection failure"); + + expect( + await database + .prepare( + `SELECT runtime_event_seq_cursor, title, + (SELECT COUNT(*) FROM session_event) AS receipt_count, + (SELECT COUNT(*) FROM session_model_call) AS model_call_count, + (SELECT COUNT(*) FROM usage_event) AS usage_event_count, + (SELECT COUNT(*) FROM native_resume_ref) AS native_ref_count + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(), + ).toEqual({ + model_call_count: 0, + native_ref_count: 0, + receipt_count: 0, + runtime_event_seq_cursor: 0, + title: null, + usage_event_count: 0, + }); } + }); - const terminalBatch = [ - ...(finalBatches.at(-1) ?? []), + test("finalizes durable usage only with the terminal Run commit", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const events = [ runtimeEvent({ - kind: "message.completed", - payload: { messageId: finalMessageId, role: "agent" }, - sourceEventId: "canary:final:completed", + kind: "usage.updated", + payload: { + inputTokens: 13, + outputTokens: 3, + source: "prompt_response", + usageContract: "openai_total_with_cached_breakdown", + }, + sourceEventId: "terminal-usage:usage", }), runtimeEvent({ kind: "run.completed", - payload: { - finalMessageId, - finalMessageText: FINAL_TEXT, - stopReason: "end_turn", - }, - sourceEventId: TERMINAL_SOURCE_EVENT_ID, + payload: { stopReason: "end_turn" }, + sourceEventId: "terminal-usage:completed", }), ]; - const failingBindings = { - ...bindings, - DB: failTerminalSessionEventInsert(database), - } as ApiBindings; - await expect(pushFreshController(failingBindings, terminalBatch)).rejects.toBeInstanceOf(Error); + await expect( + pushFreshController( + { + ...createPublicHttpTestBindings(database), + DB: failPreparedProjection(database, /completed-run:model-calls/iu), + } as ApiBindings, + events, + ), + ).rejects.toThrow("injected durable side-effect projection failure"); + expect( + await database + .prepare( + `SELECT r.status AS run_status, s.status AS session_status, + model_call.completed_at, model_call.status AS model_call_status, + (SELECT COUNT(*) FROM session_event + WHERE session_id = ? + AND event_type IN ('run.cancelled', 'run.completed', 'run.failed')) + AS terminal_count + FROM session_run AS r + JOIN session AS s ON s.id = r.session_id + JOIN session_model_call AS model_call ON model_call.session_run_id = r.id + WHERE r.id = ?`, + ) + .bind(SESSION_ID, RUN_ID) + .first(), + ).toEqual({ + completed_at: null, + model_call_status: "started", + run_status: "running", + session_status: "RUNNING", + terminal_count: 0, + }); - const completedRun = await database - .prepare("SELECT status FROM session_run WHERE id = ?") - .bind(RUN_ID) - .first<{ status: string }>(); - const projectedMessagesBeforeReplay = await database - .prepare("SELECT content_text, id FROM session_message WHERE session_run_id = ? ORDER BY seq") - .bind(RUN_ID) - .all<{ content_text: string; id: string }>(); - const finalOutputBeforeReplay = await readPublicThreadRunFinalOutput({ - database, - runId: RUN_ID, - sessionId: SESSION_ID, + await expect( + pushFreshController(createPublicHttpTestBindings(database) as ApiBindings, events), + ).resolves.toHaveLength(2); + await expect( + pushFreshController(createPublicHttpTestBindings(database) as ApiBindings, events), + ).resolves.toHaveLength(2); + expect( + await database + .prepare( + `SELECT r.status AS run_status, model_call.completed_at, + model_call.status AS model_call_status, + (SELECT COUNT(*) FROM session_event + WHERE session_id = ? AND event_type = 'run.completed') AS terminal_count + FROM session_run AS r + JOIN session_model_call AS model_call ON model_call.session_run_id = r.id + WHERE r.id = ?`, + ) + .bind(SESSION_ID, RUN_ID) + .first(), + ).toEqual({ + completed_at: expect.any(Number), + model_call_status: "completed", + run_status: "completed", + terminal_count: 1, }); - const terminalRowsBeforeReplay = await database - .prepare("SELECT source_event_id FROM session_event WHERE source_event_id = ?") - .bind(TERMINAL_SOURCE_EVENT_ID) - .all<{ source_event_id: string }>(); - expect(completedRun?.status).toBe("completed"); - expect(projectedMessagesBeforeReplay.results).toEqual([ - { content_text: FINAL_TEXT, id: finalMessageId }, - ]); - expect(finalOutputBeforeReplay?.text).toBe(FINAL_TEXT); - expect(new TextEncoder().encode(finalOutputBeforeReplay?.text)).toEqual( - new TextEncoder().encode(FINAL_TEXT), + await database + .prepare( + `UPDATE session_model_call + SET completed_at = NULL, status = 'started' + WHERE session_id = ? AND session_run_id = ?`, + ) + .bind(SESSION_ID, RUN_ID) + .run(); + const terminalEnvelope = events[1]; + + if (terminalEnvelope === undefined) { + throw new Error("Missing terminal usage test event."); + } + const canonicalTerminal = canonicalizeDriverEventEnvelope( + { event: terminalEnvelope.event, eventId: terminalEnvelope.eventId }, + { traceId: "trace-canary" }, ); - expect(terminalRowsBeforeReplay.results).toEqual([]); + await expect( + commitTerminalRunProjection(database, { + assistantMessage: null, + error: null, + runId: RUN_ID, + sessionId: SESSION_ID, + source: "driver", + targetStatus: "completed", + terminalEvent: { + event: canonicalTerminal.event, + occurredAt: Date.parse(canonicalTerminal.event.occurredAt), + sourceEventId: canonicalTerminal.event.sourceEventId ?? null, + }, + }), + ).resolves.toMatchObject({ kind: "duplicate" }); + expect( + await database + .prepare( + `SELECT completed_at, status + FROM session_model_call + WHERE session_id = ? AND session_run_id = ?`, + ) + .bind(SESSION_ID, RUN_ID) + .first(), + ).toEqual({ completed_at: expect.any(Number), status: "completed" }); + }); - const replayedFinalMessageId = createPlatformId(); - const crossBootTerminalBatch = [ - ...messageEvents({ - messageId: replayedFinalMessageId, - sourcePrefix: "canary:reconnected-final", - text: FINAL_TEXT, + test("fails closed when active hydration sees a private row in a public stream", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const messageId = createPlatformId(); + + await pushFreshController(bindings, [ + runtimeEvent({ + kind: "message.added", + payload: { content: "public", messageId, role: "agent" }, + sourceEventId: "active-mixed-visibility:public", + }), + runtimeEvent({ + kind: "message.delta", + payload: { contentDelta: "private", messageId, role: "agent" }, + sourceEventId: "active-mixed-visibility:private", + visibility: "owner_debug", }), + ]); + + await expect( + loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }), + ).rejects.toThrow("mixed-visibility"); + }); + + test("rejects an event batch after another Driver connection takes ownership", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + await database + .prepare("UPDATE session_run SET started_at = NULL, status = 'booting' WHERE id = ?") + .bind(RUN_ID) + .run(); + await database.prepare("UPDATE session SET title = NULL WHERE id = ?").bind(SESSION_ID).run(); + let replaced = false; + const replacedDatabase = new Proxy(database, { + get(target, property) { + if (property === "batch") { + return async (statements: D1PreparedStatement[]) => { + if (!replaced) { + replaced = true; + await target + .prepare("UPDATE driver_instance SET connection_id = ? WHERE id = ?") + .bind("replacement-connection", DRIVER_ID) + .run(); + } + return target.batch(statements); + }; + } + const value = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + const bindings = createPublicHttpTestBindings(replacedDatabase) as ApiBindings; + const events = [ + runtimeEvent({ + kind: "run.started", + payload: { startedAt: new Date(1).toISOString() }, + sourceEventId: "superseded-driver:run-started", + }), + runtimeEvent({ + kind: "message.added", + payload: { + content: "must not persist", + messageId: createPlatformId(), + role: "agent", + }, + sourceEventId: "superseded-driver:message", + }), + runtimeEvent({ + kind: "session.info.updated", + payload: { title: "Must not persist" }, + sourceEventId: "superseded-driver:title", + }), + ]; + + await expect(pushFreshController(bindings, events)).rejects.toThrow( + "lost its atomic session or active-run fence", + ); + expect(replaced).toBe(true); + expect( + await database + .prepare("SELECT runtime_event_seq_cursor FROM session WHERE id = ?") + .bind(SESSION_ID) + .first(), + ).toEqual({ runtime_event_seq_cursor: 0 }); + expect( + await database + .prepare( + `SELECT COUNT(*) AS count FROM session_event WHERE source_event_id IN (${events + .map(() => "?") + .join(", ")})`, + ) + .bind(...events.map((event) => event.eventId)) + .first(), + ).toEqual({ count: 0 }); + expect( + await database + .prepare("SELECT started_at, status FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first(), + ).toEqual({ started_at: null, status: "booting" }); + expect( + await database.prepare("SELECT title FROM session WHERE id = ?").bind(SESSION_ID).first(), + ).toEqual({ title: null }); + }); + + test("commits run.started and its Run transition in one receipt batch", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + await database + .prepare( + "UPDATE session_run SET started_at = NULL, status = 'booting', status_seq = 0 WHERE id = ?", + ) + .bind(RUN_ID) + .run(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const event = runtimeEvent({ + kind: "run.started", + payload: { startedAt: new Date(1).toISOString() }, + sourceEventId: "atomic-run-started", + }); + + await pushFreshController(bindings, [event]); + await pushFreshController(bindings, [event]); + + expect( + await database + .prepare( + `SELECT r.started_at, r.status, r.status_seq, s.runtime_event_seq_cursor + FROM session_run AS r + INNER JOIN session AS s ON s.id = r.session_id + WHERE r.id = ?`, + ) + .bind(RUN_ID) + .first(), + ).toEqual({ + runtime_event_seq_cursor: 1, + started_at: 1, + status: "running", + status_seq: 1, + }); + }); + + test("fences a runless Session event to the exact Driver connection", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + await database + .prepare("UPDATE session_run SET completed_at = 2, status = 'completed' WHERE id = ?") + .bind(RUN_ID) + .run(); + await database + .prepare("UPDATE session SET status = 'IDLE', title = NULL WHERE id = ?") + .bind(SESSION_ID) + .run(); + let replaced = false; + const replacedDatabase = new Proxy(database, { + get(target, property) { + if (property === "batch") { + return async (statements: D1PreparedStatement[]) => { + if (!replaced) { + replaced = true; + await target + .prepare("UPDATE driver_instance SET connection_id = ? WHERE id = ?") + .bind("replacement-connection", DRIVER_ID) + .run(); + } + return target.batch(statements); + }; + } + const value = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + const bindings = createPublicHttpTestBindings(replacedDatabase) as ApiBindings; + const event = runtimeEvent({ + kind: "session.info.updated", + payload: { title: "Must not persist" }, + runId: null, + sourceEventId: "superseded-driver:runless-title", + }); + + await expect(pushFreshController(bindings, [event])).rejects.toThrow( + "lost its atomic session or active-run fence", + ); + expect( + await database + .prepare("SELECT runtime_event_seq_cursor, title FROM session WHERE id = ?") + .bind(SESSION_ID) + .first(), + ).toEqual({ runtime_event_seq_cursor: 0, title: null }); + }); + + test("does not complete a Run after another Driver connection takes ownership", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const healthyBindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + await pushFreshController( + healthyBindings, + messageEvents({ + messageId: finalMessageId, + sourcePrefix: "superseded-terminal:message", + text: "Sealed before reconnect.", + }), + ); + + let replaced = false; + const replacedDatabase = new Proxy(database, { + get(target, property) { + if (property === "batch") { + return async (statements: D1PreparedStatement[]) => { + if (!replaced) { + replaced = true; + await target + .prepare("UPDATE driver_instance SET connection_id = ? WHERE id = ?") + .bind("replacement-connection", DRIVER_ID) + .run(); + } + return target.batch(statements); + }; + } + const value = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + const bindings = createPublicHttpTestBindings(replacedDatabase) as ApiBindings; + + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ]), + ).rejects.toThrow(); + expect(replaced).toBe(true); + expect( + await database.prepare("SELECT status FROM session_run WHERE id = ?").bind(RUN_ID).first(), + ).toEqual({ status: "running" }); + expect( + await database + .prepare("SELECT COUNT(*) AS count FROM session_event WHERE source_event_id = ?") + .bind(TERMINAL_SOURCE_EVENT_ID) + .first(), + ).toEqual({ count: 0 }); + }); + + test.each([ + { + artifactSandbox: { resolveError: new Error("injected missing sandbox") }, + name: "sandbox resolution", + }, + { + name: "conversation lookup", + removeConversation: true, + }, + { + artifactSandbox: { listError: new Error("injected output listing failure") }, + name: "output listing", + }, + { + artifactSandbox: { + files: new Map([["report.txt", "uncommitted"]]), + readError: new Error("injected output read failure"), + }, + name: "output read", + }, + { + artifactSandbox: { + files: new Map([ + ["valid.txt", "capturable"], + ["nested/tab\tline\n.txt", "unrepresentable"], + ]), + }, + name: "output filename", + }, + ] as const)("commits terminal authority when artifact $name is unavailable", async (failure) => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings( + database, + failure.artifactSandbox === undefined ? {} : { artifactSandbox: failure.artifactSandbox }, + ) as ApiBindings; + const finalMessageId = createPlatformId(); + await pushFreshController( + bindings, + messageEvents({ + messageId: finalMessageId, + sourcePrefix: `artifact-failure:${failure.name}:message`, + text: "This answer must remain canonical.", + }), + ); + if ("removeConversation" in failure) { + await database + .prepare("DELETE FROM sandbox_session WHERE session_id = ?") + .bind(SESSION_ID) + .run(); + } + const sourceEventId = `artifact-failure:${failure.name}:run-completed`; + + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId, + }), + ]), + ).resolves.toHaveLength(1); + expect(await readRuntimeArtifactManifest(database, TERMINAL_SOURCE_EVENT_ID)).toMatchObject({ + captureStatus: "omitted_runtime_unavailable", + files: [], + mode: "snapshot", + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }); + expect( + await database + .prepare("SELECT COUNT(*) AS count FROM session_artifact_head WHERE session_id = ?") + .bind(SESSION_ID) + .first(), + ).toEqual({ count: 0 }); + expect( + await database.prepare("SELECT status FROM session_run WHERE id = ?").bind(RUN_ID).first(), + ).toEqual({ status: "completed" }); + }); + + test("accepts a mixed message and unavailable file change before a later terminal event", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const files = new Map([["report.txt", "existing artifact"]]); + const bucket = new PublicApiMemoryFileBucket(); + const healthyBindings = createPublicHttpTestBindings(database, { + artifactSandbox: { files }, + fileBucket: bucket, + }) as ApiBindings; + await pushFreshController(healthyBindings, [ + runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "upsert", path: "outputs/report.txt" }] }, + sourceEventId: "artifact-unavailable:existing", + }), + ]); + const headBefore = await database + .prepare( + `SELECT file_id, runtime_event_seq, source_event_id + FROM session_artifact_head + WHERE session_id = ? AND source_path = 'outputs/report.txt'`, + ) + .bind(SESSION_ID) + .first(); + const messageId = createPlatformId(); + const messageSourceEventId = "artifact-unavailable:message"; + const artifactSourceEventId = "artifact-unavailable:file"; + files.set("report.txt", "must not replace the existing head"); + + await expect( + pushFreshController( + createPublicHttpTestBindings(database, { + artifactSandbox: { + files, + readError: new Error("injected unavailable artifact content"), + }, + fileBucket: bucket, + }) as ApiBindings, + [ + runtimeEvent({ + kind: "message.added", + payload: { content: "Canonical despite artifact loss.", messageId, role: "agent" }, + sourceEventId: messageSourceEventId, + }), + runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "upsert", path: "outputs/report.txt" }] }, + sourceEventId: artifactSourceEventId, + }), + ], + ), + ).resolves.toHaveLength(2); + + expect(await readRuntimeArtifactManifest(database, artifactSourceEventId)).toMatchObject({ + captureStatus: "omitted_runtime_unavailable", + files: [], + mode: "delta", + }); + expect( + await database + .prepare("SELECT COUNT(*) AS count FROM session_event WHERE source_event_id = ?") + .bind(messageSourceEventId) + .first(), + ).toEqual({ count: 1 }); + expect( + await database + .prepare( + `SELECT file_id, runtime_event_seq, source_event_id + FROM session_artifact_head + WHERE session_id = ? AND source_path = 'outputs/report.txt'`, + ) + .bind(SESSION_ID) + .first(), + ).toEqual(headBefore); + + await expect( + pushFreshController(healthyBindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { stopReason: "end_turn" }, + sourceEventId: "artifact-unavailable:terminal", + }), + ]), + ).resolves.toHaveLength(1); + expect( + await database.prepare("SELECT status FROM session_run WHERE id = ?").bind(RUN_ID).first(), + ).toEqual({ status: "completed" }); + }); + + test.each(["create", "claim", "put", "seal", "commit"] as const)( + "converges after artifact %s acknowledgement loss and collects losers", + async (phase) => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const sourceEventId = `artifact-ack-loss:${phase}`; + const files = new Map([["report.txt", "before-receipt"]]); + const bucket = new PublicApiMemoryFileBucket(); + const databaseFailure = + phase === "commit" + ? failAfterRuntimeArtifactCommit(database, sourceEventId) + : phase === "put" + ? { database, wasInjected: () => false } + : failAfterRuntimeArtifactDatabaseWrite(database, phase); + const putFailure = + phase === "put" + ? failAfterRuntimeArtifactPut(bucket) + : { bucket: bucket as unknown as R2Bucket, wasInjected: () => false }; + const event = runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "upsert", path: "outputs/report.txt" }] }, + sourceEventId, + }); + + const firstPush = pushFreshController( + createPublicHttpTestBindings(databaseFailure.database, { + artifactSandbox: { files }, + fileBucket: putFailure.bucket, + }) as ApiBindings, + [event], + ); + if (phase === "put") { + await expect(firstPush).resolves.toHaveLength(1); + } else { + await expect(firstPush).rejects.toThrow(`injected artifact ${phase} acknowledgement loss`); + } + expect(databaseFailure.wasInjected() || putFailure.wasInjected()).toBe(true); + + files.set("report.txt", "canonical-after-retry"); + const replayReads: string[] = []; + await expect( + pushFreshController( + createPublicHttpTestBindings(database, { + artifactSandbox: { files, onRead: (path) => replayReads.push(path) }, + fileBucket: bucket as unknown as R2Bucket, + }) as ApiBindings, + [event], + ), + ).resolves.toHaveLength(1); + + const attempts = await database + .prepare( + `SELECT accepted_event_id, id, owned_object_keys_json, status + FROM runtime_artifact_attempt + WHERE session_id = ? AND source_event_id = ? + ORDER BY id`, + ) + .bind(SESSION_ID, sourceEventId) + .all<{ + accepted_event_id: string | null; + id: string; + owned_object_keys_json: string; + status: string; + }>(); + const accepted = attempts.results.filter((attempt) => attempt.status === "accepted"); + const losers = attempts.results.filter((attempt) => attempt.status !== "accepted"); + expect(accepted).toHaveLength(1); + expect(losers).toHaveLength(phase === "commit" ? 0 : 1); + expect( + await database + .prepare("SELECT COUNT(*) AS count FROM session_event WHERE source_event_id = ?") + .bind(sourceEventId) + .first(), + ).toEqual({ count: 1 }); + expect( + await database + .prepare( + `SELECT COUNT(*) AS count + FROM session_artifact_head + WHERE session_id = ? AND source_path = 'outputs/report.txt' + AND source_event_id = ? AND file_id IS NOT NULL`, + ) + .bind(SESSION_ID, sourceEventId) + .first(), + ).toEqual({ count: phase === "put" ? 0 : 1 }); + + const acceptedManifest = await readRuntimeArtifactManifest(database, sourceEventId); + const winnerObjectKey = acceptedManifest.files[0]?.objectKey; + expect(acceptedManifest.captureStatus).toBe( + phase === "put" ? "omitted_runtime_unavailable" : "complete", + ); + if (winnerObjectKey !== undefined) { + const winner = await bucket.get(winnerObjectKey); + expect(await winner?.text()).toBe( + phase === "commit" ? "before-receipt" : "canonical-after-retry", + ); + } + expect(replayReads).toEqual( + phase === "commit" || phase === "put" ? [] : ["/workspace/outputs/report.txt"], + ); + + if (losers.length > 0) { + await database + .prepare( + `UPDATE runtime_artifact_attempt + SET created_at = 0, expires_at = 0, updated_at = 0 + WHERE id = ? AND status IN ('staging', 'staged')`, + ) + .bind(losers[0]!.id) + .run(); + await cleanupRuntimeArtifactAttempts( + createPublicHttpTestBindings(database, { + fileBucket: bucket as unknown as R2Bucket, + }) as ApiBindings, + ); + expect( + await database + .prepare("SELECT status FROM runtime_artifact_attempt WHERE id = ?") + .bind(losers[0]!.id) + .first(), + ).toEqual({ status: "deleting" }); + if (winnerObjectKey !== undefined) { + expect(await bucket.head(winnerObjectKey)).not.toBeNull(); + } + for (const objectKey of JSON.parse(losers[0]!.owned_object_keys_json) as string[]) { + expect(await bucket.head(objectKey)).toBeNull(); + } + + await database + .prepare( + `UPDATE runtime_artifact_attempt + SET delete_after = 0, updated_at = 0 + WHERE id = ? AND status = 'deleting'`, + ) + .bind(losers[0]!.id) + .run(); + await cleanupRuntimeArtifactAttempts( + createPublicHttpTestBindings(database, { + fileBucket: bucket as unknown as R2Bucket, + }) as ApiBindings, + ); + expect( + await database + .prepare("SELECT id FROM runtime_artifact_attempt WHERE id = ?") + .bind(losers[0]!.id) + .first(), + ).toBeNull(); + } + if (winnerObjectKey !== undefined) { + expect(await bucket.head(winnerObjectKey)).not.toBeNull(); + } + }, + ); + + test("accepts a 101-file completion omission without hiding prior artifacts", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const files = new Map([["existing.txt", "previous artifact"]]); + const readPaths: string[] = []; + const bindings = createPublicHttpTestBindings(database, { + artifactSandbox: { files, onRead: (path) => readPaths.push(path) }, + fileBucket: new PublicApiMemoryFileBucket(), + }) as ApiBindings; + + await pushFreshController(bindings, [ + runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "upsert", path: "outputs/existing.txt" }] }, + sourceEventId: "artifact-quota:file-before-completion", + }), + ]); + files.clear(); + for (let index = 0; index < 101; index += 1) { + files.set(`generated-${String(index).padStart(3, "0")}.txt`, "x"); + } + const sourceEventId = "artifact-quota:101-files-completed"; + + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { stopReason: "end_turn" }, + sourceEventId, + }), + ]), + ).resolves.toHaveLength(1); + + const row = await database + .prepare( + `SELECT head.file_id, head.runtime_event_seq, attempt.status AS attempt_status, + run.status AS run_status + FROM session_event AS event + INNER JOIN runtime_artifact_attempt AS attempt + ON attempt.id = event.artifact_attempt_id + INNER JOIN session_run AS run ON run.id = event.run_id + INNER JOIN session_artifact_head AS head + ON head.session_id = event.session_id + AND head.source_path = 'outputs/existing.txt' + WHERE event.source_event_id = ?`, + ) + .bind(TERMINAL_SOURCE_EVENT_ID) + .first<{ + attempt_status: string; + file_id: string | null; + run_status: string; + runtime_event_seq: number; + }>(); + expect(await readRuntimeArtifactManifest(database, TERMINAL_SOURCE_EVENT_ID)).toMatchObject({ + captureStatus: "omitted_file_limit", + files: [], + mode: "snapshot", + }); + expect(row).toMatchObject({ + attempt_status: "accepted", + file_id: expect.any(String), + run_status: "completed", + runtime_event_seq: 1, + }); + expect(readPaths).toEqual(["/workspace/outputs/existing.txt"]); + }); + + test.each([ + { + fileSizes: new Map([["large.bin", RUNTIME_SESSION_OUTPUT_MAX_FILE_BYTES + 1]]), + files: new Map([["large.bin", "must not be read"]]), + name: "single-file limit", + }, + { + fileSizes: new Map([ + ["one.bin", Math.floor(RUNTIME_SESSION_OUTPUT_MAX_TOTAL_BYTES / 2) + 1], + ["two.bin", Math.floor(RUNTIME_SESSION_OUTPUT_MAX_TOTAL_BYTES / 2) + 1], + ]), + files: new Map([ + ["one.bin", "must not be read"], + ["two.bin", "must not be read"], + ]), + name: "total-byte limit", + }, + ])("accepts a completion omitted by the artifact $name", async ({ fileSizes, files, name }) => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database, { + artifactSandbox: { + fileSizes, + files, + readError: new Error("quota omission must not read content"), + }, + }) as ApiBindings; + const sourceEventId = `artifact-quota:${name}`; + + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { stopReason: "end_turn" }, + sourceEventId, + }), + ]), + ).resolves.toHaveLength(1); + + const row = await database + .prepare( + `SELECT attempt.status AS attempt_status, run.status AS run_status + FROM session_event AS event + INNER JOIN runtime_artifact_attempt AS attempt + ON attempt.id = event.artifact_attempt_id + INNER JOIN session_run AS run ON run.id = event.run_id + WHERE event.source_event_id = ?`, + ) + .bind(TERMINAL_SOURCE_EVENT_ID) + .first<{ + attempt_status: string; + run_status: string; + }>(); + expect(await readRuntimeArtifactManifest(database, TERMINAL_SOURCE_EVENT_ID)).toMatchObject({ + captureStatus: "omitted_size_limit", + files: [], + mode: "snapshot", + }); + expect(row).toMatchObject({ attempt_status: "accepted", run_status: "completed" }); + }); + + test("accepts an oversized file change before the terminal event", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const files = new Map([["large.bin", "must not be read"]]); + const fileSizes = new Map([["large.bin", RUNTIME_SESSION_OUTPUT_MAX_FILE_BYTES + 1]]); + const bindings = createPublicHttpTestBindings(database, { + artifactSandbox: { + fileSizes, + files, + readError: new Error("quota omission must not read content"), + }, + }) as ApiBindings; + const deltaSourceEventId = "artifact-quota:oversized-delta"; + + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "upsert", path: "outputs/large.bin" }] }, + sourceEventId: deltaSourceEventId, + }), + ]), + ).resolves.toHaveLength(1); + files.clear(); + fileSizes.clear(); + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { stopReason: "end_turn" }, + sourceEventId: "artifact-quota:terminal-after-oversized-delta", + }), + ]), + ).resolves.toHaveLength(1); + + expect(await readRuntimeArtifactManifest(database, deltaSourceEventId)).toMatchObject({ + captureStatus: "omitted_size_limit", + files: [], + mode: "delta", + }); + expect( + await database.prepare("SELECT status FROM session_run WHERE id = ?").bind(RUN_ID).first(), + ).toEqual({ status: "completed" }); + }); + + test("deduplicates a missing upsert before the later delete does sandbox I/O", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database, { + artifactSandbox: { files: new Map() }, + }) as ApiBindings; + const events = [ + runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "upsert", path: "outputs/removed.txt" }] }, + sourceEventId: "artifact-missing:upsert", + }), + runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "delete", path: "outputs/removed.txt" }] }, + sourceEventId: "artifact-missing:delete", + }), + ]; + + await expect(pushFreshController(bindings, events)).resolves.toHaveLength(2); + expect( + await database + .prepare( + `SELECT file_id, source_event_id + FROM session_artifact_head + WHERE session_id = ? AND source_path = 'outputs/removed.txt'`, + ) + .bind(SESSION_ID) + .first(), + ).toEqual({ file_id: null, source_event_id: "artifact-missing:delete" }); + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { stopReason: "end_turn" }, + sourceEventId: "artifact-missing:terminal-after-delete", + }), + ]), + ).resolves.toHaveLength(1); + + expect( + await database + .prepare( + `SELECT file_id, source_event_id + FROM session_artifact_head + WHERE session_id = ? AND source_path = 'outputs/removed.txt'`, + ) + .bind(SESSION_ID) + .first(), + ).toEqual({ file_id: null, source_event_id: TERMINAL_SOURCE_EVENT_ID }); + }); + + test("accepts a file upsert whose historical sandbox source is gone", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database, { + artifactSandbox: { files: new Map() }, + }) as ApiBindings; + const sourceEventId = "artifact-missing:historical-upsert"; + + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "upsert", path: "outputs/gone.txt" }] }, + sourceEventId, + }), + ]), + ).resolves.toHaveLength(1); + expect(await readRuntimeArtifactManifest(database, sourceEventId)).toMatchObject({ + captureStatus: "omitted_source_missing", + files: [], + mode: "delta", + }); + + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { stopReason: "end_turn" }, + sourceEventId: "artifact-missing:terminal-after-historical-upsert", + }), + ]), + ).resolves.toHaveLength(1); + expect( + await database.prepare("SELECT status FROM session_run WHERE id = ?").bind(RUN_ID).first(), + ).toEqual({ status: "completed" }); + }); + + test("bounds a file that grows beyond quota between stat and read", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const files = new Map([["growing.bin", "x".repeat(RUNTIME_SESSION_OUTPUT_MAX_FILE_BYTES + 1)]]); + const bindings = createPublicHttpTestBindings(database, { + artifactSandbox: { + fileSizes: new Map([["growing.bin", 1]]), + files, + }, + }) as ApiBindings; + + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { stopReason: "end_turn" }, + sourceEventId: "artifact-quota:growing-file", + }), + ]), + ).resolves.toHaveLength(1); + + expect(await readRuntimeArtifactManifest(database, TERMINAL_SOURCE_EVENT_ID)).toMatchObject({ + captureStatus: "omitted_size_limit", + files: [], + mode: "snapshot", + }); + }); + + test("omits a source that changes below quota between stat and read", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const files = new Map([["report.txt", "old"]]); + const fileSizes = new Map(); + const bindings = createPublicHttpTestBindings(database, { + artifactSandbox: { fileSizes, files }, + fileBucket: new PublicApiMemoryFileBucket(), + }) as ApiBindings; + await pushFreshController(bindings, [ + runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "upsert", path: "outputs/report.txt" }] }, + sourceEventId: "artifact-changed:existing", + }), + ]); + const headBefore = await database + .prepare( + `SELECT file_id, runtime_event_seq, source_event_id + FROM session_artifact_head + WHERE session_id = ? AND source_path = 'outputs/report.txt'`, + ) + .bind(SESSION_ID) + .first(); + files.set("report.txt", "x"); + fileSizes.set("report.txt", 2); + const sourceEventId = "artifact-changed:delta"; + + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "upsert", path: "outputs/report.txt" }] }, + sourceEventId, + }), + ]), + ).resolves.toHaveLength(1); + expect(await readRuntimeArtifactManifest(database, sourceEventId)).toMatchObject({ + captureStatus: "omitted_source_changed", + files: [], + mode: "delta", + }); + expect( + await database + .prepare( + `SELECT file_id, runtime_event_seq, source_event_id + FROM session_artifact_head + WHERE session_id = ? AND source_path = 'outputs/report.txt'`, + ) + .bind(SESSION_ID) + .first(), + ).toEqual(headBefore); + + fileSizes.clear(); + await expect( + pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { stopReason: "end_turn" }, + sourceEventId: "artifact-changed:terminal", + }), + ]), + ).resolves.toHaveLength(1); + expect( + await database.prepare("SELECT status FROM session_run WHERE id = ?").bind(RUN_ID).first(), + ).toEqual({ status: "completed" }); + }); + + test("omits an artifact delta batch when a later same-path capture is unavailable", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const files = new Map([["report.txt", "existing"]]); + const bucket = new PublicApiMemoryFileBucket(); + const healthyBindings = createPublicHttpTestBindings(database, { + artifactSandbox: { files }, + fileBucket: bucket, + }) as ApiBindings; + await pushFreshController(healthyBindings, [ + runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "upsert", path: "outputs/report.txt" }] }, + sourceEventId: "artifact-batch-unavailable:existing", + }), + ]); + const headBefore = await database + .prepare( + `SELECT file_id, runtime_event_seq, source_event_id + FROM session_artifact_head + WHERE session_id = ? AND source_path = 'outputs/report.txt'`, + ) + .bind(SESSION_ID) + .first(); + const sourceEventIds = [ + "artifact-batch-unavailable:earlier", + "artifact-batch-unavailable:later", + ]; + + await expect( + pushFreshController( + createPublicHttpTestBindings(database, { + artifactSandbox: { + files, + resolveError: new Error("injected unavailable later capture"), + }, + fileBucket: bucket, + }) as ApiBindings, + sourceEventIds.map((sourceEventId) => + runtimeEvent({ + kind: "file.changed", + payload: { changes: [{ change: "upsert", path: "outputs/report.txt" }] }, + sourceEventId, + }), + ), + ), + ).resolves.toHaveLength(2); + expect( + await Promise.all(sourceEventIds.map((id) => readRuntimeArtifactManifest(database, id))), + ).toEqual([ + expect.objectContaining({ + captureStatus: "omitted_runtime_unavailable", + files: [], + mode: "delta", + }), + expect.objectContaining({ + captureStatus: "omitted_runtime_unavailable", + files: [], + mode: "delta", + }), + ]); + expect( + await database + .prepare( + `SELECT file_id, runtime_event_seq, source_event_id + FROM session_artifact_head + WHERE session_id = ? AND source_path = 'outputs/report.txt'`, + ) + .bind(SESSION_ID) + .first(), + ).toEqual(headBefore); + }); + + test("rejects fresh run events after terminal repair already won", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const failureSourceEventId = `session-run-terminal:${RUN_ID}:run.failed`; + await recordCanonicalSessionRunFailure(bindings, { + error: { + code: "runtime.driver_terminal", + details: {}, + message: "Driver disconnected.", + retryable: true, + }, + runId: RUN_ID, + sessionId: SESSION_ID, + source: "api", + }); + await expect( + loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }), + ).resolves.toMatchObject({ + run: { error: { retryable: true } }, + }); + const messageId = createPlatformId(); + const delta = runtimeEvent({ + kind: "message.delta", + payload: { contentDelta: "durable prefix", messageId, role: "agent" }, + sourceEventId: "recovered-message-delta", + }); + const failed = runtimeEvent({ + kind: "run.failed", + payload: { + error: { + code: "runtime.driver_terminal", + details: {}, + message: "Driver disconnected.", + retryable: true, + }, + recoverable: true, + }, + sourceEventId: "driver-run-failed", + }); + + await expect( + createController(bindings).handlePushEvents( + { driverInstanceId: DRIVER_ID, events: [delta, failed] }, + activeContext, + ), + ).rejects.toThrow("conflicts with its durable receipt"); + + const rows = await database + .prepare( + `SELECT event_type, source_event_id + FROM session_event + WHERE source_event_id IN (?, ?) + ORDER BY seq`, + ) + .bind(delta.eventId, failureSourceEventId) + .all<{ event_type: string; source_event_id: string }>(); + expect(rows.results).toEqual([ + { event_type: "run.failed", source_event_id: failureSourceEventId }, + ]); + }); + + test.each([ + ["in the terminal RPC", false], + ["before the terminal RPC", true], + ] as const)( + "persists one final assistant snapshot when the sealed stream arrives %s", + async (_arrival, persistBeforeTerminal) => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const capturedEvents: unknown[] = []; + setServerProductAnalyticsTransportForTests(async (_input, init) => { + capturedEvents.push(JSON.parse(init.body as string) as unknown); + return new Response(null, { status: 200 }); + }); + const bindings = { + ...createPublicHttpTestBindings(database), + POSTHOG_PROJECT_KEY: "phc_test", + } as ApiBindings; + const finalText = "The final answer."; + const finalMessageId = createPlatformId(); + const stream = messageEvents({ + messageId: finalMessageId, + sourcePrefix: "final-stream", + text: finalText, + }); + const terminal = runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: "final-stream:run-completed", + }); + + if (persistBeforeTerminal) { + await pushFreshController(bindings, stream); + await pushFreshController(bindings, [terminal]); + } else { + await pushFreshController(bindings, [...stream, terminal]); + } + + const rows = await database + .prepare( + `SELECT content_text, ended_at, event_type, id, occurred_at, process_status, + process_type, run_id, seq, stream_id, tokens + FROM session_event + WHERE session_id = ? AND run_id = ? + ORDER BY seq`, + ) + .bind(SESSION_ID, RUN_ID) + .all(); + const assistantMessages = createSessionProcessEventsFromSessionEventRows(rows.results).filter( + (event) => event.type === "agent.message.delta", + ); + + expect( + rows.results + .filter((row) => row.event_type === "message.added") + .map((row) => row.content_text), + ).toEqual([finalText]); + expect(assistantMessages.map((event) => event.content)).toEqual([finalText]); + expect(capturedEvents).toEqual([ + expect.objectContaining({ + event: "task_succeeded", + properties: expect.objectContaining({ + run_duration_ms: expect.any(Number), + sandbox_id: PUBLIC_API_TEST_IDS.sandbox, + sandbox_kind: "pet", + sandbox_subject_kind: "agent", + session_type: "ui", + }), + }), + ]); + }, + ); + + test("keeps equal text from distinct progress and final message streams", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const progressMessageId = createPlatformId(); + const finalMessageId = createPlatformId(); + const text = "The same text belongs to two messages."; + + await pushFreshController(bindings, [ + ...messageEvents({ + messageId: progressMessageId, + sourcePrefix: "equal-text:progress", + text, + }), + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: "equal-text:final", + text, + }), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: "equal-text:run-completed", + }), + ]); + + const rows = await database + .prepare( + `SELECT content_text, ended_at, event_type, id, occurred_at, process_status, + process_type, run_id, seq, stream_id, tokens + FROM session_event + WHERE session_id = ? AND run_id = ? + ORDER BY seq`, + ) + .bind(SESSION_ID, RUN_ID) + .all(); + const messageRows = rows.results.filter((row) => row.event_type === "message.added"); + const messages = createSessionProcessEventsFromSessionEventRows(rows.results).filter( + (event) => event.type === "agent.message.delta", + ); + + expect(messageRows.map((row) => [row.stream_id, row.content_text])).toEqual([ + [progressMessageId, text], + [finalMessageId, text], + ]); + expect(messages.map((message) => message.content)).toEqual([text, text]); + }); + + test("atomically seals a multi-megabyte final reference across hibernation and replay", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const syncedSessionIds: string[] = []; + const bindings = createPublicHttpTestBindings(database, { + sessionNamespace: createSessionSyncNamespace(syncedSessionIds), + }) as ApiBindings; + const progressMessageIds = PROGRESS_TEXTS.map(() => createPlatformId()); + const finalMessageId = createPlatformId(); + const progressEvents = [ + runtimeEvent({ + kind: "run.started", + payload: { startedAt: new Date(1).toISOString() }, + sourceEventId: "canary:run-started", + }), + ...PROGRESS_TEXTS.flatMap((text, index) => + messageEvents({ + messageId: progressMessageIds[index], + sourcePrefix: `canary:progress:${index + 1}`, + text, + }), + ), + ]; + + expect( + (await pushFreshController(bindings, progressEvents)).map((receipt) => receipt.eventId), + ).toEqual(progressEvents.map((event) => event.eventId)); + + const toolEvents = [ + runtimeEvent({ + kind: "tool.call.updated", + payload: { + parentMessageId: finalMessageId, + rawInputDelta: '{"path":', + status: "running", + title: "Create artifact", + toolCallId: "tool-canary", + }, + sourceEventId: "canary:tool:started", + }), + runtimeEvent({ + kind: "tool.call.updated", + payload: { + rawInputDelta: '"report.md"}', + rawOutputDelta: "creating ", + status: "running", + toolCallId: "tool-canary", + }, + sourceEventId: "canary:tool:streamed", + }), + ]; + expect(await pushFreshController(bindings, toolEvents)).toHaveLength(toolEvents.length); + + const finalTextChunks = Array.from( + { length: Math.ceil(LARGE_FINAL_TEXT.length / LARGE_FINAL_TEXT_CHUNK_CHARACTERS) }, + (_, index) => + LARGE_FINAL_TEXT.slice( + index * LARGE_FINAL_TEXT_CHUNK_CHARACTERS, + (index + 1) * LARGE_FINAL_TEXT_CHUNK_CHARACTERS, + ), + ); + const finalStreamEvents = [ + runtimeEvent({ + kind: "message.started", + payload: { messageId: finalMessageId, role: "agent" }, + sourceEventId: "canary:final:started", + }), + runtimeEvent({ + kind: "message.added", + payload: { content: finalTextChunks[0], messageId: finalMessageId, role: "agent" }, + sourceEventId: "canary:final:snapshot", + }), + ...finalTextChunks.slice(1).map((contentDelta, index) => + runtimeEvent({ + kind: "message.delta", + payload: { contentDelta, messageId: finalMessageId, role: "agent" }, + sourceEventId: `canary:final:delta:${index + 2}`, + }), + ), + ]; + expect(finalStreamEvents.length).toBeGreaterThan(1_000); + const finalBatches = splitIntoBatches(finalStreamEvents, 64); + + for (const batch of finalBatches.slice(0, -1)) { + expect(await pushFreshController(bindings, batch)).toHaveLength(batch.length); + } + + const terminalBatch = [ + ...(finalBatches.at(-1) ?? []), + runtimeEvent({ + kind: "tool.call.updated", + payload: { + rawInput: '{"path":"report.md"}', + rawOutput: "artifact created", + status: "completed", + toolCallId: "tool-canary", + }, + sourceEventId: "canary:tool:completed", + }), + runtimeEvent({ + kind: "message.completed", + payload: { messageId: finalMessageId, role: "agent" }, + sourceEventId: "canary:final:completed", + }), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ]; + const failingBindings = { + ...bindings, + DB: failTerminalSessionEventInsert(database), + } as ApiBindings; + + await expect(pushFreshController(failingBindings, terminalBatch)).rejects.toBeInstanceOf(Error); + + const completedRun = await database + .prepare("SELECT status FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first<{ status: string }>(); + const projectedMessagesBeforeReplay = await database + .prepare("SELECT content_text, id FROM session_message WHERE session_run_id = ? ORDER BY seq") + .bind(RUN_ID) + .all<{ content_text: string; id: string }>(); + const finalOutputBeforeReplay = await readPublicThreadRunFinalOutput({ + database, + runId: RUN_ID, + sessionId: SESSION_ID, + }); + const terminalRowsBeforeReplay = await database + .prepare("SELECT source_event_id FROM session_event WHERE source_event_id = ?") + .bind(TERMINAL_SOURCE_EVENT_ID) + .all<{ source_event_id: string }>(); + const driverReleaseBeforeReplay = await database + .prepare("SELECT status_operation_id FROM driver_instance WHERE id = ?") + .bind(DRIVER_ID) + .first<{ status_operation_id: string | null }>(); + + expect(completedRun?.status).toBe("running"); + expect(projectedMessagesBeforeReplay.results).toEqual([]); + expect(new TextEncoder().encode(LARGE_FINAL_TEXT).byteLength).toBeGreaterThan(2_000_000); + expect(finalOutputBeforeReplay).toBeNull(); + expect(terminalRowsBeforeReplay.results).toEqual([]); + expect(driverReleaseBeforeReplay).toEqual({ status_operation_id: null }); + + const crossBootTerminalBatch = [ + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ]; + const expectedCrossBootReceiptIds = crossBootTerminalBatch.map((event) => event.eventId); + const crossBootDelivery: SessionDeliveryEvent[] = []; + const structuralTerminalBindings = { + ...bindings, + DB: rejectSessionEventContentReads(database), + } as ApiBindings; + + expect( + ( + await createController(structuralTerminalBindings, (_sessionId, events) => + crossBootDelivery.push(...events), + ).handlePushEvents( + { driverInstanceId: DRIVER_ID, events: crossBootTerminalBatch }, + activeContext, + ) + ).accepted.map((receipt) => receipt.eventId), + ).toEqual(expectedCrossBootReceiptIds); + expect(crossBootDelivery).toEqual([]); + expect(syncedSessionIds).toEqual([SESSION_ID]); + expect( + (await pushFreshController(bindings, terminalBatch)).map((receipt) => receipt.eventId), + ).toEqual(terminalBatch.map((event) => event.eventId)); + expect(syncedSessionIds).toEqual([SESSION_ID, SESSION_ID]); + await expect( + database + .prepare("SELECT status_operation_id FROM driver_instance WHERE id = ?") + .bind(DRIVER_ID) + .first(), + ).resolves.toEqual({ status_operation_id: RUN_ID }); + + const projectedMessagesAfterReplay = await database + .prepare( + "SELECT content_text, id, plan_json, projection_format, segments_json FROM session_message WHERE session_run_id = ? ORDER BY seq", + ) + .bind(RUN_ID) + .all<{ + content_text: string; + id: string; + plan_json: string | null; + projection_format: string; + segments_json: string | null; + }>(); + const terminalRowsAfterReplay = await database + .prepare("SELECT source_event_id FROM session_event WHERE source_event_id = ?") + .bind(TERMINAL_SOURCE_EVENT_ID) + .all<{ source_event_id: string }>(); + const transcript = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }); + const finalTranscriptMessage = transcript.messages.find( + (message) => message.id === finalMessageId, + ); + const canonicalTranscriptMessages = transcript.messages.filter( + (message) => message.role === "assistant" && message.content === LARGE_FINAL_TEXT, + ); + const measuredRecovery = measurePreparedQueries(database); + const recoveryMessages = await getSessionRuntimeRecoveryMessages(measuredRecovery.database, { + excludeRunId: null, + sessionId: SESSION_ID, + }); + + expect(projectedMessagesAfterReplay.results).toEqual([ + { + content_text: "", + id: finalMessageId, + plan_json: null, + projection_format: "event_stream_v3", + segments_json: null, + }, + ]); + expect(terminalRowsAfterReplay.results).toEqual([ + { source_event_id: TERMINAL_SOURCE_EVENT_ID }, + ]); + expect(finalTranscriptMessage?.content).toBe(LARGE_FINAL_TEXT); + expect(canonicalTranscriptMessages.map((message) => message.id)).toEqual([finalMessageId]); + expect(finalTranscriptMessage?.content).not.toContain(PROGRESS_TEXTS.join("")); + expect(finalTranscriptMessage?.segments.filter((segment) => segment.kind !== "text")).toEqual([ + { + argsText: '{"path":"report.md"}', + kind: "tool_use", + path: null, + runId: RUN_ID, + tool: "Create artifact", + toolCallId: "tool-canary", + }, + { + kind: "tool_result", + output: "artifact created", + runId: RUN_ID, + tool: "Create artifact", + toolCallId: "tool-canary", + }, + ]); + expect(LARGE_FINAL_TEXT.split("\n")).toContain("160|中文长文本校验-Aa9-表格字符|END160"); + expect(recoveryMessages).toEqual([ + { content: LARGE_FINAL_TEXT.slice(0, 32_000), role: "assistant" }, + ]); + expect(measuredRecovery.readCount()).toBeLessThanOrEqual(10); + }); + + test("removes provider-private citations at the public final-output boundary", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + const privateCitation = "\uE200cite\uE202turn2view0\uE202turn8view0\uE201"; + const providerText = `before${privateCitation}after`; + const events = [ + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: "private-citation:final", + text: providerText, + }), + runtimeEvent({ + kind: "run.completed", + payload: { + finalMessageId, + stopReason: "end_turn", + }, + sourceEventId: "private-citation:run-completed", + }), + ]; + + await pushFreshController(bindings, events); + + const persistedMessage = await database + .prepare("SELECT content_text FROM session_message WHERE id = ?") + .bind(finalMessageId) + .first<{ content_text: string }>(); + + expect(persistedMessage?.content_text).toBe(""); + await expect( + readPublicThreadRunFinalOutput({ database, runId: RUN_ID, sessionId: SESSION_ID }), + ).resolves.toEqual({ + text: "beforeafter", + warnings: [ + { + code: "unresolved_provider_citation", + count: 1, + }, + ], + }); + }); + + test("fails closed when a v3 terminal coexists with a legacy terminal", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + await pushFreshController(bindings, [ + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: "mixed-terminal:final", + text: FINAL_TEXT, + }), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ]); + const session = await database + .prepare("SELECT runtime_event_seq_cursor FROM session WHERE id = ?") + .bind(SESSION_ID) + .first<{ runtime_event_seq_cursor: number }>(); + if (session === null) { + throw new Error("Missing Session cursor fixture."); + } + await database + .prepare( + `UPDATE session_event + SET artifact_attempt_id = NULL, + artifact_manifest_json = NULL, + artifact_manifest_sha256 = NULL, + semantic_hash = NULL, + terminal_event_json = NULL, + source_event_id = 'legacy-terminal' + WHERE source_event_id = ?`, + ) + .bind(TERMINAL_SOURCE_EVENT_ID) + .run(); + const v3TerminalSeq = session.runtime_event_seq_cursor + 1; + await insertRuntimeEvent(database, { + eventId: createPlatformId(), + kind: "run.completed", + occurredAt: Date.now(), + payload: { finalMessageId, stopReason: "end_turn" }, + runId: RUN_ID, + seq: v3TerminalSeq, + sessionId: SESSION_ID, + }); + await database + .prepare("UPDATE session SET runtime_event_seq_cursor = ? WHERE id = ?") + .bind(v3TerminalSeq, SESSION_ID) + .run(); + + await expect( + readPublicThreadRunFinalOutput({ database, runId: RUN_ID, sessionId: SESSION_ID }), + ).rejects.toThrow("not immutable"); + await expect( + loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }), + ).rejects.toThrow("terminal"); + }); + + test("budgets v3 recovery text after streaming private-citation removal", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + const visiblePrefix = "a".repeat(31_998); + const visibleTail = "VISIBLE"; + const events = [ + runtimeEvent({ + kind: "message.added", + payload: { + content: `${visiblePrefix}\uE200ci`, + messageId: finalMessageId, + role: "agent", + }, + sourceEventId: "recovery-citation:message-added", + }), + runtimeEvent({ + kind: "message.delta", + payload: { + contentDelta: `te\uE202${"private".repeat(800)}`, + messageId: finalMessageId, + role: "agent", + }, + sourceEventId: "recovery-citation:message-delta-private", + }), + runtimeEvent({ + kind: "message.delta", + payload: { + contentDelta: `\uE201${visibleTail}`, + messageId: finalMessageId, + role: "agent", + }, + sourceEventId: "recovery-citation:message-delta-tail", + }), + runtimeEvent({ + kind: "message.completed", + payload: { messageId: finalMessageId, role: "agent" }, + sourceEventId: "recovery-citation:message-completed", + }), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ]; + + await pushFreshController(bindings, events); + + const stored = await database + .prepare("SELECT projection_format FROM session_message WHERE id = ?") + .bind(finalMessageId) + .first<{ projection_format: string }>(); + const recoveryMessages = await getSessionRuntimeRecoveryMessages(database, { + excludeRunId: null, + sessionId: SESSION_ID, + }); + + expect(stored?.projection_format).toBe("event_stream_v3"); + expect(recoveryMessages).toEqual([ + { content: `${visiblePrefix}${visibleTail.slice(0, 2)}`, role: "assistant" }, + ]); + expect(recoveryMessages[0]?.content).not.toContain("\uE200"); + }); + + test.each([1_000, 2_000])( + "seals bounded recovery at an exact %i-row page boundary", + async (streamRowCount) => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + const events = [ + runtimeEvent({ + kind: "message.added", + payload: { content: "A", messageId: finalMessageId, role: "agent" }, + sourceEventId: `exact-page:${streamRowCount}:added`, + }), + ...Array.from({ length: streamRowCount - 2 }, (_, index) => + runtimeEvent({ + kind: "message.delta", + payload: { contentDelta: "b", messageId: finalMessageId, role: "agent" }, + sourceEventId: `exact-page:${streamRowCount}:delta:${index}`, + }), + ), + runtimeEvent({ + kind: "message.completed", + payload: { messageId: finalMessageId, role: "agent" }, + sourceEventId: `exact-page:${streamRowCount}:completed`, + }), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ]; + + for (const batch of splitIntoBatches(events, 64)) { + await pushFreshController(bindings, batch); + } + + const measured = measurePreparedQueries(database); + const recoveryMessages = await getSessionRuntimeRecoveryMessages(measured.database, { + excludeRunId: null, + sessionId: SESSION_ID, + }); + + expect(recoveryMessages).toEqual([ + { content: `A${"b".repeat(streamRowCount - 2)}`, role: "assistant" }, + ]); + expect(measured.readCount()).toBeLessThanOrEqual(streamRowCount === 1_000 ? 8 : 10); + }, + ); + + test("omits live-only reasoning from stored final assistant segments", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + const privateReasoningText = "Private reasoning should stay out of stored history."; + const thoughtId = "private-reasoning"; + const events = [ + runtimeEvent({ + kind: "thought.started", + payload: { messageId: finalMessageId, thoughtId }, + sourceEventId: "reasoning:started", + }), + runtimeEvent({ + kind: "thought.delta", + payload: { contentDelta: privateReasoningText, messageId: finalMessageId, thoughtId }, + sourceEventId: "reasoning:delta", + }), + runtimeEvent({ + kind: "thought.completed", + payload: { messageId: finalMessageId, thoughtId }, + sourceEventId: "reasoning:completed", + }), + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: "reasoning:final", + text: FINAL_TEXT, + }), + runtimeEvent({ + kind: "run.completed", + payload: { + finalMessageId, + stopReason: "end_turn", + }, + sourceEventId: "reasoning:run-completed", + }), + ]; + + await pushFreshController(bindings, events); + + const persistedMessage = await database + .prepare( + "SELECT content_text, projection_format, segments_json FROM session_message WHERE id = ?", + ) + .bind(finalMessageId) + .first<{ content_text: string; projection_format: string; segments_json: string | null }>(); + + expect(persistedMessage).toEqual({ + content_text: "", + projection_format: "event_stream_v3", + segments_json: null, + }); + const transcript = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }); + const finalMessage = transcript.messages.find((message) => message.id === finalMessageId); + expect(finalMessage?.segments).toEqual([{ kind: "text", text: FINAL_TEXT }]); + expect(JSON.stringify(finalMessage?.segments)).not.toContain(privateReasoningText); + }); + + test("rehydrates the shared failed-tool fallback without explicit output", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + await pushFreshController(bindings, [ + runtimeEvent({ + kind: "message.added", + payload: { content: "Finished.", messageId: finalMessageId, role: "agent" }, + sourceEventId: "failed-tool:message", + }), + runtimeEvent({ + kind: "tool.call.updated", + payload: { + parentMessageId: finalMessageId, + status: "running", + title: "Shell", + toolCallId: "failed-tool", + }, + sourceEventId: "failed-tool:running", + }), + runtimeEvent({ + kind: "tool.call.updated", + payload: { status: "failed", toolCallId: "failed-tool" }, + sourceEventId: "failed-tool:failed", + }), + runtimeEvent({ + kind: "message.completed", + payload: { messageId: finalMessageId, role: "agent" }, + sourceEventId: "failed-tool:message-completed", + }), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ]); + + const transcript = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }); + const finalMessage = transcript.messages.find((message) => message.id === finalMessageId); + expect(finalMessage?.segments.filter((segment) => segment.kind !== "text")).toEqual([ + { + argsText: "", + kind: "tool_use", + path: null, + runId: RUN_ID, + tool: "Shell", + toolCallId: "failed-tool", + }, + { + kind: "tool_result", + output: "Tool failed.", + runId: RUN_ID, + tool: "Shell", + toolCallId: "failed-tool", + }, + ]); + }); + + test("keeps a completed run's identity-free tool result after a newer run becomes latest", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const syncedSessionIds: string[] = []; + const bindings = createPublicHttpTestBindings(database, { + sessionNamespace: createSessionSyncNamespace(syncedSessionIds), + }) as ApiBindings; + const finalMessageId = createPlatformId(); + const toolCallId = "terminal-orphan-tool"; + const resultMessageId = createRuntimeToolResultMessageId({ runId: RUN_ID, toolCallId }); + const toolOutputOccurredAt = Date.parse("2026-08-30T04:00:00.000Z"); + + await pushFreshController(bindings, [ + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: "terminal-orphan:message", + text: "Finished.", + }), + runtimeEvent({ + kind: "tool.call.updated", + payload: { status: "running", title: "Shell", toolCallId }, + sourceEventId: "terminal-orphan:running", + }), + runtimeEvent({ + kind: "tool.call.updated", + occurredAt: toolOutputOccurredAt, + payload: { rawOutputDelta: "A", status: "running", toolCallId }, + sourceEventId: "terminal-orphan:output-a", + }), + runtimeEvent({ + kind: "tool.call.updated", + occurredAt: toolOutputOccurredAt + 1, + payload: { rawOutputDelta: "B", status: "completed", toolCallId }, + sourceEventId: "terminal-orphan:output-b", + }), + ]); + + const active = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }); + const activeResult = active.messages.find((message) => message.id === resultMessageId); + + const terminalDelivery: SessionDeliveryEvent[] = []; + await createController(bindings, (_sessionId, events) => + terminalDelivery.push(...events), + ).handlePushEvents( + { + driverInstanceId: DRIVER_ID, + events: [ + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ], + }, + activeContext, + ); + expect(terminalDelivery).toEqual([]); + expect(syncedSessionIds).toEqual([SESSION_ID]); + + const terminal = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }); + expect(terminal.messages.find((message) => message.id === resultMessageId)?.segments).toEqual( + activeResult?.segments, + ); + expect(activeResult?.segments).toEqual([ + { + kind: "tool_result", + output: "AB", + runId: RUN_ID, + tool: "Tool", + toolCallId, + }, + ]); + + const nextRunId = createPlatformId(); + await database + .prepare( + `INSERT INTO session_run ( + id, session_id, agent_id, created_by_account_id, deployment_version_id, + deployment_version_number, driver_instance_id, trigger, status, provider, + model, runtime_id, trace_id, started_at, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, 1, ?, 'user_prompt', 'running', 'openai', + 'gpt-5.4', 'openai-runtime', 'trace-next', 2, 2, 2)`, + ) + .bind( + nextRunId, + SESSION_ID, + PUBLIC_API_TEST_IDS.agent, + PUBLIC_API_TEST_IDS.ownerAccount, + PUBLIC_API_TEST_IDS.deployment, + DRIVER_ID, + ) + .run(); + await database + .prepare("UPDATE session SET last_run_id = ?, status = 'RUNNING' WHERE id = ?") + .bind(nextRunId, SESSION_ID) + .run(); + + const freshNextRun = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }); + const freshResults = freshNextRun.messages.filter((message) => message.id === resultMessageId); + const carrierIndex = freshNextRun.messages.findIndex( + (message) => message.id === finalMessageId, + ); + expect(freshResults).toHaveLength(1); + expect(freshResults[0]?.segments).toEqual(activeResult?.segments); + expect(freshNextRun.messages.indexOf(freshResults[0])).toBe(carrierIndex + 1); + expect(freshResults[0]?.createdAt).toBe(new Date(toolOutputOccurredAt).toISOString()); + + const repeatedFreshNextRun = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }); + expect(repeatedFreshNextRun.messages).toEqual(freshNextRun.messages); + const publicFreshNextRun = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.outsiderAccount, + }); + expect(publicFreshNextRun.messages).toEqual(freshNextRun.messages); + }); + + test("canonicalizes nonfinal and result-first tools into one terminal carrier", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const syncedSessionIds: string[] = []; + const bindings = createPublicHttpTestBindings(database, { + sessionNamespace: createSessionSyncNamespace(syncedSessionIds), + }) as ApiBindings; + const progressMessageId = createPlatformId(); + const finalMessageId = createPlatformId(); + const finalOccurredAt = Date.parse("2026-08-30T07:00:00.000Z"); + const carrierOccurredAt = finalOccurredAt + 2; + + await pushFreshController(bindings, [ + runtimeEvent({ + kind: "message.added", + occurredAt: finalOccurredAt - 2, + payload: { content: "Progress", messageId: progressMessageId, role: "agent" }, + sourceEventId: "carrier-routes:progress", + }), + runtimeEvent({ + kind: "message.completed", + occurredAt: finalOccurredAt - 1, + payload: { messageId: progressMessageId, role: "agent" }, + sourceEventId: "carrier-routes:progress-completed", + }), + runtimeEvent({ + kind: "message.added", + occurredAt: finalOccurredAt, + payload: { content: "Final", messageId: finalMessageId, role: "agent" }, + sourceEventId: "carrier-routes:final", + }), + runtimeEvent({ + kind: "message.completed", + occurredAt: finalOccurredAt + 1, + payload: { messageId: finalMessageId, role: "agent" }, + sourceEventId: "carrier-routes:final-completed", + }), + runtimeEvent({ + kind: "tool.call.updated", + occurredAt: carrierOccurredAt, + payload: { + parentMessageId: progressMessageId, + rawInput: '{"a":1}', + status: "running", + title: "Progress tool", + toolCallId: "progress-tool", + }, + sourceEventId: "carrier-routes:progress-use", + }), + runtimeEvent({ + kind: "tool.call.updated", + occurredAt: carrierOccurredAt + 1, + payload: { rawOutput: "progress output", status: "completed", toolCallId: "progress-tool" }, + sourceEventId: "carrier-routes:progress-result", + }), + runtimeEvent({ + kind: "tool.call.updated", + occurredAt: carrierOccurredAt + 2, + payload: { + parentMessageId: progressMessageId, + rawInput: '{"b":2}', + status: "running", + title: "Use only", + toolCallId: "use-only-tool", + }, + sourceEventId: "carrier-routes:use-only", + }), + runtimeEvent({ + kind: "tool.call.updated", + occurredAt: carrierOccurredAt + 3, + payload: { + rawOutput: "result first", + status: "completed", + toolCallId: "result-first-tool", + }, + sourceEventId: "carrier-routes:result-first", + }), + runtimeEvent({ + kind: "tool.call.updated", + occurredAt: carrierOccurredAt + 4, + payload: { + parentMessageId: finalMessageId, + rawInput: '{"c":3}', + status: "completed", + title: "Result first", + toolCallId: "result-first-tool", + }, + sourceEventId: "carrier-routes:result-later-parent", + }), + runtimeEvent({ + kind: "tool.call.updated", + occurredAt: carrierOccurredAt + 5, + payload: { + parentMessageId: finalMessageId, + rawInput: '{"d":4}', + status: "running", + title: "Final tool", + toolCallId: "final-tool", + }, + sourceEventId: "carrier-routes:final-use", + }), + runtimeEvent({ + kind: "tool.call.updated", + occurredAt: carrierOccurredAt + 6, + payload: { rawOutput: "final output", status: "completed", toolCallId: "final-tool" }, + sourceEventId: "carrier-routes:final-result", + }), + ]); + + const active = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }); + expect( + active.messages.find((message) => message.id === progressMessageId)?.segments, + ).toHaveLength(4); + expect(active.messages.find((message) => message.id === RUN_ID)?.segments).toHaveLength(2); + expect(active.messages.find((message) => message.id === finalMessageId)?.segments).toHaveLength( + 3, + ); + + const terminalDelivery: SessionDeliveryEvent[] = []; + await createController(bindings, (_sessionId, events) => + terminalDelivery.push(...events), + ).handlePushEvents( + { + driverInstanceId: DRIVER_ID, + events: [ + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: "carrier-routes:terminal", + }), + ], + }, + activeContext, + ); + expect(terminalDelivery).toEqual([]); + expect(syncedSessionIds).toEqual([SESSION_ID]); + + const fresh = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.ownerAccount, + }); + expect(fresh.messages.find((message) => message.id === progressMessageId)).toBeUndefined(); + expect( + fresh.messages + .filter((message) => message.id === finalMessageId || message.id === RUN_ID) + .map(({ createdAt, id }) => ({ createdAt, id })), + ).toEqual([ + { createdAt: new Date(finalOccurredAt).toISOString(), id: finalMessageId }, + { createdAt: new Date(carrierOccurredAt).toISOString(), id: RUN_ID }, + ]); + expect(fresh.messages.find((message) => message.id === finalMessageId)?.segments).toEqual([ + { kind: "text", text: "Final" }, + { + argsText: '{"d":4}', + kind: "tool_use", + path: null, + runId: RUN_ID, + tool: "Final tool", + toolCallId: "final-tool", + }, + { + kind: "tool_result", + output: "final output", + runId: RUN_ID, + tool: "Final tool", + toolCallId: "final-tool", + }, + ]); + expect(fresh.messages.find((message) => message.id === RUN_ID)?.segments).toEqual([ + { + argsText: '{"a":1}', + kind: "tool_use", + path: null, + runId: RUN_ID, + tool: "Progress tool", + toolCallId: "progress-tool", + }, + { + kind: "tool_result", + output: "progress output", + runId: RUN_ID, + tool: "Progress tool", + toolCallId: "progress-tool", + }, + { + argsText: '{"b":2}', + kind: "tool_use", + path: null, + runId: RUN_ID, + tool: "Use only", + toolCallId: "use-only-tool", + }, + { + argsText: '{"c":3}', + kind: "tool_use", + path: null, + runId: RUN_ID, + tool: "Result first", + toolCallId: "result-first-tool", + }, + { + kind: "tool_result", + output: "result first", + runId: RUN_ID, + tool: "Result first", + toolCallId: "result-first-tool", + }, + ]); + const publicFresh = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.outsiderAccount, + }); + expect(publicFresh.messages).toEqual(fresh.messages); + }); + + for (const toolFirst of [false, true]) { + test(`merges a ${toolFirst ? "tool-before-final" : "final-before-tool"} carrier when the final id is the run id`, async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = RUN_ID as unknown as SessionMessageId; + const firstOccurredAt = Date.parse("2026-08-30T08:00:00.000Z"); + const message = runtimeEvent({ + kind: "message.added", + occurredAt: firstOccurredAt + (toolFirst ? 2 : 0), + payload: { content: "Same id final", messageId: finalMessageId, role: "agent" }, + sourceEventId: `same-id:${toolFirst}:message`, + }); + const completed = runtimeEvent({ + kind: "message.completed", + occurredAt: firstOccurredAt + (toolFirst ? 3 : 1), + payload: { messageId: finalMessageId, role: "agent" }, + sourceEventId: `same-id:${toolFirst}:message-completed`, + }); + const tool = runtimeEvent({ + kind: "tool.call.updated", + occurredAt: firstOccurredAt + (toolFirst ? 0 : 2), + payload: { + rawOutput: "same id output", + status: "completed", + toolCallId: "same-id-tool", + }, + sourceEventId: `same-id:${toolFirst}:tool`, + }); + + await pushFreshController(bindings, [ + ...(toolFirst ? [tool, message, completed] : [message, completed, tool]), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: `same-id:${toolFirst}:terminal`, + }), + ]); + + const rows = await database + .prepare( + `SELECT created_at, id + FROM session_message + WHERE session_id = ? AND session_run_id = ? AND role = 'assistant' + ORDER BY seq`, + ) + .bind(SESSION_ID, RUN_ID) + .all<{ created_at: number; id: string }>(); + expect(rows.results).toEqual([{ created_at: firstOccurredAt, id: RUN_ID }]); + const fresh = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.outsiderAccount, + }); + const merged = fresh.messages.filter((candidate) => candidate.id === RUN_ID); + expect(merged).toHaveLength(1); + expect(merged[0]?.createdAt).toBe(new Date(firstOccurredAt).toISOString()); + expect(merged[0]?.segments).toEqual( + toolFirst + ? [ + { + kind: "tool_result", + output: "same id output", + runId: RUN_ID, + tool: "Tool", + toolCallId: "same-id-tool", + }, + { kind: "text", text: "Same id final" }, + ] + : [ + { kind: "text", text: "Same id final" }, + { + kind: "tool_result", + output: "same id output", + runId: RUN_ID, + tool: "Tool", + toolCallId: "same-id-tool", + }, + ], + ); + }); + } + + test("fails closed when a cross-boot replay conflicts with the persisted final snapshot", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + const terminalBatch = [ + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: "conflict:original-final", + text: FINAL_TEXT, + }), + runtimeEvent({ + kind: "run.completed", + payload: { + finalMessageId, + stopReason: "end_turn", + }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ]; + await expect(pushFreshController(bindings, terminalBatch)).resolves.toHaveLength( + terminalBatch.length, + ); + + const replayedFinalMessageId = createPlatformId(); + const conflictingText = `${FINAL_TEXT}\nCONFLICTING-REPLAY`; + const conflictingReplay = [ + ...messageEvents({ + messageId: replayedFinalMessageId, + sourcePrefix: "conflict:reconnected-final", + text: conflictingText, + }), + runtimeEvent({ + kind: "run.completed", + payload: { + finalMessageId: replayedFinalMessageId, + stopReason: "end_turn", + }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ]; + + await expect(pushFreshController(bindings, conflictingReplay)).rejects.toBeInstanceOf(Error); + + const messages = await database + .prepare("SELECT content_text, id FROM session_message WHERE session_run_id = ? ORDER BY seq") + .bind(RUN_ID) + .all<{ content_text: string; id: string }>(); + const terminalRows = await database + .prepare("SELECT source_event_id FROM session_event WHERE source_event_id = ?") + .bind(TERMINAL_SOURCE_EVENT_ID) + .all<{ source_event_id: string }>(); + + expect(messages.results).toEqual([{ content_text: "", id: finalMessageId }]); + expect(terminalRows.results).toEqual([{ source_event_id: TERMINAL_SOURCE_EVENT_ID }]); + await expect( + readPublicThreadRunFinalOutput({ database, runId: RUN_ID, sessionId: SESSION_ID }), + ).resolves.toEqual({ text: FINAL_TEXT }); + }); + + test("does not guess a progress message when the terminal RPC has no final identity", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const progressMessageId = createPlatformId(); + const progressEvents = messageEvents({ + messageId: progressMessageId, + sourcePrefix: "fallback:progress", + text: PROGRESS_TEXTS[0], + }); + + await pushFreshController(bindings, progressEvents); + await expect( + recordDriverInstanceCompletion(bindings, { + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }), + ).rejects.toThrow("Completed Session Run is missing its canonical terminal event."); + + await expect( + readPublicThreadRunFinalOutput({ + database, + runId: RUN_ID, + sessionId: SESSION_ID, + }), + ).resolves.toBeNull(); + const messages = await database + .prepare("SELECT id FROM session_message WHERE session_run_id = ?") + .bind(RUN_ID) + .all<{ id: string }>(); + + expect(messages.results).toEqual([]); + }); + + test("rejects a failure RPC after another Driver connection takes ownership", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + + await database + .prepare("UPDATE driver_instance SET connection_id = ?, generation = 1 WHERE id = ?") + .bind("replacement-connection", DRIVER_ID) + .run(); + + await expect( + recordDriverInstanceFailure(bindings, { + driverConnectionId: "canary-connection", + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + error: { + code: "driver.disconnected", + details: {}, + message: "The superseded Driver disconnected.", + retryable: true, + }, + }), + ).rejects.toThrow("lost a concurrent running race"); + + const run = await database + .prepare("SELECT status FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first<{ status: string }>(); + const terminalEvents = await database + .prepare( + "SELECT id FROM session_event WHERE run_id = ? AND event_type IN ('run.cancelled', 'run.completed', 'run.failed')", + ) + .bind(RUN_ID) + .all(); + + expect(run?.status).toBe("running"); + expect(terminalEvents.results).toEqual([]); + }); + + test("a late run-one terminal RPC cannot touch the active run-two lease", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const nextRunId = createPlatformId(); + + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + error: null, + runId: RUN_ID, + sessionId: SESSION_ID, + source: "api", + status: "completed", + }); + await database + .prepare( + `INSERT INTO session_run ( + id, session_id, agent_id, created_by_account_id, deployment_version_id, + deployment_version_number, driver_instance_id, trigger, status, provider, + model, runtime_id, trace_id, started_at, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, 1, ?, 'user_prompt', 'running', 'openai', + 'gpt-5.4', 'openai-runtime', 'trace-next', 2, 2, 2)`, + ) + .bind( + nextRunId, + SESSION_ID, + PUBLIC_API_TEST_IDS.agent, + PUBLIC_API_TEST_IDS.ownerAccount, + PUBLIC_API_TEST_IDS.deployment, + DRIVER_ID, + ) + .run(); + await database + .prepare("UPDATE session SET last_run_id = ?, status = 'RUNNING' WHERE id = ?") + .bind(nextRunId, SESSION_ID) + .run(); + + await expect( + recordDriverInstanceCompletion(bindings, { + driverConnectionId: "canary-connection", + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }), + ).rejects.toThrow("lost a concurrent completed race"); + await expect( + database.prepare("SELECT status FROM session_run WHERE id = ?").bind(nextRunId).first(), + ).resolves.toEqual({ status: "running" }); + await expect( + database + .prepare("SELECT last_run_id, status, status_operation_id FROM session WHERE id = ?") + .bind(SESSION_ID) + .first(), + ).resolves.toEqual({ + last_run_id: nextRunId, + status: "RUNNING", + status_operation_id: null, + }); + await expect( + database + .prepare("SELECT status, status_operation_id FROM driver_instance WHERE id = ?") + .bind(DRIVER_ID) + .first(), + ).resolves.toEqual({ status: "ready", status_operation_id: null }); + }); + + test("returns no final output for a canonical completion without a final message", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + + await pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { stopReason: "end_turn" }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ]); + + await expect( + readPublicThreadRunFinalOutput({ database, runId: RUN_ID, sessionId: SESSION_ID }), + ).resolves.toBeNull(); + }); + + for (const terminalCase of [ + { + kind: "run.completed" as const, + payload: { stopReason: "end_turn" }, + status: "completed", + }, + { + kind: "run.failed" as const, + payload: { + error: { + code: "driver.provider_failed", + details: {}, + message: "Provider failed after tool output.", + retryable: false, + }, + recoverable: false, + }, + status: "failed", + }, + { + kind: "run.cancelled" as const, + payload: { reason: "session.stop" }, + status: "cancelled", + }, + ]) { + test(`persists a ${terminalCase.status} run's parentless tool deltas without a final message`, async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const toolOccurredAt = Date.parse("2026-08-30T06:00:00.000Z"); + + await pushFreshController(bindings, [ + runtimeEvent({ + kind: "tool.call.updated", + occurredAt: toolOccurredAt, + payload: { + rawOutputDelta: "A", + status: "running", + toolCallId: "terminal-no-final-tool", + }, + sourceEventId: `${terminalCase.status}:tool-output-a`, + }), + runtimeEvent({ + kind: "tool.call.updated", + occurredAt: toolOccurredAt + 1, + payload: { + rawOutputDelta: "B", + status: "completed", + toolCallId: "terminal-no-final-tool", + }, + sourceEventId: `${terminalCase.status}:tool-output-b`, + }), + runtimeEvent({ + kind: terminalCase.kind, + payload: terminalCase.payload, + sourceEventId: `${terminalCase.status}:terminal`, + }), + ]); + + const rows = await database + .prepare( + `SELECT created_at, id, projection_format + FROM session_message + WHERE session_id = ? AND session_run_id = ? AND role = 'assistant' + ORDER BY seq`, + ) + .bind(SESSION_ID, RUN_ID) + .all<{ created_at: number; id: string; projection_format: string }>(); + expect(rows.results).toEqual([ + { + created_at: toolOccurredAt, + id: RUN_ID, + projection_format: "event_stream_v3", + }, + ]); + const fresh = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.outsiderAccount, + }); + expect(fresh.messages.find((message) => message.id === RUN_ID)).toEqual( + expect.objectContaining({ + content: "", + createdAt: new Date(toolOccurredAt).toISOString(), + segments: [ + { + kind: "tool_result", + output: "AB", + runId: RUN_ID, + tool: "Tool", + toolCallId: "terminal-no-final-tool", + }, + ], + }), + ); + }); + } + + test("adopts a canonical completion RPC replay without creating another terminal event", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + + await pushFreshController(bindings, [ + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: "completion-rpc-replay:final", + text: FINAL_TEXT, + }), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: "provider-completion-id", + }), + ]); + + const operationId = createPlatformId(); + await database + .prepare( + `UPDATE session + SET status = 'RESCHEDULING', status_operation_id = ?, status_seq = status_seq + 1 + WHERE id = ?`, + ) + .bind(operationId, SESSION_ID) + .run(); + const sessionBeforeReplay = await database + .prepare( + `SELECT last_run_id, message_seq_cursor, runtime_event_seq_cursor, + status, status_operation_id, status_seq + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(); + + for (let attempt = 0; attempt < 2; attempt += 1) { + await recordDriverInstanceCompletion(bindings, { + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }); + } + const sessionAfterReplay = await database + .prepare( + `SELECT last_run_id, message_seq_cursor, runtime_event_seq_cursor, + status, status_operation_id, status_seq + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(); + + const terminalEvents = await database + .prepare( + "SELECT semantic_hash, source_event_id FROM session_event WHERE run_id = ? AND event_type = 'run.completed'", + ) + .bind(RUN_ID) + .all<{ semantic_hash: string; source_event_id: string }>(); + + expect(terminalEvents.results).toEqual([ + { + semantic_hash: expect.stringMatching(/^[0-9a-f]{64}$/), + source_event_id: TERMINAL_SOURCE_EVENT_ID, + }, + ]); + expect(sessionAfterReplay).toEqual(sessionBeforeReplay); + await expect( + readPublicThreadRunFinalOutput({ database, runId: RUN_ID, sessionId: SESSION_ID }), + ).resolves.toEqual({ text: FINAL_TEXT }); + }); + + for (const rpcKind of ["completion", "failure"] as const) { + test(`rejects a stale ${rpcKind} RPC when a same-generation Driver connection replaced its terminal winner`, async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const terminalError = { + code: "driver.api_winner", + details: {}, + message: "The API terminal winner is durable.", + retryable: false, + } as const; + + if (rpcKind === "completion") { + const finalMessageId = createPlatformId(); + await pushFreshController( + bindings, + messageEvents({ + messageId: finalMessageId, + sourcePrefix: "same-generation-winner:final", + text: FINAL_TEXT, + }), + ); + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: prepareAssistantMessageProjection({ + createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + messageId: finalMessageId, + sessionId: SESSION_ID, + sessionRunId: RUN_ID, + }), + error: null, + runId: RUN_ID, + sessionId: SESSION_ID, + source: "api", + status: "completed", + }); + } else { + await recordCanonicalSessionRunFailure(bindings, { + error: terminalError, + runId: RUN_ID, + sessionId: SESSION_ID, + source: "api", + }); + } + await database + .prepare("UPDATE driver_instance SET connection_id = ? WHERE id = ?") + .bind("replacement-connection", DRIVER_ID) + .run(); + const driverBefore = await database + .prepare( + "SELECT connection_id, generation, status, status_operation_id FROM driver_instance WHERE id = ?", + ) + .bind(DRIVER_ID) + .first(); + const sessionBefore = await database + .prepare( + `SELECT message_seq_cursor, runtime_event_seq_cursor, status, + status_operation_id, status_seq, updated_at + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(); + + await expect( + rpcKind === "completion" + ? recordDriverInstanceCompletion(bindings, { + driverConnectionId: "canary-connection", + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }) + : recordDriverInstanceFailure(bindings, { + driverConnectionId: "canary-connection", + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + error: { + code: "driver.stale_rpc", + details: {}, + message: "The old connection reported terminal.", + retryable: true, + }, + }), + ).rejects.toThrow("lost a concurrent"); + + await expect( + database + .prepare( + "SELECT connection_id, generation, status, status_operation_id FROM driver_instance WHERE id = ?", + ) + .bind(DRIVER_ID) + .first(), + ).resolves.toEqual(driverBefore); + await expect( + database + .prepare( + `SELECT message_seq_cursor, runtime_event_seq_cursor, status, + status_operation_id, status_seq, updated_at + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(), + ).resolves.toEqual(sessionBefore); + await expect( + database + .prepare( + `SELECT COUNT(*) AS count + FROM session_event + WHERE run_id = ? AND event_type IN ('run.cancelled', 'run.completed', 'run.failed')`, + ) + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ count: 1 }); + }); + } + + test("atomically claims an exact terminal Driver without rewriting a newer Session lifecycle", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + await pushFreshController( + bindings, + messageEvents({ + messageId: finalMessageId, + sourcePrefix: "atomic-adoption-claim:final", + text: FINAL_TEXT, + }), + ); + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: prepareAssistantMessageProjection({ + createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + messageId: finalMessageId, + sessionId: SESSION_ID, + sessionRunId: RUN_ID, + }), + error: null, + runId: RUN_ID, + sessionId: SESSION_ID, + source: "api", + status: "completed", + }); + const operationId = createPlatformId(); + await database + .prepare( + `UPDATE session + SET status = 'RESCHEDULING', status_operation_id = ?, status_seq = status_seq + 1 + WHERE id = ?`, + ) + .bind(operationId, SESSION_ID) + .run(); + const sessionBefore = await database + .prepare( + `SELECT last_run_id, message_seq_cursor, runtime_event_seq_cursor, + status, status_operation_id, status_seq, updated_at + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(); + + await recordDriverInstanceCompletion(bindings, { + driverConnectionId: "canary-connection", + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }); + + await expect( + database + .prepare( + `SELECT last_run_id, message_seq_cursor, runtime_event_seq_cursor, + status, status_operation_id, status_seq, updated_at + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(), + ).resolves.toEqual(sessionBefore); + }); + + test("atomically claims the persisted receipt while repairing an exact failure replay", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const durableError = { + code: "driver.api_failure", + details: {}, + message: "The API failure receipt is durable.", + retryable: false, + } as const; + await recordCanonicalSessionRunFailure(bindings, { + error: durableError, + runId: RUN_ID, + sessionId: SESSION_ID, + source: "api", + }); + await database + .prepare("UPDATE session SET status = 'RUNNING', updated_at = 1 WHERE id = ?") + .bind(SESSION_ID) + .run(); + + await recordDriverInstanceFailure(bindings, { + driverConnectionId: "canary-connection", + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + error: { + code: "driver.late_failure", + details: {}, + message: "The exact Driver replay arrived later.", + retryable: true, + }, + }); + + await expect( + database.prepare("SELECT status FROM session WHERE id = ?").bind(SESSION_ID).first(), + ).resolves.toEqual({ status: "IDLE" }); + await expect( + database + .prepare("SELECT error_code, status FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ error_code: durableError.code, status: "failed" }); + await expect( + database + .prepare( + "SELECT COUNT(*) AS count FROM session_event WHERE run_id = ? AND event_type = 'run.failed'", + ) + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ count: 1 }); + }); + + test("records a fresh Driver failure while the exact connection is still connecting", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await database + .prepare("UPDATE driver_instance SET status = 'connecting' WHERE id = ?") + .bind(DRIVER_ID) + .run(); + + await recordDriverInstanceFailure(bindings, { + driverConnectionId: "canary-connection", + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + error: { + code: "driver.connect_failed", + details: {}, + message: "The Driver failed before ready.", + retryable: true, + }, + }); + + await expect( + database + .prepare("SELECT error_code, status FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ error_code: "driver.connect_failed", status: "failed" }); + await expect( + database + .prepare( + "SELECT COUNT(*) AS count FROM session_event WHERE run_id = ? AND event_type = 'run.failed'", + ) + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ count: 1 }); + }); + + test("repairs a missing final reference behind the exact terminal adoption barrier", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const finalMessageId = await insertRepairableCompletedAuthority( + database, + "adoption-barrier:exact", + ); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + + await recordDriverInstanceCompletion(bindings, { + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }); + + await expect( + database + .prepare( + "SELECT id, projection_format, seq, session_id FROM session_message WHERE session_run_id = ?", + ) + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ + id: finalMessageId, + projection_format: "event_stream_v3", + seq: 1, + session_id: SESSION_ID, + }); + }); + + for (const corruptedStream of ["null", "progress"] as const) { + test(`rejects a pre-corrupted ${corruptedStream} terminal stream before missing-final adoption`, async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const progressMessageId = createPlatformId(); + await pushFreshController( + bindings, + messageEvents({ + messageId: progressMessageId, + sourcePrefix: `adoption-corrupt-${corruptedStream}:progress`, + text: PROGRESS_TEXTS[0], + }), + ); + await insertRepairableCompletedAuthority( + database, + `adoption-corrupt-${corruptedStream}:final`, + ); + await database + .prepare( + "UPDATE session_event SET stream_id = ? WHERE run_id = ? AND event_type = 'run.completed'", + ) + .bind(corruptedStream === "null" ? null : progressMessageId, RUN_ID) + .run(); + const messagesBefore = await database + .prepare( + "SELECT content_text, id, projection_format FROM session_message WHERE session_run_id = ? AND role = 'assistant' ORDER BY seq", + ) + .bind(RUN_ID) + .all<{ content_text: string; id: string; projection_format: string }>(); + + let rejected = false; + try { + await recordDriverInstanceCompletion(bindings, { + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }); + } catch { + rejected = true; + } + const messages = await database + .prepare( + "SELECT content_text, id, projection_format FROM session_message WHERE session_run_id = ? AND role = 'assistant' ORDER BY seq", + ) + .bind(RUN_ID) + .all<{ content_text: string; id: string; projection_format: string }>(); + expect({ messages: messages.results, rejected }).toEqual({ + messages: messagesBefore.results, + rejected: true, + }); + }); + } + + test("preserves the durable TERMINATED lifecycle while adopting a partial RUNNING Session", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + await pushFreshController( + bindings, + messageEvents({ + messageId: finalMessageId, + sourcePrefix: "adoption-terminated:final", + text: FINAL_TEXT, + }), + ); + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: prepareAssistantMessageProjection({ + createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + messageId: finalMessageId, + sessionId: SESSION_ID, + sessionRunId: RUN_ID, + }), + error: null, + lifecycle: "TERMINATED", + runId: RUN_ID, + sessionId: SESSION_ID, + source: "api", + status: "completed", + }); + await database.prepare("DELETE FROM session_message WHERE id = ?").bind(finalMessageId).run(); + await database + .prepare("UPDATE session SET message_seq_cursor = 0, status = 'RUNNING' WHERE id = ?") + .bind(SESSION_ID) + .run(); + + await recordDriverInstanceCompletion(bindings, { + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }); + + await expect( + database.prepare("SELECT status FROM session WHERE id = ?").bind(SESSION_ID).first(), + ).resolves.toEqual({ status: "TERMINATED" }); + }); + + for (const receiptMutation of [ + { + field: "source", + sql: "UPDATE session_event SET source = 'system' WHERE run_id = ? AND event_type = 'run.completed'", + }, + { + field: "stream", + sql: "UPDATE session_event SET stream_id = '01J000000000000000000000ZZ' WHERE run_id = ? AND event_type = 'run.completed'", + }, + { + field: "sequence", + sql: "UPDATE session_event SET seq = seq + 1000 WHERE run_id = ? AND event_type = 'run.completed'", + }, + { + field: "visibility", + sql: "UPDATE session_event SET visibility = 'owner_debug' WHERE run_id = ? AND event_type = 'run.completed'", + }, + { + field: "artifact", + sql: `UPDATE session_event + SET artifact_manifest_sha256 = '${"0".repeat(64)}' + WHERE run_id = ? AND event_type = 'run.completed'`, + }, + ] as const) { + test(`rolls back terminal adoption when its ${receiptMutation.field} receipt field changes before the batch`, async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + await insertRepairableCompletedAuthority( + database, + `adoption-barrier:${receiptMutation.field}`, + ); + await expect( + database + .prepare( + "SELECT artifact_manifest_sha256 FROM session_event WHERE run_id = ? AND event_type = 'run.completed'", + ) + .bind(RUN_ID) + .first<{ artifact_manifest_sha256: string | null }>(), + ).resolves.toEqual({ artifact_manifest_sha256: expect.any(String) }); + const sessionBefore = await database + .prepare( + `SELECT last_message_at, message_seq_cursor, runtime_event_seq_cursor, + status, status_operation_id, status_seq, updated_at + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(); + const injected = mutateBeforeFirstDatabaseBatch(database, async (target) => { + await target.prepare(receiptMutation.sql).bind(RUN_ID).run(); + }); + const bindings = createPublicHttpTestBindings(injected.database) as ApiBindings; + + await expect( + recordDriverInstanceCompletion(bindings, { + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }), + ).rejects.toThrow(); + + expect(injected.wasInjected()).toBe(true); + await expect( + database + .prepare( + `SELECT last_message_at, message_seq_cursor, runtime_event_seq_cursor, + status, status_operation_id, status_seq, updated_at + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(), + ).resolves.toEqual(sessionBefore); + await expect( + database + .prepare( + "SELECT COUNT(*) AS count FROM session_message WHERE session_run_id = ? AND role = 'assistant'", + ) + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ count: 0 }); + await expect( + database + .prepare( + `SELECT COUNT(*) AS count + FROM session_event + WHERE run_id = ? AND event_type IN ('run.cancelled', 'run.completed', 'run.failed')`, + ) + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ count: 1 }); + }); + } + + for (const corruptedAuthority of ["final", "carrier"] as const) { + test(`rejects a canonical replay whose ${corruptedAuthority} row belongs to another session`, async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + await insertNonOwnerSession(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + const terminal = runtimeEvent({ + kind: "run.completed", + payload: + corruptedAuthority === "final" + ? { finalMessageId, stopReason: "end_turn" } + : { stopReason: "end_turn" }, + sourceEventId: `cross-session-authority:${corruptedAuthority}:terminal`, + }); + + await pushFreshController(bindings, [ + ...(corruptedAuthority === "final" + ? messageEvents({ + messageId: finalMessageId, + sourcePrefix: "cross-session-authority:final", + text: "Final authority", + }) + : [ + runtimeEvent({ + kind: "tool.call.updated", + payload: { + rawOutput: "carrier authority", + status: "completed", + toolCallId: "cross-session-carrier", + }, + sourceEventId: "cross-session-authority:carrier", + }), + ]), + terminal, + ]); + + await database + .prepare("UPDATE session_message SET session_id = ? WHERE id = ?") + .bind( + PUBLIC_API_TEST_IDS.nonOwnerSession, + corruptedAuthority === "final" ? finalMessageId : RUN_ID, + ) + .run(); + await database + .prepare( + `UPDATE session + SET status = 'RESCHEDULING', status_operation_id = ?, status_seq = status_seq + 1 + WHERE id = ?`, + ) + .bind(createPlatformId(), SESSION_ID) + .run(); + + await expect( + recordDriverInstanceCompletion(bindings, { + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }), + ).rejects.toThrow("Canonical assistant messages"); + }); + } + + test("rejects a legacy materialized final authority from another session on RPC replay", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + await insertNonOwnerSession(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + await pushFreshController(bindings, [ + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: "legacy-cross-session:final", + text: "Legacy final authority", + }), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: "legacy-cross-session:terminal", + }), + ]); + await database + .prepare( + `UPDATE session_message + SET content_text = 'Legacy final authority', projection_format = 'materialized', + session_id = ? + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession, finalMessageId) + .run(); + await database + .prepare( + `UPDATE session_event + SET artifact_attempt_id = NULL, artifact_manifest_json = NULL, + artifact_manifest_sha256 = NULL, semantic_hash = NULL, + terminal_event_json = NULL + WHERE run_id = ? AND event_type = 'run.completed'`, + ) + .bind(RUN_ID) + .run(); + + await expect( + recordDriverInstanceCompletion(bindings, { + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }), + ).rejects.toThrow("Legacy terminal assistant messages conflict"); + }); + + test("does not repair a missing terminal event across a newer Session operation", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + const terminal = runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }); + + await pushFreshController(bindings, [ + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: "partial-repair:final", + text: FINAL_TEXT, + }), + terminal, + ]); + await database + .prepare("DELETE FROM session_event WHERE source_event_id = ?") + .bind(TERMINAL_SOURCE_EVENT_ID) + .run(); + const completedRun = await getSessionRunSummary(database, RUN_ID); + if (completedRun === null) { + throw new Error("Missing completed Session Run fixture."); + } + const runEvent = createSessionRunUpdatedEvent( + completedRun, + SESSION_ID, + "IDLE", + TERMINAL_SOURCE_EVENT_ID, + ); + const repairEvent = { + ...runEvent, + payload: { ...runEvent.payload, finalMessageId, stopReason: "end_turn" }, + }; + const run = await database + .prepare("SELECT updated_at FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first<{ updated_at: number }>(); + if (run === null) { + throw new Error("Missing completed Session Run fixture."); + } + const operationId = createPlatformId(); + await database + .prepare( + `UPDATE session + SET status = 'RESCHEDULING', status_operation_id = ?, + status_seq = status_seq + 1, updated_at = ? + WHERE id = ?`, + ) + .bind(operationId, run.updated_at + 1, SESSION_ID) + .run(); + const before = await database + .prepare( + `SELECT last_run_id, message_seq_cursor, runtime_event_seq_cursor, + status, status_operation_id, status_seq, updated_at + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(); + + await expect( + commitTerminalRunProjection(database, { + assistantMessage: prepareAssistantMessageProjection({ + createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + messageId: finalMessageId, + sessionId: SESSION_ID, + sessionRunId: RUN_ID, + }), + error: null, + runId: RUN_ID, + sessionId: SESSION_ID, + source: "driver", + targetStatus: "completed", + terminalEvent: { + event: repairEvent, + occurredAt: Date.parse(repairEvent.occurredAt), + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }, + }), + ).rejects.toThrow("not safely repairable"); + + const after = await database + .prepare( + `SELECT last_run_id, message_seq_cursor, runtime_event_seq_cursor, + status, status_operation_id, status_seq, updated_at + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(); + const terminalRows = await database + .prepare("SELECT id FROM session_event WHERE source_event_id = ?") + .bind(TERMINAL_SOURCE_EVENT_ID) + .all(); + expect(after).toEqual(before); + expect(terminalRows.results).toEqual([]); + }); + + test("repairs a missing terminal event from an unchanged RUNNING Session", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + + await pushFreshController(bindings, [ + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: "partial-repair-running:final", + text: FINAL_TEXT, + }), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }), + ]); + await database + .prepare("DELETE FROM session_event WHERE source_event_id = ?") + .bind(TERMINAL_SOURCE_EVENT_ID) + .run(); + const completedRun = await getSessionRunSummary(database, RUN_ID); + if (completedRun === null) { + throw new Error("Missing completed Session Run fixture."); + } + const repairEvent = createSessionRunUpdatedEvent( + completedRun, + SESSION_ID, + "IDLE", + TERMINAL_SOURCE_EVENT_ID, + ); + await database + .prepare( + `UPDATE session + SET status = 'RUNNING', status_operation_id = NULL, updated_at = ? + WHERE id = ?`, + ) + .bind(completedRun.updatedAt, SESSION_ID) + .run(); + + await expect( + commitTerminalRunProjection(database, { + assistantMessage: prepareAssistantMessageProjection({ + createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + messageId: finalMessageId, + sessionId: SESSION_ID, + sessionRunId: RUN_ID, + }), + error: null, + runId: RUN_ID, + sessionId: SESSION_ID, + source: "driver", + targetStatus: "completed", + terminalEvent: { + event: { + ...repairEvent, + payload: { ...repairEvent.payload, finalMessageId, stopReason: "end_turn" }, + }, + occurredAt: Date.parse(repairEvent.occurredAt), + sourceEventId: TERMINAL_SOURCE_EVENT_ID, + }, + }), + ).resolves.toMatchObject({ kind: "applied" }); + + await expect( + database + .prepare( + `SELECT status, status_operation_id + FROM session WHERE id = ?`, + ) + .bind(SESSION_ID) + .first(), + ).resolves.toEqual({ status: "IDLE", status_operation_id: null }); + await expect( + readPublicThreadRunFinalOutput({ database, runId: RUN_ID, sessionId: SESSION_ID }), + ).resolves.toEqual({ text: FINAL_TEXT }); + }); + + test("rejects an explicit terminal Run view after its observed Run advances", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + await database + .prepare( + `UPDATE session_run + SET started_at = NULL, status = 'queued', status_seq = status_seq + 1 + WHERE id = ?`, + ) + .bind(RUN_ID) + .run(); + const observed = await getSessionRunSummary(database, RUN_ID); + if (observed === null) { + throw new Error("Missing Session Run race fixture."); + } + const error = { + code: "runtime.stale_projection", + details: {}, + message: "Stale terminal projection.", + retryable: true, + }; + const terminalTimestampMs = 2_000; + const staleTerminalRun = { + ...observed, + completedAt: new Date(terminalTimestampMs).toISOString(), + error, + startedAt: new Date(terminalTimestampMs).toISOString(), + status: "failed" as const, + updatedAt: new Date(terminalTimestampMs).toISOString(), + }; + const terminalEvent = createFailedSessionRunRuntimeEvent({ + lifecycle: "IDLE", + run: staleTerminalRun, + runError: error, + sessionId: SESSION_ID, + sourceEventId: `session-run-terminal:${RUN_ID}:run.failed`, + }); + + await database + .prepare( + `UPDATE session_run + SET started_at = 1500, status = 'booting', status_seq = status_seq + 1, + updated_at = 1500 + WHERE id = ?`, + ) + .bind(RUN_ID) + .run(); + + await expect( + commitTerminalRunProjection(database, { + assistantMessage: null, + error, + runId: RUN_ID, + sessionId: SESSION_ID, + source: "api", + targetStatus: "failed", + terminalEvent: { + event: terminalEvent, + occurredAt: terminalTimestampMs, + sourceEventId: `session-run-terminal:${RUN_ID}:run.failed`, + }, + timestampMs: terminalTimestampMs, + }), + ).resolves.toMatchObject({ currentStatus: "booting", kind: "stale" }); + await expect(getSessionRunSummary(database, RUN_ID)).resolves.toMatchObject({ + error: null, + startedAt: new Date(1_500).toISOString(), + status: "booting", + }); + const terminalRows = await database + .prepare("SELECT id FROM session_event WHERE run_id = ? AND event_type = 'run.failed'") + .bind(RUN_ID) + .all(); + expect(terminalRows.results).toEqual([]); + }); + + test("adopts a canonical cancellation when a later completion RPC loses the race", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + + await pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.cancelled", + payload: { reason: "session.stop" }, + sourceEventId: "provider-cancellation-id", + }), + ]); + await recordDriverInstanceCompletion(bindings, { + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }); + + const run = await database + .prepare("SELECT status FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first<{ status: string }>(); + const terminalEvents = await database + .prepare( + "SELECT event_type, source_event_id FROM session_event WHERE run_id = ? AND event_type IN ('run.cancelled', 'run.completed', 'run.failed')", + ) + .bind(RUN_ID) + .all<{ event_type: string; source_event_id: string }>(); + + expect(run).toEqual({ status: "cancelled" }); + expect(terminalEvents.results).toEqual([ + { + event_type: "run.cancelled", + source_event_id: `session-run-terminal:${RUN_ID}:run.cancelled`, + }, + ]); + }); + + test("adopts the canonical provider failure when its terminal RPC acknowledgement is lost", async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const providerError = { + code: "driver.provider_failed", + details: { phase: "stream" }, + message: "The provider stream failed.", + retryable: false, + } as const; + + await pushFreshController(bindings, [ runtimeEvent({ - kind: "run.completed", - payload: { - finalMessageId: replayedFinalMessageId, - finalMessageText: FINAL_TEXT, - stopReason: "end_turn", - }, - sourceEventId: TERMINAL_SOURCE_EVENT_ID, + kind: "run.failed", + payload: { error: providerError, recoverable: false }, + sourceEventId: "provider-failure-id", }), - ]; - - expect(await pushFreshController(bindings, crossBootTerminalBatch)).toHaveLength( - crossBootTerminalBatch.length, - ); - expect(await pushFreshController(bindings, crossBootTerminalBatch)).toHaveLength( - crossBootTerminalBatch.length, - ); + ]); + await recordDriverInstanceFailure(bindings, { + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + error: { + code: "driver.rpc_failed", + details: {}, + message: "The later terminal RPC failed.", + retryable: false, + }, + }); - const projectedMessagesAfterReplay = await database - .prepare("SELECT content_text, id FROM session_message WHERE session_run_id = ? ORDER BY seq") + const run = await database + .prepare("SELECT error_code, error_message, status FROM session_run WHERE id = ?") .bind(RUN_ID) - .all<{ content_text: string; id: string }>(); - const terminalRowsAfterReplay = await database - .prepare("SELECT source_event_id FROM session_event WHERE source_event_id = ?") - .bind(TERMINAL_SOURCE_EVENT_ID) - .all<{ source_event_id: string }>(); - const transcript = await loadSessionViewerState(database, { - sessionId: SESSION_ID, - viewerId: PUBLIC_API_TEST_IDS.ownerAccount, - }); - const finalTranscriptMessage = transcript.messages.find( - (message) => message.id === finalMessageId, - ); - const canonicalTranscriptMessages = transcript.messages.filter( - (message) => message.role === "assistant" && message.content === FINAL_TEXT, - ); + .first<{ error_code: string; error_message: string; status: string }>(); + const terminalEvents = await database + .prepare( + "SELECT content_text, source_event_id FROM session_event WHERE run_id = ? AND event_type = 'run.failed'", + ) + .bind(RUN_ID) + .all<{ content_text: string; source_event_id: string }>(); - expect(projectedMessagesAfterReplay.results).toEqual(projectedMessagesBeforeReplay.results); - expect(terminalRowsAfterReplay.results).toEqual([ - { source_event_id: TERMINAL_SOURCE_EVENT_ID }, + expect(run).toEqual({ + error_code: providerError.code, + error_message: providerError.message, + status: "failed", + }); + expect(terminalEvents.results).toEqual([ + { + content_text: providerError.message, + source_event_id: `session-run-terminal:${RUN_ID}:run.failed`, + }, ]); - expect(finalTranscriptMessage?.content).toBe(FINAL_TEXT); - expect(canonicalTranscriptMessages.map((message) => message.id)).toEqual([finalMessageId]); - expect(finalTranscriptMessage?.content).not.toContain(PROGRESS_TEXTS.join("")); - expect(FINAL_TEXT.split("\n")).toContain("160|中文长文本校验-Aa9-表格字符|END160"); }); - test("removes provider-private citations at the public final-output boundary", async () => { + for (const corruption of ["legacy hash", "noncanonical source"] as const) { + test(`rejects ${corruption} when the completion RPC adopts a terminal event`, async () => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + + await pushFreshController(bindings, [ + ...messageEvents({ + messageId: finalMessageId, + sourcePrefix: `corrupt-completion:${corruption}`, + text: FINAL_TEXT, + }), + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: "provider-completion-id", + }), + ]); + await database + .prepare( + corruption === "legacy hash" + ? `UPDATE session_event + SET artifact_attempt_id = NULL, + artifact_manifest_json = NULL, + artifact_manifest_sha256 = NULL, + semantic_hash = NULL, + terminal_event_json = NULL + WHERE run_id = ? AND event_type = 'run.completed'` + : `UPDATE session_event + SET artifact_attempt_id = NULL, + artifact_manifest_json = NULL, + artifact_manifest_sha256 = NULL, + source_event_id = 'provider-completion-id' + WHERE run_id = ? AND event_type = 'run.completed'`, + ) + .bind(RUN_ID) + .run(); + + await expect( + recordDriverInstanceCompletion(bindings, { + driverInstanceId: DRIVER_ID, + sessionRunId: RUN_ID, + }), + ).rejects.toThrow( + corruption === "legacy hash" + ? "Legacy terminal assistant messages conflict" + : "terminal semantic authority is invalid", + ); + }); + } + + test("requires an authoritative message snapshot to be sealed before run completion", async () => { const database = await createPublicHttpContractDatabase(); await insertRuntimeFixture(database); const bindings = createPublicHttpTestBindings(database) as ApiBindings; const finalMessageId = createPlatformId(); - const privateCitation = "\uE200cite\uE202turn2view0\uE202turn8view0\uE201"; - const providerText = `before${privateCitation}after`; - const events = [ - ...messageEvents({ - messageId: finalMessageId, - sourcePrefix: "private-citation:final", - text: providerText, + const terminal = runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: "unsealed:run-completed", + }); + const unsealed = [ + runtimeEvent({ + kind: "message.added", + payload: { content: "Authoritative prefix", messageId: finalMessageId, role: "agent" }, + sourceEventId: "unsealed:snapshot", }), runtimeEvent({ - kind: "run.completed", - payload: { - finalMessageId, - finalMessageText: providerText, - stopReason: "end_turn", - }, - sourceEventId: "private-citation:run-completed", + kind: "message.delta", + payload: { contentDelta: " plus delta", messageId: finalMessageId, role: "agent" }, + sourceEventId: "unsealed:delta", }), ]; - await pushFreshController(bindings, events); + await expect(pushFreshController(bindings, [...unsealed, terminal])).rejects.toThrow( + "has no sealed authoritative snapshot", + ); - const persistedMessage = await database - .prepare("SELECT content_text FROM session_message WHERE id = ?") - .bind(finalMessageId) - .first<{ content_text: string }>(); + const runBeforeSeal = await database + .prepare("SELECT status FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first<{ status: string }>(); + expect(runBeforeSeal?.status).toBe("running"); + + const completed = runtimeEvent({ + kind: "message.completed", + payload: { messageId: finalMessageId, role: "agent" }, + sourceEventId: "unsealed:completed", + }); + await pushFreshController(bindings, [completed, terminal]); - expect(persistedMessage?.content_text).toBe(providerText); await expect( readPublicThreadRunFinalOutput({ database, runId: RUN_ID, sessionId: SESSION_ID }), - ).resolves.toEqual({ - text: "beforeafter", - warnings: [ - { - code: "unresolved_provider_citation", - count: 1, - }, - ], - }); + ).resolves.toEqual({ text: "Authoritative prefix plus delta" }); }); - test("omits live-only reasoning from stored final assistant segments", async () => { + test("requires an authoritative replacement snapshot to be sealed again", async () => { const database = await createPublicHttpContractDatabase(); await insertRuntimeFixture(database); const bindings = createPublicHttpTestBindings(database) as ApiBindings; const finalMessageId = createPlatformId(); - const privateReasoningText = "Private reasoning should stay out of stored history."; - const events = [ - runtimeEvent({ - kind: "thought.started", - payload: { messageId: finalMessageId }, - sourceEventId: "reasoning:started", - }), + const originalOccurredAt = Date.parse("2026-08-30T05:00:00.000Z"); + const carrierOccurredAt = originalOccurredAt + 2; + const replacementOccurredAt = originalOccurredAt + 3; + + await pushFreshController(bindings, [ runtimeEvent({ - kind: "thought.delta", - payload: { contentDelta: privateReasoningText, messageId: finalMessageId }, - sourceEventId: "reasoning:delta", + kind: "message.added", + occurredAt: originalOccurredAt, + payload: { content: "Original answer", messageId: finalMessageId, role: "agent" }, + sourceEventId: "replacement:original:snapshot", }), runtimeEvent({ - kind: "thought.completed", - payload: { messageId: finalMessageId }, - sourceEventId: "reasoning:completed", - }), - ...messageEvents({ - messageId: finalMessageId, - sourcePrefix: "reasoning:final", - text: FINAL_TEXT, + kind: "message.completed", + occurredAt: originalOccurredAt + 1, + payload: { messageId: finalMessageId, role: "agent" }, + sourceEventId: "replacement:original:completed", }), runtimeEvent({ - kind: "run.completed", + kind: "tool.call.updated", + occurredAt: carrierOccurredAt, payload: { - finalMessageId, - finalMessageText: FINAL_TEXT, - stopReason: "end_turn", + rawOutputDelta: "carrier output", + status: "completed", + toolCallId: "replacement-orphan-tool", }, - sourceEventId: "reasoning:run-completed", + sourceEventId: "replacement:carrier-output", }), - ]; + ]); - await pushFreshController(bindings, events); + const replacement = runtimeEvent({ + kind: "message.added", + occurredAt: replacementOccurredAt, + payload: { content: "Replacement answer", messageId: finalMessageId, role: "agent" }, + sourceEventId: "replacement:snapshot", + }); + const terminal = runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: "replacement:run-completed", + }); - const persistedMessage = await database - .prepare("SELECT content_text, segments_json FROM session_message WHERE id = ?") - .bind(finalMessageId) - .first<{ content_text: string; segments_json: string }>(); + await expect(pushFreshController(bindings, [replacement, terminal])).rejects.toThrow( + "has no sealed authoritative snapshot", + ); - expect(persistedMessage?.content_text).toBe(FINAL_TEXT); - expect(JSON.parse(persistedMessage?.segments_json ?? "[]")).toEqual([ - { kind: "text", text: FINAL_TEXT }, + await pushFreshController(bindings, [ + runtimeEvent({ + kind: "message.completed", + occurredAt: replacementOccurredAt + 1, + payload: { messageId: finalMessageId, role: "agent" }, + sourceEventId: "replacement:completed", + }), + terminal, + ]); + + await expect( + readPublicThreadRunFinalOutput({ database, runId: RUN_ID, sessionId: SESSION_ID }), + ).resolves.toEqual({ text: "Replacement answer" }); + const assistantRows = await database + .prepare( + `SELECT created_at, id, seq + FROM session_message + WHERE session_id = ? AND session_run_id = ? AND role = 'assistant' + ORDER BY seq`, + ) + .bind(SESSION_ID, RUN_ID) + .all<{ created_at: number; id: string; seq: number }>(); + expect(assistantRows.results.map(({ created_at, id }) => ({ created_at, id }))).toEqual([ + { created_at: carrierOccurredAt, id: RUN_ID }, + { created_at: replacementOccurredAt, id: finalMessageId }, + ]); + expect(assistantRows.results.map(({ seq }) => seq)).toEqual([ + assistantRows.results[0]?.seq, + (assistantRows.results[0]?.seq ?? 0) + 1, + ]); + const fresh = await loadSessionViewerState(database, { + sessionId: SESSION_ID, + viewerId: PUBLIC_API_TEST_IDS.outsiderAccount, + }); + expect( + fresh.messages + .filter((message) => message.id === RUN_ID || message.id === finalMessageId) + .map(({ createdAt, id }) => ({ createdAt, id })), + ).toEqual([ + { createdAt: new Date(carrierOccurredAt).toISOString(), id: RUN_ID }, + { createdAt: new Date(replacementOccurredAt).toISOString(), id: finalMessageId }, ]); - expect(persistedMessage?.segments_json).not.toContain(privateReasoningText); }); - test("fails closed when a cross-boot replay conflicts with the persisted final snapshot", async () => { + test("rejects a final stream identity shared with owner-only message rows", async () => { const database = await createPublicHttpContractDatabase(); await insertRuntimeFixture(database); const bindings = createPublicHttpTestBindings(database) as ApiBindings; const finalMessageId = createPlatformId(); - const terminalBatch = [ + + await pushFreshController(bindings, [ ...messageEvents({ messageId: finalMessageId, - sourcePrefix: "conflict:original-final", - text: FINAL_TEXT, + sourcePrefix: "mixed-visibility:public", + text: "Public answer", }), runtimeEvent({ - kind: "run.completed", - payload: { - finalMessageId, - finalMessageText: FINAL_TEXT, - stopReason: "end_turn", - }, - sourceEventId: TERMINAL_SOURCE_EVENT_ID, - }), - ]; - const failingBindings = { - ...bindings, - DB: failTerminalSessionEventInsert(database), - } as ApiBindings; - - await expect(pushFreshController(failingBindings, terminalBatch)).rejects.toBeInstanceOf(Error); - - const replayedFinalMessageId = createPlatformId(); - const conflictingText = `${FINAL_TEXT}\nCONFLICTING-REPLAY`; - const conflictingReplay = [ - ...messageEvents({ - messageId: replayedFinalMessageId, - sourcePrefix: "conflict:reconnected-final", - text: conflictingText, + kind: "message.added", + payload: { content: "Owner-only replacement", messageId: finalMessageId, role: "agent" }, + sourceEventId: "mixed-visibility:owner:snapshot", + visibility: "owner_debug", }), runtimeEvent({ - kind: "run.completed", - payload: { - finalMessageId: replayedFinalMessageId, - finalMessageText: conflictingText, - stopReason: "end_turn", - }, - sourceEventId: TERMINAL_SOURCE_EVENT_ID, + kind: "message.completed", + payload: { messageId: finalMessageId, role: "agent" }, + sourceEventId: "mixed-visibility:owner:completed", + visibility: "owner_debug", }), - ]; - - await expect(pushFreshController(bindings, conflictingReplay)).rejects.toThrow( - "Canonical final assistant message conflicts with the persisted projection", - ); - - const messages = await database - .prepare("SELECT content_text, id FROM session_message WHERE session_run_id = ? ORDER BY seq") - .bind(RUN_ID) - .all<{ content_text: string; id: string }>(); - const terminalRows = await database - .prepare("SELECT source_event_id FROM session_event WHERE source_event_id = ?") - .bind(TERMINAL_SOURCE_EVENT_ID) - .all<{ source_event_id: string }>(); + ]); - expect(messages.results).toEqual([{ content_text: FINAL_TEXT, id: finalMessageId }]); - expect(terminalRows.results).toEqual([]); await expect( - readPublicThreadRunFinalOutput({ database, runId: RUN_ID, sessionId: SESSION_ID }), - ).resolves.toEqual({ text: FINAL_TEXT }); - }); - - test("does not guess a progress message when the terminal RPC has no final identity", async () => { - const database = await createPublicHttpContractDatabase(); - await insertRuntimeFixture(database); - const bindings = createPublicHttpTestBindings(database) as ApiBindings; - const progressMessageId = createPlatformId(); - const progressEvents = messageEvents({ - messageId: progressMessageId, - sourcePrefix: "fallback:progress", - text: PROGRESS_TEXTS[0], - }); - - await pushFreshController(bindings, progressEvents); - await recordDriverInstanceCompletion(bindings, { - driverInstanceId: DRIVER_ID, - driverReady: true, - }); + pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: "mixed-visibility:run-completed", + }), + ]), + ).rejects.toThrow("mixed visibility"); - const finalOutput = await readPublicThreadRunFinalOutput({ - database, - runId: RUN_ID, - sessionId: SESSION_ID, - }); + const run = await database + .prepare("SELECT status FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first<{ status: string }>(); const messages = await database .prepare("SELECT id FROM session_message WHERE session_run_id = ?") .bind(RUN_ID) .all<{ id: string }>(); - expect(finalOutput).toBeNull(); + expect(run?.status).toBe("running"); expect(messages.results).toEqual([]); }); - test("fails closed when run completion omits the final text snapshot", async () => { + test.each(["unknown", "user"] as const)( + "rejects a final message ID that references an %s stream", + async (boundary) => { + const database = await createPublicHttpContractDatabase(); + await insertRuntimeFixture(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const finalMessageId = createPlatformId(); + const otherMessageId = createPlatformId(); + const stream = + boundary === "unknown" + ? messageEvents({ + messageId: otherMessageId, + sourcePrefix: "invalid-final:other", + text: "Another assistant message", + }) + : [ + runtimeEvent({ + kind: "message.added", + payload: { content: "User input", messageId: finalMessageId, role: "user" }, + sourceEventId: "invalid-final:user-snapshot", + }), + runtimeEvent({ + kind: "message.completed", + payload: { messageId: finalMessageId, role: "user" }, + sourceEventId: "invalid-final:user-completed", + }), + ]; + + await expect( + pushFreshController(bindings, [ + ...stream, + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: `invalid-final:${boundary}:run-completed`, + }), + ]), + ).rejects.toThrow( + boundary === "user" ? "conflicting identity rows" : "has no sealed authoritative snapshot", + ); + + const run = await database + .prepare("SELECT status FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first<{ status: string }>(); + expect(run?.status).toBe("running"); + }, + ); + + test("rejects a sealed message stream from another run", async () => { const database = await createPublicHttpContractDatabase(); await insertRuntimeFixture(database); const bindings = createPublicHttpTestBindings(database) as ApiBindings; - const progressMessageId = createPlatformId(); - const events = [ - ...messageEvents({ - messageId: progressMessageId, - sourcePrefix: "missing-snapshot:progress", - text: PROGRESS_TEXTS[0], - }), - runtimeEvent({ - kind: "run.completed", - payload: { finalMessageId: progressMessageId, stopReason: "end_turn" }, - sourceEventId: "missing-snapshot:run-completed", - }), - ]; + const finalMessageId = createPlatformId(); + const otherRunId = createPlatformId(); - await pushFreshController(bindings, events); + await pushFreshController( + bindings, + messageEvents({ + messageId: finalMessageId, + sourcePrefix: "cross-run:message", + text: "Another run's answer", + }), + ); + await database + .prepare( + `INSERT INTO session_run ( + id, session_id, agent_id, created_by_account_id, trigger, status, + trace_id, created_at, updated_at + ) VALUES (?, ?, ?, ?, 'user_prompt', 'completed', ?, 1, 1)`, + ) + .bind( + otherRunId, + SESSION_ID, + PUBLIC_API_TEST_IDS.agent, + PUBLIC_API_TEST_IDS.ownerAccount, + "trace-other-run", + ) + .run(); + await database + .prepare("UPDATE session_event SET run_id = ? WHERE session_id = ? AND stream_id = ?") + .bind(otherRunId, SESSION_ID, finalMessageId) + .run(); await expect( - readPublicThreadRunFinalOutput({ database, runId: RUN_ID, sessionId: SESSION_ID }), - ).resolves.toBeNull(); + pushFreshController(bindings, [ + runtimeEvent({ + kind: "run.completed", + payload: { finalMessageId, stopReason: "end_turn" }, + sourceEventId: "cross-run:run-completed", + }), + ]), + ).rejects.toThrow("conflicting identity rows"); }); test("does not persist canonical output after another terminal status wins", async () => { @@ -804,14 +5251,15 @@ describe("runtime final output ingestion", () => { kind: "run.completed", payload: { finalMessageId, - finalMessageText: FINAL_TEXT, stopReason: "end_turn", }, sourceEventId: "stale-completion:run-completed", }), ]; - await pushFreshController(bindings, events); + await expect(pushFreshController(bindings, events)).rejects.toThrow( + "requires its exact active Session Run", + ); const run = await database .prepare("SELECT status FROM session_run WHERE id = ?") diff --git a/apps/api/tests/runtime-operation-ready-authority-migration.test.ts b/apps/api/tests/runtime-operation-ready-authority-migration.test.ts new file mode 100644 index 00000000..d7c90fee --- /dev/null +++ b/apps/api/tests/runtime-operation-ready-authority-migration.test.ts @@ -0,0 +1,188 @@ +import { describe, expect, test } from "bun:test"; + +import type { RuntimeOperationId, SessionId } from "@mosoo/id"; +import { + createRuntimeEventSemanticHash, + stringifyRuntimeEventSemanticValue, +} from "@mosoo/runtime-events"; +import type { RuntimeEventEnvelope } from "@mosoo/runtime-events"; + +import { createRuntimeOperationSessionEvent } from "../src/modules/runtime/application/runtime-state-operation-events"; +import type { RuntimeOperationEvent } from "../src/modules/runtime/application/runtime-state-operation-events"; +import { createSessionRuntimeEventProjection } from "../src/modules/sessions/domain/session-runtime-event-projection"; +import { applyDrizzleMigration, applyDrizzleMigrationsBefore } from "./helpers/drizzle-migrations"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const MIGRATION_TAG = "0018_runtime-operation-ready-authority"; + +const ACCOUNT_ID = "01J0000000000000000000001K"; +const AGENT_ID = "01J0000000000000000000001J"; +const APP_ID = "01J0000000000000000000001M"; +const SESSION_ID = "01J0000000000000000000001H" as SessionId; +const OPERATION_ID = "01J0000000000000000000001R" as RuntimeOperationId; +const LEGACY_OPERATION_ID = "01J0000000000000000000001S" as RuntimeOperationId; + +async function createPre0018Database(): Promise { + const database = new SqliteD1Database(); + applyDrizzleMigrationsBefore(database, MIGRATION_TAG); + return database; +} + +async function insertSession(database: SqliteD1Database): Promise { + await database + .prepare( + `INSERT INTO session ( + agent_id, created_at, creator_account_id, id, kind, model, app_id, + provider, renamed, runtime_id, status, updated_at + ) VALUES (?, 1, ?, ?, 'agent', 'gpt-5.4', ?, 'openai', 0, 'codex', 'IDLE', 1)`, + ) + .bind(AGENT_ID, ACCOUNT_ID, SESSION_ID, APP_ID) + .run(); +} + +async function insertLegacyReadyReceipt(database: SqliteD1Database): Promise { + await database + .prepare( + `INSERT INTO session_event ( + agent_id, content_text, created_at, ended_at, event_type, family, id, + occurred_at, process_status, process_type, semantic_hash, seq, session_id, + source_event_id, source, visibility + ) VALUES (?, '', 1, 1, 'agent.task.updated', 'agent', ?, 1, 'available', + 'agent_task', ?, 1, ?, ?, 'api', 'all_consumers')`, + ) + .bind( + AGENT_ID, + "01J0000000000000000000001T", + "0".repeat(64), + SESSION_ID, + `runtime-operation:${LEGACY_OPERATION_ID}:${SESSION_ID}:ready`, + ) + .run(); +} + +function operationEvent(status: RuntimeOperationEvent["status"]): RuntimeOperationEvent { + return { + agentId: AGENT_ID, + observedAt: status === "updating" ? "2026-08-30T00:00:00.000Z" : "2026-08-30T00:00:01.000Z", + operation: "restartDriver", + status, + }; +} + +async function insertCanonicalOperationEvent( + database: SqliteD1Database, + status: RuntimeOperationEvent["status"], + seq: number, +): Promise { + const event = createRuntimeOperationSessionEvent({ + event: operationEvent(status), + operationId: OPERATION_ID, + sessionId: SESSION_ID, + }); + const projection = createSessionRuntimeEventProjection(event); + const occurredAt = Date.parse(event.occurredAt); + await database + .prepare( + `INSERT INTO session_event ( + agent_id, content_text, created_at, ended_at, event_type, family, id, + occurred_at, process_status, process_type, run_id, + runtime_operation_event_json, semantic_hash, seq, session_id, + source_event_id, source, stream_id, trace_id, visibility + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + AGENT_ID, + projection.contentText, + occurredAt, + occurredAt, + projection.eventType, + projection.family, + event.id, + occurredAt, + projection.processStatus, + projection.processType, + projection.runId, + stringifyRuntimeEventSemanticValue(event), + await createRuntimeEventSemanticHash(event), + seq, + SESSION_ID, + event.sourceEventId, + projection.source, + projection.streamId, + projection.traceId, + projection.visibility, + ) + .run(); + return event; +} + +describe("runtime operation authority migration", () => { + test("preserves legacy NULL while admitting both canonical operation phases", async () => { + const database = await createPre0018Database(); + await insertSession(database); + await insertLegacyReadyReceipt(database); + + applyDrizzleMigration(database, MIGRATION_TAG); + + expect( + ( + await database.prepare("PRAGMA table_info(session_event)").all<{ name: string }>() + ).results.map(({ name }) => name), + ).toContain("runtime_operation_event_json"); + await expect( + database + .prepare( + `SELECT runtime_operation_event_json + FROM session_event + WHERE source_event_id LIKE 'runtime-operation:%:ready'`, + ) + .first(), + ).resolves.toEqual({ runtime_operation_event_json: null }); + + await insertCanonicalOperationEvent(database, "updating", 2); + await insertCanonicalOperationEvent(database, "ready", 3); + const rows = await database + .prepare( + `SELECT json_extract(runtime_operation_event_json, '$.payload.status') AS status + FROM session_event + WHERE runtime_operation_event_json IS NOT NULL + ORDER BY seq`, + ) + .all<{ status: string }>(); + expect(rows.results).toEqual([{ status: "updating" }, { status: "ready" }]); + }); + + test("rejects a non-operation phase authority carrier", async () => { + const database = await createPre0018Database(); + await insertSession(database); + applyDrizzleMigration(database, MIGRATION_TAG); + const event = await insertCanonicalOperationEvent(database, "updating", 1); + const forged = { + ...event, + id: "01J0000000000000000000001V", + payload: { ...(event.payload as Record), status: "forged" }, + sourceEventId: "forged-runtime-operation-event", + } satisfies RuntimeEventEnvelope; + + await expect( + database + .prepare( + `INSERT INTO session_event ( + agent_id, content_text, created_at, ended_at, event_type, family, id, + occurred_at, process_status, process_type, runtime_operation_event_json, + semantic_hash, seq, session_id, source_event_id, source, visibility + ) VALUES (?, '', 2, 2, 'agent.task.updated', 'agent', ?, 2, 'available', + 'agent_task', ?, ?, 2, ?, ?, 'api', 'all_consumers')`, + ) + .bind( + AGENT_ID, + forged.id, + stringifyRuntimeEventSemanticValue(forged), + await createRuntimeEventSemanticHash(forged), + SESSION_ID, + forged.sourceEventId, + ) + .run(), + ).rejects.toThrow("CHECK constraint failed"); + }); +}); diff --git a/apps/api/tests/runtime-provisioning-lease.test.ts b/apps/api/tests/runtime-provisioning-lease.test.ts new file mode 100644 index 00000000..f9672bd3 --- /dev/null +++ b/apps/api/tests/runtime-provisioning-lease.test.ts @@ -0,0 +1,799 @@ +import { describe, expect, test } from "bun:test"; + +import { createPlatformId } from "@mosoo/id"; +import type { DriverInstanceId } from "@mosoo/id"; +import { PLATFORM_ID_FIXTURES } from "@mosoo/id/testing"; + +import { + createDriverInstanceRecord, + runtimeProvisioningDriverLaunchIsOwned, +} from "../src/modules/runtime/infrastructure/driver-instance/driver-instance-record.repository"; +import { stopDriverSession } from "../src/modules/runtime/infrastructure/driver-session-stop.service"; +import { startProvisionProcessWithOwnershipFence } from "../src/modules/runtime/infrastructure/runtime-sandbox-provisioning/runtime-driver-process-cleanup"; +import { + adoptReadyRuntimeRunProvisioningLease, + claimRuntimeProvisioningDriverCleanup, + claimRuntimeRunProvisioningLease, + claimStaleRuntimeProvisioningLeases, + heartbeatRuntimeRunProvisioningLease, + readRuntimeProvisioningCleanupTargets, + recordRuntimeProvisioningConversationTarget, + releaseAbortedRuntimeProvisioningLease, + releaseReadyRuntimeRunProvisioningLease, + renewRuntimeProvisioningLeaseOwnership, +} from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-provisioning-lease-store"; +import type { RuntimeRunProvisioningLease } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-provisioning-lease-store"; +import { repairStaleRuntimeProvisioningLeases } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance.service"; +import { + ensureRuntimeSubjectId, + recordRuntimeConversationSessionClosed, +} from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store"; +import type { RuntimeProcessHandle } from "../src/modules/runtime/infrastructure/sandbox-handles"; +import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const DRIVER_ID = PLATFORM_ID_FIXTURES.driverInstance; +const NEW_CONVERSATION_ID = "01J00000000000000000000003"; +const OLD_CONVERSATION_ID = "01J00000000000000000000002"; +const RUN_ID = PLATFORM_ID_FIXTURES.sessionRun; +const SANDBOX_ID = PLATFORM_ID_FIXTURES.sandbox; +const SESSION_ID = PLATFORM_ID_FIXTURES.session; + +function createDatabase(options: { readonly insertSandbox?: boolean } = {}): SqliteD1Database { + const database = new SqliteD1Database({ foreignKeys: false }); + database.execute(` + CREATE TABLE session ( + id text PRIMARY KEY NOT NULL, + archived_at integer, + cleanup_operation_kind text, + last_run_id text, + runtime_provisioning_heartbeat_at integer, + runtime_provisioning_operation_id text, + runtime_provisioning_run_id text, + runtime_provisioning_sandbox_id text, + runtime_provisioning_sandbox_incarnation integer, + runtime_provisioning_sandbox_session_id text, + status text NOT NULL, + status_operation_id text, + CHECK ( + ( + runtime_provisioning_operation_id IS NULL + AND runtime_provisioning_run_id IS NULL + AND runtime_provisioning_sandbox_id IS NULL + AND runtime_provisioning_sandbox_incarnation IS NULL + AND runtime_provisioning_sandbox_session_id IS NULL + AND runtime_provisioning_heartbeat_at IS NULL + ) OR ( + runtime_provisioning_operation_id IS NOT NULL + AND runtime_provisioning_sandbox_id IS NOT NULL + AND runtime_provisioning_heartbeat_at IS NOT NULL + AND typeof(runtime_provisioning_heartbeat_at) = 'integer' + AND archived_at IS NULL + AND cleanup_operation_kind IS NULL + AND status_operation_id IS NULL + ) + ) + ); + CREATE TABLE session_run ( + driver_instance_id text, + id text PRIMARY KEY NOT NULL, + session_id text NOT NULL, + status text NOT NULL + ); + CREATE TABLE sandbox_session ( + cloudflare_session_id text NOT NULL DEFAULT '${OLD_CONVERSATION_ID}', + cleanup_operation_id text, + sandbox_id text NOT NULL, + sandbox_incarnation integer NOT NULL DEFAULT 1, + session_id text PRIMARY KEY NOT NULL, + status text NOT NULL, + updated_at integer NOT NULL DEFAULT 0 + ); + CREATE TABLE driver_instance ( + boot_token_expires_at integer, + boot_token_hash blob, + created_at integer NOT NULL DEFAULT 0, + expires_at integer, + generation integer NOT NULL DEFAULT 0, + heartbeat_count integer NOT NULL DEFAULT 0, + id text PRIMARY KEY NOT NULL, + protocol text NOT NULL DEFAULT 'orpc-ws', + protocol_version integer NOT NULL DEFAULT 3, + restart_count integer NOT NULL DEFAULT 0, + runtime text NOT NULL DEFAULT 'openai-runtime', + sandbox_id text NOT NULL, + sandbox_incarnation integer NOT NULL DEFAULT 1, + sandbox_session_id text NOT NULL, + status text NOT NULL, + status_changed_at integer DEFAULT 0 NOT NULL, + status_event text DEFAULT 'driver.provision' NOT NULL, + status_operation_id text, + status_seq integer DEFAULT 0 NOT NULL, + status_source text DEFAULT 'system' NOT NULL, + updated_at integer DEFAULT 0 NOT NULL + ); + CREATE TABLE sandbox ( + agent_id text, + app_id text, + bind_mount_ready integer NOT NULL DEFAULT 0, + claim_expires_at integer, + claim_owner text, + created_at integer NOT NULL DEFAULT 0, + global_mounts_json text NOT NULL DEFAULT '[]', + id text PRIMARY KEY NOT NULL, + inactive_deadline_at integer, + incarnation integer NOT NULL DEFAULT 1, + kind text, + last_backup_id text, + last_error text, + last_error_code text, + last_restore_backup_id text, + network_constraints_hash text, + operation_kind text, + owner_account_id text, + status text NOT NULL DEFAULT 'active', + status_changed_at integer NOT NULL DEFAULT 0, + status_event text NOT NULL DEFAULT 'runtime_subject.active', + status_operation_id text, + status_seq integer NOT NULL DEFAULT 0, + status_source text NOT NULL DEFAULT 'test', + subject_id text, + subject_kind text, + updated_at integer NOT NULL + ); + INSERT INTO session (id, last_run_id, status) + VALUES ('${SESSION_ID}', '${RUN_ID}', 'RUNNING'); + INSERT INTO session_run (id, session_id, status) + VALUES ('${RUN_ID}', '${SESSION_ID}', 'running'); + `); + if (options.insertSandbox !== false) { + database.execute(` + INSERT INTO sandbox ( + agent_id, app_id, id, kind, owner_account_id, subject_id, subject_kind, updated_at + ) VALUES ( + '${PLATFORM_ID_FIXTURES.agent}', '${PLATFORM_ID_FIXTURES.app}', '${SANDBOX_ID}', + 'cattle', '${PLATFORM_ID_FIXTURES.account}', '${SESSION_ID}', 'session', 0 + ); + `); + } + return database; +} + +async function insertDurableRunHandoff( + database: D1Database, + lease: RuntimeRunProvisioningLease, +): Promise { + await database + .prepare( + `INSERT INTO sandbox_session (sandbox_id, session_id, status) + VALUES (?, ?, 'active')`, + ) + .bind(SANDBOX_ID, SESSION_ID) + .run(); + await database + .prepare( + `INSERT INTO driver_instance (id, sandbox_id, sandbox_session_id, status) + VALUES (?, ?, ?, 'ready')`, + ) + .bind(DRIVER_ID, SANDBOX_ID, SESSION_ID) + .run(); + await database + .prepare("UPDATE session_run SET driver_instance_id = ? WHERE id = ?") + .bind(DRIVER_ID, RUN_ID) + .run(); + const targeted = await recordRuntimeProvisioningConversationTarget(database, { + lease, + sandboxIncarnation: 1, + sandboxSessionId: OLD_CONVERSATION_ID, + }); + if (targeted === null) { + throw new Error("Runtime provisioning fixture lost its target lease."); + } + return targeted; +} + +describe("runtime provisioning lease", () => { + test("allocates the lifecycle row before claiming initial provisioning", async () => { + const database = createDatabase({ insertSandbox: false }); + const sandboxId = await ensureRuntimeSubjectId(database, { + agentId: PLATFORM_ID_FIXTURES.agent, + appId: PLATFORM_ID_FIXTURES.app, + executionOwnerUserId: PLATFORM_ID_FIXTURES.account, + kind: "cattle", + runtimeSubjectId: SANDBOX_ID, + subjectId: SESSION_ID, + subjectKind: "session", + }); + + expect(sandboxId).toBe(SANDBOX_ID); + await expect( + ensureRuntimeSubjectId(database, { + agentId: PLATFORM_ID_FIXTURES.agent, + appId: PLATFORM_ID_FIXTURES.app, + executionOwnerUserId: "01J0000000000000000000000F", + kind: "cattle", + runtimeSubjectId: SANDBOX_ID, + subjectId: SESSION_ID, + subjectKind: "session", + }), + ).rejects.toThrow("identity does not match"); + expect( + await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId, + sessionId: SESSION_ID, + }), + ).not.toBeNull(); + }); + + test("rejects a half-written lease without a heartbeat", async () => { + const database = createDatabase(); + + await expect( + database + .prepare( + `UPDATE session + SET runtime_provisioning_operation_id = ?, + runtime_provisioning_sandbox_id = ? + WHERE id = ?`, + ) + .bind(PLATFORM_ID_FIXTURES.runtimeOperation, SANDBOX_ID, SESSION_ID) + .run(), + ).rejects.toThrow(); + }); + + test("uses stable provisioning identity across heartbeats and rejects a rotated owner", async () => { + const database = createDatabase(); + const claimed = await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(claimed).not.toBeNull(); + if (claimed === null) { + return; + } + await database + .prepare( + `INSERT INTO sandbox_session ( + cloudflare_session_id, sandbox_id, sandbox_incarnation, session_id, status + ) VALUES (?, ?, 1, ?, 'active')`, + ) + .bind(OLD_CONVERSATION_ID, SANDBOX_ID, SESSION_ID) + .run(); + const targeted = await recordRuntimeProvisioningConversationTarget(database, { + lease: claimed, + sandboxIncarnation: 1, + sandboxSessionId: OLD_CONVERSATION_ID, + }); + expect(targeted).not.toBeNull(); + if (targeted === null) { + return; + } + + expect(await heartbeatRuntimeRunProvisioningLease(database, targeted)).toBe(true); + expect( + await createDriverInstanceRecord({ DB: database } as ApiBindings, { + bootTokenHash: new Uint8Array([1]), + conflictStrategy: "insert-only", + driverInstanceId: DRIVER_ID, + runtime: "openai-runtime", + runtimeProvisioningLease: targeted, + sandboxId: SANDBOX_ID, + sandboxIncarnation: 1, + sandboxSessionId: SESSION_ID, + }), + ).toMatchObject({ generation: 0, status: "created" }); + + expect( + await runtimeProvisioningDriverLaunchIsOwned(database, { + bootTokenHash: new Uint8Array([1]), + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + lease: targeted, + }), + ).toBe(true); + + await database + .prepare("UPDATE session SET runtime_provisioning_heartbeat_at = 1 WHERE id = ?") + .bind(SESSION_ID) + .run(); + const maintenance = await claimStaleRuntimeProvisioningLeases(database, { + heartbeatAtLte: 1, + limit: 1, + }); + expect(maintenance).toHaveLength(1); + expect( + await runtimeProvisioningDriverLaunchIsOwned(database, { + bootTokenHash: new Uint8Array([1]), + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + lease: targeted, + }), + ).toBe(false); + expect( + await createDriverInstanceRecord({ DB: database } as ApiBindings, { + bootTokenHash: new Uint8Array([2]), + conflictStrategy: "insert-only", + driverInstanceId: createPlatformId(), + runtime: "openai-runtime", + runtimeProvisioningLease: targeted, + sandboxId: SANDBOX_ID, + sandboxIncarnation: 1, + sandboxSessionId: SESSION_ID, + }), + ).toEqual({ + bootTokenExpiresAt: null, + generation: null, + reason: "existing-driver", + status: "skipped", + }); + }); + + test("stops a late process after activation moves to the next incarnation", async () => { + let activeIncarnation = 1; + let disposeCalls = 0; + let killCalls = 0; + let startCalls = 0; + const process = { + [Symbol.dispose]: () => { + disposeCalls++; + }, + getStatus: async () => "running" as const, + kill: async () => { + killCalls++; + }, + } as unknown as RuntimeProcessHandle; + const assertOwned = async () => { + if (activeIncarnation !== 1) { + throw new Error("Runtime Driver provisioning lost launch ownership."); + } + }; + + await expect( + startProvisionProcessWithOwnershipFence({ + assertOwned, + context: { sandboxId: SANDBOX_ID }, + message: "test late process cleanup", + startProcess: async () => { + startCalls++; + activeIncarnation = 2; + return process; + }, + }), + ).rejects.toThrow("lost launch ownership"); + expect({ disposeCalls, killCalls, startCalls }).toEqual({ + disposeCalls: 1, + killCalls: 1, + startCalls: 1, + }); + + await expect( + startProvisionProcessWithOwnershipFence({ + assertOwned, + context: { sandboxId: SANDBOX_ID }, + message: "test pre-launch fence", + startProcess: async () => { + startCalls++; + return process; + }, + }), + ).rejects.toThrow("lost launch ownership"); + expect(startCalls).toBe(1); + }); + + test("fences cleanup until the conversation, Driver, and Run handoff are durable", async () => { + const database = createDatabase(); + let lease = await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(lease).not.toBeNull(); + if (lease === null) { + return; + } + + await expect( + database + .prepare( + "UPDATE session SET archived_at = 1, cleanup_operation_kind = 'archive' WHERE id = ?", + ) + .bind(SESSION_ID) + .run(), + ).rejects.toThrow(); + expect( + await releaseReadyRuntimeRunProvisioningLease(database, { + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + lease, + }), + ).toBe(false); + + lease = await insertDurableRunHandoff(database, lease); + expect( + await releaseReadyRuntimeRunProvisioningLease(database, { + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + lease, + }), + ).toBe(true); + expect( + await database + .prepare("SELECT runtime_provisioning_operation_id FROM session WHERE id = ?") + .bind(SESSION_ID) + .first(), + ).toEqual({ runtime_provisioning_operation_id: null }); + }); + + test("does not release a stale Driver generation handoff", async () => { + const database = createDatabase(); + let lease = await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(lease).not.toBeNull(); + if (lease === null) { + return; + } + lease = await insertDurableRunHandoff(database, lease); + await database + .prepare("UPDATE driver_instance SET generation = 1 WHERE id = ?") + .bind(DRIVER_ID) + .run(); + + expect( + await releaseReadyRuntimeRunProvisioningLease(database, { + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + lease, + }), + ).toBe(false); + expect(await heartbeatRuntimeRunProvisioningLease(database, lease)).toBe(true); + expect( + await releaseReadyRuntimeRunProvisioningLease(database, { + driverGeneration: 1, + driverInstanceId: DRIVER_ID, + lease, + }), + ).toBe(true); + }); + + test("cleanup ownership wins before provisioning without an external-work window", async () => { + const database = createDatabase(); + await database + .prepare( + `UPDATE session + SET archived_at = 1, + cleanup_operation_kind = 'archive', + status = 'RESCHEDULING', + status_operation_id = ? + WHERE id = ?`, + ) + .bind(PLATFORM_ID_FIXTURES.runtimeOperation, SESSION_ID) + .run(); + + expect( + await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }), + ).toBeNull(); + }); + + test("maintenance adopts a durable handoff instead of destroying a healthy Run", async () => { + const database = createDatabase(); + let lease = await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(lease).not.toBeNull(); + if (lease === null) { + return; + } + lease = await insertDurableRunHandoff(database, lease); + await database + .prepare("UPDATE session SET runtime_provisioning_heartbeat_at = 1 WHERE id = ?") + .bind(SESSION_ID) + .run(); + + const [maintenance] = await claimStaleRuntimeProvisioningLeases(database, { + heartbeatAtLte: 1, + limit: 1, + }); + expect(maintenance).toBeDefined(); + if (maintenance === undefined) { + return; + } + expect(await adoptReadyRuntimeRunProvisioningLease(database, maintenance)).toBe(true); + expect(await heartbeatRuntimeRunProvisioningLease(database, lease)).toBe(false); + expect(await releaseAbortedRuntimeProvisioningLease(database, lease)).toBe(false); + expect( + await database + .prepare("SELECT driver_instance_id, status FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first(), + ).toEqual({ driver_instance_id: DRIVER_ID, status: "running" }); + }); + + test("a cleanup claim makes the exact Driver generation non-assignable before remote I/O", async () => { + const database = createDatabase(); + let lease = await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(lease).not.toBeNull(); + if (lease === null) { + return; + } + lease = await insertDurableRunHandoff(database, lease); + await database + .prepare("UPDATE session SET runtime_provisioning_heartbeat_at = 1 WHERE id = ?") + .bind(SESSION_ID) + .run(); + const requestStarted = Promise.withResolvers(); + const response = Promise.withResolvers(); + const bindings = { + DB: database, + DriverConnection: { + get: () => ({ + fetch: async () => { + requestStarted.resolve(); + return response.promise; + }, + }), + idFromName: () => "driver-do-id", + }, + } as unknown as ApiBindings; + expect( + await claimRuntimeProvisioningDriverCleanup(database, { + driverGeneration: 0, + driverInstanceId: DRIVER_ID, + lease, + source: "maintenance", + }), + ).toBe(true); + const stopping = stopDriverSession(bindings, { + driverInstanceId: DRIVER_ID, + expectedDriverGeneration: 0, + expectedSessionRunId: RUN_ID, + reason: "runtime.provisioning_stale", + }); + await requestStarted.promise; + + await expect( + database + .prepare( + "SELECT status, status_operation_id, status_source FROM driver_instance WHERE id = ?", + ) + .bind(DRIVER_ID) + .first(), + ).resolves.toEqual({ + status: "stopping", + status_operation_id: expect.any(String), + status_source: "maintenance", + }); + const [maintenance] = await claimStaleRuntimeProvisioningLeases(database, { + heartbeatAtLte: 1, + limit: 1, + }); + expect(maintenance).toBeDefined(); + if (maintenance !== undefined) { + expect(await adoptReadyRuntimeRunProvisioningLease(database, maintenance)).toBe(false); + } + + response.resolve(Response.json({ error: "injected stop failure" }, { status: 500 })); + await expect(stopping).rejects.toThrow("injected stop failure"); + }); + + test("a stale cleanup cannot claim a Driver after maintenance adopts its handoff", async () => { + const database = createDatabase(); + let staleLease = await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(staleLease).not.toBeNull(); + if (staleLease === null) { + return; + } + staleLease = await insertDurableRunHandoff(database, staleLease); + const targets = await readRuntimeProvisioningCleanupTargets(database, staleLease); + const [driver] = targets?.driverInstances ?? []; + expect(driver).toBeDefined(); + if (driver === undefined) { + return; + } + await database + .prepare("UPDATE session SET runtime_provisioning_heartbeat_at = 1 WHERE id = ?") + .bind(SESSION_ID) + .run(); + const [maintenance] = await claimStaleRuntimeProvisioningLeases(database, { + heartbeatAtLte: 1, + limit: 1, + }); + expect(maintenance).toBeDefined(); + if (maintenance === undefined) { + return; + } + expect(await adoptReadyRuntimeRunProvisioningLease(database, maintenance)).toBe(true); + + expect( + await claimRuntimeProvisioningDriverCleanup(database, { + driverGeneration: driver.generation, + driverInstanceId: driver.id, + lease: staleLease, + source: "maintenance", + }), + ).toBe(false); + await expect( + database + .prepare("SELECT status, status_operation_id FROM driver_instance WHERE id = ?") + .bind(DRIVER_ID) + .first(), + ).resolves.toEqual({ status: "ready", status_operation_id: null }); + await expect( + database + .prepare("SELECT driver_instance_id, status FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ driver_instance_id: DRIVER_ID, status: "running" }); + }); + + test("the maintenance entry point releases a crashed post-handoff lease", async () => { + const database = createDatabase(); + let lease = await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(lease).not.toBeNull(); + if (lease === null) { + return; + } + lease = await insertDurableRunHandoff(database, lease); + await database + .prepare("UPDATE session SET runtime_provisioning_heartbeat_at = 1 WHERE id = ?") + .bind(SESSION_ID) + .run(); + + expect( + await repairStaleRuntimeProvisioningLeases({ DB: database } as ApiBindings, { + heartbeatAtLte: 1, + limit: 1, + }), + ).toBe(1); + expect( + await database + .prepare("SELECT runtime_provisioning_operation_id FROM session WHERE id = ?") + .bind(SESSION_ID) + .first(), + ).toEqual({ runtime_provisioning_operation_id: null }); + }); + + test("maintenance takeover revokes the old holder before clearing an incomplete attempt", async () => { + const database = createDatabase(); + const lease = await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(lease).not.toBeNull(); + if (lease === null) { + return; + } + await database + .prepare("UPDATE session SET runtime_provisioning_heartbeat_at = 1 WHERE id = ?") + .bind(SESSION_ID) + .run(); + + const [maintenance] = await claimStaleRuntimeProvisioningLeases(database, { + heartbeatAtLte: 1, + limit: 1, + }); + expect(maintenance).toBeDefined(); + if (maintenance === undefined) { + return; + } + expect(maintenance.operationId).not.toBe(lease.operationId); + expect(await heartbeatRuntimeRunProvisioningLease(database, lease)).toBe(false); + expect(await releaseAbortedRuntimeProvisioningLease(database, lease)).toBe(false); + expect(await releaseAbortedRuntimeProvisioningLease(database, maintenance)).toBe(true); + + const nextLease = await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(nextLease).not.toBeNull(); + expect(await renewRuntimeProvisioningLeaseOwnership(database, lease)).toBe(false); + if (nextLease !== null) { + expect(await heartbeatRuntimeRunProvisioningLease(database, nextLease)).toBe(true); + } + }); + + test("a late old cleanup cannot close the next owner's conversation", async () => { + const database = createDatabase(); + let lease = await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(lease).not.toBeNull(); + if (lease === null) { + return; + } + await database + .prepare( + `INSERT INTO sandbox_session + (cloudflare_session_id, sandbox_id, session_id, status) + VALUES (?, ?, ?, 'active')`, + ) + .bind(OLD_CONVERSATION_ID, SANDBOX_ID, SESSION_ID) + .run(); + const targetedLease = await recordRuntimeProvisioningConversationTarget(database, { + lease, + sandboxIncarnation: 1, + sandboxSessionId: OLD_CONVERSATION_ID, + }); + expect(targetedLease).not.toBeNull(); + if (targetedLease === null) { + return; + } + lease = targetedLease; + expect(await readRuntimeProvisioningCleanupTargets(database, lease)).toMatchObject({ + conversationSessionId: OLD_CONVERSATION_ID, + }); + await database + .prepare("UPDATE session SET runtime_provisioning_heartbeat_at = 1 WHERE id = ?") + .bind(SESSION_ID) + .run(); + + const [maintenance] = await claimStaleRuntimeProvisioningLeases(database, { + heartbeatAtLte: 1, + limit: 1, + }); + expect(maintenance).toBeDefined(); + if (maintenance === undefined) { + return; + } + expect(await releaseAbortedRuntimeProvisioningLease(database, maintenance)).toBe(true); + const nextLease = await claimRuntimeRunProvisioningLease(database, { + runId: RUN_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(nextLease).not.toBeNull(); + await database + .prepare( + `UPDATE sandbox_session + SET cloudflare_session_id = ?, status = 'active' + WHERE session_id = ?`, + ) + .bind(NEW_CONVERSATION_ID, SESSION_ID) + .run(); + + await recordRuntimeConversationSessionClosed(database, { + expectedProvisioningOperationId: lease.operationId, + inactiveDeadlineAt: 100, + now: 10, + runtimeSubjectId: SANDBOX_ID, + sandboxSessionId: OLD_CONVERSATION_ID, + sessionId: SESSION_ID, + }); + + expect( + await database + .prepare("SELECT cloudflare_session_id, status FROM sandbox_session WHERE session_id = ?") + .bind(SESSION_ID) + .first(), + ).toEqual({ cloudflare_session_id: NEW_CONVERSATION_ID, status: "active" }); + if (nextLease !== null) { + expect(await heartbeatRuntimeRunProvisioningLease(database, nextLease)).toBe(true); + } + }); +}); diff --git a/apps/api/tests/runtime-session-link.test.ts b/apps/api/tests/runtime-session-link.test.ts index 8fc36fa3..7b11ec1d 100644 --- a/apps/api/tests/runtime-session-link.test.ts +++ b/apps/api/tests/runtime-session-link.test.ts @@ -18,8 +18,11 @@ function createRuntimeSessionLinkDatabase(): SqliteD1Database { id text PRIMARY KEY NOT NULL, session_id text NOT NULL, created_by_account_id text NOT NULL, + created_at integer NOT NULL, + runtime_id text, trace_id text, - status text NOT NULL + status text NOT NULL, + updated_at integer NOT NULL ); CREATE TABLE session ( @@ -27,6 +30,7 @@ function createRuntimeSessionLinkDatabase(): SqliteD1Database { agent_id text NOT NULL, app_id text, creator_account_id text NOT NULL, + runtime_id text NOT NULL, type text DEFAULT 'preview' NOT NULL ); @@ -49,11 +53,20 @@ function createRuntimeSessionLinkDatabase(): SqliteD1Database { INSERT INTO agent (id, owner_account_id) VALUES ('01J00000000000000000000009', '01J00000000000000000000001'); - INSERT INTO session (id, agent_id, creator_account_id) - VALUES ('session-1', '01J00000000000000000000009', 'creator-1'); - - INSERT INTO session_run (driver_instance_id, id, session_id, created_by_account_id, trace_id, status) - VALUES ('driver-1', 'run-1', 'session-1', 'caller-1', 'trace-1', 'running'); + INSERT INTO session (id, agent_id, creator_account_id, runtime_id) + VALUES ('session-1', '01J00000000000000000000009', 'creator-1', 'runtime-1'); + + INSERT INTO session_run ( + driver_instance_id, + id, + session_id, + created_by_account_id, + created_at, + trace_id, + status, + updated_at + ) + VALUES ('driver-1', 'run-1', 'session-1', 'caller-1', 1, 'trace-1', 'running', 1); INSERT INTO sandbox (id, kind, subject_kind) VALUES ('01J0000000000000000000000D', 'pet', 'agent'); diff --git a/apps/api/tests/runtime-session-outputs.test.ts b/apps/api/tests/runtime-session-outputs.test.ts index 1e71ac03..32868da2 100644 --- a/apps/api/tests/runtime-session-outputs.test.ts +++ b/apps/api/tests/runtime-session-outputs.test.ts @@ -1,213 +1,138 @@ import { describe, expect, test } from "bun:test"; +import { createPlatformId } from "@mosoo/id"; +import type { RuntimeEventId } from "@mosoo/id"; import { createRuntimeEvent } from "@mosoo/runtime-events"; -import { createBaseLiveState } from "../src/modules/runtime/infrastructure/driver-instance/event-projection"; +import { fileStore } from "../src/modules/files/application/file-store"; import type { RuntimeSessionLink } from "../src/modules/runtime/infrastructure/driver-instance/event-types"; -import { appRuntimeDriverEvents } from "../src/modules/runtime/infrastructure/driver-instance/events"; +import { DriverInstanceRpcEventIngestionController } from "../src/modules/runtime/infrastructure/driver-instance/rpc-event-ingestion-controller"; import { getRuntimeSessionOutputDirectory, normalizeRuntimeSessionOutputRelativePath, - readRuntimeSessionOutputListing, + readRuntimeSessionOutputInventory, toRuntimeSessionOutputFile, } from "../src/modules/runtime/infrastructure/driver-instance/runtime-session-outputs"; -import type { SandboxHandle } from "../src/modules/runtime/infrastructure/sandbox-handles"; +import { RuntimeSessionViewCache } from "../src/modules/runtime/infrastructure/driver-instance/runtime-session-view-cache"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; -import { API_DRIVER_BOUNDARY_IDS } from "./api-driver-boundary-fixtures"; import { PUBLIC_API_TEST_IDS, PublicApiMemoryFileBucket, createPublicHttpContractDatabase, createPublicHttpTestBindings, + insertActiveSandboxSessionFixture, insertOwnerSession, nowMsForTest, } from "./helpers/public-api-http-test-fixture"; +import { createRuntimeOutputSandbox } from "./helpers/runtime-output-sandbox"; -const encoder = new TextEncoder(); - -function encodeBase64(value: string): string { - const bytes = encoder.encode(value); - let binary = ""; - - for (const byte of bytes) { - binary += String.fromCodePoint(byte); - } - - return btoa(binary); -} - -function createSandboxHandle(files: ReadonlyMap): SandboxHandle { - const unavailable = async () => { - throw new Error("Unexpected sandbox test method call."); - }; - const successfulCommand: SandboxHandle["exec"] = async (command) => { - if (command.includes("find . -type f")) { - const outputDir = "/workspace/session/outputs/"; - const output = [...files.keys()] - .filter((path) => path.startsWith(outputDir)) - .map((path) => path.slice(outputDir.length)) - .toSorted() - .join("\n"); - - return { - exitCode: 0, - stderr: "", - stdout: output.length === 0 ? "" : `${output}\n`, - success: true, - }; - } - - return { - exitCode: 0, - stderr: "", - stdout: "", - success: true, - }; - }; - const readFile: SandboxHandle["readFile"] = async (path) => { - const content = files.get(path); - - if (content === undefined) { - throw new Error(`Missing sandbox test file: ${path}`); - } - - return { - content: encodeBase64(content), - encoding: "base64", - }; - }; - - return { - configureNetworkConstraints: unavailable, - createBackup: unavailable, - createSession: unavailable, - deleteSession: unavailable, - destroy: unavailable, - exec: successfulCommand, - getSession: async () => ({ - exec: successfulCommand, - mkdir: unavailable, - readFile, - startProcess: unavailable, - watch: unavailable, - writeFile: unavailable, - }), - mkdir: unavailable, - mountBucket: unavailable, - readFile, - restoreBackup: unavailable, - setKeepAlive: unavailable, - startProcess: unavailable, - terminal: unavailable, - unmountBucket: unavailable, - watch: unavailable, - writeFile: unavailable, - wsConnect: unavailable, - }; -} - -async function insertActiveSandboxSession(database: D1Database): Promise { +async function insertActiveRuntime(database: D1Database): Promise { + const now = nowMsForTest(); await database .prepare( - ` - INSERT INTO sandbox_session ( - cloudflare_session_id, - created_at, - cwd, - origin_json, - sandbox_id, - session_id, - status, - updated_at - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) - `, + `INSERT INTO driver_instance ( + id, sandbox_id, sandbox_incarnation, sandbox_session_id, runtime, protocol, protocol_version, + status, status_changed_at, status_event, status_seq, status_source, + connection_id, command_seq_cursor, boot_token_hash, boot_token_expires_at, + generation, heartbeat_count, restart_count, expires_at, created_at, updated_at + ) VALUES (?, ?, 1, ?, 'openai-runtime', 'acp', 3, 'ready', ?, 'driver.ready', 1, + 'driver', 'connection-1', 0, X'00', ?, 1, 1, 0, ?, ?, ?)`, ) .bind( - API_DRIVER_BOUNDARY_IDS.sandboxSession, - nowMsForTest(), - "/workspace/session", - "{}", + PUBLIC_API_TEST_IDS.driverOwner, PUBLIC_API_TEST_IDS.sandbox, PUBLIC_API_TEST_IDS.ownerSession, - "active", - nowMsForTest(), + now, + now + 60_000, + now + 60_000, + now, + now, ) .run(); -} - -function createRuntimeLink(): RuntimeSessionLink { - return { - agentId: PUBLIC_API_TEST_IDS.agent, - appId: PUBLIC_API_TEST_IDS.app, - callerId: PUBLIC_API_TEST_IDS.ownerAccount, - creatorId: PUBLIC_API_TEST_IDS.ownerAccount, - executionOwnerId: PUBLIC_API_TEST_IDS.ownerAccount, - sandboxId: PUBLIC_API_TEST_IDS.sandbox, - sandboxKind: "cattle", - sandboxSubjectKind: "session", - sessionId: PUBLIC_API_TEST_IDS.ownerSession, - sessionRunId: PUBLIC_API_TEST_IDS.run, - sessionRunStatus: "running", - sessionType: "ui", - traceId: "trace-session-outputs", - }; + await database + .prepare( + `INSERT INTO session_run ( + id, session_id, agent_id, created_by_account_id, driver_instance_id, + trigger, status, provider, model, runtime_id, trace_id, created_at, + status_changed_at, status_event, status_seq, status_source, updated_at + ) VALUES (?, ?, ?, ?, ?, 'user_prompt', 'running', 'openai', 'gpt-5.4', + 'openai-runtime', 'trace-session-outputs', ?, ?, 'run.start', 1, 'driver', ?)`, + ) + .bind( + PUBLIC_API_TEST_IDS.run, + PUBLIC_API_TEST_IDS.ownerSession, + PUBLIC_API_TEST_IDS.agent, + PUBLIC_API_TEST_IDS.ownerAccount, + PUBLIC_API_TEST_IDS.driverOwner, + now, + now, + now, + ) + .run(); + await database + .prepare( + "UPDATE session SET last_run_id = ?, status = 'RUNNING', status_operation_id = NULL WHERE id = ?", + ) + .bind(PUBLIC_API_TEST_IDS.run, PUBLIC_API_TEST_IDS.ownerSession) + .run(); } function createCompletedRunEvent() { return createRuntimeEvent({ driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner, - id: API_DRIVER_BOUNDARY_IDS.runtimeEvent, + id: createPlatformId(), kind: "run.completed", - occurredAt: "2026-06-22T00:00:00.000Z", - payload: { - run: { - completedAt: "2026-06-22T00:00:00.000Z", - error: null, - startedAt: null, - status: "completed", - }, - }, + occurredAt: "2026-06-22T00:00:01.000Z", + payload: { stopReason: "end_turn" }, runId: PUBLIC_API_TEST_IDS.run, + runtimeId: "openai-runtime", sessionId: PUBLIC_API_TEST_IDS.ownerSession, }); } -function createFileChangedEvent(path = "outputs/live.txt") { +interface TestFileChange { + readonly change: "delete" | "upsert"; + readonly metadata?: { readonly contentType?: string }; + readonly path: string; +} + +function createFileChangedEvent(...changes: readonly TestFileChange[]) { return createRuntimeEvent({ driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner, - id: API_DRIVER_BOUNDARY_IDS.runtimeEvent, + id: createPlatformId(), kind: "file.changed", occurredAt: "2026-06-22T00:00:00.000Z", - payload: { - changes: [ - { - change: "upsert", - metadata: { contentType: "text/plain" }, - path, - }, - { - change: "upsert", - path: "src/temp.txt", - }, - ], - }, + payload: { changes }, runId: PUBLIC_API_TEST_IDS.run, + runtimeId: "openai-runtime", sessionId: PUBLIC_API_TEST_IDS.ownerSession, }); } -async function createBindings(input: { - database: D1Database; - files?: ReadonlyMap; -}): Promise<{ bindings: ApiBindings; bucket: PublicApiMemoryFileBucket }> { +function fileUpsert(path: string, contentType = "text/plain"): TestFileChange { + return { change: "upsert", metadata: { contentType }, path }; +} + +function createBindings(input: { database: D1Database; files?: ReadonlyMap }): { + bindings: ApiBindings; + bucket: PublicApiMemoryFileBucket; +} { const bucket = new PublicApiMemoryFileBucket(); - const sandbox = createSandboxHandle(input.files ?? new Map()); + const sandbox = createRuntimeOutputSandbox({ + files: input.files, + onExec: (command) => { + if (command.includes("find . -type f")) { + expect(command.startsWith("bash -lc ")).toBe(true); + expect(command.indexOf("head -z")).toBeLessThan(command.indexOf("sort -z")); + } + }, + root: "/workspace/session/outputs", + }); return { bindings: { ...createPublicHttpTestBindings(input.database, { - fileBucket: bucket as unknown as R2Bucket, + fileBucket: bucket, }), runtimeSubjectHandleFactory: () => sandbox, } as ApiBindings, @@ -215,41 +140,125 @@ async function createBindings(input: { }; } +const activeContext = { + assertActiveConnection: () => undefined, + connectionId: "connection-1", +} as never; + +function createController(bindings: ApiBindings): DriverInstanceRpcEventIngestionController { + const state = { + hello: { pid: 1 }, + requireDriverGeneration: () => 1, + requireDriverInstanceId: () => PUBLIC_API_TEST_IDS.driverOwner, + runtimeSessionLink: null as RuntimeSessionLink | null, + setRuntimeSessionLink(link: RuntimeSessionLink) { + this.runtimeSessionLink = link; + }, + }; + + return new DriverInstanceRpcEventIngestionController({ + env: bindings, + state, + viewCache: new RuntimeSessionViewCache(), + viewerEventDelivery: { + enqueue: () => undefined, + flush: async () => undefined, + flushSafely: async () => undefined, + requestStateSync: () => undefined, + resetAfterFlush: () => undefined, + }, + } as never); +} + +async function createRuntimeSessionOutputFixture( + files: ReadonlyMap = new Map(), +): Promise<{ + readonly bindings: ApiBindings; + readonly bucket: PublicApiMemoryFileBucket; + readonly database: D1Database; +}> { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await insertActiveSandboxSessionFixture(database, { + cwd: "/workspace/session", + ownerAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + }); + await insertActiveRuntime(database); + + return { database, ...createBindings({ database, files }) }; +} + async function dispatchRuntimeEvent(input: { bindings: ApiBindings; event: ReturnType; - link: RuntimeSessionLink; }): Promise { - await appRuntimeDriverEvents(input.bindings, { - currentLiveState: createBaseLiveState({ - callerId: input.link.callerId, - creatorId: input.link.creatorId, + const eventId = input.event.sourceEventId ?? input.event.id; + + await createController(input.bindings).handlePushEvents( + { driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner, - sessionId: input.link.sessionId, - }), - driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner, - events: [ - { - event: input.event, - eventId: "source-runtime-session-outputs", - occurredAt: 1, - }, - ], - link: input.link, - }); + events: [{ event: input.event, eventId, occurredAt: input.event.occurredAt }], + }, + activeContext, + ); } describe("runtime session outputs", () => { + test("preserves exact binary bytes in the memory R2 fixture", async () => { + const bucket = new PublicApiMemoryFileBucket(); + const source = Uint8Array.of(0xaa, 0xff, 0, 0x80, 0x41, 0xbb); + const expected = Uint8Array.of(0xff, 0, 0x80, 0x41); + const customMetadata = { contentSha256: "binary-hash", sourcePath: "outputs/data.bin" }; + + const stored = await bucket.put("binary", source.subarray(1, -1), { + customMetadata, + httpMetadata: { contentType: "application/octet-stream" }, + }); + source.fill(0); + const body = await bucket.get("binary"); + + expect(stored).toMatchObject({ customMetadata, size: expected.byteLength }); + expect(body).toMatchObject({ customMetadata, size: expected.byteLength }); + expect(new Uint8Array(await body!.arrayBuffer())).toEqual(expected); + }); + + test.each([ + ["If-None-Match", false], + ["If-Match", true], + ] as const)("serializes concurrent %s writes in the memory R2 fixture", async (name, seeded) => { + const bucket = new PublicApiMemoryFileBucket(); + const key = "conditional"; + const previous = seeded ? await bucket.put(key, "previous") : null; + const onlyIf = new Headers({ [name]: seeded ? previous!.httpEtag : "*" }); + const writes = await Promise.all([ + bucket.put(key, "first", { onlyIf }), + bucket.put(key, "second", { onlyIf }), + ]); + const winner = writes.find((object) => object !== null); + const stored = await bucket.get(key); + + expect(writes.filter((object) => object !== null)).toHaveLength(1); + expect(stored?.etag).toBe(winner?.etag); + expect(["first", "second"]).toContain(await stored?.text()); + }); + test("normalizes the session output directory contract", () => { expect(getRuntimeSessionOutputDirectory("/workspace/session")).toBe( "/workspace/session/outputs", ); - expect(normalizeRuntimeSessionOutputRelativePath("a/./b.txt")).toBe("a/b.txt"); + expect(normalizeRuntimeSessionOutputRelativePath("a/b.txt")).toBe("a/b.txt"); + expect(normalizeRuntimeSessionOutputRelativePath("a/./b.txt")).toBeNull(); expect(normalizeRuntimeSessionOutputRelativePath("../b.txt")).toBeNull(); expect(normalizeRuntimeSessionOutputRelativePath("/tmp/b.txt")).toBeNull(); - expect(readRuntimeSessionOutputListing("b.txt\n../secret.txt\nnested/a.pdf\n")).toEqual([ - "b.txt", - "nested/a.pdf", + expect( + readRuntimeSessionOutputInventory( + ["./b.txt", "12", "./nested/report.pdf", "7", ""].join("\0"), + ), + ).toEqual([ + { relativePath: "b.txt", size: 12 }, + { relativePath: "nested/report.pdf", size: 7 }, ]); expect( toRuntimeSessionOutputFile({ @@ -276,24 +285,39 @@ describe("runtime session outputs", () => { ).toBeNull(); }); - test("records files under outputs as session artifacts on run completion", async () => { - const database = await createPublicHttpContractDatabase(); - await insertOwnerSession(database); - await insertActiveSandboxSession(database); + test.each([ + ["control character", "nested/tab\tline\n.txt"], + ["trailing control character", "nested/report.txt\n"], + ["backslash", "nested/back\\slash.txt"], + ["non-canonical whitespace", "nested/ report.txt"], + ["encoded traversal", "nested/%2e%2e/report.txt"], + ["dot alias", "nested/./report.txt"], + ["empty segment", "nested//report.txt"], + ] as const)("rejects a %s that file records cannot represent", (_name, path) => { + expect(normalizeRuntimeSessionOutputRelativePath(path)).toBeNull(); + expect(() => readRuntimeSessionOutputInventory([`./${path}`, "1", ""].join("\0"))).toThrow( + "Runtime output inventory is invalid.", + ); + expect( + toRuntimeSessionOutputFile({ cwd: "/workspace/session", path: `outputs/${path}` }), + ).toBeNull(); + }); - const { bindings, bucket } = await createBindings({ - database, - files: new Map([ + test("records files under outputs as session artifacts from file changes", async () => { + const { bindings, bucket, database } = await createRuntimeSessionOutputFixture( + new Map([ ["/workspace/session/outputs/resume.txt", "Improved resume"], ["/workspace/session/outputs/nested/summary.md", "# Summary"], ]), - }); - const link = createRuntimeLink(); + ); await dispatchRuntimeEvent({ bindings, - event: createCompletedRunEvent(), - link, + event: createFileChangedEvent(fileUpsert("outputs/resume.txt")), + }); + await dispatchRuntimeEvent({ + bindings, + event: createFileChangedEvent(fileUpsert("outputs/nested/summary.md", "text/markdown")), }); const rows = await database @@ -341,23 +365,144 @@ describe("runtime session outputs", () => { size: 9, }, ]); - expect([...bucket.objects.values()]).toHaveLength(2); + expect( + [...bucket.objects.values()].filter((object) => !object.key.endsWith("/manifest.json")), + ).toHaveLength(2); }); - test("deduplicates runtime outputs by source path and content", async () => { - const database = await createPublicHttpContractDatabase(); - await insertOwnerSession(database); - await insertActiveSandboxSession(database); + test("atomically replaces delta artifact heads with the completed Run snapshot", async () => { + const files = new Map([ + ["/workspace/session/outputs/current.txt", "delta version"], + ["/workspace/session/outputs/obsolete.txt", "removed before completion"], + ]); + const { bindings, bucket, database } = await createRuntimeSessionOutputFixture(files); + + await dispatchRuntimeEvent({ + bindings, + event: createFileChangedEvent(fileUpsert("outputs/current.txt")), + }); + await dispatchRuntimeEvent({ + bindings, + event: createFileChangedEvent(fileUpsert("outputs/obsolete.txt")), + }); + + files.set("/workspace/session/outputs/current.txt", "terminal version"); + files.delete("/workspace/session/outputs/obsolete.txt"); + await dispatchRuntimeEvent({ bindings, event: createCompletedRunEvent() }); + + const heads = await database + .prepare( + `SELECT h.file_id, h.runtime_event_seq, h.source_path, f.object_key + FROM session_artifact_head AS h + LEFT JOIN file_record AS f ON f.id = h.file_id + WHERE h.session_id = ? + ORDER BY h.source_path`, + ) + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .all<{ + file_id: string | null; + object_key: string | null; + runtime_event_seq: number; + source_path: string; + }>(); + + expect( + heads.results.map(({ file_id, runtime_event_seq, source_path }) => ({ + file_id: file_id === null ? null : "present", + runtime_event_seq, + source_path, + })), + ).toEqual([ + { + file_id: "present", + runtime_event_seq: 3, + source_path: "outputs/current.txt", + }, + { + file_id: null, + runtime_event_seq: 3, + source_path: "outputs/obsolete.txt", + }, + ]); + expect(await (await bucket.get(heads.results[0]?.object_key ?? ""))?.text()).toBe( + "terminal version", + ); + expect( + await database + .prepare( + `SELECT count(*) AS count + FROM runtime_artifact_attempt + WHERE status = 'accepted' AND owned_object_keys_json = '[]'`, + ) + .first(), + ).toEqual({ count: 3 }); + }); + + test("does not resurrect a late headless legacy artifact after an empty snapshot", async () => { + const { bindings, database } = await createRuntimeSessionOutputFixture(); + const staleFileId = createPlatformId(); + await database + .prepare( + `INSERT INTO file_record ( + committed, created_at, created_by_account_id, id, name, object_key, + owner_id, owner_kind, parent_path, path, purpose, scope_id, scope_kind, + session_kind, size, status, updated_at, version + ) VALUES ( + 1, ?, ?, ?, 'stale.txt', ?, ?, 'session', ?, ?, 'session_artifact', + ?, 'session', 'artifact', 5, 'ready', ?, 1 + )`, + ) + .bind( + nowMsForTest(), + PUBLIC_API_TEST_IDS.ownerAccount, + staleFileId, + `objects/${staleFileId}`, + PUBLIC_API_TEST_IDS.ownerSession, + `runtime-output/outputs/stale.txt/${"a".repeat(64)}`, + `session-artifacts/${staleFileId}/stale.txt`, + PUBLIC_API_TEST_IDS.ownerSession, + nowMsForTest(), + ) + .run(); + await dispatchRuntimeEvent({ + bindings, + event: createCompletedRunEvent(), + }); + + await expect( + database + .prepare( + `SELECT json_extract(manifest_json, '$.captureStatus') AS capture_status, + json_extract(manifest_json, '$.mode') AS mode, + status + FROM runtime_artifact_attempt`, + ) + .first(), + ).resolves.toEqual({ capture_status: "complete", mode: "snapshot", status: "accepted" }); + expect( + await fileStore.listReadySessionFiles(database, PUBLIC_API_TEST_IDS.ownerSession), + ).toEqual([]); + expect( + await fileStore.listLatestReadySessionArtifactSources( + database, + PUBLIC_API_TEST_IDS.ownerSession, + ), + ).toEqual([]); + await expect( + database + .prepare("SELECT count(*) AS count FROM file_record WHERE id = ?") + .bind(staleFileId) + .first(), + ).resolves.toEqual({ count: 1 }); + }); + + test("deduplicates runtime outputs by source path and content", async () => { const files = new Map([ ["/workspace/session/outputs/one/report.txt", "alpha"], ["/workspace/session/outputs/two/report.txt", "bravo"], ]); - const { bindings } = await createBindings({ - database, - files, - }); - const link = createRuntimeLink(); + const { bindings, database } = await createRuntimeSessionOutputFixture(files); const readReportCount = async () => { const row = await database .prepare( @@ -370,14 +515,12 @@ describe("runtime session outputs", () => { await dispatchRuntimeEvent({ bindings, - event: createFileChangedEvent("outputs/one/report.txt"), - link, + event: createFileChangedEvent(fileUpsert("outputs/one/report.txt")), }); await dispatchRuntimeEvent({ bindings, - event: createCompletedRunEvent(), - link, + event: createFileChangedEvent(fileUpsert("outputs/two/report.txt")), }); expect(await readReportCount()).toBe(2); @@ -386,64 +529,65 @@ describe("runtime session outputs", () => { await dispatchRuntimeEvent({ bindings, - event: createFileChangedEvent("outputs/one/report.txt"), - link, + event: createFileChangedEvent(fileUpsert("outputs/one/report.txt")), }); expect(await readReportCount()).toBe(3); }); - test("records file change events only when the path is under outputs", async () => { - const database = await createPublicHttpContractDatabase(); - await insertOwnerSession(database); - await insertActiveSandboxSession(database); + test.each([ + ["relative output path", "outputs/live.txt", [{ name: "live.txt", size: 11 }]], + ["external relative path", "src/temp.txt", []], + [ + "absolute output path", + "/workspace/session/outputs/live.txt", + [{ name: "live.txt", size: 11 }], + ], + ["external absolute path", "/workspace/session/src/temp.txt", []], + ["output traversal", "outputs/../src/temp.txt", []], + ] as const)( + "keeps artifact writes inside outputs: %s", + async (_label, eventPath, expectedRows) => { + const { bindings, bucket, database } = await createRuntimeSessionOutputFixture( + new Map([ + ["/workspace/session/outputs/live.txt", "download me"], + ["/workspace/session/src/temp.txt", "ignore me"], + ]), + ); + + await dispatchRuntimeEvent({ + bindings, + event: createFileChangedEvent(fileUpsert(eventPath)), + }); + + const rows = await database + .prepare("SELECT name, size FROM file_record WHERE session_kind = 'artifact'") + .all<{ name: string; size: number }>(); + const artifactObjects = [...bucket.objects.values()].filter( + (object) => !object.key.endsWith("/manifest.json"), + ); + + expect(rows.results).toEqual(expectedRows); + expect(artifactObjects).toHaveLength(expectedRows.length); + }, + ); - const { bindings } = await createBindings({ - database, - files: new Map([ - ["/workspace/session/outputs/live.txt", "download me"], - ["/workspace/session/src/temp.txt", "ignore me"], - ]), - }); - const link = createRuntimeLink(); + test("replays an artifact receipt without reopening its sandbox path", async () => { + const files = new Map([["/workspace/session/outputs/result.txt", "R1 result"]]); + const { bindings, bucket, database } = await createRuntimeSessionOutputFixture(files); + const event = createFileChangedEvent(fileUpsert("outputs/result.txt")); - await dispatchRuntimeEvent({ - bindings, - event: createFileChangedEvent(), - link, - }); + await dispatchRuntimeEvent({ bindings, event }); + files.set("/workspace/session/outputs/result.txt", "R2 result"); + files.set("/workspace/session/outputs/later.txt", "created by R2"); + await dispatchRuntimeEvent({ bindings, event }); const rows = await database - .prepare("SELECT name, size FROM file_record WHERE session_kind = 'artifact'") - .all<{ name: string; size: number }>(); - - expect(rows.results).toEqual([ - { - name: "live.txt", - size: 11, - }, - ]); - }); - - test("skips optional output directory when it has no files", async () => { - const database = await createPublicHttpContractDatabase(); - await insertOwnerSession(database); - await insertActiveSandboxSession(database); - - const { bindings, bucket } = await createBindings({ database }); - const link = createRuntimeLink(); - - await dispatchRuntimeEvent({ - bindings, - event: createCompletedRunEvent(), - link, - }); - - const row = await database - .prepare("SELECT count(*) AS count FROM file_record") - .first<{ count: number }>(); - - expect(row?.count).toBe(0); - expect([...bucket.objects.values()]).toEqual([]); + .prepare( + "SELECT name, object_key FROM file_record WHERE session_kind = 'artifact' ORDER BY name", + ) + .all<{ name: string; object_key: string }>(); + expect(rows.results.map(({ name }) => name)).toEqual(["result.txt"]); + expect(await (await bucket.get(rows.results[0]?.object_key ?? ""))?.text()).toBe("R1 result"); }); }); diff --git a/apps/api/tests/runtime-state-operation-execution.test.ts b/apps/api/tests/runtime-state-operation-execution.test.ts index c11e2499..c1712834 100644 --- a/apps/api/tests/runtime-state-operation-execution.test.ts +++ b/apps/api/tests/runtime-state-operation-execution.test.ts @@ -39,4 +39,44 @@ describe("runtime state operation execution", () => { expect(startedRuntimeSubjectIds).toEqual(["01J0000000000000000000000D", "sandbox-2"]); expect(operationIds).toEqual(["01J0000000000000000000000R", "01J0000000000000000000000R"]); }); + + test("joins every started subject before exposing a sibling failure", async () => { + const delayed = Promise.withResolvers(); + let delayedSubjectFinished = false; + let operationRejected = false; + const plane: RuntimeStateOperationExecutionPlane = { + async recreateSubjectPreservingState() { + throw new Error("Unexpected recreate operation."); + }, + async resetSubjectAgentState() { + throw new Error("Unexpected reset operation."); + }, + async stopSubjectDrivers(_bindings, input) { + if (input.runtimeSubjectId === "failing-subject") { + throw new Error("subject failed"); + } + await delayed.promise; + delayedSubjectFinished = true; + }, + }; + const operation = executeRuntimeStateOperationSubjects({} as ApiBindings, { + executionPlane: plane, + operation: "restartDriver", + operationId: "01J0000000000000000000000R", + subjects: [ + { runtimeSubjectId: "failing-subject", targets: [] }, + { runtimeSubjectId: "delayed-subject", targets: [] }, + ], + }).catch((error: unknown) => { + operationRejected = true; + throw error; + }); + + await Promise.resolve(); + expect(operationRejected).toBe(false); + delayed.resolve(); + + await expect(operation).rejects.toThrow("subject failed"); + expect(delayedSubjectFinished).toBe(true); + }); }); diff --git a/apps/api/tests/runtime-state-operation-phases.test.ts b/apps/api/tests/runtime-state-operation-phases.test.ts index 5a785773..853dbbeb 100644 --- a/apps/api/tests/runtime-state-operation-phases.test.ts +++ b/apps/api/tests/runtime-state-operation-phases.test.ts @@ -1,28 +1,42 @@ import { describe, expect, test } from "bun:test"; +import { createSessionRunTerminalSourceId } from "@mosoo/runtime-events"; + import { completeRuntimeStateOperationPhase, failRuntimeStateOperationPhase, startRuntimeStateOperationPhase, } from "../src/modules/runtime/application/runtime-state-operation-phases"; import type { RuntimeSessionTarget } from "../src/modules/runtime/application/runtime-state-operation-target-store"; +import { recordCanonicalSessionRunTerminal } from "../src/modules/runtime/application/session-runs/session-run-terminal-failure.service"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { createPublicHttpContractDatabase, createPublicHttpTestBindings, insertNonOwnerSession, + nowMsForTest, } from "./helpers/public-api-http-test-fixture"; function createRuntimeTarget( - input: Omit & { + input: Omit< + RuntimeSessionTarget, + | "sessionRuntimeEventSeqCursor" + | "sessionStatusOperationId" + | "sessionStatusSeq" + | "sessionUpdatedAt" + > & { + readonly sessionRuntimeEventSeqCursor?: number; readonly sessionStatusOperationId?: string | null; readonly sessionStatusSeq?: number; + readonly sessionUpdatedAt?: number; }, ): RuntimeSessionTarget { return { ...input, + sessionRuntimeEventSeqCursor: input.sessionRuntimeEventSeqCursor ?? 0, sessionStatusOperationId: input.sessionStatusOperationId ?? null, sessionStatusSeq: input.sessionStatusSeq ?? 0, + sessionUpdatedAt: input.sessionUpdatedAt ?? nowMsForTest(), }; } @@ -124,7 +138,7 @@ describe("runtime state operation phases", () => { run_status: "cancelled", session_status: "IDLE", session_status_operation_id: null, - session_status_seq: 2, + session_status_seq: 3, }); }); @@ -150,10 +164,15 @@ describe("runtime state operation phases", () => { ], }); - await database - .prepare("UPDATE session_run SET status = ?, status_operation_id = ? WHERE id = ?") - .bind("cancelled", phase.operationId, "01J0000000000000000000000N") - .run(); + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + error: null, + expectedSessionOperationId: phase.operationId, + runId: "01J0000000000000000000000N", + sessionId: "01J0000000000000000000000B", + source: "driver", + status: "cancelled", + }); await completeRuntimeStateOperationPhase(bindings, { agentId: "01J00000000000000000000009", @@ -171,12 +190,12 @@ describe("runtime state operation phases", () => { ) .bind( "01J0000000000000000000000B", - `runtime-operation:${phase.operationId}:01J0000000000000000000000N:interrupted`, + createSessionRunTerminalSourceId("01J0000000000000000000000N", "run.cancelled"), ) .first<{ source_event_id: string }>(); expect(event?.source_event_id).toBe( - `runtime-operation:${phase.operationId}:01J0000000000000000000000N:interrupted`, + createSessionRunTerminalSourceId("01J0000000000000000000000N", "run.cancelled"), ); const row = await database @@ -226,10 +245,15 @@ describe("runtime state operation phases", () => { ], }); - await database - .prepare("UPDATE session_run SET status = ?, status_operation_id = ? WHERE id = ?") - .bind("cancelled", phase.operationId, "01J0000000000000000000000N") - .run(); + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + error: null, + expectedSessionOperationId: phase.operationId, + runId: "01J0000000000000000000000N", + sessionId: "01J0000000000000000000000B", + source: "driver", + status: "cancelled", + }); await failRuntimeStateOperationPhase(bindings, { agentId: "01J00000000000000000000009", @@ -247,16 +271,16 @@ describe("runtime state operation phases", () => { ) .bind( "01J0000000000000000000000B", - `runtime-operation:${phase.operationId}:01J0000000000000000000000N:interrupted`, + createSessionRunTerminalSourceId("01J0000000000000000000000N", "run.cancelled"), ) .first<{ source_event_id: string }>(); expect(event?.source_event_id).toBe( - `runtime-operation:${phase.operationId}:01J0000000000000000000000N:interrupted`, + createSessionRunTerminalSourceId("01J0000000000000000000000N", "run.cancelled"), ); }); - test("complete does not app terminal runs from another outcome as cancelled", async () => { + test("complete adopts a canonical Driver failure without rewriting its outcome", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); await insertRunningSessionRun(database); @@ -278,10 +302,20 @@ describe("runtime state operation phases", () => { ], }); - await database - .prepare("UPDATE session_run SET status = ?, status_operation_id = ? WHERE id = ?") - .bind("completed", phase.operationId, "01J0000000000000000000000N") - .run(); + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + error: { + code: "driver.failed", + details: {}, + message: "Driver failed while stopping.", + retryable: true, + }, + expectedSessionOperationId: phase.operationId, + runId: "01J0000000000000000000000N", + sessionId: "01J0000000000000000000000B", + source: "driver", + status: "failed", + }); await completeRuntimeStateOperationPhase(bindings, { agentId: "01J00000000000000000000009", @@ -299,11 +333,16 @@ describe("runtime state operation phases", () => { ) .bind( "01J0000000000000000000000B", - `runtime-operation:${phase.operationId}:01J0000000000000000000000N:interrupted`, + createSessionRunTerminalSourceId("01J0000000000000000000000N", "run.cancelled"), ) .first<{ source_event_id: string }>(); expect(event).toBeNull(); + const run = await database + .prepare("SELECT error_code, status FROM session_run WHERE id = ?") + .bind("01J0000000000000000000000N") + .first<{ error_code: string | null; status: string }>(); + expect(run).toEqual({ error_code: "driver.failed", status: "failed" }); }); test("start ignores stale targets that changed after scope resolution", async () => { @@ -318,7 +357,7 @@ describe("runtime state operation phases", () => { WHERE id = ? `, ) - .bind("run-new", "RUNNING", 1, "01J0000000000000000000000B") + .bind("01J0000000000000000000000Q", "RUNNING", 1, "01J0000000000000000000000B") .run(); const bindings = createPublicHttpTestBindings(database) as ApiBindings; @@ -352,7 +391,7 @@ describe("runtime state operation phases", () => { }>(); expect(row).toEqual({ - last_run_id: "run-new", + last_run_id: "01J0000000000000000000000Q", status: "RUNNING", status_operation_id: null, status_seq: 1, diff --git a/apps/api/tests/runtime-state-operation-scope.test.ts b/apps/api/tests/runtime-state-operation-scope.test.ts index d11660ae..f6bfdf51 100644 --- a/apps/api/tests/runtime-state-operation-scope.test.ts +++ b/apps/api/tests/runtime-state-operation-scope.test.ts @@ -16,10 +16,12 @@ function createRuntimeOperationScopeDatabase(): SqliteD1Database { agent_id text NOT NULL, creator_account_id text NOT NULL, last_run_id text, + runtime_event_seq_cursor integer DEFAULT 0 NOT NULL, status text NOT NULL, status_operation_id text, status_seq integer DEFAULT 0 NOT NULL, - archived_at integer + archived_at integer, + updated_at integer DEFAULT 0 NOT NULL ); CREATE TABLE sandbox ( @@ -134,8 +136,10 @@ function createRuntimeTarget(input: { lastRunId: null, sandboxId: input.sandboxId, sessionId: input.sessionId, + sessionRuntimeEventSeqCursor: 0, sessionStatus: "IDLE", sessionStatusOperationId: null, sessionStatusSeq: 0, + sessionUpdatedAt: 0, }; } diff --git a/apps/api/tests/runtime-state-operation-target-events.test.ts b/apps/api/tests/runtime-state-operation-target-events.test.ts index fb13c738..bdf3ebe5 100644 --- a/apps/api/tests/runtime-state-operation-target-events.test.ts +++ b/apps/api/tests/runtime-state-operation-target-events.test.ts @@ -1,22 +1,37 @@ import { describe, expect, test } from "bun:test"; +import { sandboxSessionsTable } from "@mosoo/db"; +import { + createRuntimeEventSemanticHash, + parseRuntimeEventEnvelope, + stringifyRuntimeEventSemanticValue, +} from "@mosoo/runtime-events"; + import { toRuntimeDiagnosticBaseValue } from "../src/modules/runtime/application/runtime-diagnostic-events"; import { buildRuntimeStateOperationEvents } from "../src/modules/runtime/application/runtime-state-operation-events"; import { appendRuntimeDriverRestartAttemptedEvents, appendRuntimeSubjectTerminatedEvents, - broadcastRuntimeOperationEvent, + commitRuntimeOperationReadySnapshots, writeRuntimeOperationInterruptedSnapshots, writeRuntimeOperationTimedOutSnapshots, } from "../src/modules/runtime/application/runtime-state-operation-target-events"; +import { + adoptRuntimeOperationReadyReceipt, + claimRuntimeOperationTargets, +} from "../src/modules/runtime/application/runtime-state-operation-target-store"; import type { RuntimeSessionTarget } from "../src/modules/runtime/application/runtime-state-operation-target-store"; +import { recordCanonicalSessionRunTerminal } from "../src/modules/runtime/application/session-runs/session-run-terminal-failure.service"; +import { createSessionRuntimeEventProjection } from "../src/modules/sessions/domain/session-runtime-event-projection"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { PUBLIC_API_TEST_IDS, createPublicHttpContractDatabase, createPublicHttpTestBindings, + insertActiveSandboxSessionFixture, insertNonOwnerSession, insertOwnerSession, + nowMsForTest, } from "./helpers/public-api-http-test-fixture"; async function insertRunningSessionRun( @@ -72,6 +87,82 @@ async function insertRunningSessionRun( .run(); } +async function readRuntimeTarget(database: D1Database): Promise { + const row = await database + .prepare( + `SELECT last_run_id, runtime_event_seq_cursor, status, status_operation_id, + status_seq, updated_at + FROM session + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first<{ + last_run_id: RuntimeSessionTarget["lastRunId"]; + runtime_event_seq_cursor: number; + status: RuntimeSessionTarget["sessionStatus"]; + status_operation_id: RuntimeSessionTarget["sessionStatusOperationId"]; + status_seq: number; + updated_at: number; + }>(); + if (row === null) { + throw new Error("Missing runtime operation target fixture."); + } + return createRuntimeTarget({ + agentId: PUBLIC_API_TEST_IDS.agent, + creatorAccountId: PUBLIC_API_TEST_IDS.nonOwnerAccount, + lastRunId: row.last_run_id, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + sessionRuntimeEventSeqCursor: row.runtime_event_seq_cursor, + sessionStatus: row.status, + sessionStatusOperationId: row.status_operation_id, + sessionStatusSeq: row.status_seq, + sessionUpdatedAt: row.updated_at, + }); +} + +async function rewriteTerminalAuthority( + database: D1Database, + rewrite: (value: Record) => Record, +): Promise { + const row = await database + .prepare( + `SELECT terminal_event_json + FROM session_event + WHERE run_id = ? + AND event_type IN ('run.cancelled', 'run.completed', 'run.failed')`, + ) + .bind(PUBLIC_API_TEST_IDS.run) + .first<{ terminal_event_json: string | null }>(); + if (row?.terminal_event_json === null || row === null) { + throw new Error("Missing terminal semantic authority fixture."); + } + const event = parseRuntimeEventEnvelope(rewrite(JSON.parse(row.terminal_event_json))); + const projection = createSessionRuntimeEventProjection(event); + await database + .prepare( + `UPDATE session_event + SET content_text = ?, semantic_hash = ?, stream_id = ?, terminal_event_json = ? + WHERE run_id = ? + AND event_type IN ('run.cancelled', 'run.completed', 'run.failed')`, + ) + .bind( + projection.contentText, + await createRuntimeEventSemanticHash(event), + projection.streamId, + stringifyRuntimeEventSemanticValue(event), + PUBLIC_API_TEST_IDS.run, + ) + .run(); +} + +function requireRecord(value: unknown): Record { + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new Error("Expected a record fixture value."); + } + return value as Record; +} + function createRuntimeDiagnosticTargets(): RuntimeSessionTarget[] { return [ createRuntimeTarget({ @@ -106,6 +197,7 @@ async function insertLiveDriverInstance( input: { driverInstanceId: string; sessionId: string; + tokenByte: number; }, ): Promise { await database @@ -114,6 +206,7 @@ async function insertLiveDriverInstance( INSERT INTO driver_instance ( id, sandbox_id, + sandbox_incarnation, sandbox_session_id, runtime, protocol, @@ -126,18 +219,19 @@ async function insertLiveDriverInstance( created_at, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( input.driverInstanceId, PUBLIC_API_TEST_IDS.sandbox, + 1, input.sessionId, "cloudflare-container", "driver-ws", 1, "ready", - new Uint8Array([1, 2, 3]), + new Uint8Array([input.tokenByte]), 10_000, 0, 20_000, @@ -148,12 +242,25 @@ async function insertLiveDriverInstance( } function createRuntimeTarget( - input: Omit, + input: Omit< + RuntimeSessionTarget, + | "sessionRuntimeEventSeqCursor" + | "sessionStatusOperationId" + | "sessionStatusSeq" + | "sessionUpdatedAt" + > & { + readonly sessionRuntimeEventSeqCursor?: number; + readonly sessionStatusOperationId?: RuntimeSessionTarget["sessionStatusOperationId"]; + readonly sessionStatusSeq?: number; + readonly sessionUpdatedAt?: number; + }, ): RuntimeSessionTarget { return { ...input, - sessionStatusOperationId: null, - sessionStatusSeq: 0, + sessionRuntimeEventSeqCursor: input.sessionRuntimeEventSeqCursor ?? 0, + sessionStatusOperationId: input.sessionStatusOperationId ?? null, + sessionStatusSeq: input.sessionStatusSeq ?? 0, + sessionUpdatedAt: input.sessionUpdatedAt ?? nowMsForTest(), }; } @@ -161,6 +268,122 @@ function createExistingRuntimeDiagnosticTargets(): RuntimeSessionTarget[] { return createRuntimeDiagnosticTargets().filter((target) => target.agentId !== null); } +async function createAdoptableRuntimeOperationReadyReceipt(database: D1Database): Promise<{ + readonly target: RuntimeSessionTarget; +}> { + await insertNonOwnerSession(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const [updatingEvent, readyEvent] = buildRuntimeStateOperationEvents({ + agentId: PUBLIC_API_TEST_IDS.agent, + operation: "restartDriver", + readyAt: "2026-05-08T00:00:01.000Z", + startedAt: "2026-05-08T00:00:00.000Z", + }); + const [claimed] = await claimRuntimeOperationTargets(database, { + event: updatingEvent, + operationId: PUBLIC_API_TEST_IDS.operation, + targets: createExistingRuntimeDiagnosticTargets().slice(0, 1), + }); + if (claimed === undefined) { + throw new Error("Missing runtime operation claim fixture."); + } + await commitRuntimeOperationReadySnapshots(bindings, { + event: readyEvent, + operationId: PUBLIC_API_TEST_IDS.operation, + targets: [claimed.current], + }); + await database + .prepare( + `UPDATE session + SET status = 'RESCHEDULING', + status_operation_id = ?, + status_seq = status_seq + 1 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + + return { target: await readRuntimeTarget(database) }; +} + +async function rewriteRuntimeOperationAuthority( + database: D1Database, + status: "ready" | "updating", + rewrite: (event: Record) => Record, +): Promise { + const sourceEventId = `runtime-operation:${PUBLIC_API_TEST_IDS.operation}:${PUBLIC_API_TEST_IDS.nonOwnerSession}:${status}`; + const row = await database + .prepare( + `SELECT runtime_operation_event_json + FROM session_event + WHERE session_id = ? AND source_event_id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession, sourceEventId) + .first<{ runtime_operation_event_json: string | null }>(); + if (row?.runtime_operation_event_json === null || row === null) { + throw new Error(`Missing runtime operation ${status} authority fixture.`); + } + const event = parseRuntimeEventEnvelope(rewrite(JSON.parse(row.runtime_operation_event_json))); + await database + .prepare( + `UPDATE session_event + SET runtime_operation_event_json = ?, semantic_hash = ? + WHERE session_id = ? AND source_event_id = ?`, + ) + .bind( + stringifyRuntimeEventSemanticValue(event), + await createRuntimeEventSemanticHash(event), + PUBLIC_API_TEST_IDS.nonOwnerSession, + sourceEventId, + ) + .run(); +} + +function raceRuntimeOperationReadyAdoption( + database: D1Database, + race: () => Promise, +): { readonly database: D1Database; readonly raced: () => boolean } { + let raced = false; + + function wrap(statement: D1PreparedStatement, shouldRace: boolean): D1PreparedStatement { + return new Proxy(statement, { + get(target, property, receiver) { + if (property === "bind") { + return (...values: unknown[]) => wrap(target.bind(...values), shouldRace); + } + if (property === "run" && shouldRace) { + return async () => { + if (!raced) { + raced = true; + await race(); + } + return target.run(); + }; + } + return Reflect.get(target, property, receiver); + }, + }); + } + + return { + database: new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => + wrap( + target.prepare(query), + query.includes("UPDATE session") && + query.includes("runtime_operation_event_json = ?"), + ); + } + const value = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }), + raced: () => raced, + }; +} + describe("runtime state operation target events", () => { test("operation events carry the target deployment version", () => { const events = buildRuntimeStateOperationEvents({ @@ -218,6 +441,10 @@ describe("runtime state operation target events", () => { await insertRunningSessionRun(database); const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await database + .prepare("UPDATE session SET status = ?, status_operation_id = ? WHERE id = ?") + .bind("RESCHEDULING", PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); const targets: RuntimeSessionTarget[] = [ createRuntimeTarget({ agentId: PUBLIC_API_TEST_IDS.agent, @@ -225,12 +452,14 @@ describe("runtime state operation target events", () => { lastRunId: PUBLIC_API_TEST_IDS.run, sandboxId: PUBLIC_API_TEST_IDS.sandbox, sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, - sessionStatus: "RUNNING", + sessionStatus: "RESCHEDULING", + sessionStatusOperationId: PUBLIC_API_TEST_IDS.operation, }), ]; await writeRuntimeOperationInterruptedSnapshots(bindings, { operationId: PUBLIC_API_TEST_IDS.operation, + timestampMs: nowMsForTest(), targets, }); @@ -256,6 +485,10 @@ describe("runtime state operation target events", () => { }); const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await database + .prepare("UPDATE session SET status = ?, status_operation_id = ?") + .bind("RESCHEDULING", PUBLIC_API_TEST_IDS.operation) + .run(); const targets: RuntimeSessionTarget[] = [ createRuntimeTarget({ agentId: PUBLIC_API_TEST_IDS.agent, @@ -263,7 +496,8 @@ describe("runtime state operation target events", () => { lastRunId: PUBLIC_API_TEST_IDS.run, sandboxId: PUBLIC_API_TEST_IDS.sandbox, sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, - sessionStatus: "RUNNING", + sessionStatus: "RESCHEDULING", + sessionStatusOperationId: PUBLIC_API_TEST_IDS.operation, }), createRuntimeTarget({ agentId: PUBLIC_API_TEST_IDS.agent, @@ -271,12 +505,14 @@ describe("runtime state operation target events", () => { lastRunId: PUBLIC_API_TEST_IDS.runAlt, sandboxId: PUBLIC_API_TEST_IDS.sandbox, sessionId: PUBLIC_API_TEST_IDS.ownerSession, - sessionStatus: "RUNNING", + sessionStatus: "RESCHEDULING", + sessionStatusOperationId: PUBLIC_API_TEST_IDS.operation, }), ]; await writeRuntimeOperationInterruptedSnapshots(bindings, { operationId: PUBLIC_API_TEST_IDS.operation, + timestampMs: nowMsForTest(), targets, }); @@ -310,6 +546,51 @@ describe("runtime state operation target events", () => { ]); }); + test("preserves a terminal lifecycle winner during interrupted recovery", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await insertRunningSessionRun(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + deliver: false, + error: { + code: "driver.command_failed", + details: {}, + message: "Driver failed before recovery.", + retryable: false, + }, + lifecycle: "TERMINATED", + runId: PUBLIC_API_TEST_IDS.run, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + source: "maintenance", + status: "failed", + timestampMs: 2, + }); + await database + .prepare( + `UPDATE session + SET status = 'RESCHEDULING', status_operation_id = ?, + status_seq = status_seq + 1, updated_at = 3 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + + await writeRuntimeOperationInterruptedSnapshots(bindings, { + operationId: PUBLIC_API_TEST_IDS.operation, + targets: [await readRuntimeTarget(database)], + timestampMs: 3, + }); + + await expect( + database + .prepare("SELECT status, status_operation_id FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ status: "TERMINATED", status_operation_id: null }); + }); + test("terminated subject events are written for target sessions", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); @@ -347,13 +628,40 @@ describe("runtime state operation target events", () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); await insertOwnerSession(database); + await insertActiveSandboxSessionFixture(database, { + ownerAccountId: PUBLIC_API_TEST_IDS.nonOwnerAccount, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + }); + await database + .app() + .insert(sandboxSessionsTable) + .values({ + createdAt: nowMsForTest(), + cwd: `/workspace/se/${PUBLIC_API_TEST_IDS.ownerSession}`, + originJson: JSON.stringify({ + callerUserId: PUBLIC_API_TEST_IDS.ownerAccount, + entrypoint: "api", + executionOwnerUserId: PUBLIC_API_TEST_IDS.ownerAccount, + type: "agent", + }), + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxIncarnation: 1, + sandboxSessionId: PUBLIC_API_TEST_IDS.driverOwner, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + status: "active", + updatedAt: nowMsForTest(), + }) + .run(); await insertLiveDriverInstance(database, { driverInstanceId: PUBLIC_API_TEST_IDS.driverNonOwner, sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + tokenByte: 1, }); await insertLiveDriverInstance(database, { driverInstanceId: PUBLIC_API_TEST_IDS.driverOwner, sessionId: PUBLIC_API_TEST_IDS.ownerSession, + tokenByte: 2, }); const bindings = createPublicHttpTestBindings(database) as ApiBindings; @@ -383,52 +691,484 @@ describe("runtime state operation target events", () => { expect(rows.results.every((row) => row.event_type === "runtime.driver.updated")).toBe(true); }); - test("runtime operation broadcasts events across target sessions", async () => { + test("ready snapshots atomically release operation-owned target sessions", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); await insertOwnerSession(database); const bindings = createPublicHttpTestBindings(database) as ApiBindings; - const [event] = buildRuntimeStateOperationEvents({ + const [updatingEvent, readyEvent] = buildRuntimeStateOperationEvents({ agentId: PUBLIC_API_TEST_IDS.agent, operation: "restartDriver", readyAt: "2026-05-08T00:00:01.000Z", startedAt: "2026-05-08T00:00:00.000Z", }); - await broadcastRuntimeOperationEvent(bindings, { - event, + const claimed = await claimRuntimeOperationTargets(database, { + event: updatingEvent, operationId: PUBLIC_API_TEST_IDS.operation, targets: createExistingRuntimeDiagnosticTargets(), }); + await commitRuntimeOperationReadySnapshots(bindings, { + event: readyEvent, + operationId: PUBLIC_API_TEST_IDS.operation, + targets: claimed.map((transition) => transition.current), + }); const rows = await database .prepare( ` - SELECT event_type, seq, session_id + SELECT event_type, runtime_operation_event_json, seq, session_id, source_event_id FROM session_event - ORDER BY session_id + ORDER BY session_id, seq `, ) .all<{ event_type: string; + runtime_operation_event_json: string | null; seq: number; session_id: string; + source_event_id: string; }>(); expect(rows.results.map((row) => ({ seq: row.seq, sessionId: row.session_id }))).toEqual([ { seq: 1, sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession }, + { seq: 2, sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession }, { seq: 1, sessionId: PUBLIC_API_TEST_IDS.ownerSession }, + { seq: 2, sessionId: PUBLIC_API_TEST_IDS.ownerSession }, ]); expect(rows.results.every((row) => row.event_type === "agent.task.updated")).toBe(true); + expect( + rows.results.map((row) => ({ + authority: row.runtime_operation_event_json === null ? "legacy" : "canonical", + source: row.source_event_id.endsWith(":ready") ? "ready" : "updating", + })), + ).toEqual([ + { authority: "canonical", source: "updating" }, + { authority: "canonical", source: "ready" }, + { authority: "canonical", source: "updating" }, + { authority: "canonical", source: "ready" }, + ]); + expect( + await database.prepare("SELECT status, status_operation_id FROM session ORDER BY id").all(), + ).toMatchObject({ + results: expect.arrayContaining([ + { status: "IDLE", status_operation_id: null }, + { status: "IDLE", status_operation_id: null }, + ]), + }); + }); + + test("adopts one canonical ready receipt and acknowledges its replay", async () => { + const database = await createPublicHttpContractDatabase(); + const { target } = await createAdoptableRuntimeOperationReadyReceipt(database); + const input = { operationId: PUBLIC_API_TEST_IDS.operation, target }; + + await expect(adoptRuntimeOperationReadyReceipt(database, input)).resolves.toBe("applied"); + await expect(adoptRuntimeOperationReadyReceipt(database, input)).resolves.toBe("duplicate"); + await expect( + database + .prepare("SELECT status, status_operation_id FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ status: "IDLE", status_operation_id: null }); + }); + + test("rejects an updating event forged behind the ready source receipt", async () => { + const database = await createPublicHttpContractDatabase(); + const { target } = await createAdoptableRuntimeOperationReadyReceipt(database); + const sourceEventId = `runtime-operation:${PUBLIC_API_TEST_IDS.operation}:${PUBLIC_API_TEST_IDS.nonOwnerSession}:ready`; + const row = await database + .prepare( + `SELECT runtime_operation_event_json + FROM session_event + WHERE session_id = ? AND source_event_id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession, sourceEventId) + .first<{ runtime_operation_event_json: string | null }>(); + if (row?.runtime_operation_event_json === null || row === null) { + throw new Error("Missing runtime operation ready authority fixture."); + } + const ready = parseRuntimeEventEnvelope(JSON.parse(row.runtime_operation_event_json)); + const payload = requireRecord(ready.payload); + const updating = parseRuntimeEventEnvelope({ + ...ready, + payload: { + agentId: payload["agentId"], + operation: payload["operation"], + operationId: payload["operationId"], + startedAt: ready.occurredAt, + status: "updating", + }, + }); + await database + .prepare( + `UPDATE session_event + SET runtime_operation_event_json = NULL, semantic_hash = ? + WHERE session_id = ? AND source_event_id = ?`, + ) + .bind( + await createRuntimeEventSemanticHash(updating), + PUBLIC_API_TEST_IDS.nonOwnerSession, + sourceEventId, + ) + .run(); + + await expect( + adoptRuntimeOperationReadyReceipt(database, { + operationId: PUBLIC_API_TEST_IDS.operation, + target, + }), + ).rejects.toThrow("has no semantic authority"); + await expect( + database + .prepare("SELECT status, status_operation_id FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ + status: "RESCHEDULING", + status_operation_id: PUBLIC_API_TEST_IDS.operation, + }); + }); + + test("requires the canonical updating claim before adopting ready", async () => { + const database = await createPublicHttpContractDatabase(); + const { target } = await createAdoptableRuntimeOperationReadyReceipt(database); + await database + .prepare( + `DELETE FROM session_event + WHERE session_id = ? AND source_event_id LIKE 'runtime-operation:%:updating'`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + + await expect( + adoptRuntimeOperationReadyReceipt(database, { + operationId: PUBLIC_API_TEST_IDS.operation, + target, + }), + ).rejects.toThrow("has no canonical claim"); }); - test("timed out snapshots cancel running runs and persist target events", async () => { + test("rejects a self-consistent claim payload mutation", async () => { + const database = await createPublicHttpContractDatabase(); + const { target } = await createAdoptableRuntimeOperationReadyReceipt(database); + await rewriteRuntimeOperationAuthority(database, "updating", (event) => ({ + ...event, + payload: { ...requireRecord(event["payload"]), forged: true }, + })); + + await expect( + adoptRuntimeOperationReadyReceipt(database, { + operationId: PUBLIC_API_TEST_IDS.operation, + target, + }), + ).rejects.toThrow("updating source"); + }); + + test("fails closed for self-consistent ready authority identity corruption", async () => { + const corruptions: readonly { + readonly rewrite: (event: Record) => Record; + }[] = [ + { + rewrite: (event) => ({ + ...event, + payload: { ...requireRecord(event["payload"]), forged: true }, + }), + }, + { + rewrite: (event) => ({ ...event, sourceEventId: "forged-runtime-operation-ready" }), + }, + { + rewrite: (event) => ({ ...event, sessionId: PUBLIC_API_TEST_IDS.ownerSession }), + }, + { + rewrite: (event) => ({ + ...event, + payload: { + ...requireRecord(event["payload"]), + operationId: PUBLIC_API_TEST_IDS.deployment, + }, + }), + }, + { + rewrite: (event) => ({ + ...event, + payload: { + ...requireRecord(event["payload"]), + agentId: PUBLIC_API_TEST_IDS.ownerAccount, + }, + }), + }, + { + rewrite: (event) => ({ + ...event, + payload: { + ...requireRecord(event["payload"]), + operation: "recreateSandbox", + }, + }), + }, + { + rewrite: (event) => ({ + ...event, + payload: { + ...requireRecord(event["payload"]), + deploymentVersionId: PUBLIC_API_TEST_IDS.deployment, + deploymentVersionNumber: 1, + }, + }), + }, + { + rewrite: (event) => ({ ...event, visibility: "owner_debug" }), + }, + { + rewrite: (event) => ({ ...event, actor: "system", origin: "system" }), + }, + ]; + + for (const corruption of corruptions) { + const database = await createPublicHttpContractDatabase(); + const { target } = await createAdoptableRuntimeOperationReadyReceipt(database); + await rewriteRuntimeOperationAuthority(database, "ready", corruption.rewrite); + + await expect( + adoptRuntimeOperationReadyReceipt(database, { + operationId: PUBLIC_API_TEST_IDS.operation, + target, + }), + ).rejects.toThrow("ready source"); + await expect( + database + .prepare("SELECT status, status_operation_id FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ + status: "RESCHEDULING", + status_operation_id: PUBLIC_API_TEST_IDS.operation, + }); + } + }); + + test("recomputes the ready authority hash before adopting it", async () => { + const database = await createPublicHttpContractDatabase(); + const { target } = await createAdoptableRuntimeOperationReadyReceipt(database); + await database + .prepare( + `UPDATE session_event + SET semantic_hash = ? + WHERE session_id = ? AND source_event_id LIKE 'runtime-operation:%:ready'`, + ) + .bind("0".repeat(64), PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + + await expect( + adoptRuntimeOperationReadyReceipt(database, { + operationId: PUBLIC_API_TEST_IDS.operation, + target, + }), + ).rejects.toThrow("is not canonical"); + }); + + test("does not cross either receipt mutation that races the adoption CAS", async () => { + for (const status of ["updating", "ready"] as const) { + const database = await createPublicHttpContractDatabase(); + const { target } = await createAdoptableRuntimeOperationReadyReceipt(database); + const racing = raceRuntimeOperationReadyAdoption(database, async () => { + await database + .prepare( + `UPDATE session_event + SET semantic_hash = ? + WHERE session_id = ? AND source_event_id LIKE ?`, + ) + .bind( + "f".repeat(64), + PUBLIC_API_TEST_IDS.nonOwnerSession, + `runtime-operation:%:${status}`, + ) + .run(); + }); + + await expect( + adoptRuntimeOperationReadyReceipt(racing.database, { + operationId: PUBLIC_API_TEST_IDS.operation, + target, + }), + ).rejects.toThrow("is not canonical"); + expect(racing.raced()).toBe(true); + await expect( + database + .prepare("SELECT status, status_operation_id FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ + status: "RESCHEDULING", + status_operation_id: PUBLIC_API_TEST_IDS.operation, + }); + } + }); + + test("operation start cannot cross an active runtime provisioning fence", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await database + .prepare( + `UPDATE session + SET runtime_provisioning_heartbeat_at = 1, + runtime_provisioning_operation_id = ?, + runtime_provisioning_sandbox_id = ? + WHERE id = ?`, + ) + .bind( + PUBLIC_API_TEST_IDS.operation, + PUBLIC_API_TEST_IDS.sandbox, + PUBLIC_API_TEST_IDS.nonOwnerSession, + ) + .run(); + const [updatingEvent] = buildRuntimeStateOperationEvents({ + agentId: PUBLIC_API_TEST_IDS.agent, + operation: "restartDriver", + readyAt: "2026-05-08T00:00:01.000Z", + startedAt: "2026-05-08T00:00:00.000Z", + }); + + expect( + await claimRuntimeOperationTargets(database, { + event: updatingEvent, + operationId: PUBLIC_API_TEST_IDS.operation, + targets: createExistingRuntimeDiagnosticTargets().slice(0, 1), + }), + ).toEqual([]); + expect( + await database + .prepare( + `SELECT runtime_event_seq_cursor, + runtime_provisioning_operation_id, + status, + status_operation_id + FROM session + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).toEqual({ + runtime_event_seq_cursor: 0, + runtime_provisioning_operation_id: PUBLIC_API_TEST_IDS.operation, + status: "IDLE", + status_operation_id: null, + }); + expect( + await database + .prepare("SELECT COUNT(*) AS count FROM session_event WHERE session_id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).toEqual({ count: 0 }); + }); + + test("ready wins atomically when operation timeout races its release", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await database + .prepare("UPDATE session SET updated_at = 1 WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + const [updatingEvent, readyEvent] = buildRuntimeStateOperationEvents({ + agentId: PUBLIC_API_TEST_IDS.agent, + operation: "restartDriver", + readyAt: "1970-01-01T00:00:00.002Z", + startedAt: "1970-01-01T00:00:00.001Z", + }); + const [claimed] = await claimRuntimeOperationTargets(database, { + event: updatingEvent, + operationId: PUBLIC_API_TEST_IDS.operation, + targets: [ + createRuntimeTarget({ + agentId: PUBLIC_API_TEST_IDS.agent, + creatorAccountId: PUBLIC_API_TEST_IDS.nonOwnerAccount, + lastRunId: null, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + sessionStatus: "IDLE", + sessionUpdatedAt: 1, + }), + ], + }); + if (claimed === undefined) { + throw new Error("Missing runtime operation claim fixture."); + } + const readyBindings = createPublicHttpTestBindings(database) as ApiBindings; + let raced = false; + const racingDatabase = new Proxy(database, { + get(target, property) { + if (property === "batch") { + return async (statements: D1PreparedStatement[]) => { + if (!raced) { + raced = true; + await commitRuntimeOperationReadySnapshots(readyBindings, { + event: readyEvent, + operationId: PUBLIC_API_TEST_IDS.operation, + targets: [claimed.current], + }); + } + return target.batch(statements); + }; + } + + const value = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }) as D1Database; + + await writeRuntimeOperationTimedOutSnapshots( + { + ...(createPublicHttpTestBindings(database) as ApiBindings), + DB: racingDatabase, + }, + { operationId: PUBLIC_API_TEST_IDS.operation, targets: [claimed.current] }, + ); + + expect(raced).toBe(true); + await expect( + database + .prepare("SELECT status, status_operation_id FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ status: "IDLE", status_operation_id: null }); + await expect( + database + .prepare( + `SELECT source_event_id + FROM session_event + WHERE source_event_id LIKE 'runtime-operation:%' + ORDER BY seq`, + ) + .all(), + ).resolves.toMatchObject({ + results: [ + { + source_event_id: `runtime-operation:${PUBLIC_API_TEST_IDS.operation}:${PUBLIC_API_TEST_IDS.nonOwnerSession}:updating`, + }, + { + source_event_id: `runtime-operation:${PUBLIC_API_TEST_IDS.operation}:${PUBLIC_API_TEST_IDS.nonOwnerSession}:ready`, + }, + ], + }); + }); + + test("timed out snapshots expire running runs and persist target events", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); await insertOwnerSession(database); await insertRunningSessionRun(database); const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await database + .prepare("UPDATE session SET status = ?, status_operation_id = ? WHERE id = ?") + .bind("RESCHEDULING", PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + await database + .prepare("UPDATE session SET status_operation_id = ? WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.ownerSession) + .run(); const targets: RuntimeSessionTarget[] = [ createRuntimeTarget({ agentId: PUBLIC_API_TEST_IDS.agent, @@ -436,7 +1176,8 @@ describe("runtime state operation target events", () => { lastRunId: PUBLIC_API_TEST_IDS.run, sandboxId: PUBLIC_API_TEST_IDS.sandbox, sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, - sessionStatus: "RUNNING", + sessionStatus: "RESCHEDULING", + sessionStatusOperationId: PUBLIC_API_TEST_IDS.operation, }), createRuntimeTarget({ agentId: PUBLIC_API_TEST_IDS.agent, @@ -445,6 +1186,7 @@ describe("runtime state operation target events", () => { sandboxId: PUBLIC_API_TEST_IDS.sandbox, sessionId: PUBLIC_API_TEST_IDS.ownerSession, sessionStatus: "IDLE", + sessionStatusOperationId: PUBLIC_API_TEST_IDS.operation, }), ]; @@ -457,7 +1199,7 @@ describe("runtime state operation target events", () => { .prepare("SELECT status FROM session_run WHERE id = ?") .bind(PUBLIC_API_TEST_IDS.run) .first<{ status: string }>(); - expect(run).toEqual({ status: "cancelled" }); + expect(run).toEqual({ status: "expired" }); const events = await database .prepare( ` @@ -474,4 +1216,201 @@ describe("runtime state operation target events", () => { { seq: 1, sessionId: PUBLIC_API_TEST_IDS.ownerSession }, ]); }); + + test("preserves a terminal lifecycle winner while releasing its operation fence", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await insertRunningSessionRun(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + deliver: false, + error: { + code: "driver.command_failed", + details: {}, + message: "Driver failed before the operation completed.", + retryable: false, + }, + lifecycle: "TERMINATED", + runId: PUBLIC_API_TEST_IDS.run, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + source: "maintenance", + status: "failed", + timestampMs: 2, + }); + await database + .prepare( + `UPDATE session + SET status = 'RESCHEDULING', status_operation_id = ?, + status_seq = status_seq + 1, updated_at = 3 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + + await writeRuntimeOperationTimedOutSnapshots(bindings, { + operationId: PUBLIC_API_TEST_IDS.operation, + targets: [await readRuntimeTarget(database)], + }); + + await expect( + database + .prepare("SELECT status, status_operation_id FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ status: "TERMINATED", status_operation_id: null }); + }); + + test("does not release an operation from a self-consistent missing final authority", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await insertRunningSessionRun(database); + await database + .prepare("UPDATE session SET status = ?, status_operation_id = ? WHERE id = ?") + .bind("RESCHEDULING", PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + deliver: false, + error: null, + lifecycle: "IDLE", + runId: PUBLIC_API_TEST_IDS.run, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + source: "driver", + status: "completed", + timestampMs: 2, + }); + await rewriteTerminalAuthority(database, (value) => ({ + ...value, + payload: { + ...requireRecord(value.payload), + finalMessageId: "01J0000000000000000000000M", + }, + })); + const target = await readRuntimeTarget(database); + + await expect( + writeRuntimeOperationTimedOutSnapshots(bindings, { + operationId: PUBLIC_API_TEST_IDS.operation, + targets: [target], + }), + ).rejects.toThrow("Canonical final assistant"); + await expect( + database + .prepare("SELECT status, status_operation_id FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ status: "IDLE", status_operation_id: PUBLIC_API_TEST_IDS.operation }); + }); + + test("does not release an operation from a self-consistent conflicting Run error", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await insertRunningSessionRun(database); + await database + .prepare("UPDATE session SET status = ?, status_operation_id = ? WHERE id = ?") + .bind("RESCHEDULING", PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + deliver: false, + error: { + code: "driver.command_failed", + details: {}, + message: "A", + retryable: false, + }, + lifecycle: "IDLE", + runId: PUBLIC_API_TEST_IDS.run, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + source: "driver", + status: "failed", + timestampMs: 2, + }); + const conflictingError = { + code: "driver.command_failed", + details: {}, + message: "B", + retryable: false, + }; + await rewriteTerminalAuthority(database, (value) => { + const payload = requireRecord(value.payload); + return { + ...value, + payload: { + ...payload, + error: conflictingError, + run: { ...requireRecord(payload.run), error: conflictingError }, + }, + }; + }); + const target = await readRuntimeTarget(database); + + await expect( + writeRuntimeOperationTimedOutSnapshots(bindings, { + operationId: PUBLIC_API_TEST_IDS.operation, + targets: [target], + }), + ).rejects.toThrow("persisted RunError"); + await expect( + database + .prepare("SELECT status, status_operation_id FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ status: "IDLE", status_operation_id: PUBLIC_API_TEST_IDS.operation }); + }); + + test("retries a timeout from the fresh target after losing its lifecycle CAS", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await database + .prepare("UPDATE session SET status = ?, status_operation_id = ? WHERE id = ?") + .bind("RESCHEDULING", PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + const target = await readRuntimeTarget(database); + let raced = false; + const racingDatabase = new Proxy(database, { + get(source, property) { + if (property === "batch") { + return async (statements: D1PreparedStatement[]) => { + if (!raced) { + raced = true; + await source + .prepare( + `UPDATE session + SET runtime_event_seq_cursor = runtime_event_seq_cursor + 1, + updated_at = updated_at + 1 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + } + return source.batch(statements); + }; + } + const value = Reflect.get(source, property); + return typeof value === "function" ? value.bind(source) : value; + }, + }) as D1Database; + + await writeRuntimeOperationTimedOutSnapshots( + { + ...(createPublicHttpTestBindings(database) as ApiBindings), + DB: racingDatabase, + }, + { operationId: PUBLIC_API_TEST_IDS.operation, targets: [target] }, + ); + expect(raced).toBe(true); + await expect( + database + .prepare("SELECT status, status_operation_id FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ + status: "TERMINATED", + status_operation_id: null, + }); + }); }); diff --git a/apps/api/tests/runtime-subject-activation-record.test.ts b/apps/api/tests/runtime-subject-activation-record.test.ts index 6fd1ad65..e5868bad 100644 --- a/apps/api/tests/runtime-subject-activation-record.test.ts +++ b/apps/api/tests/runtime-subject-activation-record.test.ts @@ -10,15 +10,24 @@ function createRuntimeSubjectDatabase(): SqliteD1Database { database.execute(` CREATE TABLE sandbox ( + agent_id text NOT NULL DEFAULT '01J00000000000000000000009', + app_id text NOT NULL DEFAULT '01J0000000000000000000000A', claim_expires_at integer, claim_owner text, global_mounts_json text NOT NULL, id text PRIMARY KEY NOT NULL, + incarnation integer NOT NULL DEFAULT 0, kind text NOT NULL, last_backup_id text, last_error text, last_error_code text, - status text NOT NULL + network_constraints_hash text, + operation_kind text, + owner_account_id text NOT NULL DEFAULT '01J00000000000000000000008', + status_operation_id text, + status text NOT NULL, + subject_id text NOT NULL DEFAULT '01J00000000000000000000009', + subject_kind text NOT NULL DEFAULT 'agent' ); CREATE TABLE sandbox_backup ( diff --git a/apps/api/tests/runtime-subject-lifecycle.test.ts b/apps/api/tests/runtime-subject-lifecycle.test.ts index 027df519..670f8c61 100644 --- a/apps/api/tests/runtime-subject-lifecycle.test.ts +++ b/apps/api/tests/runtime-subject-lifecycle.test.ts @@ -1,15 +1,26 @@ import { afterEach, describe, expect, test } from "bun:test"; +import { createTimeoutError } from "@mosoo/effects"; import { createPlatformId } from "@mosoo/id"; -import type { SandboxId, SessionId } from "@mosoo/id"; +import type { RuntimeOperationId, SandboxId, SessionId, SessionRunId } from "@mosoo/id"; import { decideRuntimeSubjectTransition } from "../src/modules/runtime/domain/runtime-subject-lifecycle.machine"; +import { hashSandboxNetworkConstraints } from "../src/modules/runtime/domain/sandbox-network-constraints"; import { createRuntimeSubjectLifecycleService } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-lifecycle.service"; import type { ActivateRuntimeSubjectInput } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-lifecycle.service"; -import { destroyRuntimeSubjectContainer } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-platform"; +import { + recreateRuntimeSubjectPreservingState, + resetRuntimeSubjectAgentState, + runRuntimeSubjectOperation, +} from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-operations.service"; +import { + destroyRuntimeSubjectContainer, + getRuntimeSubjectKeepAliveHandle, +} from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-platform"; import { recycleRuntimeSubject } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-recycle.service"; import { advanceRuntimeSubjectOperationStatus, + claimRuntimeSubjectOperationForRepair, markRuntimeSubjectCold, markRuntimeSubjectOperationStarted, } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store"; @@ -31,6 +42,10 @@ const RUNTIME_SUBJECT_QUOTA_SCOPE = { appId: APP_ID, executionOwnerUserId: ACCOUNT_ID, } as const; +const FULL_NETWORK_CONSTRAINTS_HASH = await hashSandboxNetworkConstraints({ + allowedHosts: [], + networkPolicy: "full", +}); afterEach(() => { setServerProductAnalyticsTransportForTests(null); @@ -50,12 +65,15 @@ function createRuntimeSubjectLifecycleDatabase(): SqliteD1Database { global_mounts_json text DEFAULT '[]' NOT NULL, id text PRIMARY KEY NOT NULL, inactive_deadline_at integer, + incarnation integer DEFAULT 0 NOT NULL, kind text NOT NULL, last_backup_id text, last_error text, last_error_code text, last_restore_backup_id text, + network_constraints_hash text, owner_account_id text, + operation_kind text, status text NOT NULL, status_changed_at integer DEFAULT 0 NOT NULL, status_event text DEFAULT 'runtime_subject.cold' NOT NULL, @@ -73,16 +91,35 @@ function createRuntimeSubjectLifecycleDatabase(): SqliteD1Database { error_message text, id text PRIMARY KEY NOT NULL, keep integer NOT NULL, + operation_id text, sandbox_id text NOT NULL, + sandbox_incarnation integer DEFAULT 0 NOT NULL, session_run_id text, status text NOT NULL, ttl_seconds integer NOT NULL, updated_at integer NOT NULL ); + CREATE TABLE sandbox_session ( + cloudflare_session_id text NOT NULL, + cleanup_operation_id text, + created_at integer NOT NULL, + cwd text NOT NULL, + origin_json text NOT NULL, + sandbox_id text NOT NULL, + sandbox_incarnation integer DEFAULT 0 NOT NULL, + session_id text PRIMARY KEY NOT NULL, + status text NOT NULL, + updated_at integer NOT NULL + ); + CREATE TABLE driver_instance ( + generation integer DEFAULT 0 NOT NULL, id text PRIMARY KEY NOT NULL, sandbox_id text NOT NULL, + sandbox_incarnation integer DEFAULT 0 NOT NULL, + sandbox_session_id text NOT NULL, + status_operation_id text, status text NOT NULL ); @@ -93,6 +130,23 @@ function createRuntimeSubjectLifecycleDatabase(): SqliteD1Database { session_id text NOT NULL, status text NOT NULL ); + + CREATE TABLE native_resume_ref ( + session_id text PRIMARY KEY NOT NULL + ); + + CREATE TABLE session ( + archived_at integer, + cleanup_operation_kind text, + id text PRIMARY KEY NOT NULL, + last_message_at integer, + runtime_provisioning_operation_id text, + runtime_provisioning_run_id text, + runtime_provisioning_sandbox_id text, + runtime_provisioning_sandbox_incarnation integer, + runtime_provisioning_sandbox_session_id text, + status text NOT NULL DEFAULT 'ready' + ); `); return database; @@ -101,6 +155,9 @@ function createRuntimeSubjectLifecycleDatabase(): SqliteD1Database { async function insertRuntimeSubject( database: D1Database, input: { + readonly incarnation?: number; + readonly kind?: "cattle" | "pet"; + readonly lastBackupId?: string | null; readonly lastError?: string | null; readonly lastErrorCode?: string | null; readonly status: string; @@ -111,16 +168,21 @@ async function insertRuntimeSubject( .prepare( ` INSERT INTO sandbox ( + agent_id, + app_id, claim_expires_at, claim_owner, created_at, id, inactive_deadline_at, + incarnation, kind, last_backup_id, last_error, last_error_code, last_restore_backup_id, + network_constraints_hash, + owner_account_id, status, status_event, status_seq, @@ -129,31 +191,93 @@ async function insertRuntimeSubject( subject_kind, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( + AGENT_ID, + APP_ID, null, null, 1, RUNTIME_SUBJECT_ID, 1, - "cattle", - null, + input.incarnation ?? (input.status === "cold" ? 0 : 1), + input.kind ?? "cattle", + input.lastBackupId ?? null, input.lastError ?? null, input.lastErrorCode ?? null, null, + FULL_NETWORK_CONSTRAINTS_HASH, + ACCOUNT_ID, input.status, `runtime_subject.${input.status === "backing_up" ? "back_up" : input.status}`, input.statusSeq ?? 0, "test", - SESSION_ID, - "session", + input.kind === "pet" ? AGENT_ID : SESSION_ID, + input.kind === "pet" ? "agent" : "session", 1, ) .run(); } +async function insertRuntimeProvisioningAuthority( + database: D1Database, + incarnation: number, +): Promise> { + const operationId = createPlatformId(); + const runId = createPlatformId(); + await database + .prepare( + `INSERT INTO session ( + id, runtime_provisioning_operation_id, runtime_provisioning_run_id, + runtime_provisioning_sandbox_id, runtime_provisioning_sandbox_incarnation, + status + ) VALUES (?, ?, ?, ?, ?, 'ready')`, + ) + .bind(SESSION_ID, operationId, runId, RUNTIME_SUBJECT_ID, incarnation) + .run(); + return { operationId, runId, sessionId: SESSION_ID }; +} + +async function insertReadySubjectBackup(database: D1Database, incarnation: number): Promise { + await database + .prepare( + `INSERT INTO sandbox_backup ( + created_at, dir, error_message, id, keep, sandbox_id, + sandbox_incarnation, status, ttl_seconds, updated_at + ) VALUES (1, '/workspace/memory', NULL, ?, 0, ?, ?, 'ready', 600, 1)`, + ) + .bind(STORED_BACKUP_ID, RUNTIME_SUBJECT_ID, incarnation) + .run(); +} + +async function insertRuntimeConversationSession( + database: D1Database, + input: { + readonly incarnation: number; + readonly sessionId: string; + readonly status: "active" | "cleanup_pending" | "closed" | "error"; + }, +): Promise { + await database + .prepare( + `INSERT INTO sandbox_session ( + cloudflare_session_id, cleanup_operation_id, created_at, cwd, origin_json, + sandbox_id, sandbox_incarnation, session_id, status, updated_at + ) VALUES (?, ?, 1, '/workspace', '{}', ?, ?, ?, ?, 1)`, + ) + .bind( + createPlatformId(), + input.status === "cleanup_pending" ? createPlatformId() : null, + RUNTIME_SUBJECT_ID, + input.incarnation, + input.sessionId, + input.status, + ) + .run(); +} + async function readRuntimeSubject(database: D1Database): Promise<{ status: string; status_event: string; @@ -188,7 +312,12 @@ function createSandboxHandle( readonly destroyError?: Error; readonly destroyPromise?: Promise; readonly onConfigureNetwork?: () => void; + readonly onDispose?: () => void; readonly onDestroy?: () => void; + readonly onExec?: (command: string) => void; + readonly inspectKind?: "healthy" | "missing" | "retired" | "stale" | "unknown"; + readonly onActivate?: (incarnation: number) => void; + readonly onReady?: (incarnation: number) => void; readonly onRestore?: (backup: { readonly dir: string; readonly id: string }) => void; readonly prepareError?: Error; } = {}, @@ -198,6 +327,10 @@ function createSandboxHandle( }; return { + [Symbol.dispose]: () => options.onDispose?.(), + activateRuntimeSubjectIncarnation: async (incarnation: number) => { + options.onActivate?.(incarnation); + }, configureNetworkConstraints: async () => { options.onConfigureNetwork?.(); if (options.configureNetworkError) { @@ -205,6 +338,7 @@ function createSandboxHandle( } }, createBackup: unavailable, + createRuntimeSubjectBackup: unavailable, createSession: unavailable, deleteSession: unavailable, destroy: async () => { @@ -216,8 +350,29 @@ function createSandboxHandle( await options.destroyPromise; }, - exec: unavailable, + destroyRuntimeSubjectIncarnation: async () => { + options.onDestroy?.(); + + if (options.destroyError) { + throw options.destroyError; + } + + await options.destroyPromise; + return { kind: "destroyed" as const }; + }, + exec: options.onExec + ? async (command) => { + options.onExec?.(command); + return { exitCode: 0, stderr: "", stdout: "", success: true }; + } + : unavailable, getSession: unavailable, + inspectRuntimeSubjectIncarnation: async () => ({ + kind: options.inspectKind ?? "healthy", + }), + markRuntimeSubjectIncarnationReady: async (incarnation: number) => { + options.onReady?.(incarnation); + }, mkdir: async () => { if (options.prepareError) { throw options.prepareError; @@ -248,9 +403,14 @@ function createBindings( readonly configureNetworkError?: Error; readonly destroyError?: Error; readonly destroyPromise?: Promise; + readonly inspectKind?: "healthy" | "missing" | "retired" | "stale" | "unknown"; + readonly onActivate?: (incarnation: number) => void; readonly onConfigureNetwork?: () => void; + readonly onDispose?: () => void; readonly onDestroy?: () => void; + readonly onExec?: (command: string) => void; readonly onRestore?: (backup: { readonly dir: string; readonly id: string }) => void; + readonly onReady?: (incarnation: number) => void; readonly prepareError?: Error; } = {}, ): ApiBindings { @@ -259,42 +419,40 @@ function createBindings( MOSOO_ACCOUNT_CONCURRENT_SANDBOX_LIMIT: options.accountConcurrentSandboxLimit ?? "5", SANDBOX_FILE_BUCKET_LOCAL: "true", runtimeSubjectHandleFactory: () => createSandboxHandle(options), - } as unknown as ApiBindings; + } as ApiBindings; } describe("runtime subject lifecycle machine", () => { - test("keeps subject operation transitions explicit", () => { - expect( - decideRuntimeSubjectTransition({ - currentStatus: "cold", - targetStatus: "restoring", - }), - ).toMatchObject({ kind: "accepted", nextStatus: "restoring" }); - expect( - decideRuntimeSubjectTransition({ - currentStatus: "restoring", - targetStatus: "backing_up", - }), - ).toMatchObject({ kind: "rejected", reason: "illegal_transition" }); - expect( - decideRuntimeSubjectTransition({ - currentStatus: "backing_up", - targetStatus: "destroying", - }), - ).toMatchObject({ kind: "accepted", nextStatus: "destroying" }); - expect( - decideRuntimeSubjectTransition({ - currentStatus: "restoring", - targetStatus: "destroying", - }), - ).toMatchObject({ kind: "accepted", nextStatus: "destroying" }); - // Confirmed teardown returns to cold; there is no `error` status. - expect( - decideRuntimeSubjectTransition({ currentStatus: "active", targetStatus: "cold" }), - ).toMatchObject({ kind: "accepted", nextStatus: "cold" }); - expect( - decideRuntimeSubjectTransition({ currentStatus: "restoring", targetStatus: "cold" }), - ).toMatchObject({ kind: "accepted", nextStatus: "cold" }); + test("keeps every subject operation transition explicit", () => { + const statuses = ["active", "backing_up", "cold", "destroying", "restoring"] as const; + const accepted = new Set([ + "active->backing_up", + "active->cold", + "active->destroying", + "backing_up->active", + "backing_up->cold", + "backing_up->destroying", + "cold->backing_up", + "cold->destroying", + "cold->restoring", + "destroying->cold", + "restoring->active", + "restoring->cold", + "restoring->destroying", + ]); + + for (const currentStatus of statuses) { + for (const targetStatus of statuses) { + const decision = decideRuntimeSubjectTransition({ currentStatus, targetStatus }); + expect(decision.kind).toBe( + currentStatus === targetStatus + ? "duplicate" + : accepted.has(`${currentStatus}->${targetStatus}`) + ? "accepted" + : "rejected", + ); + } + } }); test("rejects Pet Limited before lifecycle admission", async () => { @@ -387,16 +545,22 @@ describe("runtime subject lifecycle machine", () => { const database = createRuntimeSubjectLifecycleDatabase(); await insertRuntimeSubject(database, { status: "active" }); - await expect( - markRuntimeSubjectOperationStarted(database, { - now: 10, - runtimeSubjectId: "01J0000000000000000000000D", - status: "backing_up", - }), - ).resolves.toBe(true); + const lease = await markRuntimeSubjectOperationStarted(database, { + claimExpiresAt: 60_010, + claimOwner: "operation-test", + now: 10, + operationId: createPlatformId(), + operationKind: "hibernate", + runtimeSubjectId: "01J0000000000000000000000D", + }); + expect(lease).not.toBeNull(); + if (lease === null) { + throw new Error("Runtime subject operation lease was not created."); + } await expect( advanceRuntimeSubjectOperationStatus(database, { expectedStatus: "backing_up", + lease, runtimeSubjectId: "01J0000000000000000000000D", status: "destroying", }), @@ -404,6 +568,7 @@ describe("runtime subject lifecycle machine", () => { await markRuntimeSubjectCold(database, { clearBackups: false, expectedStatus: "destroying", + lease: { ...lease, status: "destroying" }, runtimeSubjectId: "01J0000000000000000000000D", }); @@ -415,22 +580,131 @@ describe("runtime subject lifecycle machine", () => { }); }); - test("does not let a stale operation completion overwrite a newer subject status", async () => { + test("does not let a stale operation completion mutate a newer incarnation", async () => { const database = createRuntimeSubjectLifecycleDatabase(); - await insertRuntimeSubject(database, { status: "active", statusSeq: 7 }); + await insertRuntimeSubject(database, { incarnation: 1, status: "active", statusSeq: 7 }); + await insertRuntimeConversationSession(database, { + incarnation: 1, + sessionId: SESSION_ID, + status: "active", + }); + await database + .prepare("INSERT INTO native_resume_ref (session_id) VALUES (?)") + .bind(SESSION_ID) + .run(); + const operationId = createPlatformId(); + const firstLease = await markRuntimeSubjectOperationStarted(database, { + claimExpiresAt: 10, + claimOwner: "first-worker", + now: 1, + operationId, + operationKind: "reset", + runtimeSubjectId: RUNTIME_SUBJECT_ID, + }); + expect(firstLease).not.toBeNull(); + if (firstLease === null) { + throw new Error("The first operation lease was not created."); + } + await expect( + advanceRuntimeSubjectOperationStatus(database, { + expectedStatus: "backing_up", + lease: firstLease, + runtimeSubjectId: RUNTIME_SUBJECT_ID, + status: "destroying", + }), + ).resolves.toBe(true); + const firstDestroyingLease = { ...firstLease, status: "destroying" as const }; + const batchEntered = Promise.withResolvers(); + const releaseFirstWorker = Promise.withResolvers(); + const gatedDatabase = { + batch: async (statements: D1PreparedStatement[]) => { + batchEntered.resolve(); + await releaseFirstWorker.promise; + return database.batch(statements); + }, + prepare: (query: string) => database.prepare(query), + } as D1Database; - await markRuntimeSubjectCold(database, { + const staleCompletion = markRuntimeSubjectCold(gatedDatabase, { clearBackups: false, - expectedStatus: "backing_up", - runtimeSubjectId: "01J0000000000000000000000D", + clearNativeResumeRefs: true, + expectedStatus: "destroying", + lease: firstDestroyingLease, + runtimeSubjectId: RUNTIME_SUBJECT_ID, }); + await batchEntered.promise; + + const takeoverLease = await claimRuntimeSubjectOperationForRepair(database, { + candidate: { + claimExpiresAt: 10, + claimOwner: "first-worker", + id: RUNTIME_SUBJECT_ID, + incarnation: 1, + kind: "cattle", + operationId, + operationKind: "reset", + status: "destroying", + }, + claimExpiresAt: 100, + claimOwner: "takeover-worker", + now: 11, + }); + expect(takeoverLease).not.toBeNull(); + if (takeoverLease === null) { + throw new Error("The takeover operation lease was not created."); + } + await expect( + markRuntimeSubjectCold(database, { + clearBackups: false, + clearNativeResumeRefs: true, + expectedStatus: "destroying", + lease: takeoverLease, + runtimeSubjectId: RUNTIME_SUBJECT_ID, + }), + ).resolves.toBe(true); + await database + .prepare( + `UPDATE sandbox + SET incarnation = 2, status = 'active', status_event = 'runtime_subject.active', + status_seq = status_seq + 1, status_source = 'test', updated_at = 12 + WHERE id = ?`, + ) + .bind(RUNTIME_SUBJECT_ID) + .run(); + await database + .prepare( + `UPDATE sandbox_session + SET sandbox_incarnation = 2, status = 'active', updated_at = 12 + WHERE session_id = ?`, + ) + .bind(SESSION_ID) + .run(); + await database + .prepare("INSERT INTO native_resume_ref (session_id) VALUES (?)") + .bind(SESSION_ID) + .run(); + releaseFirstWorker.resolve(); + + await expect(staleCompletion).resolves.toBe(false); await expect(readRuntimeSubject(database)).resolves.toEqual({ status: "active", status_event: "runtime_subject.active", - status_seq: 7, + status_seq: 11, status_source: "test", }); + await expect( + database + .prepare("SELECT sandbox_incarnation, status FROM sandbox_session WHERE session_id = ?") + .bind(SESSION_ID) + .first(), + ).resolves.toEqual({ sandbox_incarnation: 2, status: "active" }); + await expect( + database + .prepare("SELECT session_id FROM native_resume_ref WHERE session_id = ?") + .bind(SESSION_ID) + .first("session_id"), + ).resolves.toBe(SESSION_ID); }); test("lets interactive activation preempt best-effort prewarm activation claims", async () => { @@ -482,6 +756,71 @@ describe("runtime subject lifecycle machine", () => { }); }); + test("does not let prewarm retire an idle or shared Pet with a legacy network identity", async () => { + for (const hasActiveRun of [false, true]) { + const database = createRuntimeSubjectLifecycleDatabase(); + await insertRuntimeSubject(database, { incarnation: 4, kind: "pet", status: "active" }); + await database + .prepare("UPDATE sandbox SET network_constraints_hash = NULL WHERE id = ?") + .bind(RUNTIME_SUBJECT_ID) + .run(); + if (hasActiveRun) { + const driverInstanceId = createPlatformId(); + await database + .prepare( + `INSERT INTO driver_instance ( + id, sandbox_id, sandbox_incarnation, sandbox_session_id, status + ) VALUES (?, ?, 4, ?, 'ready')`, + ) + .bind(driverInstanceId, RUNTIME_SUBJECT_ID, SESSION_ID) + .run(); + await database + .prepare( + `INSERT INTO session_run (agent_id, driver_instance_id, id, session_id, status) + VALUES (?, ?, ?, ?, 'running')`, + ) + .bind(AGENT_ID, driverInstanceId, createPlatformId(), SESSION_ID) + .run(); + } + let destroyCalls = 0; + + await expect( + createRuntimeSubjectLifecycleService( + createBindings(database, { + onDestroy: () => { + destroyCalls++; + }, + }), + ).activate({ + ...RUNTIME_SUBJECT_QUOTA_SCOPE, + kind: "pet", + networkConstraints: { allowedHosts: [], networkPolicy: "full" }, + purpose: "prewarm", + runtimeSubjectId: RUNTIME_SUBJECT_ID, + subjectId: AGENT_ID, + subjectKind: "agent", + }), + ).rejects.toThrow("network retirement"); + + expect(destroyCalls).toBe(0); + await expect( + database + .prepare( + `SELECT claim_owner, incarnation, operation_kind, status, status_operation_id + FROM sandbox WHERE id = ?`, + ) + .bind(RUNTIME_SUBJECT_ID) + .first(), + ).resolves.toEqual({ + claim_owner: null, + incarnation: 4, + operation_kind: null, + status: "active", + status_operation_id: null, + }); + } + }); + test("captures one sandbox creation when a cold subject becomes active", async () => { const database = createRuntimeSubjectLifecycleDatabase(); await insertRuntimeSubject(database, { status: "cold" }); @@ -496,7 +835,7 @@ describe("runtime subject lifecycle machine", () => { } as ApiBindings; const service = createRuntimeSubjectLifecycleService(bindings); const activation = { - executionOwnerUserId: "01J00000000000000000000002", + ...RUNTIME_SUBJECT_QUOTA_SCOPE, kind: "cattle" as const, networkConstraints: { allowedHosts: [], networkPolicy: "full" as const }, runtimeSubjectId: RUNTIME_SUBJECT_ID, @@ -634,6 +973,120 @@ describe("runtime subject lifecycle machine", () => { expect((await readRuntimeSubject(database)).status).toBe("active"); }); + test("converges a lost recreate to cold while preserving the last ready backup", async () => { + const database = createRuntimeSubjectLifecycleDatabase(); + await insertRuntimeSubject(database, { + incarnation: 1, + kind: "pet", + lastBackupId: STORED_BACKUP_ID, + status: "active", + }); + let destroyCalls = 0; + + await expect( + recreateRuntimeSubjectPreservingState( + createBindings(database, { + inspectKind: "missing", + onDestroy: () => { + destroyCalls++; + }, + }), + { + operationId: createPlatformId(), + reason: "test missing recreate", + runtimeSubjectId: RUNTIME_SUBJECT_ID, + targets: [], + }, + ), + ).rejects.toMatchObject({ name: "RuntimeSubjectPhysicalStateLostError" }); + + expect(destroyCalls).toBe(1); + await expect( + database + .prepare( + `SELECT status, status_operation_id, claim_owner, last_backup_id, + last_error_code + FROM sandbox WHERE id = ?`, + ) + .bind(RUNTIME_SUBJECT_ID) + .first(), + ).resolves.toEqual({ + claim_owner: null, + last_backup_id: STORED_BACKUP_ID, + last_error_code: "runtime.subject_operation_failed", + status: "cold", + status_operation_id: null, + }); + }); + + test("retries a lost cold reset on a new incarnation and clears the old subject backup", async () => { + const database = createRuntimeSubjectLifecycleDatabase(); + await insertReadySubjectBackup(database, 1); + await insertRuntimeSubject(database, { + incarnation: 1, + kind: "pet", + lastBackupId: STORED_BACKUP_ID, + status: "active", + }); + const restoredBackups: Array<{ readonly dir: string; readonly id: string }> = []; + const clearCommands: string[] = []; + const bindings = createBindings(database); + bindings.runtimeSubjectHandleFactory = (physicalId) => + createSandboxHandle({ + inspectKind: physicalId === `${RUNTIME_SUBJECT_ID}-i1` ? "missing" : "healthy", + onExec: (command) => { + clearCommands.push(command); + }, + onRestore: (backup) => { + restoredBackups.push(backup); + }, + }); + + await expect( + resetRuntimeSubjectAgentState(bindings, { + operationId: createPlatformId(), + reason: "test missing reset", + runtimeSubjectId: RUNTIME_SUBJECT_ID, + targets: [], + }), + ).rejects.toMatchObject({ name: "RuntimeSubjectPhysicalStateLostError" }); + await expect( + database + .prepare("SELECT status, last_backup_id FROM sandbox WHERE id = ?") + .bind(RUNTIME_SUBJECT_ID) + .first(), + ).resolves.toEqual({ last_backup_id: STORED_BACKUP_ID, status: "cold" }); + + await expect( + resetRuntimeSubjectAgentState(bindings, { + operationId: createPlatformId(), + reason: "retry missing reset", + runtimeSubjectId: RUNTIME_SUBJECT_ID, + targets: [], + }), + ).resolves.toBeUndefined(); + + expect(restoredBackups).toEqual([{ dir: "/workspace/memory", id: CLOUDFLARE_BACKUP_ID }]); + expect(clearCommands).toHaveLength(1); + await expect( + database + .prepare( + `SELECT incarnation, status, last_backup_id, last_restore_backup_id, + last_error, last_error_code + FROM sandbox WHERE id = ?`, + ) + .bind(RUNTIME_SUBJECT_ID) + .first(), + ).resolves.toEqual({ + incarnation: 2, + last_backup_id: null, + last_error: null, + last_error_code: null, + last_restore_backup_id: null, + status: "cold", + }); + }); + test("lets interactive activation retry after a prior activation failure", async () => { const database = createRuntimeSubjectLifecycleDatabase(); // A prior activation failure leaves the subject cold (no live container) @@ -787,6 +1240,7 @@ describe("runtime subject lifecycle machine", () => { createBindings(database, { destroyPromise }), RUNTIME_SUBJECT_ID, 5, + 5, ), ).rejects.toThrow("Runtime subject destroy"); }); @@ -865,4 +1319,491 @@ describe("runtime subject lifecycle machine", () => { expect(recovered.subject).toBeTruthy(); expect((await readRuntimeSubject(database)).status).toBe("active"); }); + + test("disposes an activation handle that fails before it can be returned", async () => { + const database = createRuntimeSubjectLifecycleDatabase(); + await insertRuntimeSubject(database, { status: "cold" }); + const bindings = createBindings(database); + let createdHandles = 0; + let activationHandleDisposals = 0; + bindings.runtimeSubjectHandleFactory = () => { + const isActivationHandle = createdHandles++ === 0; + return createSandboxHandle( + isActivationHandle + ? { + onDispose: () => (activationHandleDisposals += 1), + prepareError: new Error("injected prepare failure"), + } + : {}, + ); + }; + + await expect( + createRuntimeSubjectLifecycleService(bindings).activate({ + ...RUNTIME_SUBJECT_QUOTA_SCOPE, + kind: "cattle", + networkConstraints: { allowedHosts: [], networkPolicy: "full" }, + runtimeSubjectId: RUNTIME_SUBJECT_ID, + subjectId: SESSION_ID, + subjectKind: "session", + }), + ).rejects.toThrow("injected prepare failure"); + + expect(activationHandleDisposals).toBe(1); + }); + + test("routes each incarnation to one physical sandbox id", async () => { + const physicalIds: string[] = []; + const bindings = createBindings(createRuntimeSubjectLifecycleDatabase()); + bindings.runtimeSubjectHandleFactory = (physicalId) => { + physicalIds.push(physicalId); + return createSandboxHandle(); + }; + + await getRuntimeSubjectKeepAliveHandle(bindings, RUNTIME_SUBJECT_ID, 0); + await getRuntimeSubjectKeepAliveHandle(bindings, RUNTIME_SUBJECT_ID, 36); + + expect(physicalIds).toEqual([RUNTIME_SUBJECT_ID, `${RUNTIME_SUBJECT_ID}-i10`]); + }); + + test("recovers a missing active runtime into a new physical incarnation", async () => { + const database = createRuntimeSubjectLifecycleDatabase(); + await insertRuntimeSubject(database, { status: "active" }); + await insertRuntimeConversationSession(database, { + incarnation: 4, + sessionId: SESSION_ID, + status: "active", + }); + await insertRuntimeConversationSession(database, { + incarnation: 4, + sessionId: "01J0000000000000000000000A", + status: "cleanup_pending", + }); + await database + .prepare( + `INSERT INTO sandbox_backup ( + created_at, dir, error_message, id, keep, sandbox_id, sandbox_incarnation, + status, ttl_seconds, updated_at + ) VALUES (?, ?, NULL, ?, false, ?, ?, 'ready', 600, ?)`, + ) + .bind(1, "/workspace/memory", STORED_BACKUP_ID, RUNTIME_SUBJECT_ID, 4, 1) + .run(); + await database + .prepare( + `UPDATE sandbox + SET incarnation = 4, kind = 'pet', last_backup_id = ?, subject_id = ?, subject_kind = 'agent' + WHERE id = ?`, + ) + .bind(STORED_BACKUP_ID, AGENT_ID, RUNTIME_SUBJECT_ID) + .run(); + const provisioningAuthority = await insertRuntimeProvisioningAuthority(database, 4); + + let releaseOldMutation: (() => void) | undefined; + let oldRetired = false; + const oldMutationBarrier = new Promise((resolve) => { + releaseOldMutation = resolve; + }); + const oldHandle = { + ...createSandboxHandle({ inspectKind: "missing" }), + destroyRuntimeSubjectIncarnation: async () => { + oldRetired = true; + return { kind: "destroyed" as const }; + }, + exec: async () => { + await oldMutationBarrier; + if (oldRetired) { + throw new Error("Runtime subject incarnation is retired."); + } + return { exitCode: 0, stderr: "", stdout: "", success: true }; + }, + } as SandboxHandle; + let restored = false; + let readyIncarnation: number | null = null; + const newHandle = createSandboxHandle({ + onReady: (incarnation) => { + readyIncarnation = incarnation; + }, + onRestore: () => { + restored = true; + }, + }); + const physicalIds: string[] = []; + const bindings = createBindings(database); + bindings.runtimeSubjectHandleFactory = (physicalId) => { + physicalIds.push(physicalId); + return physicalId === `${RUNTIME_SUBJECT_ID}-i4` ? oldHandle : newHandle; + }; + + const staleMutation = oldHandle.exec("stale-write"); + const lifecycle = createRuntimeSubjectLifecycleService(bindings); + const activationInput = { + ...RUNTIME_SUBJECT_QUOTA_SCOPE, + kind: "pet" as const, + networkConstraints: { allowedHosts: [], networkPolicy: "full" as const }, + provisioningAuthority, + runtimeSubjectId: RUNTIME_SUBJECT_ID, + subjectId: AGENT_ID, + subjectKind: "agent" as const, + }; + + await expect(lifecycle.activate(activationInput)).rejects.toThrow("retired"); + const activation = await lifecycle.activate(activationInput); + + releaseOldMutation?.(); + await expect(staleMutation).rejects.toThrow("retired"); + expect(activation.incarnation).toBe(5); + expect(restored).toBe(true); + expect(readyIncarnation).toBe(5); + expect(physicalIds).toContain(`${RUNTIME_SUBJECT_ID}-i4`); + expect(physicalIds).toContain(`${RUNTIME_SUBJECT_ID}-i5`); + await expect( + database + .prepare("SELECT incarnation, status FROM sandbox WHERE id = ?") + .bind(RUNTIME_SUBJECT_ID) + .first(), + ).resolves.toEqual({ incarnation: 5, status: "active" }); + await expect( + database + .prepare("SELECT sandbox_incarnation, status FROM sandbox_session WHERE session_id = ?") + .bind(SESSION_ID) + .first(), + ).resolves.toEqual({ sandbox_incarnation: 4, status: "closed" }); + await expect( + database + .prepare( + "SELECT cleanup_operation_id, sandbox_incarnation, status FROM sandbox_session WHERE session_id = ?", + ) + .bind("01J0000000000000000000000A") + .first(), + ).resolves.toEqual({ cleanup_operation_id: null, sandbox_incarnation: 4, status: "closed" }); + }); + + test("retires an ambiguously timed-out restoring incarnation before retry", async () => { + const database = createRuntimeSubjectLifecycleDatabase(); + await insertRuntimeSubject(database, { incarnation: 3, kind: "pet", status: "cold" }); + + let releaseOldMutation: (() => void) | undefined; + let oldRetired = false; + const oldMutationBarrier = new Promise((resolve) => { + releaseOldMutation = resolve; + }); + const oldHandle = { + ...createSandboxHandle({ + inspectKind: "healthy", + prepareError: createTimeoutError({ + label: "Runtime subject filesystem prepare", + timeoutMs: 15_000, + }), + }), + destroyRuntimeSubjectIncarnation: async () => { + oldRetired = true; + return { kind: "destroyed" as const }; + }, + exec: async () => { + await oldMutationBarrier; + if (oldRetired) { + throw new Error("Runtime subject incarnation is retired."); + } + return { exitCode: 0, stderr: "", stdout: "late", success: true }; + }, + } as SandboxHandle; + const newHandle = createSandboxHandle(); + const bindings = createBindings(database); + bindings.runtimeSubjectHandleFactory = (physicalId) => + physicalId === `${RUNTIME_SUBJECT_ID}-i4` ? oldHandle : newHandle; + const lifecycle = createRuntimeSubjectLifecycleService(bindings); + const activationInput = { + ...RUNTIME_SUBJECT_QUOTA_SCOPE, + kind: "pet" as const, + networkConstraints: { allowedHosts: [], networkPolicy: "full" as const }, + runtimeSubjectId: RUNTIME_SUBJECT_ID, + subjectId: AGENT_ID, + subjectKind: "agent" as const, + }; + + const staleMutation = oldHandle.exec("late-write"); + await expect(lifecycle.activate(activationInput)).rejects.toThrow("timed out"); + await expect( + database + .prepare("SELECT incarnation, status FROM sandbox WHERE id = ?") + .bind(RUNTIME_SUBJECT_ID) + .first(), + ).resolves.toEqual({ incarnation: 4, status: "cold" }); + + const recovered = await lifecycle.activate(activationInput); + releaseOldMutation?.(); + await expect(staleMutation).rejects.toThrow("retired"); + expect(recovered.incarnation).toBe(5); + }); + + test("keeps an active incarnation on an unknown health result", async () => { + const database = createRuntimeSubjectLifecycleDatabase(); + await insertRuntimeSubject(database, { status: "active" }); + await database + .prepare("UPDATE sandbox SET incarnation = 4 WHERE id = ?") + .bind(RUNTIME_SUBJECT_ID) + .run(); + let destroyCalls = 0; + let disposeCalls = 0; + + await expect( + createRuntimeSubjectLifecycleService( + createBindings(database, { + inspectKind: "unknown", + onDestroy: () => (destroyCalls += 1), + onDispose: () => (disposeCalls += 1), + }), + ).activate({ + ...RUNTIME_SUBJECT_QUOTA_SCOPE, + kind: "cattle", + networkConstraints: { allowedHosts: [], networkPolicy: "full" }, + runtimeSubjectId: RUNTIME_SUBJECT_ID, + subjectId: SESSION_ID, + subjectKind: "session", + }), + ).rejects.toThrow("health is unknown"); + + expect(destroyCalls).toBe(0); + expect(disposeCalls).toBe(1); + await expect( + database + .prepare("SELECT claim_owner, incarnation, status FROM sandbox WHERE id = ?") + .bind(RUNTIME_SUBJECT_ID) + .first(), + ).resolves.toEqual({ claim_owner: null, incarnation: 4, status: "active" }); + }); + + test("retires a cattle conversation after its missing incarnation is destroyed", async () => { + const database = createRuntimeSubjectLifecycleDatabase(); + await insertRuntimeSubject(database, { status: "active" }); + await database + .prepare("UPDATE sandbox SET incarnation = 4 WHERE id = ?") + .bind(RUNTIME_SUBJECT_ID) + .run(); + await insertRuntimeConversationSession(database, { + incarnation: 4, + sessionId: SESSION_ID, + status: "error", + }); + const provisioningAuthority = await insertRuntimeProvisioningAuthority(database, 4); + const bindings = createBindings(database, { inspectKind: "missing" }); + + const lifecycle = createRuntimeSubjectLifecycleService(bindings); + const activationInput = { + ...RUNTIME_SUBJECT_QUOTA_SCOPE, + kind: "cattle" as const, + networkConstraints: { allowedHosts: [], networkPolicy: "full" as const }, + provisioningAuthority, + runtimeSubjectId: RUNTIME_SUBJECT_ID, + subjectId: SESSION_ID, + subjectKind: "session" as const, + }; + + await expect(lifecycle.activate(activationInput)).rejects.toThrow("retired"); + const activation = await lifecycle.activate(activationInput); + + expect(activation.incarnation).toBe(5); + await expect( + database + .prepare("SELECT sandbox_incarnation, status FROM sandbox_session WHERE session_id = ?") + .bind(SESSION_ID) + .first(), + ).resolves.toEqual({ sandbox_incarnation: 4, status: "closed" }); + }); + + test("repairs active-incarnation retirement without closing a newer conversation", async () => { + const database = createRuntimeSubjectLifecycleDatabase(); + await insertRuntimeSubject(database, { status: "active" }); + await insertRuntimeConversationSession(database, { + incarnation: 4, + sessionId: SESSION_ID, + status: "active", + }); + await insertRuntimeConversationSession(database, { + incarnation: 4, + sessionId: "01J00000000000000000000008", + status: "cleanup_pending", + }); + const newerSessionId = "01J0000000000000000000000A"; + await insertRuntimeConversationSession(database, { + incarnation: 5, + sessionId: newerSessionId, + status: "active", + }); + const operationId = createPlatformId(); + const claimExpiresAt = Date.now() + 60_000; + await database + .prepare( + `UPDATE sandbox + SET claim_expires_at = ?, claim_owner = 'activation-repair', incarnation = 4, + operation_kind = 'activate', status = 'destroying', status_operation_id = ? + WHERE id = ?`, + ) + .bind(claimExpiresAt, operationId, RUNTIME_SUBJECT_ID) + .run(); + + await runRuntimeSubjectOperation(createBindings(database), { + kind: "pet", + lease: { + claimExpiresAt, + claimOwner: "activation-repair", + incarnation: 4, + kind: "activate", + operationId, + status: "destroying", + }, + reason: "repair active incarnation retirement", + runtimeSubjectId: RUNTIME_SUBJECT_ID, + }); + + await expect( + database + .prepare( + `SELECT cleanup_operation_id, sandbox_incarnation, status + FROM sandbox_session + ORDER BY session_id`, + ) + .all(), + ).resolves.toMatchObject({ + results: [ + { cleanup_operation_id: null, sandbox_incarnation: 4, status: "closed" }, + { cleanup_operation_id: null, sandbox_incarnation: 4, status: "closed" }, + { cleanup_operation_id: null, sandbox_incarnation: 5, status: "active" }, + ], + }); + expect((await readRuntimeSubject(database)).status).toBe("cold"); + }); + + test("keeps generic activation repair draining until every exact-incarnation Run is terminal", async () => { + const database = createRuntimeSubjectLifecycleDatabase(); + await insertRuntimeSubject(database, { incarnation: 4, kind: "pet", status: "active" }); + const operationId = createPlatformId(); + const driverInstanceId = createPlatformId(); + const runId = createPlatformId(); + await database + .prepare( + `UPDATE sandbox + SET claim_expires_at = 1, claim_owner = 'expired-activation', + operation_kind = 'activate', status = 'destroying', status_operation_id = ? + WHERE id = ?`, + ) + .bind(operationId, RUNTIME_SUBJECT_ID) + .run(); + await database + .prepare( + `INSERT INTO driver_instance ( + id, sandbox_id, sandbox_incarnation, sandbox_session_id, status + ) VALUES (?, ?, 4, ?, 'ready')`, + ) + .bind(driverInstanceId, RUNTIME_SUBJECT_ID, SESSION_ID) + .run(); + await database + .prepare( + `INSERT INTO session_run (agent_id, driver_instance_id, id, session_id, status) + VALUES (?, ?, ?, ?, 'running')`, + ) + .bind(AGENT_ID, driverInstanceId, runId, SESSION_ID) + .run(); + const firstRepairLease = await claimRuntimeSubjectOperationForRepair(database, { + candidate: { + claimExpiresAt: 1, + claimOwner: "expired-activation", + id: RUNTIME_SUBJECT_ID, + incarnation: 4, + kind: "pet", + operationId, + operationKind: "activate", + status: "destroying", + }, + claimExpiresAt: 60_002, + claimOwner: "first-repair", + now: 2, + }); + expect(firstRepairLease).not.toBeNull(); + if (firstRepairLease === null) { + return; + } + let destroyCalls = 0; + const bindings = createBindings(database, { + onDestroy: () => { + destroyCalls++; + }, + }); + + await expect( + runRuntimeSubjectOperation(bindings, { + kind: "pet", + lease: firstRepairLease, + reason: "test activation retirement drain", + runtimeSubjectId: RUNTIME_SUBJECT_ID, + }), + ).rejects.toThrow("still draining active Runs"); + expect(destroyCalls).toBe(0); + const waiting = await database + .prepare( + `SELECT claim_expires_at, claim_owner, operation_kind, status, status_operation_id + FROM sandbox WHERE id = ?`, + ) + .bind(RUNTIME_SUBJECT_ID) + .first<{ + claim_expires_at: number; + claim_owner: string; + operation_kind: "activate"; + status: "destroying"; + status_operation_id: RuntimeOperationId; + }>(); + expect(waiting).toMatchObject({ + claim_owner: "first-repair", + operation_kind: "activate", + status: "destroying", + status_operation_id: operationId, + }); + if (waiting === null) { + return; + } + + await database + .prepare("UPDATE session_run SET status = 'completed' WHERE id = ?") + .bind(runId) + .run(); + await database + .prepare("UPDATE driver_instance SET status = 'stopped' WHERE id = ?") + .bind(driverInstanceId) + .run(); + const nextNow = waiting.claim_expires_at + 1; + const nextRepairLease = await claimRuntimeSubjectOperationForRepair(database, { + candidate: { + claimExpiresAt: waiting.claim_expires_at, + claimOwner: waiting.claim_owner, + id: RUNTIME_SUBJECT_ID, + incarnation: 4, + kind: "pet", + operationId, + operationKind: "activate", + status: "destroying", + }, + claimExpiresAt: nextNow + 60_000, + claimOwner: "next-repair", + now: nextNow, + }); + expect(nextRepairLease).not.toBeNull(); + if (nextRepairLease === null) { + return; + } + + await runRuntimeSubjectOperation(bindings, { + kind: "pet", + lease: nextRepairLease, + reason: "test activation retirement drained", + runtimeSubjectId: RUNTIME_SUBJECT_ID, + }); + + expect(destroyCalls).toBe(1); + await expect( + database + .prepare("SELECT claim_owner, operation_kind, status FROM sandbox WHERE id = ?") + .bind(RUNTIME_SUBJECT_ID) + .first(), + ).resolves.toEqual({ claim_owner: null, operation_kind: null, status: "cold" }); + }); }); diff --git a/apps/api/tests/runtime-subject-maintenance.test.ts b/apps/api/tests/runtime-subject-maintenance.test.ts index 8160c4cc..3d8c3523 100644 --- a/apps/api/tests/runtime-subject-maintenance.test.ts +++ b/apps/api/tests/runtime-subject-maintenance.test.ts @@ -1,12 +1,26 @@ import { describe, expect, test } from "bun:test"; -import { expireStaleReschedulingSessions } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance.service"; +import { buildRuntimeStateOperationEvents } from "../src/modules/runtime/application/runtime-state-operation-events"; +import { commitRuntimeOperationReadySnapshots } from "../src/modules/runtime/application/runtime-state-operation-target-events"; +import { claimRuntimeOperationTargets } from "../src/modules/runtime/application/runtime-state-operation-target-store"; +import { recordCanonicalSessionRunTerminal } from "../src/modules/runtime/application/session-runs/session-run-terminal-failure.service"; +import { + expireStaleReschedulingSessions, + repairStaleRuntimeOperationTargets, + runSandboxMaintenance, +} from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-maintenance.service"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { + PUBLIC_API_TEST_IDS, createPublicHttpContractDatabase, createPublicHttpTestBindings, + insertActiveSandboxSessionFixture, insertNonOwnerSession, + insertOwnerSession, } from "./helpers/public-api-http-test-fixture"; +import type { SqliteD1Database } from "./helpers/sqlite-d1"; + +const RESCHEDULING_RUN_ID = "01J0000000000000000000000R"; async function insertRunningSessionRun(database: D1Database): Promise { await database @@ -30,7 +44,7 @@ async function insertRunningSessionRun(database: D1Database): Promise { `, ) .bind( - "run-rescheduling", + RESCHEDULING_RUN_ID, "01J0000000000000000000000B", "01J00000000000000000000009", "01J00000000000000000000002", @@ -46,11 +60,104 @@ async function insertRunningSessionRun(database: D1Database): Promise { .run(); await database .prepare("UPDATE session SET last_run_id = ?, status = ?, updated_at = ? WHERE id = ?") - .bind("run-rescheduling", "RESCHEDULING", 1, "01J0000000000000000000000B") + .bind(RESCHEDULING_RUN_ID, "RESCHEDULING", 1, "01J0000000000000000000000B") .run(); } +async function insertSandboxSession(database: SqliteD1Database): Promise { + await insertActiveSandboxSessionFixture(database, { + cwd: "/workspace", + ownerAccountId: PUBLIC_API_TEST_IDS.nonOwnerAccount, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxSessionId: "01J0000000000000000000000S", + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + timestampMs: 1, + }); +} + +function beforeFirstBatch(database: D1Database, action: () => Promise): D1Database { + let raced = false; + + return new Proxy(database, { + get(target, property) { + if (property === "batch") { + return async (statements: D1PreparedStatement[]) => { + if (!raced) { + raced = true; + await action(); + } + return target.batch(statements); + }; + } + + const value = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }); +} + describe("runtime subject maintenance", () => { + test("continues later maintenance after a terminal projection is poisoned", async () => { + const database = await createPublicHttpContractDatabase(); + await Promise.all([insertOwnerSession(database), insertNonOwnerSession(database)]); + await database + .prepare( + `INSERT INTO session_run ( + id, session_id, agent_id, created_by_account_id, trigger, status, + provider, model, runtime_id, trace_id, started_at, completed_at, + created_at, updated_at + ) VALUES (?, ?, ?, ?, 'user_prompt', 'completed', 'openai', 'gpt-5.4', + 'openai-runtime', 'trace-poisoned-terminal', 1, 1, 1, 1)`, + ) + .bind( + PUBLIC_API_TEST_IDS.run, + PUBLIC_API_TEST_IDS.ownerSession, + PUBLIC_API_TEST_IDS.agent, + PUBLIC_API_TEST_IDS.ownerAccount, + ) + .run(); + await database + .prepare( + `INSERT INTO session_message ( + content_text, created_at, created_by_account_id, id, plan_json, + projection_format, role, segments_json, seq, session_id, session_run_id + ) VALUES ('invalid carrier', 1, ?, ?, NULL, 'materialized', 'assistant', + NULL, 1, ?, ?)`, + ) + .bind( + PUBLIC_API_TEST_IDS.ownerAccount, + PUBLIC_API_TEST_IDS.run, + PUBLIC_API_TEST_IDS.ownerSession, + PUBLIC_API_TEST_IDS.run, + ) + .run(); + await database + .prepare("UPDATE session SET last_run_id = ?, status = 'IDLE' WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.run, PUBLIC_API_TEST_IDS.ownerSession) + .run(); + await database + .prepare("UPDATE session SET status = 'RESCHEDULING', updated_at = 1 WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + + await runSandboxMaintenance(createPublicHttpTestBindings(database) as ApiBindings); + + await expect( + database + .prepare("SELECT status FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ status: "TERMINATED" }); + await expect( + database + .prepare("SELECT terminal_reconciliation_attempted_at FROM session_run WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.run) + .first(), + ).resolves.toMatchObject({ + terminal_reconciliation_attempted_at: expect.any(Number), + }); + }); + test("expires stale rescheduling sessions", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); @@ -58,16 +165,27 @@ describe("runtime subject maintenance", () => { const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await expireStaleReschedulingSessions(bindings); await expireStaleReschedulingSessions(bindings); const run = await database .prepare("SELECT error_code, status FROM session_run WHERE id = ?") - .bind("run-rescheduling") + .bind(RESCHEDULING_RUN_ID) .first<{ error_code: string | null; status: string }>(); expect(run).toEqual({ error_code: "session.rescheduling_timeout", status: "failed", }); + const session = await database + .prepare("SELECT status FROM session WHERE id = ?") + .bind("01J0000000000000000000000B") + .first<{ status: string }>(); + const terminalEvents = await database + .prepare("SELECT COUNT(*) AS count FROM session_event WHERE run_id = ? AND event_type = ?") + .bind(RESCHEDULING_RUN_ID, "run.failed") + .first<{ count: number }>(); + expect(session).toEqual({ status: "TERMINATED" }); + expect(terminalEvents).toEqual({ count: 1 }); }); test("does not expire runtime operation owned rescheduling sessions", async () => { @@ -88,7 +206,7 @@ describe("runtime subject maintenance", () => { .first<{ status: string; status_operation_id: string | null }>(); const run = await database .prepare("SELECT error_code, status FROM session_run WHERE id = ?") - .bind("run-rescheduling") + .bind(RESCHEDULING_RUN_ID) .first<{ error_code: string | null; status: string }>(); expect(session).toEqual({ @@ -100,4 +218,381 @@ describe("runtime subject maintenance", () => { status: "running", }); }); + + test("does not append a lifecycle timeout after a no-Run Session reconnects", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await database + .prepare( + `UPDATE session + SET status = 'RESCHEDULING', status_seq = 1, updated_at = 1 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + const racingDatabase = beforeFirstBatch(database, async () => { + await database + .prepare( + `UPDATE session + SET status = 'RUNNING', status_seq = status_seq + 1, updated_at = 2 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + }); + + await expireStaleReschedulingSessions({ + ...(createPublicHttpTestBindings(database) as ApiBindings), + DB: racingDatabase, + }); + + await expect( + database + .prepare("SELECT status, status_seq, updated_at FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ status: "RUNNING", status_seq: 2, updated_at: 2 }); + await expect( + database.prepare("SELECT COUNT(*) AS count FROM session_event").first(), + ).resolves.toEqual({ count: 0 }); + }); + + test("does not fail an active Run after its stale Session reconnects", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await insertRunningSessionRun(database); + const racingDatabase = beforeFirstBatch(database, async () => { + await database + .prepare( + `UPDATE session + SET status = 'RUNNING', status_seq = status_seq + 1, updated_at = 2 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + }); + + await expireStaleReschedulingSessions({ + ...(createPublicHttpTestBindings(database) as ApiBindings), + DB: racingDatabase, + }); + + await expect( + database + .prepare("SELECT error_code, status FROM session_run WHERE id = ?") + .bind(RESCHEDULING_RUN_ID) + .first(), + ).resolves.toEqual({ error_code: null, status: "running" }); + await expect( + database + .prepare("SELECT status, status_seq, updated_at FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ status: "RUNNING", status_seq: 1, updated_at: 2 }); + await expect( + database.prepare("SELECT COUNT(*) AS count FROM session_event").first(), + ).resolves.toEqual({ count: 0 }); + }); + + test("converges an operation target after the atomic start claim crashes", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await insertRunningSessionRun(database); + await insertSandboxSession(database); + await database + .prepare("UPDATE session SET status = 'RUNNING' WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + const operationId = PUBLIC_API_TEST_IDS.operation; + const [updatingEvent] = buildRuntimeStateOperationEvents({ + agentId: PUBLIC_API_TEST_IDS.agent, + operation: "restartDriver", + readyAt: "1970-01-01T00:00:00.002Z", + startedAt: "1970-01-01T00:00:00.001Z", + }); + const claimed = await claimRuntimeOperationTargets(database, { + event: updatingEvent, + operationId, + targets: [ + { + agentId: PUBLIC_API_TEST_IDS.agent, + creatorAccountId: PUBLIC_API_TEST_IDS.nonOwnerAccount, + lastRunId: RESCHEDULING_RUN_ID, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + sessionRuntimeEventSeqCursor: 0, + sessionStatus: "RUNNING", + sessionStatusOperationId: null, + sessionStatusSeq: 0, + sessionUpdatedAt: 1, + }, + ], + }); + expect(claimed).toHaveLength(1); + expect( + await database + .prepare("SELECT status, status_operation_id, updated_at FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).toEqual({ + status: "RESCHEDULING", + status_operation_id: operationId, + updated_at: 1, + }); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + + expect( + await repairStaleRuntimeOperationTargets(bindings, { + limit: 10, + staleUpdatedAtLte: 1, + }), + ).toBe(1); + expect( + await repairStaleRuntimeOperationTargets(bindings, { + limit: 10, + staleUpdatedAtLte: 1, + }), + ).toBe(0); + + expect( + await database + .prepare( + `SELECT session.status AS session_status, + session.status_operation_id, + session_run.completed_at, + session_run.error_code, + session_run.status AS run_status + FROM session + JOIN session_run ON session_run.id = session.last_run_id + WHERE session.id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).toEqual({ + completed_at: 120_001, + error_code: "agent.runtime_state_operation_timeout", + run_status: "expired", + session_status: "TERMINATED", + status_operation_id: null, + }); + expect( + await database + .prepare( + `SELECT event_type, occurred_at + FROM session_event + WHERE run_id = ? OR event_type = 'agent.task.updated' + ORDER BY seq`, + ) + .bind(RESCHEDULING_RUN_ID) + .all(), + ).toMatchObject({ + results: [ + { event_type: "agent.task.updated", occurred_at: 1 }, + { event_type: "run.cancelled", occurred_at: 120_001 }, + ], + }); + }); + + test("leaves a durable ready winner unchanged during maintenance", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await insertSandboxSession(database); + await database + .prepare("UPDATE session SET updated_at = 1 WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + const operationId = PUBLIC_API_TEST_IDS.operation; + const [updatingEvent, readyEvent] = buildRuntimeStateOperationEvents({ + agentId: PUBLIC_API_TEST_IDS.agent, + operation: "restartDriver", + readyAt: "1970-01-01T00:00:00.002Z", + startedAt: "1970-01-01T00:00:00.001Z", + }); + const [claimed] = await claimRuntimeOperationTargets(database, { + event: updatingEvent, + operationId, + targets: [ + { + agentId: PUBLIC_API_TEST_IDS.agent, + creatorAccountId: PUBLIC_API_TEST_IDS.nonOwnerAccount, + lastRunId: null, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + sessionRuntimeEventSeqCursor: 0, + sessionStatus: "IDLE", + sessionStatusOperationId: null, + sessionStatusSeq: 0, + sessionUpdatedAt: 1, + }, + ], + }); + if (claimed === undefined) { + throw new Error("Missing runtime operation claim fixture."); + } + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await commitRuntimeOperationReadySnapshots(bindings, { + event: readyEvent, + operationId, + targets: [claimed.current], + }); + + expect( + await repairStaleRuntimeOperationTargets(bindings, { + limit: 10, + staleUpdatedAtLte: 1, + }), + ).toBe(0); + await expect( + database + .prepare( + `SELECT runtime_event_seq_cursor, status, status_operation_id, updated_at + FROM session + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).resolves.toEqual({ + runtime_event_seq_cursor: 2, + status: "IDLE", + status_operation_id: null, + updated_at: 2, + }); + await expect( + database.prepare("SELECT event_type, source_event_id FROM session_event ORDER BY seq").all(), + ).resolves.toMatchObject({ + results: [ + { + event_type: "agent.task.updated", + source_event_id: `runtime-operation:${operationId}:${PUBLIC_API_TEST_IDS.nonOwnerSession}:updating`, + }, + { + event_type: "agent.task.updated", + source_event_id: `runtime-operation:${operationId}:${PUBLIC_API_TEST_IDS.nonOwnerSession}:ready`, + }, + ], + }); + }); + + test("does not adopt archive or delete cleanup ownership as a runtime operation", async () => { + for (const cleanupOperationKind of ["archive", "delete"] as const) { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await insertSandboxSession(database); + await database + .prepare( + `UPDATE session + SET archived_at = 1, + cleanup_operation_kind = ?, + status = 'RESCHEDULING', + status_operation_id = ?, + updated_at = 1 + WHERE id = ?`, + ) + .bind( + cleanupOperationKind, + PUBLIC_API_TEST_IDS.operation, + PUBLIC_API_TEST_IDS.nonOwnerSession, + ) + .run(); + + expect( + await repairStaleRuntimeOperationTargets( + createPublicHttpTestBindings(database) as ApiBindings, + { limit: 10, staleUpdatedAtLte: 1 }, + ), + ).toBe(0); + expect( + await database + .prepare( + "SELECT cleanup_operation_kind, status, status_operation_id FROM session WHERE id = ?", + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).toEqual({ + cleanup_operation_kind: cleanupOperationKind, + status: "RESCHEDULING", + status_operation_id: PUBLIC_API_TEST_IDS.operation, + }); + } + }); + + test("releases a crashed operation fence when a canonical Driver terminal already won", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await insertRunningSessionRun(database); + await insertSandboxSession(database); + await database + .prepare("UPDATE session SET status = 'RUNNING' WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + const operationId = PUBLIC_API_TEST_IDS.operation; + const [updatingEvent] = buildRuntimeStateOperationEvents({ + agentId: PUBLIC_API_TEST_IDS.agent, + operation: "restartDriver", + readyAt: "1970-01-01T00:00:00.002Z", + startedAt: "1970-01-01T00:00:00.001Z", + }); + await claimRuntimeOperationTargets(database, { + event: updatingEvent, + operationId, + targets: [ + { + agentId: PUBLIC_API_TEST_IDS.agent, + creatorAccountId: PUBLIC_API_TEST_IDS.nonOwnerAccount, + lastRunId: RESCHEDULING_RUN_ID, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + sessionRuntimeEventSeqCursor: 0, + sessionStatus: "RUNNING", + sessionStatusOperationId: null, + sessionStatusSeq: 0, + sessionUpdatedAt: 1, + }, + ], + }); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + error: null, + expectedSessionOperationId: operationId, + runId: RESCHEDULING_RUN_ID, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + source: "driver", + status: "cancelled", + }); + + expect( + await repairStaleRuntimeOperationTargets(bindings, { + limit: 10, + staleUpdatedAtLte: Number.MAX_SAFE_INTEGER, + }), + ).toBe(1); + + expect( + await database + .prepare( + `SELECT session.status AS session_status, + session.status_operation_id, + session_run.status AS run_status, + session_run.status_source + FROM session + JOIN session_run ON session_run.id = session.last_run_id + WHERE session.id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).toEqual({ + run_status: "cancelled", + session_status: "IDLE", + status_operation_id: null, + status_source: "driver", + }); + expect( + await database + .prepare( + "SELECT COUNT(*) AS count FROM session_event WHERE event_type = 'session.lifecycle.updated'", + ) + .first(), + ).toEqual({ count: 0 }); + }); }); diff --git a/apps/api/tests/runtime-subject-network.test.ts b/apps/api/tests/runtime-subject-network.test.ts index 04b2b798..ff13e102 100644 --- a/apps/api/tests/runtime-subject-network.test.ts +++ b/apps/api/tests/runtime-subject-network.test.ts @@ -4,7 +4,7 @@ import { resolveRuntimeSubjectNetworkConstraints } from "../src/modules/runtime/ import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; function createBindings(overrides: Record = {}): ApiBindings { - return overrides as unknown as ApiBindings; + return overrides; } const LIMITED_NETWORK = { diff --git a/apps/api/tests/runtime-subject-operation-authority-migration.test.ts b/apps/api/tests/runtime-subject-operation-authority-migration.test.ts new file mode 100644 index 00000000..282e29d6 --- /dev/null +++ b/apps/api/tests/runtime-subject-operation-authority-migration.test.ts @@ -0,0 +1,959 @@ +import { describe, expect, test } from "bun:test"; + +import { + installProtocolV3CutoverSql, + INSTALL_PROTOCOL_V3_POST_MIGRATION_CUTOVER_SQL, + PROTOCOL_V3_CUTOVER_OBJECTS_SQL, + PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + REMOVE_PROTOCOL_V3_CUTOVER_SQL, +} from "../bin/protocol-v3-cutover"; +import { applyDrizzleMigration, applyDrizzleMigrationsBefore } from "./helpers/drizzle-migrations"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const MIGRATION_TAG = "0019_runtime-subject-operation-authority"; +const id = (suffix: string) => `01J0000000000000000000001${suffix}`; + +const ACCOUNT_ID = id("A"); +const APP_ID = id("B"); +const PET_AGENT_ID = id("C"); +const CATTLE_AGENT_ID = id("D"); +const PET_SESSION_ID = id("E"); +const CATTLE_SESSION_ID = id("F"); +const PET_SANDBOX_ID = id("G"); +const CATTLE_SANDBOX_ID = id("H"); +const PET_CLOUDFLARE_SESSION_ID = id("J"); +const CATTLE_CLOUDFLARE_SESSION_ID = id("K"); +const RUN_ID = id("M"); +const DRIVER_ID = id("N"); +const DRIVER_COMMAND_ID = id("P"); +const EFFECT_ID = id("Q"); +const SERVER_ID = id("R"); +const API_COMMAND_ID = id("S"); +const PET_BACKUP_ID = id("T"); +const CATTLE_BACKUP_ID = id("V"); +const PRUNED_BACKUP_ID = id("W"); +const OPERATION_ID = id("X"); +const DEPLOYMENT_ID = id("5"); + +function createLegacyDatabase(): SqliteD1Database { + const database = new SqliteD1Database(); + applyDrizzleMigrationsBefore(database, MIGRATION_TAG); + database.execute(` + INSERT INTO app ( + created_at, id, name, organization_id, owner_account_id, updated_at + ) VALUES (1, '${APP_ID}', 'Fixture', '${ACCOUNT_ID}', '${ACCOUNT_ID}', 1); + + INSERT INTO agent ( + config_json, created_at, id, kind, model, name, owner_account_id, + app_id, prompt, provider, runtime_id, status, updated_at, visibility + ) VALUES + ('{}', 1, '${PET_AGENT_ID}', 'pet', 'gpt-5.4', 'Pet', '${ACCOUNT_ID}', + '${APP_ID}', '', 'openai', 'codex', 'draft', 1, 'private'), + ('{}', 1, '${CATTLE_AGENT_ID}', 'cattle', 'gpt-5.4', 'Cattle', '${ACCOUNT_ID}', + '${APP_ID}', '', 'openai', 'codex', 'draft', 1, 'private'); + + INSERT INTO session ( + agent_id, created_at, creator_account_id, id, kind, model, app_id, + provider, renamed, runtime_id, status, updated_at + ) VALUES + ('${PET_AGENT_ID}', 1, '${ACCOUNT_ID}', '${PET_SESSION_ID}', 'pet', 'gpt-5.4', + '${APP_ID}', 'openai', 0, 'codex', 'IDLE', 1), + ('${CATTLE_AGENT_ID}', 1, '${ACCOUNT_ID}', '${CATTLE_SESSION_ID}', 'cattle', + 'gpt-5.4', '${APP_ID}', 'openai', 0, 'codex', 'IDLE', 1); + + INSERT INTO sandbox ( + agent_id, app_id, created_at, id, kind, owner_account_id, status, + subject_id, subject_kind, updated_at + ) VALUES + (NULL, NULL, 1, '${PET_SANDBOX_ID}', 'pet', NULL, 'cold', + '${PET_AGENT_ID}', 'agent', 1), + ('${CATTLE_AGENT_ID}', '${APP_ID}', 1, '${CATTLE_SANDBOX_ID}', 'cattle', + '${ACCOUNT_ID}', 'cold', '${CATTLE_SESSION_ID}', 'session', 1); + + INSERT INTO sandbox_session ( + cloudflare_session_id, created_at, cwd, origin_json, sandbox_id, + session_id, status, updated_at + ) VALUES + ('${PET_CLOUDFLARE_SESSION_ID}', 1, '/workspace', '{}', '${PET_SANDBOX_ID}', + '${PET_SESSION_ID}', 'closed', 1), + ('${CATTLE_CLOUDFLARE_SESSION_ID}', 1, '/workspace', '{}', + '${CATTLE_SANDBOX_ID}', '${CATTLE_SESSION_ID}', 'error', 1); + + INSERT INTO session_run ( + agent_id, completed_at, created_at, created_by_account_id, id, session_id, + status, status_event, trace_id, trigger, updated_at + ) VALUES ( + '${CATTLE_AGENT_ID}', 2, 1, '${ACCOUNT_ID}', '${RUN_ID}', '${CATTLE_SESSION_ID}', + 'completed', 'run.complete', 'trace', 'user', 2 + ); + + INSERT INTO driver_instance ( + boot_token_expires_at, boot_token_hash, created_at, expires_at, + generation, heartbeat_count, id, protocol, protocol_version, runtime, + sandbox_id, sandbox_session_id, status, updated_at + ) VALUES ( + 10, x'01', 1, 10, 7, 0, '${DRIVER_ID}', 'rpc', 2, 'openai-runtime', + '${CATTLE_SANDBOX_ID}', '${CATTLE_SESSION_ID}', 'failed', 2 + ); + + INSERT INTO driver_command ( + driver_instance_id, id, issued_at, kind, payload_json, seq, status + ) VALUES ('${DRIVER_ID}', '${DRIVER_COMMAND_ID}', 1, 'session.stop', '{}', 1, 'completed'); + + INSERT INTO external_tool_effect ( + command_id, created_at, driver_instance_id, id, idempotency_key, server_id, + session_run_id, status, tool_name, updated_at + ) VALUES ( + '${DRIVER_COMMAND_ID}', 1, '${DRIVER_ID}', '${EFFECT_ID}', 'fixture-effect', + '${SERVER_ID}', '${RUN_ID}', 'succeeded', 'write', 2 + ); + + INSERT INTO api_command ( + created_at, dedupe_key, id, kind, payload_json, status, updated_at + ) VALUES (1, 'fixture-command', '${API_COMMAND_ID}', 'scheduled_maintenance', '{}', + 'succeeded', 2); + + INSERT INTO sandbox_backup ( + created_at, dir, id, keep, sandbox_id, session_run_id, status, ttl_seconds, + updated_at + ) VALUES + (1, '/pet', '${PET_BACKUP_ID}', 1, '${PET_SANDBOX_ID}', NULL, 'ready', 60, 2), + (1, '/workspace', '${CATTLE_BACKUP_ID}', 0, '${CATTLE_SANDBOX_ID}', + '${RUN_ID}', 'ready', 60, 2), + (1, '/archive', '${PRUNED_BACKUP_ID}', 0, '${CATTLE_SANDBOX_ID}', NULL, + 'pruned', 60, 2); + + UPDATE sandbox + SET last_backup_id = '${PET_BACKUP_ID}' + WHERE id = '${PET_SANDBOX_ID}'; + + UPDATE sandbox + SET last_backup_id = '${CATTLE_BACKUP_ID}', + last_restore_backup_id = '${PRUNED_BACKUP_ID}' + WHERE id = '${CATTLE_SANDBOX_ID}'; + `); + return database; +} + +async function columnNames(database: SqliteD1Database, table: string): Promise { + const result = await database.prepare(`PRAGMA table_info(${table})`).all<{ name: string }>(); + return result.results.map(({ name }) => name); +} + +async function cutoverCatalog(database: SqliteD1Database) { + return ( + await database + .prepare( + `SELECT name, sql, tbl_name, type + FROM sqlite_master + WHERE name = '__protocol_v3_cutover' + OR name GLOB '__protocol_v3_cutover_*' + ORDER BY type, name`, + ) + .all() + ).results; +} + +describe("0019 runtime subject operation authority migration", () => { + test("backfills authoritative identity and preserves only durable backup lineage", async () => { + const database = createLegacyDatabase(); + + applyDrizzleMigration(database, MIGRATION_TAG); + + expect( + await database + .prepare( + `SELECT id, agent_id, app_id, owner_account_id, incarnation, + network_constraints_hash, operation_kind, status_operation_id + FROM sandbox ORDER BY id`, + ) + .all(), + ).toMatchObject({ + results: [ + { + agent_id: PET_AGENT_ID, + app_id: APP_ID, + id: PET_SANDBOX_ID, + incarnation: 0, + network_constraints_hash: null, + operation_kind: null, + owner_account_id: ACCOUNT_ID, + status_operation_id: null, + }, + { + agent_id: CATTLE_AGENT_ID, + app_id: APP_ID, + id: CATTLE_SANDBOX_ID, + incarnation: 0, + network_constraints_hash: null, + operation_kind: null, + owner_account_id: ACCOUNT_ID, + status_operation_id: null, + }, + ], + }); + + expect( + await database + .prepare( + `SELECT id, operation_id, sandbox_incarnation, session_run_id, staging_id, + workspace_session_id + FROM sandbox_backup ORDER BY id`, + ) + .all(), + ).toMatchObject({ + results: [ + { + id: PET_BACKUP_ID, + operation_id: null, + sandbox_incarnation: 0, + session_run_id: null, + staging_id: PET_BACKUP_ID, + workspace_session_id: null, + }, + { + id: CATTLE_BACKUP_ID, + operation_id: null, + sandbox_incarnation: 0, + session_run_id: RUN_ID, + staging_id: CATTLE_BACKUP_ID, + workspace_session_id: CATTLE_SESSION_ID, + }, + { + id: PRUNED_BACKUP_ID, + operation_id: null, + sandbox_incarnation: 0, + session_run_id: null, + staging_id: PRUNED_BACKUP_ID, + workspace_session_id: null, + }, + ], + }); + + expect( + await database + .prepare( + `SELECT name FROM sqlite_master + WHERE name IN ('__runtime_subject_authority_guard', '__runtime_subject_identity')`, + ) + .all(), + ).toMatchObject({ results: [] }); + expect(await columnNames(database, "driver_instance")).toContain("sandbox_incarnation"); + expect(await columnNames(database, "app_deployment")).toContain("active_script_name"); + expect(await columnNames(database, "sandbox_session")).toContain("cleanup_operation_id"); + expect(await columnNames(database, "sandbox_backup_staging")).toEqual( + expect.arrayContaining(["driver_generation", "driver_instance_id"]), + ); + expect(await columnNames(database, "session")).toContain( + "runtime_provisioning_sandbox_incarnation", + ); + expect( + await database + .prepare("SELECT count(*) AS count FROM environment_package_artifact_backup_staging") + .first(), + ).toEqual({ count: 0 }); + expect( + await database + .prepare(`SELECT delivery_generation FROM api_command WHERE id = '${API_COMMAND_ID}'`) + .first(), + ).toEqual({ delivery_generation: 1 }); + expect( + await database.prepare("SELECT count(*) AS count FROM __protocol_v3_cutover").first(), + ).toEqual({ count: 0 }); + expect((await cutoverCatalog(database)).length).toBe( + PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + ); + }); + + test("migrates legacy Pages runs to Workers Static Assets without losing run state", async () => { + const database = createLegacyDatabase(); + const deploymentRunId = id("6"); + database.execute( + `INSERT INTO app_deployment_run ( + app_id, created_at, deployment_id, error_code, error_message, + external_deployment_id, external_project_id, external_version_id, + generated_wrangler_config_json, id, mosoo_config_json, plan_json, + source_branch, source_commit_sha, status, target_kind, + target_project_name, target_script_name, updated_at, url + ) VALUES ( + '${APP_ID}', 1, '${DEPLOYMENT_ID}', 'legacy-code', 'legacy-message', + 'external-deployment', 'external-project', 'external-version', + '{"name":"fixture"}', '${deploymentRunId}', '{"build":"bun"}', + '{"output":"dist"}', 'main', '0123456789abcdef', 'success', + 'cloudflare_pages', 'legacy-project', 'candidate-script', 2, + 'https://fixture.example' + )`, + ); + + applyDrizzleMigration(database, MIGRATION_TAG); + + expect(await database.prepare("SELECT * FROM app_deployment_run").first()).toEqual({ + app_id: APP_ID, + created_at: 1, + deployment_id: DEPLOYMENT_ID, + error_code: "legacy-code", + error_message: "legacy-message", + external_deployment_id: "external-deployment", + external_project_id: "external-project", + external_version_id: "external-version", + generated_wrangler_config_json: '{"name":"fixture"}', + id: deploymentRunId, + mosoo_config_json: '{"build":"bun"}', + plan_json: '{"output":"dist"}', + source_branch: "main", + source_commit_sha: "0123456789abcdef", + status: "success", + target_kind: "cloudflare_static_assets", + target_project_name: "legacy-project", + target_script_name: "candidate-script", + updated_at: 2, + url: "https://fixture.example", + }); + expect( + ( + await database + .prepare( + `SELECT name FROM sqlite_master + WHERE type = 'index' AND tbl_name = 'app_deployment_run' AND sql IS NOT NULL + ORDER BY name`, + ) + .all<{ name: string }>() + ).results.map(({ name }) => name), + ).toEqual([ + "app_deployment_run_active_app_idx", + "app_deployment_run_app_id_idx", + "app_deployment_run_deployment_id_idx", + ]); + expect((await database.prepare("PRAGMA foreign_key_check").all()).results).toEqual([]); + expect(() => + database.execute( + `UPDATE app_deployment_run SET target_kind = 'cloudflare_pages' + WHERE id = '${deploymentRunId}'`, + ), + ).toThrow(); + }); + + test("matches the canonical inert post-migration gate exactly", async () => { + const migrated = createLegacyDatabase(); + applyDrizzleMigration(migrated, MIGRATION_TAG); + + const canonical = createLegacyDatabase(); + applyDrizzleMigration(canonical, MIGRATION_TAG); + canonical.execute(REMOVE_PROTOCOL_V3_CUTOVER_SQL); + canonical.execute(INSTALL_PROTOCOL_V3_POST_MIGRATION_CUTOVER_SQL); + + expect(await cutoverCatalog(migrated)).toEqual(await cutoverCatalog(canonical)); + }); + + test("keeps the complete post-migration schema outside the exact cutover inventory", async () => { + const database = createLegacyDatabase(); + applyDrizzleMigration(database, MIGRATION_TAG); + + expect(await database.prepare(PROTOCOL_V3_CUTOVER_OBJECTS_SQL).first()).toEqual({ + exact_object_count: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + object_count: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + }); + + database.execute(` + DROP TRIGGER app_deployment_run_target_script_update_authority; + CREATE TRIGGER app_deployment_run_target_script_update_authority + BEFORE UPDATE ON app_deployment_run WHEN 0 + BEGIN + SELECT 1; + END; + `); + expect(await database.prepare(PROTOCOL_V3_CUTOVER_OBJECTS_SQL).first()).toEqual({ + exact_object_count: 0, + object_count: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT + 1, + }); + }); + + test("preserves a completed archive as static business state", async () => { + const database = createLegacyDatabase(); + database.execute( + `UPDATE session + SET archived_at = 2, cleanup_operation_kind = 'archive' + WHERE id = '${PET_SESSION_ID}'`, + ); + + applyDrizzleMigration(database, MIGRATION_TAG); + + expect( + await database + .prepare( + `SELECT archived_at, cleanup_operation_kind, status, status_operation_id + FROM session WHERE id = '${PET_SESSION_ID}'`, + ) + .first(), + ).toEqual({ + archived_at: 2, + cleanup_operation_kind: "archive", + status: "IDLE", + status_operation_id: null, + }); + }); + + test("preserves an enabled production gate and immediately rejects old writers", async () => { + const database = createLegacyDatabase(); + const releaseTreeOid = "0123456789abcdef0123456789abcdef01234567"; + database.execute(installProtocolV3CutoverSql(releaseTreeOid)); + + applyDrizzleMigration(database, MIGRATION_TAG); + + expect( + await database + .prepare("SELECT enabled, release_tree_oid FROM __protocol_v3_cutover WHERE id = 1") + .first(), + ).toEqual({ enabled: 1, release_tree_oid: releaseTreeOid }); + expect((await cutoverCatalog(database)).length).toBe( + PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + ); + expect(() => + database.execute( + `UPDATE sandbox + SET claim_expires_at = 100, claim_owner = 'old-worker' + WHERE id = '${PET_SANDBOX_ID}'`, + ), + ).toThrow("protocol v3 cutover blocks sandbox activation"); + expect(() => + database.execute( + `INSERT INTO sandbox_backup_staging ( + created_at, dir, id, operation_id, sandbox_id, sandbox_incarnation, + ttl_seconds, updated_at + ) VALUES (1, '/workspace', '${id("6")}', '${OPERATION_ID}', '${PET_SANDBOX_ID}', + 1, 60, 1)`, + ), + ).toThrow("protocol v3 cutover blocks new sandbox backup staging"); + }); + + test("freezes identity while allowing an incarnation-fenced activation transition", () => { + const database = createLegacyDatabase(); + applyDrizzleMigration(database, MIGRATION_TAG); + + for (const assignment of [ + `id = '${id("Y")}'`, + "kind = 'cattle'", + "subject_kind = 'session'", + `subject_id = '${CATTLE_SESSION_ID}'`, + `agent_id = '${CATTLE_AGENT_ID}'`, + `app_id = '${id("Z")}'`, + `owner_account_id = '${id("0")}'`, + ]) { + expect(() => + database.execute(`UPDATE sandbox SET ${assignment} WHERE id = '${PET_SANDBOX_ID}'`), + ).toThrow("sandbox identity is immutable"); + } + + expect(() => + database.execute( + `UPDATE sandbox + SET status = 'active', network_constraints_hash = '${"0".repeat(64)}' + WHERE id = '${PET_SANDBOX_ID}'`, + ), + ).toThrow(); + + expect(() => + database.execute( + `INSERT INTO sandbox_backup_staging ( + created_at, dir, driver_generation, driver_instance_id, id, operation_id, + sandbox_id, sandbox_incarnation, ttl_seconds, updated_at + ) VALUES (1, '/workspace', 7, '${DRIVER_ID}', '${id("4")}', '${OPERATION_ID}', + '${PET_SANDBOX_ID}', 1, 60, 1)`, + ), + ).toThrow(); + expect(() => + database.execute( + `INSERT INTO sandbox_backup_staging ( + created_at, dir, id, sandbox_id, sandbox_incarnation, session_run_id, + ttl_seconds, updated_at, workspace_session_id + ) VALUES (1, '/workspace', '${id("5")}', '${CATTLE_SANDBOX_ID}', 1, + '${RUN_ID}', 60, 1, '${CATTLE_SESSION_ID}')`, + ), + ).toThrow(); + for (const [driverGeneration, driverInstanceId, suffix] of [ + ["7", "NULL", "7"], + ["NULL", `'${DRIVER_ID}'`, "8"], + ["1.5", `'${DRIVER_ID}'`, "9"], + ["9007199254740992", `'${DRIVER_ID}'`, "A"], + ] as const) { + expect(() => + database.execute( + `INSERT INTO sandbox_backup_staging ( + created_at, dir, driver_generation, driver_instance_id, id, sandbox_id, + sandbox_incarnation, session_run_id, ttl_seconds, updated_at, + workspace_session_id + ) VALUES (1, '/workspace', ${driverGeneration}, ${driverInstanceId}, '${id(suffix)}', + '${CATTLE_SANDBOX_ID}', 1, '${RUN_ID}', 60, 1, + '${CATTLE_SESSION_ID}')`, + ), + ).toThrow(); + } + database.execute( + `INSERT INTO sandbox_backup_staging ( + created_at, dir, driver_generation, driver_instance_id, id, sandbox_id, + sandbox_incarnation, session_run_id, ttl_seconds, updated_at, workspace_session_id + ) VALUES (1, '/workspace', 7, '${DRIVER_ID}', '${id("6")}', + '${CATTLE_SANDBOX_ID}', 1, '${RUN_ID}', 60, 1, '${CATTLE_SESSION_ID}')`, + ); + + database.execute( + `UPDATE sandbox + SET claim_expires_at = 100, + claim_owner = 'activation', + incarnation = incarnation + 1, + network_constraints_hash = '${"0".repeat(64)}', + operation_kind = 'activate', + status = 'restoring', + status_operation_id = '${OPERATION_ID}' + WHERE id = '${PET_SANDBOX_ID}'`, + ); + }); + + test("enforces exact incarnation and provisioning authority on new writes", () => { + const database = createLegacyDatabase(); + applyDrizzleMigration(database, MIGRATION_TAG); + + expect(() => + database.execute( + `UPDATE sandbox_session SET status = 'active' WHERE session_id = '${PET_SESSION_ID}'`, + ), + ).toThrow(); + database.execute( + `UPDATE sandbox_session + SET sandbox_incarnation = 1, status = 'active' + WHERE session_id = '${PET_SESSION_ID}'`, + ); + + expect(() => + database.execute( + `INSERT INTO sandbox_backup_staging ( + created_at, dir, id, operation_id, sandbox_id, sandbox_incarnation, + ttl_seconds, updated_at + ) VALUES (1, '/workspace', '${id("2")}', '${OPERATION_ID}', '${PET_SANDBOX_ID}', + 0, 60, 1)`, + ), + ).toThrow(); + + database.execute( + `UPDATE session + SET runtime_provisioning_heartbeat_at = 1, + runtime_provisioning_operation_id = '${OPERATION_ID}', + runtime_provisioning_sandbox_id = '${PET_SANDBOX_ID}' + WHERE id = '${PET_SESSION_ID}'`, + ); + expect(() => + database.execute( + `UPDATE session + SET runtime_provisioning_heartbeat_at = 1, + runtime_provisioning_operation_id = '${id("3")}', + runtime_provisioning_sandbox_id = '${PET_SANDBOX_ID}' + WHERE id = '${CATTLE_SESSION_ID}'`, + ), + ).toThrow(); + }); + + test("creates the environment artifact staging authority with native constraints", () => { + const database = createLegacyDatabase(); + applyDrizzleMigration(database, MIGRATION_TAG); + const pathsJson = '{"executable":[],"node":[],"python":[]}'; + const digest = "0".repeat(64); + + const insertCommand = ( + commandId: string, + commandDigest: string, + claimOwner: string, + dedupeKey: string, + ): void => { + const payloadJson = JSON.stringify({ appId: APP_ID, inputDigest: commandDigest }); + database.execute( + `INSERT INTO api_command ( + attempt_count, claim_expires_at, claim_owner, created_at, dedupe_key, + delivery_generation, id, kind, payload_json, status, updated_at + ) VALUES (1, 9007199254740991, '${claimOwner}', 1, '${dedupeKey}', 1, + '${commandId}', 'environment_package_artifact_build', '${payloadJson}', + 'running', 1)`, + ); + }; + const insertStage = ( + commandId: string, + stageDigest: string, + options: { + readonly appId?: string; + readonly attemptCount?: string; + readonly claimOwner?: string; + readonly deliveryGeneration?: string; + } = {}, + ): void => { + database.execute( + `INSERT INTO environment_package_artifact_backup_staging ( + app_id, attempt_count, claim_owner, command_id, created_at, + delivery_generation, dir, input_digest, paths_json, updated_at + ) VALUES ('${options.appId ?? APP_ID}', ${options.attemptCount ?? "1"}, + '${options.claimOwner ?? "owner"}', '${commandId}', 1, + ${options.deliveryGeneration ?? "1"}, '/artifact', '${stageDigest}', + '${pathsJson}', 1)`, + ); + }; + + expect(() => insertStage(API_COMMAND_ID, digest)).toThrow( + "environment artifact backup stage lacks command authority", + ); + database.execute( + `UPDATE api_command + SET attempt_count = 1, claim_expires_at = 1, claim_owner = 'owner', + kind = 'environment_package_artifact_build', + payload_json = '${JSON.stringify({ appId: APP_ID, inputDigest: digest })}', + status = 'running' + WHERE id = '${API_COMMAND_ID}'`, + ); + expect(() => insertStage(API_COMMAND_ID, digest)).toThrow( + "environment artifact backup stage lacks command authority", + ); + database.execute( + `UPDATE api_command SET claim_expires_at = 9007199254740991 + WHERE id = '${API_COMMAND_ID}'`, + ); + + for (const [attemptCount, claimOwner, deliveryGeneration, inputDigest] of [ + ["1", "owner", "1", "invalid"], + ["0", "owner", "1", digest], + ["1.5", "owner", "1", digest], + ["9007199254740992", "owner", "1", digest], + ["1", "", "1", digest], + ["1", "owner", "0", digest], + ["1", "owner", "1.5", digest], + ["1", "owner", "9007199254740992", digest], + ] as const) { + expect(() => + insertStage(API_COMMAND_ID, inputDigest, { + attemptCount, + claimOwner, + deliveryGeneration, + }), + ).toThrow(); + } + for (const options of [ + { attemptCount: "2" }, + { claimOwner: "other-owner" }, + { deliveryGeneration: "2" }, + { appId: id("Z") }, + ] as const) { + expect(() => insertStage(API_COMMAND_ID, digest, options)).toThrow( + "environment artifact backup stage lacks command authority", + ); + } + expect(() => insertStage(id("Y"), digest)).toThrow(); + + insertStage(API_COMMAND_ID, digest); + expect(() => + database.execute( + `UPDATE environment_package_artifact_backup_staging SET dir = '/other' + WHERE command_id = '${API_COMMAND_ID}'`, + ), + ).toThrow("environment artifact backup stage is immutable"); + database.execute( + `UPDATE environment_package_artifact_backup_staging + SET actual_backup_id = '${id("4")}', updated_at = 2 + WHERE command_id = '${API_COMMAND_ID}'`, + ); + + const duplicateIntentCommandId = id("2"); + insertCommand(duplicateIntentCommandId, digest, "owner-2", "fixture-command-2"); + expect(() => + insertStage(duplicateIntentCommandId, digest, { claimOwner: "owner-2" }), + ).toThrow(); + + const otherDigest = "1".repeat(64); + const otherCommandId = id("3"); + insertCommand(otherCommandId, otherDigest, "owner-3", "fixture-command-3"); + insertStage(otherCommandId, otherDigest, { claimOwner: "owner-3" }); + expect(() => + database.execute( + `UPDATE environment_package_artifact_backup_staging + SET actual_backup_id = '${id("4")}', updated_at = 2 + WHERE command_id = '${otherCommandId}'`, + ), + ).toThrow(); + + expect(() => + database.execute( + `UPDATE api_command SET delivery_generation = 0 WHERE id = '${API_COMMAND_ID}'`, + ), + ).toThrow(); + expect(() => + database.execute(`DELETE FROM api_command WHERE id = '${API_COMMAND_ID}'`), + ).toThrow(); + }); + + test("lets only exact pre-gate Environment command authority finish staging", async () => { + const database = createLegacyDatabase(); + applyDrizzleMigration(database, MIGRATION_TAG); + const digest = "0".repeat(64); + const payloadJson = JSON.stringify({ appId: APP_ID, inputDigest: digest }); + database.execute(` + UPDATE api_command + SET attempt_count = 1, + claim_expires_at = 9007199254740991, + claim_owner = 'pre-gate-owner', + created_at = 1, + kind = 'environment_package_artifact_build', + payload_json = '${payloadJson}', + status = 'running' + WHERE id = '${API_COMMAND_ID}'; + INSERT INTO __protocol_v3_cutover (id, enabled, release_tree_oid) + VALUES (1, 1, '${"0".repeat(40)}'); + `); + + expect(() => + database.execute( + `INSERT INTO api_command ( + attempt_count, created_at, dedupe_key, delivery_generation, id, kind, + payload_json, status, updated_at + ) VALUES ( + 0, 2, 'post-gate-environment-command', 1, '${id("Y")}', + 'environment_package_artifact_build', '${payloadJson}', 'queued', 2 + )`, + ), + ).toThrow("protocol v3 cutover blocks new nonterminal API commands"); + + database.execute( + `INSERT INTO environment_package_artifact_backup_staging ( + app_id, attempt_count, claim_owner, command_id, created_at, + delivery_generation, dir, input_digest, paths_json, updated_at + ) VALUES ( + '${APP_ID}', 1, 'pre-gate-owner', '${API_COMMAND_ID}', 1, + 1, '/artifact', '${digest}', + '{"executable":[],"node":[],"python":[]}', 1 + )`, + ); + await expect( + database + .prepare( + `SELECT command_id FROM environment_package_artifact_backup_staging + WHERE command_id = '${API_COMMAND_ID}'`, + ) + .first(), + ).resolves.toEqual({ command_id: API_COMMAND_ID }); + }); + + test("rejects public deployment traffic without durable script authority", async () => { + const database = createLegacyDatabase(); + applyDrizzleMigration(database, MIGRATION_TAG); + database.execute( + `INSERT INTO app_deployment ( + active_script_name, app_id, created_at, default_branch, deleted_at, id, + last_successful_url, latest_run_id, mosoo_subdomain, owner_account_id, + repo_name, repo_owner, repo_url, source_kind, updated_at + ) VALUES (NULL, '${APP_ID}', 1, 'main', NULL, '${DEPLOYMENT_ID}', NULL, NULL, + 'fixture-app', '${ACCOUNT_ID}', 'repo', 'owner', + 'https://example.com/repo', 'github_public', 1)`, + ); + + for (const assignment of [ + "last_successful_url = 'https://fixture.example'", + "active_script_name = 'fixture-script'", + "active_script_name = '', last_successful_url = 'https://fixture.example'", + ]) { + expect(() => + database.execute(`UPDATE app_deployment SET ${assignment} WHERE id = '${DEPLOYMENT_ID}'`), + ).toThrow(); + } + + database.execute( + `UPDATE app_deployment + SET active_script_name = NULL, deleted_at = 2, last_successful_url = NULL + WHERE id = '${DEPLOYMENT_ID}'`, + ); + + expect( + await database + .prepare( + `SELECT active_script_name, deleted_at, last_successful_url + FROM app_deployment WHERE id = '${DEPLOYMENT_ID}'`, + ) + .first(), + ).toEqual({ active_script_name: null, deleted_at: 2, last_successful_url: null }); + }); + + test("keeps an older backup incarnation when its workspace advances", async () => { + const database = createLegacyDatabase(); + applyDrizzleMigration(database, MIGRATION_TAG); + database.execute(` + UPDATE sandbox SET incarnation = 2 WHERE id = '${CATTLE_SANDBOX_ID}'; + UPDATE sandbox_session SET sandbox_incarnation = 2 + WHERE session_id = '${CATTLE_SESSION_ID}'; + INSERT INTO sandbox_backup ( + created_at, dir, id, keep, sandbox_id, sandbox_incarnation, session_run_id, + staging_id, status, ttl_seconds, updated_at, workspace_session_id + ) VALUES ( + 3, '/workspace', '${id("7")}', 0, '${CATTLE_SANDBOX_ID}', 1, '${RUN_ID}', + '${id("8")}', 'ready', 60, 3, '${CATTLE_SESSION_ID}' + ); + `); + + expect( + await database + .prepare( + `SELECT backup.sandbox_incarnation AS backup_incarnation, + workspace.sandbox_incarnation AS workspace_incarnation + FROM sandbox_backup AS backup + JOIN sandbox_session AS workspace + ON workspace.session_id = backup.workspace_session_id + WHERE backup.id = '${id("7")}'`, + ) + .first(), + ).toEqual({ backup_incarnation: 1, workspace_incarnation: 2 }); + }); + + const rollbackCases = [ + ["active sandbox", `UPDATE sandbox SET status = 'active' WHERE id = '${PET_SANDBOX_ID}'`], + [ + "partial sandbox identity", + `UPDATE sandbox SET agent_id = '${PET_AGENT_ID}' WHERE id = '${PET_SANDBOX_ID}'`, + ], + [ + "mismatched sandbox identity", + `UPDATE sandbox SET agent_id = '${CATTLE_AGENT_ID}', app_id = '${APP_ID}', + owner_account_id = '${ACCOUNT_ID}' + WHERE id = '${PET_SANDBOX_ID}'`, + ], + [ + "invalid sandbox subject", + `UPDATE sandbox SET subject_kind = 'user' WHERE id = '${PET_SANDBOX_ID}'`, + ], + ["running session", `UPDATE session SET status = 'RUNNING' WHERE id = '${PET_SESSION_ID}'`], + [ + "session status operation", + `UPDATE session SET status_operation_id = '${OPERATION_ID}' WHERE id = '${PET_SESSION_ID}'`, + ], + [ + "session delete cleanup operation", + `UPDATE session + SET archived_at = 1, + cleanup_operation_kind = 'delete', + status = 'RESCHEDULING', + status_operation_id = '${OPERATION_ID}' + WHERE id = '${PET_SESSION_ID}'`, + ], + [ + "session provisioning authority", + `UPDATE session + SET runtime_provisioning_heartbeat_at = 1, + runtime_provisioning_operation_id = '${OPERATION_ID}', + runtime_provisioning_sandbox_id = '${PET_SANDBOX_ID}' + WHERE id = '${PET_SESSION_ID}'`, + ], + [ + "sandbox claim", + `UPDATE sandbox SET claim_owner = 'legacy', claim_expires_at = 10 + WHERE id = '${PET_SANDBOX_ID}'`, + ], + [ + "active sandbox session", + `UPDATE sandbox_session SET status = 'active' WHERE session_id = '${PET_SESSION_ID}'`, + ], + [ + "cross-subject sandbox session", + `UPDATE sandbox_session SET sandbox_id = '${CATTLE_SANDBOX_ID}' + WHERE session_id = '${PET_SESSION_ID}'`, + ], + ["failed backup", `UPDATE sandbox_backup SET status = 'failed' WHERE id = '${PET_BACKUP_ID}'`], + [ + "backup error", + `UPDATE sandbox_backup SET error_message = 'failed' WHERE id = '${PET_BACKUP_ID}'`, + ], + [ + "duplicate terminal backup", + `INSERT INTO sandbox_backup ( + created_at, dir, id, keep, sandbox_id, session_run_id, status, ttl_seconds, updated_at + ) VALUES (1, '/workspace', '${id("4")}', 0, '${CATTLE_SANDBOX_ID}', '${RUN_ID}', + 'pruned', 60, 2)`, + ], + [ + "cross-subject terminal backup", + `UPDATE session_run SET session_id = '${PET_SESSION_ID}' WHERE id = '${RUN_ID}'`, + ], + [ + "cross-sandbox terminal backup", + `UPDATE sandbox SET last_backup_id = NULL WHERE id = '${CATTLE_SANDBOX_ID}'; + UPDATE sandbox_backup SET sandbox_id = '${PET_SANDBOX_ID}' + WHERE id = '${CATTLE_BACKUP_ID}'`, + ], + [ + "cross-directory terminal backup", + `UPDATE sandbox_backup SET dir = '/other' WHERE id = '${CATTLE_BACKUP_ID}'`, + ], + [ + "mismatched terminal run agent", + `UPDATE session_run SET agent_id = '${PET_AGENT_ID}' WHERE id = '${RUN_ID}'`, + ], + [ + "non-completed terminal backup", + `UPDATE session_run SET status = 'failed' WHERE id = '${RUN_ID}'`, + ], + [ + "dangling backup pointer", + `UPDATE sandbox SET last_backup_id = '${id("5")}' WHERE id = '${PET_SANDBOX_ID}'`, + ], + ["live driver", `UPDATE driver_instance SET status = 'ready' WHERE id = '${DRIVER_ID}'`], + ["unsafe driver generation", `UPDATE driver_instance SET generation = 1.5`], + ["active run", `UPDATE session_run SET status = 'running' WHERE id = '${RUN_ID}'`], + [ + "active app deployment run", + `INSERT INTO app_deployment ( + app_id, created_at, default_branch, deleted_at, id, last_successful_url, + latest_run_id, mosoo_subdomain, owner_account_id, repo_name, repo_owner, + repo_url, source_kind, updated_at + ) VALUES ('${APP_ID}', 1, 'main', NULL, '${DEPLOYMENT_ID}', NULL, + '${id("6")}', 'fixture-app', '${ACCOUNT_ID}', 'repo', 'owner', + 'https://example.com/repo', 'github_public', 1); + INSERT INTO app_deployment_run ( + app_id, created_at, deployment_id, id, source_branch, source_commit_sha, + status, updated_at + ) VALUES ('${APP_ID}', 1, '${DEPLOYMENT_ID}', '${id("6")}', 'main', + '0123456789abcdef', 'building', 1)`, + ], + ["queued driver command", `UPDATE driver_command SET driver_generation = 7, status = 'queued'`], + ["claimed external effect", `UPDATE external_tool_effect SET status = 'claimed'`], + ["queued API command", `UPDATE api_command SET status = 'queued'`], + ["running API command", `UPDATE api_command SET status = 'running'`], + [ + "legacy active app deployment traffic", + `INSERT INTO app_deployment ( + app_id, created_at, default_branch, deleted_at, id, last_successful_url, + latest_run_id, mosoo_subdomain, owner_account_id, repo_name, repo_owner, + repo_url, source_kind, updated_at + ) VALUES ('${APP_ID}', 1, 'main', NULL, '${DEPLOYMENT_ID}', + 'https://fixture.example', NULL, 'fixture-app', '${ACCOUNT_ID}', + 'repo', 'owner', 'https://example.com/repo', 'github_public', 1)`, + ], + [ + "legacy deleted app deployment traffic marker", + `INSERT INTO app_deployment ( + app_id, created_at, default_branch, deleted_at, id, last_successful_url, + latest_run_id, mosoo_subdomain, owner_account_id, repo_name, repo_owner, + repo_url, source_kind, updated_at + ) VALUES ('${APP_ID}', 1, 'main', 2, '${DEPLOYMENT_ID}', + 'https://fixture.example', NULL, 'fixture-app', '${ACCOUNT_ID}', + 'repo', 'owner', 'https://example.com/repo', 'github_public', 2)`, + ], + ] as const; + + for (const [name, mutation] of rollbackCases) { + test(`rolls back every statement for ${name}`, async () => { + const database = createLegacyDatabase(); + database.execute(mutation); + + expect(() => applyDrizzleMigration(database, MIGRATION_TAG)).toThrow(); + expect(await columnNames(database, "sandbox")).not.toContain("incarnation"); + expect( + await database + .prepare( + `SELECT count(*) AS count FROM sqlite_master + WHERE name IN ('sandbox_backup_staging', 'sandbox_identity_immutable')`, + ) + .first<{ count: number }>(), + ).toEqual({ count: 0 }); + }); + } +}); diff --git a/apps/api/tests/runtime-subject-recycle.test.ts b/apps/api/tests/runtime-subject-recycle.test.ts index ccd85298..9508808a 100644 --- a/apps/api/tests/runtime-subject-recycle.test.ts +++ b/apps/api/tests/runtime-subject-recycle.test.ts @@ -11,11 +11,16 @@ import { resumeRuntimeSubjectRecycleOperation, } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-recycle.service"; import { + claimRuntimeSubjectOperationForRepair, listInactiveRuntimeSubjects, listStaleRuntimeSubjectOperations, } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store"; +import type { RuntimeSubjectOperationLease } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-subject-store"; import { encodeSandboxBackupIdForStorage } from "../src/modules/runtime/infrastructure/sandbox-backup-id"; -import type { SandboxHandle } from "../src/modules/runtime/infrastructure/sandbox-handles"; +import type { + RuntimeSubjectIncarnationHandle, + SandboxHandle, +} from "../src/modules/runtime/infrastructure/sandbox-handles"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { SqliteD1Database } from "./helpers/sqlite-d1"; @@ -27,10 +32,15 @@ const CLOUDFLARE_BACKUP_IDS = [ ] as const; const BACKUP_ID = encodeSandboxBackupIdForStorage(CLOUDFLARE_BACKUP_ID); const CLAIM_OWNER = "scheduled-maintenance-owner"; -const OPERATION_ID = "01J0000000000000000000000R"; +const OPERATION_ID = parsePlatformId( + "01J0000000000000000000000R", + "operation id", +); const SANDBOX_ID = PLATFORM_ID_FIXTURES.sandbox; -let currentSandbox: SandboxHandle | null = null; +type RuntimeSubjectTestHandle = RuntimeSubjectIncarnationHandle & SandboxHandle; + +let currentSandbox: RuntimeSubjectTestHandle | null = null; function createRuntimeSubjectRecycleDatabase(): SqliteD1Database { const database = new SqliteD1Database({ foreignKeys: false }); @@ -39,19 +49,29 @@ function createRuntimeSubjectRecycleDatabase(): SqliteD1Database { CREATE TABLE driver_instance ( id text PRIMARY KEY NOT NULL, sandbox_id text NOT NULL, + sandbox_incarnation integer DEFAULT 1 NOT NULL, + sandbox_session_id text, generation integer DEFAULT 0 NOT NULL, - status text NOT NULL + status text NOT NULL, + status_operation_id text ); CREATE TABLE sandbox ( + agent_id text NOT NULL DEFAULT '01J00000000000000000000009', + app_id text NOT NULL DEFAULT '01J0000000000000000000000A', claim_expires_at integer, claim_owner text, id text PRIMARY KEY NOT NULL, inactive_deadline_at integer, + incarnation integer NOT NULL, kind text NOT NULL, last_backup_id text, last_error text, last_error_code text, + last_restore_backup_id text, + network_constraints_hash text, + operation_kind text, + owner_account_id text NOT NULL DEFAULT '01J00000000000000000000008', status text NOT NULL, status_changed_at integer DEFAULT 0 NOT NULL, status_event text DEFAULT 'runtime_subject.active' NOT NULL, @@ -66,27 +86,78 @@ function createRuntimeSubjectRecycleDatabase(): SqliteD1Database { CREATE TABLE sandbox_backup ( created_at integer NOT NULL, dir text NOT NULL, - error_message text, id text PRIMARY KEY NOT NULL, keep integer NOT NULL, + operation_id text, sandbox_id text NOT NULL, + sandbox_incarnation integer DEFAULT 1 NOT NULL, session_run_id text, + staging_id text NOT NULL, status text NOT NULL, ttl_seconds integer NOT NULL, - updated_at integer NOT NULL + updated_at integer NOT NULL, + workspace_session_id text + ); + + CREATE TABLE sandbox_backup_staging ( + actual_backup_id text, + claim_owner text, + created_at integer NOT NULL, + dir text NOT NULL, + driver_generation integer, + driver_instance_id text, + id text PRIMARY KEY NOT NULL, + operation_id text, + sandbox_id text NOT NULL, + sandbox_incarnation integer NOT NULL, + session_run_id text, + ttl_seconds integer NOT NULL, + updated_at integer NOT NULL, + updates_subject_backup integer NOT NULL, + workspace_session_id text + ); + + CREATE TABLE sandbox_backup_delete_intent ( + attempted_at integer, + backup_id text PRIMARY KEY NOT NULL, + created_at integer NOT NULL, + delete_after integer NOT NULL, + deleted_at integer + ); + + CREATE TABLE environment_package_artifact_backup ( + backup_id text PRIMARY KEY NOT NULL + ); + + CREATE TABLE environment_package_artifact_backup_staging ( + actual_backup_id text + ); + + CREATE TABLE native_resume_ref ( + committed_session_run_id text, + committed_value text, + observed_session_run_id text, + session_id text PRIMARY KEY NOT NULL, + value text NOT NULL ); CREATE TABLE sandbox_session ( + cleanup_operation_id text, cwd text NOT NULL, sandbox_id text NOT NULL, + sandbox_incarnation integer DEFAULT 1 NOT NULL, session_id text PRIMARY KEY NOT NULL, status text NOT NULL, updated_at integer ); CREATE TABLE session ( + archived_at integer, + cleanup_operation_kind text, id text PRIMARY KEY NOT NULL, last_message_at integer, + runtime_provisioning_operation_id text, + runtime_provisioning_sandbox_id text, status text NOT NULL ); @@ -103,10 +174,13 @@ function createRuntimeSubjectRecycleDatabase(): SqliteD1Database { claim_owner, id, inactive_deadline_at, + incarnation, kind, last_backup_id, last_error, last_error_code, + network_constraints_hash, + operation_kind, status, status_operation_id, status_seq, @@ -115,7 +189,11 @@ function createRuntimeSubjectRecycleDatabase(): SqliteD1Database { subject_kind, updated_at ) - VALUES (9999999999999, '${CLAIM_OWNER}', '${SANDBOX_ID}', 1, 'pet', NULL, NULL, NULL, 'active', NULL, 0, 'test', '01J00000000000000000000009', 'session', 1); + VALUES ( + 9999999999999, '${CLAIM_OWNER}', '${SANDBOX_ID}', 1, 1, 'pet', NULL, NULL, NULL, + '${"0".repeat(64)}', NULL, 'active', NULL, 0, 'test', + '01J00000000000000000000009', 'agent', 1 + ); `); return database; @@ -135,7 +213,7 @@ function createBindings(database: D1Database): ApiBindings { delete: async () => {}, }, Sandbox: {}, - } as unknown as ApiBindings; + }; } function requireRuntimeOperationId(value: string | null | undefined): RuntimeOperationId { @@ -146,6 +224,56 @@ function requireRuntimeOperationId(value: string | null | undefined): RuntimeOpe return parsePlatformId(value, "runtime operation id"); } +function operationLease( + operationId: RuntimeOperationId, + status: RuntimeSubjectOperationLease["status"], +): RuntimeSubjectOperationLease { + return { + claimExpiresAt: Number.MAX_SAFE_INTEGER, + claimOwner: CLAIM_OWNER, + incarnation: 1, + kind: "hibernate", + operationId, + status, + }; +} + +async function claimStaleOperation( + database: D1Database, + operationId: RuntimeOperationId, + status: RuntimeSubjectOperationLease["status"], +): Promise { + const candidates = await listStaleRuntimeSubjectOperations(database, { + limit: 10, + staleChangedAtLte: Number.MAX_SAFE_INTEGER, + }); + expect(candidates).toHaveLength(1); + const candidate = candidates[0]; + if (!candidate) { + throw new Error("Runtime subject repair candidate was not found."); + } + expect(candidate).toMatchObject({ + claimOwner: CLAIM_OWNER, + id: SANDBOX_ID, + incarnation: 1, + kind: "pet", + operationId, + operationKind: "hibernate", + status, + }); + const now = Date.now(); + const lease = await claimRuntimeSubjectOperationForRepair(database, { + candidate, + claimExpiresAt: now + 60_000, + claimOwner: "repair-owner", + now, + }); + if (lease === null) { + throw new Error("Runtime subject repair candidate could not be claimed."); + } + return lease; +} + async function readRuntimeSubjectRecycleRow(database: D1Database): Promise<{ last_backup_id: string | null; last_error: string | null; @@ -177,22 +305,30 @@ async function readRuntimeSubjectRecycleRow(database: D1Database): Promise<{ return row; } -function createSandboxHandle(): SandboxHandle { +function createSandboxHandle(): RuntimeSubjectTestHandle { const unavailable = async () => { throw new Error("Unexpected sandbox test method call."); }; return { + activateRuntimeSubjectIncarnation: async () => {}, configureNetworkConstraints: async () => {}, createBackup: async (options) => ({ dir: options.dir, id: CLOUDFLARE_BACKUP_ID, }), + createRuntimeSubjectBackup: async (_incarnation, options) => ({ + dir: options.dir, + id: CLOUDFLARE_BACKUP_ID, + }), createSession: unavailable, deleteSession: unavailable, destroy: async () => {}, + destroyRuntimeSubjectIncarnation: async () => ({ kind: "destroyed" }), exec: unavailable, getSession: unavailable, + inspectRuntimeSubjectIncarnation: async () => ({ kind: "healthy" }), + markRuntimeSubjectIncarnationReady: async () => {}, mkdir: async () => {}, mountBucket: unavailable, readFile: unavailable, @@ -204,7 +340,7 @@ function createSandboxHandle(): SandboxHandle { watch: unavailable, writeFile: unavailable, wsConnect: unavailable, - } as SandboxHandle; + }; } describe("runtime subject recycle", () => { @@ -215,7 +351,9 @@ describe("runtime subject recycle", () => { SET claim_expires_at = NULL, claim_owner = NULL, inactive_deadline_at = NULL, - kind = 'cattle' + kind = 'cattle', + subject_id = 'session-1', + subject_kind = 'session' WHERE id = '${SANDBOX_ID}'; INSERT INTO sandbox_session (cwd, sandbox_id, session_id, status, updated_at) @@ -246,7 +384,17 @@ describe("runtime subject recycle", () => { test("uses a generated operation id instead of the maintenance claim owner", async () => { const database = createRuntimeSubjectRecycleDatabase(); - currentSandbox = createSandboxHandle(); + let observedOperationId: string | null = null; + currentSandbox = { + ...createSandboxHandle(), + createRuntimeSubjectBackup: async (_incarnation, options) => { + observedOperationId = await database + .prepare("SELECT status_operation_id FROM sandbox WHERE id = ?") + .bind(SANDBOX_ID) + .first("status_operation_id"); + return { dir: options.dir, id: CLOUDFLARE_BACKUP_ID }; + }, + }; await expect( recycleRuntimeSubject(createBindings(database), { @@ -275,8 +423,9 @@ describe("runtime subject recycle", () => { expect(subject?.status).toBe("cold"); expect(subject?.last_backup_id).toBe(BACKUP_ID); - expect(subject?.status_operation_id).not.toBe(CLAIM_OWNER); - expect(isPlatformId(subject?.status_operation_id)).toBe(true); + expect(observedOperationId).not.toBe(CLAIM_OWNER); + expect(isPlatformId(observedOperationId)).toBe(true); + expect(subject?.status_operation_id).toBeNull(); }); test("hibernates one idle pet subject across sessions after checkpointing durable state", async () => { @@ -308,12 +457,11 @@ describe("runtime subject recycle", () => { ('01J0000000000000000000000U', 1, 'TERMINATED'); `); const checkpointDirs: string[] = []; - const preparedDirs: string[] = []; const lifecycleCalls: string[] = []; let backupIndex = 0; currentSandbox = { ...createSandboxHandle(), - createBackup: async (options) => { + createRuntimeSubjectBackup: async (_incarnation, options) => { checkpointDirs.push(options.dir); const id = CLOUDFLARE_BACKUP_IDS[backupIndex]; backupIndex += 1; @@ -322,11 +470,9 @@ describe("runtime subject recycle", () => { } return { dir: options.dir, id }; }, - mkdir: async (path) => { - preparedDirs.push(path); - }, - destroy: async () => { + destroyRuntimeSubjectIncarnation: async () => { lifecycleCalls.push("destroy"); + return { kind: "destroyed" }; }, setKeepAlive: async (keepAlive) => { lifecycleCalls.push(`keepAlive:${keepAlive}`); @@ -362,8 +508,7 @@ describe("runtime subject recycle", () => { "/workspace/se/session-1", "/workspace/se/session-2", ]); - expect(preparedDirs.toSorted()).toEqual(checkpointDirs.toSorted()); - expect(lifecycleCalls).toEqual(["keepAlive:false", "destroy"]); + expect(lifecycleCalls).toEqual(["destroy"]); await expect(readRuntimeSubjectRecycleRow(database)).resolves.toMatchObject({ last_error: null, last_error_code: null, @@ -389,17 +534,17 @@ describe("runtime subject recycle", () => { const database = createRuntimeSubjectRecycleDatabase(); currentSandbox = { ...createSandboxHandle(), - destroy: async () => {}, + destroyRuntimeSubjectIncarnation: async () => ({ kind: "destroyed" }), }; await database .prepare( ` UPDATE sandbox - SET status = ?, status_operation_id = ?, status_changed_at = ?, status_source = ? + SET operation_kind = ?, status = ?, status_operation_id = ?, status_changed_at = ?, status_source = ? WHERE id = ? `, ) - .bind("destroying", OPERATION_ID, 1, "maintenance", SANDBOX_ID) + .bind("hibernate", "destroying", OPERATION_ID, 1, "maintenance", SANDBOX_ID) .run(); await database .prepare( @@ -414,10 +559,9 @@ describe("runtime subject recycle", () => { await expect( resumeRuntimeSubjectRecycleOperation(createBindings(database), { kind: "pet", - operationId: OPERATION_ID, + lease: operationLease(OPERATION_ID, "destroying"), reason: "test.repair", runtimeSubjectId: SANDBOX_ID, - status: "destroying", }), ).resolves.toBe(true); @@ -441,7 +585,7 @@ describe("runtime subject recycle", () => { expect(subject).toEqual({ status: "cold", - status_operation_id: OPERATION_ID, + status_operation_id: null, }); expect(session).toEqual({ status: "closed" }); }); @@ -461,7 +605,7 @@ describe("runtime subject recycle", () => { .run(); currentSandbox = { ...createSandboxHandle(), - createBackup: async (options) => { + createRuntimeSubjectBackup: async (_incarnation, options) => { if (!backupAvailable) { backupAvailable = true; throw new Error("backup service unavailable"); @@ -472,8 +616,9 @@ describe("runtime subject recycle", () => { id: CLOUDFLARE_BACKUP_ID, }; }, - destroy: async () => { + destroyRuntimeSubjectIncarnation: async () => { lifecycleCalls.push("destroy"); + return { kind: "destroyed" }; }, setKeepAlive: async (keepAlive) => { lifecycleCalls.push(`keepAlive:${keepAlive}`); @@ -507,27 +652,14 @@ describe("runtime subject recycle", () => { .bind("01J0000000000000000000000S") .first("status"), ).resolves.toBe("active"); - await expect( - listStaleRuntimeSubjectOperations(database, { - limit: 10, - staleChangedAtLte: Number.MAX_SAFE_INTEGER, - }), - ).resolves.toEqual([ - { - id: SANDBOX_ID, - kind: "pet", - operationId, - status: "backing_up", - }, - ]); + const repairLease = await claimStaleOperation(database, operationId, "backing_up"); await expect( resumeRuntimeSubjectRecycleOperation(createBindings(database), { kind: "pet", - operationId, + lease: repairLease, reason: "test.repair", runtimeSubjectId: SANDBOX_ID, - status: "backing_up", }), ).resolves.toBe(true); @@ -536,9 +668,9 @@ describe("runtime subject recycle", () => { last_error: null, last_error_code: null, status: "cold", - status_operation_id: operationId, + status_operation_id: null, }); - expect(lifecycleCalls).toEqual(["keepAlive:false", "destroy"]); + expect(lifecycleCalls).toEqual(["destroy"]); }); test("keeps destroy failures as stale repair candidates with the recorded backup", async () => { @@ -546,17 +678,18 @@ describe("runtime subject recycle", () => { let destroyAvailable = false; currentSandbox = { ...createSandboxHandle(), - createBackup: async (options) => { + createRuntimeSubjectBackup: async (_incarnation, options) => { return { dir: options.dir, id: CLOUDFLARE_BACKUP_ID, }; }, - destroy: async () => { + destroyRuntimeSubjectIncarnation: async () => { if (!destroyAvailable) { destroyAvailable = true; throw new Error("destroy service unavailable"); } + return { kind: "destroyed" }; }, }; @@ -580,27 +713,14 @@ describe("runtime subject recycle", () => { status: "destroying", status_operation_id: operationId, }); - await expect( - listStaleRuntimeSubjectOperations(database, { - limit: 10, - staleChangedAtLte: Number.MAX_SAFE_INTEGER, - }), - ).resolves.toEqual([ - { - id: SANDBOX_ID, - kind: "pet", - operationId, - status: "destroying", - }, - ]); + const repairLease = await claimStaleOperation(database, operationId, "destroying"); await expect( resumeRuntimeSubjectRecycleOperation(createBindings(database), { kind: "pet", - operationId, + lease: repairLease, reason: "test.repair", runtimeSubjectId: SANDBOX_ID, - status: "destroying", }), ).resolves.toBe(true); @@ -609,7 +729,7 @@ describe("runtime subject recycle", () => { last_error: null, last_error_code: null, status: "cold", - status_operation_id: operationId, + status_operation_id: null, }); }); @@ -619,11 +739,11 @@ describe("runtime subject recycle", () => { .prepare( ` UPDATE sandbox - SET status = ?, status_operation_id = ?, status_changed_at = ?, status_source = ? + SET claim_expires_at = ?, operation_kind = ?, status = ?, status_operation_id = ?, status_changed_at = ?, status_source = ? WHERE id = ? `, ) - .bind("destroying", OPERATION_ID, 10, "maintenance", SANDBOX_ID) + .bind(10, "hibernate", "destroying", OPERATION_ID, 10, "maintenance", SANDBOX_ID) .run(); await expect( @@ -639,9 +759,13 @@ describe("runtime subject recycle", () => { }), ).resolves.toEqual([ { + claimExpiresAt: 10, + claimOwner: CLAIM_OWNER, id: SANDBOX_ID, + incarnation: 1, kind: "pet", operationId: OPERATION_ID, + operationKind: "hibernate", status: "destroying", }, ]); diff --git a/apps/api/tests/runtime-subject-run-lease.test.ts b/apps/api/tests/runtime-subject-run-lease.test.ts index 4435a17b..6a0195ba 100644 --- a/apps/api/tests/runtime-subject-run-lease.test.ts +++ b/apps/api/tests/runtime-subject-run-lease.test.ts @@ -47,21 +47,34 @@ function createRuntimeSubjectLeaseDatabase(): SqliteD1Database { database.execute(` CREATE TABLE driver_instance ( id text PRIMARY KEY NOT NULL, + generation integer NOT NULL, sandbox_id text NOT NULL, + sandbox_incarnation integer NOT NULL, sandbox_session_id text NOT NULL, status text NOT NULL, + status_changed_at integer DEFAULT 0 NOT NULL, + status_event text DEFAULT 'driver.provision' NOT NULL, + status_operation_id text, + status_seq integer DEFAULT 0 NOT NULL, + status_source text DEFAULT 'system' NOT NULL, updated_at integer NOT NULL ); CREATE TABLE sandbox ( + claim_owner text, id text PRIMARY KEY NOT NULL, inactive_deadline_at integer, + incarnation integer NOT NULL, kind text NOT NULL, + operation_kind text, + status text NOT NULL, + status_operation_id text, updated_at integer NOT NULL ); CREATE TABLE sandbox_session ( sandbox_id text NOT NULL, + sandbox_incarnation integer NOT NULL, session_id text PRIMARY KEY NOT NULL, status text NOT NULL ); @@ -80,20 +93,25 @@ function createRuntimeSubjectLeaseDatabase(): SqliteD1Database { WHERE driver_instance_id IS NOT NULL AND status IN ('queued', 'booting', 'running', 'waiting_input'); - INSERT INTO sandbox (id, inactive_deadline_at, kind, updated_at) - VALUES ('${SANDBOX_ID}', 1, 'cattle', 1); + INSERT INTO sandbox ( + claim_owner, id, inactive_deadline_at, incarnation, kind, operation_kind, + status, status_operation_id, updated_at + ) + VALUES (NULL, '${SANDBOX_ID}', 1, 1, 'cattle', NULL, 'active', NULL, 1); - INSERT INTO sandbox_session (sandbox_id, session_id, status) - VALUES ('${SANDBOX_ID}', '${SESSION_ID}', 'active'); + INSERT INTO sandbox_session (sandbox_id, sandbox_incarnation, session_id, status) + VALUES ('${SANDBOX_ID}', 1, '${SESSION_ID}', 'active'); INSERT INTO driver_instance ( id, + generation, sandbox_id, + sandbox_incarnation, sandbox_session_id, status, updated_at ) - VALUES ('${DRIVER_INSTANCE_ID}', '${SANDBOX_ID}', '${SESSION_ID}', 'ready', 1); + VALUES ('${DRIVER_INSTANCE_ID}', 0, '${SANDBOX_ID}', 1, '${SESSION_ID}', 'ready', 1); INSERT INTO session_run (id, session_id, status, status_seq, updated_at) VALUES ('${SESSION_RUN_ID}', '${SESSION_ID}', 'running', 0, 1); @@ -109,8 +127,10 @@ function leaseInput( } = {}, ) { return { + driverGeneration: 0, driverInstanceId: input.driverInstanceId ?? DRIVER_INSTANCE_ID, runtimeSubjectId: SANDBOX_ID, + runtimeSubjectIncarnation: 1, sessionId: SESSION_ID, sessionRunId: input.sessionRunId ?? SESSION_RUN_ID, }; @@ -128,6 +148,7 @@ describe("runtime subject run lease store", () => { await expect( recordRuntimeRunLeaseReleased(database, { driverInstanceId: DRIVER_INSTANCE_ID, + expectedDriverGeneration: 0, expectedSessionRunId: SESSION_RUN_ID, }), ).resolves.toBe(true); @@ -155,6 +176,95 @@ describe("runtime subject run lease store", () => { expect(sandbox?.inactive_deadline_at).toBeNull(); }); + test("does not let an old Driver generation release the replacement generation's lease", async () => { + const database = createRuntimeSubjectLeaseDatabase(); + await recordRuntimeRunLeaseAcquired(database, leaseInput()); + const originalBatch = database.batch.bind(database) as D1Database["batch"]; + let rotateBeforeRelease = true; + database.batch = (async (statements: D1PreparedStatement[]) => { + if (rotateBeforeRelease) { + rotateBeforeRelease = false; + database.execute( + `UPDATE driver_instance SET generation = 1 WHERE id = '${DRIVER_INSTANCE_ID}'`, + ); + } + return originalBatch(statements); + }) as D1Database["batch"]; + + await expect( + recordRuntimeRunLeaseReleasedOutcome(database, { + driverInstanceId: DRIVER_INSTANCE_ID, + expectedDriverGeneration: 0, + expectedSessionRunId: SESSION_RUN_ID, + }), + ).resolves.toEqual({ + reason: "driver_changed", + status: "stale", + transition: "release", + }); + await expect( + database + .prepare("SELECT driver_instance_id FROM session_run WHERE id = ?") + .bind(SESSION_RUN_ID) + .first("driver_instance_id"), + ).resolves.toBe(DRIVER_INSTANCE_ID); + }); + + test("does not let an old Driver generation acquire a lease for its replacement", async () => { + const database = createRuntimeSubjectLeaseDatabase(); + const originalBatch = database.batch.bind(database) as D1Database["batch"]; + let rotateBeforeAcquire = true; + database.batch = (async (statements: D1PreparedStatement[]) => { + if (rotateBeforeAcquire) { + rotateBeforeAcquire = false; + database.execute( + `UPDATE driver_instance SET generation = 1 WHERE id = '${DRIVER_INSTANCE_ID}'`, + ); + } + return originalBatch(statements); + }) as D1Database["batch"]; + + await expect(recordRuntimeRunLeaseAcquiredOutcome(database, leaseInput())).resolves.toEqual({ + reason: "driver_changed", + status: "stale", + transition: "acquire", + }); + await expect( + database + .prepare("SELECT driver_instance_id FROM session_run WHERE id = ?") + .bind(SESSION_RUN_ID) + .first("driver_instance_id"), + ).resolves.toBeNull(); + }); + + test("rolls back the run link when clearing the inactive deadline fails", async () => { + const database = createRuntimeSubjectLeaseDatabase(); + database.execute(` + CREATE TRIGGER reject_inactive_deadline_clear + BEFORE UPDATE OF inactive_deadline_at ON sandbox + WHEN NEW.inactive_deadline_at IS NULL + BEGIN + SELECT RAISE(ABORT, 'injected deadline failure'); + END; + `); + + await expect(recordRuntimeRunLeaseAcquired(database, leaseInput())).rejects.toThrow( + "injected deadline failure", + ); + + const run = await database + .prepare("SELECT driver_instance_id FROM session_run WHERE id = ?") + .bind(SESSION_RUN_ID) + .first<{ driver_instance_id: string | null }>(); + const deadline = await database + .prepare("SELECT inactive_deadline_at FROM sandbox WHERE id = ?") + .bind(SANDBOX_ID) + .first("inactive_deadline_at"); + + expect(run?.driver_instance_id).toBeNull(); + expect(deadline).toBe(1); + }); + test("arms the pet idle deadline after the final run while its conversation stays active", async () => { const database = createRuntimeSubjectLeaseDatabase(); database.execute(`UPDATE sandbox SET kind = 'pet' WHERE id = '${SANDBOX_ID}'`); @@ -166,6 +276,7 @@ describe("runtime subject run lease store", () => { await expect( recordRuntimeRunLeaseReleased(database, { driverInstanceId: DRIVER_INSTANCE_ID, + expectedDriverGeneration: 0, expectedSessionRunId: SESSION_RUN_ID, }), ).resolves.toBe(true); @@ -199,6 +310,7 @@ describe("runtime subject run lease store", () => { await expect( recordRuntimeRunLeaseReleased(database, { driverInstanceId: DRIVER_INSTANCE_ID, + expectedDriverGeneration: 0, expectedSessionRunId: SESSION_RUN_ID, }), ).resolves.toBe(true); @@ -216,6 +328,64 @@ describe("runtime subject run lease store", () => { expect(run?.driver_instance_id).toBe(DRIVER_INSTANCE_ID); }); + test("keeps a terminal Driver non-assignable until its physical close", async () => { + const database = createRuntimeSubjectLeaseDatabase(); + await recordRuntimeRunLeaseAcquired(database, leaseInput()); + database.execute(` + UPDATE session_run SET status = 'completed', status_seq = 1 + WHERE id = '${SESSION_RUN_ID}'; + UPDATE driver_instance SET status_operation_id = '${SESSION_RUN_ID}' + WHERE id = '${DRIVER_INSTANCE_ID}'; + INSERT INTO session_run (id, session_id, status, status_seq, updated_at) + VALUES ('${OTHER_SESSION_RUN_ID}', '${SESSION_ID}', 'running', 0, 2); + `); + + const terminalRelease = await recordRuntimeRunLeaseReleasedOutcome(database, { + driverInstanceId: DRIVER_INSTANCE_ID, + expectedDriverGeneration: 0, + expectedDriverOperationId: SESSION_RUN_ID, + expectedSessionRunId: SESSION_RUN_ID, + retainDriverOperationUntilTerminal: true, + }); + expect(terminalRelease).toMatchObject({ status: "applied" }); + await expect( + database + .prepare("SELECT status, status_operation_id FROM driver_instance WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ status: "stopping", status_operation_id: SESSION_RUN_ID }); + await expect( + recordRuntimeRunLeaseAcquiredOutcome( + database, + leaseInput({ sessionRunId: OTHER_SESSION_RUN_ID }), + ), + ).resolves.toEqual({ + reason: "driver_not_assignable", + status: "rejected", + transition: "acquire", + }); + + database.execute(` + UPDATE driver_instance SET status = 'stopped' + WHERE id = '${DRIVER_INSTANCE_ID}' + `); + await expect( + recordRuntimeRunLeaseReleasedOutcome(database, { + driverInstanceId: DRIVER_INSTANCE_ID, + expectedDriverGeneration: 0, + expectedDriverOperationId: SESSION_RUN_ID, + expectedSessionRunId: SESSION_RUN_ID, + retainDriverOperationUntilTerminal: true, + }), + ).resolves.toMatchObject({ status: "applied" }); + await expect( + database + .prepare("SELECT status, status_operation_id FROM driver_instance WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ status: "stopped", status_operation_id: null }); + }); + test("treats acquiring the same run as idempotent", async () => { const database = createRuntimeSubjectLeaseDatabase(); @@ -260,6 +430,21 @@ describe("runtime subject run lease store", () => { expect(run?.driver_instance_id).toBeNull(); }); + test("does not acquire while terminal cleanup owns the Driver", async () => { + const database = createRuntimeSubjectLeaseDatabase(); + database.execute(` + UPDATE driver_instance + SET status_operation_id = '${SESSION_RUN_ID}' + WHERE id = '${DRIVER_INSTANCE_ID}' + `); + + await expect(recordRuntimeRunLeaseAcquiredOutcome(database, leaseInput())).resolves.toEqual({ + reason: "driver_not_assignable", + status: "rejected", + transition: "acquire", + }); + }); + test("does not steal a run linked to another driver", async () => { const database = createRuntimeSubjectLeaseDatabase(); @@ -459,12 +644,14 @@ describe("runtime subject run lease store", () => { await expect( recordRuntimeRunLeaseReleased(database, { driverInstanceId: DRIVER_INSTANCE_ID, + expectedDriverGeneration: 0, expectedSessionRunId: UNLINKED_SESSION_RUN_ID, }), ).resolves.toBe(false); await expect( recordRuntimeRunLeaseReleasedOutcome(database, { driverInstanceId: DRIVER_INSTANCE_ID, + expectedDriverGeneration: 0, expectedSessionRunId: UNLINKED_SESSION_RUN_ID, }), ).resolves.toEqual({ diff --git a/apps/api/tests/sandbox-backup-lifecycle.test.ts b/apps/api/tests/sandbox-backup-lifecycle.test.ts new file mode 100644 index 00000000..0511d3a7 --- /dev/null +++ b/apps/api/tests/sandbox-backup-lifecycle.test.ts @@ -0,0 +1,3762 @@ +import { describe, expect, test } from "bun:test"; + +import type { ApiCommandId } from "@mosoo/db"; +import { createPlatformId, parsePlatformId } from "@mosoo/id"; +import type { AppId, RuntimeOperationId, SandboxBackupId, SandboxId, SessionId } from "@mosoo/id"; + +import { + publishEnvironmentPackageArtifactBackup, + resolveEnvironmentPackageArtifactBackup, +} from "../src/modules/environments/application/environment-package-artifact-backup"; +import { + claimEnvironmentPackageArtifactBackupActual, + commitEnvironmentPackageArtifactBackup, + getEnvironmentPackageArtifactBackupManifest, + getEnvironmentPackageArtifactBackupStage, + retireExpiredEnvironmentPackageArtifactBackups, + stageEnvironmentPackageArtifactBackup, +} from "../src/modules/environments/application/environment-package-artifact-backup-store"; +import { resolveEnvironmentPackageArtifact } from "../src/modules/environments/application/environment-package-artifact.service"; +import { + createEnvironmentPackageArtifactBackupName, + environmentPackageArtifactDir, + environmentPackageArtifactMetadataKey, + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_REFRESH_WINDOW_MS, + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS, +} from "../src/modules/environments/domain/environment-package-artifact"; +import type { + EnvironmentPackageArtifactKey, + EnvironmentPackageArtifactPaths, +} from "../src/modules/environments/domain/environment-package-artifact"; +import { encodeSandboxBackupIdForStorage } from "../src/modules/runtime/infrastructure/sandbox-backup-id"; +import { + createRuntimeSandboxBackupName, + deleteAuthorizedSandboxBackupObjects, + getSandboxBackupObjectKeys, +} from "../src/modules/runtime/infrastructure/sandbox-backup-platform"; +import { reconcileSandboxBackupPage } from "../src/modules/runtime/infrastructure/sandbox-backup-reconciliation.service"; +import { + authorizeSandboxBackupDeletion, + claimSandboxBackupStageActual, + finalizeSandboxBackupStage, + getSandboxBackupStage, + listPendingSandboxBackupDeletions, + revokeSandboxBackupStage, + revokeSandboxBackupsForSessionDelete, + stageSandboxBackupWrites, +} from "../src/modules/runtime/infrastructure/sandbox-backup-store"; +import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; +import { applyDrizzleMigrationsThrough } from "./helpers/drizzle-migrations"; +import { createApiCommandQueueStub } from "./helpers/public-api-http-test-fixture"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const SANDBOX_ID = parsePlatformId("01J0000000000000000000000D"); +const SESSION_ID = parsePlatformId("01J0000000000000000000000S"); +const OPERATION_ID = parsePlatformId("01J0000000000000000000000A"); +const CLEANUP_OPERATION_ID = parsePlatformId("01J0000000000000000000000B"); +const DIR = "/workspace/current"; +const NOW = Date.now(); +const OLD_UPLOAD = new Date(NOW - 25 * 60 * 60_000); +const ENVIRONMENT_COMMAND_ID = parsePlatformId("01J0000000000000000000000J"); +const ENVIRONMENT_APP_ID = parsePlatformId("01J0000000000000000000000K"); +const ENVIRONMENT_KEY: EnvironmentPackageArtifactKey = { + appId: ENVIRONMENT_APP_ID, + inputDigest: "a".repeat(64), +}; +const ENVIRONMENT_DIR = environmentPackageArtifactDir(ENVIRONMENT_KEY); +const ENVIRONMENT_PATHS: EnvironmentPackageArtifactPaths = { + executable: [`${ENVIRONMENT_DIR}/python/bin`], + node: [`${ENVIRONMENT_DIR}/npm/node_modules`], + python: [`${ENVIRONMENT_DIR}/python/site-packages`], +}; + +function createDatabase(): SqliteD1Database { + const database = new SqliteD1Database(); + applyDrizzleMigrationsThrough(database, "0020_sandbox-backup-object-authority"); + database.execute(` + INSERT INTO sandbox ( + agent_id, app_id, claim_expires_at, claim_owner, created_at, id, incarnation, + kind, network_constraints_hash, operation_kind, owner_account_id, status, status_operation_id, + subject_id, subject_kind, updated_at + ) VALUES ( + '01J0000000000000000000000E', '01J0000000000000000000000F', ${NOW + 60_000}, + 'owner', ${NOW}, '${SANDBOX_ID}', 1, 'pet', '${"0".repeat(64)}', 'hibernate', + '01J0000000000000000000000G', 'backing_up', '${OPERATION_ID}', + '01J0000000000000000000000E', 'agent', ${NOW} + ); + `); + return database; +} + +async function setEnvironmentArtifactCommand( + database: D1Database, + input: { + readonly attemptCount: number; + readonly claimOwner: string; + readonly commandId?: ApiCommandId; + readonly dedupeKey?: string; + readonly deliveryGeneration: number; + readonly key?: EnvironmentPackageArtifactKey; + }, +): Promise { + const commandId = input.commandId ?? ENVIRONMENT_COMMAND_ID; + const key = input.key ?? ENVIRONMENT_KEY; + await database + .prepare( + `INSERT INTO api_command ( + attempt_count, claim_expires_at, claim_owner, created_at, dedupe_key, + delivery_generation, id, kind, payload_json, status, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, 'environment_package_artifact_build', ?, 'running', ?) + ON CONFLICT(id) DO UPDATE SET + attempt_count = excluded.attempt_count, + claim_expires_at = excluded.claim_expires_at, + claim_owner = excluded.claim_owner, + delivery_generation = excluded.delivery_generation, + status = 'running', + updated_at = excluded.updated_at`, + ) + .bind( + input.attemptCount, + Date.now() + 60_000, + input.claimOwner, + NOW, + input.dedupeKey ?? `environment-artifact:${key.appId}:${key.inputDigest}`, + input.deliveryGeneration, + commandId, + JSON.stringify(key), + NOW, + ) + .run(); +} + +async function createEnvironmentArtifactStage( + database: D1Database, + input: { + readonly attemptCount?: number; + readonly claimOwner?: string; + readonly deliveryGeneration?: number; + } = {}, +) { + const authority = { + attemptCount: input.attemptCount ?? 1, + claimOwner: input.claimOwner ?? "environment-owner", + commandId: ENVIRONMENT_COMMAND_ID, + deliveryGeneration: input.deliveryGeneration ?? 1, + }; + await setEnvironmentArtifactCommand(database, authority); + const stage = await stageEnvironmentPackageArtifactBackup(database, { + ...authority, + dir: ENVIRONMENT_DIR, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + return { authority, stage }; +} + +async function succeedEnvironmentArtifactCommand( + database: D1Database, + commandId: ApiCommandId = ENVIRONMENT_COMMAND_ID, +): Promise { + await database + .prepare( + `UPDATE api_command + SET claim_expires_at = NULL, claim_owner = NULL, + completed_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), status = 'succeeded' + WHERE id = ?`, + ) + .bind(commandId) + .run(); +} + +async function expireEnvironmentArtifactManifest( + database: SqliteD1Database, + expiresAt = Date.now() - 1_000, + key: EnvironmentPackageArtifactKey = ENVIRONMENT_KEY, +): Promise { + const trigger = await database + .prepare( + `SELECT sql FROM sqlite_schema + WHERE type = 'trigger' AND name = 'environment_package_artifact_backup_rotation_authority'`, + ) + .first<{ sql: string }>(); + if (trigger === null) { + throw new Error("Rotation authority trigger is missing."); + } + database.execute("DROP TRIGGER environment_package_artifact_backup_rotation_authority"); + await database + .prepare( + `UPDATE environment_package_artifact_backup + SET committed_at = ?, expires_at = ? + WHERE app_id = ? AND input_digest = ?`, + ) + .bind(expiresAt - 25 * 60 * 60_000, expiresAt, key.appId, key.inputDigest) + .run(); + database.execute(trigger.sql); +} + +async function insertRawEnvironmentArtifactManifest( + database: D1Database, + backupId: SandboxBackupId, +): Promise { + await database + .prepare( + `INSERT INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES ( + ?, 1, ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), 1, + CAST(unixepoch('subsec') * 1000 AS INTEGER) + 315359999000, ?, 1, ? + )`, + ) + .bind( + ENVIRONMENT_APP_ID, + backupId, + ENVIRONMENT_COMMAND_ID, + ENVIRONMENT_KEY.inputDigest, + JSON.stringify(ENVIRONMENT_PATHS), + ) + .run(); +} + +interface StoredObject { + readonly body: string; + readonly uploaded: Date; +} + +class MemoryR2Bucket { + readonly gets: string[] = []; + readonly objects = new Map(); + readonly #cursorKeys = new Map(); + #cursorSequence = 0; + failAfterDeletingKey: string | null = null; + + async delete(keys: string | string[]): Promise { + for (const key of typeof keys === "string" ? [keys] : keys) { + this.objects.delete(key); + if (this.failAfterDeletingKey === key) { + this.failAfterDeletingKey = null; + throw new Error("Simulated partial R2 deletion failure."); + } + } + } + + async get(key: string): Promise { + this.gets.push(key); + const object = this.objects.get(key); + return object === undefined + ? null + : ({ + text: async () => object.body, + uploaded: object.uploaded, + } as R2ObjectBody); + } + + async head(key: string): Promise { + const object = this.objects.get(key); + return object === undefined ? null : ({ key, uploaded: object.uploaded } as R2Object); + } + + async list(options: R2ListOptions): Promise { + const keys = [...this.objects.keys()] + .filter((key) => key.startsWith(options.prefix ?? "")) + .toSorted(); + const cursorKey = + options.cursor === undefined ? undefined : this.#cursorKeys.get(options.cursor); + if (options.cursor !== undefined && cursorKey === undefined) { + throw new Error("Unknown R2 test cursor."); + } + const start = cursorKey === undefined ? 0 : keys.findIndex((key) => key > cursorKey); + const normalizedStart = start === -1 ? keys.length : start; + const end = Math.min(normalizedStart + (options.limit ?? 1_000), keys.length); + const pageKeys = keys.slice(normalizedStart, end); + const cursor = `opaque-${++this.#cursorSequence}`; + const lastKey = pageKeys.at(-1); + if (lastKey !== undefined) { + this.#cursorKeys.set(cursor, lastKey); + } + return { + cursor, + delimitedPrefixes: [], + objects: pageKeys.map((key) => ({ + key, + uploaded: this.objects.get(key)!.uploaded, + })) as R2Object[], + truncated: end < keys.length, + }; + } + + async put( + key: string, + body = "", + options: Date | R2PutOptions = OLD_UPLOAD, + ): Promise { + const uploaded = options instanceof Date ? options : new Date(); + this.objects.set(key, { body, uploaded }); + return { key, uploaded } as R2Object; + } + + putBackup(input: { + readonly createdAt?: Date; + readonly dir?: string; + readonly metadataId?: string; + readonly name?: string | null; + readonly platformId: string; + readonly stagingId?: SandboxBackupId; + readonly uploaded?: Date; + }): SandboxBackupId { + const backupId = encodeSandboxBackupIdForStorage(input.platformId); + const [dataKey, metadataKey] = getSandboxBackupObjectKeys(backupId); + const createdAt = input.createdAt ?? new Date(); + void this.put(dataKey, "archive", input.uploaded ?? OLD_UPLOAD); + const name = + input.name !== undefined + ? input.name + : input.stagingId === undefined + ? (() => { + throw new Error("A runtime backup stage ID is required."); + })() + : createRuntimeSandboxBackupName(input.stagingId); + void this.put( + metadataKey, + JSON.stringify({ + dir: input.dir ?? DIR, + id: input.metadataId ?? input.platformId, + name, + createdAt: createdAt.toISOString(), + ttl: ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS, + }), + input.uploaded ?? OLD_UPLOAD, + ); + return backupId; + } + + backupExpiresAt(backupId: SandboxBackupId): number { + const [, metadataKey] = getSandboxBackupObjectKeys(backupId); + const metadata = JSON.parse(this.objects.get(metadataKey)!.body) as { createdAt: string }; + return Date.parse(metadata.createdAt) + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS * 1_000; + } +} + +function createBindings(database: D1Database, bucket: MemoryR2Bucket): ApiBindings { + return { + DB: database, + SANDBOX_STATE_BUCKET: bucket, + } as ApiBindings; +} + +function loseRunAcknowledgementOnce( + database: D1Database, + queryNeedle: string, +): { readonly database: D1Database; readonly wasLost: () => boolean } { + let lost = false; + const wrap = (statement: D1PreparedStatement): D1PreparedStatement => + new Proxy(statement, { + get(target, property, receiver) { + if (property === "bind") { + return (...values: unknown[]) => wrap(Reflect.apply(target.bind, target, values)); + } + if (property === "run") { + return async (...args: unknown[]) => { + const result = await Reflect.apply(target.run, target, args); + if (!lost) { + lost = true; + throw new Error("Simulated D1 write acknowledgement loss."); + } + return result; + }; + } + return Reflect.get(target, property, receiver); + }, + }); + return { + database: new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => { + const statement = target.prepare(query); + return query.includes(queryNeedle) ? wrap(statement) : statement; + }; + } + const value = Reflect.get(target, property, target) as unknown; + return typeof value === "function" ? value.bind(target) : value; + }, + }), + wasLost: () => lost, + }; +} + +async function createOperationStage( + database: D1Database, + input: { readonly updateSubjectBackup?: boolean; readonly workspace?: boolean } = {}, +) { + if (input.workspace === true) { + await database + .prepare( + `INSERT INTO session ( + agent_id, app_id, created_at, creator_account_id, id, kind, last_message_at, + model, provider, renamed, runtime_id, status, updated_at + ) VALUES (?, ?, ?, ?, ?, 'pet', 1, 'gpt-5.4', 'openai', 0, 'openai-runtime', 'IDLE', ?)`, + ) + .bind( + "01J0000000000000000000000E", + "01J0000000000000000000000F", + NOW, + "01J0000000000000000000000G", + SESSION_ID, + NOW, + ) + .run(); + await database + .prepare( + `INSERT INTO sandbox_session ( + cloudflare_session_id, created_at, cwd, origin_json, sandbox_id, + sandbox_incarnation, session_id, status, updated_at + ) VALUES (?, ?, ?, '{}', ?, 1, ?, 'active', ?)`, + ) + .bind("01J0000000000000000000000H", NOW, DIR, SANDBOX_ID, SESSION_ID, NOW) + .run(); + } + const [write] = await stageSandboxBackupWrites(database, { + admission: { + kind: "operation", + lease: { + claimExpiresAt: NOW + 60_000, + claimOwner: "owner", + incarnation: 1, + kind: "hibernate", + operationId: OPERATION_ID, + status: "backing_up", + }, + }, + sandboxId: SANDBOX_ID, + targets: [ + { + dir: DIR, + updateSandboxLastBackup: input.updateSubjectBackup ?? false, + workspaceSessionId: input.workspace === true ? SESSION_ID : null, + }, + ], + ttlSeconds: 100, + }); + if (write?.kind !== "staged") { + throw new Error("Test stage was not created."); + } + return write.stage; +} + +async function readyRows(database: D1Database): Promise> { + return ( + await database + .prepare("SELECT id, staging_id FROM sandbox_backup WHERE status = 'ready' ORDER BY id") + .all<{ id: string; staging_id: string }>() + ).results; +} + +describe("sandbox backup lifecycle", () => { + test("scanner adopts a complete object after the create ACK is lost", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const stage = await createOperationStage(database); + const actualId = bucket.putBackup({ + platformId: "550e8400-e29b-41d4-a716-446655440001", + stagingId: stage.id, + }); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + + expect(await readyRows(database)).toEqual([{ id: actualId, staging_id: stage.id }]); + expect(await getSandboxBackupStage(database, stage.id)).toBeNull(); + }); + + test("two candidates converge on one ready row and the scanner deletes the loser", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const stage = await createOperationStage(database); + const candidates = [ + bucket.putBackup({ + platformId: "550e8400-e29b-41d4-a716-446655440002", + stagingId: stage.id, + }), + bucket.putBackup({ + platformId: "550e8400-e29b-41d4-a716-446655440003", + stagingId: stage.id, + }), + ]; + + await Promise.all( + candidates.map((actualBackupId) => + claimSandboxBackupStageActual(database, { + actualBackupId, + dir: DIR, + sandboxIncarnation: 1, + stagingId: stage.id, + }), + ), + ); + const winner = (await getSandboxBackupStage(database, stage.id))?.actualBackupId; + if (winner === null || winner === undefined) { + throw new Error("Concurrent claims produced no winner."); + } + await finalizeSandboxBackupStage(database, { actualBackupId: winner, stagingId: stage.id }); + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + + expect(await readyRows(database)).toEqual([{ id: winner, staging_id: stage.id }]); + const loser = candidates.find((candidate) => candidate !== winner)!; + expect(getSandboxBackupObjectKeys(winner).every((key) => bucket.objects.has(key))).toBe(true); + expect(getSandboxBackupObjectKeys(loser).every((key) => !bucket.objects.has(key))).toBe(true); + }); + + test("claim keeps its winner when another worker finalizes before the caller resumes", async () => { + const database = createDatabase(); + const stage = await createOperationStage(database); + const actualId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440009"); + let intercepted = false; + const interceptClaim = (statement: D1PreparedStatement): D1PreparedStatement => + new Proxy(statement, { + get(target, property, receiver) { + if (property === "bind") { + return (...values: unknown[]) => + interceptClaim(Reflect.apply(target.bind, target, values)); + } + if (property === "first") { + return async (...args: unknown[]) => { + const result = await Reflect.apply(target.first, target, args); + intercepted = true; + await finalizeSandboxBackupStage(database, { + actualBackupId: actualId, + stagingId: stage.id, + }); + return result; + }; + } + return Reflect.get(target, property, receiver); + }, + }); + const racingDatabase = new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => { + const statement = target.prepare(query); + return query.includes("UPDATE sandbox_backup_staging AS stage") && + query.includes("RETURNING actual_backup_id") + ? interceptClaim(statement) + : statement; + }; + } + const value = Reflect.get(target, property, target) as unknown; + return typeof value === "function" ? value.bind(target) : value; + }, + }); + + await expect( + claimSandboxBackupStageActual(racingDatabase, { + actualBackupId: actualId, + dir: DIR, + sandboxIncarnation: 1, + stagingId: stage.id, + }), + ).resolves.toEqual({ actualBackupId: actualId }); + expect(intercepted).toBe(true); + expect(await readyRows(database)).toEqual([{ id: actualId, staging_id: stage.id }]); + expect(await getSandboxBackupStage(database, stage.id)).toBeNull(); + }); + + test("a stale collector cannot tombstone an object after finalization wins", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const stage = await createOperationStage(database); + const actualId = bucket.putBackup({ + platformId: "550e8400-e29b-41d4-a716-44665544000a", + stagingId: stage.id, + }); + expect( + await claimSandboxBackupStageActual(database, { + actualBackupId: actualId, + dir: DIR, + sandboxIncarnation: 1, + stagingId: stage.id, + }), + ).toEqual({ actualBackupId: actualId }); + + // The collector observed no finalized row, then the authoritative finalizer won D1. + expect(await readyRows(database)).toEqual([]); + await finalizeSandboxBackupStage(database, { + actualBackupId: actualId, + stagingId: stage.id, + }); + expect( + await authorizeSandboxBackupDeletion(database, { + authority: { kind: "runtime_invalid", stagingId: stage.id }, + backupId: actualId, + }), + ).toBe(false); + + expect(await readyRows(database)).toEqual([{ id: actualId, staging_id: stage.id }]); + expect(getSandboxBackupObjectKeys(actualId).every((key) => bucket.objects.has(key))).toBe(true); + }); + + test("an old lease owner cannot borrow a successor lease to finalize its backup", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const stageA = await createOperationStage(database); + const delayedA = bucket.putBackup({ + platformId: "550e8400-e29b-41d4-a716-446655440012", + stagingId: stageA.id, + }); + await database + .prepare( + `UPDATE sandbox + SET claim_owner = 'successor-owner', + claim_expires_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) + 60000 + WHERE id = ?`, + ) + .bind(SANDBOX_ID) + .run(); + + await expect( + claimSandboxBackupStageActual(database, { + actualBackupId: delayedA, + dir: DIR, + sandboxIncarnation: 1, + stagingId: stageA.id, + }), + ).resolves.toBeNull(); + const [writeB] = await stageSandboxBackupWrites(database, { + admission: { + kind: "operation", + lease: { + claimExpiresAt: NOW + 60_000, + claimOwner: "successor-owner", + incarnation: 1, + kind: "hibernate", + operationId: OPERATION_ID, + status: "backing_up", + }, + }, + sandboxId: SANDBOX_ID, + targets: [{ dir: DIR, updateSandboxLastBackup: false, workspaceSessionId: null }], + ttlSeconds: 100, + }); + expect(writeB?.kind).toBe("staged"); + if (writeB?.kind !== "staged") { + throw new Error("Successor lease did not create a replacement backup stage."); + } + expect(writeB.stage).toMatchObject({ claimOwner: "successor-owner" }); + expect(writeB.stage.id).not.toBe(stageA.id); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { cursor: null }); + expect(getSandboxBackupObjectKeys(delayedA).every((key) => !bucket.objects.has(key))).toBe( + true, + ); + expect(await getSandboxBackupStage(database, writeB.stage.id)).not.toBeNull(); + expect(await readyRows(database)).toEqual([]); + }); + + test("a handoff race never returns a third owner's backup stage", async () => { + const database = createDatabase(); + const stageA = await createOperationStage(database); + await database + .prepare( + `UPDATE sandbox + SET claim_owner = 'owner-b', + claim_expires_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) + 60000 + WHERE id = ?`, + ) + .bind(SANDBOX_ID) + .run(); + let stageC: Awaited> = null; + let raced = false; + const interceptDelete = (statement: D1PreparedStatement): D1PreparedStatement => + new Proxy(statement, { + get(target, property, receiver) { + if (property === "bind") { + return (...values: unknown[]) => + interceptDelete(Reflect.apply(target.bind, target, values)); + } + if (property === "first") { + return async (...args: unknown[]) => { + if (!raced) { + raced = true; + await database + .prepare( + `UPDATE sandbox + SET claim_owner = 'owner-c', + claim_expires_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) + 60000 + WHERE id = ?`, + ) + .bind(SANDBOX_ID) + .run(); + await revokeSandboxBackupStage(database, { + onlyIfStale: true, + stagingId: stageA.id, + }); + const [writeC] = await stageSandboxBackupWrites(database, { + admission: { + kind: "operation", + lease: { + claimExpiresAt: NOW + 60_000, + claimOwner: "owner-c", + incarnation: 1, + kind: "hibernate", + operationId: OPERATION_ID, + status: "backing_up", + }, + }, + sandboxId: SANDBOX_ID, + targets: [{ dir: DIR, updateSandboxLastBackup: false, workspaceSessionId: null }], + ttlSeconds: 100, + }); + stageC = writeC?.kind === "staged" ? writeC.stage : null; + } + return Reflect.apply(target.first, target, args); + }; + } + return Reflect.get(target, property, receiver); + }, + }); + const racingDatabase = new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => { + const statement = target.prepare(query); + return query.includes("DELETE FROM sandbox_backup_staging AS stage") && + query.includes("AND NOT") + ? interceptDelete(statement) + : statement; + }; + } + const value = Reflect.get(target, property, target) as unknown; + return typeof value === "function" ? value.bind(target) : value; + }, + }); + + await expect( + stageSandboxBackupWrites(racingDatabase, { + admission: { + kind: "operation", + lease: { + claimExpiresAt: NOW + 60_000, + claimOwner: "owner-b", + incarnation: 1, + kind: "hibernate", + operationId: OPERATION_ID, + status: "backing_up", + }, + }, + sandboxId: SANDBOX_ID, + targets: [{ dir: DIR, updateSandboxLastBackup: false, workspaceSessionId: null }], + ttlSeconds: 100, + }), + ).rejects.toThrow("another admission authority"); + expect(raced).toBe(true); + expect(stageC).toMatchObject({ claimOwner: "owner-c" }); + }); + + test("an expired D1 lease cannot claim a backup stage", async () => { + const database = createDatabase(); + const stage = await createOperationStage(database); + await database + .prepare( + `UPDATE sandbox + SET claim_expires_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) - 1 + WHERE id = ?`, + ) + .bind(SANDBOX_ID) + .run(); + + await expect( + claimSandboxBackupStageActual(database, { + actualBackupId: encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440013"), + dir: DIR, + sandboxIncarnation: 1, + stagingId: stage.id, + }), + ).resolves.toBeNull(); + }); + + test("a partial finalization cannot borrow a successor owner to update the subject", async () => { + const database = createDatabase(); + const stage = await createOperationStage(database, { updateSubjectBackup: true }); + const actualId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440017"); + await claimSandboxBackupStageActual(database, { + actualBackupId: actualId, + dir: DIR, + sandboxIncarnation: 1, + stagingId: stage.id, + }); + await expect( + database + .prepare( + `INSERT OR REPLACE INTO sandbox_backup ( + created_at, dir, id, keep, operation_id, sandbox_id, sandbox_incarnation, + staging_id, status, ttl_seconds, updated_at + ) SELECT created_at, dir, actual_backup_id, 0, operation_id, sandbox_id, + sandbox_incarnation, id, 'pruned', ttl_seconds, + CAST(unixepoch('subsec') * 1000 AS INTEGER) + FROM sandbox_backup_staging WHERE id = ?`, + ) + .bind(stage.id) + .run(), + ).rejects.toThrow("already referenced"); + await database + .prepare( + `INSERT INTO sandbox_backup ( + created_at, dir, id, keep, operation_id, sandbox_id, sandbox_incarnation, + staging_id, status, ttl_seconds, updated_at + ) SELECT created_at, dir, actual_backup_id, 0, operation_id, sandbox_id, + sandbox_incarnation, id, 'ready', ttl_seconds, + CAST(unixepoch('subsec') * 1000 AS INTEGER) + FROM sandbox_backup_staging WHERE id = ?`, + ) + .bind(stage.id) + .run(); + await expect( + database + .prepare("UPDATE sandbox_backup_staging SET actual_backup_id = NULL WHERE id = ?") + .bind(stage.id) + .run(), + ).rejects.toThrow("already owned"); + await expect( + database + .prepare("UPDATE sandbox_backup SET status = 'pruned' WHERE id = ?") + .bind(actualId) + .run(), + ).rejects.toThrow("already referenced"); + await database + .prepare( + `UPDATE sandbox + SET claim_owner = 'successor-owner', + claim_expires_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) + 60000 + WHERE id = ?`, + ) + .bind(SANDBOX_ID) + .run(); + + await expect( + finalizeSandboxBackupStage(database, { + actualBackupId: actualId, + stagingId: stage.id, + }), + ).resolves.toMatchObject({ candidateAccepted: true, complete: false }); + expect( + await database + .prepare("SELECT last_backup_id FROM sandbox WHERE id = ?") + .bind(SANDBOX_ID) + .first(), + ).toEqual({ last_backup_id: null }); + expect(await getSandboxBackupStage(database, stage.id)).not.toBeNull(); + }); + + test("the tombstone guard independently rejects a ready row ID change", async () => { + const database = createDatabase(); + const stage = await createOperationStage(database); + const readyId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-44665544000d"); + const tombstonedId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-44665544000e"); + expect( + await authorizeSandboxBackupDeletion(database, { + authority: { kind: "unattributed" }, + backupId: tombstonedId, + }), + ).toBe(true); + await claimSandboxBackupStageActual(database, { + actualBackupId: readyId, + dir: DIR, + sandboxIncarnation: 1, + stagingId: stage.id, + }); + await finalizeSandboxBackupStage(database, { + actualBackupId: readyId, + stagingId: stage.id, + }); + database.execute("DROP TRIGGER sandbox_backup_identity_immutable"); + + await expect( + database + .prepare("UPDATE sandbox_backup SET id = ? WHERE id = ?") + .bind(tombstonedId, readyId) + .run(), + ).rejects.toThrow("sandbox backup object is tombstoned or already referenced"); + expect(await readyRows(database)).toEqual([{ id: readyId, staging_id: stage.id }]); + }); + + test("runtime final identity cannot use OR REPLACE to delete another owner", async () => { + const database = createDatabase(); + const actualA = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440020"); + const actualB = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440021"); + const actualC = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440022"); + const actualD = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440027"); + const stageA = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440023"); + const stageB = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440024"); + const stageC = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440028"); + const insert = (actualId: SandboxBackupId, stagingId: SandboxBackupId, dir: string) => + database + .prepare( + `INSERT INTO sandbox_backup ( + created_at, dir, id, keep, operation_id, sandbox_id, sandbox_incarnation, + session_run_id, staging_id, status, ttl_seconds, updated_at, workspace_session_id + ) VALUES (?, ?, ?, 0, ?, ?, 1, NULL, ?, 'ready', 100, ?, NULL)`, + ) + .bind(NOW, dir, actualId, OPERATION_ID, SANDBOX_ID, stagingId, NOW) + .run(); + await insert(actualA, stageA, "/workspace/a"); + await insert(actualB, stageB, "/workspace/b"); + await database.prepare("PRAGMA recursive_triggers = OFF").run(); + + for (const [actualId, stagingId, dir] of [ + [actualC, stageA, "/workspace/c"], + [actualD, stageC, "/workspace/b"], + ] as const) { + await expect( + database + .prepare( + `INSERT OR REPLACE INTO sandbox_backup ( + created_at, dir, id, keep, operation_id, sandbox_id, sandbox_incarnation, + session_run_id, staging_id, status, ttl_seconds, updated_at, workspace_session_id + ) VALUES (?, ?, ?, 0, ?, ?, 1, NULL, ?, 'ready', 100, ?, NULL)`, + ) + .bind(NOW, dir, actualId, OPERATION_ID, SANDBOX_ID, stagingId, NOW) + .run(), + ).rejects.toThrow("already referenced"); + } + await expect( + database + .prepare("UPDATE OR REPLACE sandbox_backup SET dir = '/workspace/b' WHERE id = ?") + .bind(actualA) + .run(), + ).rejects.toThrow("identity is immutable"); + await expect( + database.prepare("DELETE FROM sandbox_backup WHERE id = ?").bind(actualA).run(), + ).rejects.toThrow("record is permanent"); + expect( + (await database.prepare("SELECT id, staging_id FROM sandbox_backup ORDER BY id").all()) + .results, + ).toEqual([ + { id: actualA, staging_id: stageA }, + { id: actualB, staging_id: stageB }, + ]); + }); + + test("ready and pruned runtime records prevent a second stage from claiming their ID", async () => { + const database = createDatabase(); + const stageA = await createOperationStage(database); + const actualId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440025"); + await claimSandboxBackupStageActual(database, { + actualBackupId: actualId, + dir: DIR, + sandboxIncarnation: 1, + stagingId: stageA.id, + }); + await finalizeSandboxBackupStage(database, { + actualBackupId: actualId, + stagingId: stageA.id, + }); + await expect( + database + .prepare( + `INSERT OR REPLACE INTO sandbox_backup_staging ( + actual_backup_id, claim_owner, created_at, dir, driver_generation, + driver_instance_id, id, operation_id, sandbox_id, sandbox_incarnation, + session_run_id, ttl_seconds, updated_at, updates_subject_backup, + workspace_session_id + ) VALUES ( + NULL, 'owner', ?, '/workspace/resurrected', NULL, NULL, ?, ?, ?, 1, + NULL, 100, ?, 0, NULL + )`, + ) + .bind(NOW, stageA.id, OPERATION_ID, SANDBOX_ID, NOW) + .run(), + ).rejects.toThrow("already owned"); + const [write] = await stageSandboxBackupWrites(database, { + admission: { + kind: "operation", + lease: { + claimExpiresAt: NOW + 60_000, + claimOwner: "owner", + incarnation: 1, + kind: "hibernate", + operationId: OPERATION_ID, + status: "backing_up", + }, + }, + sandboxId: SANDBOX_ID, + targets: [ + { + dir: "/workspace/other", + updateSandboxLastBackup: false, + workspaceSessionId: null, + }, + ], + ttlSeconds: 100, + }); + if (write?.kind !== "staged") { + throw new Error("Second runtime backup stage was not created."); + } + + await expect( + claimSandboxBackupStageActual(database, { + actualBackupId: actualId, + dir: write.stage.dir, + sandboxIncarnation: 1, + stagingId: write.stage.id, + }), + ).resolves.toBeNull(); + await database + .prepare("UPDATE sandbox_backup SET status = 'pruned' WHERE id = ?") + .bind(actualId) + .run(); + await expect( + database + .prepare("UPDATE sandbox_backup SET status = 'ready' WHERE id = ?") + .bind(actualId) + .run(), + ).rejects.toThrow("cannot become ready"); + await expect( + claimSandboxBackupStageActual(database, { + actualBackupId: actualId, + dir: write.stage.dir, + sandboxIncarnation: 1, + stagingId: write.stage.id, + }), + ).resolves.toBeNull(); + await expect( + database + .prepare("UPDATE OR REPLACE sandbox_backup_staging SET actual_backup_id = ? WHERE id = ?") + .bind(actualId, write.stage.id) + .run(), + ).rejects.toThrow("already owned"); + await expect( + database + .prepare( + `INSERT OR REPLACE INTO sandbox_backup_staging ( + actual_backup_id, claim_owner, created_at, dir, driver_generation, + driver_instance_id, id, operation_id, sandbox_id, sandbox_incarnation, + session_run_id, ttl_seconds, updated_at, updates_subject_backup, + workspace_session_id + ) VALUES ( + ?, 'owner', ?, '/workspace/third', NULL, NULL, ?, ?, ?, 1, + NULL, 100, ?, 0, NULL + )`, + ) + .bind( + actualId, + NOW, + encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440026"), + OPERATION_ID, + SANDBOX_ID, + NOW, + ) + .run(), + ).rejects.toThrow("already owned"); + expect( + await authorizeSandboxBackupDeletion(database, { + authority: { kind: "pruned" }, + backupId: actualId, + }), + ).toBe(true); + expect(await getSandboxBackupStage(database, write.stage.id)).toMatchObject({ + actualBackupId: null, + }); + }); + + test("a crash after delete intent is retried to physical completion", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const actualId = bucket.putBackup({ + platformId: "550e8400-e29b-41d4-a716-44665544000b", + stagingId: parsePlatformId("01J0000000000000000000000C"), + }); + expect( + await authorizeSandboxBackupDeletion(database, { + authority: { kind: "unattributed" }, + backupId: actualId, + }), + ).toBe(true); + expect( + await database + .prepare("SELECT deleted_at FROM sandbox_backup_delete_intent WHERE backup_id = ?") + .bind(actualId) + .first(), + ).toEqual({ deleted_at: null }); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + + expect(getSandboxBackupObjectKeys(actualId).every((key) => !bucket.objects.has(key))).toBe( + true, + ); + expect( + await database + .prepare("SELECT deleted_at FROM sandbox_backup_delete_intent WHERE backup_id = ?") + .bind(actualId) + .first<{ deleted_at: number | null }>(), + ).toMatchObject({ deleted_at: expect.any(Number) }); + await expect( + database + .prepare("DELETE FROM sandbox_backup_delete_intent WHERE backup_id = ?") + .bind(actualId) + .run(), + ).rejects.toThrow("deletion intent is permanent"); + await database.prepare("PRAGMA recursive_triggers = OFF").run(); + await expect( + database + .prepare( + `INSERT OR REPLACE INTO sandbox_backup_delete_intent ( + attempted_at, backup_id, created_at, delete_after, deleted_at + ) VALUES ( + NULL, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), + CAST(unixepoch('subsec') * 1000 AS INTEGER), NULL + )`, + ) + .bind(actualId) + .run(), + ).rejects.toThrow("deletion lacks D1 authority"); + bucket.putBackup({ + platformId: "550e8400-e29b-41d4-a716-44665544000b", + stagingId: parsePlatformId("01J0000000000000000000000C"), + }); + expect( + await authorizeSandboxBackupDeletion(database, { + authority: { kind: "unattributed" }, + backupId: actualId, + }), + ).toBe(true); + await deleteAuthorizedSandboxBackupObjects(createBindings(database, bucket), [actualId]); + expect(getSandboxBackupObjectKeys(actualId).every((key) => !bucket.objects.has(key))).toBe( + true, + ); + }); + + test.each([ + ["data", "550e8400-e29b-41d4-a716-446655440019"], + ["metadata", "550e8400-e29b-41d4-a716-44665544001c"], + ] as const)( + "reconciliation resumes after deleting %s when the R2 ACK is lost", + async (kind, id) => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const actualId = bucket.putBackup({ + platformId: id, + stagingId: parsePlatformId("01J0000000000000000000000C"), + uploaded: new Date(), + }); + await authorizeSandboxBackupDeletion(database, { + authority: { kind: "unattributed" }, + backupId: actualId, + }); + const [dataKey, metadataKey] = getSandboxBackupObjectKeys(actualId); + bucket.failAfterDeletingKey = kind === "data" ? dataKey : metadataKey; + + await reconcileSandboxBackupPage(createBindings(database, bucket), { cursor: null }); + + expect(bucket.objects.has(dataKey)).toBe(false); + expect(bucket.objects.has(metadataKey)).toBe(kind === "data"); + expect( + await database + .prepare( + `SELECT attempted_at, deleted_at + FROM sandbox_backup_delete_intent WHERE backup_id = ?`, + ) + .bind(actualId) + .first(), + ).toEqual({ attempted_at: expect.any(Number), deleted_at: null }); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { cursor: null }); + + expect(bucket.objects.has(metadataKey)).toBe(false); + expect( + await database + .prepare("SELECT deleted_at FROM sandbox_backup_delete_intent WHERE backup_id = ?") + .bind(actualId) + .first(), + ).toEqual({ deleted_at: expect.any(Number) }); + }, + ); + + test("a lost deletion completion acknowledgement remains idempotently complete", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const actualId = bucket.putBackup({ + platformId: "550e8400-e29b-41d4-a716-44665544001a", + stagingId: parsePlatformId("01J0000000000000000000000C"), + }); + await authorizeSandboxBackupDeletion(database, { + authority: { kind: "unattributed" }, + backupId: actualId, + }); + const fault = loseRunAcknowledgementOnce(database, "SET deleted_at = coalesce"); + + await reconcileSandboxBackupPage(createBindings(fault.database, bucket), { cursor: null }); + + expect(fault.wasLost()).toBe(true); + expect(getSandboxBackupObjectKeys(actualId).every((key) => !bucket.objects.has(key))).toBe( + true, + ); + expect( + await database + .prepare("SELECT deleted_at FROM sandbox_backup_delete_intent WHERE backup_id = ?") + .bind(actualId) + .first(), + ).toEqual({ deleted_at: expect.any(Number) }); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { cursor: null }); + expect(await listPendingSandboxBackupDeletions(database, 1)).toEqual([]); + }); + + test("a deletion cannot be completed before its first physical attempt", async () => { + const database = createDatabase(); + const backupId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-44665544001b"); + await authorizeSandboxBackupDeletion(database, { + authority: { kind: "unattributed" }, + backupId, + }); + + await expect( + database + .prepare( + `UPDATE sandbox_backup_delete_intent + SET deleted_at = CAST(unixepoch('subsec') * 1000 AS INTEGER) + WHERE backup_id = ?`, + ) + .bind(backupId) + .run(), + ).rejects.toThrow("completion must use D1 time and is irreversible"); + }); + + test("an unattempted deletion wins an otherwise exact retry ordering tie", async () => { + const database = createDatabase(); + const attempted = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-44665544001d"); + const unattempted = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-44665544001e"); + database.execute("DROP TRIGGER sandbox_backup_delete_intent_authority"); + await database + .prepare( + `INSERT INTO sandbox_backup_delete_intent ( + attempted_at, backup_id, created_at, delete_after, deleted_at + ) VALUES (1, ?, 1, 1, NULL), (NULL, ?, 1, 1, NULL)`, + ) + .bind(attempted, unattempted) + .run(); + + expect(await listPendingSandboxBackupDeletions(database, 2)).toEqual([unattempted, attempted]); + }); + + test("cleanup admitted between SDK creation and D1 claim leaves no ready row", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const stage = await createOperationStage(database, { workspace: true }); + const actualId = bucket.putBackup({ + platformId: "550e8400-e29b-41d4-a716-446655440004", + stagingId: stage.id, + }); + await database + .prepare( + `UPDATE session + SET archived_at = ?, cleanup_operation_kind = 'delete', + status = 'RESCHEDULING', status_operation_id = ? + WHERE id = ?`, + ) + .bind(NOW, CLEANUP_OPERATION_ID, SESSION_ID) + .run(); + + const claim = await claimSandboxBackupStageActual(database, { + actualBackupId: actualId, + dir: DIR, + sandboxIncarnation: 1, + stagingId: stage.id, + }); + expect(claim).toBeNull(); + expect( + await revokeSandboxBackupsForSessionDelete(database, { + cwd: DIR, + operationId: CLEANUP_OPERATION_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }), + ).toEqual([]); + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + + expect(await readyRows(database)).toEqual([]); + expect(await getSandboxBackupStage(database, stage.id)).toBeNull(); + expect(getSandboxBackupObjectKeys(actualId).every((key) => !bucket.objects.has(key))).toBe( + true, + ); + }); + + for (const mismatch of ["dir", "id", "scope"] as const) { + test(`scanner rejects ${mismatch} metadata that does not match its stage`, async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const stage = await createOperationStage(database); + const platformId = "550e8400-e29b-41d4-a716-446655440005"; + const actualId = bucket.putBackup({ + ...(mismatch === "dir" ? { dir: "/workspace/wrong" } : {}), + ...(mismatch === "id" ? { metadataId: "550e8400-e29b-41d4-a716-446655440006" } : {}), + ...(mismatch === "scope" + ? { + name: createRuntimeSandboxBackupName( + parsePlatformId("01J0000000000000000000000C"), + ), + } + : {}), + platformId, + stagingId: stage.id, + }); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + + expect(await readyRows(database)).toEqual([]); + expect((await getSandboxBackupStage(database, stage.id))?.actualBackupId).toBeNull(); + expect(getSandboxBackupObjectKeys(actualId).every((key) => !bucket.objects.has(key))).toBe( + true, + ); + }); + } + + test("reconciliation remains bounded and resumes with an opaque page cursor", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + for (let index = 0; index < 65; index += 1) { + const backupId = encodeSandboxBackupIdForStorage( + `550e8400-e29b-4000-8000-${index.toString(16).padStart(12, "0")}`, + ); + const [dataKey] = getSandboxBackupObjectKeys(backupId); + await bucket.put(dataKey, "orphan"); + } + + const first = await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + const second = await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: first.nextCursor, + }); + + expect(first).toEqual({ hasMore: true, nextCursor: "opaque-1", processed: 64 }); + expect(second).toEqual({ hasMore: false, nextCursor: null, processed: 1 }); + expect(bucket.objects.size).toBe(0); + }); + + test("database-backed orphan grace retains fresh uploads", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const backupId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-44665544000c"); + const [dataKey] = getSandboxBackupObjectKeys(backupId); + await bucket.put(dataKey, "in-flight", new Date()); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + + expect(bucket.objects.has(dataKey)).toBe(true); + }); +}); + +describe("environment artifact backup lifecycle", () => { + test("a pruned runtime backup ID cannot be claimed before its tombstone exists", async () => { + const database = createDatabase(); + const runtimeStage = await createOperationStage(database); + const { authority } = await createEnvironmentArtifactStage(database); + const sharedId = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440040"); + + expect( + await claimSandboxBackupStageActual(database, { + actualBackupId: sharedId, + dir: DIR, + sandboxIncarnation: 1, + stagingId: runtimeStage.id, + }), + ).toEqual({ actualBackupId: sharedId }); + await finalizeSandboxBackupStage(database, { + actualBackupId: sharedId, + stagingId: runtimeStage.id, + }); + await database + .prepare("UPDATE sandbox_backup SET status = 'pruned' WHERE id = ?") + .bind(sharedId) + .run(); + + await expect( + claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: sharedId, + authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }), + ).resolves.toBeNull(); + await expect( + database + .prepare( + `UPDATE OR REPLACE environment_package_artifact_backup_staging + SET actual_backup_id = ? WHERE command_id = ?`, + ) + .bind(sharedId, ENVIRONMENT_COMMAND_ID) + .run(), + ).rejects.toThrow("already owned"); + expect( + await database + .prepare("SELECT id, status FROM sandbox_backup WHERE id = ?") + .bind(sharedId) + .first(), + ).toEqual({ id: sharedId, status: "pruned" }); + expect( + await getEnvironmentPackageArtifactBackupStage(database, ENVIRONMENT_COMMAND_ID), + ).toMatchObject({ actualBackupId: null }); + }); + + test("an environment artifact backup ID cannot be claimed by runtime", async () => { + const database = createDatabase(); + const runtimeStage = await createOperationStage(database); + const { authority } = await createEnvironmentArtifactStage(database); + const sharedId = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440041"); + + expect( + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: sharedId, + authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }), + ).toEqual({ actualBackupId: sharedId }); + expect( + await commitEnvironmentPackageArtifactBackup(database, { + actualBackupId: sharedId, + ...authority, + expiresAt: Date.now() + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS * 1_000 - 1_000, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }), + ).toBe(true); + + await expect( + claimSandboxBackupStageActual(database, { + actualBackupId: sharedId, + dir: DIR, + sandboxIncarnation: 1, + stagingId: runtimeStage.id, + }), + ).resolves.toBeNull(); + await expect( + database + .prepare( + `UPDATE OR REPLACE sandbox_backup_staging + SET actual_backup_id = ? WHERE id = ?`, + ) + .bind(sharedId, runtimeStage.id) + .run(), + ).rejects.toThrow("already owned"); + expect(await getSandboxBackupStage(database, runtimeStage.id)).toMatchObject({ + actualBackupId: null, + }); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: sharedId }); + await database + .prepare("DELETE FROM environment_package_artifact_backup_staging WHERE command_id = ?") + .bind(ENVIRONMENT_COMMAND_ID) + .run(); + const successor = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + await expect( + claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: sharedId, + authority: successor.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }), + ).resolves.toBeNull(); + await expect( + database + .prepare( + `UPDATE OR REPLACE environment_package_artifact_backup_staging + SET actual_backup_id = ? WHERE command_id = ?`, + ) + .bind(sharedId, ENVIRONMENT_COMMAND_ID) + .run(), + ).rejects.toThrow("already owned"); + }); + + test("bounded reconciliation removes a terminal stage before any R2 object exists", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + await createEnvironmentArtifactStage(database); + database.execute( + `UPDATE api_command SET status = 'failed' WHERE id = '${ENVIRONMENT_COMMAND_ID}'`, + ); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + + expect( + await getEnvironmentPackageArtifactBackupStage(database, ENVIRONMENT_COMMAND_ID), + ).toBeNull(); + }); + + test("an empty R2 page continues until every terminal stage is revoked", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + for (let index = 0; index < 65; index += 1) { + const commandId = createPlatformId(); + const inputDigest = index.toString(16).padStart(64, "0"); + const dir = `/workspace/.mosoo/environment-artifacts/${inputDigest}`; + await database + .prepare( + `INSERT INTO api_command ( + attempt_count, claim_expires_at, claim_owner, created_at, dedupe_key, + delivery_generation, id, kind, payload_json, status, updated_at + ) VALUES (1, ?, 'owner', ?, ?, 1, ?, 'environment_package_artifact_build', ?, + 'running', ?)`, + ) + .bind( + Date.now() + 60_000, + NOW, + `terminal-environment-stage:${commandId}`, + commandId, + JSON.stringify({ appId: ENVIRONMENT_APP_ID, inputDigest }), + NOW, + ) + .run(); + await database + .prepare( + `INSERT INTO environment_package_artifact_backup_staging ( + actual_backup_id, app_id, attempt_count, claim_owner, command_id, created_at, + delivery_generation, dir, input_digest, paths_json, updated_at + ) VALUES (NULL, ?, 1, 'owner', ?, ?, 1, ?, ?, ?, ?)`, + ) + .bind( + ENVIRONMENT_APP_ID, + commandId, + NOW, + dir, + inputDigest, + JSON.stringify({ executable: [`${dir}/bin/tool`], node: [], python: [] }), + NOW, + ) + .run(); + } + await database + .prepare( + "UPDATE api_command SET status = 'failed' WHERE dedupe_key LIKE 'terminal-environment-stage:%'", + ) + .run(); + + const first = await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + expect(first).toEqual({ hasMore: true, nextCursor: null, processed: 0 }); + await expect( + database + .prepare("SELECT count(*) AS count FROM environment_package_artifact_backup_staging") + .first(), + ).resolves.toEqual({ count: 1 }); + + const second = await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + expect(second).toEqual({ hasMore: false, nextCursor: null, processed: 0 }); + await expect( + database + .prepare("SELECT count(*) AS count FROM environment_package_artifact_backup_staging") + .first(), + ).resolves.toEqual({ count: 0 }); + }); + + test("scanner commits a complete backup to D1 without writing a projection", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-446655440001", + }); + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + expect( + await getEnvironmentPackageArtifactBackupStage(database, ENVIRONMENT_COMMAND_ID), + ).toBeNull(); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualId, manifestGeneration: 1, paths: ENVIRONMENT_PATHS }); + expect(bucket.objects.has(environmentPackageArtifactMetadataKey(ENVIRONMENT_KEY))).toBe(false); + expect(getSandboxBackupObjectKeys(actualId).every((key) => bucket.objects.has(key))).toBe(true); + }); + + test("a D1 manifest resolves without an R2 projection", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-446655440009", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualId, + authority, + commandId: authority.commandId, + dir: ENVIRONMENT_DIR, + }); + expect( + await commitEnvironmentPackageArtifactBackup(database, { + actualBackupId: actualId, + ...authority, + expiresAt: bucket.backupExpiresAt(actualId), + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }), + ).toBe(true); + expect(bucket.objects.has(environmentPackageArtifactMetadataKey(ENVIRONMENT_KEY))).toBe(false); + const manifest = await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY); + expect(manifest?.expiresAt).toBe(bucket.backupExpiresAt(actualId)); + expect(manifest?.expiresAt).not.toBe( + OLD_UPLOAD.getTime() + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS * 1_000, + ); + + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toMatchObject({ + backupId: "650e8400-e29b-41d4-a716-446655440009", + paths: ENVIRONMENT_PATHS, + }); + }); + + test.each([ + ["data", 0], + ["metadata", 1], + ] as const)( + "a current manifest resolves null when its %s object is missing", + async (_kind, keyIndex) => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-446655440019", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualId, + authority, + commandId: authority.commandId, + dir: ENVIRONMENT_DIR, + }); + await commitEnvironmentPackageArtifactBackup(database, { + actualBackupId: actualId, + ...authority, + expiresAt: bucket.backupExpiresAt(actualId), + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + await bucket.delete(getSandboxBackupObjectKeys(actualId)[keyIndex]); + + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toBeNull(); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualId, manifestGeneration: 1 }); + }, + ); + + test("a near-expiry current manifest resolves null without reading R2", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualId = bucket.putBackup({ + createdAt: new Date( + Date.now() - + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS * 1_000 + + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_REFRESH_WINDOW_MS / 2, + ), + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-44665544001a", + }); + const expiresAt = bucket.backupExpiresAt(actualId); + database.execute("DROP TRIGGER environment_package_artifact_backup_authority"); + await database + .prepare( + `INSERT INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES (?, 1, ?, ?, ?, 1, ?, ?, 1, ?)`, + ) + .bind( + ENVIRONMENT_APP_ID, + actualId, + ENVIRONMENT_COMMAND_ID, + expiresAt - 25 * 60 * 60_000, + expiresAt, + ENVIRONMENT_KEY.inputDigest, + JSON.stringify(ENVIRONMENT_PATHS), + ) + .run(); + + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toBeNull(); + expect(bucket.gets).not.toContain(getSandboxBackupObjectKeys(actualId)[1]); + }); + + test("default resolution automatically requeues a succeeded near-expiry artifact", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const buildQueue = createApiCommandQueueStub(); + const bindings = { + ...createBindings(database, bucket), + API_COMMAND_QUEUE: createApiCommandQueueStub(), + ENVIRONMENT_ARTIFACT_BUILD_QUEUE: buildQueue, + } as ApiBindings; + const packages = [{ manager: "pip", packages: ["requests==2.32.4"] }] as const; + const initial = await resolveEnvironmentPackageArtifact(bindings, ENVIRONMENT_APP_ID, packages); + if (initial === null) { + throw new Error("Environment package artifact resolution returned no key."); + } + const command = await database + .prepare( + `SELECT id FROM api_command + WHERE kind = 'environment_package_artifact_build'`, + ) + .first<{ id: ApiCommandId }>(); + if (command === null) { + throw new Error("Environment package artifact command was not enqueued."); + } + await database + .prepare( + `UPDATE api_command + SET attempt_count = 1, claim_expires_at = NULL, claim_owner = NULL, + completed_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), status = 'succeeded' + WHERE id = ?`, + ) + .bind(command.id) + .run(); + const dir = environmentPackageArtifactDir(initial.key); + const paths: EnvironmentPackageArtifactPaths = { + executable: [`${dir}/python/bin`], + node: [`${dir}/npm/node_modules`], + python: [`${dir}/python/site-packages`], + }; + const backupId = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-44665544002b"); + await database + .prepare( + `INSERT INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES ( + ?, 1, ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), 1, + CAST(unixepoch('subsec') * 1000 AS INTEGER) + 315359999000, ?, 1, ? + )`, + ) + .bind(initial.key.appId, backupId, command.id, initial.key.inputDigest, JSON.stringify(paths)) + .run(); + await expireEnvironmentArtifactManifest( + database, + Date.now() + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_REFRESH_WINDOW_MS / 2, + initial.key, + ); + + await expect( + resolveEnvironmentPackageArtifact(bindings, ENVIRONMENT_APP_ID, packages), + ).resolves.toMatchObject({ key: initial.key, metadata: null }); + expect(buildQueue.sent).toHaveLength(2); + expect( + await database + .prepare("SELECT delivery_generation, status FROM api_command WHERE id = ?") + .bind(command.id) + .first(), + ).toEqual({ delivery_generation: 2, status: "queued" }); + }); + + test("a delete intent that wins D1 prevents a late environment projection", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-44665544000a", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualId, + authority, + commandId: authority.commandId, + dir: ENVIRONMENT_DIR, + }); + database.execute("DROP TRIGGER sandbox_backup_delete_intent_authority"); + await database + .prepare( + `INSERT INTO sandbox_backup_delete_intent ( + attempted_at, backup_id, created_at, delete_after, deleted_at + ) VALUES ( + NULL, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), + CAST(unixepoch('subsec') * 1000 AS INTEGER), NULL + )`, + ) + .bind(actualId) + .run(); + database.execute("DROP TRIGGER environment_package_artifact_backup_authority"); + + await expect( + database + .prepare( + `INSERT INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES ( + ?, 1, ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), 1, + ?, ?, 1, ? + )`, + ) + .bind( + ENVIRONMENT_APP_ID, + actualId, + ENVIRONMENT_COMMAND_ID, + bucket.backupExpiresAt(actualId), + ENVIRONMENT_KEY.inputDigest, + JSON.stringify(ENVIRONMENT_PATHS), + ) + .run(), + ).rejects.toThrow("sandbox backup object is tombstoned or already referenced"); + expect(await getEnvironmentPackageArtifactBackupStage(database, authority.commandId)).toEqual( + expect.objectContaining({ actualBackupId: actualId }), + ); + await deleteAuthorizedSandboxBackupObjects(createBindings(database, bucket), [actualId]); + expect(bucket.objects.has(environmentPackageArtifactMetadataKey(ENVIRONMENT_KEY))).toBe(false); + expect(getSandboxBackupObjectKeys(actualId).every((key) => !bucket.objects.has(key))).toBe( + true, + ); + }); + + test("a verified named R2 projection is atomically adopted into D1", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const attemptA = await createEnvironmentArtifactStage(database); + const actualId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptA.authority), + platformId: "650e8400-e29b-41d4-a716-44665544000b", + }); + const legacyKey = environmentPackageArtifactMetadataKey(ENVIRONMENT_KEY); + await bucket.put( + legacyKey, + JSON.stringify({ + backupId: "650e8400-e29b-41d4-a716-44665544000b", + paths: ENVIRONMENT_PATHS, + }), + ); + await succeedEnvironmentArtifactCommand(database); + + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toMatchObject({ paths: ENVIRONMENT_PATHS }); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualId, paths: ENVIRONMENT_PATHS }); + await bucket.delete(legacyKey); + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toMatchObject({ paths: ENVIRONMENT_PATHS }); + expect(bucket.gets.filter((key) => key === legacyKey)).toHaveLength(1); + expect( + await authorizeSandboxBackupDeletion(database, { + authority: { kind: "unattributed" }, + backupId: actualId, + }), + ).toBe(false); + }); + + test("a real legacy nameless R2 projection is atomically adopted into D1", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + await createEnvironmentArtifactStage(database); + const actualId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: null, + platformId: "650e8400-e29b-41d4-a716-446655440010", + }); + await bucket.put( + environmentPackageArtifactMetadataKey(ENVIRONMENT_KEY), + JSON.stringify({ + backupId: "650e8400-e29b-41d4-a716-446655440010", + paths: ENVIRONMENT_PATHS, + }), + ); + await succeedEnvironmentArtifactCommand(database); + + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toMatchObject({ paths: ENVIRONMENT_PATHS }); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ + attemptCount: 1, + backupId: actualId, + commandId: ENVIRONMENT_COMMAND_ID, + deliveryGeneration: 1, + paths: ENVIRONMENT_PATHS, + }); + await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toMatchObject({ paths: ENVIRONMENT_PATHS }); + }); + + test("a D1 manifest ignores a stale legacy R2 projection", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-44665544000c", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualId, + authority, + commandId: authority.commandId, + dir: ENVIRONMENT_DIR, + }); + await commitEnvironmentPackageArtifactBackup(database, { + actualBackupId: actualId, + ...authority, + expiresAt: bucket.backupExpiresAt(actualId), + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + await bucket.put( + environmentPackageArtifactMetadataKey(ENVIRONMENT_KEY), + JSON.stringify({ + backupId: "650e8400-e29b-41d4-a716-44665544000d", + paths: ENVIRONMENT_PATHS, + }), + ); + + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toEqual({ + backupId: "650e8400-e29b-41d4-a716-44665544000c", + paths: ENVIRONMENT_PATHS, + }); + expect(getSandboxBackupObjectKeys(actualId).every((key) => bucket.objects.has(key))).toBe(true); + }); + + test("a D1 manifest rejects mismatched object authority metadata", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-44665544000e", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualId, + authority, + commandId: authority.commandId, + dir: ENVIRONMENT_DIR, + }); + await commitEnvironmentPackageArtifactBackup(database, { + actualBackupId: actualId, + ...authority, + expiresAt: bucket.backupExpiresAt(actualId), + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + await bucket.put( + environmentPackageArtifactMetadataKey(ENVIRONMENT_KEY), + JSON.stringify({ + backupId: "650e8400-e29b-41d4-a716-44665544000e", + paths: ENVIRONMENT_PATHS, + }), + ); + const [, backupMetadataKey] = getSandboxBackupObjectKeys(actualId); + const backupMetadata = JSON.parse(bucket.objects.get(backupMetadataKey)!.body) as Record< + string, + unknown + >; + for (const tamperedAuthority of [ + { ...authority, commandId: "01J0000000000000000000000L" }, + { ...authority, deliveryGeneration: authority.deliveryGeneration + 1 }, + { ...authority, attemptCount: authority.attemptCount + 1 }, + ]) { + await bucket.put( + backupMetadataKey, + JSON.stringify({ + ...backupMetadata, + name: createEnvironmentPackageArtifactBackupName(tamperedAuthority), + }), + ); + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toBeNull(); + } + for (const tamperedMetadata of [ + { ...backupMetadata, ttl: ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS - 1 }, + { ...backupMetadata, createdAt: "2026-01-01" }, + ]) { + await bucket.put(backupMetadataKey, JSON.stringify(tamperedMetadata)); + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toBeNull(); + } + }); + + test("a current manifest rejects direct, replacement, and malformed writes", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-446655440014", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualId, + authority, + commandId: authority.commandId, + dir: ENVIRONMENT_DIR, + }); + await commitEnvironmentPackageArtifactBackup(database, { + actualBackupId: actualId, + ...authority, + expiresAt: bucket.backupExpiresAt(actualId), + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + await database.prepare("PRAGMA recursive_triggers = OFF").run(); + await expect( + database + .prepare( + `INSERT OR REPLACE INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES (?, ?, ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), ?, ?, ?, 1, ?)`, + ) + .bind( + ENVIRONMENT_APP_ID, + authority.attemptCount, + actualId, + authority.commandId, + authority.deliveryGeneration, + bucket.backupExpiresAt(actualId), + ENVIRONMENT_KEY.inputDigest, + JSON.stringify(ENVIRONMENT_PATHS), + ) + .run(), + ).rejects.toThrow("backup lacks D1 authority"); + + const malformedDatabase = createDatabase(); + const malformedStage = await createEnvironmentArtifactStage(malformedDatabase); + const malformedId = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440015"); + await claimEnvironmentPackageArtifactBackupActual(malformedDatabase, { + actualBackupId: malformedId, + authority: malformedStage.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + const duplicatePaths = `{"executable":[],"executable":[],"node":[],"python":[]}`; + malformedDatabase.execute("DROP TRIGGER environment_package_artifact_backup_staging_immutable"); + await malformedDatabase + .prepare( + "UPDATE environment_package_artifact_backup_staging SET paths_json = ? WHERE command_id = ?", + ) + .bind(duplicatePaths, ENVIRONMENT_COMMAND_ID) + .run(); + await expect( + malformedDatabase + .prepare( + `INSERT INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES (?, 1, ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), 1, + CAST(unixepoch('subsec') * 1000 AS INTEGER) + 315360000000, ?, 1, ?)`, + ) + .bind( + ENVIRONMENT_APP_ID, + malformedId, + ENVIRONMENT_COMMAND_ID, + ENVIRONMENT_KEY.inputDigest, + duplicatePaths, + ) + .run(), + ).rejects.toThrow("backup lacks D1 authority"); + + const invalidClockDatabase = createDatabase(); + const invalidClockStage = await createEnvironmentArtifactStage(invalidClockDatabase); + const invalidClockId = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440016"); + await claimEnvironmentPackageArtifactBackupActual(invalidClockDatabase, { + actualBackupId: invalidClockId, + authority: invalidClockStage.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await expect( + invalidClockDatabase + .prepare( + `INSERT INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES (?, 1, ?, ?, 0, 1, 315360000000, ?, 1, ?)`, + ) + .bind( + ENVIRONMENT_APP_ID, + invalidClockId, + ENVIRONMENT_COMMAND_ID, + ENVIRONMENT_KEY.inputDigest, + JSON.stringify(ENVIRONMENT_PATHS), + ) + .run(), + ).rejects.toThrow("backup lacks D1 authority"); + }); + + test("a raw manifest insert accepts an exact live staged command", async () => { + const database = createDatabase(); + const { authority } = await createEnvironmentArtifactStage(database); + const backupId = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440017"); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: backupId, + authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + + await insertRawEnvironmentArtifactManifest(database, backupId); + + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId }); + }); + + test.each(["queued", "failed", "dead_lettered"] as const)( + "an exact live stage cannot borrow a %s command", + async (status) => { + const database = createDatabase(); + const { authority } = await createEnvironmentArtifactStage(database); + const backupId = encodeSandboxBackupIdForStorage( + `650e8400-e29b-4000-8000-${status.length.toString(16).padStart(12, "0")}`, + ); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: backupId, + authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await database + .prepare("UPDATE api_command SET status = ? WHERE id = ?") + .bind(status, ENVIRONMENT_COMMAND_ID) + .run(); + + await expect(insertRawEnvironmentArtifactManifest(database, backupId)).rejects.toThrow( + "backup lacks D1 authority", + ); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toBeNull(); + }, + ); + + test("an exact staged running command requires a live lease", async () => { + const database = createDatabase(); + const { authority } = await createEnvironmentArtifactStage(database); + const backupId = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440018"); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: backupId, + authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await database + .prepare("UPDATE api_command SET claim_expires_at = 0 WHERE id = ?") + .bind(ENVIRONMENT_COMMAND_ID) + .run(); + + await expect(insertRawEnvironmentArtifactManifest(database, backupId)).rejects.toThrow( + "backup lacks D1 authority", + ); + }); + + test("a succeeded command can authorize an unstaged legacy manifest", async () => { + const database = createDatabase(); + await setEnvironmentArtifactCommand(database, { + attemptCount: 1, + claimOwner: "environment-owner", + deliveryGeneration: 1, + }); + await succeedEnvironmentArtifactCommand(database); + const backupId = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-44665544001f"); + + await insertRawEnvironmentArtifactManifest(database, backupId); + + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId }); + }); + + test.each([ + ["terminal status", "UPDATE api_command SET status = 'failed' WHERE id = ?"], + ["completion clock", "UPDATE api_command SET completed_at = NULL WHERE id = ?"], + ["claim owner", "UPDATE api_command SET claim_owner = 'stale-owner' WHERE id = ?"], + ["claim lease", "UPDATE api_command SET claim_expires_at = 9007199254740991 WHERE id = ?"], + ] as const)( + "an unstaged legacy manifest requires an exact succeeded command %s", + async (_, sql) => { + const database = createDatabase(); + await setEnvironmentArtifactCommand(database, { + attemptCount: 1, + claimOwner: "environment-owner", + deliveryGeneration: 1, + }); + await succeedEnvironmentArtifactCommand(database); + await database.prepare(sql).bind(ENVIRONMENT_COMMAND_ID).run(); + const backupId = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440047"); + + await expect(insertRawEnvironmentArtifactManifest(database, backupId)).rejects.toThrow( + "backup lacks D1 authority", + ); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toBeNull(); + }, + ); + + test("OR REPLACE cannot move a live environment candidate to another key", async () => { + const database = createDatabase(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualId = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440042"); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualId, + authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + + const appB = parsePlatformId("01J0000000000000000000000M"); + const commandB = parsePlatformId("01J0000000000000000000000N"); + const keyB: EnvironmentPackageArtifactKey = { appId: appB, inputDigest: "b".repeat(64) }; + const dirB = environmentPackageArtifactDir(keyB); + const pathsB: EnvironmentPackageArtifactPaths = { + executable: [`${dirB}/python/bin`], + node: [`${dirB}/npm/node_modules`], + python: [`${dirB}/python/site-packages`], + }; + await database + .prepare( + `INSERT INTO api_command ( + attempt_count, claim_expires_at, claim_owner, created_at, dedupe_key, + delivery_generation, id, kind, payload_json, status, updated_at + ) VALUES ( + 1, CAST(unixepoch('subsec') * 1000 AS INTEGER) + 60000, 'owner-b', ?, ?, 1, ?, + 'environment_package_artifact_build', ?, 'running', ? + )`, + ) + .bind( + NOW, + `environment-artifact:${appB}:${keyB.inputDigest}`, + commandB, + JSON.stringify(keyB), + NOW, + ) + .run(); + await database.prepare("PRAGMA recursive_triggers = OFF").run(); + + await expect( + database + .prepare( + `INSERT OR REPLACE INTO environment_package_artifact_backup_staging ( + actual_backup_id, app_id, attempt_count, claim_owner, command_id, created_at, + delivery_generation, dir, input_digest, paths_json, updated_at + ) VALUES ( + ?, ?, 1, 'owner-b', ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), + 1, ?, ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER) + )`, + ) + .bind(actualId, appB, commandB, dirB, keyB.inputDigest, JSON.stringify(pathsB)) + .run(), + ).rejects.toThrow("already owned"); + + await database + .prepare( + `UPDATE api_command + SET claim_expires_at = NULL, claim_owner = NULL, + completed_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), status = 'succeeded' + WHERE id = ?`, + ) + .bind(commandB) + .run(); + await expect( + database + .prepare( + `INSERT OR REPLACE INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES ( + ?, 1, ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), 1, + CAST(unixepoch('subsec') * 1000 AS INTEGER) + 315359999000, ?, 1, ? + )`, + ) + .bind(appB, actualId, commandB, keyB.inputDigest, JSON.stringify(pathsB)) + .run(), + ).rejects.toThrow("backup lacks D1 authority"); + + expect( + await getEnvironmentPackageArtifactBackupStage(database, ENVIRONMENT_COMMAND_ID), + ).toMatchObject({ actualBackupId: actualId }); + expect(await getEnvironmentPackageArtifactBackupManifest(database, keyB)).toBeNull(); + }); + + test("OR REPLACE cannot replace an occupied environment stage key with a new backup", async () => { + const database = createDatabase(); + await createEnvironmentArtifactStage(database); + const commandB = parsePlatformId("01J0000000000000000000000P"); + const actualB = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440048"); + await setEnvironmentArtifactCommand(database, { + attemptCount: 1, + claimOwner: "owner-b", + commandId: commandB, + dedupeKey: `environment-artifact-replacement:${commandB}`, + deliveryGeneration: 1, + key: ENVIRONMENT_KEY, + }); + await database.prepare("PRAGMA recursive_triggers = OFF").run(); + + await expect( + database + .prepare( + `INSERT OR REPLACE INTO environment_package_artifact_backup_staging ( + actual_backup_id, app_id, attempt_count, claim_owner, command_id, created_at, + delivery_generation, dir, input_digest, paths_json, updated_at + ) VALUES ( + ?, ?, 1, 'owner-b', ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), + 1, ?, ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER) + )`, + ) + .bind( + actualB, + ENVIRONMENT_APP_ID, + commandB, + ENVIRONMENT_DIR, + ENVIRONMENT_KEY.inputDigest, + JSON.stringify(ENVIRONMENT_PATHS), + ) + .run(), + ).rejects.toThrow("already owned"); + + expect( + await getEnvironmentPackageArtifactBackupStage(database, ENVIRONMENT_COMMAND_ID), + ).not.toBeNull(); + expect(await getEnvironmentPackageArtifactBackupStage(database, commandB)).toBeNull(); + }); + + test("OR REPLACE cannot replace an occupied environment manifest key with a new backup", async () => { + const database = createDatabase(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualA = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440049"); + const actualB = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-44665544004a"); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualA, + authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await insertRawEnvironmentArtifactManifest(database, actualA); + const commandB = parsePlatformId("01J0000000000000000000000Q"); + await setEnvironmentArtifactCommand(database, { + attemptCount: 1, + claimOwner: "owner-b", + commandId: commandB, + dedupeKey: `environment-artifact-replacement:${commandB}`, + deliveryGeneration: 1, + key: ENVIRONMENT_KEY, + }); + await succeedEnvironmentArtifactCommand(database, commandB); + await database.prepare("PRAGMA recursive_triggers = OFF").run(); + + await expect( + database + .prepare( + `INSERT OR REPLACE INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES ( + ?, 1, ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), 1, + CAST(unixepoch('subsec') * 1000 AS INTEGER) + 315359999000, ?, 1, ? + )`, + ) + .bind( + ENVIRONMENT_APP_ID, + actualB, + commandB, + ENVIRONMENT_KEY.inputDigest, + JSON.stringify(ENVIRONMENT_PATHS), + ) + .run(), + ).rejects.toThrow("backup lacks D1 authority"); + + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualA }); + }); + + test("OR REPLACE cannot move a live runtime candidate between stages", async () => { + const database = createDatabase(); + const stageA = await createOperationStage(database); + const actualA = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440043"); + const actualB = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440044"); + const stageB = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440045"); + const stageC = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440046"); + await claimSandboxBackupStageActual(database, { + actualBackupId: actualA, + dir: DIR, + sandboxIncarnation: 1, + stagingId: stageA.id, + }); + const insertStage = (stagingId: SandboxBackupId, actualBackupId: SandboxBackupId) => + database + .prepare( + `INSERT OR REPLACE INTO sandbox_backup_staging ( + actual_backup_id, claim_owner, created_at, dir, driver_generation, + driver_instance_id, id, operation_id, sandbox_id, sandbox_incarnation, + session_run_id, ttl_seconds, updated_at, updates_subject_backup, + workspace_session_id + ) VALUES ( + ?, 'owner', CAST(unixepoch('subsec') * 1000 AS INTEGER), '/workspace/other', + NULL, NULL, ?, ?, ?, 1, NULL, 100, + CAST(unixepoch('subsec') * 1000 AS INTEGER), 0, NULL + )`, + ) + .bind(actualBackupId, stagingId, OPERATION_ID, SANDBOX_ID) + .run(); + await database.prepare("PRAGMA recursive_triggers = OFF").run(); + + await expect(insertStage(stageC, actualA)).rejects.toThrow("already owned"); + await insertStage(stageB, actualB); + await expect( + database + .prepare("UPDATE OR REPLACE sandbox_backup_staging SET actual_backup_id = ? WHERE id = ?") + .bind(actualA, stageB) + .run(), + ).rejects.toThrow("already owned"); + + expect(await getSandboxBackupStage(database, stageA.id)).toMatchObject({ + actualBackupId: actualA, + }); + expect(await getSandboxBackupStage(database, stageB)).toMatchObject({ + actualBackupId: actualB, + }); + }); + + test("a wrong-dir stage cannot authorize a manifest", async () => { + const database = createDatabase(); + await setEnvironmentArtifactCommand(database, { + attemptCount: 1, + claimOwner: "environment-owner", + deliveryGeneration: 1, + }); + const backupId = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-44665544002c"); + await database + .prepare( + `INSERT INTO environment_package_artifact_backup_staging ( + actual_backup_id, app_id, attempt_count, claim_owner, command_id, created_at, + delivery_generation, dir, input_digest, paths_json, updated_at + ) VALUES ( + ?, ?, 1, 'environment-owner', ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), + 1, '/workspace/wrong', ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER) + )`, + ) + .bind( + backupId, + ENVIRONMENT_APP_ID, + ENVIRONMENT_COMMAND_ID, + ENVIRONMENT_KEY.inputDigest, + JSON.stringify(ENVIRONMENT_PATHS), + ) + .run(); + const expiresAt = Date.now() + ENVIRONMENT_PACKAGE_ARTIFACT_BACKUP_TTL_SECONDS * 1_000 - 1_000; + + await expect( + commitEnvironmentPackageArtifactBackup(database, { + actualBackupId: backupId, + attemptCount: 1, + claimOwner: "environment-owner", + commandId: ENVIRONMENT_COMMAND_ID, + deliveryGeneration: 1, + expiresAt, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }), + ).resolves.toBe(false); + await expect( + database + .prepare( + `INSERT INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES ( + ?, 1, ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), 1, ?, ?, 1, ? + )`, + ) + .bind( + ENVIRONMENT_APP_ID, + backupId, + ENVIRONMENT_COMMAND_ID, + expiresAt, + ENVIRONMENT_KEY.inputDigest, + JSON.stringify(ENVIRONMENT_PATHS), + ) + .run(), + ).rejects.toThrow("backup lacks D1 authority"); + }); + + test("UPDATE OR REPLACE cannot steal another environment key's backup", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualA = bucket.putBackup({ + createdAt: new Date(Date.now() - 60_000), + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-446655440027", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualA, + authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...authority, + backupId: actualA, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + + const appB = parsePlatformId("01J0000000000000000000000M"); + const commandB = parsePlatformId("01J0000000000000000000000N"); + const keyB: EnvironmentPackageArtifactKey = { appId: appB, inputDigest: "b".repeat(64) }; + const dirB = environmentPackageArtifactDir(keyB); + const pathsB: EnvironmentPackageArtifactPaths = { + executable: [`${dirB}/python/bin`], + node: [`${dirB}/npm/node_modules`], + python: [`${dirB}/python/site-packages`], + }; + const actualB = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440028"); + await database + .prepare( + `INSERT INTO api_command ( + attempt_count, claim_expires_at, claim_owner, completed_at, created_at, + dedupe_key, delivery_generation, id, kind, payload_json, status, updated_at + ) VALUES ( + 1, NULL, NULL, CAST(unixepoch('subsec') * 1000 AS INTEGER), ?, + ?, 1, ?, 'environment_package_artifact_build', ?, 'succeeded', ? + )`, + ) + .bind( + NOW, + `environment-artifact:${appB}:${keyB.inputDigest}`, + commandB, + JSON.stringify(keyB), + NOW, + ) + .run(); + await database + .prepare( + `INSERT INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES ( + ?, 1, ?, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), 1, + CAST(unixepoch('subsec') * 1000 AS INTEGER) + 315359999000, ?, 1, ? + )`, + ) + .bind(appB, actualB, commandB, keyB.inputDigest, JSON.stringify(pathsB)) + .run(); + await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "replacement-owner", + deliveryGeneration: 2, + }); + const stageGuard = await database + .prepare( + `SELECT sql FROM sqlite_schema + WHERE type = 'trigger' + AND name = 'sandbox_backup_delete_intent_blocks_environment_stage_update'`, + ) + .first<{ sql: string }>(); + if (stageGuard === null) { + throw new Error("Environment backup stage ownership trigger is missing."); + } + database.execute("DROP TRIGGER sandbox_backup_delete_intent_blocks_environment_stage_update"); + await database + .prepare( + `UPDATE environment_package_artifact_backup_staging + SET actual_backup_id = ? WHERE command_id = ?`, + ) + .bind(actualB, ENVIRONMENT_COMMAND_ID) + .run(); + database.execute(stageGuard.sql); + await database.prepare("PRAGMA recursive_triggers = OFF").run(); + + await expect( + database + .prepare( + `UPDATE OR REPLACE environment_package_artifact_backup + SET attempt_count = 2, backup_id = ?, command_id = ?, + committed_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), + delivery_generation = 2, expires_at = ?, manifest_generation = 2, + paths_json = ? + WHERE app_id = ? AND input_digest = ?`, + ) + .bind( + actualB, + ENVIRONMENT_COMMAND_ID, + bucket.backupExpiresAt(actualA) + 30_000, + JSON.stringify(ENVIRONMENT_PATHS), + ENVIRONMENT_APP_ID, + ENVIRONMENT_KEY.inputDigest, + ) + .run(), + ).rejects.toThrow("rotation lacks D1 authority"); + expect( + ( + await database + .prepare( + `SELECT app_id, backup_id, input_digest + FROM environment_package_artifact_backup ORDER BY app_id`, + ) + .all() + ).results, + ).toEqual([ + { app_id: ENVIRONMENT_APP_ID, backup_id: actualA, input_digest: ENVIRONMENT_KEY.inputDigest }, + { app_id: appB, backup_id: actualB, input_digest: keyB.inputDigest }, + ]); + }); + + test("raw rotation requires a new backup, exact next generation, and longer expiry", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const attemptA = await createEnvironmentArtifactStage(database); + const actualA = bucket.putBackup({ + createdAt: new Date(Date.now() - 60_000), + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptA.authority), + platformId: "650e8400-e29b-41d4-a716-44665544002d", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualA, + authority: attemptA.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptA.authority, + backupId: actualA, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + const attemptB = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + const actualB = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptB.authority), + platformId: "650e8400-e29b-41d4-a716-44665544002e", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualB, + authority: attemptB.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + + for (const [backupId, manifestGeneration, expiresAt] of [ + [actualB, 3, bucket.backupExpiresAt(actualB)], + [actualB, 2, bucket.backupExpiresAt(actualA)], + ] as const) { + await expect( + database + .prepare( + `UPDATE environment_package_artifact_backup + SET attempt_count = 2, backup_id = ?, command_id = ?, + committed_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), + delivery_generation = 2, expires_at = ?, manifest_generation = ?, + paths_json = ? + WHERE app_id = ? AND input_digest = ?`, + ) + .bind( + backupId, + ENVIRONMENT_COMMAND_ID, + expiresAt, + manifestGeneration, + JSON.stringify(ENVIRONMENT_PATHS), + ENVIRONMENT_APP_ID, + ENVIRONMENT_KEY.inputDigest, + ) + .run(), + ).rejects.toThrow("rotation lacks D1 authority"); + } + database.execute("DROP TRIGGER sandbox_backup_delete_intent_blocks_environment_stage_update"); + await database + .prepare( + `UPDATE environment_package_artifact_backup_staging + SET actual_backup_id = ? WHERE command_id = ?`, + ) + .bind(actualA, ENVIRONMENT_COMMAND_ID) + .run(); + await expect( + database + .prepare( + `UPDATE environment_package_artifact_backup + SET attempt_count = 2, backup_id = ?, command_id = ?, + committed_at = CAST(unixepoch('subsec') * 1000 AS INTEGER), + delivery_generation = 2, expires_at = ?, manifest_generation = 2, + paths_json = ? + WHERE app_id = ? AND input_digest = ?`, + ) + .bind( + actualA, + ENVIRONMENT_COMMAND_ID, + bucket.backupExpiresAt(actualB), + JSON.stringify(ENVIRONMENT_PATHS), + ENVIRONMENT_APP_ID, + ENVIRONMENT_KEY.inputDigest, + ) + .run(), + ).rejects.toThrow("rotation lacks D1 authority"); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualA, manifestGeneration: 1 }); + }); + + test("a failed old-backup tombstone rolls back the whole manifest swap", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const attemptA = await createEnvironmentArtifactStage(database); + const actualA = bucket.putBackup({ + createdAt: new Date(Date.now() - 60_000), + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptA.authority), + platformId: "650e8400-e29b-41d4-a716-44665544002f", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualA, + authority: attemptA.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptA.authority, + backupId: actualA, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + const attemptB = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + const actualB = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptB.authority), + platformId: "650e8400-e29b-41d4-a716-446655440030", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualB, + authority: attemptB.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + database.execute("DROP TRIGGER sandbox_backup_delete_intent_authority"); + await database + .prepare( + `INSERT INTO sandbox_backup_delete_intent ( + attempted_at, backup_id, created_at, delete_after, deleted_at + ) VALUES ( + NULL, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), + CAST(unixepoch('subsec') * 1000 AS INTEGER), NULL + )`, + ) + .bind(actualA) + .run(); + + await expect( + publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptB.authority, + backupId: actualB, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }), + ).rejects.toThrow(); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualA, manifestGeneration: 1 }); + expect( + await database + .prepare("SELECT count(*) AS count FROM sandbox_backup_delete_intent WHERE backup_id = ?") + .bind(actualA) + .first(), + ).toEqual({ count: 1 }); + }); + + test.each([ + ["traversal", { ...ENVIRONMENT_PATHS, python: [`${ENVIRONMENT_DIR}/python/../escape`] }], + ["outside", { ...ENVIRONMENT_PATHS, python: ["/workspace/escape"] }], + ["relative", { ...ENVIRONMENT_PATHS, python: ["python/site-packages"] }], + ["root itself", { ...ENVIRONMENT_PATHS, python: [ENVIRONMENT_DIR] }], + ["prefix collision", { ...ENVIRONMENT_PATHS, python: [`${ENVIRONMENT_DIR}-evil/path`] }], + ["empty segment", { ...ENVIRONMENT_PATHS, python: [`${ENVIRONMENT_DIR}//python`] }], + ["dot segment", { ...ENVIRONMENT_PATHS, python: [`${ENVIRONMENT_DIR}/./python`] }], + ["extra key", { ...ENVIRONMENT_PATHS, extra: [`${ENVIRONMENT_DIR}/extra`] }], + ["duplicate", { ...ENVIRONMENT_PATHS, python: [ENVIRONMENT_PATHS.executable[0]] }], + ["NUL", { ...ENVIRONMENT_PATHS, python: [`${ENVIRONMENT_DIR}/python/\0escape`] }], + ["colon", { ...ENVIRONMENT_PATHS, python: [`${ENVIRONMENT_DIR}/python:escape`] }], + ])("a legacy projection rejects %s artifact paths", async (_case, paths) => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-44665544000f", + }); + await bucket.put( + environmentPackageArtifactMetadataKey(ENVIRONMENT_KEY), + JSON.stringify({ + backupId: "650e8400-e29b-41d4-a716-44665544000f", + paths, + }), + ); + + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toBeNull(); + expect(await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY)).toBeNull(); + }); + + test("a legacy projection rejects extra metadata keys", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-446655440011", + }); + await bucket.put( + environmentPackageArtifactMetadataKey(ENVIRONMENT_KEY), + JSON.stringify({ + backupId: "650e8400-e29b-41d4-a716-446655440011", + extra: true, + paths: ENVIRONMENT_PATHS, + }), + ); + + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toBeNull(); + expect(await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY)).toBeNull(); + }); + + test("an existing D1 manifest cannot complete a successor stage", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const attemptA = await createEnvironmentArtifactStage(database); + const actualA = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptA.authority), + platformId: "650e8400-e29b-41d4-a716-446655440002", + }); + expect( + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualA, + authority: attemptA.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }), + ).toEqual({ actualBackupId: actualA }); + expect( + await commitEnvironmentPackageArtifactBackup(database, { + actualBackupId: actualA, + ...attemptA.authority, + expiresAt: bucket.backupExpiresAt(actualA), + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }), + ).toBe(true); + + const attemptB = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + attemptCount: attemptA.authority.attemptCount, + backupId: actualA, + claimOwner: attemptA.authority.claimOwner, + commandId: attemptA.authority.commandId, + deliveryGeneration: attemptA.authority.deliveryGeneration, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + + expect( + await getEnvironmentPackageArtifactBackupStage(database, ENVIRONMENT_COMMAND_ID), + ).toEqual(attemptB.stage); + expect(getSandboxBackupObjectKeys(actualA).every((key) => bucket.objects.has(key))).toBe(true); + expect(bucket.objects.has(environmentPackageArtifactMetadataKey(ENVIRONMENT_KEY))).toBe(false); + }); + + test("rotation keeps old objects until every issued reference has expired", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const attemptA = await createEnvironmentArtifactStage(database); + const actualA = bucket.putBackup({ + createdAt: new Date(Date.now() - 60_000), + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptA.authority), + platformId: "650e8400-e29b-41d4-a716-446655440008", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualA, + authority: attemptA.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptA.authority, + backupId: actualA, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + + const attemptB = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + const actualB = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptB.authority), + platformId: "650e8400-e29b-41d4-a716-446655440018", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualB, + authority: attemptB.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptB.authority, + backupId: actualB, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualB, manifestGeneration: 2 }); + expect( + await database + .prepare( + `SELECT attempted_at, delete_after, deleted_at + FROM sandbox_backup_delete_intent WHERE backup_id = ?`, + ) + .bind(actualA) + .first(), + ).toEqual({ + attempted_at: null, + delete_after: bucket.backupExpiresAt(actualA), + deleted_at: null, + }); + expect(await listPendingSandboxBackupDeletions(database, 64)).not.toContain(actualA); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + + expect(getSandboxBackupObjectKeys(actualA).every((key) => bucket.objects.has(key))).toBe(true); + expect(getSandboxBackupObjectKeys(actualB).every((key) => bucket.objects.has(key))).toBe(true); + }); + + test("an expired manifest rotates atomically and makes its old objects immediately due", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const attemptA = await createEnvironmentArtifactStage(database); + const actualA = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptA.authority), + platformId: "650e8400-e29b-41d4-a716-446655440020", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualA, + authority: attemptA.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptA.authority, + backupId: actualA, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + + await expireEnvironmentArtifactManifest(database); + const [, oldMetadataKey] = getSandboxBackupObjectKeys(actualA); + const oldMetadata = JSON.parse(bucket.objects.get(oldMetadataKey)!.body) as Record< + string, + unknown + >; + await bucket.put(oldMetadataKey, JSON.stringify({ ...oldMetadata, name: null })); + + const attemptB = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + const actualB = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptB.authority), + platformId: "650e8400-e29b-41d4-a716-446655440021", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualB, + authority: attemptB.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptB.authority, + backupId: actualB, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualB, manifestGeneration: 2 }); + expect(await listPendingSandboxBackupDeletions(database, 64)).toContain(actualA); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { cursor: null }); + expect(getSandboxBackupObjectKeys(actualA).every((key) => !bucket.objects.has(key))).toBe(true); + expect(getSandboxBackupObjectKeys(actualB).every((key) => bucket.objects.has(key))).toBe(true); + }); + + test("reconciliation retires an unused expired manifest and its R2 objects", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const platformId = "650e8400-e29b-41d4-a716-446655440029"; + const actualId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId, + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualId, + authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...authority, + backupId: actualId, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + await reconcileSandboxBackupPage(createBindings(database, bucket), { cursor: null }); + + await expect( + database + .prepare( + `DELETE FROM environment_package_artifact_backup + WHERE app_id = ? AND input_digest = ?`, + ) + .bind(ENVIRONMENT_APP_ID, ENVIRONMENT_KEY.inputDigest) + .run(), + ).rejects.toThrow("manifest has not expired"); + await expireEnvironmentArtifactManifest(database); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { cursor: null }); + + expect(await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY)).toBeNull(); + expect(getSandboxBackupObjectKeys(actualId).every((key) => !bucket.objects.has(key))).toBe( + true, + ); + expect( + await database + .prepare( + `SELECT attempted_at, deleted_at + FROM sandbox_backup_delete_intent WHERE backup_id = ?`, + ) + .bind(actualId) + .first(), + ).toEqual({ attempted_at: expect.any(Number), deleted_at: expect.any(Number) }); + + bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId, + }); + await bucket.put( + environmentPackageArtifactMetadataKey(ENVIRONMENT_KEY), + JSON.stringify({ backupId: platformId, paths: ENVIRONMENT_PATHS }), + ); + await succeedEnvironmentArtifactCommand(database); + await expect( + resolveEnvironmentPackageArtifactBackup(createBindings(database, bucket), ENVIRONMENT_KEY), + ).resolves.toBeNull(); + expect(await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY)).toBeNull(); + + const successor = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + const successorId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(successor.authority), + platformId: "650e8400-e29b-41d4-a716-44665544002a", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: successorId, + authority: successor.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...successor.authority, + backupId: successorId, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: successorId, manifestGeneration: 1 }); + }); + + test("a rotation retries as generation one when retirement wins its CAS", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const attemptA = await createEnvironmentArtifactStage(database); + const actualA = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptA.authority), + platformId: "650e8400-e29b-41d4-a716-446655440031", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualA, + authority: attemptA.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptA.authority, + backupId: actualA, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + await reconcileSandboxBackupPage(createBindings(database, bucket), { cursor: null }); + await expireEnvironmentArtifactManifest(database); + + const attemptB = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + const actualB = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptB.authority), + platformId: "650e8400-e29b-41d4-a716-446655440032", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualB, + authority: attemptB.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + let retired = false; + const intercept = (statement: D1PreparedStatement): D1PreparedStatement => + new Proxy(statement, { + get(target, property, receiver) { + if (property === "bind") { + return (...values: unknown[]) => intercept(Reflect.apply(target.bind, target, values)); + } + if (property === "run") { + return async (...args: unknown[]) => { + if (!retired) { + retired = true; + await retireExpiredEnvironmentPackageArtifactBackups(database, 64); + } + return Reflect.apply(target.run, target, args); + }; + } + return Reflect.get(target, property, receiver); + }, + }); + const racingDatabase = new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => { + const statement = target.prepare(query); + return query.includes("UPDATE environment_package_artifact_backup\n SET") + ? intercept(statement) + : statement; + }; + } + const value = Reflect.get(target, property, target) as unknown; + return typeof value === "function" ? value.bind(target) : value; + }, + }); + + await publishEnvironmentPackageArtifactBackup(createBindings(racingDatabase, bucket), { + ...attemptB.authority, + backupId: actualB, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + + expect(retired).toBe(true); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualB, manifestGeneration: 1 }); + expect(await listPendingSandboxBackupDeletions(database, 64)).toContain(actualA); + }); + + test("an initial manifest survives a lost D1 acknowledgement", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const actualId = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-446655440043", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualId, + authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + const fault = loseRunAcknowledgementOnce( + database, + "INSERT INTO environment_package_artifact_backup", + ); + + await expect( + publishEnvironmentPackageArtifactBackup(createBindings(fault.database, bucket), { + ...authority, + backupId: actualId, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }), + ).rejects.toThrow("Simulated D1 write acknowledgement loss"); + expect(fault.wasLost()).toBe(true); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualId, manifestGeneration: 1 }); + + await expect( + publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...authority, + backupId: actualId, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }), + ).resolves.toBeUndefined(); + }); + + test("a committed manifest swap survives a lost D1 acknowledgement", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const attemptA = await createEnvironmentArtifactStage(database); + const actualA = bucket.putBackup({ + createdAt: new Date(Date.now() - 60_000), + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptA.authority), + platformId: "650e8400-e29b-41d4-a716-446655440022", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualA, + authority: attemptA.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptA.authority, + backupId: actualA, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + + const attemptB = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + const actualB = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptB.authority), + platformId: "650e8400-e29b-41d4-a716-446655440023", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualB, + authority: attemptB.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + const fault = loseRunAcknowledgementOnce( + database, + "UPDATE environment_package_artifact_backup\n SET", + ); + + await expect( + publishEnvironmentPackageArtifactBackup(createBindings(fault.database, bucket), { + ...attemptB.authority, + backupId: actualB, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }), + ).rejects.toThrow("Simulated D1 write acknowledgement loss"); + + expect(fault.wasLost()).toBe(true); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualB, manifestGeneration: 2 }); + expect( + await database + .prepare("SELECT count(*) AS count FROM sandbox_backup_delete_intent WHERE backup_id = ?") + .bind(actualA) + .first(), + ).toEqual({ count: 1 }); + await expect( + publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptB.authority, + backupId: actualB, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }), + ).resolves.toBeUndefined(); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualB, manifestGeneration: 2 }); + }); + + test("a stale rotation CAS loses without disturbing its winner", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const attemptA = await createEnvironmentArtifactStage(database); + const actualA = bucket.putBackup({ + createdAt: new Date(Date.now() - 120_000), + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptA.authority), + platformId: "650e8400-e29b-41d4-a716-446655440024", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualA, + authority: attemptA.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptA.authority, + backupId: actualA, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + + const attemptB = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "second-owner", + deliveryGeneration: 2, + }); + const actualB = bucket.putBackup({ + createdAt: new Date(Date.now() - 60_000), + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptB.authority), + platformId: "650e8400-e29b-41d4-a716-446655440025", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualB, + authority: attemptB.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + + let actualC: SandboxBackupId | null = null; + let intercepted = false; + const intercept = (statement: D1PreparedStatement): D1PreparedStatement => + new Proxy(statement, { + get(target, property, receiver) { + if (property === "bind") { + return (...values: unknown[]) => intercept(Reflect.apply(target.bind, target, values)); + } + if (property === "run") { + return async (...args: unknown[]) => { + if (!intercepted) { + intercepted = true; + const attemptC = await createEnvironmentArtifactStage(database, { + attemptCount: 3, + claimOwner: "third-owner", + deliveryGeneration: 3, + }); + actualC = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptC.authority), + platformId: "650e8400-e29b-41d4-a716-446655440026", + }); + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: actualC, + authority: attemptC.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }); + await publishEnvironmentPackageArtifactBackup(createBindings(database, bucket), { + ...attemptC.authority, + backupId: actualC, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + } + return Reflect.apply(target.run, target, args); + }; + } + return Reflect.get(target, property, receiver); + }, + }); + const racingDatabase = new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => { + const statement = target.prepare(query); + return query.includes("UPDATE environment_package_artifact_backup\n SET") + ? intercept(statement) + : statement; + }; + } + const value = Reflect.get(target, property, target) as unknown; + return typeof value === "function" ? value.bind(target) : value; + }, + }); + + await publishEnvironmentPackageArtifactBackup(createBindings(racingDatabase, bucket), { + ...attemptB.authority, + backupId: actualB, + key: ENVIRONMENT_KEY, + paths: ENVIRONMENT_PATHS, + }); + + expect(intercepted).toBe(true); + expect(actualC).not.toBeNull(); + expect( + await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY), + ).toMatchObject({ backupId: actualC, manifestGeneration: 2 }); + expect(await listPendingSandboxBackupDeletions(database, 64)).toContain(actualB); + expect(getSandboxBackupObjectKeys(actualB).every((key) => bucket.objects.has(key))).toBe(true); + + bucket.failAfterDeletingKey = getSandboxBackupObjectKeys(actualB)[0]; + await reconcileSandboxBackupPage(createBindings(database, bucket), { cursor: null }); + await reconcileSandboxBackupPage(createBindings(database, bucket), { cursor: null }); + expect(getSandboxBackupObjectKeys(actualB).every((key) => !bucket.objects.has(key))).toBe(true); + expect(getSandboxBackupObjectKeys(actualA).every((key) => bucket.objects.has(key))).toBe(true); + expect( + actualC !== null && + getSandboxBackupObjectKeys(actualC).every((key) => bucket.objects.has(key)), + ).toBe(true); + }); + + test("scanner cannot claim an old attempt candidate into a rotated stage", async () => { + const database = createDatabase(); + const attemptA = await createEnvironmentArtifactStage(database); + const candidateA = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440003"); + + const attemptB = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + expect( + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: candidateA, + authority: attemptA.authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }), + ).toBeNull(); + expect( + await getEnvironmentPackageArtifactBackupStage(database, ENVIRONMENT_COMMAND_ID), + ).toEqual(attemptB.stage); + }); + + test("old attempt objects are collected after grace without disturbing the current stage", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const attemptA = await createEnvironmentArtifactStage(database); + const actualA = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(attemptA.authority), + platformId: "650e8400-e29b-41d4-a716-446655440004", + }); + const attemptB = await createEnvironmentArtifactStage(database, { + attemptCount: 2, + claimOwner: "successor-owner", + deliveryGeneration: 2, + }); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + + expect(getSandboxBackupObjectKeys(actualA).every((key) => !bucket.objects.has(key))).toBe(true); + expect( + await getEnvironmentPackageArtifactBackupStage(database, ENVIRONMENT_COMMAND_ID), + ).toEqual(attemptB.stage); + }); + + test("partial objects use bounded grace even while another artifact writer is active", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const metaOnly = bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: "650e8400-e29b-41d4-a716-446655440005", + }); + expect( + await claimEnvironmentPackageArtifactBackupActual(database, { + actualBackupId: metaOnly, + authority, + commandId: ENVIRONMENT_COMMAND_ID, + dir: ENVIRONMENT_DIR, + }), + ).toEqual({ actualBackupId: metaOnly }); + const [metaOnlyDataKey, metaOnlyMetadataKey] = getSandboxBackupObjectKeys(metaOnly); + await bucket.delete(metaOnlyDataKey); + + const oldDataOnly = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440006"); + const freshDataOnly = encodeSandboxBackupIdForStorage("650e8400-e29b-41d4-a716-446655440007"); + const [oldDataKey] = getSandboxBackupObjectKeys(oldDataOnly); + const [freshDataKey] = getSandboxBackupObjectKeys(freshDataOnly); + await bucket.put(oldDataKey, "old-partial"); + await bucket.put(freshDataKey, "fresh-partial", new Date(NOW)); + + await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + + expect(bucket.objects.has(metaOnlyMetadataKey)).toBe(false); + expect(bucket.objects.has(oldDataKey)).toBe(false); + expect(bucket.objects.has(freshDataKey)).toBe(true); + expect( + await getEnvironmentPackageArtifactBackupStage(database, ENVIRONMENT_COMMAND_ID), + ).toMatchObject({ actualBackupId: null }); + }); + + test("opaque pagination converges many complete candidates on one D1 manifest", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + const { authority } = await createEnvironmentArtifactStage(database); + const candidates: SandboxBackupId[] = []; + for (let index = 0; index < 65; index += 1) { + candidates.push( + bucket.putBackup({ + dir: ENVIRONMENT_DIR, + name: createEnvironmentPackageArtifactBackupName(authority), + platformId: `650e8400-e29b-4000-8000-${index.toString(16).padStart(12, "0")}`, + }), + ); + } + + let cursor: string | null = null; + let pageCount = 0; + do { + const page = await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor, + }); + cursor = page.nextCursor; + pageCount += 1; + if (!page.hasMore) { + break; + } + } while (pageCount < 10); + + const winner = (await getEnvironmentPackageArtifactBackupManifest(database, ENVIRONMENT_KEY)) + ?.backupId; + if (winner === undefined) { + throw new Error("Environment artifact candidates produced no D1 winner."); + } + expect(candidates).toContain(winner); + expect(pageCount).toBeGreaterThan(1); + expect( + [...bucket.objects.keys()].filter((key) => key.startsWith("backups/")).toSorted(), + ).toEqual([...getSandboxBackupObjectKeys(winner)].toSorted()); + }); + + test("an empty R2 page continues until every expired D1 manifest is retired", async () => { + const database = createDatabase(); + const bucket = new MemoryR2Bucket(); + await createEnvironmentArtifactStage(database); + database.execute("DROP TRIGGER environment_package_artifact_backup_authority"); + const committedAt = Date.now() - 2 * 86_400_000; + const expiresAt = committedAt + 86_400_001; + for (let index = 0; index < 65; index += 1) { + const inputDigest = index.toString(16).padStart(64, "0"); + const dir = `/workspace/.mosoo/environment-artifacts/${inputDigest}`; + const backupId = encodeSandboxBackupIdForStorage( + `650e8400-e29b-4000-8001-${index.toString(16).padStart(12, "0")}`, + ); + await database + .prepare( + `INSERT INTO environment_package_artifact_backup ( + app_id, attempt_count, backup_id, command_id, committed_at, + delivery_generation, expires_at, input_digest, manifest_generation, paths_json + ) VALUES (?, 1, ?, ?, ?, 1, ?, ?, 1, ?)`, + ) + .bind( + ENVIRONMENT_APP_ID, + backupId, + ENVIRONMENT_COMMAND_ID, + committedAt, + expiresAt, + inputDigest, + JSON.stringify({ executable: [`${dir}/bin/tool`], node: [], python: [] }), + ) + .run(); + } + + const first = await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + expect(first).toEqual({ hasMore: true, nextCursor: null, processed: 0 }); + expect( + await database + .prepare("SELECT count(*) AS count FROM environment_package_artifact_backup") + .first(), + ).toEqual({ count: 1 }); + + const second = await reconcileSandboxBackupPage(createBindings(database, bucket), { + cursor: null, + }); + expect(second).toEqual({ hasMore: false, nextCursor: null, processed: 0 }); + expect( + await database + .prepare("SELECT count(*) AS count FROM environment_package_artifact_backup") + .first(), + ).toEqual({ count: 0 }); + }); + + test("hidden rowid cannot replace permanent backup authority", async () => { + const database = createDatabase(); + const stage = await createOperationStage(database); + const readyId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440030"); + const replacementId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440031"); + const replacementStageId = encodeSandboxBackupIdForStorage( + "550e8400-e29b-41d4-a716-446655440032", + ); + const intentId = encodeSandboxBackupIdForStorage("550e8400-e29b-41d4-a716-446655440033"); + const replacementIntentId = encodeSandboxBackupIdForStorage( + "550e8400-e29b-41d4-a716-446655440034", + ); + + await claimSandboxBackupStageActual(database, { + actualBackupId: readyId, + dir: DIR, + sandboxIncarnation: 1, + stagingId: stage.id, + }); + await finalizeSandboxBackupStage(database, { + actualBackupId: readyId, + stagingId: stage.id, + }); + expect( + await authorizeSandboxBackupDeletion(database, { + authority: { kind: "unattributed" }, + backupId: intentId, + }), + ).toBe(true); + + await database.prepare("PRAGMA recursive_triggers = OFF").run(); + + expect( + ( + await database + .prepare( + `SELECT name, wr FROM pragma_table_list + WHERE name IN ( + 'environment_package_artifact_backup', + 'environment_package_artifact_backup_staging', + 'sandbox_backup', + 'sandbox_backup_delete_intent', + 'sandbox_backup_staging' + ) + ORDER BY name`, + ) + .all() + ).results, + ).toEqual([ + { name: "environment_package_artifact_backup", wr: 1 }, + { name: "environment_package_artifact_backup_staging", wr: 1 }, + { name: "sandbox_backup", wr: 1 }, + { name: "sandbox_backup_delete_intent", wr: 1 }, + { name: "sandbox_backup_staging", wr: 1 }, + ]); + + await expect( + database + .prepare( + `INSERT OR REPLACE INTO sandbox_backup ( + rowid, created_at, dir, id, keep, operation_id, sandbox_id, + sandbox_incarnation, session_run_id, staging_id, status, + ttl_seconds, updated_at, workspace_session_id + ) VALUES ( + 1, ?, '/workspace/rowid-replacement', ?, 0, ?, ?, 1, + NULL, ?, 'ready', 100, ?, NULL + )`, + ) + .bind(NOW, replacementId, OPERATION_ID, SANDBOX_ID, replacementStageId, NOW) + .run(), + ).rejects.toThrow(/rowid/i); + + await expect( + database + .prepare( + `INSERT OR REPLACE INTO sandbox_backup_delete_intent ( + rowid, attempted_at, backup_id, created_at, delete_after, deleted_at + ) VALUES ( + 1, NULL, ?, CAST(unixepoch('subsec') * 1000 AS INTEGER), + CAST(unixepoch('subsec') * 1000 AS INTEGER), NULL + )`, + ) + .bind(replacementIntentId) + .run(), + ).rejects.toThrow(/rowid/i); + + expect(await readyRows(database)).toEqual([{ id: readyId, staging_id: stage.id }]); + expect( + ( + await database + .prepare("SELECT backup_id FROM sandbox_backup_delete_intent ORDER BY backup_id") + .all() + ).results, + ).toEqual([{ backup_id: intentId }]); + }); +}); diff --git a/apps/api/tests/sandbox-backup-pruning.test.ts b/apps/api/tests/sandbox-backup-pruning.test.ts index 56016445..f8da52f8 100644 --- a/apps/api/tests/sandbox-backup-pruning.test.ts +++ b/apps/api/tests/sandbox-backup-pruning.test.ts @@ -1,216 +1,183 @@ import { describe, expect, test } from "bun:test"; +import { createPlatformId } from "@mosoo/id"; +import type { RuntimeOperationId, SandboxBackupId } from "@mosoo/id"; + import { selectSandboxBackupPruneIds } from "../src/modules/runtime/infrastructure/sandbox-backup-pruning"; import { listReadySandboxBackupsForPruning, markSandboxBackupsPruned, - recordCreatedSandboxBackups, } from "../src/modules/runtime/infrastructure/sandbox-backup-store"; +import { applyDrizzleMigrationsThrough } from "./helpers/drizzle-migrations"; import { SqliteD1Database } from "./helpers/sqlite-d1"; -const BACKUP_ID_1 = "01J000000000000000000000H1"; -const BACKUP_ID_2 = "01J000000000000000000000H2"; -const BACKUP_ID_3 = "01J000000000000000000000H3"; -const BACKUP_ID_4 = "01J000000000000000000000H4"; -const BACKUP_ID_5 = "01J000000000000000000000H5"; -const KEEP_BACKUP_ID = "01J000000000000000000000HZ"; -const MEMORY_NEW_BACKUP_ID = "01J000000000000000000000H6"; -const MEMORY_OLD_BACKUP_ID = "01J000000000000000000000H7"; -const SESSION_BACKUP_ID = "01J000000000000000000000H8"; - -function createSandboxBackupDatabase(input: { maxBoundParams?: number } = {}): SqliteD1Database { - const database = new SqliteD1Database(input); - - database.execute(` - CREATE TABLE sandbox ( - id text PRIMARY KEY NOT NULL, - last_backup_id text, - status text NOT NULL, - status_changed_at integer DEFAULT 0 NOT NULL, - status_event text DEFAULT 'runtime_subject.cold' NOT NULL, - status_operation_id text, - status_seq integer DEFAULT 0 NOT NULL, - status_source text DEFAULT 'system' NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE sandbox_backup ( - created_at integer NOT NULL, - dir text NOT NULL, - error_message text, - id text PRIMARY KEY NOT NULL, - keep integer DEFAULT 0 NOT NULL, - sandbox_id text NOT NULL, - session_run_id text, - status text NOT NULL, - ttl_seconds integer NOT NULL, - updated_at integer NOT NULL - ); - `); +const SANDBOX_ID = "01J0000000000000000000000D"; +const BACKUP_IDS = [..."12345Z"].map((suffix) => `01J000000000000000000000H${suffix}`); +function createDatabase(): SqliteD1Database { + const database = new SqliteD1Database(); + applyDrizzleMigrationsThrough(database, "0020_sandbox-backup-object-authority"); return database; } -async function insertSandbox(database: D1Database): Promise { +async function insertActiveSandbox(database: D1Database): Promise { + await database + .prepare( + `INSERT INTO sandbox ( + agent_id, app_id, created_at, id, incarnation, kind, network_constraints_hash, + owner_account_id, status, subject_id, subject_kind, updated_at + ) VALUES (?, ?, 1, ?, 1, 'pet', ?, ?, 'active', ?, 'agent', 1)`, + ) + .bind( + "01J0000000000000000000000E", + "01J0000000000000000000000F", + SANDBOX_ID, + "0".repeat(64), + "01J0000000000000000000000G", + "01J0000000000000000000000E", + ) + .run(); +} + +async function insertWorkspaceSession(database: D1Database, sessionId: string): Promise { await database - .prepare("INSERT INTO sandbox (id, last_backup_id, status, updated_at) VALUES (?, ?, ?, ?)") - .bind("01J0000000000000000000000D", null, "backing_up", 1) + .prepare( + `INSERT INTO session ( + agent_id, app_id, created_at, creator_account_id, id, kind, model, + provider, renamed, runtime_id, status, updated_at + ) VALUES (?, ?, 1, ?, ?, 'pet', 'gpt-5.4', 'openai', 0, 'openai-runtime', 'IDLE', 1)`, + ) + .bind( + "01J0000000000000000000000E", + "01J0000000000000000000000F", + "01J0000000000000000000000G", + sessionId, + ) .run(); } async function insertBackup( database: D1Database, - input: { - createdAt: number; - id: string; - keep?: boolean; - }, + input: { createdAt: number; id: string; keep?: boolean; workspaceSessionId?: string }, ): Promise { await database .prepare( - ` - INSERT INTO sandbox_backup ( - created_at, - dir, - error_message, - id, - keep, - sandbox_id, - status, - ttl_seconds, - updated_at - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) - `, + `INSERT INTO sandbox_backup ( + created_at, dir, id, keep, operation_id, sandbox_id, sandbox_incarnation, + session_run_id, staging_id, status, ttl_seconds, updated_at, workspace_session_id + ) VALUES (?, '/workspace', ?, ?, ?, ?, 1, NULL, ?, 'ready', 100, ?, ?)`, ) .bind( input.createdAt, - "/workspace", - null, input.id, input.keep === true ? 1 : 0, - "01J0000000000000000000000D", - "ready", - 100, + createPlatformId(), + SANDBOX_ID, + createPlatformId(), input.createdAt, + input.workspaceSessionId ?? null, ) .run(); } describe("sandbox backup pruning", () => { - test("selects pruned backup ids and marks records", async () => { - const database = createSandboxBackupDatabase(); - - await insertBackup(database, { createdAt: 1, id: BACKUP_ID_1 }); - await insertBackup(database, { createdAt: 2, id: BACKUP_ID_2 }); - await insertBackup(database, { createdAt: 3, id: BACKUP_ID_3 }); - await insertBackup(database, { createdAt: 4, id: BACKUP_ID_4 }); - await insertBackup(database, { createdAt: 5, id: BACKUP_ID_5 }); - await insertBackup(database, { createdAt: 0, id: KEEP_BACKUP_ID, keep: true }); - - const backups = await listReadySandboxBackupsForPruning(database, "01J0000000000000000000000D"); - const pruneIds = selectSandboxBackupPruneIds(backups); - - await markSandboxBackupsPruned(database, pruneIds); - - expect(pruneIds).toEqual([BACKUP_ID_2, BACKUP_ID_1]); + test("marks only stable overflow after D1 protection is rechecked", async () => { + const database = createDatabase(); + await insertActiveSandbox(database); + for (const [index, id] of BACKUP_IDS.entries()) { + await insertBackup(database, { + createdAt: index + 1, + id, + keep: index === BACKUP_IDS.length - 1, + }); + } + await database + .prepare("UPDATE sandbox SET last_backup_id = ? WHERE id = ?") + .bind(BACKUP_IDS[0], SANDBOX_ID) + .run(); + + const backups = await listReadySandboxBackupsForPruning(database, SANDBOX_ID); + const candidates = selectSandboxBackupPruneIds(backups); + const pruned = await markSandboxBackupsPruned(database, candidates); + + expect(candidates).toEqual([BACKUP_IDS[1]]); + expect(pruned).toEqual([BACKUP_IDS[1]]); const rows = await database - .prepare("SELECT id, status FROM sandbox_backup ORDER BY id") + .prepare("SELECT id, status FROM sandbox_backup ORDER BY created_at") .all<{ id: string; status: string }>(); - expect(rows.results).toEqual([ - { id: BACKUP_ID_1, status: "pruned" }, - { id: BACKUP_ID_2, status: "pruned" }, - { id: BACKUP_ID_3, status: "ready" }, - { id: BACKUP_ID_4, status: "ready" }, - { id: BACKUP_ID_5, status: "ready" }, - { id: KEEP_BACKUP_ID, status: "ready" }, - ]); - }); - - test("records checkpoint backup batch and stores latest subject checkpoint", async () => { - const database = createSandboxBackupDatabase(); - await insertSandbox(database); - - await recordCreatedSandboxBackups(database, { - backups: [ - { - backup: { dir: "/workspace/one", id: SESSION_BACKUP_ID }, - updateSandboxLastBackup: false, - }, - { - backup: { dir: "/memory", id: MEMORY_OLD_BACKUP_ID }, - updateSandboxLastBackup: true, - }, - { - backup: { dir: "/memory", id: MEMORY_NEW_BACKUP_ID }, - updateSandboxLastBackup: true, - }, - ], - sandboxId: "01J0000000000000000000000D", - ttlSeconds: 100, - }); - - const backupRows = await database - .prepare("SELECT id, dir, status, ttl_seconds FROM sandbox_backup ORDER BY id") - .all<{ dir: string; id: string; status: string; ttl_seconds: number }>(); - expect(backupRows.results).toEqual([ - { - dir: "/memory", - id: MEMORY_NEW_BACKUP_ID, - status: "ready", - ttl_seconds: 100, - }, - { - dir: "/memory", - id: MEMORY_OLD_BACKUP_ID, - status: "ready", - ttl_seconds: 100, - }, - { - dir: "/workspace/one", - id: SESSION_BACKUP_ID, - status: "ready", - ttl_seconds: 100, - }, + expect(rows.results.map((row) => row.status)).toEqual([ + "ready", + "pruned", + "ready", + "ready", + "ready", + "ready", ]); - - const sandbox = await database - .prepare("SELECT last_backup_id, status, status_seq FROM sandbox WHERE id = ?") - .bind("01J0000000000000000000000D") - .first<{ last_backup_id: string; status: string; status_seq: number }>(); - expect(sandbox).toEqual({ - last_backup_id: MEMORY_NEW_BACKUP_ID, - status: "backing_up", - status_seq: 0, - }); }); - test("records more backups than fit in one D1 statement", async () => { - const database = createSandboxBackupDatabase({ maxBoundParams: 100 }); - await insertSandbox(database); - const suffixes = [..."123456789ABCDEFG"]; - - await recordCreatedSandboxBackups(database, { - backups: suffixes.map((suffix, index) => ({ - backup: { - dir: index === suffixes.length - 1 ? "/memory" : `/workspace/${index}`, - id: `01J000000000000000000000H${suffix}`, - }, - updateSandboxLastBackup: index === suffixes.length - 1, - })), - sandboxId: "01J0000000000000000000000D", - ttlSeconds: 100, + for (const protection of ["last_backup", "restoring", "provisioning"] as const) { + test(`rechecks ${protection} protection after pruning candidates were listed`, async () => { + const database = createDatabase(); + const sessionId = "01J0000000000000000000000S"; + await insertActiveSandbox(database); + if (protection === "provisioning") { + await insertWorkspaceSession(database, sessionId); + } + for (const [index, id] of BACKUP_IDS.slice(0, 5).entries()) { + await insertBackup(database, { + createdAt: index + 1, + id, + workspaceSessionId: protection === "provisioning" && index === 1 ? sessionId : undefined, + }); + } + + const listed = await listReadySandboxBackupsForPruning(database, SANDBOX_ID); + const candidates = selectSandboxBackupPruneIds(listed); + expect(candidates).toEqual([BACKUP_IDS[1], BACKUP_IDS[0]]); + const protectedId = candidates[0]; + switch (protection) { + case "last_backup": { + await database + .prepare("UPDATE sandbox SET last_backup_id = ? WHERE id = ?") + .bind(protectedId, SANDBOX_ID) + .run(); + break; + } + case "restoring": { + await database + .prepare( + `UPDATE sandbox + SET claim_expires_at = ?, claim_owner = 'restore-owner', + last_restore_backup_id = ?, operation_kind = 'activate', + status = 'restoring', status_operation_id = ? + WHERE id = ?`, + ) + .bind(Date.now() + 60_000, protectedId, "01J0000000000000000000000A", SANDBOX_ID) + .run(); + break; + } + case "provisioning": { + await database + .prepare( + `UPDATE session + SET runtime_provisioning_heartbeat_at = 1, + runtime_provisioning_operation_id = ?, + runtime_provisioning_sandbox_id = ? + WHERE id = ?`, + ) + .bind("01J0000000000000000000000A", SANDBOX_ID, sessionId) + .run(); + break; + } + } + + expect(await markSandboxBackupsPruned(database, candidates)).toEqual([candidates[1]]); + await expect( + database + .prepare("SELECT status FROM sandbox_backup WHERE id = ?") + .bind(protectedId) + .first(), + ).resolves.toEqual({ status: "ready" }); }); - - const backupCount = await database - .prepare("SELECT COUNT(*) AS count FROM sandbox_backup") - .first<{ count: number }>(); - const sandbox = await database - .prepare("SELECT last_backup_id FROM sandbox WHERE id = ?") - .bind("01J0000000000000000000000D") - .first<{ last_backup_id: string }>(); - - expect(backupCount?.count).toBe(16); - expect(sandbox?.last_backup_id).toBe("01J000000000000000000000HG"); - }); + } }); diff --git a/apps/api/tests/sandbox-backup-staging.test.ts b/apps/api/tests/sandbox-backup-staging.test.ts new file mode 100644 index 00000000..2bd2adea --- /dev/null +++ b/apps/api/tests/sandbox-backup-staging.test.ts @@ -0,0 +1,365 @@ +import { expect, test } from "bun:test"; + +import { createPlatformId, parsePlatformId } from "@mosoo/id"; +import type { RuntimeOperationId, SandboxBackupId, SandboxId, SessionId } from "@mosoo/id"; + +import { + deferSandboxBackupStageRepair, + listSandboxBackupStages, + listSandboxSessionBackupCandidates, + revokeSandboxBackupsForSessionDelete, + stageSandboxBackupWrites, +} from "../src/modules/runtime/infrastructure/sandbox-backup-store"; +import { applyDrizzleMigrationsThrough } from "./helpers/drizzle-migrations"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +const SANDBOX_ID = parsePlatformId("01J0000000000000000000000D"); +const SESSION_ID = parsePlatformId("01J0000000000000000000000S"); +const NEXT_SESSION_ID = parsePlatformId("01J0000000000000000000000T"); +const OPERATION_ID = parsePlatformId("01J0000000000000000000000A"); + +function createDatabase(): SqliteD1Database { + const database = new SqliteD1Database(); + applyDrizzleMigrationsThrough(database, "0020_sandbox-backup-object-authority"); + return database; +} + +async function insertBackingUpSandbox( + database: D1Database, + sandboxId: SandboxId = SANDBOX_ID, +): Promise { + const now = Date.now(); + await database + .prepare( + `INSERT INTO sandbox ( + agent_id, app_id, claim_expires_at, claim_owner, created_at, id, incarnation, + kind, network_constraints_hash, operation_kind, owner_account_id, status, status_operation_id, + subject_id, subject_kind, updated_at + ) VALUES (?, ?, ?, 'owner', ?, ?, 1, 'pet', ?, 'hibernate', ?, 'backing_up', ?, ?, 'agent', ?)`, + ) + .bind( + sandboxId, + "01J0000000000000000000000F", + now + 60_000, + now, + sandboxId, + "0".repeat(64), + "01J0000000000000000000000G", + OPERATION_ID, + sandboxId, + now, + ) + .run(); +} + +async function insertSession( + database: D1Database, + input: { + readonly cleanup?: boolean; + readonly id: SessionId; + readonly lastMessageAt: number; + }, +): Promise { + const now = Date.now(); + await database + .prepare( + `INSERT INTO session ( + agent_id, app_id, archived_at, cleanup_operation_kind, created_at, + creator_account_id, id, kind, last_message_at, model, provider, renamed, + runtime_id, status, status_operation_id, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, 'pet', ?, 'gpt-5.4', 'openai', 0, + 'openai-runtime', 'IDLE', ?, ?)`, + ) + .bind( + "01J0000000000000000000000E", + "01J0000000000000000000000F", + input.cleanup === true ? now : null, + input.cleanup === true ? "delete" : null, + now, + "01J0000000000000000000000G", + input.id, + input.lastMessageAt, + input.cleanup === true ? OPERATION_ID : null, + now, + ) + .run(); +} + +async function insertWorkspace( + database: D1Database, + input: { + readonly cwd: string; + readonly incarnation: number; + readonly sandboxId?: SandboxId; + readonly sessionId: SessionId; + readonly status: "active" | "closed"; + }, +): Promise { + const now = Date.now(); + await database + .prepare( + `INSERT INTO sandbox_session ( + cloudflare_session_id, created_at, cwd, origin_json, sandbox_id, + sandbox_incarnation, session_id, status, updated_at + ) VALUES (?, ?, ?, '{}', ?, ?, ?, ?, ?)`, + ) + .bind( + input.sessionId, + now, + input.cwd, + input.sandboxId ?? SANDBOX_ID, + input.incarnation, + input.sessionId, + input.status, + now, + ) + .run(); +} + +test("workspace staging binds the exact session scope in SQLite", async () => { + const database = createDatabase(); + await insertBackingUpSandbox(database); + await insertSession(database, { id: SESSION_ID, lastMessageAt: 1 }); + await insertWorkspace(database, { + cwd: "/workspace", + incarnation: 1, + sessionId: SESSION_ID, + status: "active", + }); + + const [write] = await stageSandboxBackupWrites(database, { + admission: { + kind: "operation", + lease: { + claimExpiresAt: Date.now() + 60_000, + claimOwner: "owner", + incarnation: 1, + kind: "hibernate", + operationId: OPERATION_ID, + status: "backing_up", + }, + }, + sandboxId: SANDBOX_ID, + targets: [ + { + dir: "/workspace", + updateSandboxLastBackup: false, + workspaceSessionId: SESSION_ID, + }, + ], + ttlSeconds: 100, + }); + + expect(write?.kind).toBe("staged"); + if (write?.kind === "staged") { + expect(write.stage.workspaceSessionId).toBe(SESSION_ID); + expect(write.stage.sandboxIncarnation).toBe(1); + } +}); + +test("session cleanup fences new workspace staging", async () => { + const database = createDatabase(); + await insertBackingUpSandbox(database); + await insertSession(database, { cleanup: true, id: SESSION_ID, lastMessageAt: 1 }); + await insertWorkspace(database, { + cwd: "/workspace", + incarnation: 1, + sessionId: SESSION_ID, + status: "active", + }); + + await expect( + stageSandboxBackupWrites(database, { + admission: { + kind: "operation", + lease: { + claimExpiresAt: Date.now() + 60_000, + claimOwner: "owner", + incarnation: 1, + kind: "hibernate", + operationId: OPERATION_ID, + status: "backing_up", + }, + }, + sandboxId: SANDBOX_ID, + targets: [ + { + dir: "/workspace", + updateSandboxLastBackup: false, + workspaceSessionId: SESSION_ID, + }, + ], + ttlSeconds: 100, + }), + ).rejects.toThrow("lost its exact lifecycle authority"); +}); + +test("checkpoint candidates exclude a closed workspace from the prior incarnation", async () => { + const database = createDatabase(); + await insertSession(database, { id: SESSION_ID, lastMessageAt: 1 }); + await insertSession(database, { id: NEXT_SESSION_ID, lastMessageAt: 2 }); + await insertWorkspace(database, { + cwd: "/workspace/old", + incarnation: 1, + sessionId: SESSION_ID, + status: "closed", + }); + await insertWorkspace(database, { + cwd: "/workspace/current", + incarnation: 2, + sessionId: NEXT_SESSION_ID, + status: "active", + }); + + const candidates = await listSandboxSessionBackupCandidates(database, SANDBOX_ID, 2); + + expect(candidates.map(({ cwd, sessionId }) => ({ cwd, sessionId }))).toEqual([ + { cwd: "/workspace/current", sessionId: NEXT_SESSION_ID }, + ]); +}); + +test("failed repair pages rotate instead of starving the next stage", async () => { + const database = createDatabase(); + const ids = Array.from({ length: 65 }, () => createPlatformId()); + for (const [index, id] of ids.entries()) { + await database + .prepare( + `INSERT INTO sandbox_backup_staging ( + claim_owner, created_at, dir, id, operation_id, sandbox_id, sandbox_incarnation, + ttl_seconds, updated_at, updates_subject_backup + ) VALUES ('stale-owner', 1, ?, ?, ?, ?, 1, 100, 1, 0)`, + ) + .bind(`/workspace/${index.toString().padStart(2, "0")}`, id, OPERATION_ID, SANDBOX_ID) + .run(); + } + + const firstPage = await listSandboxBackupStages(database, 64); + expect(firstPage).toHaveLength(64); + expect(firstPage.map(({ id }) => id)).not.toContain(ids[64]); + for (const stage of firstPage) { + expect(await deferSandboxBackupStageRepair(database, stage)).toBe(true); + } + + const secondPage = await listSandboxBackupStages(database, 64); + expect(secondPage[0]?.id).toBe(ids[64]); +}); + +test("session delete revokes every exact workspace incarnation and preserves other owners", async () => { + const database = createDatabase(); + const otherSandboxId = parsePlatformId("01J0000000000000000000000E"); + const sameSandboxOtherSessionId = parsePlatformId("01J0000000000000000000000V"); + await insertBackingUpSandbox(database); + await insertBackingUpSandbox(database, otherSandboxId); + await insertSession(database, { cleanup: true, id: SESSION_ID, lastMessageAt: 1 }); + await insertSession(database, { id: NEXT_SESSION_ID, lastMessageAt: 2 }); + await insertSession(database, { id: sameSandboxOtherSessionId, lastMessageAt: 3 }); + await insertWorkspace(database, { + cwd: "/workspace/shared", + incarnation: 2, + sessionId: SESSION_ID, + status: "closed", + }); + await insertWorkspace(database, { + cwd: "/workspace/shared", + incarnation: 1, + sandboxId: otherSandboxId, + sessionId: NEXT_SESSION_ID, + status: "closed", + }); + await insertWorkspace(database, { + cwd: "/workspace/shared", + incarnation: 2, + sessionId: sameSandboxOtherSessionId, + status: "closed", + }); + + const exactDeletedPriorId = createPlatformId(); + const exactDeletedCurrentId = createPlatformId(); + const sameSandboxOtherSessionBackupId = createPlatformId(); + const otherSandboxBackupId = createPlatformId(); + const legacyDeletedSandboxId = createPlatformId(); + const legacyOtherSandboxId = createPlatformId(); + for (const [backupId, sandboxId, workspaceSessionId, incarnation] of [ + [exactDeletedPriorId, SANDBOX_ID, SESSION_ID, 1], + [exactDeletedCurrentId, SANDBOX_ID, SESSION_ID, 2], + [sameSandboxOtherSessionBackupId, SANDBOX_ID, sameSandboxOtherSessionId, 2], + [otherSandboxBackupId, otherSandboxId, NEXT_SESSION_ID, 1], + [legacyDeletedSandboxId, SANDBOX_ID, null, 0], + [legacyOtherSandboxId, otherSandboxId, null, 0], + ] as const) { + await database + .prepare( + `INSERT INTO sandbox_backup ( + created_at, dir, id, keep, operation_id, sandbox_id, sandbox_incarnation, + staging_id, status, ttl_seconds, updated_at, workspace_session_id + ) VALUES (1, '/workspace/shared', ?, 0, ?, ?, ?, ?, 'ready', 100, 1, ?)`, + ) + .bind( + backupId, + incarnation === 0 ? null : createPlatformId(), + sandboxId, + incarnation, + incarnation === 0 ? backupId : createPlatformId(), + workspaceSessionId, + ) + .run(); + } + + const stagedActualIds = [ + createPlatformId(), + createPlatformId(), + ]; + for (const [incarnation, actualBackupId] of [ + [1, stagedActualIds[0]], + [2, stagedActualIds[1]], + ] as const) { + await database + .prepare( + `INSERT INTO sandbox_backup_staging ( + actual_backup_id, claim_owner, created_at, dir, id, operation_id, sandbox_id, + sandbox_incarnation, ttl_seconds, updated_at, updates_subject_backup, + workspace_session_id + ) VALUES (?, 'owner', 1, '/workspace/shared', ?, ?, ?, ?, 100, 1, 0, ?)`, + ) + .bind( + actualBackupId, + createPlatformId(), + createPlatformId(), + SANDBOX_ID, + incarnation, + SESSION_ID, + ) + .run(); + } + + const revoked = await revokeSandboxBackupsForSessionDelete(database, { + cwd: "/workspace/shared", + operationId: OPERATION_ID, + sandboxId: SANDBOX_ID, + sessionId: SESSION_ID, + }); + expect(new Set(revoked)).toEqual( + new Set([exactDeletedPriorId, exactDeletedCurrentId, stagedActualIds[0], stagedActualIds[1]]), + ); + await expect( + database + .prepare( + "SELECT COUNT(*) AS count FROM sandbox_backup_staging WHERE workspace_session_id = ?", + ) + .bind(SESSION_ID) + .first(), + ).resolves.toEqual({ count: 0 }); + const rows = await database + .prepare("SELECT id, status FROM sandbox_backup ORDER BY id") + .all<{ id: string; status: string }>(); + expect(rows.results).toEqual( + [ + { id: exactDeletedCurrentId, status: "pruned" }, + { id: exactDeletedPriorId, status: "pruned" }, + { id: legacyDeletedSandboxId, status: "ready" }, + { id: legacyOtherSandboxId, status: "ready" }, + { id: otherSandboxBackupId, status: "ready" }, + { id: sameSandboxOtherSessionBackupId, status: "ready" }, + ].toSorted((left, right) => left.id.localeCompare(right.id)), + ); +}); diff --git a/apps/api/tests/sandbox-conversation-session.test.ts b/apps/api/tests/sandbox-conversation-session.test.ts index cfda4c5d..cadce3f1 100644 --- a/apps/api/tests/sandbox-conversation-session.test.ts +++ b/apps/api/tests/sandbox-conversation-session.test.ts @@ -11,8 +11,14 @@ import type { SandboxHandle, } from "../src/modules/runtime/infrastructure/sandbox-handles"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; -import { PublicApiMemoryFileBucket } from "./helpers/public-api-http-test-fixture"; -import { SqliteD1Database } from "./helpers/sqlite-d1"; +import { + PUBLIC_API_TEST_IDS, + PublicApiMemoryFileBucket, + createPublicHttpContractDatabase, + insertActiveSandboxSessionFixture, + insertOwnerSession, +} from "./helpers/public-api-http-test-fixture"; +import type { SqliteD1Database } from "./helpers/sqlite-d1"; mock.module("@cloudflare/sandbox", () => ({ getSandbox: () => { @@ -20,18 +26,25 @@ mock.module("@cloudflare/sandbox", () => ({ }, })); -const { closeIdleCattleConversationSession, ensureSandboxConversationSession } = +const { + closeIdleCattleConversationSession, + ensureSandboxConversationSession, + repairPendingSandboxConversationSessionCleanups, +} = await import("../src/modules/runtime/infrastructure/sandbox-session/sandbox-conversation-session.service"); const ORIGIN = { - callerUserId: "01J00000000000000000000001", + callerUserId: PUBLIC_API_TEST_IDS.ownerAccount, entrypoint: "api", - executionOwnerUserId: "01J00000000000000000000001", + executionOwnerUserId: PUBLIC_API_TEST_IDS.ownerAccount, type: "agent", } as const; const CLOUDFLARE_BACKUP_ID = "550e8400-e29b-41d4-a716-446655440000"; const STORED_BACKUP_ID = encodeSandboxBackupIdForStorage(CLOUDFLARE_BACKUP_ID); const ARTIFACT_SHA256 = "a".repeat(64); +const SANDBOX_SESSION_ID = "01J00000000000000000000001"; +const SANDBOX_INCARNATION = 1; +const SESSION_ID = PUBLIC_API_TEST_IDS.ownerSession; function commandResult(): RuntimeCommandResultHandle { return { @@ -51,80 +64,23 @@ function failedCommandResult(): RuntimeCommandResultHandle { }; } -function createConversationSessionDatabase(kind: AgentKind = "pet"): SqliteD1Database { - const database = new SqliteD1Database(); - - database.execute(` - CREATE TABLE sandbox ( - id text PRIMARY KEY NOT NULL, - inactive_deadline_at integer, - kind text NOT NULL, - status text DEFAULT 'active' NOT NULL, - status_changed_at integer DEFAULT 0 NOT NULL, - status_event text DEFAULT 'runtime_subject.active' NOT NULL, - status_operation_id text, - status_seq integer DEFAULT 0 NOT NULL, - status_source text DEFAULT 'system' NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE sandbox_session ( - cloudflare_session_id text NOT NULL, - created_at integer NOT NULL, - cwd text NOT NULL, - origin_json text NOT NULL, - sandbox_id text NOT NULL, - session_id text PRIMARY KEY NOT NULL, - status text NOT NULL, - updated_at integer NOT NULL - ); - - CREATE TABLE sandbox_backup ( - created_at integer NOT NULL, - dir text NOT NULL, - id text PRIMARY KEY NOT NULL, - sandbox_id text NOT NULL, - session_run_id text, - status text NOT NULL - ); - - CREATE TABLE session ( - agent_id text, - id text PRIMARY KEY NOT NULL, - kind text DEFAULT '${kind}' NOT NULL, - last_run_id text, - workspace_checkpoint_required integer DEFAULT 0 NOT NULL - ); +async function createConversationSessionDatabase( + kind: AgentKind = "pet", +): Promise { + const database = await createPublicHttpContractDatabase(); - CREATE TABLE driver_instance ( - id text PRIMARY KEY NOT NULL, - sandbox_id text NOT NULL - ); - - CREATE TABLE session_run ( - driver_instance_id text, - id text PRIMARY KEY NOT NULL, - status text NOT NULL - ); - - CREATE TABLE file_record ( - id text PRIMARY KEY NOT NULL, - created_at integer NOT NULL, - name text NOT NULL, - object_key text NOT NULL, - parent_path text NOT NULL, - scope_id text, - scope_kind text NOT NULL, - session_kind text, - size integer NOT NULL, - status text NOT NULL - ); - `); - - database.execute(` - INSERT INTO sandbox (id, inactive_deadline_at, kind, updated_at) - VALUES ('01J0000000000000000000000D', 123, '${kind}', 1); - `); + await insertOwnerSession(database); + await database.prepare("UPDATE session SET kind = ? WHERE id = ?").bind(kind, SESSION_ID).run(); + await insertActiveSandboxSessionFixture(database, { + inactiveDeadlineAt: 123, + kind, + ownerAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxSessionId: SANDBOX_SESSION_ID, + sessionId: SESSION_ID, + timestampMs: 1, + }); + await database.prepare("DELETE FROM sandbox_session WHERE session_id = ?").bind(SESSION_ID).run(); return database; } @@ -135,10 +91,6 @@ async function insertConversationSession( readonly status: SandboxSessionStatus; }, ): Promise { - await database - .prepare("INSERT INTO session (agent_id, id) VALUES (?, ?) ON CONFLICT (id) DO NOTHING") - .bind(null, "session-1") - .run(); await database .prepare( ` @@ -148,20 +100,22 @@ async function insertConversationSession( cwd, origin_json, sandbox_id, + sandbox_incarnation, session_id, status, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( - "01J00000000000000000000001", + SANDBOX_SESSION_ID, 1, - "/workspace/se/session-1", + `/workspace/se/${SESSION_ID}`, JSON.stringify(ORIGIN), - "01J0000000000000000000000D", - "session-1", + PUBLIC_API_TEST_IDS.sandbox, + SANDBOX_INCARNATION, + SESSION_ID, input.status, 1, ) @@ -170,7 +124,11 @@ async function insertConversationSession( async function insertConversationBackup( database: D1Database, - input: { createdAt?: number; dir?: string } = {}, + input: { + createdAt?: number; + dir?: string; + workspaceSessionId?: string | null; + } = {}, ): Promise { await database .prepare( @@ -179,18 +137,32 @@ async function insertConversationBackup( created_at, dir, id, + keep, + operation_id, sandbox_id, - status + sandbox_incarnation, + staging_id, + status, + ttl_seconds, + updated_at, + workspace_session_id ) - VALUES (?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( input.createdAt ?? 1, - input.dir ?? "/workspace/se/session-1", + input.dir ?? `/workspace/se/${SESSION_ID}`, + STORED_BACKUP_ID, + 0, + PUBLIC_API_TEST_IDS.operation, + PUBLIC_API_TEST_IDS.sandbox, + SANDBOX_INCARNATION, STORED_BACKUP_ID, - "01J0000000000000000000000D", "ready", + 86_400, + input.createdAt ?? 1, + input.workspaceSessionId === undefined ? SESSION_ID : input.workspaceSessionId, ) .run(); } @@ -208,7 +180,7 @@ async function readConversationSession(database: D1Database): Promise<{ WHERE session_id = ? `, ) - .bind("session-1") + .bind(SESSION_ID) .first<{ cloudflare_session_id: string; cwd: string; @@ -225,7 +197,7 @@ async function readConversationSession(database: D1Database): Promise<{ async function readInactiveDeadline(database: D1Database): Promise { return database .prepare("SELECT inactive_deadline_at FROM sandbox WHERE id = ?") - .bind("01J0000000000000000000000D") + .bind(PUBLIC_API_TEST_IDS.sandbox) .first("inactive_deadline_at"); } @@ -325,25 +297,26 @@ async function setWorkspaceCheckpointRequired( ): Promise { await database .prepare("UPDATE session SET workspace_checkpoint_required = ? WHERE id = ?") - .bind(required ? 1 : 0, "session-1") + .bind(required ? 1 : 0, SESSION_ID) .run(); } function createInput(sandbox: SandboxHandle, kind: AgentKind = "pet") { return { - agentId: "01J00000000000000000000009", + agentId: PUBLIC_API_TEST_IDS.agent, kind, mountSessionResources: false, origin: ORIGIN, sandbox, - sandboxId: "01J0000000000000000000000D", - sessionId: "session-1", + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxIncarnation: SANDBOX_INCARNATION, + sessionId: SESSION_ID, }; } describe("ensureSandboxConversationSession", () => { test("reuses an active session without preparing directories", async () => { - const database = createConversationSessionDatabase(); + const database = await createConversationSessionDatabase(); await insertConversationSession(database, { status: "active" }); const sandbox = createSandbox(); @@ -352,17 +325,17 @@ describe("ensureSandboxConversationSession", () => { createInput(sandbox), ); - expect(result.sandboxSessionId).toBe("01J00000000000000000000001"); + expect(result.sandboxSessionId).toBe(SANDBOX_SESSION_ID); await expect(readConversationSession(database)).resolves.toEqual({ - cloudflare_session_id: "01J00000000000000000000001", - cwd: "/workspace/se/session-1", + cloudflare_session_id: SANDBOX_SESSION_ID, + cwd: `/workspace/se/${SESSION_ID}`, status: "active", }); await expect(readInactiveDeadline(database)).resolves.toBe(123); }); test("creates a missing conversation session record", async () => { - const database = createConversationSessionDatabase(); + const database = await createConversationSessionDatabase(); const sandbox = createSandbox(); const result = await ensureSandboxConversationSession( @@ -379,7 +352,7 @@ describe("ensureSandboxConversationSession", () => { }); test("arms an idle deadline when a legacy Pet session has none", async () => { - const database = createConversationSessionDatabase(); + const database = await createConversationSessionDatabase(); database.execute("UPDATE sandbox SET inactive_deadline_at = NULL"); const sandbox = createSandbox(); const startedAt = Date.now(); @@ -392,7 +365,7 @@ describe("ensureSandboxConversationSession", () => { }); test("continues a warm closed cattle session with a new execution session id", async () => { - const database = createConversationSessionDatabase("cattle"); + const database = await createConversationSessionDatabase("cattle"); await insertConversationSession(database, { status: "closed" }); const sandbox = createSandbox(); @@ -401,7 +374,7 @@ describe("ensureSandboxConversationSession", () => { createInput(sandbox, "cattle"), ); - expect(result.sandboxSessionId).not.toBe("01J00000000000000000000001"); + expect(result.sandboxSessionId).not.toBe(SANDBOX_SESSION_ID); expect(isPlatformId(result.sandboxSessionId)).toBe(true); await expect(readConversationSession(database)).resolves.toMatchObject({ @@ -412,7 +385,7 @@ describe("ensureSandboxConversationSession", () => { }); test("restores a cold cattle session from a 20-day-old committed checkpoint", async () => { - const database = createConversationSessionDatabase("cattle"); + const database = await createConversationSessionDatabase("cattle"); await insertConversationSession(database, { status: "closed" }); await setWorkspaceCheckpointRequired(database, true); await insertConversationBackup(database, { @@ -432,13 +405,13 @@ describe("ensureSandboxConversationSession", () => { ); expect(restoredBackup).toEqual({ - dir: "/workspace/se/session-1", + dir: `/workspace/se/${SESSION_ID}`, id: CLOUDFLARE_BACKUP_ID, }); }); test("fails cold cattle continuation when its exact Thread checkpoint is missing", async () => { - const database = createConversationSessionDatabase("cattle"); + const database = await createConversationSessionDatabase("cattle"); await insertConversationSession(database, { status: "closed" }); await setWorkspaceCheckpointRequired(database, true); await insertConversationBackup(database, { dir: "/workspace/se/another-session" }); @@ -449,8 +422,22 @@ describe("ensureSandboxConversationSession", () => { ).rejects.toThrow("has no committed workspace checkpoint"); }); + test("does not consume a legacy checkpoint without exact Thread ownership", async () => { + const database = await createConversationSessionDatabase("cattle"); + await insertConversationSession(database, { status: "closed" }); + await setWorkspaceCheckpointRequired(database, true); + await insertConversationBackup(database, { workspaceSessionId: null }); + + await expect( + ensureSandboxConversationSession( + createBindings(database), + createInput(createSandbox({ cwdHasContent: false }), "cattle"), + ), + ).rejects.toThrow("has no committed workspace checkpoint"); + }); + test("reports an actionable error for a corrupt cattle checkpoint", async () => { - const database = createConversationSessionDatabase("cattle"); + const database = await createConversationSessionDatabase("cattle"); await insertConversationSession(database, { status: "closed" }); await setWorkspaceCheckpointRequired(database, true); await insertConversationBackup(database); @@ -464,17 +451,28 @@ describe("ensureSandboxConversationSession", () => { ).rejects.toThrow("workspace checkpoint could not be restored. Retry the continuation"); }); - test("restores recorded artifacts for a pre-rollout cattle Thread", async () => { - const database = createConversationSessionDatabase("cattle"); + test("restores backfilled artifacts for a pre-rollout cattle Thread", async () => { + const database = await createConversationSessionDatabase("cattle"); await insertConversationSession(database, { status: "closed" }); database.execute(` INSERT INTO file_record ( - id, created_at, name, object_key, parent_path, scope_id, scope_kind, - session_kind, size, status + committed, created_at, created_by_account_id, id, name, object_key, + owner_id, owner_kind, parent_path, path, purpose, runtime_event_seq, + scope_id, scope_kind, session_kind, size, status, updated_at, version ) VALUES ( - 'artifact-legacy', 1, 'legacy.txt', 'artifacts/legacy.txt', - 'runtime-output/outputs/legacy.txt/${ARTIFACT_SHA256}', 'session-1', 'session', - 'artifact', 14, 'ready' + 1, 1, '${PUBLIC_API_TEST_IDS.ownerAccount}', '${PUBLIC_API_TEST_IDS.file}', + 'legacy.txt', 'artifacts/legacy.txt', '${SESSION_ID}', 'session', + 'runtime-output/outputs/legacy.txt/${ARTIFACT_SHA256}', + 'session-artifacts/${PUBLIC_API_TEST_IDS.file}/legacy.txt', 'session_artifact', 0, + '${SESSION_ID}', 'session', 'artifact', 14, 'ready', 1, 1 + ); + + INSERT INTO session_artifact_head ( + file_id, runtime_event_seq, session_id, source_event_id, source_path, updated_at + ) VALUES ( + '${PUBLIC_API_TEST_IDS.file}', 0, '${SESSION_ID}', + 'legacy-file:${PUBLIC_API_TEST_IDS.file}', + 'outputs/legacy.txt', 1 ); `); const bucket = new PublicApiMemoryFileBucket(); @@ -486,15 +484,15 @@ describe("ensureSandboxConversationSession", () => { }); await ensureSandboxConversationSession( - createBindings(database, undefined, bucket as unknown as R2Bucket), + createBindings(database, undefined, bucket), createInput(sandbox, "cattle"), ); - expect(restoredPaths).toContain("/workspace/se/session-1/outputs/legacy.txt"); + expect(restoredPaths).toContain(`/workspace/se/${SESSION_ID}/outputs/legacy.txt`); }); test("continues a closed pet session through the stable restore path", async () => { - const database = createConversationSessionDatabase(); + const database = await createConversationSessionDatabase(); await insertConversationSession(database, { status: "closed" }); await insertConversationBackup(database); let restoredBackup: { readonly dir: string; readonly id: string } | null = null; @@ -510,13 +508,32 @@ describe("ensureSandboxConversationSession", () => { createInput(sandbox, "pet"), ); - expect(result.sandboxSessionId).toBe("01J00000000000000000000001"); + expect(result.sandboxSessionId).toBe(SANDBOX_SESSION_ID); expect(restoredBackup).toEqual({ - dir: "/workspace/se/session-1", + dir: `/workspace/se/${SESSION_ID}`, id: CLOUDFLARE_BACKUP_ID, }); await expect(readConversationSession(database)).resolves.toMatchObject({ - cloudflare_session_id: "01J00000000000000000000001", + cloudflare_session_id: SANDBOX_SESSION_ID, + status: "active", + }); + }); + + test("replaces a closed pet execution session for a fenced Run provisioning", async () => { + const database = await createConversationSessionDatabase(); + await insertConversationSession(database, { status: "closed" }); + await insertConversationBackup(database); + const sandbox = createSandbox({ cwdHasContent: false }); + + const result = await ensureSandboxConversationSession(createBindings(database), { + ...createInput(sandbox, "pet"), + replaceClosedExecutionSession: true, + }); + + expect(result.sandboxSessionId).not.toBe(SANDBOX_SESSION_ID); + expect(isPlatformId(result.sandboxSessionId)).toBe(true); + await expect(readConversationSession(database)).resolves.toMatchObject({ + cloudflare_session_id: result.sandboxSessionId, status: "active", }); }); @@ -524,26 +541,66 @@ describe("ensureSandboxConversationSession", () => { describe("closeIdleCattleConversationSession", () => { test("arms subject reclamation when the remote session is already absent", async () => { - const database = createConversationSessionDatabase("cattle"); + const database = await createConversationSessionDatabase("cattle"); await insertConversationSession(database, { status: "active" }); database.execute("UPDATE sandbox SET inactive_deadline_at = NULL"); - const sandboxSessionId = "01J00000000000000000000001"; const sandbox = createSandbox({ - deleteSessionError: new Error(`Session '${sandboxSessionId}' not found`), + deleteSessionError: new Error(`Session '${SANDBOX_SESSION_ID}' not found`), }); const startedAt = Date.now(); await expect( closeIdleCattleConversationSession(createBindings(database, sandbox), { idleSinceLte: 1, - sandboxId: "01J0000000000000000000000D", - sessionId: "session-1", + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sessionId: SESSION_ID, }), - ).rejects.toThrow(`Session '${sandboxSessionId}' not found`); + ).resolves.toBe(true); await expect(readConversationSession(database)).resolves.toMatchObject({ status: "closed" }); + await expect( + database + .prepare( + `SELECT event_type, family, source, visibility + FROM session_event + WHERE session_id = ? AND event_type = 'runtime.sandbox.updated'`, + ) + .bind(SESSION_ID) + .first(), + ).resolves.toEqual({ + event_type: "runtime.sandbox.updated", + family: "sandbox", + source: "system", + visibility: "owner_debug", + }); const deadline = await readInactiveDeadline(database); expect(deadline).toBeGreaterThanOrEqual(startedAt + 5 * 60_000); expect(deadline).toBeLessThanOrEqual(Date.now() + 5 * 60_000); }); + + test("repairs cleanup against its retired physical incarnation", async () => { + const database = await createConversationSessionDatabase("cattle"); + await insertConversationSession(database, { status: "active" }); + database.execute(` + UPDATE sandbox SET incarnation = 5; + UPDATE sandbox_session + SET cleanup_operation_id = '${PUBLIC_API_TEST_IDS.operation}', + sandbox_incarnation = 4, + status = 'cleanup_pending'; + `); + const physicalIds: string[] = []; + const sandbox = createSandbox(); + const bindings = { + DB: database, + runtimeSubjectHandleFactory: (physicalId: string) => { + physicalIds.push(physicalId); + return sandbox; + }, + } as ApiBindings; + + await expect(repairPendingSandboxConversationSessionCleanups(bindings, 10)).resolves.toBe(1); + + expect(physicalIds).toEqual([`${PUBLIC_API_TEST_IDS.sandbox}-i4`]); + await expect(readConversationSession(database)).resolves.toMatchObject({ status: "closed" }); + }); }); diff --git a/apps/api/tests/sandbox-runtime-incarnation.test.ts b/apps/api/tests/sandbox-runtime-incarnation.test.ts new file mode 100644 index 00000000..102ba344 --- /dev/null +++ b/apps/api/tests/sandbox-runtime-incarnation.test.ts @@ -0,0 +1,742 @@ +import { describe, expect, mock, spyOn, test } from "bun:test"; + +import { hashSandboxNetworkConstraints } from "../src/modules/runtime/domain/sandbox-network-constraints"; + +const FULL_NETWORK_CONSTRAINTS = { allowedHosts: [], networkPolicy: "full" } as const; +const FULL_NETWORK_CONSTRAINTS_HASH = await hashSandboxNetworkConstraints(FULL_NETWORK_CONSTRAINTS); + +class SandboxDelegateMock { + alarmBarrier: Promise | null = null; + backupCalls = 0; + backupOptions: unknown[] = []; + destroyBarrier: Promise | null = null; + destroyCalls = 0; + destroyError: Error | null = null; + enableInternet = true; + envVars: Record = {}; + getPlacementError: Error | null = null; + interceptHttps = false; + lastSessionKeys: PropertyKey[] = []; + mkdirCalls = 0; + placement: string | null | undefined = "placement-1"; + readonly plainDto = { nested: { value: "unchanged" }, success: true }; + processActionCalls = 0; + readFileBarrier: Promise | null = null; + readFileError: Error | null = null; + startProcessBarrier: Promise | null = null; + terminalBarrier: Promise | null = null; + terminalCalls = 0; + terminalError: Error | null = null; + streamOpenCalls = 0; + streamReadCalls = 0; + writeFileBarrier: Promise | null = null; + onAlarmStart: (() => void) | null = null; + onMkdir: (() => void) | null = null; + onReadFileStart: (() => void) | null = null; + onStartProcessStart: (() => void) | null = null; + onTerminalStart: (() => void) | null = null; + onWriteFileStart: (() => void) | null = null; + readonly files = new Map(); + + async alarm(): Promise { + this.onAlarmStart?.(); + await this.alarmBarrier; + } + + async createBackup(options: unknown): Promise<{ dir: string; id: string }> { + this.backupCalls += 1; + this.backupOptions.push(options); + return { + dir: Reflect.get(options as object, "dir") as string, + id: "550e8400-e29b-41d4-a716-446655440001", + }; + } + + async createSession(options: { id: string }) { + void options; + const session = { + execStream: () => this.execStream(), + getProcess: () => this.getProcess(), + id: "user-session", + listProcesses: () => this.listProcesses(), + readFileStream: () => this.readFileStream(), + readFile: async (path: string) => { + const content = this.files.get(path); + if (content === undefined) { + throw new Error("missing file"); + } + return { content }; + }, + startProcess: ( + command: string, + startOptions?: Parameters[1], + ) => this.startProcess(command, startOptions), + streamProcessLogs: () => this.streamProcessLogs(), + terminal: async () => { + this.onTerminalStart?.(); + await this.terminalBarrier; + if (this.terminalError !== null) { + throw this.terminalError; + } + this.terminalCalls += 1; + return new Response("terminal"); + }, + upstreamOnlyMethod: async () => "preserved", + }; + this.lastSessionKeys = Reflect.ownKeys(session); + return session; + } + + async deleteSession(): Promise {} + + async destroy(): Promise { + this.destroyCalls += 1; + await this.destroyBarrier; + if (this.destroyError !== null) { + throw this.destroyError; + } + } + + async exec(): Promise { + return this.plainDto; + } + + async exists(path: string) { + return { + exists: this.files.has(path), + path, + success: true, + timestamp: new Date(0).toISOString(), + }; + } + + async execStream(): Promise> { + this.streamOpenCalls += 1; + return this.streamHandle(); + } + + async getContainerPlacementId(): Promise { + if (this.getPlacementError !== null) { + throw this.getPlacementError; + } + return this.placement; + } + + async getProcess(): Promise> { + return this.processHandle(); + } + + async getSession(): ReturnType { + return this.createSession({ id: "user-session" }); + } + + async mkdir(): Promise { + this.mkdirCalls += 1; + this.onMkdir?.(); + } + + async listProcesses(): Promise>> { + return [this.processHandle()]; + } + + processHandle() { + const action = async () => { + this.processActionCalls += 1; + }; + return { + command: "driver", + endTime: undefined, + exitCode: undefined, + getLogs: action, + getStatus: action, + id: "process-1", + kill: action, + pid: 1, + sessionId: "session-1", + startTime: new Date(0), + status: "running", + waitForExit: action, + waitForLog: action, + waitForPort: action, + }; + } + + async startProcess( + _command: string, + options?: { onStart?: (process: ReturnType) => void }, + ): Promise> { + const process = this.processHandle(); + options?.onStart?.(process); + this.onStartProcessStart?.(); + await this.startProcessBarrier; + return process; + } + + streamHandle() { + return { + [Symbol.asyncIterator]() { + return this; + }, + next: async () => { + this.streamReadCalls += 1; + return { done: true as const, value: undefined }; + }, + }; + } + + async readFile(path: string): Promise<{ content: string }> { + this.onReadFileStart?.(); + await this.readFileBarrier; + if (this.readFileError !== null) { + throw this.readFileError; + } + const content = this.files.get(path); + if (content === undefined) { + throw new Error("missing file"); + } + return { content }; + } + + async readFileStream(): Promise> { + this.streamOpenCalls += 1; + return this.streamHandle(); + } + + async setAllowedHosts(): Promise {} + + async setKeepAlive(): Promise {} + + async streamProcessLogs(): Promise> { + this.streamOpenCalls += 1; + return this.streamHandle(); + } + + async writeFile(path: string, content: string): Promise { + this.onWriteFileStart?.(); + await this.writeFileBarrier; + this.files.set(path, content); + } +} + +let delegate: SandboxDelegateMock; +let abortCalls = 0; +let abortRetryAlarm = false; +let alarmWrites = 0; + +function registerDelegate(value: SandboxDelegateMock): void { + delegate = value; +} + +mock.module("cloudflare:workers", () => ({ + DurableObject: class { + constructor( + protected readonly ctx: unknown, + protected readonly env: unknown, + ) {} + }, +})); + +mock.module("@cloudflare/sandbox", () => ({ + Sandbox: class extends SandboxDelegateMock { + constructor() { + super(); + registerDelegate(this); + } + }, +})); + +const { Sandbox } = await import("../src/adapters/durable-objects/sandbox.do"); + +function createBarrier(): { readonly promise: Promise; readonly release: () => void } { + let release = () => {}; + const promise = new Promise((resolve) => { + release = resolve; + }); + return { promise, release }; +} + +function createSandbox(): Sandbox { + const values = new Map(); + const ctx = { + abort: (_reason?: string, options?: { retryAlarm?: boolean }) => { + abortCalls += 1; + abortRetryAlarm = options?.retryAlarm === true; + throw new Error("durable object aborted"); + }, + blockConcurrencyWhile: (action: () => Promise) => action(), + container: { running: true }, + storage: { + get: async (key: string) => values.get(key) as T | undefined, + put: async (key: string | Record, value?: unknown) => { + if (typeof key === "string") { + values.set(key, value); + return; + } + for (const [entryKey, entryValue] of Object.entries(key)) { + values.set(entryKey, entryValue); + } + }, + setAlarm: async () => { + alarmWrites += 1; + }, + }, + }; + abortCalls = 0; + abortRetryAlarm = false; + alarmWrites = 0; + return new Sandbox(ctx, { SANDBOX_FILE_BUCKET_LOCAL: "false" }); +} + +async function createReadySandbox(): Promise { + const sandbox = createSandbox(); + await sandbox.activateRuntimeSubjectIncarnation(4, FULL_NETWORK_CONSTRAINTS_HASH); + await sandbox.markRuntimeSubjectIncarnationReady(4, FULL_NETWORK_CONSTRAINTS_HASH); + return sandbox; +} + +describe("sandbox runtime incarnation health", () => { + test("fails closed before SDK backup creation when legacy persistent auth exists", async () => { + const sandbox = await createReadySandbox(); + const authPath = "/workspace/se/session/.state/openai-runtime/auth.json"; + delegate.files.set(authPath, "secret"); + const handle = sandbox as unknown as { + createRuntimeSubjectBackup( + incarnation: number, + options: { + dir: string; + excludes: string[]; + forbiddenPaths: string[]; + name: string; + }, + ): Promise<{ dir: string; id: string }>; + }; + const options = { + dir: "/workspace/se/session", + excludes: [".mosoo-session-files-session"], + forbiddenPaths: [authPath], + name: "mosoo:runtime-backup:v1:01J0000000000000000000000B", + }; + + await expect(handle.createRuntimeSubjectBackup(4, options)).rejects.toThrow( + "legacy persistent credentials", + ); + expect(delegate.backupCalls).toBe(0); + expect(delegate.mkdirCalls).toBe(0); + + delegate.files.delete(authPath); + await expect(handle.createRuntimeSubjectBackup(4, options)).resolves.toEqual({ + dir: options.dir, + id: "550e8400-e29b-41d4-a716-446655440001", + }); + expect(delegate.mkdirCalls).toBe(1); + expect(delegate.backupCalls).toBe(1); + expect(delegate.backupOptions).toEqual([ + { + dir: options.dir, + excludes: options.excludes, + name: options.name, + }, + ]); + }); + + test("keeps incarnation and network policy as one immutable identity", async () => { + const sandbox = createSandbox(); + await sandbox.activateRuntimeSubjectIncarnation(4, FULL_NETWORK_CONSTRAINTS_HASH); + await expect( + sandbox.activateRuntimeSubjectIncarnation(4, FULL_NETWORK_CONSTRAINTS_HASH), + ).resolves.toBeUndefined(); + await expect(sandbox.activateRuntimeSubjectIncarnation(4, "0".repeat(64))).rejects.toThrow( + "identity does not match", + ); + }); + + test("rejects network mutation that differs from the admitted identity", async () => { + const sandbox = createSandbox(); + await sandbox.activateRuntimeSubjectIncarnation(4, FULL_NETWORK_CONSTRAINTS_HASH); + + await expect( + sandbox.configureNetworkConstraints({ + allowedHosts: ["example.com"], + networkPolicy: "limited", + }), + ).rejects.toThrow("do not match the admitted incarnation"); + }); + + test("keeps a ready incarnation when the placement handshake is unavailable", async () => { + const sandbox = await createReadySandbox(); + delegate.getPlacementError = new Error("control plane unavailable"); + + await expect( + (sandbox as unknown as { mkdir(path: string): Promise }).mkdir("/workspace"), + ).rejects.toThrow("health is unknown"); + expect(delegate.destroyCalls).toBe(0); + }); + + test("retires a ready incarnation after a refreshed placement mismatch", async () => { + const sandbox = await createReadySandbox(); + delegate.placement = "placement-2"; + + await expect( + (sandbox as unknown as { mkdir(path: string): Promise }).mkdir("/workspace"), + ).rejects.toThrow("replaced or stopped"); + expect(delegate.destroyCalls).toBe(1); + await expect( + sandbox.inspectRuntimeSubjectIncarnation(4, FULL_NETWORK_CONSTRAINTS_HASH), + ).resolves.toEqual({ + kind: "retired", + }); + }); + + test("refreshes cached placement before running a business mutation", async () => { + const sandbox = await createReadySandbox(); + delegate.onReadFileStart = () => { + delegate.placement = "placement-2"; + }; + + await expect( + (sandbox as unknown as { mkdir(path: string): Promise }).mkdir("/workspace"), + ).rejects.toThrow("replaced or stopped"); + expect(delegate.mkdirCalls).toBe(0); + expect(delegate.destroyCalls).toBe(1); + }); + + test("retires when a missing sentinel refreshes the cached placement", async () => { + const sandbox = await createReadySandbox(); + delegate.onReadFileStart = () => { + delegate.files.delete("/tmp/.mosoo-runtime-subject-incarnation"); + delegate.placement = "placement-2"; + }; + + await expect( + (sandbox as unknown as { mkdir(path: string): Promise }).mkdir("/workspace"), + ).rejects.toThrow("replaced or stopped"); + expect(delegate.mkdirCalls).toBe(0); + expect(delegate.destroyCalls).toBe(1); + }); + + test("retires a structured missing sentinel without placement ids", async () => { + const sandbox = createSandbox(); + await sandbox.activateRuntimeSubjectIncarnation(4, FULL_NETWORK_CONSTRAINTS_HASH); + delegate.placement = null; + await sandbox.markRuntimeSubjectIncarnationReady(4, FULL_NETWORK_CONSTRAINTS_HASH); + delegate.readFileError = Object.assign(new Error("missing sentinel"), { + code: "FILE_NOT_FOUND", + }); + + await expect( + (sandbox as unknown as { mkdir(path: string): Promise }).mkdir("/workspace"), + ).rejects.toThrow("replaced or stopped"); + expect(delegate.mkdirCalls).toBe(0); + expect(delegate.destroyCalls).toBe(1); + }); + + test("re-destroys a retired incarnation after a late local sentinel read", async () => { + const sandbox = createSandbox(); + await sandbox.activateRuntimeSubjectIncarnation(4, FULL_NETWORK_CONSTRAINTS_HASH); + await sandbox.configureNetworkConstraints(FULL_NETWORK_CONSTRAINTS); + delegate.placement = null; + await sandbox.markRuntimeSubjectIncarnationReady(4, FULL_NETWORK_CONSTRAINTS_HASH); + const read = createBarrier(); + const started = createBarrier(); + delegate.readFileBarrier = read.promise; + delegate.onReadFileStart = started.release; + + const inspection = sandbox.inspectRuntimeSubjectIncarnation(4, FULL_NETWORK_CONSTRAINTS_HASH); + await started.promise; + await sandbox.destroyRuntimeSubjectIncarnation(4); + await sandbox.alarm(); + expect(delegate.destroyCalls).toBe(2); + + read.release(); + await expect(inspection).resolves.toEqual({ kind: "retired" }); + expect(delegate.destroyCalls).toBe(3); + }); + + test("rejects and retires a mutation whose container placement changes mid-RPC", async () => { + const sandbox = await createReadySandbox(); + delegate.onMkdir = () => { + delegate.placement = "placement-2"; + }; + + await expect( + (sandbox as unknown as { mkdir(path: string): Promise }).mkdir("/workspace"), + ).rejects.toThrow("replaced or stopped"); + expect(delegate.mkdirCalls).toBe(1); + expect(delegate.destroyCalls).toBe(1); + }); + + test("re-destroys a retired incarnation after a late readiness write", async () => { + const sandbox = createSandbox(); + await sandbox.activateRuntimeSubjectIncarnation(4, FULL_NETWORK_CONSTRAINTS_HASH); + await sandbox.configureNetworkConstraints(FULL_NETWORK_CONSTRAINTS); + const write = createBarrier(); + const started = createBarrier(); + delegate.writeFileBarrier = write.promise; + delegate.onWriteFileStart = started.release; + + const readiness = sandbox.markRuntimeSubjectIncarnationReady(4, FULL_NETWORK_CONSTRAINTS_HASH); + await started.promise; + await sandbox.destroyRuntimeSubjectIncarnation(4); + await sandbox.alarm(); + expect(delegate.destroyCalls).toBe(2); + + write.release(); + await expect(readiness).rejects.toThrow("retired"); + expect(delegate.destroyCalls).toBe(3); + }); + + test("fences terminal sessions returned by create and get", async () => { + const sandbox = await createReadySandbox(); + const handle = sandbox as unknown as { + createSession(options: { id: string }): Promise<{ + terminal(request: Request): Promise; + }>; + getSession(id: string): Promise<{ terminal(request: Request): Promise }>; + }; + const created = await handle.createSession({ id: "created-session" }); + const fetched = await handle.getSession("created-session"); + + const byPropertyKey = (left: PropertyKey, right: PropertyKey) => + String(left).localeCompare(String(right)); + expect(Reflect.ownKeys(created).toSorted(byPropertyKey)).toEqual( + delegate.lastSessionKeys.toSorted(byPropertyKey), + ); + + await expect(created.terminal(new Request("https://sandbox.test"))).resolves.toBeInstanceOf( + Response, + ); + await expect(fetched.terminal(new Request("https://sandbox.test"))).resolves.toBeInstanceOf( + Response, + ); + expect(delegate.terminalCalls).toBe(2); + }); + + test("re-destroys after a late terminal RPC resumes past the consumed alarm", async () => { + const sandbox = await createReadySandbox(); + const handle = sandbox as unknown as { + createSession(options: { id: string }): Promise<{ + terminal(request: Request): Promise; + }>; + }; + const session = await handle.createSession({ id: "terminal-session" }); + const terminal = createBarrier(); + const started = createBarrier(); + delegate.terminalBarrier = terminal.promise; + delegate.onTerminalStart = started.release; + + const response = session.terminal(new Request("https://sandbox.test")); + await started.promise; + await sandbox.destroyRuntimeSubjectIncarnation(4); + await sandbox.alarm(); + expect(delegate.destroyCalls).toBe(2); + + terminal.release(); + await expect(response).rejects.toThrow("retired during the operation"); + expect(delegate.destroyCalls).toBe(3); + }); + + test.each([ + ["successful cleanup", null], + ["failed cleanup", new Error("destroy failed")], + ] as const)("preserves a task failure after retirement with %s", async (_, destroyFailure) => { + const sandbox = await createReadySandbox(); + const handle = sandbox as unknown as { + createSession(options: { id: string }): Promise<{ + terminal(request: Request): Promise; + }>; + }; + const session = await handle.createSession({ id: "terminal-session" }); + const terminal = createBarrier(); + const started = createBarrier(); + const taskFailure = new Error("terminal failed"); + delegate.terminalBarrier = terminal.promise; + delegate.terminalError = taskFailure; + delegate.onTerminalStart = started.release; + + const response = session.terminal(new Request("https://sandbox.test")); + await started.promise; + await sandbox.destroyRuntimeSubjectIncarnation(4); + delegate.destroyError = destroyFailure; + terminal.release(); + + try { + await response; + throw new Error("The retired RPC unexpectedly succeeded."); + } catch (error) { + if (destroyFailure === null) { + expect(error).toBe(taskFailure); + } else { + if (!(error instanceof AggregateError)) { + throw error; + } + expect(error.errors).toEqual([taskFailure, destroyFailure]); + } + } + expect(delegate.destroyCalls).toBe(2); + }); + + test("rejects a late process start after the next incarnation retires its owner", async () => { + const sandbox = await createReadySandbox(); + const handle = sandbox as unknown as { + startProcess(command: string): Promise>; + }; + const process = createBarrier(); + const started = createBarrier(); + delegate.startProcessBarrier = process.promise; + delegate.onStartProcessStart = started.release; + + const lateStart = handle.startProcess("driver"); + await started.promise; + await sandbox.destroyRuntimeSubjectIncarnation(4); + await sandbox.alarm(); + expect(delegate.destroyCalls).toBe(2); + + process.release(); + await expect(lateStart).rejects.toThrow("retired during the operation"); + expect(delegate.destroyCalls).toBe(3); + }); + + test("fences every Process source and method after retirement", async () => { + const sandbox = await createReadySandbox(); + type GuardedProcess = ReturnType; + type GuardedSession = { + getProcess(id: string): Promise; + listProcesses(): Promise; + startProcess( + command: string, + options?: { onStart(process: GuardedProcess): void }, + ): Promise; + }; + let callbackProcess: GuardedProcess | null = null; + let sessionCallbackProcess: GuardedProcess | null = null; + const handle = sandbox as unknown as { + createSession(options: { id: string }): Promise; + getProcess(id: string): Promise; + listProcesses(): Promise; + startProcess( + command: string, + options?: { onStart(process: GuardedProcess): void }, + ): Promise; + }; + const started = await handle.startProcess("driver", { + onStart: (processHandle) => { + callbackProcess = processHandle; + }, + }); + expect(Object.keys(started).toSorted()).toEqual( + Object.keys(delegate.processHandle()).toSorted(), + ); + const fetched = await handle.getProcess("process-1"); + const [listed] = await handle.listProcesses(); + const session = await handle.createSession({ id: "process-session" }); + const sessionStarted = await session.startProcess("driver", { + onStart: (processHandle) => { + sessionCallbackProcess = processHandle; + }, + }); + const sessionFetched = await session.getProcess("process-1"); + const [sessionListed] = await session.listProcesses(); + if ( + callbackProcess === null || + sessionCallbackProcess === null || + listed === undefined || + sessionListed === undefined + ) { + throw new Error("Every Process source must return a handle."); + } + const processes = [ + started, + fetched, + listed, + callbackProcess, + sessionStarted, + sessionFetched, + sessionListed, + sessionCallbackProcess, + ]; + const methods = [ + ["getLogs", []], + ["getStatus", []], + ["kill", []], + ["waitForExit", []], + ["waitForLog", ["ready"]], + ["waitForPort", [8080]], + ] as const; + + await sandbox.destroyRuntimeSubjectIncarnation(4); + await sandbox.alarm(); + for (const process of processes) { + for (const [method, args] of methods) { + await expect(Reflect.apply(process[method], process, args)).rejects.toThrow("retired"); + } + } + expect(delegate.processActionCalls).toBe(0); + expect(delegate.destroyCalls).toBe(2 + processes.length * methods.length); + }); + + test("leaves DTOs and established streams untouched without reopening them", async () => { + const sandbox = await createReadySandbox(); + type Stream = ReturnType; + const handle = sandbox as unknown as { + createSession(options: { id: string }): Promise<{ execStream(): Promise }>; + exec(): Promise; + }; + const dto = await handle.exec(); + const session = await handle.createSession({ id: "stream-session" }); + const stream = await session.execStream(); + + expect(dto).toBe(delegate.plainDto); + expect(dto.nested).toBe(delegate.plainDto.nested); + expect(delegate.streamOpenCalls).toBe(1); + + await sandbox.destroyRuntimeSubjectIncarnation(4); + await sandbox.alarm(); + await expect(session.execStream()).rejects.toThrow("retired"); + expect(delegate.streamOpenCalls).toBe(1); + await expect(stream.next()).resolves.toEqual({ done: true, value: undefined }); + expect(delegate.streamReadCalls).toBe(1); + }); + + test("re-destroys after a late delegate alarm resumes", async () => { + const sandbox = await createReadySandbox(); + const alarm = createBarrier(); + const started = createBarrier(); + delegate.alarmBarrier = alarm.promise; + delegate.onAlarmStart = started.release; + + const lateAlarm = sandbox.alarm(); + await started.promise; + await sandbox.destroyRuntimeSubjectIncarnation(4); + expect(delegate.destroyCalls).toBe(1); + + alarm.release(); + await lateAlarm; + expect(delegate.destroyCalls).toBe(2); + }); + + test("aborts a wedged retired destroy while preserving its retry alarm", async () => { + const sandbox = await createReadySandbox(); + delegate.destroyBarrier = new Promise(() => {}); + const timeout = spyOn(globalThis, "setTimeout").mockImplementation((( + callback: TimerHandler, + ) => { + if (typeof callback === "function") { + queueMicrotask(callback); + } + return 1; + }) as typeof setTimeout); + + try { + await expect(sandbox.destroyRuntimeSubjectIncarnation(4)).rejects.toThrow( + "durable object aborted", + ); + } finally { + timeout.mockRestore(); + } + + expect(delegate.destroyCalls).toBe(1); + expect(alarmWrites).toBe(2); + expect(abortCalls).toBe(1); + expect(abortRetryAlarm).toBe(true); + }); +}); diff --git a/apps/api/tests/send-agent-session-events.test.ts b/apps/api/tests/send-agent-session-events.test.ts index 589cb054..8c3738ce 100644 --- a/apps/api/tests/send-agent-session-events.test.ts +++ b/apps/api/tests/send-agent-session-events.test.ts @@ -110,6 +110,10 @@ async function insertQueuedRunFixture( input.nowMs, ) .run(); + await database + .prepare("UPDATE session SET last_run_id = ?, status = ? WHERE id = ?") + .bind(input.runId, "RUNNING", PUBLIC_API_TEST_IDS.ownerSession) + .run(); } async function insertSessionFileRecord( @@ -531,6 +535,10 @@ describe("send agent session events", () => { nowMs, ) .run(); + await database + .prepare("UPDATE session SET last_run_id = ?, status = ? WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.run, "RUNNING", PUBLIC_API_TEST_IDS.ownerSession) + .run(); await database .prepare( ` diff --git a/apps/api/tests/session-event-stream-fold.test.ts b/apps/api/tests/session-event-stream-fold.test.ts index 2523d068..348d6319 100644 --- a/apps/api/tests/session-event-stream-fold.test.ts +++ b/apps/api/tests/session-event-stream-fold.test.ts @@ -2,100 +2,183 @@ import { describe, expect, test } from "bun:test"; import type { RuntimeEventId, SessionRunId } from "@mosoo/id"; -import { foldStreamedSessionEventRows } from "../src/modules/sessions/domain/session-event-stream-fold"; +import { + createMessageStreamLifecycle, + findLeftIncompleteSessionEventStreamKeys, + foldStreamedSessionEventRows, + reduceMessageStreamLifecycle, + resolveSealedMessageStream, +} from "../src/modules/sessions/domain/session-event-stream-fold"; import type { StreamFoldableSessionEventRow } from "../src/modules/sessions/domain/session-event-stream-fold"; const RUN_ID = "run-1" as SessionRunId; +interface TestSessionEventRow extends StreamFoldableSessionEventRow { + process_status: "available" | "error"; +} + function row(input: { content: string; eventType: string; id: string; + processStatus?: TestSessionEventRow["process_status"]; processType?: string; runId?: SessionRunId | null; seq: number; -}): StreamFoldableSessionEventRow { + streamId?: string | null; +}): TestSessionEventRow { return { content_text: input.content, ended_at: input.seq * 1000, event_type: input.eventType, id: input.id as RuntimeEventId, occurred_at: input.seq * 1000, + process_status: input.processStatus ?? "available", process_type: input.processType ?? "agent.message.delta", run_id: input.runId === undefined ? RUN_ID : input.runId, seq: input.seq, + stream_id: input.streamId === undefined ? "stream-1" : input.streamId, tokens: null, }; } describe("session event stream folding", () => { - test("folds a streamed assistant message into one row", () => { - const folded = foldStreamedSessionEventRows( - [ - row({ content: "Message updated.", eventType: "message.started", id: "m-start", seq: 1 }), - row({ content: "你", eventType: "message.delta", id: "m-1", seq: 2 }), - row({ content: "好", eventType: "message.delta", id: "m-2", seq: 3 }), - row({ content: "，世界", eventType: "message.delta", id: "m-3", seq: 4 }), - row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 5 }), - ], - { flushOpenStreams: true }, - ); - - expect(folded.openStreamRows).toEqual([]); - expect(folded.rows).toHaveLength(1); - expect(folded.rows[0]).toMatchObject({ - content_text: "你好，世界", - event_type: "message.completed", - id: "m-end", - occurred_at: 1000, - seq: 1, - }); + test.each([ + [[], false, false], + [["message.added"], true, false], + [["message.added", "message.completed"], true, true], + [["message.added", "message.completed", "message.delta"], true, false], + [["message.added", "message.completed", "message.started"], false, false], + [["message.added", "message.cancelled", "message.completed"], false, false], + [["message.added", "message.failed", "message.completed"], false, false], + [["message.started", "message.completed"], false, false], + ] as const)( + "reduces message lifecycle %j to authoritative=%s sealed=%s", + (eventTypes, authoritative, sealed) => { + let state = createMessageStreamLifecycle(); + for (const eventType of eventTypes) { + state = reduceMessageStreamLifecycle(state, eventType); + } + expect(state).toEqual({ authoritative, sealed }); + }, + ); + + test("rejects a streamed row without identity", () => { + expect(() => + foldStreamedSessionEventRows([ + row({ + content: "orphan", + eventType: "message.delta", + id: "orphan", + seq: 1, + streamId: null, + }), + ]), + ).toThrow("has no stream identity"); }); - test("appends repeated delta fragments without deduplicating prefixes", () => { - const folded = foldStreamedSessionEventRows( - [ - row({ content: "ha", eventType: "message.delta", id: "m-1", seq: 1 }), - row({ content: "ha", eventType: "message.delta", id: "m-2", seq: 2 }), - row({ content: "h", eventType: "message.delta", id: "m-3", seq: 3 }), - row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 4 }), - ], - { flushOpenStreams: true }, - ); - - expect(folded.rows.map((entry) => entry.content_text)).toEqual(["hahah"]); + test.each([ + ["message.completed", "available"], + ["message.cancelled", "available"], + ["message.failed", "error"], + ] as const)("folds a streamed assistant message closed by %s", (eventType, processStatus) => { + const folded = foldStreamedSessionEventRows([ + row({ content: "Partial ", eventType: "message.delta", id: "m-1", seq: 1 }), + row({ content: "answer", eventType: "message.delta", id: "m-2", seq: 2 }), + row({ + content: "Message updated.", + eventType, + id: "m-end", + processStatus, + seq: 3, + }), + ]); + + expect(folded).toEqual([ + expect.objectContaining({ + content_text: "Partial answer", + event_type: eventType, + id: "m-end", + process_status: processStatus, + seq: 1, + }), + ]); }); - test("prefers a closing snapshot that extends the streamed prefix", () => { - const folded = foldStreamedSessionEventRows( - [ - row({ content: "Final ans", eventType: "message.delta", id: "m-1", seq: 1 }), - row({ content: "Final answer.", eventType: "message.added", id: "m-added", seq: 2 }), - ], - { flushOpenStreams: true }, - ); - - expect(folded.rows).toHaveLength(1); - expect(folded.rows[0]).toMatchObject({ - content_text: "Final answer.", - id: "m-added", - }); + test("folds a cancelled thought stream", () => { + const folded = foldStreamedSessionEventRows([ + row({ + content: "Inspect", + eventType: "thought.delta", + id: "th-1", + processType: "agent.thinking.delta", + seq: 1, + }), + row({ + content: "Agent thinking updated.", + eventType: "thought.cancelled", + id: "th-end", + processType: "agent.thinking.delta", + seq: 2, + }), + ]); + + expect(folded).toEqual([ + expect.objectContaining({ content_text: "Inspect", event_type: "thought.cancelled" }), + ]); + }); + + test("appends deltas after an authoritative message snapshot", () => { + const folded = foldStreamedSessionEventRows([ + row({ content: "Final ", eventType: "message.added", id: "m-added", seq: 1 }), + row({ content: "answer.", eventType: "message.delta", id: "m-delta", seq: 2 }), + row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 3 }), + ]); + + expect(folded).toEqual([ + expect.objectContaining({ content_text: "Final answer.", id: "m-end" }), + ]); }); - test("keeps standalone message.added rows untouched", () => { - const first = row({ content: "First message.", eventType: "message.added", id: "m-1", seq: 1 }); + test("appends repeated delta fragments without prefix guessing", () => { + const folded = foldStreamedSessionEventRows([ + row({ content: "ha", eventType: "message.delta", id: "m-1", seq: 1 }), + row({ content: "ha", eventType: "message.delta", id: "m-2", seq: 2 }), + row({ content: "h", eventType: "message.delta", id: "m-3", seq: 3 }), + row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 4 }), + ]); + + expect(folded.map((entry) => entry.content_text)).toEqual(["hahah"]); + }); + + test("keeps standalone snapshots for distinct identities", () => { + const first = row({ + content: "First message.", + eventType: "message.added", + id: "m-1", + seq: 1, + streamId: "message-1", + }); const second = row({ content: "Second message.", eventType: "message.added", id: "m-2", seq: 2, + streamId: "message-2", }); - const folded = foldStreamedSessionEventRows([first, second], { flushOpenStreams: true }); - expect(folded.rows).toEqual([first, second]); + expect(foldStreamedSessionEventRows([first, second])).toEqual([first, second]); + }); + + test("treats an authoritative snapshot as a complete reverse-scan boundary", () => { + expect( + findLeftIncompleteSessionEventStreamKeys([ + row({ content: "Complete snapshot", eventType: "message.added", id: "m-1", seq: 1 }), + ]), + ).toEqual(new Set()); }); - test("keeps interleaved non-stream rows and folds around them", () => { + test("keeps timeline order around interleaved non-stream rows", () => { const toolRow = row({ content: "Read file", eventType: "tool.call.updated", @@ -103,57 +186,97 @@ describe("session event stream folding", () => { processType: "tool.use.started", seq: 3, }); - const folded = foldStreamedSessionEventRows( - [ - row({ content: "部分", eventType: "message.delta", id: "m-1", seq: 1 }), - row({ content: "回答", eventType: "message.delta", id: "m-2", seq: 2 }), - toolRow, - row({ content: "。", eventType: "message.delta", id: "m-3", seq: 4 }), - row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 5 }), - ], - { flushOpenStreams: true }, - ); - - expect(folded.rows).toHaveLength(2); - expect(folded.rows[0]).toEqual(toolRow); - expect(folded.rows[1]).toMatchObject({ content_text: "部分回答。", seq: 1 }); + const folded = foldStreamedSessionEventRows([ + row({ content: "部分", eventType: "message.delta", id: "m-1", seq: 1 }), + row({ content: "回答", eventType: "message.delta", id: "m-2", seq: 2 }), + toolRow, + row({ content: "。", eventType: "message.delta", id: "m-3", seq: 4 }), + row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 5 }), + ]); + + expect(folded).toHaveLength(2); + expect(folded[0]).toMatchObject({ content_text: "部分回答。", seq: 1 }); + expect(folded[1]).toEqual(toolRow); }); - test("folds message and thought streams independently", () => { - const folded = foldStreamedSessionEventRows( - [ - row({ - content: "思考", - eventType: "thought.delta", - id: "th-1", - processType: "agent.thinking.delta", - seq: 1, - }), - row({ content: "回答", eventType: "message.delta", id: "m-1", seq: 2 }), - row({ - content: "中", - eventType: "thought.delta", - id: "th-2", - processType: "agent.thinking.delta", - seq: 3, - }), - row({ - content: "Agent thinking updated.", - eventType: "thought.completed", - id: "th-end", - processType: "agent.thinking.delta", - seq: 4, - }), - row({ content: "完毕", eventType: "message.delta", id: "m-2", seq: 5 }), - row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 6 }), - ], - { flushOpenStreams: true }, - ); + test("separates message and thought streams with the same identity", () => { + const folded = foldStreamedSessionEventRows([ + row({ + content: "思考", + eventType: "thought.delta", + id: "th-1", + processType: "agent.thinking.delta", + seq: 1, + }), + row({ content: "回答", eventType: "message.delta", id: "m-1", seq: 2 }), + row({ + content: "Agent thinking updated.", + eventType: "thought.completed", + id: "th-end", + processType: "agent.thinking.delta", + seq: 3, + }), + row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 4 }), + ]); - expect(folded.rows.map((entry) => entry.content_text)).toEqual(["思考中", "回答完毕"]); + expect(folded.map((entry) => entry.content_text)).toEqual(["思考", "回答"]); }); - test("flushes interrupted streams when the run reaches a terminal event", () => { + test("separates user and assistant messages with the same stream identity", () => { + const folded = foldStreamedSessionEventRows([ + row({ + content: "User", + eventType: "message.delta", + id: "user-delta", + processType: "user.message", + seq: 1, + }), + row({ content: "Assistant", eventType: "message.delta", id: "agent-delta", seq: 2 }), + row({ + content: "Message updated.", + eventType: "message.completed", + id: "user-end", + processType: "user.message", + seq: 3, + }), + row({ + content: "Message updated.", + eventType: "message.completed", + id: "agent-end", + seq: 4, + }), + ]); + + expect(folded.map((entry) => [entry.process_type, entry.content_text])).toEqual([ + ["user.message", "User"], + ["agent.message.delta", "Assistant"], + ]); + }); + + test("separates the same stream identity across runs", () => { + const folded = foldStreamedSessionEventRows([ + row({ content: "run one", eventType: "message.delta", id: "r1-delta", seq: 1 }), + row({ + content: "run two", + eventType: "message.delta", + id: "r2-delta", + runId: "run-2" as SessionRunId, + seq: 2, + }), + row({ + content: "Message updated.", + eventType: "message.completed", + id: "r2-end", + runId: "run-2" as SessionRunId, + seq: 3, + }), + row({ content: "Message updated.", eventType: "message.completed", id: "r1-end", seq: 4 }), + ]); + + expect(folded.map((entry) => entry.content_text)).toEqual(["run one", "run two"]); + }); + + test("flushes interrupted streams when their run terminates", () => { const failedRow = row({ content: "Run failed.", eventType: "run.failed", @@ -161,215 +284,227 @@ describe("session event stream folding", () => { processType: "run.failed", seq: 3, }); - const folded = foldStreamedSessionEventRows( - [ - row({ content: "写到一", eventType: "message.delta", id: "m-1", seq: 1 }), - row({ content: "半", eventType: "message.delta", id: "m-2", seq: 2 }), - failedRow, - ], - { flushOpenStreams: false }, - ); - - expect(folded.openStreamRows).toEqual([]); - expect(folded.rows.map((entry) => entry.content_text)).toEqual(["写到一半", "Run failed."]); - }); + const folded = foldStreamedSessionEventRows([ + row({ content: "写到一", eventType: "message.delta", id: "m-1", seq: 1 }), + row({ content: "半", eventType: "message.delta", id: "m-2", seq: 2 }), + failedRow, + ]); - test("closes the previous stream when a new one starts without a completed row", () => { - const folded = foldStreamedSessionEventRows( - [ - row({ content: "第一条", eventType: "message.delta", id: "m-1", seq: 1 }), - row({ content: "Message updated.", eventType: "message.started", id: "m-start", seq: 2 }), - row({ content: "第二条", eventType: "message.delta", id: "m-2", seq: 3 }), - row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 4 }), - ], - { flushOpenStreams: true }, - ); - - expect(folded.rows.map((entry) => entry.content_text)).toEqual(["第一条", "第二条"]); + expect(folded.map((entry) => entry.content_text)).toEqual(["写到一半", "Run failed."]); }); - test("withholds open streams so callers can carry them across reads", () => { - const fragments = [ - row({ content: "流式", eventType: "message.delta", id: "m-1", seq: 1 }), - row({ content: "输出", eventType: "message.delta", id: "m-2", seq: 2 }), - ]; - const firstFold = foldStreamedSessionEventRows(fragments, { flushOpenStreams: false }); - - expect(firstFold.rows).toEqual([]); - expect(firstFold.openStreamRows).toEqual(fragments); - - const secondFold = foldStreamedSessionEventRows( - [ - ...firstFold.openStreamRows, - row({ content: "完成", eventType: "message.delta", id: "m-3", seq: 3 }), - row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 4 }), - ], - { flushOpenStreams: false }, - ); - - expect(secondFold.openStreamRows).toEqual([]); - expect(secondFold.rows).toHaveLength(1); - expect(secondFold.rows[0]).toMatchObject({ content_text: "流式输出完成", id: "m-end" }); - }); + test("reconciles lossless message rows persisted after a repaired run terminal", () => { + const folded = foldStreamedSessionEventRows([ + row({ content: "Message updated.", eventType: "message.started", id: "m-start", seq: 1 }), + row({ content: "Draft", eventType: "message.delta", id: "m-delta", seq: 2 }), + row({ content: "Run failed.", eventType: "run.failed", id: "run-failed", seq: 3 }), + row({ content: "Final ", eventType: "message.added", id: "m-added", seq: 4 }), + row({ content: "world", eventType: "message.delta", id: "m-late", seq: 5 }), + ]); - test("supersedes a closed stream with its trailing snapshot instead of duplicating it", () => { - // Streamed replies persist message.completed (at message_stop) before the - // aggregated assistant snapshot arrives, so the snapshot lands after its - // stream already closed and must replace the folded row, not repeat it. - const folded = foldStreamedSessionEventRows( - [ - row({ content: "Message updated.", eventType: "message.started", id: "m-start", seq: 1 }), - row({ content: "P", eventType: "message.delta", id: "m-1", seq: 2 }), - row({ - content: "ong. What would you like to work on?", - eventType: "message.delta", - id: "m-2", - seq: 3, - }), - row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 4 }), - row({ - content: "Pong. What would you like to work on?", - eventType: "message.added", - id: "m-added", - seq: 5, - }), - ], - { flushOpenStreams: true }, - ); - - expect(folded.rows).toHaveLength(1); - expect(folded.rows[0]).toMatchObject({ - content_text: "Pong. What would you like to work on?", - event_type: "message.added", - id: "m-added", - seq: 1, - }); + expect(folded.map((entry) => entry.content_text)).toEqual(["Final world", "Run failed."]); }); - test("collapses a fractured stream whose snapshot matches the concatenated fragments", () => { - // YEF-884: a dropped message_start fractures one reply into per-fragment - // messages, so the rows arrive as started/delta pairs per fragment plus a - // final full snapshot. The snapshot equals the fragment concatenation and - // must fold everything into a single timeline entry. - const folded = foldStreamedSessionEventRows( - [ - row({ content: "Message updated.", eventType: "message.started", id: "s-1", seq: 1 }), - row({ content: "P", eventType: "message.delta", id: "m-1", seq: 2 }), - row({ content: "Message updated.", eventType: "message.started", id: "s-2", seq: 3 }), - row({ - content: "ong. What would you like to work on?", - eventType: "message.delta", - id: "m-2", - seq: 4, - }), - row({ content: "Message updated.", eventType: "message.started", id: "s-3", seq: 5 }), - row({ - content: "Pong. What would you like to work on?", - eventType: "message.delta", - id: "m-3", - seq: 6, - }), - row({ - content: "Pong. What would you like to work on?", - eventType: "message.added", - id: "m-added", - seq: 7, - }), - row({ content: "Message updated.", eventType: "message.completed", id: "c-3", seq: 8 }), - row({ content: "Message updated.", eventType: "message.completed", id: "c-1", seq: 9 }), - row({ content: "Message updated.", eventType: "message.completed", id: "c-2", seq: 10 }), - row({ - content: "run-1", - eventType: "run.completed", - id: "r-done", - processType: "run.completed", - seq: 11, - }), - ], - { flushOpenStreams: true }, - ); + test("uses a trailing authoritative snapshot without replacing the terminal row", () => { + const folded = foldStreamedSessionEventRows([ + row({ content: "Message updated.", eventType: "message.started", id: "m-start", seq: 1 }), + row({ content: "corrupt preview", eventType: "message.delta", id: "m-delta", seq: 2 }), + row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 3 }), + row({ + content: "Authoritative answer.", + eventType: "message.added", + id: "m-added", + seq: 4, + }), + ]); - expect(folded.rows.map((entry) => entry.content_text)).toEqual([ - "Pong. What would you like to work on?", - "run-1", + expect(folded).toEqual([ + expect.objectContaining({ + content_text: "Authoritative answer.", + event_type: "message.completed", + id: "m-end", + seq: 1, + }), ]); - expect(folded.rows[0]).toMatchObject({ id: "m-added", seq: 1 }); }); - test("supersedes a truncated stream with the longer prefix-matching snapshot", () => { - const folded = foldStreamedSessionEventRows( - [ - row({ content: "Final ans", eventType: "message.delta", id: "m-1", seq: 1 }), - row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 2 }), - row({ content: "Final answer.", eventType: "message.added", id: "m-added", seq: 3 }), - ], - { flushOpenStreams: true }, - ); - - expect(folded.rows).toHaveLength(1); - expect(folded.rows[0]).toMatchObject({ - content_text: "Final answer.", - id: "m-added", - seq: 1, + test("folds interleaved message streams by identity", () => { + const folded = foldStreamedSessionEventRows([ + row({ content: "A1", eventType: "message.delta", id: "a-1", seq: 1, streamId: "a" }), + row({ content: "B1", eventType: "message.delta", id: "b-1", seq: 2, streamId: "b" }), + row({ content: "A2", eventType: "message.delta", id: "a-2", seq: 3, streamId: "a" }), + row({ + content: "Message updated.", + eventType: "message.completed", + id: "b-end", + seq: 4, + streamId: "b", + }), + row({ + content: "Message updated.", + eventType: "message.completed", + id: "a-end", + seq: 5, + streamId: "a", + }), + ]); + + expect( + Object.fromEntries(folded.map((entry) => [entry.stream_id, entry.content_text])), + ).toEqual({ + a: "A1A2", + b: "B1", }); }); - test("keeps a snapshot that does not extend the closed stream as its own message", () => { - const folded = foldStreamedSessionEventRows( - [ - row({ content: "第一条进度", eventType: "message.delta", id: "m-1", seq: 1 }), - row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 2 }), - row({ content: "另一条最终回复", eventType: "message.added", id: "m-added", seq: 3 }), - ], - { flushOpenStreams: true }, - ); - - expect(folded.rows.map((entry) => entry.content_text)).toEqual([ - "第一条进度", - "另一条最终回复", + test("keeps a snapshot for a different identity as its own message", () => { + const folded = foldStreamedSessionEventRows([ + row({ + content: "第一条进度", + eventType: "message.delta", + id: "m-1", + seq: 1, + streamId: "message-1", + }), + row({ + content: "Message updated.", + eventType: "message.completed", + id: "m-end", + seq: 2, + streamId: "message-1", + }), + row({ + content: "另一条最终回复", + eventType: "message.added", + id: "m-added", + seq: 3, + streamId: "message-2", + }), ]); + + expect(folded.map((entry) => entry.content_text)).toEqual(["第一条进度", "另一条最终回复"]); }); - test("consumes fragments per snapshot across a multi-message turn", () => { - const folded = foldStreamedSessionEventRows( - [ - row({ content: "进度说明", eventType: "message.delta", id: "m-1", seq: 1 }), - row({ content: "Message updated.", eventType: "message.completed", id: "c-1", seq: 2 }), - row({ content: "进度说明", eventType: "message.added", id: "a-1", seq: 3 }), - row({ content: "最终回复", eventType: "message.delta", id: "m-2", seq: 4 }), - row({ content: "Message updated.", eventType: "message.completed", id: "c-2", seq: 5 }), - row({ content: "最终回复", eventType: "message.added", id: "a-2", seq: 6 }), - ], - { flushOpenStreams: true }, - ); - - expect(folded.rows.map((entry) => entry.content_text)).toEqual(["进度说明", "最终回复"]); - expect(folded.rows.map((entry) => entry.id)).toEqual(["a-1", "a-2"]); + test("associates trailing snapshots with their stream in a multi-message turn", () => { + const folded = foldStreamedSessionEventRows([ + row({ content: "进度", eventType: "message.delta", id: "m-1", seq: 1, streamId: "one" }), + row({ + content: "Message updated.", + eventType: "message.completed", + id: "c-1", + seq: 2, + streamId: "one", + }), + row({ content: "进度说明", eventType: "message.added", id: "a-1", seq: 3, streamId: "one" }), + row({ content: "最终", eventType: "message.delta", id: "m-2", seq: 4, streamId: "two" }), + row({ + content: "Message updated.", + eventType: "message.completed", + id: "c-2", + seq: 5, + streamId: "two", + }), + row({ content: "最终回复", eventType: "message.added", id: "a-2", seq: 6, streamId: "two" }), + ]); + + expect(folded.map((entry) => entry.content_text)).toEqual(["进度说明", "最终回复"]); + expect(folded.map((entry) => entry.id)).toEqual(["c-1", "c-2"]); }); - test("drops streams that never carried text", () => { - const folded = foldStreamedSessionEventRows( - [ - row({ content: "Message updated.", eventType: "message.started", id: "m-start", seq: 1 }), - row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 2 }), - ], - { flushOpenStreams: true }, - ); + test("keeps a terminal stream row with empty content", () => { + const folded = foldStreamedSessionEventRows([ + row({ content: "Message updated.", eventType: "message.started", id: "m-start", seq: 1 }), + row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 2 }), + ]); - expect(folded.rows).toEqual([]); + expect(folded).toEqual([ + expect.objectContaining({ content_text: "", event_type: "message.completed", id: "m-end" }), + ]); }); test("re-folding folded rows is a no-op", () => { - const folded = foldStreamedSessionEventRows( - [ - row({ content: "你", eventType: "message.delta", id: "m-1", seq: 1 }), - row({ content: "好", eventType: "message.delta", id: "m-2", seq: 2 }), - row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 3 }), - ], - { flushOpenStreams: true }, - ); - const refolded = foldStreamedSessionEventRows(folded.rows, { flushOpenStreams: true }); - - expect(refolded.rows).toEqual(folded.rows); + const folded = foldStreamedSessionEventRows([ + row({ content: "你", eventType: "message.delta", id: "m-1", seq: 1 }), + row({ content: "好", eventType: "message.delta", id: "m-2", seq: 2 }), + row({ content: "Message updated.", eventType: "message.completed", id: "m-end", seq: 3 }), + ]); + + expect(foldStreamedSessionEventRows(folded)).toEqual(folded); + }); + + test("resolves only the latest sealed authoritative message snapshot", () => { + const firstCompletion = [ + row({ content: "Draft", eventType: "message.added", id: "m-draft", seq: 1 }), + row({ content: "Message updated.", eventType: "message.completed", id: "m-end-1", seq: 2 }), + ]; + const replacement = [ + ...firstCompletion, + row({ content: "Final ", eventType: "message.added", id: "m-added", seq: 3 }), + row({ content: "answer", eventType: "message.delta", id: "m-final", seq: 4 }), + ]; + + expect(resolveSealedMessageStream(firstCompletion)).toEqual({ text: "Draft" }); + expect(resolveSealedMessageStream(replacement)).toBeNull(); + expect( + resolveSealedMessageStream([ + ...replacement, + row({ + content: "Message updated.", + eventType: "message.completed", + id: "m-end-2", + seq: 5, + }), + ]), + ).toEqual({ text: "Final answer" }); + }); + + test("does not treat failed, cancelled, or restarted messages as sealed final output", () => { + for (const eventType of ["message.failed", "message.cancelled"] as const) { + expect( + resolveSealedMessageStream([ + row({ content: "partial", eventType: "message.delta", id: "m-delta", seq: 1 }), + row({ content: "Message updated.", eventType, id: "m-terminal", seq: 2 }), + ]), + ).toBeNull(); + } + + expect( + resolveSealedMessageStream([ + row({ content: "old", eventType: "message.added", id: "m-delta", seq: 1 }), + row({ + content: "Message updated.", + eventType: "message.completed", + id: "m-terminal", + seq: 2, + }), + row({ content: "Message updated.", eventType: "message.started", id: "m-restart", seq: 3 }), + ]), + ).toBeNull(); + + expect( + resolveSealedMessageStream([ + row({ content: "Message updated.", eventType: "message.started", id: "m-start", seq: 1 }), + row({ content: "best effort", eventType: "message.delta", id: "m-delta", seq: 2 }), + row({ + content: "Message updated.", + eventType: "message.completed", + id: "m-terminal", + seq: 3, + }), + ]), + ).toBeNull(); + + expect( + resolveSealedMessageStream([ + row({ content: "old", eventType: "message.added", id: "m-added", seq: 1 }), + row({ content: "Message updated.", eventType: "message.failed", id: "m-failed", seq: 2 }), + row({ + content: "Message updated.", + eventType: "message.completed", + id: "m-terminal", + seq: 3, + }), + ]), + ).toBeNull(); }); }); diff --git a/apps/api/tests/session-lifecycle-mutation.test.ts b/apps/api/tests/session-lifecycle-mutation.test.ts index 5381f65e..0725bdd9 100644 --- a/apps/api/tests/session-lifecycle-mutation.test.ts +++ b/apps/api/tests/session-lifecycle-mutation.test.ts @@ -3,6 +3,8 @@ import { describe, expect, test } from "bun:test"; import type { RuntimeOperationId } from "@mosoo/id"; import type { AuthenticatedViewer } from "../src/modules/auth/application/viewer-auth.service"; +import type { SandboxHandle } from "../src/modules/runtime/infrastructure/sandbox-handles"; +import { setSessionRunStatus } from "../src/modules/runtime/infrastructure/session-runs/session-run-store.repository"; import { deleteSessionCascade, repairStaleSessionDeleteCleanups, @@ -10,16 +12,20 @@ import { import { archiveAgentSession, deleteAgentSession, + repairStaleSessionArchiveCleanups, unarchiveAgentSession, } from "../src/modules/sessions/application/session-lifecycle-mutation.service"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; +import { applyDrizzleMigration } from "./helpers/drizzle-migrations"; import { PUBLIC_API_TEST_IDS, createPublicHttpContractDatabase, createPublicHttpTestBindings, + insertActiveSandboxSessionFixture, insertOwnerSession, insertNonOwnerSession, } from "./helpers/public-api-http-test-fixture"; +import { SqliteD1Database } from "./helpers/sqlite-d1"; const OWNER_VIEWER: AuthenticatedViewer = { email: "owner@example.com", @@ -64,11 +70,17 @@ function withDriverConnection(bindings: ApiBindings, paths: string[]): ApiBindin }; } -function createSessionLifecycleBinding(paths: string[], options: { destroyError?: Error } = {}) { +function createSessionLifecycleBinding( + paths: string[], + options: { closeError?: Error; destroyError?: Error } = {}, +) { return { get: () => ({ closeViewers: async (_sessionId: string, reason: string) => { paths.push(`close:${reason}`); + if (options.closeError !== undefined) { + throw options.closeError; + } }, destroy: async (_sessionId: string, reason: string) => { paths.push(`destroy:${reason}`); @@ -87,108 +99,65 @@ function createSessionLifecycleBinding(paths: string[], options: { destroyError? function withSessionLifecycleBinding( bindings: ApiBindings, paths: string[] = [], - options: { destroyError?: Error } = {}, + options: { closeError?: Error; destroyError?: Error } = {}, ): ApiBindings { return { ...bindings, + runtimeSubjectHandleFactory: () => createCleanupSandboxHandle(), Session: createSessionLifecycleBinding(paths, options) as ApiBindings["Session"], }; } -async function ensureRuntimeLifecycleTables(database: D1Database): Promise { - await database - .prepare( - ` - CREATE TABLE IF NOT EXISTS sandbox ( - id text PRIMARY KEY NOT NULL, - inactive_deadline_at integer, - kind text NOT NULL, - status text DEFAULT 'active' NOT NULL, - updated_at integer DEFAULT 1 NOT NULL - ) - `, - ) - .run(); - await database - .prepare( - ` - CREATE TABLE IF NOT EXISTS sandbox_session ( - cloudflare_session_id text NOT NULL, - created_at integer NOT NULL, - cwd text NOT NULL, - origin_json text NOT NULL, - sandbox_id text NOT NULL, - session_id text PRIMARY KEY NOT NULL, - status text NOT NULL, - updated_at integer NOT NULL - ) - `, - ) - .run(); - await database - .prepare( - ` - CREATE TABLE IF NOT EXISTS sandbox_backup ( - id text PRIMARY KEY NOT NULL, - dir text NOT NULL, - sandbox_id text NOT NULL, - status text NOT NULL, - created_at integer NOT NULL, - updated_at integer NOT NULL - ) - `, - ) - .run(); +function createCleanupSandboxHandle(): SandboxHandle { + const unavailable = async (): Promise => { + throw new Error("Unexpected sandbox test method call."); + }; + + return { + configureNetworkConstraints: unavailable, + createBackup: unavailable, + createSession: unavailable, + deleteSession: async (sessionId) => ({ + sessionId, + success: true, + timestamp: new Date(0).toISOString(), + }), + destroy: unavailable, + exec: unavailable, + getSession: unavailable, + mkdir: unavailable, + mountBucket: unavailable, + readFile: unavailable, + restoreBackup: unavailable, + setKeepAlive: unavailable, + startProcess: unavailable, + terminal: unavailable, + unmountBucket: unavailable, + watch: unavailable, + writeFile: unavailable, + wsConnect: unavailable, + }; } async function insertSandboxSession( - database: D1Database, + database: SqliteD1Database, sessionId: string = PUBLIC_API_TEST_IDS.nonOwnerSession, ): Promise { - await ensureRuntimeLifecycleTables(database); - await database - .prepare( - ` - INSERT INTO sandbox ( - id, - kind, - subject_kind, - subject_id, - status, - created_at, - updated_at - ) - VALUES (?, ?, ?, ?, ?, ?, ?) - `, - ) - .bind(PUBLIC_API_TEST_IDS.sandbox, "pet", "agent", PUBLIC_API_TEST_IDS.agent, "active", 1, 1) - .run(); + const ownerAccountId = + sessionId === PUBLIC_API_TEST_IDS.ownerSession + ? PUBLIC_API_TEST_IDS.ownerAccount + : PUBLIC_API_TEST_IDS.nonOwnerAccount; + await insertActiveSandboxSessionFixture(database, { + cwd: "/workspace/session-cwd", + ownerAccountId, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxSessionId: "01J0000000000000000000000S", + sessionId, + timestampMs: 1, + }); await database - .prepare( - ` - INSERT INTO sandbox_session ( - cloudflare_session_id, - created_at, - cwd, - origin_json, - sandbox_id, - session_id, - status, - updated_at - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) - `, - ) - .bind( - "01J0000000000000000000000S", - 1, - "session-cwd", - "{}", - PUBLIC_API_TEST_IDS.sandbox, - sessionId, - "closed", - 1, - ) + .prepare("UPDATE sandbox_session SET status = 'closed' WHERE session_id = ?") + .bind(sessionId) .run(); } @@ -260,6 +229,7 @@ async function insertDriverInstance( INSERT INTO driver_instance ( id, sandbox_id, + sandbox_incarnation, sandbox_session_id, runtime, protocol, @@ -272,12 +242,13 @@ async function insertDriverInstance( created_at, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( input.driverId, PUBLIC_API_TEST_IDS.sandbox, + 1, input.sandboxSessionId, "cloudflare-container", "driver-ws", @@ -357,15 +328,286 @@ async function insertSessionFileForCleanup(database: D1Database): Promise } describe("session lifecycle mutations", () => { + test("cleanup migration preserves data without enqueueing external cleanup", async () => { + const database = new SqliteD1Database(); + database.execute(` + CREATE TABLE session ( + id text PRIMARY KEY, + archived_at integer, + status text NOT NULL, + status_operation_id text, + status_seq integer NOT NULL, + updated_at integer NOT NULL + ); + CREATE TABLE child ( + session_id text REFERENCES session(id) ON DELETE CASCADE + ); + INSERT INTO session VALUES ('session-1', NULL, 'IDLE', NULL, 0, 1); + INSERT INTO session VALUES ('legacy-delete', 2, 'TERMINATED', 'delete-op', 4, 2); + INSERT INTO session VALUES ('legacy-archive', 3, 'IDLE', NULL, 7, 3); + INSERT INTO child VALUES ('session-1'); + `); + applyDrizzleMigration(database, "0015_session-cleanup-operation"); + + expect(await database.prepare("SELECT COUNT(*) AS count FROM child").first()).toEqual({ + count: 1, + }); + expect( + await database + .prepare( + "SELECT cleanup_operation_kind, status, status_operation_id, status_seq FROM session WHERE id = 'legacy-delete'", + ) + .first(), + ).toEqual({ + cleanup_operation_kind: null, + status: "TERMINATED", + status_operation_id: "delete-op", + status_seq: 4, + }); + expect( + await database + .prepare( + "SELECT cleanup_operation_kind, status, status_operation_id, status_seq FROM session WHERE id = 'legacy-archive'", + ) + .first(), + ).toEqual({ + cleanup_operation_kind: null, + status: "IDLE", + status_operation_id: null, + status_seq: 7, + }); + await expect( + database + .prepare("UPDATE session SET cleanup_operation_kind = 'archive' WHERE id = 'session-1'") + .run(), + ).rejects.toThrow("session_cleanup_operation_kind_check"); + await expect( + database + .prepare( + `UPDATE session + SET runtime_provisioning_operation_id = '01J0000000000000000000000R', + runtime_provisioning_sandbox_id = '01J0000000000000000000000S' + WHERE id = 'session-1'`, + ) + .run(), + ).rejects.toThrow("session_runtime_provisioning_lease_check"); + await expect( + database + .prepare( + `UPDATE session + SET runtime_provisioning_operation_id = '01J0000000000000000000000R', + runtime_provisioning_sandbox_id = '01J0000000000000000000000S', + runtime_provisioning_heartbeat_at = 'not-a-timestamp' + WHERE id = 'session-1'`, + ) + .run(), + ).rejects.toThrow("session_runtime_provisioning_lease_check"); + await database + .prepare( + `UPDATE session + SET archived_at = 1, + cleanup_operation_kind = 'archive', + status = 'RESCHEDULING', + status_operation_id = 'operation-1' + WHERE id = 'session-1'`, + ) + .run(); + await database + .prepare( + `UPDATE session + SET cleanup_operation_kind = 'archive', + status = 'IDLE', + status_operation_id = NULL + WHERE id = 'session-1'`, + ) + .run(); + }); + + test("maintenance adopts cleanup claims written by an old worker after migration", async () => { + const archiveDatabase = await createPublicHttpContractDatabase(); + await insertOwnerSession(archiveDatabase); + await archiveDatabase + .prepare( + `UPDATE session + SET archived_at = 1, + cleanup_operation_kind = NULL, + status = 'RESCHEDULING', + status_operation_id = ?, + updated_at = 1 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.ownerSession) + .run(); + + expect( + await repairStaleSessionArchiveCleanups( + withSessionLifecycleBinding(createPublicHttpTestBindings(archiveDatabase) as ApiBindings), + { limit: 10, staleUpdatedAtLte: 1 }, + ), + ).toBe(1); + expect( + await archiveDatabase + .prepare( + "SELECT cleanup_operation_kind, status, status_operation_id FROM session WHERE id = ?", + ) + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first(), + ).toEqual({ + cleanup_operation_kind: "archive", + status: "IDLE", + status_operation_id: null, + }); + + const deleteDatabase = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(deleteDatabase); + await deleteDatabase + .prepare( + `UPDATE session + SET archived_at = 1, + cleanup_operation_kind = NULL, + status = 'TERMINATED', + status_operation_id = ?, + updated_at = 1 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + + expect( + await repairStaleSessionDeleteCleanups( + withSessionLifecycleBinding(createPublicHttpTestBindings(deleteDatabase) as ApiBindings), + { limit: 10, staleUpdatedAtLte: 1 }, + ), + ).toBe(1); + expect( + await deleteDatabase + .prepare("SELECT id FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first(), + ).toBeNull(); + }); + + test("archive and delete cannot pass an active runtime provisioning fence", async () => { + for (const action of ["archive", "delete"] as const) { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await database + .prepare( + `UPDATE session + SET runtime_provisioning_heartbeat_at = 1, + runtime_provisioning_operation_id = ?, + runtime_provisioning_sandbox_id = ? + WHERE id = ?`, + ) + .bind( + PUBLIC_API_TEST_IDS.operation, + PUBLIC_API_TEST_IDS.sandbox, + PUBLIC_API_TEST_IDS.ownerSession, + ) + .run(); + const bindings = withSessionLifecycleBinding( + createPublicHttpTestBindings(database) as ApiBindings, + ); + + const mutation = + action === "archive" + ? archiveAgentSession({ + bindings, + appId: PUBLIC_API_TEST_IDS.app, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + viewer: OWNER_VIEWER, + }) + : deleteSessionCascade(bindings, PUBLIC_API_TEST_IDS.ownerSession); + await expect(mutation).rejects.toThrow(); + expect( + await database + .prepare( + "SELECT archived_at, cleanup_operation_kind, runtime_provisioning_operation_id FROM session WHERE id = ?", + ) + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first(), + ).toEqual({ + archived_at: null, + cleanup_operation_kind: null, + runtime_provisioning_operation_id: PUBLIC_API_TEST_IDS.operation, + }); + } + }); + + test("failed archive repairs move behind older pending work", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await insertOwnerSession(database); + await database + .prepare( + `UPDATE session + SET archived_at = 1, + cleanup_operation_kind = 'archive', + status = 'RESCHEDULING', + status_operation_id = ?, + updated_at = 1 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.nonOwnerSession) + .run(); + await database + .prepare( + `UPDATE session + SET archived_at = 2, + cleanup_operation_kind = 'archive', + status = 'RESCHEDULING', + status_operation_id = ?, + updated_at = 2 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.ownerSession) + .run(); + const bindings = withSessionLifecycleBinding( + createPublicHttpTestBindings(database) as ApiBindings, + [], + { closeError: new Error("viewer cleanup interrupted") }, + ); + + expect( + await repairStaleSessionArchiveCleanups(bindings, { + limit: 1, + staleUpdatedAtLte: 2, + }), + ).toBe(1); + const firstAttempt = await database + .prepare("SELECT updated_at FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.nonOwnerSession) + .first<{ updated_at: number }>(); + expect(firstAttempt?.updated_at).toBeGreaterThan(2); + expect( + await database + .prepare("SELECT updated_at FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first(), + ).toEqual({ updated_at: 2 }); + + expect( + await repairStaleSessionArchiveCleanups(bindings, { + limit: 1, + staleUpdatedAtLte: 2, + }), + ).toBe(1); + const secondAttempt = await database + .prepare("SELECT updated_at FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first<{ updated_at: number }>(); + expect(secondAttempt?.updated_at).toBeGreaterThan(2); + }); + test("delete cascade removes live and terminal driver instances associated with the session", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); await insertSandboxSession(database); - await insertSessionRun(database, { runId: PUBLIC_API_TEST_IDS.run }); await insertSessionRun(database, { runId: PUBLIC_API_TEST_IDS.runAlt, status: "completed", }); + await insertSessionRun(database, { runId: PUBLIC_API_TEST_IDS.run }); await insertDriverInstance(database, { driverId: PUBLIC_API_TEST_IDS.driverNonOwner, sandboxSessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, @@ -466,7 +708,7 @@ describe("session lifecycle mutations", () => { expect(anchoredSession?.archived_at).toBeNumber(); expect(anchoredSession).toMatchObject({ - status: "TERMINATED", + status: "IDLE", status_operation_id: operationId, }); @@ -493,7 +735,7 @@ describe("session lifecycle mutations", () => { const bucket = new RetriableSessionFileBucket(); const bindings = withSessionLifecycleBinding( createPublicHttpTestBindings(database, { - fileBucket: bucket as unknown as R2Bucket, + fileBucket: bucket as R2Bucket, }) as ApiBindings, ); @@ -513,7 +755,7 @@ describe("session lifecycle mutations", () => { .first<{ status: string }>(); expect(bucket.deletedKeys).toEqual([`session-files/${SESSION_FILE_ID}/notes.txt`]); - expect(anchoredSession).toEqual({ status: "TERMINATED" }); + expect(anchoredSession).toEqual({ status: "IDLE" }); expect(retainedFile).toEqual({ status: "deleting" }); bucket.failDelete = false; @@ -542,7 +784,6 @@ describe("session lifecycle mutations", () => { test("archive cancels active runs and exposes an idle archived session", async () => { const database = await createPublicHttpContractDatabase(); await insertOwnerSession(database); - await ensureRuntimeLifecycleTables(database); await insertSandboxSession(database, PUBLIC_API_TEST_IDS.ownerSession); await insertSessionRun(database, { createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, @@ -586,7 +827,120 @@ describe("session lifecycle mutations", () => { expect(outcomes.every((outcome) => outcome.status !== "failed")).toBe(true); }); - test("unarchive normalizes stale rescheduling state before exposing the session", async () => { + test("maintenance resumes an admitted archive after external cleanup fails", async () => { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + const baseBindings = createPublicHttpTestBindings(database) as ApiBindings; + const failingBindings = withSessionLifecycleBinding(baseBindings, [], { + closeError: new Error("viewer cleanup interrupted"), + }); + + await expect( + archiveAgentSession({ + bindings: failingBindings, + appId: PUBLIC_API_TEST_IDS.app, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + viewer: OWNER_VIEWER, + }), + ).rejects.toThrow("viewer cleanup interrupted"); + + const admitted = await database + .prepare( + "SELECT archived_at, cleanup_operation_kind, status, status_operation_id FROM session WHERE id = ?", + ) + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first<{ + archived_at: number | null; + cleanup_operation_kind: string | null; + status: string; + status_operation_id: string | null; + }>(); + expect(admitted?.archived_at).toBeNumber(); + expect(admitted).toMatchObject({ + cleanup_operation_kind: "archive", + status: "RESCHEDULING", + }); + expect(admitted?.status_operation_id).toBeString(); + + const repaired = await repairStaleSessionArchiveCleanups( + withSessionLifecycleBinding(baseBindings), + { limit: 10, staleUpdatedAtLte: Date.now() }, + ); + const archived = await database + .prepare( + "SELECT archived_at, cleanup_operation_kind, status, status_operation_id FROM session WHERE id = ?", + ) + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first<{ + archived_at: number | null; + cleanup_operation_kind: string | null; + status: string; + status_operation_id: string | null; + }>(); + + expect(repaired).toBe(1); + expect(archived?.archived_at).toBeNumber(); + expect(archived).toMatchObject({ + cleanup_operation_kind: "archive", + status: "IDLE", + status_operation_id: null, + }); + }); + + for (const cleanupOperationKind of ["archive", "delete"] as const) { + test(`${cleanupOperationKind} ownership rejects duplicate and advancing active Run projections`, async () => { + for (const initialRunStatus of ["booting", "queued"] as const) { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await insertSessionRun(database, { + createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + runId: PUBLIC_API_TEST_IDS.run, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + status: initialRunStatus, + }); + const racingDatabase = claimCleanupAfterLifecycleRead(database, cleanupOperationKind); + + const outcome = await setSessionRunStatus(racingDatabase, { + runId: PUBLIC_API_TEST_IDS.run, + status: "booting", + }); + + expect(outcome.kind).toBe(initialRunStatus === "booting" ? "duplicate" : "stale"); + expect(await readCleanupOwnedRunProjection(database)).toEqual({ + archived_at: 2, + cleanup_operation_kind: cleanupOperationKind, + run_status: initialRunStatus, + run_status_seq: 0, + session_status: "RESCHEDULING", + session_status_operation_id: PUBLIC_API_TEST_IDS.operation, + session_status_seq: 1, + }); + } + }); + + test(`${cleanupOperationKind} ownership rejects a duplicate active projection repair`, async () => { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await insertSessionRun(database, { + createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + runId: PUBLIC_API_TEST_IDS.run, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + status: "booting", + }); + await claimCleanup(database, cleanupOperationKind); + const before = await readCleanupOwnedRunProjection(database); + + const outcome = await setSessionRunStatus(database, { + runId: PUBLIC_API_TEST_IDS.run, + status: "booting", + }); + + expect(outcome.kind).toBe("repair_needed"); + expect(await readCleanupOwnedRunProjection(database)).toEqual(before); + }); + } + + test("unarchive refuses to clear an in-progress cleanup fence", async () => { const database = await createPublicHttpContractDatabase(); await insertOwnerSession(database); await insertSessionRun(database, { @@ -599,25 +953,35 @@ describe("session lifecycle mutations", () => { ` UPDATE session SET archived_at = ?, + cleanup_operation_kind = ?, status = ?, status_operation_id = ? WHERE id = ? `, ) - .bind(1, "RESCHEDULING", PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.ownerSession) + .bind( + 1, + "delete", + "RESCHEDULING", + PUBLIC_API_TEST_IDS.operation, + PUBLIC_API_TEST_IDS.ownerSession, + ) .run(); - await unarchiveAgentSession({ - database, - appId: PUBLIC_API_TEST_IDS.app, - sessionId: PUBLIC_API_TEST_IDS.ownerSession, - viewer: OWNER_VIEWER, - }); + await expect( + unarchiveAgentSession({ + database, + appId: PUBLIC_API_TEST_IDS.app, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + viewer: OWNER_VIEWER, + }), + ).rejects.toThrow("cleanup is still in progress"); const row = await database .prepare( ` SELECT session.archived_at, + session.cleanup_operation_kind, session.status AS session_status, session.status_operation_id, session_run.status AS run_status @@ -629,17 +993,44 @@ describe("session lifecycle mutations", () => { .bind(PUBLIC_API_TEST_IDS.ownerSession) .first<{ archived_at: number | null; + cleanup_operation_kind: string | null; run_status: string; session_status: string; status_operation_id: string | null; }>(); expect(row).toEqual({ - archived_at: null, - run_status: "cancelled", - session_status: "IDLE", - status_operation_id: null, + archived_at: 1, + cleanup_operation_kind: "delete", + run_status: "running", + session_status: "RESCHEDULING", + status_operation_id: PUBLIC_API_TEST_IDS.operation, + }); + }); + + test("unarchive exposes a stable archived Session with no cleanup owner", async () => { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await database + .prepare( + "UPDATE session SET archived_at = 1, cleanup_operation_kind = 'archive' WHERE id = ?", + ) + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .run(); + + await unarchiveAgentSession({ + database, + appId: PUBLIC_API_TEST_IDS.app, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + viewer: OWNER_VIEWER, }); + + expect( + await database + .prepare("SELECT archived_at FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first(), + ).toEqual({ archived_at: null }); }); test("lifecycle mutations reject attributed participants who are not session creators", async () => { @@ -690,3 +1081,95 @@ describe("session lifecycle mutations", () => { ).rejects.toThrow(); }); }); + +async function readCleanupOwnedRunProjection(database: D1Database) { + return database + .prepare( + `SELECT session.archived_at, + session.cleanup_operation_kind, + session.status AS session_status, + session.status_operation_id AS session_status_operation_id, + session.status_seq AS session_status_seq, + session_run.status AS run_status, + session_run.status_seq AS run_status_seq + FROM session + JOIN session_run ON session_run.id = session.last_run_id + WHERE session.id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first(); +} + +function claimCleanupAfterLifecycleRead( + database: D1Database, + cleanupOperationKind: "archive" | "delete", +): D1Database { + let claimed = false; + let interceptedLifecycleRead = false; + + async function claim(): Promise { + if (claimed) { + return; + } + claimed = true; + await claimCleanup(database, cleanupOperationKind); + } + + function injectAfterRead(statement: D1PreparedStatement): D1PreparedStatement { + return new Proxy(statement, { + get(target, property, receiver) { + if (property === "bind") { + return (...values: unknown[]) => injectAfterRead(target.bind(...values)); + } + if (property === "all" || property === "first" || property === "raw") { + return async (...args: unknown[]) => { + const method = Reflect.get(target, property, receiver) as ( + ...values: unknown[] + ) => unknown; + const result = await method.apply(target, args); + await claim(); + return result; + }; + } + const value = Reflect.get(target, property, receiver); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + } + + return new Proxy(database, { + get(target, property, receiver) { + if (property === "prepare") { + return (query: string) => { + const statement = target.prepare(query); + if (interceptedLifecycleRead) { + return statement; + } + interceptedLifecycleRead = true; + return injectAfterRead(statement); + }; + } + const value = Reflect.get(target, property, receiver); + return typeof value === "function" ? value.bind(target) : value; + }, + }); +} + +async function claimCleanup( + database: D1Database, + cleanupOperationKind: "archive" | "delete", +): Promise { + await database + .prepare( + `UPDATE session + SET archived_at = 2, + cleanup_operation_kind = ?, + status = 'RESCHEDULING', + status_operation_id = ?, + status_seq = status_seq + 1, + updated_at = 2 + WHERE id = ?`, + ) + .bind(cleanupOperationKind, PUBLIC_API_TEST_IDS.operation, PUBLIC_API_TEST_IDS.ownerSession) + .run(); +} diff --git a/apps/api/tests/session-message-store.test.ts b/apps/api/tests/session-message-store.test.ts index 92a875be..ce3a0fad 100644 --- a/apps/api/tests/session-message-store.test.ts +++ b/apps/api/tests/session-message-store.test.ts @@ -21,11 +21,23 @@ function createSessionMessageStoreDatabase(): SqliteD1Database { created_by_account_id text NOT NULL, id text PRIMARY KEY NOT NULL, plan_json text, + projection_format text DEFAULT 'materialized' NOT NULL, role text NOT NULL, segments_json text, seq integer NOT NULL, session_id text NOT NULL, - session_run_id text + session_run_id text, + CHECK (projection_format IN ('materialized', 'event_stream_v3')), + CHECK ( + projection_format <> 'event_stream_v3' + OR ( + role = 'assistant' + AND session_run_id IS NOT NULL + AND content_text = '' + AND plan_json IS NULL + AND segments_json IS NULL + ) + ) ); CREATE UNIQUE INDEX session_message_session_seq_idx diff --git a/apps/api/tests/session-model-call-identity.test.ts b/apps/api/tests/session-model-call-identity.test.ts index 27cabc0c..9ed5f1d9 100644 --- a/apps/api/tests/session-model-call-identity.test.ts +++ b/apps/api/tests/session-model-call-identity.test.ts @@ -37,6 +37,7 @@ const DRIVER_INSTANCE_ID = parsePlatformId( "01J00000000000000000000017", "driver instance ID", ); +const USAGE_CREATED_AT_MS = 1_500; interface IdentityProjection { metadata_json: string | null; @@ -86,6 +87,7 @@ function createSessionModelCallDatabase(): SqliteD1Database { CREATE TABLE session_run ( agent_id text NOT NULL, completed_at integer, + created_at integer NOT NULL, created_by_account_id text NOT NULL, deployment_version_id text, id text PRIMARY KEY NOT NULL, @@ -94,6 +96,7 @@ function createSessionModelCallDatabase(): SqliteD1Database { runtime_id text, session_id text NOT NULL, started_at integer, + status text NOT NULL, trigger text NOT NULL ); @@ -114,6 +117,7 @@ function createSessionModelCallDatabase(): SqliteD1Database { native_call_id text, output_tokens integer, provider text NOT NULL, + source_event_seq integer DEFAULT 0 NOT NULL, session_id text NOT NULL, session_run_id text NOT NULL, started_at integer, @@ -148,10 +152,18 @@ function createSessionModelCallDatabase(): SqliteD1Database { session_run_id text, source text NOT NULL, source_event_id text NOT NULL, + source_event_seq integer DEFAULT 0 NOT NULL, total_cost_usd_micros integer NOT NULL, usage_contract text NOT NULL, UNIQUE (source, source_event_id) ); + + CREATE TABLE usage_event_rollup_receipt ( + rolled_up_at integer NOT NULL, + source text NOT NULL, + source_event_id text NOT NULL, + PRIMARY KEY (source, source_event_id) + ); `); return database; @@ -257,6 +269,7 @@ async function seedRunIdentity( INSERT INTO session_run ( agent_id, completed_at, + created_at, created_by_account_id, deployment_version_id, id, @@ -265,9 +278,10 @@ async function seedRunIdentity( runtime_id, session_id, started_at, + status, trigger ) - VALUES (?, 1800, ?, ?, ?, 'gpt-5.4', 'openai', 'openai-runtime', ?, 1200, 'user_prompt') + VALUES (?, 1800, 1000, ?, ?, ?, 'gpt-5.4', 'openai', 'openai-runtime', ?, 1200, 'completed', 'user_prompt') `, ) .bind(AGENT_ID, createdByAccountId, deploymentVersionId, SESSION_RUN_ID, SESSION_ID) @@ -293,10 +307,11 @@ describe("session model call identity", () => { } satisfies SessionUsageSummary; await upsertSessionModelCallUsage(database, { + createdAtMs: USAGE_CREATED_AT_MS, driverInstanceId: DRIVER_INSTANCE_ID, sessionId: SESSION_ID, sessionRunId: SESSION_RUN_ID, - status: "completed", + sourceEventSeq: 5, traceId: "trace-1", usage, }); @@ -377,10 +392,11 @@ describe("session model call identity", () => { await expect( upsertSessionModelCallUsage(database, { + createdAtMs: USAGE_CREATED_AT_MS, driverInstanceId: DRIVER_INSTANCE_ID, sessionId: SESSION_ID, sessionRunId: SESSION_RUN_ID, - status: "completed", + sourceEventSeq: 5, traceId: "trace-wrong-app", usage, }), @@ -398,10 +414,11 @@ describe("session model call identity", () => { usageContract: "openai_total_with_cached_breakdown", } satisfies SessionUsageSummary; const input = { + createdAtMs: USAGE_CREATED_AT_MS, driverInstanceId: DRIVER_INSTANCE_ID, sessionId: SESSION_ID, sessionRunId: SESSION_RUN_ID, - status: "completed" as const, + sourceEventSeq: 5, traceId: "trace-atomic-ledger", usage, }; @@ -435,4 +452,247 @@ describe("session model call identity", () => { .first<{ count: number }>(), ).toEqual({ count: 1 }); }); + + test.each([ + ["running", null, "started"], + ["waiting_input", null, "started"], + ["completed", 1_800, "completed"], + ["failed", 1_800, "failed"], + ["cancelled", 1_800, "failed"], + ["expired", 1_800, "failed"], + ] as const)( + "derives model-call status from durable Run status %s", + async (runStatus, completedAt, expectedStatus) => { + const database = createSessionModelCallDatabase(); + await seedRunIdentity(database); + await database + .prepare("UPDATE session_run SET completed_at = ?, status = ? WHERE id = ?") + .bind(completedAt, runStatus, SESSION_RUN_ID) + .run(); + + await upsertSessionModelCallUsage(database, { + createdAtMs: USAGE_CREATED_AT_MS, + driverInstanceId: DRIVER_INSTANCE_ID, + sessionId: SESSION_ID, + sessionRunId: SESSION_RUN_ID, + sourceEventSeq: 5, + traceId: "trace-status", + usage: { + callId: "status-call", + inputTokens: 10, + source: "prompt_response", + usageContract: "openai_total_with_cached_breakdown", + }, + }); + + expect( + await database + .prepare("SELECT completed_at, status FROM session_model_call") + .first<{ completed_at: number | null; status: string }>(), + ).toEqual({ completed_at: completedAt, status: expectedStatus }); + }, + ); + + test("converges model-call and usage rows by durable event seq", async () => { + const database = createSessionModelCallDatabase(); + await seedRunIdentity(database); + const input = { + createdAtMs: USAGE_CREATED_AT_MS, + driverInstanceId: DRIVER_INSTANCE_ID, + sessionId: SESSION_ID, + sessionRunId: SESSION_RUN_ID, + sourceEventSeq: 5, + traceId: "trace-sequenced", + usage: { + callId: "sequenced-call", + inputTokens: 10, + outputTokens: 5, + source: "prompt_response" as const, + usageContract: "openai_total_with_cached_breakdown" as const, + }, + }; + + await upsertSessionModelCallUsage(database, input); + await upsertSessionModelCallUsage(database, { + ...input, + sourceEventSeq: 4, + usage: { ...input.usage, inputTokens: 1 }, + }); + await upsertSessionModelCallUsage(database, input); + await expect( + upsertSessionModelCallUsage(database, { + ...input, + usage: { ...input.usage, inputTokens: 20 }, + }), + ).rejects.toThrow("replayed with conflicting content"); + await upsertSessionModelCallUsage(database, { + ...input, + sourceEventSeq: 6, + usage: { ...input.usage, inputTokens: 30 }, + }); + + expect( + await database + .prepare( + "SELECT input_tokens, source_event_seq, status FROM session_model_call WHERE call_key = ?", + ) + .bind("model_call:sequenced-call") + .first<{ input_tokens: number; source_event_seq: number; status: string }>(), + ).toEqual({ input_tokens: 30, source_event_seq: 6, status: "completed" }); + expect( + await database + .prepare("SELECT input_tokens, source_event_seq FROM usage_event WHERE source_event_id = ?") + .bind(`${DRIVER_INSTANCE_ID}:sequenced-call`) + .first<{ input_tokens: number; source_event_seq: number }>(), + ).toEqual({ input_tokens: 30, source_event_seq: 6 }); + }); + + test("merges higher partial usage identically into the model call and ledger", async () => { + const database = createSessionModelCallDatabase(); + await seedRunIdentity(database); + const input = { + createdAtMs: USAGE_CREATED_AT_MS, + driverInstanceId: DRIVER_INSTANCE_ID, + sessionId: SESSION_ID, + sessionRunId: SESSION_RUN_ID, + sourceEventSeq: 5, + traceId: "trace-partial", + usage: { + callId: "partial-call", + inputTokens: 10, + outputTokens: 5, + source: "prompt_response" as const, + usageContract: "openai_total_with_cached_breakdown" as const, + }, + }; + + const partialInput = { + ...input, + createdAtMs: USAGE_CREATED_AT_MS + 100, + sourceEventSeq: 6, + usage: { + callId: "partial-call", + outputTokens: 7, + source: "prompt_response", + usageContract: "openai_total_with_cached_breakdown", + }, + }; + + await upsertSessionModelCallUsage(database, input); + await upsertSessionModelCallUsage(database, partialInput); + await upsertSessionModelCallUsage(database, partialInput); + + expect( + await database + .prepare( + `SELECT created_at, input_tokens, output_tokens, source_event_seq + FROM session_model_call WHERE call_key = ?`, + ) + .bind("model_call:partial-call") + .first(), + ).toEqual({ + created_at: USAGE_CREATED_AT_MS, + input_tokens: 10, + output_tokens: 7, + source_event_seq: 6, + }); + expect( + await database + .prepare( + `SELECT created_at, input_tokens, output_tokens, source_event_seq + FROM usage_event WHERE source_event_id = ?`, + ) + .bind(`${DRIVER_INSTANCE_ID}:partial-call`) + .first(), + ).toEqual({ + created_at: USAGE_CREATED_AT_MS, + input_tokens: 10, + output_tokens: 7, + source_event_seq: 6, + }); + + await database + .prepare("UPDATE session_run SET model = 'gpt-5.5' WHERE id = ?") + .bind(SESSION_RUN_ID) + .run(); + await expect( + upsertSessionModelCallUsage(database, { + ...partialInput, + sourceEventSeq: 7, + usage: { ...partialInput.usage, outputTokens: 9 }, + }), + ).rejects.toThrow(); + + expect( + await database + .prepare( + `SELECT model, output_tokens, source_event_seq + FROM session_model_call WHERE call_key = ?`, + ) + .bind("model_call:partial-call") + .first(), + ).toEqual({ model: "gpt-5.4", output_tokens: 7, source_event_seq: 6 }); + expect( + await database + .prepare( + `SELECT model, output_tokens, source_event_seq + FROM usage_event WHERE source_event_id = ?`, + ) + .bind(`${DRIVER_INSTANCE_ID}:partial-call`) + .first(), + ).toEqual({ model: "gpt-5.4", output_tokens: 7, source_event_seq: 6 }); + }); + + test("never replaces raw usage after its durable call was rolled up", async () => { + const database = createSessionModelCallDatabase(); + await seedRunIdentity(database); + const input = { + createdAtMs: USAGE_CREATED_AT_MS, + driverInstanceId: DRIVER_INSTANCE_ID, + sessionId: SESSION_ID, + sessionRunId: SESSION_RUN_ID, + sourceEventSeq: 5, + traceId: "trace-rolled", + usage: { + callId: "rolled-call", + inputTokens: 10, + outputTokens: 5, + source: "prompt_response" as const, + usageContract: "openai_total_with_cached_breakdown" as const, + }, + }; + + await upsertSessionModelCallUsage(database, input); + await database + .prepare( + "INSERT INTO usage_event_rollup_receipt (source, source_event_id, rolled_up_at) VALUES (?, ?, ?)", + ) + .bind("runtime_driver", `${DRIVER_INSTANCE_ID}:rolled-call`, 2_000) + .run(); + await database + .prepare("DELETE FROM usage_event WHERE source_event_id = ?") + .bind(`${DRIVER_INSTANCE_ID}:rolled-call`) + .run(); + + await upsertSessionModelCallUsage(database, input); + await expect( + upsertSessionModelCallUsage(database, { + ...input, + sourceEventSeq: 6, + usage: { ...input.usage, inputTokens: 20 }, + }), + ).rejects.toThrow("already rolled up"); + + expect( + await database + .prepare("SELECT source_event_seq FROM session_model_call WHERE call_key = ?") + .bind("model_call:rolled-call") + .first<{ source_event_seq: number }>(), + ).toEqual({ source_event_seq: 5 }); + expect( + await database + .prepare("SELECT COUNT(*) AS count FROM usage_event") + .first<{ count: number }>(), + ).toEqual({ count: 0 }); + }); }); diff --git a/apps/api/tests/session-process-events.test.ts b/apps/api/tests/session-process-events.test.ts index 123ac20d..b93e89e0 100644 --- a/apps/api/tests/session-process-events.test.ts +++ b/apps/api/tests/session-process-events.test.ts @@ -96,6 +96,7 @@ function createProcessEventQueryDatabase(): SqliteD1Database { run_id text, seq integer NOT NULL, session_id text NOT NULL, + stream_id text, tokens integer, visibility text NOT NULL ); @@ -200,12 +201,18 @@ async function insertSessionProcessEvent( runId?: string | null; seq: number; sessionId?: string; + streamId?: string | null; tokens?: number | null; visibility?: SessionRuntimeEventVisibility; }, ): Promise { const occurredAt = input.occurredAt ?? input.seq * 1000; const processType = input.processType ?? "run.started"; + const eventType = input.eventType ?? processType; + const streamId = + input.streamId === undefined && /^(?:message|thought)\./u.test(eventType) + ? "stream-1" + : (input.streamId ?? null); await database .prepare( @@ -221,22 +228,24 @@ async function insertSessionProcessEvent( run_id, seq, session_id, + stream_id, tokens, visibility - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( input.id, input.content ?? `run-${input.seq}`, input.endedAt ?? occurredAt, - input.eventType ?? processType, + eventType, occurredAt, input.processStatus ?? "available", processType, input.runId ?? null, input.seq, input.sessionId ?? SESSION_ID, + streamId, input.tokens ?? null, input.visibility ?? "all_consumers", ) @@ -498,6 +507,464 @@ describe("session process event projection", () => { }); }); + test("pages raw stream rows until the latest logical events are complete", async () => { + const database = createProcessEventQueryDatabase(); + const fragmentCount = 1_005; + await insertSessionProcessEvent(database, { + content: "run-1", + eventType: "run.started", + id: "event-run-started", + runId: "run-1", + seq: 1, + }); + + for (let index = 0; index < fragmentCount; index += 1) { + await insertSessionProcessEvent(database, { + content: "x", + eventType: "message.delta", + id: `event-message-delta-${index}`, + processType: "agent.message.delta", + runId: "run-1", + seq: index + 2, + }); + } + + await insertSessionProcessEvent(database, { + content: "Message updated.", + eventType: "message.completed", + id: "event-message-completed", + processType: "agent.message.delta", + runId: "run-1", + seq: fragmentCount + 2, + }); + await insertSessionProcessEvent(database, { + content: "run-1", + eventType: "run.completed", + id: "event-run-completed", + processType: "run.completed", + runId: "run-1", + seq: fragmentCount + 3, + }); + + const events = await getThreadSessionProcessEvents( + database, + VIEWER, + { appId: APP_ID, sessionId: SESSION_ID }, + { limit: 2 }, + ); + + expect(events.map((event) => event.type)).toEqual([ + "session.status", + "agent.message.delta", + "run.completed", + ]); + expect(events[1]).toMatchObject({ + content: "x".repeat(fragmentCount), + id: "event-message-completed", + }); + }); + + test("does not cut a retained stream at a page boundary", async () => { + const database = createProcessEventQueryDatabase(); + const fragmentCount = 550; + await insertSessionProcessEvent(database, { + content: "run-1", + eventType: "run.started", + id: "event-run-started", + runId: "run-1", + seq: 1, + }); + await insertSessionProcessEvent(database, { + content: "Agent thinking updated.", + eventType: "thought.started", + id: "event-thought-a-started", + processType: "agent.thinking.delta", + runId: "run-1", + seq: 2, + streamId: "thought-a", + }); + await insertSessionProcessEvent(database, { + content: "Agent thinking updated.", + eventType: "thought.started", + id: "event-thought-b-started", + processType: "agent.thinking.delta", + runId: "run-1", + seq: 3, + streamId: "thought-b", + }); + + for (let index = 0; index < fragmentCount; index += 1) { + await insertSessionProcessEvent(database, { + content: "a", + eventType: "thought.delta", + id: `event-thought-a-delta-${index}`, + processType: "agent.thinking.delta", + runId: "run-1", + seq: index * 2 + 4, + streamId: "thought-a", + }); + await insertSessionProcessEvent(database, { + content: "b", + eventType: "thought.delta", + id: `event-thought-b-delta-${index}`, + processType: "agent.thinking.delta", + runId: "run-1", + seq: index * 2 + 5, + streamId: "thought-b", + }); + } + + const firstTerminalSeq = fragmentCount * 2 + 4; + await insertSessionProcessEvent(database, { + content: "Agent thinking updated.", + eventType: "thought.completed", + id: "event-thought-a-completed", + processType: "agent.thinking.delta", + runId: "run-1", + seq: firstTerminalSeq, + streamId: "thought-a", + }); + await insertSessionProcessEvent(database, { + content: "Agent thinking updated.", + eventType: "thought.completed", + id: "event-thought-b-completed", + processType: "agent.thinking.delta", + runId: "run-1", + seq: firstTerminalSeq + 1, + streamId: "thought-b", + }); + await insertSessionProcessEvent(database, { + content: "run-1", + eventType: "run.completed", + id: "event-run-completed", + processType: "run.completed", + runId: "run-1", + seq: firstTerminalSeq + 2, + }); + + const events = await getThreadSessionProcessEvents( + database, + VIEWER, + { appId: APP_ID, sessionId: SESSION_ID }, + { limit: 2 }, + ); + + expect(events.map((event) => event.type)).toEqual([ + "session.status", + "agent.thinking.delta", + "run.completed", + ]); + expect(events[1]).toMatchObject({ + content: "b".repeat(fragmentCount), + id: "event-thought-b-completed", + }); + }); + + test("does not expose an incomplete stream when the raw scan ceiling is reached", async () => { + const database = createProcessEventQueryDatabase(); + const fragmentCount = 20_001; + await insertSessionProcessEvent(database, { + content: "Agent thinking updated.", + eventType: "thought.started", + id: "event-thought-started", + processType: "agent.thinking.delta", + runId: "run-1", + seq: 1, + streamId: "ceiling-thought", + }); + + for (let index = 0; index < fragmentCount; index += 1) { + await insertSessionProcessEvent(database, { + content: "x", + eventType: "thought.delta", + id: `event-thought-delta-${index}`, + processType: "agent.thinking.delta", + runId: "run-1", + seq: index + 2, + streamId: "ceiling-thought", + }); + } + + await insertSessionProcessEvent(database, { + content: "Agent thinking updated.", + eventType: "thought.completed", + id: "event-thought-completed", + processType: "agent.thinking.delta", + runId: "run-1", + seq: fragmentCount + 2, + streamId: "ceiling-thought", + }); + await insertSessionProcessEvent(database, { + content: "Complete final answer", + eventType: "message.added", + id: "event-final-message", + processType: "agent.message.delta", + runId: "run-1", + seq: fragmentCount + 3, + streamId: "complete-final-message", + }); + await insertSessionProcessEvent(database, { + content: "run-1", + eventType: "run.completed", + id: "event-run-completed", + processType: "run.completed", + runId: "run-1", + seq: fragmentCount + 4, + }); + + const events = await getThreadSessionProcessEvents( + database, + VIEWER, + { appId: APP_ID, sessionId: SESSION_ID }, + { limit: 2 }, + ); + + expect(events.map((event) => event.type)).toEqual([ + "session.status", + "agent.message.delta", + "run.completed", + ]); + expect(events[0]?.content).toContain("Earlier runtime events are hidden"); + expect(events[1]?.content).toBe("Complete final answer"); + }, 15_000); + + test("recognizes the database start at the exact raw scan ceiling", async () => { + const database = createProcessEventQueryDatabase(); + const fillerCount = 19_997; + + for (let seq = 1; seq <= fillerCount; seq += 1) { + await insertSessionProcessEvent(database, { + content: `filler-${seq}`, + id: `exact-ceiling-filler-${seq}`, + processType: "usage.updated", + seq, + }); + } + + await insertSessionProcessEvent(database, { + content: "retained", + eventType: "message.delta", + id: "exact-ceiling-message-delta", + processType: "agent.message.delta", + runId: null, + seq: fillerCount + 1, + streamId: "exact-ceiling-message", + }); + await insertSessionProcessEvent(database, { + content: "Message updated.", + eventType: "message.completed", + id: "exact-ceiling-message-completed", + processType: "agent.message.delta", + runId: null, + seq: fillerCount + 2, + streamId: "exact-ceiling-message", + }); + await insertSessionProcessEvent(database, { + content: "run-1", + eventType: "run.completed", + id: "exact-ceiling-run-completed", + processType: "run.completed", + runId: "run-1", + seq: fillerCount + 3, + }); + + const events = await getThreadSessionProcessEvents( + database, + VIEWER, + { appId: APP_ID, sessionId: SESSION_ID }, + { limit: 2 }, + ); + + expect(events.map((event) => event.type)).toEqual([ + "session.status", + "agent.message.delta", + "run.completed", + ]); + expect(events[1]?.content).toBe("retained"); + }, 15_000); + + test("orders paged events by durable sequence rather than driver time", async () => { + const database = createProcessEventQueryDatabase(); + const fragmentCount = 1_002; + await insertSessionProcessEvent(database, { + content: "run-1", + eventType: "run.started", + id: "out-of-order-run-started", + occurredAt: 1_000, + processType: "run.started", + runId: "run-1", + seq: 1, + }); + await insertSessionProcessEvent(database, { + content: "Message updated.", + eventType: "message.started", + id: "out-of-order-message-started", + occurredAt: 100_000, + processType: "agent.message.delta", + runId: "run-1", + seq: 2, + streamId: "out-of-order-message", + }); + + for (let index = 0; index < fragmentCount; index += 1) { + await insertSessionProcessEvent(database, { + content: "x", + eventType: "message.delta", + id: `out-of-order-message-delta-${index}`, + occurredAt: 101_000 + index, + processType: "agent.message.delta", + runId: "run-1", + seq: index + 3, + streamId: "out-of-order-message", + }); + } + + await insertSessionProcessEvent(database, { + content: "Read file", + eventType: "tool.call.updated", + id: "out-of-order-tool", + occurredAt: 2_000, + processType: "tool.use.started", + runId: "run-1", + seq: fragmentCount + 3, + }); + await insertSessionProcessEvent(database, { + content: "run-1", + eventType: "run.completed", + id: "out-of-order-run-completed", + occurredAt: 3_000, + processType: "run.completed", + runId: "run-1", + seq: fragmentCount + 4, + }); + + const events = await getThreadSessionProcessEvents( + database, + VIEWER, + { appId: APP_ID, sessionId: SESSION_ID }, + { limit: 2 }, + ); + + expect(events.map((event) => event.type)).toEqual([ + "session.status", + "tool.use.started", + "run.completed", + ]); + expect(events[1]?.id).toBe("out-of-order-tool"); + }); + + test.each([ + ["message.cancelled", "available"], + ["message.failed", "error"], + ] as const)("folds persisted streams closed by %s", async (eventType, processStatus) => { + const database = createProcessEventQueryDatabase(); + await insertSessionProcessEvent(database, { + content: "Partial answer", + eventType: "message.delta", + id: "event-message-delta", + processType: "agent.message.delta", + runId: "run-1", + seq: 1, + }); + await insertSessionProcessEvent(database, { + content: "Message updated.", + eventType, + id: "event-message-end", + processStatus, + processType: "agent.message.delta", + runId: "run-1", + seq: 2, + }); + + const events = await getThreadSessionProcessEvents( + database, + VIEWER, + { appId: APP_ID, sessionId: SESSION_ID }, + { limit: 100 }, + ); + + expect(events).toEqual([ + expect.objectContaining({ + content: "Partial answer", + id: "event-message-end", + status: processStatus, + type: "agent.message.delta", + }), + ]); + }); + + test("keeps a terminal message without deltas", async () => { + const database = createProcessEventQueryDatabase(); + await insertSessionProcessEvent(database, { + content: "Message updated.", + eventType: "message.completed", + id: "event-message-completed", + processType: "agent.message.delta", + runId: "run-1", + seq: 1, + }); + + const events = await getThreadSessionProcessEvents( + database, + VIEWER, + { appId: APP_ID, sessionId: SESSION_ID }, + { limit: 100 }, + ); + + expect(events).toEqual([ + expect.objectContaining({ + content: "", + id: "event-message-completed", + type: "agent.message.delta", + }), + ]); + }); + + test("folds interleaved messages by stream identity", async () => { + const database = createProcessEventQueryDatabase(); + + for (const event of [ + { content: "A1", id: "a-1", seq: 1, streamId: "message-a" }, + { content: "B1", id: "b-1", seq: 2, streamId: "message-b" }, + { content: "A2", id: "a-2", seq: 3, streamId: "message-a" }, + ]) { + await insertSessionProcessEvent(database, { + ...event, + eventType: "message.delta", + processType: "agent.message.delta", + runId: "run-1", + }); + } + + await insertSessionProcessEvent(database, { + content: "Message updated.", + eventType: "message.completed", + id: "b-end", + processType: "agent.message.delta", + runId: "run-1", + seq: 4, + streamId: "message-b", + }); + await insertSessionProcessEvent(database, { + content: "Message updated.", + eventType: "message.completed", + id: "a-end", + processType: "agent.message.delta", + runId: "run-1", + seq: 5, + streamId: "message-a", + }); + + const events = await getThreadSessionProcessEvents( + database, + VIEWER, + { appId: APP_ID, sessionId: SESSION_ID }, + { limit: 100 }, + ); + + expect(events.map((event) => event.content)).toEqual(["A1A2", "B1"]); + }); + test("shows an in-flight streamed message as a single folded process event", async () => { const database = createProcessEventQueryDatabase(); await insertSessionProcessEvent(database, { @@ -546,11 +1013,20 @@ describe("session process event projection", () => { seq: 1, visibility: "owner_debug", }); + await insertSessionProcessEvent(innerDatabase, { + content: "1 background task active.", + eventType: "agent.tasks.replaced", + id: "event-task-state", + occurredAt: 950, + processType: "session.status", + seq: 2, + visibility: "owner_debug", + }); await insertSessionProcessEvent(innerDatabase, { content: "run-1", id: "event-public", occurredAt: 1000, - seq: 2, + seq: 3, }); const events = await getThreadSessionProcessEvents( diff --git a/apps/api/tests/session-resource-files.test.ts b/apps/api/tests/session-resource-files.test.ts index 20315008..d334943c 100644 --- a/apps/api/tests/session-resource-files.test.ts +++ b/apps/api/tests/session-resource-files.test.ts @@ -12,13 +12,6 @@ import { listSessionResources, } from "../src/modules/sessions/application/session-resource.service"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; -import { - PUBLIC_API_TEST_IDS, - PublicApiMemoryFileBucket, - createPublicHttpContractDatabase, - createPublicHttpTestBindings, - insertOwnerSession, -} from "./helpers/public-api-http-test-fixture"; import { SqliteD1Database } from "./helpers/sqlite-d1"; const OWNER_ID = parsePlatformId("01J00000000000000000000001", "owner ID"); @@ -99,6 +92,7 @@ function createSessionResourceDatabase(input: { includeFile?: boolean } = {}): S parent_path text NOT NULL, path text NOT NULL, purpose text NOT NULL, + runtime_event_seq integer, size integer NOT NULL, updated_at integer NOT NULL, version integer NOT NULL @@ -573,55 +567,6 @@ describe("session resource files", () => { expect(artifacts.files.map((file) => file.id)).toEqual([ARTIFACT_FILE_ID]); }); - test("records runtime outputs as session-scoped artifacts", async () => { - const database = await createPublicHttpContractDatabase(); - await insertOwnerSession(database); - const bucket = new PublicApiMemoryFileBucket(); - const ownerViewer: AuthenticatedViewer = { - email: "owner@example.com", - emailVerified: true, - id: PUBLIC_API_TEST_IDS.ownerAccount, - imageUrl: null, - name: "Owner", - }; - const file = await fileStore.recordRuntimeOutput({ - bindings: createPublicHttpTestBindings(database, { - fileBucket: bucket as unknown as R2Bucket, - }) as ApiBindings, - body: new TextEncoder().encode("runtime summary"), - contentType: "text/markdown", - createdBy: PUBLIC_API_TEST_IDS.ownerAccount, - path: "outputs/reports/summary.md", - sessionId: PUBLIC_API_TEST_IDS.ownerSession, - }); - - expect(file.owner).toEqual({ - id: PUBLIC_API_TEST_IDS.ownerSession, - kind: "session", - }); - expect(file.purpose).toBe("session_artifact"); - expect(file.scope).toEqual({ - id: PUBLIC_API_TEST_IDS.ownerSession, - kind: "session", - }); - expect(file.sessionKind).toBe("artifact"); - expect(file.sourcePath).toBe("outputs/reports/summary.md"); - - const resources = await listSessionResources(database, ownerViewer, { - appId: PUBLIC_API_TEST_IDS.app, - sessionId: PUBLIC_API_TEST_IDS.ownerSession, - }); - - expect(resources).toEqual([ - expect.objectContaining({ - id: file.id, - kind: "artifact", - name: "summary.md", - path: `session-artifacts/${file.id}/summary.md`, - }), - ]); - }); - test("lists session artifacts but does not treat them as removable attachments", async () => { const database = createSessionResourceDatabase(); insertSessionArtifact(database); diff --git a/apps/api/tests/session-resource-mount.test.ts b/apps/api/tests/session-resource-mount.test.ts index ac4a7df2..83d5a621 100644 --- a/apps/api/tests/session-resource-mount.test.ts +++ b/apps/api/tests/session-resource-mount.test.ts @@ -75,7 +75,7 @@ function createSandbox(input: { async destroy() {}, async exec(command) { calls.exec.push(command); - return commandResult(input.pathExists); + return commandResult(command.includes("readlink") || input.pathExists); }, async getSession() { throw new Error("getSession is not used in session resource mount tests."); @@ -123,7 +123,8 @@ describe("ensureSessionResourcesMounted", () => { sessionId: "session-local", }); - expect(sandbox.calls.exec).toHaveLength(1); + expect(sandbox.calls.exec).toHaveLength(2); + expect(sandbox.calls.exec[1]).toContain("../../.mosoo/session-files/session-local"); expect(sandbox.calls.mkdir).toEqual([]); expect(sandbox.calls.mountBucket).toEqual([]); }); @@ -137,7 +138,8 @@ describe("ensureSessionResourcesMounted", () => { sessionId: "session-remote", }); - expect(sandbox.calls.exec).toHaveLength(1); + expect(sandbox.calls.exec).toHaveLength(2); + expect(sandbox.calls.exec[1]).toContain("../../.mosoo/session-files/session-remote"); expect(sandbox.calls.mkdir).toHaveLength(1); expect(sandbox.calls.mountBucket).toHaveLength(1); expect(sandbox.calls.mountBucket[0]).toEqual({ @@ -170,7 +172,8 @@ describe("ensureSessionResourcesMounted", () => { sessionId: "session-remote", }); - expect(sandbox.calls.exec).toHaveLength(2); + expect(sandbox.calls.exec).toHaveLength(3); + expect(sandbox.calls.exec[2]).toContain("../../.mosoo/session-files/session-remote"); expect(sandbox.calls.mkdir).toHaveLength(1); expect(sandbox.calls.mountBucket.map((call) => call.mountPath)).toEqual(sandbox.calls.mkdir); }); diff --git a/apps/api/tests/session-run-admission-atomicity.test.ts b/apps/api/tests/session-run-admission-atomicity.test.ts index 9be8b6dd..5c484648 100644 --- a/apps/api/tests/session-run-admission-atomicity.test.ts +++ b/apps/api/tests/session-run-admission-atomicity.test.ts @@ -1,12 +1,21 @@ import { describe, expect, test } from "bun:test"; -import { parsePlatformId } from "@mosoo/id"; -import type { AgentDeploymentVersionId, SessionId, SessionRunId } from "@mosoo/id"; +import { createPlatformId, parsePlatformId } from "@mosoo/id"; +import type { + AgentDeploymentVersionId, + SessionId, + SessionMessageId, + SessionRunId, +} from "@mosoo/id"; +import { createRuntimeEventSemanticHash } from "@mosoo/runtime-events"; import { API_COMMAND_QUEUE_SEND_FAILED_CODE } from "../src/modules/api-command/application/api-command-ledger"; import { getAccountViewer } from "../src/modules/auth/application/public-api-caller.service"; import type { AuthenticatedViewer } from "../src/modules/auth/application/viewer-auth.service"; import { queueSessionRun } from "../src/modules/runtime/application/session-run.service"; +import { recordCanonicalSessionRunTerminal } from "../src/modules/runtime/application/session-runs/session-run-terminal-failure.service"; +import { createQueuedSessionRunRuntimeEvents } from "../src/modules/runtime/application/session-runs/session-run-view-events.service"; +import { prepareAssistantMessageProjection } from "../src/modules/runtime/infrastructure/driver-instance/assistant-message-projection"; import { setSessionRunStatus } from "../src/modules/runtime/infrastructure/session-runs/session-run-store.repository"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { API_ERROR_CODE } from "../src/platform/errors"; @@ -18,6 +27,7 @@ import { insertOwnerSession, } from "./helpers/public-api-http-test-fixture"; import type { ApiCommandQueueStub, SqliteD1Database } from "./helpers/public-api-http-test-fixture"; +import { insertRuntimeEvent } from "./public-thread-api-fixtures"; interface AdmissionCounts { apiCommand: number; @@ -186,8 +196,12 @@ async function createFixture() { return { database, viewer }; } -async function completeRun(database: D1Database, runId: SessionRunId): Promise { - for (const status of ["booting", "running", "completed"] as const) { +async function completeRun( + bindings: ApiBindings, + database: SqliteD1Database, + runId: SessionRunId, +): Promise { + for (const status of ["booting", "running"] as const) { const outcome = await setSessionRunStatus(database, { runId, source: "driver", @@ -195,74 +209,200 @@ async function completeRun(database: D1Database, runId: SessionRunId): Promise(); + await insertRuntimeEvent(database, { + kind: "message.added", + occurredAt: 3, + payload: { content: "done", messageId: finalMessageId, role: "agent" }, + runId, + seq: 3, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + }); + await insertRuntimeEvent(database, { + kind: "message.completed", + occurredAt: 4, + payload: { messageId: finalMessageId, role: "agent" }, + runId, + seq: 4, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + }); + await database + .prepare("UPDATE session SET runtime_event_seq_cursor = 4 WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .run(); + await recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: prepareAssistantMessageProjection({ + createdByAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + messageId: finalMessageId, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + sessionRunId: runId, + }), + error: null, + runId, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + source: "driver", + status: "completed", + timestampMs: 5, + }); } describe("Session Run atomic admission", () => { - test("blocks a cattle follow-up until the previous completed Run has a ready checkpoint", async () => { + test("hashes the exact sourced admission events and keeps replay cursors stable", async () => { const { database, viewer } = await createFixture(); - const apiCommandQueue = createApiCommandQueueStub(); - const bindings = createPublicHttpTestBindings(database, { apiCommandQueue }) as ApiBindings; - await database - .prepare("UPDATE session SET kind = 'cattle' WHERE id = ?") - .bind(PUBLIC_API_TEST_IDS.ownerSession) - .run(); - const first = await queueOwnerRun({ - bindings, - clientRequestId: "checkpoint-run-a", - viewer, + const bindings = createPublicHttpTestBindings(database, { + apiCommandQueue: createApiCommandQueueStub(), + }) as ApiBindings; + const clientRequestId = "semantic-admission-request"; + const prompt = "Admit this request atomically."; + const admitted = await queueOwnerRun({ bindings, clientRequestId, viewer }); + const messageId = await database + .prepare("SELECT id FROM session_message WHERE session_run_id = ?") + .bind(admitted.run.id) + .first("id"); + const storedEvents = await database + .prepare( + `SELECT id, occurred_at, semantic_hash, source_event_id + FROM session_event + WHERE run_id = ? + ORDER BY seq`, + ) + .bind(admitted.run.id) + .all<{ + id: string; + occurred_at: number; + semantic_hash: string; + source_event_id: string; + }>(); + + if (messageId === null || storedEvents.results.length !== 2) { + throw new Error("Admission did not persist its canonical message and events."); + } + + const expectedEvents = createQueuedSessionRunRuntimeEvents({ + prompt, + run: admitted.run, + sessionId: parsePlatformId(PUBLIC_API_TEST_IDS.ownerSession, "fixture session"), + sessionMessageId: parsePlatformId(messageId, "fixture message"), }); - await completeRun(database, first.run.id); + const expectedHashes = await Promise.all( + expectedEvents.map((event, index) => { + const stored = storedEvents.results[index]; + if (stored === undefined) { + throw new Error("Admission event identity is missing."); + } + return createRuntimeEventSemanticHash({ + ...event, + id: parsePlatformId(stored.id, "fixture runtime event"), + occurredAt: new Date(stored.occurred_at).toISOString(), + sourceEventId: stored.source_event_id, + }); + }), + ); - await expect( - database - .prepare("SELECT workspace_checkpoint_required FROM session WHERE id = ?") - .bind(PUBLIC_API_TEST_IDS.ownerSession) - .first("workspace_checkpoint_required"), - ).resolves.toBe(1); + expect(storedEvents.results.map((event) => event.semantic_hash)).toEqual(expectedHashes); + expect(storedEvents.results.map((event) => event.source_event_id)).toEqual([ + clientRequestId, + storedEvents.results[1]?.id, + ]); + const beforeReplay = await readSessionState(database); - await expect( - queueOwnerRun({ bindings, clientRequestId: "checkpoint-run-b", viewer }), - ).rejects.toMatchObject({ - code: API_ERROR_CODE.sessionRunCheckpointPending, - message: expect.stringContaining("still committing its previous workspace checkpoint"), + await expect(queueOwnerRun({ bindings, clientRequestId, viewer })).rejects.toMatchObject({ + code: API_ERROR_CODE.sessionRunClientRequestDuplicate, status: 409, }); + await expect(readSessionState(database)).resolves.toEqual(beforeReplay); + await expect(readAdmissionCounts(database)).resolves.toEqual({ + apiCommand: 1, + event: 2, + message: 1, + run: 1, + }); + }); - database.execute(` + test.each([ + ["rejects a checkpoint owned by another workspace", PUBLIC_API_TEST_IDS.nonOwnerSession, false], + ["accepts its ready checkpoint", PUBLIC_API_TEST_IDS.ownerSession, true], + ] as const)( + "%s before admitting a cattle follow-up", + async (_name, workspaceSessionId, ready) => { + const { database, viewer } = await createFixture(); + const apiCommandQueue = createApiCommandQueueStub(); + const bindings = createPublicHttpTestBindings(database, { apiCommandQueue }) as ApiBindings; + await database + .prepare("UPDATE session SET kind = 'cattle' WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .run(); + const first = await queueOwnerRun({ + bindings, + clientRequestId: "checkpoint-run-a", + viewer, + }); + await completeRun(bindings, database, first.run.id); + + await expect( + database + .prepare("SELECT workspace_checkpoint_required FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first("workspace_checkpoint_required"), + ).resolves.toBe(1); + + await expect( + queueOwnerRun({ bindings, clientRequestId: "checkpoint-run-b", viewer }), + ).rejects.toMatchObject({ + code: API_ERROR_CODE.sessionRunCheckpointPending, + message: expect.stringContaining("still committing its previous workspace checkpoint"), + status: 409, + }); + + database.execute(` INSERT INTO sandbox ( - id, kind, subject_kind, subject_id, status, bind_mount_ready, + agent_id, app_id, id, incarnation, kind, network_constraints_hash, + owner_account_id, subject_kind, subject_id, status, bind_mount_ready, global_mounts_json, created_at, updated_at ) VALUES ( - '${PUBLIC_API_TEST_IDS.sandbox}', 'cattle', 'session', '${PUBLIC_API_TEST_IDS.ownerSession}', + '${PUBLIC_API_TEST_IDS.agent}', '${PUBLIC_API_TEST_IDS.app}', + '${PUBLIC_API_TEST_IDS.sandbox}', 1, 'cattle', '${"0".repeat(64)}', + '${PUBLIC_API_TEST_IDS.ownerAccount}', 'session', '${PUBLIC_API_TEST_IDS.ownerSession}', 'active', 1, '[]', 1, 1 ); INSERT INTO sandbox_session ( cloudflare_session_id, created_at, cwd, origin_json, sandbox_id, - session_id, status, updated_at + sandbox_incarnation, session_id, status, updated_at ) VALUES ( '01J0000000000000000000000Z', 1, '/workspace/se/${PUBLIC_API_TEST_IDS.ownerSession}', '{}', - '${PUBLIC_API_TEST_IDS.sandbox}', '${PUBLIC_API_TEST_IDS.ownerSession}', 'active', 1 + '${PUBLIC_API_TEST_IDS.sandbox}', 1, '${PUBLIC_API_TEST_IDS.ownerSession}', 'active', 1 ); INSERT INTO sandbox_backup ( - created_at, dir, id, keep, sandbox_id, session_run_id, status, ttl_seconds, updated_at + created_at, dir, id, keep, sandbox_id, sandbox_incarnation, session_run_id, + staging_id, status, ttl_seconds, updated_at, workspace_session_id ) VALUES ( 1, '/workspace/se/${PUBLIC_API_TEST_IDS.ownerSession}', '${PUBLIC_API_TEST_IDS.operation}', - 0, '${PUBLIC_API_TEST_IDS.sandbox}', '${first.run.id}', 'ready', 315360000, 1 + 0, '${PUBLIC_API_TEST_IDS.sandbox}', 1, '${first.run.id}', + '${PUBLIC_API_TEST_IDS.operation}', 'ready', 315360000, 1, + '${workspaceSessionId}' ); `); - const second = await queueOwnerRun({ - bindings, - clientRequestId: "checkpoint-run-b", - viewer, - }); - expect(second.run.status).toBe("queued"); - }); + const followUp = queueOwnerRun({ + bindings, + clientRequestId: "checkpoint-run-b", + viewer, + }); + if (!ready) { + await expect(followUp).rejects.toMatchObject({ + code: API_ERROR_CODE.sessionRunCheckpointPending, + }); + return; + } + + await expect(followUp).resolves.toMatchObject({ run: { status: "queued" } }); + }, + ); test("grandfathers a completed cattle Run from before the checkpoint rollout", async () => { const { database, viewer } = await createFixture(); @@ -379,6 +519,44 @@ describe("Session Run atomic admission", () => { expect(apiCommandQueue.sent).toHaveLength(1); }); + test("does not admit a Run while provisioning owns the Session", async () => { + const { database, viewer } = await createFixture(); + const apiCommandQueue = createApiCommandQueueStub(); + const bindings = createPublicHttpTestBindings(database, { apiCommandQueue }) as ApiBindings; + + await database + .prepare( + `UPDATE session + SET runtime_provisioning_heartbeat_at = ?, + runtime_provisioning_operation_id = ?, + runtime_provisioning_sandbox_id = ? + WHERE id = ?`, + ) + .bind( + 1, + PUBLIC_API_TEST_IDS.operation, + PUBLIC_API_TEST_IDS.sandbox, + PUBLIC_API_TEST_IDS.ownerSession, + ) + .run(); + + await expect(queueOwnerRun({ bindings, viewer })).rejects.toThrow(); + await expect(readAdmissionCounts(database)).resolves.toEqual({ + apiCommand: 0, + event: 0, + message: 0, + run: 0, + }); + await expect(readSessionState(database)).resolves.toEqual({ + lastMessageAt: null, + lastRunId: null, + messageSeqCursor: 0, + runtimeEventSeqCursor: 0, + status: "IDLE", + }); + expect(apiCommandQueue.sent).toHaveLength(0); + }); + test("classifies a completed client request replay without creating a second Run", async () => { const { database, viewer } = await createFixture(); const apiCommandQueue = createApiCommandQueueStub(); diff --git a/apps/api/tests/session-run-cancel.test.ts b/apps/api/tests/session-run-cancel.test.ts index e6fbd213..b690a99f 100644 --- a/apps/api/tests/session-run-cancel.test.ts +++ b/apps/api/tests/session-run-cancel.test.ts @@ -6,14 +6,20 @@ import type { AccountId, DriverInstanceId, SandboxId, SessionId, SessionRunId } import type { AuthenticatedViewer } from "../src/modules/auth/application/viewer-auth.service"; import { cancelRun } from "../src/modules/runtime/application/session-runs/cancel-run.service"; import { resolvePermissionRequest } from "../src/modules/runtime/application/session-runs/resolve-permission-request.service"; +import { createSessionRunUpdatedEvent } from "../src/modules/runtime/application/session-runs/session-run-view-events.service"; +import { createSessionRunTerminalSourceId } from "../src/modules/runtime/domain/session-run-terminal-event-id"; +import { commitTerminalRunProjection } from "../src/modules/runtime/infrastructure/driver-instance/completed-run-commit.repository"; import { recordRuntimeRunLeaseAcquired } from "../src/modules/runtime/infrastructure/runtime-subject-lifecycle/runtime-run-lease-store"; +import { getSessionRunSummary } from "../src/modules/runtime/infrastructure/session-runs/session-run-store.repository"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { PUBLIC_API_TEST_IDS, createPublicHttpContractDatabase, createPublicHttpTestBindings, + insertActiveSandboxSessionFixture, insertOwnerSession, } from "./helpers/public-api-http-test-fixture"; +import type { SqliteD1Database } from "./helpers/sqlite-d1"; const OWNER_ACCOUNT_ID = parsePlatformId( "01J00000000000000000000001", @@ -38,7 +44,7 @@ const ownerViewer: AuthenticatedViewer = { name: "Owner", }; -function createDriverConnectionBinding(requests: unknown[]) { +function createDriverConnectionBinding(requests: unknown[], onSend?: () => Promise) { return { get: () => ({ fetch: async (request: Request) => { @@ -47,6 +53,7 @@ function createDriverConnectionBinding(requests: unknown[]) { body, path: new URL(request.url).pathname, }); + await onSend?.(); return Response.json({ ok: true }); }, }), @@ -54,58 +61,49 @@ function createDriverConnectionBinding(requests: unknown[]) { }; } -function withDriverConnection(bindings: ApiBindings, requests: unknown[]): ApiBindings { +function withDriverConnection( + bindings: ApiBindings, + requests: unknown[], + onSend?: () => Promise, +): ApiBindings { return { ...bindings, - DriverConnection: createDriverConnectionBinding(requests) as ApiBindings["DriverConnection"], + DriverConnection: createDriverConnectionBinding( + requests, + onSend, + ) as ApiBindings["DriverConnection"], }; } -async function ensureRuntimeLeaseTables(database: D1Database): Promise { - await database - .prepare( - ` - CREATE TABLE IF NOT EXISTS driver_command ( - acked_at integer, - completed_at integer, - delivery_connection_id text, - driver_instance_id text NOT NULL, - error_json text, - expires_at integer, - id text PRIMARY KEY NOT NULL, - issued_at integer NOT NULL, - kind text NOT NULL, - payload_json text NOT NULL, - result_json text, - seq integer NOT NULL, - status text NOT NULL - ) - `, - ) - .run(); - await database - .prepare( - ` - CREATE TABLE IF NOT EXISTS sandbox ( - id text PRIMARY KEY NOT NULL, - inactive_deadline_at integer, - kind text NOT NULL, - updated_at integer NOT NULL - ) - `, - ) - .run(); - await database - .prepare( - ` - CREATE TABLE IF NOT EXISTS sandbox_session ( - sandbox_id text NOT NULL, - session_id text PRIMARY KEY NOT NULL, - status text NOT NULL - ) - `, - ) - .run(); +async function commitDriverTerminalRun( + database: D1Database, + status: "cancelled" | "completed", + timestampMs = Date.now(), +): Promise { + const current = await getSessionRunSummary(database, RUN_ID); + if (current === null) throw new Error("Missing test Session Run."); + const timestamp = new Date(timestampMs).toISOString(); + const run = { + ...current, + completedAt: timestamp, + startedAt: current.startedAt ?? timestamp, + status, + updatedAt: timestamp, + }; + const kind = status === "completed" ? "run.completed" : "run.cancelled"; + const sourceEventId = createSessionRunTerminalSourceId(RUN_ID, kind); + const event = createSessionRunUpdatedEvent(run, OWNER_SESSION_ID, "IDLE", sourceEventId); + + await commitTerminalRunProjection(database, { + assistantMessage: null, + error: null, + runId: RUN_ID, + sessionId: OWNER_SESSION_ID, + source: "driver", + targetStatus: status, + terminalEvent: { event, occurredAt: timestampMs, sourceEventId }, + timestampMs, + }); } async function insertRunningSessionRun( @@ -157,15 +155,28 @@ async function insertRunningSessionRun( } async function insertRunDriverInstance( - database: D1Database, - input: { bindRun?: boolean; sessionId: SessionId; status?: string }, + database: SqliteD1Database, + input: { + bindRun?: boolean; + kind?: "cattle" | "pet"; + sessionId: SessionId; + status?: string; + }, ): Promise { + await insertActiveSandboxSessionFixture(database, { + kind: input.kind, + ownerAccountId: OWNER_ACCOUNT_ID, + sandboxId: SANDBOX_ID, + sessionId: input.sessionId, + }); await database .prepare( ` INSERT INTO driver_instance ( id, + connection_id, sandbox_id, + sandbox_incarnation, sandbox_session_id, runtime, protocol, @@ -179,12 +190,14 @@ async function insertRunDriverInstance( created_at, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( DRIVER_INSTANCE_ID, + "driver-connection-1", SANDBOX_ID, + 1, input.sessionId, "cloudflare-container", "driver-ws", @@ -244,62 +257,19 @@ describe("session run cancel", () => { createdByAccountId: OWNER_ACCOUNT_ID, sessionId: OWNER_SESSION_ID, }); - await ensureRuntimeLeaseTables(database); - await database - .prepare( - ` - INSERT INTO sandbox ( - id, - inactive_deadline_at, - kind, - subject_kind, - subject_id, - status, - created_at, - updated_at - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) - `, - ) - .bind(SANDBOX_ID, 1, "cattle", "session", OWNER_SESSION_ID, "active", 1, 1) - .run(); - await database - .prepare( - ` - INSERT INTO sandbox_session ( - cloudflare_session_id, - created_at, - cwd, - origin_json, - sandbox_id, - session_id, - status, - updated_at - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) - `, - ) - .bind( - "cloudflare-session-1", - 1, - "/workspace", - "{}", - SANDBOX_ID, - OWNER_SESSION_ID, - "active", - 1, - ) - .run(); await insertRunDriverInstance(database, { bindRun: false, + kind: "cattle", sessionId: OWNER_SESSION_ID, status: "provisioning", }); await expect( recordRuntimeRunLeaseAcquired(database, { + driverGeneration: 0, driverInstanceId: DRIVER_INSTANCE_ID, runtimeSubjectId: SANDBOX_ID, + runtimeSubjectIncarnation: 1, sessionId: OWNER_SESSION_ID, sessionRunId: RUN_ID, }), @@ -315,6 +285,7 @@ describe("session run cancel", () => { const bindings = withDriverConnection( createPublicHttpTestBindings(database) as ApiBindings, driverRequests, + () => commitDriverTerminalRun(database, "cancelled"), ); const result = await cancelRun(bindings, ownerViewer, { @@ -327,6 +298,194 @@ describe("session run cancel", () => { expect(driverRequests).toHaveLength(1); }); + test("adopts a Driver completion that wins the cancellation race", async () => { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await insertRunningSessionRun(database, { + createdByAccountId: OWNER_ACCOUNT_ID, + sessionId: OWNER_SESSION_ID, + }); + await insertRunDriverInstance(database, { + bindRun: true, + sessionId: OWNER_SESSION_ID, + }); + const bindings = withDriverConnection( + createPublicHttpTestBindings(database) as ApiBindings, + [], + () => commitDriverTerminalRun(database, "completed"), + ); + + const result = await cancelRun(bindings, ownerViewer, { + appId: PUBLIC_API_TEST_IDS.app, + runId: RUN_ID, + sessionId: OWNER_SESSION_ID, + }); + + expect(result.run.status).toBe("completed"); + const terminals = await database + .prepare( + "SELECT event_type FROM session_event WHERE run_id = ? AND event_type IN ('run.cancelled', 'run.completed', 'run.failed')", + ) + .bind(RUN_ID) + .all<{ event_type: string }>(); + expect(terminals.results).toEqual([{ event_type: "run.completed" }]); + }); + + test("adopts one canonical legacy terminal projection without advancing cursors", async () => { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await insertRunningSessionRun(database, { + createdByAccountId: OWNER_ACCOUNT_ID, + sessionId: OWNER_SESSION_ID, + }); + await commitDriverTerminalRun(database, "cancelled"); + await database + .prepare( + "UPDATE session_event SET semantic_hash = NULL, terminal_event_json = NULL WHERE run_id = ?", + ) + .bind(RUN_ID) + .run(); + const before = await database + .prepare( + "SELECT message_seq_cursor, runtime_event_seq_cursor, status, status_operation_id, status_seq FROM session WHERE id = ?", + ) + .bind(OWNER_SESSION_ID) + .first(); + + await expect(commitDriverTerminalRun(database, "cancelled")).resolves.toBeUndefined(); + + const after = await database + .prepare( + "SELECT message_seq_cursor, runtime_event_seq_cursor, status, status_operation_id, status_seq FROM session WHERE id = ?", + ) + .bind(OWNER_SESSION_ID) + .first(); + const terminals = await database + .prepare("SELECT semantic_hash, source_event_id FROM session_event WHERE run_id = ?") + .bind(RUN_ID) + .all(); + expect(after).toEqual(before); + expect(terminals.results).toEqual([ + { + semantic_hash: null, + source_event_id: createSessionRunTerminalSourceId(RUN_ID, "run.cancelled"), + }, + ]); + }); + + test("does not move the Session lease timestamp backwards when a terminal event is older", async () => { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await insertRunningSessionRun(database, { + createdByAccountId: OWNER_ACCOUNT_ID, + sessionId: OWNER_SESSION_ID, + }); + await database + .prepare("UPDATE session SET updated_at = ? WHERE id = ?") + .bind(2_000, OWNER_SESSION_ID) + .run(); + + await commitDriverTerminalRun(database, "cancelled", 1_000); + + await expect( + database + .prepare("SELECT updated_at FROM session WHERE id = ?") + .bind(OWNER_SESSION_ID) + .first("updated_at"), + ).resolves.toBe(2_000); + }); + + test("rejects a noncanonical legacy terminal projection", async () => { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await insertRunningSessionRun(database, { + createdByAccountId: OWNER_ACCOUNT_ID, + sessionId: OWNER_SESSION_ID, + }); + await commitDriverTerminalRun(database, "cancelled"); + await database + .prepare( + "UPDATE session_event SET semantic_hash = NULL, terminal_event_json = NULL, source_event_id = 'provider-terminal' WHERE run_id = ?", + ) + .bind(RUN_ID) + .run(); + + await expect(commitDriverTerminalRun(database, "cancelled")).rejects.toThrow( + "Legacy terminal projection conflicts", + ); + }); + + test("synthesizes cancellation only when the Driver control socket is confirmed missing", async () => { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await insertRunningSessionRun(database, { + createdByAccountId: OWNER_ACCOUNT_ID, + sessionId: OWNER_SESSION_ID, + }); + await insertRunDriverInstance(database, { + bindRun: true, + sessionId: OWNER_SESSION_ID, + }); + await database + .prepare("UPDATE driver_instance SET connection_id = NULL WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .run(); + const driverRequests: string[] = []; + const bindings = { + ...(createPublicHttpTestBindings(database) as ApiBindings), + DriverConnection: { + get: () => ({ + fetch: async (request: Request) => { + const path = new URL(request.url).pathname; + driverRequests.push(path); + if (path === "/control/send") { + throw new Error("Runtime driver control socket is not connected."); + } + if (path === "/control/fail") { + await database + .prepare("UPDATE driver_instance SET status = 'failed' WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .run(); + } + return Response.json({ ok: true }); + }, + }), + idFromName: (name: string) => name, + } as ApiBindings["DriverConnection"], + }; + + const result = await cancelRun(bindings, ownerViewer, { + appId: PUBLIC_API_TEST_IDS.app, + runId: RUN_ID, + sessionId: OWNER_SESSION_ID, + }); + const run = await database + .prepare("SELECT status, status_source FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first<{ status: string; status_source: string }>(); + + expect(result.run.status).toBe("cancelled"); + expect(run).toEqual({ status: "cancelled", status_source: "viewer" }); + expect(driverRequests).toEqual([ + "/control/send", + "/control/send", + "/control/fail", + "/wait/close", + ]); + await expect( + database + .prepare("SELECT status, status_operation_id FROM driver_instance WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ status: "failed", status_operation_id: null }); + await expect( + database + .prepare("SELECT driver_instance_id FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ driver_instance_id: DRIVER_INSTANCE_ID }); + }); + test("resolves permission requests for the App owner through App ownership", async () => { const database = await createPublicHttpContractDatabase(); await insertOwnerSession(database); @@ -350,6 +509,7 @@ describe("session run cancel", () => { driverInstanceId: DRIVER_INSTANCE_ID, appId: PUBLIC_API_TEST_IDS.app, requestId: "permission-1", + runId: RUN_ID, sessionId: OWNER_SESSION_ID, }), ).resolves.toBeUndefined(); diff --git a/apps/api/tests/session-run-lifecycle.test.ts b/apps/api/tests/session-run-lifecycle.test.ts index a64ffb5e..a69030d5 100644 --- a/apps/api/tests/session-run-lifecycle.test.ts +++ b/apps/api/tests/session-run-lifecycle.test.ts @@ -1,6 +1,5 @@ import { describe, expect, test } from "bun:test"; -import { updateSessionLastRun } from "../src/modules/runtime/infrastructure/session-runs/session-run-session.repository"; import { createSessionRunRecordIfSessionIdle, setSessionRunStatus, @@ -8,16 +7,13 @@ import { import { createPublicHttpContractDatabase, insertNonOwnerSession, + PUBLIC_API_TEST_IDS, } from "./helpers/public-api-http-test-fixture"; -async function insertSessionRun( - database: D1Database, - input: { - runId: string; - sessionId?: string; - status: string; - }, -): Promise { +const SESSION_ID = PUBLIC_API_TEST_IDS.nonOwnerSession; +const RUN_ID = "01J0000000000000000000000T"; + +async function insertRunningSessionRun(database: D1Database): Promise { await database .prepare( ` @@ -39,453 +35,141 @@ async function insertSessionRun( `, ) .bind( - input.runId, - input.sessionId ?? "01J0000000000000000000000B", - "01J00000000000000000000009", - "01J00000000000000000000002", + RUN_ID, + SESSION_ID, + PUBLIC_API_TEST_IDS.agent, + PUBLIC_API_TEST_IDS.nonOwnerAccount, "user_prompt", - input.status, + "running", "openai", "gpt-5.4", "openai-runtime", - `trace-${input.runId}`, + `trace-${RUN_ID}`, 1, 1, ) .run(); await database - .prepare("UPDATE session SET last_run_id = ?, status = ?, updated_at = ? WHERE id = ?") - .bind( - input.runId, - input.status === "completed" ? "IDLE" : "RUNNING", - 1, - "01J0000000000000000000000B", - ) + .prepare("UPDATE session SET last_run_id = ?, status = 'RUNNING', updated_at = 1 WHERE id = ?") + .bind(RUN_ID, SESSION_ID) .run(); } -function failSessionProjectionStatementInBatch(database: D1Database): D1Database { - return new Proxy(database, { - get(target, property, receiver) { - if (property === "batch") { - return async (statements: D1PreparedStatement[]) => { - const firstStatement = statements[0]; - - if (firstStatement === undefined) { - throw new Error("Expected a Run status statement in the D1 batch."); - } - - const failingStatement = new Proxy(target.prepare("SELECT 1"), { - get(statement, statementProperty, statementReceiver) { - if (statementProperty === "run") { - return async () => { - throw new Error("injected Session lifecycle projection failure"); - }; - } - - const value = Reflect.get(statement, statementProperty, statementReceiver); - return typeof value === "function" ? value.bind(statement) : value; - }, - }); - - return target.batch([firstStatement, failingStatement]); - }; - } - - const value = Reflect.get(target, property, receiver); - return typeof value === "function" ? value.bind(target) : value; - }, - }); -} - -function advanceRunBeforeBatch( - database: D1Database, - input: { - readonly runId: string; - }, -): D1Database { - let advanced = false; - - return new Proxy(database, { - get(target, property, receiver) { - if (property === "batch") { - return async (statements: D1PreparedStatement[]) => { - if (!advanced) { - advanced = true; - await setSessionRunStatus(target, { - runId: input.runId, - source: "driver", - status: "running", - }); - } - - return target.batch(statements); - }; - } - - const value = Reflect.get(target, property, receiver); - return typeof value === "function" ? value.bind(target) : value; - }, - }); +function createRunInput() { + return { + agentId: PUBLIC_API_TEST_IDS.agent, + createdBy: PUBLIC_API_TEST_IDS.nonOwnerAccount, + model: "gpt-5.4", + provider: "openai", + runtimeId: "openai-runtime", + sessionId: SESSION_ID, + status: "queued" as const, + trigger: "user_prompt" as const, + }; } -describe("session run lifecycle", () => { - test("emits one structured business log for an applied terminal transition", async () => { +describe("Session Run non-terminal lifecycle writer", () => { + test("fails closed when an untyped caller attempts a terminal transition", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); - await insertSessionRun(database, { - runId: "run-terminal-log", - status: "running", - }); - - const originalConsoleInfo = console.info; - const output: string[] = []; - console.info = (...values: unknown[]) => output.push(values.map(String).join(" ")); - - try { - await setSessionRunStatus(database, { - runId: "run-terminal-log", - source: "driver", - status: "completed", - }); - await setSessionRunStatus(database, { - runId: "run-terminal-log", - source: "driver", - status: "completed", - }); - } finally { - console.info = originalConsoleInfo; - } + await insertRunningSessionRun(database); - const terminalEntries = output - .map((entry) => JSON.parse(entry)) - .filter((entry) => entry.message === "session.run.terminal"); - - expect(terminalEntries).toHaveLength(1); - expect(terminalEntries[0]).toMatchObject({ - level: "info", - metadata: { - errorCode: null, - runId: "run-terminal-log", - runtimeId: "openai-runtime", - sessionType: "ui", + await expect( + setSessionRunStatus(database, { + runId: RUN_ID, source: "driver", status: "completed", - traceId: "trace-run-terminal-log", - trigger: "user_prompt", - }, - namespace: "api", - }); - expect(terminalEntries[0]?.metadata.durationMs).toBeGreaterThanOrEqual(0); + } as never), + ).rejects.toThrow("atomic terminal projection"); + + expect( + await database + .prepare("SELECT status, status_seq FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first(), + ).toEqual({ status: "running", status_seq: 0 }); }); - test("does not let a stale terminal event revive or overwrite a completed run", async () => { + test("keeps duplicate active transitions idempotent", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); - await insertSessionRun(database, { - runId: "run-terminal", - status: "running", - }); + await insertRunningSessionRun(database); - await setSessionRunStatus(database, { - runId: "run-terminal", - source: "driver", - status: "completed", - }); - await setSessionRunStatus(database, { - error: { - code: "runtime.late_failure", - details: {}, - message: "Late failure.", - retryable: false, - }, - runId: "run-terminal", - source: "driver", - status: "failed", - }); - - const row = await database - .prepare( - ` - SELECT error_code, status - FROM session_run - WHERE id = ? - `, - ) - .bind("run-terminal") - .first<{ - error_code: string | null; - status: string; - }>(); - expect(row).toEqual({ - error_code: null, - status: "completed", - }); - }); - - test("leaves duplicate transitions idempotent", async () => { - const database = await createPublicHttpContractDatabase(); - await insertNonOwnerSession(database); - await insertSessionRun(database, { - runId: "run-duplicate", - status: "running", - }); - - await setSessionRunStatus(database, { - runId: "run-duplicate", + const outcome = await setSessionRunStatus(database, { + runId: RUN_ID, source: "driver", status: "running", }); - const row = await database - .prepare("SELECT status FROM session_run WHERE id = ?") - .bind("run-duplicate") - .first<{ status: string }>(); - expect(row).toEqual({ status: "running" }); + expect(outcome.kind).toBe("duplicate"); + expect( + await database + .prepare("SELECT status, status_seq FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first(), + ).toEqual({ status: "running", status_seq: 0 }); }); - test("rejects new runs after the owning session is terminated", async () => { + test("does not advance a current Run after its Session is terminated", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); + await insertRunningSessionRun(database); await database - .prepare("UPDATE session SET status = ? WHERE id = ?") - .bind("TERMINATED", "01J0000000000000000000000B") + .prepare("UPDATE session_run SET status = 'booting' WHERE id = ?") + .bind(RUN_ID) .run(); - - await expect( - createSessionRunRecordIfSessionIdle(database, { - agentId: "01J00000000000000000000009", - createdBy: "01J00000000000000000000002", - model: "gpt-5.4", - provider: "openai", - runtimeId: "openai-runtime", - sessionId: "01J0000000000000000000000B", - status: "queued", - trigger: "user_prompt", - }), - ).rejects.toThrow(); - }); - - test("rejects new runs while a runtime operation owns the session", async () => { - const database = await createPublicHttpContractDatabase(); - await insertNonOwnerSession(database); await database - .prepare( - ` - UPDATE session - SET status = ?, status_operation_id = ? - WHERE id = ? - `, - ) - .bind("RESCHEDULING", "01J0000000000000000000000R", "01J0000000000000000000000B") + .prepare("UPDATE session SET status = 'TERMINATED' WHERE id = ?") + .bind(SESSION_ID) .run(); - await expect( - createSessionRunRecordIfSessionIdle(database, { - agentId: "01J00000000000000000000009", - createdBy: "01J00000000000000000000002", - model: "gpt-5.4", - provider: "openai", - runtimeId: "openai-runtime", - sessionId: "01J0000000000000000000000B", - status: "queued", - trigger: "user_prompt", - }), - ).rejects.toThrow(); - }); - - test("session run projections expose the session as idle after completion", async () => { - const database = await createPublicHttpContractDatabase(); - await insertNonOwnerSession(database); - const run = await createSessionRunRecordIfSessionIdle(database, { - agentId: "01J00000000000000000000009", - createdBy: "01J00000000000000000000002", - model: "gpt-5.4", - provider: "openai", - runtimeId: "openai-runtime", - sessionId: "01J0000000000000000000000B", - status: "running", - trigger: "user_prompt", - }); - if (run.createdRun === null) { - throw new Error("Expected session run creation."); - } - - await setSessionRunStatus(database, { - runId: run.createdRun.id, + const outcome = await setSessionRunStatus(database, { + runId: RUN_ID, source: "driver", - status: "completed", - }); - - const row = await database - .prepare( - ` - SELECT status, status_operation_id - FROM session - WHERE id = ? - `, - ) - .bind("01J0000000000000000000000B") - .first<{ - status: string; - status_operation_id: string | null; - }>(); - - expect(row).toEqual({ - status: "IDLE", - status_operation_id: null, - }); - }); - - test("rolls back a terminal Run transition when its Session projection fails", async () => { - const database = await createPublicHttpContractDatabase(); - await insertNonOwnerSession(database); - await insertSessionRun(database, { - runId: "run-atomic-terminal-projection", status: "running", }); - await expect( - setSessionRunStatus(failSessionProjectionStatementInBatch(database), { - runId: "run-atomic-terminal-projection", - source: "driver", - status: "completed", - }), - ).rejects.toThrow("injected Session lifecycle projection failure"); - - const interrupted = await database - .prepare( - ` - SELECT session.status AS session_status, session_run.status AS run_status - FROM session - INNER JOIN session_run ON session_run.id = session.last_run_id - WHERE session.id = ? - `, - ) - .bind("01J0000000000000000000000B") - .first<{ run_status: string; session_status: string }>(); - - expect(interrupted).toEqual({ - run_status: "running", - session_status: "RUNNING", - }); - - await setSessionRunStatus(database, { - runId: "run-atomic-terminal-projection", - source: "driver", - status: "completed", - }); - - const admitted = await createSessionRunRecordIfSessionIdle(database, { - agentId: "01J00000000000000000000009", - createdBy: "01J00000000000000000000002", - model: "gpt-5.4", - provider: "openai", - runtimeId: "openai-runtime", - sessionId: "01J0000000000000000000000B", - status: "queued", - trigger: "user_prompt", + expect(outcome.kind).toBe("stale"); + expect( + await database + .prepare( + `SELECT session.status AS session_status, + session.status_seq AS session_status_seq, + session_run.status AS run_status, + session_run.status_seq AS run_status_seq + FROM session + JOIN session_run ON session_run.id = session.last_run_id + WHERE session.id = ?`, + ) + .bind(SESSION_ID) + .first(), + ).toEqual({ + run_status: "booting", + run_status_seq: 0, + session_status: "TERMINATED", + session_status_seq: 0, }); - - expect(admitted.createdRun).not.toBeNull(); }); - test("does not project a stale terminal transition onto a newer Run state", async () => { + test("rejects admission after the Session is terminated", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); - await insertSessionRun(database, { - runId: "run-stale-terminal-projection", - status: "booting", - }); - - const outcome = await setSessionRunStatus( - advanceRunBeforeBatch(database, { - runId: "run-stale-terminal-projection", - }), - { - error: { - code: "runtime.stale_terminal", - details: {}, - message: "The stale terminal transition must not update the Session.", - retryable: false, - }, - runId: "run-stale-terminal-projection", - source: "driver", - status: "failed", - }, - ); - - expect(outcome).toMatchObject({ - kind: "stale", - reason: "concurrent_transition", - }); - const current = await database - .prepare( - ` - SELECT session.status AS session_status, session_run.status AS run_status - FROM session - INNER JOIN session_run ON session_run.id = session.last_run_id - WHERE session.id = ? - `, - ) - .bind("01J0000000000000000000000B") - .first<{ run_status: string; session_status: string }>(); + await database + .prepare("UPDATE session SET status = 'TERMINATED' WHERE id = ?") + .bind(SESSION_ID) + .run(); - expect(current).toEqual({ - run_status: "running", - session_status: "RUNNING", - }); + await expect(createSessionRunRecordIfSessionIdle(database, createRunInput())).rejects.toThrow(); }); - test("does not revive terminated sessions from stale run projections", async () => { + test("rejects admission while another operation owns the Session", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); - await insertSessionRun(database, { - runId: "run-stale-session", - status: "running", - }); await database - .prepare("UPDATE session SET status = ? WHERE id = ?") - .bind("TERMINATED", "01J0000000000000000000000B") + .prepare("UPDATE session SET status = 'RESCHEDULING', status_operation_id = ? WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.operation, SESSION_ID) .run(); - await expect( - updateSessionLastRun(database, { - model: "gpt-5.4", - provider: "openai", - runId: "run-stale-session", - sessionId: "01J0000000000000000000000B", - timestampMs: 2, - }), - ).resolves.toBe(false); - await setSessionRunStatus(database, { - error: { - code: "runtime.stale_session", - details: {}, - message: "Stale session.", - retryable: false, - }, - preserveSessionLifecycle: true, - runId: "run-stale-session", - source: "maintenance", - status: "failed", - }); - - const row = await database - .prepare( - ` - SELECT session.status AS session_status, session_run.status AS run_status - FROM session - INNER JOIN session_run ON session_run.id = session.last_run_id - WHERE session.id = ? - `, - ) - .bind("01J0000000000000000000000B") - .first<{ run_status: string; session_status: string }>(); - - expect(row).toEqual({ - run_status: "failed", - session_status: "TERMINATED", - }); + await expect(createSessionRunRecordIfSessionIdle(database, createRunInput())).rejects.toThrow(); }); }); diff --git a/apps/api/tests/session-run-read.test.ts b/apps/api/tests/session-run-read.test.ts index 39b4958f..19fb8a24 100644 --- a/apps/api/tests/session-run-read.test.ts +++ b/apps/api/tests/session-run-read.test.ts @@ -7,6 +7,7 @@ import { import { createPublicHttpContractDatabase, insertNonOwnerSession, + PUBLIC_API_TEST_IDS, } from "./helpers/public-api-http-test-fixture"; async function insertQueuedSessionRun( @@ -16,7 +17,7 @@ async function insertQueuedSessionRun( id?: string; } = {}, ): Promise { - const id = input.id ?? "run-active-probe"; + const id = input.id ?? PUBLIC_API_TEST_IDS.run; const createdAt = input.createdAt ?? 1; await database @@ -68,11 +69,11 @@ describe("session run reads", () => { test("loads the latest active run id", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); - await insertQueuedSessionRun(database, { createdAt: 1, id: "run-active-probe-old" }); - await insertQueuedSessionRun(database, { createdAt: 2, id: "run-active-probe-latest" }); + await insertQueuedSessionRun(database, { createdAt: 1, id: PUBLIC_API_TEST_IDS.run }); + await insertQueuedSessionRun(database, { createdAt: 2, id: PUBLIC_API_TEST_IDS.runAlt }); await expect(getActiveSessionRunId(database, "01J0000000000000000000000B")).resolves.toBe( - "run-active-probe-latest", + PUBLIC_API_TEST_IDS.runAlt, ); }); }); diff --git a/apps/api/tests/session-run-reconciliation.test.ts b/apps/api/tests/session-run-reconciliation.test.ts index 90aea293..19126974 100644 --- a/apps/api/tests/session-run-reconciliation.test.ts +++ b/apps/api/tests/session-run-reconciliation.test.ts @@ -9,7 +9,9 @@ import { RUNTIME_SOCKET_TIMEOUT_MS, } from "../src/modules/runtime/domain/runtime-config"; import { + PUBLIC_API_TEST_IDS, createPublicHttpContractDatabase, + insertActiveSandboxSessionFixture, insertNonOwnerSession, } from "./helpers/public-api-http-test-fixture"; @@ -17,6 +19,11 @@ describe("session run reconciliation", () => { test("keeps connecting runs alive for the cold ready budget", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); + await insertActiveSandboxSessionFixture(database, { + ownerAccountId: PUBLIC_API_TEST_IDS.nonOwnerAccount, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + }); const driverId = "01J0000000000000000000000E"; const runId = "01J0000000000000000000000N"; @@ -26,6 +33,7 @@ describe("session run reconciliation", () => { INSERT INTO driver_instance ( id, sandbox_id, + sandbox_incarnation, sandbox_session_id, runtime, protocol, @@ -39,12 +47,13 @@ describe("session run reconciliation", () => { created_at, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( driverId, "01J0000000000000000000000D", + 1, "01J0000000000000000000000B", "cloudflare-container", "driver-ws", @@ -140,7 +149,7 @@ describe("session run reconciliation", () => { `, ) .bind( - "run-stale", + PUBLIC_API_TEST_IDS.run, "01J0000000000000000000000B", "01J00000000000000000000009", "01J00000000000000000000002", @@ -156,7 +165,7 @@ describe("session run reconciliation", () => { .run(); await database .prepare("UPDATE session SET last_run_id = ?, status = ? WHERE id = ?") - .bind("run-stale", "RUNNING", "01J0000000000000000000000B") + .bind(PUBLIC_API_TEST_IDS.run, "RUNNING", "01J0000000000000000000000B") .run(); await expect( @@ -165,7 +174,7 @@ describe("session run reconciliation", () => { const run = await database .prepare("SELECT error_code, status FROM session_run WHERE id = ?") - .bind("run-stale") + .bind(PUBLIC_API_TEST_IDS.run) .first<{ error_code: string | null; status: string }>(); expect(run).toMatchObject({ status: "failed", @@ -173,6 +182,108 @@ describe("session run reconciliation", () => { expect(run?.error_code).toBeString(); }); + test("does not fail a stale candidate after its Driver heartbeat recovers", async () => { + const database = await createPublicHttpContractDatabase(); + await insertNonOwnerSession(database); + await insertActiveSandboxSessionFixture(database, { + ownerAccountId: PUBLIC_API_TEST_IDS.nonOwnerAccount, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sessionId: PUBLIC_API_TEST_IDS.nonOwnerSession, + }); + const driverId = "01J0000000000000000000000E"; + const runId = PUBLIC_API_TEST_IDS.run; + const staleAt = Date.now() - RUNTIME_SOCKET_TIMEOUT_MS - 1_000; + + await database + .prepare( + `INSERT INTO driver_instance ( + id, sandbox_id, sandbox_incarnation, sandbox_session_id, runtime, protocol, protocol_version, status, + boot_token_hash, boot_token_expires_at, generation, heartbeat_count, + last_heartbeat_at, expires_at, created_at, updated_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + driverId, + "01J0000000000000000000000D", + 1, + "01J0000000000000000000000B", + "cloudflare-container", + "driver-ws", + 1, + "ready", + new Uint8Array([1]), + Date.now() + 10_000, + 0, + 1, + staleAt, + Date.now() + 20_000, + 1, + staleAt, + ) + .run(); + await database + .prepare( + `INSERT INTO session_run ( + id, session_id, agent_id, created_by_account_id, trigger, status, provider, model, + runtime_id, trace_id, driver_instance_id, created_at, updated_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + runId, + "01J0000000000000000000000B", + "01J00000000000000000000009", + "01J00000000000000000000002", + "user_prompt", + "running", + "openai", + "gpt-5.4", + "openai-runtime", + "trace-heartbeat-race", + driverId, + 1, + 1, + ) + .run(); + await database + .prepare("UPDATE session SET last_run_id = ?, status = ? WHERE id = ?") + .bind(runId, "RUNNING", "01J0000000000000000000000B") + .run(); + + let raced = false; + const racingDatabase = new Proxy(database, { + get(target, property) { + if (property === "batch") { + return async (statements: D1PreparedStatement[]) => { + if (!raced) { + raced = true; + const recoveredAt = Date.now(); + await target + .prepare( + "UPDATE driver_instance SET heartbeat_count = heartbeat_count + 1, last_heartbeat_at = ?, updated_at = ? WHERE id = ?", + ) + .bind(recoveredAt, recoveredAt, driverId) + .run(); + } + return target.batch(statements); + }; + } + + const value = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }) as D1Database; + + await expect( + reconcileStaleActiveSessionRun(racingDatabase, "01J0000000000000000000000B"), + ).resolves.toBe(false); + expect(raced).toBe(true); + await expect( + database.prepare("SELECT status FROM session_run WHERE id = ?").bind(runId).first(), + ).resolves.toEqual({ status: "running" }); + }); + test("reconciles stale active runs in batches", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); @@ -198,7 +309,7 @@ describe("session run reconciliation", () => { `, ) .bind( - "run-stale", + PUBLIC_API_TEST_IDS.run, "01J0000000000000000000000B", "01J00000000000000000000009", "01J00000000000000000000002", @@ -214,7 +325,7 @@ describe("session run reconciliation", () => { .run(); await database .prepare("UPDATE session SET last_run_id = ?, status = ? WHERE id = ?") - .bind("run-stale", "RUNNING", "01J0000000000000000000000B") + .bind(PUBLIC_API_TEST_IDS.run, "RUNNING", "01J0000000000000000000000B") .run(); await expect( @@ -222,13 +333,13 @@ describe("session run reconciliation", () => { limit: 10, }), ).resolves.toEqual({ - reconciledRunIds: ["run-stale"], + reconciledRunIds: [PUBLIC_API_TEST_IDS.run], reconciledSessionIds: ["01J0000000000000000000000B"], }); const run = await database .prepare("SELECT error_code, status FROM session_run WHERE id = ?") - .bind("run-stale") + .bind(PUBLIC_API_TEST_IDS.run) .first<{ error_code: string | null; status: string }>(); expect(run).toMatchObject({ status: "failed", diff --git a/apps/api/tests/session-run-skill-persistence.test.ts b/apps/api/tests/session-run-skill-persistence.test.ts index f84425bb..4ff0317a 100644 --- a/apps/api/tests/session-run-skill-persistence.test.ts +++ b/apps/api/tests/session-run-skill-persistence.test.ts @@ -4,9 +4,10 @@ import { persistSessionRunSkills } from "../src/modules/runtime/application/sess import { createPublicHttpContractDatabase, insertNonOwnerSession, + PUBLIC_API_TEST_IDS, } from "./helpers/public-api-http-test-fixture"; -const RUN_ID = "run-skill-persistence"; +const RUN_ID = PUBLIC_API_TEST_IDS.run; async function insertQueuedSessionRun(database: D1Database): Promise { await database diff --git a/apps/api/tests/session-run-state.test.ts b/apps/api/tests/session-run-state.test.ts index 3cdb5263..18660f29 100644 --- a/apps/api/tests/session-run-state.test.ts +++ b/apps/api/tests/session-run-state.test.ts @@ -1,14 +1,14 @@ import { describe, expect, test } from "bun:test"; -import { - acquireSessionRunDispatch, - updateSessionRunStatusIfActive, -} from "../src/modules/runtime/application/session-runs/session-run-state.repository"; +import { acquireSessionRunDispatch } from "../src/modules/runtime/application/session-runs/session-run-state.repository"; import { createPublicHttpContractDatabase, insertNonOwnerSession, + PUBLIC_API_TEST_IDS, } from "./helpers/public-api-http-test-fixture"; +const RUN_ID = PUBLIC_API_TEST_IDS.run; + async function insertQueuedSessionRun(database: D1Database): Promise { await database .prepare( @@ -31,7 +31,7 @@ async function insertQueuedSessionRun(database: D1Database): Promise { `, ) .bind( - "run-active-transition", + RUN_ID, "01J0000000000000000000000B", "01J00000000000000000000009", "01J00000000000000000000002", @@ -47,46 +47,26 @@ async function insertQueuedSessionRun(database: D1Database): Promise { .run(); await database .prepare("UPDATE session SET last_run_id = ?, status = ?, updated_at = ? WHERE id = ?") - .bind("run-active-transition", "RUNNING", 1, "01J0000000000000000000000B") + .bind(RUN_ID, "RUNNING", 1, "01J0000000000000000000000B") .run(); } describe("session run state", () => { - test("returns active status transitions", async () => { - const database = await createPublicHttpContractDatabase(); - await insertNonOwnerSession(database); - await insertQueuedSessionRun(database); - - const run = await updateSessionRunStatusIfActive(database, { - runId: "run-active-transition", - status: "booting", - }); - - expect(run?.id).toBe("run-active-transition"); - expect(run?.status).toBe("booting"); - - const stored = await database - .prepare("SELECT status FROM session_run WHERE id = ?") - .bind("run-active-transition") - .first<{ status: string }>(); - expect(stored).toEqual({ status: "booting" }); - }); - test("only the first dispatch acquire can continue a run", async () => { const database = await createPublicHttpContractDatabase(); await insertNonOwnerSession(database); await insertQueuedSessionRun(database); - const first = await acquireSessionRunDispatch(database, "run-active-transition"); - const second = await acquireSessionRunDispatch(database, "run-active-transition"); + const first = await acquireSessionRunDispatch(database, RUN_ID); + const second = await acquireSessionRunDispatch(database, RUN_ID); - expect(first?.id).toBe("run-active-transition"); + expect(first?.id).toBe(RUN_ID); expect(first?.status).toBe("booting"); expect(second).toBeNull(); const stored = await database .prepare("SELECT status FROM session_run WHERE id = ?") - .bind("run-active-transition") + .bind(RUN_ID) .first<{ status: string }>(); expect(stored).toEqual({ status: "booting" }); }); diff --git a/apps/api/tests/session-run-terminal-failure.test.ts b/apps/api/tests/session-run-terminal-failure.test.ts index 54b023a0..e251b419 100644 --- a/apps/api/tests/session-run-terminal-failure.test.ts +++ b/apps/api/tests/session-run-terminal-failure.test.ts @@ -1,23 +1,35 @@ import { describe, expect, test } from "bun:test"; -import type { DriverInstanceId, SessionRunId } from "@mosoo/id"; +import type { DriverInstanceId, SessionMessageId, SessionRunId } from "@mosoo/id"; -import { recordCanonicalSessionRunFailure } from "../src/modules/runtime/application/session-runs/session-run-terminal-failure.service"; -import { reconcileTerminalSessionRuns } from "../src/modules/runtime/application/session-runs/terminal-run-reconciliation.service"; +import { + recordCanonicalSessionRunFailure, + recordCanonicalSessionRunTerminal, +} from "../src/modules/runtime/application/session-runs/session-run-terminal-failure.service"; +import { createSessionRunUpdatedEvent } from "../src/modules/runtime/application/session-runs/session-run-view-events.service"; +import { + assertCanonicalTerminalSessionRunProjection, + reconcileTerminalSessionRuns, +} from "../src/modules/runtime/application/session-runs/terminal-run-reconciliation.service"; +import { createSessionRunTerminalSourceId } from "../src/modules/runtime/domain/session-run-terminal-event-id"; +import { commitTerminalRunProjection } from "../src/modules/runtime/infrastructure/driver-instance/completed-run-commit.repository"; import { getRuntimeSessionLink } from "../src/modules/runtime/infrastructure/driver-instance/session-link.repository"; import { recordDriverInstanceFailure } from "../src/modules/runtime/infrastructure/driver-instance/terminal-driver-events"; -import { setSessionRunStatus } from "../src/modules/runtime/infrastructure/session-runs/session-run-store.repository"; +import { getSessionRunSummary } from "../src/modules/runtime/infrastructure/session-runs/session-run-store.repository"; import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; import { createPublicHttpContractDatabase, createPublicHttpTestBindings, + insertActiveSandboxSessionFixture, insertOwnerSession, PUBLIC_API_TEST_IDS, } from "./helpers/public-api-http-test-fixture"; import type { SqliteD1Database } from "./helpers/public-api-http-test-fixture"; +import { insertRuntimeEvent } from "./public-thread-api-fixtures"; const RUN_ID = PUBLIC_API_TEST_IDS.run as SessionRunId; const DRIVER_ID = PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId; +const FINAL_MESSAGE_ID = "01J0000000000000000000000M" as SessionMessageId; const CANONICAL_FAILURE_SOURCE_ID = `session-run-terminal:${RUN_ID}:run.failed`; const DRIVER_ERROR = { code: "driver.command_failed", @@ -38,11 +50,62 @@ interface FailureEventRow { source_event_id: string; } +function advanceRunAfterFirstSummaryRead( + database: SqliteD1Database, + advance: () => Promise, +): { database: D1Database; advanced: () => boolean } { + let advanced = false; + + function wrap(statement: D1PreparedStatement): D1PreparedStatement { + return new Proxy(statement, { + get(target, property, receiver) { + if (property === "bind") { + return (...values: unknown[]) => wrap(target.bind(...values)); + } + if (property === "raw") { + return async () => { + const rows = await target.raw(); + if (!advanced) { + advanced = true; + await advance(); + } + return rows; + }; + } + return Reflect.get(target, property, receiver); + }, + }); + } + + return { + advanced: () => advanced, + database: new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => + query.includes('from "session_run"') + ? wrap(target.prepare(query)) + : target.prepare(query); + } + + const value = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }) as D1Database, + }; +} + async function insertLinkedRunFixture( database: SqliteD1Database, status: "booting" | "completed" | "cancelled" | "failed" = "booting", ): Promise { await insertOwnerSession(database); + await insertActiveSandboxSessionFixture(database, { + ownerAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + timestampMs: 1, + }); await database .prepare( ` @@ -57,11 +120,12 @@ async function insertLinkedRunFixture( protocol_version, runtime, sandbox_id, + sandbox_incarnation, sandbox_session_id, status, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( @@ -75,6 +139,7 @@ async function insertLinkedRunFixture( 1, "openai-runtime", PUBLIC_API_TEST_IDS.sandbox, + 1, PUBLIC_API_TEST_IDS.ownerSession, "ready", 1, @@ -102,10 +167,11 @@ async function insertLinkedRunFixture( error_code, error_message, error_details_json, + error_retryable, created_at, updated_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) `, ) .bind( @@ -127,16 +193,48 @@ async function insertLinkedRunFixture( status === "failed" ? DRIVER_ERROR.code : null, status === "failed" ? DRIVER_ERROR.message : null, status === "failed" ? "{}" : null, + status === "failed" ? 0 : null, 1, 1, ) .run(); await database - .prepare("UPDATE session SET last_run_id = ?, status = ? WHERE id = ?") + .prepare("UPDATE session SET last_run_id = ?, status = ?, updated_at = 1 WHERE id = ?") .bind(RUN_ID, status === "booting" ? "RUNNING" : "IDLE", PUBLIC_API_TEST_IDS.ownerSession) .run(); } +async function insertHealthyFailedRunFixture(database: SqliteD1Database): Promise { + const runId = PUBLIC_API_TEST_IDS.runAlt as SessionRunId; + await database + .prepare( + `INSERT INTO session_run ( + id, session_id, agent_id, created_by_account_id, deployment_version_id, + deployment_version_number, driver_instance_id, trigger, status, provider, + model, runtime_id, trace_id, started_at, completed_at, error_code, + error_message, error_details_json, error_retryable, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, 1, ?, 'user_prompt', 'failed', 'openai', + 'gpt-5.4', 'openai-runtime', 'trace-healthy-terminal-repair', 2, 2, + ?, ?, '{}', 0, 2, 2)`, + ) + .bind( + runId, + PUBLIC_API_TEST_IDS.ownerSession, + PUBLIC_API_TEST_IDS.agent, + PUBLIC_API_TEST_IDS.ownerAccount, + PUBLIC_API_TEST_IDS.deployment, + DRIVER_ID, + PROVISION_ERROR.code, + PROVISION_ERROR.message, + ) + .run(); + await database + .prepare("UPDATE session SET last_run_id = ?, status = 'RUNNING', updated_at = 2 WHERE id = ?") + .bind(runId, PUBLIC_API_TEST_IDS.ownerSession) + .run(); + return runId; +} + async function readFailureEvents(database: SqliteD1Database): Promise { return database .prepare( @@ -152,7 +250,255 @@ async function readFailureEvents(database: SqliteD1Database): Promise result.results ?? []); } +async function insertSealedAssistantAuthority(database: SqliteD1Database): Promise { + await insertRuntimeEvent(database, { + kind: "message.added", + occurredAt: 1, + payload: { content: "done", messageId: FINAL_MESSAGE_ID, role: "agent" }, + runId: RUN_ID, + seq: 1, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + }); + await insertRuntimeEvent(database, { + kind: "message.completed", + occurredAt: 2, + payload: { messageId: FINAL_MESSAGE_ID, role: "agent" }, + runId: RUN_ID, + seq: 2, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + }); + await database + .prepare( + ` + INSERT INTO session_message ( + content_text, + created_at, + created_by_account_id, + id, + plan_json, + projection_format, + role, + segments_json, + seq, + session_id, + session_run_id + ) + VALUES ('', 1, ?, ?, NULL, 'event_stream_v3', 'assistant', NULL, 1, ?, ?) + `, + ) + .bind( + PUBLIC_API_TEST_IDS.ownerAccount, + FINAL_MESSAGE_ID, + PUBLIC_API_TEST_IDS.ownerSession, + RUN_ID, + ) + .run(); + await database + .prepare("UPDATE session SET message_seq_cursor = 1, runtime_event_seq_cursor = 2 WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .run(); +} + +async function insertParentlessToolOutput( + database: SqliteD1Database, + eventSeq: number, +): Promise { + await insertRuntimeEvent(database, { + kind: "tool.call.updated", + occurredAt: eventSeq, + payload: { + rawOutput: "parentless output", + status: "completed", + title: "Shell", + toolCallId: "parentless-tool", + }, + runId: RUN_ID, + seq: eventSeq, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + }); + await database + .prepare( + `UPDATE session + SET runtime_event_seq_cursor = MAX(runtime_event_seq_cursor, ?) + WHERE id = ?`, + ) + .bind(eventSeq, PUBLIC_API_TEST_IDS.ownerSession) + .run(); +} + +async function insertParentlessToolCarrier( + database: SqliteD1Database, + input: { readonly eventSeq: number; readonly messageSeq: number }, +): Promise { + await insertParentlessToolOutput(database, input.eventSeq); + await database + .prepare( + `INSERT INTO session_message ( + content_text, created_at, created_by_account_id, id, plan_json, + projection_format, role, segments_json, seq, session_id, session_run_id + ) VALUES ('', ?, ?, ?, NULL, 'event_stream_v3', 'assistant', NULL, ?, ?, ?)`, + ) + .bind( + input.eventSeq, + PUBLIC_API_TEST_IDS.ownerAccount, + RUN_ID, + input.messageSeq, + PUBLIC_API_TEST_IDS.ownerSession, + RUN_ID, + ) + .run(); + await database + .prepare( + `UPDATE session + SET message_seq_cursor = MAX(message_seq_cursor, ?) + WHERE id = ?`, + ) + .bind(input.messageSeq, PUBLIC_API_TEST_IDS.ownerSession) + .run(); +} + +async function commitCompletedWithoutFinalStream(database: SqliteD1Database): Promise { + const current = await getSessionRunSummary(database, RUN_ID); + if (current === null) { + throw new Error("Missing Session Run fixture."); + } + const completedAt = new Date(2).toISOString(); + const run = { + ...current, + completedAt, + startedAt: current.startedAt ?? completedAt, + status: "completed" as const, + updatedAt: completedAt, + }; + const sourceEventId = createSessionRunTerminalSourceId(RUN_ID, "run.completed"); + const event = createSessionRunUpdatedEvent( + run, + PUBLIC_API_TEST_IDS.ownerSession, + "IDLE", + sourceEventId, + ); + await commitTerminalRunProjection(database, { + assistantMessage: null, + error: null, + runId: RUN_ID, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + source: "driver", + targetStatus: "completed", + terminalEvent: { event, occurredAt: 2, sourceEventId }, + timestampMs: 2, + }); +} + +async function insertLegacyCompletedAuthorityWithCarrier( + database: SqliteD1Database, +): Promise { + await insertParentlessToolCarrier(database, { eventSeq: 1, messageSeq: 2 }); + await database + .prepare( + `INSERT INTO session_message ( + content_text, created_at, created_by_account_id, id, plan_json, + projection_format, role, segments_json, seq, session_id, session_run_id + ) VALUES ('legacy final', 1, ?, ?, NULL, 'materialized', 'assistant', NULL, 1, ?, ?)`, + ) + .bind( + PUBLIC_API_TEST_IDS.ownerAccount, + FINAL_MESSAGE_ID, + PUBLIC_API_TEST_IDS.ownerSession, + RUN_ID, + ) + .run(); + await insertRuntimeEvent(database, { + kind: "run.completed", + occurredAt: 2, + payload: { finalMessageId: FINAL_MESSAGE_ID, stopReason: "end_turn" }, + runId: RUN_ID, + seq: 2, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + }); + await database + .prepare( + `UPDATE session_event SET semantic_hash = NULL, terminal_event_json = NULL + WHERE run_id = ? AND event_type = 'run.completed'`, + ) + .bind(RUN_ID) + .run(); + await database + .prepare( + `UPDATE session + SET message_seq_cursor = 2, runtime_event_seq_cursor = 2 + WHERE id = ?`, + ) + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .run(); +} + describe("canonical session run terminal failure", () => { + test("returns stale when the Run advances after its authoritative read", async () => { + const database = await createPublicHttpContractDatabase(); + await insertLinkedRunFixture(database); + await database + .prepare( + `UPDATE session_run + SET started_at = NULL, + status = 'queued', + status_seq = status_seq + 1, + updated_at = 1 + WHERE id = ?`, + ) + .bind(RUN_ID) + .run(); + const race = advanceRunAfterFirstSummaryRead(database, async () => { + await database + .prepare( + `UPDATE session_run + SET started_at = 2, + status = 'booting', + status_seq = status_seq + 1, + updated_at = 2 + WHERE id = ?`, + ) + .bind(RUN_ID) + .run(); + }); + const bindings = { + ...(createPublicHttpTestBindings(database) as ApiBindings), + DB: race.database, + }; + + await expect( + recordCanonicalSessionRunFailure(bindings, { + error: PROVISION_ERROR, + runId: RUN_ID, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + source: "api", + }), + ).resolves.toEqual({ kind: "not_failed", status: "booting" }); + + expect(race.advanced()).toBe(true); + await expect( + database + .prepare( + `SELECT completed_at, error_code, started_at, status + FROM session_run + WHERE id = ?`, + ) + .bind(RUN_ID) + .first(), + ).resolves.toEqual({ + completed_at: null, + error_code: null, + started_at: 2, + status: "booting", + }); + await expect( + database + .prepare("SELECT status FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first(), + ).resolves.toEqual({ status: "RUNNING" }); + expect(await readFailureEvents(database)).toEqual([]); + }); + test("treats a concurrent failed transition as canonical success", async () => { const database = await createPublicHttpContractDatabase(); await insertLinkedRunFixture(database); @@ -173,7 +519,7 @@ describe("canonical session run terminal failure", () => { }), ]); - expect(outcomes.map((outcome) => outcome.kind)).toEqual(["failed", "failed"]); + expect(outcomes.map((outcome) => outcome.kind).toSorted()).toEqual(["failed", "not_failed"]); expect(await readFailureEvents(database)).toHaveLength(1); }); @@ -185,6 +531,7 @@ describe("canonical session run terminal failure", () => { await recordDriverInstanceFailure(bindings, { driverInstanceId: DRIVER_ID, error: DRIVER_ERROR, + sessionRunId: RUN_ID, }); await recordCanonicalSessionRunFailure(bindings, { error: PROVISION_ERROR, @@ -226,6 +573,7 @@ describe("canonical session run terminal failure", () => { await recordDriverInstanceFailure(bindings, { driverInstanceId: DRIVER_ID, error: DRIVER_ERROR, + sessionRunId: RUN_ID, }); const run = await database @@ -253,16 +601,16 @@ describe("canonical session run terminal failure", () => { const bindings = createPublicHttpTestBindings(database) as ApiBindings; await recordCanonicalSessionRunFailure(bindings, { - error: PROVISION_ERROR, + error: DRIVER_ERROR, runId: RUN_ID, sessionId: PUBLIC_API_TEST_IDS.ownerSession, - source: "api", + source: "system", }); await recordCanonicalSessionRunFailure(bindings, { - error: PROVISION_ERROR, + error: DRIVER_ERROR, runId: RUN_ID, sessionId: PUBLIC_API_TEST_IDS.ownerSession, - source: "api", + source: "system", }); expect(await readFailureEvents(database)).toEqual([ @@ -274,24 +622,24 @@ describe("canonical session run terminal failure", () => { ]); }); - test("repairs an API failure event from the Driver cached run link", async () => { + test("keeps the API terminal winner when the Driver uses a cached run link", async () => { const database = await createPublicHttpContractDatabase(); await insertLinkedRunFixture(database); const bindings = createPublicHttpTestBindings(database) as ApiBindings; const link = await getRuntimeSessionLink(database, DRIVER_ID); - await setSessionRunStatus(database, { + await recordCanonicalSessionRunFailure(bindings, { error: PROVISION_ERROR, runId: RUN_ID, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, source: "api", - status: "failed", }); - expect(await readFailureEvents(database)).toEqual([]); await recordDriverInstanceFailure(bindings, { driverInstanceId: DRIVER_ID, error: DRIVER_ERROR, link, + sessionRunId: RUN_ID, }); expect(await readFailureEvents(database)).toEqual([ @@ -308,12 +656,29 @@ describe("canonical session run terminal failure", () => { await insertLinkedRunFixture(database); const bindings = createPublicHttpTestBindings(database) as ApiBindings; - await setSessionRunStatus(database, { - error: PROVISION_ERROR, - runId: RUN_ID, - source: "api", - status: "failed", - }); + await database + .prepare( + ` + UPDATE session_run + SET completed_at = 2, + error_code = ?, + error_details_json = '{}', + error_message = ?, + error_retryable = 0, + status = 'failed', + status_operation_id = NULL, + status_seq = status_seq + 1, + status_source = 'api', + updated_at = 2 + WHERE id = ? + `, + ) + .bind(PROVISION_ERROR.code, PROVISION_ERROR.message, RUN_ID) + .run(); + await database + .prepare("UPDATE session SET status = 'IDLE', updated_at = 2 WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .run(); await database .prepare("UPDATE driver_instance SET status = ? WHERE id = ?") @@ -322,8 +687,12 @@ describe("canonical session run terminal failure", () => { expect(await readFailureEvents(database)).toEqual([]); - const firstRepair = await reconcileTerminalSessionRuns(bindings, { limit: 10 }); - const secondRepair = await reconcileTerminalSessionRuns(bindings, { limit: 10 }); + const firstRepair = await reconcileTerminalSessionRuns(bindings, { + limit: 10, + }); + const secondRepair = await reconcileTerminalSessionRuns(bindings, { + limit: 10, + }); expect(firstRepair.reconciledRunIds).toEqual([RUN_ID]); expect(secondRepair.reconciledRunIds).toEqual([]); @@ -336,9 +705,43 @@ describe("canonical session run terminal failure", () => { ]); }); - test("repairs an inherited terminal Run whose Session and completion event are both stale", async () => { + test("repairs a completed Run that has no final assistant", async () => { + const database = await createPublicHttpContractDatabase(); + await insertLinkedRunFixture(database, "completed"); + await insertParentlessToolCarrier(database, { eventSeq: 1, messageSeq: 1 }); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + + await database + .prepare("UPDATE driver_instance SET status = ? WHERE id = ?") + .bind("stopped", DRIVER_ID) + .run(); + + const repaired = await reconcileTerminalSessionRuns(bindings, { + limit: 10, + }); + + expect(repaired.reconciledRunIds).toEqual([RUN_ID]); + expect( + await database + .prepare( + "SELECT stream_id FROM session_event WHERE run_id = ? AND event_type = 'run.completed'", + ) + .bind(RUN_ID) + .first(), + ).toEqual({ stream_id: null }); + expect( + await database + .prepare("SELECT id FROM session_message WHERE session_run_id = ? AND role = 'assistant'") + .bind(RUN_ID) + .all(), + ).toMatchObject({ results: [{ id: RUN_ID }] }); + }); + + test("repairs a completed Run from its unique sealed assistant authority", async () => { const database = await createPublicHttpContractDatabase(); await insertLinkedRunFixture(database, "completed"); + await insertSealedAssistantAuthority(database); + await insertParentlessToolCarrier(database, { eventSeq: 3, messageSeq: 2 }); const bindings = createPublicHttpTestBindings(database) as ApiBindings; await database @@ -350,7 +753,16 @@ describe("canonical session run terminal failure", () => { .bind("stopped", DRIVER_ID) .run(); - const repaired = await reconcileTerminalSessionRuns(bindings, { limit: 10 }); + const repaired = await reconcileTerminalSessionRuns(bindings, { + limit: 10, + }); + await expect( + assertCanonicalTerminalSessionRunProjection(bindings, { + runId: RUN_ID, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + status: "completed", + }), + ).resolves.toBeUndefined(); const session = await database .prepare("SELECT status FROM session WHERE id = ?") .bind(PUBLIC_API_TEST_IDS.ownerSession) @@ -376,9 +788,343 @@ describe("canonical session run terminal failure", () => { ]); }); - test("repairs an old terminal receipt without changing a newer Run", async () => { + test("recognizes a legacy materialized assistant beside a parentless tool carrier", async () => { const database = await createPublicHttpContractDatabase(); await insertLinkedRunFixture(database, "completed"); + await insertLegacyCompletedAuthorityWithCarrier(database); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + + await expect( + assertCanonicalTerminalSessionRunProjection(bindings, { + runId: RUN_ID, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + status: "completed", + }), + ).resolves.toBeUndefined(); + }); + + test("adopts a canonical completed Run that intentionally has no final stream", async () => { + const database = await createPublicHttpContractDatabase(); + await insertLinkedRunFixture(database); + await insertParentlessToolOutput(database, 1); + await commitCompletedWithoutFinalStream(database); + await database + .prepare("UPDATE session SET status = 'RUNNING', updated_at = 1 WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .run(); + await database + .prepare("UPDATE driver_instance SET status = 'stopped' WHERE id = ?") + .bind(DRIVER_ID) + .run(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + + const repaired = await reconcileTerminalSessionRuns(bindings, { + limit: 10, + }); + + expect(repaired.reconciledSessionIds).toEqual([PUBLIC_API_TEST_IDS.ownerSession]); + expect( + await database + .prepare("SELECT status FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first(), + ).toEqual({ status: "IDLE" }); + expect( + await database + .prepare( + "SELECT COUNT(*) AS count FROM session_message WHERE session_run_id = ? AND role = 'assistant'", + ) + .bind(RUN_ID) + .first(), + ).toEqual({ count: 1 }); + expect( + await database + .prepare("SELECT id FROM session_message WHERE session_run_id = ? AND role = 'assistant'") + .bind(RUN_ID) + .first(), + ).toEqual({ id: RUN_ID }); + }); + + test("does not report a newer stale Session projection as reconciled", async () => { + const database = await createPublicHttpContractDatabase(); + await insertLinkedRunFixture(database); + await commitCompletedWithoutFinalStream(database); + const run = await database + .prepare("SELECT updated_at FROM session_run WHERE id = ?") + .bind(RUN_ID) + .first<{ updated_at: number }>(); + if (run === null) { + throw new Error("Missing terminal Run fixture."); + } + await database + .prepare("UPDATE session SET status = 'RUNNING', updated_at = ? WHERE id = ?") + .bind(run.updated_at + 1, PUBLIC_API_TEST_IDS.ownerSession) + .run(); + await database + .prepare("UPDATE driver_instance SET status = 'stopped' WHERE id = ?") + .bind(DRIVER_ID) + .run(); + + const result = await reconcileTerminalSessionRuns( + createPublicHttpTestBindings(database) as ApiBindings, + { limit: 10 }, + ); + + expect(result).toMatchObject({ + failures: [ + { + message: expect.stringContaining("did not converge its current Session"), + runId: RUN_ID, + }, + ], + reconciledRunIds: [], + reconciledSessionIds: [], + }); + await expect( + database + .prepare("SELECT status FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first(), + ).resolves.toEqual({ status: "RUNNING" }); + }); + + test("retries a repaired authority without changing terminal Run semantics", async () => { + const database = await createPublicHttpContractDatabase(); + await insertLinkedRunFixture(database); + await commitCompletedWithoutFinalStream(database); + await database + .prepare( + `INSERT INTO session_message ( + content_text, created_at, created_by_account_id, id, plan_json, + projection_format, role, segments_json, seq, session_id, session_run_id + ) VALUES ('not a carrier', 2, ?, ?, NULL, 'materialized', 'assistant', NULL, 1, ?, ?)`, + ) + .bind(PUBLIC_API_TEST_IDS.ownerAccount, RUN_ID, PUBLIC_API_TEST_IDS.ownerSession, RUN_ID) + .run(); + await database + .prepare("UPDATE session SET status = 'RUNNING', updated_at = 1 WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .run(); + await database + .prepare("UPDATE driver_instance SET status = 'stopped' WHERE id = ?") + .bind(DRIVER_ID) + .run(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + const semanticRun = await database + .prepare( + `SELECT completed_at, error_code, error_details_json, error_message, + error_retryable, status, status_seq, status_source, updated_at + FROM session_run + WHERE id = ?`, + ) + .bind(RUN_ID) + .first(); + + await expect(reconcileTerminalSessionRuns(bindings, { limit: 10 })).resolves.toMatchObject({ + failures: [ + { + message: expect.stringContaining("Canonical assistant messages"), + runId: RUN_ID, + }, + ], + }); + const claimedRun = await database + .prepare( + `SELECT completed_at, error_code, error_details_json, error_message, + error_retryable, status, status_seq, status_source, + terminal_reconciliation_attempted_at, updated_at + FROM session_run + WHERE id = ?`, + ) + .bind(RUN_ID) + .first<{ terminal_reconciliation_attempted_at: number }>(); + expect(claimedRun).toEqual({ + ...semanticRun, + terminal_reconciliation_attempted_at: expect.any(Number), + }); + const attemptedAt = claimedRun?.terminal_reconciliation_attempted_at; + if (attemptedAt === undefined) { + throw new Error("Missing terminal reconciliation attempt marker."); + } + await database.prepare("DELETE FROM session_message WHERE id = ?").bind(RUN_ID).run(); + await database + .prepare("UPDATE session_run SET terminal_reconciliation_attempted_at = ? WHERE id = ?") + .bind(attemptedAt - 10 * 60_000 - 1, RUN_ID) + .run(); + + await expect(reconcileTerminalSessionRuns(bindings, { limit: 10 })).resolves.toMatchObject({ + failures: [], + reconciledSessionIds: [PUBLIC_API_TEST_IDS.ownerSession], + }); + await expect( + database + .prepare("SELECT status FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first(), + ).resolves.toEqual({ status: "IDLE" }); + }); + + test("rotates a poisoned candidate behind the next repair batch", async () => { + const database = await createPublicHttpContractDatabase(); + await insertLinkedRunFixture(database, "completed"); + const healthyRunId = await insertHealthyFailedRunFixture(database); + await database + .prepare( + `INSERT INTO session_message ( + content_text, created_at, created_by_account_id, id, plan_json, + projection_format, role, segments_json, seq, session_id, session_run_id + ) VALUES ('not a carrier', 1, ?, ?, NULL, 'materialized', 'assistant', NULL, 1, ?, ?)`, + ) + .bind(PUBLIC_API_TEST_IDS.ownerAccount, RUN_ID, PUBLIC_API_TEST_IDS.ownerSession, RUN_ID) + .run(); + await database + .prepare("UPDATE driver_instance SET status = 'stopped' WHERE id = ?") + .bind(DRIVER_ID) + .run(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + + await expect(reconcileTerminalSessionRuns(bindings, { limit: 1 })).resolves.toMatchObject({ + failures: [ + { + message: expect.stringContaining("invalid parentless tool carrier"), + runId: RUN_ID, + }, + ], + reconciledRunIds: [], + }); + await expect( + database + .prepare( + "SELECT source_event_id FROM session_event WHERE run_id = ? AND event_type = 'run.failed'", + ) + .bind(healthyRunId) + .first(), + ).resolves.toBeNull(); + await expect(reconcileTerminalSessionRuns(bindings, { limit: 1 })).resolves.toMatchObject({ + failures: [], + reconciledRunIds: [healthyRunId], + }); + await expect( + database + .prepare( + "SELECT source_event_id FROM session_event WHERE run_id = ? AND event_type = 'run.failed'", + ) + .bind(healthyRunId) + .first(), + ).resolves.toEqual({ + source_event_id: createSessionRunTerminalSourceId(healthyRunId, "run.failed"), + }); + await expect( + database + .prepare("SELECT last_run_id, status FROM session WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .first(), + ).resolves.toEqual({ last_run_id: healthyRunId, status: "IDLE" }); + }); + + test("rotates a malformed terminal Run before parsing the next batch", async () => { + const database = await createPublicHttpContractDatabase(); + await insertLinkedRunFixture(database, "failed"); + const healthyRunId = await insertHealthyFailedRunFixture(database); + await database + .prepare("UPDATE session_run SET error_details_json = '{' WHERE id = ?") + .bind(RUN_ID) + .run(); + await database + .prepare("UPDATE driver_instance SET status = 'stopped' WHERE id = ?") + .bind(DRIVER_ID) + .run(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + + await expect(reconcileTerminalSessionRuns(bindings, { limit: 1 })).resolves.toMatchObject({ + failures: [{ runId: RUN_ID }], + reconciledRunIds: [], + }); + await expect(reconcileTerminalSessionRuns(bindings, { limit: 1 })).resolves.toMatchObject({ + failures: [], + reconciledRunIds: [healthyRunId], + }); + }); + + test("rejects assistant rows that contradict a canonical no-final completion", async () => { + const database = await createPublicHttpContractDatabase(); + await insertLinkedRunFixture(database); + await commitCompletedWithoutFinalStream(database); + await database + .prepare( + `INSERT INTO session_message ( + content_text, created_at, created_by_account_id, id, plan_json, + projection_format, role, segments_json, seq, session_id, session_run_id + ) VALUES ('', 2, ?, ?, NULL, 'event_stream_v3', 'assistant', NULL, 1, ?, ?)`, + ) + .bind( + PUBLIC_API_TEST_IDS.ownerAccount, + FINAL_MESSAGE_ID, + PUBLIC_API_TEST_IDS.ownerSession, + RUN_ID, + ) + .run(); + await database + .prepare("UPDATE session SET status = 'RUNNING', updated_at = 1 WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .run(); + await database + .prepare("UPDATE driver_instance SET status = 'stopped' WHERE id = ?") + .bind(DRIVER_ID) + .run(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + + await expect(reconcileTerminalSessionRuns(bindings, { limit: 10 })).resolves.toMatchObject({ + failures: [ + { + message: expect.stringContaining("Canonical assistant messages"), + runId: RUN_ID, + }, + ], + }); + }); + + test("rejects a terminal receipt whose semantic stream authority changed", async () => { + const database = await createPublicHttpContractDatabase(); + await insertLinkedRunFixture(database); + await commitCompletedWithoutFinalStream(database); + await database + .prepare("UPDATE session_event SET stream_id = 'corrupt-progress' WHERE run_id = ?") + .bind(RUN_ID) + .run(); + const bindings = createPublicHttpTestBindings(database) as ApiBindings; + await expect( + recordCanonicalSessionRunTerminal(bindings, { + assistantMessage: null, + error: null, + runId: RUN_ID, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + source: "driver", + status: "completed", + }), + ).rejects.toThrow("terminal semantic authority is invalid"); + await database + .prepare("UPDATE session SET status = 'RUNNING', updated_at = 1 WHERE id = ?") + .bind(PUBLIC_API_TEST_IDS.ownerSession) + .run(); + await database + .prepare("UPDATE driver_instance SET status = 'stopped' WHERE id = ?") + .bind(DRIVER_ID) + .run(); + await expect(reconcileTerminalSessionRuns(bindings, { limit: 10 })).resolves.toMatchObject({ + failures: [ + { + message: expect.stringContaining("terminal semantic authority is invalid"), + runId: RUN_ID, + }, + ], + }); + }); + + test("refuses to invent an old terminal receipt once a newer Run owns the Session", async () => { + const database = await createPublicHttpContractDatabase(); + await insertLinkedRunFixture(database, "completed"); + await insertSealedAssistantAuthority(database); const bindings = createPublicHttpTestBindings(database) as ApiBindings; await database @@ -425,7 +1171,14 @@ describe("canonical session run terminal failure", () => { .bind("stopped", DRIVER_ID) .run(); - await reconcileTerminalSessionRuns(bindings, { limit: 10 }); + await expect(reconcileTerminalSessionRuns(bindings, { limit: 10 })).resolves.toMatchObject({ + failures: [ + { + message: expect.stringContaining("not safely repairable"), + runId: RUN_ID, + }, + ], + }); const session = await database .prepare("SELECT last_run_id, status FROM session WHERE id = ?") @@ -443,9 +1196,7 @@ describe("canonical session run terminal failure", () => { .bind(RUN_ID) .first<{ source_event_id: string }>(); - expect(completedEvent).toEqual({ - source_event_id: `session-run-terminal:${RUN_ID}:run.completed`, - }); + expect(completedEvent).toBeNull(); expect(await readFailureEvents(database)).toEqual([]); }); diff --git a/apps/api/tests/session-runtime-event-store.test.ts b/apps/api/tests/session-runtime-event-store.test.ts index 73ddd860..59820be1 100644 --- a/apps/api/tests/session-runtime-event-store.test.ts +++ b/apps/api/tests/session-runtime-event-store.test.ts @@ -1,17 +1,18 @@ import { describe, expect, test } from "bun:test"; -import { createRuntimeEvent } from "@mosoo/runtime-events"; +import { createRuntimeEvent, createRuntimeEventSemanticHash } from "@mosoo/runtime-events"; import type { RuntimeEventEnvelope, RuntimeEventKind, RuntimeEventOrigin, } from "@mosoo/runtime-events"; -import { RuntimeEventPersistenceCompactor } from "../src/modules/runtime/infrastructure/driver-instance/runtime-event-persistence-compactor"; +import { loadSessionAgentTaskSnapshot } from "../src/modules/sessions/infrastructure/session-agent-task-snapshot.repository"; import { persistOneRuntimeEventPerSession, persistSessionRuntimeEvents, } from "../src/modules/sessions/infrastructure/session-runtime-event-store.repository"; +import { applyDrizzleMigration } from "./helpers/drizzle-migrations"; import { SqliteD1Database } from "./helpers/sqlite-d1"; function runtimeEvent(input: { @@ -39,6 +40,63 @@ function runtimeEvent(input: { }); } +function agentTasksEvent(input: { + driverInstanceId: string; + id: string; + occurredAtMs: number; + runId: string; + taskId?: string; +}): RuntimeEventEnvelope { + return runtimeEvent({ + driverInstanceId: input.driverInstanceId, + id: input.id, + kind: "agent.tasks.replaced", + occurredAtMs: input.occurredAtMs, + payload: { tasks: input.taskId === undefined ? [] : [{ taskId: input.taskId }] }, + runId: input.runId, + }); +} + +function activateRun( + database: SqliteD1Database, + input: { driverInstanceId: string; runId: string }, +): void { + database + .prepare("UPDATE session SET last_run_id = ?, status = 'RUNNING' WHERE id = 'session-1'") + .bind(input.runId) + .run(); + database + .prepare("UPDATE session_run SET driver_instance_id = ?, status = 'running' WHERE id = ?") + .bind(input.driverInstanceId, input.runId) + .run(); +} + +async function insertConcurrentMessageReceipt( + database: SqliteD1Database, + input: { event: RuntimeEventEnvelope; sourceEventId: string }, +): Promise { + const semanticHash = await createRuntimeEventSemanticHash(input.event); + database.execute("UPDATE session SET runtime_event_seq_cursor = 1 WHERE id = 'session-1'"); + await database + .prepare( + `INSERT INTO session_event ( + agent_id, content_text, created_at, ended_at, event_type, family, id, + occurred_at, process_status, process_type, run_id, semantic_hash, seq, + session_id, source_event_id, source, stream_id, visibility + ) VALUES (?, ?, 1000, 1000, 'message.delta', 'message', ?, 1000, + 'available', 'agent.message.delta', 'run-1', ?, 1, 'session-1', + ?, 'driver', 'message-1', 'all_consumers')`, + ) + .bind( + "01J00000000000000000000009", + (input.event.payload as { contentDelta: string }).contentDelta, + `winner:${input.sourceEventId}`, + semanticHash, + input.sourceEventId, + ) + .run(); +} + function createRuntimeEventStoreDatabase( input: { maxBoundParams?: number } = {}, ): SqliteD1Database { @@ -52,37 +110,120 @@ function createRuntimeEventStoreDatabase( id text PRIMARY KEY NOT NULL, agent_id text NOT NULL, archived_at integer, + last_run_id text, status text NOT NULL, runtime_event_seq_cursor integer DEFAULT 0 NOT NULL ); CREATE TABLE session_run ( + completed_at integer, id text PRIMARY KEY NOT NULL, - session_id text NOT NULL + driver_instance_id text, + error_code text, + error_details_json text, + error_message text, + error_retryable integer, + session_id text NOT NULL, + status text DEFAULT 'running' NOT NULL, + status_changed_at integer NOT NULL DEFAULT 0, + status_event text NOT NULL DEFAULT 'run.start', + status_operation_id text, + status_seq integer NOT NULL DEFAULT 0, + status_source text NOT NULL DEFAULT 'driver', + started_at integer, + updated_at integer NOT NULL DEFAULT 0 + ); + + CREATE TABLE session_agent_task_snapshot ( + driver_instance_id text NOT NULL, + run_id text NOT NULL, + seq integer NOT NULL, + session_id text PRIMARY KEY NOT NULL, + tasks_json text NOT NULL ); CREATE TABLE session_event ( agent_id text NOT NULL, + artifact_attempt_id text, + artifact_manifest_json text, + artifact_manifest_sha256 text, content_text text NOT NULL, created_at integer NOT NULL, ended_at integer NOT NULL, event_type text NOT NULL, family text NOT NULL, id text PRIMARY KEY NOT NULL, + mcp_command_id text, occurred_at integer NOT NULL, process_status text NOT NULL, process_type text NOT NULL, run_id text, + semantic_hash text CHECK ( + semantic_hash IS NULL OR ( + length(semantic_hash) = 64 + AND semantic_hash = lower(semantic_hash) + AND semantic_hash NOT GLOB '*[^0-9a-f]*' + ) + ), seq integer NOT NULL, session_id text NOT NULL, source_event_id text NOT NULL, source text NOT NULL, + stream_id text, + terminal_event_json text, tool_call_id text, + tool_input_delta_json text, tool_input_json text, tool_name text, + tool_output_delta_text text, + tool_output_text text, + tool_parent_message_id text, + tool_result_message_id text, + tool_status text, tokens integer, trace_id text, - visibility text NOT NULL + visibility text NOT NULL, + CHECK (tool_input_delta_json IS NULL OR tool_input_json IS NULL), + CHECK (tool_output_delta_text IS NULL OR tool_output_text IS NULL), + CHECK ( + (terminal_event_json IS NULL AND NOT (semantic_hash IS NOT NULL AND event_type IN ('run.cancelled', 'run.completed', 'run.failed'))) + OR (terminal_event_json IS NOT NULL AND json_valid(terminal_event_json) = 1 AND semantic_hash IS NOT NULL AND event_type IN ('run.cancelled', 'run.completed', 'run.failed')) + ), + CHECK ( + (artifact_attempt_id IS NULL AND artifact_manifest_json IS NULL AND artifact_manifest_sha256 IS NULL) + OR ( + artifact_attempt_id IS NOT NULL + AND artifact_manifest_json IS NOT NULL + AND json_valid(artifact_manifest_json) = 1 + AND json_extract(artifact_manifest_json, '$.version') IS 1 + AND json_type(artifact_manifest_json, '$.captureStatus') IS 'text' + AND json_extract(artifact_manifest_json, '$.captureStatus') IN ('complete', 'omitted_file_limit', 'omitted_runtime_unavailable', 'omitted_size_limit', 'omitted_source_changed', 'omitted_source_missing') + AND json_type(artifact_manifest_json, '$.mode') IS 'text' + AND json_extract(artifact_manifest_json, '$.mode') IN ('delta', 'snapshot') + AND (json_extract(artifact_manifest_json, '$.captureStatus') = 'complete' OR json_array_length(artifact_manifest_json, '$.files') = 0) + AND json_extract(artifact_manifest_json, '$.sourceEventId') IS source_event_id + AND json_extract(artifact_manifest_json, '$.semanticHash') IS semantic_hash + AND json_type(artifact_manifest_json, '$.files') IS 'array' + AND artifact_manifest_sha256 IS NOT NULL + AND length(artifact_manifest_sha256) = 64 + AND artifact_manifest_sha256 = lower(artifact_manifest_sha256) + AND artifact_manifest_sha256 NOT GLOB '*[^0-9a-f]*' + AND semantic_hash IS NOT NULL + AND event_type IN ('file.change.updated', 'file.changed', 'run.completed') + ) + ), + CHECK (tool_status IS NULL OR tool_status IN ('running', 'completed', 'failed', 'cancelled')), + CHECK ( + mcp_command_id IS NULL OR ( + mcp_command_id = upper(mcp_command_id) + AND length(mcp_command_id) = 26 + AND substr(mcp_command_id, 1, 1) GLOB '[0-7]' + AND mcp_command_id NOT GLOB '*[^0-9A-HJKMNP-TV-Z]*' + AND event_type = 'tool.call.updated' + AND tool_status IS NOT NULL + AND tool_status IN ('completed', 'failed', 'cancelled') + ) + ) ); CREATE UNIQUE INDEX session_event_session_seq_idx @@ -91,29 +232,25 @@ function createRuntimeEventStoreDatabase( CREATE UNIQUE INDEX session_event_session_source_idx ON session_event (session_id, source_event_id); - CREATE TRIGGER session_event_tool_identity_consistency - BEFORE INSERT ON session_event - WHEN NEW.tool_call_id IS NOT NULL AND EXISTS ( - SELECT 1 - FROM session_event AS existing - WHERE existing.session_id = NEW.session_id - AND existing.tool_call_id = NEW.tool_call_id - AND ( - ( - NEW.tool_name IS NOT NULL - AND existing.tool_name IS NOT NULL - AND NEW.tool_name <> existing.tool_name - ) - OR ( - NEW.tool_input_json IS NOT NULL - AND existing.tool_input_json IS NOT NULL - AND NEW.tool_input_json <> existing.tool_input_json - ) - ) - ) - BEGIN - SELECT RAISE(ABORT, 'session_event tool identity conflict'); - END; + CREATE UNIQUE INDEX session_event_artifact_attempt_idx + ON session_event (artifact_attempt_id) + WHERE artifact_attempt_id IS NOT NULL; + + CREATE UNIQUE INDEX session_event_mcp_terminal_winner_idx + ON session_event (session_id, mcp_command_id) + WHERE mcp_command_id IS NOT NULL; + + CREATE UNIQUE INDEX session_event_run_terminal_winner_idx + ON session_event (session_id, run_id) + WHERE semantic_hash IS NOT NULL + AND run_id IS NOT NULL + AND event_type IN ('run.cancelled', 'run.completed', 'run.failed'); + + CREATE INDEX session_event_run_stream_process_seq_idx + ON session_event (run_id, stream_id, process_type, seq); + + CREATE INDEX session_event_run_tool_call_seq_idx + ON session_event (run_id, tool_call_id, seq); CREATE TABLE session_permission_request ( created_at integer NOT NULL, @@ -142,11 +279,341 @@ function createRuntimeEventStoreDatabase( ('run-1', 'session-1'), ('run-2', 'session-2'); `); + applyDrizzleMigration(database, "0018_runtime-operation-ready-authority"); return database; } describe("session runtime event store", () => { + test("atomically persists only the latest current-run task snapshot", async () => { + const database = createRuntimeEventStoreDatabase(); + activateRun(database, { driverInstanceId: "driver-1", runId: "run-1" }); + + const result = await persistSessionRuntimeEvents(database, { + records: [ + { + event: agentTasksEvent({ + driverInstanceId: "driver-1", + id: "tasks-1", + occurredAtMs: 1_000, + runId: "run-1", + taskId: "first", + }), + occurredAt: 1_000, + sourceEventId: "source-tasks-1", + }, + { + event: agentTasksEvent({ + driverInstanceId: "driver-1", + id: "tasks-2", + occurredAtMs: 1_001, + runId: "run-1", + taskId: "latest", + }), + occurredAt: 1_001, + sourceEventId: "source-tasks-2", + }, + ], + sessionId: "session-1", + }); + const snapshot = await database + .prepare( + "SELECT driver_instance_id, run_id, seq, tasks_json FROM session_agent_task_snapshot", + ) + .first<{ + driver_instance_id: string; + run_id: string; + seq: number; + tasks_json: string; + }>(); + const eventRows = await database + .prepare("SELECT content_text, visibility FROM session_event ORDER BY seq") + .all<{ content_text: string; visibility: string }>(); + + expect(result.persistedCount).toBe(2); + expect(snapshot).toEqual({ + driver_instance_id: "driver-1", + run_id: "run-1", + seq: 2, + tasks_json: JSON.stringify({ tasks: [{ taskId: "latest" }] }), + }); + expect(eventRows.results).toEqual([ + { content_text: "1 background task active.", visibility: "owner_debug" }, + { content_text: "1 background task active.", visibility: "owner_debug" }, + ]); + }); + + test("rolls back the task event receipt when snapshot persistence fails", async () => { + const database = createRuntimeEventStoreDatabase(); + activateRun(database, { driverInstanceId: "driver-1", runId: "run-1" }); + database.execute(` + CREATE TRIGGER reject_agent_task_snapshot + BEFORE INSERT ON session_agent_task_snapshot + BEGIN + SELECT RAISE(ABORT, 'forced snapshot failure'); + END; + `); + + await expect( + persistSessionRuntimeEvents(database, { + records: [ + { + event: agentTasksEvent({ + driverInstanceId: "driver-1", + id: "tasks-1", + occurredAtMs: 1_000, + runId: "run-1", + taskId: "task-1", + }), + occurredAt: 1_000, + sourceEventId: "source-tasks-1", + }, + ], + sessionId: "session-1", + }), + ).rejects.toThrow("forced snapshot failure"); + + expect(await database.prepare("SELECT COUNT(*) AS count FROM session_event").first()).toEqual({ + count: 0, + }); + expect( + await database + .prepare("SELECT runtime_event_seq_cursor FROM session WHERE id = 'session-1'") + .first(), + ).toEqual({ runtime_event_seq_cursor: 0 }); + }); + + test("rolls back a permission receipt when its viewer projection fails", async () => { + const database = createRuntimeEventStoreDatabase(); + const event = runtimeEvent({ + driverInstanceId: "driver-1", + id: "permission-atomic-event", + kind: "permission.requested", + occurredAtMs: 1_100, + payload: { + requestId: "permission-atomic", + targetItemId: "tool-call-1", + title: "Approve command", + toolCall: { kind: "shell", toolCallId: "tool-call-1" }, + }, + runId: "run-1", + }); + const persist = () => + persistSessionRuntimeEvents(database, { + records: [{ event, occurredAt: 1_100, sourceEventId: "permission-atomic-source" }], + sessionId: "session-1", + }); + database.execute(` + CREATE TRIGGER reject_permission_projection + BEFORE INSERT ON session_permission_request + BEGIN + SELECT RAISE(ABORT, 'forced permission projection failure'); + END; + `); + + await expect(persist()).rejects.toThrow("forced permission projection failure"); + expect(await database.prepare("SELECT COUNT(*) AS count FROM session_event").first()).toEqual({ + count: 0, + }); + expect( + await database + .prepare("SELECT runtime_event_seq_cursor FROM session WHERE id = 'session-1'") + .first(), + ).toEqual({ runtime_event_seq_cursor: 0 }); + + database.execute("DROP TRIGGER reject_permission_projection"); + await expect(persist()).resolves.toMatchObject({ persistedCount: 1 }); + await expect(persist()).resolves.toMatchObject({ persistedCount: 0 }); + expect( + await database + .prepare("SELECT request_id FROM session_permission_request") + .all<{ request_id: string }>(), + ).toMatchObject({ results: [{ request_id: "permission-atomic" }] }); + expect( + await database + .prepare("SELECT runtime_event_seq_cursor FROM session WHERE id = 'session-1'") + .first(), + ).toEqual({ runtime_event_seq_cursor: 1 }); + }); + + test("atomically rejects the cursor, receipt, and snapshot when archive wins", async () => { + const database = createRuntimeEventStoreDatabase(); + activateRun(database, { driverInstanceId: "driver-1", runId: "run-1" }); + const persistBatch = database.batch.bind(database) as D1Database["batch"]; + let archivedBeforeBatch = false; + + database.batch = async (statements: D1PreparedStatement[]) => { + if (!archivedBeforeBatch) { + archivedBeforeBatch = true; + database.execute("UPDATE session SET archived_at = 2 WHERE id = 'session-1'"); + } + + return persistBatch(statements); + }; + + await expect( + persistSessionRuntimeEvents(database, { + records: [ + { + event: agentTasksEvent({ + driverInstanceId: "driver-1", + id: "tasks-after-archive", + occurredAtMs: 1_000, + runId: "run-1", + taskId: "must-not-land", + }), + occurredAt: 1_000, + sourceEventId: "source-after-archive", + }, + ], + sessionId: "session-1", + }), + ).rejects.toThrow("not writable for runtime events"); + const session = await database + .prepare("SELECT archived_at, runtime_event_seq_cursor FROM session WHERE id = 'session-1'") + .first<{ archived_at: number | null; runtime_event_seq_cursor: number }>(); + + expect(archivedBeforeBatch).toBe(true); + expect(session).toEqual({ archived_at: 2, runtime_event_seq_cursor: 0 }); + expect(await database.prepare("SELECT COUNT(*) AS count FROM session_event").first()).toEqual({ + count: 0, + }); + expect( + await database.prepare("SELECT COUNT(*) AS count FROM session_agent_task_snapshot").first(), + ).toEqual({ count: 0 }); + expect(await loadSessionAgentTaskSnapshot(database, "session-1")).toBeNull(); + }); + + test("persists an explicit empty snapshot and hides it after archive", async () => { + const database = createRuntimeEventStoreDatabase(); + activateRun(database, { driverInstanceId: "driver-1", runId: "run-1" }); + + for (const [index, taskId] of ["task-1", undefined].entries()) { + await persistSessionRuntimeEvents(database, { + records: [ + { + event: agentTasksEvent({ + driverInstanceId: "driver-1", + id: `tasks-${index + 1}`, + occurredAtMs: 1_000 + index, + runId: "run-1", + ...(taskId === undefined ? {} : { taskId }), + }), + occurredAt: 1_000 + index, + sourceEventId: `source-tasks-${index + 1}`, + }, + ], + sessionId: "session-1", + }); + } + + expect( + await database + .prepare("SELECT tasks_json FROM session_agent_task_snapshot WHERE session_id = ?") + .bind("session-1") + .first<{ tasks_json: string }>(), + ).toEqual({ tasks_json: JSON.stringify({ tasks: [] }) }); + expect(await loadSessionAgentTaskSnapshot(database, "session-1")).toEqual({ + driverInstanceId: "driver-1", + runId: "run-1", + tasks: [], + }); + + database.execute("UPDATE session SET archived_at = 2 WHERE id = 'session-1'"); + expect(await loadSessionAgentTaskSnapshot(database, "session-1")).toBeNull(); + }); + + test("does not let exact replay receipts or stale run and driver snapshots replace current state", async () => { + const database = createRuntimeEventStoreDatabase(); + activateRun(database, { driverInstanceId: "driver-1", runId: "run-1" }); + + await persistSessionRuntimeEvents(database, { + records: [ + { + event: agentTasksEvent({ + driverInstanceId: "driver-1", + id: "tasks-1", + occurredAtMs: 1_000, + runId: "run-1", + taskId: "first", + }), + occurredAt: 1_000, + sourceEventId: "source-tasks-1", + }, + ], + sessionId: "session-1", + }); + await persistSessionRuntimeEvents(database, { + records: [ + { + event: agentTasksEvent({ + driverInstanceId: "driver-1", + id: "tasks-replay", + occurredAtMs: 1_001, + runId: "run-1", + taskId: "first", + }), + occurredAt: 1_001, + sourceEventId: "source-tasks-1", + }, + ], + sessionId: "session-1", + }); + + database.execute( + "INSERT INTO session_run (id, driver_instance_id, session_id, status) VALUES ('run-3', 'driver-2', 'session-1', 'running')", + ); + activateRun(database, { driverInstanceId: "driver-2", runId: "run-3" }); + await persistSessionRuntimeEvents(database, { + records: [ + { + event: agentTasksEvent({ + driverInstanceId: "driver-2", + id: "tasks-current", + occurredAtMs: 1_002, + runId: "run-3", + taskId: "current", + }), + occurredAt: 1_002, + sourceEventId: "source-current", + }, + { + event: agentTasksEvent({ + driverInstanceId: "driver-1", + id: "tasks-old-run", + occurredAtMs: 1_003, + runId: "run-1", + taskId: "old-run", + }), + occurredAt: 1_003, + sourceEventId: "source-old-run", + }, + { + event: agentTasksEvent({ + driverInstanceId: "driver-1", + id: "tasks-old-driver", + occurredAtMs: 1_004, + runId: "run-3", + taskId: "old-driver", + }), + occurredAt: 1_004, + sourceEventId: "source-old-driver", + }, + ], + sessionId: "session-1", + }); + + expect( + await database + .prepare("SELECT driver_instance_id, run_id, tasks_json FROM session_agent_task_snapshot") + .first(), + ).toEqual({ + driver_instance_id: "driver-2", + run_id: "run-3", + tasks_json: JSON.stringify({ tasks: [{ taskId: "current" }] }), + }); + }); + test("keeps runtime event inserts within D1's bound parameter limit", async () => { const database = createRuntimeEventStoreDatabase({ maxBoundParams: 100 }); const records = Array.from({ length: 6 }, (_, index) => { @@ -206,6 +673,11 @@ describe("session runtime event store", () => { .first<{ count: number }>(); expect(count?.count).toBe(0); + expect( + await failingDatabase + .prepare("SELECT runtime_event_seq_cursor FROM session WHERE id = 'session-1'") + .first(), + ).toEqual({ runtime_event_seq_cursor: 0 }); }); test("persists mixed source ids and skips source replays before allocating sequence", async () => { @@ -287,7 +759,7 @@ describe("session runtime event store", () => { { event: runtimeEvent({ id: "event-3", - kind: "run.completed", + kind: "run.waiting", occurredAtMs: 1_200, origin: "driver", runId: "run-1", @@ -319,7 +791,7 @@ describe("session runtime event store", () => { expect(rows.results.map((row) => row.event_type)).toEqual([ "run.started", "runtime.timing.recorded", - "run.completed", + "run.waiting", ]); expect(rows.results.map((row) => row.seq)).toEqual([1, 2, 3]); expect(rows.results.map((row) => row.source_event_id)).toEqual([ @@ -389,7 +861,95 @@ describe("session runtime event store", () => { }); }); - test("persists compacted semantic columns", async () => { + test.each([ + ["adopts an exact", false], + ["rejects a changed", true], + ] as const)("%s source winner that races the atomic batch", async (_name, changed) => { + const database = createRuntimeEventStoreDatabase(); + const candidate = runtimeEvent({ + id: "candidate", + kind: "message.delta", + occurredAtMs: 1_000, + payload: { contentDelta: changed ? "changed" : "winner", messageId: "message-1" }, + runId: "run-1", + }); + const winner = runtimeEvent({ + id: "winner", + kind: "message.delta", + occurredAtMs: 999, + payload: { contentDelta: "winner", messageId: "message-1" }, + runId: "run-1", + }); + const sourceEventId = "source-race"; + const originalBatch = database.batch.bind(database) as D1Database["batch"]; + let injected = false; + database.batch = async (statements: D1PreparedStatement[]) => { + if (!injected) { + injected = true; + await insertConcurrentMessageReceipt(database, { event: winner, sourceEventId }); + } + return originalBatch(statements); + }; + + const persistence = persistSessionRuntimeEvents(database, { + records: [{ event: candidate, occurredAt: 1_000, sourceEventId }], + sessionId: "session-1", + }); + if (changed) { + await expect(persistence).rejects.toThrow("conflicts with its durable receipt"); + } else { + await expect(persistence).resolves.toMatchObject({ persistedCount: 0 }); + } + + expect( + await database + .prepare("SELECT runtime_event_seq_cursor FROM session WHERE id = 'session-1'") + .first(), + ).toEqual({ runtime_event_seq_cursor: 1 }); + expect(await database.prepare("SELECT COUNT(*) AS count FROM session_event").first()).toEqual({ + count: 1, + }); + }); + + test("adopts an exact receipt after commit succeeds but its ACK is lost", async () => { + const database = createRuntimeEventStoreDatabase(); + const event = runtimeEvent({ + id: "commit-before-ack", + kind: "message.delta", + occurredAtMs: 1_000, + payload: { contentDelta: "durable", messageId: "message-1" }, + runId: "run-1", + }); + const input = { + records: [{ event, occurredAt: 1_000, sourceEventId: "source-commit-before-ack" }], + sessionId: "session-1" as const, + }; + const originalBatch = database.batch.bind(database) as D1Database["batch"]; + let disconnected = false; + database.batch = async (statements: D1PreparedStatement[]) => { + const result = await originalBatch(statements); + if (!disconnected) { + disconnected = true; + throw new Error("injected ACK loss"); + } + return result; + }; + + await expect(persistSessionRuntimeEvents(database, input)).rejects.toThrow("injected ACK loss"); + await expect(persistSessionRuntimeEvents(database, input)).resolves.toMatchObject({ + persistedCount: 0, + }); + expect( + await database + .prepare("SELECT runtime_event_seq_cursor FROM session WHERE id = 'session-1'") + .first(), + ).toEqual({ runtime_event_seq_cursor: 1 }); + expect(await database.prepare("SELECT COUNT(*) AS count FROM session_event").first()).toEqual({ + count: 1, + }); + }); + + test("persists terminal tool semantic columns", async () => { const database = createRuntimeEventStoreDatabase(); await persistSessionRuntimeEvents(database, { @@ -411,6 +971,21 @@ describe("session runtime event store", () => { occurredAt: 2_000, sourceEventId: "source-tool-completed", }, + { + event: runtimeEvent({ + id: "tool-cancelled", + kind: "tool.call.updated", + occurredAtMs: 2_100, + payload: { + status: "cancelled", + title: "Search", + toolCallId: "tool-2", + }, + runId: "run-1", + }), + occurredAt: 2_100, + sourceEventId: "source-tool-cancelled", + }, ], sessionId: "session-1", }); @@ -427,6 +1002,7 @@ describe("session runtime event store", () => { e.tool_input_json, e.tool_name FROM session_event e + ORDER BY e.seq `, ) .all<{ @@ -439,19 +1015,98 @@ describe("session runtime event store", () => { tool_name: string | null; }>(); - expect(rows.results).toHaveLength(1); + expect(rows.results).toHaveLength(2); expect(rows.results[0]).toMatchObject({ - content_text: "Search result: ok", + content_text: "ok", event_type: "tool.call.updated", process_status: "available", process_type: "tool.use.completed", tool_call_id: "tool-1", tool_input_json: '{"query":"mosoo"}', - tool_name: null, + tool_name: "Search", + }); + expect(rows.results[1]).toMatchObject({ + content_text: "", + event_type: "tool.call.updated", + process_status: "available", + process_type: "tool.use.completed", + tool_call_id: "tool-2", + tool_input_json: null, + tool_name: "Search", }); }); - test("accepts streamed arguments but rejects terminal tool input reuse", async () => { + test.each([ + [ + "message.cancelled", + { messageId: "message-1", role: "agent" }, + "Message updated.", + "available", + "agent.message.delta", + "message-1", + ], + [ + "message.failed", + { + error: { code: "runtime.failed", message: "Runtime failed." }, + messageId: "message-1", + role: "agent", + }, + "Message updated.", + "error", + "agent.message.delta", + "message-1", + ], + [ + "thought.cancelled", + { thoughtId: "thought-1" }, + "Agent thinking updated.", + "available", + "agent.thinking.delta", + "thought-1", + ], + ] as const)( + "persists %s as a terminal semantic row", + async (kind, payload, contentText, processStatus, processType, streamId) => { + const database = createRuntimeEventStoreDatabase(); + + await persistSessionRuntimeEvents(database, { + records: [ + { + event: runtimeEvent({ + id: "terminal-event", + kind, + occurredAtMs: 2_000, + payload, + runId: "run-1", + }), + occurredAt: 2_000, + sourceEventId: "source-terminal-event", + }, + ], + sessionId: "session-1", + }); + + expect( + await database + .prepare( + ` + SELECT content_text, event_type, process_status, process_type, stream_id + FROM session_event + `, + ) + .first(), + ).toMatchObject({ + content_text: contentText, + event_type: kind, + process_status: processStatus, + process_type: processType, + stream_id: streamId, + }); + }, + ); + + test("persists replacement tool input snapshots through terminal enrichment", async () => { const database = createRuntimeEventStoreDatabase(); const persistToolEvent = (input: { id: string; @@ -504,20 +1159,18 @@ describe("session runtime event store", () => { title: "Command exited 0", }); - await expect( - persistToolEvent({ - id: "tool-conflict", - rawInput: '{"command":"python other.py","cwd":"/workspace"}', - sourceEventId: "source-tool-conflict", - status: "completed", - title: "Command exited 0", - }), - ).rejects.toThrow("session_event tool identity conflict"); + await persistToolEvent({ + id: "tool-late-enrichment", + rawInput: '{"command":"python other.py","cwd":"/workspace"}', + sourceEventId: "source-tool-late-enrichment", + status: "completed", + title: "Command exited 0", + }); const count = await database .prepare("SELECT COUNT(*) AS count FROM session_event") .first<{ count: number }>(); - expect(count?.count).toBe(3); + expect(count?.count).toBe(4); }); test("rejects runtime event batches for a different envelope session", async () => { @@ -572,73 +1225,6 @@ describe("session runtime event store", () => { ).rejects.toThrow(); }); - test("persists compacted stream fragments as semantic rows", async () => { - const database = createRuntimeEventStoreDatabase(); - const compactor = new RuntimeEventPersistenceCompactor(); - const fragments = [ - runtimeEvent({ - id: "message-start", - kind: "message.started", - occurredAtMs: 3_000, - payload: { messageId: "message-1", role: "agent" }, - runId: "run-1", - }), - runtimeEvent({ - id: "message-delta-1", - kind: "message.delta", - occurredAtMs: 3_010, - payload: { contentDelta: "Hello ", messageId: "message-1", role: "agent" }, - runId: "run-1", - }), - runtimeEvent({ - id: "message-delta-2", - kind: "message.delta", - occurredAtMs: 3_020, - payload: { contentDelta: "world", messageId: "message-1", role: "agent" }, - runId: "run-1", - }), - runtimeEvent({ - id: "message-end", - kind: "message.completed", - occurredAtMs: 3_030, - payload: { messageId: "message-1", role: "agent" }, - runId: "run-1", - }), - ]; - const compacted = compactor.compact( - fragments.map((event) => ({ - event, - occurredAt: Date.parse(event.occurredAt), - sourceEventId: `source-${event.id}`, - })), - ); - - await persistSessionRuntimeEvents(database, { - records: compacted, - sessionId: "session-1", - }); - - const rows = await database - .prepare( - ` - SELECT content_text, event_type, process_type - FROM session_event - `, - ) - .all<{ - content_text: string; - event_type: string; - process_type: string; - }>(); - - expect(rows.results).toHaveLength(1); - expect(rows.results[0]).toMatchObject({ - content_text: "Hello world", - event_type: "message.added", - process_type: "agent.message.delta", - }); - }); - test("updates viewer projections only for inserted runtime events", async () => { const database = createRuntimeEventStoreDatabase(); const permissionEvent = runtimeEvent({ @@ -720,6 +1306,11 @@ describe("session runtime event store", () => { title: "Approve command", }, ]); + expect( + await database + .prepare("SELECT status, status_seq FROM session_run WHERE id = 'run-1'") + .first(), + ).toEqual({ status: "waiting_input", status_seq: 1 }); expect(readiness === null ? null : JSON.parse(readiness.readiness_json)).toEqual({ checkedAt: "2026-05-08T00:00:04.010Z", issues: [], @@ -736,6 +1327,7 @@ describe("session runtime event store", () => { payload: { requestId: "permission-1", }, + runId: "run-1", }), occurredAt: 4_020, sourceEventId: "permission-resolved-source", @@ -749,6 +1341,11 @@ describe("session runtime event store", () => { .all<{ request_id: string }>(); expect(remainingPermissions.results).toEqual([]); + expect( + await database + .prepare("SELECT status, status_seq FROM session_run WHERE id = 'run-1'") + .first(), + ).toEqual({ status: "running", status_seq: 2 }); }); test("rejects late event batches for archived sessions without allocating sequence", async () => { @@ -890,6 +1487,42 @@ describe("session runtime event store", () => { persistedCount: 0, skippedSessionIds: ["session-1"], }); + expect( + await database + .prepare("SELECT runtime_event_seq_cursor FROM session WHERE id = 'session-1'") + .first(), + ).toEqual({ runtime_event_seq_cursor: 1 }); + }); + + test("does not consume one-event sequence numbers when the active-run fence rejects", async () => { + const database = createRuntimeEventStoreDatabase(); + database.execute("UPDATE session_run SET status = 'failed' WHERE id = 'run-1'"); + + await expect( + persistOneRuntimeEventPerSession(database, { + records: [ + { + event: runtimeEvent({ + id: "late-one-event", + kind: "agent.task.updated", + occurredAtMs: 4_050, + payload: { status: "completed" }, + runId: "run-1", + }), + occurredAt: 4_050, + sessionId: "session-1", + }, + ], + }), + ).rejects.toThrow("atomic session or active-run fence"); + expect( + await database + .prepare("SELECT runtime_event_seq_cursor FROM session WHERE id = 'session-1'") + .first(), + ).toEqual({ runtime_event_seq_cursor: 0 }); + expect(await database.prepare("SELECT COUNT(*) AS count FROM session_event").first()).toEqual({ + count: 0, + }); }); test("rejects one-event-per-session records for a different envelope session", async () => { diff --git a/apps/api/tests/session-title-mutation.test.ts b/apps/api/tests/session-title-mutation.test.ts index e1ca1559..19099bc1 100644 --- a/apps/api/tests/session-title-mutation.test.ts +++ b/apps/api/tests/session-title-mutation.test.ts @@ -4,6 +4,7 @@ import type { AuthenticatedViewer } from "../src/modules/auth/application/viewer import type { SessionSummaryRow } from "../src/modules/sessions/application/session-summary-query.service"; import { hydrateSessionSummariesFromRows } from "../src/modules/sessions/application/session-summary-query.service"; import { + applyDurableSessionAutoTitle, autoTitleSession, renameSession, } from "../src/modules/sessions/application/session-title.service"; @@ -29,6 +30,7 @@ function createSessionTitleMutationDatabase(): SqliteD1Database { id text PRIMARY KEY NOT NULL, agent_id text NOT NULL, archived_at integer, + auto_title_event_seq integer, attributed_user_id text, created_at integer NOT NULL, creator_account_id text NOT NULL, @@ -191,6 +193,62 @@ describe("session title mutations", () => { expect(session.title).toBe("Auto title"); }); + test("applies durable auto-titles by monotonic event seq", async () => { + const database = createSessionTitleMutationDatabase(); + const input = { + creatorAccountId: VIEWER.id, + eventSeq: 5, + sessionId: SESSION_ID, + title: "Title five", + }; + + await applyDurableSessionAutoTitle(database, input); + await applyDurableSessionAutoTitle(database, { + ...input, + eventSeq: 4, + title: "Stale title", + }); + await applyDurableSessionAutoTitle(database, input); + await expect( + applyDurableSessionAutoTitle(database, { ...input, title: "Conflicting title" }), + ).rejects.toThrow("replayed with conflicting content"); + await applyDurableSessionAutoTitle(database, { + ...input, + eventSeq: 6, + title: "Title six", + }); + + expect( + await database + .prepare("SELECT auto_title_event_seq, title FROM session WHERE id = ?") + .bind(SESSION_ID) + .first<{ auto_title_event_seq: number; title: string }>(), + ).toEqual({ auto_title_event_seq: 6, title: "Title six" }); + }); + + test("does not replace a title chosen outside the durable Driver stream", async () => { + const database = createSessionTitleMutationDatabase(); + + await autoTitleSession(database, VIEWER, { + appId: APP_ID, + sessionId: SESSION_ID, + title: "Prompt title", + }); + await applyDurableSessionAutoTitle(database, { + creatorAccountId: VIEWER.id, + eventSeq: 10, + sessionId: SESSION_ID, + title: "Driver title", + }); + + expect( + await database + .prepare("SELECT auto_title_event_seq, title FROM session WHERE id = ?") + .bind(SESSION_ID) + .first<{ auto_title_event_seq: number | null; title: string }>(), + ).toEqual({ auto_title_event_seq: null, title: "Prompt title" }); + }); + test("hydrates updated summary rows with their last runs", async () => { const database = createSessionTitleMutationDatabase(); diff --git a/apps/api/tests/session-viewer-event-delivery-buffer.test.ts b/apps/api/tests/session-viewer-event-delivery-buffer.test.ts index a887253b..fff9f18c 100644 --- a/apps/api/tests/session-viewer-event-delivery-buffer.test.ts +++ b/apps/api/tests/session-viewer-event-delivery-buffer.test.ts @@ -14,9 +14,49 @@ interface Deferred { interface PublishedRequest { events: AgUiSessionEvent[]; + previousRuntimeEventSeqCursor?: number; + runtimeEventSeqCursor?: number; sessionId: string; } +const largeTaskMetadata = "x".repeat(4_096); +const maxSerializedBatchBytes = 1_536 * 1024; + +function createLargeTaskSnapshot(input: { + driverInstanceId: string; + marker: string; + runId: string; +}): AgUiSessionEvent { + return createServerCustomEvent(MOSOO_CUSTOM_EVENT.sessionTasksReplaced.name, { + driverInstanceId: input.driverInstanceId, + runId: input.runId, + tasks: Array.from({ length: 120 }, (_, index) => ({ + taskId: `${input.marker}-${index}`, + taskType: largeTaskMetadata, + title: largeTaskMetadata, + })), + }); +} + +function createTerminalEvent(runId = "run-1"): AgUiSessionEvent { + return createServerCustomEvent(MOSOO_CUSTOM_EVENT.sessionRunUpdated.name, { + driverInstanceId: null, + lifecycle: "IDLE", + run: { + completedAt: "2026-04-30T00:00:01.000Z", + error: null, + id: runId, + startedAt: "2026-04-30T00:00:00.000Z", + status: "completed", + traceId: null, + }, + }); +} + +function serializedEventBytes(events: AgUiSessionEvent[]): number { + return new TextEncoder().encode(JSON.stringify(events)).byteLength; +} + function createDeferred(): Deferred { let rejectDeferred: (reason?: unknown) => void = () => {}; let resolveDeferred: (value: T) => void = () => {}; @@ -33,20 +73,61 @@ function createDeferred(): Deferred { } function createBufferHarness(): { + batchPublishCount: () => number; buffer: SessionViewerEventDeliveryBuffer; published: PublishedRequest[]; + pushAfterResponse: (callback: () => void) => void; pushResponse: (response: Promise | Response) => void; + stateSyncCount: () => number; waitForPublish: () => Promise; waitForWaitUntil: () => Promise; } { const publishWaiters: Array<() => void> = []; const published: PublishedRequest[] = []; + const afterResponseCallbacks: Array<() => void> = []; const responses: Array | Response> = []; const waitUntilTasks: Promise[] = []; + let batchPublishCount = 0; + let stateSyncCount = 0; const sessionStub = { - async publishEvents(sessionId: string, events: AgUiSessionEvent[]): Promise { + async publishEventBatches( + sessionId: string, + batches: Array<{ + events: AgUiSessionEvent[]; + previousRuntimeEventSeqCursor: number | null; + runtimeEventSeqCursor: number | null; + }>, + ): Promise { + batchPublishCount += 1; + for (const batch of batches) { + published.push({ + events: batch.events, + ...(batch.previousRuntimeEventSeqCursor === null + ? {} + : { previousRuntimeEventSeqCursor: batch.previousRuntimeEventSeqCursor }), + ...(batch.runtimeEventSeqCursor === null + ? {} + : { runtimeEventSeqCursor: batch.runtimeEventSeqCursor }), + sessionId, + }); + } + publishWaiters.shift()?.(); + const response = await (responses.shift() ?? new Response(null, { status: 204 })); + if (!response.ok) { + throw new Error(`Session event publish failed with status ${response.status}.`); + } + afterResponseCallbacks.shift()?.(); + }, + async publishEvents( + sessionId: string, + events: AgUiSessionEvent[], + runtimeEventSeqCursor: number | null, + previousRuntimeEventSeqCursor: number | null, + ): Promise { published.push({ events, + ...(previousRuntimeEventSeqCursor === null ? {} : { previousRuntimeEventSeqCursor }), + ...(runtimeEventSeqCursor === null ? {} : { runtimeEventSeqCursor }), sessionId, }); publishWaiters.shift()?.(); @@ -56,6 +137,11 @@ function createBufferHarness(): { if (!response.ok) { throw new Error(`Session event publish failed with status ${response.status}.`); } + + afterResponseCallbacks.shift()?.(); + }, + async syncViewers(): Promise { + stateSyncCount += 1; }, }; const env = { @@ -77,11 +163,16 @@ function createBufferHarness(): { }); return { + batchPublishCount: () => batchPublishCount, buffer, published, + pushAfterResponse: (callback) => { + afterResponseCallbacks.push(callback); + }, pushResponse: (response) => { responses.push(response); }, + stateSyncCount: () => stateSyncCount, waitForPublish: () => new Promise((resolve) => { publishWaiters.push(resolve); @@ -139,17 +230,7 @@ describe("SessionViewerEventDeliveryBuffer", () => { test("flushes terminal events immediately", async () => { const { buffer, published, waitForWaitUntil } = createBufferHarness(); - const terminalEvent = createServerCustomEvent(MOSOO_CUSTOM_EVENT.sessionRunUpdated.name, { - lifecycle: "IDLE", - run: { - completedAt: "2026-04-30T00:00:01.000Z", - error: null, - id: "run-1", - startedAt: "2026-04-30T00:00:00.000Z", - status: "completed", - traceId: null, - }, - }); + const terminalEvent = createTerminalEvent(); buffer.enqueue("session-1", [ { delta: "done", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }, @@ -163,6 +244,95 @@ describe("SessionViewerEventDeliveryBuffer", () => { ]); }); + test("keeps a delayed flush alive until the buffered events are published", async () => { + const { buffer, published, waitForWaitUntil } = createBufferHarness(); + + buffer.enqueue("session-1", [ + { delta: "batched", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }, + ]); + await waitForWaitUntil(); + + expect(published).toEqual([ + { + events: [{ delta: "batched", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }], + sessionId: "session-1", + }, + ]); + }); + + test("keeps different durable cursor generations in separate deliveries", async () => { + const { batchPublishCount, buffer, published } = createBufferHarness(); + + buffer.enqueue( + "session-1", + [{ delta: "one", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }], + 1, + 0, + ); + buffer.enqueue( + "session-1", + [{ delta: "two", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }], + 2, + 1, + ); + await buffer.flush(); + + expect(batchPublishCount()).toBe(1); + expect( + published.map(({ events, previousRuntimeEventSeqCursor, runtimeEventSeqCursor }) => ({ + events, + previousRuntimeEventSeqCursor, + runtimeEventSeqCursor, + })), + ).toEqual([ + { + events: [{ delta: "one", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }], + previousRuntimeEventSeqCursor: 0, + runtimeEventSeqCursor: 1, + }, + { + events: [{ delta: "two", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }], + previousRuntimeEventSeqCursor: 1, + runtimeEventSeqCursor: 2, + }, + ]); + }); + + test("keeps one durable cursor generation in one atomic delivery", async () => { + const { batchPublishCount, buffer, published } = createBufferHarness(); + const events: AgUiSessionEvent[] = Array.from({ length: 128 }, (_, index) => ({ + delta: [{ op: "replace", path: `/commands/${index}`, value: index }], + type: "STATE_DELTA", + })); + + buffer.enqueue("session-1", events, 128, 0); + await buffer.flush(); + + expect(batchPublishCount()).toBe(1); + expect(published).toHaveLength(1); + expect(published[0]).toMatchObject({ + events, + previousRuntimeEventSeqCursor: 0, + runtimeEventSeqCursor: 128, + }); + }); + + test("falls back to an authoritative state sync for an oversized durable generation", async () => { + const { batchPublishCount, buffer, published, stateSyncCount } = createBufferHarness(); + const oversized = createServerCustomEvent(MOSOO_CUSTOM_EVENT.sessionTasksReplaced.name, { + driverInstanceId: "driver-1", + runId: "run-1", + tasks: [{ taskId: "oversized", title: "x".repeat(maxSerializedBatchBytes + 1) }], + }); + + buffer.enqueue("session-1", [oversized], 1, 0); + await buffer.flush(); + + expect(batchPublishCount()).toBe(0); + expect(published).toEqual([]); + expect(stateSyncCount()).toBe(1); + }); + test("flushes the first delta of a run immediately, then batches the rest", async () => { const { buffer, published, waitForWaitUntil } = createBufferHarness(); const runStarted = { @@ -193,28 +363,144 @@ describe("SessionViewerEventDeliveryBuffer", () => { buffer.enqueue("session-1", [ { delta: "there", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }, ]); - await waitForWaitUntil(); expect(published).toHaveLength(1); await buffer.flush(); expect(published).toHaveLength(2); }); - test("requeues failed deliveries before events enqueued during the failed publish", async () => { + test("does not strand a terminal batch queued as the prior delivery settles", async () => { + const { buffer, published, pushAfterResponse, waitForWaitUntil } = createBufferHarness(); + const terminalEvent = createTerminalEvent(); + + pushAfterResponse(() => { + // Cross the Session stub, client, and drain continuations so this lands + // after the drain resolves but before a chained cleanup reaction can run. + queueMicrotask(() => { + queueMicrotask(() => { + queueMicrotask(() => { + buffer.enqueue("session-1", [terminalEvent]); + }); + }); + }); + }); + buffer.enqueue("session-1", [ + { delta: "done", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }, + ]); + await buffer.flush(); + await waitForWaitUntil(); + + expect(published.map((request) => request.events)).toEqual([ + [{ delta: "done", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }], + [terminalEvent], + ]); + }); + + test("coalesces a same-generation large task snapshot flood to the latest event", async () => { + const { buffer, published } = createBufferHarness(); + let latestSnapshot: AgUiSessionEvent | null = null; + + for (let index = 0; index < 34; index += 1) { + latestSnapshot = createLargeTaskSnapshot({ + driverInstanceId: "driver-1", + marker: `snapshot-${index}`, + runId: "run-1", + }); + buffer.enqueue("session-1", [latestSnapshot]); + } + await buffer.flush(); + + expect(published).toHaveLength(1); + expect(published[0]?.events).toEqual([latestSnapshot]); + }); + + test("keeps only the in-flight and latest same-generation snapshot during a slow publish", async () => { const { buffer, published, pushResponse, waitForPublish } = createBufferHarness(); + const response = createDeferred(); + const inFlightSnapshot = createLargeTaskSnapshot({ + driverInstanceId: "driver-1", + marker: "in-flight", + runId: "run-1", + }); + const pendingSnapshots = Array.from({ length: 5 }, (_, index) => + createLargeTaskSnapshot({ + driverInstanceId: "driver-1", + marker: `pending-${index}`, + runId: "run-1", + }), + ); + pushResponse(response.promise); + + buffer.enqueue("session-1", [inFlightSnapshot]); + const firstPublish = waitForPublish(); + const flush = buffer.flush(); + await firstPublish; + + for (const [index, snapshot] of pendingSnapshots.entries()) { + buffer.enqueue("session-1", [snapshot, createTerminalEvent(`terminal-${index}`)]); + } + + expect(published.map((request) => request.events)).toEqual([[inFlightSnapshot]]); + + response.resolve(new Response(null, { status: 204 })); + await flush; + + const deliveredSnapshots = published + .flatMap((request) => request.events) + .filter( + (event) => + event.type === "CUSTOM" && event.name === MOSOO_CUSTOM_EVENT.sessionTasksReplaced.name, + ); + expect(deliveredSnapshots).toEqual([inFlightSnapshot, pendingSnapshots.at(-1)]); + }); + + test("bounds a cross-generation backlog behind one authoritative state sync", async () => { + const { buffer, published, stateSyncCount } = createBufferHarness(); + const snapshots = Array.from({ length: 3 }, (_, index) => + createLargeTaskSnapshot({ + driverInstanceId: `driver-${index}`, + marker: `snapshot-${index}`, + runId: `run-${index}`, + }), + ); + + buffer.enqueue("session-1", [snapshots[0]]); + expect(published).toHaveLength(0); + buffer.enqueue("session-1", [snapshots[1]]); + expect(published[0]?.events).toEqual([snapshots[0]]); + buffer.enqueue("session-1", [snapshots[2]]); + await buffer.flush(); + + expect(published).toHaveLength(1); + expect(published.map((request) => request.events.length)).toEqual([1]); + expect( + published.every((request) => serializedEventBytes(request.events) <= maxSerializedBatchBytes), + ).toBe(true); + expect(stateSyncCount()).toBe(1); + }); + + test("retries failed bounded batches before events enqueued during the failed publish", async () => { + const { buffer, published, pushResponse, stateSyncCount, waitForPublish } = + createBufferHarness(); const failedResponse = createDeferred(); + const failedSnapshot = createLargeTaskSnapshot({ + driverInstanceId: "driver-1", + marker: "failed", + runId: "run-1", + }); + const nextSnapshot = createLargeTaskSnapshot({ + driverInstanceId: "driver-2", + marker: "next", + runId: "run-2", + }); pushResponse(failedResponse.promise); - buffer.enqueue("session-1", [ - { delta: "A", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }, - ]); + buffer.enqueue("session-1", [failedSnapshot]); const firstPublish = waitForPublish(); const failedFlush = buffer.flush().catch((error: unknown) => error); await firstPublish; - buffer.enqueue("session-1", [ - { delta: "B", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }, - ]); + buffer.enqueue("session-1", [nextSnapshot]); failedResponse.resolve( new Response(JSON.stringify({ error: "publish failed" }), { headers: { "content-type": "application/json" }, @@ -225,15 +511,42 @@ describe("SessionViewerEventDeliveryBuffer", () => { expect(await failedFlush).toBeInstanceOf(Error); await buffer.flush(); - expect(published).toEqual([ - { - events: [{ delta: "A", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }], - sessionId: "session-1", - }, - { - events: [{ delta: "AB", messageId: "assistant-1", type: "TEXT_MESSAGE_CONTENT" }], - sessionId: "session-1", - }, + expect(published.map((request) => request.events)).toEqual([ + [failedSnapshot], + [failedSnapshot], + ]); + expect(stateSyncCount()).toBe(1); + expect( + published.every((request) => serializedEventBytes(request.events) <= maxSerializedBatchBytes), + ).toBe(true); + }); + + test("drops derived payloads after delivery failure instead of hot-looping", async () => { + const { buffer, published, pushResponse } = createBufferHarness(); + const snapshots = Array.from({ length: 5 }, (_, index) => + createLargeTaskSnapshot({ + driverInstanceId: "driver-1", + marker: `snapshot-${index}`, + runId: "run-1", + }), + ); + + for (const snapshot of snapshots) { + pushResponse(new Response(null, { status: 500 })); + buffer.enqueue("session-1", [snapshot]); + await buffer.flushSafely(); + } + await buffer.flush(); + + expect(published.map((request) => request.events)).toEqual([ + [snapshots[0]], + [snapshots[1]], + [snapshots[2]], + [snapshots[3]], + [snapshots[4]], ]); + expect( + published.every((request) => serializedEventBytes(request.events) <= maxSerializedBatchBytes), + ).toBe(true); }); }); diff --git a/apps/api/tests/session-viewer-socket-state-order.test.ts b/apps/api/tests/session-viewer-socket-state-order.test.ts new file mode 100644 index 00000000..d0ff3865 --- /dev/null +++ b/apps/api/tests/session-viewer-socket-state-order.test.ts @@ -0,0 +1,754 @@ +import { describe, expect, test } from "bun:test"; + +import { + applyAgUiEventsToSessionLiveState, + createInitialSessionLiveState, + createServerCustomEvent, + createViewerCustomEvent, + MOSOO_CUSTOM_EVENT, + parseAgUiSessionEventJson, +} from "@mosoo/ag-ui-session"; +import { createPromiseDeferred } from "@mosoo/effects"; + +import type { AuthenticatedViewer } from "../src/modules/auth/application/viewer-auth.service"; +import { loadSessionViewerStateSnapshot } from "../src/modules/sessions/infrastructure/session-viewer-live-snapshot.repository"; +import { writeSessionViewerSocketHeaders } from "../src/modules/sessions/infrastructure/session/socket-headers"; +import { SessionViewerSocketHub } from "../src/modules/sessions/infrastructure/session/viewer-socket-hub"; +import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; +import { + createPublicHttpContractDatabase, + createPublicHttpTestBindings, + insertOwnerSession, + PUBLIC_API_TEST_IDS, +} from "./helpers/public-api-http-test-fixture"; +import type { SqliteD1Database } from "./helpers/sqlite-d1"; +import { insertRuntimeEvent } from "./public-thread-api-fixtures"; + +const DRIVER_ID = PUBLIC_API_TEST_IDS.driverOwner; +const RUN_ID = PUBLIC_API_TEST_IDS.run; +const SESSION_ID = PUBLIC_API_TEST_IDS.ownerSession; +const VIEWER: AuthenticatedViewer = { + email: "owner@example.com", + emailVerified: true, + id: PUBLIC_API_TEST_IDS.ownerAccount, + imageUrl: null, + name: "Owner", +}; + +class TestSocket { + attachment: unknown = null; + readonly frames: string[] = []; + readyState = WebSocket.OPEN; + + close(): void { + this.readyState = WebSocket.CLOSED; + } + + deserializeAttachment(): unknown { + return this.attachment; + } + + send(frame: string): void { + this.frames.push(frame); + } + + serializeAttachment(attachment: unknown): void { + this.attachment = attachment; + } +} + +function createContext(): { + ctx: DurableObjectState; + pending: Promise[]; +} { + const accepted: { socket: TestSocket; tags: string[] }[] = []; + const pending: Promise[] = []; + const stored = new Map(); + const storage = { + delete: async (key: string) => stored.delete(key), + deleteAlarm: async () => {}, + get: async (key: string) => stored.get(key) as T | undefined, + put: async (key: string, value: unknown) => { + stored.set(key, value); + }, + setAlarm: async () => {}, + }; + const ctx = { + acceptWebSocket(socket: TestSocket, tags: string[]) { + accepted.push({ socket, tags }); + }, + getWebSockets(tag?: string) { + return accepted + .filter( + ({ socket, tags }) => + socket.readyState === WebSocket.OPEN && (tag === undefined || tags.includes(tag)), + ) + .map(({ socket }) => socket); + }, + storage, + waitUntil(promise: Promise) { + pending.push(promise); + }, + } as unknown as DurableObjectState; + + return { ctx, pending }; +} + +function deferFirstTaskSnapshotRead(database: SqliteD1Database): { + database: D1Database; + readStarted: Promise; + releaseRead: () => void; +} { + const readStarted = createPromiseDeferred(); + const releaseRead = createPromiseDeferred(); + let deferred = false; + + function wrapStatement(statement: D1PreparedStatement): D1PreparedStatement { + return new Proxy(statement, { + get(target, property) { + if (property === "all" || property === "first" || property === "raw") { + return async (...args: unknown[]) => { + const read = Reflect.get(target, property) as ( + ...values: unknown[] + ) => Promise; + const result = await read.apply(target, args); + + if (!deferred) { + deferred = true; + readStarted.resolve(); + await releaseRead.promise; + } + + return result; + }; + } + + if (property === "bind") { + return (...values: unknown[]) => wrapStatement(target.bind(...values)); + } + + const value: unknown = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + } + + return { + database: new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => { + const statement = target.prepare(query); + return query.includes("session_agent_task_snapshot") + ? wrapStatement(statement) + : statement; + }; + } + + const value: unknown = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }), + readStarted: readStarted.promise, + releaseRead: () => { + releaseRead.resolve(); + }, + }; +} + +function mutateAfterFirstViewerSnapshotRead( + database: SqliteD1Database, + mutate: () => Promise, +): { database: D1Database; readCount: () => number } { + let mutated = false; + let readCount = 0; + + function wrapStatement(statement: D1PreparedStatement, query: string): D1PreparedStatement { + return new Proxy(statement, { + get(target, property) { + if (property === "all" || property === "first" || property === "raw") { + return async (...args: unknown[]) => { + const read = Reflect.get(target, property) as ( + ...values: unknown[] + ) => Promise; + const result = await read.apply(target, args); + if (query.includes("session_agent_task_snapshot")) { + readCount += 1; + if (!mutated) { + mutated = true; + await mutate(); + } + } + return result; + }; + } + if (property === "bind") { + return (...values: unknown[]) => wrapStatement(target.bind(...values), query); + } + const value: unknown = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + } + + return { + database: new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => wrapStatement(target.prepare(query), query); + } + const value: unknown = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }), + readCount: () => readCount, + }; +} + +function failFirstViewerSnapshotRead(database: SqliteD1Database): D1Database { + let failed = false; + + function wrapStatement(statement: D1PreparedStatement, query: string): D1PreparedStatement { + return new Proxy(statement, { + get(target, property) { + if (property === "all" || property === "first" || property === "raw") { + return async (...args: unknown[]) => { + if (!failed && query.includes("session_agent_task_snapshot")) { + failed = true; + throw new Error("injected initial viewer snapshot failure"); + } + + const read = Reflect.get(target, property) as ( + ...values: unknown[] + ) => Promise; + return read.apply(target, args); + }; + } + if (property === "bind") { + return (...values: unknown[]) => wrapStatement(target.bind(...values), query); + } + const value: unknown = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }); + } + + return new Proxy(database, { + get(target, property) { + if (property === "prepare") { + return (query: string) => wrapStatement(target.prepare(query), query); + } + const value: unknown = Reflect.get(target, property); + return typeof value === "function" ? value.bind(target) : value; + }, + }); +} + +function installTestWebSocketPair(): { restore(): void; sockets: TestSocket[] } { + const sockets: TestSocket[] = []; + const originalWebSocketPair = Reflect.get(globalThis, "WebSocketPair"); + Reflect.set(globalThis, "WebSocketPair", function TestWebSocketPair() { + const client = new TestSocket(); + const server = new TestSocket(); + sockets.push(client, server); + return [client, server]; + }); + + return { + restore() { + if (originalWebSocketPair === undefined) { + Reflect.deleteProperty(globalThis, "WebSocketPair"); + } else { + Reflect.set(globalThis, "WebSocketPair", originalWebSocketPair); + } + }, + sockets, + }; +} + +function createViewerSocketRequest(): Request { + const headers = new Headers({ upgrade: "websocket" }); + writeSessionViewerSocketHeaders(headers, { + publicOrigin: "https://mosoo.ai", + appId: PUBLIC_API_TEST_IDS.app, + sessionId: SESSION_ID, + viewer: VIEWER, + }); + return new Request("https://session.internal/viewer/ws", { headers }); +} + +async function createDatabase(): Promise { + const database = await createPublicHttpContractDatabase(); + const now = Date.now(); + await insertOwnerSession(database); + database.execute(` + UPDATE session + SET last_run_id = '${RUN_ID}', status = 'RUNNING' + WHERE id = '${SESSION_ID}'; + + INSERT INTO session_run ( + agent_id, + created_at, + created_by_account_id, + driver_instance_id, + id, + model, + provider, + session_id, + started_at, + status, + trace_id, + trigger, + updated_at + ) VALUES ( + '${PUBLIC_API_TEST_IDS.agent}', + ${now}, + '${VIEWER.id}', + '${DRIVER_ID}', + '${RUN_ID}', + 'gpt-5.4', + 'openai', + '${SESSION_ID}', + ${now}, + 'running', + 'trace-1', + 'user_message', + ${now} + ); + + INSERT INTO session_agent_task_snapshot ( + driver_instance_id, + run_id, + seq, + session_id, + tasks_json + ) VALUES ( + '${DRIVER_ID}', + '${RUN_ID}', + 1, + '${SESSION_ID}', + '{"tasks":[{"taskId":"stale"}]}' + ); + `); + return database; +} + +describe("session viewer socket state ordering", () => { + test("retries a viewer snapshot when its durable cursor changes mid-read", async () => { + const database = await createDatabase(); + const messageId = "01J0000000000000000000002E"; + const guarded = mutateAfterFirstViewerSnapshotRead(database, async () => { + await insertRuntimeEvent(database, { + eventId: "01J0000000000000000000002F", + kind: "message.added", + occurredAt: Date.now(), + payload: { content: "Committed during snapshot.", messageId, role: "agent" }, + runId: RUN_ID, + seq: 1, + sessionId: SESSION_ID, + }); + await database + .prepare("UPDATE session SET runtime_event_seq_cursor = 1 WHERE id = ?") + .bind(SESSION_ID) + .run(); + }); + + const snapshot = await loadSessionViewerStateSnapshot(guarded.database, { + sessionId: SESSION_ID, + viewerId: VIEWER.id, + }); + + expect(guarded.readCount()).toBe(2); + expect(snapshot.runtimeEventSeqCursor).toBe(1); + expect(snapshot.state.messages).toContainEqual( + expect.objectContaining({ content: "Committed during snapshot.", id: messageId }), + ); + }); + + test("linearizes a cold initial snapshot before a concurrent task broadcast", async () => { + const database = await createDatabase(); + const deferredRead = deferFirstTaskSnapshotRead(database); + const bindings = { + ...createPublicHttpTestBindings(database), + DB: deferredRead.database, + } as ApiBindings; + const { ctx, pending } = createContext(); + const hub = new SessionViewerSocketHub({ + ctx, + env: bindings, + getSessionId: () => SESSION_ID, + rememberSessionId: () => {}, + withSessionLogContext: (operation) => operation(), + }); + const sockets: TestSocket[] = []; + const originalWebSocketPair = Reflect.get(globalThis, "WebSocketPair"); + Reflect.set(globalThis, "WebSocketPair", function TestWebSocketPair() { + const client = new TestSocket(); + const server = new TestSocket(); + sockets.push(client, server); + return [client, server]; + }); + + try { + const headers = new Headers({ upgrade: "websocket" }); + writeSessionViewerSocketHeaders(headers, { + publicOrigin: "https://mosoo.ai", + appId: PUBLIC_API_TEST_IDS.app, + sessionId: SESSION_ID, + viewer: VIEWER, + }); + expect( + hub.connect(new Request("https://session.internal/viewer/ws", { headers })).status, + ).toBe(101); + const server = sockets[1]; + + if (!server) { + throw new Error("Expected a server websocket."); + } + + await deferredRead.readStarted; + database.execute(` + UPDATE session_agent_task_snapshot + SET seq = 2, tasks_json = '{"tasks":[{"taskId":"latest"}]}' + WHERE session_id = '${SESSION_ID}'; + `); + const latestTaskEvent = createServerCustomEvent( + MOSOO_CUSTOM_EVENT.sessionTasksReplaced.name, + { + driverInstanceId: DRIVER_ID, + runId: RUN_ID, + tasks: [{ taskId: "latest" }], + }, + ); + const broadcast = hub.broadcastEvents([latestTaskEvent]); + + await Promise.resolve(); + expect(server.frames).toEqual([]); + + deferredRead.releaseRead(); + await broadcast; + await Promise.all(pending); + + const firstEvents = server.frames.map(parseAgUiSessionEventJson); + expect(firstEvents.map((event) => event.type)).toEqual(["STATE_SNAPSHOT", "CUSTOM"]); + const state = applyAgUiEventsToSessionLiveState( + createInitialSessionLiveState({ sessionId: SESSION_ID, title: null, viewerId: VIEWER.id }), + firstEvents, + ); + expect(state.taskSnapshot?.tasks).toEqual([{ taskId: "latest" }]); + + const durableMessageId = "01J0000000000000000000002A"; + await insertRuntimeEvent(database, { + eventId: "01J0000000000000000000002B", + kind: "message.added", + occurredAt: Date.now(), + payload: { + content: "Durable replay content.", + messageId: durableMessageId, + role: "agent", + }, + runId: RUN_ID, + seq: 1, + sessionId: SESSION_ID, + }); + database.execute(` + UPDATE session + SET runtime_event_seq_cursor = 1 + WHERE id = '${SESSION_ID}'; + `); + + expect( + hub.connect( + new Request("https://session.internal/viewer/ws", { + headers: new Headers(headers), + }), + ).status, + ).toBe(101); + const secondServer = sockets[3]; + if (!secondServer) { + throw new Error("Expected a second server websocket."); + } + await Promise.all(pending); + await hub.broadcastStateSync(); + + const durableSnapshot = parseAgUiSessionEventJson(server.frames.at(-1) ?? ""); + const secondDurableSnapshot = parseAgUiSessionEventJson(secondServer.frames.at(-1) ?? ""); + expect(secondDurableSnapshot).toEqual(durableSnapshot); + expect(durableSnapshot.type).toBe("STATE_SNAPSHOT"); + if (durableSnapshot.type !== "STATE_SNAPSHOT") { + throw new Error("Expected durable state snapshot."); + } + expect(durableSnapshot.snapshot.messages).toEqual([ + expect.objectContaining({ + content: "Durable replay content.", + id: durableMessageId, + }), + ]); + expect(durableSnapshot.snapshot.viewerId).toBe(VIEWER.id); + + const framesBeforeReplay = server.frames.length; + await hub.broadcastEvents( + [ + { + delta: " must-not-repeat", + messageId: durableMessageId, + type: "TEXT_MESSAGE_CONTENT", + }, + ], + 1, + 0, + ); + expect(server.frames).toHaveLength(framesBeforeReplay); + + await insertRuntimeEvent(database, { + eventId: "01J0000000000000000000002C", + kind: "message.delta", + occurredAt: Date.now(), + payload: { + content: " Appended once.", + messageId: durableMessageId, + }, + runId: RUN_ID, + seq: 2, + sessionId: SESSION_ID, + }); + database.execute(` + UPDATE session + SET runtime_event_seq_cursor = 2 + WHERE id = '${SESSION_ID}'; + `); + await hub.broadcastEvents( + [ + { + delta: " Appended once.", + messageId: durableMessageId, + type: "TEXT_MESSAGE_CONTENT", + }, + ], + 2, + 1, + ); + + await insertRuntimeEvent(database, { + eventId: "01J0000000000000000000002D", + kind: "message.delta", + occurredAt: Date.now(), + payload: { + content: " Gap recovered.", + messageId: durableMessageId, + }, + runId: RUN_ID, + seq: 3, + sessionId: SESSION_ID, + }); + database.execute(` + UPDATE session + SET runtime_event_seq_cursor = 3 + WHERE id = '${SESSION_ID}'; + `); + await hub.broadcastEvents( + [ + { + delta: " Gap recovered.", + messageId: durableMessageId, + type: "TEXT_MESSAGE_CONTENT", + }, + ], + 3, + 0, + ); + const recoveredGapSnapshot = parseAgUiSessionEventJson(server.frames.at(-1) ?? ""); + expect(recoveredGapSnapshot).toMatchObject({ + snapshot: { + messages: [ + expect.objectContaining({ + content: "Durable replay content. Appended once. Gap recovered.", + id: durableMessageId, + }), + ], + }, + type: "STATE_SNAPSHOT", + }); + + await hub.handleSocketMessage( + server as unknown as WebSocket, + JSON.stringify( + createViewerCustomEvent(MOSOO_CUSTOM_EVENT.sessionSyncRequest.name, { + reason: "reconnect", + }), + ), + ); + + const reconnectSnapshot = parseAgUiSessionEventJson(server.frames.at(-1) ?? ""); + expect(reconnectSnapshot.type).toBe("STATE_SNAPSHOT"); + if (reconnectSnapshot.type !== "STATE_SNAPSHOT") { + throw new Error("Expected reconnect state snapshot."); + } + expect(reconnectSnapshot.snapshot.taskSnapshot?.tasks).toEqual([{ taskId: "latest" }]); + expect(reconnectSnapshot.snapshot.messages).toEqual([ + expect.objectContaining({ + content: "Durable replay content. Appended once. Gap recovered.", + id: durableMessageId, + }), + ]); + + await insertRuntimeEvent(database, { + eventId: "01J0000000000000000000002G", + kind: "message.delta", + occurredAt: Date.now(), + payload: { + content: " Alarm recovered.", + messageId: durableMessageId, + }, + runId: RUN_ID, + seq: 4, + sessionId: SESSION_ID, + }); + database.execute(` + UPDATE session + SET runtime_event_seq_cursor = 4 + WHERE id = '${SESSION_ID}'; + `); + await hub.handleAlarm(); + + const alarmSnapshot = parseAgUiSessionEventJson(server.frames.at(-1) ?? ""); + expect(alarmSnapshot).toMatchObject({ + snapshot: { + messages: [ + expect.objectContaining({ + content: "Durable replay content. Appended once. Gap recovered. Alarm recovered.", + id: durableMessageId, + }), + ], + }, + type: "STATE_SNAPSHOT", + }); + expect(server.attachment).toMatchObject({ runtimeEventSeqCursor: 4 }); + } finally { + if (originalWebSocketPair === undefined) { + Reflect.deleteProperty(globalThis, "WebSocketPair"); + } else { + Reflect.set(globalThis, "WebSocketPair", originalWebSocketPair); + } + } + }); + + test("retries the authoritative snapshot before delivering a delta after initial sync fails", async () => { + const database = await createDatabase(); + const { ctx, pending } = createContext(); + const pair = installTestWebSocketPair(); + const hub = new SessionViewerSocketHub({ + ctx, + env: { + ...createPublicHttpTestBindings(database), + DB: failFirstViewerSnapshotRead(database), + } as ApiBindings, + getSessionId: () => SESSION_ID, + rememberSessionId: () => {}, + withSessionLogContext: (operation) => operation(), + }); + + try { + expect(hub.connect(createViewerSocketRequest()).status).toBe(101); + await Promise.allSettled(pending); + const server = pair.sockets[1]; + if (!server) { + throw new Error("Expected a server websocket."); + } + expect(server.frames).toEqual([]); + expect(server.attachment).not.toHaveProperty("runtimeEventSeqCursor"); + + const messageId = "01J0000000000000000000002H"; + await insertRuntimeEvent(database, { + eventId: "01J0000000000000000000002J", + kind: "message.added", + occurredAt: Date.now(), + payload: { content: "Recovered base.", messageId, role: "agent" }, + runId: RUN_ID, + seq: 1, + sessionId: SESSION_ID, + }); + database.execute(` + UPDATE session + SET runtime_event_seq_cursor = 1 + WHERE id = '${SESSION_ID}'; + `); + + await hub.broadcastEvents( + [{ delta: "Recovered base.", messageId, type: "TEXT_MESSAGE_CONTENT" }], + 1, + 0, + ); + + expect(server.frames).toHaveLength(1); + expect(parseAgUiSessionEventJson(server.frames[0] ?? "")).toMatchObject({ + snapshot: { + messages: [expect.objectContaining({ content: "Recovered base.", id: messageId })], + }, + type: "STATE_SNAPSHOT", + }); + expect(server.attachment).toMatchObject({ runtimeEventSeqCursor: 1 }); + } finally { + pair.restore(); + } + }); + + test("recovers the Session identity from socket attachments after hibernation", async () => { + const database = await createDatabase(); + const { ctx, pending } = createContext(); + const pair = installTestWebSocketPair(); + const initialHub = new SessionViewerSocketHub({ + ctx, + env: createPublicHttpTestBindings(database) as ApiBindings, + getSessionId: () => SESSION_ID, + rememberSessionId: () => {}, + withSessionLogContext: (operation) => operation(), + }); + + try { + expect(initialHub.connect(createViewerSocketRequest()).status).toBe(101); + await Promise.all(pending); + const server = pair.sockets[1]; + if (!server) { + throw new Error("Expected a server websocket."); + } + + const messageId = "01J0000000000000000000002K"; + await insertRuntimeEvent(database, { + eventId: "01J0000000000000000000002M", + kind: "message.added", + occurredAt: Date.now(), + payload: { content: "Alarm recovery.", messageId, role: "agent" }, + runId: RUN_ID, + seq: 1, + sessionId: SESSION_ID, + }); + database.execute(` + UPDATE session + SET runtime_event_seq_cursor = 1 + WHERE id = '${SESSION_ID}'; + `); + + const recoveredSessionIds: string[] = []; + const hibernatedHub = new SessionViewerSocketHub({ + ctx, + env: createPublicHttpTestBindings(database) as ApiBindings, + getSessionId: () => null, + rememberSessionId: (sessionId) => recoveredSessionIds.push(sessionId), + withSessionLogContext: (operation) => operation(), + }); + await hibernatedHub.handleAlarm(); + + expect(recoveredSessionIds).toContain(SESSION_ID); + expect(parseAgUiSessionEventJson(server.frames.at(-1) ?? "")).toMatchObject({ + snapshot: { + messages: [expect.objectContaining({ content: "Alarm recovery.", id: messageId })], + }, + type: "STATE_SNAPSHOT", + }); + expect(server.attachment).toMatchObject({ runtimeEventSeqCursor: 1 }); + } finally { + pair.restore(); + } + }); +}); diff --git a/apps/api/tests/session-viewer-state.test.ts b/apps/api/tests/session-viewer-state.test.ts index 4de2257a..5deb2864 100644 --- a/apps/api/tests/session-viewer-state.test.ts +++ b/apps/api/tests/session-viewer-state.test.ts @@ -1,15 +1,26 @@ import { describe, expect, test } from "bun:test"; import { loadSessionViewerState } from "../src/modules/sessions/infrastructure/session-viewer-live-snapshot.repository"; +import { loadViewerLiveState } from "../src/modules/sessions/infrastructure/session/viewer-live-state"; import { SqliteD1Database } from "./helpers/sqlite-d1"; +const VIEWER = { + email: "viewer@example.com", + emailVerified: true, + id: "viewer-1", + imageUrl: null, + name: "Viewer", +}; + function createSessionViewerStateDatabase(): SqliteD1Database { const database = new SqliteD1Database({ foreignKeys: false }); database.execute(` CREATE TABLE session ( + archived_at integer, id text PRIMARY KEY NOT NULL, last_run_id text, + runtime_event_seq_cursor integer NOT NULL DEFAULT 0, status text NOT NULL, title text, updated_at integer NOT NULL @@ -20,9 +31,11 @@ function createSessionViewerStateDatabase(): SqliteD1Database { created_at integer NOT NULL, deployment_version_id text, deployment_version_number integer, + driver_instance_id text, error_code text, error_details_json text, error_message text, + error_retryable integer, id text PRIMARY KEY NOT NULL, model text, provider text, @@ -34,15 +47,25 @@ function createSessionViewerStateDatabase(): SqliteD1Database { updated_at integer NOT NULL ); + CREATE TABLE session_agent_task_snapshot ( + driver_instance_id text NOT NULL, + run_id text NOT NULL, + seq integer NOT NULL, + session_id text PRIMARY KEY NOT NULL, + tasks_json text NOT NULL + ); + CREATE TABLE session_message ( content_text text NOT NULL, created_at integer NOT NULL, id text PRIMARY KEY NOT NULL, plan_json text, + projection_format text NOT NULL DEFAULT 'materialized', role text NOT NULL, segments_json text, seq integer NOT NULL, - session_id text NOT NULL + session_id text NOT NULL, + session_run_id text ); CREATE TABLE file_record ( @@ -60,6 +83,7 @@ function createSessionViewerStateDatabase(): SqliteD1Database { parent_path text NOT NULL, path text NOT NULL, purpose text NOT NULL, + runtime_event_seq integer, scope_id text NOT NULL, scope_kind text NOT NULL, session_kind text, @@ -69,6 +93,16 @@ function createSessionViewerStateDatabase(): SqliteD1Database { version integer NOT NULL ); + CREATE TABLE session_artifact_head ( + file_id text, + runtime_event_seq integer NOT NULL, + session_id text NOT NULL, + source_event_id text NOT NULL, + source_path text NOT NULL, + updated_at integer NOT NULL, + UNIQUE (session_id, source_path) + ); + CREATE TABLE session_permission_request ( created_at integer NOT NULL, driver_instance_id text NOT NULL, @@ -101,6 +135,7 @@ function createSessionViewerStateDatabase(): SqliteD1Database { INSERT INTO session_run ( completed_at, created_at, + driver_instance_id, id, model, provider, @@ -114,6 +149,7 @@ function createSessionViewerStateDatabase(): SqliteD1Database { VALUES ( NULL, 10, + 'driver-1', 'run-1', 'gpt-5.4', 'openai', @@ -199,6 +235,7 @@ describe("session viewer state", () => { expect(state.run.id).toBe("run-1"); expect(state.run.status).toBe("running"); + expect(state.infra.driverInstanceId).toBe("driver-1"); expect(state.files).toHaveLength(1); expect(state.messages).toHaveLength(1); expect(state.title).toBe("Investigate issue"); @@ -217,6 +254,271 @@ describe("session viewer state", () => { expect(state.messages).toHaveLength(1); }); + test("loads the schema-validated current task snapshot from the same canonical read", async () => { + const database = createSessionViewerStateDatabase(); + database.execute(` + INSERT INTO session_agent_task_snapshot ( + driver_instance_id, + run_id, + seq, + session_id, + tasks_json + ) + VALUES ( + 'driver-1', + 'run-1', + 7, + 'session-1', + '{"tasks":[{"taskId":"task-1","title":"Inspect"}]}' + ); + `); + + const state = await loadSessionViewerState(database, { + sessionId: "session-1", + viewerId: "viewer-1", + }); + + expect(state.taskSnapshot).toEqual({ + driverInstanceId: "driver-1", + runId: "run-1", + tasks: [{ taskId: "task-1", title: "Inspect" }], + }); + }); + + test("refreshes the independently durable task snapshot when the live-state cache is stale", async () => { + const database = createSessionViewerStateDatabase(); + database.execute(` + INSERT INTO session_agent_task_snapshot ( + driver_instance_id, + run_id, + seq, + session_id, + tasks_json + ) + VALUES ('driver-1', 'run-1', 7, 'session-1', '{"tasks":[{"taskId":"stale"}]}'); + `); + const cachedState = await loadSessionViewerState(database, { + sessionId: "session-1", + viewerId: VIEWER.id, + }); + database.execute(` + UPDATE session_agent_task_snapshot + SET seq = 8, tasks_json = '{"tasks":[{"taskId":"latest"}]}' + WHERE session_id = 'session-1'; + `); + + const state = await loadViewerLiveState({ + cachedState, + database, + reconciledStaleRun: false, + sessionId: "session-1", + viewer: VIEWER, + }); + + expect(state.taskSnapshot?.tasks).toEqual([{ taskId: "latest" }]); + expect(state.messages).toEqual(cachedState.messages); + }); + + test("reloads the canonical viewer generation when a newer run broadcast was missed", async () => { + const database = createSessionViewerStateDatabase(); + database.execute(` + INSERT INTO session_agent_task_snapshot ( + driver_instance_id, + run_id, + seq, + session_id, + tasks_json + ) + VALUES ('driver-1', 'run-1', 7, 'session-1', '{"tasks":[{"taskId":"run-1"}]}'); + `); + const cachedState = await loadSessionViewerState(database, { + sessionId: "session-1", + viewerId: VIEWER.id, + }); + database.execute(` + INSERT INTO session_run ( + completed_at, + created_at, + driver_instance_id, + id, + model, + provider, + session_id, + started_at, + status, + trace_id, + trigger, + updated_at + ) + VALUES ( + NULL, + 31, + 'driver-2', + 'run-2', + 'gpt-5.4', + 'openai', + 'session-1', + 32, + 'running', + 'trace-2', + 'user_message', + 33 + ); + UPDATE session + SET last_run_id = 'run-2', updated_at = 34 + WHERE id = 'session-1'; + UPDATE session_agent_task_snapshot + SET + driver_instance_id = 'driver-2', + run_id = 'run-2', + seq = 8, + tasks_json = '{"tasks":[{"taskId":"run-2"}]}' + WHERE session_id = 'session-1'; + `); + + const state = await loadViewerLiveState({ + cachedState, + database, + reconciledStaleRun: false, + sessionId: "session-1", + viewer: VIEWER, + }); + + expect(state.run).toMatchObject({ id: "run-2", status: "running", traceId: "trace-2" }); + expect(state.infra.driverInstanceId).toBe("driver-2"); + expect(state.taskSnapshot).toEqual({ + driverInstanceId: "driver-2", + runId: "run-2", + tasks: [{ taskId: "run-2" }], + }); + }); + + test("reloads canonical terminal state when its live broadcast was missed", async () => { + const database = createSessionViewerStateDatabase(); + database.execute(` + INSERT INTO session_agent_task_snapshot ( + driver_instance_id, + run_id, + seq, + session_id, + tasks_json + ) + VALUES ('driver-1', 'run-1', 7, 'session-1', '{"tasks":[{"taskId":"running"}]}'); + `); + const cachedState = await loadSessionViewerState(database, { + sessionId: "session-1", + viewerId: VIEWER.id, + }); + database.execute(` + UPDATE session + SET status = 'IDLE', updated_at = 40 + WHERE id = 'session-1'; + UPDATE session_run + SET completed_at = 39, status = 'completed', updated_at = 39 + WHERE id = 'run-1'; + `); + + const state = await loadViewerLiveState({ + cachedState, + database, + reconciledStaleRun: false, + sessionId: "session-1", + viewer: VIEWER, + }); + + expect(state.lifecycle).toBe("IDLE"); + expect(state.run).toMatchObject({ id: "run-1", status: "completed" }); + expect(state.infra.driverInstanceId).toBeNull(); + expect(state.taskSnapshot).toBeNull(); + }); + + test.each([ + ["canonical row is removed", "DELETE FROM session_agent_task_snapshot"], + [ + "canonical run becomes terminal", + "UPDATE session_run SET status = 'completed' WHERE id = 'run-1'", + ], + ])("clears a cached task snapshot when the %s", async (_label, boundarySql) => { + const database = createSessionViewerStateDatabase(); + database.execute(` + INSERT INTO session_agent_task_snapshot ( + driver_instance_id, + run_id, + seq, + session_id, + tasks_json + ) + VALUES ('driver-1', 'run-1', 7, 'session-1', '{"tasks":[{"taskId":"stale"}]}'); + `); + const cachedState = await loadSessionViewerState(database, { + sessionId: "session-1", + viewerId: VIEWER.id, + }); + database.execute(boundarySql); + + const state = await loadViewerLiveState({ + cachedState, + database, + reconciledStaleRun: false, + sessionId: "session-1", + viewer: VIEWER, + }); + + expect(state.taskSnapshot).toBeNull(); + }); + + test.each([ + ["terminal run", "UPDATE session_run SET status = 'completed' WHERE id = 'run-1'"], + ["rescheduling session", "UPDATE session SET status = 'RESCHEDULING' WHERE id = 'session-1'"], + ["archived running session", "UPDATE session SET archived_at = 9 WHERE id = 'session-1'"], + ["new run", "UPDATE session SET last_run_id = 'run-2' WHERE id = 'session-1'"], + [ + "new driver generation", + "UPDATE session_run SET driver_instance_id = 'driver-2' WHERE id = 'run-1'", + ], + ])("does not leak task history across a %s boundary", async (_label, boundarySql) => { + const database = createSessionViewerStateDatabase(); + database.execute(` + INSERT INTO session_agent_task_snapshot ( + driver_instance_id, + run_id, + seq, + session_id, + tasks_json + ) + VALUES ('driver-1', 'run-1', 7, 'session-1', '{"tasks":[{"taskId":"stale"}]}'); + ${boundarySql}; + `); + + const state = await loadSessionViewerState(database, { + sessionId: "session-1", + viewerId: "viewer-1", + }); + + expect(state.taskSnapshot).toBeNull(); + }); + + test("fails closed when persisted task JSON is malformed", async () => { + const database = createSessionViewerStateDatabase(); + database.execute(` + INSERT INTO session_agent_task_snapshot ( + driver_instance_id, + run_id, + seq, + session_id, + tasks_json + ) + VALUES ('driver-1', 'run-1', 7, 'session-1', '{"tasks":"invalid"}'); + `); + + const state = await loadSessionViewerState(database, { + sessionId: "session-1", + viewerId: "viewer-1", + }); + + expect(state.taskSnapshot).toBeNull(); + }); + test("loads active permissions and readiness projections", async () => { const database = createSessionViewerStateDatabase(); database.execute(` diff --git a/apps/api/tests/skill-package-snapshot.test.ts b/apps/api/tests/skill-package-snapshot.test.ts index e8e7123e..433dbe75 100644 --- a/apps/api/tests/skill-package-snapshot.test.ts +++ b/apps/api/tests/skill-package-snapshot.test.ts @@ -117,7 +117,7 @@ describe("publishSkillSnapshot", () => { const limitedDatabase = new BindLimitedD1Database(sqlite, 100); const bucket = new PublicApiMemoryFileBucket(); const bindings = createPublicHttpTestBindings(limitedDatabase, { - fileBucket: bucket as unknown as R2Bucket, + fileBucket: bucket, }) as ApiBindings; const published = await publishSkillSnapshot( diff --git a/apps/api/tests/sqlite-d1.test.ts b/apps/api/tests/sqlite-d1.test.ts new file mode 100644 index 00000000..3b481a52 --- /dev/null +++ b/apps/api/tests/sqlite-d1.test.ts @@ -0,0 +1,24 @@ +import { expect, test } from "bun:test"; + +import { SqliteD1Database } from "./helpers/sqlite-d1"; + +test("serialized SQLite clones isolate data and transactions", async () => { + const template = new SqliteD1Database(); + template.execute("CREATE TABLE item (id integer PRIMARY KEY)"); + const bytes = template.serialize(); + const first = new SqliteD1Database({ serialized: bytes }); + const second = new SqliteD1Database({ serialized: bytes }); + + await expect( + first.batch([ + first.prepare("INSERT INTO item (id) VALUES (1)"), + first.prepare("INSERT INTO item (id) VALUES (1)"), + ]), + ).rejects.toThrow(); + await first.prepare("INSERT INTO item (id) VALUES (2)").run(); + + await expect(first.prepare("SELECT id FROM item").all()).resolves.toMatchObject({ + results: [{ id: 2 }], + }); + await expect(second.prepare("SELECT id FROM item").all()).resolves.toMatchObject({ results: [] }); +}); diff --git a/apps/api/tests/terminal-run-release-recovery.test.ts b/apps/api/tests/terminal-run-release-recovery.test.ts new file mode 100644 index 00000000..e3ff1d5c --- /dev/null +++ b/apps/api/tests/terminal-run-release-recovery.test.ts @@ -0,0 +1,325 @@ +import { describe, expect, test } from "bun:test"; + +import type { DriverInstanceId, SessionRunId } from "@mosoo/id"; + +import { recordCanonicalSessionRunFailure } from "../src/modules/runtime/application/session-runs/session-run-terminal-failure.service"; +import { classifyReclaim } from "../src/modules/runtime/domain/session-run-reclaim-recovery"; +import { createSessionRunTerminalFailureSourceId } from "../src/modules/runtime/domain/session-run-terminal-event-id"; +import { recordDriverInstanceFailure } from "../src/modules/runtime/infrastructure/driver-instance/terminal-driver-events"; +import { + releaseTerminalDriverInstanceSessionRun, + repairFinalizedTerminalDriverRunState, +} from "../src/modules/runtime/infrastructure/driver-instance/terminal-run-release"; +import type { ApiBindings } from "../src/platform/cloudflare/worker-types"; +import { + createPublicHttpContractDatabase, + createPublicHttpTestBindings, + insertActiveSandboxSessionFixture, + insertOwnerSession, + PUBLIC_API_TEST_IDS, +} from "./helpers/public-api-http-test-fixture"; +import type { SqliteD1Database } from "./helpers/sqlite-d1"; + +const DRIVER_INSTANCE_ID = PUBLIC_API_TEST_IDS.driverOwner as DriverInstanceId; +const SESSION_RUN_ID = "01J0000000000000000000000T" as SessionRunId; +const SUCCESSOR_SESSION_RUN_ID = "01J0000000000000000000000V" as SessionRunId; + +async function createTerminalReleaseFixture(): Promise<{ + bindings: ApiBindings; + database: SqliteD1Database; +}> { + const database = await createPublicHttpContractDatabase(); + await insertOwnerSession(database); + await insertActiveSandboxSessionFixture(database, { + ownerAccountId: PUBLIC_API_TEST_IDS.ownerAccount, + sandboxId: PUBLIC_API_TEST_IDS.sandbox, + sandboxSessionId: "01J0000000000000000000000W", + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + timestampMs: 1, + }); + database.execute(` + INSERT INTO driver_instance ( + id, boot_token_expires_at, boot_token_hash, connection_id, created_at, + expires_at, heartbeat_count, protocol, protocol_version, runtime, + sandbox_id, sandbox_incarnation, sandbox_session_id, status, updated_at + ) VALUES ( + '${DRIVER_INSTANCE_ID}', 1, X'01', 'terminal-recovery-connection', 1, + 1, 0, 'orpc-ws', 1, 'openai-runtime', '${PUBLIC_API_TEST_IDS.sandbox}', + 1, '${PUBLIC_API_TEST_IDS.ownerSession}', 'stopped', 1 + ); + + INSERT INTO session_run ( + id, session_id, agent_id, created_by_account_id, deployment_version_id, + deployment_version_number, driver_instance_id, trigger, status, provider, + model, runtime_id, trace_id, started_at, created_at, updated_at + ) VALUES ( + '${SESSION_RUN_ID}', '${PUBLIC_API_TEST_IDS.ownerSession}', '${PUBLIC_API_TEST_IDS.agent}', + '${PUBLIC_API_TEST_IDS.ownerAccount}', '${PUBLIC_API_TEST_IDS.deployment}', 1, + '${DRIVER_INSTANCE_ID}', 'user_prompt', 'running', 'openai', 'gpt-5.4', + 'openai-runtime', 'trace-terminal-recovery', 1, 1, 1 + ); + + UPDATE session + SET last_run_id = '${SESSION_RUN_ID}', status = 'RUNNING' + WHERE id = '${PUBLIC_API_TEST_IDS.ownerSession}'; + `); + + return { + bindings: createPublicHttpTestBindings(database) as ApiBindings, + database, + }; +} + +async function recordReclaimFailure(bindings: ApiBindings): Promise { + const outcome = await recordCanonicalSessionRunFailure(bindings, { + error: classifyReclaim({ + driverInstanceId: DRIVER_INSTANCE_ID, + driverTerminalStatus: "stopped", + reclaimReason: "socket_closed", + }), + runId: SESSION_RUN_ID, + sessionId: PUBLIC_API_TEST_IDS.ownerSession, + source: "driver", + }); + + expect(outcome.kind).toBe("failed"); +} + +async function readReleaseState(database: D1Database): Promise<{ + inactiveDeadlineAt: number | null; + sandboxUpdatedAt: number; +}> { + const row = await database + .prepare( + "SELECT inactive_deadline_at AS inactiveDeadlineAt, updated_at AS sandboxUpdatedAt FROM sandbox WHERE id = ?", + ) + .bind(PUBLIC_API_TEST_IDS.sandbox) + .first<{ + inactiveDeadlineAt: number | null; + sandboxUpdatedAt: number; + }>(); + + if (row === null) { + throw new Error("Sandbox release state was not found."); + } + + return row; +} + +describe("terminal Run release recovery", () => { + test("releases a Driver failure reported before the connection becomes ready", async () => { + const { bindings, database } = await createTerminalReleaseFixture(); + database.execute(` + UPDATE driver_instance + SET status = 'connecting', status_operation_id = NULL + WHERE id = '${DRIVER_INSTANCE_ID}'; + `); + + await recordDriverInstanceFailure(bindings, { + driverConnectionId: "terminal-recovery-connection", + driverGeneration: 0, + driverInstanceId: DRIVER_INSTANCE_ID, + sessionRunId: SESSION_RUN_ID, + error: { + code: "driver.startup_failed", + details: {}, + message: "Driver failed before ready.", + retryable: true, + }, + }); + await expect( + database + .prepare("SELECT status, status_operation_id FROM driver_instance WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ status: "stopping", status_operation_id: SESSION_RUN_ID }); + await expect( + database + .prepare("SELECT error_code, status FROM session_run WHERE id = ?") + .bind(SESSION_RUN_ID) + .first(), + ).resolves.toEqual({ error_code: "driver.startup_failed", status: "failed" }); + }); + + test("does not release a replacement connection from an old terminal RPC", async () => { + const { bindings, database } = await createTerminalReleaseFixture(); + await recordReclaimFailure(bindings); + const beforeRelease = await readReleaseState(database); + const commandId = "01J0000000000000000000000X"; + + await database + .prepare( + `INSERT INTO driver_command ( + driver_generation, driver_instance_id, id, issued_at, kind, + payload_json, seq, status + ) VALUES (?, ?, ?, ?, 'input.start', ?, ?, 'accepted')`, + ) + .bind( + 0, + DRIVER_INSTANCE_ID, + commandId, + 1, + JSON.stringify({ + commandId, + input: { text: "continue" }, + kind: "input.start", + requestId: "replacement-command", + runId: SESSION_RUN_ID, + }), + 1, + ) + .run(); + database.execute(` + UPDATE driver_instance + SET connection_id = 'replacement-connection', status_operation_id = NULL + WHERE id = '${DRIVER_INSTANCE_ID}'; + `); + + await expect( + releaseTerminalDriverInstanceSessionRun(bindings, { + expectedDriverConnectionId: "terminal-recovery-connection", + driverGeneration: 0, + driverInstanceId: DRIVER_INSTANCE_ID, + sessionRunId: SESSION_RUN_ID, + }), + ).rejects.toThrow("lost its exact Driver ownership"); + await expect( + database + .prepare("SELECT connection_id, status_operation_id FROM driver_instance WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ + connection_id: "replacement-connection", + status_operation_id: null, + }); + await expect( + database + .prepare("SELECT error_json, result_json, status FROM driver_command WHERE id = ?") + .bind(commandId) + .first(), + ).resolves.toEqual({ error_json: null, result_json: null, status: "accepted" }); + expect(await readReleaseState(database)).toEqual(beforeRelease); + }); + + test("does not let an old terminal Run claim a Driver owned by an active successor", async () => { + const { bindings, database } = await createTerminalReleaseFixture(); + await recordReclaimFailure(bindings); + database.execute(` + INSERT INTO session_run ( + id, session_id, agent_id, created_by_account_id, deployment_version_id, + deployment_version_number, driver_instance_id, trigger, status, provider, + model, runtime_id, trace_id, started_at, created_at, updated_at + ) VALUES ( + '${SUCCESSOR_SESSION_RUN_ID}', '${PUBLIC_API_TEST_IDS.ownerSession}', + '${PUBLIC_API_TEST_IDS.agent}', '${PUBLIC_API_TEST_IDS.ownerAccount}', + '${PUBLIC_API_TEST_IDS.deployment}', 1, '${DRIVER_INSTANCE_ID}', + 'user_prompt', 'running', 'openai', 'gpt-5.4', 'openai-runtime', + 'trace-terminal-successor', 2, 2, 2 + ); + + UPDATE session + SET last_run_id = '${SUCCESSOR_SESSION_RUN_ID}', status = 'RUNNING' + WHERE id = '${PUBLIC_API_TEST_IDS.ownerSession}'; + + UPDATE driver_instance + SET status = 'ready', status_operation_id = NULL + WHERE id = '${DRIVER_INSTANCE_ID}'; + `); + + await expect( + releaseTerminalDriverInstanceSessionRun(bindings, { + driverGeneration: 0, + driverInstanceId: DRIVER_INSTANCE_ID, + sessionRunId: SESSION_RUN_ID, + }), + ).rejects.toThrow("lost its exact Driver ownership"); + await expect( + database + .prepare("SELECT status, status_operation_id FROM driver_instance WHERE id = ?") + .bind(DRIVER_INSTANCE_ID) + .first(), + ).resolves.toEqual({ status: "ready", status_operation_id: null }); + await expect( + database + .prepare("SELECT driver_instance_id, status FROM session_run WHERE id = ?") + .bind(SUCCESSOR_SESSION_RUN_ID) + .first(), + ).resolves.toEqual({ driver_instance_id: DRIVER_INSTANCE_ID, status: "running" }); + }); + + test("repairs the lease from a complete canonical terminal projection", async () => { + const { bindings, database } = await createTerminalReleaseFixture(); + await recordReclaimFailure(bindings); + const sourceEventId = createSessionRunTerminalFailureSourceId(SESSION_RUN_ID); + + await expect( + database + .prepare("SELECT COUNT(*) AS count FROM session_event WHERE source_event_id = ?") + .bind(sourceEventId) + .first(), + ).resolves.toEqual({ count: 1 }); + + await expect( + repairFinalizedTerminalDriverRunState(bindings, { + driverGeneration: 0, + driverInstanceId: DRIVER_INSTANCE_ID, + sessionRunId: SESSION_RUN_ID, + status: "stopped", + }), + ).resolves.toMatchObject({ released: true }); + const firstRelease = await readReleaseState(database); + + await expect( + repairFinalizedTerminalDriverRunState(bindings, { + driverGeneration: 0, + driverInstanceId: DRIVER_INSTANCE_ID, + sessionRunId: SESSION_RUN_ID, + status: "stopped", + }), + ).resolves.toMatchObject({ released: true }); + await expect( + database + .prepare("SELECT COUNT(*) AS count FROM session_event WHERE source_event_id = ?") + .bind(sourceEventId) + .first(), + ).resolves.toEqual({ count: 1 }); + expect(await readReleaseState(database)).toEqual(firstRelease); + }); + + test("converges after the atomic lease release commits but its acknowledgement is lost", async () => { + const { bindings, database } = await createTerminalReleaseFixture(); + await recordReclaimFailure(bindings); + const originalBatch = database.batch; + let loseAcknowledgement = true; + database.batch = (async (statements: D1PreparedStatement[]) => { + const results = await originalBatch.call(database, statements); + + if (loseAcknowledgement) { + loseAcknowledgement = false; + throw new Error("lease release acknowledgement lost"); + } + + return results as D1Result[]; + }) as D1Database["batch"]; + + await expect( + releaseTerminalDriverInstanceSessionRun(bindings, { + driverGeneration: 0, + driverInstanceId: DRIVER_INSTANCE_ID, + sessionRunId: SESSION_RUN_ID, + }), + ).rejects.toThrow("lease release acknowledgement lost"); + const committedRelease = await readReleaseState(database); + expect(committedRelease.inactiveDeadlineAt).not.toBeNull(); + + database.batch = originalBatch; + await expect( + releaseTerminalDriverInstanceSessionRun(bindings, { + driverGeneration: 0, + driverInstanceId: DRIVER_INSTANCE_ID, + sessionRunId: SESSION_RUN_ID, + }), + ).resolves.toMatchObject({ released: true }); + expect(await readReleaseState(database)).toEqual(committedRelease); + }); +}); diff --git a/apps/api/tests/wrangler-d1-migration-atomicity.test.ts b/apps/api/tests/wrangler-d1-migration-atomicity.test.ts new file mode 100644 index 00000000..85a6dc2a --- /dev/null +++ b/apps/api/tests/wrangler-d1-migration-atomicity.test.ts @@ -0,0 +1,510 @@ +import { expect, test } from "bun:test"; +import { spawnSync } from "node:child_process"; +import { copyFileSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { fileURLToPath } from "node:url"; + +import { parseD1JsonResults } from "../bin/d1-json"; +import { + createProdSchemaCatalogFromIntrospectionRows, + DRIZZLE_MIGRATION_PROD_SCHEMA_CATALOG, + createProdSchemaIntrospectionStatements, + findProdSchemaDifferences, +} from "../bin/prod-schema-guard"; +import { + installProtocolV3CutoverSql, + installProtocolV3PostMigrationCutoverSql, + PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + PROTOCOL_V3_CUTOVER_OBJECTS_SQL, + PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, +} from "../bin/protocol-v3-cutover"; +import { drizzleMigrations, getDrizzleMigration } from "./helpers/drizzle-migrations"; + +const MIGRATION_TAG = "0019_runtime-subject-operation-authority"; +const SCHEMA_MIGRATION_TAG = "0020_sandbox-backup-object-authority"; +const RELEASE_TREE_OID = "0123456789abcdef0123456789abcdef01234567"; +const migrationsSource = fileURLToPath(new URL("../../../pkgs/db/drizzle/", import.meta.url)); +const wrangler = fileURLToPath(new URL("../node_modules/.bin/wrangler", import.meta.url)); + +interface WranglerResult { + readonly output: string; + readonly status: number | null; +} + +function runWrangler(root: string, args: readonly string[]): WranglerResult { + const result = spawnSync(wrangler, args, { + cwd: root, + encoding: "utf8", + env: { + ...process.env, + CI: "1", + NO_UPDATE_NOTIFIER: "1", + TMPDIR: join(root, "tmp"), + WRANGLER_SEND_METRICS: "false", + WRANGLER_LOG_PATH: join(root, "logs"), + WRANGLER_WRITE_LOGS: "false", + XDG_CACHE_HOME: join(root, "cache"), + XDG_CONFIG_HOME: join(root, "config"), + }, + timeout: 120_000, + }); + if (result.error) throw result.error; + return { + output: `${result.stdout ?? ""}\n${result.stderr ?? ""}`, + status: result.status, + }; +} + +function runWranglerSuccessfully(root: string, args: readonly string[]): string { + const result = runWrangler(root, args); + if (result.status !== 0) { + throw new Error(`wrangler ${args.join(" ")} failed:\n${result.output}`); + } + return result.output; +} + +function localD1Args(state: string, action: "execute" | "migrations", tail: string[]): string[] { + return [ + "d1", + action, + ...(action === "migrations" ? ["apply"] : []), + "DB", + "--local", + "--persist-to", + state, + ...tail, + ]; +} + +interface MigrationStateExpectation { + readonly authorityColumnCount: number; + readonly backupAuthorityWithoutRowidCount: number; + readonly gateObjectCount: number; + readonly lastMigration: string; + readonly stagingClaimOwnerCount: number; + readonly stagingIndexCount: number; + readonly stagingTableCount: number; + readonly targetApplied: number; +} + +interface MigrationAtomicityFixture { + readonly failurePoint: string; + readonly gateSql: string; + readonly initialGateObjectCount: number; + readonly recovered: MigrationStateExpectation; + readonly rolledBack: MigrationStateExpectation; + readonly verifyLatestCatalog: boolean; +} + +const MIGRATION_ATOMICITY_CASES = [ + [ + MIGRATION_TAG, + { + failurePoint: "ALTER TABLE `driver_instance` ADD `sandbox_incarnation`", + gateSql: installProtocolV3CutoverSql(RELEASE_TREE_OID), + initialGateObjectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + recovered: { + authorityColumnCount: 1, + backupAuthorityWithoutRowidCount: 0, + gateObjectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + lastMigration: `${MIGRATION_TAG}.sql`, + stagingClaimOwnerCount: 0, + stagingIndexCount: 4, + stagingTableCount: 2, + targetApplied: 1, + }, + rolledBack: { + authorityColumnCount: 0, + backupAuthorityWithoutRowidCount: 0, + gateObjectCount: PROTOCOL_V3_CUTOVER_OBJECT_COUNT, + lastMigration: "0018_runtime-operation-ready-authority.sql", + stagingClaimOwnerCount: 0, + stagingIndexCount: 0, + stagingTableCount: 0, + targetApplied: 0, + }, + verifyLatestCatalog: false, + }, + ], + [ + SCHEMA_MIGRATION_TAG, + { + failurePoint: "CREATE TRIGGER `sandbox_backup_delete_intent_authority`", + gateSql: installProtocolV3PostMigrationCutoverSql(RELEASE_TREE_OID), + initialGateObjectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + recovered: { + authorityColumnCount: 1, + backupAuthorityWithoutRowidCount: 5, + gateObjectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + lastMigration: `${SCHEMA_MIGRATION_TAG}.sql`, + stagingClaimOwnerCount: 1, + stagingIndexCount: 4, + stagingTableCount: 2, + targetApplied: 1, + }, + rolledBack: { + authorityColumnCount: 1, + backupAuthorityWithoutRowidCount: 0, + gateObjectCount: PROTOCOL_V3_POST_MIGRATION_CUTOVER_OBJECT_COUNT, + lastMigration: `${MIGRATION_TAG}.sql`, + stagingClaimOwnerCount: 0, + stagingIndexCount: 4, + stagingTableCount: 2, + targetApplied: 0, + }, + verifyLatestCatalog: true, + }, + ], +] as const satisfies readonly (readonly [string, MigrationAtomicityFixture])[]; + +const SCHEMA_MIGRATION_STAGING_SEED_SQL = ` +INSERT INTO sandbox ( + agent_id, app_id, claim_expires_at, claim_owner, created_at, id, incarnation, + kind, network_constraints_hash, operation_kind, owner_account_id, status, + status_operation_id, subject_id, subject_kind, updated_at +) VALUES ( + '01J0000000000000000000000B', '01J0000000000000000000000C', 2000000000000, + 'live-owner', 1, '01J0000000000000000000000A', 2, 'pet', '${"0".repeat(64)}', + 'hibernate', '01J0000000000000000000000D', 'backing_up', + '01J0000000000000000000000E', '01J0000000000000000000000B', 'agent', 2 +); +INSERT INTO sandbox_backup_staging ( + actual_backup_id, created_at, dir, driver_generation, driver_instance_id, id, + operation_id, sandbox_id, sandbox_incarnation, session_run_id, ttl_seconds, + updated_at, updates_subject_backup, workspace_session_id +) VALUES + ('01J0000000000000000000000M', 10, '/workspace/live', NULL, NULL, + '01J0000000000000000000000G', '01J0000000000000000000000E', + '01J0000000000000000000000A', 2, NULL, 60, 11, 1, NULL), + ('01J0000000000000000000000N', 20, '/workspace/stale', NULL, NULL, + '01J0000000000000000000000H', '01J0000000000000000000000F', + '01J0000000000000000000000A', 2, NULL, 90, 21, 0, NULL), + ('01J0000000000000000000000P', 30, '/workspace/terminal', 3, + '01J0000000000000000000000K', '01J0000000000000000000000J', NULL, + '01J0000000000000000000000A', 2, '01J0000000000000000000000Q', + 120, 31, 0, '01J0000000000000000000000R'); +INSERT INTO sandbox_backup ( + created_at, dir, id, keep, operation_id, sandbox_id, sandbox_incarnation, + session_run_id, staging_id, status, ttl_seconds, updated_at, workspace_session_id +) VALUES ( + 40, '/workspace/ready', '01J0000000000000000000000V', 1, + '01J0000000000000000000000E', '01J0000000000000000000000A', 2, NULL, + '01J0000000000000000000000W', 'ready', 180, 41, NULL +); +INSERT INTO api_command ( + attempt_count, claim_expires_at, claim_owner, created_at, dedupe_key, + delivery_generation, id, kind, payload_json, status, updated_at +) VALUES ( + 4, 2000000000000, 'environment-owner', 42, 'environment-artifact-migration-seed', + 3, '01J0000000000000000000000X', 'environment_package_artifact_build', + '{"appId":"01J0000000000000000000000Y","inputDigest":"${"c".repeat(64)}"}', + 'running', 43 +); +INSERT INTO environment_package_artifact_backup_staging ( + actual_backup_id, app_id, attempt_count, claim_owner, command_id, created_at, + delivery_generation, dir, input_digest, paths_json, updated_at +) VALUES ( + '01J0000000000000000000000Z', '01J0000000000000000000000Y', 4, + 'environment-owner', '01J0000000000000000000000X', 42, 3, + '/workspace/.mosoo/environment-artifacts/${"c".repeat(64)}', '${"c".repeat(64)}', + '{"executable":["/workspace/.mosoo/environment-artifacts/${"c".repeat(64)}/bin"],"node":[],"python":[]}', + 43 +); +`; + +function executeLocalD1Json(root: string, state: string, sql: string) { + return parseD1JsonResults( + runWranglerSuccessfully( + root, + localD1Args(state, "execute", ["--command", sql, "--json", "--yes"]), + ), + ); +} + +function expectExactCutoverGate(root: string, state: string, expectedObjectCount: number): void { + expect(executeLocalD1Json(root, state, PROTOCOL_V3_CUTOVER_OBJECTS_SQL)).toEqual([ + [{ exact_object_count: expectedObjectCount, object_count: expectedObjectCount }], + ]); +} + +function expectMigrationState( + root: string, + state: string, + migrationTag: string, + expected: MigrationStateExpectation, +): void { + expect( + executeLocalD1Json( + root, + state, + `SELECT value FROM rollback_probe WHERE id = 1; + SELECT count(*) AS count FROM pragma_table_info('sandbox') WHERE name = 'incarnation'; + SELECT count(*) AS count FROM pragma_table_info('api_command') + WHERE name = 'delivery_generation'; + SELECT count(*) AS count FROM pragma_table_info('app_deployment') + WHERE name = 'active_script_name'; + SELECT count(*) AS count FROM sqlite_master + WHERE type = 'table' + AND name IN ('sandbox_backup_staging', 'environment_package_artifact_backup_staging'); + SELECT count(*) AS count FROM pragma_table_info('sandbox_backup_staging') + WHERE name = 'claim_owner'; + SELECT count(*) AS count FROM sqlite_master + WHERE type = 'index' + AND tbl_name = 'sandbox_backup_staging' + AND name IN ('sandbox_backup_staging_updated_idx', + 'sandbox_backup_staging_actual_idx', + 'sandbox_backup_staging_terminal_checkpoint_idx', + 'sandbox_backup_staging_operation_checkpoint_idx'); + SELECT count(*) AS count FROM pragma_table_list + WHERE name IN ('environment_package_artifact_backup', + 'environment_package_artifact_backup_staging', + 'sandbox_backup', + 'sandbox_backup_delete_intent', + 'sandbox_backup_staging') + AND wr = 1; + ${PROTOCOL_V3_CUTOVER_OBJECTS_SQL} + SELECT enabled, release_tree_oid FROM __protocol_v3_cutover WHERE id = 1; + SELECT count(*) AS count FROM d1_migrations WHERE name = '${migrationTag}.sql'; + SELECT name FROM d1_migrations ORDER BY id DESC LIMIT 1;`, + ), + ).toEqual([ + [{ value: 1 }], + [{ count: expected.authorityColumnCount }], + [{ count: expected.authorityColumnCount }], + [{ count: expected.authorityColumnCount }], + [{ count: expected.stagingTableCount }], + [{ count: expected.stagingClaimOwnerCount }], + [{ count: expected.stagingIndexCount }], + [{ count: expected.backupAuthorityWithoutRowidCount }], + [ + { + exact_object_count: expected.gateObjectCount, + object_count: expected.gateObjectCount, + }, + ], + [{ enabled: 1, release_tree_oid: RELEASE_TREE_OID }], + [{ count: expected.targetApplied }], + [{ name: expected.lastMigration }], + ]); +} + +function expectSchemaMigrationAuthorityRows(root: string, state: string, migrated: boolean): void { + const claimOwner = migrated ? "claim_owner" : "NULL AS claim_owner"; + expect( + executeLocalD1Json( + root, + state, + `SELECT actual_backup_id, ${claimOwner}, created_at, dir, driver_generation, + driver_instance_id, id, operation_id, sandbox_id, sandbox_incarnation, + session_run_id, ttl_seconds, updated_at, updates_subject_backup, + workspace_session_id + FROM sandbox_backup_staging ORDER BY id; + SELECT created_at, dir, id, keep, operation_id, sandbox_id, sandbox_incarnation, + session_run_id, staging_id, status, ttl_seconds, updated_at, workspace_session_id + FROM sandbox_backup ORDER BY id; + SELECT actual_backup_id, app_id, attempt_count, claim_owner, command_id, created_at, + delivery_generation, dir, input_digest, paths_json, updated_at + FROM environment_package_artifact_backup_staging ORDER BY command_id; + PRAGMA foreign_key_check`, + ), + ).toEqual([ + [ + { + actual_backup_id: "01J0000000000000000000000M", + claim_owner: migrated ? "live-owner" : null, + created_at: 10, + dir: "/workspace/live", + driver_generation: null, + driver_instance_id: null, + id: "01J0000000000000000000000G", + operation_id: "01J0000000000000000000000E", + sandbox_id: "01J0000000000000000000000A", + sandbox_incarnation: 2, + session_run_id: null, + ttl_seconds: 60, + updated_at: 11, + updates_subject_backup: 1, + workspace_session_id: null, + }, + { + actual_backup_id: "01J0000000000000000000000N", + claim_owner: migrated ? "__legacy_stale__" : null, + created_at: 20, + dir: "/workspace/stale", + driver_generation: null, + driver_instance_id: null, + id: "01J0000000000000000000000H", + operation_id: "01J0000000000000000000000F", + sandbox_id: "01J0000000000000000000000A", + sandbox_incarnation: 2, + session_run_id: null, + ttl_seconds: 90, + updated_at: 21, + updates_subject_backup: 0, + workspace_session_id: null, + }, + { + actual_backup_id: "01J0000000000000000000000P", + claim_owner: null, + created_at: 30, + dir: "/workspace/terminal", + driver_generation: 3, + driver_instance_id: "01J0000000000000000000000K", + id: "01J0000000000000000000000J", + operation_id: null, + sandbox_id: "01J0000000000000000000000A", + sandbox_incarnation: 2, + session_run_id: "01J0000000000000000000000Q", + ttl_seconds: 120, + updated_at: 31, + updates_subject_backup: 0, + workspace_session_id: "01J0000000000000000000000R", + }, + ], + [ + { + created_at: 40, + dir: "/workspace/ready", + id: "01J0000000000000000000000V", + keep: 1, + operation_id: "01J0000000000000000000000E", + sandbox_id: "01J0000000000000000000000A", + sandbox_incarnation: 2, + session_run_id: null, + staging_id: "01J0000000000000000000000W", + status: "ready", + ttl_seconds: 180, + updated_at: 41, + workspace_session_id: null, + }, + ], + [ + { + actual_backup_id: "01J0000000000000000000000Z", + app_id: "01J0000000000000000000000Y", + attempt_count: 4, + claim_owner: "environment-owner", + command_id: "01J0000000000000000000000X", + created_at: 42, + delivery_generation: 3, + dir: `/workspace/.mosoo/environment-artifacts/${"c".repeat(64)}`, + input_digest: "c".repeat(64), + paths_json: JSON.stringify({ + executable: [`/workspace/.mosoo/environment-artifacts/${"c".repeat(64)}/bin`], + node: [], + python: [], + }), + updated_at: 43, + }, + ], + [], + ]); +} + +function expectLatestSchemaCatalog(root: string, state: string): void { + runWranglerSuccessfully( + root, + localD1Args(state, "execute", [ + "--command", + "DROP TABLE rollback_probe; DROP TABLE __protocol_v3_migration_intent", + "--yes", + ]), + ); + const tableNames = DRIZZLE_MIGRATION_PROD_SCHEMA_CATALOG.tables.map(({ name }) => name); + const catalogOutput = runWranglerSuccessfully( + root, + localD1Args(state, "execute", [ + "--command", + createProdSchemaIntrospectionStatements(tableNames).join(";\n"), + "--json", + "--yes", + ]), + ); + expect( + findProdSchemaDifferences( + DRIZZLE_MIGRATION_PROD_SCHEMA_CATALOG, + createProdSchemaCatalogFromIntrospectionRows(parseD1JsonResults(catalogOutput), tableNames), + ), + ).toEqual([]); +} + +test.each(MIGRATION_ATOMICITY_CASES)( + "Wrangler rolls back and can reapply %s after a middle failure", + (migrationTag, fixture) => { + const root = mkdtempSync(join(tmpdir(), "mosoo-wrangler-migration-")); + try { + const migrations = join(root, "migrations"); + const state = join(root, "state"); + for (const directory of ["cache", "config", "migrations", "tmp"]) { + mkdirSync(join(root, directory)); + } + writeFileSync(join(root, "worker.ts"), "export default {};\n"); + writeFileSync( + join(root, "wrangler.toml"), + `name = "migration-atomicity" +main = "worker.ts" +compatibility_date = "2026-08-30" + +[[d1_databases]] +binding = "DB" +database_name = "migration-atomicity" +database_id = "00000000-0000-0000-0000-000000000000" +migrations_dir = "migrations" +`, + ); + + const target = getDrizzleMigration(migrationTag); + for (const migration of drizzleMigrations.slice(0, target.index)) { + copyFileSync( + join(migrationsSource, `${migration.tag}.sql`), + join(migrations, `${migration.tag}.sql`), + ); + } + runWranglerSuccessfully(root, localD1Args(state, "migrations", [])); + + const setup = join(root, "setup.sql"); + writeFileSync( + setup, + `${migrationTag === SCHEMA_MIGRATION_TAG ? SCHEMA_MIGRATION_STAGING_SEED_SQL : ""} +${fixture.gateSql} +CREATE TABLE rollback_probe (id integer PRIMARY KEY, value integer NOT NULL); +INSERT INTO rollback_probe (id, value) VALUES (1, 1); +`, + ); + runWranglerSuccessfully(root, localD1Args(state, "execute", ["--file", setup, "--yes"])); + expectExactCutoverGate(root, state, fixture.initialGateObjectCount); + + const source = readFileSync(join(migrationsSource, `${migrationTag}.sql`), "utf8"); + expect(source.split(fixture.failurePoint)).toHaveLength(2); + writeFileSync( + join(migrations, `${migrationTag}.sql`), + source.replace( + fixture.failurePoint, + `UPDATE rollback_probe SET value = 2 WHERE id = 1;--> statement-breakpoint +SELECT * FROM __intentional_mid_migration_failure;--> statement-breakpoint +${fixture.failurePoint}`, + ), + ); + + const failed = runWrangler(root, localD1Args(state, "migrations", [])); + expect(failed.status).not.toBe(0); + expect(failed.output).toContain("__intentional_mid_migration_failure"); + expectMigrationState(root, state, migrationTag, fixture.rolledBack); + if (migrationTag === SCHEMA_MIGRATION_TAG) { + expectSchemaMigrationAuthorityRows(root, state, false); + } + + writeFileSync(join(migrations, `${migrationTag}.sql`), source); + runWranglerSuccessfully(root, localD1Args(state, "migrations", [])); + expectMigrationState(root, state, migrationTag, fixture.recovered); + if (migrationTag === SCHEMA_MIGRATION_TAG) { + expectSchemaMigrationAuthorityRows(root, state, true); + } + if (fixture.verifyLatestCatalog) expectLatestSchemaCatalog(root, state); + } finally { + rmSync(root, { force: true, recursive: true }); + } + }, + 120_000, +); diff --git a/apps/api/wrangler.toml b/apps/api/wrangler.toml index b7ad6d57..f09d929c 100644 --- a/apps/api/wrangler.toml +++ b/apps/api/wrangler.toml @@ -31,6 +31,11 @@ POSTHOG_API_HOST = "https://us.i.posthog.com" [[send_email]] name = "AUTH_EMAIL" +[[workflows]] +binding = "APP_DEPLOYMENT_WORKFLOW" +name = "mosoo-app-deployment-dev" +class_name = "AppDeploymentWorkflow" + [secrets] required = [ "BETTER_AUTH_SECRET", @@ -137,6 +142,7 @@ BACKUP_BUCKET_NAME = "mosoo-stage-sandbox-state" SANDBOX_STATE_BUCKET_NAME = "mosoo-stage-sandbox-state" SANDBOX_FILE_BUCKET_LOCAL = "false" MOSOO_APP_DEPLOYMENT_DOMAIN = "apps-stage.mosoo.ai" +MOSOO_APP_DISPATCH_NAMESPACE = "mosoo-app-deployments-stage" MOSOO_ACCOUNT_CONCURRENT_SANDBOX_LIMIT = "5" MOSOO_DEPLOYMENT_MODE = "cloud" MOSOO_ENVIRONMENT = "production" @@ -159,6 +165,19 @@ required = [ "GOOGLE_OAUTH_CLIENT_SECRET", ] +[[env.stage.routes]] +zone_name = "mosoo.ai" +pattern = "*.apps-stage.mosoo.ai/*" + +[[env.stage.dispatch_namespaces]] +binding = "APP_DEPLOYMENT_DISPATCHER" +namespace = "mosoo-app-deployments-stage" + +[[env.stage.workflows]] +binding = "APP_DEPLOYMENT_WORKFLOW" +name = "mosoo-app-deployment-stage" +class_name = "AppDeploymentWorkflow" + [[env.stage.durable_objects.bindings]] name = "DriverConnection" class_name = "DriverConnection" @@ -254,6 +273,7 @@ BACKUP_BUCKET_NAME = "mosoo-sandbox-state" SANDBOX_STATE_BUCKET_NAME = "mosoo-sandbox-state" SANDBOX_FILE_BUCKET_LOCAL = "false" MOSOO_APP_DEPLOYMENT_DOMAIN = "apps.mosoo.ai" +MOSOO_APP_DISPATCH_NAMESPACE = "mosoo-app-deployments-prod" MOSOO_ACCOUNT_CONCURRENT_SANDBOX_LIMIT = "5" MOSOO_DEPLOYMENT_MODE = "cloud" MOSOO_ENVIRONMENT = "production" @@ -284,6 +304,19 @@ pattern = "cloud.mosoo.ai/api/*" zone_name = "mosoo.ai" pattern = "try.mosoo.ai/api/*" +[[env.prod.routes]] +zone_name = "mosoo.ai" +pattern = "*.apps.mosoo.ai/*" + +[[env.prod.dispatch_namespaces]] +binding = "APP_DEPLOYMENT_DISPATCHER" +namespace = "mosoo-app-deployments-prod" + +[[env.prod.workflows]] +binding = "APP_DEPLOYMENT_WORKFLOW" +name = "mosoo-app-deployment-prod" +class_name = "AppDeploymentWorkflow" + [[env.prod.durable_objects.bindings]] name = "DriverConnection" class_name = "DriverConnection" diff --git a/apps/driver b/apps/driver index 50fde116..5e661259 160000 --- a/apps/driver +++ b/apps/driver @@ -1 +1 @@ -Subproject commit 50fde116e301efeeee83418f43fc9ee6759ac82d +Subproject commit 5e661259e853663e561fed167aa495e792450045 diff --git a/apps/web/package.json b/apps/web/package.json index e0bf3cce..8ee407ae 100644 --- a/apps/web/package.json +++ b/apps/web/package.json @@ -6,15 +6,15 @@ "dev": "vp dev", "build": "vp build", "preview": "vp preview", - "deploy": "vp run build && vp exec wrangler deploy --env prod", + "deploy": "vp run build && vp exec wrangler deploy --env prod --experimental-provision=false", "lint": "vp lint .", "tc": "vp exec tsc --noEmit", - "test": "vp exec bun test tests" + "test": "vp exec bun test tests --isolate" }, "dependencies": { "@assistant-ui/react": "0.15.14", "@base-ui/react": "^1.7.0", - "@cloudflare/sandbox": "0.12.6", + "@cloudflare/sandbox": "0.12.9", "@hugeicons/core-free-icons": "^4.2.3", "@hugeicons/react": "^1.1.9", "@lobehub/icons-static-svg": "^1.90.0", @@ -30,11 +30,9 @@ "@xterm/addon-fit": "^0.11.0", "@xterm/xterm": "^6.0.0", "better-auth": "^1.6.28", - "boring-avatars": "^2.0.4", "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", "lucide-react": "^1.31.0", - "motion": "^12.40.0", "react": "^19.2.8", "react-dom": "^19.2.8", "react-markdown": "^10.1.0", @@ -55,7 +53,7 @@ "jsdom": "^29.1.1", "typescript": "^6.0.3", "vite": "^8.2.1", - "vite-plus": "^0.1.24", - "wrangler": "^4.123.0" + "vite-plus": "^0.3.0", + "wrangler": "^4.127.1" } } diff --git a/apps/web/src/domains/runtime/session-stream/session-stream-actions.ts b/apps/web/src/domains/runtime/session-stream/session-stream-actions.ts index 64c48302..36ad90ca 100644 --- a/apps/web/src/domains/runtime/session-stream/session-stream-actions.ts +++ b/apps/web/src/domains/runtime/session-stream/session-stream-actions.ts @@ -4,6 +4,7 @@ import type { SessionPermissionRequestView, SessionRunView, } from "@mosoo/ag-ui-session"; +import type { AgentTask } from "@mosoo/contracts/session"; import { useCallback, useMemo } from "react"; import type { MutableRefObject } from "react"; @@ -37,6 +38,7 @@ interface UseSessionStreamActionsInput { } export function useSessionStreamActions(input: UseSessionStreamActionsInput): { + activeTasks: AgentTask[]; messages: SessionLiveState["messages"]; lifecycle: SessionLiveState["lifecycle"]; permissionRequests: SessionPermissionRequestView[]; @@ -61,22 +63,23 @@ export function useSessionStreamActions(input: UseSessionStreamActionsInput): { streaming: boolean; syncSession: () => Promise; } { + const { activeSessionIdRef, sendViewerEvent } = input; const syncSession = useCallback(async (): Promise => { - const activeSessionId = input.activeSessionIdRef.current; + const activeSessionId = activeSessionIdRef.current; if (!isTruthy(activeSessionId)) { return; } const options: SendViewerEventOptions = { maxAttempts: 2 }; - await input.sendViewerEvent( + await sendViewerEvent( activeSessionId, createViewerCustomEvent("mosoo.session.sync.request", { reason: "manual", }), options, ); - }, [input.activeSessionIdRef, input.sendViewerEvent]); + }, [activeSessionIdRef, sendViewerEvent]); const sendUserMessage = useCallback( async (message: { attachmentIds?: string[]; @@ -150,6 +153,7 @@ export function useSessionStreamActions(input: UseSessionStreamActionsInput): { const streaming = useMemo(() => isSessionStreamStreaming(input.liveState), [input.liveState]); return { + activeTasks: input.liveState?.taskSnapshot?.tasks ?? [], lifecycle: input.liveState?.lifecycle ?? "IDLE", messages, permissionRequests, diff --git a/apps/web/src/domains/runtime/session-stream/session-stream-socket.ts b/apps/web/src/domains/runtime/session-stream/session-stream-socket.ts index e47434eb..6b0ac112 100644 --- a/apps/web/src/domains/runtime/session-stream/session-stream-socket.ts +++ b/apps/web/src/domains/runtime/session-stream/session-stream-socket.ts @@ -11,7 +11,7 @@ import type { MosooViewerCustomEvent, } from "@mosoo/ag-ui-session"; import { createPromiseDeferred, ignorePromiseRejection } from "@mosoo/effects"; -import { useCallback, useEffect, useRef, useState } from "react"; +import { useCallback, useEffect, useLayoutEffect, useRef, useState } from "react"; import type { MutableRefObject } from "react"; import { isTruthy } from "../../../shared/lib/truthiness"; @@ -78,16 +78,22 @@ export function useSessionStreamSocket( ); const activeSessionIdRef = useRef(sessionId); const liveStateRef = useRef(snapshot.liveState); + const reconnectTimeoutRef = useRef | null>(null); const renderSchedulerRef = useRef(null); const socketRef = useRef(null); - activeSessionIdRef.current = sessionId; + useLayoutEffect(() => { + activeSessionIdRef.current = sessionId; + + return () => { + activeSessionIdRef.current = null; + }; + }, [sessionId]); let scopedSnapshot = snapshot; if (snapshot.sessionId !== sessionId) { scopedSnapshot = createSessionStreamSnapshot(sessionId); - liveStateRef.current = scopedSnapshot.liveState; setSnapshot(scopedSnapshot); } @@ -115,19 +121,21 @@ export function useSessionStreamSocket( [], ); - const queueSocketEvents = useCallback( - (targetSessionId: string, events: AgUiSessionEvent[]) => { - if (events.length === 0) { - return; - } + const queueSocketEvents = useCallback((targetSessionId: string, events: AgUiSessionEvent[]) => { + if (events.length === 0) { + return; + } - renderSchedulerRef.current ??= new SessionStreamRenderScheduler(applyScheduledEvents); - renderSchedulerRef.current.enqueueMany(targetSessionId, events); - }, - [applyScheduledEvents], - ); + renderSchedulerRef.current ??= new SessionStreamRenderScheduler(applyScheduledEvents); + renderSchedulerRef.current.enqueueMany(targetSessionId, events); + }, []); const closeSocket = useCallback((reason: string) => { + if (reconnectTimeoutRef.current !== null) { + globalThis.clearTimeout(reconnectTimeoutRef.current); + reconnectTimeoutRef.current = null; + } + const { current } = socketRef; if (!current) { @@ -153,7 +161,7 @@ export function useSessionStreamSocket( }, []); const connectToSession = useCallback( - async (targetSessionId: string): Promise => { + async function connectToSession(targetSessionId: string): Promise { if (!isTruthy(appId)) { throw new Error("App id is required to open a session stream."); } @@ -232,20 +240,24 @@ export function useSessionStreamSocket( rescheduleStartedAt: new Date().toISOString(), }), ]); - globalThis.setTimeout(() => { + reconnectTimeoutRef.current = globalThis.setTimeout(() => { + reconnectTimeoutRef.current = null; + if (activeSessionIdRef.current !== targetSessionId) { return; } - void connectToSession(targetSessionId).then((nextSocket) => { - nextSocket.send( - JSON.stringify( - createViewerCustomEvent("mosoo.session.sync.request", { - reason: "reconnect", - }), - ), - ); - }); + void connectToSession(targetSessionId) + .then((nextSocket) => { + nextSocket.send( + JSON.stringify( + createViewerCustomEvent("mosoo.session.sync.request", { + reason: "reconnect", + }), + ), + ); + }) + .catch(ignorePromiseRejection); }, 350); } }); @@ -313,6 +325,7 @@ export function useSessionStreamSocket( void connectToSession(sessionId).catch(ignorePromiseRejection); return () => { + renderSchedulerRef.current?.clear(); closeSocket("session.effect.cleanup"); }; }, [closeSocket, connectToSession, appId, sessionId]); diff --git a/apps/web/src/domains/session/api/agent-session-retrieve.ts b/apps/web/src/domains/session/api/agent-session-retrieve.ts index 91359223..a74bdd03 100644 --- a/apps/web/src/domains/session/api/agent-session-retrieve.ts +++ b/apps/web/src/domains/session/api/agent-session-retrieve.ts @@ -20,6 +20,14 @@ const THREAD_AGENT_SESSION_RETRIEVE_QUERY = graphql(/* GraphQL */ ` reason status } + taskSnapshot { + runId + tasks { + taskId + taskType + title + } + } session { agentId archivedAt diff --git a/apps/web/src/features/help/help-menu.tsx b/apps/web/src/features/help/help-menu.tsx index 283a723e..e0f7aabd 100644 --- a/apps/web/src/features/help/help-menu.tsx +++ b/apps/web/src/features/help/help-menu.tsx @@ -1,5 +1,5 @@ import { HelpCircle } from "lucide-react"; -import { lazy, Suspense, useEffect, useState } from "react"; +import { lazy, Suspense, useCallback, useEffect, useState } from "react"; import { useTranslation } from "@/shared/i18n"; import { cn } from "@/shared/lib/class-names"; @@ -37,10 +37,10 @@ export function HelpMenu({ // fetched on a normal page load. const [hasOpened, setHasOpened] = useState(false); - function openHelp(): void { + const openHelp = useCallback((): void => { setHasOpened(true); setOpen(true); - } + }, []); function handleOpenChange(nextOpen: boolean): void { if (nextOpen) { @@ -72,7 +72,7 @@ export function HelpMenu({ return () => { globalThis.removeEventListener("keydown", handleKeyDown); }; - }, [shortcutEnabled]); + }, [openHelp, shortcutEnabled]); const button = (