diff --git a/docs/api-reference.md b/docs/api-reference.md index 8726e70..a1bb137 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -235,6 +235,8 @@ POST /api/v1/conversations Creates a new conversation, optionally with initial messages. +Message writes are limited to **100 messages per request**, **64 KiB (65,536 UTF-8 bytes) per content string**, and **1 MiB (1,048,576 UTF-8 bytes) of combined content**. These limits also apply to appending messages, trace ingestion (`observations`), and the MCP `store_conversation` tool. REST rejects violations before writing with `400 BAD_REQUEST`; MCP returns an `isError` tool result containing `INVALID_PARAMS`. Empty initial message arrays remain valid; append and trace ingestion require at least one item. + **Request body:** | Field | Type | Required | Description | @@ -242,14 +244,14 @@ Creates a new conversation, optionally with initial messages. | `external_id` | string | No | Caller-provided identifier. Must be unique within the project. | | `title` | string | No | Conversation title. | | `metadata` | object | No | Arbitrary key-value pairs. | -| `messages` | array | No | Initial messages to attach. | +| `messages` | array | No | Initial messages to attach (maximum 100). Combined message content must not exceed 1 MiB (1,048,576 UTF-8 bytes). | Each message object: | Field | Type | Required | Description | |-------|------|----------|-------------| | `role` | string | Yes | One of: `system`, `user`, `assistant`, `tool`. | -| `content` | string | Yes | Message content (min 1 character). | +| `content` | string | Yes | Non-empty message content, at most 64 KiB (65,536 UTF-8 bytes). | | `metadata` | object | No | Arbitrary key-value pairs. | | `token_count` | integer | No | Non-negative integer. Defaults to 0. | @@ -469,7 +471,7 @@ Add one or more messages to an existing conversation. Automatically updates the | Field | Type | Required | Description | |-------|------|----------|-------------| -| `messages` | array | Yes | At least 1 message. Same schema as [Create Conversation](#create-conversation) messages. | +| `messages` | array | Yes | 1–100 messages; each content string at most 64 KiB and combined content at most 1 MiB (UTF-8 bytes). Same schema as [Create Conversation](#create-conversation) messages. | **Response:** `201 Created` @@ -988,6 +990,8 @@ POST /v1/conversations Creates a new conversation, optionally with initial messages. +Message writes are limited to **100 messages per request**, **64 KiB (65,536 UTF-8 bytes) per content string**, and **1 MiB (1,048,576 UTF-8 bytes) of combined content**. These limits also apply to appending messages, trace ingestion (`observations`), and the MCP `store_conversation` tool. REST rejects violations before writing with `400 BAD_REQUEST`; MCP returns an `isError` tool result containing `INVALID_PARAMS`. Empty initial message arrays remain valid; append and trace ingestion require at least one item. + **Request body:** | Field | Type | Required | Description | @@ -995,14 +999,14 @@ Creates a new conversation, optionally with initial messages. | `external_id` | string | No | Caller-provided identifier. Must be unique within the project. | | `title` | string | No | Conversation title. | | `metadata` | object | No | Arbitrary key-value pairs. | -| `messages` | array | No | Initial messages to attach. | +| `messages` | array | No | Initial messages to attach (maximum 100). Combined message content must not exceed 1 MiB (1,048,576 UTF-8 bytes). | Each message object: | Field | Type | Required | Description | |-------|------|----------|-------------| | `role` | string | Yes | One of: `system`, `user`, `assistant`, `tool`. | -| `content` | string | Yes | Message content (min 1 character). | +| `content` | string | Yes | Non-empty message content, at most 64 KiB (65,536 UTF-8 bytes). | | `metadata` | object | No | Arbitrary key-value pairs. | | `token_count` | integer | No | Non-negative integer. Defaults to 0. | @@ -1365,7 +1369,7 @@ Add one or more messages to an existing conversation. Automatically updates the | Field | Type | Required | Description | |-------|------|----------|-------------| -| `messages` | array | Yes | At least 1 message. Same schema as [Create Conversation](#create-conversation) messages. | +| `messages` | array | Yes | 1–100 messages; each content string at most 64 KiB and combined content at most 1 MiB (UTF-8 bytes). Same schema as [Create Conversation](#create-conversation) messages. | **Response:** `201 Created` diff --git a/packages/api/src/lib/validation.ts b/packages/api/src/lib/validation.ts index b24d592..e48c69e 100644 --- a/packages/api/src/lib/validation.ts +++ b/packages/api/src/lib/validation.ts @@ -32,9 +32,24 @@ export const SLUG_PATTERN = /^[a-z0-9][a-z0-9-]*[a-z0-9]$|^[a-z0-9]$/; // Message schemas // --------------------------------------------------------------------------- +export const MAX_MESSAGES_PER_REQUEST = 100; +export const MAX_MESSAGE_CONTENT_BYTES = 64 * 1024; +export const MAX_MESSAGES_CONTENT_BYTES = 1024 * 1024; + +const utf8Encoder = new TextEncoder(); + +export const MessageContentSchema = z + .string() + .min(1) + .max(MAX_MESSAGE_CONTENT_BYTES) + .refine((content) => utf8Encoder.encode(content).byteLength <= MAX_MESSAGE_CONTENT_BYTES, { + message: `Message content must not exceed ${MAX_MESSAGE_CONTENT_BYTES} UTF-8 bytes`, + }) + .describe("Non-empty content, at most 65536 UTF-8 bytes (64 KiB)"); + export const MessageInputSchema = z.object({ role: z.enum(MESSAGE_ROLES), - content: z.string().min(1), + content: MessageContentSchema, metadata: z.record(z.string(), z.unknown()).optional(), token_count: z.number().int().nonnegative().optional(), model: z.string().max(100).optional(), @@ -50,6 +65,26 @@ export const MessageInputSchema = z.object({ }); export type MessageInput = z.infer; +export function boundedMessagesSchema>(itemSchema: T) { + return z + .array(itemSchema) + .max(MAX_MESSAGES_PER_REQUEST) + .refine( + (items) => { + let bytes = 0; + for (const item of items) { + bytes += utf8Encoder.encode(item.content).byteLength; + if (bytes > MAX_MESSAGES_CONTENT_BYTES) return false; + } + return true; + }, + { message: `Combined message content must not exceed ${MAX_MESSAGES_CONTENT_BYTES} UTF-8 bytes` }, + ) + .describe("At most 100 messages; combined content at most 1048576 UTF-8 bytes (1 MiB)"); +} + +export const MessagesInputSchema = boundedMessagesSchema(MessageInputSchema); + // --------------------------------------------------------------------------- // Conversation schemas // --------------------------------------------------------------------------- @@ -58,7 +93,7 @@ export const CreateConversationSchema = z.object({ external_id: z.string().optional(), title: z.string().optional(), metadata: z.record(z.string(), z.unknown()).optional(), - messages: z.array(MessageInputSchema).optional(), + messages: MessagesInputSchema.optional(), }); export type CreateConversationInput = z.infer; @@ -69,7 +104,7 @@ export const UpdateConversationSchema = z.object({ export type UpdateConversationInput = z.infer; export const AppendMessagesSchema = z.object({ - messages: z.array(MessageInputSchema).min(1), + messages: MessagesInputSchema.min(1), }); export type AppendMessagesInput = z.infer; @@ -277,7 +312,7 @@ export type CreateClaimInput = z.infer; const ObservationInputSchema = z.object({ role: z.enum(MESSAGE_ROLES).optional().default("assistant"), - content: z.string().min(1), + content: MessageContentSchema, parent_message_id: z.string().optional(), observation_type: z.enum(OBSERVATION_TYPES), metadata: z.record(z.string(), z.unknown()).optional(), @@ -298,6 +333,6 @@ export const IngestTraceSchema = z.object({ title: z.string().optional(), metadata: z.record(z.string(), z.unknown()).optional(), }), - observations: z.array(ObservationInputSchema).min(1).max(100), + observations: boundedMessagesSchema(ObservationInputSchema).min(1), }); export type IngestTraceInput = z.infer; diff --git a/packages/api/src/routes/mcp/tools.ts b/packages/api/src/routes/mcp/tools.ts index 3035cf1..61b45f2 100644 --- a/packages/api/src/routes/mcp/tools.ts +++ b/packages/api/src/routes/mcp/tools.ts @@ -4,7 +4,11 @@ import { z } from "zod"; import { conversations as conversationsTable, messages as messagesTable } from "../../db/schema"; import { GrantableScopeSchema, scopesSatisfyAll } from "../../lib/scopes"; import { deserializeConversationFull, deserializeMessage } from "../../lib/serialization"; -import { CapabilityScopeSchema } from "../../lib/validation"; +import { + boundedMessagesSchema, + CapabilityScopeSchema, + MessageContentSchema, +} from "../../lib/validation"; import * as capabilityTokensService from "../../services/capability-tokens"; import * as claimsService from "../../services/claims"; import * as keysService from "../../services/keys"; @@ -46,7 +50,7 @@ export class ToolError extends Error { const messageSchema = z.object({ role: z.enum(["user", "assistant", "system", "tool"]), - content: z.string(), + content: MessageContentSchema, metadata: z.record(z.string(), z.unknown()).optional(), token_count: z.number().int().optional(), }); @@ -82,7 +86,7 @@ const storeConversationSchema = z.object({ external_id: z.string().optional().describe("Optional external identifier for deduplication"), title: z.string().optional().describe("Human-readable title for the conversation"), metadata: z.record(z.string(), z.unknown()).optional().describe("Arbitrary JSON metadata"), - messages: z.array(messageSchema).optional().describe("Initial messages to include"), + messages: boundedMessagesSchema(messageSchema).optional(), }); const recallConversationSchema = z.object({ diff --git a/packages/api/src/services/conversations.ts b/packages/api/src/services/conversations.ts index ce2496b..e15554e 100644 --- a/packages/api/src/services/conversations.ts +++ b/packages/api/src/services/conversations.ts @@ -14,6 +14,7 @@ import { serializeMetadata } from "../lib/serialization"; import { TagSchema } from "../lib/validation"; import { sendWebhookWithRetry } from "../lib/webhook"; import * as webhooksService from "../services/webhooks"; +import { insertMessageRows } from "./messages"; // --------------------------------------------------------------------------- // Field Selection Types and Utilities @@ -238,7 +239,7 @@ export async function createConversation( createdAt: now, })); - await db.insert(messages).values(rows); + await insertMessageRows(db, rows); messageRows = rows as (typeof messages.$inferSelect)[]; } diff --git a/packages/api/src/services/mcp-conversations.ts b/packages/api/src/services/mcp-conversations.ts index ae24231..dbb74cd 100644 --- a/packages/api/src/services/mcp-conversations.ts +++ b/packages/api/src/services/mcp-conversations.ts @@ -17,6 +17,7 @@ import { invalidateAnalyticsCache } from "../lib/analytics-cache"; import { CACHE_COUNT_TTL_S } from "../lib/config"; import { generateId } from "../lib/id"; import { serializeMetadata } from "../lib/serialization"; +import { insertMessageRows } from "./messages"; // --------------------------------------------------------------------------- // Types @@ -206,7 +207,7 @@ export async function createConversation( createdAt: now, })); - await db.insert(messages).values(rows); + await insertMessageRows(db, rows); } invalidateAnalyticsCache(cache, executionCtx, projectId); diff --git a/packages/api/src/services/messages.ts b/packages/api/src/services/messages.ts index cfa93b9..7805d1d 100644 --- a/packages/api/src/services/messages.ts +++ b/packages/api/src/services/messages.ts @@ -6,6 +6,23 @@ import { generateId } from "../lib/id"; import { serializeMetadata } from "../lib/serialization"; import type { MessageInput } from "../lib/validation"; +/** D1 permits 100 bound parameters per statement; a message has 17 columns. */ +export const MESSAGE_INSERT_CHUNK_SIZE = 5; + +export async function insertMessageRows( + db: DrizzleD1Database, + rows: (typeof messages.$inferInsert)[], +): Promise { + if (rows.length === 0) return; + const first = db.insert(messages).values(rows.slice(0, MESSAGE_INSERT_CHUNK_SIZE)); + const remaining = []; + for (let index = MESSAGE_INSERT_CHUNK_SIZE; index < rows.length; index += MESSAGE_INSERT_CHUNK_SIZE) { + remaining.push(db.insert(messages).values(rows.slice(index, index + MESSAGE_INSERT_CHUNK_SIZE))); + } + // D1 batch is transactional, so a failed chunk cannot leave a partial append. + await db.batch([first, ...remaining]); +} + /** * Append messages to a conversation. * @@ -44,7 +61,7 @@ export async function appendMessages( createdAt: now, })); - await db.insert(messages).values(messageRows); + await insertMessageRows(db, messageRows); const addedTokens = inputMessages.reduce((sum, m) => sum + (m.token_count ?? 0), 0); const addedCost = inputMessages.reduce((sum, m) => sum + (m.cost_microdollars ?? 0), 0); diff --git a/packages/api/src/services/traces.ts b/packages/api/src/services/traces.ts index 81c0d26..2229ee0 100644 --- a/packages/api/src/services/traces.ts +++ b/packages/api/src/services/traces.ts @@ -9,6 +9,7 @@ import { } from "../lib/serialization"; import { buildObservationTree } from "../lib/trace-tree"; import type { IngestTraceInput } from "../lib/validation"; +import { insertMessageRows } from "./messages"; // --------------------------------------------------------------------------- // Types @@ -104,7 +105,7 @@ export async function ingestTrace( } } - await db.insert(messages).values(messageRows); + await insertMessageRows(db, messageRows); const conversation = { id: conversationId, diff --git a/packages/api/test/message-limits.test.ts b/packages/api/test/message-limits.test.ts new file mode 100644 index 0000000..aff3c5b --- /dev/null +++ b/packages/api/test/message-limits.test.ts @@ -0,0 +1,393 @@ +import { env, SELF } from "cloudflare:test"; +import { drizzle } from "drizzle-orm/d1"; +import { beforeAll, describe, expect, it, vi } from "vitest"; +import { + AppendMessagesSchema, + CreateConversationSchema, + IngestTraceSchema, + MAX_MESSAGE_CONTENT_BYTES, + MAX_MESSAGES_CONTENT_BYTES, + MAX_MESSAGES_PER_REQUEST, +} from "../src/lib/validation"; +import { insertMessageRows, serializeMessageRows } from "../src/services/messages"; +import { applyMigrations, authHeaders, seedProject } from "./setup"; + +const message = (content = "hello") => ({ role: "user" as const, content, token_count: 2 }); +type InputMessage = ReturnType; +type ObservationInput = InputMessage & { + observation_type: "span"; + parent_message_id?: string; +}; +interface StoredMessage { + id: string; + content: string; + parent_message_id: string | null; +} +interface ConversationResponse { + id: string; + message_count: number; + messages: StoredMessage[]; +} +interface TraceResponse { + conversation: { id: string; message_count: number; token_count: number }; + observations: StoredMessage[]; +} + +const create = (messages?: InputMessage[]) => + SELF.fetch("http://localhost/api/v1/conversations", { + method: "POST", + headers: authHeaders(), + body: JSON.stringify({ messages }), + }); +const append = (id: string, messages: InputMessage[]) => + SELF.fetch(`http://localhost/api/v1/conversations/${id}/messages`, { + method: "POST", + headers: authHeaders(), + body: JSON.stringify({ messages }), + }); +const observations = (messages: InputMessage[]): ObservationInput[] => + messages.map((m) => ({ ...m, observation_type: "span" })); +const ingest = (observations: ObservationInput[]) => + SELF.fetch("http://localhost/api/v1/conversations/traces/ingest", { + method: "POST", + headers: authHeaders(), + body: JSON.stringify({ trace: { title: "Message limits trace" }, observations }), + }); +const store = async (messages?: InputMessage[]) => { + const response = await SELF.fetch("http://localhost/api/mcp", { + method: "POST", + headers: { ...authHeaders(), Accept: "application/json, text/event-stream" }, + body: JSON.stringify({ + jsonrpc: "2.0", + id: 1, + method: "tools/call", + params: { name: "store_conversation", arguments: { messages } }, + }), + }); + expect(response.status).toBe(200); + return response.json<{ + result: { isError?: boolean; content: { text: string }[] }; + }>(); +}; + +// Check both tables: validation must not leave an orphan conversation or partial messages. +const writeSnapshot = () => + env.DB.prepare(`SELECT + (SELECT COUNT(*) FROM conversations) AS conversations, + (SELECT COUNT(*) FROM messages) AS messages, + (SELECT COALESCE(SUM(message_count), 0) FROM conversations) AS message_count, + (SELECT COALESCE(SUM(token_count), 0) FROM conversations) AS token_count`).first(); +const conversationRow = (id: string) => + env.DB.prepare("SELECT * FROM conversations WHERE id = ?").bind(id).first(); +const messageCount = (id: string) => + env.DB.prepare("SELECT COUNT(*) AS n FROM messages WHERE conversation_id = ?").bind(id).first(); +const contentBytes = (id: string) => + env.DB.prepare( + "SELECT COUNT(*) AS n, SUM(length(CAST(content AS BLOB))) AS bytes FROM messages WHERE conversation_id = ?", + ) + .bind(id) + .first(); + +const invalidBatches = [ + { + name: "101 messages", + messages: Array.from({ length: MAX_MESSAGES_PER_REQUEST + 1 }, () => message()), + }, + { name: "empty content", messages: [message("")] }, + { + name: "oversized ASCII content", + messages: [message("a".repeat(MAX_MESSAGE_CONTENT_BYTES + 1))], + }, + { + name: "oversized UTF-8 content", + messages: [message(`${"é".repeat(MAX_MESSAGE_CONTENT_BYTES / 2)}a`)], + }, + { + name: "combined content above 1 MiB", + messages: [ + ...Array.from({ length: MAX_MESSAGES_CONTENT_BYTES / MAX_MESSAGE_CONTENT_BYTES }, () => + message("é".repeat(MAX_MESSAGE_CONTENT_BYTES / 2)), + ), + message("a"), + ], + }, +]; + +describe("message request limits", () => { + beforeAll(async () => { + await applyMigrations(); + await seedProject(); + }); + + it("accepts batches immediately below count and UTF-8 byte limits", () => { + const batches = [ + Array.from({ length: MAX_MESSAGES_PER_REQUEST - 1 }, () => message()), + [message(`${"é".repeat(MAX_MESSAGE_CONTENT_BYTES / 2 - 1)}a`)], + [ + ...Array.from({ length: MAX_MESSAGES_CONTENT_BYTES / MAX_MESSAGE_CONTENT_BYTES - 1 }, () => + message("é".repeat(MAX_MESSAGE_CONTENT_BYTES / 2)), + ), + message(`${"é".repeat(MAX_MESSAGE_CONTENT_BYTES / 2 - 1)}a`), + ], + ]; + for (const messages of batches) { + expect(CreateConversationSchema.safeParse({ messages }).success).toBe(true); + expect(AppendMessagesSchema.safeParse({ messages }).success).toBe(true); + expect( + IngestTraceSchema.safeParse({ trace: {}, observations: observations(messages) }).success, + ).toBe(true); + } + }); + + it("retains omitted/empty-create semantics for REST and MCP, but rejects empty append/trace", async () => { + for (const messages of [undefined, []]) { + const response = await create(messages); + expect(response.status).toBe(201); + const conversation = await response.json(); + expect(conversation.message_count).toBe(0); + expect(conversation.messages).toEqual([]); + expect(await messageCount(conversation.id)).toEqual({ n: 0 }); + const before = await writeSnapshot(); + const row = await conversationRow(conversation.id); + const appended = await append(conversation.id, []); + expect(appended.status).toBe(400); + await appended.json(); + expect(await writeSnapshot()).toEqual(before); + expect(await conversationRow(conversation.id)).toEqual(row); + + const stored = await store(messages); + expect(stored.result.isError).not.toBe(true); + const result = JSON.parse(stored.result.content[0].text) as ConversationResponse; + expect(result.message_count).toBe(0); + expect(await messageCount(result.id)).toEqual({ n: 0 }); + } + const before = await writeSnapshot(); + const response = await ingest([]); + expect(response.status).toBe(400); + await response.json(); + expect(await writeSnapshot()).toEqual(before); + }); + + for (const { name, messages } of invalidBatches) { + it(`rejects ${name} on REST create without writes`, async () => { + const before = await writeSnapshot(); + const response = await create(messages); + expect(response.status).toBe(400); + const body = await response.json<{ error: { code: string; message: string } }>(); + expect(body.error.code).toBe("BAD_REQUEST"); + expect(body.error.message.length).toBeGreaterThan(0); + expect(await writeSnapshot()).toEqual(before); + }); + + it(`rejects ${name} on append without writes or counter changes`, async () => { + const response = await create([message()]); + expect(response.status).toBe(201); + const conversation = await response.json(); + const before = await writeSnapshot(); + const row = await conversationRow(conversation.id); + const appended = await append(conversation.id, messages); + expect(appended.status).toBe(400); + expect((await appended.json<{ error: { code: string } }>()).error.code).toBe("BAD_REQUEST"); + expect(await writeSnapshot()).toEqual(before); + expect(await conversationRow(conversation.id)).toEqual(row); + expect(await messageCount(conversation.id)).toEqual({ n: 1 }); + }); + + it(`rejects ${name} on MCP with INVALID_PARAMS and no writes`, async () => { + const before = await writeSnapshot(); + const body = await store(messages); + expect(body.result.isError).toBe(true); + expect(body.result.content[0].text).toContain("INVALID_PARAMS"); + expect(await writeSnapshot()).toEqual(before); + }); + + it(`rejects ${name} on trace ingestion without writes`, async () => { + const before = await writeSnapshot(); + const response = await ingest(observations(messages)); + expect(response.status).toBe(400); + expect((await response.json<{ error: { code: string } }>()).error.code).toBe("BAD_REQUEST"); + expect(await writeSnapshot()).toEqual(before); + }); + } + + it("creates and appends 100 messages across D1 chunks, preserving order and counts", async () => { + const messages = Array.from({ length: MAX_MESSAGES_PER_REQUEST }, (_, i) => + message(`message ${i}`), + ); + const response = await create(messages); + expect(response.status).toBe(201); + const conversation = await response.json(); + expect(conversation.messages.map((m) => m.content)).toEqual(messages.map((m) => m.content)); + const appended = await append(conversation.id, messages); + expect(appended.status).toBe(201); + const body = await appended.json<{ messages: StoredMessage[] }>(); + expect(body.messages.map((m) => m.content)).toEqual(messages.map((m) => m.content)); + expect(await conversationRow(conversation.id)).toMatchObject({ + message_count: 200, + token_count: 400, + }); + expect(await messageCount(conversation.id)).toEqual({ n: 200 }); + }); + + it("accepts exactly 64 KiB UTF-8 per message and 1 MiB combined via create and append", async () => { + const messages = Array.from( + { length: MAX_MESSAGES_CONTENT_BYTES / MAX_MESSAGE_CONTENT_BYTES }, + () => message("é".repeat(MAX_MESSAGE_CONTENT_BYTES / 2)), + ); + const response = await create(messages); + expect(response.status).toBe(201); + const conversation = await response.json(); + expect(conversation.message_count).toBe(16); + expect(await contentBytes(conversation.id)).toEqual({ + n: 16, + bytes: MAX_MESSAGES_CONTENT_BYTES, + }); + const appended = await append(conversation.id, messages); + expect(appended.status).toBe(201); + const body = await appended.json<{ messages: StoredMessage[] }>(); + expect(body.messages.map((m) => m.content)).toEqual(messages.map((m) => m.content)); + expect(await contentBytes(conversation.id)).toEqual({ + n: 32, + bytes: 2 * MAX_MESSAGES_CONTENT_BYTES, + }); + expect(await conversationRow(conversation.id)).toMatchObject({ + message_count: 32, + token_count: 64, + }); + }); + + it("accepts exactly 64 KiB UTF-8 per message and 1 MiB combined via MCP and trace ingestion", async () => { + const messages = Array.from( + { length: MAX_MESSAGES_CONTENT_BYTES / MAX_MESSAGE_CONTENT_BYTES }, + () => message("é".repeat(MAX_MESSAGE_CONTENT_BYTES / 2)), + ); + const body = await store(messages); + expect(body.result.isError).not.toBe(true); + const conversation = JSON.parse(body.result.content[0].text) as ConversationResponse; + expect(conversation.message_count).toBe(16); + expect(await contentBytes(conversation.id)).toEqual({ + n: 16, + bytes: MAX_MESSAGES_CONTENT_BYTES, + }); + const response = await ingest(observations(messages)); + expect(response.status).toBe(201); + const trace = await response.json(); + expect(trace.conversation.message_count).toBe(16); + expect(trace.observations.map((o) => o.content)).toEqual(messages.map((m) => m.content)); + expect(await contentBytes(trace.conversation.id)).toEqual({ + n: 16, + bytes: MAX_MESSAGES_CONTENT_BYTES, + }); + }); + + it("stores 100 messages through MCP across D1 chunks", async () => { + const body = await store(Array.from({ length: MAX_MESSAGES_PER_REQUEST }, () => message())); + expect(body.result.isError).not.toBe(true); + const conversation = JSON.parse(body.result.content[0].text) as ConversationResponse; + expect(conversation.message_count).toBe(100); + expect(await conversationRow(conversation.id)).toMatchObject({ + message_count: 100, + token_count: 200, + }); + expect(await messageCount(conversation.id)).toEqual({ n: 100 }); + }); + + it("ingests 100 observations with forward and backward parent references across chunks", async () => { + const input = observations( + Array.from({ length: MAX_MESSAGES_PER_REQUEST }, (_, i) => message(`observation ${i}`)), + ).map((o, i) => ({ + ...o, + // Second observation references the last chunk; every fifth references the previous chunk. + parent_message_id: i === 0 ? undefined : i === 1 ? "$100" : i % 5 === 0 ? `$${i}` : "$1", + })); + const response = await ingest(input); + expect(response.status).toBe(201); + const body = await response.json(); + expect(body.conversation).toMatchObject({ message_count: 100, token_count: 200 }); + expect(body.observations.map((o) => o.content)).toEqual(input.map((o) => o.content)); + const expected = body.observations.map((o, i) => ({ + id: o.id, + content: input[i].content, + parent_message_id: input[i].parent_message_id + ? body.observations[Number(input[i].parent_message_id?.slice(1)) - 1].id + : null, + })); + expect( + body.observations.map(({ id, content, parent_message_id }) => ({ + id, + content, + parent_message_id, + })), + ).toEqual(expected); + const persisted = await env.DB.prepare( + "SELECT id, content, parent_message_id FROM messages WHERE conversation_id = ? ORDER BY id", + ) + .bind(body.conversation.id) + .all(); + expect(persisted.results).toEqual( + expected.sort((a, b) => (a.id < b.id ? -1 : a.id > b.id ? 1 : 0)), + ); + expect(await conversationRow(body.conversation.id)).toMatchObject({ + message_count: 100, + token_count: 200, + }); + }); + + it("keeps every insert statement within D1's 100 bound parameters", async () => { + const db = drizzle(env.DB); + const rows = serializeMessageRows( + "bind-limit-conversation", + Array.from({ length: MAX_MESSAGES_PER_REQUEST }, () => ({ + ...message(), + metadata: { source: "bind-limit-test" }, + model: "test-model", + input_tokens: 1, + output_tokens: 1, + cost_microdollars: 1, + parent_message_id: "parent", + observation_type: "span" as const, + start_time: 1, + end_time: 2, + status: "success" as const, + level: "default" as const, + })), + ); + const batch = vi.spyOn(db, "batch").mockResolvedValue([]); + try { + await insertMessageRows(db, rows); + expect(batch).toHaveBeenCalledTimes(1); + const statements = batch.mock.calls[0][0]; + expect(statements.length).toBeGreaterThan(1); + for (const statement of statements) { + expect(statement.toSQL().params.length).toBeLessThanOrEqual(100); + } + } finally { + batch.mockRestore(); + } + }); + + it("rolls back the first insert chunk when the second chunk fails", async () => { + const response = await create([message("existing message")]); + expect(response.status).toBe(201); + const conversation = await response.json(); + const db = drizzle(env.DB); + const rows = serializeMessageRows( + conversation.id, + Array.from({ length: 6 }, (_, i) => message(`rollback ${i}`)), + ); + // Five valid rows fit in chunk one; a duplicate primary key fails chunk two. + rows[5].id = rows[0].id; + const before = await writeSnapshot(); + const row = await conversationRow(conversation.id); + await expect(insertMessageRows(db, rows)).rejects.toThrow(/UNIQUE constraint failed/); + expect(await writeSnapshot()).toEqual(before); + expect(await conversationRow(conversation.id)).toEqual(row); + expect(await messageCount(conversation.id)).toEqual({ n: 1 }); + + // These same first five rows are valid and must be reusable after the rollback. + await insertMessageRows(db, rows.slice(0, 5)); + expect(await messageCount(conversation.id)).toEqual({ n: 6 }); + await insertMessageRows(db, []); + expect(await messageCount(conversation.id)).toEqual({ n: 6 }); + }); +});