diff --git a/headroom/ccr/response_handler.py b/headroom/ccr/response_handler.py index 7cdea7a49..a4acfeb3d 100644 --- a/headroom/ccr/response_handler.py +++ b/headroom/ccr/response_handler.py @@ -571,11 +571,23 @@ class StreamingCCRBuffer: chunks: list[bytes] = field(default_factory=list) detected_ccr: bool = False complete_response: dict[str, Any] | None = None + provider: str = "anthropic" - # Patterns to detect tool_use in stream + # Wire markers for the start of a tool call. Anthropic streams + # `"type":"tool_use"` content blocks; OpenAI-compatible streams carry a + # `"tool_calls"` array inside `choices[].delta` and never emit the + # Anthropic marker, so scanning only for the latter meant CCR was never + # detected on an OpenAI stream. _tool_use_start: bytes = b'"type":"tool_use"' + _openai_tool_use_start: bytes = b'"tool_calls"' _ccr_tool_pattern: bytes = f'"{CCR_TOOL_NAME}"'.encode() + def _tool_call_marker(self) -> bytes: + """The provider's on-the-wire marker for the start of a tool call.""" + if self.provider == "anthropic": + return self._tool_use_start + return self._openai_tool_use_start + def add_chunk(self, chunk: bytes) -> bool: """Add a chunk and check for CCR tool calls. @@ -587,7 +599,7 @@ class StreamingCCRBuffer: # Quick check: does accumulated content contain CCR tool? accumulated = b"".join(self.chunks) - if self._tool_use_start in accumulated and self._ccr_tool_pattern in accumulated: + if self._tool_call_marker() in accumulated and self._ccr_tool_pattern in accumulated: self.detected_ccr = True return True @@ -622,7 +634,7 @@ class StreamingCCRHandler: ) -> None: self.response_handler = response_handler self.provider = provider - self.buffer = StreamingCCRBuffer() + self.buffer = StreamingCCRBuffer(provider=provider) async def process_stream( self, @@ -648,8 +660,12 @@ class StreamingCCRHandler: Response chunks (possibly from continuation response). """ # Phase 1: Initial detection - # Buffer chunks until we can determine if there's a CCR call - detection_complete = False + # Buffer chunks until we can determine if there's a CCR call. + # + # The end-of-stream marker is provider-specific. Anthropic signals the + # terminal state with `stop_reason` in `message_delta`; OpenAI-compatible + # streams have no such field and terminate with the `[DONE]` sentinel. + end_marker = b'"stop_reason"' if self.provider == "anthropic" else b"data: [DONE]" async for chunk in stream_iterator: self.buffer.add_chunk(chunk) @@ -660,9 +676,7 @@ class StreamingCCRHandler: accumulated = self.buffer.get_accumulated() # Look for stream end markers - if b'"stop_reason"' in accumulated: - detection_complete = True - + if end_marker in accumulated: if self.buffer.detected_ccr: # CCR detected - need to handle break @@ -679,13 +693,15 @@ class StreamingCCRHandler: yield buffered_chunk self.buffer.clear() - # Continue streaming rest of response - if not detection_complete and not self.buffer.detected_ccr: - async for chunk in stream_iterator: - if self.buffer.detected_ccr: - self.buffer.add_chunk(chunk) - else: - yield chunk + # The end marker is not guaranteed to arrive: upstream can truncate, a + # provider can omit the sentinel, or the stream can be a shape this + # detector does not recognise. Anything still buffered once the source + # iterator is exhausted is real response data the client has never + # seen, so flush it instead of dropping it. + if not self.buffer.detected_ccr and self.buffer.chunks: + for buffered_chunk in self.buffer.chunks: + yield buffered_chunk + self.buffer.clear() # Phase 2: Handle CCR if detected if self.buffer.detected_ccr: @@ -903,13 +919,38 @@ class StreamingCCRHandler: } tool_calls_map: dict[int, dict[str, Any]] = {} + finish_reason: str | None = None + envelope: dict[str, Any] = {} + usage: Any = None for event in events: - choices = event.get("choices", []) - if not choices: + # Carry the chunk envelope through. Dropping it left the + # reconstructed body without `id`, `model`, `created` or `usage`, + # which downstream middleware reads for routing and metering. + for key in ("id", "created", "model", "system_fingerprint"): + value = event.get(key) + if value is not None: + envelope[key] = value + if event.get("usage") is not None: + usage = event["usage"] + + choices = event.get("choices") + if not isinstance(choices, list) or not choices: + continue + choice = choices[0] + if not isinstance(choice, dict): continue - delta = choices[0].get("delta", {}) + # `finish_reason` is null on every chunk but the last, so keep the + # most recent non-null value rather than the first one seen. + if choice.get("finish_reason") is not None: + finish_reason = choice["finish_reason"] + + # Some OpenAI-compatible providers send `"delta": null` on the + # terminal chunk instead of an empty object. + delta = choice.get("delta") + if not isinstance(delta, dict): + delta = {} if "content" in delta and delta["content"]: message["content"] = (message.get("content") or "") + delta["content"] @@ -944,14 +985,94 @@ class StreamingCCRHandler: tc["function"]["arguments"] += fn["arguments"] message["tool_calls"] = [tool_calls_map[i] for i in sorted(tool_calls_map.keys())] - if not message["tool_calls"]: + has_tool_calls = bool(message["tool_calls"]) + if not has_tool_calls: del message["tool_calls"] if not message["content"]: message["content"] = None - return { - "choices": [{"message": message, "finish_reason": "stop"}], + # OpenAI requires `finish_reason: "tool_calls"` whenever the message + # carries tool calls. This was hardcoded to "stop", which tells any + # client that drives its agent loop off `finish_reason` that the turn + # is over, so the reconstructed tool calls were never executed. + if has_tool_calls: + finish_reason = "tool_calls" + elif finish_reason is None: + finish_reason = "stop" + + response: dict[str, Any] = { + "object": "chat.completion", + **envelope, + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], } + if usage is not None: + response["usage"] = usage + return response + + def _openai_response_to_chunks(self, response: dict[str, Any]) -> list[bytes]: + """Split a non-streaming ``chat.completion`` body into SSE chunk frames. + + A streaming client reads ``choices[].delta``, not ``choices[].message``. + Serialising the reconstructed non-streaming body into a single SSE frame + produced a stream in which both the text and the tool calls were + invisible to the client. + """ + choices = response.get("choices") + choice = choices[0] if isinstance(choices, list) and choices else {} + if not isinstance(choice, dict): + choice = {} + message = choice.get("message") + if not isinstance(message, dict): + message = {} + finish_reason = choice.get("finish_reason") or "stop" + + base: dict[str, Any] = {"object": "chat.completion.chunk"} + for key in ("id", "created", "model", "system_fingerprint"): + if response.get(key) is not None: + base[key] = response[key] + + def frame(delta: dict[str, Any], reason: str | None) -> bytes: + payload = { + **base, + "choices": [{"index": 0, "delta": delta, "finish_reason": reason}], + } + return f"data: {json.dumps(payload)}\n\n".encode() + + frames = [frame({"role": message.get("role") or "assistant"}, None)] + + content = message.get("content") + if content: + frames.append(frame({"content": content}, None)) + + tool_calls = message.get("tool_calls") + if isinstance(tool_calls, list): + for index, tool_call in enumerate(tool_calls): + if not isinstance(tool_call, dict): + continue + function = tool_call.get("function") + if not isinstance(function, dict): + function = {} + frames.append( + frame( + { + "tool_calls": [ + { + "index": index, + "id": tool_call.get("id", ""), + "type": tool_call.get("type", "function"), + "function": { + "name": function.get("name", ""), + "arguments": function.get("arguments", ""), + }, + } + ] + }, + None, + ) + ) + + frames.append(frame({}, finish_reason)) + return frames async def _response_to_sse( self, @@ -968,6 +1089,7 @@ class StreamingCCRHandler: for chunk in StreamingMixin()._response_to_sse(response, "anthropic"): yield chunk else: - # OpenAI SSE format - yield f"data: {json.dumps(response)}\n\n".encode() + # OpenAI SSE format: `chat.completion.chunk` frames, then [DONE]. + for chunk in self._openai_response_to_chunks(response): + yield chunk yield b"data: [DONE]\n\n" diff --git a/tests/test_ccr_response_handler_extra.py b/tests/test_ccr_response_handler_extra.py index 1e8d97d93..c2754c413 100644 --- a/tests/test_ccr_response_handler_extra.py +++ b/tests/test_ccr_response_handler_extra.py @@ -445,7 +445,15 @@ async def test_streaming_handler_falls_back_to_buffer_on_processing_error( lambda data: (_ for _ in ()).throw(RuntimeError("parse failed")), ) - chunks = [b'{"type":"tool_use","name":"headroom_retrieve"', b',"stop_reason":"tool_use"}'] + # Real OpenAI wire shape: a `tool_calls` delta naming the CCR tool, then + # the `[DONE]` sentinel. This test previously fed Anthropic-shaped bytes to + # an ``openai`` handler, so it never reached the OpenAI detection path. + chunks = [ + b'data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,' + b'"id":"call_1","function":{"name":"headroom_retrieve",' + b'"arguments":"{}"}}]}}]}\n\n', + b"data: [DONE]\n\n", + ] streamed = [ chunk async for chunk in handler.process_stream(_async_iter(chunks), [], None, lambda m, t: None) @@ -462,7 +470,14 @@ async def test_response_to_sse_formats() -> None: openai = StreamingCCRHandler(CCRResponseHandler(), provider="openai") openai_chunks = [chunk async for chunk in openai._response_to_sse({"choices": []})] - assert openai_chunks == [b'data: {"choices": []}\n\n', b"data: [DONE]\n\n"] + # An empty body still produces well-formed chunk frames (role, then a + # terminal frame carrying finish_reason) rather than a single non-streaming + # body a streaming client cannot read. + assert openai_chunks[-1] == b"data: [DONE]\n\n" + frames = [json.loads(chunk.decode()[len("data: ") :]) for chunk in openai_chunks[:-1]] + assert [frame["object"] for frame in frames] == ["chat.completion.chunk"] * 2 + assert frames[0]["choices"][0]["delta"] == {"role": "assistant"} + assert frames[-1]["choices"][0]["finish_reason"] == "stop" @pytest.mark.asyncio @@ -533,3 +548,229 @@ def test_reconstruct_server_tool_use_input_from_partial_json() -> None: assert block["type"] == "server_tool_use" assert block["input"] == {"query": "x"} assert "_partial_json" not in block + + +def _openai_chunk(delta: dict[str, Any], finish_reason: str | None = None) -> bytes: + """One `chat.completion.chunk` SSE frame in the shape a real backend sends.""" + payload = { + "id": "chatcmpl_1", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(payload)}\n\n".encode() + + +def test_streaming_buffer_detects_ccr_in_openai_tool_calls_delta() -> None: + # An OpenAI-compatible stream never emits Anthropic's `"type":"tool_use"` + # marker; its tool calls arrive as a `tool_calls` array inside + # `choices[].delta`. Scanning only for the Anthropic marker meant CCR was + # never detected on this provider. + buffer = StreamingCCRBuffer(provider="openai") + + assert buffer.add_chunk(_openai_chunk({"role": "assistant"})) is False + detected = buffer.add_chunk( + _openai_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "function": {"name": CCR_TOOL_NAME, "arguments": ""}, + } + ] + } + ) + ) + + assert detected is True + assert buffer.detected_ccr is True + + # A non-CCR tool call on the same provider must not trip detection. + other = StreamingCCRBuffer(provider="openai") + assert ( + other.add_chunk( + _openai_chunk( + {"tool_calls": [{"index": 0, "id": "c", "function": {"name": "other_tool"}}]} + ) + ) + is False + ) + assert other.detected_ccr is False + + +@pytest.mark.asyncio +async def test_openai_stream_without_ccr_yields_every_chunk() -> None: + # A short OpenAI stream with no CCR call must pass through byte for byte. + # End-of-stream was detected by scanning for Anthropic's `stop_reason`, + # which an OpenAI stream never contains, so nothing was ever flushed and + # the client received an empty response. + handler = StreamingCCRHandler(CCRResponseHandler(), provider="openai") + chunks = [ + _openai_chunk({"role": "assistant"}), + _openai_chunk({"content": "hello "}), + _openai_chunk({"content": "world"}), + _openai_chunk({}, finish_reason="stop"), + b"data: [DONE]\n\n", + ] + + streamed = [ + chunk + async for chunk in handler.process_stream(_async_iter(chunks), [], None, lambda m, t: None) + ] + + assert streamed == chunks + + +@pytest.mark.asyncio +async def test_openai_stream_past_flush_threshold_keeps_the_tail() -> None: + # Past 10 000 buffered bytes the handler flushes in batches. Without an + # end-of-stream match the final sub-threshold batch was never flushed, so + # a long response visibly stopped mid-sentence. + handler = StreamingCCRHandler(CCRResponseHandler(), provider="openai") + chunks = [ + _openai_chunk({"role": "assistant"}), + _openai_chunk({"content": "x" * 11000}), + _openai_chunk({"content": "the tail that used to be dropped"}), + _openai_chunk({}, finish_reason="stop"), + b"data: [DONE]\n\n", + ] + + streamed = [ + chunk + async for chunk in handler.process_stream(_async_iter(chunks), [], None, lambda m, t: None) + ] + + assert streamed == chunks + assert b"the tail that used to be dropped" in b"".join(streamed) + + +@pytest.mark.asyncio +async def test_openai_stream_without_done_sentinel_still_flushes() -> None: + # Upstream can truncate before `[DONE]`, and some gateways omit it. Bytes + # left in the buffer when the source iterator is exhausted are real + # response data, so they are flushed rather than discarded. + handler = StreamingCCRHandler(CCRResponseHandler(), provider="openai") + chunks = [_openai_chunk({"role": "assistant"}), _openai_chunk({"content": "partial answer"})] + + streamed = [ + chunk + async for chunk in handler.process_stream(_async_iter(chunks), [], None, lambda m, t: None) + ] + + assert streamed == chunks + + +def test_reconstruct_openai_response_marks_tool_calls_finish_reason() -> None: + # OpenAI requires `finish_reason: "tool_calls"` when the message carries + # tool calls. It was hardcoded to "stop", so a client driving its agent + # loop off `finish_reason` ended the turn instead of running the tools. + handler = StreamingCCRHandler(CCRResponseHandler(), provider="openai") + + parsed = handler._reconstruct_openai_response( + [ + {"id": "chatcmpl_1", "model": "gpt-4o-mini", "created": 1700000000}, + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "function": { + "name": CCR_TOOL_NAME, + "arguments": '{"hash":"abc"}', + }, + } + ] + }, + "finish_reason": None, + } + ] + }, + { + "choices": [{"delta": None, "finish_reason": "tool_calls"}], + "usage": {"prompt_tokens": 12, "completion_tokens": 3}, + }, + ] + ) + + assert parsed["choices"][0]["finish_reason"] == "tool_calls" + # The chunk envelope is carried through so the reconstructed body is a + # valid `chat.completion` rather than a bare `choices` list. + assert parsed["object"] == "chat.completion" + assert parsed["id"] == "chatcmpl_1" + assert parsed["model"] == "gpt-4o-mini" + assert parsed["created"] == 1700000000 + assert parsed["usage"] == {"prompt_tokens": 12, "completion_tokens": 3} + + +def test_reconstruct_openai_response_keeps_upstream_finish_reason() -> None: + # With no tool calls, the upstream reason is preserved instead of being + # rewritten to "stop": a truncated turn must stay reported as truncated. + handler = StreamingCCRHandler(CCRResponseHandler(), provider="openai") + + parsed = handler._reconstruct_openai_response( + [ + {"choices": [{"delta": {"content": "half an ans"}, "finish_reason": None}]}, + {"choices": [{"delta": {}, "finish_reason": "length"}]}, + ] + ) + + assert parsed["choices"][0]["finish_reason"] == "length" + assert parsed["choices"][0]["message"]["content"] == "half an ans" + + # And an absent reason still defaults to "stop". + defaulted = handler._reconstruct_openai_response( + [{"choices": [{"delta": {"content": "hi"}}]}], + ) + assert defaulted["choices"][0]["finish_reason"] == "stop" + + +@pytest.mark.asyncio +async def test_response_to_sse_emits_openai_chunk_frames() -> None: + # A streaming client reads `choices[].delta`. Re-serialising the + # reconstructed non-streaming body (`choices[].message`) into one SSE frame + # made both the content and the tool calls invisible to it. + handler = StreamingCCRHandler(CCRResponseHandler(), provider="openai") + response = { + "id": "chatcmpl_2", + "object": "chat.completion", + "created": 1700000001, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "done", + "tool_calls": [ + { + "id": "call_9", + "type": "function", + "function": {"name": "do_thing", "arguments": '{"a":1}'}, + } + ], + }, + "finish_reason": "tool_calls", + } + ], + } + + chunks = [chunk async for chunk in handler._response_to_sse(response)] + + assert chunks[-1] == b"data: [DONE]\n\n" + frames = [json.loads(chunk.decode()[len("data: ") :]) for chunk in chunks[:-1]] + assert all(frame["object"] == "chat.completion.chunk" for frame in frames) + assert all("delta" in frame["choices"][0] for frame in frames) + assert all(frame["id"] == "chatcmpl_2" for frame in frames) + + deltas = [frame["choices"][0]["delta"] for frame in frames] + assert deltas[0] == {"role": "assistant"} + assert deltas[1] == {"content": "done"} + assert deltas[2]["tool_calls"][0]["id"] == "call_9" + assert deltas[2]["tool_calls"][0]["index"] == 0 + assert deltas[2]["tool_calls"][0]["function"] == {"name": "do_thing", "arguments": '{"a":1}'} + assert frames[-1]["choices"][0]["finish_reason"] == "tool_calls"