From c588606a9bf30d5d616e16e803710da4bb266bc4 Mon Sep 17 00:00:00 2001 From: Josh Lehman Date: Wed, 24 Jun 2026 06:15:09 -0700 Subject: [PATCH] refactor: route checkpoint mutations through accessor (#96222) --- scripts/check-session-accessor-boundary.mjs | 1 + src/config/sessions/session-accessor.test.ts | 130 +++++++++++++ src/config/sessions/session-accessor.ts | 174 +++++++++++++++++- .../session-compaction-checkpoints.test.ts | 31 ++++ src/gateway/session-compaction-checkpoints.ts | 149 ++++++--------- .../check-session-accessor-boundary.test.ts | 1 + 6 files changed, 390 insertions(+), 96 deletions(-) diff --git a/scripts/check-session-accessor-boundary.mjs b/scripts/check-session-accessor-boundary.mjs index 0bc11c4d24ac..4daa0f62ad7a 100644 --- a/scripts/check-session-accessor-boundary.mjs +++ b/scripts/check-session-accessor-boundary.mjs @@ -164,6 +164,7 @@ export const migratedSessionAccessorWriteFiles = new Set([ "src/commands/tasks.ts", "src/config/sessions/cleanup-service.ts", "src/gateway/server-node-events.ts", + "src/gateway/session-compaction-checkpoints.ts", "src/plugins/host-hook-cleanup.ts", "src/plugins/host-hook-state.ts", "src/tui/embedded-backend.ts", diff --git a/src/config/sessions/session-accessor.test.ts b/src/config/sessions/session-accessor.test.ts index 794c3646ec66..b00fc3a7e2e3 100644 --- a/src/config/sessions/session-accessor.test.ts +++ b/src/config/sessions/session-accessor.test.ts @@ -11,6 +11,7 @@ import { appendTranscriptEvent, applySessionEntryLifecycleMutation, applySessionPatchProjection, + branchSessionFromCompactionCheckpoint, canonicalizeSessionEntryAliases, cleanupSessionLifecycleArtifacts, commitReplySessionInitialization, @@ -28,6 +29,7 @@ import { readSessionUpdatedAt, replaceSessionEntry, resolveSessionEntryAccessTarget, + restoreSessionFromCompactionCheckpoint, resolveSessionTranscriptReadTarget, resolveSessionTranscriptRuntimeReadTarget, resolveSessionTranscriptRuntimeTarget, @@ -887,6 +889,134 @@ describe("session accessor file-backed seam", () => { }); }); + it("branches checkpoint sessions without exposing mutable store rows", async () => { + const sourceSessionId = "11111111-1111-4111-8111-111111111111"; + const branchSessionId = "22222222-2222-4222-8222-222222222222"; + const branchPath = path.join(tempDir, "branch.jsonl"); + const now = Date.now(); + fs.writeFileSync(transcriptPath, `{"type":"session","id":"${sourceSessionId}"}\n`, "utf8"); + fs.writeFileSync(branchPath, `{"type":"session","id":"${branchSessionId}"}\n`, "utf8"); + const checkpoint = { + checkpointId: "checkpoint-1", + sessionKey: "agent:main:main", + sessionId: sourceSessionId, + createdAt: now, + reason: "manual", + preCompaction: { + sessionId: sourceSessionId, + sessionFile: transcriptPath, + leafId: "leaf-1", + }, + postCompaction: { sessionId: "33333333-3333-4333-8333-333333333333" }, + } satisfies NonNullable[number]; + fs.writeFileSync( + storePath, + JSON.stringify( + { + main: { + label: "Main", + sessionFile: transcriptPath, + sessionId: sourceSessionId, + updatedAt: now, + compactionCheckpoints: [checkpoint], + }, + } satisfies Record, + null, + 2, + ), + "utf8", + ); + + const result = await branchSessionFromCompactionCheckpoint({ + storePath, + sourceKey: "agent:main:main", + sourceStoreKey: "main", + nextKey: "agent:main:branch", + checkpointId: "checkpoint-1", + forkTranscriptFromCheckpoint: async (selectedCheckpoint) => { + expect(selectedCheckpoint).toEqual(checkpoint); + return { + status: "created", + transcript: { + sessionFile: branchPath, + sessionId: branchSessionId, + totalTokens: 42, + }, + }; + }, + buildEntry: ({ currentEntry, forkedTranscript }) => ({ + ...currentEntry, + sessionFile: forkedTranscript.sessionFile, + sessionId: forkedTranscript.sessionId, + totalTokens: forkedTranscript.totalTokens, + updatedAt: now + 1, + }), + }); + + expect(result).toMatchObject({ + status: "created", + key: "agent:main:branch", + entry: { + sessionFile: branchPath, + sessionId: branchSessionId, + totalTokens: 42, + }, + }); + expect(loadSessionStore(storePath)).toEqual({ + main: expect.objectContaining({ sessionId: sourceSessionId }), + "agent:main:branch": expect.objectContaining({ + sessionFile: branchPath, + sessionId: branchSessionId, + totalTokens: 42, + }), + }); + }); + + it("does not persist checkpoint restores when the transcript boundary is missing", async () => { + fs.writeFileSync( + storePath, + JSON.stringify( + { + "agent:main:main": { + sessionId: "session-1", + updatedAt: 10, + compactionCheckpoints: [ + { + checkpointId: "checkpoint-1", + sessionKey: "agent:main:main", + sessionId: "session-1", + createdAt: 20, + reason: "manual", + preCompaction: { + sessionId: "session-1", + leafId: "leaf-1", + }, + postCompaction: { sessionId: "session-2" }, + }, + ], + }, + } satisfies Record, + null, + 2, + ), + "utf8", + ); + + const before = fs.readFileSync(storePath, "utf8"); + const result = await restoreSessionFromCompactionCheckpoint({ + storePath, + sessionKey: "agent:main:main", + checkpointId: "checkpoint-1", + forkTranscriptFromCheckpoint: async () => ({ status: "missing-boundary" }), + buildEntry: () => { + throw new Error("missing boundary should skip entry replacement"); + }, + }); + + expect(result).toEqual({ status: "missing-boundary" }); + expect(fs.readFileSync(storePath, "utf8")).toBe(before); + }); + it("cleans scoped lifecycle entries and unreferenced transcript artifacts", async () => { const nowMs = Date.now(); const oldDate = new Date(nowMs - 600_000); diff --git a/src/config/sessions/session-accessor.ts b/src/config/sessions/session-accessor.ts index 444208e743cd..f9529f249ad7 100644 --- a/src/config/sessions/session-accessor.ts +++ b/src/config/sessions/session-accessor.ts @@ -98,7 +98,7 @@ import { resolveOwnedSessionTranscriptWriteLockRunner, withOwnedSessionTranscriptWrites, } from "./transcript-write-context.js"; -import type { SessionEntry } from "./types.js"; +import type { SessionCompactionCheckpoint, SessionEntry } from "./types.js"; /** * Session access API for callers that need entries or transcripts without @@ -466,6 +466,81 @@ export type RestartRecoveryLifecycleUpdate = { replacements?: Iterable; }; +/** File-backed checkpoint transcript fork produced by the checkpoint storage boundary. */ +export type SessionCompactionCheckpointForkedTranscript = { + sessionFile: string; + sessionId: string; + totalTokens?: number; +}; + +/** Result of resolving and copying checkpoint transcript content for branch/restore. */ +export type SessionCompactionCheckpointTranscriptForkResult = + | { status: "created"; transcript: SessionCompactionCheckpointForkedTranscript } + | { status: "missing-boundary" } + | { status: "failed" }; + +/** Result of applying a checkpoint branch or restore mutation to session storage. */ +export type SessionCompactionCheckpointMutationResult = + | { + status: "created"; + key: string; + checkpoint: SessionCompactionCheckpoint; + entry: SessionEntry; + } + | { status: "missing-session" } + | { status: "missing-checkpoint" } + | { status: "missing-boundary" } + | { status: "failed" }; + +export type SessionCompactionCheckpointEntryBuildContext = { + /** Checkpoint row selected from the current persisted session entry. */ + checkpoint: SessionCompactionCheckpoint; + /** Persisted entry that owns the selected checkpoint. */ + currentEntry: SessionEntry; + /** Forked transcript identity created from the stored checkpoint boundary. */ + forkedTranscript: SessionCompactionCheckpointForkedTranscript; +}; + +export type SessionCompactionCheckpointTranscriptForker = ( + checkpoint: SessionCompactionCheckpoint, +) => Promise; + +export type SessionCompactionCheckpointEntryBuilder = ( + context: SessionCompactionCheckpointEntryBuildContext, +) => Promise | SessionEntry; + +export type BranchSessionFromCompactionCheckpointParams = { + /** Checkpoint id stored on the source session entry. */ + checkpointId: string; + /** Builds the branched session entry from the forked transcript. */ + buildEntry: SessionCompactionCheckpointEntryBuilder; + /** Copies transcript content through the stored checkpoint boundary. */ + forkTranscriptFromCheckpoint: SessionCompactionCheckpointTranscriptForker; + /** Persisted key for the new checkpoint branch. */ + nextKey: string; + /** Canonical key used as the branch parent. */ + sourceKey: string; + /** Actual persisted key to read when a legacy alias still owns the row. */ + sourceStoreKey?: string; + /** Explicit store target for file-backed stores and SQLite migration adapters. */ + storePath: string; +}; + +export type RestoreSessionFromCompactionCheckpointParams = { + /** Checkpoint id stored on the current session entry. */ + checkpointId: string; + /** Builds the restored session entry from the forked transcript. */ + buildEntry: SessionCompactionCheckpointEntryBuilder; + /** Copies transcript content through the stored checkpoint boundary. */ + forkTranscriptFromCheckpoint: SessionCompactionCheckpointTranscriptForker; + /** Canonical key to replace with the restored checkpoint state. */ + sessionKey: string; + /** Actual persisted key to read when a legacy alias still owns the row. */ + sessionStoreKey?: string; + /** Explicit store target for file-backed stores and SQLite migration adapters. */ + storePath: string; +}; + export type SessionEntryCreateWithTranscriptContext = { /** Current entry under the requested key before creation, if any. */ existingEntry?: SessionEntry; @@ -1164,6 +1239,103 @@ function applySessionAbortCutoff( entry.abortCutoffTimestamp = cutoff?.timestamp; } +function findSessionCompactionCheckpoint(params: { + checkpointId: string; + entry: SessionEntry; +}): SessionCompactionCheckpoint | undefined { + const checkpointId = params.checkpointId.trim(); + if (!checkpointId || !Array.isArray(params.entry.compactionCheckpoints)) { + return undefined; + } + return [...params.entry.compactionCheckpoints] + .toSorted((a, b) => b.createdAt - a.createdAt) + .find((checkpoint) => checkpoint.checkpointId === checkpointId); +} + +type ApplySessionCompactionCheckpointMutationParams = { + buildEntry: SessionCompactionCheckpointEntryBuilder; + checkpointId: string; + forkTranscriptFromCheckpoint: SessionCompactionCheckpointTranscriptForker; + readKey: string; + storePath: string; + writeKey: string; +}; + +async function applySessionCompactionCheckpointMutation( + params: ApplySessionCompactionCheckpointMutationParams, +): Promise { + return await updateSessionStore( + params.storePath, + async (store) => { + const currentEntry = store[params.readKey]; + if (!currentEntry?.sessionId) { + return { status: "missing-session" }; + } + const checkpoint = findSessionCompactionCheckpoint({ + entry: currentEntry, + checkpointId: params.checkpointId, + }); + if (!checkpoint) { + return { status: "missing-checkpoint" }; + } + const forkedSession = await params.forkTranscriptFromCheckpoint(checkpoint); + if (forkedSession.status !== "created") { + return forkedSession; + } + + const nextEntry = await params.buildEntry({ + checkpoint, + currentEntry, + forkedTranscript: forkedSession.transcript, + }); + store[params.writeKey] = nextEntry; + return { + status: "created", + key: params.writeKey, + checkpoint, + entry: nextEntry, + }; + }, + { skipSaveWhenResult: (result) => result.status !== "created" }, + ); +} + +/** + * Forks checkpoint transcript content and persists a new branch entry in one + * storage-sized mutation. SQLite adapters implement the transcript row copy + * and `session_entries.entry_json` insert inside the same write transaction. + */ +export async function branchSessionFromCompactionCheckpoint( + params: BranchSessionFromCompactionCheckpointParams, +): Promise { + return await applySessionCompactionCheckpointMutation({ + buildEntry: params.buildEntry, + checkpointId: params.checkpointId, + forkTranscriptFromCheckpoint: params.forkTranscriptFromCheckpoint, + readKey: params.sourceStoreKey ?? params.sourceKey, + storePath: params.storePath, + writeKey: params.nextKey, + }); +} + +/** + * Forks checkpoint transcript content and replaces the current entry in one + * storage-sized mutation. SQLite adapters implement the transcript row copy + * and `session_entries.entry_json` update inside the same write transaction. + */ +export async function restoreSessionFromCompactionCheckpoint( + params: RestoreSessionFromCompactionCheckpointParams, +): Promise { + return await applySessionCompactionCheckpointMutation({ + buildEntry: params.buildEntry, + checkpointId: params.checkpointId, + forkTranscriptFromCheckpoint: params.forkTranscriptFromCheckpoint, + readKey: params.sessionStoreKey ?? params.sessionKey, + storePath: params.storePath, + writeKey: params.sessionKey, + }); +} + /** * Applies a session patch projection through the accessor boundary. * The resolver sees a read-only snapshot and names the persisted key set; the diff --git a/src/gateway/session-compaction-checkpoints.test.ts b/src/gateway/session-compaction-checkpoints.test.ts index 1c7cfb44525f..18f1bef67514 100644 --- a/src/gateway/session-compaction-checkpoints.test.ts +++ b/src/gateway/session-compaction-checkpoints.test.ts @@ -930,6 +930,37 @@ describe("session-compaction-checkpoints", () => { expect(nextStore[MAIN_SESSION_KEY]?.compactionCheckpoints).toBeUndefined(); }); + test("persist skips malformed session rows without synthesizing a session id", async () => { + const { storePath, sessionId, sessionKey, now } = await makeTempSessionStore( + "openclaw-checkpoint-malformed-row-", + ); + await writeSessionStore(storePath, sessionKey, { + sessionId: "", + updatedAt: now, + }); + + const stored = await persistMainCheckpoint(storePath, { + sessionId, + snapshot: { + sessionId, + leafId: "pre-leaf", + }, + postSessionFile: path.join(path.dirname(storePath), "sess.compacted.jsonl"), + postLeafId: "post-leaf", + createdAt: now, + }); + + expect(stored).toBeNull(); + const nextStore = await readSessionStore<{ + compactionCheckpoints?: unknown[]; + sessionId?: string; + }>(storePath); + expect(nextStore[MAIN_SESSION_KEY]).toEqual({ + sessionId: "", + updatedAt: now, + }); + }); + test("persist trims retained checkpoint snapshots by total byte budget", async () => { const { dir, storePath, sessionId, sessionKey, now } = await makeTempSessionStore( "openclaw-checkpoint-byte-trim-", diff --git a/src/gateway/session-compaction-checkpoints.ts b/src/gateway/session-compaction-checkpoints.ts index b24d2a269791..90898ea01910 100644 --- a/src/gateway/session-compaction-checkpoints.ts +++ b/src/gateway/session-compaction-checkpoints.ts @@ -8,7 +8,6 @@ import { SessionManager, type FileEntry as SessionFileEntry, } from "../agents/sessions/session-manager.js"; -import { updateSessionStore } from "../config/sessions.js"; import type { SessionCompactionCheckpoint, SessionCompactionCheckpointReason, @@ -16,6 +15,12 @@ import type { } from "../config/sessions.js"; import { isCompactionCheckpointTranscriptFileName } from "../config/sessions/artifacts.js"; import { readFileRangeAsync } from "../config/sessions/file-range.js"; +import { + branchSessionFromCompactionCheckpoint, + restoreSessionFromCompactionCheckpoint, + type SessionCompactionCheckpointMutationResult, + updateSessionEntry, +} from "../config/sessions/session-accessor.js"; import { streamSessionTranscriptLines } from "../config/sessions/transcript-stream.js"; import { scanSessionTranscriptTree } from "../config/sessions/transcript-tree.js"; import { CURRENT_SESSION_VERSION } from "../config/sessions/version.js"; @@ -66,17 +71,7 @@ export type CompactionCheckpointTranscriptForkResult = | { status: "missing-boundary" } | { status: "failed" }; -export type CompactionCheckpointSessionMutationResult = - | { - status: "created"; - key: string; - checkpoint: SessionCompactionCheckpoint; - entry: SessionEntry; - } - | { status: "missing-session" } - | { status: "missing-checkpoint" } - | { status: "missing-boundary" } - | { status: "failed" }; +export type CompactionCheckpointSessionMutationResult = SessionCompactionCheckpointMutationResult; export type BranchCheckpointSessionParams = { storePath: string; @@ -583,30 +578,19 @@ function cloneCheckpointSessionEntry(params: { async function branchCheckpointSessionFromStoredBoundary( params: BranchCheckpointSessionParams, ): Promise { - return await updateSessionStore( - params.storePath, - async (store) => { - const currentEntry = store[params.sourceStoreKey ?? params.sourceKey]; - if (!currentEntry?.sessionId) { - return { status: "missing-session" }; - } - const checkpoint = getSessionCompactionCheckpoint({ - entry: currentEntry, - checkpointId: params.checkpointId, - }); - if (!checkpoint) { - return { status: "missing-checkpoint" }; - } - const forkedSession = await forkCheckpointTranscriptFromStoredBoundary({ checkpoint }); - if (forkedSession.status !== "created") { - return forkedSession; - } - - const forkedTranscript = forkedSession.transcript; + return await branchSessionFromCompactionCheckpoint({ + storePath: params.storePath, + sourceKey: params.sourceKey, + nextKey: params.nextKey, + checkpointId: params.checkpointId, + ...(params.sourceStoreKey ? { sourceStoreKey: params.sourceStoreKey } : {}), + forkTranscriptFromCheckpoint: async (checkpoint) => + await forkCheckpointTranscriptFromStoredBoundary({ checkpoint }), + buildEntry: ({ currentEntry, forkedTranscript }) => { const label = currentEntry.label?.trim() ? `${currentEntry.label.trim()} (checkpoint)` : "Checkpoint branch"; - const nextEntry = cloneCheckpointSessionEntry({ + return cloneCheckpointSessionEntry({ currentEntry, nextSessionId: forkedTranscript.sessionId, nextSessionFile: forkedTranscript.sessionFile, @@ -614,58 +598,29 @@ async function branchCheckpointSessionFromStoredBoundary( parentSessionKey: params.sourceKey, totalTokens: forkedTranscript.totalTokens, }); - store[params.nextKey] = nextEntry; - return { - status: "created", - key: params.nextKey, - checkpoint, - entry: nextEntry, - }; }, - { skipSaveWhenResult: (result) => result.status !== "created" }, - ); + }); } async function restoreCheckpointSessionFromStoredBoundary( params: RestoreCheckpointSessionParams, ): Promise { - return await updateSessionStore( - params.storePath, - async (store) => { - const currentEntry = store[params.sessionStoreKey ?? params.sessionKey]; - if (!currentEntry?.sessionId) { - return { status: "missing-session" }; - } - const checkpoint = getSessionCompactionCheckpoint({ - entry: currentEntry, - checkpointId: params.checkpointId, - }); - if (!checkpoint) { - return { status: "missing-checkpoint" }; - } - const restoredSession = await forkCheckpointTranscriptFromStoredBoundary({ checkpoint }); - if (restoredSession.status !== "created") { - return restoredSession; - } - - const restoredTranscript = restoredSession.transcript; - const nextEntry = cloneCheckpointSessionEntry({ + return await restoreSessionFromCompactionCheckpoint({ + storePath: params.storePath, + sessionKey: params.sessionKey, + checkpointId: params.checkpointId, + ...(params.sessionStoreKey ? { sessionStoreKey: params.sessionStoreKey } : {}), + forkTranscriptFromCheckpoint: async (checkpoint) => + await forkCheckpointTranscriptFromStoredBoundary({ checkpoint }), + buildEntry: ({ currentEntry, forkedTranscript }) => + cloneCheckpointSessionEntry({ currentEntry, - nextSessionId: restoredTranscript.sessionId, - nextSessionFile: restoredTranscript.sessionFile, - totalTokens: restoredTranscript.totalTokens, + nextSessionId: forkedTranscript.sessionId, + nextSessionFile: forkedTranscript.sessionFile, + totalTokens: forkedTranscript.totalTokens, preserveCompactionCheckpoints: true, - }); - store[params.sessionKey] = nextEntry; - return { - status: "created", - key: params.sessionKey, - checkpoint, - entry: nextEntry, - }; - }, - { skipSaveWhenResult: (result) => result.status !== "created" }, - ); + }), + }); } /** @@ -818,31 +773,35 @@ export async function persistSessionCompactionCheckpoint( }, }; - let stored = false; let trimmedCheckpoints: | { kept: SessionCompactionCheckpoint[] | undefined; removed: SessionCompactionCheckpoint[]; } | undefined; - await updateSessionStore(target.storePath, async (store) => { - const existing = store[target.canonicalKey]; - if (!existing?.sessionId) { - return; - } - const checkpoints = sessionStoreCheckpoints(existing); - checkpoints.push(checkpoint); - const snapshotBytesByPath = await statCheckpointSnapshotBytes(checkpoints); - trimmedCheckpoints = trimSessionCheckpoints(checkpoints, snapshotBytesByPath); - store[target.canonicalKey] = { - ...existing, - updatedAt: Math.max(existing.updatedAt ?? 0, createdAt), - compactionCheckpoints: trimmedCheckpoints.kept, - }; - stored = true; - }); + let stored = false; + const updatedEntry = await updateSessionEntry( + { + storePath: target.storePath, + sessionKey: target.canonicalKey, + }, + async (existing) => { + if (!existing.sessionId) { + return null; + } + const checkpoints = sessionStoreCheckpoints(existing); + checkpoints.push(checkpoint); + const snapshotBytesByPath = await statCheckpointSnapshotBytes(checkpoints); + trimmedCheckpoints = trimSessionCheckpoints(checkpoints, snapshotBytesByPath); + stored = true; + return { + updatedAt: Math.max(existing.updatedAt ?? 0, createdAt), + compactionCheckpoints: trimmedCheckpoints.kept, + }; + }, + ); - if (!stored) { + if (!updatedEntry || !stored) { log.warn("skipping compaction checkpoint persist: session not found", { sessionKey: params.sessionKey, }); diff --git a/test/scripts/check-session-accessor-boundary.test.ts b/test/scripts/check-session-accessor-boundary.test.ts index 2eba8e16279c..7177ad6bc05c 100644 --- a/test/scripts/check-session-accessor-boundary.test.ts +++ b/test/scripts/check-session-accessor-boundary.test.ts @@ -122,6 +122,7 @@ describe("session accessor boundary guard", () => { "src/commands/tasks.ts", "src/config/sessions/cleanup-service.ts", "src/gateway/server-node-events.ts", + "src/gateway/session-compaction-checkpoints.ts", "src/plugins/host-hook-cleanup.ts", "src/plugins/host-hook-state.ts", "src/tui/embedded-backend.ts",