Prefer tool roots in streamed snapshots

Keep the streaming early-stop caller untouched by making extract_json_root_string prefer the first complete JSON root that normalizes as a tool request while preserving its first-root fallback for non-tool JSON.

Add regression coverage for incidental JSON before a streamed tool-call envelope.

helpers/extract_tools.py

@@ -10,12 +10,9 @@ def json_parse_dirty(json: str) -> dict[str, Any] | None:
 
     parsed_candidates: list[dict[str, Any]] = []
     for ext_json in extract_json_root_strings(json.strip()):
-        try:
-            data = DirtyJson.parse_string(ext_json)
-            if isinstance(data, dict):
-                parsed_candidates.append(data)
-        except Exception:
-            continue
+        data = _parse_json_root_object(ext_json)
+        if data is not None:
+            parsed_candidates.append(data)
     for data in parsed_candidates:
         try:
             normalize_tool_request(data)
@@ -53,9 +50,17 @@ def normalize_tool_request(tool_request: Any) -> tuple[str, dict]:
 
 
 def extract_json_root_string(content: str) -> str | None:
-    for root in extract_json_root_strings(content):
+    roots = extract_json_root_strings(content)
+    for root in roots:
+        data = _parse_json_root_object(root)
+        if data is None:
+            continue
+        try:
+            normalize_tool_request(data)
+        except ValueError:
+            continue
         return root
-    return None
+    return roots[0] if roots else None
 
 
 def extract_json_root_strings(content: str) -> list[str]:
@@ -83,6 +88,14 @@ def extract_json_root_strings(content: str) -> list[str]:
     return roots
 
 
+def _parse_json_root_object(root: str) -> dict[str, Any] | None:
+    try:
+        data = DirtyJson.parse_string(root)
+    except Exception:
+        return None
+    return data if isinstance(data, dict) else None
+
+
 def extract_json_object_string(content):
     start = content.find("{")
     if start == -1:

helpers/extract_tools.py.dox.md

@@ -25,6 +25,7 @@
 - Update this file whenever public functions, classes, persistence behavior, path/security assumptions, side effects, or cross-module contracts change.
 - 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.
 - Imported dependency areas include: `dirty_json`, `helpers.modules`, `re`, `regex`, `typing`.
 
 ## Key Concepts

tests/test_stream_tool_early_stop.py

@@ -90,6 +90,20 @@ def test_json_parse_dirty_prefers_valid_tool_request_after_preamble_object():
     }
 
 
+def test_extract_json_root_string_prefers_valid_tool_request():
+    text = (
+        'I will call the tool after this note {"note":"not the tool"}.\n'
+        '{"tool_name":"response","tool_args":{"text":"ok"}} trailing text'
+    )
+
+    assert extract_tools.extract_json_root_string(text) == (
+        '{"tool_name":"response","tool_args":{"text":"ok"}}'
+    )
+    assert extract_tools.extract_json_root_string(
+        'Only a note {"note":"not the tool"}'
+    ) == '{"note":"not the tool"}'
+
+
 def test_litellm_global_kwargs_merge_defaults_and_config(monkeypatch):
     monkeypatch.setattr(
         models.settings,
@@ -242,6 +256,76 @@ async def test_unified_call_stops_after_canonical_root_snapshot(monkeypatch):
     assert seen[0][1] == '{"tool_name":"response","tool_args":{"text":"hello"}} trailing text'
 
 
+@pytest.mark.asyncio
+async def test_unified_call_stops_after_tool_root_with_incidental_json(monkeypatch):
+    stream = _AsyncChunkStream(
+        [
+            {"type": "response.created"},
+            _response_event('Preamble {"note":"not the tool"}.\n'),
+            _response_event(
+                '{"tool_name":"response","tool_args":{"text":"ok"}} trailing text'
+            ),
+            _response_event(" unreachable"),
+        ]
+    )
+
+    async def fake_aresponses(*args, **kwargs):
+        assert kwargs["stream"] is True
+        assert kwargs["input"] == ""
+        assert kwargs["store"] is True
+        return stream
+
+    async def fake_rate_limiter(*args, **kwargs):
+        return None
+
+    monkeypatch.setattr(litellm_transport, "aresponses", fake_aresponses)
+    monkeypatch.setattr(models, "apply_rate_limiter", fake_rate_limiter)
+    monkeypatch.setattr(
+        models.settings,
+        "get_settings",
+        lambda: {"litellm_global_kwargs": {}},
+    )
+
+    wrapper = models.LiteLLMChatWrapper(
+        model="test-model",
+        provider="openai",
+        model_config=None,
+    )
+
+    seen: list[tuple[str, str]] = []
+
+    async def response_callback(chunk: str, full: str):
+        seen.append((chunk, full))
+        snapshot = extract_tools.extract_json_root_string(full)
+        if not snapshot:
+            return None
+        parsed_snapshot = extract_tools.json_parse_dirty(snapshot)
+        if parsed_snapshot is None:
+            return None
+        try:
+            extract_tools.normalize_tool_request(parsed_snapshot)
+        except ValueError:
+            return None
+        return snapshot
+
+    response, reasoning = await wrapper.unified_call(
+        messages=[],
+        response_callback=response_callback,
+    )
+
+    assert response == '{"tool_name":"response","tool_args":{"text":"ok"}}'
+    assert reasoning == ""
+    assert stream.index == 3
+    assert stream.closed is True
+    assert len(seen) == 2
+    assert seen[0][1] == 'Preamble {"note":"not the tool"}.\n'
+    assert (
+        seen[1][1]
+        == 'Preamble {"note":"not the tool"}.\n'
+        '{"tool_name":"response","tool_args":{"text":"ok"}} trailing text'
+    )
+
+
 @pytest.mark.asyncio
 async def test_unified_call_closes_responses_stream_when_callback_raises(monkeypatch):
     stream = _AsyncChunkStream([_response_event("interrupt me")])