Skip to content

Commit f03dfd9

Browse files
kathiekiwikathiekiwi
committed
fix(dashboard-agent): normalize message JSON to well-formed UTF-16 at persist
1 parent ac029e5 commit f03dfd9

4 files changed

Lines changed: 249 additions & 10 deletions

File tree

‎apps/webapp/test/dashboardAgentSurrogatePersist.test.ts‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
import {
2+
appendChatMessageOnce,
3+
createChat,
4+
createDashboardAgentDb,
5+
getChatMessages,
6+
persistMessages,
7+
persistTurn,
8+
type DashboardAgentDb,
9+
type DashboardAgentDbClient,
10+
} from "@internal/dashboard-agent-db";
11+
import { postgresTest } from "@internal/testcontainers";
12+
import type { PrismaClient } from "@trigger.dev/database";
13+
import { readdirSync, readFileSync } from "node:fs";
14+
import path from "node:path";
15+
import { afterEach, describe, expect } from "vitest";
16+
17+
/**
18+
* jsonb rejects a lone UTF-16 surrogate, and a message body carries strings we never
19+
* authored — tool inputs, filenames, urls — from transports that don't pass the webapp
20+
* routes. `storeChatMessages` and `appendOneMessage` are where they have to be made storable.
21+
*/
22+
23+
let agentDb: DashboardAgentDb;
24+
let agentDbClient: DashboardAgentDbClient | undefined;
25+
26+
const MIGRATIONS = path.resolve(__dirname, "../../../internal-packages/dashboard-agent-db/drizzle");
27+
28+
async function applyAgentSchema(prisma: PrismaClient) {
29+
for (const name of readdirSync(MIGRATIONS)
30+
.filter((file) => file.endsWith(".sql"))
31+
.sort()) {
32+
const sql = readFileSync(path.join(MIGRATIONS, name), "utf8");
33+
for (const statement of sql.split("--> statement-breakpoint")) {
34+
const trimmed = statement.trim();
35+
if (trimmed.length > 0) await prisma.$executeRawUnsafe(trimmed);
36+
}
37+
}
38+
}
39+
40+
const ORG = "org_surrogate";
41+
const USER = "user_surrogate";
42+
43+
afterEach(async () => {
44+
await agentDbClient?.close();
45+
agentDbClient = undefined;
46+
});
47+
48+
async function boot(prisma: PrismaClient, connectionUri: string, chatId: string) {
49+
await applyAgentSchema(prisma);
50+
agentDbClient = createDashboardAgentDb(connectionUri, { max: 4 });
51+
agentDb = agentDbClient.db;
52+
await createChat(agentDb, { id: chatId, organizationId: ORG, userId: USER });
53+
}
54+
55+
async function transcript(chatId: string): Promise<unknown[]> {
56+
return getChatMessages(agentDb, { chatId, organizationId: ORG, userId: USER });
57+
}
58+
59+
describe("a lone surrogate anywhere in a message body is storable", () => {
60+
postgresTest(
61+
"persists a tool input carrying a lone surrogate",
62+
async ({ prisma, postgresContainer }) => {
63+
const chatId = "chat_surrogate";
64+
await boot(prisma, postgresContainer.getConnectionUri(), chatId);
65+
66+
await persistMessages(agentDb, {
67+
chatId,
68+
messages: [
69+
{
70+
id: "u1",
71+
role: "user",
72+
parts: [
73+
{
74+
type: "tool-search_docs",
75+
state: "input-available",
76+
toolCallId: "u1_call",
77+
input: { query: "how do i \ud83d", filename: "\udc00.png" },
78+
},
79+
],
80+
},
81+
],
82+
});
83+
84+
const stored = (await transcript(chatId)) as {
85+
parts: { input: { query: string; filename: string } }[];
86+
}[];
87+
88+
expect(stored).toHaveLength(1);
89+
expect(stored[0]!.parts[0]!.input.query).toBe("how do i �");
90+
expect(stored[0]!.parts[0]!.input.filename).toBe("�.png");
91+
},
92+
30_000
93+
);
94+
95+
postgresTest(
96+
"finalises a tool output carrying a lone surrogate",
97+
async ({ prisma, postgresContainer }) => {
98+
const chatId = "chat_surrogate_turn";
99+
await boot(prisma, postgresContainer.getConnectionUri(), chatId);
100+
101+
const call = (state: string, extra: Record<string, unknown>) => ({
102+
id: "a1",
103+
role: "assistant",
104+
parts: [{ type: "tool-search_docs", state, toolCallId: "a1_call", input: {}, ...extra }],
105+
});
106+
107+
// Stored mid-flight by `onTurnStart`, then rewritten in place when the turn completes.
108+
await persistMessages(agentDb, { chatId, messages: [call("input-available", {})] });
109+
await persistTurn(agentDb, {
110+
chatId,
111+
messages: [call("output-available", { output: { text: "the page says \ud83d" } })],
112+
finalizeMessageIds: ["a1"],
113+
session: { publicAccessToken: "pat", lastEventId: "1", runId: "run" },
114+
});
115+
116+
const stored = (await transcript(chatId)) as {
117+
parts: { state: string; output: { text: string } }[];
118+
}[];
119+
120+
expect(stored).toHaveLength(1);
121+
expect(stored[0]!.parts[0]!.state).toBe("output-available");
122+
expect(stored[0]!.parts[0]!.output.text).toBe("the page says �");
123+
},
124+
30_000
125+
);
126+
127+
postgresTest(
128+
"appends a wake message carrying a lone surrogate",
129+
async ({ prisma, postgresContainer }) => {
130+
const chatId = "chat_surrogate_append";
131+
await boot(prisma, postgresContainer.getConnectionUri(), chatId);
132+
133+
const appended = await appendChatMessageOnce(agentDb, {
134+
chatId,
135+
userId: USER,
136+
organizationId: ORG,
137+
message: {
138+
id: "w1",
139+
role: "assistant",
140+
parts: [{ type: "text", text: "the queue \udc00 backed up" }],
141+
} as { id: string; role: string },
142+
});
143+
144+
expect(appended).toBe(true);
145+
const stored = (await transcript(chatId)) as { parts: { text: string }[] }[];
146+
expect(stored[0]!.parts[0]!.text).toBe("the queue � backed up");
147+
},
148+
30_000
149+
);
150+
});

‎internal-packages/dashboard-agent-contracts/src/well-formed.test.ts‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { describe, expect, it } from "vitest";
2-
import { sliceWellFormed } from "./well-formed.js";
2+
import { sliceWellFormed, toWellFormedDeep } from "./well-formed.js";
33

44
const emoji = "😀"; // one surrogate pair
55

@@ -35,3 +35,43 @@ describe("sliceWellFormed", () => {
3535
expect(sliceWellFormed(`${emoji}x${emoji}yz`, 5)).toBe(`${emoji}x${emoji}`);
3636
});
3737
});
38+
39+
describe("toWellFormedDeep", () => {
40+
it("replaces a lone surrogate nested in a tool input", () => {
41+
const message = {
42+
id: "msg_1",
43+
role: "assistant",
44+
parts: [{ type: "tool-search", input: { query: "cat \ud83d", limit: 5 } }],
45+
};
46+
const result = toWellFormedDeep(message);
47+
expect(result.parts[0].input.query).toBe("cat �");
48+
expect(result.parts[0].input.limit).toBe(5);
49+
expect(result).not.toBe(message);
50+
});
51+
52+
it("returns the same reference when nothing changed", () => {
53+
const message = { id: "msg_1", parts: [{ text: `hello ${emoji}` }], meta: null };
54+
expect(toWellFormedDeep(message)).toBe(message);
55+
});
56+
57+
it("replaces a lone surrogate in a key", () => {
58+
const result: Record<string, unknown> = toWellFormedDeep({ "k\ud800": "v" });
59+
expect(Object.keys(result)).toEqual(["k�"]);
60+
expect(result["k�"]).toBe("v");
61+
});
62+
63+
it("walks arrays", () => {
64+
const messages = [{ text: "ok" }, { text: "\udc00bad" }];
65+
const result = toWellFormedDeep(messages);
66+
expect(result[1].text).toBe("�bad");
67+
expect(result[0]).toBe(messages[0]);
68+
});
69+
70+
it("leaves non-string primitives and non-plain objects alone", () => {
71+
const date = new Date(0);
72+
const value = { n: 1, b: true, nil: null, u: undefined, date };
73+
const result = toWellFormedDeep(value);
74+
expect(result).toBe(value);
75+
expect(result.date).toBe(date);
76+
});
77+
});

‎internal-packages/dashboard-agent-contracts/src/well-formed.ts‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,3 +7,47 @@ export function sliceWellFormed(s: string, n: number): string {
77
if (last >= 0xd800 && last <= 0xdbff) return cut.slice(0, -1);
88
return cut;
99
}
10+
11+
// `toWellFormed` is ES2024; this package targets ES2022.
12+
interface WellFormable {
13+
toWellFormed(): string;
14+
}
15+
16+
function wellFormed(s: string): string {
17+
return (s as unknown as WellFormable).toWellFormed();
18+
}
19+
20+
function isPlainObject(value: object): boolean {
21+
const proto = Object.getPrototypeOf(value);
22+
return proto === Object.prototype || proto === null;
23+
}
24+
25+
/**
26+
* Every string in a JSON-ish value made well-formed, keys included, so a lone surrogate
27+
* anywhere — tool input, filename, url — can't reach jsonb. Anything unchanged is returned
28+
* as it was, and anything with a custom prototype (a Date, a class instance) is untouched.
29+
*/
30+
export function toWellFormedDeep<T>(value: T): T {
31+
if (typeof value === "string") return wellFormed(value) as T;
32+
if (Array.isArray(value)) {
33+
let changed = false;
34+
const next = value.map((item) => {
35+
const fixed = toWellFormedDeep(item);
36+
if (fixed !== item) changed = true;
37+
return fixed;
38+
});
39+
return (changed ? next : value) as T;
40+
}
41+
if (value !== null && typeof value === "object" && isPlainObject(value)) {
42+
let changed = false;
43+
const next: Record<string, unknown> = {};
44+
for (const [key, item] of Object.entries(value)) {
45+
const fixedKey = wellFormed(key);
46+
const fixed = toWellFormedDeep(item);
47+
if (fixed !== item || fixedKey !== key) changed = true;
48+
next[fixedKey] = fixed;
49+
}
50+
return (changed ? next : value) as T;
51+
}
52+
return value;
53+
}

‎internal-packages/dashboard-agent-db/src/queries.ts‎

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import {
22
investigationBlockSchema,
3+
toWellFormedDeep,
34
VIEW_BLOCK_VERSION,
45
WATCH_REQUEST_MESSAGE_ID_PREFIX,
56
} from "@internal/dashboard-agent-contracts";
@@ -256,7 +257,7 @@ export async function createChat(
256257
organizationId: params.organizationId,
257258
userId: params.userId,
258259
title: params.title ?? DEFAULT_CHAT_TITLE,
259-
metadata: params.metadata ?? {},
260+
metadata: toWellFormedDeep(params.metadata ?? {}),
260261
})
261262
.onConflictDoNothing();
262263
}
@@ -446,8 +447,9 @@ async function storeChatMessages(
446447
tx: DashboardAgentDbOrTx,
447448
params: { chatId: string; messages: unknown[]; finalizable?: ReadonlySet<string> }
448449
): Promise<void> {
450+
// Every batch write lands here, so this is where a lone surrogate stops before jsonb.
449451
const deduped = new Map<string, unknown>();
450-
for (const message of params.messages) {
452+
for (const message of toWellFormedDeep(params.messages)) {
451453
const id = messageIdOf(params.chatId, message);
452454
if (deduped.has(id)) {
453455
throw new Error(`Chat ${params.chatId} was handed message id ${id} twice in one batch`);
@@ -568,7 +570,9 @@ async function appendOneMessage(
568570
db: DashboardAgentDbOrTx,
569571
params: { chatId: string; message: unknown; scope: SQL[] }
570572
): Promise<boolean> {
571-
const messageId = messageIdOf(params.chatId, params.message);
573+
// Single-message appends land here — normalize like storeChatMessages.
574+
const message = toWellFormedDeep(params.message);
575+
const messageId = messageIdOf(params.chatId, message);
572576
const rows = await db.execute<{ message_id: string }>(sql`
573577
with reserved as (
574578
update ${chats}
@@ -585,8 +589,8 @@ async function appendOneMessage(
585589
returning "next_message_position" - 1 as "position"
586590
)
587591
insert into ${chatMessages} ("chat_id", "message_id", "position", "role", "message")
588-
select ${params.chatId}, ${messageId}, reserved."position", ${messageRoleOf(params.chatId, params.message)},
589-
${JSON.stringify(params.message)}::jsonb
592+
select ${params.chatId}, ${messageId}, reserved."position", ${messageRoleOf(params.chatId, message)},
593+
${JSON.stringify(message)}::jsonb
590594
from reserved
591595
on conflict ("chat_id", "message_id") do nothing
592596
returning "message_id"
@@ -740,7 +744,7 @@ export async function persistTurn(
740744

741745
/** Idempotent on `(chatId, turn)`: a retried eval task can't write a second row. */
742746
export async function insertTurnEval(db: DashboardAgentDb, row: NewChatTurnEval): Promise<void> {
743-
await db.insert(chatTurnEvals).values(row).onConflictDoNothing();
747+
await db.insert(chatTurnEvals).values(toWellFormedDeep(row)).onConflictDoNothing();
744748
}
745749

746750
/**
@@ -864,6 +868,7 @@ export async function upsertInvestigationRevision(
864868
state: unknown;
865869
}
866870
): Promise<UpsertInvestigationResult> {
871+
const state = toWellFormedDeep(params.state);
867872
if (!params.id) {
868873
const id = generateInvestigationId();
869874
await db.insert(investigations).values({
@@ -872,15 +877,15 @@ export async function upsertInvestigationRevision(
872877
projectRef: params.projectRef,
873878
environmentRef: params.environmentRef,
874879
revision: 0,
875-
state: params.state,
880+
state,
876881
});
877882
return { ok: true, id, revision: 0, created: true };
878883
}
879884

880885
const rows = await db
881886
.update(investigations)
882887
.set({
883-
state: params.state,
888+
state,
884889
revision: sql`${investigations.revision} + 1`,
885890
updatedAt: sql`now()`,
886891
})
@@ -938,7 +943,7 @@ export async function seedInvestigation(
938943
projectRef: params.projectRef,
939944
environmentRef: params.environmentRef,
940945
revision: 0,
941-
state: params.state,
946+
state: toWellFormedDeep(params.state),
942947
})
943948
.onConflictDoNothing({ target: investigations.id })
944949
.returning({ id: investigations.id });

0 commit comments

Comments
 (0)