diff --git a/core/tools/systemMessageTools/buildToolsSystemMessage.ts b/core/tools/systemMessageTools/buildToolsSystemMessage.ts index 750e6bb430d..5948b044565 100644 --- a/core/tools/systemMessageTools/buildToolsSystemMessage.ts +++ b/core/tools/systemMessageTools/buildToolsSystemMessage.ts @@ -40,7 +40,7 @@ export const generateToolsSystemMessage = ( `\nAlso, these additional tool definitions show other tools you can call with the same syntax:`, ); - for (const tool of tools) { + for (const tool of withDynamicMessage) { try { const definition = framework.toolToSystemToolDefinition(tool); instructions.push(`\n${definition}`); diff --git a/core/tools/systemMessageTools/interceptSystemToolCalls.ts b/core/tools/systemMessageTools/interceptSystemToolCalls.ts index a89fbdcfbb5..bc5d8bee3ea 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 @@ -96,12 +96,13 @@ export async function* interceptSystemToolCalls( }, ]; } - } else { - // Prevent content after tool calls for now - if (parseState) { - continue; + // Completed tool calls should not terminate parsing for subsequent + // chunks/messages; reset state so normal content (or another tool + // call) can be handled. + if (parseState.done) { + parseState = undefined; } - + } else { // Yield normal assistant message yield [ { diff --git a/core/tools/systemMessageTools/toolCodeblocks/buildSystemMessage.vitest.ts b/core/tools/systemMessageTools/toolCodeblocks/buildSystemMessage.vitest.ts index 669d1c10313..4742559ef63 100644 --- a/core/tools/systemMessageTools/toolCodeblocks/buildSystemMessage.vitest.ts +++ b/core/tools/systemMessageTools/toolCodeblocks/buildSystemMessage.vitest.ts @@ -191,6 +191,11 @@ describe("generateToolsSystemMessage", () => { const hasDynamicToolsSection = /additional tool definitions/i.test(result); expect(hasDynamicToolsSection).toBe(true); + + // Dynamic definitions should not duplicate tools that already have predefined messages. + expect(result.match(/TOOL_NAME: tool_with_description/g)?.length ?? 0).toBe( + 1, + ); }); it("includes example tool definition and call", () => { diff --git a/core/tools/systemMessageTools/toolCodeblocks/interceptSystemToolCalls.vitest.ts b/core/tools/systemMessageTools/toolCodeblocks/interceptSystemToolCalls.vitest.ts index a636cc6a9da..ab49c06028b 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( @@ -265,9 +265,79 @@ describe("interceptSystemToolCalls", () => { result = await generator.next(); } - // The content after the tool call should be ignored + // The content after the tool call should be 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 a tool call that appears mid-message and preserves trailing content", async () => { + const messages: ChatMessage[][] = [ + [ + { + role: "assistant", + content: + "Before tool\n```tool\nTOOL_NAME: test_tool\nBEGIN_ARG: arg1\nvalue1\nEND_ARG\n```\nAfter tool", + }, + ], + ]; + + const generator = interceptSystemToolCalls( + createAsyncGenerator(messages), + abortController, + framework, + ); + + let result = await generator.next(); + expect(result.value).toEqual([ + { + role: "assistant", + content: [{ type: "text", text: "Before tool" }], + }, + ]); + + result = await generator.next(); + expect(result.value).toEqual([ + { + role: "assistant", + content: [{ type: "text", text: "\n" }], + }, + ]); + + result = await generator.next(); + expect( + (result.value as AssistantChatMessage[])[0].toolCalls?.[0].function?.name, + ).toBe("test_tool"); + + result = await generator.next(); + expect( + (result.value as AssistantChatMessage[])[0].toolCalls?.[0].function + ?.arguments, + ).toContain('{"arg1":'); + + result = await generator.next(); + expect( + (result.value as AssistantChatMessage[])[0].toolCalls?.[0].function + ?.arguments, + ).toBe('"value1"'); + + result = await generator.next(); + expect( + (result.value as AssistantChatMessage[])[0].toolCalls?.[0].function + ?.arguments, + ).toBe("}"); + + result = await generator.next(); + expect(result.value).toEqual([ + { + role: "assistant", + content: [{ type: "text", text: "After tool" }], + }, + ]); }); it("stops processing when aborted", async () => {