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