diff --git a/apps/server/src/orchestration-v2/ProjectionStore.test.ts b/apps/server/src/orchestration-v2/ProjectionStore.test.ts index aadd9ecc8be0..4076e9194212 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.test.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.test.ts @@ -2097,6 +2097,69 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { byteIds, Array.from({ length: 70 }, (_, index) => `turn-item:bounded-sql-history:${index + 1}`), ); + + // Selected message reads recover ordinals for only the loaded messages, while 69 + // unloaded assistant messages remain in the thread's history. + const firstMessageId = MessageId.make("message:bounded-sql-history:1"); + const lastMessageId = MessageId.make("message:bounded-sql-history:1000"); + const insertMessageItem = ( + suffix: string, + itemThreadId: string, + ordinal: number, + messageId: string | null | undefined, + ) => sql` + INSERT INTO orchestration_v2_projection_turn_items ( + turn_item_id, thread_id, run_id, node_id, provider_thread_id, provider_turn_id, + parent_item_id, ordinal, type, status, updated_at, payload_json + ) VALUES ( + ${`turn-item:bounded-sql-history:ordinal-${suffix}`}, ${itemThreadId}, NULL, NULL, NULL, + NULL, NULL, ${ordinal}, 'assistant_message', 'completed', ${nowIso}, + ${encodeUnknownJsonString({ messageId })} + ) + `; + yield* insertMessageItem("duplicate", threadId, 0, lastMessageId); + yield* insertMessageItem( + "other-thread", + "thread:bounded-sql-history:other", + -1, + firstMessageId, + ); + yield* insertMessageItem("null-id", threadId, -2, null); + yield* insertMessageItem("missing-id", threadId, -3, undefined); + const readMessageIds = (filter?: ProjectionStore.ProjectionRecordFilter) => + projectionStore + .getThreadRecords(threadId, ["messages"], filter) + .pipe(Effect.map(({ messages }) => messages.map((message) => message.id))); + assert.deepEqual(yield* readMessageIds(), [lastMessageId, firstMessageId]); + assert.deepEqual(yield* readMessageIds({ messageRoles: ["user", "assistant"] }), [ + lastMessageId, + firstMessageId, + ]); + for (const filter of [ + { + messageIds: [ + firstMessageId, + lastMessageId, + MessageId.make("message:bounded-sql-history:absent"), + ], + }, + { messageRoles: ["user" as const] }, + { + messageRunIds: [ + RunId.make("run:bounded-sql-history:1"), + RunId.make("run:bounded-sql-history:1000"), + ], + }, + ]) { + assert.deepEqual(yield* readMessageIds(filter), [lastMessageId, firstMessageId]); + } + assert.deepEqual(yield* readMessageIds({ messageIds: [] }), []); + assert.deepEqual(yield* readMessageIds({ messageIds: [firstMessageId] }), [firstMessageId]); + yield* sql` + DELETE FROM orchestration_v2_projection_turn_items + WHERE turn_item_id LIKE 'turn-item:bounded-sql-history:ordinal-%' + `; + yield* sql` UPDATE orchestration_v2_projection_turn_items SET type = 'command_execution', payload_json = json_set(payload_json, '$.type', 'command_execution') diff --git a/apps/server/src/orchestration-v2/ProjectionStore.ts b/apps/server/src/orchestration-v2/ProjectionStore.ts index 975ce7b0b9c5..4d22c7b60ae3 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -3602,10 +3602,19 @@ export const layer: Layer.Layer = ]); const messageOrdinals = new Map(); if (fields !== undefined && fields.includes("messages") && messages.length > 1) { + // Broad role reads keep the original query to avoid per-ID lookup overhead. + const loadedMessageIds = + filter?.messageIds === undefined && + filter?.messageRunIds === undefined && + (filter?.messageRoles === undefined || + (filter.messageRoles.includes("user") && filter.messageRoles.includes("assistant"))) + ? undefined + : messages.map((message) => message.id); const rows = yield* sql<{ message_id: string; ordinal: number }>` SELECT json_extract(payload_json, '$.messageId') AS message_id, MIN(ordinal) AS ordinal FROM orchestration_v2_projection_turn_items WHERE thread_id = ${threadId} AND type IN ('user_message', 'assistant_message') + ${loadedMessageIds === undefined ? sql`` : sql`AND json_extract(payload_json, '$.messageId') IN (SELECT value FROM json_each(${encodeIdList(loadedMessageIds)}))`} GROUP BY json_extract(payload_json, '$.messageId') `; for (const row of rows) messageOrdinals.set(row.message_id, row.ordinal);