diff --git a/config/scripts/session-search-retention-benchmark.ts b/config/scripts/session-search-retention-benchmark.ts new file mode 100644 index 00000000000..095eb5f29b0 --- /dev/null +++ b/config/scripts/session-search-retention-benchmark.ts @@ -0,0 +1,148 @@ +import assert from 'node:assert/strict' +import { mkdtemp, rm, stat } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { setImmediate as yieldToEventLoop } from 'node:timers/promises' +import { + syntheticCandidate, + syntheticSession, + userMessages +} from '../../src/main/ai-vault-search/session-search-index-test-fixture' +import { SessionSearchStore } from '../../src/main/ai-vault-search/session-search-store' +import SyncDatabase from '../../src/main/sqlite/sync-database' + +// Bundle with esbuild --bundle --platform=node, then run on the host under test. +// Every mode seeds through SessionSearchStore so the three arms are comparable; +// only `whole-file` leaves the shipped path, because it is the baseline the +// batched purge exists to replace. + +const ROWS = 60_000 + +/** The purge yields with `setImmediate` between chunks, so a peer chain samples each gap. */ +async function sampleLoopStalls(running: () => boolean, intervals: number[]): Promise { + let previous = performance.now() + while (running()) { + await yieldToEventLoop() + const now = performance.now() + intervals.push(now - previous) + previous = now + } +} + +/** What a search would still return: rows whose session row is still there. */ +function visibleRows(db: SyncDatabase): number { + return ( + db + .prepare(`SELECT count(*) AS n FROM messages m JOIN sessions s ON s.id = m.session_row_id`) + .get() as { n: number } + ).n +} + +const root = await mkdtemp(join(tmpdir(), 'orca-search-retention-bench-')) +try { + for (const mode of ['whole-file', 'batched', 'batched-pinned-reader']) { + const path = join(root, `${mode}.sqlite`) + const errors: unknown[] = [] + const store = new SessionSearchStore(path, (error) => errors.push(error)) + let reader: SyncDatabase | null = null + try { + const write = store.beginWrite(syntheticCandidate(), 'replace', 0)! + for (const message of userMessages( + 'synthetic benchmark needle repeated context for a representative coding conversation with commands and paths src/example.ts', + ROWS + )) { + write.add(message) + } + assert.equal( + write.commit({ + session: syntheticSession(), + byteOffset: 4096, + incomplete: false + }), + true + ) + assert.deepEqual(errors, []) + // Truncating first is what makes walBytes below the purge's own growth. + const checkpoint = new SyncDatabase(path) + checkpoint.pragma('wal_checkpoint(TRUNCATE)') + checkpoint.close() + if (mode === 'batched-pinned-reader') { + reader = new SyncDatabase(path, { readonly: true }) + reader.exec('BEGIN') + reader.prepare('SELECT count(*) FROM messages').get() + } + const probe = new SyncDatabase(path, { readonly: true }) + const intervals: number[] = [] + const started = performance.now() + if (mode === 'whole-file') { + const raw = new SyncDatabase(path) + try { + raw.exec('BEGIN IMMEDIATE') + const ids = raw.prepare('SELECT id FROM messages').all() as { + id: number + }[] + for (const { id } of ids) { + raw.prepare('DELETE FROM messages_fts WHERE rowid=?').run(id) + } + raw.exec('DELETE FROM messages; DELETE FROM sessions; DELETE FROM files; COMMIT') + } finally { + raw.close() + } + intervals.push(performance.now() - started) + } else { + let purging = true + const purge = store.purgeOlderThan(Date.now() + 60_000) + // Hiding is immediate: cutting the session loose from its file is the + // first transaction, so a read one turn in already sees nothing, long + // before the rows are gone. + const hiddenEarly = yieldToEventLoop().then(() => visibleRows(probe)) + const sampler = sampleLoopStalls(() => purging, intervals) + await purge + purging = false + await sampler + assert.equal(await hiddenEarly, 0) + assert.deepEqual(errors, []) + } + const wallMs = performance.now() - started + probe.close() + reader?.exec('COMMIT') + reader?.close() + reader = null + const after = new SyncDatabase(path, { readonly: true }) + try { + for (const table of ['messages_fts']) { + assert.equal( + ( + after.prepare(`SELECT count(*) AS n FROM ${table}`).get() as { + n: number + } + ).n, + 0 + ) + } + } finally { + after.close() + } + const walBytes = (await stat(`${path}-wal`)).size + intervals.sort((a, b) => a - b) + console.log( + JSON.stringify({ + mode, + platform: process.platform, + node: process.version, + rows: ROWS, + wallMs: Math.round(wallMs), + samples: intervals.length, + maxStepMs: Math.round(intervals.at(-1) ?? 0), + p95StepMs: Math.round(intervals[Math.floor(intervals.length * 0.95)] ?? 0), + walBytes + }) + ) + } finally { + reader?.close() + store.close() + } + } +} finally { + await rm(root, { recursive: true, force: true }) +} diff --git a/config/scripts/session-search-write-benchmark.ts b/config/scripts/session-search-write-benchmark.ts new file mode 100644 index 00000000000..db09275cd98 --- /dev/null +++ b/config/scripts/session-search-write-benchmark.ts @@ -0,0 +1,266 @@ +import assert from 'node:assert/strict' +import { rm, stat } from 'node:fs/promises' +import { join } from 'node:path' +import { setImmediate as yieldToEventLoop } from 'node:timers/promises' +import { + createSessionParseStats, + parseAgentSessionFileCached, + resetSessionParseCacheForTests +} from '../../src/main/ai-vault/session-scanner-parse-cache' +import { resetTranscriptConsumersForTests } from '../../src/main/ai-vault/session-transcript-consumers' +import { requestWholeTranscriptRead } from '../../src/main/ai-vault/session-transcript-reader' +import { registerSessionSearchIndexConsumer } from '../../src/main/ai-vault-search/session-search-index-consumer' +import { SessionSearchStore } from '../../src/main/ai-vault-search/session-search-store' +import { writeSyntheticTranscriptCorpus } from '../../src/main/ai-vault-search/session-search-synthetic-corpus' +import { sessionCandidate } from '../../src/main/ai-vault-search/session-search-transcript-fixtures' +import SyncDatabase from '../../src/main/sqlite/sync-database' + +// Measures the real transcript reader and search store over a synthetic corpus. +// Never point this at a real transcript tree. + +/** + * How long the longest single transaction held the process. + * + * With one transaction per file that is the whole stall a file costs, so it is + * the number the commit ceiling exists to bound. Measured by wrapping `exec`, + * because the writer's transactions are the only ones this benchmark runs. + */ +function recordTransactionDurations(durations: number[]): () => void { + const exec = SyncDatabase.prototype.exec + let started = 0 + SyncDatabase.prototype.exec = function (this: SyncDatabase, sql: string): void { + if (sql === 'BEGIN IMMEDIATE') { + started = performance.now() + } + exec.call(this, sql) + if (sql === 'COMMIT' && started > 0) { + durations.push(performance.now() - started) + started = 0 + } + } + return () => { + SyncDatabase.prototype.exec = exec + } +} + +/** The writer commits synchronously, so a peer chain samples the gap each read leaves. */ +async function sampleLoopStalls(running: () => boolean, stalls: number[]): Promise { + let previous = performance.now() + while (running()) { + await yieldToEventLoop() + const now = performance.now() + stalls.push(now - previous) + previous = now + } +} + +function tableBytes(db: SyncDatabase): Record { + const rows = db.prepare('SELECT name, sum(pgsize) AS bytes FROM dbstat GROUP BY name').all() as { + name: string + bytes: number + }[] + const group = (prefix: string): number => + rows + .filter((row) => row.name === prefix || row.name.startsWith(`${prefix}_`)) + .reduce((sum, row) => sum + row.bytes, 0) + return { + messagesFts: group('messages_fts'), + messages: group('messages') - group('messages_fts'), + sessions: group('sessions'), + total: rows.reduce((sum, row) => sum + row.bytes, 0) + } +} + +function assertIndexedMessages(db: SyncDatabase, expected: number): number { + const { n } = db + .prepare('SELECT count(*) AS n FROM messages m JOIN sessions s ON s.id = m.session_row_id') + .get() as { n: number } + assert.equal(n, expected, 'indexed message count') + return n +} + +async function checkpointedFileBytes(db: SyncDatabase, path: string): Promise { + // Flush committed WAL pages before reporting the final database footprint. + const [checkpoint] = db.pragma('wal_checkpoint(TRUNCATE)') as { busy: number }[] + assert.equal(checkpoint?.busy, 0, 'storage measurement requires a completed checkpoint') + return (await stat(path)).size +} + +// The default corpus puts tool output at about half the message text; set this +// far higher to price the tool-row cap against the real 80-97 % band. +const toolResultWords = Number(process.env.ORCA_SEARCH_BENCH_TOOL_WORDS ?? 200) +const corpus = await writeSyntheticTranscriptCorpus({ toolResultWords }) +const indexPath = join(corpus.root, 'index.sqlite') +try { + const errors: unknown[] = [] + const store = new SessionSearchStore(indexPath, (error) => errors.push(error)) + const unregister = registerSessionSearchIndexConsumer(store) + const stalls: number[] = [] + const transactions: number[] = [] + const restoreExec = recordTransactionDurations(transactions) + let indexing = true + try { + const stats = createSessionParseStats() + const started = performance.now() + const sampler = sampleLoopStalls(() => indexing, stalls) + for (const path of corpus.files) { + await parseAgentSessionFileCached( + await sessionCandidate('claude', path), + process.platform, + stats + ) + } + indexing = false + await sampler + restoreExec() + const rebuildMs = performance.now() - started + assert.deepEqual(errors, []) + + const reader = new SyncDatabase(indexPath, { readonly: true }) + try { + const rows = assertIndexedMessages(reader, corpus.messageCount) + const sessions = ( + reader.prepare('SELECT count(*) AS n FROM sessions').get() as { + n: number + } + ).n + assert.equal(sessions, corpus.files.length) + const bytes = tableBytes(reader) + const perMb = (value: number): number => + Math.round((value / (corpus.transcriptBytes / (1024 * 1024))) * 10) / 10 + const fileBytes = await checkpointedFileBytes(store.connection, indexPath) + stalls.sort((a, b) => a - b) + transactions.sort((a, b) => a - b) + console.log( + JSON.stringify( + { + platform: process.platform, + node: process.version, + transcriptMb: Math.round((corpus.transcriptBytes / (1024 * 1024)) * 100) / 100, + toolResultWords, + sessions, + rows, + rebuildMs: Math.round(rebuildMs), + rowsPerSecond: Math.round(rows / (rebuildMs / 1000)), + transcriptMbPerSecond: + Math.round((corpus.transcriptBytes / (1024 * 1024) / (rebuildMs / 1000)) * 100) / 100, + bytesPerTranscriptMb: { + messagesFts: perMb(bytes.messagesFts), + messages: perMb(bytes.messages), + sessions: perMb(bytes.sessions), + total: perMb(bytes.total) + }, + writeAmplification: Math.round((bytes.total / corpus.transcriptBytes) * 100) / 100, + fileWriteAmplification: Math.round((fileBytes / corpus.transcriptBytes) * 100) / 100, + transactions: transactions.length, + maxTransactionMs: Math.round((transactions.at(-1) ?? 0) * 100) / 100, + maxLoopStallMs: Math.round(stalls.at(-1) ?? 0), + p95LoopStallMs: Math.round(stalls[Math.floor(stalls.length * 0.95)] ?? 0), + loopStallSamples: stalls.length, + parseStats: stats + }, + null, + 2 + ) + ) + } finally { + reader.close() + } + } finally { + indexing = false + restoreExec() + unregister() + resetTranscriptConsumersForTests() + resetSessionParseCacheForTests() + store.close() + } +} finally { + await rm(corpus.root, { recursive: true, force: true }) +} + +// Phase two: one transcript far larger than any real one, to price the ceiling +// that decides whether a file commits once or in chunks. +const largeTurns = Number(process.env.ORCA_SEARCH_BENCH_LARGE_TURNS ?? 23_000) +const large = await writeSyntheticTranscriptCorpus({ + sessions: 1, + turnsPerSession: largeTurns, + seed: 2 +}) +const largeIndexPath = join(large.root, 'index.sqlite') +try { + const errors: unknown[] = [] + const store = new SessionSearchStore(largeIndexPath, (error) => errors.push(error)) + const unregister = registerSessionSearchIndexConsumer(store) + const transactions: number[] = [] + const restoreExec = recordTransactionDurations(transactions) + try { + const stats = createSessionParseStats() + const started = performance.now() + await parseAgentSessionFileCached( + await sessionCandidate('claude', large.files[0]!), + process.platform, + stats + ) + const indexMs = performance.now() - started + restoreExec() + assert.deepEqual(errors, []) + assertIndexedMessages(store.connection, large.messageCount) + transactions.sort((a, b) => a - b) + + // The same file again, over a generation the index already holds. That is + // the pass a growing transcript really costs, and the one whose transaction + // used to be sized by the old session rather than by the chunk being + // written. The drain that reclaims the cut-loose generation runs after the + // commit, so its bounded batches are in `replaceTransactions` too. + const replaceTransactions: number[] = [] + const restoreReplaceExec = recordTransactionDurations(replaceTransactions) + requestWholeTranscriptRead(large.files[0]!) + const replaceStarted = performance.now() + await parseAgentSessionFileCached( + await sessionCandidate('claude', large.files[0]!), + process.platform, + stats + ) + const replaceMs = performance.now() - replaceStarted + // Finishes whatever the scheduled drain has not reached, so the reclaim is + // priced rather than left half done under the next measurement. + const reclaimStarted = performance.now() + await store.purgeOlderThan(null) + const reclaimMs = performance.now() - reclaimStarted + restoreReplaceExec() + assert.deepEqual(errors, []) + assertIndexedMessages(store.connection, large.messageCount) + replaceTransactions.sort((a, b) => a - b) + + console.log( + JSON.stringify( + { + phase: 'single-large-file', + transcriptMb: Math.round((large.transcriptBytes / (1024 * 1024)) * 100) / 100, + indexMs: Math.round(indexMs), + transactions: transactions.length, + maxTransactionMs: Math.round(transactions.at(-1) ?? 0), + replaceMs: Math.round(replaceMs), + replaceTransactions: replaceTransactions.length, + maxReplaceTransactionMs: Math.round(replaceTransactions.at(-1) ?? 0), + reclaimMs: Math.round(reclaimMs), + indexMb: + Math.round( + ((await checkpointedFileBytes(store.connection, largeIndexPath)) / (1024 * 1024)) * + 100 + ) / 100 + }, + null, + 2 + ) + ) + } finally { + restoreExec() + unregister() + resetTranscriptConsumersForTests() + resetSessionParseCacheForTests() + store.close() + } +} finally { + await rm(large.root, { recursive: true, force: true }) +} diff --git a/src/main/ai-vault-search/session-search-content-hash.test.ts b/src/main/ai-vault-search/session-search-content-hash.test.ts new file mode 100644 index 00000000000..1e7232fc855 --- /dev/null +++ b/src/main/ai-vault-search/session-search-content-hash.test.ts @@ -0,0 +1,48 @@ +import { expect, it } from 'vitest' +import { + EMPTY_CONTENT_HASH, + foldContentHash, + isCollapsibleContentHash +} from './session-search-content-hash' +import { userMessages } from './session-search-index-test-fixture' + +it('reaches the same digest whether the prefix arrives whole or in two appends', () => { + const messages = userMessages('turn', 5) + const whole = foldContentHash(EMPTY_CONTENT_HASH, messages) + const resumed = foldContentHash( + foldContentHash(EMPTY_CONTENT_HASH, messages.slice(0, 2)), + messages.slice(2) + ) + + expect(resumed).toEqual(whole) + expect(whole.count).toBe(5) +}) + +it('freezes once the prefix limit is reached so later appends cannot move it', () => { + // Found rather than imported: the limit is the module's business, and a test + // that reads it off the export cannot notice the fold ignoring it. + const capped = foldContentHash(EMPTY_CONTENT_HASH, userMessages('turn', 64)) + expect(capped.count).toBeLessThan(64) + expect(foldContentHash(capped, userMessages('later', 20))).toEqual(capped) +}) + +it('separates two conversations that share an opening prompt', () => { + const shared = userMessages('same opening', 1) + const first = foldContentHash(EMPTY_CONTENT_HASH, [ + ...shared, + { role: 'user', text: 'left', timestamp: null } + ]) + const second = foldContentHash(EMPTY_CONTENT_HASH, [ + ...shared, + { role: 'user', text: 'right', timestamp: null } + ]) + expect(first.hash).not.toBe(second.hash) +}) + +it('refuses to collapse on a prefix too short to mean anything', () => { + const one = foldContentHash(EMPTY_CONTENT_HASH, userMessages('only turn', 1)) + expect(isCollapsibleContentHash(one.hash, one.count)).toBe(false) + const two = foldContentHash(EMPTY_CONTENT_HASH, userMessages('two turns', 2)) + expect(isCollapsibleContentHash(two.hash, two.count)).toBe(true) + expect(isCollapsibleContentHash(null, 9)).toBe(false) +}) diff --git a/src/main/ai-vault-search/session-search-content-hash.ts b/src/main/ai-vault-search/session-search-content-hash.ts new file mode 100644 index 00000000000..a9d2ab229da --- /dev/null +++ b/src/main/ai-vault-search/session-search-content-hash.ts @@ -0,0 +1,45 @@ +import { createHash } from 'node:crypto' +import type { TranscriptMessage } from '../ai-vault/session-transcript-consumers' + +// Why: Claude `--resume` and Codex fork copy the parent transcript into a new +// file under a new session id, so one conversation lands N times in results. +// The shared opening prefix is what identifies the copy; the tail diverges. +const CONTENT_HASH_MESSAGE_LIMIT = 8 +// One shared opening prompt is not evidence of a fork; two turns is. +const CONTENT_HASH_MIN_MESSAGES = 2 + +export type SessionContentHash = { hash: string | null; count: number } + +export const EMPTY_CONTENT_HASH: SessionContentHash = { hash: null, count: 0 } + +/** + * Chained digest over the first `CONTENT_HASH_MESSAGE_LIMIT` messages. Chaining + * (rather than hashing one joined string) makes it resumable, so an `append` + * can finish a prefix a short `replace` started; once the limit is reached the + * value is frozen and later appends leave it untouched. + */ +export function foldContentHash( + previous: SessionContentHash, + messages: readonly TranscriptMessage[] +): SessionContentHash { + let { hash, count } = previous + for (const message of messages) { + if (count >= CONTENT_HASH_MESSAGE_LIMIT) { + break + } + hash = createHash('sha256') + .update(hash ?? '') + .update('\0') + .update(message.role) + .update('\0') + .update(message.text) + .digest('hex') + count += 1 + } + return { hash, count } +} + +/** Sessions collapse only on a hash that covers enough turns to mean anything. */ +export function isCollapsibleContentHash(hash: string | null, count: number): hash is string { + return hash !== null && count >= CONTENT_HASH_MIN_MESSAGES +} diff --git a/src/main/ai-vault-search/session-search-cwd-key.test.ts b/src/main/ai-vault-search/session-search-cwd-key.test.ts new file mode 100644 index 00000000000..24d915eedc9 --- /dev/null +++ b/src/main/ai-vault-search/session-search-cwd-key.test.ts @@ -0,0 +1,37 @@ +import { expect, it } from 'vitest' +import { folderGroupKey } from '../../shared/ai-vault-session-filters' +import { cwdKey } from './session-search-file-records' + +// The sidebar groups sessions by `folderGroupKey`, which is the shared +// normalizer under a `folder:` prefix. A hit's `cwd_key` has to be the same +// string, or joining an indexed hit to a sidebar group returns nothing. +const CASES: [name: string, cwd: string][] = [ + ['a POSIX path', '/repo/app'], + ['a trailing slash', '/repo/app/'], + ['a Windows drive', 'C:\\Users\\me\\repo'], + ['a WSL interop mount', '/mnt/c/Users/me/repo'], + ['a WSL UNC path', '\\\\wsl.localhost\\Ubuntu\\home\\me\\repo'], + ['the wsl$ alias for the same path', '//wsl$/Ubuntu/home/me/repo'], + ['a Linux path from inside WSL', '/home/me/repo'], + ['the filesystem root', '/'] +] + +it.each(CASES)('keys %s exactly as the sidebar does', (_name, cwd) => { + expect(`folder:${cwdKey(cwd)}`).toBe(folderGroupKey(cwd)) +}) + +it('keeps the root as a path rather than collapsing it to nothing', () => { + // An empty key is indistinguishable from "no cwd", and the scope filter builds + // its child prefix as `key + '/'`, which would be `//` for an empty key. + expect(cwdKey('/')).toBe('/') +}) + +it('has no key for a session whose cwd the transcript never recorded', () => { + expect(cwdKey(null)).toBeNull() +}) + +it('folds the two WSL UNC aliases onto one key', () => { + expect(cwdKey('\\\\wsl.localhost\\Ubuntu\\home\\me\\repo')).toBe( + cwdKey('//wsl$/ubuntu/home/me/repo') + ) +}) diff --git a/src/main/ai-vault-search/session-search-file-cursor.ts b/src/main/ai-vault-search/session-search-file-cursor.ts new file mode 100644 index 00000000000..2ba978b00ae --- /dev/null +++ b/src/main/ai-vault-search/session-search-file-cursor.ts @@ -0,0 +1,41 @@ +import type { FileWithMtime } from '../ai-vault/session-scanner-types' + +// Why the index keeps its own cursor: the parse cache's cursor answers "what +// does the session list already show", which is a different question from "what +// bytes of this file are already rows". They diverge the moment either side +// declines a read, so neither may consult the other. + +/** Filesystem identity, when discovery could prove it. */ +export type SessionSearchFileIdentity = { dev: number; ino: number } | null + +/** + * What the index holds for one transcript. + * + * A null `byteOffset` is a file the index holds rows for and cannot continue: + * a chunked read committed a prefix, and the reader only hands out an offset + * when a read finishes. Null rather than a flag because every caller that does + * arithmetic on the offset then has to say what it means here, at compile time, + * instead of ignoring a boolean it did not know to read. + */ +export type SessionSearchIndexedFile = { + byteOffset: number | null + mtimeMs: number + sizeBytes: number | null +} + +/** + * Whether this file has to be read from the start, whatever its stat says. + * + * The mtime and size are the real ones, so a freshness check that compares only + * those would call a half-written file current and never re-read it. Every such + * check must start here. + */ +export function requiresWholeRead(indexed: SessionSearchIndexedFile | null): boolean { + return indexed !== null && indexed.byteOffset === null +} + +export function fileIdentity(file: FileWithMtime): SessionSearchFileIdentity { + return typeof file.dev === 'number' && typeof file.ino === 'number' + ? { dev: file.dev, ino: file.ino } + : null +} diff --git a/src/main/ai-vault-search/session-search-file-records.ts b/src/main/ai-vault-search/session-search-file-records.ts new file mode 100644 index 00000000000..2ee79cbdc13 --- /dev/null +++ b/src/main/ai-vault-search/session-search-file-records.ts @@ -0,0 +1,134 @@ +import { fileIdentity } from './session-search-file-cursor' +import type { AiVaultSession } from '../../shared/ai-vault-types' +import type { TranscriptSessionIdentity } from '../ai-vault/session-transcript-consumers' +import type { SessionFileCandidate } from '../ai-vault/session-scanner-types' +import type SyncDatabase from '../sqlite/sync-database' +import { EMPTY_CONTENT_HASH, type SessionContentHash } from './session-search-content-hash' +import { normalizeRuntimePathForComparison } from '../../shared/cross-platform-path' + +/** + * The stored comparison key for a session's working directory. + * + * Why the shared normalizer verbatim: the sidebar already groups sessions by + * `folderGroupKey`, which is this function under a prefix. A second spelling + * here means any later join between an indexed hit and a sidebar group returns + * nothing. An earlier version qualified a WSL cwd with its distro so two + * distros could not collide at `/home/me/repo`; that is a real collision, but it + * is one every SSH host has too, neither key qualifies for SSH, and the fix for + * it is a column that names the execution host, not a path key that only some + * hosts spell differently. + */ +export function cwdKey(cwd: string | null): string | null { + return cwd ? normalizeRuntimePathForComparison(cwd) : null +} + +export class SessionSearchFileRecords { + constructor(private readonly db: SyncDatabase) {} + /** + * The row a read hangs its messages off, before the parser has said what the + * session is. The same transaction fills it in: from the decoded session when + * the read finished, and from `updateProvisionalSession` when this is a chunk + * of one that has not. + */ + createSessionRow(candidate: SessionFileCandidate): number { + return Number( + this.db + .prepare( + `INSERT INTO sessions(agent,session_id,file_path,title,resume_command) + VALUES (?,'',?,'','')` + ) + .run(candidate.agent, candidate.file.path).lastInsertRowid + ) + } + + /** + * Writes what the parser knows so far onto a session a chunk is committing. + * + * Rows a chunk commits answer searches the moment they land, so the session + * they hang off has to be nameable before the read producing it ends — and it + * may never end, because a crash between chunks leaves exactly this row. That + * is why the identity is required rather than optional: a read that has none + * does not chunk at all. The final commit overwrites all of it from the + * decoded session; until then the title in particular is provisional. + */ + updateProvisionalSession(rowId: number, identity: TranscriptSessionIdentity): void { + this.db + .prepare( + `UPDATE sessions SET session_id = ?, title = ?, cwd = ?, cwd_key = ?, + created_at = ?, updated_at = ? WHERE id = ?` + ) + .run( + identity.sessionId, + identity.title ?? '', + identity.cwd, + cwdKey(identity.cwd), + identity.createdAt, + identity.updatedAt, + rowId + ) + } + + contentHash(rowId: number): SessionContentHash { + const row = this.db + .prepare('SELECT content_hash, content_hash_count FROM sessions WHERE id = ?') + .get(rowId) as { content_hash: string | null; content_hash_count: number } | undefined + return row ? { hash: row.content_hash, count: row.content_hash_count } : EMPTY_CONTENT_HASH + } + + updateSession(session: AiVaultSession, rowId: number, contentHash: SessionContentHash): void { + const values = [ + session.agent, + session.sessionId, + session.filePath, + session.codexHome, + session.title, + session.cwd, + cwdKey(session.cwd), + session.branch, + session.createdAt, + session.updatedAt, + session.messageCount, + session.resumeCommand, + contentHash.hash, + contentHash.count + ] + this.db + .prepare( + `UPDATE sessions SET agent = ?, session_id = ?, file_path = ?, codex_home = ?, title = ?, + cwd = ?, cwd_key = ?, branch = ?, created_at = ?, updated_at = ?, message_count = ?, resume_command = ?, + content_hash = ?, content_hash_count = ? WHERE id = ?` + ) + .run(...values, rowId) + } + + upsertFile( + candidate: SessionFileCandidate, + byteOffset: number, + sessionRowId: number | null + ): void { + const { file } = candidate + const identity = fileIdentity(file) + this.db + .prepare( + `INSERT INTO files(path, dev, ino, byte_offset, mtime_ms, size_bytes, session_row_id) + VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(path) DO UPDATE SET + -- Partial observations must never create a pair that no stat proved. + dev = CASE WHEN excluded.dev IS NOT NULL AND excluded.ino IS NOT NULL + THEN excluded.dev ELSE files.dev END, + ino = CASE WHEN excluded.dev IS NOT NULL AND excluded.ino IS NOT NULL + THEN excluded.ino ELSE files.ino END, + byte_offset = excluded.byte_offset, mtime_ms = excluded.mtime_ms, + size_bytes = excluded.size_bytes, session_row_id = excluded.session_row_id` + ) + .run( + file.path, + identity?.dev ?? null, + identity?.ino ?? null, + byteOffset, + file.mtimeMs, + file.sizeBytes ?? null, + sessionRowId + ) + } +} diff --git a/src/main/ai-vault-search/session-search-file-write.test.ts b/src/main/ai-vault-search/session-search-file-write.test.ts new file mode 100644 index 00000000000..942565e18d6 --- /dev/null +++ b/src/main/ai-vault-search/session-search-file-write.test.ts @@ -0,0 +1,614 @@ +import { afterEach, beforeEach, expect, it, vi } from 'vitest' +import { resetTranscriptConsumersForTests } from '../ai-vault/session-transcript-consumers' +import SyncDatabase from '../sqlite/sync-database' +import { registerSessionSearchIndexConsumer } from './session-search-index-consumer' +import { cwdKey } from './session-search-file-records' +import { requiresWholeRead } from './session-search-file-cursor' +import { SessionSearchIndexWriter } from './session-search-index-writer' +import { deleteExpiredSearchFiles } from './session-search-retention-delete' +import { + openSessionSearchIndexFile, + replayTranscriptRead, + syntheticCandidate, + syntheticSession, + SYNTHETIC_TRANSCRIPT, + userMessages, + type SessionSearchIndexFile +} from './session-search-index-test-fixture' +import { SessionSearchStore } from './session-search-store' + +// Every assertion here reads through `index.db`, a second connection to the same +// file. That is the whole consistency model: one transaction per file in WAL +// mode, so another handle sees the last committed state and never a session part +// way through being rewritten. + +let index: SessionSearchIndexFile +let store: SessionSearchStore +let errors: unknown[] + +beforeEach(async () => { + index = await openSessionSearchIndexFile('ss-file-write') + errors = [] + store = new SessionSearchStore(index.path, (error) => errors.push(error)) + registerSessionSearchIndexConsumer(store) +}) + +afterEach(async () => { + vi.restoreAllMocks() + resetTranscriptConsumersForTests() + store.close() + await index.close() +}) + +function matches(db: SyncDatabase, table: string, term: string): number { + return ( + db + .prepare( + `SELECT count(*) AS n FROM ${table} JOIN messages m ON m.id = ${table}.rowid + JOIN sessions s ON s.id = m.session_row_id WHERE ${table} MATCH ?` + ) + .get(term) as { n: number } + ).n +} + +/** Fails the nth statement matching `pick`, wherever the writer prepares it. */ +function failOnStatement(pick: (sql: string) => boolean, nth: number): void { + const prepare = SyncDatabase.prototype.prepare + let seen = 0 + vi.spyOn(SyncDatabase.prototype, 'prepare').mockImplementation(function ( + this: SyncDatabase, + sql: string + ) { + if (pick(sql) && ++seen === nth) { + throw new Error('index write crashed mid transaction') + } + return prepare.call(this, sql) + }) +} + +function counts(db: SyncDatabase): Record { + const one = (sql: string): number => (db.prepare(sql).get() as { n: number }).n + return { + sessions: one('SELECT count(*) AS n FROM sessions'), + messages: one('SELECT count(*) AS n FROM messages'), + files: one('SELECT count(*) AS n FROM files'), + full: one('SELECT count(*) AS n FROM messages_fts') + } +} + +it('writes a whole read in one transaction', () => { + replayTranscriptRead({ messages: userMessages('needle text', 300) }) + + const after = counts(index.db) + expect(after.sessions).toBe(1) + expect(after.messages).toBe(300) + expect(after.full).toBe(300) + expect(errors).toEqual([]) +}) + +it('files every row in one FTS table, under the column its role owns', () => { + replayTranscriptRead({ + messages: [ + { role: 'user', text: 'alpha question', timestamp: null }, + { role: 'assistant', text: 'beta answer', timestamp: null }, + { role: 'tool', text: 'gamma tool output', timestamp: null } + ] + }) + + // One table carries all three; the conversation scope is a column filter over + // it, which is what the second table used to be. + expect(counts(index.db).full).toBe(3) + expect(matches(index.db, 'messages_fts', 'gamma')).toBe(1) + expect(matches(index.db, 'messages_fts', '{user_text assistant_text}: gamma')).toBe(0) + expect(matches(index.db, 'messages_fts', '{user_text assistant_text}: beta')).toBe(1) +}) + +it('leaves the index exactly as it found it when a read never finishes', () => { + const write = store.beginWrite(syntheticCandidate(), 'replace', 0)! + for (const message of userMessages('neverfinished', 200)) { + write.add(message) + } + // The process dies here: the rows only ever existed in this buffer. + expect(counts(index.db)).toMatchObject({ + sessions: 0, + messages: 0, + files: 0 + }) +}) + +it('rolls a whole file back when a write throws part way through its transaction', () => { + replayTranscriptRead({ + messages: userMessages('firstgeneration', 3), + outcome: { byteOffset: 40 } + }) + const before = counts(index.db) + + failOnStatement((sql) => sql.startsWith('INSERT INTO messages('), 50) + replayTranscriptRead({ + messages: userMessages('crashedgeneration', 100), + outcome: { byteOffset: 900 } + }) + vi.restoreAllMocks() + + // Not one of the 49 rows that were already inserted survived, the previous + // generation is untouched, and the cursor still describes what is really here. + expect(counts(index.db)).toEqual(before) + expect(matches(index.db, 'messages_fts', 'crashedgeneration')).toBe(0) + expect(matches(index.db, 'messages_fts', 'firstgeneration')).toBe(3) + expect(store.indexedFile(SYNTHETIC_TRANSCRIPT, null)?.byteOffset).toBe(40) + expect(errors).toHaveLength(1) + // The file is owed a re-read, which is the only reason anything was lost. + expect(store.takeStale().map((candidate) => candidate.file.path)).toEqual([SYNTHETIC_TRANSCRIPT]) + + // And the connection is usable again: a transaction left open by the failure + // would take down every write after it, not just the one that threw. + replayTranscriptRead({ + messages: userMessages('afterthecrash', 2), + outcome: { byteOffset: 900 } + }) + expect(matches(index.db, 'messages_fts', 'afterthecrash')).toBe(2) + expect(store.indexedFile(SYNTHETIC_TRANSCRIPT, null)?.byteOffset).toBe(900) +}) + +it('takes the rows back when recording the cursor is what fails', () => { + replayTranscriptRead({ + messages: userMessages('firstgeneration', 3), + outcome: { byteOffset: 40 } + }) + + // The cursor is written last, so this is the crash point that would leave rows + // no cursor describes: a later append would continue from an offset those rows + // already cover, and index the same span twice. + failOnStatement((sql) => sql.startsWith('INSERT INTO files('), 1) + replayTranscriptRead({ + messages: userMessages('crashedgeneration', 5), + outcome: { byteOffset: 900 } + }) + vi.restoreAllMocks() + + expect(counts(index.db)).toMatchObject({ sessions: 1, messages: 3 }) + expect(matches(index.db, 'messages_fts', 'firstgeneration')).toBe(3) + expect(matches(index.db, 'messages_fts', 'crashedgeneration')).toBe(0) + expect(store.indexedFile(SYNTHETIC_TRANSCRIPT, null)?.byteOffset).toBe(40) +}) + +it('shows a reader on another handle one generation or the other, never a mixture', async () => { + replayTranscriptRead({ + messages: userMessages('firstgeneration', 3), + outcome: { byteOffset: 40 } + }) + expect(counts(index.db).messages).toBe(3) + + const write = store.beginWrite(syntheticCandidate(), 'replace', 0)! + for (const message of userMessages('secondgeneration', 7)) { + write.add(message) + // Every point at which the other handle could issue a query mid-read. + expect(counts(index.db).messages).toBe(3) + expect(matches(index.db, 'messages_fts', 'secondgeneration')).toBe(0) + } + expect( + write.commit({ + session: syntheticSession(), + byteOffset: 900, + incomplete: false + }) + ).toBe(true) + + expect(matches(index.db, 'messages_fts', 'firstgeneration')).toBe(0) + expect(matches(index.db, 'messages_fts', 'secondgeneration')).toBe(7) + // The old three are cut loose, not deleted, so they are still on disk and + // already unreachable; the drain the store scheduled hands them back. + expect(counts(index.db).messages).toBe(10) + await vi.waitFor(() => { + expect(counts(index.db).messages).toBe(7) + }) +}) + +// Four of these fill the 400-char ceiling the two tests below construct. +const CHUNKED_MESSAGE = `chunkedneedle ${'filler '.repeat(12)}nd` + +const PROVISIONAL_IDENTITY = { + sessionId: 'aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee', + cwd: '/repo/app', + title: 'provisional title', + createdAt: '2026-05-01T10:00:00.000Z', + updatedAt: '2026-05-01T10:05:00.000Z' +} + +// Only a read that can name its session chunks at all, so every test below that +// wants a chunk has to supply one. +const named = (): typeof PROVISIONAL_IDENTITY => PROVISIONAL_IDENTITY + +it('leaves the session consistent after every chunk of a file too large for one transaction', () => { + expect(CHUNKED_MESSAGE.length).toBe(100) + const writer = new SessionSearchIndexWriter(index.db, 400) + const write = writer.beginWrite(syntheticCandidate(), 'replace', 0, named)! + for (const [position, message] of userMessages(CHUNKED_MESSAGE, 10).entries()) { + write.add(message) + const rows = counts(index.db).messages + // Four messages per chunk, and nothing else reaches the file between them. + expect(rows).toBe(Math.floor((position + 1) / 4) * 4) + // Whatever landed is a coherent prefix of this session and answers searches. + expect(matches(index.db, 'messages_fts', 'chunkedneedle')).toBe(rows) + if (rows > 0) { + // The cursor a chunk leaves refuses every append rather than inventing an + // offset the reader never gave it. + expect(requiresWholeRead(writer.indexedFile(SYNTHETIC_TRANSCRIPT, null))).toBe(true) + expect(writer.beginWrite(syntheticCandidate(), 'append', 0)).toBeNull() + } + } + expect(counts(index.db).messages).toBe(8) + + expect( + write.commit({ + session: syntheticSession(), + byteOffset: 4096, + incomplete: false + }) + ).toBe(true) + expect(counts(index.db)).toMatchObject({ + sessions: 1, + messages: 10, + full: 10 + }) + expect(writer.indexedFile(SYNTHETIC_TRANSCRIPT, null)?.byteOffset).toBe(4096) +}) + +it('holds the ceiling against a single message larger than it', () => { + const writer = new SessionSearchIndexWriter(index.db, 8000) + const write = writer.beginWrite(syntheticCandidate(), 'replace', 0, named)! + const exec = SyncDatabase.prototype.exec + let opened = 0 + vi.spyOn(SyncDatabase.prototype, 'exec').mockImplementation(function ( + this: SyncDatabase, + sql: string + ) { + if (sql === 'BEGIN IMMEDIATE') { + opened += 1 + } + exec.call(this, sql) + }) + + // One conversation turn, three times the ceiling. Checked once per message, + // this commits all 24,000 characters in a single transaction — the ceiling + // bounds nothing that a message can exceed on its own. + write.add({ role: 'assistant', text: 'a'.repeat(24_000), timestamp: null }) + vi.restoreAllMocks() + + expect(opened).toBe(3) + expect(counts(index.db).messages).toBe(3) + expect( + write.commit({ + session: syntheticSession(), + byteOffset: 4096, + incomplete: false + }) + ).toBe(true) + expect(counts(index.db)).toMatchObject({ sessions: 1, messages: 3 }) +}) + +it('names a session on its first chunk, not only when the read ends', () => { + const writer = new SessionSearchIndexWriter(index.db, 400) + const write = writer.beginWrite(syntheticCandidate(), 'replace', 0, named)! + for (const message of userMessages(CHUNKED_MESSAGE, 10)) { + write.add(message) + } + + // The chunks that landed already answer searches, so the session they hang + // off has to be nameable on another handle before the read ends. This is also + // the whole record a crash between chunks leaves behind. + expect(counts(index.db).messages).toBe(8) + expect( + index.db.prepare('SELECT session_id, cwd, cwd_key, title, created_at FROM sessions').get() + ).toEqual({ + session_id: PROVISIONAL_IDENTITY.sessionId, + cwd: '/repo/app', + cwd_key: cwdKey('/repo/app'), + title: 'provisional title', + created_at: '2026-05-01T10:00:00.000Z' + }) + + // And the decoded session still wins at the end: the mid-read title is + // provisional, never a value the final commit has to defer to. + expect( + write.commit({ + session: syntheticSession({ title: 'the settled title' }), + byteOffset: 4096, + incomplete: false + }) + ).toBe(true) + expect(index.db.prepare('SELECT title FROM sessions').get()).toEqual({ + title: 'the settled title' + }) +}) + +it('commits a whole-file read over the ceiling in one transaction, never a chunk', () => { + // The whole-file readers (Grok, Cursor, Gemini, OpenCode) pass no identity: + // their formats are rewritten in place and have no resumable state to ask. + const writer = new SessionSearchIndexWriter(index.db, 400) + const write = writer.beginWrite(syntheticCandidate(), 'replace', 0)! + const exec = SyncDatabase.prototype.exec + let opened = 0 + vi.spyOn(SyncDatabase.prototype, 'exec').mockImplementation(function ( + this: SyncDatabase, + sql: string + ) { + if (sql === 'BEGIN IMMEDIATE') { + opened += 1 + } + exec.call(this, sql) + }) + + for (const message of userMessages(CHUNKED_MESSAGE, 10)) { + write.add(message) + // Chunking here would publish rows under a session with an empty id, an + // empty title and a null cwd, and an interrupted read would leave that + // prefix answering searches for good. + expect(counts(index.db)).toMatchObject({ sessions: 0, messages: 0, files: 0 }) + } + expect(write.commit({ session: syntheticSession(), byteOffset: 4096, incomplete: false })).toBe( + true + ) + vi.restoreAllMocks() + + expect(opened).toBe(1) + expect(counts(index.db)).toMatchObject({ sessions: 1, messages: 10, full: 10 }) + // And a real cursor, not the partial sentinel a chunk would have left. + expect(writer.indexedFile(SYNTHETIC_TRANSCRIPT, null)?.byteOffset).toBe(4096) +}) + +it('starts chunking only once the parser has an id to name the session with', () => { + const writer = new SessionSearchIndexWriter(index.db, 400) + let decoded: typeof PROVISIONAL_IDENTITY | null = null + const write = writer.beginWrite(syntheticCandidate(), 'replace', 0, () => decoded)! + for (const message of userMessages(CHUNKED_MESSAGE, 4)) { + write.add(message) + } + // Past the ceiling, but the parser has decoded nothing: the buffer keeps + // growing rather than naming a session it cannot name. + expect(counts(index.db).messages).toBe(0) + + decoded = PROVISIONAL_IDENTITY + write.add(userMessages(CHUNKED_MESSAGE, 1)[0]!) + + // Everything held goes with the first chunk that can say what it is. + expect(counts(index.db).messages).toBe(5) + expect(index.db.prepare('SELECT session_id, cwd FROM sessions').get()).toEqual({ + session_id: PROVISIONAL_IDENTITY.sessionId, + cwd: '/repo/app' + }) +}) + +it('reports a chunk-partial file as held, and as one that must be read whole', () => { + const writer = new SessionSearchIndexWriter(index.db, 400) + const write = writer.beginWrite(syntheticCandidate(), 'replace', 0, named)! + for (const message of userMessages(CHUNKED_MESSAGE, 10)) { + write.add(message) + } + + // Held, with no cursor to continue. Reporting nothing here reads as "never + // indexed", so a caller asks for whatever read the parse cache offers, the + // reader picks append, and only a decline heals it a cycle later. + const held = writer.indexedFile(SYNTHETIC_TRANSCRIPT, null) + expect(held).not.toBeNull() + expect(held?.byteOffset).toBeNull() + expect(requiresWholeRead(held)).toBe(true) + expect(held?.mtimeMs).toBe(syntheticCandidate().file.mtimeMs) + + // A file this index has never seen is still the other answer, so the two + // states a caller has to tell apart are distinguishable. + expect(writer.indexedFile('/never-seen.jsonl', null)).toBeNull() + expect(requiresWholeRead(null)).toBe(false) + + // And no offset continues it, including the one the chunk recorded. + for (const offset of [0, -1, 400, 1000]) { + expect(writer.beginWrite(syntheticCandidate(), 'append', offset)).toBeNull() + } +}) + +it('re-reads a chunked file whole when its writer died between chunks', async () => { + const writer = new SessionSearchIndexWriter(index.db, 400) + const abandoned = writer.beginWrite(syntheticCandidate(), 'replace', 0, named)! + for (const message of userMessages(CHUNKED_MESSAGE, 10)) { + abandoned.add(message) + } + expect(counts(index.db).messages).toBe(8) + + // Nothing can continue that prefix, so the only way forward is a whole re-read, + // and that replaces every row the dead writer left. + expect(requiresWholeRead(writer.indexedFile(SYNTHETIC_TRANSCRIPT, null))).toBe(true) + const replacement = writer.beginWrite(syntheticCandidate(), 'replace', 0)! + replacement.add(userMessages('wholereread', 1)[0]!) + expect( + replacement.commit({ + session: syntheticSession(), + byteOffset: 4096, + incomplete: false + }) + ).toBe(true) + // The eight stranded rows stop answering the moment the replace commits, and + // the drain hands them back after it rather than inside it. + expect(counts(index.db)).toMatchObject({ sessions: 1, messages: 9 }) + expect(matches(index.db, 'messages_fts', 'chunkedneedle')).toBe(0) + await deleteExpiredSearchFiles(index.db, null, () => false) + expect(counts(index.db)).toMatchObject({ sessions: 1, messages: 1 }) +}) + +it('stops a chunked read whose file was removed between its chunks', () => { + const writer = new SessionSearchIndexWriter(index.db, 400) + const write = writer.beginWrite(syntheticCandidate(), 'replace', 0, named)! + const messages = userMessages(CHUNKED_MESSAGE, 10) + for (const message of messages.slice(0, 4)) { + write.add(message) + } + expect(counts(index.db).messages).toBe(4) + + writer.removeFile(SYNTHETIC_TRANSCRIPT) + const exec = SyncDatabase.prototype.exec + let opened = 0 + vi.spyOn(SyncDatabase.prototype, 'exec').mockImplementation(function ( + this: SyncDatabase, + sql: string + ) { + if (sql === 'BEGIN IMMEDIATE') { + opened += 1 + } + exec.call(this, sql) + }) + for (const message of messages.slice(4)) { + write.add(message) + } + expect(write.commit({ session: syntheticSession(), byteOffset: 4096, incomplete: false })).toBe( + false + ) + vi.restoreAllMocks() + + // Not one row of the removed source came back. The read stopped at the first + // refusal rather than reopening a transaction it already knows will roll back, + // once for every message left in a file that may be a hundred megabytes. + expect(opened).toBe(1) + expect(counts(index.db)).toMatchObject({ sessions: 0, messages: 0, files: 0, full: 0 }) +}) + +it('fences a first-ever read whose file was removed before it committed', () => { + const candidate = syntheticCandidate({ path: '/never-indexed.jsonl' }) + const write = store.beginWrite(candidate, 'replace', 0)! + for (const message of userMessages('removedbeforefirstcommit', 3)) { + write.add(message) + } + // The path was never indexed, so there is no cursor for the removal to move. + // PR 3's retirement sweep removes exactly these: paths the index deferred over + // budget and never wrote, while the registered consumer is fed concurrently. + store.removeFile('/never-indexed.jsonl') + + expect(write.commit({ session: syntheticSession(), byteOffset: 300, incomplete: false })).toBe( + false + ) + expect(counts(index.db)).toMatchObject({ sessions: 0, messages: 0, files: 0, full: 0 }) +}) + +it('replaces the previous generation without ever showing both', async () => { + replayTranscriptRead({ messages: userMessages('firstgeneration', 10) }) + replayTranscriptRead({ messages: userMessages('secondgeneration', 10) }) + + expect(matches(index.db, 'messages_fts', 'firstgeneration')).toBe(0) + expect(matches(index.db, 'messages_fts', 'secondgeneration')).toBe(10) + await vi.waitFor(() => { + expect(counts(index.db)).toMatchObject({ sessions: 1, messages: 10, full: 10 }) + }) +}) + +it('replaces a generation by cutting the old one loose, not by deleting it inline', async () => { + const writer = new SessionSearchIndexWriter(index.db) + const first = writer.beginWrite(syntheticCandidate(), 'replace', 0)! + for (const message of userMessages('firstgeneration', 200)) { + first.add(message) + } + expect(first.commit({ session: syntheticSession(), byteOffset: 100, incomplete: false })).toBe( + true + ) + const before = index.db.prepare('SELECT id FROM sessions').get() as { id: number } + expect(counts(index.db).messages).toBe(200) + + const second = writer.beginWrite(syntheticCandidate(), 'replace', 0)! + for (const message of userMessages('secondgeneration', 3)) { + second.add(message) + } + expect(second.commit({ session: syntheticSession(), byteOffset: 200, incomplete: false })).toBe( + true + ) + + // The transaction inserted three rows and deleted one, rather than deleting + // two hundred: all 203 are still on disk, and the old 200 already answer + // nothing, because every retrieval joins `sessions`. + expect(counts(index.db)).toMatchObject({ sessions: 1, messages: 203, full: 203 }) + expect(matches(index.db, 'messages_fts', 'firstgeneration')).toBe(0) + expect(matches(index.db, 'messages_fts', 'secondgeneration')).toBe(3) + + // A new session row, with `files` repointed at it in that same transaction. + // AUTOINCREMENT never hands the freed id back while orphans still name it. + const after = index.db.prepare('SELECT id FROM sessions').get() as { id: number } + expect(after.id).toBeGreaterThan(before.id) + expect(index.db.prepare('SELECT session_row_id FROM files').get()).toEqual({ + session_row_id: after.id + }) + + await deleteExpiredSearchFiles(index.db, null, () => false) + expect(counts(index.db)).toMatchObject({ sessions: 1, messages: 3, full: 3 }) +}) + +it('drains what a replace cut loose without being asked', async () => { + replayTranscriptRead({ messages: userMessages('firstgeneration', 200) }) + replayTranscriptRead({ messages: userMessages('secondgeneration', 3) }) + + // The store schedules the reclaim the way it schedules retention's. Hiding a + // generation and never reclaiming it would grow the file by every re-read. + expect(matches(index.db, 'messages_fts', 'firstgeneration')).toBe(0) + await vi.waitFor(() => { + expect(counts(index.db)).toMatchObject({ sessions: 1, messages: 3, full: 3 }) + }) + expect(errors).toEqual([]) +}) + +it('continues a session across an append rather than replaying it', () => { + replayTranscriptRead({ + messages: userMessages('openingturn', 3), + outcome: { byteOffset: 40 } + }) + replayTranscriptRead({ + messages: userMessages('laterturn', 2), + mode: 'append', + previousByteOffset: 40, + outcome: { byteOffset: 90 } + }) + + expect(counts(index.db)).toMatchObject({ sessions: 1, messages: 5 }) + expect(matches(index.db, 'messages_fts', 'openingturn')).toBe(3) + expect(matches(index.db, 'messages_fts', 'laterturn')).toBe(2) + expect(store.indexedFile(SYNTHETIC_TRANSCRIPT, null)?.byteOffset).toBe(90) +}) + +it('stops answering for a removed file the moment it is removed', () => { + replayTranscriptRead({ messages: userMessages('removedneedle', 3) }) + store.removeFile(SYNTHETIC_TRANSCRIPT) + + expect(counts(index.db)).toMatchObject({ + sessions: 0, + messages: 0, + files: 0, + full: 0 + }) + expect(matches(index.db, 'messages_fts', 'removedneedle')).toBe(0) +}) + +it('writes nothing for an incomplete read and owes the file a whole re-read', () => { + replayTranscriptRead({ + messages: userMessages('incompleteread', 300), + outcome: { incomplete: true } + }) + + expect(counts(index.db)).toMatchObject({ + sessions: 0, + messages: 0, + files: 0, + full: 0 + }) + expect(store.pendingFileCount).toBe(1) + expect(errors).toEqual([]) +}) + +it('exposes the handle a composed reader queries through', () => { + replayTranscriptRead({ messages: userMessages('composedreader', 3) }) + + // PR 4's engine reads through this rather than opening a second connection, + // so it sees a write the moment the transaction commits. + expect(store.connection.prepare('SELECT count(*) AS n FROM messages').get()).toEqual({ n: 3 }) +}) + +it('closes twice without turning the second call into an error', () => { + store.close() + // node:sqlite throws ERR_INVALID_STATE on a second close of one handle, and a + // store is closed both by whoever owns it and by a teardown that cannot know. + expect(() => store.close()).not.toThrow() + store = new SessionSearchStore(index.path, (error) => errors.push(error)) +}) diff --git a/src/main/ai-vault-search/session-search-identifier-split.test.ts b/src/main/ai-vault-search/session-search-identifier-split.test.ts new file mode 100644 index 00000000000..24f8a67d5c3 --- /dev/null +++ b/src/main/ai-vault-search/session-search-identifier-split.test.ts @@ -0,0 +1,29 @@ +import { expect, it } from 'vitest' +import { identifierShadowTerms, identifierShadowText } from './session-search-identifier-split' + +it('splits a camel-case symbol into its pieces and keeps the whole', () => { + expect(identifierShadowTerms('call resolveTerminalPath here')).toEqual([ + 'resolveterminalpath', + 'resolve', + 'terminal', + 'path' + ]) +}) + +it('splits a path into its segments and extension', () => { + // The whole path already tokenizes on its own; only the pieces need shadowing. + expect(identifierShadowText('src/main/foo-bar.ts')).toBe('src main foo bar ts') +}) + +it('leaves ordinary prose alone', () => { + expect(identifierShadowTerms('the quick brown fox')).toEqual([]) +}) + +it('shadows a screaming-case constant', () => { + expect(identifierShadowTerms('MAX_RETRIES')).toEqual(['max', 'retries']) +}) + +it('stops at the term limit rather than growing with the message', () => { + const text = Array.from({ length: 50 }, (_unused, index) => `alpha_beta${index}`).join(' ') + expect(identifierShadowTerms(text, 10)).toHaveLength(10) +}) diff --git a/src/main/ai-vault-search/session-search-identifier-split.ts b/src/main/ai-vault-search/session-search-identifier-split.ts new file mode 100644 index 00000000000..e2df1822cfa --- /dev/null +++ b/src/main/ai-vault-search/session-search-identifier-split.ts @@ -0,0 +1,54 @@ +// Identifier shadow terms: `resolveTerminalPath` → `resolve terminal path`, +// `src/main/foo-bar.ts` → `src main foo bar ts`. Stored in a separate FTS5 +// column so a partial identifier still matches; the largest single accuracy +// win measured in the retrieval shoot-out (MRR 0.50 → 0.55). + +const RAW_TOKEN = /[A-Za-z0-9_./-]+/g +const CAMEL_PIECE = /[A-Z]+(?![a-z])|[A-Z][a-z0-9]*|[a-z0-9]+/g +const SEPARATOR = /[_./-]+/ +// Worth shadowing: has a separator, a camel boundary, or is SCREAMING_CASE. +const INTERESTING = /[_./-]|[a-z0-9][A-Z]|^[A-Z]{2,}[0-9_]*$/ +const MIN_TOKEN = 3 +const MAX_TOKEN = 120 +const MIN_PIECE = 2 + +function hasMixedCase(piece: string): boolean { + return /[a-z]/.test(piece) && /[A-Z]/.test(piece) +} + +export function identifierShadowTerms(text: string, limit = 4000): string[] { + const out: string[] = [] + const seen = new Set() + for (const match of text.matchAll(RAW_TOKEN)) { + const token = match[0] + if (token.length < MIN_TOKEN || token.length > MAX_TOKEN || !INTERESTING.test(token)) { + continue + } + const parts: string[] = [] + for (const piece of token.split(SEPARATOR)) { + if (!piece) { + continue + } + parts.push(piece) + if (hasMixedCase(piece)) { + parts.push(...(piece.match(CAMEL_PIECE) ?? [])) + } + } + for (const part of parts) { + const lowered = part.toLowerCase() + if (lowered.length < MIN_PIECE || seen.has(lowered)) { + continue + } + seen.add(lowered) + out.push(lowered) + if (out.length >= limit) { + return out + } + } + } + return out +} + +export function identifierShadowText(text: string, limit?: number): string { + return identifierShadowTerms(text, limit).join(' ') +} diff --git a/src/main/ai-vault-search/session-search-index-consumer.test.ts b/src/main/ai-vault-search/session-search-index-consumer.test.ts new file mode 100644 index 00000000000..ca02333cf0b --- /dev/null +++ b/src/main/ai-vault-search/session-search-index-consumer.test.ts @@ -0,0 +1,373 @@ +import { afterEach, beforeEach, expect, it } from 'vitest' +import { resetTranscriptConsumersForTests } from '../ai-vault/session-transcript-consumers' +import { registerSessionSearchIndexConsumer } from './session-search-index-consumer' +import { + openSessionSearchIndexFile, + replayTranscriptRead, + syntheticCandidate, + syntheticSession, + SYNTHETIC_TRANSCRIPT, + userMessages, + type SessionSearchIndexFile +} from './session-search-index-test-fixture' +import { SessionSearchStore, STALE_PATH_LIMIT } from './session-search-store' + +let index: SessionSearchIndexFile +let store: SessionSearchStore +let errors: unknown[] + +beforeEach(async () => { + index = await openSessionSearchIndexFile('ss-index-consumer') + errors = [] + store = new SessionSearchStore(index.path, (error) => errors.push(error)) + registerSessionSearchIndexConsumer(store) +}) + +afterEach(async () => { + resetTranscriptConsumersForTests() + store.close() + await index.close() +}) + +function indexedMessages(): number { + return ( + index.db.prepare('SELECT count(*) AS n FROM messages').get() as { + n: number + } + ).n +} + +function cursor(): number | null | undefined { + return store.indexedFile(SYNTHETIC_TRANSCRIPT, null)?.byteOffset +} + +it('appends onto its own cursor and carries the content hash forward', async () => { + replayTranscriptRead({ + messages: userMessages('first half', 3), + outcome: { byteOffset: 100 } + }) + const first = index.db + .prepare('SELECT content_hash AS hash, content_hash_count AS count FROM sessions') + .get() as { hash: string; count: number } + + replayTranscriptRead({ + mode: 'append', + previousByteOffset: 100, + messages: userMessages('second half', 2), + outcome: { byteOffset: 220 } + }) + + expect(indexedMessages()).toBe(5) + expect(cursor()).toBe(220) + const second = index.db + .prepare('SELECT content_hash AS hash, content_hash_count AS count FROM sessions') + .get() as { hash: string; count: number } + expect(second.count).toBe(first.count + 2) + expect(second.hash).not.toBe(first.hash) + expect(store.takeStale()).toEqual([]) +}) + +it('appends onto a file it read through and decoded no session from', async () => { + // An excluded Codex worker transcript: read through, nothing to index, and + // still growing. Its cursor is sound, so a re-read of the whole file every + // pass buys nothing. + replayTranscriptRead({ + messages: userMessages('excluded span', 3), + outcome: { session: null, byteOffset: 100 } + }) + expect(cursor()).toBe(100) + expect(store.takeStale()).toEqual([]) + + replayTranscriptRead({ + mode: 'append', + previousByteOffset: 100, + messages: userMessages('decoded at last', 2), + outcome: { byteOffset: 220 } + }) + + expect(indexedMessages()).toBe(2) + expect(cursor()).toBe(220) + expect(store.takeStale()).toEqual([]) +}) + +it('declines an append that starts past its own cursor and records the file', async () => { + replayTranscriptRead({ + messages: userMessages('indexed span', 3), + outcome: { byteOffset: 100 } + }) + + // The session list read further than this index did, so the appended span + // continues from bytes the index never saw. + replayTranscriptRead({ + mode: 'append', + previousByteOffset: 900, + messages: userMessages('unseen span', 4), + outcome: { byteOffset: 1200 } + }) + + expect(indexedMessages()).toBe(3) + expect(cursor()).toBe(100) + expect(store.takeStale().map((candidate) => candidate.file.path)).toEqual([SYNTHETIC_TRANSCRIPT]) +}) + +it('declines a file whose identity changed under the same path', async () => { + const original = syntheticCandidate({ dev: 1, ino: 10 }) + replayTranscriptRead({ + candidate: original, + messages: userMessages('original file', 2), + outcome: { byteOffset: 100 } + }) + + replayTranscriptRead({ + candidate: syntheticCandidate({ dev: 1, ino: 77 }), + mode: 'append', + previousByteOffset: 100, + messages: userMessages('replacement file', 2), + outcome: { byteOffset: 200 } + }) + + expect(indexedMessages()).toBe(2) + expect(store.takeStale()).toHaveLength(1) +}) + +it('never advances the cursor for an incomplete read', async () => { + replayTranscriptRead({ + messages: userMessages('complete span', 3), + outcome: { byteOffset: 100 } + }) + + replayTranscriptRead({ + mode: 'append', + previousByteOffset: 100, + messages: userMessages('partial span', 5), + outcome: { byteOffset: 400, incomplete: true } + }) + + expect(indexedMessages()).toBe(3) + expect(cursor()).toBe(100) + expect( + ( + index.db.prepare('SELECT count(*) AS n FROM messages').get() as { + n: number + } + ).n + ).toBe(3) + expect(store.takeStale()).toHaveLength(1) +}) + +it('indexes nothing at all from a read that was incomplete from the start', async () => { + replayTranscriptRead({ + messages: userMessages('unreachable', 4), + outcome: { byteOffset: 0, incomplete: true } + }) + + expect(index.db.prepare('SELECT count(*) AS n FROM sessions').get()).toEqual({ + n: 0 + }) + expect(index.db.prepare('SELECT count(*) AS n FROM messages').get()).toEqual({ + n: 0 + }) + expect(cursor()).toBeUndefined() +}) + +it('drops a file whose parser returned no session', async () => { + replayTranscriptRead({ + messages: userMessages('was indexed', 3), + outcome: { byteOffset: 100 } + }) + + replayTranscriptRead({ + messages: userMessages('now rejected', 2), + outcome: { session: null, byteOffset: 300 } + }) + + expect(index.db.prepare('SELECT count(*) AS n FROM sessions').get()).toEqual({ + n: 0 + }) + expect(index.db.prepare('SELECT count(*) AS n FROM messages').get()).toEqual({ + n: 0 + }) + // The file is still read through, so a later scan does not re-read it. + expect(cursor()).toBe(300) +}) + +it('writes nothing for a source whose parser cannot reach the channel', async () => { + // An OpenCode SQLite candidate decodes in a worker, so every read of it is + // incomplete, and no re-read would help. + const candidate = { + ...syntheticCandidate({ path: '/opencode/opencode.db#session-1' }), + agent: 'opencode' as const + } + replayTranscriptRead({ + candidate, + messages: [], + outcome: { byteOffset: 0, incomplete: true } + }) + + expect(index.db.prepare('SELECT count(*) AS n FROM sessions').get()).toEqual({ + n: 0 + }) + expect(store.takeStale()).toEqual([]) +}) + +it('ignores a candidate older than the retention cutoff', async () => { + store.setRetentionCutoffMs(Date.now()) + replayTranscriptRead({ messages: userMessages('too old', 3) }) + + expect(index.db.prepare('SELECT count(*) AS n FROM sessions').get()).toEqual({ + n: 0 + }) + expect(store.takeStale()).toEqual([]) +}) + +it('stops writing while the store refuses writes, but remembers what it skipped', async () => { + store.setAcceptingWrites(false) + replayTranscriptRead({ messages: userMessages('paused', 3) }) + + expect(index.db.prepare('SELECT count(*) AS n FROM sessions').get()).toEqual({ + n: 0 + }) + expect(errors).toEqual([]) + // A pause is exactly the window in which every read is declined. Forgetting + // them would leave the whole paused span unindexed with nothing to replay it. + expect(store.takeStale().map((candidate) => candidate.file.path)).toEqual([SYNTHETIC_TRANSCRIPT]) +}) + +it('keeps the paused re-read set when the retention window is reconfigured', async () => { + store.setAcceptingWrites(false) + replayTranscriptRead({ messages: userMessages('paused', 2) }) + expect(store.pendingFileCount).toBe(1) + + // The set records what still has to be read, not what is worth keeping. A + // window that now excludes this file is enforced where the re-read is + // dispatched, so nothing is written and the file leaves the set there. + store.setRetentionCutoffMs(Date.now()) + expect(store.pendingFileCount).toBe(1) + + store.setAcceptingWrites(true) + expect(store.takeStale()).toHaveLength(1) + replayTranscriptRead({ messages: userMessages('outside the window now', 2) }) + expect(index.db.prepare('SELECT count(*) AS n FROM sessions').get()).toEqual({ + n: 0 + }) + expect(store.pendingFileCount).toBe(0) +}) + +it('drops the oldest record rather than growing without a bound, and says so', () => { + store.setAcceptingWrites(false) + for (let index = 0; index < STALE_PATH_LIMIT + 5; index++) { + store.markStale(syntheticCandidate({ path: `/transcript-${index}.jsonl` })) + } + + expect(store.pendingFileCount).toBe(STALE_PATH_LIMIT) + expect(store.droppedPendingFileCount).toBe(5) + const kept = store.takeStale().map((candidate) => candidate.file.path) + expect(kept).not.toContain('/transcript-0.jsonl') + expect(kept).toContain(`/transcript-${STALE_PATH_LIMIT + 4}.jsonl`) +}) + +it('keeps the session list running when the index write fails', async () => { + replayTranscriptRead({ + messages: userMessages('healthy', 2), + outcome: { byteOffset: 100 } + }) + index.db.exec('DROP TABLE messages_fts') + + expect(() => + replayTranscriptRead({ + mode: 'append', + previousByteOffset: 100, + messages: userMessages('broken', 400), + outcome: { byteOffset: 500 } + }) + ).not.toThrow() + expect(errors.length).toBeGreaterThan(0) + expect(store.takeStale()).toHaveLength(1) +}) + +it('unregisters cleanly, leaving later reads unindexed', async () => { + resetTranscriptConsumersForTests() + replayTranscriptRead({ messages: userMessages('after unregister', 3) }) + + expect(index.db.prepare('SELECT count(*) AS n FROM sessions').get()).toEqual({ + n: 0 + }) +}) + +it('drops a removed source and keeps its cursor gone', async () => { + replayTranscriptRead({ + messages: userMessages('present', 3), + outcome: { byteOffset: 100 } + }) + store.removeFile(SYNTHETIC_TRANSCRIPT) + + expect(cursor()).toBeUndefined() + expect(index.db.prepare('SELECT count(*) AS n FROM sessions').get()).toEqual({ + n: 0 + }) + expect(index.db.prepare('SELECT count(*) AS n FROM messages').get()).toEqual({ + n: 0 + }) +}) + +it('writes the session metadata the read decoded', async () => { + replayTranscriptRead({ + messages: userMessages('metadata', 1), + outcome: { + session: syntheticSession({ + sessionId: 'abc-123', + title: 'a titled session', + cwd: '/repo/app', + branch: 'main', + messageCount: 1, + resumeCommand: 'claude --resume abc-123' + }), + byteOffset: 42 + } + }) + + expect( + index.db + .prepare('SELECT session_id, title, cwd, cwd_key, branch, resume_command FROM sessions') + .get() + ).toEqual({ + session_id: 'abc-123', + title: 'a titled session', + cwd: '/repo/app', + cwd_key: '/repo/app', + branch: 'main', + resume_command: 'claude --resume abc-123' + }) +}) + +it('keeps a proven file identity when a later read cannot stat it', async () => { + const withIdentity = syntheticCandidate({ dev: 1, ino: 10 }) + replayTranscriptRead({ + candidate: withIdentity, + messages: userMessages('first', 2), + outcome: { byteOffset: 100 } + }) + + // A host that cannot prove identity re-reads the same file. + replayTranscriptRead({ + candidate: syntheticCandidate(), + mode: 'append', + previousByteOffset: 100, + messages: userMessages('second', 2), + outcome: { byteOffset: 200 } + }) + expect(indexedMessages()).toBe(4) + + // The stored identity survived, so a rename-replace is still detectable. + replayTranscriptRead({ + candidate: syntheticCandidate({ dev: 1, ino: 99 }), + mode: 'append', + previousByteOffset: 200, + messages: userMessages('replacement', 2), + outcome: { byteOffset: 300 } + }) + + expect(indexedMessages()).toBe(4) + expect(cursor()).toBe(200) + expect(store.takeStale()).toHaveLength(1) +}) diff --git a/src/main/ai-vault-search/session-search-index-consumer.ts b/src/main/ai-vault-search/session-search-index-consumer.ts new file mode 100644 index 00000000000..e4fa6da4a8e --- /dev/null +++ b/src/main/ai-vault-search/session-search-index-consumer.ts @@ -0,0 +1,118 @@ +import { parserPublishesMessages } from '../ai-vault/session-scanner-agent-parser' +import { + registerTranscriptConsumer, + type TranscriptConsumer, + type TranscriptMessage, + type TranscriptReadConsumer, + type TranscriptReadOutcome, + type TranscriptReadStart +} from '../ai-vault/session-transcript-consumers' +import { fileIdentity } from './session-search-file-cursor' +import type { SessionSearchFileWrite } from './session-search-index-writer' +import type { SessionSearchStore } from './session-search-store' + +/** + * The search index as a consumer of the transcript reader. + * + * It keeps its own cursor in the `files` table and never consults the parse + * cache: the two answer different questions and diverge the moment either + * declines a read. Three refusals, each of which leaves the cursor where it + * was and records the file for a later whole re-read: + * + * - `beginRead` returns null when this index's cursor is behind the offset an + * `append` continues from, or when the file's identity changed. + * - a buffering failure stops the read's rows without failing the session list. + * - an `incomplete` outcome never commits; those rows are not the whole span. + */ +export class SessionSearchIndexConsumer implements TranscriptConsumer { + constructor(private readonly store: SessionSearchStore) {} + + beginRead(start: TranscriptReadStart): TranscriptReadConsumer | null { + const { candidate } = start + if (!this.store.acceptsCandidate(candidate)) { + // A pause is a reason not to write now, not a reason to forget the read. + // `markStale` applies the retention rule itself, so a candidate that is + // out of scope rather than merely paused is still dropped here. + this.store.markStale(candidate) + return null + } + // A parser that decodes where the channel cannot reach it reports every read + // as incomplete. Declining here is not the same as being behind: no re-read + // would help, so the file is not recorded either. + if (!parserPublishesMessages(candidate)) { + return null + } + if (start.mode === 'append') { + const cursor = this.store.indexedFile(candidate.file.path, fileIdentity(candidate.file)) + if (!cursor || cursor.byteOffset !== start.previousByteOffset) { + // This index never saw the span before `previousByteOffset`; appending + // here would leave a hole no later read can fill. A null cursor is the + // file a chunked read left half written, which no offset continues. + this.store.markStale(candidate) + return null + } + } + const write = this.store.beginWrite( + candidate, + start.mode, + start.previousByteOffset, + start.identity + ) + if (!write) { + this.store.markStale(candidate) + return null + } + return new SessionSearchReadConsumer(this.store, start, write) + } +} + +class SessionSearchReadConsumer implements TranscriptReadConsumer { + private failed = false + + constructor( + private readonly store: SessionSearchStore, + private readonly start: TranscriptReadStart, + private readonly write: SessionSearchFileWrite + ) {} + + message(message: TranscriptMessage): void { + if (this.failed) { + return + } + try { + this.write.add(message) + } catch (error) { + // Never throws back into the reader: the channel would drop this consumer + // for the rest of the read and `finish` would never run. Failing here + // keeps the whole read on one path — the buffer is dropped and the file is + // re-read. + this.failed = true + this.store.reportWriteFailure(error) + } + } + + finish(outcome: TranscriptReadOutcome): void { + const { candidate } = this.start + let committed = false + try { + // An incomplete read's rows are not the whole span, so the cursor must not + // move past them; the file is re-read whole instead. + committed = !this.failed && !outcome.incomplete && this.write.commit(outcome) + } catch (error) { + this.store.reportWriteFailure(error) + } + if (committed) { + this.store.writeCommitted(candidate) + return + } + this.store.markStale(candidate) + } +} + +/** + * Registers the index with the reader and returns the unregister function. + * Nothing in production calls this yet: PR 3 owns when the index is live. + */ +export function registerSessionSearchIndexConsumer(store: SessionSearchStore): () => void { + return registerTranscriptConsumer(new SessionSearchIndexConsumer(store)) +} diff --git a/src/main/ai-vault-search/session-search-index-test-fixture.ts b/src/main/ai-vault-search/session-search-index-test-fixture.ts new file mode 100644 index 00000000000..baa3e2976fe --- /dev/null +++ b/src/main/ai-vault-search/session-search-index-test-fixture.ts @@ -0,0 +1,124 @@ +import { mkdtemp } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { removeTree } from '../../shared/windows-transient-lock-removal' +import type { AiVaultSession } from '../../shared/ai-vault-types' +import type { SessionFileCandidate } from '../ai-vault/session-scanner-types' +import { TranscriptMessageChannel } from '../ai-vault/session-transcript-channel' +import type { + TranscriptMessage, + TranscriptReadOutcome, + TranscriptReadStart +} from '../ai-vault/session-transcript-consumers' +import type SyncDatabase from '../sqlite/sync-database' +import { openSessionSearchDatabase } from './session-search-schema' + +export const SYNTHETIC_TRANSCRIPT = 'synthetic-transcript' + +export function syntheticCandidate( + overrides: Partial = {} +): SessionFileCandidate { + const at = new Date(1740000000000) + return { + agent: 'claude', + codexHome: null, + file: { + path: SYNTHETIC_TRANSCRIPT, + mtimeMs: at.getTime(), + modifiedAt: at.toISOString(), + sizeBytes: 4096, + ...overrides + } + } +} + +export function syntheticSession(overrides: Partial = {}): AiVaultSession { + const at = new Date(1740000000000).toISOString() + return { + id: 'fixture', + executionHostId: 'local', + agent: 'claude', + sessionId: 'fixture', + title: 'fixture session', + cwd: '/fixture', + branch: null, + model: null, + filePath: SYNTHETIC_TRANSCRIPT, + codexHome: null, + createdAt: at, + updatedAt: at, + modifiedAt: at, + messageCount: 0, + totalTokens: 0, + previewMessages: [], + queuedMessageCount: 0, + subagentTranscriptCount: 0, + resumeCommand: '', + subagent: null, + ...overrides + } +} + +export function userMessages(text: string, count: number): TranscriptMessage[] { + return Array.from({ length: count }, (_unused, index) => ({ + role: 'user' as const, + text, + timestamp: new Date(1740000000000 + index * 1000).toISOString() + })) +} + +/** + * Drives one read through the real fan-out channel, so a test exercises the + * registration path the transcript reader uses rather than the consumer alone. + */ +export function replayTranscriptRead(args: { + candidate?: SessionFileCandidate + mode?: TranscriptReadStart['mode'] + previousByteOffset?: number + messages: TranscriptMessage[] + outcome?: Partial +}): void { + const candidate = args.candidate ?? syntheticCandidate() + const mode = args.mode ?? 'replace' + const channel = new TranscriptMessageChannel() + channel.beginRead({ + candidate, + mode, + previousByteOffset: args.previousByteOffset ?? 0 + }) + for (const message of args.messages) { + channel.push(message) + } + channel.finishRead({ + session: syntheticSession(), + byteOffset: 4096, + incomplete: false, + ...args.outcome + }) +} + +export type SessionSearchIndexFile = { + path: string + /** The store keeps its own connection private, so row assertions need this one. */ + db: SyncDatabase + close: () => Promise +} + +/** An on-disk index: `:memory:` is per-connection, so a second reader needs a real file. */ +export async function openSessionSearchIndexFile(name: string): Promise { + const root = await mkdtemp(join(tmpdir(), `${name}-`)) + const path = join(root, 'index.sqlite') + const db = openSessionSearchDatabase(path) + let open = true + return { + path, + db, + close: async () => { + if (open) { + open = false + db.close() + } + await removeTree(root) + } + } +} diff --git a/src/main/ai-vault-search/session-search-index-writer.test.ts b/src/main/ai-vault-search/session-search-index-writer.test.ts new file mode 100644 index 00000000000..46d147dcce0 --- /dev/null +++ b/src/main/ai-vault-search/session-search-index-writer.test.ts @@ -0,0 +1,228 @@ +import { afterEach, beforeEach, expect, it } from 'vitest' +import { SessionSearchIndexConsumer } from './session-search-index-consumer' +import { + openSessionSearchIndexFile, + syntheticCandidate, + syntheticSession, + SYNTHETIC_TRANSCRIPT, + userMessages, + type SessionSearchIndexFile +} from './session-search-index-test-fixture' +import { SessionSearchStore } from './session-search-store' + +// The store is driven directly here. Every guard below is also shadowed by the +// consumer's own check, so a test that goes through the consumer proves nothing +// about which of the two is holding. + +let index: SessionSearchIndexFile +let store: SessionSearchStore +let errors: unknown[] + +beforeEach(async () => { + index = await openSessionSearchIndexFile('ss-index-writer') + errors = [] + store = new SessionSearchStore(index.path, (error) => errors.push(error)) +}) + +afterEach(async () => { + store.close() + await index.close() +}) + +function count(table: string): number { + return ( + index.db.prepare(`SELECT count(*) AS n FROM ${table}`).get() as { + n: number + } + ).n +} + +function indexRead(previousByteOffset: number, byteOffset: number, text: string): boolean { + const write = store.beginWrite( + syntheticCandidate(), + previousByteOffset === 0 ? 'replace' : 'append', + previousByteOffset + ) + if (!write) { + return false + } + for (const message of userMessages(text, 2)) { + write.add(message) + } + return write.commit({ + session: syntheticSession(), + byteOffset, + incomplete: false + }) +} + +it('refuses an append whose predecessor offset is not the committed cursor', () => { + expect(indexRead(0, 100, 'first')).toBe(true) + + expect(store.beginWrite(syntheticCandidate(), 'append', 900)).toBeNull() + expect(store.beginWrite(syntheticCandidate(), 'append', 99)).toBeNull() + // The one offset that does continue the committed span is accepted. + expect(store.beginWrite(syntheticCandidate(), 'append', 100)).not.toBeNull() +}) + +it('refuses to commit a write whose cursor moved underneath it', () => { + const stale = store.beginWrite(syntheticCandidate(), 'replace', 0)! + for (const message of userMessages('stalegeneration', 40)) { + stale.add(message) + } + // A second read of the same path finishes first. Without the parse file lane + // this is the overlap that would otherwise resurrect the stale rows. + expect(indexRead(0, 200, 'winninggeneration')).toBe(true) + + expect( + stale.commit({ + session: syntheticSession(), + byteOffset: 100, + incomplete: false + }) + ).toBe(false) + expect(store.indexedFile(SYNTHETIC_TRANSCRIPT, null)?.byteOffset).toBe(200) + expect(count('sessions')).toBe(1) + expect(count('messages')).toBe(2) + expect(errors).toEqual([]) +}) + +it('refuses to commit a write whose file was removed mid-read', () => { + expect(indexRead(0, 100, 'firstgeneration')).toBe(true) + const write = store.beginWrite(syntheticCandidate(), 'append', 100)! + for (const message of userMessages('afterremoval', 10)) { + write.add(message) + } + store.removeFile(SYNTHETIC_TRANSCRIPT) + + // Committing here would put a source back that its owner proved was deleted. + expect( + write.commit({ + session: syntheticSession(), + byteOffset: 300, + incomplete: false + }) + ).toBe(false) + expect(store.indexedFile(SYNTHETIC_TRANSCRIPT, null)).toBeNull() + expect(count('sessions')).toBe(0) + expect(count('messages')).toBe(0) + expect(count('files')).toBe(0) +}) + +it('declines a behind cursor in beginRead before it ever reaches the store', () => { + const attempted: number[] = [] + const stub = { + acceptsCandidate: () => true, + indexedFile: () => ({ byteOffset: 100, mtimeMs: 1, sizeBytes: 1 }), + beginWrite: (_candidate: unknown, _mode: unknown, previousByteOffset: number) => { + attempted.push(previousByteOffset) + return { add: () => undefined, commit: () => true } + }, + markStale: () => undefined + } as unknown as SessionSearchStore + const consumer = new SessionSearchIndexConsumer(stub) + + expect( + consumer.beginRead({ + candidate: syntheticCandidate(), + mode: 'append', + previousByteOffset: 900 + }) + ).toBeNull() + // The store was never asked, so the writer's own guard cannot be what refused. + expect(attempted).toEqual([]) + expect( + consumer.beginRead({ + candidate: syntheticCandidate(), + mode: 'append', + previousByteOffset: 100 + }) + ).not.toBeNull() + expect(attempted).toEqual([100]) +}) + +it("hands the read's identity accessor to the store", () => { + const captured: unknown[] = [] + const stub = { + acceptsCandidate: () => true, + indexedFile: () => null, + beginWrite: ( + _candidate: unknown, + _mode: unknown, + _previousByteOffset: unknown, + identity: unknown + ) => { + captured.push(identity) + return { add: () => undefined, commit: () => true } + }, + markStale: () => undefined + } as unknown as SessionSearchStore + const identity = (): null => null + + new SessionSearchIndexConsumer(stub).beginRead({ + candidate: syntheticCandidate(), + mode: 'replace', + previousByteOffset: 0, + identity + }) + + // Dropped here, a chunked read writes rows under a session with no id and no + // cwd for as long as the read lasts, and for ever if it crashes first. + expect(captured).toEqual([identity]) +}) + +it('treats half a recorded identity as no identity at all', () => { + // New partial observations are not stored as identities. + const partial = { + ...syntheticCandidate({ dev: 7 }), + agent: 'claude' as const + } + const write = store.beginWrite(partial, 'replace', 0)! + for (const message of userMessages('halfidentity', 2)) { + write.add(message) + } + write.commit({ + session: syntheticSession(), + byteOffset: 100, + incomplete: false + }) + expect(index.db.prepare('SELECT dev, ino FROM files').get()).toEqual({ + dev: null, + ino: null + }) + // Older indexes may still carry a half-pair. + index.db.exec('UPDATE files SET dev = 7') + + // One matching number is not proof of sameness, and one mismatching number is + // not proof of replacement. Neither compares, so neither declines. + expect(store.indexedFile(SYNTHETIC_TRANSCRIPT, { dev: 7, ino: 99 })?.byteOffset).toBe(100) + expect(store.indexedFile(SYNTHETIC_TRANSCRIPT, { dev: 8, ino: 99 })?.byteOffset).toBe(100) + expect(store.beginWrite(syntheticCandidate({ dev: 8, ino: 99 }), 'append', 100)).not.toBeNull() +}) + +it.each([ + [null, { dev: null, ino: null }], + [ + { dev: 7, ino: 11 }, + { dev: 7, ino: 11 } + ] +])('never combines partial stats with the previous identity %j', (initial, expected) => { + const observations = [initial ?? {}, { dev: 9 }, { ino: 13 }, { dev: 17, ino: 19 }] + for (const [position, identity] of observations.entries()) { + const write = store.beginWrite( + syntheticCandidate(identity), + position ? 'append' : 'replace', + position * 100 + )! + expect( + write.commit({ + session: syntheticSession(), + byteOffset: (position + 1) * 100, + incomplete: false + }) + ).toBe(true) + expect(index.db.prepare('SELECT dev, ino FROM files').get()).toEqual( + position === 3 ? { dev: 17, ino: 19 } : expected + ) + } +}) diff --git a/src/main/ai-vault-search/session-search-index-writer.ts b/src/main/ai-vault-search/session-search-index-writer.ts new file mode 100644 index 00000000000..a5c981b5bb2 --- /dev/null +++ b/src/main/ai-vault-search/session-search-index-writer.ts @@ -0,0 +1,359 @@ +import type SyncDatabase from '../sqlite/sync-database' +import type { SessionFileCandidate } from '../ai-vault/session-scanner-types' +import type { + TranscriptMessage, + TranscriptReadOutcome, + TranscriptSessionIdentity +} from '../ai-vault/session-transcript-consumers' +import { EMPTY_CONTENT_HASH, foldContentHash } from './session-search-content-hash' +import type { + SessionSearchFileIdentity, + SessionSearchIndexedFile +} from './session-search-file-cursor' +import { SessionSearchFileRecords } from './session-search-file-records' +import { + deleteSearchMessages, + insertSearchMessage, + searchMessageRows +} from './session-search-message-rows' + +/** + * How much decoded text one transaction may carry. + * + * A file's rows are buffered in memory and written in one transaction, so the + * whole read is either in the index or not. The ceiling is what keeps that + * promise affordable: at the measured 26 MB of transcript per second it caps a + * single commit near a second and the WAL it produces near 64 MB, and it is far + * above the largest real transcript (the 40-session benchmark corpus is 10.5 MB + * in total), so an ordinary file never reaches it. Above the ceiling the read is + * cut into chunks that each leave the index consistent — but only a read that + * can name its session chunks at all. See `add`. + */ +export const SESSION_SEARCH_COMMIT_CHARS = 32 * 1024 * 1024 + +/** + * The cursor of a file whose rows are a prefix, written by a chunk of a read + * that has not reached the end of the file. + * + * The reader hands out byte offsets only when a read finishes, so a chunk has + * no honest offset to record. This one is unusable on purpose: `indexedFile` + * reports no cursor for it, so an append is declined and the file is re-read + * whole. The rows are still a coherent prefix of that session and answer + * searches until the re-read replaces them. + */ +const PARTIAL_FILE_CURSOR = -1 + +type FileRow = { + dev: number | null + ino: number | null + byte_offset: number + mtime_ms: number + size_bytes: number | null + session_row_id: number | null +} + +type FileCursor = Pick + +export type SessionSearchFileWrite = { + /** + * Buffers one message, committing a chunk when the buffer reaches the ceiling + * — and only while this read can name the session it is writing. + * + * A chunk's rows answer searches the moment they land, so a read with no + * `identity` would publish them under a session with an empty id, an empty + * title and a null cwd, and an interrupted read would leave that prefix + * behind for good. The readers that supply no identity are the whole-file + * ones (Grok, Cursor, Gemini, OpenCode), whose formats are rewritten in place + * and have no resumable state to ask; they are also small — the largest on + * the author's machine is 5 MB — so buffering one to the end and committing + * it whole costs nothing. Chunking stays reserved for the readers that can + * say which session this is before the read ends. + */ + add(message: TranscriptMessage): void + /** + * Writes this file's rows, its session and its cursor in one transaction. + * False when the file's record changed under this read — it was removed, or + * another writer moved the cursor these rows continue from. A read that never + * calls this leaves the index exactly as it found it, unless it chunked. + */ + commit(outcome: TranscriptReadOutcome): boolean +} + +export class SessionSearchIndexWriter { + private readonly records: SessionSearchFileRecords + // Removals per path, so a write can prove its source was not dropped under it + // rather than infer it from the cursor. In memory is enough: one process owns + // the index, and a removal only has to fence writes this process opened. + private readonly removals = new Map() + + constructor( + private readonly db: SyncDatabase, + private readonly commitChars: number = SESSION_SEARCH_COMMIT_CHARS, + /** + * Called after a transaction that left a session's messages with no session + * row, so the owner can start the bounded drain that reclaims them. + * Synchronous work here would put the cost back where it was taken from. + */ + private readonly onOrphanedRows: () => void = () => undefined + ) { + this.records = new SessionSearchFileRecords(db) + } + + /** + * What the index holds for this file, or null when it holds nothing usable: + * an unknown path, or one whose recorded identity no longer matches. + * + * A file a chunked read left half written is reported, with a null cursor. + * Reporting nothing for it would read as "never indexed", so the caller would + * ask for whatever read the parse cache offers, the reader would pick append, + * and the decline would be the only thing that ever forced the whole read. + */ + indexedFile(path: string, identity: SessionSearchFileIdentity): SessionSearchIndexedFile | null { + const row = this.db + .prepare( + 'SELECT dev, ino, byte_offset, mtime_ms, size_bytes, session_row_id FROM files WHERE path = ?' + ) + .get(path) as FileRow | undefined + if (!row) { + return null + } + // Older indexes can carry half-pairs; only a complete identity can prove replacement. + if (identity && row.dev !== null && row.ino !== null) { + if (row.dev !== identity.dev || row.ino !== identity.ino) { + return null + } + } + return { + byteOffset: row.byte_offset === PARTIAL_FILE_CURSOR ? null : row.byte_offset, + mtimeMs: row.mtime_ms, + sizeBytes: row.size_bytes + } + } + + /** + * Opens a buffered write for one read, or returns null when the read cannot + * extend what the index holds: an `append` whose predecessor byte offset is + * not this index's own cursor covers a span the index never saw. + */ + beginWrite( + candidate: SessionFileCandidate, + mode: 'replace' | 'append', + previousByteOffset: number, + identity?: () => TranscriptSessionIdentity | null + ): SessionSearchFileWrite | null { + const path = candidate.file.path + const cursor = this.cursor(path) + if (mode === 'append') { + // The partial sentinel is not a byte offset, so nothing continues it — + // including a caller that reads it back off the row and passes it in. + if (cursor === undefined || cursor.byte_offset === PARTIAL_FILE_CURSOR) { + return null + } + if (cursor.byte_offset !== previousByteOffset) { + return null + } + } + // A file the index read through and decoded no session from still has a + // cursor worth continuing: it has no session row to hang new rows off, so + // this read makes one. Declining instead would force a whole re-read of + // that file on every pass for as long as it grows. + return this.buffered(candidate, cursor, mode === 'append', identity) + } + + /** + * Drops a source: its session, its rows and its file record, in one + * transaction. Unbounded on purpose — the caller has proven this one file is + * gone and expects it out of results when the call returns, and a read of it + * that is still in flight is fenced by the cursor its commit re-reads. + */ + removeFile(path: string): void { + this.removals.set(path, (this.removals.get(path) ?? 0) + 1) + const cursor = this.cursor(path) + this.db.exec('BEGIN IMMEDIATE') + try { + this.dropSession(cursor?.session_row_id ?? null) + this.db.prepare('DELETE FROM files WHERE path = ?').run(path) + this.db.exec('COMMIT') + } catch (error) { + this.db.exec('ROLLBACK') + throw error + } + } + + private cursor(path: string): FileCursor | undefined { + return this.db + .prepare('SELECT session_row_id,byte_offset FROM files WHERE path = ?') + .get(path) as FileCursor | undefined + } + + private buffered( + candidate: SessionFileCandidate, + opened: FileCursor | undefined, + append: boolean, + identity?: () => TranscriptSessionIdentity | null + ): SessionSearchFileWrite { + const db = this.db + const path = candidate.file.path + const buffer: TranscriptMessage[] = [] + let bufferedChars = 0 + // What this write believes the file record holds. Re-read inside every + // transaction: a `removeFile` or another writer between two chunks means + // these rows no longer continue anything, and committing on top of that + // would resurrect a deleted source or duplicate a span. + let expected = opened + const removalsAtStart = this.removals.get(path) ?? 0 + // The session row is reused across re-reads of one file, so a `replace` + // swaps a session's rows rather than minting a second generation of it. + let session = opened?.session_row_id ?? null + let hash = append && session !== null ? this.records.contentHash(session) : EMPTY_CONTENT_HASH + // A replace owns the session's whole row set, so the old generation goes in + // the same transaction as the first of the new one. Chunk two onwards must + // not repeat it. + // + // It goes by being cut loose, not by being deleted. Deleting every old row + // inline sizes the transaction by the session being replaced rather than by + // the chunk being written: 1,286 ms against 720 ms fresh on the 100 MB + // corpus, and it grows with the history. Instead the first transaction + // mints a new session row, points `files` at it and deletes the one old + // `sessions` row. Every retrieval joins `sessions`, so the old generation + // stops answering the moment that commits, and its messages are reclaimed + // afterwards by the same bounded drain retention uses — which is where the + // old rows would have ended up had the process died here anyway. + // `sessions.id` is AUTOINCREMENT, so the freed id is never handed to + // another session while those rows still name it (round 8). + let replaced = append + // Set by the transaction that cut a generation loose; read once it commits. + let orphaned = false + // Set when the file record moved under this read. Nothing this write holds + // can land after that, so it stops buffering rather than reopening a + // transaction it already knows will roll back, once per remaining message. + let fenced = false + + // Why a counter and not the cursor alone: on a path this index never wrote, + // `expected` and the absent row are both undefined, so the cursor compare + // reads a removal as no change and the write recreates the source. + const current = (): boolean => { + if ((this.removals.get(path) ?? 0) !== removalsAtStart) { + return false + } + const row = this.cursor(path) + return ( + row?.session_row_id === expected?.session_row_id && + row?.byte_offset === expected?.byte_offset + ) + } + + /** + * `outcome` is null for a chunk of a read that has not reached the file's + * end, and `named` is what that chunk writes onto its session row. + */ + const write = ( + outcome: TranscriptReadOutcome | null, + named: TranscriptSessionIdentity | null + ): boolean => { + const decoded = outcome?.session ?? null + db.exec('BEGIN IMMEDIATE') + try { + if (!current()) { + db.exec('ROLLBACK') + return false + } + if (outcome && !decoded) { + // Read through, but nothing to search: the cursor advances so the file + // is not re-read whole on every pass, and whatever generation was here + // — including this read's own committed chunks — goes with it. + this.dropSession(session) + session = null + this.records.upsertFile(candidate, outcome.byteOffset, null) + } else { + if (replaced) { + session ??= this.records.createSessionRow(candidate) + } else { + const previous = session + session = this.records.createSessionRow(candidate) + if (previous !== null) { + db.prepare('DELETE FROM sessions WHERE id = ?').run(previous) + orphaned = true + } + replaced = true + } + for (const row of buffer) { + insertSearchMessage(db, session, row) + } + if (decoded) { + this.records.updateSession(decoded, session, hash) + } else if (named) { + // A chunk's rows answer searches as soon as they land, so the + // session they hang off is written with whatever the parser has + // decoded rather than left empty until a read that may never end. + // `add` refuses to chunk without this, so it is never absent here. + this.records.updateProvisionalSession(session, named) + } + this.records.upsertFile( + candidate, + outcome ? outcome.byteOffset : PARTIAL_FILE_CURSOR, + session + ) + } + db.exec('COMMIT') + } catch (error) { + db.exec('ROLLBACK') + throw error + } + // After the transaction that cut them loose is durable, never before: a + // rollback leaves the old session row standing and nothing to reclaim. + if (orphaned) { + orphaned = false + this.onOrphanedRows() + } + expected = { + session_row_id: session, + byte_offset: outcome ? outcome.byteOffset : PARTIAL_FILE_CURSOR + } + buffer.length = 0 + bufferedChars = 0 + return true + } + + return { + add: (message) => { + if (fenced) { + return + } + hash = foldContentHash(hash, [message]) + // The ceiling is checked per row, not per message: one message is a whole + // conversation turn and may be megabytes, so checking it after the whole + // message had been buffered let a single one carry a transaction as far + // past the ceiling as it was large. + for (const row of searchMessageRows([message])) { + buffer.push(row) + bufferedChars += row.text.length + if (bufferedChars < this.commitChars) { + continue + } + // Publishing a chunk under a session nothing can identify is worse + // than holding the buffer: the rows answer searches at once, and an + // interrupted read leaves that prefix for good. A read with nothing + // to name it keeps buffering and commits whole at `finish`. + const named = identity?.() ?? null + if (named && !write(null, named)) { + fenced = true + buffer.length = 0 + bufferedChars = 0 + return + } + } + }, + commit: (outcome) => !fenced && write(outcome, null) + } + } + + /** Caller's transaction: drops a session and every row that hangs off it. */ + private dropSession(sessionRowId: number | null): void { + if (sessionRowId === null) { + return + } + deleteSearchMessages(this.db, sessionRowId) + this.db.prepare('DELETE FROM sessions WHERE id = ?').run(sessionRowId) + } +} diff --git a/src/main/ai-vault-search/session-search-live-transcript.test.ts b/src/main/ai-vault-search/session-search-live-transcript.test.ts new file mode 100644 index 00000000000..7f37ca662cb --- /dev/null +++ b/src/main/ai-vault-search/session-search-live-transcript.test.ts @@ -0,0 +1,208 @@ +import { mkdtemp, rm, writeFile, appendFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, expect, it } from 'vitest' +import { resetSessionParseCacheForTests } from '../ai-vault/session-scanner-parse-cache' +import { + registerTranscriptConsumer, + resetTranscriptConsumersForTests, + type TranscriptSessionIdentity +} from '../ai-vault/session-transcript-consumers' +import { requestWholeTranscriptRead } from '../ai-vault/session-transcript-reader' +import SyncDatabase from '../sqlite/sync-database' +import { registerSessionSearchIndexConsumer } from './session-search-index-consumer' +import { SessionSearchStore } from './session-search-store' +import { + assistantRecord, + CLAUDE_SESSION_ID as SESSION_ID, + CODEX_ROLLOUT_FILE, + CODEX_SESSION_ID, + codexRolloutLines, + parseTranscript, + userRecord +} from './session-search-transcript-fixtures' + +let tempRoots: string[] = [] +let store: SessionSearchStore +// The store keeps its connection private, so row assertions need a second one. +let reader: SyncDatabase +let errors: unknown[] + +beforeEach(async () => { + resetSessionParseCacheForTests() + resetTranscriptConsumersForTests() + errors = [] + const path = join(await makeTempDir(), 'index.sqlite') + store = new SessionSearchStore(path, (error) => errors.push(error)) + registerSessionSearchIndexConsumer(store) + reader = new SyncDatabase(path, { readonly: true }) +}) + +afterEach(async () => { + resetTranscriptConsumersForTests() + reader.close() + store.close() + await Promise.all(tempRoots.map((root) => rm(root, { recursive: true, force: true }))) + tempRoots = [] +}) + +async function makeTempDir(): Promise { + const root = await mkdtemp(join(tmpdir(), 'orca-session-search-live-')) + tempRoots.push(root) + return root +} + +/** Sessions a query would return for one FTS term, read on a second handle. */ +function sessionsMatching(term: string): string[] { + return ( + reader + .prepare( + `SELECT DISTINCT s.session_id AS id FROM messages_fts + JOIN messages m ON m.id = messages_fts.rowid + JOIN sessions s ON s.id = m.session_row_id + WHERE messages_fts MATCH ? ORDER BY s.session_id` + ) + .all(term) as { id: string }[] + ).map((row) => row.id) +} + +it('indexes a Claude transcript through the reader and resumes on append', async () => { + const root = await makeTempDir() + const path = join(root, `${SESSION_ID}.jsonl`) + await writeFile( + path, + `${[ + userRecord(0, 'find the flaky terminal reattach'), + assistantRecord(1, 'look at resolveTerminalPath first') + ].join('\n')}\n` + ) + await parseTranscript(path) + expect(errors).toEqual([]) + expect(sessionsMatching('reattach')).toEqual([SESSION_ID]) + // The identifier column shadows a camel-case symbol into its pieces. + expect(sessionsMatching('terminal')).toEqual([SESSION_ID]) + + await appendFile(path, `${assistantRecord(2, 'the zygomorphic follow-up landed')}\n`) + const resumed = await parseTranscript(path) + // The reader resumed, so the index saw an `append`, not a whole re-read. + expect(resumed.stats).toMatchObject({ incremental: 1, fullParses: 0 }) + expect(errors).toEqual([]) + expect(sessionsMatching('zygomorphic')).toEqual([SESSION_ID]) + // An append extends one session rather than creating a second. + expect(reader.prepare('SELECT count(*) AS n FROM sessions').get()).toEqual({ + n: 1 + }) +}) + +it('keeps a tool result searchable but out of the conversation half', async () => { + const root = await makeTempDir() + const codexHome = await makeTempDir() + const path = join(root, CODEX_ROLLOUT_FILE) + await writeFile( + path, + `${codexRolloutLines( + ['rg', 'pericardium'], + `outputonly ${'padding '.repeat(600)}tailonly`, + 'promptonly search for the module' + ).join('\n')}\n` + ) + await parseTranscript(path, 'codex', codexHome) + expect(errors).toEqual([]) + + expect(sessionsMatching('pericardium')).toHaveLength(1) + // The prompt is conversation; the command output is not, and the column + // filter is what tells them apart. + expect(sessionsMatching('outputonly')).toHaveLength(1) + expect(sessionsMatching('tailonly')).toHaveLength(0) + expect(sessionsMatching('rg')).toHaveLength(1) + expect(sessionsMatching('{user_text assistant_text}: promptonly')).toHaveLength(1) + expect(sessionsMatching('{user_text assistant_text}: outputonly')).toHaveLength(0) + expect(sessionsMatching('{user_text assistant_text}: rg')).toHaveLength(0) +}) + +/** What `start.identity()` returns at each message of one read. */ +function recordIdentityPerMessage(): (TranscriptSessionIdentity | null)[] { + const seen: (TranscriptSessionIdentity | null)[] = [] + registerTranscriptConsumer({ + beginRead: (start) => ({ + message: () => { + seen.push(start.identity?.() ?? null) + }, + finish: () => undefined + }) + }) + return seen +} + +it('names the session mid-read, before the reader has finished the file', async () => { + const root = await makeTempDir() + const path = join(root, `${SESSION_ID}.jsonl`) + await writeFile( + path, + `${[ + userRecord(0, 'find the flaky terminal reattach'), + assistantRecord(1, 'look at resolveTerminalPath first') + ].join('\n')}\n` + ) + const seen = recordIdentityPerMessage() + await parseTranscript(path) + + // A chunked read commits partway through a file this size or larger, so what + // it can name the session with is exactly this. + expect(seen.length).toBeGreaterThan(0) + expect(seen[0]).toMatchObject({ + sessionId: SESSION_ID, + cwd: '/repo/app', + createdAt: expect.any(String) + }) +}) + +it('names a Codex session mid-read from its own opening record', async () => { + const root = await makeTempDir() + const codexHome = await makeTempDir() + const path = join(root, CODEX_ROLLOUT_FILE) + await writeFile( + path, + `${codexRolloutLines(['rg', 'pericardium'], 'src/main/pericardium.ts:12: match', 'search for the pericardium module').join('\n')}\n` + ) + const seen = recordIdentityPerMessage() + await parseTranscript(path, 'codex', codexHome) + + // Codex builds its own resumable state rather than the shared accumulator + // fold, so it is the other half of the surface a chunked commit depends on. + expect(seen[0]).toMatchObject({ + sessionId: CODEX_SESSION_ID, + cwd: '/repo/app' + }) +}) + +it('indexes a file the session list already read past, once a whole read is asked for', async () => { + const root = await makeTempDir() + const path = join(root, `${SESSION_ID}.jsonl`) + await writeFile(path, `${userRecord(0, 'the opening prompt')}\n`) + + // The state on first enablement inside a running app: the session list has + // read this file, so the parse cache is warm, while the index is empty. + resetTranscriptConsumersForTests() + await parseTranscript(path) + registerSessionSearchIndexConsumer(store) + + await appendFile(path, `${assistantRecord(1, 'a zygomorphic reply')}\n`) + const appended = await parseTranscript(path) + expect(appended.stats).toMatchObject({ incremental: 1, fullParses: 0 }) + // The append continued from a byte offset the index never saw, so it declined. + expect(sessionsMatching('zygomorphic')).toEqual([]) + + const behind = store.takeStale() + expect(behind.map((candidate) => candidate.file.path)).toEqual([path]) + for (const candidate of behind) { + requestWholeTranscriptRead(candidate.file.path) + } + + const reread = await parseTranscript(path) + expect(reread.stats).toMatchObject({ incremental: 0, fullParses: 1 }) + expect(errors).toEqual([]) + expect(sessionsMatching('zygomorphic')).toEqual([SESSION_ID]) + expect(sessionsMatching('opening')).toEqual([SESSION_ID]) + expect(store.takeStale()).toEqual([]) +}) diff --git a/src/main/ai-vault-search/session-search-message-rows.test.ts b/src/main/ai-vault-search/session-search-message-rows.test.ts new file mode 100644 index 00000000000..d2cf2fbca61 --- /dev/null +++ b/src/main/ai-vault-search/session-search-message-rows.test.ts @@ -0,0 +1,249 @@ +import { expect, it } from 'vitest' +import type { TranscriptMessage } from '../ai-vault/session-transcript-consumers' +import { insertSearchMessage, searchMessageRows } from './session-search-message-rows' +import { + openSessionSearchIndexFile, + type SessionSearchIndexFile +} from './session-search-index-test-fixture' + +/** Every column of the FTS table, so an assertion cannot miss the shadow terms. */ +async function indexedColumns( + index: SessionSearchIndexFile, + message: TranscriptMessage +): Promise { + for (const row of searchMessageRows([message])) { + insertSearchMessage(index.db, 1, row) + } + const full = index.db + .prepare('SELECT user_text, assistant_text, tool_text, identifiers FROM messages_fts') + .all() as Record[] + return full.flatMap((row) => Object.values(row)) +} + +it('splits an oversized message on a line boundary and keeps every character', () => { + const line = `${'padding '.repeat(11)}word\n` + const text = line.repeat(400) + const chunks = [...searchMessageRows([{ role: 'user', text, timestamp: null }])].map( + (row) => row.text + ) + + expect(chunks.length).toBeGreaterThan(1) + expect(chunks.join('')).toBe(text) + for (const chunk of chunks) { + expect(chunk.length).toBeLessThanOrEqual(8000) + expect(chunk.endsWith('\n')).toBe(true) + } +}) + +it('cuts at whitespace rather than through the word on the boundary', async () => { + const index = await openSessionSearchIndexFile('ss-rows-whitespace') + try { + // The 8,000th character lands inside `pericardium`. Cutting at the target + // would file `per` under one row and `icardium` under another, and the word + // the user types would match neither. + const text = `${' '.repeat(7997)}pericardium` + const chunks = [...searchMessageRows([{ role: 'user', text, timestamp: null }])] + expect(chunks.map((row) => row.text).join('')).toBe(text) + for (const row of chunks) { + insertSearchMessage(index.db, 1, row) + } + + expect( + index.db + .prepare('SELECT count(*) AS n FROM messages_fts WHERE messages_fts MATCH ?') + .get('pericardium') + ).toEqual({ n: 1 }) + } finally { + await index.close() + } +}) + +it.each(['/repo/pericardium.ts', 'PROJ-12345', 'C++', 'cafe\u0301ine'])( + 'preserves the exact FTS token %s at a chunk boundary', + async (token) => { + const index = await openSessionSearchIndexFile('ss-rows-tokenchars') + try { + const text = ' '.repeat(7998) + token + const chunks = [...searchMessageRows([{ role: 'user', text, timestamp: null }])] + expect(chunks.map((row) => row.text).join('')).toBe(text) + for (const row of chunks) { + insertSearchMessage(index.db, 1, row) + } + expect( + index.db + .prepare('SELECT count(*) AS n FROM messages_fts WHERE messages_fts MATCH ?') + .get(`"${token}"`) + ).toEqual({ n: 1 }) + } finally { + await index.close() + } + } +) + +it.each(['\u0305', '\u030d', '\u0332'])( + 'cuts at a combining mark unicode61 treats as a separator: %s', + async (mark) => { + const index = await openSessionSearchIndexFile('ss-rows-unicode-separator') + try { + const text = `${'x'.repeat(7997)}${mark}pericardium` + for (const row of searchMessageRows([{ role: 'user', text, timestamp: null }])) { + insertSearchMessage(index.db, 1, row) + } + expect( + index.db + .prepare("SELECT count(*) AS n FROM messages_fts WHERE messages_fts MATCH 'pericardium'") + .get() + ).toEqual({ n: 1 }) + } finally { + await index.close() + } + } +) + +it('backs up to any whitespace, not only a newline', () => { + // An ideographic space separates words in a CJK transcript exactly as a + // space does here, and a newline-only backoff tears the token after it. + const text = `${'\u4e00'.repeat(7000)}\u3000${'\u4e8c'.repeat(2000)}` + const chunks = [...searchMessageRows([{ role: 'user', text, timestamp: null }])].map( + (row) => row.text + ) + + expect(chunks[0]).toBe(`${'\u4e00'.repeat(7000)}\u3000`) + expect(chunks.join('')).toBe(text) +}) + +it('cuts at punctuation when the window holds no whitespace at all', async () => { + const index = await openSessionSearchIndexFile('ss-rows-minified') + try { + // Valid minified JSON, the shape a tool result carries: 8,000 characters + // without a single space. The 8,000th lands inside `pericardium`, and a + // whitespace-only backoff has nothing in the window to back up to, so it + // files `perica` under one row and `rdium` under the next. + const text = `{"pad":"${'x'.repeat(7976)}","note":"pericardium"}` + expect(JSON.parse(text)).toEqual({ pad: 'x'.repeat(7976), note: 'pericardium' }) + expect(text.slice(7994, 8005)).toBe('pericardium') + expect(/\s/.test(text)).toBe(false) + + const chunks = [...searchMessageRows([{ role: 'user', text, timestamp: null }])] + expect(chunks.map((row) => row.text).join('')).toBe(text) + for (const row of chunks) { + insertSearchMessage(index.db, 1, row) + } + + expect( + index.db + .prepare('SELECT count(*) AS n FROM messages_fts WHERE messages_fts MATCH ?') + .get('pericardium') + ).toEqual({ n: 1 }) + } finally { + await index.close() + } +}) + +it('keeps a 9,000-character identifier whole rather than cutting at its underscores', () => { + // `_` sits inside a token for this tokenizer, so it is not a boundary. A + // snake_case name that long holds none at all, and the target itself is the + // honest cut — backing up to every `_` would file the name in pieces. + const text = 'ab_'.repeat(3000) + expect(text.length).toBe(9000) + const chunks = [...searchMessageRows([{ role: 'user', text, timestamp: null }])].map( + (row) => row.text + ) + + expect(chunks.map((chunk) => chunk.length)).toEqual([8000, 1000]) + expect(chunks.join('')).toBe(text) +}) + +it('still chunks a message that holds no whitespace at all', () => { + // A 20,000-character token is not a word, so the target itself is the cut and + // the message is still bounded. + const chunks = [ + ...searchMessageRows([{ role: 'user', text: 'a'.repeat(20_000), timestamp: null }]) + ] + expect(chunks.map((row) => row.text.length)).toEqual([8000, 8000, 4000]) +}) + +it('leaves a message that fits as a single row', () => { + const rows = [...searchMessageRows([{ role: 'user', text: 'short enough', timestamp: null }])] + expect(rows.map((row) => row.text)).toEqual(['short enough']) +}) + +it('caps a tool row at its head and never caps the conversation', async () => { + const index = await openSessionSearchIndexFile('ss-rows-tool-cap') + try { + // The reader hands over untruncated text (its own bound is 256 KB per + // message and a consumer may be handed more); the cap is this module's. + const output = `pericardium ${'padding '.repeat(140_000)}` + expect(output.length).toBeGreaterThan(1024 * 1024) + + const toolRows = [...searchMessageRows([{ role: 'tool', text: output, timestamp: null }])] + expect(toolRows).toHaveLength(1) + expect(toolRows[0]!.text.length).toBe(3072) + // The head is what identifies what ran, so it is what survives. + expect(toolRows[0]!.text.startsWith('pericardium ')).toBe(true) + + // The same text as an assistant turn is conversation, and keeps every byte. + const assistantRows = [ + ...searchMessageRows([{ role: 'assistant', text: output, timestamp: null }]) + ] + expect(assistantRows.map((row) => row.text).join('')).toBe(output) + expect(assistantRows.length).toBeGreaterThan(100) + + for (const row of toolRows) { + insertSearchMessage(index.db, 1, row) + } + expect( + index.db + .prepare('SELECT count(*) AS n FROM messages_fts WHERE messages_fts MATCH ?') + .get('pericardium') + ).toEqual({ n: 1 }) + } finally { + await index.close() + } +}) + +it('files a tool row under the tool column alone', async () => { + const index = await openSessionSearchIndexFile('ss-message-rows-tool') + try { + for (const row of searchMessageRows([ + { role: 'tool', text: 'rg pericardium', timestamp: null } + ])) { + insertSearchMessage(index.db, 1, row) + } + expect(index.db.prepare('SELECT count(*) AS n FROM messages_fts').get()).toEqual({ n: 1 }) + // What makes a conversation-scoped search exclude it: the column filter, not + // a second table. + expect( + index.db + .prepare('SELECT count(*) AS n FROM messages_fts WHERE messages_fts MATCH ?') + .get('{user_text assistant_text}: pericardium') + ).toEqual({ n: 0 }) + expect( + index.db + .prepare('SELECT count(*) AS n FROM messages_fts WHERE messages_fts MATCH ?') + .get('{tool_text}: pericardium') + ).toEqual({ n: 1 }) + } finally { + await index.close() + } +}) + +it('stores a chunk exactly as the transcript wrote it', async () => { + const index = await openSessionSearchIndexFile('ss-rows-verbatim') + try { + const text = 'deploy with AKIAIOSFODNN7EXAMPLE and the resolveTerminalPath fix' + const stored = await indexedColumns(index, { + role: 'assistant', + text, + timestamp: null + }) + + // The index is a second copy of content the user already holds in plaintext, + // so it neither rewrites nor drops any of it. + expect(stored).toContain(text) + // Identifier shadow terms come off that same raw chunk. + expect(stored.some((column) => column.includes('resolve terminal path'))).toBe(true) + } finally { + await index.close() + } +}) diff --git a/src/main/ai-vault-search/session-search-message-rows.ts b/src/main/ai-vault-search/session-search-message-rows.ts new file mode 100644 index 00000000000..21254c183aa --- /dev/null +++ b/src/main/ai-vault-search/session-search-message-rows.ts @@ -0,0 +1,135 @@ +import type SyncDatabase from '../sqlite/sync-database' +import type { TranscriptMessage } from '../ai-vault/session-transcript-consumers' +import { sliceAtCodeUnitLimit } from '../ai-vault/session-scanner-text-normalization' +import { identifierShadowText } from './session-search-identifier-split' + +const CHUNK_TARGET_CHARS = 8000 + +/** + * How much of one tool output is indexed. Its head: a command, its arguments and + * the first lines of what it printed are what a user searches for, while the + * tail is the padding that makes these messages large in the first place. + * + * Tool output is 80-97 % of a transcript's bytes, and a single one can be a + * quarter of a megabyte (the reader's own per-message bound). Without this the + * index, the in-memory buffer a read holds and the transaction it commits are + * all sized by how much a tool printed rather than by how much is worth + * searching. 3 KB was the accuracy/size sweet spot in the original design + * measurement. User and assistant text is never capped: it is the conversation, + * and it is small. + */ +const TOOL_ROW_CHARS = 3072 + +// Keep unicode61's tokenchars intact, including before an available space. +// SQLite ext/fts5/fts5_unicode2.c: sqlite3Fts5UnicodeIsdiacritic, with remove_diacritics=1. +const FOLDED_DIACRITIC = + /[\u0300-\u0304\u0306-\u030c\u030f\u0311\u031b\u0323-\u0328\u032d-\u032e\u0330-\u0331]/ +const TOKEN_BOUNDARY = /[^\p{L}\p{N}\p{Co}_.\-/+\uD800-\uDFFF]/u + +/** + * Index just past the last token boundary in `[floor, end)`, or -1 when the + * window holds none. Not only a newline: a wrapped paragraph, a CJK transcript + * separated by ideographic spaces and a minified log all chunk on a boundary a + * tokenizer would have picked anyway. + */ +function lastTokenBoundaryEnd(text: string, floor: number, end: number): number { + for (let at = end - 1; at >= floor; at--) { + if (TOKEN_BOUNDARY.test(text[at]!) && !FOLDED_DIACRITIC.test(text[at]!)) { + return at + 1 + } + } + return -1 +} + +/** + * Splits an oversized message into rows of at most `CHUNK_TARGET_CHARS`, cutting + * on a token boundary so no token is torn in half and every word stays + * searchable. A phrase that straddles two chunks is not matched: chunks are + * separate FTS rows and FTS5 cannot span them. + */ +function* textChunks(text: string): Generator { + if (text.length <= CHUNK_TARGET_CHARS) { + yield text + return + } + let start = 0 + while (start < text.length) { + let end = Math.min(text.length, start + CHUNK_TARGET_CHARS) + if (end < text.length) { + // Only the second half of the window: backing up further would trade a + // torn token for chunks half the size. No boundary at all in 4,000 + // characters is not a word, so the target itself is the honest cut. + const split = lastTokenBoundaryEnd(text, start + CHUNK_TARGET_CHARS / 2, end) + if (split > start) { + end = split + } + } + yield text.slice(start, end) + start = end + } +} + +/** + * The row policy for one message: a `tool` message becomes one capped row, and + * anything else becomes N chunks, because FTS5 ranks a short row far better + * than a huge one. + */ +export function* searchMessageRows( + messages: Iterable +): Generator { + for (const message of messages) { + if (message.role === 'tool') { + yield { + ...message, + text: sliceAtCodeUnitLimit(message.text, TOOL_ROW_CHARS) + } + continue + } + for (const text of textChunks(message.text)) { + yield { ...message, text } + } + } +} + +/** + * Writes one row into `messages` and `messages_fts` in the caller's + * transaction, so a message is never present in one and absent from the other. + * A conversation-scoped query filters the columns rather than reading a second + * table (see the schema). + */ +export function insertSearchMessage( + db: SyncDatabase, + sessionId: number, + message: TranscriptMessage +): void { + const text = message.text + const id = db + .prepare('INSERT INTO messages(session_row_id, role, ts) VALUES (?, ?, ?)') + .run(sessionId, message.role, message.timestamp).lastInsertRowid + const user = message.role === 'user' ? text : '' + const assistant = message.role === 'assistant' ? text : '' + const tool = message.role === 'tool' ? text : '' + db.prepare( + 'INSERT INTO messages_fts(rowid,user_text,assistant_text,tool_text,identifiers) VALUES (?,?,?,?,?)' + ).run(id, user, assistant, tool, identifierShadowText(text)) +} + +/** + * Deletes up to `limit` of a session's rows from `messages` and `messages_fts`, + * in the caller's transaction, and reports how many went. Bounded + * because a retention sweep must not hold one transaction over a whole + * session; a replace passes no limit, since its rows and their replacements + * have to land together. + */ +export function deleteSearchMessages(db: SyncDatabase, sessionId: number, limit = -1): number { + const ids = db + .prepare('SELECT id FROM messages WHERE session_row_id = ? LIMIT ?') + .all(sessionId, limit) as { id: number }[] + const full = db.prepare('DELETE FROM messages_fts WHERE rowid = ?') + const message = db.prepare('DELETE FROM messages WHERE id = ?') + for (const { id } of ids) { + full.run(id) + message.run(id) + } + return ids.length +} diff --git a/src/main/ai-vault-search/session-search-retention-delete.test.ts b/src/main/ai-vault-search/session-search-retention-delete.test.ts new file mode 100644 index 00000000000..c21c8d7fe2b --- /dev/null +++ b/src/main/ai-vault-search/session-search-retention-delete.test.ts @@ -0,0 +1,188 @@ +import { expect, it } from 'vitest' +import type SyncDatabase from '../sqlite/sync-database' +import { + deleteExpiredSearchFiles, + RETENTION_DELETE_ROWS_PER_STEP +} from './session-search-retention-delete' +import { openSessionSearchIndexFile } from './session-search-index-test-fixture' +import { SessionSearchStore } from './session-search-store' + +function seed(db: SyncDatabase, id: number, rows: number, mtime: number): void { + db.prepare( + `INSERT INTO sessions(id,agent,session_id,file_path,title,cwd,cwd_key,resume_command) + VALUES (?, 'claude', ?, ?, 'synthetic retention', '/fixture', '/fixture', '')` + ).run(id, String(id), String(id)) + db.prepare('INSERT INTO files(path,byte_offset,mtime_ms,session_row_id) VALUES (?,1,?,?)').run( + String(id), + mtime, + id + ) + db.exec('BEGIN') + for (let i = 0; i < rows; i++) { + const row = db + .prepare("INSERT INTO messages(session_row_id,role) VALUES (?,'user')") + .run(id).lastInsertRowid + db.prepare('INSERT INTO messages_fts(rowid,user_text) VALUES (?,?)').run(row, 'retentionneedle') + } + db.exec('COMMIT') +} + +/** + * Sessions a search would still return. Every retrieval joins a message to its + * session, which is what makes cutting the session loose enough to hide the + * whole thing while its rows are still being reclaimed. + */ +function visibleSessionIds(db: SyncDatabase): string[] { + return ( + db + .prepare( + `SELECT DISTINCT s.session_id AS id FROM messages_fts + JOIN messages m ON m.id = messages_fts.rowid + JOIN sessions s ON s.id = m.session_row_id + WHERE messages_fts MATCH 'retentionneedle' ORDER BY s.session_id` + ) + .all() as { id: string }[] + ).map((row) => row.id) +} + +function count(db: SyncDatabase, table: string): number { + return (db.prepare(`SELECT count(*) AS n FROM ${table}`).get() as { n: number }).n +} + +it('seeks the expiring end of the file list instead of scanning it', async () => { + const index = await openSessionSearchIndexFile('ss-retention-plan') + try { + seed(index.db, 1, 1, 1) + const plan = ( + index.db + .prepare('EXPLAIN QUERY PLAN SELECT path FROM files WHERE mtime_ms < ? ORDER BY mtime_ms') + .all(100) as { detail: string }[] + ) + .map((row) => row.detail) + .join(' ') + // Without files_mtime this is "SCAN files" plus a "USE TEMP B-TREE FOR ORDER BY". + expect(plan).toContain('files_mtime') + expect(plan).not.toContain('TEMP B-TREE') + } finally { + await index.close() + } +}) + +it('hides an expiring session at once, then reclaims its rows in bounded steps', async () => { + const index = await openSessionSearchIndexFile('ss-retention-yield') + seed(index.db, 1, 1025, 1) + seed(index.db, 2, 1, 200) + let previous = 1025 + const steps: number[] = [] + try { + await deleteExpiredSearchFiles( + index.db, + 100, + () => false, + async () => { + const left = count(index.db, 'messages WHERE session_row_id=1') + steps.push(previous - left) + previous = left + // Cut loose in the very first transaction, so no query ever sees it with + // some of its messages already gone. + expect(visibleSessionIds(index.db)).toEqual(['2']) + } + ) + // The file transaction, then one bounded batch per step until the rows are gone. + expect(steps).toEqual([0, RETENTION_DELETE_ROWS_PER_STEP, 256, 256, 256, 1]) + expect(count(index.db, 'messages_fts')).toBe(1) + expect(count(index.db, 'sessions')).toBe(1) + } finally { + await index.close() + } +}) + +it('finishes an interrupted deletion after reopening', async () => { + const index = await openSessionSearchIndexFile('ss-retention-resume') + let store = new SessionSearchStore(index.path) + let closed = false + let steps = 0 + try { + seed(index.db, 1, 513, 1) + await deleteExpiredSearchFiles( + index.db, + 100, + () => closed, + async () => { + if (++steps === 2) { + store.close() + closed = true + } + } + ) + // Some rows went, the rest did not, and nothing recorded that anywhere. + const stranded = count(index.db, 'messages') + expect(stranded).toBeGreaterThan(0) + expect(stranded).toBeLessThan(513) + expect(visibleSessionIds(index.db)).toEqual([]) + + store = new SessionSearchStore(index.path) + closed = false + // Rows nothing points at are the whole record of unfinished work, so the + // rest goes even with retention now unlimited. + await store.purgeOlderThan(null) + expect(count(index.db, 'messages')).toBe(0) + expect(count(index.db, 'messages_fts')).toBe(0) + } finally { + if (!closed) { + store.close() + } + await index.close() + } +}) + +it('cancels retention between batches and resumes without exposing a partial session', async () => { + const index = await openSessionSearchIndexFile('ss-retention-cancel') + const store = new SessionSearchStore(index.path) + try { + seed(index.db, 1, 1025, 1) + const controller = new AbortController() + const purge = store.purgeOlderThan(100, controller.signal) + setImmediate(() => controller.abort()) + await purge + const remaining = count(index.db, 'messages') + expect(remaining).toBeGreaterThan(0) + expect(remaining).toBeLessThan(1025) + expect(visibleSessionIds(index.db)).toEqual([]) + await store.purgeOlderThan(null) + expect(count(index.db, 'messages')).toBe(0) + } finally { + store.close() + await index.close() + } +}) + +it('keeps a file a read refreshed after the expiry list was taken', async () => { + const index = await openSessionSearchIndexFile('ss-retention-refreshed') + try { + seed(index.db, 1, 2, 1) + seed(index.db, 2, 2, 2) + let refreshed = false + // The scan of `files` happens once, up front. A read of the second transcript + // lands while the first is being deleted, which makes it new enough to keep. + await deleteExpiredSearchFiles( + index.db, + 100, + () => false, + async () => { + if (!refreshed) { + refreshed = true + index.db.prepare('UPDATE files SET mtime_ms = 500 WHERE path = ?').run('2') + } + } + ) + + // Only the per-file transaction re-reading the mtime it is about to act on + // keeps that session; the list it came from says both should go. + expect(count(index.db, 'files')).toBe(1) + expect(visibleSessionIds(index.db)).toEqual(['2']) + expect(count(index.db, 'messages')).toBe(2) + } finally { + await index.close() + } +}) diff --git a/src/main/ai-vault-search/session-search-retention-delete.ts b/src/main/ai-vault-search/session-search-retention-delete.ts new file mode 100644 index 00000000000..e8d407f8be9 --- /dev/null +++ b/src/main/ai-vault-search/session-search-retention-delete.ts @@ -0,0 +1,102 @@ +import { setImmediate as yieldToEventLoop } from 'node:timers/promises' +import type SyncDatabase from '../sqlite/sync-database' +import { deleteSearchMessages } from './session-search-message-rows' + +export const RETENTION_DELETE_ROWS_PER_STEP = 256 +// Why in step with the deletes rather than one sweep at the end: `auto_vacuum = +// INCREMENTAL` holds every freed page until something asks for it back, and +// asking for a whole purge's worth at once is one long stall (40 ms per 22 MB +// freed, measured) instead of many short ones. +const RECLAIM_PAGES_PER_STEP = 2000 + +/** + * Drops every file older than the cutoff, then hands its rows back in bounded + * steps. + * + * The two halves are separate on purpose. Cutting a session loose from its file + * is one small transaction, and it is what makes the session stop answering + * searches — every read joins `sessions`, so a row whose session is gone is + * already unreachable. Reclaiming those rows is the expensive half, and it can + * be paused, interrupted or resumed at any point without a reader ever seeing a + * session that is half deleted. A crash in the middle leaves rows nothing + * points at, and `drainOrphanedMessages` finds them on the next pass. + */ +export async function deleteExpiredSearchFiles( + db: SyncDatabase, + cutoffMs: number | null, + closed: () => boolean, + yieldStep: () => Promise = yieldToEventLoop +): Promise { + if (cutoffMs !== null) { + const expired = db + .prepare('SELECT path FROM files WHERE mtime_ms < ? ORDER BY mtime_ms') + .all(cutoffMs) as { path: string }[] + for (const { path } of expired) { + if (closed()) { + return + } + db.exec('BEGIN IMMEDIATE') + try { + // Re-read under the lock: a read of this file may have landed since the + // list was taken, which makes it new enough to keep. + const file = db + .prepare('SELECT session_row_id FROM files WHERE path = ? AND mtime_ms < ?') + .get(path, cutoffMs) as { session_row_id: number | null } | undefined + if (file) { + db.prepare('DELETE FROM sessions WHERE id = ?').run(file.session_row_id) + db.prepare('DELETE FROM files WHERE path = ?').run(path) + } + db.exec('COMMIT') + } catch (error) { + db.exec('ROLLBACK') + throw error + } + await yieldStep() + } + } + await drainOrphanedMessages(db, closed, yieldStep) +} + +/** + * Deletes rows whose session no longer exists, a bounded batch per transaction. + * + * That set is exactly what retention, a replace that cut its old generation + * loose, a removed source and an interrupted earlier drain leave behind, so the + * index needs no record of unfinished work beyond the rows themselves. + * + * Exported for the store, which runs it after a replace commits for the same + * reason retention runs it after its own small transaction: cutting a session + * loose is what hides it, and reclaiming its rows is the half that must not + * hold one transaction. + */ +export async function drainOrphanedMessages( + db: SyncDatabase, + closed: () => boolean, + yieldStep: () => Promise = yieldToEventLoop +): Promise { + // Ordered by session so one call to this walks a session's rows to the end + // before paying for the scan that finds the next one. + const nextOrphan = db.prepare( + `SELECT session_row_id FROM messages + WHERE session_row_id NOT IN (SELECT id FROM sessions) LIMIT 1` + ) + let orphan = (nextOrphan.get() as { session_row_id: number } | undefined)?.session_row_id + while (orphan !== undefined && !closed()) { + db.exec('BEGIN IMMEDIATE') + let deleted = 0 + try { + deleted = deleteSearchMessages(db, orphan, RETENTION_DELETE_ROWS_PER_STEP) + db.exec('COMMIT') + } catch (error) { + db.exec('ROLLBACK') + throw error + } + db.pragma(`incremental_vacuum(${RECLAIM_PAGES_PER_STEP})`) + if (deleted < RETENTION_DELETE_ROWS_PER_STEP) { + orphan = (nextOrphan.get() as { session_row_id: number } | undefined)?.session_row_id + } + await yieldStep() + } + // A `removeFile` frees its pages outside this loop and may leave none to drain. + db.pragma(`incremental_vacuum(${RECLAIM_PAGES_PER_STEP})`) +} diff --git a/src/main/ai-vault-search/session-search-row-identity.test.ts b/src/main/ai-vault-search/session-search-row-identity.test.ts new file mode 100644 index 00000000000..4adc655b584 --- /dev/null +++ b/src/main/ai-vault-search/session-search-row-identity.test.ts @@ -0,0 +1,116 @@ +import { afterEach, beforeEach, expect, it } from 'vitest' +import type SyncDatabase from '../sqlite/sync-database' +import { deleteExpiredSearchFiles } from './session-search-retention-delete' +import { + openSessionSearchIndexFile, + syntheticCandidate, + syntheticSession, + userMessages, + type SessionSearchIndexFile +} from './session-search-index-test-fixture' +import { SessionSearchStore } from './session-search-store' + +// A session row id outlives the row: it names the rows in `messages` until a +// retention drain has walked all of them, which takes many transactions. These +// tests are about what may be handed that id in the meantime. + +let index: SessionSearchIndexFile +let store: SessionSearchStore +let errors: unknown[] + +beforeEach(async () => { + index = await openSessionSearchIndexFile('ss-row-identity') + errors = [] + store = new SessionSearchStore(index.path, (error) => errors.push(error)) +}) + +afterEach(async () => { + store.close() + await index.close() +}) + +const OLD_MTIME = 1_000 +const LIVE_MTIME = 1_000_000 +const LIVE_PATH = '/live.jsonl' + +function count(db: SyncDatabase, table: string): number { + return (db.prepare(`SELECT count(*) AS n FROM ${table}`).get() as { n: number }).n +} + +/** Rows a search would return for a term: the join every retrieval makes. */ +function matches(db: SyncDatabase, term: string): number { + return ( + db + .prepare( + `SELECT count(*) AS n FROM messages_fts JOIN messages m ON m.id = messages_fts.rowid + JOIN sessions s ON s.id = m.session_row_id WHERE messages_fts MATCH ?` + ) + .get(term) as { n: number } + ).n +} + +function indexFile(path: string, mtimeMs: number, text: string, rows: number): void { + const write = store.beginWrite(syntheticCandidate({ path, mtimeMs }), 'replace', 0)! + for (const message of userMessages(text, rows)) { + write.add(message) + } + expect(write.commit({ session: syntheticSession(), byteOffset: 50, incomplete: false })).toBe( + true + ) +} + +it('never hands a live session the rows of a purged one', async () => { + // Two expiring transcripts, each large enough that reclaiming their rows takes + // several transactions, and one live transcript the parser decoded no session + // from — so it holds a cursor and no session row of its own. + indexFile('/old-a.jsonl', OLD_MTIME, 'purgedneedle', 400) + indexFile('/old-b.jsonl', OLD_MTIME, 'purgedneedle', 400) + const live = syntheticCandidate({ path: LIVE_PATH, mtimeMs: LIVE_MTIME }) + const opening = store.beginWrite(live, 'replace', 0)! + opening.add(userMessages('excluded', 1)[0]!) + expect(opening.commit({ session: null, byteOffset: 50, incomplete: false })).toBe(true) + + let appended = false + await deleteExpiredSearchFiles( + index.db, + LIVE_MTIME, + () => false, + async () => { + // The window: both expiring sessions are cut loose, most of their rows are + // still on disk, and the live transcript grows. The append is legitimate — + // it continues this index's own cursor — and it needs a session row. + if (appended || count(index.db, 'sessions') > 0) { + return + } + appended = true + const write = store.beginWrite(live, 'append', 50)! + for (const message of userMessages('liveneedle', 2)) { + write.add(message) + } + expect( + write.commit({ session: syntheticSession(), byteOffset: 120, incomplete: false }) + ).toBe(true) + } + ) + + expect(appended).toBe(true) + // Reusing a freed id would adopt whatever of that session's rows the drain had + // not reached, and put them behind a live session no purge will visit again. + expect(matches(index.db, 'purgedneedle')).toBe(0) + expect(matches(index.db, 'liveneedle')).toBe(2) + expect(count(index.db, 'messages')).toBe(2) + expect(errors).toEqual([]) +}) + +it('never reissues a session row id a delete freed', () => { + for (const path of ['/a.jsonl', '/b.jsonl', '/c.jsonl']) { + indexFile(path, OLD_MTIME, 'seeded', 1) + } + const before = (index.db.prepare('SELECT max(id) AS id FROM sessions').get() as { id: number }).id + index.db.exec('DELETE FROM sessions') + + indexFile('/d.jsonl', OLD_MTIME, 'seeded', 1) + expect((index.db.prepare('SELECT id FROM sessions').get() as { id: number }).id).toBeGreaterThan( + before + ) +}) diff --git a/src/main/ai-vault-search/session-search-schema.test.ts b/src/main/ai-vault-search/session-search-schema.test.ts new file mode 100644 index 00000000000..b23f36cde55 --- /dev/null +++ b/src/main/ai-vault-search/session-search-schema.test.ts @@ -0,0 +1,350 @@ +import type * as NodeFs from 'node:fs' +import { mkdtemp, readFile, stat, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, describe, expect, it, vi } from 'vitest' +import { + removeTree, + WINDOWS_RM_MAX_RETRIES, + WINDOWS_RM_RETRY_DELAY_MS +} from '../../shared/windows-transient-lock-removal' +import SyncDatabase from '../sqlite/sync-database' +import { + SESSION_SEARCH_SCHEMA_VERSION, + openSessionSearchDatabase, + removeSessionSearchDatabase +} from './session-search-schema' + +const recordedRmSync = vi.hoisted(() => vi.fn()) +vi.mock('node:fs', async () => { + const actual = await vi.importActual('node:fs') + return { + ...actual, + rmSync: (...args: Parameters) => { + recordedRmSync(...args) + return actual.rmSync(...args) + } + } +}) + +let roots: string[] = [] + +afterEach(async () => { + await Promise.all(roots.map((root) => removeTree(root))) + roots = [] +}) + +async function tempDatabasePath(): Promise { + const root = await mkdtemp(join(tmpdir(), 'orca-session-search-schema-')) + roots.push(root) + return join(root, 'index.sqlite') +} + +function schemaVersion(db: SyncDatabase): string | undefined { + return ( + db.prepare("SELECT value FROM meta WHERE key = 'schema_version'").get() as + | { value: string } + | undefined + )?.value +} + +describe('openSessionSearchDatabase', () => { + it('keeps a current-version index and its rows', async () => { + const path = await tempDatabasePath() + const first = openSessionSearchDatabase(path) + first.prepare("INSERT INTO files(path,byte_offset,mtime_ms) VALUES ('a',1,1)").run() + first.close() + + const second = openSessionSearchDatabase(path) + expect(schemaVersion(second)).toBe(String(SESSION_SEARCH_SCHEMA_VERSION)) + expect(second.prepare('SELECT COUNT(*) AS c FROM files').get()).toEqual({ + c: 1 + }) + second.close() + }) + + it('carries one FTS table and throws away an index that carries two', async () => { + const path = await tempDatabasePath() + const fresh = openSessionSearchDatabase(path) + const tables = (): string[] => + ( + fresh + .prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name LIKE '%_fts'") + .all() as { name: string }[] + ).map((row) => row.name) + expect(tables()).toEqual(['messages_fts']) + + // What an index written before this bump looks like: the second table, and + // rows in it. `CREATE TABLE IF NOT EXISTS` would leave both in place, so + // only the version bump makes that file go. + fresh.exec('CREATE VIRTUAL TABLE conversation_fts USING fts5(user_text, assistant_text)') + fresh.prepare("INSERT INTO files(path,byte_offset,mtime_ms) VALUES ('a',1,1)").run() + fresh.prepare("UPDATE meta SET value = '3' WHERE key = 'schema_version'").run() + fresh.close() + + const rebuilt = openSessionSearchDatabase(path) + expect(schemaVersion(rebuilt)).toBe(String(SESSION_SEARCH_SCHEMA_VERSION)) + expect( + rebuilt + .prepare("SELECT count(*) AS n FROM sqlite_master WHERE name = 'conversation_fts'") + .get() + ).toEqual({ n: 0 }) + expect(rebuilt.prepare('SELECT COUNT(*) AS c FROM files').get()).toEqual({ c: 0 }) + rebuilt.close() + }) + + it('replaces the file on a version mismatch instead of dropping tables in place', async () => { + const path = await tempDatabasePath() + const stale = openSessionSearchDatabase(path) + stale.prepare("INSERT INTO files(path,byte_offset,mtime_ms) VALUES ('a',1,1)").run() + stale + .prepare("UPDATE meta SET value = ? WHERE key = 'schema_version'") + .run(String(SESSION_SEARCH_SCHEMA_VERSION + 1)) + stale.close() + // Why: a stale sidecar must go with the main file, or SQLite replays it into the new one. + await writeFile(`${path}-wal`, 'stale wal bytes') + const before = await stat(path) + + const fresh = openSessionSearchDatabase(path) + expect(schemaVersion(fresh)).toBe(String(SESSION_SEARCH_SCHEMA_VERSION)) + expect(fresh.prepare('SELECT COUNT(*) AS c FROM files').get()).toEqual({ + c: 0 + }) + fresh.close() + // Why not inode: ext4 hands a freed inode straight back to the next create. + // The planted sidecar is gone (a fresh WAL is checkpointed away on close). + await expect(stat(`${path}-wal`)).rejects.toMatchObject({ code: 'ENOENT' }) + expect((await stat(path)).mtimeMs).toBeGreaterThanOrEqual(before.mtimeMs) + }) + + it('removes the database with every sidecar', async () => { + const path = await tempDatabasePath() + openSessionSearchDatabase(path).close() + await writeFile(`${path}-shm`, '') + removeSessionSearchDatabase(path) + for (const suffix of ['', '-wal', '-shm']) { + await expect(stat(`${path}${suffix}`)).rejects.toMatchObject({ + code: 'ENOENT' + }) + } + }) +}) + +it('rebuilds a file too corrupt to open instead of refusing forever', async () => { + const path = await tempDatabasePath() + const healthy = openSessionSearchDatabase(path) + healthy.prepare("INSERT INTO files(path,byte_offset,mtime_ms) VALUES ('a',1,1)").run() + healthy.close() + // A torn page, not a truncation: SQLite opens the header and fails on the read. + const bytes = await readFile(path) + bytes.fill(0x7f, 4096, Math.min(bytes.length, 12_288)) + await writeFile(path, bytes) + + const rebuilt = openSessionSearchDatabase(path) + try { + expect(schemaVersion(rebuilt)).toBe(String(SESSION_SEARCH_SCHEMA_VERSION)) + expect(rebuilt.prepare('SELECT COUNT(*) AS c FROM files').get()).toEqual({ + c: 0 + }) + } finally { + rebuilt.close() + } +}) + +it('rebuilds a file that is not a database at all', async () => { + const path = await tempDatabasePath() + await writeFile(path, 'not a SQLite database') + + const rebuilt = openSessionSearchDatabase(path) + try { + expect(schemaVersion(rebuilt)).toBe(String(SESSION_SEARCH_SCHEMA_VERSION)) + } finally { + rebuilt.close() + } +}) + +it('gives up rather than looping when a fresh file still cannot be opened', async () => { + const path = await tempDatabasePath() + await writeFile(path, 'not a SQLite database') + // Every open of this path fails, so the one permitted retry is exhausted. + const open = vi.spyOn(SyncDatabase.prototype, 'pragma').mockImplementation(() => { + throw Object.assign(new Error('database disk image is malformed'), { + code: 'SQLITE_CORRUPT' + }) + }) + try { + expect(() => openSessionSearchDatabase(path)).toThrow(/malformed/) + } finally { + open.mockRestore() + } +}) + +it('surfaces the unlink failure itself when a stale index cannot be removed', async () => { + const path = await tempDatabasePath() + const stale = openSessionSearchDatabase(path) + stale + .prepare("UPDATE meta SET value = ? WHERE key = 'schema_version'") + .run(String(SESSION_SEARCH_SCHEMA_VERSION + 1)) + stale.close() + recordedRmSync.mockReset() + recordedRmSync.mockImplementation(() => { + throw Object.assign(new Error('EPERM: operation not permitted, unlink'), { + code: 'EPERM' + }) + }) + try { + // The stale handle is closed before the unlink, so the failure path must not + // close it again: ERR_INVALID_STATE would bury the cause and would not be + // classified as worth a rebuild. + expect(() => openSessionSearchDatabase(path)).toThrow(/EPERM/) + expect(() => openSessionSearchDatabase(path)).not.toThrow(/not open/) + } finally { + recordedRmSync.mockReset() + } +}) + +it('creates the directory the index lives in', async () => { + const root = await mkdtemp(join(tmpdir(), 'orca-session-search-mkdir-')) + roots.push(root) + // The real layout: `/ai-vault-search/index.sqlite`, where nothing + // has made that folder yet. SQLite would fail with `unable to open database + // file`, which is correctly not treated as corruption, so it never retries. + const db = openSessionSearchDatabase(join(root, 'ai-vault-search', 'index.sqlite')) + try { + expect(schemaVersion(db)).toBe(String(SESSION_SEARCH_SCHEMA_VERSION)) + } finally { + db.close() + } +}) + +it('rebuilds a newer index rather than reading a schema it does not know', async () => { + const path = await tempDatabasePath() + const newer = openSessionSearchDatabase(path) + newer.prepare("INSERT INTO files(path,byte_offset,mtime_ms) VALUES ('a',1,1)").run() + newer + .prepare("UPDATE meta SET value = ? WHERE key = 'schema_version'") + .run(String(SESSION_SEARCH_SCHEMA_VERSION + 1)) + newer.close() + + const rebuilt = openSessionSearchDatabase(path) + try { + expect(schemaVersion(rebuilt)).toBe(String(SESSION_SEARCH_SCHEMA_VERSION)) + expect(rebuilt.prepare('SELECT COUNT(*) AS c FROM files').get()).toEqual({ + c: 0 + }) + } finally { + rebuilt.close() + } +}) + +it('rebuilds when meta exists but its version row is gone', async () => { + const path = await tempDatabasePath() + const damaged = openSessionSearchDatabase(path) + damaged.prepare("INSERT INTO files(path,byte_offset,mtime_ms) VALUES ('a',1,1)").run() + // A meta table with no version is a damaged index, never a fresh one: seeding + // the current version over it would keep whatever the old schema left behind. + damaged.prepare("DELETE FROM meta WHERE key = 'schema_version'").run() + damaged.close() + + const rebuilt = openSessionSearchDatabase(path) + try { + expect(schemaVersion(rebuilt)).toBe(String(SESSION_SEARCH_SCHEMA_VERSION)) + expect(rebuilt.prepare('SELECT COUNT(*) AS c FROM files').get()).toEqual({ + c: 0 + }) + } finally { + rebuilt.close() + } +}) + +it('opens with the pragmas the write path depends on', async () => { + const db = openSessionSearchDatabase(await tempDatabasePath()) + try { + // auto_vacuum=2 is INCREMENTAL, and only takes on an empty file: without it + // a purge cannot hand pages back in bounded steps. + expect(Number(db.pragma('auto_vacuum', { simple: true }))).toBe(2) + expect(String(db.pragma('journal_mode', { simple: true })).toLowerCase()).toBe('wal') + expect(Number(db.pragma('synchronous', { simple: true }))).toBe(1) + // A WAL with no size limit never hands its space back after a large write. + expect(Number(db.pragma('journal_size_limit', { simple: true }))).toBe(8388608) + // Zero here turns every contended write into an immediate SQLITE_BUSY. + expect(Number(db.pragma('busy_timeout', { simple: true }))).toBe(5000) + } finally { + db.close() + } +}) + +it("walks a session's rows through an index rather than scanning the table", async () => { + const db = openSessionSearchDatabase(await tempDatabasePath()) + try { + // The replace delete and the orphan drain both take this path, once per file. + const plan = ( + db + .prepare('EXPLAIN QUERY PLAN SELECT id FROM messages WHERE session_row_id = ? LIMIT ?') + .all(1, 1) as { detail: string }[] + ) + .map((row) => row.detail) + .join(' ') + expect(plan).toContain('messages_session') + } finally { + db.close() + } +}) + +it('keeps only the session indexes a retrieval query can seek', async () => { + const db = openSessionSearchDatabase(await tempDatabasePath()) + try { + const names = ( + db + .prepare("SELECT name FROM sqlite_master WHERE type='index' AND tbl_name='sessions'") + .all() as { name: string }[] + ) + .map((row) => row.name) + .sort() + // One per shape PR 4's retrieval seeks: the agent filter, the newest-first + // order and date window, and the folder-prefix range scan. Fork folding reads + // `content_hash` off rows it already holds, so that column is not indexed. + expect(names).toEqual(['sessions_agent', 'sessions_cwd_key', 'sessions_updated_at']) + } finally { + db.close() + } +}) + +it("retries a Windows lock that outlives rmSync's own retries", async () => { + const path = await tempDatabasePath() + openSessionSearchDatabase(path).close() + vi.spyOn(process, 'platform', 'get').mockReturnValue('win32') + recordedRmSync.mockReset() + const locked = Object.assign(new Error('EPERM: operation not permitted'), { + code: 'EPERM' + }) + recordedRmSync.mockImplementationOnce(() => { + throw locked + }) + try { + expect(() => removeSessionSearchDatabase(path)).not.toThrow() + expect(recordedRmSync.mock.calls.length).toBe(5) + await expect(stat(path)).rejects.toMatchObject({ code: 'ENOENT' }) + } finally { + recordedRmSync.mockReset() + vi.restoreAllMocks() + } +}) + +it('gives Windows the shared retry options for a late handle release', async () => { + const path = await tempDatabasePath() + vi.spyOn(process, 'platform', 'get').mockReturnValue('win32') + recordedRmSync.mockClear() + try { + removeSessionSearchDatabase(path) + expect(recordedRmSync).toHaveBeenCalled() + for (const [, options] of recordedRmSync.mock.calls) { + expect(options).toMatchObject({ + maxRetries: WINDOWS_RM_MAX_RETRIES, + retryDelay: WINDOWS_RM_RETRY_DELAY_MS + }) + } + } finally { + vi.restoreAllMocks() + } +}) diff --git a/src/main/ai-vault-search/session-search-schema.ts b/src/main/ai-vault-search/session-search-schema.ts new file mode 100644 index 00000000000..2da1e64bc11 --- /dev/null +++ b/src/main/ai-vault-search/session-search-schema.ts @@ -0,0 +1,198 @@ +import { mkdirSync } from 'node:fs' +import { dirname } from 'node:path' +import SyncDatabase from '../sqlite/sync-database' +import { removeTreeSync } from '../../shared/windows-transient-lock-removal' + +// The index stores transcript content as written, with no redaction. A secret in +// a transcript is already plaintext under the user's home directory and is +// treated as compromised; this is a second copy of content the user already +// holds. What a snippet may carry once it leaves this machine is a transport +// policy, decided where the wire is. + +// Bump to drop and rebuild: the index is a cache over the transcripts, never a source. +export const SESSION_SEARCH_SCHEMA_VERSION = 5 + +// unicode61 keeps `_ . - /` inside tokens so paths and identifiers match exactly; +// the `identifiers` column carries the split form (see session-search-identifier-split). +// Why: `+` keeps `C++` a token of its own instead of the letter `c`; `#` is +// left out so `#123` still answers a search for `123`. +const TOKENIZER = `tokenize="unicode61 tokenchars '_.-/+'"` + +const SCHEMA_SQL = ` +CREATE TABLE IF NOT EXISTS meta(key TEXT PRIMARY KEY, value TEXT NOT NULL); +CREATE TABLE IF NOT EXISTS sessions( + -- AUTOINCREMENT, because this id names rows in the messages table for longer + -- than the row itself lives: retention cuts a session loose in one + -- transaction and reclaims its messages over many. A plain rowid is reissued + -- as max+1, so a session created inside that window would be handed a freed + -- id and adopt whatever of the purged conversation the drain had not reached, + -- behind a live session no later purge visits. + id INTEGER PRIMARY KEY AUTOINCREMENT, + agent TEXT NOT NULL, + session_id TEXT NOT NULL, + -- The transcript this session was decoded from. Not unique: OpenCode's SQLite + -- sessions all report the store's own path here, while files.path holds the + -- synthetic db#sessionId candidate that really is one per session. + file_path TEXT NOT NULL, + codex_home TEXT, + title TEXT NOT NULL, + cwd TEXT, + cwd_key TEXT, + branch TEXT, + created_at TEXT, + updated_at TEXT, + message_count INTEGER NOT NULL DEFAULT 0, + resume_command TEXT NOT NULL, + -- Chained digest of the first N messages; forks of one conversation share it. + content_hash TEXT, + content_hash_count INTEGER NOT NULL DEFAULT 0 +); +CREATE INDEX IF NOT EXISTS sessions_agent ON sessions(agent); +CREATE INDEX IF NOT EXISTS sessions_updated_at ON sessions(updated_at); +CREATE INDEX IF NOT EXISTS sessions_cwd_key ON sessions(cwd_key); +CREATE TABLE IF NOT EXISTS files( + path TEXT PRIMARY KEY, + dev INTEGER, + ino INTEGER, + byte_offset INTEGER NOT NULL, + mtime_ms REAL NOT NULL, + size_bytes INTEGER, + session_row_id INTEGER +); +-- Retention walks the expiring end of this column; without it that is a full scan and a sort. +CREATE INDEX IF NOT EXISTS files_mtime ON files(mtime_ms); +CREATE TABLE IF NOT EXISTS messages( + id INTEGER PRIMARY KEY, + session_row_id INTEGER NOT NULL, + role TEXT NOT NULL, + ts TEXT +); +-- Both the replace delete and the orphan drain walk a session's rows through this. +CREATE INDEX IF NOT EXISTS messages_session ON messages(session_row_id); +-- One FTS table, not two. A conversation-scoped search is a column filter on +-- this one — 'MATCH {user_text assistant_text}: q' with bm25 weights that zero +-- the other two — and PR 4 measured that at 1.16-1.36x the p95 of a dedicated +-- second table on a 105 MB corpus, under the 2x bar the decision was set at. +CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5( + user_text, assistant_text, tool_text, identifiers, ${TOKENIZER}, detail=full +); +` + +/** + * Opens the index, rebuilding it whenever what is on disk cannot be trusted: + * a different schema version, a version SQLite cannot report, or a file torn + * badly enough that opening or recovery fails. The index is a cache over the + * transcripts, so throwing away a bad one costs a re-scan and nothing else; + * refusing to open would strand the feature until a human deleted the file. + */ +export function openSessionSearchDatabase(path: string): SyncDatabase { + // SQLite will not create the directory, and its failure is `unable to open + // database file`, which is correctly not corruption — so without this the + // feature strands on a profile that has never held an index. + if (path !== ':memory:') { + mkdirSync(dirname(path), { recursive: true }) + } + try { + return openExisting(path) + } catch (error) { + if (!isUnusableDatabaseError(error)) { + throw error + } + // One retry only: a second failure on a file we just created is not corruption. + removeSessionSearchDatabase(path) + return openExisting(path) + } +} + +function openExisting(path: string): SyncDatabase { + // Nulled while no handle is open, because closing an already-closed handle + // throws ERR_INVALID_STATE, which would replace whatever really failed — + // an unlink refused by a virus scanner or a second Orca holding the file — + // with an error nothing classifies as worth rebuilding for. + let db: SyncDatabase | null = openWithPragmas(path) + try { + if (isStaleSchema(db)) { + // Why: DROP TABLE on a multi-GB FTS index takes minutes and runs inside the + // scanner service's init, past its ready timeout; unlinking is instant. + db.close() + db = null + removeSessionSearchDatabase(path) + db = openWithPragmas(path) + } + db.exec(SCHEMA_SQL) + db.prepare('INSERT OR REPLACE INTO meta(key, value) VALUES (?, ?)').run( + 'schema_version', + String(SESSION_SEARCH_SCHEMA_VERSION) + ) + return db + } catch (error) { + db?.close() + throw error + } +} + +// SQLite reports a torn file at the first statement that has to read a page, so +// this has to match on the message as well as the code. +const UNUSABLE_DATABASE = + /SQLITE_CORRUPT|SQLITE_NOTADB|file is not a database|database disk image is malformed/i + +function isUnusableDatabaseError(error: unknown): boolean { + if (!(error instanceof Error)) { + return false + } + const code = (error as { code?: unknown }).code + return ( + (typeof code === 'string' && UNUSABLE_DATABASE.test(code)) || + UNUSABLE_DATABASE.test(error.message) + ) +} + +function openWithPragmas(path: string): SyncDatabase { + const db = new SyncDatabase(path) + try { + // Why: only takes effect on an empty file; it is what lets a purge hand pages + // back in bounded steps instead of a full VACUUM. Set before any table exists. + db.pragma('auto_vacuum = INCREMENTAL') + // The whole consistency model: a file's rows and its cursor land in one + // transaction, and a reader on another handle sees the last committed state + // of the index rather than a session half way through being rewritten. + db.pragma('journal_mode = WAL') + db.pragma('synchronous = NORMAL') + db.pragma('journal_size_limit = 8388608') + db.pragma('busy_timeout = 5000') + return db + } catch (error) { + db?.close() + throw error + } +} + +export function removeSessionSearchDatabase(path: string): void { + if (path === ':memory:') { + return + } + for (const suffix of ['', '-wal', '-shm', '-journal']) { + removeTreeSync(`${path}${suffix}`) + } +} + +/** + * Whether what is on disk has to be thrown away. No `meta` table at all is a + * file with nothing in it to throw away, and removing it would make the first + * open of every new profile a create-remove-create. A meta table whose version + * row is missing or unparseable is a damaged index rather than a new one: + * seeding the current version over it would keep whatever rows the old schema + * left. + */ +function isStaleSchema(db: SyncDatabase): boolean { + const table = db + .prepare("SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'meta'") + .get() + if (!table) { + return false + } + const row = db.prepare("SELECT value FROM meta WHERE key = 'schema_version'").get() as + | { value: string } + | undefined + return (row ? Number(row.value) : Number.NaN) !== SESSION_SEARCH_SCHEMA_VERSION +} diff --git a/src/main/ai-vault-search/session-search-store.ts b/src/main/ai-vault-search/session-search-store.ts new file mode 100644 index 00000000000..da9f4d6f731 --- /dev/null +++ b/src/main/ai-vault-search/session-search-store.ts @@ -0,0 +1,253 @@ +import type SyncDatabase from '../sqlite/sync-database' +import type { SessionFileCandidate } from '../ai-vault/session-scanner-types' +import type { TranscriptSessionIdentity } from '../ai-vault/session-transcript-consumers' +import type { + SessionSearchFileIdentity, + SessionSearchIndexedFile +} from './session-search-file-cursor' +import { + SESSION_SEARCH_COMMIT_CHARS, + SessionSearchIndexWriter, + type SessionSearchFileWrite +} from './session-search-index-writer' +import { deleteExpiredSearchFiles, drainOrphanedMessages } from './session-search-retention-delete' +import { openSessionSearchDatabase } from './session-search-schema' + +// A paused store keeps recording what it declined, so the set needs a ceiling. +// Above it the oldest record goes and the drop is counted, because a re-read set +// that silently forgets is worse than one that says it is incomplete. +export const STALE_PATH_LIMIT = 20_000 + +/** + * Owns the index database. PR 2 scope: the write half only — the transcript + * consumer writes through it and nothing reads from it yet. Lifecycle (who + * indexes, when, and how the re-read set is drained) belongs to the service. + */ +export class SessionSearchStore { + private readonly db: SyncDatabase + private readonly writer: SessionSearchIndexWriter + private closed = false + private acceptingWrites = true + private retentionCutoffMs: number | null = null + // Files this index knows it is behind on. Filled by a declined or abandoned + // read; PR 3's indexer drains it. Nothing here schedules the re-read. + private readonly stale = new Map() + private droppedStalePaths = 0 + // One drain at a time. A replace that commits while one is running asks for + // another pass rather than starting a second walk of the same rows. + private draining = false + private drainRequested = false + + constructor( + path: string, + private readonly onError: (error: unknown) => void = (error) => + console.warn( + '[ai-vault-search] index write failed:', + error instanceof Error ? error.name : 'IndexError' + ) + ) { + this.db = openSessionSearchDatabase(path) + this.writer = new SessionSearchIndexWriter(this.db, SESSION_SEARCH_COMMIT_CHARS, () => + this.scheduleOrphanDrain() + ) + } + + /** + * Reclaims the rows a replace cut loose, once its transaction has committed. + * + * The same split retention makes, for the same reason: deleting the old + * session row is what stops it answering, because every retrieval joins + * `sessions`, and handing its messages back is the expensive half that must + * not hold one transaction. Nothing records the work: rows whose session row + * is gone are the whole record, so a crash before or during a drain is found + * by the next one. + */ + private scheduleOrphanDrain(): void { + this.drainRequested = true + if (this.draining || this.closed) { + return + } + this.draining = true + // Off the committing stack. An async function runs synchronously up to its + // first `await`, so calling the drain here would put its first batch back + // inside the call that committed the replace — the cost this took out. + void Promise.resolve().then(() => this.runOrphanDrain()) + } + + private async runOrphanDrain(): Promise { + try { + while (this.drainRequested && !this.closed) { + this.drainRequested = false + await drainOrphanedMessages(this.db, () => this.closed) + } + } catch (error) { + if (!this.closed) { + this.onError(error) + } + } finally { + this.draining = false + } + } + + /** + * The index handle, for a reader composed over this store (PR 4's engine). + * + * Two rules come with it, both measured in this PR. **Never hold a read + * transaction across an `await`**: a checkpoint cannot pass an open read + * snapshot, so a paginated read that opened `BEGIN` and yielded between pages + * takes the WAL from 10 MB to 266 MB and it does not come back. And **no + * `.iterate()` that outlives its statement**, which is the same pin by + * another name. Every retrieval a single synchronous statement is the whole + * contract. + */ + get connection(): SyncDatabase { + return this.db + } + + setAcceptingWrites(accept: boolean): void { + this.acceptingWrites = accept + } + + /** The oldest transcript mtime worth indexing; PR 3 derives it from the retention setting. */ + setRetentionCutoffMs(cutoffMs: number | null): void { + this.retentionCutoffMs = cutoffMs + } + + /** Whether this candidate is new enough to be worth holding rows for at all. */ + private withinRetention(candidate: SessionFileCandidate): boolean { + return this.retentionCutoffMs === null || candidate.file.mtimeMs >= this.retentionCutoffMs + } + + /** Whether a write for this candidate may start right now. */ + acceptsCandidate(candidate: SessionFileCandidate): boolean { + return !this.closed && this.acceptingWrites && this.withinRetention(candidate) + } + + indexedFile(path: string, identity: SessionSearchFileIdentity): SessionSearchIndexedFile | null { + try { + return this.writer.indexedFile(path, identity) + } catch (error) { + this.onError(error) + return null + } + } + + /** Null when this read cannot extend the index, or when the store refuses writes. */ + beginWrite( + candidate: SessionFileCandidate, + mode: 'replace' | 'append', + previousByteOffset: number, + identity?: () => TranscriptSessionIdentity | null + ): SessionSearchFileWrite | null { + if (!this.acceptsCandidate(candidate)) { + return null + } + try { + return this.writer.beginWrite(candidate, mode, previousByteOffset, identity) + } catch (error) { + this.reportWriteFailure(error) + return null + } + } + + writeCommitted(candidate: SessionFileCandidate): void { + // Why: a list scan queues every file the backfill has not reached yet; once + // one lands, a later pass must not re-read the whole queue. + this.stale.delete(candidate.file.path) + } + + reportWriteFailure(error: unknown): void { + this.onError(error) + } + + /** + * Records a file whose content the index is behind on, for a later whole + * re-read. Recorded while paused too: a pause is exactly the window in which + * reads are declined, so refusing to remember them would lose every file the + * pause covered. + */ + markStale(candidate: SessionFileCandidate): void { + if (this.closed || !this.withinRetention(candidate)) { + return + } + // Re-inserting moves the path to the end, so the oldest record is the one + // dropped when a long pause overruns the bound. + this.stale.delete(candidate.file.path) + this.stale.set(candidate.file.path, candidate) + while (this.stale.size > STALE_PATH_LIMIT) { + const oldest = this.stale.keys().next() + if (oldest.done) { + break + } + this.stale.delete(oldest.value) + this.droppedStalePaths += 1 + } + } + + /** + * Files the index knew it was behind on and could not keep a record of. A + * non-zero count means the re-read set is incomplete, so coverage cannot be + * reported as whole until a full pass runs. + */ + get droppedPendingFileCount(): number { + return this.droppedStalePaths + } + + /** + * Hands the re-read set to its scheduler and clears it. + * + * These paths are behind, not merely dirty: the index declined their last read + * because it covered a span the index never saw. Re-dispatching a scan is not + * enough on its own, because the reader picks `append` from the session list's + * resume point and the consumer will decline again. The caller must pass each + * path to `requestWholeTranscriptRead` first. + */ + takeStale(): SessionFileCandidate[] { + const candidates = [...this.stale.values()] + this.stale.clear() + return candidates + } + + get pendingFileCount(): number { + return this.stale.size + } + + /** + * Drops a source's rows. Only a proven deletion may call this: an unreadable + * source is `unverifiable`, not `missing`, and keeps its rows + * (docs/reference/ssh-execution-boundary.md). + */ + removeFile(path: string): void { + this.stale.delete(path) + try { + this.writer.removeFile(path) + } catch (error) { + this.onError(error) + } + } + + /** Cuts expired sessions loose at once, then reclaims their rows in resumable batches. */ + async purgeOlderThan(cutoffMs: number | null, signal?: AbortSignal): Promise { + try { + await deleteExpiredSearchFiles( + this.db, + cutoffMs, + () => this.closed || signal?.aborted === true + ) + } catch (error) { + if (!this.closed) { + this.onError(error) + } + } + } + + close(): void { + // node:sqlite throws ERR_INVALID_STATE on a second close, and a store is + // closed both by its owner and by a test's teardown. + if (this.closed) { + return + } + this.closed = true + this.db.close() + } +} diff --git a/src/main/ai-vault-search/session-search-synthetic-corpus.test.ts b/src/main/ai-vault-search/session-search-synthetic-corpus.test.ts new file mode 100644 index 00000000000..9b6a48beb4e --- /dev/null +++ b/src/main/ai-vault-search/session-search-synthetic-corpus.test.ts @@ -0,0 +1,44 @@ +import { rm } from 'node:fs/promises' +import { join } from 'node:path' +import { expect, it } from 'vitest' +import { resetSessionParseCacheForTests } from '../ai-vault/session-scanner-parse-cache' +import { resetTranscriptConsumersForTests } from '../ai-vault/session-transcript-consumers' +import { registerSessionSearchIndexConsumer } from './session-search-index-consumer' +import { SessionSearchStore } from './session-search-store' +import { writeSyntheticTranscriptCorpus } from './session-search-synthetic-corpus' +import { parseTranscript } from './session-search-transcript-fixtures' + +it.each([Infinity, -Infinity, Number.NaN, -1, 1.5])( + 'rejects invalid corpus loop bounds: %s', + async (value) => { + for (const field of ['sessions', 'turnsPerSession', 'toolResultWords']) { + await expect(writeSyntheticTranscriptCorpus({ [field]: value })).rejects.toThrow(RangeError) + } + } +) + +it.each([0, 200, 2000])( + 'counts the indexed messages with %s tool words', + async (toolResultWords) => { + const corpus = await writeSyntheticTranscriptCorpus({ + sessions: 1, + turnsPerSession: 1, + toolResultWords + }) + const store = new SessionSearchStore(join(corpus.root, 'index.sqlite')) + const unregister = registerSessionSearchIndexConsumer(store) + try { + await parseTranscript(corpus.files[0]!) + expect(corpus.messageCount).toBe(toolResultWords === 0 ? 3 : 4) + expect(store.connection.prepare('SELECT count(*) AS n FROM messages').get()).toEqual({ + n: corpus.messageCount + }) + } finally { + unregister() + resetTranscriptConsumersForTests() + resetSessionParseCacheForTests() + store.close() + await rm(corpus.root, { recursive: true, force: true }) + } + } +) diff --git a/src/main/ai-vault-search/session-search-synthetic-corpus.ts b/src/main/ai-vault-search/session-search-synthetic-corpus.ts new file mode 100644 index 00000000000..14e8a27fe5f --- /dev/null +++ b/src/main/ai-vault-search/session-search-synthetic-corpus.ts @@ -0,0 +1,154 @@ +import { mkdtemp, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' + +// Why synthetic and in-repo: the cost model has to be reproducible on any host +// and must never read a real transcript. The shapes here mirror what a Claude +// JSONL transcript actually holds — prose turns, a pasted diff, tool calls and +// their output — because the index's disk cost tracks the mix, not the size. + +const WORDS = [ + 'terminal', + 'reattach', + 'worktree', + 'resolveTerminalPath', + 'src/main/ai-vault/session-transcript-reader.ts', + 'the', + 'index', + 'cursor', + 'byteOffset', + 'publish', + 'staged', + 'transaction', + 'MAX_RETRIES', + 'relay', + 'daemon', + 'pty', + 'snapshot', + 'because' +] + +/** Deterministic: the same seed gives the same corpus on every host and run. */ +function mulberry32(seed: number): () => number { + let state = seed >>> 0 + return () => { + state = (state + 0x6d2b79f5) >>> 0 + let t = Math.imul(state ^ (state >>> 15), 1 | state) + t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t + return ((t ^ (t >>> 14)) >>> 0) / 4294967296 + } +} + +function words(random: () => number, count: number): string { + const out: string[] = [] + for (let index = 0; index < count; index++) { + out.push(WORDS[Math.floor(random() * WORDS.length)]) + } + return out.join(' ') +} + +export type SyntheticCorpus = { + root: string + files: string[] + /** Total bytes of transcript written, the denominator of write amplification. */ + transcriptBytes: number + messageCount: number +} + +export type SyntheticCorpusOptions = { + sessions?: number + turnsPerSession?: number + seed?: number + /** + * Words per tool result. The default keeps tool output at about half the + * message text; the real distribution is 80-97 %, which is what prices the + * tool-row cap, so the benchmark runs a second arm well above the default. + */ + toolResultWords?: number +} + +/** Writes a corpus of Claude JSONL transcripts and reports what it cost on disk. */ +export async function writeSyntheticTranscriptCorpus( + options: SyntheticCorpusOptions = {} +): Promise { + const sessions = options.sessions ?? 40 + const turns = options.turnsPerSession ?? 60 + const toolWords = options.toolResultWords ?? 200 + for (const [name, value] of Object.entries({ + sessions, + turnsPerSession: turns, + toolResultWords: toolWords + })) { + if (!Number.isSafeInteger(value) || value < 0) { + throw new RangeError(`${name} must be a finite non-negative safe integer`) + } + } + const random = mulberry32(options.seed ?? 1) + const root = await mkdtemp(join(tmpdir(), 'orca-search-corpus-')) + const files: string[] = [] + let transcriptBytes = 0 + let messageCount = 0 + + for (let session = 0; session < sessions; session++) { + const sessionId = `00000000-0000-4000-8000-${String(session).padStart(12, '0')}` + const lines: string[] = [] + for (let turn = 0; turn < turns; turn++) { + const at = new Date(1740000000000 + turn * 60_000).toISOString() + lines.push( + JSON.stringify({ + type: 'user', + sessionId, + timestamp: at, + cwd: `/repo/app-${session % 7}`, + gitBranch: 'main', + message: { role: 'user', content: words(random, 40) } + }) + ) + lines.push( + JSON.stringify({ + type: 'assistant', + sessionId, + timestamp: at, + message: { + role: 'assistant', + model: 'claude-fable-5', + content: [ + { type: 'text', text: words(random, 120) }, + { + type: 'tool_use', + name: 'Bash', + input: { command: `rg ${words(random, 3)}` } + } + ] + } + }) + ) + lines.push( + JSON.stringify({ + type: 'user', + sessionId, + timestamp: at, + message: { + role: 'user', + content: [ + { + type: 'tool_result', + tool_use_id: 'toolu_1', + content: words(random, toolWords) + } + ] + } + }) + ) + // Empty tool results emit no searchable message. + messageCount += toolWords === 0 ? 3 : 4 + } + const path = join(root, `${sessionId}.jsonl`) + const body = `${lines.join('\n')}\n` + await writeFile(path, body) + transcriptBytes += Buffer.byteLength(body) + files.push(path) + } + + return { root, files, transcriptBytes, messageCount } +} diff --git a/src/main/ai-vault-search/session-search-transcript-fixtures.ts b/src/main/ai-vault-search/session-search-transcript-fixtures.ts new file mode 100644 index 00000000000..bc7eda9a8ff --- /dev/null +++ b/src/main/ai-vault-search/session-search-transcript-fixtures.ts @@ -0,0 +1,118 @@ +import { stat } from 'node:fs/promises' +import { + createSessionParseStats, + parseAgentSessionFileCached, + type SessionParseStats +} from '../ai-vault/session-scanner-parse-cache' +import type { SessionFileCandidate } from '../ai-vault/session-scanner-types' + +// Transcript builders shared by the session-search store tests; each file owns +// its temp directories, this module only shapes records and drives the parser. + +export const CLAUDE_SESSION_ID = 'aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee' +export const CODEX_SESSION_ID = '019f0000-1111-7222-8333-444444444444' +export const CODEX_ROLLOUT_FILE = `rollout-2026-05-01T10-00-00-${CODEX_SESSION_ID}.jsonl` + +const RECORD_EPOCH_MS = 1740000000000 + +export function recordTimestamp(index: number): string { + return new Date(RECORD_EPOCH_MS + index * 60_000).toISOString() +} + +export function userRecord( + index: number, + content: unknown, + sessionId = CLAUDE_SESSION_ID, + cwd = '/repo/app' +): string { + return JSON.stringify({ + type: 'user', + sessionId, + timestamp: recordTimestamp(index), + cwd, + gitBranch: 'main', + message: { role: 'user', content } + }) +} + +export function assistantRecord( + index: number, + content: unknown, + sessionId = CLAUDE_SESSION_ID +): string { + return JSON.stringify({ + type: 'assistant', + sessionId, + timestamp: recordTimestamp(index), + message: { role: 'assistant', model: 'claude-fable-5', content } + }) +} + +export async function sessionCandidate( + agent: SessionFileCandidate['agent'], + path: string, + codexHome: string | null = null +): Promise { + const fileStat = await stat(path) + return { + agent, + codexHome, + file: { + path, + mtimeMs: fileStat.mtimeMs, + modifiedAt: fileStat.mtime.toISOString(), + sizeBytes: fileStat.size, + dev: fileStat.dev, + ino: fileStat.ino + } + } +} + +export async function parseTranscript( + path: string, + agent: SessionFileCandidate['agent'] = 'claude', + codexHome: string | null = null +): Promise<{ stats: SessionParseStats }> { + const stats = createSessionParseStats() + await parseAgentSessionFileCached( + await sessionCandidate(agent, path, codexHome), + process.platform, + stats + ) + return { stats } +} + +function codexLine(record: Record): string { + return JSON.stringify(record) +} + +/** Minimal Codex rollout: meta, one user message, one completed shell command. */ +export function codexRolloutLines(command: string[], output: string, prompt: string): string[] { + return [ + codexLine({ + timestamp: recordTimestamp(0), + type: 'session_meta', + payload: { id: CODEX_SESSION_ID, cwd: '/repo/app', git: { branch: 'main' } } + }), + codexLine({ + timestamp: recordTimestamp(1), + type: 'response_item', + payload: { type: 'message', role: 'user', content: prompt } + }), + codexLine({ + timestamp: recordTimestamp(2), + type: 'response_item', + payload: { + type: 'function_call', + call_id: 'call-1', + name: 'shell', + arguments: JSON.stringify({ command }) + } + }), + codexLine({ + timestamp: recordTimestamp(3), + type: 'response_item', + payload: { type: 'function_call_output', call_id: 'call-1', output } + }) + ] +} diff --git a/src/main/ai-vault/session-scanner-accumulator.ts b/src/main/ai-vault/session-scanner-accumulator.ts index 88e09627eb7..18f273d6be3 100644 --- a/src/main/ai-vault/session-scanner-accumulator.ts +++ b/src/main/ai-vault/session-scanner-accumulator.ts @@ -23,7 +23,11 @@ import { normalizePreviewText, timestampMs } from './session-scanner-values' -import { NO_TRANSCRIPT_MESSAGES, type TranscriptMessageSink } from './session-transcript-consumers' +import { + NO_TRANSCRIPT_MESSAGES, + type TranscriptMessageSink, + type TranscriptSessionIdentity +} from './session-transcript-consumers' import { boundedText, transcriptMessageRole, @@ -64,6 +68,28 @@ export function createAccumulator(args: { } } +/** + * The session identity a fold holds right now. Null until it has an id, which + * every supported format writes in the opening lines of the transcript. + */ +export function accumulatorSessionIdentity( + accumulator: SessionAccumulator +): TranscriptSessionIdentity | null { + const sessionId = accumulator.sessionId.trim() + if (!sessionId) { + return null + } + return { + sessionId, + cwd: accumulator.cwd, + // The generated fallback is `finalizeSession`'s, not this one's: a title + // that is still absent mid-read is better said to be absent. + title: accumulator.title ?? accumulator.fallbackTitle, + createdAt: accumulator.createdAt, + updatedAt: accumulator.updatedAt + } +} + export function cloneSessionAccumulator(accumulator: SessionAccumulator): SessionAccumulator { return { ...accumulator, previewMessages: [...accumulator.previewMessages] } } @@ -77,6 +103,7 @@ export function accumulatorFoldResumeState( ): ResumableSessionParseState { return { consumeLine: (line) => consumeRecordLine(accumulator, line), + identity: () => accumulatorSessionIdentity(accumulator), clone: () => accumulatorFoldResumeState(cloneSessionAccumulator(accumulator), consumeRecordLine), touchFile: (file) => { diff --git a/src/main/ai-vault/session-scanner-codex-message-records.ts b/src/main/ai-vault/session-scanner-codex-message-records.ts index 5aa739275ff..a8a5c3c9d13 100644 --- a/src/main/ai-vault/session-scanner-codex-message-records.ts +++ b/src/main/ai-vault/session-scanner-codex-message-records.ts @@ -1,3 +1,7 @@ +import { + publishCodexResponseTool, + publishCodexCompletedTool +} from './session-scanner-codex-tool-records' import { normalizePromptField } from '../../shared/agent-status-field-normalization' import { addPreviewContent } from './session-scanner-accumulator' import type { SessionAccumulator } from './session-scanner-types' @@ -8,6 +12,10 @@ export function consumeCodexResponseMessage( payload: Record, timestamp: unknown ): boolean { + publishCodexResponseTool(accumulator, payload, timestamp) + if (payload.type !== 'message') { + return false + } accumulator.messageCount++ const role = payload.role === 'assistant' ? 'assistant' : payload.role === 'user' ? 'user' : 'unknown' @@ -24,6 +32,7 @@ export function consumeCodexCompletedMessage( payload: Record, timestamp: unknown ): boolean { + publishCodexCompletedTool(accumulator, payload, timestamp) const item = asRecord(payload.item) if (!item) { return false diff --git a/src/main/ai-vault/session-scanner-codex-parser.ts b/src/main/ai-vault/session-scanner-codex-parser.ts index 02a385400de..b3571d4a337 100644 --- a/src/main/ai-vault/session-scanner-codex-parser.ts +++ b/src/main/ai-vault/session-scanner-codex-parser.ts @@ -4,6 +4,7 @@ import type { AiVaultSession } from '../../shared/ai-vault-types' import { readCodexSessionIndexTitle } from './session-scanner-codex-title-index' import type { ExecutionHostId } from '../../shared/execution-host' import { + accumulatorSessionIdentity, cloneSessionAccumulator, createAccumulator, finalizeSession, @@ -153,19 +154,13 @@ function consumeCodexRecordLine(state: CodexSessionParseState, line: string): vo accumulator.title = metadataTitle state.titleSource = 'meta' } - const cwd = extractString(payload.cwd) - if (cwd) { - accumulator.cwd = cwd - } + accumulator.cwd = extractString(payload.cwd) ?? accumulator.cwd accumulator.branch = extractGitBranch(payload.git) ?? accumulator.branch return } if (record.type === 'turn_context' && payload) { - const cwd = extractString(payload.cwd) - if (cwd) { - accumulator.cwd = cwd - } + accumulator.cwd = extractString(payload.cwd) ?? accumulator.cwd const model = extractModel(payload) if (model) { accumulator.model = model @@ -177,7 +172,7 @@ function consumeCodexRecordLine(state: CodexSessionParseState, line: string): vo return } - if (record.type === 'response_item' && payload.type === 'message') { + if (record.type === 'response_item') { if (state.historyMode === 'paginated') { return } @@ -285,7 +280,10 @@ function codexResumeStateFromParseState( return { consumeLine: (line) => consumeCodexRecordLine(state, line), consumeLineBytes: (line) => { - const timelineOnlyRecord = readCodexTimelineOnlyRecord(line) + const timelineOnlyRecord = readCodexTimelineOnlyRecord( + line, + state.accumulator.messages.active && state.historyMode !== 'paginated' + ) if (timelineOnlyRecord) { updateTimeline(state.accumulator, timelineOnlyRecord.timestamp) } else { @@ -293,6 +291,7 @@ function codexResumeStateFromParseState( } }, shouldStop: () => state.rejectedWorkerSession, + identity: () => accumulatorSessionIdentity(state.accumulator), clone: () => codexResumeStateFromParseState(cloneCodexParseState(state), codexHome, titleReader), touchFile: (file) => { diff --git a/src/main/ai-vault/session-scanner-codex-record-fast-path.ts b/src/main/ai-vault/session-scanner-codex-record-fast-path.ts index 1c4322f89f6..852ec174e95 100644 --- a/src/main/ai-vault/session-scanner-codex-record-fast-path.ts +++ b/src/main/ai-vault/session-scanner-codex-record-fast-path.ts @@ -1,3 +1,5 @@ +import { CODEX_TOOL_RESPONSE_TYPES } from './session-scanner-codex-tool-records' + // Records below this size are decoded and parsed exactly: JSON.parse on a // kilobyte costs less than the risk of a prefix heuristic, and the scan cost // this path exists to remove is entirely in megabyte-scale records. @@ -22,7 +24,10 @@ const PARSED_EVENT_TYPES = new Set([ ]) /** Returns the timestamp only when the record cannot affect other visible session fields. */ -export function readCodexTimelineOnlyRecord(line: Buffer): { timestamp: string } | null { +export function readCodexTimelineOnlyRecord( + line: Buffer, + includeTools = false +): { timestamp: string } | null { if (line.length <= CODEX_RECORD_PREFIX_LIMIT) { return null } @@ -41,6 +46,13 @@ export function readCodexTimelineOnlyRecord(line: Buffer): { timestamp: string } if (!payloadType) { return null } + if ( + includeTools && + recordType === 'response_item' && + CODEX_TOOL_RESPONSE_TYPES.has(payloadType) + ) { + return null + } const parsedPayloadTypes = recordType === 'response_item' ? PARSED_RESPONSE_ITEM_TYPES : PARSED_EVENT_TYPES return parsedPayloadTypes.has(payloadType) ? null : { timestamp } diff --git a/src/main/ai-vault/session-scanner-codex-tool-records.test.ts b/src/main/ai-vault/session-scanner-codex-tool-records.test.ts new file mode 100644 index 00000000000..a5aa909bfa1 --- /dev/null +++ b/src/main/ai-vault/session-scanner-codex-tool-records.test.ts @@ -0,0 +1,125 @@ +import { expect, it } from 'vitest' +import { createCodexSessionResumeState } from './session-scanner-codex-parser' +import type { TranscriptMessage } from './session-transcript-consumers' +import { readCodexTimelineOnlyRecord } from './session-scanner-codex-record-fast-path' + +const timestamp = '2026-05-01T10:00:00.000Z' +const file = { + path: '/fixture/rollout.jsonl', + mtimeMs: Date.parse(timestamp), + modifiedAt: timestamp +} +const record = (type: string, payload: Record): Buffer => + Buffer.from(JSON.stringify({ timestamp, type, payload })) + +it.each(['function_call_output', 'custom_tool_call_output'])( + 'reads large %s records only when a consumer needs them', + (type) => { + const line = record('response_item', { type, output: 'outputonly '.repeat(300) }) + expect(readCodexTimelineOnlyRecord(line)).toEqual({ timestamp }) + expect(readCodexTimelineOnlyRecord(line, true)).toBeNull() + const messages: TranscriptMessage[] = [] + const state = createCodexSessionResumeState(file, null, { + active: true, + push: (message) => messages.push(message) + }) + state.consumeLineBytes!(line) + expect(messages).toEqual([{ role: 'tool', text: 'outputonly '.repeat(300), timestamp }]) + } +) + +it.each([false, true])( + 'uses one tool representation across append when paginated=%s', + async (paginated) => { + const messages: TranscriptMessage[] = [] + let state = createCodexSessionResumeState(file, null, { + active: true, + push: (message) => messages.push(message) + }) + const consume = (type: string, payload: Record) => + state.consumeLineBytes!(record(type, payload)) + consume('session_meta', { id: 'session-1', history_mode: paginated ? 'paginated' : 'full' }) + consume('response_item', { type: 'message', role: 'user', content: 'promptonly' }) + consume('event_msg', { + type: 'item_completed', + item: { type: 'UserMessage', content: [{ type: 'text', text: 'promptonly' }] } + }) + consume('response_item', { + type: 'function_call', + name: 'shell', + arguments: '{"command":"commandonly"}' + }) + // The next scan resumes between the call and its output. + state = state.clone() + consume('response_item', { type: 'function_call_output', output: 'outputonly' }) + consume('event_msg', { + type: 'item_completed', + item: { type: 'CommandExecution', command: ['commandonly'], aggregated_output: 'outputonly' } + }) + expect(messages.filter((message) => message.text.includes('commandonly'))).toHaveLength(1) + expect(messages.filter((message) => message.text === 'outputonly')).toEqual([ + { role: 'tool', text: 'outputonly', timestamp } + ]) + expect(messages.filter((message) => message.role === 'user')).toHaveLength(1) + expect(await state.finalize(process.platform)).toMatchObject({ messageCount: 1 }) + } +) + +it.each([ + { type: 'add', content: '+ addedneedle' }, + { type: 'delete', content: '+ addedneedle' }, + { type: 'update', unified_diff: '+ addedneedle', move_path: null } +])('publishes paginated $type file changes', (change) => { + const messages: TranscriptMessage[] = [] + const state = createCodexSessionResumeState(file, null, { + active: true, + push: (message) => messages.push(message) + }) + state.consumeLineBytes!(record('session_meta', { id: 'session-1', history_mode: 'paginated' })) + state.consumeLineBytes!( + record('event_msg', { + type: 'item_completed', + item: { type: 'FileChange', changes: { 'src/changed.ts': change } } + }) + ) + expect(messages.map((message) => [message.role, message.text])).toEqual([ + ['tool', 'apply_patch: src/changed.ts'], + ['tool', '+ addedneedle'] + ]) +}) + +it('normalizes custom calls and structured results through the existing content reader', () => { + const messages: TranscriptMessage[] = [] + const state = createCodexSessionResumeState(file, null, { + active: true, + push: (message) => messages.push(message) + }) + state.consumeLineBytes!( + record('response_item', { type: 'custom_tool_call', name: 'apply_patch', input: 'patchneedle' }) + ) + state.consumeLineBytes!( + record('response_item', { + type: 'custom_tool_call_output', + output: { content: [{ type: 'text', text: 'resultneedle' }] } + }) + ) + expect(messages.map((message) => message.text)).toEqual([ + 'apply_patch: patchneedle', + 'resultneedle' + ]) +}) + +it('keeps local shell argv searchable', () => { + const messages: TranscriptMessage[] = [] + const state = createCodexSessionResumeState(file, null, { + active: true, + push: (message) => messages.push(message) + }) + state.consumeLineBytes!( + record('response_item', { + type: 'local_shell_call', + action: { type: 'exec', command: ['rg', 'argvneedle'] } + }) + ) + expect(messages.map((message) => message.text)).toEqual(['tool: rg argvneedle']) +}) diff --git a/src/main/ai-vault/session-scanner-codex-tool-records.ts b/src/main/ai-vault/session-scanner-codex-tool-records.ts new file mode 100644 index 00000000000..976d5eff36d --- /dev/null +++ b/src/main/ai-vault/session-scanner-codex-tool-records.ts @@ -0,0 +1,95 @@ +import { timestampIso } from './session-scanner-accumulator' +import { asRecord } from './session-scanner-record-value' +import type { SessionAccumulator } from './session-scanner-types' +import { transcriptMessagesFromContent } from './session-transcript-message-content' + +export const CODEX_TOOL_RESPONSE_TYPES = new Set([ + 'function_call', + 'local_shell_call', + 'custom_tool_call', + 'function_call_output', + 'custom_tool_call_output' +]) + +function publishToolContent( + accumulator: SessionAccumulator, + content: unknown, + timestamp: unknown +): void { + for (const message of transcriptMessagesFromContent('tool', content, timestampIso(timestamp))) { + accumulator.messages.push(message) + } +} + +export function publishCodexResponseTool( + accumulator: SessionAccumulator, + payload: Record, + timestamp: unknown +): void { + if (!accumulator.messages.active || !CODEX_TOOL_RESPONSE_TYPES.has(String(payload.type))) { + return + } + if (payload.type === 'function_call_output' || payload.type === 'custom_tool_call_output') { + const output = asRecord(payload.output) + publishToolContent( + accumulator, + [{ type: 'tool_result', content: output?.content ?? output?.output ?? payload.output }], + timestamp + ) + return + } + const input = payload.arguments ?? payload.input ?? payload.action + const action = asRecord(input) + const normalizedInput = + action && Array.isArray(action.command) + ? { ...action, command: action.command.filter((part) => typeof part === 'string').join(' ') } + : input + publishToolContent( + accumulator, + [ + { + type: 'tool_use', + name: payload.name ?? 'tool', + input: normalizedInput + } + ], + timestamp + ) +} + +export function publishCodexCompletedTool( + accumulator: SessionAccumulator, + payload: Record, + timestamp: unknown +): void { + if (!accumulator.messages.active) { + return + } + const item = asRecord(payload.item) + if (item?.type === 'CommandExecution' || item?.type === 'command_execution') { + const command = Array.isArray(item.command) + ? item.command.filter((part) => typeof part === 'string').join(' ') + : item.command + publishToolContent( + accumulator, + [ + { type: 'tool_use', name: 'shell', input: command }, + { type: 'tool_result', content: item.aggregated_output ?? item.aggregatedOutput } + ], + timestamp + ) + } else if (item?.type === 'FileChange' || item?.type === 'file_change') { + const changes = asRecord(item.changes) ?? {} + for (const [path, value] of Object.entries(changes)) { + const change = asRecord(value) + publishToolContent( + accumulator, + [ + { type: 'tool_use', name: 'apply_patch', input: { path } }, + { type: 'tool_result', content: change?.unified_diff ?? change?.content } + ], + timestamp + ) + } + } +} diff --git a/src/main/ai-vault/session-scanner-omp-subagent-transcripts.ts b/src/main/ai-vault/session-scanner-omp-subagent-transcripts.ts index 57cadebeee0..07f3646b537 100644 --- a/src/main/ai-vault/session-scanner-omp-subagent-transcripts.ts +++ b/src/main/ai-vault/session-scanner-omp-subagent-transcripts.ts @@ -101,6 +101,7 @@ export function withOmpSubagentTranscriptCount( ): ResumableSessionParseState { return { consumeLine: (line) => state.consumeLine(line), + identity: () => state.identity?.() ?? null, clone: () => withOmpSubagentTranscriptCount(state.clone(), transcriptFilePath), touchFile: (file) => state.touchFile(file), finalize: async (platform, options) => { diff --git a/src/main/ai-vault/session-scanner-primary-parsers.ts b/src/main/ai-vault/session-scanner-primary-parsers.ts index 9783f8a0520..62149920b5a 100644 --- a/src/main/ai-vault/session-scanner-primary-parsers.ts +++ b/src/main/ai-vault/session-scanner-primary-parsers.ts @@ -12,6 +12,7 @@ import type { } from './session-scanner-types' import type { TranscriptMessageSink } from './session-transcript-consumers' import { + accumulatorSessionIdentity, addPreviewContent, createAccumulator, finalizeSession, @@ -205,6 +206,7 @@ function claudeResumeStateFromParseState( ): ResumableSessionParseState { return { consumeLine: (line) => consumeClaudeSessionLine(state, line), + identity: () => accumulatorSessionIdentity(state.accumulator), clone: () => claudeResumeStateFromParseState(cloneClaudeSessionParseState(state)), touchFile: (file) => { state.accumulator.modifiedAt = file.modifiedAt diff --git a/src/main/ai-vault/session-scanner-types.ts b/src/main/ai-vault/session-scanner-types.ts index b1d480aa944..f4b60270544 100644 --- a/src/main/ai-vault/session-scanner-types.ts +++ b/src/main/ai-vault/session-scanner-types.ts @@ -5,7 +5,10 @@ import type { AiVaultSessionPreviewMessage } from '../../shared/ai-vault-types' import type { ExecutionHostId } from '../../shared/execution-host' -import type { TranscriptMessageSink } from './session-transcript-consumers' +import type { + TranscriptMessageSink, + TranscriptSessionIdentity +} from './session-transcript-consumers' import type { SessionSidecarObservation } from './session-sidecar-stat' export type AiVaultScanOptions = { @@ -103,6 +106,9 @@ export type ResumableSessionParseState = { consumeLineBytes?(line: Buffer): void // Lets a parser terminate an excluded transcript without draining the file. shouldStop?(): boolean + // What the fold knows about the session right now, for a consumer that has to + // commit before the read ends (see TranscriptSessionIdentity). + identity?(): TranscriptSessionIdentity | null clone(): ResumableSessionParseState // Refresh per-scan file metadata (mtime display string) without re-parsing. touchFile(file: FileWithMtime): void diff --git a/src/main/ai-vault/session-transcript-consumers.ts b/src/main/ai-vault/session-transcript-consumers.ts index 6707b298586..7aff8dcf87b 100644 --- a/src/main/ai-vault/session-transcript-consumers.ts +++ b/src/main/ai-vault/session-transcript-consumers.ts @@ -27,12 +27,34 @@ export const NO_TRANSCRIPT_MESSAGES: TranscriptMessageSink = { push: () => undefined } +/** + * What a parser has decoded about the session so far, mid-read. + * + * Provisional by construction: it is read before the file ends, so a title can + * still change and a timestamp can still move. Every field the transcript + * formats put in their opening lines, which is what a consumer that has to + * commit before the read finishes needs to name what it is holding. + */ +export type TranscriptSessionIdentity = { + sessionId: string + cwd: string | null + title: string | null + createdAt: string | null + updatedAt: string | null +} + export type TranscriptReadStart = { candidate: SessionFileCandidate /** `replace`: the whole file is being re-read; `append`: a resumed read. */ mode: 'replace' | 'append' /** Byte offset the messages of this read continue from. */ previousByteOffset: number + /** + * The session identity decoded so far, or null before the parser has an id. + * Called during the read, never here: nothing is decoded yet when a read + * begins. Absent when the read has no resumable parse state to ask. + */ + identity?: () => TranscriptSessionIdentity | null } export type TranscriptReadOutcome = { diff --git a/src/main/ai-vault/session-transcript-reader.ts b/src/main/ai-vault/session-transcript-reader.ts index 228b0c832a3..3fc0ddcc146 100644 --- a/src/main/ai-vault/session-transcript-reader.ts +++ b/src/main/ai-vault/session-transcript-reader.ts @@ -3,7 +3,10 @@ import type { AiVaultSession } from '../../shared/ai-vault-types' import { parseAgentSessionFile, parserPublishesMessages } from './session-scanner-agent-parser' import { consumeCompleteJsonlLines } from './session-scanner-jsonl-reader' import type { ResumableSessionParseState, SessionFileCandidate } from './session-scanner-types' -import type { SessionParseResumePoint } from './session-parse-cache-store' +import { + invalidateSessionParseCacheEntry, + type SessionParseResumePoint +} from './session-parse-cache-store' import { TranscriptMessageChannel } from './session-transcript-channel' const NEWLINE_BYTE = 0x0a @@ -29,6 +32,24 @@ export type ResumableTranscriptRead = { resume: SessionParseResumePoint } +/** + * Ask for the next read of `path` to be a whole-file `replace`. + * + * Why this lives here: a consumer never chooses its own mode. The reader picks + * `append` or `replace` from the resume point the session list left behind, so a + * consumer that declined an append has no way to get the span it missed — with + * an empty index and a warm parse cache, every read arrives as `append`, every + * one is declined, and nothing is ever indexed. Dropping the resume point is the + * one lever that changes the next read's mode, and only the reader's own cache + * owns it. + * + * The cost is a re-parse for the session list too. That is the honest price of a + * second consumer being behind, and it is paid once per file rather than per scan. + */ +export function requestWholeTranscriptRead(path: string): void { + invalidateSessionParseCacheEntry(path) +} + /** * Read an append-only transcript, resuming from `resume` when the file only * grew and the recorded offset still sits on a line boundary. Anything else @@ -70,7 +91,10 @@ export async function readResumableTranscript(args: { channel.beginRead({ candidate: args.candidate, mode: canResume ? 'append' : 'replace', - previousByteOffset: startOffset + previousByteOffset: startOffset, + // Read by a consumer during the read, not here: the fold has decoded + // nothing yet at this point of a whole-file read. + identity: () => state.identity?.() ?? null }) try { const readResult = await consumeCompleteJsonlLines({