Skip to content

Commit dd06745

Browse files
committed
fix(chat): report a drained steer from a turn that fails
The error path built newUIMessages from the wire message and the partial only, never from the per-turn list the drain appends to, so a turn that failed after a steer handed onTurnComplete everything except the steer. Seeded from the per-turn list at both construction sites, deduped by id. Reproduced with a stream that rejects mid-answer; an AI SDK error part takes the normal completion path and was never affected.
1 parent bf82221 commit dd06745

2 files changed

Lines changed: 189 additions & 5 deletions

File tree

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

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9206,10 +9206,26 @@ function chatAgent<
92069206
i === partialIdx ? partialResponse! : m
92079207
) as TUIMessage[]);
92089208

9209-
let erroredNewUIMessages: TUIMessage[] = erroredWireMessage ? [erroredWireMessage] : [];
9210-
if (includePartial) {
9211-
erroredNewUIMessages.push(partialResponse!);
9212-
}
9209+
/**
9210+
* Seeded from the per-turn list, not just the wire message and the
9211+
* partial, so a steering message the drain consumed is reported too.
9212+
* An app persisting from `newUIMessages` would otherwise lose the
9213+
* instruction whenever the turn it steered went on to fail.
9214+
*/
9215+
const buildErroredNew = (): TUIMessage[] => {
9216+
const out: TUIMessage[] = [];
9217+
const addUnique = (m?: TUIMessage) => {
9218+
if (m && !out.some((existing) => existing.id === m.id)) out.push(m);
9219+
};
9220+
addUnique(erroredWireMessage);
9221+
for (const m of (locals.get(chatTurnNewUIMessagesKey) ?? []) as TUIMessage[]) {
9222+
addUnique(m);
9223+
}
9224+
if (includePartial) addUnique(partialResponse!);
9225+
return out;
9226+
};
9227+
9228+
let erroredNewUIMessages: TUIMessage[] = buildErroredNew();
92139229

92149230
let erroredNewModelMessages: ModelMessage[] = [];
92159231

@@ -9237,7 +9253,7 @@ function chatAgent<
92379253
} catch {
92389254
erroredNewModelMessages = [];
92399255
erroredUIMessagesWithPartial = erroredUIMessages;
9240-
erroredNewUIMessages = erroredWireMessage ? [erroredWireMessage] : [];
9256+
erroredNewUIMessages = buildErroredNew().filter((m) => m !== partialResponse);
92419257
}
92429258
}
92439259

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
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 { stepCountIs, streamText, tool } from "ai";
6+
import { MockLanguageModelV3 } from "ai/test";
7+
import type { UIMessage } from "ai";
8+
import { describe, expect, it } from "vitest";
9+
import { z } from "zod";
10+
import { chat } from "../src/v3/ai.js";
11+
12+
/**
13+
* A turn that drains a steering message and then fails.
14+
*
15+
* The error path builds its own `newUIMessages` from the wire message and the
16+
* partial response, so a steer the drain consumed is not in it. That is the
17+
* same append-only persistence hole the steering fix exists to close: the app
18+
* stores what `onTurnComplete` hands it, the failed turn hands it everything
19+
* except the steer, and the instruction is gone.
20+
*/
21+
22+
const USAGE = {
23+
inputTokens: { total: 1, noCache: 1, cacheRead: undefined, cacheWrite: undefined },
24+
outputTokens: { total: 1, text: 1, reasoning: undefined },
25+
};
26+
27+
function userMessage(text: string, id: string) {
28+
return { id, role: "user" as const, parts: [{ type: "text" as const, text }] };
29+
}
30+
31+
function deferred() {
32+
let resolve!: () => void;
33+
const promise = new Promise<void>((res) => {
34+
resolve = res;
35+
});
36+
return { promise, resolve };
37+
}
38+
39+
async function waitFor(check: () => boolean, label = "condition", timeoutMs = 8_000) {
40+
const start = Date.now();
41+
while (Date.now() - start < timeoutMs) {
42+
if (check()) return;
43+
await new Promise((r) => setTimeout(r, 10));
44+
}
45+
throw new Error(`waitFor timed out: ${label}`);
46+
}
47+
48+
function toolCallChunks(callId: string): LanguageModelV3StreamPart[] {
49+
return [
50+
{ type: "tool-call", toolCallId: callId, toolName: "gate", input: JSON.stringify({ q: "go" }) },
51+
{ type: "finish", finishReason: { unified: "tool-calls", raw: "tool-calls" }, usage: USAGE },
52+
];
53+
}
54+
55+
/**
56+
* Emits a partial then errors, one chunk per pull so the queue isn't reset by
57+
* erroring in the same tick as the enqueue.
58+
*/
59+
function erroringStream(): ReadableStream<LanguageModelV3StreamPart> {
60+
const chunks: LanguageModelV3StreamPart[] = [
61+
{ type: "text-start", id: "t1" },
62+
{ type: "text-delta", id: "t1", delta: "partial" },
63+
];
64+
let i = 0;
65+
return new ReadableStream({
66+
pull(controller) {
67+
if (i < chunks.length) {
68+
controller.enqueue(chunks[i++]!);
69+
return;
70+
}
71+
controller.error(new Error("UND_ERR_BODY_TIMEOUT"));
72+
},
73+
});
74+
}
75+
76+
type SeqReader = { lastSeqNum: (chatId: string, dir: "in" | "out") => number | undefined };
77+
78+
async function sendAndLand(
79+
harness: { sendMessage: (m: ReturnType<typeof userMessage>) => Promise<unknown> },
80+
chatId: string,
81+
text: string,
82+
id: string
83+
) {
84+
const seqs = sessionStreams as unknown as SeqReader;
85+
const before = seqs.lastSeqNum(chatId, "in") ?? -1;
86+
void harness.sendMessage(userMessage(text, id));
87+
await waitFor(() => (seqs.lastSeqNum(chatId, "in") ?? -1) > before, `append ${id}`);
88+
}
89+
90+
describe("chat.agent steering on a turn that fails", () => {
91+
it("still reports the steer in the error path's newUIMessages", async () => {
92+
const chatId = "steer-error-path";
93+
const toolGate = deferred();
94+
let toolEntered = false;
95+
const events: { newUIMessages: UIMessage[]; finishReason?: string }[] = [];
96+
const promptsSawSteer: boolean[] = [];
97+
98+
const gateTool = tool({
99+
description: "blocks until the test opens it",
100+
inputSchema: z.object({ q: z.string() }),
101+
execute: async () => {
102+
toolEntered = true;
103+
await toolGate.promise;
104+
return "ok";
105+
},
106+
});
107+
108+
let step = 0;
109+
const model = new MockLanguageModelV3({
110+
doStream: async ({ prompt }) => {
111+
promptsSawSteer.push(JSON.stringify(prompt).includes("steer-me"));
112+
// Step 1 calls the tool, step 2 fails mid-stream.
113+
return step++ === 0
114+
? {
115+
stream: new ReadableStream({
116+
start(c) {
117+
for (const ch of toolCallChunks("tc-1")) c.enqueue(ch);
118+
c.close();
119+
},
120+
}),
121+
}
122+
: { stream: erroringStream() };
123+
},
124+
});
125+
126+
const agent = chat.agent({
127+
id: "steer-error-path",
128+
pendingMessages: { shouldInject: () => true },
129+
onTurnComplete: async ({ newUIMessages, finishReason }) => {
130+
events.push({ newUIMessages: [...(newUIMessages ?? [])], finishReason });
131+
},
132+
run: async ({ messages, signal }) =>
133+
streamText({
134+
model,
135+
messages,
136+
abortSignal: signal,
137+
...chat.toStreamTextOptions(),
138+
tools: { gate: gateTool },
139+
stopWhen: stepCountIs(5),
140+
}),
141+
});
142+
143+
const harness = mockChatAgent(agent, { chatId });
144+
try {
145+
void harness.sendMessage(userMessage("m1", "u-1"));
146+
await waitFor(() => toolEntered, "tool entered");
147+
await sendAndLand(harness, chatId, "steer-me", "u-2");
148+
toolGate.resolve();
149+
150+
await waitFor(() => events.length >= 1, "turn complete fired");
151+
152+
// The turn really did fail, and the steer really did reach the model,
153+
// so the lane check below is about persistence and nothing else.
154+
expect(events[0]!.finishReason).toBe("error");
155+
expect(promptsSawSteer.some(Boolean)).toBe(true);
156+
157+
const texts = events[0]!.newUIMessages.flatMap((m) =>
158+
((m.parts ?? []) as { type: string; text?: string }[])
159+
.filter((p) => p.type === "text")
160+
.map((p) => p.text ?? "")
161+
);
162+
expect(texts).toContain("steer-me");
163+
} finally {
164+
toolGate.resolve();
165+
await harness.close();
166+
}
167+
});
168+
});

0 commit comments

Comments
 (0)