diff --git a/proxy.py b/proxy.py
index 268c27b..5d69bdf 100644
--- a/proxy.py
+++ b/proxy.py
@@ -815,13 +815,28 @@ def _stringify_message_content(content: object) -> str:
class SSEChunkLogger:
+ # Markers that can start a tool-call XML block
+ _START_MARKERS = ("", "", "function": ""}
+
def __init__(self, wrapped, chat_logger: ChatLogger) -> None:
self._wrapped = wrapped
self._chat_logger = chat_logger
self._buffer = b""
+ # --- existing logging state ---
self._current_kind: str | None = None
self._current_text = ""
self._tool_calls: dict[int, dict[str, str]] = {}
+ # --- rescue state machine ---
+ self._rescue_capturing: bool = False
+ self._rescue_buf: str = ""
+ self._rescue_kind: str | None = None # "tool_call" or "function"
+ self._reasoning_holdback: str = ""
+ self._rescued_any: bool = False
+ self._rescue_index: int = 0
+ # cache upstream chunk id for synthesised events
+ self._last_chunk_id: str = "rescued"
async def _flush_text(self) -> None:
if self._current_kind and self._current_text:
@@ -845,12 +860,93 @@ async def _flush_all(self) -> None:
await self._flush_text()
await self._flush_tool_calls()
+ # ---- rescue helpers ----
+
+ @staticmethod
+ def _split_safe_prefix(text: str, markers: tuple[str, ...]) -> tuple[str, str]:
+ """Return (emit, holdback) where holdback is the longest suffix of *text*
+ that is a proper prefix of any *marker*."""
+ for length in range(len(text), 0, -1):
+ suffix = text[len(text) - length:]
+ for marker in markers:
+ if len(suffix) < len(marker) and marker.startswith(suffix):
+ return text[: len(text) - length], suffix
+ return text, ""
+
+ @staticmethod
+ def _parse_tool_call_xml(block: str) -> dict | None:
+ """Parse a tool-call XML block. Returns {name, arguments} or None."""
+ # Find
+ import re as _re
+ m = _re.search(r"\s]+)", block)
+ if not m:
+ return None
+ name: str = m.group(1)
+ # Find all VALUE
+ args: dict[str, object] = {}
+ for pm in _re.finditer(r"(.*?)", block, _re.DOTALL):
+ key = pm.group(1)
+ value = pm.group(2).strip()
+ # Coerce: try JSON parse
+ try:
+ value = json.loads(value)
+ except (json.JSONDecodeError, ValueError):
+ pass
+ args[key] = value
+ return {"name": name, "arguments": args}
+
+ def _build_synthesized_event(self, parsed: dict) -> bytes:
+ """Build a synthesised tool_calls SSE event from a parsed tool call."""
+ import secrets as _secrets
+ event: dict = {
+ "id": self._last_chunk_id,
+ "object": "chat.completion.chunk",
+ "choices": [
+ {
+ "index": 0,
+ "delta": {
+ "tool_calls": [
+ {
+ "index": self._rescue_index,
+ "id": f"call_{_secrets.token_hex(4)}",
+ "type": "function",
+ "function": {
+ "name": parsed["name"],
+ "arguments": json.dumps(parsed["arguments"]),
+ },
+ }
+ ]
+ },
+ "finish_reason": None,
+ }
+ ],
+ }
+ self._rescue_index += 1
+ body = json.dumps(event, ensure_ascii=False)
+ return f"data: {body}\r\n\r\n".encode()
+
+ # ---- main read loop ----
+
async def readany(self) -> bytes:
+ # Loop so we only ever return b"" at true EOF — the consumer treats an
+ # empty return as end-of-stream. A single upstream chunk may not complete
+ # an SSE event, in which case _readany_once returns None and we read more.
+ while True:
+ out = await self._readany_once()
+ if out is not None:
+ return out
+
+ async def _readany_once(self) -> bytes | None:
data = await self._wrapped.content.readany()
if not data:
await self._flush_all()
- return data
+ if self._buffer:
+ leftover = self._buffer
+ self._buffer = b""
+ return leftover
+ return b""
self._buffer += data
+ outbound: list[bytes] = []
while True:
crlf_idx = self._buffer.find(b"\r\n\r\n")
lf_idx = self._buffer.find(b"\n\n")
@@ -860,53 +956,159 @@ async def readany(self) -> bytes:
idx, sep_len = crlf_idx, 4
else:
idx, sep_len = lf_idx, 2
- event = self._buffer[:idx]
+ raw_event = self._buffer[: idx + sep_len]
self._buffer = self._buffer[idx + sep_len:]
- text = event.decode("utf-8", errors="replace").strip()
- if not text:
+ event_text = raw_event.decode("utf-8", errors="replace").strip()
+ if not event_text:
+ outbound.append(raw_event)
continue
+ # Extract payload line
payload = ""
- for line in text.splitlines():
+ for line in event_text.splitlines():
if line.startswith("data:"):
payload = line[5:].strip()
+ # Non-data lines, comments, [DONE] → pass through
if not payload:
+ outbound.append(raw_event)
continue
if payload == "[DONE]":
await self._flush_all()
await self._chat_logger.log_response("[DONE]", True)
+ outbound.append(raw_event)
continue
try:
obj = json.loads(payload)
except json.JSONDecodeError:
+ outbound.append(raw_event)
continue
+ # --- (a) existing logging on original delta ---
choices = obj.get("choices") or []
- if not choices:
- continue
- delta = choices[0].get("delta") or {}
- reasoning = delta.get("reasoning_content")
- content = delta.get("content")
- tool_calls = delta.get("tool_calls")
- if reasoning:
- if self._current_kind != "thinking":
- await self._flush_all()
- self._current_kind = "thinking"
- self._current_text += reasoning
- if content:
- if self._current_kind != "content":
- await self._flush_all()
- self._current_kind = "content"
- self._current_text += content
- if tool_calls:
- await self._flush_text()
- for tc in tool_calls:
- i = tc.get("index", 0)
- slot = self._tool_calls.setdefault(i, {"name": "", "arguments": ""})
- fn = tc.get("function") or {}
- if fn.get("name"):
- slot["name"] = fn["name"]
- if fn.get("arguments"):
- slot["arguments"] += fn["arguments"]
- return data
+ if choices:
+ delta = choices[0].get("delta") or {}
+ reasoning = delta.get("reasoning_content")
+ content = delta.get("content")
+ tool_calls = delta.get("tool_calls")
+ if reasoning:
+ if self._current_kind != "thinking":
+ await self._flush_all()
+ self._current_kind = "thinking"
+ self._current_text += reasoning
+ if content:
+ if self._current_kind != "content":
+ await self._flush_all()
+ self._current_kind = "content"
+ self._current_text += content
+ if tool_calls:
+ await self._flush_text()
+ for tc in tool_calls:
+ i = tc.get("index", 0)
+ slot = self._tool_calls.setdefault(i, {"name": "", "arguments": ""})
+ fn = tc.get("function") or {}
+ if fn.get("name"):
+ slot["name"] = fn["name"]
+ if fn.get("arguments"):
+ slot["arguments"] += fn["arguments"]
+ # Track chunk id for synthesised events
+ cid = obj.get("id")
+ if cid:
+ self._last_chunk_id = cid
+ # --- (b) build outbound bytes with rescue transform ---
+ outbound_event = self._transform_event(obj)
+ outbound.append(outbound_event)
+ return b"".join(outbound) if outbound else None
+
+ def _transform_event(self, obj: dict) -> bytes:
+ """Transform a single parsed event dict into outbound SSE bytes,
+ applying the rescue state machine."""
+ choices = obj.get("choices") or []
+ if not choices:
+ # No choices — pass through
+ body = json.dumps(obj, ensure_ascii=False)
+ return f"data: {body}\r\n\r\n".encode()
+
+ delta = choices[0].get("delta") or {}
+ reasoning = delta.get("reasoning_content")
+
+ # If no reasoning_content, just rewrite finish_reason if needed
+ if not reasoning:
+ if self._rescued_any and choices[0].get("finish_reason") == "stop":
+ choices[0]["finish_reason"] = "tool_calls"
+ body = json.dumps(obj, ensure_ascii=False)
+ return f"data: {body}\r\n\r\n".encode()
+
+ # Run rescue state machine on reasoning_content
+ work = self._reasoning_holdback + reasoning
+ self._reasoning_holdback = ""
+ prose_parts: list[str] = []
+ synthesized: list[bytes] = []
+
+ while work:
+ if not self._rescue_capturing:
+ # Look for earliest start marker
+ earliest_pos = len(work)
+ earliest_marker: str | None = None
+ for marker in self._START_MARKERS:
+ pos = work.find(marker)
+ if pos != -1 and pos < earliest_pos:
+ earliest_pos = pos
+ earliest_marker = marker
+
+ if earliest_marker is None:
+ # No marker found — apply split_safe_prefix
+ emit, holdback = self._split_safe_prefix(work, self._START_MARKERS)
+ prose_parts.append(emit)
+ self._reasoning_holdback = holdback
+ work = ""
+ else:
+ # A complete start marker is present, so everything before it
+ # is safe prose — no partial-marker holdback needed here.
+ prose_parts.append(work[:earliest_pos])
+ # Start capturing
+ self._rescue_capturing = True
+ self._rescue_kind = (
+ "tool_call" if earliest_marker == "" else "function"
+ )
+ self._rescue_buf = earliest_marker
+ work = work[earliest_pos + len(earliest_marker):]
+ else:
+ # Capturing — look for end marker
+ end_marker = self._END_MARKERS[self._rescue_kind]
+ end_pos = work.find(end_marker)
+ if end_pos != -1:
+ self._rescue_buf += work[: end_pos + len(end_marker)]
+ # Parse the block
+ parsed = self._parse_tool_call_xml(self._rescue_buf)
+ if parsed:
+ synthesized.append(self._build_synthesized_event(parsed))
+ self._rescued_any = True
+ self._rescue_capturing = False
+ self._rescue_buf = ""
+ self._rescue_kind = None
+ work = work[end_pos + len(end_marker):]
+ else:
+ # End marker not found — keep all of work in buffer
+ self._rescue_buf += work
+ work = ""
+
+ # Build the outbound event
+ forwarded_reasoning = "".join(prose_parts)
+ if forwarded_reasoning:
+ delta["reasoning_content"] = forwarded_reasoning
+ else:
+ delta.pop("reasoning_content", None)
+ # If delta is now empty and has no other keys, we still emit the event
+ # (the caller handles skipping if needed)
+
+ # Rewrite finish_reason if rescued
+ if self._rescued_any and choices[0].get("finish_reason") == "stop":
+ choices[0]["finish_reason"] = "tool_calls"
+
+ body = json.dumps(obj, ensure_ascii=False)
+ result = f"data: {body}\r\n\r\n".encode()
+ # Append any synthesised events after the reasoning event
+ for syn in synthesized:
+ result += syn
+ return result
def __getattr__(self, name: str) -> object:
return getattr(self._wrapped, name)
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_toolcall_rescue.py b/tests/test_toolcall_rescue.py
new file mode 100644
index 0000000..be3eab0
--- /dev/null
+++ b/tests/test_toolcall_rescue.py
@@ -0,0 +1,487 @@
+"""Tests for tool-call XML rescue from reasoning_content deltas."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import unittest
+
+import proxy
+
+
+# ---------------------------------------------------------------------------
+# Fakes
+# ---------------------------------------------------------------------------
+
+class FakeContent:
+ """Yields preset byte chunks then b"" (EOF)."""
+
+ def __init__(self, chunks: list[bytes]) -> None:
+ self._chunks = list(chunks)
+ self._idx = 0
+
+ async def readany(self) -> bytes:
+ if self._idx >= len(self._chunks):
+ return b""
+ chunk = self._chunks[self._idx]
+ self._idx += 1
+ return chunk
+
+
+class ChunkedContent:
+ """Takes a single bytes blob and yields it in fixed-size slices, then b"" (EOF).
+ Simulates real network chunk boundaries that split SSE events mid-stream."""
+
+ def __init__(self, blob: bytes, slice_size: int = 13) -> None:
+ self._blob = blob
+ self._slice_size = slice_size
+ self._pos = 0
+
+ async def readany(self) -> bytes:
+ if self._pos >= len(self._blob):
+ return b""
+ chunk = self._blob[self._pos : self._pos + self._slice_size]
+ self._pos += self._slice_size
+ return chunk
+
+
+class FakeUpstream:
+ """Minimal upstream response with .content and .headers."""
+
+ def __init__(self, chunks: list[bytes]) -> None:
+ self.content = FakeContent(chunks)
+ self.headers = {}
+
+
+class FakeChatLogger:
+ """Records log_response calls."""
+
+ def __init__(self) -> None:
+ self.calls: list[tuple[str, bool]] = []
+
+ async def log_response(self, data: str, is_done: bool) -> None:
+ self.calls.append((data, is_done))
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _sse_event(payload: dict, extra_headers: str = "") -> bytes:
+ """Build a raw SSE event byte string from a JSON-serialisable dict."""
+ body = json.dumps(payload, ensure_ascii=False)
+ return f"{extra_headers}data: {body}\r\n\r\n".encode()
+
+
+def _split_sse_events(raw: bytes) -> list[dict]:
+ """Split raw SSE bytes into parsed JSON payloads (data: lines only)."""
+ events: list[dict] = []
+ # Normalise line endings
+ text = raw.decode("utf-8", errors="replace")
+ for block in text.split("\r\n\r\n"):
+ if not block.strip():
+ continue
+ for line in block.splitlines():
+ if line.startswith("data:"):
+ payload = line[5:].strip()
+ if payload:
+ try:
+ events.append(json.loads(payload))
+ except json.JSONDecodeError:
+ pass # skip [DONE] and other non-JSON payloads
+ return events
+
+
+def _make_logger(chunks: list[bytes]) -> proxy.SSEChunkLogger:
+ return proxy.SSEChunkLogger(FakeUpstream(chunks), FakeChatLogger())
+
+
+async def _drive(chunks: list[bytes]) -> tuple[str, FakeChatLogger]:
+ """Run SSEChunkLogger through all chunks, return (combined output, logger)."""
+ logger = FakeChatLogger()
+ wrapped = proxy.SSEChunkLogger(FakeUpstream(chunks), logger)
+ out: list[bytes] = []
+ while True:
+ piece = await wrapped.readany()
+ if not piece:
+ break
+ out.append(piece)
+ return b"".join(out).decode("utf-8", errors="replace"), logger
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+class TestToolCallRescue(unittest.TestCase):
+
+ # 1. Passthrough — normal stream, nothing rescued
+ def test_passthrough_normal_stream(self) -> None:
+ """Real tool_calls delta is preserved; finish_reason stays 'stop'."""
+ chunks = [
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
+ }),
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": "Let me think..."}, "finish_reason": None}],
+ }),
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"content": "Hello"}, "finish_reason": None}],
+ }),
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {
+ "tool_calls": [{"index": 0, "id": "call_abc", "type": "function",
+ "function": {"name": "read", "arguments": '{"path":"/x"}'}}],
+ }, "finish_reason": None}],
+ }),
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+ }),
+ ]
+ raw, _ = asyncio.run(_drive(chunks))
+ events = _split_sse_events(raw.encode())
+
+ # Find the tool_calls event
+ tc_events = [e for e in events if e.get("choices", [{}])[0].get("delta", {}).get("tool_calls")]
+ self.assertEqual(len(tc_events), 1)
+ tc = tc_events[0]["choices"][0]["delta"]["tool_calls"][0]
+ self.assertEqual(tc["function"]["name"], "read")
+
+ # finish_reason should stay "stop" (nothing rescued)
+ stop_events = [e for e in events if e.get("choices", [{}])[0].get("finish_reason") == "stop"]
+ self.assertEqual(len(stop_events), 1)
+
+ # 2. Stuck call split across chunks AND events
+ def test_rescue_split_across_chunks(self) -> None:
+ """Tool-call XML split across multiple SSE events/chunks is rescued."""
+ # Build the XML tool call split across several events
+ xml_parts = [
+ "\n",
+ "\n",
+ "\n",
+ "/x\n",
+ "\n",
+ "\n",
+ "",
+ ]
+ chunks: list[bytes] = []
+ for part in xml_parts:
+ chunks.append(_sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": part}, "finish_reason": None}],
+ }))
+ chunks.append(_sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+ }))
+
+ raw, logger = asyncio.run(_drive(chunks))
+ events = _split_sse_events(raw.encode())
+
+ # Should have a tool_calls event
+ tc_events = [e for e in events if e.get("choices", [{}])[0].get("delta", {}).get("tool_calls")]
+ self.assertEqual(len(tc_events), 1)
+ tc = tc_events[0]["choices"][0]["delta"]["tool_calls"][0]
+ self.assertEqual(tc["function"]["name"], "read")
+ args = json.loads(tc["function"]["arguments"])
+ self.assertEqual(args, {"path": "/x"})
+
+ # No reasoning_content should contain any XML markers
+ for e in events:
+ rc = e.get("choices", [{}])[0].get("delta", {}).get("reasoning_content")
+ if rc:
+ self.assertNotIn("", rc)
+
+ # Combined reasoning must also be clean (catches leaked etc.)
+ combined_reasoning = "".join(
+ e["choices"][0]["delta"].get("reasoning_content", "")
+ for e in events
+ if e.get("choices", [{}])[0].get("delta", {}).get("reasoning_content")
+ )
+ self.assertNotIn("", combined_reasoning)
+ self.assertNotIn("", combined_reasoning)
+ self.assertNotIn("", combined_reasoning)
+
+ # finish_reason rewritten to "tool_calls"
+ stop_events = [e for e in events if e.get("choices", [{}])[0].get("finish_reason") == "stop"]
+ self.assertEqual(len(stop_events), 0)
+ tc_finish = [e for e in events if e.get("choices", [{}])[0].get("finish_reason") == "tool_calls"]
+ self.assertEqual(len(tc_finish), 1)
+
+ # 3. Prose around call
+ def test_prose_around_call(self) -> None:
+ """Prose before and after XML block is forwarded; XML stripped."""
+ xml_block = "\n\n\n/x\n\n\n"
+ chunks = [
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": "Let me read it."}, "finish_reason": None}],
+ }),
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": xml_block}, "finish_reason": None}],
+ }),
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": " done"}, "finish_reason": None}],
+ }),
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+ }),
+ ]
+ raw, _ = asyncio.run(_drive(chunks))
+ events = _split_sse_events(raw.encode())
+
+ # Collect all reasoning_content
+ reasonings = [
+ e["choices"][0]["delta"].get("reasoning_content")
+ for e in events
+ if e.get("choices", [{}])[0].get("delta", {}).get("reasoning_content")
+ ]
+ combined_reasoning = "".join(reasonings)
+ self.assertEqual(combined_reasoning, "Let me read it. done")
+ self.assertNotIn("", combined_reasoning)
+ self.assertNotIn("", combined_reasoning)
+ self.assertNotIn("", combined_reasoning)
+
+ # Tool call rescued
+ tc_events = [e for e in events if e.get("choices", [{}])[0].get("delta", {}).get("tool_calls")]
+ self.assertEqual(len(tc_events), 1)
+
+ # 4. No wrapper (no )
+ def test_no_wrapper(self) -> None:
+ """ without is rescued."""
+ xml = "ls"
+ chunks = [
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": xml}, "finish_reason": None}],
+ }),
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+ }),
+ ]
+ raw, _ = asyncio.run(_drive(chunks))
+ events = _split_sse_events(raw.encode())
+
+ tc_events = [e for e in events if e.get("choices", [{}])[0].get("delta", {}).get("tool_calls")]
+ self.assertEqual(len(tc_events), 1)
+ tc = tc_events[0]["choices"][0]["delta"]["tool_calls"][0]
+ self.assertEqual(tc["function"]["name"], "bash")
+ args = json.loads(tc["function"]["arguments"])
+ self.assertEqual(args, {"command": "ls"})
+
+ # 5. Bare empty
+ def test_bare_empty(self) -> None:
+ """Bare empty is stripped; no tool_calls event; finish_reason stays stop."""
+ chunks = [
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": ""}, "finish_reason": None}],
+ }),
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+ }),
+ ]
+ raw, _ = asyncio.run(_drive(chunks))
+ events = _split_sse_events(raw.encode())
+
+ tc_events = [e for e in events if e.get("choices", [{}])[0].get("delta", {}).get("tool_calls")]
+ self.assertEqual(len(tc_events), 0)
+
+ # finish_reason should stay "stop"
+ stop_events = [e for e in events if e.get("choices", [{}])[0].get("finish_reason") == "stop"]
+ self.assertEqual(len(stop_events), 1)
+
+ # 6. Arg coercion
+ def test_arg_coercion(self) -> None:
+ """Numeric args are coerced to int; string args stay strings."""
+ xml = "\n\n\n/some/path\n\n200\n\n"
+ chunks = [
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": xml}, "finish_reason": None}],
+ }),
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+ }),
+ ]
+ raw, _ = asyncio.run(_drive(chunks))
+ events = _split_sse_events(raw.encode())
+
+ tc_events = [e for e in events if e.get("choices", [{}])[0].get("delta", {}).get("tool_calls")]
+ self.assertEqual(len(tc_events), 1)
+ tc = tc_events[0]["choices"][0]["delta"]["tool_calls"][0]
+ args = json.loads(tc["function"]["arguments"])
+ self.assertEqual(args["limit"], 200)
+ self.assertIsInstance(args["limit"], int)
+ self.assertEqual(args["path"], "/some/path")
+ self.assertIsInstance(args["path"], str)
+
+
+ # 7. Two tool calls in one turn
+ def test_multiple_calls_one_turn(self) -> None:
+ """Two tool-call blocks back-to-back are both rescued with correct indices."""
+ xml_block = (
+ "/a\n"
+ "ls"
+ )
+ chunks = [
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": xml_block}, "finish_reason": None}],
+ }),
+ _sse_event({
+ "id": "req1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+ }),
+ ]
+ raw, _ = asyncio.run(_drive(chunks))
+ events = _split_sse_events(raw.encode())
+
+ tc_events = [e for e in events if e.get("choices", [{}])[0].get("delta", {}).get("tool_calls")]
+ self.assertEqual(len(tc_events), 2)
+
+ # First call: read with path=/a, index 0
+ tc0 = tc_events[0]["choices"][0]["delta"]["tool_calls"][0]
+ self.assertEqual(tc0["index"], 0)
+ self.assertEqual(tc0["function"]["name"], "read")
+ args0 = json.loads(tc0["function"]["arguments"])
+ self.assertEqual(args0, {"path": "/a"})
+
+ # Second call: bash with command=ls, index 1
+ tc1 = tc_events[1]["choices"][0]["delta"]["tool_calls"][0]
+ self.assertEqual(tc1["index"], 1)
+ self.assertEqual(tc1["function"]["name"], "bash")
+ args1 = json.loads(tc1["function"]["arguments"])
+ self.assertEqual(args1, {"command": "ls"})
+
+ # finish_reason rewritten to "tool_calls"
+ stop_events = [e for e in events if e.get("choices", [{}])[0].get("finish_reason") == "stop"]
+ self.assertEqual(len(stop_events), 0)
+ tc_finish = [e for e in events if e.get("choices", [{}])[0].get("finish_reason") == "tool_calls"]
+ self.assertEqual(len(tc_finish), 1)
+
+
+ # 8. Sub-event chunk boundaries — normal stream fed in tiny slices
+ def test_subevent_chunk_boundaries(self) -> None:
+ """A normal multi-event stream assembled into one blob, fed in 13-byte
+ slices, still produces correct output (readany loops until complete
+ events or true EOF)."""
+ blob = b""
+ blob += _sse_event({"id": "r1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]})
+ blob += _sse_event({"id": "r1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": "Let me "}, "finish_reason": None}]})
+ blob += _sse_event({"id": "r1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": "think"}, "finish_reason": None}]})
+ blob += _sse_event({"id": "r1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": " here."}, "finish_reason": None}]})
+ blob += _sse_event({"id": "r1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"content": "The "}, "finish_reason": None}]})
+ blob += _sse_event({"id": "r1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"content": "answer."}, "finish_reason": None}]})
+ blob += _sse_event({"id": "r1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]})
+ blob += b"data: [DONE]\n\n"
+
+ content = ChunkedContent(blob, slice_size=13)
+ wrapped = proxy.SSEChunkLogger(
+ FakeUpstream([]), # headers unused; we swap content
+ FakeChatLogger(),
+ )
+ # Replace the wrapped content with our chunked source
+ wrapped._wrapped.content = content
+
+ out: list[bytes] = []
+ while True:
+ piece = asyncio.run(wrapped.readany())
+ if not piece:
+ break
+ out.append(piece)
+
+ raw = b"".join(out).decode("utf-8", errors="replace")
+ self.assertTrue(len(raw) > 0, "output should be non-empty")
+
+ events = _split_sse_events(raw.encode())
+ reasoning_parts = [
+ e["choices"][0]["delta"].get("reasoning_content", "")
+ for e in events
+ if e.get("choices", [{}])[0].get("delta", {}).get("reasoning_content")
+ ]
+ self.assertEqual("".join(reasoning_parts), "Let me think here.")
+
+ content_parts = [
+ e["choices"][0]["delta"].get("content", "")
+ for e in events
+ if e.get("choices", [{}])[0].get("delta", {}).get("content")
+ ]
+ self.assertEqual("".join(content_parts), "The answer.")
+
+ stop_events = [e for e in events if e.get("choices", [{}])[0].get("finish_reason") == "stop"]
+ self.assertEqual(len(stop_events), 1)
+
+ self.assertIn("[DONE]", raw)
+
+ # 9. Rescue survives arbitrary sub-event chunk boundaries
+ def test_rescue_subevent_chunks(self) -> None:
+ """Stuck-tool-call XML assembled into one blob and fed in 13-byte slices
+ is still rescued (name "read", args {"path":"/x"}), finish_reason
+ rewritten to "tool_calls"."""
+ xml_block = "\n\n/x\n\n"
+ blob = b""
+ blob += _sse_event({"id": "r1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {"reasoning_content": f"\n{xml_block}\n"}, "finish_reason": None}]})
+ blob += _sse_event({"id": "r1", "object": "chat.completion.chunk",
+ "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]})
+ blob += b"data: [DONE]\n\n"
+
+ content = ChunkedContent(blob, slice_size=13)
+ wrapped = proxy.SSEChunkLogger(
+ FakeUpstream([]),
+ FakeChatLogger(),
+ )
+ wrapped._wrapped.content = content
+
+ out: list[bytes] = []
+ while True:
+ piece = asyncio.run(wrapped.readany())
+ if not piece:
+ break
+ out.append(piece)
+
+ raw = b"".join(out).decode("utf-8", errors="replace")
+ self.assertTrue(len(raw) > 0, "output should be non-empty")
+
+ events = _split_sse_events(raw.encode())
+
+ # Tool call rescued
+ tc_events = [e for e in events if e.get("choices", [{}])[0].get("delta", {}).get("tool_calls")]
+ self.assertEqual(len(tc_events), 1)
+ tc = tc_events[0]["choices"][0]["delta"]["tool_calls"][0]
+ self.assertEqual(tc["function"]["name"], "read")
+ args = json.loads(tc["function"]["arguments"])
+ self.assertEqual(args, {"path": "/x"})
+
+ # finish_reason rewritten to "tool_calls"
+ stop_events = [e for e in events if e.get("choices", [{}])[0].get("finish_reason") == "stop"]
+ self.assertEqual(len(stop_events), 0)
+ tc_finish = [e for e in events if e.get("choices", [{}])[0].get("finish_reason") == "tool_calls"]
+ self.assertEqual(len(tc_finish), 1)
+
+
+if __name__ == "__main__":
+ unittest.main()