fix: fence companion context snapshots

This commit is contained in:
Shakker
2026-08-12 03:41:55 +02:00
parent 0e78a64df5
commit 05d0d980d8
7 changed files with 217 additions and 39 deletions
@@ -350,6 +350,34 @@ describe("SQLite active transcript event projection", () => {
expect(readSessionTranscriptMessageEventById(scope, "old")).toBeDefined();
});
it("fails closed when the latest indexed reset payload is malformed", async () => {
await persistSessionTranscriptTurn(scope, {
messages: [
{ eventId: "old", parentId: null, message: { role: "user", content: "old" } },
{
eventId: "kept",
parentId: "old",
message: { role: "assistant", content: "kept" },
},
],
touchSessionEntry: false,
});
await appendTranscriptEvent(scope, {
type: "reset",
id: "reset-boundary",
parentId: "kept",
timestamp: "2026-08-12T00:00:00.000Z",
reason: "new",
firstKeptEntryId: "kept",
});
const database = openOpenClawAgentDatabase({ agentId: scope.agentId, env: scope.env });
database.db
.prepare("UPDATE transcript_events SET event_json = '{' WHERE session_id = ? AND seq = 3")
.run(scope.sessionId);
expect(() => readSessionTranscriptMessageEventCount(scope)).toThrow();
});
it("recomputes a cached reset window after a branch-changing message", async () => {
await persistSessionTranscriptTurn(scope, {
messages: [
@@ -14,6 +14,7 @@ import type {
TranscriptEvent,
} from "./session-accessor.sqlite-contract.js";
import {
readTranscriptProjectionGeneration,
readVisibleMessageRange,
resolveVisibleMessagePositionRange,
resolveVisibleMessagePositions,
@@ -58,6 +59,10 @@ export type SessionTranscriptMessageAnchorPage = SessionTranscriptMessageEventPa
export type SessionTranscriptBoundedMessageTailPage = SessionTranscriptMessageEventPage & {
scannedMessages: number;
serializedBytes: number;
snapshot: {
generation?: string;
indexedSeq: number;
};
};
function parseMessageEventRow(row: {
@@ -446,6 +451,10 @@ export function readSessionTranscriptBoundedMessageTailPage(
): SessionTranscriptBoundedMessageTailPage {
return withCurrentProjectionSnapshot(scope, (projection) => {
const visible = resolveVisibleMessagePositions(projection);
const snapshot = {
generation: readTranscriptProjectionGeneration(projection),
indexedSeq: projection.state.indexedSeq,
};
const totalMessages = visible.total;
const offset = Math.min(
Math.max(0, Math.floor(Number.isFinite(options.offset) ? options.offset : 0)),
@@ -468,6 +477,7 @@ export function readSessionTranscriptBoundedMessageTailPage(
events: [],
scannedMessages: positions.length,
serializedBytes: 0,
snapshot,
totalMessages,
};
}
@@ -521,6 +531,7 @@ export function readSessionTranscriptBoundedMessageTailPage(
events,
scannedMessages: positions.length,
serializedBytes,
snapshot,
totalMessages,
};
});
@@ -96,20 +96,13 @@ function readMessageRange(
).rows.map(parseMessageEventRow);
}
function parseTranscriptEventType(eventJson: string): string | undefined {
try {
const parsed = JSON.parse(eventJson) as { type?: unknown };
return typeof parsed.type === "string" ? parsed.type : undefined;
} catch {
return undefined;
}
}
function resetMessageWindowCacheKey(projection: ResetWindowProjection): string {
return `${projection.database.path}\0${projection.resolved.sessionId}`;
}
function readTranscriptGeneration(projection: ResetWindowProjection): string | undefined {
export function readTranscriptProjectionGeneration(
projection: ResetWindowProjection,
): string | undefined {
return executeSqliteQueryTakeFirstSync(
projection.database.db,
getResetWindowKysely(projection.database)
@@ -125,34 +118,70 @@ function cacheResetMessageWindow(key: string, entry: ResetMessageWindowCacheEntr
pruneMapToMaxSize(resetMessageWindowCache, MAX_RESET_MESSAGE_WINDOW_CACHE);
}
function readLatestActiveBoundaryMetadataByType(
projection: ResetWindowProjection,
eventType: "compaction" | "reset",
) {
const db = getResetWindowKysely(projection.database);
return executeSqliteQueryTakeFirstSync(
projection.database.db,
db
.selectFrom("session_transcript_active_events as active")
.innerJoin("transcript_event_identities as identity", (join) =>
join
.onRef("identity.session_id", "=", "active.session_id")
.onRef("identity.seq", "=", "active.event_seq"),
)
.select(["active.active_position", "identity.event_type", "identity.seq"])
.where("active.session_id", "=", projection.resolved.sessionId)
.where("identity.event_type", "=", eventType)
.orderBy("identity.seq", "desc")
.limit(1),
);
}
function readLatestActiveBoundaryMetadata(projection: ResetWindowProjection) {
const reset = readLatestActiveBoundaryMetadataByType(projection, "reset");
const compaction = readLatestActiveBoundaryMetadataByType(projection, "compaction");
if (!reset) {
return compaction;
}
if (!compaction) {
return reset;
}
return reset.seq > compaction.seq ? reset : compaction;
}
function readResetBoundary(projection: ResetWindowProjection, seq: number) {
const row = executeSqliteQueryTakeFirstSync(
projection.database.db,
getResetWindowKysely(projection.database)
.selectFrom("transcript_events")
.select("event_json")
.where("session_id", "=", projection.resolved.sessionId)
.where("seq", "=", seq)
.limit(1),
);
if (!row) {
throw new Error("Active transcript reset boundary is missing");
}
const parsed = JSON.parse(row.event_json) as { firstKeptEntryId?: unknown; type?: unknown };
if (parsed.type !== "reset") {
throw new Error("Active transcript reset boundary has invalid payload");
}
return parsed;
}
function findLatestResetMessageWindow(
projection: ResetWindowProjection,
generation: string | undefined,
): ResetMessageWindow | null {
const db = getResetWindowKysely(projection.database);
const nonMessageRows = executeSqliteQuerySync(
projection.database.db,
db
.selectFrom("session_transcript_active_events as active")
.innerJoin("transcript_events as event", (join) =>
join
.onRef("event.session_id", "=", "active.session_id")
.onRef("event.seq", "=", "active.event_seq"),
)
.select(["active.active_position", "event.event_json"])
.where("active.session_id", "=", projection.resolved.sessionId)
.where("active.message_position", "is", null)
.orderBy("active.active_position", "desc"),
).rows;
const latestBoundaryRow = nonMessageRows.find((row) => {
const type = parseTranscriptEventType(row.event_json);
return type === "reset" || type === "compaction";
});
if (!latestBoundaryRow || parseTranscriptEventType(latestBoundaryRow.event_json) !== "reset") {
const latestBoundary = readLatestActiveBoundaryMetadata(projection);
if (!latestBoundary || latestBoundary.event_type !== "reset") {
return null;
}
const resetRow = latestBoundaryRow;
const reset = JSON.parse(resetRow.event_json) as { firstKeptEntryId?: unknown };
const reset = readResetBoundary(projection, latestBoundary.seq);
const postBoundaryMessagePosition =
executeSqliteQueryTakeFirstSync(
projection.database.db,
@@ -160,7 +189,7 @@ function findLatestResetMessageWindow(
.selectFrom("session_transcript_active_events")
.select("message_position")
.where("session_id", "=", projection.resolved.sessionId)
.where("active_position", ">", resetRow.active_position)
.where("active_position", ">", latestBoundary.active_position)
.where("message_position", "is not", null)
.orderBy("active_position", "asc")
.limit(1),
@@ -180,7 +209,7 @@ function findLatestResetMessageWindow(
.where("identity.session_id", "=", projection.resolved.sessionId)
.where("identity.event_id", "=", reset.firstKeptEntryId),
);
if (firstKept && firstKept.active_position < resetRow.active_position) {
if (firstKept && firstKept.active_position < latestBoundary.active_position) {
keptMessagePositions = executeSqliteQuerySync(
projection.database.db,
db
@@ -193,7 +222,7 @@ function findLatestResetMessageWindow(
.select(["active.message_position", "event.event_json"])
.where("active.session_id", "=", projection.resolved.sessionId)
.where("active.active_position", ">=", firstKept.active_position)
.where("active.active_position", "<", resetRow.active_position)
.where("active.active_position", "<", latestBoundary.active_position)
.where("active.message_position", "is not", null)
.orderBy("active.active_position", "asc"),
).rows.flatMap((row) => {
@@ -221,7 +250,7 @@ function findLatestResetMessageWindow(
function resolveResetMessageWindow(projection: ResetWindowProjection): ResetMessageWindow | null {
const key = resetMessageWindowCacheKey(projection);
const cached = resetMessageWindowCache.get(key);
const generation = readTranscriptGeneration(projection);
const generation = readTranscriptProjectionGeneration(projection);
if (cached) {
if (cached.generation === generation && cached.indexedSeq === projection.state.indexedSeq) {
return cached.window;
@@ -5,6 +5,7 @@ import {
persistSessionTranscriptTurn,
upsertSessionEntryCore,
} from "../config/sessions/session-accessor.js";
import * as activeTranscriptEvents from "../config/sessions/session-accessor.sqlite-active-events.js";
import {
closeOpenClawAgentDatabasesForTest,
openOpenClawAgentDatabase,
@@ -195,6 +196,44 @@ describe("session companion context", () => {
});
});
it("rejects context assembled across different transcript snapshots", async () => {
const scope = createScope("companion-context-snapshot-fence");
await upsertSessionEntryCore(scope, { sessionId: scope.sessionId, updatedAt: 1 });
const page = vi
.spyOn(activeTranscriptEvents, "readSessionTranscriptBoundedMessageTailPage")
.mockReturnValueOnce({
activeLeafEntryId: "leaf-1",
events: [
{
event: {
type: "message",
id: "message-1",
parentId: null,
message: { role: "user", content: "stable context", timestamp: 1 },
},
seq: 1,
},
],
scannedMessages: 1,
serializedBytes: 128,
snapshot: { generation: "generation-1", indexedSeq: 1 },
totalMessages: 1,
})
.mockReturnValueOnce({
activeLeafEntryId: "leaf-1",
events: [],
scannedMessages: 0,
serializedBytes: 0,
snapshot: { generation: "generation-2", indexedSeq: 1 },
totalMessages: 1,
});
await expect(defaultSessionCompanionContextReader.read(scope)).resolves.toEqual({
kind: "unavailable",
});
expect(page).toHaveBeenCalledTimes(2);
});
it("keeps transcript-visible messages across compaction", async () => {
const scope = createScope("companion-context-compaction");
await upsertSessionEntryCore(scope, { sessionId: scope.sessionId, updatedAt: 1 });
+38
View File
@@ -142,6 +142,14 @@ async function readSessionCompanionContext(params: {
let rawBytes = 0;
let scannedMessages = 0;
let totalMessages = 0;
let snapshot:
| {
activeLeafEntryId?: string | null;
generation?: string;
indexedSeq: number;
totalMessages: number;
}
| undefined;
let contextMessages: SessionCompanionContextMessage[] = [];
while (
contextMessages.length < CONTEXT_MAX_MESSAGES &&
@@ -158,6 +166,21 @@ async function readSessionCompanionContext(params: {
if (params.signal?.aborted || page.events.length !== page.scannedMessages) {
return { kind: "unavailable" };
}
const pageSnapshot = {
activeLeafEntryId: page.activeLeafEntryId,
generation: page.snapshot.generation,
indexedSeq: page.snapshot.indexedSeq,
totalMessages: page.totalMessages,
};
snapshot ??= pageSnapshot;
if (
pageSnapshot.activeLeafEntryId !== snapshot.activeLeafEntryId ||
pageSnapshot.generation !== snapshot.generation ||
pageSnapshot.indexedSeq !== snapshot.indexedSeq ||
pageSnapshot.totalMessages !== snapshot.totalMessages
) {
return { kind: "unavailable" };
}
totalMessages = page.totalMessages;
rawBytes += page.serializedBytes;
scannedMessages += page.scannedMessages;
@@ -173,6 +196,21 @@ async function readSessionCompanionContext(params: {
if (contextMessages.length < CONTEXT_MAX_MESSAGES && offset < totalMessages) {
return { kind: "unavailable" };
}
const fence = readSessionTranscriptBoundedMessageTailPage(scope, {
maxBytes: 0,
maxMessages: 0,
offset: 0,
});
if (
params.signal?.aborted ||
!snapshot ||
fence.activeLeafEntryId !== snapshot.activeLeafEntryId ||
fence.snapshot.generation !== snapshot.generation ||
fence.snapshot.indexedSeq !== snapshot.indexedSeq ||
fence.totalMessages !== snapshot.totalMessages
) {
return { kind: "unavailable" };
}
return {
kind: "ready",
context: {
+16 -1
View File
@@ -27,9 +27,14 @@ export type ChatSessionCompanionThread = {
};
type MutableCompanionThread = ChatSessionCompanionThread & {
failedQuestionKnownExchanges: ReadonlySet<string> | null;
revision: number;
};
function exchangeKey(exchange: SessionCompanionExchange): string {
return JSON.stringify([exchange.question, exchange.answer, exchange.ts]);
}
function errorDetailCode(error: unknown): string | null {
if (!error || typeof error !== "object") {
return null;
@@ -65,6 +70,7 @@ function createThread(): MutableCompanionThread {
exchanges: [],
pendingQuestion: null,
failedQuestion: null,
failedQuestionKnownExchanges: null,
hint: null,
retryable: false,
draft: "",
@@ -118,9 +124,14 @@ export class ChatSessionCompanionThreads {
}));
if (
thread.failedQuestion &&
thread.exchanges.some((exchange) => exchange.question === thread.failedQuestion)
thread.exchanges.some(
(exchange) =>
exchange.question === thread.failedQuestion &&
!thread.failedQuestionKnownExchanges?.has(exchangeKey(exchange)),
)
) {
thread.failedQuestion = null;
thread.failedQuestionKnownExchanges = null;
thread.hint = null;
thread.retryable = false;
}
@@ -152,11 +163,13 @@ export class ChatSessionCompanionThreads {
}
thread.pendingQuestion = normalized;
thread.failedQuestion = null;
thread.failedQuestionKnownExchanges = null;
thread.hint = null;
thread.retryable = false;
thread.draft = "";
thread.revision += 1;
const token = Symbol(key);
const knownExchanges = new Set(thread.exchanges.map(exchangeKey));
this.submissionTokens.set(key, token);
this.notify();
try {
@@ -168,11 +181,13 @@ export class ChatSessionCompanionThreads {
...thread.exchanges,
{ question: normalized, answer: result.answer, ts: result.ts },
].slice(-MAX_COMPANION_EXCHANGES);
thread.failedQuestionKnownExchanges = null;
} catch (error) {
if (this.submissionTokens.get(key) !== token) {
return;
}
thread.failedQuestion = normalized;
thread.failedQuestionKnownExchanges = knownExchanges;
const reason = errorDetailReason(error);
thread.hint =
errorDetailCode(error) === COMPANION_BUSY_DETAIL_CODE
+21 -3
View File
@@ -339,8 +339,11 @@ describe("ChatSessionCompanionThreads", () => {
});
});
it("hydrates an answer committed before a disconnect response was lost", async () => {
it("hydrates only a newly committed repeated question after a lost response", async () => {
const threads = new ChatSessionCompanionThreads();
await threads.hydrate("one", async () => ({
exchanges: [{ question: "What changed?", answer: "Earlier answer.", ts: 1 }],
}));
await threads.submit("one", "What changed?", async () => {
throw new Error("socket closed");
});
@@ -351,14 +354,29 @@ describe("ChatSessionCompanionThreads", () => {
});
await threads.hydrate("one", async () => ({
exchanges: [{ question: "What changed?", answer: "The fix committed.", ts: 4 }],
exchanges: [{ question: "What changed?", answer: "Earlier answer.", ts: 1 }],
}));
expect(threads.view("one")).toMatchObject({
failedQuestion: "What changed?",
hint: "unavailable",
retryable: true,
});
await threads.hydrate("one", async () => ({
exchanges: [
{ question: "What changed?", answer: "Earlier answer.", ts: 1 },
{ question: "What changed?", answer: "The fix committed.", ts: 4 },
],
}));
expect(threads.view("one")).toMatchObject({
failedQuestion: null,
hint: null,
retryable: false,
exchanges: [{ question: "What changed?", answer: "The fix committed.", ts: 4 }],
exchanges: [
{ question: "What changed?", answer: "Earlier answer.", ts: 1 },
{ question: "What changed?", answer: "The fix committed.", ts: 4 },
],
});
});