diff --git a/packages/core/src/query/__tests__/message-query-functions.test.ts b/packages/core/src/query/__tests__/message-query-functions.test.ts index 918b9804..ce784d7e 100644 --- a/packages/core/src/query/__tests__/message-query-functions.test.ts +++ b/packages/core/src/query/__tests__/message-query-functions.test.ts @@ -7,6 +7,8 @@ import assert from 'node:assert/strict' import { describe, it } from 'node:test' import Database from 'better-sqlite3' +import { getMessageContext, getSearchMessageContext } from '../message-queries' +import { SqliteTestAdapter } from './sqlite-test-adapter' import { fetchAllRecentMessages, @@ -142,17 +144,28 @@ describe('message search and context', () => { it('loads chronological context instead of insertion order', async () => { const db = createBackfilledMessageDb() try { - const messages = await fetchMessageContext(createSqliteExecutor(db), 1, 2) - assert.deepEqual( - messages.map((message) => message.id), - [3, 4, 1, 2] - ) + const adapter = new SqliteTestAdapter(db) + const executor = createSqliteExecutor(db) + for (const messages of [getMessageContext(adapter, [1], 2), await fetchMessageContext(executor, 1, 2)]) { + assert.deepEqual( + messages.map((message) => message.id), + [3, 4, 1, 2] + ) + } + db.exec(`INSERT INTO message (id, sender_id, ts, type, content) VALUES + (5, 1, 300, 0, 'same-second backfill'), (6, 1, 300, 0, 'same-second later');`) + for (const messages of [getMessageContext(adapter, [5], 1), await fetchMessageContext(executor, 5, 1)]) { + assert.deepEqual( + messages.map((message) => message.id), + [1, 5, 6] + ) + } } finally { db.close() } }) - it('falls back to neighboring message ids without a message_context table', async () => { + it('falls back to chronological neighbors without a message_context table', async () => { const db = createMessageDb() try { const result = await fetchSearchMessageContext(createSqliteExecutor(db), [3], 1, 1) @@ -164,6 +177,43 @@ describe('message search and context', () => { db.close() } }) + + for (const indexed of [false, true]) { + it(`loads backfilled search context in time order and respects segments (indexed=${indexed})`, async () => { + const db = createBackfilledMessageDb() + try { + db.exec(`INSERT INTO message (id, sender_id, ts, type, content) VALUES + (5, 1, 300, 0, 'same-second backfill'), (6, 1, 300, 0, 'same-second later'), + (7, 1, 5000, 0, 'another segment');`) + if (indexed) { + db.exec(`CREATE TABLE message_context (message_id INTEGER PRIMARY KEY, segment_id INTEGER); + INSERT INTO message_context SELECT id, CASE WHEN id = 7 THEN 2 ELSE 1 END FROM message;`) + } + const adapter = new SqliteTestAdapter(db) + const executor = createSqliteExecutor(db) + for (const { ids, before, after, expected } of [ + { ids: [1], before: 1, after: 1, expected: [4, 1, 5] }, + { ids: [5], before: 1, after: 1, expected: [1, 5, 6] }, + { ids: [6], before: 1, after: 1, expected: [5, 6, 2] }, + { ids: [2], before: 1, after: 1, expected: indexed ? [6, 2] : [6, 2, 7] }, + { ids: [1, 5], before: 1, after: 1, expected: [4, 1, 5, 6] }, + { ids: [1], before: 0, after: 1, expected: [1, 5] }, + ]) { + for (const messages of [ + getSearchMessageContext(adapter, ids, before, after), + await fetchSearchMessageContext(executor, ids, before, after), + ]) { + assert.deepEqual( + messages.map((message) => message.id), + expected + ) + } + } + } finally { + db.close() + } + }) + } }) describe('recent messages', () => { diff --git a/packages/core/src/query/message-queries.ts b/packages/core/src/query/message-queries.ts index 295f6c5a..cc7eac10 100644 --- a/packages/core/src/query/message-queries.ts +++ b/packages/core/src/query/message-queries.ts @@ -592,14 +592,14 @@ function hydrateMessagesByIds(db: DatabaseAdapter, ids: number[]): MappedMessage if (ids.length === 0) return [] const placeholders = ids.map(() => '?').join(', ') const rows = db - .prepare(`${FULL_MSG_SELECT} WHERE msg.id IN (${placeholders}) ORDER BY msg.id ASC`) + .prepare(`${FULL_MSG_SELECT} WHERE msg.id IN (${placeholders}) ORDER BY msg.ts ASC, msg.id ASC`) .all(...ids) as unknown as FullMessageRow[] return rows.map(mapMessageRow) } /** * Get surrounding context messages for given message IDs. - * Uses simple id-based ordering (not session-aware). + * Uses chronological (timestamp, id) ordering (not session-aware). */ export function getMessageContext( db: DatabaseAdapter, @@ -614,12 +614,22 @@ export function getMessageContext( contextIds.add(messageId) const beforeRows = db - .prepare('SELECT id FROM message WHERE id < ? ORDER BY id DESC LIMIT ?') + .prepare( + `SELECT msg.id FROM message msg + JOIN message anchor ON anchor.id = ? + WHERE msg.ts < anchor.ts OR (msg.ts = anchor.ts AND msg.id < anchor.id) + ORDER BY msg.ts DESC, msg.id DESC LIMIT ?` + ) .all(messageId, contextSize) as { id: number }[] beforeRows.forEach((r) => contextIds.add(r.id)) const afterRows = db - .prepare('SELECT id FROM message WHERE id > ? ORDER BY id ASC LIMIT ?') + .prepare( + `SELECT msg.id FROM message msg + JOIN message anchor ON anchor.id = ? + WHERE msg.ts > anchor.ts OR (msg.ts = anchor.ts AND msg.id > anchor.id) + ORDER BY msg.ts ASC, msg.id ASC LIMIT ?` + ) .all(messageId, contextSize) as { id: number }[] afterRows.forEach((r) => contextIds.add(r.id)) } @@ -629,7 +639,7 @@ export function getMessageContext( /** * Get context messages around search results. - * Session-aware when message_context table is available, falls back to id-based ordering. + * Uses chronological neighbors within the indexed segment, or across the chat when unindexed. */ export function getSearchMessageContext( db: DatabaseAdapter, @@ -657,21 +667,25 @@ export function getSearchMessageContext( if (contextBefore > 0) { const rows = db .prepare( - `SELECT mc.message_id as id FROM message_context mc - WHERE mc.segment_id = ? AND mc.message_id < ? - ORDER BY mc.message_id DESC LIMIT ?` + `SELECT msg.id FROM message_context mc + JOIN message msg ON msg.id = mc.message_id + JOIN message anchor ON anchor.id = ? + WHERE mc.segment_id = ? AND (msg.ts < anchor.ts OR (msg.ts = anchor.ts AND msg.id < anchor.id)) + ORDER BY msg.ts DESC, msg.id DESC LIMIT ?` ) - .all(sessionRow.segment_id, messageId, contextBefore) as { id: number }[] + .all(messageId, sessionRow.segment_id, contextBefore) as { id: number }[] rows.forEach((r) => contextIds.add(r.id)) } if (contextAfter > 0) { const rows = db .prepare( - `SELECT mc.message_id as id FROM message_context mc - WHERE mc.segment_id = ? AND mc.message_id > ? - ORDER BY mc.message_id ASC LIMIT ?` + `SELECT msg.id FROM message_context mc + JOIN message msg ON msg.id = mc.message_id + JOIN message anchor ON anchor.id = ? + WHERE mc.segment_id = ? AND (msg.ts > anchor.ts OR (msg.ts = anchor.ts AND msg.id > anchor.id)) + ORDER BY msg.ts ASC, msg.id ASC LIMIT ?` ) - .all(sessionRow.segment_id, messageId, contextAfter) as { id: number }[] + .all(messageId, sessionRow.segment_id, contextAfter) as { id: number }[] rows.forEach((r) => contextIds.add(r.id)) } continue @@ -680,13 +694,23 @@ export function getSearchMessageContext( if (contextBefore > 0) { const rows = db - .prepare('SELECT id FROM message WHERE id < ? ORDER BY id DESC LIMIT ?') + .prepare( + `SELECT msg.id FROM message msg + JOIN message anchor ON anchor.id = ? + WHERE msg.ts < anchor.ts OR (msg.ts = anchor.ts AND msg.id < anchor.id) + ORDER BY msg.ts DESC, msg.id DESC LIMIT ?` + ) .all(messageId, contextBefore) as { id: number }[] rows.forEach((r) => contextIds.add(r.id)) } if (contextAfter > 0) { const rows = db - .prepare('SELECT id FROM message WHERE id > ? ORDER BY id ASC LIMIT ?') + .prepare( + `SELECT msg.id FROM message msg + JOIN message anchor ON anchor.id = ? + WHERE msg.ts > anchor.ts OR (msg.ts = anchor.ts AND msg.id > anchor.id) + ORDER BY msg.ts ASC, msg.id ASC LIMIT ?` + ) .all(messageId, contextAfter) as { id: number }[] rows.forEach((r) => contextIds.add(r.id)) } diff --git a/packages/core/src/query/message-query-functions.ts b/packages/core/src/query/message-query-functions.ts index b4f44d05..67fe47e0 100644 --- a/packages/core/src/query/message-query-functions.ts +++ b/packages/core/src/query/message-query-functions.ts @@ -177,7 +177,7 @@ export async function fetchMessageContext( /** * Get context messages around search results. - * Session-aware when message_context table is available, falls back to id-based. + * Uses chronological neighbors within the indexed segment, or across the chat when unindexed. */ export async function fetchSearchMessageContext( executor: AsyncSqlExecutor, @@ -211,19 +211,23 @@ export async function fetchSearchMessageContext( if (sessionRow) { if (contextBefore > 0) { const rows = await executor.all<{ id: number }>( - `SELECT mc.message_id as id FROM message_context mc - WHERE mc.segment_id = ? AND mc.message_id < ? - ORDER BY mc.message_id DESC LIMIT ?`, - [sessionRow.segment_id, messageId, contextBefore] + `SELECT msg.id FROM message_context mc + JOIN message msg ON msg.id = mc.message_id + JOIN message anchor ON anchor.id = ? + WHERE mc.segment_id = ? AND (msg.ts < anchor.ts OR (msg.ts = anchor.ts AND msg.id < anchor.id)) + ORDER BY msg.ts DESC, msg.id DESC LIMIT ?`, + [messageId, sessionRow.segment_id, contextBefore] ) rows.forEach((r) => contextIds.add(r.id)) } if (contextAfter > 0) { const rows = await executor.all<{ id: number }>( - `SELECT mc.message_id as id FROM message_context mc - WHERE mc.segment_id = ? AND mc.message_id > ? - ORDER BY mc.message_id ASC LIMIT ?`, - [sessionRow.segment_id, messageId, contextAfter] + `SELECT msg.id FROM message_context mc + JOIN message msg ON msg.id = mc.message_id + JOIN message anchor ON anchor.id = ? + WHERE mc.segment_id = ? AND (msg.ts > anchor.ts OR (msg.ts = anchor.ts AND msg.id > anchor.id)) + ORDER BY msg.ts ASC, msg.id ASC LIMIT ?`, + [messageId, sessionRow.segment_id, contextAfter] ) rows.forEach((r) => contextIds.add(r.id)) } @@ -232,17 +236,23 @@ export async function fetchSearchMessageContext( } if (contextBefore > 0) { - const rows = await executor.all<{ id: number }>('SELECT id FROM message WHERE id < ? ORDER BY id DESC LIMIT ?', [ - messageId, - contextBefore, - ]) + const rows = await executor.all<{ id: number }>( + `SELECT msg.id FROM message msg + JOIN message anchor ON anchor.id = ? + WHERE msg.ts < anchor.ts OR (msg.ts = anchor.ts AND msg.id < anchor.id) + ORDER BY msg.ts DESC, msg.id DESC LIMIT ?`, + [messageId, contextBefore] + ) rows.forEach((r) => contextIds.add(r.id)) } if (contextAfter > 0) { - const rows = await executor.all<{ id: number }>('SELECT id FROM message WHERE id > ? ORDER BY id ASC LIMIT ?', [ - messageId, - contextAfter, - ]) + const rows = await executor.all<{ id: number }>( + `SELECT msg.id FROM message msg + JOIN message anchor ON anchor.id = ? + WHERE msg.ts > anchor.ts OR (msg.ts = anchor.ts AND msg.id > anchor.id) + ORDER BY msg.ts ASC, msg.id ASC LIMIT ?`, + [messageId, contextAfter] + ) rows.forEach((r) => contextIds.add(r.id)) } }