diff --git a/app/extension/src/__tests__/sessionStoragePersistence.test.ts b/app/extension/src/__tests__/sessionStoragePersistence.test.ts index dde8e58..b484ecc 100644 --- a/app/extension/src/__tests__/sessionStoragePersistence.test.ts +++ b/app/extension/src/__tests__/sessionStoragePersistence.test.ts @@ -249,6 +249,79 @@ function createRunningSession(text: string) { }; } +function createTwoTurnSession(secondQuestion = "Follow-up") { + return { + ...createSession("Question", "First answer"), + messages: [ + { + id: "user-1", + role: "user" as const, + parts: [{ type: "text" as const, text: "Question" }], + status: "complete" as const, + }, + { + id: "assistant-1", + role: "assistant" as const, + parts: [{ type: "text" as const, text: "First answer" }], + status: "complete" as const, + }, + { + id: "user-2", + role: "user" as const, + parts: [{ type: "text" as const, text: secondQuestion }], + status: "complete" as const, + }, + { + id: "assistant-2", + role: "assistant" as const, + parts: [{ type: "text" as const, text: "Second answer" }], + status: "complete" as const, + }, + ], + }; +} + +function createTwoTurnSessionWithAttachment(secondAnswer = "Second answer") { + return { + ...createSession("Question", "First answer"), + messages: [ + { + id: "user-1", + role: "user" as const, + parts: [ + { type: "text" as const, text: "Question" }, + { + type: "file" as const, + filename: "note.txt", + mediaType: "text/plain", + dataUrl: "data:text/plain;base64,SGVsbG8=", + size: 5, + }, + ], + status: "complete" as const, + }, + { + id: "assistant-1", + role: "assistant" as const, + parts: [{ type: "text" as const, text: "First answer" }], + status: "complete" as const, + }, + { + id: "user-2", + role: "user" as const, + parts: [{ type: "text" as const, text: "Follow-up" }], + status: "complete" as const, + }, + { + id: "assistant-2", + role: "assistant" as const, + parts: [{ type: "text" as const, text: secondAnswer }], + status: "complete" as const, + }, + ], + }; +} + describe("sessionStorage persistence layout", () => { let fakeIndexedDB: FakeIndexedDB; @@ -275,7 +348,15 @@ describe("sessionStorage persistence layout", () => { expect(sessionRecord.schemaVersion).toBe(4); expect(sessionRecord.messages).toBeUndefined(); expect(sessionRecord.messageRefs).toHaveLength(2); + expect(sessionRecord.messageRefs).toEqual([ + expect.not.objectContaining({ signature: expect.anything() }), + expect.not.objectContaining({ signature: expect.anything() }), + ]); expect(messageRecords.size).toBe(2); + expect(Array.from(messageRecords.values())).toEqual([ + expect.not.objectContaining({ signature: expect.anything() }), + expect.not.objectContaining({ signature: expect.anything() }), + ]); const restored = await getSession("session-1"); expect(restored?.messages.map((message) => message.id)).toEqual([ @@ -284,7 +365,7 @@ describe("sessionStorage persistence layout", () => { ]); }); - it("only rewrites changed message rows on later saves", async () => { + it("only rewrites the latest message row on streaming saves", async () => { const { saveSession } = await import("../sidepanel/sessionStorage"); await saveSession(createSession("Question", "First chunk")); @@ -305,6 +386,90 @@ describe("sessionStorage persistence layout", () => { ).toBe("Second chunk"); }); + it("rewrites only the divergent suffix after a history edit", async () => { + const { saveSession } = await import("../sidepanel/sessionStorage"); + + await saveSession(createTwoTurnSession()); + + const db = fakeIndexedDB.databases.get("huntly-agent")!; + const messageStore = db.stores.get("session-messages")!; + messageStore.putCount = 0; + + const edited = createTwoTurnSession("Edited follow-up"); + edited.messages[2] = { + ...edited.messages[2], + id: "user-2-edited", + }; + edited.messages[3] = { + ...edited.messages[3], + id: "assistant-2-edited", + }; + + await saveSession(edited); + + expect(messageStore.putCount).toBe(2); + const storedIds = Array.from(messageStore.records.values()).map( + (record) => (record as { message: { id: string } }).message.id + ); + expect(storedIds).toEqual([ + "user-1", + "assistant-1", + "user-2-edited", + "assistant-2-edited", + ]); + }); + + it("keeps earlier attachment refs stable when only the latest message rewrites", async () => { + const { saveSession } = await import("../sidepanel/sessionStorage"); + const originalFetch = globalThis.fetch; + + globalThis.fetch = jest.fn(async () => ({ + ok: true, + blob: async () => + ({ size: 5, type: "text/plain" }) as unknown as Blob, + })) as unknown as typeof fetch; + + try { + await saveSession(createTwoTurnSessionWithAttachment()); + + const db = fakeIndexedDB.databases.get("huntly-agent")!; + const messageStore = db.stores.get("session-messages")!; + const attachmentStore = db.stores.get("session-attachments")!; + const storedUserMessage = messageStore.records.get( + "session-1\u001fuser-1" + ) as { + message: { parts: Array<{ type: string; attachmentId?: string }> }; + }; + const firstAttachmentId = storedUserMessage.message.parts[1].attachmentId; + + expect(firstAttachmentId).toBeDefined(); + expect(attachmentStore.records.has(firstAttachmentId!)).toBe(true); + + messageStore.putCount = 0; + + await saveSession( + createTwoTurnSessionWithAttachment("Updated second answer") + ); + + expect(messageStore.putCount).toBe(1); + + const updatedUserMessage = messageStore.records.get( + "session-1\u001fuser-1" + ) as { + message: { parts: Array<{ type: string; attachmentId?: string }> }; + }; + const updatedAttachmentId = + updatedUserMessage.message.parts[1].attachmentId; + + expect(updatedAttachmentId).toBe(firstAttachmentId); + expect(Array.from(attachmentStore.records.keys())).toEqual([ + firstAttachmentId, + ]); + } finally { + globalThis.fetch = originalFetch; + } + }); + it("keeps history metadata and stored messages after a run completes", async () => { const { getSession, listSessionMetadata, saveSession } = await import( "../sidepanel/sessionStorage" @@ -328,43 +493,81 @@ describe("sessionStorage persistence layout", () => { expect(restored?.messages[1].parts[0].text).toBe("Final answer"); }); - it("resets older chat stores before writing the current schema", async () => { + it("migrates older chat stores without dropping persisted data", async () => { const legacyDb = new FakeDatabase(); - legacyDb.version = 3; + legacyDb.version = 2; legacyDb.createObjectStore("sessions", { keyPath: "id" }); legacyDb.createObjectStore("session-metadata", { keyPath: "id" }); + legacyDb.createObjectStore("session-attachments", { keyPath: "id" }); legacyDb.stores.get("sessions")!.records.set("legacy-session", { + ...createSession("Legacy question", "Legacy answer"), id: "legacy-session", - messages: [], + title: "Legacy chat", + createdAt: "2026-04-24T08:00:00.000Z", + updatedAt: "2026-04-24T08:00:01.000Z", + lastMessageAt: "2026-04-24T08:00:01.000Z", + lastOpenedAt: "2026-04-24T08:00:01.000Z", }); legacyDb.stores.get("session-metadata")!.records.set("legacy-session", { id: "legacy-session", title: "Legacy chat", createdAt: "2026-04-24T08:00:00.000Z", - updatedAt: "2026-04-24T08:00:00.000Z", - messageCount: 0, + updatedAt: "2026-04-24T08:00:01.000Z", + messageCount: 2, preview: "", currentModelId: null, }); + legacyDb.stores + .get("session-attachments")! + .records.set("legacy-attachment", { + id: "legacy-attachment", + sessionId: "legacy-session", + blob: { size: 17, type: "text/plain" } as unknown as Blob, + createdAt: "2026-04-24T08:00:01.000Z", + mediaType: "text/plain", + size: 17, + }); fakeIndexedDB.databases.set("huntly-agent", legacyDb); - const { listSessionMetadata, saveSession } = await import( + const { getSession, listSessionMetadata, saveSession } = await import( "../sidepanel/sessionStorage" ); - await saveSession(createSession("Question", "Answer")); - const db = fakeIndexedDB.databases.get("huntly-agent")!; + const restoredLegacy = await getSession("legacy-session"); + expect(db.version).toBe(4); - expect(db.stores.get("sessions")!.records.has("legacy-session")).toBe( - false - ); + const migratedSession = db.stores + .get("sessions")! + .records.get("legacy-session") as Record; + expect(migratedSession.messages).toBeUndefined(); + expect(migratedSession.messageRefs).toHaveLength(2); + expect(db.stores.get("session-messages")!.records.size).toBe(2); expect( - db.stores.get("session-metadata")!.records.has("legacy-session") - ).toBe(false); + db.stores.get("session-attachments")!.records.has("legacy-attachment") + ).toBe(true); + + expect(restoredLegacy?.title).toBe("Legacy chat"); + expect(restoredLegacy?.messages.map((message) => message.id)).toEqual([ + "user-1", + "assistant-1", + ]); + expect(restoredLegacy?.messages[1].parts[0].text).toBe("Legacy answer"); const metadata = await listSessionMetadata(); - expect(metadata.map((session) => session.id)).toEqual(["session-1"]); - expect(metadata[0].preview).toBe("Question\nAnswer"); + expect(metadata).toHaveLength(1); + expect(metadata[0].id).toBe("legacy-session"); + expect(metadata[0].preview).toBe("Legacy question\nLegacy answer"); + + await saveSession(createSession("Question", "Answer")); + + const mergedMetadata = await listSessionMetadata(); + expect(mergedMetadata).toHaveLength(2); + expect(mergedMetadata.map((session) => session.id)).toEqual( + expect.arrayContaining(["legacy-session", "session-1"]) + ); + expect( + (await getSession("legacy-session"))?.messages[1].parts[0].text + ).toBe("Legacy answer"); }); }); diff --git a/app/extension/src/sidepanel/sessionStorage.ts b/app/extension/src/sidepanel/sessionStorage.ts index b33ebbd..2de8c29 100644 --- a/app/extension/src/sidepanel/sessionStorage.ts +++ b/app/extension/src/sidepanel/sessionStorage.ts @@ -23,6 +23,7 @@ import { const DB_NAME = "huntly-agent"; const DB_VERSION = 4; +const SPLIT_STORES_DB_VERSION = 4; const SESSION_RECORD_SCHEMA_VERSION = 4; const SESSIONS_STORE = "sessions"; const SESSION_MESSAGES_STORE = "session-messages"; @@ -35,7 +36,6 @@ type StoredMessageRef = { storageKey: string; messageId: string; order: number; - signature: string; }; type StoredSessionRecord = Omit & { @@ -50,7 +50,6 @@ type StoredMessageRecord = { sessionId: string; messageId: string; order: number; - signature: string; message: ChatMessage; updatedAt: string; }; @@ -66,13 +65,6 @@ type StoredAttachmentRecord = { let databasePromise: Promise | null = null; -const CHAT_STORE_NAMES = [ - SESSIONS_STORE, - SESSION_MESSAGES_STORE, - METADATA_STORE, - ATTACHMENTS_STORE, -] as const; - function requestToPromise(request: IDBRequest): Promise { return new Promise((resolve, reject) => { request.onsuccess = () => resolve(request.result); @@ -81,12 +73,6 @@ function requestToPromise(request: IDBRequest): Promise { }); } -function deleteObjectStoreIfExists(db: IDBDatabase, storeName: string): void { - if (db.objectStoreNames.contains(storeName)) { - db.deleteObjectStore(storeName); - } -} - function ensureObjectStores( db: IDBDatabase, transaction: IDBTransaction @@ -129,14 +115,13 @@ function openDatabase(): Promise { request.onupgradeneeded = (event) => { const db = request.result; const oldVersion = event.oldVersion; + const transaction = request.transaction!; - if (oldVersion > 0 && oldVersion < DB_VERSION) { - CHAT_STORE_NAMES.forEach((storeName) => { - deleteObjectStoreIfExists(db, storeName); - }); - } + ensureObjectStores(db, transaction); - ensureObjectStores(db, request.transaction!); + if (oldVersion > 0 && oldVersion < SPLIT_STORES_DB_VERSION) { + migrateLegacySessionRecords(transaction); + } }; request.onsuccess = () => resolve(request.result); @@ -210,8 +195,12 @@ function getMessageStorageKey(sessionId: string, messageId: string): string { return `${sessionId}\u001f${messageId}`; } -function getMessageSignature(message: ChatMessage): string { - return JSON.stringify(message); +function getFallbackAttachmentId( + sessionId: string, + messageId: string, + partIndex: number +): string { + return `${sessionId}\u001f${messageId}\u001fattachment-${partIndex}`; } function createMessageRef( @@ -224,7 +213,6 @@ function createMessageRef( storageKey: getMessageStorageKey(sessionId, messageId), messageId, order, - signature: getMessageSignature(message), }; } @@ -457,13 +445,19 @@ async function serializeSessionAttachments( const referencedAttachmentIds = new Set(); const messages = await Promise.all( - session.messages.map(async (message) => ({ + session.messages.map(async (message, messageOrder) => ({ ...message, parts: await Promise.all( - message.parts.map(async (part) => { + message.parts.map(async (part, partIndex) => { if (part.type !== "file") return part; - const attachmentId = part.attachmentId || crypto.randomUUID(); + const attachmentId = + part.attachmentId || + getFallbackAttachmentId( + session.id, + getMessageId(message, messageOrder), + partIndex + ); referencedAttachmentIds.add(attachmentId); if (part.dataUrl && !attachmentRecords.has(attachmentId)) { @@ -636,39 +630,119 @@ export function buildSessionMetadata(session: SessionData): SessionMetadata { }; } -// --------------------------------------------------------------------------- -// CRUD operations -// --------------------------------------------------------------------------- +function isLegacySessionRecord(record: unknown): record is SessionData { + return Boolean( + record && + typeof record === "object" && + Array.isArray((record as { messages?: unknown }).messages) + ); +} -export async function saveSession(session: SessionData): Promise { - const normalizedSession = normalizeSessionTiming(session); - const serialized = await serializeSessionAttachments(normalizedSession); - const safeMessages = toJsonSafe(serialized.session.messages); - const safeSessionForMetadata: SessionData = { - ...serialized.session, +function buildStoredSessionArtifacts(session: SessionData): { + sessionRecord: StoredSessionRecord; + metadata: SessionMetadata; + messageRecords: StoredMessageRecord[]; +} { + const safeMessages = toJsonSafe(session.messages); + const safeSession: SessionData = { + ...session, messages: safeMessages, }; const messageRecords = safeMessages.map((message, order) => - createMessageRecord( - serialized.session.id, - message, - order, - serialized.session.updatedAt - ) + createMessageRecord(safeSession.id, message, order, safeSession.updatedAt) ); const messageRefs = messageRecords.map( - ({ storageKey, messageId, order, signature }) => ({ + ({ storageKey, messageId, order }) => ({ storageKey, messageId, order, - signature, }) ); - const sessionRecord = createStoredSessionRecord( - safeSessionForMetadata, - messageRefs + + return { + sessionRecord: createStoredSessionRecord(safeSession, messageRefs), + metadata: buildSessionMetadata(safeSession), + messageRecords, + }; +} + +function getFirstMessageRefDivergenceIndex( + existingRefs: StoredMessageRef[], + messageRecords: StoredMessageRecord[] +): number { + const comparableLength = Math.min(existingRefs.length, messageRecords.length); + + for (let index = 0; index < comparableLength; index += 1) { + const existingRef = existingRefs[index]; + const messageRecord = messageRecords[index]; + if ( + existingRef.storageKey !== messageRecord.storageKey || + existingRef.messageId !== messageRecord.messageId || + existingRef.order !== messageRecord.order + ) { + return index; + } + } + + return existingRefs.length === messageRecords.length ? -1 : comparableLength; +} + +function getMessageRecordsToPut( + existingRefs: StoredMessageRef[], + messageRecords: StoredMessageRecord[] +): StoredMessageRecord[] { + if (messageRecords.length === 0) { + return []; + } + + const divergenceIndex = getFirstMessageRefDivergenceIndex( + existingRefs, + messageRecords ); - const metadata = buildSessionMetadata(safeSessionForMetadata); + + if (divergenceIndex !== -1) { + return messageRecords.slice(divergenceIndex); + } + + return [messageRecords[messageRecords.length - 1]]; +} + +function migrateLegacySessionRecords(transaction: IDBTransaction): void { + const sessionStore = transaction.objectStore(SESSIONS_STORE); + const messageStore = transaction.objectStore(SESSION_MESSAGES_STORE); + const metadataStore = transaction.objectStore(METADATA_STORE); + const getAllSessionsRequest = sessionStore.getAll(); + + getAllSessionsRequest.onsuccess = () => { + const storedSessions = (getAllSessionsRequest.result || []) as unknown[]; + + storedSessions.forEach((storedSession) => { + if (!isLegacySessionRecord(storedSession)) { + return; + } + + const artifacts = buildStoredSessionArtifacts( + normalizeSessionTiming(storedSession) + ); + + sessionStore.put(toJsonSafe(artifacts.sessionRecord)); + metadataStore.put(toJsonSafe(artifacts.metadata)); + artifacts.messageRecords.forEach((messageRecord) => { + messageStore.put(toJsonSafe(messageRecord)); + }); + }); + }; +} + +// --------------------------------------------------------------------------- +// CRUD operations +// --------------------------------------------------------------------------- + +export async function saveSession(session: SessionData): Promise { + const normalizedSession = normalizeSessionTiming(session); + const serialized = await serializeSessionAttachments(normalizedSession); + const { messageRecords, metadata, sessionRecord } = + buildStoredSessionArtifacts(serialized.session); const safeSessionRecord = toJsonSafe(sessionRecord); const safeMetadata = toJsonSafe(metadata); const safeMessageRecords = toJsonSafe(messageRecords); @@ -681,24 +755,15 @@ export async function saveSession(session: SessionData): Promise { stores[SESSIONS_STORE].get(safeSessionRecord.id) )) as StoredSessionRecord | null) || null; const existingRefs = existingSession ? existingSession.messageRefs : []; - const existingRefsByKey = new Map( - existingRefs.map((ref) => [ref.storageKey, ref]) - ); const nextMessageKeys = new Set( safeMessageRecords.map((messageRecord) => messageRecord.storageKey) ); const removedMessageKeys = existingRefs .map((ref) => ref.storageKey) .filter((key) => !nextMessageKeys.has(key)); - const changedMessageRecords = safeMessageRecords.filter( - (messageRecord) => { - const existingRef = existingRefsByKey.get(messageRecord.storageKey); - return ( - !existingRef || - existingRef.order !== messageRecord.order || - existingRef.signature !== messageRecord.signature - ); - } + const messageRecordsToPut = getMessageRecordsToPut( + existingRefs, + safeMessageRecords ); const existingAttachmentIds = await getAttachmentIdsForSession( @@ -712,7 +777,7 @@ export async function saveSession(session: SessionData): Promise { await Promise.all([ requestToPromise(stores[SESSIONS_STORE].put(safeSessionRecord)), requestToPromise(stores[METADATA_STORE].put(safeMetadata)), - ...changedMessageRecords.map((messageRecord) => + ...messageRecordsToPut.map((messageRecord) => requestToPromise(stores[SESSION_MESSAGES_STORE].put(messageRecord)) ), ...removedMessageKeys.map((messageKey) =>