" in content_events[0]["text"]
+ assert "I need to search" in content_events[0]["text"]
+
+ # Tool should be executed
+ tool_starts = [e for e in events if e["type"] == "tool_start"]
+ assert len(tool_starts) == 1, f"Expected tool_start: {events}"
+ assert tool_starts[0]["tool_name"] == "web_search"
+
+ print("PASS: test_reasoning_then_tool_call")
+
+
+def test_empty_response():
+ """Model returns empty stream (just role + [DONE]). Should not crash."""
+ backend = _make_backend()
+
+ chunks = [
+ _make_chunk({"role": "assistant"}),
+ _make_chunk({}, finish_reason="stop"),
+ ]
+ sse = _build_sse_stream(chunks)
+ events = _collect_events(backend, sse)
+
+ # Should not crash, just return with no content
+ error_events = [e for e in events if e.get("type") == "error"]
+ assert len(error_events) == 0, f"Should not error: {error_events}"
+
+ print("PASS: test_empty_response")
+
+
+def test_buffer_prefix_timeout():
+ """Content starts with '<' but is not a tool call (e.g., 'Hello
').
+ Buffer should hold briefly then flush when no prefix match at 32 chars."""
+ backend = _make_backend()
+
+ content = "This is a paragraph of HTML content that is not a tool call
"
+ chunks = [_make_chunk({"role": "assistant"})]
+ # Stream char by char
+ for char in content:
+ chunks.append(_make_chunk({"content": char}))
+ chunks.append(_make_chunk({}, finish_reason="stop"))
+
+ sse = _build_sse_stream(chunks)
+ events = _collect_events(backend, sse)
+
+ content_events = [e for e in events if e["type"] == "content"]
+ assert len(content_events) >= 1, f"Should have content events: {events}"
+
+ # No tool calls should be detected
+ tool_starts = [e for e in events if e["type"] == "tool_start"]
+ assert len(tool_starts) == 0, f"Should not detect tools in HTML: {tool_starts}"
+
+ # Final content should contain the HTML
+ final = content_events[-1]["text"]
+ assert "" in final, f"HTML content should pass through: {final}"
+
+ print("PASS: test_buffer_prefix_timeout")
+
+
+def test_buffer_resolves_to_streaming_on_non_xml_first_char():
+ """First content char is not '<' and not whitespace.
+ Should immediately transition to STREAMING."""
+ backend = _make_backend()
+
+ chunks = [
+ _make_chunk({"role": "assistant"}),
+ _make_chunk({"content": "H"}), # 'H' is not '<', instant STREAMING
+ _make_chunk({"content": "ello"}),
+ ]
+ sse = _build_sse_stream(chunks)
+ events = _collect_events(backend, sse)
+
+ content_events = [e for e in events if e["type"] == "content"]
+ # First content event should appear immediately with just "H"
+ assert len(content_events) >= 1
+ assert "H" in content_events[0]["text"]
+
+ print("PASS: test_buffer_resolves_to_streaming_on_non_xml_first_char")
+
+
+def test_draining_false_positive():
+ """Buffer detects 'Use a screwdriver"}),
+ _make_chunk({}, finish_reason="stop"),
+ ]
+ sse = _build_sse_stream(chunks)
+ events = _collect_events(backend, sse)
+
+ # "" so it enters BUFFERING.
+ # Then "_tip>" does NOT match "" since the buffer becomes
+ # "..." which doesn't start with "" or "32 chars the buffer should flush.
+ # No tool should be executed.
+ tool_starts = [e for e in events if e["type"] == "tool_start"]
+ assert len(tool_starts) == 0, f"Should not detect tool in : {tool_starts}"
+
+ print("PASS: test_draining_false_positive")
+
+
+def test_structured_tool_args_json_parsing():
+ """Verify that arguments streamed across multiple chunks get reassembled
+ and parsed correctly as JSON."""
+ backend = _make_backend()
+
+ # Arguments split across 4 chunks
+ arg_parts = ['{"qu', 'ery":', ' "wha', 't is python?"}']
+
+ chunks = [
+ _make_chunk({"role": "assistant"}),
+ _make_chunk({"tool_calls": [{"index": 0, "id": "call_abc",
+ "function": {"name": "web_search", "arguments": ""}}]}),
+ ]
+ for part in arg_parts:
+ chunks.append(_make_chunk({"tool_calls": [{"index": 0,
+ "function": {"arguments": part}}]}))
+ chunks.append(_make_chunk({}, finish_reason="tool_calls"))
+
+ sse = _build_sse_stream(chunks)
+ events = _collect_events(backend, sse)
+
+ tool_starts = [e for e in events if e["type"] == "tool_start"]
+ assert len(tool_starts) == 1
+ assert tool_starts[0]["arguments"] == {"query": "what is python?"}, \
+ f"Arguments not reassembled correctly: {tool_starts[0]['arguments']}"
+ assert tool_starts[0]["tool_call_id"] == "call_abc"
+
+ print("PASS: test_structured_tool_args_json_parsing")
+
+
+def test_auto_heal_disabled():
+ """When auto_heal_tool_calls=False, XML tool calls in content should NOT
+ be parsed -- only structured tool_calls are honored."""
+ backend = _make_backend()
+
+ tc_json = json.dumps({"name": "web_search", "arguments": {"query": "test"}})
+ content = f"{tc_json}"
+
+ chunks = [_make_chunk({"role": "assistant"})]
+ # Send as one big content chunk
+ chunks.append(_make_chunk({"content": content}))
+ chunks.append(_make_chunk({}, finish_reason="stop"))
+
+ sse = _build_sse_stream(chunks)
+ events = _collect_events(backend, sse, auto_heal_tool_calls=False)
+
+ # With auto_heal disabled, the XML should NOT be parsed as a tool call
+ tool_starts = [e for e in events if e["type"] == "tool_start"]
+ assert len(tool_starts) == 0, \
+ f"auto_heal_tool_calls=False should not parse XML tools: {tool_starts}"
+
+ print("PASS: test_auto_heal_disabled")
+
+
+def test_metrics_accumulation_across_tool_iterations():
+ """When tools are called, metrics from the tool iteration should be
+ accumulated and included in the final metadata."""
+ backend = _make_backend()
+
+ # First iteration: tool call
+ tool_chunks = [
+ _make_chunk({"role": "assistant"}),
+ _make_chunk({"tool_calls": [{"index": 0, "id": "call_0",
+ "function": {"name": "web_search", "arguments": '{"query": "test"}'}}]}),
+ _make_chunk({}, finish_reason="tool_calls"),
+ ]
+ tool_usage = {"prompt_tokens": 10, "completion_tokens": 5}
+ tool_timings = {"predicted_ms": 50, "predicted_n": 5}
+ tool_sse = _build_sse_stream(tool_chunks, final_usage=tool_usage, final_timings=tool_timings)
+
+ # Second iteration: plain text response (synthesis)
+ synth_chunks = [
+ _make_chunk({"role": "assistant"}),
+ _make_chunk({"content": "Based on my search, the answer is X."}),
+ _make_chunk({}, finish_reason="stop"),
+ ]
+ synth_usage = {"prompt_tokens": 20, "completion_tokens": 8}
+ synth_timings = {"predicted_ms": 100, "predicted_n": 8}
+ synth_sse = _build_sse_stream(synth_chunks, final_usage=synth_usage, final_timings=synth_timings)
+
+ # We need to return different SSE streams for each iteration
+ call_count = [0]
+ original_sse = [tool_sse, synth_sse]
+
+ fake_responses = [FakeResponse(tool_sse), FakeResponse(synth_sse)]
+
+ @contextlib.contextmanager
+ def fake_stream_with_retry(client, url, payload, cancel_event, headers=None):
+ idx = min(call_count[0], len(fake_responses) - 1)
+ call_count[0] += 1
+ yield fake_responses[idx]
+
+ def fake_execute_tool(tool_name, arguments, cancel_event=None, timeout=None, session_id=None):
+ return "Search result: success"
+
+ backend._stream_with_retry = fake_stream_with_retry
+
+ events = []
+ with patch("core.inference.tools.execute_tool", fake_execute_tool, create=True):
+ for event in backend.generate_chat_completion_with_tools(
+ messages=[{"role": "user", "content": "Search for test"}],
+ tools=[{"type": "function", "function": {"name": "web_search",
+ "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}}}],
+ ):
+ events.append(event)
+
+ meta_events = [e for e in events if e["type"] == "metadata"]
+ assert len(meta_events) == 1, f"Expected exactly 1 metadata event, got: {meta_events}"
+
+ meta = meta_events[0]
+ # completion_tokens should be accumulated: 5 (tool iter) + 8 (synthesis) = 13
+ assert meta["usage"]["completion_tokens"] == 13, \
+ f"Expected 13 total completion tokens, got: {meta['usage']['completion_tokens']}"
+
+ # predicted_ms and predicted_n should also accumulate
+ assert meta["timings"]["predicted_n"] == 13, \
+ f"Expected 13 predicted_n, got: {meta['timings']['predicted_n']}"
+
+ print("PASS: test_metrics_accumulation_across_tool_iterations")
+
+
+# ── Run all tests ────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+ tests = [
+ test_no_tool_call_plain_text,
+ test_structured_tool_calls,
+ test_xml_tool_call_at_start,
+ test_xml_function_tag_at_start,
+ test_whitespace_before_tool_xml,
+ test_content_then_tool_xml_safety_net,
+ test_multiple_structured_tool_calls,
+ test_reasoning_tokens_stream_immediately,
+ test_reasoning_then_tool_call,
+ test_empty_response,
+ test_buffer_prefix_timeout,
+ test_buffer_resolves_to_streaming_on_non_xml_first_char,
+ test_draining_false_positive,
+ test_structured_tool_args_json_parsing,
+ test_auto_heal_disabled,
+ test_metrics_accumulation_across_tool_iterations,
+ ]
+
+ passed = 0
+ failed = 0
+ errors = []
+
+ for test_fn in tests:
+ try:
+ test_fn()
+ passed += 1
+ except Exception as e:
+ failed += 1
+ errors.append((test_fn.__name__, str(e)))
+ import traceback
+ print(f"FAIL: {test_fn.__name__}: {e}")
+ traceback.print_exc()
+ print()
+
+ print(f"\n{'='*60}")
+ print(f"Results: {passed} passed, {failed} failed, {len(tests)} total")
+ if errors:
+ print(f"\nFailed tests:")
+ for name, err in errors:
+ print(f" - {name}: {err}")
+ print(f"{'='*60}")