diff --git a/src/config/sessions.cache.test.ts b/src/config/sessions.cache.test.ts index 68bc62f100e..ff27cae0010 100644 --- a/src/config/sessions.cache.test.ts +++ b/src/config/sessions.cache.test.ts @@ -7,10 +7,12 @@ import { createSuiteTempRootTracker } from "../test-helpers/temp-dir.js"; import { getSerializedSessionStore, getSerializedSessionStoreCacheStatsForTest, + getSerializedSessionStorePromptRefs, getSessionStoreSnapshotCacheStatsForTest, getSessionStoreStringInternStatsForTest, readSessionStoreCache, setSerializedSessionStore, + setSerializedSessionStorePromptRefs, writeSessionStoreCache, } from "./sessions/store-cache.js"; import { @@ -27,6 +29,7 @@ import { updateLastRoute, } from "./sessions/store.js"; import type { SessionEntry } from "./sessions/types.js"; +import type { SessionSkillPromptRef } from "./sessions/types.js"; function createSessionEntry(overrides: Partial = {}): SessionEntry { return { @@ -102,6 +105,24 @@ describe("Session Store Cache", () => { expect(getSerializedSessionStoreCacheStatsForTest().entries).toBe(maxEntries); }); + it("keeps serialized prompt refs on the serialized cache entry lifecycle", () => { + const promptRef: SessionSkillPromptRef = { + version: 1, + algorithm: "sha256", + hash: "a".repeat(64), + bytes: 123, + }; + const refs = new Map([["session:1", promptRef]]); + + setSerializedSessionStore("store:refs", "{}"); + setSerializedSessionStorePromptRefs("store:refs", refs); + + expect(getSerializedSessionStorePromptRefs("store:refs")).toBe(refs); + + setSerializedSessionStore("store:refs", "{}"); + expect(getSerializedSessionStorePromptRefs("store:refs")).toBeUndefined(); + }); + it("should load session store from disk on first call", async () => { const testStore = createSingleSessionStore(); diff --git a/src/config/sessions/store-cache.ts b/src/config/sessions/store-cache.ts index 68761dfe8a6..cc00450a9d2 100644 --- a/src/config/sessions/store-cache.ts +++ b/src/config/sessions/store-cache.ts @@ -1,7 +1,7 @@ import { parseStrictNonNegativeInteger } from "../../infra/parse-finite-number.js"; import { createExpiringMapCache, isCacheEnabled, resolveCacheTtlMs } from "../cache-utils.js"; import { clearSessionSkillPromptRefCache } from "./skill-prompt-blobs.js"; -import type { SessionEntry } from "./types.js"; +import type { SessionEntry, SessionSkillPromptRef } from "./types.js"; export type DeepReadonly = T extends (...args: never[]) => unknown ? T @@ -35,6 +35,7 @@ type SessionStoreSnapshotCacheEntry = { type SerializedSessionStoreCacheEntry = { serialized: string; sizeBytes: number; + promptRefs?: ReadonlyMap; }; const DEFAULT_SESSION_STORE_TTL_MS = 45_000; // 45 seconds (between 30-60s) @@ -315,10 +316,30 @@ export function getSerializedSessionStore(storePath: string): string | undefined return SESSION_STORE_SERIALIZED_CACHE.get(storePath)?.serialized; } +export function getSerializedSessionStorePromptRefs( + storePath: string, +): ReadonlyMap | undefined { + pruneSerializedSessionStoreCache(); + return SESSION_STORE_SERIALIZED_CACHE.get(storePath)?.promptRefs; +} + +export function setSerializedSessionStorePromptRefs( + storePath: string, + promptRefs: ReadonlyMap, +): void { + pruneSerializedSessionStoreCache(); + const cached = SESSION_STORE_SERIALIZED_CACHE.get(storePath); + if (!cached) { + return; + } + cached.promptRefs = promptRefs; +} + export function setSerializedSessionStore( storePath: string, serialized?: string, sizeBytesHint?: number, + promptRefs?: ReadonlyMap, ): void { deleteSerializedSessionStore(storePath); if (serialized === undefined) { @@ -333,7 +354,7 @@ export function setSerializedSessionStore( if (maxEntries <= 0 || maxBytes <= 0 || sizeBytes > maxBytes) { return; } - SESSION_STORE_SERIALIZED_CACHE.set(storePath, { serialized, sizeBytes }); + SESSION_STORE_SERIALIZED_CACHE.set(storePath, { serialized, sizeBytes, promptRefs }); sessionStoreSerializedCacheBytes += sizeBytes; pruneSerializedSessionStoreCache(); } @@ -422,6 +443,7 @@ export function writeSessionStoreCache(params: { mtimeMs?: number; sizeBytes?: number; serialized?: string; + serializedPromptRefs?: ReadonlyMap; cloneSerialized?: string; takeOwnership?: boolean; }): void { @@ -437,5 +459,10 @@ export function writeSessionStoreCache(params: { sizeBytes: params.sizeBytes, serialized: params.cloneSerialized, }); - setSerializedSessionStore(params.storePath, params.serialized, params.sizeBytes); + setSerializedSessionStore( + params.storePath, + params.serialized, + params.sizeBytes, + params.serializedPromptRefs, + ); } diff --git a/src/config/sessions/store.ts b/src/config/sessions/store.ts index 43d3d00b3fa..ef02e1f6bc4 100644 --- a/src/config/sessions/store.ts +++ b/src/config/sessions/store.ts @@ -28,9 +28,11 @@ import { dropSessionStoreObjectCache, dropSessionStoreSnapshotCache, getSerializedSessionStore, + getSerializedSessionStorePromptRefs, getSessionStoreCacheVersion, invalidateSessionStoreCache, isSessionStoreCacheEnabled, + setSerializedSessionStorePromptRefs, setSerializedSessionStore, takeMutableSessionStoreCache, writeSessionStoreCache, @@ -234,11 +236,17 @@ function updateSessionStoreWriteCaches(params: { storePath: string; store: Record; serialized: string; + serializedPromptRefs?: ReadonlyMap; cloneSerialized?: string; takeOwnership?: boolean; }): void { const fileStat = getFileStatSnapshot(params.storePath); - setSerializedSessionStore(params.storePath, params.serialized, fileStat?.sizeBytes); + setSerializedSessionStore( + params.storePath, + params.serialized, + fileStat?.sizeBytes, + params.serializedPromptRefs, + ); if (!isSessionStoreCacheEnabled()) { dropSessionStoreObjectCache(params.storePath); dropSessionStoreSnapshotCache(params.storePath); @@ -250,6 +258,7 @@ function updateSessionStoreWriteCaches(params: { mtimeMs: fileStat?.mtimeMs, sizeBytes: fileStat?.sizeBytes, serialized: params.serialized, + serializedPromptRefs: params.serializedPromptRefs, cloneSerialized: params.cloneSerialized, takeOwnership: params.takeOwnership, }); @@ -273,18 +282,26 @@ function restoreUnchangedSessionStoreCache( return; } const serialized = getSerializedSessionStore(storePath); + const serializedPromptRefs = + serialized !== undefined ? getSerializedSessionStorePromptRefs(storePath) : undefined; writeSessionStoreCache({ storePath, store, mtimeMs: loadedFileStat?.mtimeMs, sizeBytes: loadedFileStat?.sizeBytes, serialized, + serializedPromptRefs, takeOwnership: true, }); if (serialized !== undefined) { // Keep hydrated blob prompts in the object cache, but preserve the disk JSON // comparison string so repeated no-op saves do not rewrite sessions.json. - setSerializedSessionStore(storePath, serialized, loadedFileStat?.sizeBytes); + setSerializedSessionStore( + storePath, + serialized, + loadedFileStat?.sizeBytes, + serializedPromptRefs, + ); } } @@ -333,11 +350,16 @@ function indentTopLevelEntryJson(json: string): string { function buildSingleEntrySerializedStore(params: { storePath: string; patch: SingleEntryPersistencePatch; -}): { serialized: string; promptBlobs: SessionSkillPromptBlobProjection[] } | null { +}): { + serialized: string; + promptBlobs: SessionSkillPromptBlobProjection[]; + promptRefs: ReadonlyMap; +} | null { const currentSerialized = getSerializedSessionStore(params.storePath); if (currentSerialized === undefined) { return null; } + const currentPromptRefs = getSerializedPromptRefs(params.storePath, currentSerialized); const marker = `\n ${JSON.stringify(params.patch.sessionKey)}: `; const markerIndex = currentSerialized.indexOf(marker); if (markerIndex < 0) { @@ -360,10 +382,18 @@ function buildSingleEntrySerializedStore(params: { return null; } const entryJson = indentTopLevelEntryJson(JSON.stringify(projectedEntry, null, 2)); + const promptRefs = new Map(currentPromptRefs); + const promptRef = projectedEntry.skillsSnapshot?.promptRef; + if (promptRef) { + promptRefs.set(params.patch.sessionKey, promptRef); + } else { + promptRefs.delete(params.patch.sessionKey); + } return { serialized: currentSerialized.slice(0, valueStart) + entryJson + currentSerialized.slice(valueEnd), promptBlobs: [...projected.promptBlobs.values()], + promptRefs, }; } @@ -383,15 +413,42 @@ function collectSerializedPromptRefs(serialized: string): Map, +): Map { + const refs = new Map(); + for (const [key, entry] of Object.entries(store)) { + const ref = entry?.skillsSnapshot?.promptRef; + if (ref) { + refs.set(key, ref); + } + } + return refs; +} + +function getSerializedPromptRefs( + storePath: string, + serialized: string, +): ReadonlyMap { + const cached = getSerializedSessionStorePromptRefs(storePath); + if (cached) { + return cached; + } + const refs = collectSerializedPromptRefs(serialized); + setSerializedSessionStorePromptRefs(storePath, refs); + return refs; +} + function storeHasUnsafeUntouchedHydratedSkillPrompts( storePath: string, store: Record, changedSessionKey: string, ): boolean { const currentSerialized = getSerializedSessionStore(storePath); - const serializedPromptRefs = currentSerialized - ? collectSerializedPromptRefs(currentSerialized) - : undefined; + const serializedPromptRefs = + currentSerialized !== undefined + ? getSerializedPromptRefs(storePath, currentSerialized) + : undefined; for (const [key, entry] of Object.entries(store)) { if (key === changedSessionKey || typeof entry.skillsSnapshot?.prompt !== "string") { continue; @@ -684,6 +741,7 @@ async function saveSessionStoreUnlocked( storePath, store, serialized: singleEntrySerialized.serialized, + serializedPromptRefs: singleEntrySerialized.promptRefs, promptBlobs: singleEntrySerialized.promptBlobs, takeOwnership: opts?.takeCacheOwnership, }); @@ -692,6 +750,7 @@ async function saveSessionStoreUnlocked( } const persisted = projectSessionStoreForPersistence({ storePath, store }); const promptBlobs = [...persisted.promptBlobs.values()]; + const promptRefs = collectStorePromptRefs(persisted.store); const json = JSON.stringify(persisted.store, null, 2); const cloneSerialized = persisted.changed ? undefined : json; if (getSerializedSessionStore(storePath) === json) { @@ -703,6 +762,7 @@ async function saveSessionStoreUnlocked( storePath, store, serialized: json, + serializedPromptRefs: promptRefs, cloneSerialized, takeOwnership: opts?.takeCacheOwnership, }); @@ -717,6 +777,7 @@ async function saveSessionStoreUnlocked( storePath, store, serialized: json, + serializedPromptRefs: promptRefs, cloneSerialized, promptBlobs, takeOwnership: opts?.takeCacheOwnership, @@ -744,6 +805,7 @@ async function saveSessionStoreUnlocked( storePath, store, serialized: json, + serializedPromptRefs: promptRefs, cloneSerialized, promptBlobs, takeOwnership: opts?.takeCacheOwnership, @@ -759,6 +821,7 @@ async function saveSessionStoreUnlocked( storePath, store, serialized: json, + serializedPromptRefs: promptRefs, cloneSerialized, promptBlobs, takeOwnership: opts?.takeCacheOwnership, @@ -882,6 +945,7 @@ async function writeSessionStoreAtomic(params: { storePath: string; store: Record; serialized: string; + serializedPromptRefs?: ReadonlyMap; cloneSerialized?: string; promptBlobs: Iterable; takeOwnership?: boolean; @@ -904,6 +968,7 @@ async function writeSessionStoreAtomic(params: { storePath: params.storePath, store: params.store, serialized: params.serialized, + serializedPromptRefs: params.serializedPromptRefs, cloneSerialized: params.cloneSerialized, takeOwnership: params.takeOwnership, });