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.
Alessandro committed
Jul 1, 2026 at 16:15 UTC
0f7429daef11d8b0d51bc7bd24550d1a017fe04e
3 files changed
+58
-4
helpers/extract_tools.py
+31
-4
@@ -64,10 +64,7 @@ def extract_json_root_strings(content: str) -> list[str]:
64
return []
65
66
roots: list[str] = []
67
- for start, char in enumerate(content):
68
- if char != "{":
69
- continue
70
-
67
+ for start in _json_root_object_starts(content):
68
parser = DirtyJson()
69
try:
70
parser.parse(content[start:])
@@ -81,6 +78,36 @@ def extract_json_root_strings(content: str) -> list[str]:
78
return roots
79
80
81
+def _json_root_object_starts(content: str) -> list[int]:
82
+ starts: list[int] = []
83
+ depth = 0
84
+ quote: str | None = None
85
+ escaped = False
86
+
87
+ for index, char in enumerate(content):
88
+ if quote:
89
+ if escaped:
90
+ escaped = False
91
+ elif char == "\\":
92
+ escaped = True
93
+ elif char == quote:
94
+ quote = None
95
+ continue
96
+
97
+ if depth and char in ['"', "'", "`"]:
98
+ quote = char
99
+ elif char == "{":
100
+ if depth == 0:
101
+ starts.append(index)
102
+ depth += 1
103
+ elif depth and char == "[":
104
+ depth += 1
105
+ elif depth and char in ["}", "]"]:
106
+ depth -= 1
107
+
108
+ return starts
109
+
110
+
111
def _parse_json_root_object(root: str) -> dict[str, Any] | None:
112
try:
113
data = DirtyJson.parse_string(root)
helpers/extract_tools.py.dox.md
+1
@@ -26,6 +26,7 @@
26
- Observed side-effect areas: settings/state persistence.
27
- 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.
28
- 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.
29
+- 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.
30
- Imported dependency areas include: `dirty_json`, `helpers.modules`, `re`, `regex`, `typing`.
31
32
## Key Concepts
tests/test_stream_tool_early_stop.py
+26
@@ -112,6 +112,32 @@ def test_extract_json_root_string_prefers_valid_tool_request():
112
) == '{"note":"not the tool"}'
113
114
115
+def test_extract_json_root_string_waits_for_complete_parallel_parent():
116
+ partial = (
117
+ '{"tool_name":"parallel","tool_args":{"tool_calls":['
118
+ '{"tool_name":"code_execution_tool","tool_args":{"code":"first"}}'
119
+ )
120
+
121
+ assert extract_tools.extract_json_root_string(partial) is None
122
+
123
+ full = (
124
+ partial
125
+ + ',{"tool_name":"code_execution_tool","tool_args":{"code":"second"}}'
126
+ '],"wait":true}} trailing text'
127
+ )
128
+
129
+ root = extract_tools.extract_json_root_string(full)
130
+ assert root == (
131
+ '{"tool_name":"parallel","tool_args":{"tool_calls":['
132
+ '{"tool_name":"code_execution_tool","tool_args":{"code":"first"}},'
133
+ '{"tool_name":"code_execution_tool","tool_args":{"code":"second"}}'
134
+ '],"wait":true}}'
135
+ )
136
+ parsed = extract_tools.json_parse_dirty(root)
137
+ assert parsed["tool_name"] == "parallel"
138
+ assert len(parsed["tool_args"]["tool_calls"]) == 2
139
+
140
+
141
def test_litellm_global_kwargs_merge_defaults_and_config(monkeypatch):
142
monkeypatch.setattr(
143
models.settings,