Skip to content

Commit 72efbcc

Browse files
committed
fix(sdk): tighten chat.history read primitives from review feedback
- Clarify getPendingToolCalls JSDoc: walks back to most-recent assistant (not strict tail), so a trailing user message does not change the result. Document why approval-* states are not surfaced. - extractNewToolResults now collapses duplicate toolCallIds within a single message to one emission (within-message dedup). - Collapse getChain to a true alias delegating to all() so the bodies cannot drift. - Tests: add real-stream chain dedup, output-error via the runtime merger, multi-tool single message with one new id (plus within- message duplicate), and pending-after-trailing-user semantic.
1 parent e88d185 commit 72efbcc

2 files changed

Lines changed: 323 additions & 10 deletions

File tree

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

Lines changed: 27 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1517,6 +1517,10 @@ function getResolvedToolCallsFromHistory(messages: UIMessage[]): ChatToolCallRef
15171517
* Pure helper: tool parts in `message` that have a fresh result not
15181518
* already represented by the resolved toolCallIds in `messages`. The
15191519
* `errorText` field is present only for `output-error` parts.
1520+
*
1521+
* Within a single `message`, duplicate `toolCallId`s emit only once
1522+
* (first occurrence wins). This guards against malformed assistants
1523+
* with repeated tool parts.
15201524
* @internal
15211525
*/
15221526
function extractNewToolResultsFromHistory(
@@ -1526,10 +1530,13 @@ function extractNewToolResultsFromHistory(
15261530
const resolved = new Set(
15271531
getResolvedToolCallsFromHistory(messages).map((r) => r.toolCallId)
15281532
);
1533+
const seen = new Set<string>();
15291534
const out: ChatNewToolResult[] = [];
15301535
for (const { part, toolCallId, toolName, state } of iterateToolParts(message)) {
15311536
if (!isResolvedToolState(state)) continue;
15321537
if (resolved.has(toolCallId)) continue;
1538+
if (seen.has(toolCallId)) continue;
1539+
seen.add(toolCallId);
15331540
if (state === "output-error") {
15341541
out.push({ toolCallId, toolName, output: undefined, errorText: part.errorText });
15351542
} else {
@@ -1556,12 +1563,11 @@ const chatHistory = {
15561563
},
15571564

15581565
/**
1559-
* Read the current chain as an ordered `UIMessage[]`. Alias for `all()` —
1560-
* use this when working alongside parent-aware APIs (TRI-9120) where
1561-
* "chain" disambiguates from "graph".
1566+
* Read the current chain as an ordered `UIMessage[]`. Identical to
1567+
* `all()`; use whichever name reads better in context.
15621568
*/
15631569
getChain(): UIMessage[] {
1564-
return [...getChatHistoryState()];
1570+
return chatHistory.all();
15651571
},
15661572

15671573
/**
@@ -1573,13 +1579,22 @@ const chatHistory = {
15731579
},
15741580

15751581
/**
1576-
* Tool calls on the leaf assistant message still waiting on an answer
1577-
* (`input-available` state). Use this to gate fresh user turns during
1578-
* HITL flows: if `getPendingToolCalls().length > 0`, an `addToolOutput`
1579-
* is expected before any new user message.
1582+
* Tool calls on the *most recent* assistant message that are still in
1583+
* `input-available` state (waiting on an `addToolOutput` answer). The
1584+
* scan walks back from the tail and stops at the first assistant
1585+
* message it finds, so a trailing user message does not change the
1586+
* result — pending tool calls remain pending until they're resolved
1587+
* on that assistant or the assistant is removed.
1588+
*
1589+
* Use this to gate fresh user turns or actions during HITL flows: if
1590+
* `getPendingToolCalls().length > 0`, an `addToolOutput` is expected.
1591+
*
1592+
* Returns `[]` if there is no assistant message yet, or if the most
1593+
* recent assistant has no pending tool calls.
15801594
*
1581-
* Returns `[]` if there is no assistant message yet, or if the leaf
1582-
* assistant has no pending tool calls.
1595+
* Approval flows (`approval-requested` / `approval-responded` states)
1596+
* are not surfaced here. Those are about the user authorizing a tool
1597+
* to run; "pending" is about the user *answering* a tool call.
15831598
*/
15841599
getPendingToolCalls(): ChatToolCallRef[] {
15851600
return getPendingToolCallsFromHistory(getChatHistoryState());
@@ -1599,6 +1614,8 @@ const chatHistory = {
15991614
* not already represented in the current chain. Use this when
16001615
* persisting tool results to your own store: each call surfaces only
16011616
* the *new* answers, so writes stay idempotent across re-streams.
1617+
* Duplicate `toolCallId`s within `message` itself are also collapsed
1618+
* to a single entry.
16021619
*/
16031620
extractNewToolResults(message: UIMessage): ChatNewToolResult[] {
16041621
return extractNewToolResultsFromHistory(message, getChatHistoryState());

packages/trigger-sdk/test/mockChatAgent.test.ts

Lines changed: 296 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -904,6 +904,302 @@ describe("mockChatAgent", () => {
904904
}
905905
});
906906

907+
it("extractNewToolResults dedups against a real-stream-built chain", async () => {
908+
// Build the chain through real model streams (no chat.history.set seed)
909+
// and assert extractNewToolResults compares against the post-merge state.
910+
const { z } = await import("zod");
911+
const { tool } = await import("ai");
912+
const askUser = tool({
913+
description: "Ask the user.",
914+
inputSchema: z.object({ q: z.string() }),
915+
});
916+
const TC = "tc_real_chain_1";
917+
918+
let callIdx = 0;
919+
const model = new MockLanguageModelV3({
920+
doStream: async () => ({
921+
stream:
922+
callIdx++ === 0
923+
? toolCallStream({ toolCallId: TC, toolName: "askUser", input: { q: "?" } })
924+
: textStream("done"),
925+
}),
926+
});
927+
928+
let extractedAgainstRealChain: any;
929+
const agent = chat.agent({
930+
id: "mockChatAgent.history.extract-real",
931+
tools: { askUser },
932+
onTurnComplete: async () => {
933+
// After the HITL answer turn, the chain has TC resolved. An
934+
// incoming "echo" message carrying TC again should yield [].
935+
// A second new TC should yield exactly one entry.
936+
const incoming = {
937+
id: "echo",
938+
role: "assistant" as const,
939+
parts: [
940+
{
941+
type: "tool-askUser",
942+
toolCallId: TC,
943+
state: "output-available" as const,
944+
input: { q: "?" },
945+
output: { answer: "hi" },
946+
},
947+
{
948+
type: "tool-askUser",
949+
toolCallId: "tc_real_chain_2",
950+
state: "output-available" as const,
951+
input: { q: "second" },
952+
output: { answer: "yes" },
953+
},
954+
],
955+
};
956+
extractedAgainstRealChain = chat.history.extractNewToolResults(incoming as any);
957+
},
958+
run: async ({ messages, signal }) => {
959+
return streamText({ model, messages, tools: { askUser }, abortSignal: signal });
960+
},
961+
});
962+
963+
const harness = mockChatAgent(agent, { chatId: "test-history-extract-real" });
964+
try {
965+
await harness.sendMessage(userMessage("hi"));
966+
await new Promise((r) => setTimeout(r, 50));
967+
// HITL answer for TC, lands via the runtime merger.
968+
const toolAnswer = {
969+
id: "ai-sdk-fresh-id-real",
970+
role: "assistant" as const,
971+
parts: [
972+
{
973+
type: "tool-askUser",
974+
toolCallId: TC,
975+
state: "output-available" as const,
976+
input: { q: "?" },
977+
output: { answer: "hi" },
978+
},
979+
],
980+
};
981+
await harness.sendMessage(toolAnswer as any);
982+
await new Promise((r) => setTimeout(r, 50));
983+
984+
expect(extractedAgainstRealChain).toEqual([
985+
{ toolCallId: "tc_real_chain_2", toolName: "askUser", output: { answer: "yes" } },
986+
]);
987+
} finally {
988+
await harness.close();
989+
}
990+
});
991+
992+
it("extractNewToolResults surfaces output-error parts via the runtime merger", async () => {
993+
// The runtime merges incoming tool-answer messages onto the head
994+
// assistant via the toolCallId map. Here we send an answer in
995+
// `output-error` state and verify (a) getResolvedToolCalls reports
996+
// it, and (b) extractNewToolResults emits it with errorText set.
997+
const { z } = await import("zod");
998+
const { tool } = await import("ai");
999+
const search = tool({
1000+
description: "Search.",
1001+
inputSchema: z.object({ q: z.string() }),
1002+
});
1003+
const TC = "tc_err_via_merger";
1004+
1005+
let callIdx = 0;
1006+
const model = new MockLanguageModelV3({
1007+
doStream: async () => ({
1008+
stream:
1009+
callIdx++ === 0
1010+
? toolCallStream({ toolCallId: TC, toolName: "search", input: { q: "x" } })
1011+
: textStream("noted"),
1012+
}),
1013+
});
1014+
1015+
let resolved: any;
1016+
let extracted: any;
1017+
const agent = chat.agent({
1018+
id: "mockChatAgent.history.extract-error",
1019+
tools: { search },
1020+
onTurnComplete: async () => {
1021+
resolved = chat.history.getResolvedToolCalls();
1022+
// An echo carrying the same error toolCallId — should NOT surface
1023+
// as new because it's already resolved on the chain.
1024+
const echo = {
1025+
id: "echo-err",
1026+
role: "assistant" as const,
1027+
parts: [
1028+
{
1029+
type: "tool-search",
1030+
toolCallId: TC,
1031+
state: "output-error" as const,
1032+
input: { q: "x" },
1033+
errorText: "boom",
1034+
},
1035+
],
1036+
};
1037+
extracted = chat.history.extractNewToolResults(echo as any);
1038+
},
1039+
run: async ({ messages, signal }) => {
1040+
return streamText({ model, messages, tools: { search }, abortSignal: signal });
1041+
},
1042+
});
1043+
1044+
const harness = mockChatAgent(agent, { chatId: "test-history-extract-error" });
1045+
try {
1046+
await harness.sendMessage(userMessage("kick"));
1047+
await new Promise((r) => setTimeout(r, 50));
1048+
// HITL answer arriving as output-error.
1049+
const errAnswer = {
1050+
id: "ai-sdk-err-fresh",
1051+
role: "assistant" as const,
1052+
parts: [
1053+
{
1054+
type: "tool-search",
1055+
toolCallId: TC,
1056+
state: "output-error" as const,
1057+
input: { q: "x" },
1058+
errorText: "boom",
1059+
},
1060+
],
1061+
};
1062+
await harness.sendMessage(errAnswer as any);
1063+
await new Promise((r) => setTimeout(r, 50));
1064+
1065+
expect(resolved).toHaveLength(1);
1066+
expect(resolved[0]).toMatchObject({ toolCallId: TC, toolName: "search" });
1067+
// Echo of the same error toolCallId is already resolved → []
1068+
expect(extracted).toEqual([]);
1069+
} finally {
1070+
await harness.close();
1071+
}
1072+
});
1073+
1074+
it("extractNewToolResults handles a multi-tool message where only one is new", async () => {
1075+
// Pure-helper edge: incoming message has two tool parts with the
1076+
// same toolName but different toolCallIds — one already resolved
1077+
// on the chain, one fresh. Only the fresh one should surface.
1078+
let extracted: any;
1079+
const agent = chat.agent({
1080+
id: "mockChatAgent.history.extract-multi",
1081+
run: async ({ messages, signal }) => {
1082+
chat.history.set([
1083+
{
1084+
id: "a-seed",
1085+
role: "assistant",
1086+
parts: [
1087+
{
1088+
type: "tool-search",
1089+
toolCallId: "tc-old",
1090+
state: "output-available",
1091+
input: { q: "old" },
1092+
output: { hits: 1 },
1093+
},
1094+
],
1095+
} as any,
1096+
{ id: "u-1", role: "user", parts: [{ type: "text", text: "u" }] } as any,
1097+
]);
1098+
1099+
const incoming = {
1100+
id: "a-incoming",
1101+
role: "assistant" as const,
1102+
parts: [
1103+
// Same tool, already-resolved id — should be filtered.
1104+
{
1105+
type: "tool-search",
1106+
toolCallId: "tc-old",
1107+
state: "output-available" as const,
1108+
input: { q: "old" },
1109+
output: { hits: 1 },
1110+
},
1111+
// Same tool, fresh id — should surface.
1112+
{
1113+
type: "tool-search",
1114+
toolCallId: "tc-new",
1115+
state: "output-available" as const,
1116+
input: { q: "new" },
1117+
output: { hits: 9 },
1118+
},
1119+
// Duplicate of tc-new in the same message — must collapse
1120+
// to a single emission (within-message dedup).
1121+
{
1122+
type: "tool-search",
1123+
toolCallId: "tc-new",
1124+
state: "output-available" as const,
1125+
input: { q: "new" },
1126+
output: { hits: 9 },
1127+
},
1128+
],
1129+
};
1130+
extracted = chat.history.extractNewToolResults(incoming as any);
1131+
1132+
return streamText({
1133+
model: new MockLanguageModelV3({
1134+
doStream: async () => ({ stream: textStream("ok") }),
1135+
}),
1136+
messages,
1137+
abortSignal: signal,
1138+
});
1139+
},
1140+
});
1141+
1142+
const harness = mockChatAgent(agent, { chatId: "test-history-extract-multi" });
1143+
try {
1144+
await harness.sendMessage(userMessage("kick"));
1145+
await new Promise((r) => setTimeout(r, 50));
1146+
1147+
expect(extracted).toEqual([
1148+
{ toolCallId: "tc-new", toolName: "search", output: { hits: 9 } },
1149+
]);
1150+
} finally {
1151+
await harness.close();
1152+
}
1153+
});
1154+
1155+
it("getPendingToolCalls still returns the assistant's pending calls when a user message follows", async () => {
1156+
// Edge: the chain is [assistant(input-available), user]. The most
1157+
// recent assistant is the one with the pending tool call, even
1158+
// though the strict tail is a user message. The walk-back semantic
1159+
// means pending stays pending until the assistant is mutated.
1160+
let pendingAfterUser: any;
1161+
const agent = chat.agent({
1162+
id: "mockChatAgent.history.pending-after-user",
1163+
run: async ({ messages, signal }) => {
1164+
chat.history.set([
1165+
{
1166+
id: "a-pending",
1167+
role: "assistant",
1168+
parts: [
1169+
{
1170+
type: "tool-askUser",
1171+
toolCallId: "tc-still-pending",
1172+
state: "input-available",
1173+
input: { q: "?" },
1174+
},
1175+
],
1176+
} as any,
1177+
{ id: "u-after", role: "user", parts: [{ type: "text", text: "anyway..." }] } as any,
1178+
]);
1179+
pendingAfterUser = chat.history.getPendingToolCalls();
1180+
return streamText({
1181+
model: new MockLanguageModelV3({
1182+
doStream: async () => ({ stream: textStream("ok") }),
1183+
}),
1184+
messages,
1185+
abortSignal: signal,
1186+
});
1187+
},
1188+
});
1189+
1190+
const harness = mockChatAgent(agent, { chatId: "test-history-pending-after-user" });
1191+
try {
1192+
await harness.sendMessage(userMessage("kick"));
1193+
await new Promise((r) => setTimeout(r, 50));
1194+
1195+
expect(pendingAfterUser).toEqual([
1196+
{ toolCallId: "tc-still-pending", toolName: "askUser", messageId: "a-pending" },
1197+
]);
1198+
} finally {
1199+
await harness.close();
1200+
}
1201+
});
1202+
9071203
it("getChain returns a defensive copy parallel to all()", async () => {
9081204
let chainCopy: any;
9091205
let allCopy: any;

0 commit comments

Comments
 (0)