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.
This commit is contained in:
Alessandro 2026-07-01 16:15:49 +02:00
parent b805773c91
commit 0f7429daef
3 changed files with 58 additions and 4 deletions

View file

@ -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)

View file

@ -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

View file

@ -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,