Skip to content

feat(sdk): persist compaction and injected context through the transcript storage state #4894

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ericallam wants to merge 1 commit into from
+611 −16
Draft

feat(sdk): persist compaction and injected context through the transcript storage state #4894

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 91 additions & 3 deletions packages/trigger-sdk/src/v3/ai.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,10 +76,29 @@ import {
createTranscriptShadow,
defaultStorage,
diffTranscript,
parseTranscriptRuntimeState,
prefixFingerprint,
restoreModelLane,
type TranscriptChange,
type TranscriptChangeReason,
type TranscriptRuntimeState,
type TranscriptShadow,
type TranscriptStorage,
type TranscriptStorageContext,
} from "./transcriptStorage.js";

let transcriptStorageOverride: TranscriptStorage<unknown> | undefined;

/**
* Test-only override for the storage `chat.agent` persists through, so a
* test can capture the exact changesets the runtime produces.
* @internal
*/
export function __setTranscriptStorageForTests(
storage: TranscriptStorage<unknown> | undefined
): void {
transcriptStorageOverride = storage;
}
import {
type ChatInputChunk,
type ChatTaskWirePayload,
Expand Down Expand Up @@ -6870,8 +6889,18 @@ function chatAgent<
// collectively cost ~600ms on every first-message TTFC. Both reads
// swallow errors internally; the agent stays available either way.
const sessionIdForSnapshot = payload.sessionId ?? payload.chatId;
const transcriptStorage = defaultStorage;
const transcriptStorage = transcriptStorageOverride ?? defaultStorage;
let transcriptShadow: TranscriptShadow = createTranscriptShadow([]);
let bootTranscriptState: unknown = null;
/**
* True while the model lane holds a compaction summary, so it cannot be
* rebuilt from the transcript and has to be persisted as state. Reset
* wherever the lane is reconverted from the UI lane.
*/
let laneCompacted = false;
/** Conversational `chat.inject` messages in the lane, anchored to the transcript. */
let laneInjections: NonNullable<TranscriptRuntimeState["injections"]> = [];
let persistedStateSet = false;
let bootSnapshot:
| { messages: TUIMessage[]; lastOutEventId?: string; lastInEventId?: string }
| undefined;
Expand Down Expand Up @@ -6919,6 +6948,23 @@ function chatAgent<
const { changes, shadow } = diffTranscript(transcriptShadow, opts.messages, {
nonFinalIds: opts.nonFinalIds,
});
const lastId = opts.messages.at(-1)?.id;
const runtimeState: TranscriptRuntimeState | null =
laneCompacted && lastId !== undefined
? {
v: 1,
compaction: {
modelMessages: accumulatedMessages,
throughId: lastId,
fingerprint: prefixFingerprint(shadow, lastId),
},
}
: laneInjections.length > 0
? { v: 1, injections: laneInjections }
: null;
if (runtimeState !== null || persistedStateSet) {
changes.push({ op: "state", value: runtimeState } satisfies TranscriptChange);
}
const inCursor = chatInputRouter().resumeFloor();
await transcriptStorage.save(
{
Expand All @@ -6939,6 +6985,7 @@ function chatAgent<
}
);
transcriptShadow = shadow;
persistedStateSet = runtimeState !== null;
};

/**
Expand Down Expand Up @@ -7023,6 +7070,8 @@ function chatAgent<
clientData: bootClientData,
});
transcriptShadow = createTranscriptShadow(loaded.messages);
bootTranscriptState = loaded.state;
persistedStateSet = loaded.state !== null && loaded.state !== undefined;
bootSnapshot = {
messages: loaded.messages,
lastOutEventId: loaded.cursors?.lastOutEventId,
Expand Down Expand Up @@ -7367,7 +7416,14 @@ function chatAgent<
}
}
try {
accumulatedMessages = await toModelMessages(accumulatedUIMessages);
const restored = await restoreModelLane(
accumulatedUIMessages,
parseTranscriptRuntimeState(bootTranscriptState),
(messages) => toModelMessages(messages)
);
accumulatedMessages = restored.messages;
laneCompacted = restored.compacted;
laneInjections = restored.injections;
} catch (error) {
logger.warn("chat.agent: toModelMessages failed at boot; starting empty", {
error: error instanceof Error ? error.message : String(error),
Expand Down Expand Up @@ -8039,6 +8095,8 @@ function chatAgent<
);
accumulatedUIMessages = [...hydrated] as TUIMessage[];
accumulatedMessages = await toModelMessages(hydrated);
laneCompacted = false;
laneInjections = [];
locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
}

Expand Down Expand Up @@ -8076,6 +8134,8 @@ function chatAgent<
locals.set(chatOverrideMessagesKey, undefined);
accumulatedUIMessages = [...actionOverride] as TUIMessage[];
accumulatedMessages = await toModelMessages(actionOverride);
laneCompacted = false;
laneInjections = [];
locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);

actionChangedHistory = true;
Expand Down Expand Up @@ -8208,6 +8268,8 @@ function chatAgent<

accumulatedUIMessages = merged;
accumulatedMessages = await toModelMessages(merged);
laneCompacted = false;
laneInjections = [];
locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);

// Track new messages for onTurnComplete.newUIMessages.
Expand Down Expand Up @@ -8257,6 +8319,8 @@ function chatAgent<
accumulatedUIMessages.pop();
}
accumulatedMessages = await toModelMessages(accumulatedUIMessages);
laneCompacted = false;
laneInjections = [];
} else if (cleanedUIMessages.length > 0) {
// Submit-message (and the special-cased
// handover-prepare → submit-message rewrite earlier in
Expand Down Expand Up @@ -8310,6 +8374,8 @@ function chatAgent<
"chat.agent: replaced message not found at the model lane tail; reconverting the lane"
);
accumulatedMessages = await toModelMessages(accumulatedUIMessages);
laneCompacted = false;
laneInjections = [];
}
} else {
const incomingModelMessages = await toModelMessages(cleanedUIMessages);
Expand Down Expand Up @@ -8491,6 +8557,8 @@ function chatAgent<
locals.set(chatOverrideMessagesKey, undefined);
accumulatedUIMessages = [...turnStartOverride] as TUIMessage[];
accumulatedMessages = await toModelMessages(turnStartOverride);
laneCompacted = false;
laneInjections = [];
locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
}
},
Expand Down Expand Up @@ -8556,7 +8624,12 @@ function chatAgent<
const lastAccumulated = accumulatedMessages[accumulatedMessages.length - 1];
const bgQueue = locals.get(chatBackgroundQueueKey);
if (bgQueue && bgQueue.length > 0 && lastAccumulated?.role !== "tool") {
accumulatedMessages.push(...bgQueue.splice(0));
const injected = bgQueue.splice(0);
accumulatedMessages.push(...injected);
laneInjections.push({
afterId: accumulatedUIMessages.at(-1)?.id ?? "",
messages: injected,
});
Comment on lines +8627 to +8632

@coderabbitai coderabbitai Bot Sep 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift

Background injections drained through the inner-loop prepareStep path are never tracked, so they are silently lost from accumulatedMessages and never persisted.

This drain only runs when the last accumulated model message is not a tool message. The comment on this block states the fallback: when the tail is a tool message, the queued messages are picked up instead by toStreamTextOptions()'s own auto-injected prepareStep (the background-context-injection step further up in this file). That sibling drain returns the injected content only inside resultMessages for the current streamText step. It does not push into accumulatedMessages and does not push an entry into laneInjections.

Two consequences follow from the missing sync:

  • Within the same run, accumulatedMessages no longer reflects what the model actually saw for that step. Outer-loop compaction, onTurnComplete.messages/newMessages, and any later reconversion checkpoint all operate on a lane that is missing the injected content.
  • On a continuation boot, restoreModelLane has no record of that injection (it was never added to laneInjections, and it is not part of the transcript, so it cannot be recovered by reconverting UI messages either). This directly contradicts the PR's stated goal that "Injected messages are stored with anchors to the transcript messages they followed and reinserted on continuation" — for this specific, already-documented fallback case, they are not.

This is reachable whenever a chat.inject() call queues conversational content while the accumulator's last model message is a tool message (for example, right after a HITL tool-result exchange), and the agent's turn involves multiple streamText steps (tool calls).

Fix this by reconciling the inner-loop background-injection drain the same way drainSteeringQueue/reconcilePendingSteer reconcile steering messages: record what was actually drained and its anchor, then merge it back into accumulatedMessages and laneInjections once the turn's userRun/streamText call returns (before the turn's snapshot is written). For example, in the toStreamTextOptions() prepareStep step handling background context (around the block that does const injected = bgQueue.splice(0); resultMessages = [...(resultMessages ?? messages), ...injected];), also stash injected alongside the current tail id so the caller can merge it back:

// toStreamTextOptions()'s prepareStep, background-context step:
const bgQueue = locals.get(chatBackgroundQueueKey);
if (bgQueue && bgQueue.length > 0) {
  const injected = bgQueue.splice(0);
  resultMessages = [...(resultMessages ?? messages), ...injected];
  const pendingBg = locals.get(chatPendingBackgroundInjectionKey) ?? [];
  pendingBg.push({ afterId: <current UI tail id>, messages: injected });
  locals.set(chatPendingBackgroundInjectionKey, pendingBg);
}

Then, after userRun/streamText returns, merge chatPendingBackgroundInjectionKey into accumulatedMessages and laneInjections the same way the existing pre-run drain does at lines 8627-8632.

None of the new tests exercise this fallback (they only cover the pre-run drain, triggered from onTurnComplete). Consider adding a test that forces the accumulator's tail to be a tool message before queuing a chat.inject() call, to lock in the fix.

}

if (isHeadStartFinalTurn) {
Expand Down Expand Up @@ -8749,6 +8822,8 @@ function chatAgent<
accumulatedMessages = await toModelMessages(
runOverride.filter((m) => !pendingIds.has(m.id))
);
laneCompacted = false;
laneInjections = [];
locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
}

Expand All @@ -8774,6 +8849,8 @@ function chatAgent<
accumulatedMessages = taskCompactionConfig?.compactModelMessages
? await taskCompactionConfig.compactModelMessages(compactEvent)
: modelOnlyOverride;
laneCompacted = true;
laneInjections = [];

// Apply UI messages: callback or default (preserve all)
if (taskCompactionConfig?.compactUIMessages) {
Expand Down Expand Up @@ -8866,6 +8943,8 @@ function chatAgent<
"chat.agent: replaced response not found at the model lane tail; reconverting the lane"
);
accumulatedMessages = await toModelMessages(accumulatedUIMessages);
laneCompacted = false;
laneInjections = [];
}
} else {
accumulatedMessages.push(...responseModelMessages);
Expand Down Expand Up @@ -8985,6 +9064,9 @@ function chatAgent<
},
];

laneCompacted = true;
laneInjections = [];

// UI messages: callback or default (preserve all)
if (outerCompaction.compactUIMessages) {
accumulatedUIMessages = (await outerCompaction.compactUIMessages(
Expand Down Expand Up @@ -9089,6 +9171,8 @@ function chatAgent<
locals.set(chatOverrideMessagesKey, undefined);
accumulatedUIMessages = [...override] as TUIMessage[];
accumulatedMessages = await toModelMessages(override);
laneCompacted = false;
laneInjections = [];
locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
// Update event so onTurnComplete sees compacted messages
turnCompleteEvent.messages = accumulatedMessages;
Expand Down Expand Up @@ -9148,6 +9232,8 @@ function chatAgent<
locals.set(chatOverrideMessagesKey, undefined);
accumulatedUIMessages = [...turnCompleteOverride] as TUIMessage[];
accumulatedMessages = await toModelMessages(turnCompleteOverride);
laneCompacted = false;
laneInjections = [];
locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
}
},
Expand Down Expand Up @@ -9512,6 +9598,8 @@ function chatAgent<
"chat.agent: replaced partial not found at the model lane tail; reconverting the lane"
);
accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial);
laneCompacted = false;
laneInjections = [];
}
}
accumulatedUIMessages = erroredUIMessagesWithPartial;
Expand Down
Loading
Loading