diff --git a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/chatStateStorage.ts b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/chatStateStorage.ts index 977ddf935fe..c9b65645611 100644 --- a/workspaces/ballerina/ballerina-extension/src/views/ai-panel/chatStateStorage.ts +++ b/workspaces/ballerina/ballerina-extension/src/views/ai-panel/chatStateStorage.ts @@ -284,6 +284,48 @@ export class ChatStateStorage { } } + /** + * Persist a single generation incrementally (append-only). + * Appends only the changed header and/or the model messages not yet + * written — instead of rewriting the whole thread. Use for every + * per-generation mutation. Relies on the store's per-generation tracking, + * which is primed by loadThread()/saveThread() (thread creation + load). + */ + private persistGeneration(projectRootPath: string, threadId: string, generation: Generation): void { + const thread = this.storage.get(projectRootPath)?.threads.get(threadId); + if (!thread) { + return; + } + try { + this.persistenceStore.appendGeneration( + projectRootPath, + threadId, + toPersistedGeneration(generation), + thread.updatedAt + ); + } catch (err) { + console.error(`[ChatStateStorage] Failed to append generation ${generation.id}:`, err); + } + } + + /** Persist a generation removal (append-only tombstone). */ + private persistGenerationRemoval(projectRootPath: string, threadId: string, generationId: string): void { + try { + this.persistenceStore.removeGenerationRecord(projectRootPath, threadId, generationId); + } catch (err) { + console.error(`[ChatStateStorage] Failed to persist removal of generation ${generationId}:`, err); + } + } + + /** Persist a truncation from a generation onward (append-only). */ + private persistTruncation(projectRootPath: string, threadId: string, fromGenerationId: string): void { + try { + this.persistenceStore.truncateFromGeneration(projectRootPath, threadId, fromGenerationId); + } catch (err) { + console.error(`[ChatStateStorage] Failed to persist truncation from ${fromGenerationId}:`, err); + } + } + /** * Flush workspace metadata to disk. */ @@ -549,8 +591,8 @@ export class ChatStateStorage { thread.generations.push(generation); thread.updatedAt = Date.now(); - // Persist immediately - this.flushThread(projectRootPath, threadId); + // Persist immediately (append-only) + this.persistGeneration(projectRootPath, threadId, generation); console.log(`[ChatStateStorage] Added generation: ${generation.id} to thread: ${threadId}`); // Capture checkpoint for this generation asynchronously (skip for synthetic compacted generations) @@ -618,8 +660,8 @@ export class ChatStateStorage { Object.assign(generation, updates); thread.updatedAt = Date.now(); - // Persist immediately - this.flushThread(projectRootPath, threadId); + // Persist immediately (append-only: only changed header / new messages) + this.persistGeneration(projectRootPath, threadId, generation); console.log(`[ChatStateStorage] Updated generation: ${generationId}`); } @@ -635,8 +677,8 @@ export class ChatStateStorage { if (index !== -1) { thread.generations.splice(index, 1); thread.updatedAt = Date.now(); - // Persist immediately - this.flushThread(projectRootPath, threadId); + // Persist immediately (append-only tombstone) + this.persistGenerationRemoval(projectRootPath, threadId, generationId); // Also clean up checkpoint file if it exists this.persistenceStore.deleteCheckpoint(projectRootPath, threadId, generationId); console.log(`[ChatStateStorage] Removed generation: ${generationId}`); @@ -779,8 +821,8 @@ export class ChatStateStorage { Object.assign(generation.reviewState, state); thread.updatedAt = Date.now(); - // Persist immediately - this.flushThread(projectRootPath, threadId); + // Persist immediately (append-only: reviewState is part of the header) + this.persistGeneration(projectRootPath, threadId, generation); console.log(`[ChatStateStorage] Updated review state for generation: ${generationId}, status: ${generation.reviewState.status}`); } @@ -793,21 +835,23 @@ export class ChatStateStorage { */ acceptAllReviews(projectRootPath: string, threadId: string): void { const thread = this.getOrCreateThread(projectRootPath, threadId); - let count = 0; + const changed: Generation[] = []; for (const generation of thread.generations) { if (generation.reviewState.status === 'under_review') { generation.reviewState.status = 'accepted'; generation.reviewState.affectedPackagePaths = []; - count++; + changed.push(generation); } } - if (count > 0) { + if (changed.length > 0) { thread.updatedAt = Date.now(); - this.flushThread(projectRootPath, threadId); + for (const generation of changed) { + this.persistGeneration(projectRootPath, threadId, generation); + } } - console.log(`[ChatStateStorage] Accepted ${count} review(s) in thread: ${threadId}`); + console.log(`[ChatStateStorage] Accepted ${changed.length} review(s) in thread: ${threadId}`); } /** @@ -819,22 +863,24 @@ export class ChatStateStorage { */ declineAllReviews(projectRootPath: string, threadId: string): void { const thread = this.getOrCreateThread(projectRootPath, threadId); - let count = 0; + const changed: Generation[] = []; for (const generation of thread.generations) { if (generation.reviewState.status === 'under_review') { generation.reviewState.status = 'error'; generation.reviewState.errorMessage = 'Declined by user'; generation.reviewState.affectedPackagePaths = []; - count++; + changed.push(generation); } } - if (count > 0) { + if (changed.length > 0) { thread.updatedAt = Date.now(); - this.flushThread(projectRootPath, threadId); + for (const generation of changed) { + this.persistGeneration(projectRootPath, threadId, generation); + } } - console.log(`[ChatStateStorage] Declined ${count} review(s) in thread: ${threadId}`); + console.log(`[ChatStateStorage] Declined ${changed.length} review(s) in thread: ${threadId}`); } // ============================================ @@ -928,8 +974,10 @@ export class ChatStateStorage { console.error(`[ChatStateStorage] Failed to persist checkpoint ${checkpoint.id}:`, err); } - // Enforce checkpoint limit (evicts oldest) and flush thread once at the end + // Enforce checkpoint limit (evicts oldest, persisting each evicted generation) await this.enforceCheckpointLimit(projectRootPath, threadId); + // Persist this generation's header (hasCheckpoint flag now reflects the checkpoint) + this.persistGeneration(projectRootPath, threadId, generation); } /** @@ -943,8 +991,7 @@ export class ChatStateStorage { const config = getCheckpointConfig(); if (!config.enabled) { - // Still flush the thread to persist the newly added checkpoint - this.flushThread(projectRootPath, threadId); + // Nothing to evict; caller persists the newly-checkpointed generation. return; } @@ -955,15 +1002,17 @@ export class ChatStateStorage { .map((gen, index) => ({ generation: gen, index })) .filter(item => item.generation.checkpoint !== undefined); - // If we're within the limit, just flush and return + // If we're within the limit, nothing to evict; caller persists the added generation. if (generationsWithCheckpoints.length <= config.maxCount) { - this.flushThread(projectRootPath, threadId); return; } // Calculate how many checkpoints to remove const checkpointsToRemove = generationsWithCheckpoints.length - config.maxCount; + // Bump updatedAt before persisting so the appended records carry the new timestamp. + thread.updatedAt = Date.now(); + // Remove checkpoints from oldest generations (keep the most recent maxCount) for (let i = 0; i < checkpointsToRemove; i++) { const { generation } = generationsWithCheckpoints[i]; @@ -974,11 +1023,10 @@ export class ChatStateStorage { // Delete the persisted checkpoint file this.persistenceStore.deleteCheckpoint(projectRootPath, threadId, generation.id); - } - thread.updatedAt = Date.now(); - // Persist thread with updated checkpoint flags - this.flushThread(projectRootPath, threadId); + // Persist the evicted generation's header (hasCheckpoint flips to false) + this.persistGeneration(projectRootPath, threadId, generation); + } } /** @@ -1010,6 +1058,7 @@ export class ChatStateStorage { } // Clean up checkpoint files for removed generations + const fromGenerationId = thread.generations[checkpointGenerationIndex].id; for (let i = checkpointGenerationIndex; i < thread.generations.length; i++) { this.persistenceStore.deleteCheckpoint(projectRootPath, threadId, thread.generations[i].id); } @@ -1020,8 +1069,8 @@ export class ChatStateStorage { thread.generations = thread.generations.slice(0, checkpointGenerationIndex); thread.updatedAt = Date.now(); - // Persist immediately - this.flushThread(projectRootPath, threadId); + // Persist truncation (append-only): removes fromGenerationId and everything after it + this.persistTruncation(projectRootPath, threadId, fromGenerationId); return true; } diff --git a/workspaces/common-libs/copilot-utilities/src/chat-persistence/__tests__/persistence-store.test.ts b/workspaces/common-libs/copilot-utilities/src/chat-persistence/__tests__/persistence-store.test.ts index 320543857bf..e2b2bfdc1ac 100644 --- a/workspaces/common-libs/copilot-utilities/src/chat-persistence/__tests__/persistence-store.test.ts +++ b/workspaces/common-libs/copilot-utilities/src/chat-persistence/__tests__/persistence-store.test.ts @@ -529,16 +529,23 @@ describe('CopilotPersistenceStore', () => { assert.equal(tmpFiles.length, 0, `Found leftover tmp files: ${tmpFiles.join(', ')}`); }); - it('should produce valid JSON on disk', () => { + it('should produce a valid JSONL log on disk (one JSON object per line)', () => { store.saveThread(WORKSPACE_PATH, 'default', makeThread()); - const threadFile = path.join( - store.getWorkspaceDir(WORKSPACE_PATH), 'threads', 'default', 'thread.json' - ); - const raw = fs.readFileSync(threadFile, 'utf8'); - // Should not throw - const parsed = JSON.parse(raw); - assert.equal(parsed.id, 'default'); + const threadDir = path.join(store.getWorkspaceDir(WORKSPACE_PATH), 'threads', 'default'); + // The append-only log replaces the legacy whole-file snapshot. + assert.ok(fs.existsSync(path.join(threadDir, 'thread.jsonl'))); + assert.equal(fs.existsSync(path.join(threadDir, 'thread.json')), false); + + const raw = fs.readFileSync(path.join(threadDir, 'thread.jsonl'), 'utf8'); + const lines = raw.split('\n').filter(l => l.trim().length > 0); + // Each line must parse as JSON on its own. + const records = lines.map(l => JSON.parse(l)); + // Compact form: head, meta, then one gen per generation. + assert.equal(records[0].t, 'head'); + assert.equal(records[0].id, 'default'); + assert.equal(records[1].t, 'meta'); + assert.ok(records.some(r => r.t === 'gen')); }); }); diff --git a/workspaces/common-libs/copilot-utilities/src/chat-persistence/__tests__/thread-log.test.ts b/workspaces/common-libs/copilot-utilities/src/chat-persistence/__tests__/thread-log.test.ts new file mode 100644 index 00000000000..93b87e94c9b --- /dev/null +++ b/workspaces/common-libs/copilot-utilities/src/chat-persistence/__tests__/thread-log.test.ts @@ -0,0 +1,630 @@ +/** + * Copyright (c) 2026, WSO2 LLC. (https://www.wso2.com) All Rights Reserved. + * + * WSO2 LLC. licenses this file to you under the Apache License, + * Version 2.0 (the "License"); you may not use this file except + * in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { describe, it, beforeEach, afterEach } from 'node:test'; +import * as assert from 'node:assert/strict'; +import * as fs from 'fs'; +import * as path from 'path'; +import * as os from 'os'; + +import { CopilotPersistenceStore } from '../persistence-store'; +import { appendLineSync, readJsonlSync } from '../file-utils'; +import { CURRENT_THREAD_SCHEMA_VERSION } from '../schema-migration'; +import { PersistedGeneration, PersistedThread, ThreadLogRecord } from '../types'; + +// ============================================ +// Test Helpers +// ============================================ + +let tmpDir: string; +let store: CopilotPersistenceStore; + +const WORKSPACE_PATH = '/Users/test/projects/my-ballerina-project'; + +function createTmpDir(): string { + return fs.mkdtempSync(path.join(os.tmpdir(), 'copilot-threadlog-test-')); +} + +function cleanupTmpDir(dir: string): void { + if (fs.existsSync(dir)) { + fs.rmSync(dir, { recursive: true, force: true }); + } +} + +let genCounter = 0; +function makeGeneration(overrides: Partial = {}): PersistedGeneration { + genCounter++; + return { + id: `gen-${genCounter}`, + userPrompt: 'Create a REST API', + modelMessages: [ + { role: 'user', content: 'Create a REST API' }, + { + role: 'assistant', + content: [ + { type: 'text', text: 'Working on it.' }, + { type: 'tool-call', toolCallId: 'tc1', toolName: 'writeFile', args: { path: 'main.bal' } }, + ], + }, + ], + uiResponse: 'Done.', + timestamp: 1000 + genCounter, + currentTaskIndex: -1, + reviewState: { status: 'accepted', modifiedFiles: ['main.bal'] }, + metadata: { isPlanMode: false, generationType: 'agent' }, + hasCheckpoint: false, + ...overrides, + }; +} + +function makeThread(overrides: Partial = {}): PersistedThread { + return { + schemaVersion: CURRENT_THREAD_SCHEMA_VERSION, + id: 'default', + name: 'Default Thread', + createdAt: 1000, + updatedAt: 1000, + generations: [], + ...overrides, + }; +} + +function threadDir(threadId: string): string { + return path.join(store.getWorkspaceDir(WORKSPACE_PATH), 'threads', threadId); +} + +function logPath(threadId: string): string { + return path.join(threadDir(threadId), 'thread.jsonl'); +} + +function legacyJsonPath(threadId: string): string { + return path.join(threadDir(threadId), 'thread.json'); +} + +function countLines(threadId: string): number { + const raw = fs.readFileSync(logPath(threadId), 'utf8'); + return raw.split('\n').filter(l => l.trim().length > 0).length; +} + +// ============================================ +// Tests +// ============================================ + +describe('Thread append-only log (JSONL)', () => { + beforeEach(() => { + tmpDir = createTmpDir(); + store = new CopilotPersistenceStore({ baseDir: tmpDir }); + genCounter = 0; + }); + + afterEach(() => { + cleanupTmpDir(tmpDir); + }); + + // --- Append → replay round-trip --- + + describe('append and replay', () => { + it('should reconstruct generations appended after a save-initialized thread', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + const g1 = makeGeneration({ id: 'g1', userPrompt: 'first' }); + const g2 = makeGeneration({ id: 'g2', userPrompt: 'second' }); + store.appendGeneration(WORKSPACE_PATH, 'default', g1); + store.appendGeneration(WORKSPACE_PATH, 'default', g2); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.equal(loaded.generations.length, 2); + assert.deepEqual(loaded.generations.map(g => g.id), ['g1', 'g2']); + assert.equal(loaded.generations[0].userPrompt, 'first'); + assert.deepEqual(loaded.generations[1].modelMessages, g2.modelMessages); + }); + + it('should treat the latest upsert for a generation id as authoritative (in place)', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1', uiResponse: 'v1' })); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g2', uiResponse: 'other' })); + // Update g1 in place (simulates AgentExecutor persisting after each step). + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1', uiResponse: 'v2-final' })); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + // Order preserved (g1 keeps its original position), content is latest. + assert.deepEqual(loaded.generations.map(g => g.id), ['g1', 'g2']); + assert.equal(loaded.generations[0].uiResponse, 'v2-final'); + }); + + it('should apply a tombstone (del) removing a generation', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1' })); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g2' })); + store.removeGenerationRecord(WORKSPACE_PATH, 'default', 'g1'); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.deepEqual(loaded.generations.map(g => g.id), ['g2']); + }); + + it('should ignore a del for an unknown generation id', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1' })); + store.removeGenerationRecord(WORKSPACE_PATH, 'default', 'does-not-exist'); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.deepEqual(loaded.generations.map(g => g.id), ['g1']); + }); + + it('should truncate a generation and everything appended after it', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1' })); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g2' })); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g3' })); + // Restore-to-checkpoint at g2 drops g2 and g3. + store.truncateFromGeneration(WORKSPACE_PATH, 'default', 'g2'); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.deepEqual(loaded.generations.map(g => g.id), ['g1']); + }); + + it('should ignore a trunc for an unknown generation id', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1' })); + store.truncateFromGeneration(WORKSPACE_PATH, 'default', 'unknown'); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.deepEqual(loaded.generations.map(g => g.id), ['g1']); + }); + + it('should re-add a generation after it was tombstoned (position resets to append order)', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1' })); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g2' })); + store.removeGenerationRecord(WORKSPACE_PATH, 'default', 'g1'); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1', uiResponse: 'again' })); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.deepEqual(loaded.generations.map(g => g.id), ['g2', 'g1']); + assert.equal(loaded.generations[1].uiResponse, 'again'); + }); + + it('should apply thread metadata updates (name, sessionId)', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + store.updateThreadMeta(WORKSPACE_PATH, 'default', { name: 'Renamed', sessionId: 'sess-42' }); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.equal(loaded.name, 'Renamed'); + assert.equal(loaded.sessionId, 'sess-42'); + }); + + it('should derive updatedAt from the latest record', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread({ updatedAt: 1000 })); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1' }), 5000); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.equal(loaded.updatedAt, 5000); + }); + + it('should round-trip an empty thread (head + meta, no generations)', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread({ generations: [] })); + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.deepEqual(loaded.generations, []); + assert.equal(loaded.name, 'Default Thread'); + }); + + function records(threadId: string): any[] { + return fs.readFileSync(logPath(threadId), 'utf8') + .split('\n').filter(l => l.trim()).map(l => JSON.parse(l)); + } + + it('should write each model message exactly once across a multi-step turn (no O(S^2) duplication)', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread({ generations: [] })); + + // Simulate the agent hot path: same generation, unchanged header, the + // modelMessages array grows by one message per step. + const S = 30; + const msgs: unknown[] = []; + for (let step = 1; step <= S; step++) { + msgs.push({ role: 'assistant', content: `step-${step}` }); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ + id: 'g1', + timestamp: 5, // fixed so the header never changes + uiResponse: 'fixed', + modelMessages: [...msgs], + })); + } + + const recs = records('default'); + const msgRecords = recs.filter(r => r.t === 'msg').length; + const genRecords = recs.filter(r => r.t === 'gen').length; + // Each message persisted once (S), NOT 1+2+...+S. This is the fix. + assert.equal(msgRecords, S, `expected ${S} msg records, got ${msgRecords}`); + // The header is written once, since it never changed. + assert.equal(genRecords, 1, `expected a single gen header, got ${genRecords}`); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.equal(loaded.generations[0].modelMessages.length, S); + assert.deepEqual(loaded.generations[0].modelMessages[S - 1], { role: 'assistant', content: `step-${S}` }); + }); + + it('should append a gen header only when a non-message field changes', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread({ generations: [] })); + const base = { id: 'g1', timestamp: 5, modelMessages: [{ role: 'user', content: 'x' }] }; + // Same header three times -> one header record. + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ ...base, uiResponse: 'a' })); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ ...base, uiResponse: 'a' })); + // uiResponse change -> a second header record. + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ ...base, uiResponse: 'b' })); + + const recs = records('default'); + assert.equal(recs.filter(r => r.t === 'gen').length, 2); + assert.equal(recs.filter(r => r.t === 'msg').length, 1); // one message, written once + assert.equal(store.loadThread(WORKSPACE_PATH, 'default')!.generations[0].uiResponse, 'b'); + }); + + it('should emit a msgs reset when modelMessages are rewritten rather than extended', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread({ generations: [] })); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ + id: 'g1', timestamp: 5, + modelMessages: [{ role: 'user', content: 'a' }, { role: 'assistant', content: 'b' }], + })); + // Same length but the last message changed (e.g. server-side rewrite). + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ + id: 'g1', timestamp: 5, + modelMessages: [{ role: 'user', content: 'a' }, { role: 'assistant', content: 'REWRITTEN' }], + })); + + const recs = records('default'); + assert.ok(recs.some(r => r.t === 'msgs'), 'expected a msgs reset record'); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.deepEqual(loaded.generations[0].modelMessages, [ + { role: 'user', content: 'a' }, { role: 'assistant', content: 'REWRITTEN' }, + ]); + }); + }); + + // --- Crash / corruption resilience --- + + describe('resilience', () => { + it('should recover all complete records when the trailing line is torn', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1' })); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g2' })); + // Simulate a crash mid-append: a partial JSON line with no newline. + fs.appendFileSync(logPath('default'), '{"t":"gen","updatedAt":9999,"gen":{"id":"g3","userPr'); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.deepEqual(loaded.generations.map(g => g.id), ['g1', 'g2']); + }); + + it('should skip a corrupt line in the middle and keep the rest', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1' })); + appendLineSync(logPath('default'), 'this is not json at all'); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g2' })); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.deepEqual(loaded.generations.map(g => g.id), ['g1', 'g2']); + }); + + it('should return null when the log exists but contains no recognizable records', () => { + fs.mkdirSync(threadDir('garbage'), { recursive: true }); + fs.writeFileSync(logPath('garbage'), 'not json\n{also not\n'); + assert.equal(store.loadThread(WORKSPACE_PATH, 'garbage'), null); + }); + + it('should return null for an empty log file', () => { + fs.mkdirSync(threadDir('empty-file'), { recursive: true }); + fs.writeFileSync(logPath('empty-file'), ''); + assert.equal(store.loadThread(WORKSPACE_PATH, 'empty-file'), null); + }); + + it('should reconstruct without data loss when the head record is missing', () => { + // Simulate gen header + message records that reached disk before an + // interrupted init wrote the thread `head` record. + const { modelMessages: m1, ...h1 } = makeGeneration({ id: 'g1' }); + const { modelMessages: _m2, ...h2 } = makeGeneration({ id: 'g2' }); + fs.mkdirSync(threadDir('headless'), { recursive: true }); + appendLineSync(logPath('headless'), JSON.stringify({ t: 'gen', updatedAt: 2000, gen: h1 } as ThreadLogRecord)); + appendLineSync(logPath('headless'), JSON.stringify({ t: 'msg', updatedAt: 2000, genId: 'g1', message: m1[0] } as ThreadLogRecord)); + appendLineSync(logPath('headless'), JSON.stringify({ t: 'gen', updatedAt: 2500, gen: h2 } as ThreadLogRecord)); + + const loaded = store.loadThread(WORKSPACE_PATH, 'headless'); + assert.ok(loaded); + assert.deepEqual(loaded.generations.map(g => g.id), ['g1', 'g2']); + assert.deepEqual(loaded.generations[0].modelMessages, [m1[0]]); // message reconstructed + assert.deepEqual(loaded.generations[1].modelMessages, []); + assert.equal(loaded.id, 'headless'); + assert.equal(loaded.createdAt, 2000); // earliest record timestamp + }); + + it('should return null for a non-existent thread (no log, no legacy json)', () => { + assert.equal(store.loadThread(WORKSPACE_PATH, 'nope'), null); + }); + }); + + // --- Legacy migration --- + + describe('legacy thread.json migration', () => { + it('should migrate a v1 thread.json to thread.jsonl and delete the legacy file', () => { + // Hand-write a legacy v1 whole-file snapshot. + const legacy = { + schemaVersion: 1, + id: 'default', + name: 'Legacy Thread', + createdAt: 111, + updatedAt: 222, + generations: [makeGeneration({ id: 'g1', userPrompt: 'legacy prompt' })], + }; + fs.mkdirSync(threadDir('default'), { recursive: true }); + fs.writeFileSync(legacyJsonPath('default'), JSON.stringify(legacy), 'utf8'); + assert.equal(fs.existsSync(logPath('default')), false); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.equal(loaded.name, 'Legacy Thread'); + assert.equal(loaded.schemaVersion, CURRENT_THREAD_SCHEMA_VERSION); + assert.equal(loaded.generations[0].userPrompt, 'legacy prompt'); + + // The log now exists and the legacy file is gone. + assert.ok(fs.existsSync(logPath('default'))); + assert.equal(fs.existsSync(legacyJsonPath('default')), false); + + // Second load reads from the log and is identical. + const reloaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.deepEqual(reloaded, loaded); + }); + + it('should prefer thread.jsonl and remove a stale thread.json if both exist', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread({ name: 'From Log' })); + // A stale legacy file lingers (e.g. an interrupted prior migration). + fs.writeFileSync(legacyJsonPath('default'), JSON.stringify({ schemaVersion: 1, id: 'default', name: 'Stale', createdAt: 1, updatedAt: 1, generations: [] }), 'utf8'); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.equal(loaded.name, 'From Log'); + assert.equal(fs.existsSync(legacyJsonPath('default')), false); + }); + + it('should fall back to a legacy thread.json when the jsonl is empty/unrecoverable', () => { + // Defensive: an externally-truncated (0-byte) log must not shadow a + // still-present legacy snapshot and report the thread as gone. + fs.mkdirSync(threadDir('both'), { recursive: true }); + fs.writeFileSync(logPath('both'), ''); + const legacy = { + schemaVersion: 1, id: 'both', name: 'Recovered', createdAt: 1, updatedAt: 2, + generations: [makeGeneration({ id: 'g1' })], + }; + fs.writeFileSync(legacyJsonPath('both'), JSON.stringify(legacy), 'utf8'); + + const loaded = store.loadThread(WORKSPACE_PATH, 'both'); + assert.ok(loaded); + assert.equal(loaded.name, 'Recovered'); + assert.deepEqual(loaded.generations.map(g => g.id), ['g1']); + // Migration rewrote the log and removed the legacy file. + assert.equal(fs.existsSync(legacyJsonPath('both')), false); + }); + + it('should return null for a corrupt legacy thread.json', () => { + fs.mkdirSync(threadDir('corrupt'), { recursive: true }); + fs.writeFileSync(legacyJsonPath('corrupt'), '{ invalid json', 'utf8'); + assert.equal(store.loadThread(WORKSPACE_PATH, 'corrupt'), null); + }); + + it('should return null when a legacy thread.json has an unmigratable schema version', () => { + // Parses as JSON but has no migration path (version 0 -> no 0->1 + // migration exists) so migrateThread throws and loadThread reports + // the thread as unreadable rather than crashing. + const unmigratable = { schemaVersion: 0, id: 'default', name: 'X', createdAt: 1, updatedAt: 1, generations: [] }; + fs.mkdirSync(threadDir('ancient'), { recursive: true }); + fs.writeFileSync(legacyJsonPath('ancient'), JSON.stringify(unmigratable), 'utf8'); + assert.equal(store.loadThread(WORKSPACE_PATH, 'ancient'), null); + }); + }); + + // --- Compaction --- + + describe('compaction', () => { + it('compactThread should collapse superseded records while preserving the replayed thread', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + // Update the same generation several times (mid-stream persistence). + // Kept under the on-load compaction threshold so loadThread does not + // pre-compact and we can isolate compactThread's effect. + for (let i = 0; i < 10; i++) { + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1', uiResponse: `step-${i}` })); + } + const before = store.loadThread(WORKSPACE_PATH, 'default'); + const linesBefore = countLines('default'); + assert.ok(linesBefore > 3); + + store.compactThread(WORKSPACE_PATH, 'default'); + + const after = store.loadThread(WORKSPACE_PATH, 'default'); + const linesAfter = countLines('default'); + + // Replayed thread is unchanged by compaction. + assert.deepEqual(after, before); + // The log shrank to head + meta + 1 gen header + its messages. + const msgCount = after!.generations[0].modelMessages.length; + assert.equal(linesAfter, 3 + msgCount); + assert.ok(linesBefore > linesAfter); + assert.equal(after!.generations[0].uiResponse, 'step-9'); + }); + + it('should preserve sessionId through a compacted save/replay', () => { + // Exercises the compacted-record builder's sessionId branch. + store.saveThread(WORKSPACE_PATH, 'default', makeThread({ sessionId: 'sess-xyz', generations: [makeGeneration({ id: 'g1' })] })); + store.compactThread(WORKSPACE_PATH, 'default'); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.equal(loaded.sessionId, 'sess-xyz'); + assert.deepEqual(loaded.generations.map(g => g.id), ['g1']); + }); + + it('compactThread should be a no-op when the thread has no log', () => { + // Should not throw and should not create a file. + store.compactThread(WORKSPACE_PATH, 'ghost'); + assert.equal(fs.existsSync(logPath('ghost')), false); + }); + + it('should auto-compact once the size-aware append threshold is reached', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + // One live generation -> threshold is the floor (64). Append well past + // it with in-place updates so compaction must fire without any explicit + // compactThread call. + for (let i = 0; i < 120; i++) { + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1', uiResponse: `v${i}` })); + } + const lines = countLines('default'); + assert.ok(lines < 70, `expected auto-compaction to bound growth, got ${lines} lines`); + + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.equal(loaded.generations.length, 1); + assert.equal(loaded.generations[0].uiResponse, 'v119'); + }); + + it('should scale the compaction threshold with live generation count', () => { + // 30 live generations -> threshold = max(64, 3*30) = 90. + const gens = Array.from({ length: 30 }, (_, i) => makeGeneration({ id: `g${i}` })); + store.saveThread(WORKSPACE_PATH, 'big', makeThread({ id: 'big', generations: gens })); + + // 80 in-place updates of one generation stay UNDER the size-aware + // threshold (90), so NO auto-compaction fires. A fixed 64-append + // interval would have compacted here — this asserts the amortized, + // size-proportional behaviour that keeps per-append cost ~O(1). + for (let i = 0; i < 80; i++) { + store.appendGeneration(WORKSPACE_PATH, 'big', makeGeneration({ id: 'g0', uiResponse: `v${i}` })); + } + const lines = countLines('big'); + assert.ok(lines > 100, `expected no compaction below the size-aware threshold, got ${lines} lines`); + + const loaded = store.loadThread(WORKSPACE_PATH, 'big'); + assert.ok(loaded); + assert.equal(loaded.generations.length, 30); + }); + + it('should not fail load when on-load compaction cannot write (read-only dir)', () => { + store.saveThread(WORKSPACE_PATH, 'ro', makeThread({ id: 'ro' })); + for (let i = 0; i < 50; i++) { + store.appendGeneration(WORKSPACE_PATH, 'ro', makeGeneration({ id: 'g1', uiResponse: `v${i}` })); + } + // 52 lines: past the on-load compaction trigger. Make the thread dir + // read-only so the compaction's atomic write (a new tmp file) fails. + const dir = threadDir('ro'); + fs.chmodSync(dir, 0o555); + try { + // Must NOT throw and must return the already-replayed thread even + // though compaction could not persist. (If the test runs as root, + // the write succeeds; either way load returns the thread.) + const loaded = store.loadThread(WORKSPACE_PATH, 'ro'); + assert.ok(loaded, 'load must still return the replayed thread'); + assert.equal(loaded.generations[0].uiResponse, 'v49'); + } finally { + fs.chmodSync(dir, 0o755); + } + }); + + it('should compact on load when the log has grown well beyond its compact size', () => { + store.saveThread(WORKSPACE_PATH, 'default', makeThread()); + // Stay under the auto-compaction interval so the file stays large. + for (let i = 0; i < 50; i++) { + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'g1', uiResponse: `v${i}` })); + } + assert.ok(countLines('default') > 40); + + // Loading triggers opportunistic compaction. + const loaded = store.loadThread(WORKSPACE_PATH, 'default'); + assert.ok(loaded); + assert.equal(loaded.generations[0].uiResponse, 'v49'); + // head + meta + 1 gen header + its messages. + assert.equal(countLines('default'), 3 + loaded.generations[0].modelMessages.length); + }); + }); + + // --- Isolation --- + + describe('log isolation', () => { + it('should keep append logs separate across threads and workspaces', () => { + const wsB = '/Users/test/other-project'; + store.saveThread(WORKSPACE_PATH, 'default', makeThread({ name: 'A' })); + store.saveThread(wsB, 'default', makeThread({ name: 'B' })); + store.appendGeneration(WORKSPACE_PATH, 'default', makeGeneration({ id: 'a1' })); + store.appendGeneration(wsB, 'default', makeGeneration({ id: 'b1' })); + + const a = store.loadThread(WORKSPACE_PATH, 'default'); + const b = store.loadThread(wsB, 'default'); + assert.equal(a!.name, 'A'); + assert.equal(b!.name, 'B'); + assert.deepEqual(a!.generations.map(g => g.id), ['a1']); + assert.deepEqual(b!.generations.map(g => g.id), ['b1']); + }); + }); +}); + +// ============================================ +// File utilities (JSONL primitives) +// ============================================ + +describe('file-utils JSONL primitives', () => { + beforeEach(() => { tmpDir = createTmpDir(); }); + afterEach(() => { cleanupTmpDir(tmpDir); }); + + it('appendLineSync should create parent directories and append newline-terminated lines', () => { + const p = path.join(tmpDir, 'nested', 'deep', 'log.jsonl'); + appendLineSync(p, JSON.stringify({ a: 1 })); + appendLineSync(p, JSON.stringify({ a: 2 })); + const raw = fs.readFileSync(p, 'utf8'); + assert.equal(raw, '{"a":1}\n{"a":2}\n'); + }); + + it('readJsonlSync should return null for a missing file', () => { + assert.equal(readJsonlSync(path.join(tmpDir, 'missing.jsonl')), null); + }); + + it('readJsonlSync should parse valid lines and skip blank/corrupt ones', () => { + const p = path.join(tmpDir, 'mixed.jsonl'); + fs.writeFileSync(p, '{"a":1}\n\n \nnot-json\n{"a":2}\n'); + const records = readJsonlSync<{ a: number }>(p); + assert.deepEqual(records, [{ a: 1 }, { a: 2 }]); + }); + + it('readJsonlSync should return an empty array for an empty file', () => { + const p = path.join(tmpDir, 'empty.jsonl'); + fs.writeFileSync(p, ''); + assert.deepEqual(readJsonlSync(p), []); + }); + + it('readJsonlSync should return null on an I/O error (path is a directory)', () => { + const p = path.join(tmpDir, 'a-directory'); + fs.mkdirSync(p); + // existsSync is true but readFileSync throws EISDIR -> outer catch. + assert.equal(readJsonlSync(p), null); + }); +}); diff --git a/workspaces/common-libs/copilot-utilities/src/chat-persistence/file-utils.ts b/workspaces/common-libs/copilot-utilities/src/chat-persistence/file-utils.ts index 23bad50f832..fb4742dc324 100644 --- a/workspaces/common-libs/copilot-utilities/src/chat-persistence/file-utils.ts +++ b/workspaces/common-libs/copilot-utilities/src/chat-persistence/file-utils.ts @@ -92,6 +92,53 @@ export function writeJsonSync(filePath: string, data: unknown): void { atomicWriteSync(filePath, JSON.stringify(data)); } +/** + * Append a single line to a file (creating it and any parent directories if + * needed). The newline terminator is added automatically. + * + * This is the core primitive behind the append-only (JSONL) thread log: adding + * a record costs O(size of the record) regardless of how large the file already + * is, unlike a full-file rewrite which costs O(size of the whole file). + */ +export function appendLineSync(filePath: string, line: string): void { + ensureDirSync(path.dirname(filePath)); + fs.appendFileSync(filePath, line + '\n', 'utf8'); +} + +/** + * Read a newline-delimited JSON (JSONL) file and parse each line. + * + * Returns `null` if the file does not exist (so callers can distinguish + * "no log" from "empty log"). Blank lines are skipped. Individual lines that + * fail to parse are skipped rather than aborting the whole read — this makes + * replay resilient to a torn trailing line left behind by a crash mid-append + * (`fs.appendFileSync` is not atomic across a crash) and to any single corrupt + * record, without losing the rest of the history. + */ +export function readJsonlSync(filePath: string): T[] | null { + try { + if (!fs.existsSync(filePath)) { + return null; + } + const raw = fs.readFileSync(filePath, 'utf8'); + const records: T[] = []; + for (const line of raw.split('\n')) { + const trimmed = line.trim(); + if (trimmed.length === 0) { + continue; + } + try { + records.push(JSON.parse(trimmed) as T); + } catch { + // Skip a torn/corrupt line; keep replaying the rest. + } + } + return records; + } catch { + return null; + } +} + /** * Write data as gzip-compressed JSON atomically. * Used for large checkpoint snapshots. diff --git a/workspaces/common-libs/copilot-utilities/src/chat-persistence/index.ts b/workspaces/common-libs/copilot-utilities/src/chat-persistence/index.ts index f5c0b2af460..0c808c0cc4a 100644 --- a/workspaces/common-libs/copilot-utilities/src/chat-persistence/index.ts +++ b/workspaces/common-libs/copilot-utilities/src/chat-persistence/index.ts @@ -31,6 +31,8 @@ export type { PersistedLinePosition, WorkspaceMetadata, ThreadSummary, + ThreadLogRecord, + PersistedGenerationHeader, PersistenceStoreConfig, } from './types'; diff --git a/workspaces/common-libs/copilot-utilities/src/chat-persistence/persistence-store.ts b/workspaces/common-libs/copilot-utilities/src/chat-persistence/persistence-store.ts index 63c6cd580cd..10cf9a4798e 100644 --- a/workspaces/common-libs/copilot-utilities/src/chat-persistence/persistence-store.ts +++ b/workspaces/common-libs/copilot-utilities/src/chat-persistence/persistence-store.ts @@ -24,6 +24,9 @@ import { PersistenceStoreConfig, PersistedThread, PersistedCheckpoint, + PersistedGeneration, + PersistedGenerationHeader, + ThreadLogRecord, WorkspaceMetadata, ThreadSummary, } from './types'; @@ -36,6 +39,9 @@ import { removeDirSync, listSubdirectoriesSync, listFilesBySuffixSync, + atomicWriteSync, + appendLineSync, + readJsonlSync, } from './file-utils'; import { computeWorkspaceHash } from './workspace-hash'; import { @@ -53,8 +59,31 @@ const THREADS_DIR = 'threads'; const CHECKPOINTS_DIR = 'checkpoints'; const WORKSPACE_META_FILE = 'workspace.meta.json'; const THREAD_FILE = 'thread.json'; +const THREAD_LOG_FILE = 'thread.jsonl'; const CHECKPOINT_SUFFIX = '.snapshot.gz'; +/** + * Compaction tuning for the append-only thread log. + * + * Live-session compaction is **size-proportional** so that per-append cost stays + * genuinely amortized O(1) (not just smaller): a thread is compacted once it has + * received more than `max(COMPACT_MIN_APPENDS, COMPACT_GROWTH_FACTOR × liveGens)` + * appends since its last compaction. A compaction costs O(liveGens) and happens + * at most once per that many appends, so the amortized cost per append does not + * grow with the conversation length — this is what keeps the format from + * regressing to the O(n²)-per-session behaviour of full-file rewrites. Using a + * fixed interval instead would make each compaction O(liveGens) every N appends, + * i.e. still superlinear for large threads. + * + * `COMPACT_LOAD_FACTOR` / `COMPACT_LOAD_MIN_LINES` bound startup replay cost: + * on load, if the log holds more than `FACTOR ×` the records a compact log + * would (and more than the minimum), it is compacted opportunistically. + */ +const COMPACT_MIN_APPENDS = 64; +const COMPACT_GROWTH_FACTOR = 3; +const COMPACT_LOAD_FACTOR = 3; +const COMPACT_LOAD_MIN_LINES = 16; + /** * File-based persistence store for copilot chat threads and checkpoints. * @@ -62,10 +91,52 @@ const CHECKPOINT_SUFFIX = '.snapshot.gz'; * All thread reads/writes are synchronous (files are typically 1-2MB). * Checkpoint writes offer an async variant for large snapshots. */ +/** + * Build the compact record sequence for a thread: a `head`, a `meta`, then for + * each generation in order a `gen` header followed by one `msg` per model + * message. Replaying this sequence reproduces `thread`, and every message + * appears exactly once (no duplication to compact away later). + */ +function buildCompactedRecords(thread: PersistedThread): ThreadLogRecord[] { + const records: ThreadLogRecord[] = [ + { t: 'head', v: thread.schemaVersion, id: thread.id, createdAt: thread.createdAt }, + { + t: 'meta', + updatedAt: thread.updatedAt, + name: thread.name, + ...(thread.sessionId !== undefined ? { sessionId: thread.sessionId } : {}), + }, + ]; + for (const gen of thread.generations) { + const { modelMessages, ...header } = gen; + records.push({ t: 'gen', updatedAt: thread.updatedAt, gen: header }); + for (const message of modelMessages ?? []) { + records.push({ t: 'msg', updatedAt: thread.updatedAt, genId: gen.id, message }); + } + } + return records; +} + export class CopilotPersistenceStore { private readonly baseDir: string; private readonly workspaceIdResolver: (workspacePath: string) => string; + // Per-thread count of appends since the last compaction (runtime-only) and + // the size-proportional threshold at which the next compaction fires. Both + // are keyed by `${workspaceHash}/${threadId}` and live for the process + // lifetime. Drives amortized (size-aware) compaction. + private readonly appendsSinceCompaction: Map = new Map(); + private readonly compactionThreshold: Map = new Map(); + + // Per-generation append tracking (runtime-only), keyed by + // `${workspaceHash}/${threadId}#${generationId}`. Lets appendGeneration + // persist only what changed: the header when it differs, and only the model + // messages not yet written. `count` = messages already persisted; + // `lastJson` = JSON of the last persisted message (used to detect whether a + // new modelMessages array extends what we have or was rewritten). + private readonly lastHeaderJson: Map = new Map(); + private readonly msgState: Map = new Map(); + constructor(config: PersistenceStoreConfig = {}) { this.baseDir = config.baseDir ?? DEFAULT_BASE_DIR; this.workspaceIdResolver = config.workspaceIdResolver ?? ((p) => path.resolve(p)); @@ -86,11 +157,26 @@ export class CopilotPersistenceStore { return path.join(this.getWorkspaceDir(workspacePath), THREADS_DIR, threadId); } - /** Resolve the path to the thread JSON file. */ + /** Resolve the path to the legacy whole-file thread JSON (pre-v2). */ private getThreadFilePath(workspacePath: string, threadId: string): string { return path.join(this.getThreadDir(workspacePath, threadId), THREAD_FILE); } + /** Resolve the path to the append-only thread log (`thread.jsonl`). */ + private getThreadLogFilePath(workspacePath: string, threadId: string): string { + return path.join(this.getThreadDir(workspacePath, threadId), THREAD_LOG_FILE); + } + + /** Key used to track per-thread append counts for amortized compaction. */ + private compactionKey(workspacePath: string, threadId: string): string { + return `${computeWorkspaceHash(this.workspaceIdResolver(workspacePath))}/${threadId}`; + } + + /** Key used to track per-generation persisted header/messages. */ + private genKey(workspacePath: string, threadId: string, generationId: string): string { + return `${this.compactionKey(workspacePath, threadId)}#${generationId}`; + } + /** Resolve the path to a checkpoint snapshot file. */ private getCheckpointFilePath(workspacePath: string, threadId: string, generationId: string): string { return path.join( @@ -166,6 +252,17 @@ export class CopilotPersistenceStore { */ deleteWorkspace(workspacePath: string): void { removeDirSync(this.getWorkspaceDir(workspacePath)); + const prefix = `${computeWorkspaceHash(this.workspaceIdResolver(workspacePath))}/`; + this.clearRuntimeStateByPrefix(prefix); + } + + /** Drop all in-memory tracking whose key starts with `prefix`. */ + private clearRuntimeStateByPrefix(prefix: string): void { + for (const map of [this.appendsSinceCompaction, this.compactionThreshold, this.lastHeaderJson, this.msgState]) { + for (const key of map.keys()) { + if (key.startsWith(prefix)) { map.delete(key); } + } + } } // ============================================ @@ -205,20 +302,46 @@ export class CopilotPersistenceStore { } /** - * Load a thread from disk. Applies schema migrations if needed. + * Load a thread from disk. + * + * Threads persist as an append-only log (`thread.jsonl`, schema v2). The log + * is replayed to rebuild the thread. If only a legacy whole-file snapshot + * (`thread.json`, schema v1) exists, it is read, migrated, rewritten as a + * `thread.jsonl`, and the legacy file is deleted (one-time migration). + * * Returns `null` if the thread does not exist or is corrupt. */ loadThread(workspacePath: string, threadId: string): PersistedThread | null { + const logPath = this.getThreadLogFilePath(workspacePath, threadId); + const records = readJsonlSync(logPath); + + if (records !== null) { + // Append-only log is the source of truth once it exists. + const thread = this.replayThread(threadId, records); + if (thread) { + this.compactOnLoadIfNeeded(workspacePath, threadId, thread, records.length); + // Prime per-generation tracking so subsequent appends persist only + // deltas (idempotent if on-load compaction already rebuilt it). + this.rebuildGenTracking(workspacePath, threadId, thread.generations); + // A stale legacy snapshot may linger if a previous migration was + // interrupted after writing the log but before deleting the JSON. + this.deleteLegacyThreadFile(workspacePath, threadId); + return thread; + } + // The log exists but yielded nothing recoverable (empty or every line + // corrupt). Fall through to the legacy snapshot if one still exists, + // rather than reporting the thread as gone. + } + + // Legacy fallback: whole-file thread.json (schema v1). const raw = readJsonSync>(this.getThreadFilePath(workspacePath, threadId)); if (!raw) { return null; } try { const migrated = migrateThread(raw); - // Re-save if migration changed the schema version - if ((raw.schemaVersion as number) !== CURRENT_THREAD_SCHEMA_VERSION) { - this.saveThread(workspacePath, threadId, migrated); - } + // Convert to the append-only format and remove the legacy file. + this.saveThread(workspacePath, threadId, migrated); return migrated; } catch (err) { console.error(`[CopilotPersistenceStore] Failed to migrate thread ${threadId}:`, err); @@ -227,15 +350,392 @@ export class CopilotPersistenceStore { } /** - * Save a thread to disk atomically. - * Injects the current schemaVersion automatically. + * Persist a thread by rewriting its log in compact form (one record per + * live generation), atomically. Used for thread creation, one-time + * migration from the legacy format, and compaction. + * + * This does NOT append — for incremental mutations use {@link appendGeneration}, + * {@link removeGenerationRecord}, {@link truncateFromGeneration} and + * {@link updateThreadMeta}, which append a single record instead of + * rewriting the whole thread. */ saveThread(workspacePath: string, threadId: string, thread: Omit): void { const data: PersistedThread = { ...thread, schemaVersion: CURRENT_THREAD_SCHEMA_VERSION, }; - writeJsonSync(this.getThreadFilePath(workspacePath, threadId), data); + const records = buildCompactedRecords(data); + const content = records.map(r => JSON.stringify(r)).join('\n') + '\n'; + atomicWriteSync(this.getThreadLogFilePath(workspacePath, threadId), content); + // The compacted log fully supersedes any legacy whole-file snapshot. + this.deleteLegacyThreadFile(workspacePath, threadId); + // Reset the append counter and set the next compaction threshold + // proportional to the live generation count. + const key = this.compactionKey(workspacePath, threadId); + this.appendsSinceCompaction.set(key, 0); + this.compactionThreshold.set( + key, + Math.max(COMPACT_MIN_APPENDS, COMPACT_GROWTH_FACTOR * data.generations.length) + ); + // Rebuild per-generation tracking to match what we just wrote. + this.rebuildGenTracking(workspacePath, threadId, data.generations); + } + + /** + * Reset per-generation header/message tracking to match a known-persisted + * set of generations (after a full write, compaction, or load). Clears any + * stale entries for the thread so removed generations don't linger. + */ + private rebuildGenTracking( + workspacePath: string, + threadId: string, + generations: PersistedGeneration[] + ): void { + const prefix = `${this.compactionKey(workspacePath, threadId)}#`; + for (const k of this.lastHeaderJson.keys()) { + if (k.startsWith(prefix)) { this.lastHeaderJson.delete(k); } + } + for (const k of this.msgState.keys()) { + if (k.startsWith(prefix)) { this.msgState.delete(k); } + } + for (const gen of generations) { + const gKey = this.genKey(workspacePath, threadId, gen.id); + const { modelMessages, ...header } = gen; + const msgs = modelMessages ?? []; + this.lastHeaderJson.set(gKey, JSON.stringify(header)); + this.msgState.set(gKey, { + count: msgs.length, + lastJson: msgs.length ? JSON.stringify(msgs[msgs.length - 1]) : '', + }); + } + } + + // ============================================ + // Thread Log — Append API (incremental persistence) + // ============================================ + + /** Append one record to the log without touching the compaction counter. */ + private rawAppendRecord(workspacePath: string, threadId: string, record: ThreadLogRecord): void { + appendLineSync(this.getThreadLogFilePath(workspacePath, threadId), JSON.stringify(record)); + } + + /** + * Append a single record to a thread's log, then run amortized compaction. + */ + appendThreadRecord(workspacePath: string, threadId: string, record: ThreadLogRecord): void { + this.rawAppendRecord(workspacePath, threadId, record); + this.noteAppend(workspacePath, threadId, 1); + } + + /** + * Persist a generation incrementally (covers both add and in-place update). + * + * Writes only what actually changed: + * - a `gen` header record iff the header (everything except `modelMessages`) + * differs from what was last persisted for this generation; + * - one `msg` record for each model message not yet persisted, when the new + * `modelMessages` extend what we already have (the normal per-step case); + * - a single `msgs` reset record when `modelMessages` were rewritten rather + * than extended (e.g. server-side context compaction shortened them). + * + * This is what makes a turn cost O(messages) to persist instead of + * O(messages²): a message is written once, not re-written on every step. + */ + appendGeneration( + workspacePath: string, + threadId: string, + generation: PersistedGeneration, + updatedAt: number = Date.now() + ): void { + const gKey = this.genKey(workspacePath, threadId, generation.id); + const { modelMessages, ...header } = generation; + const msgs = modelMessages ?? []; + let written = 0; + + // 1) Header: append only when it changed. + const headerJson = JSON.stringify(header); + if (this.lastHeaderJson.get(gKey) !== headerJson) { + this.rawAppendRecord(workspacePath, threadId, { t: 'gen', updatedAt, gen: header }); + this.lastHeaderJson.set(gKey, headerJson); + written++; + } + + // 2) Messages: append only the new ones, or reset if rewritten. + const state = this.msgState.get(gKey) ?? { count: 0, lastJson: '' }; + const extendsPrevious = + msgs.length >= state.count && + (state.count === 0 || JSON.stringify(msgs[state.count - 1]) === state.lastJson); + if (extendsPrevious) { + for (let i = state.count; i < msgs.length; i++) { + this.rawAppendRecord(workspacePath, threadId, { t: 'msg', updatedAt, genId: generation.id, message: msgs[i] }); + written++; + } + } else { + // Prefix changed or list shrank — the incremental model no longer + // holds; replace the whole message list for this generation. + this.rawAppendRecord(workspacePath, threadId, { t: 'msgs', updatedAt, genId: generation.id, messages: msgs }); + written++; + } + this.msgState.set(gKey, { + count: msgs.length, + lastJson: msgs.length ? JSON.stringify(msgs[msgs.length - 1]) : '', + }); + + // Trigger amortized compaction once for the whole batch (never mid-batch, + // so the log is always in a consistent state when compaction reads it). + if (written > 0) { + this.noteAppend(workspacePath, threadId, written); + } + } + + /** + * Append a tombstone removing a generation by id. + */ + removeGenerationRecord( + workspacePath: string, + threadId: string, + generationId: string, + updatedAt: number = Date.now() + ): void { + const gKey = this.genKey(workspacePath, threadId, generationId); + this.lastHeaderJson.delete(gKey); + this.msgState.delete(gKey); + this.appendThreadRecord(workspacePath, threadId, { t: 'del', updatedAt, id: generationId }); + } + + /** + * Append a truncation removing `fromGenerationId` and every generation + * appended after it (restore-to-checkpoint). + */ + truncateFromGeneration( + workspacePath: string, + threadId: string, + fromGenerationId: string, + updatedAt: number = Date.now() + ): void { + this.appendThreadRecord(workspacePath, threadId, { t: 'trunc', updatedAt, fromId: fromGenerationId }); + } + + /** + * Append a thread-level metadata update (name / sessionId). + */ + updateThreadMeta( + workspacePath: string, + threadId: string, + meta: { name?: string; sessionId?: string }, + updatedAt: number = Date.now() + ): void { + this.appendThreadRecord(workspacePath, threadId, { t: 'meta', updatedAt, ...meta }); + } + + /** + * Rewrite a thread's log in compact form. Safe to call at any time; a no-op + * if the thread has no log yet. + */ + compactThread(workspacePath: string, threadId: string): void { + const records = readJsonlSync(this.getThreadLogFilePath(workspacePath, threadId)); + if (records === null) { + return; + } + const thread = this.replayThread(threadId, records); + if (thread) { + this.saveThread(workspacePath, threadId, thread); + } + } + + // ============================================ + // Thread Log — internal helpers + // ============================================ + + /** + * Rebuild a {@link PersistedThread} by replaying log records in order. + * + * Returns `null` only when there is nothing to reconstruct (empty file, or a + * file whose every line was unparseable). A `head` record supplies the + * thread's precise `id`/`createdAt`; if it is missing but other records + * exist (e.g. appends reached disk before an interrupted init), the thread + * is still reconstructed using the directory's `threadId` and the earliest + * record timestamp — so no committed data is lost. + */ + private replayThread(threadId: string, records: ThreadLogRecord[]): PersistedThread | null { + let hasHead = false; + let recognized = 0; + let minTimestamp = Number.POSITIVE_INFINITY; + let id = threadId; + let name = ''; + let sessionId: string | undefined; + let createdAt = 0; + let updatedAt = 0; + + // Preserve first-seen order of generations; headers and messages are + // tracked separately and stitched together at the end. + const order: string[] = []; + const headers = new Map(); + const messages = new Map(); + const seen = (genId: string): boolean => headers.has(genId) || messages.has(genId); + const noteGen = (genId: string): void => { if (!seen(genId)) { order.push(genId); } }; + const removeGen = (genId: string): boolean => { + const existed = headers.delete(genId); + const existedM = messages.delete(genId); + if (existed || existedM) { + const idx = order.indexOf(genId); + if (idx !== -1) { order.splice(idx, 1); } + return true; + } + return false; + }; + + for (const record of records) { + switch (record.t) { + case 'head': + hasHead = true; + recognized++; + id = record.id; + createdAt = record.createdAt; + updatedAt = Math.max(updatedAt, record.createdAt); + minTimestamp = Math.min(minTimestamp, record.createdAt); + break; + case 'meta': + recognized++; + if (record.name !== undefined) { name = record.name; } + if (record.sessionId !== undefined) { sessionId = record.sessionId; } + updatedAt = Math.max(updatedAt, record.updatedAt); + minTimestamp = Math.min(minTimestamp, record.updatedAt); + break; + case 'gen': + recognized++; + noteGen(record.gen.id); + headers.set(record.gen.id, record.gen); + updatedAt = Math.max(updatedAt, record.updatedAt); + minTimestamp = Math.min(minTimestamp, record.updatedAt); + break; + case 'msg': { + recognized++; + noteGen(record.genId); + const arr = messages.get(record.genId) ?? []; + arr.push(record.message); + messages.set(record.genId, arr); + updatedAt = Math.max(updatedAt, record.updatedAt); + minTimestamp = Math.min(minTimestamp, record.updatedAt); + break; + } + case 'msgs': + recognized++; + noteGen(record.genId); + messages.set(record.genId, [...record.messages]); + updatedAt = Math.max(updatedAt, record.updatedAt); + minTimestamp = Math.min(minTimestamp, record.updatedAt); + break; + case 'del': + recognized++; + removeGen(record.id); + updatedAt = Math.max(updatedAt, record.updatedAt); + minTimestamp = Math.min(minTimestamp, record.updatedAt); + break; + case 'trunc': { + recognized++; + const idx = order.indexOf(record.fromId); + if (idx !== -1) { + for (const genId of order.slice(idx)) { + headers.delete(genId); + messages.delete(genId); + } + order.splice(idx); + } + updatedAt = Math.max(updatedAt, record.updatedAt); + minTimestamp = Math.min(minTimestamp, record.updatedAt); + break; + } + } + } + + if (recognized === 0) { + // Nothing recoverable — empty file or every line was corrupt. + return null; + } + if (!hasHead && Number.isFinite(minTimestamp)) { + // Reconstruct createdAt from the earliest record when head is absent. + createdAt = minTimestamp; + } + + const generations: PersistedGeneration[] = []; + for (const genId of order) { + const header = headers.get(genId); + // A generation needs its header to be reconstructable. In practice the + // header is always written before its messages; skip if it is missing. + if (!header) { continue; } + generations.push({ ...header, modelMessages: messages.get(genId) ?? [] }); + } + const thread: PersistedThread = { + schemaVersion: CURRENT_THREAD_SCHEMA_VERSION, + id, + name, + createdAt, + updatedAt, + generations, + }; + if (sessionId !== undefined) { + thread.sessionId = sessionId; + } + return thread; + } + + /** + * Increment the append counter and compact once the size-proportional + * threshold is reached. Compaction is a pure optimization — a write failure + * (e.g. transient disk error) must never propagate out of an append, so it + * is caught and the counter is reset to back off until the next interval. + */ + private noteAppend(workspacePath: string, threadId: string, count: number): void { + const key = this.compactionKey(workspacePath, threadId); + const next = (this.appendsSinceCompaction.get(key) ?? 0) + count; + const threshold = this.compactionThreshold.get(key) ?? COMPACT_MIN_APPENDS; + if (next >= threshold) { + try { + // compactThread -> saveThread resets the counter + threshold. + this.compactThread(workspacePath, threadId); + } catch (err) { + console.error(`[CopilotPersistenceStore] Compaction failed for thread ${threadId}:`, err); + this.appendsSinceCompaction.set(key, 0); // back off; retry after another interval + } + } else { + this.appendsSinceCompaction.set(key, next); + } + } + + /** + * Compact on load when the log has grown well beyond its compact size. + * Purely an optimization: `loadThread` has already replayed the thread, so a + * write failure here must be swallowed rather than fail the (previously + * read-only) load and abort the whole workspace restore. + */ + private compactOnLoadIfNeeded( + workspacePath: string, + threadId: string, + thread: PersistedThread, + lineCount: number + ): void { + const compactSize = thread.generations.length + 2; // head + meta + gens + if (lineCount > COMPACT_LOAD_MIN_LINES && lineCount > COMPACT_LOAD_FACTOR * compactSize) { + try { + this.saveThread(workspacePath, threadId, thread); + } catch (err) { + console.error(`[CopilotPersistenceStore] On-load compaction failed for thread ${threadId}:`, err); + // Keep the already-replayed thread; compaction will retry later. + } + } + } + + /** Remove a superseded legacy `thread.json`, if present. */ + private deleteLegacyThreadFile(workspacePath: string, threadId: string): void { + const legacyPath = this.getThreadFilePath(workspacePath, threadId); + try { + if (fs.existsSync(legacyPath)) { + fs.unlinkSync(legacyPath); + } + } catch (err) { + console.error(`[CopilotPersistenceStore] Failed to remove legacy thread file ${threadId}:`, err); + } } /** @@ -243,6 +743,11 @@ export class CopilotPersistenceStore { */ deleteThread(workspacePath: string, threadId: string): void { removeDirSync(this.getThreadDir(workspacePath, threadId)); + const key = this.compactionKey(workspacePath, threadId); + this.appendsSinceCompaction.delete(key); + this.compactionThreshold.delete(key); + // Clear per-generation tracking for this thread (keys are `${key}#genId`). + this.clearRuntimeStateByPrefix(`${key}#`); } // ============================================ diff --git a/workspaces/common-libs/copilot-utilities/src/chat-persistence/schema-migration.ts b/workspaces/common-libs/copilot-utilities/src/chat-persistence/schema-migration.ts index e1f30764f59..4fb39544501 100644 --- a/workspaces/common-libs/copilot-utilities/src/chat-persistence/schema-migration.ts +++ b/workspaces/common-libs/copilot-utilities/src/chat-persistence/schema-migration.ts @@ -22,7 +22,12 @@ import { PersistedThread, PersistedCheckpoint, WorkspaceMetadata } from './types // Current Schema Versions // ============================================ -export const CURRENT_THREAD_SCHEMA_VERSION = 1; +// v2: threads persist as an append-only JSONL log (`thread.jsonl`) instead of a +// single whole-file JSON snapshot (`thread.json`). The generation data shape is +// unchanged between v1 and v2 — only the on-disk storage format differs — so the +// v1 -> v2 migration is a pure identity on the data. Legacy `thread.json` files +// are read once, rewritten as `thread.jsonl`, and then removed. +export const CURRENT_THREAD_SCHEMA_VERSION = 2; export const CURRENT_WORKSPACE_SCHEMA_VERSION = 1; export const CURRENT_CHECKPOINT_SCHEMA_VERSION = 1; @@ -68,8 +73,21 @@ function applyMigrations( // ============================================ // Add future migrations here: -// { fromVersion: 1, toVersion: 2, migrate: (data) => { ... } } -const threadMigrations: SchemaMigration[] = []; +// { fromVersion: 2, toVersion: 3, migrate: (data) => { ... } } +const threadMigrations: SchemaMigration[] = [ + { + // v1 (whole-file thread.json) -> v2 (append-only thread.jsonl). + // Data shape is identical; only the storage format changes. The store + // handles the file-format conversion, so this migration just stamps the + // new version onto the in-memory object. + fromVersion: 1, + toVersion: 2, + migrate: (data) => ({ + ...(data as Record), + schemaVersion: 2, + }) as unknown as PersistedThread, + }, +]; /** * Migrate a raw thread object to the current schema version. diff --git a/workspaces/common-libs/copilot-utilities/src/chat-persistence/types.ts b/workspaces/common-libs/copilot-utilities/src/chat-persistence/types.ts index 591c0bdd8d2..de32c146d7e 100644 --- a/workspaces/common-libs/copilot-utilities/src/chat-persistence/types.ts +++ b/workspaces/common-libs/copilot-utilities/src/chat-persistence/types.ts @@ -141,6 +141,51 @@ export interface PersistedThread { generations: PersistedGeneration[]; } +// ============================================ +// Thread Log Records (append-only JSONL format) +// ============================================ + +/** + * A generation's persisted fields WITHOUT its (potentially large, append-only) + * `modelMessages`. Stored in a `gen` log record; the messages are stored + * separately as individual `msg` records so they are each written exactly once. + */ +export type PersistedGenerationHeader = Omit; + +/** + * A single record in a thread's append-only log (`thread.jsonl`). + * + * Instead of rewriting the whole thread on every mutation, each change appends + * exactly what changed. Crucially, a generation's `modelMessages` grow one + * message at a time as the agent runs, so each message is persisted **once** as + * its own `msg` record rather than re-writing the whole (growing) generation on + * every step — that is what keeps per-turn write cost O(messages) instead of + * O(messages²). On load the records are replayed in order to rebuild an + * identical {@link PersistedThread}. + * + * Replay semantics: + * - `head` — thread identity + schema version. Written once as the first line. + * - `meta` — thread-level mutable fields (name/sessionId). Last write wins. + * - `gen` — upsert a generation's header (everything except `modelMessages`) + * by `gen.id`. First occurrence fixes its position; later + * occurrences replace the header in place (last write wins). + * - `msg` — append one model message to a generation's `modelMessages`. + * - `msgs` — replace a generation's entire `modelMessages` list. Used only when + * the messages are rewritten rather than extended (e.g. server-side + * context compaction), which the incremental `msg` path can't model. + * - `del` — remove a generation by id (tombstone). + * - `trunc` — remove the generation `fromId` and every generation that was + * appended after it (restore-to-checkpoint). + */ +export type ThreadLogRecord = + | { t: 'head'; v: number; id: string; createdAt: number } + | { t: 'meta'; updatedAt: number; name?: string; sessionId?: string } + | { t: 'gen'; updatedAt: number; gen: PersistedGenerationHeader } + | { t: 'msg'; updatedAt: number; genId: string; message: unknown } + | { t: 'msgs'; updatedAt: number; genId: string; messages: unknown[] } + | { t: 'del'; updatedAt: number; id: string } + | { t: 'trunc'; updatedAt: number; fromId: string }; + // ============================================ // Persisted Checkpoint // ============================================