diff --git a/chat_logger.py b/chat_logger.py index 687bd1a..09d9c66 100644 --- a/chat_logger.py +++ b/chat_logger.py @@ -19,6 +19,7 @@ RAW_BODY_CAP = 1024 * 1024 +THINK_LEAK_WINDOW = 4096 # bytes of accumulated content text class ChatLogger: @@ -156,6 +157,12 @@ def __init__(self, wrapped, chat_logger: ChatLogger) -> None: self._rescue_index: int = 0 # cache upstream chunk id for synthesised events self._last_chunk_id: str = "rescued" + # --- content-side think-leak buffer (Fix 1) --- + self._content_buffer_active: bool = False + self._content_buf: str = "" + self._content_buffer_decided: bool = False + # --- buffered finish event (Fix 2 — defer until truncated rescue resolves) --- + self._pending_finish: bytes | None = None async def _flush_text(self) -> None: if self._current_kind and self._current_text: @@ -259,14 +266,47 @@ async def readany(self) -> bytes: # ── read loop helpers ── + def _resolve_stream_end(self) -> bytes: + """Resolve all held state at end of stream ([DONE] or EOF): flush the + content buffer, attempt the truncated tool-call rescue, and release or + drop the held finish event. Idempotent — safe to call from both the + [DONE] path and _handle_eof.""" + outbound: list[bytes] = [] + + # Flush content buffer (Fix 1) + if self._content_buffer_active and not self._content_buffer_decided and self._content_buf: + flush_bytes = self._flush_content_buffer() + if flush_bytes: + outbound.append(flush_bytes) + + # Try truncated rescue (Fix 2) + rescued_succeeded = False + if self._rescue_capturing and self._rescue_buf: + synthesized = self._try_rescue_truncated() + if synthesized: + outbound.append(synthesized) + rescued_succeeded = True + + # Emit the held finish event only if the truncated rescue didn't + # replace it with its own finish_reason=tool_calls event. + if self._pending_finish: + if not rescued_succeeded: + outbound.append(self._pending_finish) + self._pending_finish = None + + return b"".join(outbound) + async def _handle_eof(self) -> bytes | None: """Handle EOF: flush buffers and return any leftover.""" await self._flush_all() + outbound: list[bytes] = [self._resolve_stream_end()] + if self._buffer: leftover = self._buffer self._buffer = b"" - return leftover - return b"" + outbound.append(leftover) + + return b"".join(outbound) @staticmethod def _find_event_boundary(buffer: bytes): @@ -310,10 +350,11 @@ async def _log_delta(self, obj: dict) -> None: 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 not self._content_buffer_active: + 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: @@ -352,6 +393,11 @@ async def _readany_once(self) -> bytes | None: if payload == "[DONE]": await self._flush_all() await self._chat_logger.log_response("[DONE]", True) + # Resolve any held state BEFORE forwarding [DONE]: clients + # stop reading at [DONE], so buffered content, a pending + # truncated-tool-call rescue, and a held finish event must + # all be emitted ahead of it. + outbound.append(self._resolve_stream_end()) outbound.append(raw_event) continue try: @@ -364,7 +410,12 @@ async def _readany_once(self) -> bytes | None: # --- (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 + # readany()'s contract: b"" means EOF. A transform may legitimately + # swallow an event (held finish, buffered content) and produce no + # bytes — report None so the read loop keeps going instead of + # signalling a premature end-of-stream. + joined = b"".join(outbound) + return joined if joined else None # ── transform helpers ── @@ -467,17 +518,183 @@ def _build_outbound_event(self, obj: dict, prose_parts: list[str], synthesized: result += syn return result + # ── Fix 1: content-side think-leak buffer ── + + @staticmethod + def _find_closing_tag_line_index(text: str) -> int | None: + """Find the line index of a standalone "" line outside code fences. + + Returns the 0-based line index or None if not found. + The tag line must be complete (terminated by \n or end-of-string). + """ + lines = text.split("\n") + in_fence = False + for i, line in enumerate(lines): + stripped = line.strip() + if stripped.startswith("```"): + in_fence = not in_fence + continue + if not in_fence and stripped == "": + return i + return None + + def _flush_content_buffer(self) -> bytes: + """Flush the content buffer, possibly splitting at "" tag. + + Returns outbound SSE event bytes (reasoning event, content event, or both). + """ + buf = self._content_buf + self._content_buf = "" + self._content_buffer_decided = True + + tag_line_idx = self._find_closing_tag_line_index(buf) + if tag_line_idx is not None: + lines = buf.split("\n") + before_lines = lines[:tag_line_idx] + after_lines = lines[tag_line_idx + 1:] + + reasoning_text = "\n".join(before_lines) + if reasoning_text and not reasoning_text.endswith("\n"): + reasoning_text += "\n" + + after_text = "\n".join(after_lines) + after_text = after_text.lstrip("\n") + + reasoning_event = self._build_content_buffer_event( + reasoning_content=reasoning_text + ) + + if after_text: + content_event = self._build_content_buffer_event( + content=after_text + ) + return reasoning_event + content_event + return reasoning_event + else: + if buf: + return self._build_content_buffer_event(content=buf) + return b"" + + def _build_content_buffer_event(self, reasoning_content="", content="") -> bytes: + """Build a synthetic SSE event for content buffer output.""" + obj = { + "id": self._last_chunk_id, + "object": "chat.completion.chunk", + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": None, + }], + } + delta = obj["choices"][0]["delta"] + if reasoning_content: + delta["reasoning_content"] = reasoning_content + if content: + delta["content"] = content + body = json.dumps(obj, ensure_ascii=False) + return ("data: " + body + "\r\n\r\n").encode() + + def _handle_content_delta(self, content_text: str) -> bytes | None: + """Process a content delta through the think-leak buffer. + + Returns outbound SSE event bytes if the buffer was flushed, + or None if still buffering. + """ + if self._content_buffer_decided: + return None + + self._content_buf += content_text + buf_bytes = len(self._content_buf.encode("utf-8")) + + if buf_bytes >= THINK_LEAK_WINDOW: + return self._flush_content_buffer() + + if self._find_closing_tag_line_index(self._content_buf) is not None: + return self._flush_content_buffer() + + return None + + # ── Fix 2: truncated tool-call rescue at EOF ── + + def _try_rescue_truncated(self) -> bytes | None: + """Try to synthesize a tool call from an incomplete capture at EOF. + + Returns synthesized SSE event bytes if successful, None otherwise. + """ + buf = self._rescue_buf + import re as _re + m = _re.search(r"\s]+)", buf) + if not m: + return None + + fn_end = buf.find("") + if fn_end == -1: + return None + + residue = buf[fn_end + len(""):] + trimmed_residue = residue.strip() + closing_tag = "" + if not closing_tag.startswith(trimmed_residue): + return None + + block = buf[:fn_end + len("")] + parsed = self._parse_tool_call_xml(block) + if not parsed: + return None + + synthesized = self._build_synthesized_event(parsed) + self._rescued_any = True + + self._rescue_capturing = False + self._rescue_buf = "" + self._rescue_kind = None + + finish_obj = { + "id": self._last_chunk_id, + "object": "chat.completion.chunk", + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": "tool_calls", + }], + } + finish_bytes = ("data: " + json.dumps(finish_obj, ensure_ascii=False) + "\r\n\r\n").encode() + + return synthesized + finish_bytes + def _transform_event(self, obj: dict) -> bytes: """Transform a single parsed event dict into outbound SSE bytes, - applying the rescue state machine.""" + applying the rescue state machine and content-side think-leak buffer.""" choices = obj.get("choices") or [] if not choices: return self._transform_pass_through(obj) delta = choices[0].get("delta") or {} + content = delta.get("content") reasoning = delta.get("reasoning_content") + + if content is not None and not self._content_buffer_decided: + self._content_buffer_active = True + result = self._handle_content_delta(content) + if result is not None: + return result + # Still buffering — swallow the content from this event; anything + # else it carries (finish_reason etc.) falls through below. + delta.pop("content", None) + if not reasoning: - return self._transform_no_reasoning(obj) + finish_reason = choices[0].get("finish_reason") + # A finish event resolves the content buffer first, so rerouted or + # flushed text is emitted before the finish_reason reaches clients. + prefix = b"" + if finish_reason and self._content_buffer_active and not self._content_buffer_decided: + prefix = self._flush_content_buffer() + # Hold finish events while a rescue capture is pending — the + # truncated-rescue at EOF may need to rewrite the finish_reason. + if finish_reason and self._rescue_capturing: + self._pending_finish = self._transform_no_reasoning(obj) + return prefix + return prefix + self._transform_no_reasoning(obj) prose_parts, synthesized = self._run_rescue_loop(reasoning) return self._build_outbound_event(obj, prose_parts, synthesized) diff --git a/tests/_helpers.py b/tests/_helpers.py new file mode 100644 index 0000000..50d2b29 --- /dev/null +++ b/tests/_helpers.py @@ -0,0 +1,112 @@ +"""Shared test helpers for SSEChunkLogger unit tests. + +Extracted from tests/test_toolcall_rescue.py to avoid duplication across +test modules. +""" + +from __future__ import annotations + +import asyncio +import json +from typing import List + +from chat_logger import SSEChunkLogger + + +# --------------------------------------------------------------------------- +# 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] = [] + + 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] = [] + 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]) -> SSEChunkLogger: + return SSEChunkLogger(FakeUpstream(chunks), FakeChatLogger()) + + +async def _drive(chunks: List[bytes]) -> tuple: + """Run SSEChunkLogger through all chunks, return (combined output str, logger).""" + logger = FakeChatLogger() + wrapped = 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 diff --git a/tests/test_think_leak_replication.py b/tests/test_think_leak_replication.py new file mode 100644 index 0000000..d8d9620 --- /dev/null +++ b/tests/test_think_leak_replication.py @@ -0,0 +1,505 @@ +"""Replication test harness for think-tag leaks and thinking-only stalls (issue #8). + +Two categories: + Category 1 — REPLICATION tests: replicate real production failure streams + and assert the corrected behaviour (implemented in chat_logger.py). + Category 2 — CONTROL tests: guard against regressions / over-eager fixes + (legit tag mentions, fences, pure stalls, normal streams). +""" + +from __future__ import annotations + +import asyncio +import json +import re +import unittest + +from tests._helpers import ( + ChunkedContent, + FakeChatLogger, + FakeUpstream, + _drive, + _sse_event, + _split_sse_events, +) + + +# --------------------------------------------------------------------------- +# Shared fixture builders +# --------------------------------------------------------------------------- + +def _base_chunk(id_: str = "req1") -> dict: + """Minimal chat.completion.chunk envelope.""" + return { + "id": id_, + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {}, "finish_reason": None}], + } + + +def _reasoning_chunk(text: str, id_: str = "req1") -> bytes: + """SSE event with a reasoning_content delta.""" + obj = _base_chunk(id_) + obj["choices"][0]["delta"]["reasoning_content"] = text + return _sse_event(obj) + + +def _content_chunk(text: str, id_: str = "req1") -> bytes: + """SSE event with a content delta.""" + obj = _base_chunk(id_) + obj["choices"][0]["delta"]["content"] = text + return _sse_event(obj) + + +def _finish_chunk(finish_reason: str = "stop", id_: str = "req1") -> bytes: + """SSE event with finish_reason and empty delta.""" + obj = _base_chunk(id_) + obj["choices"][0]["delta"] = {} + obj["choices"][0]["finish_reason"] = finish_reason + return _sse_event(obj) + + +def _collect_reasoning(events: list) -> str: + """Concatenate all reasoning_content deltas from *events*.""" + parts: list[str] = [] + for e in events: + rc = e.get("choices", [{}])[0].get("delta", {}).get("reasoning_content") + if rc: + parts.append(rc) + return "".join(parts) + + +def _collect_content(events: list) -> str: + """Concatenate all content deltas from *events*.""" + parts: list[str] = [] + for e in events: + c = e.get("choices", [{}])[0].get("delta", {}).get("content") + if c: + parts.append(c) + return "".join(parts) + + +def _has_tool_calls(events: list) -> list: + """Return all tool_calls entries found across events.""" + results: list = [] + for e in events: + tcs = e.get("choices", [{}])[0].get("delta", {}).get("tool_calls") + if tcs: + results.extend(tcs) + return results + + +def _finish_reasons(events: list) -> list: + """Return all non-None finish_reason values.""" + reasons: list[str] = [] + for e in events: + fr = e.get("choices", [{}])[0].get("finish_reason") + if fr: + reasons.append(fr) + return reasons + + +# =================================================================== +# Category 1: REPLICATION tests (real production failure streams) +# =================================================================== + +class TestThinkLeakReplication(unittest.TestCase): + """Tests replicating real production think-tag leaks and thinking-only + stalls, asserting the corrected stream behaviour.""" + + # --- Shape A: double (leak) --- + + def test_shapeA_draft_rerouted_to_reasoning(self) -> None: + """Content carrying text before a standalone line should be + rerouted: pre-tag text → reasoning_content, tag dropped, + post-tag text → normal content. + + Real production seam (session 8d62feae): + reasoning ends "...ically check the abort signal while waiting for user input.\\n" + content = "\\nNow I have the complete picture. The freeze is a **deadlock**:\\n\\n\\nGot it — this is a deadlock. Let me update the issue draft with the precise root cause.\\n\\n" + """ + reasoning_tail = ( + "Let me think through this carefully. " + "I need to check the abort signal while waiting for user input.\n" + ) + # The leaked content: draft answer + standalone + final answer + leaked_content = ( + "\nNow I have the complete picture. The freeze is a **deadlock**:\n" + "\n\n" + "Got it — this is a deadlock. Let me update the issue draft with the precise root cause.\n\n" + ) + + chunks = [ + _reasoning_chunk(reasoning_tail), + _content_chunk(leaked_content), + _finish_chunk(), + ] + + raw, _ = asyncio.run(_drive(chunks)) + events = _split_sse_events(raw.encode()) + + # Desired: reasoning = original thinking + draft (text before ) + combined_reasoning = _collect_reasoning(events) + expected_reasoning = reasoning_tail + "\nNow I have the complete picture. The freeze is a **deadlock**:\n" + self.assertEqual(combined_reasoning, expected_reasoning) + + # Desired: content = final answer only, NO literal + combined_content = _collect_content(events) + expected_content = "Got it — this is a deadlock. Let me update the issue draft with the precise root cause.\n\n" + self.assertEqual(combined_content, expected_content) + self.assertNotIn("", combined_content) + + # --- Shape B: mention consumed as real tag (leak) --- + + def test_shapeB_mention_leak_rerouted(self) -> None: + """Same standalone-line rule applied to Shape B: everything in content + up to and including the standalone line goes to reasoning; + the rest is content. + + Real production seam: the model's thinking discusses `` inside + backticks. llama-server's parser closes reasoning at the mention and + eats the mention text. + """ + # Reasoning ends mid-code-span (parser ate the rest) + reasoning_part = ( + "I need to check if there's a " + ) + # Content starts where reasoning was cut, has the real close tag + # appear literally later, then the final answer + content_part = ( + "`` variant (without brackets) that might be what they mean\n" + "Let me extract them.\n" + "\n\n" + "Here are the extracted items.\n\n" + ) + + chunks = [ + _reasoning_chunk(reasoning_part), + _content_chunk(content_part), + _finish_chunk(), + ] + + raw, _ = asyncio.run(_drive(chunks)) + events = _split_sse_events(raw.encode()) + + # Desired: reasoning = original + everything up to standalone + combined_reasoning = _collect_reasoning(events) + expected_reasoning = ( + reasoning_part + + "`` variant (without brackets) that might be what they mean\n" + "Let me extract them.\n" + ) + self.assertEqual(combined_reasoning, expected_reasoning) + + # Desired: content = text after the standalone only + combined_content = _collect_content(events) + expected_content = "Here are the extracted items.\n\n" + self.assertEqual(combined_content, expected_content) + self.assertNotIn("", combined_content) + + # --- Shape A split across chunks --- + + def test_shapeA_tag_split_across_chunks(self) -> None: + """Same as test_shapeA but the literal is split across + content deltas and SSE event boundaries.""" + reasoning_text = "Analyzing the problem step by step.\n" + + # The tag is split: "" in the next + draft_before_tag = "\nDraft answer here.\n) + combined_reasoning = _collect_reasoning(events) + expected_reasoning = reasoning_text + "\nDraft answer here.\n" + self.assertEqual(combined_reasoning, expected_reasoning) + + # Desired: content = final answer only, no + combined_content = _collect_content(events) + expected_content = "Final answer.\n\n" + self.assertEqual(combined_content, expected_content) + self.assertNotIn("", combined_content) + + # --- Shape C1: trapped tool-call rescued --- + # NOTE: this is NOT expectedFailure — the existing rescue state machine + # in SSEChunkLogger already handles tool-call XML (...) in + # reasoning_content. It passes today and must keep passing. + + def test_shapeC1_trapped_toolcall_rescued(self) -> None: + """Reasoning-only stream whose reasoning ends with a complete + ... block and finish_reason=stop with empty content: the + client should receive a synthesized tool_call and the XML + stripped from reasoning_content tail. + + 37 of 94 real stalls look like this — the model wrote its tool + call inside the unclosed think block. + """ + reasoning_prose = "Let me check the directory listing.\n" + xml_block = ( + "\n" + "\n" + "\n" + "ls C:\\some\\path\n" + "\n" + "\n" + "" + ) + + chunks = [ + _reasoning_chunk(reasoning_prose), + _reasoning_chunk(xml_block), + _finish_chunk(), + ] + + raw, _ = asyncio.run(_drive(chunks)) + events = _split_sse_events(raw.encode()) + + # Desired: tool call is synthesized + tc_list = _has_tool_calls(events) + self.assertEqual(len(tc_list), 1) + tc = tc_list[0] + self.assertEqual(tc["function"]["name"], "bash") + args = json.loads(tc["function"]["arguments"]) + self.assertEqual(args["command"], "ls C:\\some\\path") + + # Desired: reasoning = prose only, XML stripped + combined_reasoning = _collect_reasoning(events) + self.assertEqual(combined_reasoning, reasoning_prose) + self.assertNotIn("", combined_reasoning) + self.assertNotIn("", combined_reasoning) + self.assertNotIn(" None: + """Reasoning-only stream whose reasoning ends with a tool-call XML + block whose closing tag is truncated (missing final '>') + and finish_reason=stop with empty content: the tool call should still + be synthesized, the XML (including truncated tail) stripped from + reasoning_content, and finish_reason rewritten to tool_calls. + + Real production seam — the model stalled inside an unclosed think + block and the stream ended with the tool-call XML cut off mid + closing tag. + """ + reasoning_prose = "Let me list the docs directory.\n" + xml_block = ( + "\n" + "\n" + "\n" + 'ls "C:\\Users\\HTK\\AppData\\Roaming\\npm\\node_modules\\@earendil-works\\pi-coding-agent\\docs"\n' + "\n" + "\n" + "' + ) + + chunks = [ + _reasoning_chunk(reasoning_prose), + _reasoning_chunk(xml_block), + _finish_chunk(), + ] + + raw, _ = asyncio.run(_drive(chunks)) + events = _split_sse_events(raw.encode()) + + # Desired: tool call is synthesized despite truncated closing tag + tc_list = _has_tool_calls(events) + self.assertEqual(len(tc_list), 1) + tc = tc_list[0] + self.assertEqual(tc["function"]["name"], "bash") + args = json.loads(tc["function"]["arguments"]) + self.assertIn("command", args) + self.assertIn("ls", args["command"]) + self.assertIn("pi-coding-agent", args["command"]) + + # Desired: reasoning = prose only, XML stripped (including truncated tail) + combined_reasoning = _collect_reasoning(events) + self.assertEqual(combined_reasoning, reasoning_prose) + self.assertNotIn("", combined_reasoning) + self.assertNotIn(" None: + """A normal stream whose content legitimately mentions the tag inline + in backticks with no standalone-line tag: must pass through byte-identical. + Guard against over-eager fixes.""" + reasoning_text = "I've analyzed the problem.\n" + content_text = ( + "Some models leak `` into output. " + "This is just a mention of the closing tag, not an actual one.\n" + ) + + chunks = [ + _reasoning_chunk(reasoning_text), + _content_chunk(content_text), + _finish_chunk(), + ] + + raw, _ = asyncio.run(_drive(chunks)) + events = _split_sse_events(raw.encode()) + + # Content must be byte-identical + combined_content = _collect_content(events) + self.assertEqual(combined_content, content_text) + + # Reasoning must be intact + combined_reasoning = _collect_reasoning(events) + self.assertEqual(combined_reasoning, reasoning_text) + + # --- 6. Legit fenced block mention untouched --- + + def test_legit_fenced_block_mention_untouched(self) -> None: + """Content containing a fenced code block in which sits alone + on a line must pass through unmodified. The standalone-line rule must + not fire inside a code fence. + + NOTE: today's passthrough makes this green; a naive fix would break it. + """ + reasoning_text = "Here is an example.\n" + content_text = ( + "Example of the bug:\n" + "```text\n" + "thinking...\n" + "\n" + "```\n" + "That was the example.\n" + ) + + chunks = [ + _reasoning_chunk(reasoning_text), + _content_chunk(content_text), + _finish_chunk(), + ] + + raw, _ = asyncio.run(_drive(chunks)) + events = _split_sse_events(raw.encode()) + + combined_content = _collect_content(events) + self.assertEqual(combined_content, content_text) + + combined_reasoning = _collect_reasoning(events) + self.assertEqual(combined_reasoning, reasoning_text) + + # --- 7. Shape C2: pure stall passthrough --- + + def test_shapeC2_pure_stall_passthrough(self) -> None: + """Reasoning-only stream with NO tool-call XML: passes through + unchanged (we do not invent content). Green today, must stay green.""" + reasoning_text = ( + "Let me think about this problem. " + "There are several approaches we could take. " + "I need more information to proceed.\n" + ) + + chunks = [ + _reasoning_chunk(reasoning_text), + _finish_chunk(), + ] + + raw, _ = asyncio.run(_drive(chunks)) + events = _split_sse_events(raw.encode()) + + # Reasoning passes through intact + combined_reasoning = _collect_reasoning(events) + self.assertEqual(combined_reasoning, reasoning_text) + + # No content invented + combined_content = _collect_content(events) + self.assertEqual(combined_content, "") + + # No tool calls synthesized + tc_list = _has_tool_calls(events) + self.assertEqual(len(tc_list), 0) + + # finish_reason stays "stop" + reasons = _finish_reasons(events) + self.assertIn("stop", reasons) + + # --- 8. Normal stream with reasoning untouched --- + + def test_normal_stream_with_reasoning_untouched(self) -> None: + """Ordinary reasoning + content + native tool_calls stream passes + through unchanged (extension of existing passthrough test but with + reasoning deltas).""" + chunks = [ + _sse_event({ + "id": "req1", "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], + }), + _reasoning_chunk("Let me think about this carefully.\n"), + _reasoning_chunk("I have a plan.\n"), + _content_chunk("Here is the answer.\n"), + _content_chunk("It should be correct.\n"), + _sse_event({ + "id": "req1", "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": { + "tool_calls": [{"index": 0, "id": "call_xyz", "type": "function", + "function": {"name": "write", "arguments": '{"path":"/out"}'}}], + }, "finish_reason": None}], + }), + _finish_chunk(), + ] + + raw, _ = asyncio.run(_drive(chunks)) + events = _split_sse_events(raw.encode()) + + combined_reasoning = _collect_reasoning(events) + self.assertEqual(combined_reasoning, "Let me think about this carefully.\nI have a plan.\n") + + combined_content = _collect_content(events) + self.assertEqual(combined_content, "Here is the answer.\nIt should be correct.\n") + + # Tool call preserved + tc_list = _has_tool_calls(events) + self.assertEqual(len(tc_list), 1) + self.assertEqual(tc_list[0]["function"]["name"], "write") + + # finish_reason stays "stop" + reasons = _finish_reasons(events) + self.assertIn("stop", reasons) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_toolcall_rescue.py b/tests/test_toolcall_rescue.py index 2b7b733..727bc7f 100644 --- a/tests/test_toolcall_rescue.py +++ b/tests/test_toolcall_rescue.py @@ -7,106 +7,15 @@ import unittest from chat_logger import SSEChunkLogger - - -# --------------------------------------------------------------------------- -# 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]) -> SSEChunkLogger: - return 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 = 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 +from tests._helpers import ( + ChunkedContent, + FakeChatLogger, + FakeUpstream, + _drive, + _make_logger, + _sse_event, + _split_sse_events, +) # ---------------------------------------------------------------------------