From 0f7429daef11d8b0d51bc7bd24550d1a017fe04e Mon Sep 17 00:00:00 2001
From: Alessandro <155005371+3clyp50@users.noreply.github.com>
Date: Wed, 1 Jul 2026 16:15:49 +0200
Subject: [PATCH] Fix streamed parallel tool extraction

Only treat top-level JSON objects as tool roots during streaming, so a complete nested tool_calls item cannot end the stream before the parallel wrapper closes.

Add regression coverage for partial parallel wrapper snapshots.
---
 helpers/extract_tools.py             | 35 ++++++++++++++++++++++++----
 helpers/extract_tools.py.dox.md      |  1 +
 tests/test_stream_tool_early_stop.py | 26 +++++++++++++++++++++
 3 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/helpers/extract_tools.py b/helpers/extract_tools.py
index 8e1124d5b..0c49df968 100644
--- a/helpers/extract_tools.py
+++ b/helpers/extract_tools.py
@@ -64,10 +64,7 @@ def extract_json_root_strings(content: str) -> list[str]:
         return []
 
     roots: list[str] = []
-    for start, char in enumerate(content):
-        if char != "{":
-            continue
-
+    for start in _json_root_object_starts(content):
         parser = DirtyJson()
         try:
             parser.parse(content[start:])
@@ -81,6 +78,36 @@ def extract_json_root_strings(content: str) -> list[str]:
     return roots
 
 
+def _json_root_object_starts(content: str) -> list[int]:
+    starts: list[int] = []
+    depth = 0
+    quote: str | None = None
+    escaped = False
+
+    for index, char in enumerate(content):
+        if quote:
+            if escaped:
+                escaped = False
+            elif char == "\\":
+                escaped = True
+            elif char == quote:
+                quote = None
+            continue
+
+        if depth and char in ['"', "'", "`"]:
+            quote = char
+        elif char == "{":
+            if depth == 0:
+                starts.append(index)
+            depth += 1
+        elif depth and char == "[":
+            depth += 1
+        elif depth and char in ["}", "]"]:
+            depth -= 1
+
+    return starts
+
+
 def _parse_json_root_object(root: str) -> dict[str, Any] | None:
     try:
         data = DirtyJson.parse_string(root)
diff --git a/helpers/extract_tools.py.dox.md b/helpers/extract_tools.py.dox.md
index 0d5f7fdd1..d932d79c2 100644
--- a/helpers/extract_tools.py.dox.md
+++ b/helpers/extract_tools.py.dox.md
@@ -26,6 +26,7 @@
 - Observed side-effect areas: settings/state persistence.
 - Dirty parsing scans complete JSON object roots in prose and prefers the first object that normalizes as a valid tool request, so a leading text preamble or incidental non-tool object does not force a misformat warning when a valid tool call follows.
 - Streaming tool snapshots use the same valid-tool preference through `extract_json_root_string`, while preserving the first complete object fallback when no valid tool-call object is present.
+- Root extraction ignores objects nested inside an open parent object, so streamed wrapper tools such as `parallel` cannot stop early on the first nested `tool_calls` item.
 - Imported dependency areas include: `dirty_json`, `helpers.modules`, `re`, `regex`, `typing`.
 
 ## Key Concepts
diff --git a/tests/test_stream_tool_early_stop.py b/tests/test_stream_tool_early_stop.py
index 5caff077c..bdc9e5acb 100644
--- a/tests/test_stream_tool_early_stop.py
+++ b/tests/test_stream_tool_early_stop.py
@@ -112,6 +112,32 @@ def test_extract_json_root_string_prefers_valid_tool_request():
     ) == '{"note":"not the tool"}'
 
 
+def test_extract_json_root_string_waits_for_complete_parallel_parent():
+    partial = (
+        '{"tool_name":"parallel","tool_args":{"tool_calls":['
+        '{"tool_name":"code_execution_tool","tool_args":{"code":"first"}}'
+    )
+
+    assert extract_tools.extract_json_root_string(partial) is None
+
+    full = (
+        partial
+        + ',{"tool_name":"code_execution_tool","tool_args":{"code":"second"}}'
+        '],"wait":true}} trailing text'
+    )
+
+    root = extract_tools.extract_json_root_string(full)
+    assert root == (
+        '{"tool_name":"parallel","tool_args":{"tool_calls":['
+        '{"tool_name":"code_execution_tool","tool_args":{"code":"first"}},'
+        '{"tool_name":"code_execution_tool","tool_args":{"code":"second"}}'
+        '],"wait":true}}'
+    )
+    parsed = extract_tools.json_parse_dirty(root)
+    assert parsed["tool_name"] == "parallel"
+    assert len(parsed["tool_args"]["tool_calls"]) == 2
+
+
 def test_litellm_global_kwargs_merge_defaults_and_config(monkeypatch):
     monkeypatch.setattr(
         models.settings,