From 653adb1772b2134ee099d08a061ff446cbe48903 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Fri, 6 Mar 2026 03:03:55 +0800 Subject: [PATCH] fix(core): allow multiple system tool codeblock calls per response --- .../interceptSystemToolCalls.ts | 13 ++- .../toolCodeblocks/index.ts | 2 +- .../interceptSystemToolCalls.vitest.ts | 87 +++++++++++++++++-- 3 files changed, 87 insertions(+), 15 deletions(-) diff --git a/core/tools/systemMessageTools/interceptSystemToolCalls.ts b/core/tools/systemMessageTools/interceptSystemToolCalls.ts index a89fbdcfbb5..9e0931e8c18 100644 --- a/core/tools/systemMessageTools/interceptSystemToolCalls.ts +++ b/core/tools/systemMessageTools/interceptSystemToolCalls.ts @@ -46,7 +46,7 @@ export async function* interceptSystemToolCalls( return result.value; } else { for await (const message of result.value) { - if (abortController.signal.aborted || parseState?.done) { + if (abortController.signal.aborted) { break; } // Skip non-assistant messages or messages with native tool calls @@ -69,6 +69,10 @@ export async function* interceptSystemToolCalls( for (const chunk of chunks) { buffer += chunk; + if (parseState?.done) { + parseState = undefined; + } + if (!parseState) { const { isInPartialStart, isInToolCall, modifiedBuffer } = detectToolCallStart(buffer, systemToolFramework); @@ -82,7 +86,7 @@ export async function* interceptSystemToolCalls( } } - if (parseState && !parseState.done) { + if (parseState) { const delta = systemToolFramework.handleToolCallBuffer( buffer, parseState, @@ -97,11 +101,6 @@ export async function* interceptSystemToolCalls( ]; } } else { - // Prevent content after tool calls for now - if (parseState) { - continue; - } - // Yield normal assistant message yield [ { diff --git a/core/tools/systemMessageTools/toolCodeblocks/index.ts b/core/tools/systemMessageTools/toolCodeblocks/index.ts index 97baa19fa06..a68ae21e226 100644 --- a/core/tools/systemMessageTools/toolCodeblocks/index.ts +++ b/core/tools/systemMessageTools/toolCodeblocks/index.ts @@ -71,7 +71,7 @@ To use a tool, respond with a tool code block (\`\`\`tool) using the syntax show systemMessageSuffix = `If it seems like the User's request could be solved with one of the tools, choose the BEST one for the job based on the user's request and the tool descriptions Then send the \`\`\`tool codeblock (YOU call the tool, not the user). Always start the codeblock on a new line. Do not perform actions with/for hypothetical files. Ask the user or use tools to deduce which files are relevant. -You can only call ONE tool at at time. The tool codeblock should be the last thing you say; stop your response after the tool codeblock.`; +You may call multiple tools if needed. Put each call in its own tool codeblock and continue with normal text whenever helpful.`; exampleDynamicToolDefinition = ` \`\`\`tool_definition diff --git a/core/tools/systemMessageTools/toolCodeblocks/interceptSystemToolCalls.vitest.ts b/core/tools/systemMessageTools/toolCodeblocks/interceptSystemToolCalls.vitest.ts index a636cc6a9da..619c4ca244a 100644 --- a/core/tools/systemMessageTools/toolCodeblocks/interceptSystemToolCalls.vitest.ts +++ b/core/tools/systemMessageTools/toolCodeblocks/interceptSystemToolCalls.vitest.ts @@ -242,7 +242,7 @@ describe("interceptSystemToolCalls", () => { ).toBe("}"); }); - it("ignores content after a tool call", async () => { + it("preserves content after a tool call", async () => { const messages: ChatMessage[][] = [ [{ role: "assistant", content: "```tool\n" }], [{ role: "assistant", content: "TOOL_NAME: test_tool\n" }], @@ -250,7 +250,7 @@ describe("interceptSystemToolCalls", () => { [{ role: "assistant", content: "value1\n" }], [{ role: "assistant", content: "END_ARG\n" }], [{ role: "assistant", content: "```\n" }], - [{ role: "assistant", content: "This content should be ignored" }], + [{ role: "assistant", content: "This content should be preserved" }], ]; const generator = interceptSystemToolCalls( @@ -259,15 +259,88 @@ describe("interceptSystemToolCalls", () => { framework, ); - let result; - // Process through all the tool call parts + // Process through tool call parts for (let i = 0; i < 6; i++) { - result = await generator.next(); + await generator.next(); } - // The content after the tool call should be ignored + // Newline after closing codeblock is now surfaced as text + let result = await generator.next(); + expect(result.value).toEqual([ + { + role: "assistant", + content: [{ type: "text", text: "\n" }], + }, + ]); + + // Content after tool call is preserved result = await generator.next(); - expect(result.value).toBeUndefined(); + expect(result.value).toEqual([ + { + role: "assistant", + content: [{ type: "text", text: "This content should be preserved" }], + }, + ]); + }); + + it("parses multiple tool calls in a single response stream", async () => { + const messages: ChatMessage[][] = [ + [{ role: "assistant", content: "```tool\n" }], + [{ role: "assistant", content: "TOOL_NAME: first_tool\n" }], + [{ role: "assistant", content: "BEGIN_ARG: arg1\n" }], + [{ role: "assistant", content: "value1\n" }], + [{ role: "assistant", content: "END_ARG\n" }], + [{ role: "assistant", content: "```\n" }], + [{ role: "assistant", content: "Now running another tool.\n" }], + [{ role: "assistant", content: "```tool\n" }], + [{ role: "assistant", content: "TOOL_NAME: second_tool\n" }], + [{ role: "assistant", content: "BEGIN_ARG: arg2\n" }], + [{ role: "assistant", content: "value2\n" }], + [{ role: "assistant", content: "END_ARG\n" }], + [{ role: "assistant", content: "```" }], + ]; + + const generator = interceptSystemToolCalls( + createAsyncGenerator(messages), + abortController, + framework, + ); + + // First tool + await generator.next(); // first start token + const firstToolName = await generator.next(); + expect( + (firstToolName.value as AssistantChatMessage[])[0].toolCalls?.[0] + .function?.name, + ).toBe("first_tool"); + + await generator.next(); // begin arg + await generator.next(); // arg value + await generator.next(); // end arg + + const transitionText = await generator.next(); // newline after closing ``` + expect(transitionText.value).toEqual([ + { + role: "assistant", + content: [{ type: "text", text: "\n" }], + }, + ]); + + const bridgeText = await generator.next(); + expect(bridgeText.value).toEqual([ + { + role: "assistant", + content: [{ type: "text", text: "Now running another tool." }], + }, + ]); + + await generator.next(); // trailing newline from bridge text + await generator.next(); // second start token + const secondToolName = await generator.next(); + expect( + (secondToolName.value as AssistantChatMessage[])[0].toolCalls?.[0] + .function?.name, + ).toBe("second_tool"); }); it("stops processing when aborted", async () => {