Skip to content

Commit ab1fa5a

Browse files
jalehmanjalehman
committed
refactor: add sqlite checkpoint session fork
1 parent 831e2f5 commit ab1fa5a

2 files changed

Lines changed: 494 additions & 2 deletions

File tree

‎src/config/sessions/session-accessor.conformance.test.ts‎

Lines changed: 140 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import {
3232
import {
3333
appendSqliteTranscriptEvent,
3434
appendSqliteTranscriptMessage,
35+
branchSqliteCompactionCheckpointSession,
3536
listSqliteSessionEntries,
3637
loadExactSqliteSessionEntry,
3738
loadSqliteSessionEntry,
@@ -41,11 +42,13 @@ import {
4142
persistSqliteSessionTranscriptTurn,
4243
publishSqliteTranscriptUpdate,
4344
readSqliteSessionUpdatedAt,
45+
replaceSqliteTranscriptEvents,
4446
replaceSqliteSessionEntry,
47+
restoreSqliteCompactionCheckpointSession,
4548
updateSqliteSessionEntry,
4649
upsertSqliteSessionEntry,
4750
} from "./session-accessor.sqlite.js";
48-
import type { SessionEntry } from "./types.js";
51+
import type { SessionCompactionCheckpoint, SessionEntry } from "./types.js";
4952

5053
type AccessorAdapter = {
5154
name: string;
@@ -770,4 +773,140 @@ describe("sqlite transcript turn persistence", () => {
770773
}),
771774
]);
772775
});
776+
777+
it("branches a checkpoint by copying rows and creating the branch entry atomically", async () => {
778+
const sourceScope = sqliteAdapter.transcriptScope(paths, "source-session");
779+
const sourceEntryScope = sqliteAdapter.entryScope(paths);
780+
const branchKey = "agent:main:checkpoint-branch";
781+
const checkpoint: SessionCompactionCheckpoint = {
782+
checkpointId: "checkpoint-1",
783+
sessionKey: sourceEntryScope.sessionKey,
784+
sessionId: "source-session",
785+
createdAt: Date.parse("2026-01-01T00:00:00.000Z"),
786+
reason: "manual",
787+
tokensBefore: 42,
788+
tokensAfter: 84,
789+
preCompaction: {
790+
sessionId: "pre-compaction-session",
791+
leafId: "pre-msg",
792+
},
793+
postCompaction: {
794+
sessionId: "source-session",
795+
entryId: "msg-2",
796+
},
797+
};
798+
799+
await replaceSqliteTranscriptEvents(sourceScope, [
800+
{ type: "session", id: "source-session", cwd: paths.tempDir },
801+
{ type: "message", id: "msg-1", parentId: null, message: { content: "one" } },
802+
{ type: "message", id: "msg-2", parentId: "msg-1", message: { content: "two" } },
803+
{ type: "message", id: "msg-3", parentId: "msg-2", message: { content: "three" } },
804+
]);
805+
await upsertSqliteSessionEntry(sourceEntryScope, {
806+
label: "Source",
807+
sessionId: "source-session",
808+
updatedAt: 10,
809+
compactionCheckpoints: [checkpoint],
810+
});
811+
812+
const result = await branchSqliteCompactionCheckpointSession({
813+
agentId: "main",
814+
env: sourceScope.env,
815+
storePath: paths.sqlitePath,
816+
sourceKey: sourceEntryScope.sessionKey,
817+
nextKey: branchKey,
818+
checkpointId: checkpoint.checkpointId,
819+
});
820+
if (result.status !== "created") {
821+
throw new Error(`expected branch creation, got ${result.status}`);
822+
}
823+
824+
const branchScope = {
825+
...sourceScope,
826+
sessionId: result.entry.sessionId,
827+
sessionKey: branchKey,
828+
};
829+
expect(loadSqliteSessionEntry({ ...sourceEntryScope, sessionKey: branchKey })).toEqual(
830+
result.entry,
831+
);
832+
expect(result.entry).toEqual(
833+
expect.objectContaining({
834+
label: "Source (checkpoint)",
835+
parentSessionKey: sourceEntryScope.sessionKey,
836+
sessionFile: expect.stringContaining(result.entry.sessionId),
837+
totalTokens: 84,
838+
totalTokensFresh: true,
839+
}),
840+
);
841+
await expect(loadSqliteTranscriptEvents(branchScope)).resolves.toEqual([
842+
expect.objectContaining({ type: "session", id: result.entry.sessionId }),
843+
expect.objectContaining({ id: "msg-1", type: "message" }),
844+
expect.objectContaining({ id: "msg-2", type: "message" }),
845+
]);
846+
});
847+
848+
it("restores a checkpoint by copying rows and replacing the session entry atomically", async () => {
849+
const sourceScope = sqliteAdapter.transcriptScope(paths, "source-session");
850+
const sourceEntryScope = sqliteAdapter.entryScope(paths);
851+
const checkpoint: SessionCompactionCheckpoint = {
852+
checkpointId: "checkpoint-restore",
853+
sessionKey: sourceEntryScope.sessionKey,
854+
sessionId: "current-session",
855+
createdAt: Date.parse("2026-01-01T00:00:00.000Z"),
856+
reason: "manual",
857+
tokensBefore: 12,
858+
tokensAfter: 24,
859+
preCompaction: {
860+
sessionId: "pre-compaction-session",
861+
leafId: "pre-msg",
862+
},
863+
postCompaction: {
864+
sessionId: "source-session",
865+
entryId: "msg-1",
866+
},
867+
};
868+
869+
await replaceSqliteTranscriptEvents(sourceScope, [
870+
{ type: "session", id: "source-session", cwd: paths.tempDir },
871+
{ type: "message", id: "msg-1", parentId: null, message: { content: "restore" } },
872+
{ type: "message", id: "msg-2", parentId: "msg-1", message: { content: "skip" } },
873+
]);
874+
await upsertSqliteSessionEntry(sourceEntryScope, {
875+
label: "Current",
876+
sessionId: "current-session",
877+
sessionFile: "sqlite:main:current-session",
878+
updatedAt: 10,
879+
compactionCheckpoints: [checkpoint],
880+
});
881+
882+
const result = await restoreSqliteCompactionCheckpointSession({
883+
agentId: "main",
884+
env: sourceScope.env,
885+
storePath: paths.sqlitePath,
886+
sessionKey: sourceEntryScope.sessionKey,
887+
checkpointId: checkpoint.checkpointId,
888+
});
889+
if (result.status !== "created") {
890+
throw new Error(`expected restore creation, got ${result.status}`);
891+
}
892+
893+
const restoredScope = {
894+
...sourceScope,
895+
sessionId: result.entry.sessionId,
896+
};
897+
expect(loadSqliteSessionEntry(sourceEntryScope)).toEqual(result.entry);
898+
expect(result.entry).toEqual(
899+
expect.objectContaining({
900+
label: "Current",
901+
compactionCheckpoints: [checkpoint],
902+
sessionFile: expect.stringContaining(result.entry.sessionId),
903+
totalTokens: 24,
904+
totalTokensFresh: true,
905+
}),
906+
);
907+
await expect(loadSqliteTranscriptEvents(restoredScope)).resolves.toEqual([
908+
expect.objectContaining({ type: "session", id: result.entry.sessionId }),
909+
expect.objectContaining({ id: "msg-1", type: "message" }),
910+
]);
911+
});
773912
});

0 commit comments

Comments
 (0)