diff --git a/packages/chat/src/migrations.ts b/packages/chat/src/migrations.ts index ace9b7a3f..d9bae76be 100644 --- a/packages/chat/src/migrations.ts +++ b/packages/chat/src/migrations.ts @@ -414,6 +414,60 @@ export const chatMigrations: readonly ChatMigration[] = [ ADD COLUMN IF NOT EXISTS "sources_digest" text; `, }, + { + // CL-7130: `ensureRootThread`/`anchoredReplyThread` used to + // select-then-insert with no unique constraint backing the read, + // so concurrent first writers could each insert a thread for the + // same (workbench, kind='root') or (workbench, parent_message_id, + // kind='reply') key. Dedupe existing duplicates — keeping the + // oldest row per key — before the partial unique indexes below + // make a repeat impossible: every reference to a dropped + // duplicate's id (thread membership, message thread pointers, and + // any reply thread anchored under a dropped duplicate root) is + // repointed at the kept row first, so nothing is left dangling. + name: "0025_workbench_threads_unique_key", + sql: ` + CREATE TEMP TABLE "thread_dedupe_map" ON COMMIT DROP AS + SELECT "id" AS "drop_id", "keep_id" FROM ( + SELECT + "id", + first_value("id") OVER ( + PARTITION BY "tenant_id", "workbench_id", "kind", "parent_message_id" + ORDER BY "created_at", "id" + ) AS "keep_id" + FROM "chat"."workbench_threads" + WHERE "kind" IN ('root', 'reply') + ) "ranked" + WHERE "id" <> "keep_id"; + + UPDATE "chat"."workbench_thread_messages" "wtm" + SET "thread_id" = "m"."keep_id" + FROM "thread_dedupe_map" "m" + WHERE "wtm"."thread_id" = "m"."drop_id"; + + UPDATE "chat"."workbench_messages" "wm" + SET "thread_id" = "m"."keep_id" + FROM "thread_dedupe_map" "m" + WHERE "wm"."thread_id" = "m"."drop_id"; + + UPDATE "chat"."workbench_threads" "wt" + SET "parent_thread_id" = "m"."keep_id" + FROM "thread_dedupe_map" "m" + WHERE "wt"."parent_thread_id" = "m"."drop_id"; + + DELETE FROM "chat"."workbench_threads" "wt" + USING "thread_dedupe_map" "m" + WHERE "wt"."id" = "m"."drop_id"; + + CREATE UNIQUE INDEX IF NOT EXISTS "workbench_threads_root_key" + ON "chat"."workbench_threads" ("tenant_id", "workbench_id") + WHERE "kind" = 'root'; + + CREATE UNIQUE INDEX IF NOT EXISTS "workbench_threads_reply_key" + ON "chat"."workbench_threads" ("tenant_id", "workbench_id", "parent_message_id") + WHERE "kind" = 'reply'; + `, + }, ]; /** diff --git a/packages/chat/src/schema.ts b/packages/chat/src/schema.ts index a8883afc0..372cb12e1 100644 --- a/packages/chat/src/schema.ts +++ b/packages/chat/src/schema.ts @@ -6,6 +6,7 @@ // docs/package-migrations.md. `tenantId`/`principalId` are plain text // identifiers, not foreign keys, so referencing platform tenant/principal // ids works identically from a named schema. +import { sql } from "drizzle-orm"; import { index, integer, @@ -15,6 +16,7 @@ import { text, timestamp, unique, + uniqueIndex, } from "drizzle-orm/pg-core"; export const chatSchema = pgSchema("chat"); @@ -229,6 +231,12 @@ export const workbenchThreads = chatSchema.table( table.tenantId, table.workbenchId, ), + uniqueIndex("workbench_threads_root_key") + .on(table.tenantId, table.workbenchId) + .where(sql`${table.kind} = 'root'`), + uniqueIndex("workbench_threads_reply_key") + .on(table.tenantId, table.workbenchId, table.parentMessageId) + .where(sql`${table.kind} = 'reply'`), ], ); diff --git a/packages/chat/src/threads.ts b/packages/chat/src/threads.ts index 72d8f07d9..f865f9c1e 100644 --- a/packages/chat/src/threads.ts +++ b/packages/chat/src/threads.ts @@ -385,11 +385,11 @@ function mapThreadRow( export function createDrizzleThreadStore< TSchema extends Record, >(db: ThreadDb): ThreadStore { - async function ensureRootThread( + async function selectRootThread( tenantId: string, workbenchId: string, - ): Promise { - const existing = await db + ): Promise { + const rows = await db .select() .from(workbenchThreads) .where( @@ -399,8 +399,22 @@ export function createDrizzleThreadStore< eq(workbenchThreads.kind, "root"), ), ) + .orderBy(asc(workbenchThreads.createdAt)) .limit(1); - if (existing[0]) return mapThreadRow(existing[0]); + return rows[0] ? mapThreadRow(rows[0]) : undefined; + } + + async function ensureRootThread( + tenantId: string, + workbenchId: string, + ): Promise { + const existing = await selectRootThread(tenantId, workbenchId); + if (existing) return existing; + // Insert-first, not select-then-insert: two concurrent first + // writers both attempt the insert, the partial unique index + // (tenant_id, workbench_id) WHERE kind = 'root' serializes them, + // and the loser gets an empty `returning()` rather than a + // duplicate root thread — re-select picks up the winner's row. const id = newThreadId(); const inserted = await db .insert(workbenchThreads) @@ -414,8 +428,18 @@ export function createDrizzleThreadStore< runRef: null, title: null, }) + .onConflictDoNothing({ + target: [workbenchThreads.tenantId, workbenchThreads.workbenchId], + where: eq(workbenchThreads.kind, "root"), + }) .returning(); - return mapThreadRow(requireReturningRow(inserted, "root thread")); + const row = inserted[0]; + if (row) return mapThreadRow(row); + const reselected = await selectRootThread(tenantId, workbenchId); + if (!reselected) { + throw new Error("expected root thread row after conflicting insert"); + } + return reselected; } /** The thread a message currently lives in, or `root` if unassigned. */ @@ -451,23 +475,37 @@ export function createDrizzleThreadStore< return containerRows[0] ? mapThreadRow(containerRows[0]) : root; } - async function anchoredReplyThread( - input: OpenReplyThreadInput, - mode: "reply" | "fork", - ): Promise { - const existing = await db + async function selectReplyThread( + tenantId: string, + workbenchId: string, + parentMessageId: string, + ): Promise { + const rows = await db .select() .from(workbenchThreads) .where( and( - eq(workbenchThreads.tenantId, input.tenantId), - eq(workbenchThreads.workbenchId, input.workbenchId), + eq(workbenchThreads.tenantId, tenantId), + eq(workbenchThreads.workbenchId, workbenchId), eq(workbenchThreads.kind, "reply"), - eq(workbenchThreads.parentMessageId, input.parentMessageId), + eq(workbenchThreads.parentMessageId, parentMessageId), ), ) + .orderBy(asc(workbenchThreads.createdAt)) .limit(1); - if (existing[0]) return mapThreadRow(existing[0]); + return rows[0] ? mapThreadRow(rows[0]) : undefined; + } + + async function anchoredReplyThread( + input: OpenReplyThreadInput, + mode: "reply" | "fork", + ): Promise { + const existing = await selectReplyThread( + input.tenantId, + input.workbenchId, + input.parentMessageId, + ); + if (existing) return existing; const root = await ensureRootThread(input.tenantId, input.workbenchId); const container = await containerThreadFor( input.tenantId, @@ -477,6 +515,12 @@ export function createDrizzleThreadStore< ); const anchor = resolveThreadAnchor(root, container); if (anchor.blocked && mode === "reply") throw new ThreadDepthCapError(); + // Insert-first, not select-then-insert: two concurrent first + // repliers to the same message both attempt the insert, the + // partial unique index (tenant_id, workbench_id, + // parent_message_id) WHERE kind = 'reply' serializes them, and + // the loser's empty `returning()` re-selects the winner's row + // rather than creating a duplicate reply thread. const id = newThreadId(); const inserted = await db .insert(workbenchThreads) @@ -490,8 +534,26 @@ export function createDrizzleThreadStore< runRef: null, title: input.title ?? null, }) + .onConflictDoNothing({ + target: [ + workbenchThreads.tenantId, + workbenchThreads.workbenchId, + workbenchThreads.parentMessageId, + ], + where: eq(workbenchThreads.kind, "reply"), + }) .returning(); - return mapThreadRow(requireReturningRow(inserted, `${mode} thread`)); + const row = inserted[0]; + if (row) return mapThreadRow(row); + const reselected = await selectReplyThread( + input.tenantId, + input.workbenchId, + input.parentMessageId, + ); + if (!reselected) { + throw new Error(`expected ${mode} thread row after conflicting insert`); + } + return reselected; } return { diff --git a/packages/chat/test/migrations.test.ts b/packages/chat/test/migrations.test.ts index b386b0fbf..1df9b6426 100644 --- a/packages/chat/test/migrations.test.ts +++ b/packages/chat/test/migrations.test.ts @@ -45,6 +45,7 @@ const migrationNames = [ "0022_agent_turns", "0023_drop_workbench_host_arm", "0024_workbench_launch_sources_digest", + "0025_workbench_threads_unique_key", ]; describeIfDb("applyChatMigrations", () => { @@ -159,6 +160,19 @@ describeIfDb("applyChatMigrations", () => { expect(pinIndexes.map((row) => String(row["indexname"]))).toContain( "pinned_messages_workbench_idx", ); + + // CL-7130: `ensureRootThread`/`anchoredReplyThread` insert-then- + // reselect on conflict, and these partial unique indexes are + // what makes the conflict possible instead of a silent + // duplicate row. + const threadIndexes = await sql.unsafe( + `SELECT indexname FROM pg_indexes WHERE schemaname = 'chat' AND tablename = 'workbench_threads'`, + ); + const threadIndexNames = threadIndexes.map((row) => + String(row["indexname"]), + ); + expect(threadIndexNames).toContain("workbench_threads_root_key"); + expect(threadIndexNames).toContain("workbench_threads_reply_key"); } finally { await sql.end(); } @@ -213,3 +227,161 @@ describeIfDb("applyChatMigrations", () => { } }); }); + +describeIfDb("0025_workbench_threads_unique_key dedupe", () => { + const scratchUrl = scratchUrlFor( + databaseUrl ?? "postgres://localhost:5432/unused", + ).replace("_chat_migrations_test", "_chat_migrations_dedupe_test"); + const scratchTarget = new URL(scratchUrl); + const scratchDatabase = scratchTarget.pathname.replace(/^\//, ""); + + beforeAll(async () => { + const maintenanceUrl = new URL(scratchUrl); + maintenanceUrl.pathname = "/postgres"; + const maintenance = postgres(maintenanceUrl.toString(), { + max: 1, + onnotice: () => undefined, + }); + try { + await maintenance.unsafe(`DROP DATABASE IF EXISTS "${scratchDatabase}"`); + await maintenance.unsafe(`CREATE DATABASE "${scratchDatabase}"`); + } finally { + await maintenance.end(); + } + }, 20000); + + afterAll(async () => { + const maintenanceUrl = new URL(scratchUrl); + maintenanceUrl.pathname = "/postgres"; + const maintenance = postgres(maintenanceUrl.toString(), { + max: 1, + onnotice: () => undefined, + }); + try { + await maintenance.unsafe(`DROP DATABASE IF EXISTS "${scratchDatabase}"`); + } finally { + await maintenance.end(); + } + }, 20000); + + test("collapses duplicate root and reply threads, repointing thread_id and parent_thread_id references at the kept row", async () => { + // Replay every migration through 0023 by hand, recording each in + // the same ledger table `applyChatMigrations` reads, so the two + // duplicate roots below can be seeded *before* 0025 exists to + // forbid them — then let `applyChatMigrations` run 0025 for real. + const preDedupeMigrations = chatMigrations.filter( + (migration) => migration.name !== "0025_workbench_threads_unique_key", + ); + const seed = postgres(scratchUrl, { max: 1, onnotice: () => undefined }); + try { + await seed.unsafe(`CREATE SCHEMA IF NOT EXISTS "chat"`); + await seed.unsafe( + `CREATE TABLE IF NOT EXISTS "chat"."chat_migrations" ` + + `(name text PRIMARY KEY, applied_at timestamptz NOT NULL DEFAULT now())`, + ); + for (const migration of preDedupeMigrations) { + await seed.begin(async (tx) => { + await tx.unsafe(migration.sql); + await tx.unsafe( + `INSERT INTO "chat"."chat_migrations" (name) VALUES ($1)`, + [migration.name], + ); + }); + } + + await seed.unsafe(` + INSERT INTO "chat"."workbench_threads" + (id, tenant_id, workbench_id, kind, parent_message_id, created_at) + VALUES + ('thr_root_old', 'tnt_dedupe', 'wb_dedupe', 'root', NULL, now() - interval '1 hour'), + ('thr_root_new', 'tnt_dedupe', 'wb_dedupe', 'root', NULL, now()), + ('thr_reply_old', 'tnt_dedupe', 'wb_dedupe', 'reply', 'msg_parent', now() - interval '1 hour'), + ('thr_reply_new', 'tnt_dedupe', 'wb_dedupe', 'reply', 'msg_parent', now()) + `); + await seed.unsafe(` + INSERT INTO "chat"."workbench_thread_messages" + (tenant_id, workbench_id, thread_id, message_id) + VALUES ('tnt_dedupe', 'wb_dedupe', 'thr_root_new', 'msg_dedupe') + `); + // A message parked in the dropped reply duplicate — proves + // `workbench_messages.thread_id` is repointed, not just + // `workbench_thread_messages.thread_id`. + await seed.unsafe(` + INSERT INTO "chat"."workbench_messages" + (id, tenant_id, workbench_id, sender_address, thread_id, parts) + VALUES ( + 'msg_in_dropped_reply', 'tnt_dedupe', 'wb_dedupe', 'addr_1', + 'thr_reply_new', '[]'::jsonb + ) + `); + // A depth-2 thread anchored off the dropped root duplicate — + // proves `workbench_threads.parent_thread_id` is repointed too, + // not just message-facing references. + await seed.unsafe(` + INSERT INTO "chat"."workbench_threads" + (id, tenant_id, workbench_id, kind, parent_message_id, parent_thread_id, created_at) + VALUES ( + 'thr_child_of_dropped_root', 'tnt_dedupe', 'wb_dedupe', 'reply', + 'msg_child', 'thr_root_new', now() + ) + `); + } finally { + await seed.end(); + } + + const report = await applyChatMigrations(scratchUrl); + expect(report.applied).toEqual(["0025_workbench_threads_unique_key"]); + + const verify = postgres(scratchUrl, { max: 1, onnotice: () => undefined }); + try { + const roots = await verify.unsafe( + `SELECT id FROM "chat"."workbench_threads" ` + + `WHERE tenant_id = 'tnt_dedupe' AND workbench_id = 'wb_dedupe' AND kind = 'root'`, + ); + expect(roots.map((row) => String(row["id"]))).toEqual(["thr_root_old"]); + + const membership = await verify.unsafe( + `SELECT thread_id FROM "chat"."workbench_thread_messages" WHERE message_id = 'msg_dedupe'`, + ); + expect(String(membership[0]?.["thread_id"])).toBe("thr_root_old"); + + const replies = await verify.unsafe( + `SELECT id FROM "chat"."workbench_threads" ` + + `WHERE tenant_id = 'tnt_dedupe' AND workbench_id = 'wb_dedupe' ` + + `AND kind = 'reply' AND parent_message_id = 'msg_parent'`, + ); + expect(replies.map((row) => String(row["id"]))).toEqual([ + "thr_reply_old", + ]); + + const messageThreadId = await verify.unsafe( + `SELECT thread_id FROM "chat"."workbench_messages" WHERE id = 'msg_in_dropped_reply'`, + ); + expect(String(messageThreadId[0]?.["thread_id"])).toBe("thr_reply_old"); + + const childParentThreadId = await verify.unsafe( + `SELECT parent_thread_id FROM "chat"."workbench_threads" WHERE id = 'thr_child_of_dropped_root'`, + ); + expect(String(childParentThreadId[0]?.["parent_thread_id"])).toBe( + "thr_root_old", + ); + + // The reply unique index holds: a second reply row for the same + // (tenant, workbench, parent_message_id) key is now rejected + // rather than silently accepted as a second duplicate. + // postgres.js queries are lazy thenables, not Promises; wrap so + // `expect(...).rejects` actually runs the statement. + await expect( + (async () => { + await verify.unsafe( + `INSERT INTO "chat"."workbench_threads" ` + + `(id, tenant_id, workbench_id, kind, parent_message_id) ` + + `VALUES ('thr_reply_conflict', 'tnt_dedupe', 'wb_dedupe', 'reply', 'msg_parent')`, + ); + })(), + ).rejects.toThrow(); + } finally { + await verify.end(); + } + }, 120000); +}); diff --git a/packages/chat/test/threads.drizzle.test.ts b/packages/chat/test/threads.drizzle.test.ts new file mode 100644 index 000000000..2a3ed34cb --- /dev/null +++ b/packages/chat/test/threads.drizzle.test.ts @@ -0,0 +1,125 @@ +// DB-gated: skipped when no DATABASE_URL is reachable (a fresh +// checkout still runs the unit gates), mirroring +// `reactions.drizzle.test.ts`. Runs against its own scratch database. +// +// `threads.test.ts` proves `ensureRootThread`/`openReplyThread`'s +// idempotency against the in-memory store, which can never actually +// race (no `await` between its read and write). This exercises the +// real `createDrizzleThreadStore` path, where two concurrent first +// writers for the same root or reply key really do race at the +// database: proves the fix (insert with `onConflictDoNothing` backed +// by the partial unique index, then re-select on conflict — never +// select-then-insert) never throws a raw unique-violation and both +// callers converge on the same thread row (CL-7130). +import { afterAll, beforeAll, describe, expect, test } from "bun:test"; +import { drizzle } from "drizzle-orm/postgres-js"; +import postgres from "postgres"; + +import { e2eDatabaseUrl } from "../../../scripts/e2e/harness"; +import { applyChatMigrations } from "../src/migrations"; +import { createDrizzleThreadStore } from "../src/threads"; + +function scratchUrlFor(e2eUrl: string): string { + const url = new URL(e2eUrl); + const database = url.pathname.replace(/^\//, ""); + url.pathname = `/${database}_chat_threads_drizzle_test`; + return url.toString(); +} + +const databaseUrl = e2eDatabaseUrl(); +const describeIfDb = databaseUrl === undefined ? describe.skip : describe; + +const TENANT = "tnt_1"; +const WORKBENCH = "run_workbench1"; + +describeIfDb("createDrizzleThreadStore: concurrent first writers", () => { + const scratchUrl = scratchUrlFor( + databaseUrl ?? "postgres://localhost:5432/unused", + ); + const scratchTarget = new URL(scratchUrl); + const scratchDatabase = scratchTarget.pathname.replace(/^\//, ""); + + beforeAll(async () => { + const maintenanceUrl = new URL(scratchUrl); + maintenanceUrl.pathname = "/postgres"; + const maintenance = postgres(maintenanceUrl.toString(), { + max: 1, + onnotice: () => undefined, + }); + try { + await maintenance.unsafe(`DROP DATABASE IF EXISTS "${scratchDatabase}"`); + await maintenance.unsafe(`CREATE DATABASE "${scratchDatabase}"`); + } finally { + await maintenance.end(); + } + await applyChatMigrations(scratchUrl); + }); + + afterAll(async () => { + const maintenanceUrl = new URL(scratchUrl); + maintenanceUrl.pathname = "/postgres"; + const maintenance = postgres(maintenanceUrl.toString(), { + max: 1, + onnotice: () => undefined, + }); + try { + await maintenance.unsafe(`DROP DATABASE IF EXISTS "${scratchDatabase}"`); + } finally { + await maintenance.end(); + } + }); + + test("two concurrent ensureRootThread calls for the same workbench never throw a unique-violation and agree on one row", async () => { + // `max: 5` — a real connection pool, so the two calls below issue + // genuinely overlapping queries rather than being serialized onto + // one connection before either can race the other. + const sql = postgres(scratchUrl, { max: 5, onnotice: () => undefined }); + try { + const store = createDrizzleThreadStore(drizzle(sql)); + + const [first, second] = await Promise.all([ + store.ensureRootThread(TENANT, WORKBENCH), + store.ensureRootThread(TENANT, WORKBENCH), + ]); + + expect(first.id).toBe(second.id); + + const rows = await sql.unsafe( + `SELECT id FROM "chat"."workbench_threads" ` + + `WHERE tenant_id = $1 AND workbench_id = $2 AND kind = 'root'`, + [TENANT, WORKBENCH], + ); + expect(rows).toHaveLength(1); + } finally { + await sql.end(); + } + }); + + test("two concurrent openReplyThread calls for the same parent message never throw a unique-violation and agree on one row", async () => { + const sql = postgres(scratchUrl, { max: 5, onnotice: () => undefined }); + try { + const store = createDrizzleThreadStore(drizzle(sql)); + const input = { + tenantId: TENANT, + workbenchId: WORKBENCH, + parentMessageId: "msg_race", + }; + + const [first, second] = await Promise.all([ + store.openReplyThread(input), + store.openReplyThread(input), + ]); + + expect(first.id).toBe(second.id); + + const rows = await sql.unsafe( + `SELECT id FROM "chat"."workbench_threads" ` + + `WHERE tenant_id = $1 AND workbench_id = $2 AND kind = 'reply' AND parent_message_id = $3`, + [TENANT, WORKBENCH, "msg_race"], + ); + expect(rows).toHaveLength(1); + } finally { + await sql.end(); + } + }); +});