Skip to content

Commit 5625960

Browse files
steipetegaliniliev
steipete
and
galiniliev
authored
fix(agent-core): ignore truncated tool calls (#97140)
* fix(agent-core): ignore truncated tool calls Co-authored-by: Galin Iliev <[email protected]> * fix(agent-core): require explicit tool-call terminals --------- Co-authored-by: Galin Iliev <[email protected]>
1 parent 552ec2b commit 5625960

12 files changed

Lines changed: 440 additions & 32 deletions

‎extensions/google/transport-stream.test.ts‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -476,6 +476,45 @@ describe("google transport stream", () => {
476476
expect(result.content[2]).toHaveProperty("thoughtSignature", "Y2FsbF9zaWdfMQ==");
477477
});
478478

479+
it("preserves MAX_TOKENS when the partial response contains a function call", async () => {
480+
guardedFetchMock.mockResolvedValueOnce(
481+
buildSseResponse([
482+
{
483+
candidates: [
484+
{
485+
content: {
486+
parts: [{ functionCall: { name: "lookup", args: { q: "hello" } } }],
487+
},
488+
finishReason: "MAX_TOKENS",
489+
},
490+
],
491+
},
492+
]),
493+
);
494+
495+
const streamFn = createGoogleGenerativeAiTransportStreamFn();
496+
const stream = await Promise.resolve(
497+
streamFn(
498+
buildGeminiModel(),
499+
{
500+
messages: [{ role: "user", content: "hello", timestamp: 0 }],
501+
tools: [
502+
{
503+
name: "lookup",
504+
description: "Look up a value",
505+
parameters: { type: "object" },
506+
},
507+
],
508+
} as Parameters<typeof streamFn>[1],
509+
{ apiKey: "gemini-api-key" } as Parameters<typeof streamFn>[2],
510+
),
511+
);
512+
const result = await stream.result();
513+
514+
expect(result.stopReason).toBe("length");
515+
expect(result.content).toEqual([expect.objectContaining({ type: "toolCall", name: "lookup" })]);
516+
});
517+
479518
it("strips redundant google provider prefixes from Gemini API model paths", async () => {
480519
guardedFetchMock.mockResolvedValueOnce(buildSseResponse([]));
481520

‎extensions/google/transport-stream.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1404,7 +1404,12 @@ function createGoogleTransportStreamFn(kind: CanonicalGoogleTransportApi): Strea
14041404
}
14051405
if (typeof candidate?.finishReason === "string") {
14061406
output.stopReason = mapStopReasonString(candidate.finishReason);
1407-
if (output.content.some((block) => block.type === "toolCall")) {
1407+
// MAX_TOKENS can leave a complete-looking partial call. Only a normal
1408+
// Google stop may promote parsed calls into an executable tool-use turn.
1409+
if (
1410+
output.stopReason === "stop" &&
1411+
output.content.some((block) => block.type === "toolCall")
1412+
) {
14081413
output.stopReason = "toolUse";
14091414
}
14101415
}

‎extensions/ollama/src/stream.test.ts‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,13 +98,13 @@ describe("buildAssistantMessage", () => {
9898
expect(msg.stopReason).toBe("length");
9999
});
100100

101-
it("keeps tool use authoritative over a length stop", () => {
101+
it("keeps a length stop authoritative over complete-looking tool calls", () => {
102102
const response = makeOllamaResponse({
103103
done_reason: "length",
104104
tool_calls: [{ function: { name: "read", arguments: { path: "README.md" } } }],
105105
});
106106
const msg = buildAssistantMessage(response, MODEL_INFO);
107-
expect(msg.stopReason).toBe("toolUse");
107+
expect(msg.stopReason).toBe("length");
108108
});
109109
});
110110

@@ -282,6 +282,32 @@ describe("createOllamaStreamFn thinking events", () => {
282282
expect(done.message?.stopReason).toBe("length");
283283
});
284284

285+
it("preserves a native length stop when the partial response contains tool calls", async () => {
286+
const events = await streamOllamaEvents(
287+
[
288+
makeOllamaResponse({
289+
done_reason: "length",
290+
tool_calls: [{ function: { name: "read", arguments: { path: "README.md" } } }],
291+
}),
292+
],
293+
{},
294+
{
295+
messages: [{ role: "user", content: "test" }],
296+
tools: [{ name: "read", description: "Read files", parameters: { type: "object" } }],
297+
} as never,
298+
);
299+
300+
const done = events.find((event) => event.type === "done") as {
301+
reason?: string;
302+
message?: { content?: Array<Record<string, unknown>>; stopReason?: string };
303+
};
304+
expect(done.reason).toBe("length");
305+
expect(done.message?.stopReason).toBe("length");
306+
expect(done.message?.content).toEqual([
307+
expect.objectContaining({ type: "toolCall", name: "read" }),
308+
]);
309+
});
310+
285311
it("uses generic stream timeout for Ollama request timeout", async () => {
286312
await streamOllamaEvents([makeOllamaResponse({ content: "ok" })], { timeoutMs: 2500 });
287313

‎extensions/ollama/src/stream.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -656,10 +656,15 @@ function estimateTokensFromChars(chars: number): number {
656656
}
657657

658658
function resolveOllamaStopReason(response: OllamaChatResponse) {
659+
// Ollama's length terminal means generation hit its token limit, even when
660+
// the partial response already contains a complete-looking tool call.
661+
if (response.done_reason === "length") {
662+
return "length" as const;
663+
}
659664
if (response.message.tool_calls?.length) {
660665
return "toolUse" as const;
661666
}
662-
return response.done_reason === "length" ? ("length" as const) : ("stop" as const);
667+
return "stop" as const;
663668
}
664669

665670
function estimateOllamaPromptTokens(params: {

‎packages/agent-core/src/agent-loop.test.ts‎

Lines changed: 122 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,126 @@ describe("agentLoop streaming updates", () => {
164164
expect(update.assistantMessageEvent).not.toHaveProperty("partial");
165165
}
166166
});
167+
168+
it("does not execute tool calls from a max-token-truncated assistant turn", async () => {
169+
const execute = vi.fn(
170+
async (): Promise<AgentToolResult<unknown>> => ({
171+
content: [{ type: "text", text: "should not run" }],
172+
details: {},
173+
}),
174+
);
175+
const contexts: Context[] = [];
176+
let streamCalls = 0;
177+
const streamFn: StreamFn = async (_model, context) => {
178+
contexts.push(context);
179+
streamCalls += 1;
180+
const stream = createAssistantMessageEventStream();
181+
if (streamCalls > 1) {
182+
const message: AssistantMessage = {
183+
role: "assistant",
184+
content: [{ type: "text", text: "continued" }],
185+
api: model.api,
186+
provider: model.provider,
187+
model: model.id,
188+
usage: TEST_USAGE,
189+
stopReason: "stop",
190+
timestamp: 2,
191+
};
192+
queueMicrotask(() => {
193+
stream.push({ type: "done", reason: "stop", message });
194+
});
195+
return stream;
196+
}
197+
const toolCall = {
198+
type: "toolCall" as const,
199+
id: "call-truncated-spawn",
200+
name: "sessions_spawn",
201+
arguments: {},
202+
};
203+
const message: AssistantMessage = {
204+
role: "assistant",
205+
content: [{ type: "text", text: "spawning" }, toolCall],
206+
api: model.api,
207+
provider: model.provider,
208+
model: model.id,
209+
usage: TEST_USAGE,
210+
stopReason: "length",
211+
timestamp: 1,
212+
};
213+
214+
queueMicrotask(() => {
215+
stream.push({ type: "start", partial: { ...message, content: [] } });
216+
stream.push({ type: "toolcall_start", contentIndex: 1, partial: message });
217+
stream.push({
218+
type: "toolcall_end",
219+
contentIndex: 1,
220+
toolCall,
221+
partial: message,
222+
});
223+
stream.push({ type: "done", reason: "length", message });
224+
});
225+
226+
return stream;
227+
};
228+
229+
const stream = agentLoop(
230+
[{ role: "user", content: "spawn specialists", timestamp: 1 }],
231+
{
232+
systemPrompt: "",
233+
messages: [],
234+
tools: [
235+
{
236+
name: "sessions_spawn",
237+
label: "sessions_spawn",
238+
description: "Spawn a child session",
239+
parameters: Type.Object({}, { additionalProperties: false }),
240+
execute,
241+
},
242+
],
243+
},
244+
{
245+
...config,
246+
getFollowUpMessages: async () =>
247+
streamCalls === 1 ? [{ role: "user", content: "continue", timestamp: 2 }] : [],
248+
},
249+
undefined,
250+
streamFn,
251+
);
252+
253+
const events = await collectEvents(stream);
254+
const messages = await stream.result();
255+
const truncatedMessageEnd = events.find(
256+
(event): event is Extract<AgentEvent, { type: "message_end" }> =>
257+
event.type === "message_end" &&
258+
event.message.role === "assistant" &&
259+
event.message.stopReason === "length",
260+
);
261+
const replayedTruncatedMessage = contexts[1]?.messages[1];
262+
263+
if (!truncatedMessageEnd || !replayedTruncatedMessage) {
264+
throw new Error("expected the truncated assistant message to be emitted and replayed");
265+
}
266+
267+
expect(execute).not.toHaveBeenCalled();
268+
expect(events.some((event) => event.type === "tool_execution_start")).toBe(false);
269+
expect(messages.map((message) => message.role)).toEqual([
270+
"user",
271+
"assistant",
272+
"user",
273+
"assistant",
274+
]);
275+
expect(messages[1]).toMatchObject({ role: "assistant", stopReason: "length" });
276+
expect(messages[1]).not.toMatchObject({
277+
content: expect.arrayContaining([expect.objectContaining({ type: "toolCall" })]),
278+
});
279+
expect(truncatedMessageEnd.message).not.toMatchObject({
280+
content: expect.arrayContaining([expect.objectContaining({ type: "toolCall" })]),
281+
});
282+
expect(replayedTruncatedMessage).toMatchObject({ role: "assistant", stopReason: "length" });
283+
expect(replayedTruncatedMessage).not.toMatchObject({
284+
content: expect.arrayContaining([expect.objectContaining({ type: "toolCall" })]),
285+
});
286+
});
167287
});
168288

169289
describe("runAgentLoop deferred tool hydration", () => {
@@ -936,7 +1056,7 @@ describe("agentLoop thinking state", () => {
9361056
totalTokens: 0,
9371057
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
9381058
},
939-
stopReason: "stop",
1059+
stopReason: content.some((item) => item.type === "toolCall") ? "toolUse" : "stop",
9401060
timestamp: 1,
9411061
};
9421062
}
@@ -968,7 +1088,7 @@ describe("agentLoop thinking state", () => {
9681088
: [{ type: "text", text: "done" }];
9691089
stream.push({
9701090
type: "done",
971-
reason: "stop",
1091+
reason: content.some((item) => item.type === "toolCall") ? "toolUse" : "stop",
9721092
message: makeAssistantMessage(activeModel, content),
9731093
});
9741094
stream.end();

‎packages/agent-core/src/agent-loop.ts‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,14 @@ function resolveAssistantMessageUpdate(
8080
return currentMessage;
8181
}
8282

83+
function removeNonExecutableToolCalls(message: AssistantMessage): AssistantMessage {
84+
if (message.stopReason === "toolUse") {
85+
return message;
86+
}
87+
const content = message.content.filter((item) => item.type !== "toolCall");
88+
return content.length === message.content.length ? message : { ...message, content };
89+
}
90+
8391
/**
8492
* Start an agent loop with a new prompt message.
8593
* The prompt is added to the context and events are emitted for it.
@@ -342,12 +350,12 @@ async function runLoop(
342350
return;
343351
}
344352

345-
// Check for tool calls
353+
// Only completed toolUse turns dispatch; length/stop can carry partial stream blocks.
346354
const toolCalls = message.content.filter((c) => c.type === "toolCall");
347355

348356
const toolResults: ToolResultMessage[] = [];
349357
hasMoreToolCalls = false;
350-
if (toolCalls.length > 0) {
358+
if (message.stopReason === "toolUse" && toolCalls.length > 0) {
351359
const executedToolBatch = await executeToolCalls(
352360
currentContext,
353361
message,
@@ -508,7 +516,7 @@ async function streamAssistantResponse(
508516

509517
case "done":
510518
case "error": {
511-
const finalMessage = await response.result();
519+
const finalMessage = removeNonExecutableToolCalls(await response.result());
512520
if (addedPartial) {
513521
context.messages[context.messages.length - 1] = finalMessage;
514522
} else {
@@ -523,7 +531,7 @@ async function streamAssistantResponse(
523531
}
524532
}
525533

526-
const finalMessage = await response.result();
534+
const finalMessage = removeNonExecutableToolCalls(await response.result());
527535
if (addedPartial) {
528536
context.messages[context.messages.length - 1] = finalMessage;
529537
} else {

0 commit comments

Comments
 (0)