perf: cache serialized session prompt refs

This commit is contained in:
Peter Steinberger
2026-05-30 22:43:59 +01:00
parent 0be3ef5a38
commit 71b3bc87ca
3 changed files with 122 additions and 9 deletions

View File

@@ -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> = {}): 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();

View File

@@ -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> = T extends (...args: never[]) => unknown
? T
@@ -35,6 +35,7 @@ type SessionStoreSnapshotCacheEntry = {
type SerializedSessionStoreCacheEntry = {
serialized: string;
sizeBytes: number;
promptRefs?: ReadonlyMap<string, SessionSkillPromptRef>;
};
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<string, SessionSkillPromptRef> | undefined {
pruneSerializedSessionStoreCache();
return SESSION_STORE_SERIALIZED_CACHE.get(storePath)?.promptRefs;
}
export function setSerializedSessionStorePromptRefs(
storePath: string,
promptRefs: ReadonlyMap<string, SessionSkillPromptRef>,
): 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<string, SessionSkillPromptRef>,
): 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<string, SessionSkillPromptRef>;
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,
);
}

View File

@@ -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<string, SessionEntry>;
serialized: string;
serializedPromptRefs?: ReadonlyMap<string, SessionSkillPromptRef>;
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<string, SessionSkillPromptRef>;
} | 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<string, SessionSki
return refs;
}
function collectStorePromptRefs(
store: Record<string, SessionEntry>,
): Map<string, SessionSkillPromptRef> {
const refs = new Map<string, SessionSkillPromptRef>();
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<string, SessionSkillPromptRef> {
const cached = getSerializedSessionStorePromptRefs(storePath);
if (cached) {
return cached;
}
const refs = collectSerializedPromptRefs(serialized);
setSerializedSessionStorePromptRefs(storePath, refs);
return refs;
}
function storeHasUnsafeUntouchedHydratedSkillPrompts(
storePath: string,
store: Record<string, SessionEntry>,
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<string, SessionEntry>;
serialized: string;
serializedPromptRefs?: ReadonlyMap<string, SessionSkillPromptRef>;
cloneSerialized?: string;
promptBlobs: Iterable<SessionSkillPromptBlobProjection>;
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,
});