From d8fda075da4762642019549361b59763eb07a782 Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Thu, 6 Aug 2026 16:28:00 +0000 Subject: [PATCH] fix(webapp): stop an ordinary transcript write from rewriting a stored message MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `storeChatMessages` ended in `onConflictDoUpdate`, so `persistMessages` and `persistTurn` — which are handed a whole snapshot — treated any differing body under an existing message id as a deliberate finalisation. A stale snapshot carrying `wake:watch_1:fired`, the watch consent record, the deterministic confirmation or an investigation settlement card with a different body would overwrite the durable row that was already recorded. The proxy caps body size and metadata but does not rewrite message ids, so this was not an internal-bug-only exposure. The same clause updated only the `message` JSONB and never the `role` column, so `chat_messages.role` could end up disagreeing with `message.role` — and the UI reads one while the quota query reads the other. Ordinary transcript writes are now insert-only. Changing a stored message is its own operation, `finalizeChatMessage`, guarded on chat id, message id and role. `role` is verified rather than updated, and verified on both sides: the stored column must match `expectedRole` and so must the incoming body's own `role`, so the two cannot drift. A finalisation that matches nothing returns false; one whose body contradicts `expectedRole` throws. No production caller depended on the implicit finalisation. Every existing finalisation-shaped path already writes through an insert-only append: `settleInvestigationAndCloseCard`, `settleInvestigationStateAndCloseCard` and the watch request/confirmation/refusal records all use `appendChatMessageOnce(ByChatId)`. Also: re-sending a snapshot no longer reserves positions for messages that are already stored. The chat row is held, the missing ids are read under that lock, and only those get slots. A 40-message chat grown one turn at a time used to burn 1+2+…+40 = 820 slots for its 40 rows; it now burns 40. Deltas would be the proper fix, but that reaches into the agent's turn hooks and is a larger change than this pass. Two smaller repairs in the same file: `messageIdOf`/`messageRoleOf` now fail fast and name the chat and the offending message instead of casting unchecked and surfacing a `NOT NULL` violation from the driver; and a batch carrying the same message id twice throws instead of silently keeping the first, since that is an impossible state and a silent pick is how the upstream bug would stay invisible. The comment on `reserveMessagePositions` claiming the row lock is "released with the statement" was wrong — Postgres holds it to commit — and now says what is true. --- .../dashboardAgentTranscriptStore.test.ts | 253 +++++++++++++++++- .../dashboard-agent-db/README.md | 25 +- .../dashboard-agent-db/src/queries.ts | 125 +++++++-- 3 files changed, 369 insertions(+), 34 deletions(-) diff --git a/apps/webapp/test/dashboardAgentTranscriptStore.test.ts b/apps/webapp/test/dashboardAgentTranscriptStore.test.ts index e3b3feec5..10f516a9f 100644 --- a/apps/webapp/test/dashboardAgentTranscriptStore.test.ts +++ b/apps/webapp/test/dashboardAgentTranscriptStore.test.ts @@ -3,6 +3,7 @@ import { countUserMessages, createChat, createDashboardAgentDb, + finalizeChatMessage, getChatMessages, getInvestigation, investigationSettlementMessageId, @@ -92,6 +93,35 @@ async function rows(prisma: PrismaClient, chatId: string): Promise ); } +/** The position allocator itself: what a wasted reservation is visible in. */ +async function nextPosition(prisma: PrismaClient, chatId: string): Promise { + const rows = await prisma.$queryRawUnsafe<{ next_message_position: number }[]>( + `select next_message_position from trigger_dashboard_agent.chats where id = $1`, + chatId + ); + return rows[0]!.next_message_position; +} + +async function chatStamps( + prisma: PrismaClient, + chatId: string +): Promise<{ last_message_at: Date | null; updated_at: Date }[]> { + return prisma.$queryRawUnsafe( + `select last_message_at, updated_at from trigger_dashboard_agent.chats where id = $1`, + chatId + ); +} + +/** The structural column, which the JSONB payload must never be able to contradict. */ +async function roleOf(prisma: PrismaClient, chatId: string, messageId: string): Promise { + const rows = await prisma.$queryRawUnsafe<{ role: string }[]>( + `select role from trigger_dashboard_agent.chat_messages where chat_id = $1 and message_id = $2`, + chatId, + messageId + ); + return rows[0]!.role; +} + function openState(): InvestigationState { return investigationStateSchema.parse({ outcome: "in_progress", @@ -231,9 +261,94 @@ describe("invariant 2: concurrent different messages get distinct positions", () ); }); -describe("invariant 3: a controlled update changes the body and nothing else", () => { +describe("invariant 3: an ordinary transcript write can never change a stored message", () => { postgresTest( - "finalising a message keeps its id, its position and its row", + "a differing body under an existing id leaves the durable row exactly as it was", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_no_implicit_update"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + await persistMessages(agentDb, { chatId, messages: [textMessage("u1")] }); + // A durable event: the wake that actually fired. + await appendChatMessageOnceByChatId(agentDb, { + chatId, + message: textMessage("wake:watch_1:fired", "The watch on send-order-receipt resolved."), + }); + const before = await rows(prisma, chatId); + + // A stale snapshot carrying the same id with a different body. `persistMessages` is + // not a finalisation, so it must not be able to rewrite it. + await persistMessages(agentDb, { + chatId, + messages: [textMessage("u1"), textMessage("wake:watch_1:fired", "something else entirely")], + }); + + expect(await rows(prisma, chatId)).toEqual(before); + }, + 30_000 + ); + + postgresTest( + "persistTurn cannot rewrite a stored message either", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_no_implicit_update_turn"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + await persistMessages(agentDb, { chatId, messages: [textMessage("a1", "the answer")] }); + const before = await rows(prisma, chatId); + + await persistTurn(agentDb, { + chatId, + messages: [textMessage("a1", "a different answer")], + session: { publicAccessToken: "pat_store" }, + }); + + expect(await rows(prisma, chatId)).toEqual(before); + }, + 30_000 + ); + + postgresTest( + "a batch carrying the same id twice is refused rather than silently picking one", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_dup_in_batch"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + await expect( + persistMessages(agentDb, { + chatId, + messages: [textMessage("a1", "first"), textMessage("a1", "second")], + }) + ).rejects.toThrow(/message id a1 twice in one batch/); + + // And nothing landed: the throw is before any reservation. + expect(await rows(prisma, chatId)).toHaveLength(0); + expect(await nextPosition(prisma, chatId)).toBe(1); + }, + 30_000 + ); + + postgresTest( + "a message with no id or no role is refused by name, not by a NOT NULL violation", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_malformed"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + await expect( + persistMessages(agentDb, { chatId, messages: [{ role: "user", parts: [] }] }) + ).rejects.toThrow(/Chat chat_malformed was handed a message with no id: .*"role":"user"/); + + await expect( + persistMessages(agentDb, { chatId, messages: [{ id: "a1", parts: [] }] }) + ).rejects.toThrow(/Chat chat_malformed was handed a message with no role: .*"id":"a1"/); + }, + 30_000 + ); +}); + +describe("invariant 4: a controlled finalisation changes the body and nothing else", () => { + postgresTest( + "finalising a message keeps its id, its position and its role", async ({ prisma, postgresContainer }) => { const chatId = "chat_finalise"; await boot(prisma, postgresContainer.getConnectionUri(), chatId); @@ -244,19 +359,143 @@ describe("invariant 3: a controlled update changes the body and nothing else", ( }); const before = await rows(prisma, chatId); - await persistMessages(agentDb, { - chatId, - messages: [textMessage("u1"), textMessage("a1", "here is the answer")], - }); + expect( + await finalizeChatMessage(agentDb, { + chatId, + messageId: "a1", + expectedRole: "assistant", + message: textMessage("a1", "here is the answer"), + }) + ).toBe(true); const after = await rows(prisma, chatId); expect(after).toHaveLength(2); expect(after.map((row) => [row.message_id, row.position])).toEqual( before.map((row) => [row.message_id, row.position]) ); - // Only the one message that changed changed. + // Only the one message named changed. expect(after[0]!.message).toEqual(before[0]!.message); expect(after[1]!.message).toMatchObject({ parts: [{ text: "here is the answer" }] }); + expect(await roleOf(prisma, chatId, "a1")).toBe("assistant"); + }, + 30_000 + ); + + postgresTest( + "a finalisation aimed at the wrong role writes nothing", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_finalise_role"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + await persistMessages(agentDb, { chatId, messages: [textMessage("a1", "still working")] }); + const before = await rows(prisma, chatId); + + // The stored row is an assistant message, so a user finalisation is not its own. + expect( + await finalizeChatMessage(agentDb, { + chatId, + messageId: "a1", + expectedRole: "user", + message: { id: "a1", role: "user", parts: [{ type: "text", text: "hijacked" }] }, + }) + ).toBe(false); + + expect(await rows(prisma, chatId)).toEqual(before); + expect(await roleOf(prisma, chatId, "a1")).toBe("assistant"); + }, + 30_000 + ); + + postgresTest( + "the row's role and the body's role cannot be made to disagree", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_finalise_drift"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + await persistMessages(agentDb, { chatId, messages: [textMessage("a1")] }); + + // The column says assistant, the body would say user. Refused outright rather + // than stored as a row whose column and payload disagree. + await expect( + finalizeChatMessage(agentDb, { + chatId, + messageId: "a1", + expectedRole: "assistant", + message: { id: "a1", role: "user", parts: [] }, + }) + ).rejects.toThrow(/expected role assistant but its body carries user/); + + expect(await roleOf(prisma, chatId, "a1")).toBe("assistant"); + expect((await rows(prisma, chatId))[0]!.message).toMatchObject({ role: "assistant" }); + }, + 30_000 + ); + + postgresTest( + "a finalisation of a message that isn't there writes nothing", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_finalise_missing"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + expect( + await finalizeChatMessage(agentDb, { + chatId, + messageId: "never-stored", + expectedRole: "assistant", + message: textMessage("never-stored"), + }) + ).toBe(false); + expect(await rows(prisma, chatId)).toHaveLength(0); + }, + 30_000 + ); +}); + +describe("invariant 5: re-sending a snapshot is free", () => { + postgresTest( + "a re-sent snapshot reserves no position, touches no row and writes no timestamp", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_snapshot_free"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + const snapshot = Array.from({ length: 6 }, (_, i) => textMessage(`m${i}`)); + await persistMessages(agentDb, { chatId, messages: snapshot }); + + const before = await rows(prisma, chatId); + const positionBefore = await nextPosition(prisma, chatId); + const chatBefore = await chatStamps(prisma, chatId); + + await persistMessages(agentDb, { chatId, messages: snapshot }); + await persistTurn(agentDb, { + chatId, + messages: snapshot, + session: { publicAccessToken: "pat_store" }, + }); + + expect(await rows(prisma, chatId)).toEqual(before); + // The allocator is the observable cost: a re-send that reserved slots would grow it. + expect(await nextPosition(prisma, chatId)).toBe(positionBefore); + expect(await chatStamps(prisma, chatId)).toEqual(chatBefore); + }, + 30_000 + ); + + postgresTest( + "a transcript grown by re-sent snapshots spends one position per message", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_snapshot_slots"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + // The real write pattern: every turn hands over the whole transcript again. With + // the old insert-everything path this cost 1+2+…+40 = 820 slots for 40 rows. + const snapshot: ReturnType[] = []; + for (let i = 0; i < 40; i++) { + snapshot.push(textMessage(`m${i}`)); + await persistMessages(agentDb, { chatId, messages: [...snapshot] }); + } + + expect(await rows(prisma, chatId)).toHaveLength(40); + expect(await nextPosition(prisma, chatId)).toBe(41); }, 30_000 ); diff --git a/internal-packages/dashboard-agent-db/README.md b/internal-packages/dashboard-agent-db/README.md index f58688a30..df0a4cf68 100644 --- a/internal-packages/dashboard-agent-db/README.md +++ b/internal-packages/dashboard-agent-db/README.md @@ -27,10 +27,27 @@ of truth. ## Tables -- `chats` — one row per conversation: org/user scope, title, a `messages` JSONB - display copy of the transcript, and `metadata` (the project/env context the chat - ran in). Soft-deleted via `deleted_at`, pinned via `pinned_at`, read-marked via - `last_read_at` (NULL = never read, so every watch wake in it counts as unread). +- `chats` — one row per conversation: org/user scope, title, `metadata` (the + project/env context the chat ran in), and `next_message_position`, the allocator the + transcript's ordering comes from. No transcript of its own. Soft-deleted via + `deleted_at`, pinned via `pinned_at`, read-marked via `last_read_at` (NULL = never + read, so every watch wake in it counts as unread). +- `chat_messages` — the transcript, one row per message. Identity is + `(chat_id, message_id)` and order is `position`, unique per chat and reserved from + `chats.next_message_position` by the same single statement that reads it, so + concurrent writers get disjoint contiguous ranges. `role` is lifted out of the + payload so the message-quota count is an index scan. + + Three write modes, and only the third may change a message the chat already holds: + a new message is a plain insert; a redelivered durable event (a watch wake, a + settlement card) is `ON CONFLICT DO NOTHING` on `(chat_id, message_id)`, so it + leaves the recorded row untouched; a deliberate finalisation is + `finalizeChatMessage`, which rewrites one body under a verified `role` and never + moves the id or the position. So re-sending a whole turn snapshot is a no-op. + + Positions are monotonic, not gapless: a reservation whose insert then conflicts, + or a batch that rolls back, leaves the slot unused. Only the relative order + matters, so a gap is expected and harmless. - `chat_sessions` — live transport state keyed by `chat_id`: the session-scoped `public_access_token` and `last_event_id` for resume. Separate table so the secret token is isolated from list queries and the hot per-turn write stays off diff --git a/internal-packages/dashboard-agent-db/src/queries.ts b/internal-packages/dashboard-agent-db/src/queries.ts index 82ad5cbbd..27e168f2b 100644 --- a/internal-packages/dashboard-agent-db/src/queries.ts +++ b/internal-packages/dashboard-agent-db/src/queries.ts @@ -3,7 +3,7 @@ import { VIEW_BLOCK_VERSION, WATCH_REQUEST_MESSAGE_ID_PREFIX, } from "@internal/dashboard-agent-contracts"; -import { and, desc, eq, ne, notLike, sql, isNull, type SQL } from "drizzle-orm"; +import { and, desc, eq, inArray, ne, notLike, sql, isNull, type SQL } from "drizzle-orm"; import type { DashboardAgentDb } from "./client.js"; import { generateInvestigationId } from "./ids.js"; import { lockChatForWatches, type DashboardAgentDbOrTx } from "./internal.js"; @@ -280,22 +280,41 @@ export async function softDeleteChat( }); } -/** Row identity. No fallback: `message_id` is `NOT NULL`, so the database rejects a message with no id. */ -function messageIdOf(message: unknown): string { - return (message as { id: string }).id; +/** Enough of the payload to recognise it in an error, without logging a whole transcript. */ +function describeMessage(message: unknown): string { + try { + return JSON.stringify(message)?.slice(0, 200) ?? String(message); + } catch { + return String(message); + } +} + +/** Row identity. Checked here so a malformed message names itself, not a `NOT NULL` violation. */ +function messageIdOf(chatId: string, message: unknown): string { + const id = (message as { id?: unknown } | null | undefined)?.id; + if (typeof id !== "string" || id.length === 0) { + throw new Error(`Chat ${chatId} was handed a message with no id: ${describeMessage(message)}`); + } + return id; } /** Lifted out of the payload so the quota count never opens the JSONB. */ -function messageRoleOf(message: unknown): string { - return (message as { role: string }).role; +function messageRoleOf(chatId: string, message: unknown): string { + const role = (message as { role?: unknown } | null | undefined)?.role; + if (typeof role !== "string" || role.length === 0) { + throw new Error( + `Chat ${chatId} was handed a message with no role: ${describeMessage(message)}` + ); + } + return role; } /** * Reserve `count` contiguous positions on the chat, or null when there is no such chat. * - * One statement, so two writers are handed disjoint ranges and the row lock is released - * with the statement rather than held across a round trip. `scope` is the caller's - * tenancy check, applied here because this is the statement that has the chat row. + * One statement, so two writers are handed disjoint ranges and the row lock it takes is + * held for the rest of this short transaction rather than across a round trip. `scope` is + * the caller's tenancy check, applied here because this is the statement that has the row. */ async function reserveMessagePositions( tx: DashboardAgentDbOrTx, @@ -316,13 +335,16 @@ async function reserveMessagePositions( } /** - * Store a batch of messages: insert the ones that aren't there, and update one that is - * only when its content genuinely changed. Nothing is ever deleted and no message - * outside `messages` is touched, so a wake or a terminal card that landed mid-turn - * cannot be written away — the property the old read-modify-write merge only approximated. + * Store a batch of messages, insert-only. A message id already in the chat is left exactly + * as it was recorded — body, position and role — so a stale snapshot cannot overwrite a + * durable event, and nothing outside `messages` is touched either. Changing a message that + * is already stored is a different operation: {@link finalizeChatMessage}. * * The batch keeps its incoming order, and a message already stored keeps the position it * was first given, which is why a mid-turn append sits before the turn's later messages. + * + * The already-stored ids are dropped before any position is reserved, so re-sending a whole + * snapshot costs one slot per genuinely new message instead of one per message sent. */ async function storeChatMessages( tx: DashboardAgentDbOrTx, @@ -330,11 +352,36 @@ async function storeChatMessages( ): Promise { const deduped = new Map(); for (const message of params.messages) { - const id = messageIdOf(message); - if (!deduped.has(id)) deduped.set(id, message); + const id = messageIdOf(params.chatId, message); + if (deduped.has(id)) { + throw new Error(`Chat ${params.chatId} was handed message id ${id} twice in one batch`); + } + deduped.set(id, message); } if (deduped.size === 0) return; + // Hold the chat row before reading which ids are missing, so a concurrent batch can't + // reserve a slot for a message this one is about to insert. + const locked = await tx + .select({ id: chats.id }) + .from(chats) + .where(and(eq(chats.id, params.chatId), isNull(chats.deletedAt))) + .limit(1) + .for("update"); + if (locked.length === 0) return; + + const stored = await tx + .select({ messageId: chatMessages.messageId }) + .from(chatMessages) + .where( + and( + eq(chatMessages.chatId, params.chatId), + inArray(chatMessages.messageId, [...deduped.keys()]) + ) + ); + for (const row of stored) deduped.delete(row.messageId); + if (deduped.size === 0) return; + const start = await reserveMessagePositions(tx, { chatId: params.chatId, count: deduped.size }); if (start === null) return; @@ -345,15 +392,47 @@ async function storeChatMessages( chatId: params.chatId, messageId, position: start + offset, - role: messageRoleOf(message), + role: messageRoleOf(params.chatId, message), message, })) ) - .onConflictDoUpdate({ - target: [chatMessages.chatId, chatMessages.messageId], - set: { message: sql`excluded.message` }, - setWhere: sql`${chatMessages.message} is distinct from excluded.message`, - }); + .onConflictDoNothing({ target: [chatMessages.chatId, chatMessages.messageId] }); +} + +/** + * Rewrite one already-stored message's body, deliberately. The only operation that may + * change a message the transcript already holds; every other write is insert-only. + * + * `expectedRole` is verified rather than updated, on both sides: the stored row's `role` + * column has to match, and so does the incoming body's own `role`. So the structural + * column and the payload cannot drift apart, and a finalisation aimed at the wrong + * message — or at a message some other lane has since replaced — writes nothing and says + * so by returning false. Position and id are never touched. + */ +export async function finalizeChatMessage( + db: DashboardAgentDbOrTx, + params: { chatId: string; messageId: string; expectedRole: string; message: unknown } +): Promise { + const role = messageRoleOf(params.chatId, params.message); + if (role !== params.expectedRole) { + throw new Error( + `Chat ${params.chatId} finalisation of ${params.messageId} expected role ${params.expectedRole} but its body carries ${role}` + ); + } + + const rows = await db + .update(chatMessages) + .set({ message: params.message }) + .where( + and( + eq(chatMessages.chatId, params.chatId), + eq(chatMessages.messageId, params.messageId), + eq(chatMessages.role, params.expectedRole) + ) + ) + .returning({ messageId: chatMessages.messageId }); + + return rows.length > 0; } /** No session state, unlike {@link persistTurn}. */ @@ -376,7 +455,7 @@ async function appendOneMessage( db: DashboardAgentDbOrTx, params: { chatId: string; message: unknown; scope: SQL[] } ): Promise { - const messageId = messageIdOf(params.message); + const messageId = messageIdOf(params.chatId, params.message); const rows = await db.execute<{ message_id: string }>(sql` with reserved as ( update ${chats} @@ -393,7 +472,7 @@ async function appendOneMessage( returning "next_message_position" - 1 as "position" ) insert into ${chatMessages} ("chat_id", "message_id", "position", "role", "message") - select ${params.chatId}, ${messageId}, reserved."position", ${messageRoleOf(params.message)}, + select ${params.chatId}, ${messageId}, reserved."position", ${messageRoleOf(params.chatId, params.message)}, ${JSON.stringify(params.message)}::jsonb from reserved on conflict ("chat_id", "message_id") do nothing