Skip to content

Commit bf82221

Browse files
committed
fix(chat): append a steer to the model lane instead of rebuilding it
Reconciling the model lane by reconverting the UI lane assumed the UI lane is a superset of it. Compaction breaks that by design: it replaces the model lane with a summary and deliberately leaves the UI lane whole, so any reconversion restored every message the summary had replaced. Reproduced on both surfaces: the next turn was sent the full transcript with no summary, while the steer itself was present, which is what made a steer-presence check pass. The model lane is now only ever appended to. The stale flag becomes the claimed messages themselves, reconciled before the response is appended so the order stays steer-then-answer, and onto whatever the lane holds, summary included. absorbSteering appends to both lanes the same way.
1 parent 6683e4b commit bf82221

3 files changed

Lines changed: 285 additions & 36 deletions

File tree

‎packages/trigger-sdk/src/v3/ai.ts‎

Lines changed: 47 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -3563,17 +3563,22 @@ const chatSteeringQueueKey = locals.create<SteeringQueueEntry[]>("chat.steeringQ
35633563
const chatTurnNewUIMessagesKey = locals.create<UIMessage[]>("chat.turnNewUIMessages");
35643564

35653565
/**
3566-
* Set when a steering drain appended to the UI accumulator, so the model
3567-
* accumulator gets rebuilt from it at the end of the turn.
3566+
* Steering messages a drain consumed that the model accumulator has not been
3567+
* given yet.
35683568
*
35693569
* The two accumulators are maintained separately, and the model one is
35703570
* normally advanced by appending each turn's delta. A drained message is
35713571
* appended to the UI one but reaches the model only through the `prepareStep`
3572-
* return value, which is per-step: without a rebuild the model lane never
3573-
* learns the message exists and every later turn of the run answers without
3574-
* it, while the browser, the snapshot and `chat.history.*` all still show it.
3572+
* return value, which is per-step: without this the model lane never learns
3573+
* the message exists and every later turn of the run answers without it,
3574+
* while the browser, the snapshot and `chat.history.*` all still show it.
3575+
*
3576+
* Held as the messages rather than a "rebuild me" flag because the model lane
3577+
* can only be appended to, never reconstructed. Compaction replaces it with a
3578+
* summary and deliberately leaves the UI lane whole, so reconverting the UI
3579+
* lane restores every message the summary replaced.
35753580
*/
3576-
const chatModelLaneStaleKey = locals.create<boolean>("chat.modelLaneStale");
3581+
const chatPendingSteerKey = locals.create<UIMessage[]>("chat.pendingSteer");
35773582
/** @internal — IDs of messages that were successfully injected via prepareStep */
35783583
const chatInjectedMessageIdsKey = locals.create<Set<string>>("chat.injectedMessageIds");
35793584
/** @internal — non-transient data parts queued via chat.response or writer.write() for accumulation into the response message */
@@ -4315,7 +4320,11 @@ async function drainSteeringQueue(
43154320
}
43164321
}
43174322
if (claimedUIMessages.length > 0 && currentUIMessages) {
4318-
locals.set(chatModelLaneStaleKey, true);
4323+
const pendingSteer = locals.get(chatPendingSteerKey) ?? [];
4324+
for (const m of claimedUIMessages) {
4325+
if (!pendingSteer.some((existing) => existing.id === m.id)) pendingSteer.push(m);
4326+
}
4327+
locals.set(chatPendingSteerKey, pendingSteer);
43194328
}
43204329

43214330
// Write injection confirmation chunk to the stream so the frontend
@@ -8543,6 +8552,27 @@ function chatAgent<
85438552
// Determine if the user stopped generation this turn (not a full run cancel).
85448553
const wasStopped = stopController.signal.aborted && !runSignal.aborted;
85458554

8555+
// Give the model accumulator the steering messages the drain
8556+
// consumed. Appended, never reconverted from the UI lane, so a
8557+
// model-only compaction summary set just above survives; and done
8558+
// before the response is appended so the order stays
8559+
// steer-then-answer. Outside the `capturedResponseMessage`
8560+
// branches below, so a turn that captured no response is covered.
8561+
const pendingSteer = locals.get(chatPendingSteerKey);
8562+
if (pendingSteer && pendingSteer.length > 0) {
8563+
locals.set(chatPendingSteerKey, []);
8564+
try {
8565+
accumulatedMessages.push(
8566+
...(await toModelMessages(pendingSteer.map(stripProviderMetadata)))
8567+
);
8568+
} catch (error) {
8569+
logger.warn(
8570+
"chat.agent: toModelMessages failed for an injected message; it will be missing from the next turn",
8571+
{ error: error instanceof Error ? error.message : String(error) }
8572+
);
8573+
}
8574+
}
8575+
85468576
// Append the assistant's response (partial or complete) to the accumulator.
85478577
// The onFinish callback fires even on abort/stop, so partial responses
85488578
// from stopped generation are captured correctly.
@@ -8631,23 +8661,6 @@ function chatAgent<
86318661
turnBufferedChunks.length = 0;
86328662
}
86338663

8634-
// Bring the model accumulator back in line with the UI one
8635-
// after a steering drain. Placed after response accumulation
8636-
// and before compaction reads `accumulatedMessages`, and
8637-
// outside the `capturedResponseMessage` branches so a turn
8638-
// that captured no response is covered too.
8639-
if (locals.get(chatModelLaneStaleKey)) {
8640-
locals.set(chatModelLaneStaleKey, false);
8641-
try {
8642-
accumulatedMessages = await toModelMessages(accumulatedUIMessages);
8643-
} catch (error) {
8644-
logger.warn(
8645-
"chat.agent: toModelMessages failed rebuilding after an injection; the injected message will be missing from the next turn",
8646-
{ error: error instanceof Error ? error.message : String(error) }
8647-
);
8648-
}
8649-
}
8650-
86518664
if (runSignal.aborted) return "exit";
86528665

86538666
// Await deferred background work (e.g. DB writes from onTurnStart)
@@ -10669,20 +10682,18 @@ class ChatMessageAccumulator {
1066910682
* The drain only puts them in this step's prompt, so without this they
1067010683
* shape one answer and then exist in neither lane: not in `uiMessages`,
1067110684
* which is what an app persists from, and not in `modelMessages`, which is
10672-
* what every later turn sends. The UI lane is authoritative and the model
10673-
* lane is its conversion, matching how `chat.agent` reconciles the two.
10685+
* what every later turn sends.
10686+
*
10687+
* Both lanes are appended to. The model lane is never reconverted from the
10688+
* UI lane, because `compactIfNeeded` replaces it with a summary and leaves
10689+
* the UI lane whole: a reconversion would restore everything the summary
10690+
* replaced.
1067410691
*/
1067510692
async absorbSteering(claimed: UIMessage[]): Promise<void> {
10676-
if (claimed.length === 0) return;
10677-
let added = false;
10678-
for (const m of claimed) {
10679-
if (!this.uiMessages.some((existing) => existing.id === m.id)) {
10680-
this.uiMessages.push(m);
10681-
added = true;
10682-
}
10683-
}
10684-
if (!added) return;
10685-
this.modelMessages = await toModelMessages(this.uiMessages);
10693+
const fresh = claimed.filter((m) => !this.uiMessages.some((e) => e.id === m.id));
10694+
if (fresh.length === 0) return;
10695+
this.uiMessages.push(...fresh);
10696+
this.modelMessages.push(...(await toModelMessages(fresh)));
1068610697
}
1068710698

1068810699
/**
Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
import type { UIMessage } from "ai";
2+
import { describe, expect, it } from "vitest";
3+
import { chat } from "../src/v3/ai.js";
4+
5+
/**
6+
* `chat.MessageAccumulator` compaction is model-only: it replaces
7+
* `modelMessages` with a summary and leaves `uiMessages` whole so the chat can
8+
* still display the conversation. Recording a steer by reconverting
9+
* `modelMessages` from `uiMessages` therefore restores everything the summary
10+
* replaced, on the next steer after any compaction.
11+
*
12+
* Asserted directly on the accumulator: no run, no model, no harness, because
13+
* the whole question is which of its two lanes gets written and how.
14+
*/
15+
16+
const USAGE = {
17+
inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined },
18+
outputTokens: { total: 5, text: 5, reasoning: undefined },
19+
totalTokens: 15,
20+
};
21+
22+
const userMessage = (text: string, id: string): UIMessage =>
23+
({ id, role: "user", parts: [{ type: "text", text }] }) as UIMessage;
24+
25+
const assistantMessage = (text: string, id: string): UIMessage =>
26+
({ id, role: "assistant", parts: [{ type: "text", text }] }) as UIMessage;
27+
28+
const flatten = (messages: { content: unknown }[]) => JSON.stringify(messages);
29+
30+
describe("chat.MessageAccumulator steering after compaction", () => {
31+
it("keeps the summary in the model lane when a later steer is absorbed", async () => {
32+
const conversation = new chat.MessageAccumulator({
33+
compaction: {
34+
shouldCompact: () => true,
35+
summarize: async () => "SUMMARY-OF-EVERYTHING",
36+
},
37+
});
38+
39+
await conversation.addIncoming([userMessage("EARLY-SENTINEL", "u-1")], "submit-message", 0);
40+
await conversation.addResponse(assistantMessage("first answer", "a-1"));
41+
42+
const didCompact = await conversation.compactIfNeeded(USAGE as never);
43+
expect(didCompact).toBe(true);
44+
45+
// Compaction is model-only, so the two lanes deliberately disagree here.
46+
expect(flatten(conversation.modelMessages)).toContain("SUMMARY-OF-EVERYTHING");
47+
expect(flatten(conversation.modelMessages)).not.toContain("EARLY-SENTINEL");
48+
expect(JSON.stringify(conversation.uiMessages)).toContain("EARLY-SENTINEL");
49+
50+
await conversation.absorbSteering([userMessage("steer-me", "u-2")]);
51+
52+
// The steer has to land in both lanes.
53+
expect(JSON.stringify(conversation.uiMessages)).toContain("steer-me");
54+
expect(flatten(conversation.modelMessages)).toContain("steer-me");
55+
// And the compaction has to survive it. Reconverting from the UI lane
56+
// brings the compacted message back and drops the summary.
57+
expect(flatten(conversation.modelMessages)).toContain("SUMMARY-OF-EVERYTHING");
58+
expect(flatten(conversation.modelMessages)).not.toContain("EARLY-SENTINEL");
59+
});
60+
});
Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
import { mockChatAgent } from "../src/v3/test/index.js";
2+
3+
import { sessionStreams } from "@trigger.dev/core/v3";
4+
import type { LanguageModelV3StreamPart } from "@ai-sdk/provider";
5+
import { simulateReadableStream, stepCountIs, streamText, tool } from "ai";
6+
import { MockLanguageModelV3 } from "ai/test";
7+
import { describe, expect, it } from "vitest";
8+
import { z } from "zod";
9+
import { chat } from "../src/v3/ai.js";
10+
11+
/**
12+
* Steering and compaction in the same turn.
13+
*
14+
* Compaction is model-only by design: it replaces the model messages with a
15+
* summary and deliberately leaves the UI messages whole, so the chat still
16+
* displays the full conversation. Reconciling the model lane by rebuilding it
17+
* from the UI lane therefore un-compacts it, and the next turn is sent the
18+
* entire pre-compaction transcript.
19+
*
20+
* The assertion that catches this is the absence of an early message, not the
21+
* presence of the steer: a rebuild puts the steer in the prompt too, so a
22+
* steer-presence check passes while compaction has been silently undone.
23+
*/
24+
25+
const USAGE = {
26+
inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined },
27+
outputTokens: { total: 1, text: 1, reasoning: undefined },
28+
};
29+
30+
function userMessage(text: string, id: string) {
31+
return { id, role: "user" as const, parts: [{ type: "text" as const, text }] };
32+
}
33+
34+
function deferred() {
35+
let resolve!: () => void;
36+
const promise = new Promise<void>((res) => {
37+
resolve = res;
38+
});
39+
return { promise, resolve };
40+
}
41+
42+
async function waitFor(check: () => boolean, label = "condition", timeoutMs = 8_000) {
43+
const start = Date.now();
44+
while (Date.now() - start < timeoutMs) {
45+
if (check()) return;
46+
await new Promise((r) => setTimeout(r, 10));
47+
}
48+
throw new Error(`waitFor timed out: ${label}`);
49+
}
50+
51+
function textChunks(text: string): LanguageModelV3StreamPart[] {
52+
return [
53+
{ type: "text-start", id: "t1" },
54+
{ type: "text-delta", id: "t1", delta: text },
55+
{ type: "text-end", id: "t1" },
56+
{ type: "finish", finishReason: { unified: "stop", raw: "stop" }, usage: USAGE },
57+
];
58+
}
59+
60+
function toolCallChunks(callId: string): LanguageModelV3StreamPart[] {
61+
return [
62+
{ type: "tool-call", toolCallId: callId, toolName: "gate", input: JSON.stringify({ q: "go" }) },
63+
{ type: "finish", finishReason: { unified: "tool-calls", raw: "tool-calls" }, usage: USAGE },
64+
];
65+
}
66+
67+
type SeqReader = { lastSeqNum: (chatId: string, dir: "in" | "out") => number | undefined };
68+
69+
async function sendAndLand(
70+
harness: { sendMessage: (m: ReturnType<typeof userMessage>) => Promise<unknown> },
71+
chatId: string,
72+
text: string,
73+
id: string
74+
) {
75+
const seqs = sessionStreams as unknown as SeqReader;
76+
const before = seqs.lastSeqNum(chatId, "in") ?? -1;
77+
void harness.sendMessage(userMessage(text, id));
78+
await waitFor(() => (seqs.lastSeqNum(chatId, "in") ?? -1) > before, `append ${id}`);
79+
}
80+
81+
const flatUserTexts = (prompt: { role: string; content: unknown }[]) =>
82+
prompt
83+
.filter((m) => m.role === "user")
84+
.flatMap((m) =>
85+
Array.isArray(m.content)
86+
? (m.content as { type: string; text?: string }[])
87+
.filter((c) => c.type === "text")
88+
.map((c) => c.text ?? "")
89+
: []
90+
);
91+
92+
describe("chat.agent steering with compaction in the same turn", () => {
93+
it("keeps the summary and adds the steer, rather than restoring the transcript", async () => {
94+
const chatId = "steer-compaction";
95+
const toolGate = deferred();
96+
let toolEntered = false;
97+
const prompts: string[][] = [];
98+
const allPrompts: string[] = [];
99+
100+
const gateTool = tool({
101+
description: "blocks until the test opens it",
102+
inputSchema: z.object({ q: z.string() }),
103+
execute: async () => {
104+
toolEntered = true;
105+
await toolGate.promise;
106+
return "ok";
107+
},
108+
});
109+
110+
let step = 0;
111+
const model = new MockLanguageModelV3({
112+
doStream: async ({ prompt }) => {
113+
prompts.push(flatUserTexts(prompt));
114+
allPrompts.push(JSON.stringify(prompt));
115+
const isToolStep = step++ % 2 === 0;
116+
return {
117+
stream: simulateReadableStream({
118+
chunks: isToolStep ? toolCallChunks(`tc-${step}`) : textChunks("done"),
119+
initialDelayInMs: 10,
120+
chunkDelayInMs: 2,
121+
}),
122+
};
123+
},
124+
});
125+
126+
let compacted = 0;
127+
const agent = chat.agent({
128+
id: "steer-compaction",
129+
pendingMessages: { shouldInject: () => true },
130+
compaction: {
131+
// Compact once, at the first step boundary of turn 1.
132+
shouldCompact: () => compacted === 0,
133+
summarize: async () => {
134+
compacted++;
135+
return "SUMMARY-OF-EVERYTHING";
136+
},
137+
},
138+
run: async ({ messages, signal }) =>
139+
streamText({
140+
model,
141+
messages,
142+
abortSignal: signal,
143+
...chat.toStreamTextOptions(),
144+
tools: { gate: gateTool },
145+
stopWhen: stepCountIs(5),
146+
}),
147+
});
148+
149+
const harness = mockChatAgent(agent, { chatId });
150+
try {
151+
const first = harness.sendMessage(userMessage("EARLY-SENTINEL", "u-1"));
152+
await waitFor(() => toolEntered, "tool entered");
153+
await sendAndLand(harness, chatId, "steer-me", "u-2");
154+
toolGate.resolve();
155+
await first;
156+
157+
await waitFor(() => compacted > 0, "compaction ran");
158+
const promptsAfterTurn1 = prompts.length;
159+
160+
await harness.sendMessage(userMessage("m3", "u-3"));
161+
await waitFor(() => prompts.length > promptsAfterTurn1, "turn 2 prompt built");
162+
163+
const turn2 = prompts[promptsAfterTurn1]!;
164+
const turn2Raw = allPrompts[promptsAfterTurn1]!;
165+
166+
// The steer has to survive.
167+
expect(turn2).toContain("steer-me");
168+
// And so does the compaction: the summary is what the model gets...
169+
expect(turn2Raw).toContain("SUMMARY-OF-EVERYTHING");
170+
// ...instead of the message the summary replaced. This is the assertion a
171+
// rebuild-based reconciliation fails.
172+
expect(turn2).not.toContain("EARLY-SENTINEL");
173+
} finally {
174+
toolGate.resolve();
175+
await harness.close();
176+
}
177+
});
178+
});

0 commit comments

Comments
 (0)