Fix system tool parser dropping text and blocking follow-up tool calls

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 [
               {

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

core/tools/systemMessageTools/toolCodeblocks/interceptSystemToolCalls.vitest.ts

@@ -242,15 +242,15 @@ 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" }],
       [{ 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: "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 () => {