diff --git a/doc/DATABASE.md b/doc/DATABASE.md index f69fb4eca8a..3f1c105a130 100644 --- a/doc/DATABASE.md +++ b/doc/DATABASE.md @@ -746,6 +746,43 @@ blocks concurrent task or run updates. It allows audit inserts to retain their foreign-key `KEY SHARE` locks without waiting on identity acquisition. The audit foreign keys and their deletion behavior remain enforced. +## Execution admission accounting + +The execution-admission module (`server/src/modules/execution-admission/`) owns +three tables. Nothing writes admissions or caps in production yet, and no claim +path enforces a project maximum. + +- `project_task_execution_policies` stores a project's maximum concurrent tasks + apart from `projects`, so project readers never return an unenforced cap. No + row means never configured. A row with a null maximum means an explicit "no + limit", which a later legacy migration must not overwrite. Maximums are + positive. The row is deleted with its project. +- `execution_admissions` stores one admission attempt of a run. A run can be + released before dispatch and admitted again, so `(run_id, attempt)` is unique + and a partial unique index allows one unsettled attempt per run. The row keeps + the context revision, the nullable primary project as provenance (no foreign + key) and the settlement identity, which is set together with `settled_at`. +- `execution_admission_project_charges` stores one row per distinct admitted + project, including unlimited projects. An admission without projects has no + rows. Composite `(company_id, …)` foreign keys keep admissions and projects in + one company. Rows exist only while the admission is unsettled: settlement + deletes the whole set in one transaction, and `(company_id, project_id)` is + indexed for active counts, so lookups scale with active work, not history. + +Archiving a project changes no charge. Deleting a project with a charge returns +`409 project_execution_admission_pending`; the foreign key is the backstop. +Deleting an agent refuses unsettled admissions on its runs and removes settled +ones. Company deletion removes the company's admissions before its runs and +projects. + +The proposed global lock order for the admission transaction is: company row +(`FOR NO KEY UPDATE`), native finalization recovery owner, the whole project set +in ascending ID order (`FOR NO KEY UPDATE`, one statement), issues, wake +requests, the run, the external-admission advisory lock, the admission row and +its charges, then budget reservations. Policy updates take the same project +lock. Settlement takes no project lock. Only the project-set step is implemented +so far; claim paths keep their current locks until they are integrated. + ## Switching between modes The database mode is controlled by `DATABASE_URL`: diff --git a/packages/db/src/execution-admission-migration.test.ts b/packages/db/src/execution-admission-migration.test.ts new file mode 100644 index 00000000000..6fefa53f9bd --- /dev/null +++ b/packages/db/src/execution-admission-migration.test.ts @@ -0,0 +1,53 @@ +import { createHash, randomUUID } from "node:crypto"; +import { readFile } from "node:fs/promises"; +import postgres from "postgres"; +import { describe, expect, it } from "vitest"; +import { applyPendingMigrations, inspectMigrations } from "./client.js"; +import { EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./test-embedded-postgres.js"; + +const migrationName = "0426_absent_guardian.sql"; +const migrationHash = createHash("sha256") + .update(await readFile(new URL(`./migrations/${migrationName}`, import.meta.url), "utf8")).digest("hex"); +const support = await getEmbeddedPostgresTestSupport(); +const describePostgres = support.supported ? describe : describe.skip; + +describePostgres("execution admission migration", () => { + it("adds empty admission storage without changing existing projects or runs", async () => { + const database = await startEmbeddedPostgresTestDatabase("paperclip-execution-admission-migration-"); + const sql = postgres(database.connectionString, { max: 1, onnotice: () => {} }); + try { + await sql`DROP TABLE execution_admission_project_charges, execution_admissions, project_task_execution_policies`; + await sql`DELETE FROM drizzle.__drizzle_migrations WHERE hash = ${migrationHash}`; + const companyId = randomUUID(), agentId = randomUUID(), projectId = randomUUID(); + await sql`INSERT INTO companies (id, name, issue_prefix) VALUES (${companyId}, 'Upgrade', 'UPG')`; + await sql`INSERT INTO agents (id, company_id, name, adapter_type) VALUES (${agentId}, ${companyId}, 'Agent', 'process')`; + await sql`INSERT INTO projects (id, company_id, name, execution_workspace_policy, archived_at) VALUES + (${projectId}, ${companyId}, 'Legacy', '{"enabled":true,"sharedWorkspaceConcurrency":"serialize"}', now())`; + await sql`INSERT INTO heartbeat_runs (company_id, agent_id, status, context_snapshot) VALUES + (${companyId}, ${agentId}, 'running', ${sql.json({ projectId })})`; + const contents = async () => ({ + projects: await sql`SELECT * FROM projects ORDER BY id`, + runs: await sql`SELECT * FROM heartbeat_runs ORDER BY id`, + }); + const before = await contents(); + expect(await inspectMigrations(database.connectionString)) + .toMatchObject({ status: "needsMigrations", pendingMigrations: [migrationName] }); + + await applyPendingMigrations(database.connectionString); + + expect(await contents()).toEqual(before); + for (const table of ["project_task_execution_policies", "execution_admissions", "execution_admission_project_charges"]) { + expect(await sql.unsafe(`SELECT count(*)::int AS count FROM "${table}"`)).toEqual([{ count: 0 }]); + } + expect((await sql`SELECT conname FROM pg_constraint WHERE conrelid = 'execution_admission_project_charges'::regclass + AND contype = 'f' ORDER BY conname`).map(row => row.conname)).toEqual([ + "execution_admission_project_charges_company_admission_fk", + "execution_admission_project_charges_company_project_fk", + ]); + expect(await inspectMigrations(database.connectionString)).toMatchObject({ status: "upToDate" }); + } finally { + await sql.end(); + await database.cleanup(); + } + }, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS); +}); diff --git a/packages/db/src/migrations/0426_absent_guardian.sql b/packages/db/src/migrations/0426_absent_guardian.sql new file mode 100644 index 00000000000..2bee4c7874f --- /dev/null +++ b/packages/db/src/migrations/0426_absent_guardian.sql @@ -0,0 +1,39 @@ +CREATE TABLE "execution_admission_project_charges" ( + "admission_id" uuid NOT NULL, + "company_id" uuid NOT NULL, + "project_id" uuid NOT NULL, + CONSTRAINT "execution_admission_project_charges_admission_id_project_id_pk" PRIMARY KEY("admission_id","project_id") +); +--> statement-breakpoint +CREATE TABLE "execution_admissions" ( + "id" uuid PRIMARY KEY NOT NULL, + "company_id" uuid NOT NULL, + "run_id" uuid NOT NULL, + "attempt" integer NOT NULL, + "primary_project_id" uuid, + "context_revision" text NOT NULL, + "settlement_id" uuid, + "settled_at" timestamp with time zone, + "created_at" timestamp with time zone DEFAULT now() NOT NULL, + CONSTRAINT "execution_admissions_attempt_check" CHECK ("execution_admissions"."attempt" > 0), + CONSTRAINT "execution_admissions_settlement_check" CHECK (("execution_admissions"."settlement_id" is null) = ("execution_admissions"."settled_at" is null)) +); +--> statement-breakpoint +CREATE TABLE "project_task_execution_policies" ( + "project_id" uuid PRIMARY KEY NOT NULL, + "company_id" uuid NOT NULL, + "max_concurrent_tasks" integer, + "created_at" timestamp with time zone DEFAULT now() NOT NULL, + "updated_at" timestamp with time zone DEFAULT now() NOT NULL, + CONSTRAINT "project_task_execution_policies_maximum_check" CHECK ("project_task_execution_policies"."max_concurrent_tasks" is null or "project_task_execution_policies"."max_concurrent_tasks" > 0) +); +--> statement-breakpoint +CREATE INDEX "execution_admission_project_charges_project_idx" ON "execution_admission_project_charges" USING btree ("company_id","project_id");--> statement-breakpoint +CREATE UNIQUE INDEX "execution_admissions_company_id_uq" ON "execution_admissions" USING btree ("company_id","id");--> statement-breakpoint +CREATE UNIQUE INDEX "execution_admissions_run_attempt_uq" ON "execution_admissions" USING btree ("run_id","attempt");--> statement-breakpoint +CREATE UNIQUE INDEX "execution_admissions_live_run_uq" ON "execution_admissions" USING btree ("run_id") WHERE "execution_admissions"."settled_at" is null;--> statement-breakpoint +ALTER TABLE "execution_admission_project_charges" ADD CONSTRAINT "execution_admission_project_charges_company_admission_fk" FOREIGN KEY ("company_id","admission_id") REFERENCES "public"."execution_admissions"("company_id","id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint +ALTER TABLE "execution_admission_project_charges" ADD CONSTRAINT "execution_admission_project_charges_company_project_fk" FOREIGN KEY ("company_id","project_id") REFERENCES "public"."projects"("company_id","id") ON DELETE no action ON UPDATE no action;--> statement-breakpoint +ALTER TABLE "execution_admissions" ADD CONSTRAINT "execution_admissions_company_id_companies_id_fk" FOREIGN KEY ("company_id") REFERENCES "public"."companies"("id") ON DELETE no action ON UPDATE no action;--> statement-breakpoint +ALTER TABLE "execution_admissions" ADD CONSTRAINT "execution_admissions_run_id_heartbeat_runs_id_fk" FOREIGN KEY ("run_id") REFERENCES "public"."heartbeat_runs"("id") ON DELETE no action ON UPDATE no action;--> statement-breakpoint +ALTER TABLE "project_task_execution_policies" ADD CONSTRAINT "project_task_execution_policies_company_project_fk" FOREIGN KEY ("company_id","project_id") REFERENCES "public"."projects"("company_id","id") ON DELETE cascade ON UPDATE no action; \ No newline at end of file diff --git a/packages/db/src/migrations/meta/0420_snapshot.json b/packages/db/src/migrations/meta/0426_snapshot.json similarity index 99% rename from packages/db/src/migrations/meta/0420_snapshot.json rename to packages/db/src/migrations/meta/0426_snapshot.json index 1f538a50635..0f939fcb071 100644 --- a/packages/db/src/migrations/meta/0420_snapshot.json +++ b/packages/db/src/migrations/meta/0426_snapshot.json @@ -1,6 +1,6 @@ { - "id": "8f75a1bb-32d9-4aed-9f93-134a6ae71f79", - "prevId": "52cf22f3-3bd9-4b0f-8d23-b422bb0d2540", + "id": "a1fcb4c6-40e7-47a6-b3f2-c43cd9bf332b", + "prevId": "8d9f4abf-57c7-443a-aa8b-3e3fd4f09dae", "version": "7", "dialect": "postgresql", "tables": { @@ -2233,6 +2233,74 @@ "checkConstraints": {}, "isRLSEnabled": false }, + "public.agent_file_roots": { + "name": "agent_file_roots", + "schema": "", + "columns": { + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "agent_id": { + "name": "agent_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "generation": { + "name": "generation", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "entries": { + "name": "entries", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": {}, + "foreignKeys": { + "agent_file_roots_company_id_agent_id_agents_company_id_id_fk": { + "name": "agent_file_roots_company_id_agent_id_agents_company_id_id_fk", + "tableFrom": "agent_file_roots", + "tableTo": "agents", + "columnsFrom": [ + "company_id", + "agent_id" + ], + "columnsTo": [ + "company_id", + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": { + "agent_file_roots_company_id_agent_id_pk": { + "name": "agent_file_roots_company_id_agent_id_pk", + "columns": [ + "company_id", + "agent_id" + ] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, "public.agent_identity_keys": { "name": "agent_identity_keys", "schema": "", @@ -28757,6 +28825,327 @@ "checkConstraints": {}, "isRLSEnabled": false }, + "public.execution_admission_project_charges": { + "name": "execution_admission_project_charges", + "schema": "", + "columns": { + "admission_id": { + "name": "admission_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "project_id": { + "name": "project_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + } + }, + "indexes": { + "execution_admission_project_charges_project_idx": { + "name": "execution_admission_project_charges_project_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "project_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "execution_admission_project_charges_company_admission_fk": { + "name": "execution_admission_project_charges_company_admission_fk", + "tableFrom": "execution_admission_project_charges", + "tableTo": "execution_admissions", + "columnsFrom": [ + "company_id", + "admission_id" + ], + "columnsTo": [ + "company_id", + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "execution_admission_project_charges_company_project_fk": { + "name": "execution_admission_project_charges_company_project_fk", + "tableFrom": "execution_admission_project_charges", + "tableTo": "projects", + "columnsFrom": [ + "company_id", + "project_id" + ], + "columnsTo": [ + "company_id", + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": { + "execution_admission_project_charges_admission_id_project_id_pk": { + "name": "execution_admission_project_charges_admission_id_project_id_pk", + "columns": [ + "admission_id", + "project_id" + ] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.execution_admissions": { + "name": "execution_admissions", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "run_id": { + "name": "run_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "attempt": { + "name": "attempt", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "primary_project_id": { + "name": "primary_project_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "context_revision": { + "name": "context_revision", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "settlement_id": { + "name": "settlement_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "settled_at": { + "name": "settled_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "execution_admissions_company_id_uq": { + "name": "execution_admissions_company_id_uq", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "execution_admissions_run_attempt_uq": { + "name": "execution_admissions_run_attempt_uq", + "columns": [ + { + "expression": "run_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "attempt", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "execution_admissions_live_run_uq": { + "name": "execution_admissions_live_run_uq", + "columns": [ + { + "expression": "run_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "\"execution_admissions\".\"settled_at\" is null", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "execution_admissions_company_id_companies_id_fk": { + "name": "execution_admissions_company_id_companies_id_fk", + "tableFrom": "execution_admissions", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "execution_admissions_run_id_heartbeat_runs_id_fk": { + "name": "execution_admissions_run_id_heartbeat_runs_id_fk", + "tableFrom": "execution_admissions", + "tableTo": "heartbeat_runs", + "columnsFrom": [ + "run_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": { + "execution_admissions_attempt_check": { + "name": "execution_admissions_attempt_check", + "value": "\"execution_admissions\".\"attempt\" > 0" + }, + "execution_admissions_settlement_check": { + "name": "execution_admissions_settlement_check", + "value": "(\"execution_admissions\".\"settlement_id\" is null) = (\"execution_admissions\".\"settled_at\" is null)" + } + }, + "isRLSEnabled": false + }, + "public.project_task_execution_policies": { + "name": "project_task_execution_policies", + "schema": "", + "columns": { + "project_id": { + "name": "project_id", + "type": "uuid", + "primaryKey": true, + "notNull": true + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "max_concurrent_tasks": { + "name": "max_concurrent_tasks", + "type": "integer", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": {}, + "foreignKeys": { + "project_task_execution_policies_company_project_fk": { + "name": "project_task_execution_policies_company_project_fk", + "tableFrom": "project_task_execution_policies", + "tableTo": "projects", + "columnsFrom": [ + "company_id", + "project_id" + ], + "columnsTo": [ + "company_id", + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": { + "project_task_execution_policies_maximum_check": { + "name": "project_task_execution_policies_maximum_check", + "value": "\"project_task_execution_policies\".\"max_concurrent_tasks\" is null or \"project_task_execution_policies\".\"max_concurrent_tasks\" > 0" + } + }, + "isRLSEnabled": false + }, "public.execution_workspace_repositories": { "name": "execution_workspace_repositories", "schema": "", @@ -29705,6 +30094,241 @@ "checkConstraints": {}, "isRLSEnabled": false }, + "public.execution_workspace_snapshots": { + "name": "execution_workspace_snapshots", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "execution_workspace_id": { + "name": "execution_workspace_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "source_operation_id": { + "name": "source_operation_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "source_run_id": { + "name": "source_run_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "parent_snapshot_id": { + "name": "parent_snapshot_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "generation": { + "name": "generation", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "epoch": { + "name": "epoch", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'pending'" + }, + "format": { + "name": "format", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'paperclip.remote-workspace-archive.v1'" + }, + "object_key": { + "name": "object_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "sha256": { + "name": "sha256", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "byte_size": { + "name": "byte_size", + "type": "bigint", + "primaryKey": false, + "notNull": false + }, + "key": { + "name": "key", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "manifest_digest": { + "name": "manifest_digest", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "metadata": { + "name": "metadata", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "delete_after": { + "name": "delete_after", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_error": { + "name": "last_error", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "execution_workspace_snapshots_operation_uniq": { + "name": "execution_workspace_snapshots_operation_uniq", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "execution_workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "source_operation_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "execution_workspace_snapshots_generation_uniq": { + "name": "execution_workspace_snapshots_generation_uniq", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "execution_workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "generation", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "execution_workspace_snapshots_workspace_idx": { + "name": "execution_workspace_snapshots_workspace_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "execution_workspace_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "execution_workspace_snapshots_cleanup_idx": { + "name": "execution_workspace_snapshots_cleanup_idx", + "columns": [ + { + "expression": "status", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "delete_after", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, "public.execution_workspaces": { "name": "execution_workspaces", "schema": "", @@ -50177,6 +50801,12 @@ "primaryKey": false, "notNull": false }, + "created_by_run_id": { + "name": "created_by_run_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, "dispatch_fingerprint": { "name": "dispatch_fingerprint", "type": "text", @@ -62139,6 +62769,12 @@ "primaryKey": false, "notNull": false }, + "cleanup_verified_at": { + "name": "cleanup_verified_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, "created_at": { "name": "created_at", "type": "timestamp with time zone", diff --git a/packages/db/src/migrations/meta/_journal.json b/packages/db/src/migrations/meta/_journal.json index b42ac91d174..b8a5f20dd9a 100644 --- a/packages/db/src/migrations/meta/_journal.json +++ b/packages/db/src/migrations/meta/_journal.json @@ -2969,6 +2969,13 @@ "when": 1793724587619, "tag": "0425_awesome_living_lightning", "breakpoints": true + }, + { + "idx": 426, + "version": "7", + "when": 1793733446322, + "tag": "0426_absent_guardian", + "breakpoints": true } ] } \ No newline at end of file diff --git a/packages/db/src/schema/execution_admissions.ts b/packages/db/src/schema/execution_admissions.ts new file mode 100644 index 00000000000..2a9e9fb641f --- /dev/null +++ b/packages/db/src/schema/execution_admissions.ts @@ -0,0 +1,59 @@ +import { sql } from "drizzle-orm"; +import { check, foreignKey, index, integer, pgTable, primaryKey, text, timestamp, uniqueIndex, uuid } from "drizzle-orm/pg-core"; +import { companies } from "./companies.js"; +import { heartbeatRuns } from "./heartbeat_runs.js"; +import { projects } from "./projects.js"; + +/** + * Project task-execution policy, kept apart from `projects` so unenforced caps + * never reach project readers. No row: never configured. Null maximum: explicit + * unlimited. Owned by the execution-admission module. + */ +export const projectTaskExecutionPolicies = pgTable("project_task_execution_policies", { + projectId: uuid("project_id").primaryKey(), + companyId: uuid("company_id").notNull(), + maxConcurrentTasks: integer("max_concurrent_tasks"), + createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), + updatedAt: timestamp("updated_at", { withTimezone: true }).notNull().defaultNow(), +}, t => ({ + project: foreignKey({ columns: [t.companyId, t.projectId], foreignColumns: [projects.companyId, projects.id], + name: "project_task_execution_policies_company_project_fk" }).onDelete("cascade"), + maximum: check("project_task_execution_policies_maximum_check", sql`${t.maxConcurrentTasks} is null or ${t.maxConcurrentTasks} > 0`), +})); + +/** + * One admission attempt of a run. A run can be admitted again after a settled + * attempt, so the run ID is not the admission identity. Primary project is + * provenance only; the charges below are the admitted project set. + */ +export const executionAdmissions = pgTable("execution_admissions", { + id: uuid("id").primaryKey(), + companyId: uuid("company_id").notNull().references(() => companies.id), + runId: uuid("run_id").notNull().references(() => heartbeatRuns.id), + attempt: integer("attempt").notNull(), + primaryProjectId: uuid("primary_project_id"), + contextRevision: text("context_revision").notNull(), + settlementId: uuid("settlement_id"), + settledAt: timestamp("settled_at", { withTimezone: true }), + createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), +}, t => ({ + companyIdentity: uniqueIndex("execution_admissions_company_id_uq").on(t.companyId, t.id), + runAttempt: uniqueIndex("execution_admissions_run_attempt_uq").on(t.runId, t.attempt), + liveRun: uniqueIndex("execution_admissions_live_run_uq").on(t.runId).where(sql`${t.settledAt} is null`), + attemptCheck: check("execution_admissions_attempt_check", sql`${t.attempt} > 0`), + settlementCheck: check("execution_admissions_settlement_check", sql`(${t.settlementId} is null) = (${t.settledAt} is null)`), +})); + +/** One active charge per distinct admitted project. Settlement deletes the whole set; a charged project cannot be deleted. */ +export const executionAdmissionProjectCharges = pgTable("execution_admission_project_charges", { + admissionId: uuid("admission_id").notNull(), + companyId: uuid("company_id").notNull(), + projectId: uuid("project_id").notNull(), +}, t => ({ + identity: primaryKey({ columns: [t.admissionId, t.projectId] }), + admission: foreignKey({ columns: [t.companyId, t.admissionId], foreignColumns: [executionAdmissions.companyId, executionAdmissions.id], + name: "execution_admission_project_charges_company_admission_fk" }).onDelete("cascade"), + project: foreignKey({ columns: [t.companyId, t.projectId], foreignColumns: [projects.companyId, projects.id], + name: "execution_admission_project_charges_company_project_fk" }), + activeProject: index("execution_admission_project_charges_project_idx").on(t.companyId, t.projectId), +})); diff --git a/packages/db/src/schema/index.ts b/packages/db/src/schema/index.ts index 1af7fbcfdf2..f12096d60cb 100644 --- a/packages/db/src/schema/index.ts +++ b/packages/db/src/schema/index.ts @@ -274,3 +274,5 @@ export { workflows } from "./workflows.js"; export { userWorkspacePreferences } from "./user_workspace_preferences.js"; export { issueThreadInteractionIdempotencyKeys } from "./issue_thread_interaction_idempotency_keys.js"; + +export { projectTaskExecutionPolicies, executionAdmissions, executionAdmissionProjectCharges } from "./execution_admissions.js"; diff --git a/server/src/modules/README.md b/server/src/modules/README.md index 6c3bb456c49..1da9c9b7c67 100644 --- a/server/src/modules/README.md +++ b/server/src/modules/README.md @@ -56,3 +56,9 @@ Ordinary configuration updates use the root lifecycle command. The [agent lifecycle module](../../../doc/AGENT-LIFECYCLE.md) owns agent creation, pause, resume, and termination. Use its commands for these changes. + +The execution admission module owns project task-execution policy and the +durable admission records that charge project capacity. Today it exposes only +internal policy commands and the deletion guards that project, agent, and +company deletion call. Execution claims still belong to `services/heartbeat/queue.ts` +and native review dispatch. See [execution admission accounting](../../../doc/DATABASE.md#execution-admission-accounting). diff --git a/server/src/modules/execution-admission/adapters/company-deletion.ts b/server/src/modules/execution-admission/adapters/company-deletion.ts new file mode 100644 index 00000000000..5e421170738 --- /dev/null +++ b/server/src/modules/execution-admission/adapters/company-deletion.ts @@ -0,0 +1,11 @@ +import { executionAdmissions } from "@paperclipai/db"; +import { eq } from "drizzle-orm"; +import type { CompanyDeletionParticipant } from "../../../lib/company-deletion.js"; + +/** Removes every admission, settled or not, before the company's runs and projects are deleted. */ +export const executionAdmissionCompanyDeletion: CompanyDeletionParticipant = { + async deleteCompanyData(tx, companyId) { + if (!("nestedIndex" in tx)) throw new Error("Company deletion requires a database transaction"); + await tx.delete(executionAdmissions).where(eq(executionAdmissions.companyId, companyId)); + }, +}; diff --git a/server/src/modules/execution-admission/adapters/postgres.test.ts b/server/src/modules/execution-admission/adapters/postgres.test.ts new file mode 100644 index 00000000000..ecdfef3cf07 --- /dev/null +++ b/server/src/modules/execution-admission/adapters/postgres.test.ts @@ -0,0 +1,239 @@ +import { randomUUID } from "node:crypto"; +import { + agents, + companies, + createDb, + executionAdmissionProjectCharges as charges, + executionAdmissions as admissions, + heartbeatRuns, + projectTaskExecutionPolicies as policies, + projects, +} from "@paperclipai/db"; +import { asc, eq, sql } from "drizzle-orm"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../../__tests__/helpers/embedded-postgres.js"; +import { withAccountingTransaction } from "../../../services/accounting-transaction.js"; +import { deleteAgentDependencies } from "../../../services/agent-deletion.js"; +import { deleteCompany } from "../../../services/company-deletion.js"; +import { projectService } from "../../../services/projects.js"; +import { prepareAdmissionRecord, type AdmissionRecord } from "../domain/admission.js"; +import { createExecutionAdmission } from "../index.js"; +import { countActiveProjectCharges, recordAdmission, settleAdmission } from "./postgres.js"; + +const support = await getEmbeddedPostgresTestSupport(); +type TestDb = ReturnType; + +function record(primaryProjectId: string | null, referencedProjectIds: string[], contextRevision = "issue:1") { + const result = prepareAdmissionRecord({ primaryProjectId, referencedProjectIds, contextRevision }); + if (result.kind !== "ok") throw new Error(result.kind); + return result.record; +} + +describe.skipIf(!support.supported)("execution admission storage", () => { + let database: Awaited>; + let db: TestDb; + + beforeAll(async () => { + database = await startEmbeddedPostgresTestDatabase("execution-admission-"); + db = createDb(database.connectionString); + }, 60_000); + + afterAll(async () => { + await db?.$client.end(); + await database?.cleanup(); + }); + + async function fixture() { + const [company] = await db.insert(companies).values({ name: "Admission fixture", issuePrefix: `A${randomUUID().slice(0, 6)}` }).returning(); + const [agent] = await db.insert(agents).values({ companyId: company.id, name: "Admission agent", adapterType: "process" }).returning(); + const projectRows = await db.insert(projects).values(["Design", "Docs", "Unlimited"].map(name => ({ companyId: company.id, name }))) + .returning(); + const [run] = await db.insert(heartbeatRuns).values({ companyId: company.id, agentId: agent.id }).returning(); + return { companyId: company.id, agentId: agent.id, projectIds: projectRows.map(row => row.id).sort(), runId: run.id }; + } + + const inTransaction = (work: (tx: TestDb) => Promise) => db.transaction(tx => work(tx as unknown as TestDb)); + const admit = (input: { companyId: string; runId: string; admissionId?: string; record: AdmissionRecord }) => + inTransaction(tx => recordAdmission(tx, { ...input, admissionId: input.admissionId ?? randomUUID() })); + const settle = (companyId: string, admissionId: string, settlementId = randomUUID()) => + inTransaction(tx => settleAdmission(tx, { companyId, admissionId, settlementId })); + const chargedProjects = async (admissionId: string) => (await db.select({ projectId: charges.projectId }).from(charges) + .where(eq(charges.admissionId, admissionId)).orderBy(asc(charges.projectId))).map(row => row.projectId); + + it("stores absent, explicit unlimited and positive maximums with partial updates", async () => { + const { companyId, projectIds: [projectId] } = await fixture(); + const policy = createExecutionAdmission(db); + const scope = { companyId, projectId }; + expect(await policy.readTaskExecutionPolicy(scope)).toEqual({ kind: "found", policy: { state: "absent", maxConcurrentTasks: null } }); + expect(await policy.updateTaskExecutionPolicy(scope, {})).toMatchObject({ kind: "unchanged", policy: { state: "absent" } }); + expect(await db.select().from(policies)).not.toContainEqual(expect.objectContaining({ projectId })); + + expect(await policy.updateTaskExecutionPolicy(scope, { maxConcurrentTasks: 2 })) + .toEqual({ kind: "updated", policy: { state: "configured", maxConcurrentTasks: 2 } }); + expect(await policy.updateTaskExecutionPolicy(scope, undefined)).toMatchObject({ kind: "unchanged", policy: { maxConcurrentTasks: 2 } }); + expect(await policy.updateTaskExecutionPolicy(scope, { maxConcurrentTasks: 0 })).toMatchObject({ kind: "invalid" }); + expect(await policy.updateTaskExecutionPolicy(scope, null)) + .toEqual({ kind: "updated", policy: { state: "configured", maxConcurrentTasks: null } }); + expect(await policy.readTaskExecutionPolicy(scope)).toEqual({ kind: "found", policy: { state: "configured", maxConcurrentTasks: null } }); + + const foreign = { companyId: (await fixture()).companyId, projectId }; + expect(await policy.readTaskExecutionPolicy(foreign)).toEqual({ kind: "project_not_found" }); + expect(await policy.updateTaskExecutionPolicy(foreign, { maxConcurrentTasks: 1 })).toEqual({ kind: "project_not_found" }); + expect(await policy.readTaskExecutionPolicy(scope)).toMatchObject({ policy: { maxConcurrentTasks: null } }); + }); + + it("enforces policy constraints and company scope in the database", async () => { + const { companyId, projectIds: [projectId] } = await fixture(); + const other = await fixture(); + await expect(db.insert(policies).values({ companyId, projectId, maxConcurrentTasks: 0 })) + .rejects.toMatchObject({ cause: { constraint_name: "project_task_execution_policies_maximum_check" } }); + await expect(db.insert(policies).values({ companyId: other.companyId, projectId, maxConcurrentTasks: 1 })) + .rejects.toMatchObject({ cause: { constraint_name: "project_task_execution_policies_company_project_fk" } }); + }); + + it("charges every distinct project, including unlimited and primaryless sets", async () => { + const { companyId, runId, projectIds: [a, b, c] } = await fixture(); + await createExecutionAdmission(db).updateTaskExecutionPolicy({ companyId, projectId: c }, null); + const admissionId = randomUUID(); + const admitted = record(a, [c, b, a]); + expect(await admit({ companyId, runId, admissionId, record: admitted })).toEqual({ kind: "recorded", attempt: 1 }); + expect(await chargedProjects(admissionId)).toEqual([a, b, c]); + const [epoch] = await db.select().from(admissions).where(eq(admissions.id, admissionId)); + expect(epoch).toMatchObject({ companyId, runId, attempt: 1, primaryProjectId: a, contextRevision: "issue:1", settlementId: null }); + + const other = await fixture(); + const primaryless = randomUUID(); + expect(await admit({ ...other, admissionId: primaryless, record: record(null, [other.projectIds[1]]) })) + .toMatchObject({ kind: "recorded" }); + expect(await chargedProjects(primaryless)).toEqual([other.projectIds[1]]); + + const empty = await fixture(); + const emptyId = randomUUID(); + expect(await admit({ ...empty, admissionId: emptyId, record: record(null, []) })).toMatchObject({ kind: "recorded" }); + expect(await chargedProjects(emptyId)).toEqual([]); + expect([...(await countActiveProjectCharges(db, empty.companyId, empty.projectIds)).values()]).toEqual([0, 0, 0]); + }); + + it("replays one admission identity and refuses a second live epoch for the run", async () => { + const { companyId, runId, projectIds: [a, b] } = await fixture(); + const admissionId = randomUUID(); + const admitted = record(a, [b]); + await admit({ companyId, runId, admissionId, record: admitted }); + expect(await admit({ companyId, runId, admissionId, record: admitted })).toEqual({ kind: "replayed", attempt: 1 }); + expect(await admit({ companyId, runId, admissionId, record: record(a, []) })).toEqual({ kind: "admission_conflict" }); + expect(await admit({ companyId, runId, admissionId, record: record(a, [b], "issue:2") })).toEqual({ kind: "admission_conflict" }); + expect(await admit({ companyId, runId, record: admitted })).toEqual({ kind: "run_already_admitted", admissionId }); + expect(await db.select({ id: admissions.id }).from(admissions).where(eq(admissions.runId, runId))).toHaveLength(1); + await expect(db.insert(admissions).values({ id: randomUUID(), companyId, runId, attempt: 2, contextRevision: "x" })) + .rejects.toMatchObject({ cause: { constraint_name: "execution_admissions_live_run_uq" } }); + }); + + it("settles the whole set once and starts a new attempt for the same run", async () => { + const { companyId, runId, projectIds: [a, b] } = await fixture(); + const first = randomUUID(); + await admit({ companyId, runId, admissionId: first, record: record(a, [b]) }); + const settlementId = randomUUID(); + expect(await settle(companyId, first, settlementId)).toEqual({ kind: "settled", releasedProjectIds: [a, b] }); + expect(await settle(companyId, first, settlementId)).toEqual({ kind: "already_settled" }); + expect(await settle(companyId, first)).toEqual({ kind: "settled_by_other" }); + expect(await settle(randomUUID(), first)).toEqual({ kind: "not_found" }); + expect(await chargedProjects(first)).toEqual([]); + expect(await admit({ companyId, runId, admissionId: first, record: record(a, [b]) })).toEqual({ kind: "admission_settled" }); + + const second = randomUUID(); + expect(await admit({ companyId, runId, admissionId: second, record: record(a, [], "issue:2") })).toEqual({ kind: "recorded", attempt: 2 }); + expect(await db.select({ attempt: admissions.attempt, settled: sql`${admissions.settledAt} is not null` }) + .from(admissions).where(eq(admissions.runId, runId)).orderBy(asc(admissions.attempt))) + .toEqual([{ attempt: 1, settled: true }, { attempt: 2, settled: false }]); + await expect(db.update(admissions).set({ settlementId: randomUUID() }).where(eq(admissions.id, second))) + .rejects.toMatchObject({ cause: { constraint_name: "execution_admissions_settlement_check" } }); + }); + + it("keeps company isolation and rolls back the whole record", async () => { + const { companyId, runId, projectIds: [a] } = await fixture(); + const other = await fixture(); + expect(await admit({ companyId, runId, record: record(a, [other.projectIds[0]]) })) + .toEqual({ kind: "projects_unavailable", projectIds: [other.projectIds[0]] }); + expect(await admit({ companyId, runId: other.runId, record: record(a, []) })).toEqual({ kind: "run_not_found" }); + const admissionId = randomUUID(); + await admit({ companyId, runId, admissionId, record: record(a, []) }); + expect(await admit({ companyId: other.companyId, runId: other.runId, admissionId, record: record(null, []) })) + .toEqual({ kind: "admission_conflict" }); + await expect(db.insert(charges).values({ admissionId, companyId, projectId: other.projectIds[0] })) + .rejects.toMatchObject({ cause: { constraint_name: "execution_admission_project_charges_company_project_fk" } }); + await expect(db.insert(charges).values({ admissionId, companyId, projectId: a })) + .rejects.toMatchObject({ cause: { constraint_name: "execution_admission_project_charges_admission_id_project_id_pk" } }); + + const rolledBack = randomUUID(); + await expect(inTransaction(async tx => { + await recordAdmission(tx, { companyId: other.companyId, runId: other.runId, admissionId: rolledBack, record: record(null, other.projectIds) }); + throw new Error("claim lost"); + })).rejects.toThrow("claim lost"); + expect(await db.select().from(admissions).where(eq(admissions.id, rolledBack))).toEqual([]); + expect(await chargedProjects(rolledBack)).toEqual([]); + await expect(recordAdmission(db, { companyId, runId, admissionId: randomUUID(), record: record(a, []) })) + .rejects.toThrow("requires a database transaction"); + }); + + it("counts active charges per project through the project index", async () => { + const first = await fixture(); + const [a, b, c] = first.projectIds; + const [secondRun] = await db.insert(heartbeatRuns).values({ companyId: first.companyId, agentId: first.agentId }).returning(); + const held = randomUUID(); + await admit({ ...first, admissionId: held, record: record(a, [b]) }); + await admit({ companyId: first.companyId, runId: secondRun.id, record: record(null, [a]) }); + expect(Object.fromEntries(await countActiveProjectCharges(db, first.companyId, [a, b, c]))).toEqual({ [a]: 2, [b]: 1, [c]: 0 }); + await settle(first.companyId, held); + expect(Object.fromEntries(await countActiveProjectCharges(db, first.companyId, [a, b, c]))).toEqual({ [a]: 1, [b]: 0, [c]: 0 }); + expect(Object.fromEntries(await countActiveProjectCharges(db, randomUUID(), [a]))).toEqual({ [a]: 0 }); + + const plan = await inTransaction(async tx => { + await tx.execute(sql`set local enable_seqscan = off`); + return tx.execute(sql`explain select project_id, count(*) from execution_admission_project_charges + where company_id = ${first.companyId} and project_id in (${a}, ${b}) group by project_id`); + }); + expect(JSON.stringify(plan)).toContain("execution_admission_project_charges_project_idx"); + }); + + it("preserves unsettled charges through project archive and delete", async () => { + const { companyId, runId, projectIds: [a, b] } = await fixture(); + const admissionId = randomUUID(); + await admit({ companyId, runId, admissionId, record: record(a, [b]) }); + const projectsService = projectService(db); + await db.update(projects).set({ archivedAt: new Date() }).where(eq(projects.id, b)); + expect(await chargedProjects(admissionId)).toEqual([a, b]); + await expect(projectsService.remove(a)).rejects.toMatchObject({ status: 409, details: { code: "project_execution_admission_pending" } }); + await expect(db.delete(projects).where(eq(projects.id, b))) + .rejects.toMatchObject({ cause: { constraint_name: "execution_admission_project_charges_company_project_fk" } }); + expect(await chargedProjects(admissionId)).toEqual([a, b]); + + await createExecutionAdmission(db).updateTaskExecutionPolicy({ companyId, projectId: a }, { maxConcurrentTasks: 1 }); + await settle(companyId, admissionId); + expect(await projectsService.remove(a)).toMatchObject({ id: a }); + expect(await db.select().from(policies).where(eq(policies.projectId, a))).toEqual([]); + expect(await db.select({ primaryProjectId: admissions.primaryProjectId }).from(admissions).where(eq(admissions.id, admissionId))) + .toEqual([{ primaryProjectId: a }]); + }); + + it("guards agent deletion while unsettled and removes admissions with the company", async () => { + const { companyId, agentId, runId, projectIds: [a] } = await fixture(); + const admissionId = randomUUID(); + await admit({ companyId, runId, admissionId, record: record(a, []) }); + await expect(withAccountingTransaction(db, companyId, tx => deleteAgentDependencies(tx, companyId, agentId))) + .rejects.toMatchObject({ status: 409, details: { code: "agent_execution_admission_pending" } }); + expect(await chargedProjects(admissionId)).toEqual([a]); + await settle(companyId, admissionId); + await withAccountingTransaction(db, companyId, tx => deleteAgentDependencies(tx, companyId, agentId)); + expect(await db.select().from(admissions).where(eq(admissions.runId, runId))).toEqual([]); + + const deleted = await fixture(); + const retained = await fixture(); + const retainedAdmission = randomUUID(); + await admit({ ...retained, admissionId: retainedAdmission, record: record(retained.projectIds[0], []) }); + await admit({ ...deleted, record: record(deleted.projectIds[0], deleted.projectIds) }); + await db.update(agents).set({ lifecycleState: "terminated", status: "terminated" }).where(eq(agents.id, deleted.agentId)); + expect(await deleteCompany(db, deleted.companyId)).toMatchObject({ id: deleted.companyId }); + expect(await db.select().from(admissions).where(eq(admissions.companyId, deleted.companyId))).toEqual([]); + expect(await chargedProjects(retainedAdmission)).toEqual([retained.projectIds[0]]); + }, 60_000); +}); diff --git a/server/src/modules/execution-admission/adapters/postgres.ts b/server/src/modules/execution-admission/adapters/postgres.ts new file mode 100644 index 00000000000..937eef72e14 --- /dev/null +++ b/server/src/modules/execution-admission/adapters/postgres.ts @@ -0,0 +1,165 @@ +import { + executionAdmissionProjectCharges as charges, + executionAdmissions as admissions, + heartbeatRuns, + projectTaskExecutionPolicies as policies, + projects, + type Db, +} from "@paperclipai/db"; +import { and, asc, eq, inArray, isNull, sql } from "drizzle-orm"; +import type { TaskExecutionPolicyStore } from "../application/policy.js"; +import { decideSettlement, sameAdmissionRecord, type AdmissionRecord } from "../domain/admission.js"; +import { ABSENT_TASK_EXECUTION_POLICY, type TaskExecutionPolicy } from "../domain/policy.js"; + +function assertTransaction(tx: Db) { + if (!("nestedIndex" in tx)) throw new Error("Execution admission storage requires a database transaction"); +} + +/** + * Global lock order step 3: locks the whole project set in ascending ID order + * in one statement. Returns the IDs that are not projects of the company. + */ +export async function lockAdmissionProjects(tx: Db, companyId: string, projectIds: readonly string[]) { + assertTransaction(tx); + if (projectIds.length === 0) return []; + const rows = await tx.select({ id: projects.id }).from(projects) + .where(and(eq(projects.companyId, companyId), inArray(projects.id, [...projectIds]))) + .orderBy(asc(projects.id)).for("no key update"); + const found = new Set(rows.map(row => row.id)); + return projectIds.filter(id => !found.has(id)); +} + +function toPolicy(row: { maxConcurrentTasks: number | null } | undefined): TaskExecutionPolicy { + return row ? { state: "configured", maxConcurrentTasks: row.maxConcurrentTasks } : ABSENT_TASK_EXECUTION_POLICY; +} + +export function createPostgresTaskExecutionPolicyStore(db: Db): TaskExecutionPolicyStore { + return { + async read({ companyId, projectId }) { + const [row] = await db.select({ configured: policies.projectId, maxConcurrentTasks: policies.maxConcurrentTasks }) + .from(projects).leftJoin(policies, eq(policies.projectId, projects.id)) + .where(and(eq(projects.companyId, companyId), eq(projects.id, projectId))); + if (!row) return null; + return row.configured ? toPolicy(row) : ABSENT_TASK_EXECUTION_POLICY; + }, + + update: ({ companyId, projectId }, change) => db.transaction(async transaction => { + const tx = transaction as unknown as Db; + if ((await lockAdmissionProjects(tx, companyId, [projectId])).length > 0) return null; + const [row] = await tx.select({ maxConcurrentTasks: policies.maxConcurrentTasks }).from(policies) + .where(eq(policies.projectId, projectId)); + const current = toPolicy(row); + const next = change(current); + if (next === current) return current; + if (next.state === "absent") throw new Error("A stored task execution policy cannot become absent"); + await tx.insert(policies).values({ projectId, companyId, maxConcurrentTasks: next.maxConcurrentTasks }) + .onConflictDoUpdate({ target: policies.projectId, set: { maxConcurrentTasks: next.maxConcurrentTasks, updatedAt: sql`now()` } }); + return next; + }), + }; +} + +export type RecordAdmissionInput = { companyId: string; runId: string; admissionId: string; record: AdmissionRecord }; + +export type RecordAdmissionOutcome = + | { kind: "recorded" | "replayed"; attempt: number } + | { kind: "projects_unavailable"; projectIds: string[] } + | { kind: "run_not_found" } + | { kind: "run_already_admitted"; admissionId: string } + | { kind: "admission_settled" } + | { kind: "admission_conflict" }; + +/** + * Storage transition only: persists one admission attempt and its charges in + * the caller's transaction. It neither checks capacity nor claims work; the + * composed admission command must do both under the same locks. Takes lock + * steps 3 (projects) and 6 (run); composed callers take them earlier. + */ +export async function recordAdmission(tx: Db, input: RecordAdmissionInput): Promise { + assertTransaction(tx); + const { companyId, runId, admissionId, record } = input; + const missing = await lockAdmissionProjects(tx, companyId, record.projectIds); + if (missing.length > 0) return { kind: "projects_unavailable", projectIds: missing }; + const [run] = await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns) + .where(and(eq(heartbeatRuns.id, runId), eq(heartbeatRuns.companyId, companyId))).for("no key update"); + if (!run) return { kind: "run_not_found" }; + + const [existing] = await tx.select().from(admissions).where(eq(admissions.id, admissionId)); + if (existing) { + if (existing.companyId !== companyId || existing.runId !== runId) return { kind: "admission_conflict" }; + if (existing.settledAt) return { kind: "admission_settled" }; + const charged = await tx.select({ projectId: charges.projectId }).from(charges) + .where(and(eq(charges.companyId, companyId), eq(charges.admissionId, admissionId))).orderBy(asc(charges.projectId)); + const stored = { primaryProjectId: existing.primaryProjectId, contextRevision: existing.contextRevision, + projectIds: charged.map(row => row.projectId) }; + return sameAdmissionRecord(stored, record) ? { kind: "replayed", attempt: existing.attempt } : { kind: "admission_conflict" }; + } + const [live] = await tx.select({ id: admissions.id }).from(admissions) + .where(and(eq(admissions.runId, runId), isNull(admissions.settledAt))); + if (live) return { kind: "run_already_admitted", admissionId: live.id }; + + const [{ attempt }] = await tx.select({ attempt: sql`coalesce(max(${admissions.attempt}), 0)::int + 1` }) + .from(admissions).where(eq(admissions.runId, runId)); + await tx.insert(admissions).values({ id: admissionId, companyId, runId, attempt, + primaryProjectId: record.primaryProjectId, contextRevision: record.contextRevision }); + if (record.projectIds.length > 0) { + await tx.insert(charges).values(record.projectIds.map(projectId => ({ admissionId, companyId, projectId }))); + } + return { kind: "recorded", attempt }; +} + +export type SettleAdmissionOutcome = + | { kind: "settled"; releasedProjectIds: string[] } + | { kind: "already_settled" | "settled_by_other" | "not_found" }; + +/** + * Storage transition only: releases the admission's whole charge set once. + * Callers must first verify authoritative lifecycle facts. Takes no project lock. + */ +export async function settleAdmission(tx: Db, input: { companyId: string; admissionId: string; settlementId: string }): +Promise { + assertTransaction(tx); + const { companyId, admissionId, settlementId } = input; + const [row] = await tx.select({ settlementId: admissions.settlementId }).from(admissions) + .where(and(eq(admissions.companyId, companyId), eq(admissions.id, admissionId))).for("no key update"); + if (!row) return { kind: "not_found" }; + const decision = decideSettlement(row.settlementId, settlementId); + if (decision !== "settle") return { kind: decision }; + const released = await tx.delete(charges) + .where(and(eq(charges.companyId, companyId), eq(charges.admissionId, admissionId))) + .returning({ projectId: charges.projectId }); + await tx.update(admissions).set({ settlementId, settledAt: sql`now()` }).where(eq(admissions.id, admissionId)); + return { kind: "settled", releasedProjectIds: released.map(charge => charge.projectId).sort() }; +} + +/** Active charge counts per project. Every stored charge is unsettled. */ +export async function countActiveProjectCharges(db: Db, companyId: string, projectIds: readonly string[]) { + const counts = new Map(projectIds.map(id => [id, 0])); + if (projectIds.length === 0) return counts; + const rows = await db.select({ projectId: charges.projectId, active: sql`count(*)::int` }).from(charges) + .where(and(eq(charges.companyId, companyId), inArray(charges.projectId, [...projectIds]))) + .groupBy(charges.projectId); + for (const row of rows) counts.set(row.projectId, row.active); + return counts; +} + +/** Project deletion holds the project row lock, so no charge can be added concurrently. */ +export async function hasUnsettledProjectCharges(db: Db, scope: { companyId: string; projectId: string }) { + const [row] = await db.select({ admissionId: charges.admissionId }).from(charges) + .where(and(eq(charges.companyId, scope.companyId), eq(charges.projectId, scope.projectId))).limit(1); + return Boolean(row); +} + +/** Removes settled history of an agent's runs before they are deleted; refuses while any admission is unsettled. */ +export async function deleteAgentAdmissionHistory(tx: Db, scope: { companyId: string; agentId: string }) { + assertTransaction(tx); + const agentRuns = tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns) + .where(and(eq(heartbeatRuns.companyId, scope.companyId), eq(heartbeatRuns.agentId, scope.agentId))); + const [live] = await tx.select({ id: admissions.id }).from(admissions) + .where(and(eq(admissions.companyId, scope.companyId), inArray(admissions.runId, agentRuns), isNull(admissions.settledAt))) + .limit(1); + if (live) return { kind: "unsettled" as const }; + await tx.delete(admissions).where(and(eq(admissions.companyId, scope.companyId), inArray(admissions.runId, agentRuns))); + return { kind: "deleted" as const }; +} + diff --git a/server/src/modules/execution-admission/application/policy.ts b/server/src/modules/execution-admission/application/policy.ts new file mode 100644 index 00000000000..3498b31102a --- /dev/null +++ b/server/src/modules/execution-admission/application/policy.ts @@ -0,0 +1,50 @@ +import { + applyTaskExecutionPolicyPatch, + parseTaskExecutionPolicyPatch, + type TaskExecutionPolicy, +} from "../domain/policy.js"; + +export type ProjectScope = { companyId: string; projectId: string }; + +export interface TaskExecutionPolicyStore { + /** Null when the project is not in the company. */ + read(scope: ProjectScope): Promise; + /** + * Locks the project in admission lock order and stores `change(current)`. + * Returning `current` itself writes nothing. Null when the project is not in the company. + */ + update(scope: ProjectScope, change: (current: TaskExecutionPolicy) => TaskExecutionPolicy): Promise; +} + +export type PolicyReadResult = { kind: "found"; policy: TaskExecutionPolicy } | { kind: "project_not_found" }; + +export type PolicyUpdateResult = + | { kind: "updated" | "unchanged"; policy: TaskExecutionPolicy } + | { kind: "invalid"; field: "policy" | "maxConcurrentTasks"; reason: string } + | { kind: "project_not_found" }; + +/** + * Internal policy commands. No route, tool or import calls `update` until + * admission enforces the maximum; a cap written earlier would mislead users. + */ +export function createTaskExecutionPolicyCommands(store: TaskExecutionPolicyStore) { + return { + async readTaskExecutionPolicy(scope: ProjectScope): Promise { + const policy = await store.read(scope); + return policy ? { kind: "found", policy } : { kind: "project_not_found" }; + }, + + async updateTaskExecutionPolicy(scope: ProjectScope, input: unknown): Promise { + const parsed = parseTaskExecutionPolicyPatch(input); + if (parsed.kind === "invalid") return parsed; + let changed = false; + const policy = await store.update(scope, current => { + const next = applyTaskExecutionPolicyPatch(current, parsed.patch); + changed = next !== current; + return next; + }); + if (!policy) return { kind: "project_not_found" }; + return { kind: changed ? "updated" : "unchanged", policy }; + }, + }; +} diff --git a/server/src/modules/execution-admission/company-deletion.ts b/server/src/modules/execution-admission/company-deletion.ts new file mode 100644 index 00000000000..cfaf4097f2a --- /dev/null +++ b/server/src/modules/execution-admission/company-deletion.ts @@ -0,0 +1 @@ +export { executionAdmissionCompanyDeletion } from "./adapters/company-deletion.js"; diff --git a/server/src/modules/execution-admission/domain/admission.ts b/server/src/modules/execution-admission/domain/admission.ts new file mode 100644 index 00000000000..4177e9c4c06 --- /dev/null +++ b/server/src/modules/execution-admission/domain/admission.ts @@ -0,0 +1,53 @@ +/** Referenced projects are rejected beyond this bound instead of admitting a subset. */ +export const MAX_ADMISSION_PROJECTS = 100; +const MAX_CONTEXT_REVISION_LENGTH = 256; +const UUID = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/; + +export type AdmissionRecordInput = { + primaryProjectId: string | null; + referencedProjectIds: readonly string[]; + contextRevision: string; +}; + +/** + * `projectIds` is the complete distinct set in canonical lock order. Lowercase + * UUID text order matches PostgreSQL uuid order. An empty set is valid and + * charges no project. + */ +export type AdmissionRecord = { + primaryProjectId: string | null; + projectIds: string[]; + contextRevision: string; +}; + +export type AdmissionRecordResult = + | { kind: "ok"; record: AdmissionRecord } + | { kind: "invalid_project_id" } + | { kind: "too_many_projects"; count: number; limit: number } + | { kind: "invalid_context_revision" }; + +export function prepareAdmissionRecord(input: AdmissionRecordInput): AdmissionRecordResult { + const contextRevision = input.contextRevision.trim(); + if (!contextRevision || contextRevision.length > MAX_CONTEXT_REVISION_LENGTH) return { kind: "invalid_context_revision" }; + const primaryProjectId = input.primaryProjectId?.toLowerCase() ?? null; + const candidates = [...input.referencedProjectIds, ...(primaryProjectId === null ? [] : [primaryProjectId])].map(id => id.toLowerCase()); + if (candidates.some(id => !UUID.test(id))) return { kind: "invalid_project_id" }; + const projectIds = [...new Set(candidates)].sort(); + if (projectIds.length > MAX_ADMISSION_PROJECTS) { + return { kind: "too_many_projects", count: projectIds.length, limit: MAX_ADMISSION_PROJECTS }; + } + return { kind: "ok", record: { primaryProjectId, projectIds, contextRevision } }; +} + +export function sameAdmissionRecord(left: AdmissionRecord, right: AdmissionRecord) { + return left.primaryProjectId === right.primaryProjectId && left.contextRevision === right.contextRevision + && left.projectIds.length === right.projectIds.length && left.projectIds.every((id, index) => id === right.projectIds[index]); +} + +export type SettlementDecision = "settle" | "already_settled" | "settled_by_other"; + +/** A settlement identity can be replayed; a second, different settlement cannot take over. */ +export function decideSettlement(existingSettlementId: string | null, settlementId: string): SettlementDecision { + if (existingSettlementId === null) return "settle"; + return existingSettlementId === settlementId ? "already_settled" : "settled_by_other"; +} diff --git a/server/src/modules/execution-admission/domain/policy.test.ts b/server/src/modules/execution-admission/domain/policy.test.ts new file mode 100644 index 00000000000..0c5ce8ebd6d --- /dev/null +++ b/server/src/modules/execution-admission/domain/policy.test.ts @@ -0,0 +1,88 @@ +import { describe, expect, it } from "vitest"; +import { decideSettlement, MAX_ADMISSION_PROJECTS, prepareAdmissionRecord } from "./admission.js"; +import { ABSENT_TASK_EXECUTION_POLICY, applyTaskExecutionPolicyPatch, parseTaskExecutionPolicyPatch, type TaskExecutionPolicy } from "./policy.js"; + +const configured = (maxConcurrentTasks: number | null): TaskExecutionPolicy => ({ state: "configured", maxConcurrentTasks }); +const apply = (current: TaskExecutionPolicy, input: unknown) => { + const parsed = parseTaskExecutionPolicyPatch(input); + if (parsed.kind !== "ok") throw new Error(parsed.reason); + return applyTaskExecutionPolicyPatch(current, parsed.patch); +}; + +describe("task execution policy", () => { + it("keeps never-configured distinct from explicit unlimited", () => { + expect(apply(ABSENT_TASK_EXECUTION_POLICY, null)).toEqual(configured(null)); + expect(apply(ABSENT_TASK_EXECUTION_POLICY, { maxConcurrentTasks: null })).toEqual(configured(null)); + expect(apply(ABSENT_TASK_EXECUTION_POLICY, { maxConcurrentTasks: 3 })).toEqual(configured(3)); + }); + + it("preserves the current value for omitted policies and fields", () => { + for (const current of [ABSENT_TASK_EXECUTION_POLICY, configured(null), configured(2)]) { + expect(apply(current, undefined)).toBe(current); + expect(apply(current, {})).toBe(current); + expect(apply(current, { maxConcurrentTasks: undefined })).toBe(current); + } + const two = configured(2); + expect(apply(two, { maxConcurrentTasks: 2 })).toBe(two); + const unlimited = configured(null); + expect(apply(unlimited, null)).toBe(unlimited); + }); + + it("clears or replaces a configured maximum", () => { + expect(apply(configured(2), null)).toEqual(configured(null)); + expect(apply(configured(2), { maxConcurrentTasks: null })).toEqual(configured(null)); + expect(apply(configured(null), { maxConcurrentTasks: 1 })).toEqual(configured(1)); + }); + + it.each([0, -1, 1.5, Number.NaN, Number.POSITIVE_INFINITY, 2_147_483_648, "2", true, {}])( + "rejects maximum %s", value => { + expect(parseTaskExecutionPolicyPatch({ maxConcurrentTasks: value })).toMatchObject({ kind: "invalid", field: "maxConcurrentTasks" }); + }, + ); + + it("rejects malformed policies and unknown fields", () => { + for (const input of [1, "unlimited", [], { maxConcurrentTasks: 1, perTaskOverride: 3 }]) { + expect(parseTaskExecutionPolicyPatch(input)).toMatchObject({ kind: "invalid", field: "policy" }); + } + expect(parseTaskExecutionPolicyPatch({ maxConcurrentTasks: 2_147_483_647 })).toEqual({ kind: "ok", patch: { maxConcurrentTasks: 2_147_483_647 } }); + }); +}); + +describe("admission record", () => { + const a = "0b9d8e52-6f2e-4bb0-9a4a-1f0a7f2e0a01"; + const b = "7c1f3d9a-2b4e-4d8f-8a6b-3e5c7d9f1b02"; + const c = "f3a2b1c0-9d8e-4f7a-8b6c-5d4e3f2a1b03"; + + it("unions the primary with referenced projects in canonical order", () => { + expect(prepareAdmissionRecord({ primaryProjectId: c, referencedProjectIds: [b, a, b.toUpperCase(), c], contextRevision: " r1 " })) + .toEqual({ kind: "ok", record: { primaryProjectId: c, projectIds: [a, b, c], contextRevision: "r1" } }); + }); + + it("keeps a primaryless set and an empty set without a project bucket", () => { + expect(prepareAdmissionRecord({ primaryProjectId: null, referencedProjectIds: [b], contextRevision: "r" })) + .toEqual({ kind: "ok", record: { primaryProjectId: null, projectIds: [b], contextRevision: "r" } }); + expect(prepareAdmissionRecord({ primaryProjectId: null, referencedProjectIds: [], contextRevision: "r" })) + .toEqual({ kind: "ok", record: { primaryProjectId: null, projectIds: [], contextRevision: "r" } }); + }); + + it("rejects invalid references, oversized sets and missing context identity instead of narrowing", () => { + expect(prepareAdmissionRecord({ primaryProjectId: null, referencedProjectIds: [a, "Design"], contextRevision: "r" })) + .toEqual({ kind: "invalid_project_id" }); + expect(prepareAdmissionRecord({ primaryProjectId: "", referencedProjectIds: [a], contextRevision: "r" })) + .toEqual({ kind: "invalid_project_id" }); + const many = Array.from({ length: MAX_ADMISSION_PROJECTS + 1 }, (_, index) => + `00000000-0000-4000-8000-${index.toString(16).padStart(12, "0")}`); + expect(prepareAdmissionRecord({ primaryProjectId: null, referencedProjectIds: many, contextRevision: "r" })) + .toEqual({ kind: "too_many_projects", count: MAX_ADMISSION_PROJECTS + 1, limit: MAX_ADMISSION_PROJECTS }); + expect(prepareAdmissionRecord({ primaryProjectId: a, referencedProjectIds: [], contextRevision: " " })) + .toEqual({ kind: "invalid_context_revision" }); + expect(prepareAdmissionRecord({ primaryProjectId: a, referencedProjectIds: [], contextRevision: "x".repeat(257) })) + .toEqual({ kind: "invalid_context_revision" }); + }); + + it("settles once and replays only the same settlement", () => { + expect(decideSettlement(null, a)).toBe("settle"); + expect(decideSettlement(a, a)).toBe("already_settled"); + expect(decideSettlement(a, b)).toBe("settled_by_other"); + }); +}); diff --git a/server/src/modules/execution-admission/domain/policy.ts b/server/src/modules/execution-admission/domain/policy.ts new file mode 100644 index 00000000000..e171011d3b2 --- /dev/null +++ b/server/src/modules/execution-admission/domain/policy.ts @@ -0,0 +1,47 @@ +/** + * A project's task-execution policy. `absent` was never configured, so a later + * legacy migration may still choose a value. `configured` with a null maximum + * is an explicit "no limit" that such a migration must preserve. + */ +export type TaskExecutionPolicy = + | { state: "absent"; maxConcurrentTasks: null } + | { state: "configured"; maxConcurrentTasks: number | null }; + +/** Null clears the cap explicitly. An omitted field preserves its value. */ +export type TaskExecutionPolicyPatch = { maxConcurrentTasks?: number | null } | null; + +export type PolicyPatchResult = + | { kind: "ok"; patch: TaskExecutionPolicyPatch } + | { kind: "invalid"; field: "policy" | "maxConcurrentTasks"; reason: string }; + +export const ABSENT_TASK_EXECUTION_POLICY: TaskExecutionPolicy = { state: "absent", maxConcurrentTasks: null }; + +// The column is a PostgreSQL integer. +const MAX_STORED_MAXIMUM = 2_147_483_647; + +/** Undefined (an omitted policy) preserves everything; null clears the cap. */ +export function parseTaskExecutionPolicyPatch(input: unknown): PolicyPatchResult { + if (input === undefined) return { kind: "ok", patch: {} }; + if (input === null) return { kind: "ok", patch: null }; + if (typeof input !== "object" || Array.isArray(input)) { + return { kind: "invalid", field: "policy", reason: "Policy must be an object or null" }; + } + const unknownKey = Object.keys(input).find(key => key !== "maxConcurrentTasks"); + if (unknownKey) return { kind: "invalid", field: "policy", reason: `Unknown policy field ${unknownKey}` }; + if (!("maxConcurrentTasks" in input)) return { kind: "ok", patch: {} }; + const value = (input as { maxConcurrentTasks: unknown }).maxConcurrentTasks; + if (value === undefined) return { kind: "ok", patch: {} }; + if (value === null) return { kind: "ok", patch: { maxConcurrentTasks: null } }; + if (typeof value !== "number" || !Number.isInteger(value) || value < 1 || value > MAX_STORED_MAXIMUM) { + return { kind: "invalid", field: "maxConcurrentTasks", reason: "Maximum concurrent tasks must be a positive whole number" }; + } + return { kind: "ok", patch: { maxConcurrentTasks: value } }; +} + +/** Applies a validated patch. Returns the current policy when nothing changes. */ +export function applyTaskExecutionPolicyPatch(current: TaskExecutionPolicy, patch: TaskExecutionPolicyPatch): TaskExecutionPolicy { + if (patch !== null && patch.maxConcurrentTasks === undefined) return current; + const maxConcurrentTasks = patch === null ? null : patch.maxConcurrentTasks ?? null; + if (current.state === "configured" && current.maxConcurrentTasks === maxConcurrentTasks) return current; + return { state: "configured", maxConcurrentTasks }; +} diff --git a/server/src/modules/execution-admission/index.ts b/server/src/modules/execution-admission/index.ts new file mode 100644 index 00000000000..85878c38a05 --- /dev/null +++ b/server/src/modules/execution-admission/index.ts @@ -0,0 +1,17 @@ +import type { Db } from "@paperclipai/db"; +import { createPostgresTaskExecutionPolicyStore } from "./adapters/postgres.js"; +import { createTaskExecutionPolicyCommands } from "./application/policy.js"; + +export type { TaskExecutionPolicy, TaskExecutionPolicyPatch } from "./domain/policy.js"; +export type { PolicyReadResult, PolicyUpdateResult, ProjectScope } from "./application/policy.js"; +// Lifecycle guards run inside the caller's deletion transaction. +export { deleteAgentAdmissionHistory, hasUnsettledProjectCharges } from "./adapters/postgres.js"; + +/** + * Execution admission owns project task-execution policy and the durable + * admission epochs that charge project capacity. Admission and settlement + * commands are not composed yet; nothing enforces the maximum. + */ +export function createExecutionAdmission(db: Db) { + return createTaskExecutionPolicyCommands(createPostgresTaskExecutionPolicyStore(db)); +} diff --git a/server/src/services/agent-deletion.ts b/server/src/services/agent-deletion.ts index ca18e8f02ec..245919fff3a 100644 --- a/server/src/services/agent-deletion.ts +++ b/server/src/services/agent-deletion.ts @@ -14,6 +14,7 @@ import { agents, agentApiKeys, agentRuntimeState, agentTaskSessions, agentWakeup caseHandoffs, principalPermissionGrants, type Db } from "@paperclipai/db"; import { conflict } from "../errors.js"; import { issueThreadInteractionService } from "./issue-thread-interactions.js"; +import { deleteAgentAdmissionHistory } from "../modules/execution-admission/index.js"; export async function deleteAgentDependencies(tx: Db, companyId: string, id: string) { const [retainedHandoff] = await tx.select({ id: caseHandoffs.id }).from(caseHandoffs).where(and( @@ -63,6 +64,11 @@ export async function deleteAgentDependencies(tx: Db, companyId: string, id: str await notifyInstructionCollection(tx, { companyId: companyId }); await notifyCancelledWorkspacePreservation(tx, { companyId: companyId }); await notifyUnsafeWorkspaceExportRecovery(tx, { companyId }); + if ((await deleteAgentAdmissionHistory(tx, { companyId, agentId: id })).kind === "unsettled") { + throw conflict("Wait for this agent's admitted task executions to settle before deleting it", { + code: "agent_execution_admission_pending", + }); + } await tx.delete(heartbeatRuns).where(eq(heartbeatRuns.agentId, id)); await tx.delete(agentWakeupRequests).where(eq(agentWakeupRequests.agentId, id)); await tx.delete(agentApiKeys).where(eq(agentApiKeys.agentId, id)); diff --git a/server/src/services/company-deletion.ts b/server/src/services/company-deletion.ts index 66e81df7673..2d341072d2f 100644 --- a/server/src/services/company-deletion.ts +++ b/server/src/services/company-deletion.ts @@ -20,6 +20,7 @@ import { eq, inArray } from "drizzle-orm"; import { createAgentLifecycle, processAgentLifecycle } from "./agent-lifecycle.js"; import { requestRestorePointDeletion } from "./git-restore-point-ledger.js"; import { agentLifecycleCompanyDeletion } from "../modules/agent-lifecycle/company-deletion.js"; +import { executionAdmissionCompanyDeletion } from "../modules/execution-admission/company-deletion.js"; import { companies, agents, @@ -117,6 +118,7 @@ export async function deleteCompany(db: Db, id: string) { await notifyInstructionCollection(tx, { companyId: id }); await notifyCancelledWorkspacePreservation(tx, { companyId: id }); await notifyUnsafeWorkspaceExportRecovery(tx, { companyId: id }); + await executionAdmissionCompanyDeletion.deleteCompanyData(tx, id); await tx.delete(heartbeatRuns).where(eq(heartbeatRuns.companyId, id)); await tx.delete(agentWakeupRequests).where(eq(agentWakeupRequests.companyId, id)); await tx.delete(agentApiKeys).where(eq(agentApiKeys.companyId, id)); diff --git a/server/src/services/projects.ts b/server/src/services/projects.ts index 45891a218a6..5af27d351a7 100644 --- a/server/src/services/projects.ts +++ b/server/src/services/projects.ts @@ -43,6 +43,7 @@ import { parseProjectExecutionWorkspacePolicy } from "./execution-workspace-poli import { mergeProjectWorkspaceRuntimeConfig, readProjectWorkspaceRuntimeConfig } from "./project-workspace-runtime-config.js"; import { resolveManagedProjectWorkspaceDir } from "../home-paths.js"; import { recordProjectLifecycleEvent, recordResourceCreationEvent } from "./resource-lifecycle-events.js"; +import { hasUnsettledProjectCharges } from "../modules/execution-admission/index.js"; type ProjectTransaction = Parameters[0]>[0]; @@ -1085,6 +1086,11 @@ export function projectService(db: Db, options: { captureWorkspaceUpdates?: bool const [retained] = await tx.select({ id: executionWorkspaces.id }).from(executionWorkspaces) .where(eq(executionWorkspaces.projectId, id)).limit(1); if (retained) throw conflict("This project still supplies access policy for retained task workspaces. Archive the project or explicitly clean up its workspaces before deleting it."); + if (await hasUnsettledProjectCharges(tx as unknown as Db, { companyId: project.companyId, projectId: id })) { + throw conflict("Task executions admitted to this project have not finished saving. Wait for them to settle before deleting the project.", { + code: "project_execution_admission_pending", + }); + } await tx.select({ id: projectWorkspaces.id }).from(projectWorkspaces).where(eq(projectWorkspaces.projectId, id)).for("update"); await recordProjectLifecycleEvent(tx as unknown as Db, project.companyId, id, "delete"); await notifyDeliveryWork(tx, DELIVERY_QUEUES.skillSourceSync);