diff --git a/packages/db/src/agent-instruction-migration.test.ts b/packages/db/src/agent-instruction-migration.test.ts new file mode 100644 index 00000000000..18df341de91 --- /dev/null +++ b/packages/db/src/agent-instruction-migration.test.ts @@ -0,0 +1,65 @@ +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 { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./test-embedded-postgres.js"; + +const files = ["0285_legal_thunderbolts.sql", "0286_amusing_shiver_man.sql"]; +const migrations = await Promise.all(files.map(async file => ({ + file, sql: await readFile(new URL(`./migrations/${file}`, import.meta.url), "utf8"), +}))); +const support = await getEmbeddedPostgresTestSupport(); +const describePostgres = support.supported ? describe : describe.skip; + +describePostgres("instruction revision migrations", () => { + it("replays over existing revisions and pending copies without losing bytes or company constraints", async () => { + const database = await startEmbeddedPostgresTestDatabase("instruction-migration-replay-"); + const sql = postgres(database.connectionString, { max: 1, onnotice: () => {} }); + try { + const companyId = randomUUID(), otherCompanyId = randomUUID(), agentId = randomUUID(); + const runId = randomUUID(), revisionId = randomUUID(); + const bytes = Buffer.from("\uFEFF# Saved\r\nExact bytes ☃\n"); + const contentBase64 = bytes.toString("base64"); + const contentHash = createHash("sha256").update(bytes).digest("hex"); + await sql`INSERT INTO companies (id, name, issue_prefix) VALUES + (${companyId}, 'Instructions', 'INS'), (${otherCompanyId}, 'Other', 'OTH')`; + await sql`INSERT INTO agents (id, company_id, name) VALUES (${agentId}, ${companyId}, 'Writer')`; + await sql`INSERT INTO heartbeat_runs (id, company_id, agent_id) VALUES (${runId}, ${companyId}, ${agentId})`; + await sql`INSERT INTO agent_instruction_revisions + (id, company_id, agent_id, entry_file, content_base64, content_hash, byte_length, source, source_run_id) + VALUES (${revisionId}, ${companyId}, ${agentId}, 'AGENTS.md', ${contentBase64}, ${contentHash}, ${bytes.length}, 'cleanup', ${runId})`; + await sql`INSERT INTO agent_instruction_heads (company_id, agent_id, entry_file, revision_id) + VALUES (${companyId}, ${agentId}, 'AGENTS.md', ${revisionId})`; + await sql`INSERT INTO agent_instruction_working_copies + (run_id, company_id, agent_id, responsible_user_id, entry_file, base_revision_id, base_hash, local_root, execution_root, location, state, candidate_base64, candidate_hash) + VALUES (${runId}, ${companyId}, ${agentId}, 'editor', 'AGENTS.md', ${revisionId}, ${contentHash}, '/private/copy', '/private/copy', 'local', 'pending_commit', ${contentBase64}, ${contentHash})`; + const before = { + revisions: await sql`SELECT * FROM agent_instruction_revisions`, + heads: await sql`SELECT * FROM agent_instruction_heads`, + copies: await sql`SELECT * FROM agent_instruction_working_copies`, + }; + for (const migration of migrations) { + await sql`DELETE FROM drizzle.__drizzle_migrations WHERE hash = ${createHash("sha256").update(migration.sql).digest("hex")}`; + } + expect(await inspectMigrations(database.connectionString)).toMatchObject({ status: "needsMigrations", pendingMigrations: files }); + await applyPendingMigrations(database.connectionString); + // Also replay the SQL directly with every object already present. + await sql.begin(async tx => { + for (const migration of migrations) for (const statement of migration.sql.split("--> statement-breakpoint")) { + if (statement.trim()) await tx.unsafe(statement); + } + }); + expect(await sql`SELECT * FROM agent_instruction_revisions`).toEqual(before.revisions); + expect(await sql`SELECT * FROM agent_instruction_heads`).toEqual(before.heads); + expect(await sql`SELECT * FROM agent_instruction_working_copies`).toEqual(before.copies); + await expect(sql`UPDATE agent_instruction_heads SET company_id = ${otherCompanyId} WHERE revision_id = ${revisionId}`) + .rejects.toMatchObject({ code: "23503" }); + await expect(sql`UPDATE agent_instruction_working_copies SET company_id = ${otherCompanyId} WHERE run_id = ${runId}`) + .rejects.toMatchObject({ code: "23503" }); + } finally { + await sql.end(); + await database.cleanup(); + } + }, 30_000); +}); diff --git a/packages/db/src/migrations/0285_legal_thunderbolts.sql b/packages/db/src/migrations/0285_legal_thunderbolts.sql index 27522ed65a9..4835c499bdf 100644 --- a/packages/db/src/migrations/0285_legal_thunderbolts.sql +++ b/packages/db/src/migrations/0285_legal_thunderbolts.sql @@ -1,4 +1,5 @@ -CREATE TABLE "agent_instruction_heads" ( +-- Keep pre-release or renumbered migration replay safe without replacing stored instruction bytes. +CREATE TABLE IF NOT EXISTS "agent_instruction_heads" ( "company_id" uuid NOT NULL, "agent_id" uuid NOT NULL, "entry_file" text NOT NULL, @@ -7,7 +8,7 @@ CREATE TABLE "agent_instruction_heads" ( CONSTRAINT "agent_instruction_heads_identity_uq" UNIQUE("company_id","agent_id","entry_file") ); --> statement-breakpoint -CREATE TABLE "agent_instruction_revisions" ( +CREATE TABLE IF NOT EXISTS "agent_instruction_revisions" ( "id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL, "company_id" uuid NOT NULL, "agent_id" uuid NOT NULL, @@ -27,6 +28,12 @@ CREATE TABLE "agent_instruction_revisions" ( CONSTRAINT "agent_instruction_revisions_identity_uq" UNIQUE("company_id","agent_id","entry_file","id") ); --> statement-breakpoint -ALTER TABLE "agent_instruction_heads" ADD CONSTRAINT "agent_instruction_heads_company_id_agent_id_entry_file_revision_id_agent_instruction_revisions_company_id_agent_id_entry_file_id_fk" FOREIGN KEY ("company_id","agent_id","entry_file","revision_id") REFERENCES "public"."agent_instruction_revisions"("company_id","agent_id","entry_file","id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint -ALTER TABLE "agent_instruction_revisions" ADD CONSTRAINT "agent_instruction_revisions_company_id_agent_id_agents_company_id_id_fk" FOREIGN KEY ("company_id","agent_id") REFERENCES "public"."agents"("company_id","id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint -CREATE INDEX "agent_instruction_revisions_history_idx" ON "agent_instruction_revisions" USING btree ("company_id","agent_id","entry_file","created_at","id"); \ No newline at end of file +DO $$ BEGIN +ALTER TABLE "agent_instruction_heads" ADD CONSTRAINT "agent_instruction_heads_company_id_agent_id_entry_file_revision_id_agent_instruction_revisions_company_id_agent_id_entry_file_id_fk" FOREIGN KEY ("company_id","agent_id","entry_file","revision_id") REFERENCES "public"."agent_instruction_revisions"("company_id","agent_id","entry_file","id") ON DELETE cascade ON UPDATE no action; +EXCEPTION WHEN duplicate_object THEN NULL; +END $$;--> statement-breakpoint +DO $$ BEGIN +ALTER TABLE "agent_instruction_revisions" ADD CONSTRAINT "agent_instruction_revisions_company_id_agent_id_agents_company_id_id_fk" FOREIGN KEY ("company_id","agent_id") REFERENCES "public"."agents"("company_id","id") ON DELETE cascade ON UPDATE no action; +EXCEPTION WHEN duplicate_object THEN NULL; +END $$;--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "agent_instruction_revisions_history_idx" ON "agent_instruction_revisions" USING btree ("company_id","agent_id","entry_file","created_at","id"); diff --git a/packages/db/src/migrations/0286_amusing_shiver_man.sql b/packages/db/src/migrations/0286_amusing_shiver_man.sql index e50df4b3053..6bededa36fe 100644 --- a/packages/db/src/migrations/0286_amusing_shiver_man.sql +++ b/packages/db/src/migrations/0286_amusing_shiver_man.sql @@ -1,4 +1,5 @@ -CREATE TABLE "agent_instruction_working_copies" ( +-- Keep pre-release or renumbered migration replay safe without replacing stored instruction bytes. +CREATE TABLE IF NOT EXISTS "agent_instruction_working_copies" ( "run_id" uuid PRIMARY KEY NOT NULL, "company_id" uuid NOT NULL, "agent_id" uuid NOT NULL, @@ -22,7 +23,13 @@ CREATE TABLE "agent_instruction_working_copies" ( "updated_at" timestamp with time zone DEFAULT now() NOT NULL ); --> statement-breakpoint -ALTER TABLE "agent_instruction_working_copies" ADD CONSTRAINT "agent_instruction_working_copies_run_id_heartbeat_runs_id_fk" FOREIGN KEY ("run_id") REFERENCES "public"."heartbeat_runs"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint -ALTER TABLE "agent_instruction_working_copies" ADD CONSTRAINT "agent_instruction_working_copies_company_id_agent_id_agents_company_id_id_fk" FOREIGN KEY ("company_id","agent_id") REFERENCES "public"."agents"("company_id","id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint -CREATE INDEX "agent_instruction_copies_pending_idx" ON "agent_instruction_working_copies" USING btree ("state","next_attempt_at");--> statement-breakpoint -CREATE INDEX "agent_instruction_copies_agent_idx" ON "agent_instruction_working_copies" USING btree ("company_id","agent_id","created_at"); \ No newline at end of file +DO $$ BEGIN +ALTER TABLE "agent_instruction_working_copies" ADD CONSTRAINT "agent_instruction_working_copies_run_id_heartbeat_runs_id_fk" FOREIGN KEY ("run_id") REFERENCES "public"."heartbeat_runs"("id") ON DELETE cascade ON UPDATE no action; +EXCEPTION WHEN duplicate_object THEN NULL; +END $$;--> statement-breakpoint +DO $$ BEGIN +ALTER TABLE "agent_instruction_working_copies" ADD CONSTRAINT "agent_instruction_working_copies_company_id_agent_id_agents_company_id_id_fk" FOREIGN KEY ("company_id","agent_id") REFERENCES "public"."agents"("company_id","id") ON DELETE cascade ON UPDATE no action; +EXCEPTION WHEN duplicate_object THEN NULL; +END $$;--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "agent_instruction_copies_pending_idx" ON "agent_instruction_working_copies" USING btree ("state","next_attempt_at");--> statement-breakpoint +CREATE INDEX IF NOT EXISTS "agent_instruction_copies_agent_idx" ON "agent_instruction_working_copies" USING btree ("company_id","agent_id","created_at");