fix(db): make instruction migrations safe to replay

Co-Authored-By: Paperclip <noreply@paperclip.ing>
This commit is contained in:
DottaandPaperclip committed 2026-09-27 18:58:49 -05:00
1 parent 69549d4248
commit 97df59cb3a
3 files changed
+89 -10

No files matched your search

@@ -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);
});
@@ -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");
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");
@@ -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");
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");