diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 6291a1cf6..a0fe245a8 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2824,726 +2824,861 @@ class AnthropicHandlerMixin: session_key=session_key, ) else: - async with stage_timer.measure("upstream_connect"): - response = await self._retry_request( - "POST", - url, - headers, - body, - original_body_bytes=original_body_bytes, - body_mutated=body_mutation_tracker.mutated, - mutation_reasons=body_mutation_tracker.reasons, - request_id=request_id, - forwarder_name="anthropic_messages", - path_for_log="/v1/messages", - timeout=self._anthropic_buffered_request_timeout(), - ) - self.pipeline_extensions.emit( - PipelineStage.POST_SEND, - operation="proxy.request", - request_id=request_id, - provider=pipeline_provider, - model=model, - messages=body["messages"], - tools=tools, - response=response, - metadata={ - "path": pipeline_path, - "stream": False, - "client_stream": buffered_stream_ccr, - "ccr_stream_buffered": buffered_stream_ccr, - "status_code": response.status_code, - }, - ) - self.pipeline_extensions.emit( - PipelineStage.RESPONSE_RECEIVED, - operation="proxy.request", - request_id=request_id, - provider=pipeline_provider, - model=model, - response=response, - metadata={ - "path": pipeline_path, - "stream": False, - "client_stream": buffered_stream_ccr, - "ccr_stream_buffered": buffered_stream_ccr, - "status_code": response.status_code, - }, - ) - if ( - "upstream_first_byte" not in stage_timer - and "upstream_connect" in stage_timer - ): - stage_timer.record( - "upstream_first_byte", - stage_timer.summary()["upstream_connect"], - ) - await _finalize_pre_upstream() - # Full diagnostic dump on upstream errors. - # Writes pre/post compression messages, tools, and error - # to ~/.headroom/logs/debug_400/ for offline analysis. - if response.status_code >= 400: - try: - err_body = response.json() - err_msg = err_body.get("error", {}).get("message", "") - err_type = err_body.get("error", {}).get("type", "") - except Exception: - err_body = {"raw": response.text[:2000]} - err_msg = str(response.text[:500]) - err_type = "parse_error" - logger.warning( - f"[{request_id}] UPSTREAM_ERROR " - f"status={response.status_code} " - f"error_type={err_type} " - f"error_msg={err_msg!r} " - f"model={model} " - f"compressed={'yes' if transforms_applied else 'no'} " - f"transforms={transforms_applied} " - f"original_tokens={original_tokens} " - f"optimized_tokens={optimized_tokens} " - f"message_count={len(body.get('messages', []))} " - f"stream={stream}" - ) - - # Diagnostic dump of the full upstream-error request. - # OFF by default: it can contain cleartext prompt / tool / - # system content. Opt in with HEADROOM_DEBUG_DUMP=1 - # (redacted: structure + lengths only) or =full (content). - # Never written in stateless mode. - dump_mode = _debug_dump_mode(self.config) - if dump_mode != "off": - try: - from headroom import paths as _hr_paths - - debug_dir = _hr_paths.debug_400_dir() - debug_dir.mkdir(parents=True, exist_ok=True) - ts = datetime.now().strftime("%Y%m%d_%H%M%S") - debug_file = debug_dir / f"{ts}_{request_id}.json" - - # Sanitize headers (redact API keys) - safe_headers = {} - _sensitive_header_names = {"x-api-key", "authorization"} | { - k.lower() for k in (self.config.anthropic_extra_headers or {}) - } - for k, v in headers.items(): - if k.lower() in _sensitive_header_names: - safe_headers[k] = v[:12] + "..." if v else "" - else: - safe_headers[k] = v - - # In redacted mode, elide prompt/tool/system - # content but keep structure, roles, and lengths. - redact = dump_mode == "redacted" - messages_sent = body.get("messages") - original_dump: Any = ( - original_messages - if original_messages is not body.get("messages") - else "__same_as_sent__" - ) - tools_sent = body.get("tools") - system_prompt = body.get("system") - if redact: - messages_sent = _redact_debug_value(messages_sent) - if original_dump != "__same_as_sent__": - original_dump = _redact_debug_value(original_dump) - tools_sent = _redact_debug_value(tools_sent) - system_prompt = _redact_debug_value(system_prompt) - - debug_payload = { - "request_id": request_id, - "timestamp": datetime.now().isoformat(), - "dump_mode": dump_mode, - "status_code": response.status_code, - "error_response": err_body, - "model": model, - "stream": stream, - "headers": safe_headers, - "compression": { - "was_compressed": bool(transforms_applied), - "transforms": transforms_applied, - "original_tokens": original_tokens, - "optimized_tokens": optimized_tokens, - "tokens_saved": tokens_saved, - "compression_failed": _compression_failed, - }, - "tools_sent": tools_sent, - "tool_count": len(body.get("tools") or []), - "original_tool_count": len(_original_tools or []), - "messages_sent": messages_sent, - "message_count": len(body.get("messages", [])), - "original_messages": original_dump, - "original_message_count": len(original_messages), - "system_prompt": system_prompt, - } - - with open(debug_file, "w") as f: - json.dump(debug_payload, f, indent=2, default=str) - - logger.warning( - f"[{request_id}] Debug dump ({dump_mode}): {debug_file}" - ) - except Exception as dump_err: - logger.error( - f"[{request_id}] Failed to write debug dump: {dump_err}" - ) - - # Parse response for CCR handling - resp_json = None - try: - resp_json = response.json() - except (json.JSONDecodeError, ValueError) as e: - logger.debug( - f"[{request_id}] Failed to parse response JSON for CCR handling: {e}" - ) - - # CCR Response Handling: Handle headroom_retrieve tool calls automatically - if ( - self.ccr_response_handler - and resp_json - and response.status_code == 200 - and self.ccr_response_handler.has_ccr_tool_calls(resp_json, "anthropic") - ): - logger.info( - f"[{request_id}] CCR: Detected retrieval tool call, handling..." - ) - - # Create API call function for continuation - # Use a fresh client to avoid potential decompression state issues - async def api_call_fn( - msgs: list[dict], tls: list[dict] | None - ) -> dict[str, Any]: - continuation_body = { - **body, - "messages": msgs, - } - if tls is not None: - continuation_body["tools"] = tls - - # Use clean headers for continuation - continuation_headers = { - k: v - for k, v in headers.items() - if k.lower() - not in ( - "content-encoding", - "transfer-encoding", - "accept-encoding", - "content-length", - ) - } - - # Reuse main client for CCR continuations (connection pooling) - logger.info( - f"CCR: Making continuation request with {len(msgs)} messages" - ) - assert self.http_client is not None, "HTTP client not initialized" - # Byte-faithful (PR-A3, fixes P0-2). The CCR - # continuation body is synthesized by Headroom - # so it is treated as mutated and goes through - # the canonical serializer. - from headroom.proxy.body_forwarding import ( - prepare_outbound_body_bytes, - ) - from headroom.proxy.helpers import log_outbound_request - - ccr_outbound_bytes, ccr_outbound_source = prepare_outbound_body_bytes( - body=continuation_body, - original_body_bytes=None, - body_mutated=True, - ) - ccr_outbound_headers = { - **continuation_headers, - "content-type": "application/json", - } - log_outbound_request( - forwarder="anthropic_ccr_continuation", - method="POST", - path=url, - body_bytes_count=len(ccr_outbound_bytes), - body_mutated=True, - mutation_reasons=["ccr_continuation"], + async def _buffered_ccr_operation(): + async with stage_timer.measure("upstream_connect"): + response = await self._retry_request( + "POST", + url, + headers, + body, + original_body_bytes=original_body_bytes, + body_mutated=body_mutation_tracker.mutated, + mutation_reasons=body_mutation_tracker.reasons, request_id=request_id, - source=ccr_outbound_source, + forwarder_name="anthropic_messages", + path_for_log="/v1/messages", + timeout=self._anthropic_buffered_request_timeout(), ) + self.pipeline_extensions.emit( + PipelineStage.POST_SEND, + operation="proxy.request", + request_id=request_id, + provider=pipeline_provider, + model=model, + messages=body["messages"], + tools=tools, + response=response, + metadata={ + "path": pipeline_path, + "stream": False, + "client_stream": buffered_stream_ccr, + "ccr_stream_buffered": buffered_stream_ccr, + "status_code": response.status_code, + }, + ) + self.pipeline_extensions.emit( + PipelineStage.RESPONSE_RECEIVED, + operation="proxy.request", + request_id=request_id, + provider=pipeline_provider, + model=model, + response=response, + metadata={ + "path": pipeline_path, + "stream": False, + "client_stream": buffered_stream_ccr, + "ccr_stream_buffered": buffered_stream_ccr, + "status_code": response.status_code, + }, + ) + if ( + "upstream_first_byte" not in stage_timer + and "upstream_connect" in stage_timer + ): + stage_timer.record( + "upstream_first_byte", + stage_timer.summary()["upstream_connect"], + ) + await _finalize_pre_upstream() + # Full diagnostic dump on upstream errors. + # Writes pre/post compression messages, tools, and error + # to ~/.headroom/logs/debug_400/ for offline analysis. + if response.status_code >= 400: try: - cont_response = await self.http_client.post( - url, - content=ccr_outbound_bytes, - headers=ccr_outbound_headers, - timeout=self._anthropic_buffered_request_timeout(), - ) - logger.info( - f"CCR: Got response status={cont_response.status_code}, " - f"content-encoding={cont_response.headers.get('content-encoding')}" - ) - result: dict[str, Any] = cont_response.json() - logger.info("CCR: Parsed JSON successfully") - return result - except Exception as e: - resp_headers: str | dict[str, str] = "N/A" + err_body = response.json() + err_msg = err_body.get("error", {}).get("message", "") + err_type = err_body.get("error", {}).get("type", "") + except Exception: + err_body = {"raw": response.text[:2000]} + err_msg = str(response.text[:500]) + err_type = "parse_error" + + logger.warning( + f"[{request_id}] UPSTREAM_ERROR " + f"status={response.status_code} " + f"error_type={err_type} " + f"error_msg={err_msg!r} " + f"model={model} " + f"compressed={'yes' if transforms_applied else 'no'} " + f"transforms={transforms_applied} " + f"original_tokens={original_tokens} " + f"optimized_tokens={optimized_tokens} " + f"message_count={len(body.get('messages', []))} " + f"stream={stream}" + ) + + # Diagnostic dump of the full upstream-error request. + # OFF by default: it can contain cleartext prompt / tool / + # system content. Opt in with HEADROOM_DEBUG_DUMP=1 + # (redacted: structure + lengths only) or =full (content). + # Never written in stateless mode. + dump_mode = _debug_dump_mode(self.config) + if dump_mode != "off": try: - resp_headers = dict(cont_response.headers) - except Exception: - pass + from headroom import paths as _hr_paths + + debug_dir = _hr_paths.debug_400_dir() + debug_dir.mkdir(parents=True, exist_ok=True) + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + debug_file = debug_dir / f"{ts}_{request_id}.json" + + # Sanitize headers (redact API keys) + safe_headers = {} + _sensitive_header_names = {"x-api-key", "authorization"} | { + k.lower() + for k in (self.config.anthropic_extra_headers or {}) + } + for k, v in headers.items(): + if k.lower() in _sensitive_header_names: + safe_headers[k] = v[:12] + "..." if v else "" + else: + safe_headers[k] = v + + # In redacted mode, elide prompt/tool/system + # content but keep structure, roles, and lengths. + redact = dump_mode == "redacted" + messages_sent = body.get("messages") + original_dump: Any = ( + original_messages + if original_messages is not body.get("messages") + else "__same_as_sent__" + ) + tools_sent = body.get("tools") + system_prompt = body.get("system") + if redact: + messages_sent = _redact_debug_value(messages_sent) + if original_dump != "__same_as_sent__": + original_dump = _redact_debug_value(original_dump) + tools_sent = _redact_debug_value(tools_sent) + system_prompt = _redact_debug_value(system_prompt) + + debug_payload = { + "request_id": request_id, + "timestamp": datetime.now().isoformat(), + "dump_mode": dump_mode, + "status_code": response.status_code, + "error_response": err_body, + "model": model, + "stream": stream, + "headers": safe_headers, + "compression": { + "was_compressed": bool(transforms_applied), + "transforms": transforms_applied, + "original_tokens": original_tokens, + "optimized_tokens": optimized_tokens, + "tokens_saved": tokens_saved, + "compression_failed": _compression_failed, + }, + "tools_sent": tools_sent, + "tool_count": len(body.get("tools") or []), + "original_tool_count": len(_original_tools or []), + "messages_sent": messages_sent, + "message_count": len(body.get("messages", [])), + "original_messages": original_dump, + "original_message_count": len(original_messages), + "system_prompt": system_prompt, + } + + with open(debug_file, "w") as f: + json.dump(debug_payload, f, indent=2, default=str) + + logger.warning( + f"[{request_id}] Debug dump ({dump_mode}): {debug_file}" + ) + except Exception as dump_err: + logger.error( + f"[{request_id}] Failed to write debug dump: {dump_err}" + ) + + # Parse response for CCR handling + resp_json = None + try: + resp_json = response.json() + except (json.JSONDecodeError, ValueError) as e: + logger.debug( + f"[{request_id}] Failed to parse response JSON for CCR handling: {e}" + ) + + # CCR Response Handling: Handle headroom_retrieve tool calls automatically + if ( + self.ccr_response_handler + and resp_json + and response.status_code == 200 + and self.ccr_response_handler.has_ccr_tool_calls(resp_json, "anthropic") + ): + logger.info( + f"[{request_id}] CCR: Detected retrieval tool call, handling..." + ) + + # Create API call function for continuation + # Use a fresh client to avoid potential decompression state issues + async def api_call_fn( + msgs: list[dict], tls: list[dict] | None + ) -> dict[str, Any]: + continuation_body = { + **body, + "messages": msgs, + } + if tls is not None: + continuation_body["tools"] = tls + + # Use clean headers for continuation + continuation_headers = { + k: v + for k, v in headers.items() + if k.lower() + not in ( + "content-encoding", + "transfer-encoding", + "accept-encoding", + "content-length", + ) + } + + # Reuse main client for CCR continuations (connection pooling) + logger.info( + f"CCR: Making continuation request with {len(msgs)} messages" + ) + assert self.http_client is not None, "HTTP client not initialized" + # Byte-faithful (PR-A3, fixes P0-2). The CCR + # continuation body is synthesized by Headroom + # so it is treated as mutated and goes through + # the canonical serializer. + from headroom.proxy.body_forwarding import ( + prepare_outbound_body_bytes, + ) + from headroom.proxy.helpers import log_outbound_request + + ccr_outbound_bytes, ccr_outbound_source = ( + prepare_outbound_body_bytes( + body=continuation_body, + original_body_bytes=None, + body_mutated=True, + ) + ) + ccr_outbound_headers = { + **continuation_headers, + "content-type": "application/json", + } + log_outbound_request( + forwarder="anthropic_ccr_continuation", + method="POST", + path=url, + body_bytes_count=len(ccr_outbound_bytes), + body_mutated=True, + mutation_reasons=["ccr_continuation"], + request_id=request_id, + source=ccr_outbound_source, + ) + try: + cont_response = await self.http_client.post( + url, + content=ccr_outbound_bytes, + headers=ccr_outbound_headers, + timeout=self._anthropic_buffered_request_timeout(), + ) + logger.info( + f"CCR: Got response status={cont_response.status_code}, " + f"content-encoding={cont_response.headers.get('content-encoding')}" + ) + result: dict[str, Any] = cont_response.json() + logger.info("CCR: Parsed JSON successfully") + return result + except Exception as e: + resp_headers: str | dict[str, str] = "N/A" + try: + resp_headers = dict(cont_response.headers) + except Exception: + pass + logger.error( + f"CCR: API call failed: {e}, response headers: {resp_headers}" + ) + raise + + # Handle CCR tool calls + try: + final_resp_json = await self.ccr_response_handler.handle_response( + resp_json, + optimized_messages, + tools, + api_call_fn, + provider="anthropic", + ) + # Update response content with final response + resp_json = final_resp_json + # Turn hooks (opt-in extensions) may inspect the turn or + # re-drive the model before we hand back the response. + # Inert when no hook is registered. + from headroom.proxy.turn_hooks import ( + TurnContext, + run_response_hooks, + ) + + final_resp_json = await run_response_hooks( + TurnContext( + provider="anthropic", + model=str(model), + messages=optimized_messages, + tools=tools, + config=self.config, + ), + final_resp_json, + api_call_fn, + ) + resp_json = final_resp_json + # Remove encoding headers since content is now uncompressed JSON + ccr_response_headers = { + k: v + for k, v in response.headers.items() + if k.lower() not in ("content-encoding", "content-length") + } + try: + ccr_content = json.dumps(final_resp_json).encode() + except (TypeError, ValueError) as json_err: + logger.warning( + f"[{request_id}] CCR: JSON serialization failed: {json_err}" + ) + ccr_content = json.dumps(resp_json).encode() + response = httpx.Response( + status_code=200, + content=ccr_content, + headers=ccr_response_headers, + ) + # Only claim success when no headroom_retrieve remains. + # On an intentional mixed-tool skip (#839) the response + # still carries headroom_retrieve for the client to + # resolve — logging "handled successfully" there is + # misleading. Classify via the shared, provider-generic + # residual-CCR signal. + from headroom.ccr.response_handler import ( + RESIDUAL_CCR_SKIPPED_MIXED, + ) + + residual_status = self.ccr_response_handler.residual_ccr_status( + final_resp_json, "anthropic" + ) + if residual_status == RESIDUAL_CCR_SKIPPED_MIXED: + logger.info( + f"[{request_id}] CCR: Skipped retrieval — " + "headroom_retrieve returned alongside a client " + "tool for the client to resolve" + ) + else: + logger.info( + f"[{request_id}] CCR: Retrieval handled successfully" + ) + except Exception as e: + import traceback + logger.error( - f"CCR: API call failed: {e}, response headers: {resp_headers}" + f"[{request_id}] CCR: Response handling failed: {e}\n" + f"Traceback: {traceback.format_exc()}" ) raise - # Handle CCR tool calls - try: - final_resp_json = await self.ccr_response_handler.handle_response( - resp_json, - optimized_messages, - tools, - api_call_fn, - provider="anthropic", - ) - # Update response content with final response - resp_json = final_resp_json - # Turn hooks (opt-in extensions) may inspect the turn or - # re-drive the model before we hand back the response. - # Inert when no hook is registered. - from headroom.proxy.turn_hooks import ( - TurnContext, - run_response_hooks, + # Memory: Handle memory tool calls in response + if ( + self.memory_handler + and memory_user_id + and resp_json + and response.status_code == 200 + and self.memory_handler.has_memory_tool_calls(resp_json, "anthropic") + ): + logger.info( + f"[{request_id}] Memory: Detected memory tool call, handling..." ) - final_resp_json = await run_response_hooks( - TurnContext( - provider="anthropic", - model=str(model), - messages=optimized_messages, - tools=tools, - config=self.config, - ), - final_resp_json, - api_call_fn, - ) - resp_json = final_resp_json - # Remove encoding headers since content is now uncompressed JSON - ccr_response_headers = { - k: v - for k, v in response.headers.items() - if k.lower() not in ("content-encoding", "content-length") - } try: - ccr_content = json.dumps(final_resp_json).encode() - except (TypeError, ValueError) as json_err: + # Execute memory tool calls + tool_results = await self.memory_handler.handle_memory_tool_calls( + resp_json, + memory_user_id, + "anthropic", + request_context=memory_request_ctx, + ) + + if tool_results: + # Create continuation messages + assistant_msg = { + "role": "assistant", + "content": resp_json.get("content", []), + } + user_msg = { + "role": "user", + "content": tool_results, + } + + continuation_messages = optimized_messages + [ + assistant_msg, + user_msg, + ] + + # Make continuation API call + continuation_body = {**body, "messages": continuation_messages} + if tools: + continuation_body["tools"] = tools + + cont_response = await self._retry_request( + "POST", + url, + headers, + continuation_body, + timeout=self._anthropic_buffered_request_timeout(), + ) + + # Update response with continuation + resp_json = cont_response.json() + response = cont_response + logger.info( + f"[{request_id}] Memory: Tool calls handled, continuation complete" + ) + + except Exception as e: logger.warning( - f"[{request_id}] CCR: JSON serialization failed: {json_err}" + f"[{request_id}] Memory: Tool call handling failed: {e}" ) - ccr_content = json.dumps(resp_json).encode() - response = httpx.Response( - status_code=200, - content=ccr_content, - headers=ccr_response_headers, - ) - # Only claim success when no headroom_retrieve remains. - # On an intentional mixed-tool skip (#839) the response - # still carries headroom_retrieve for the client to - # resolve — logging "handled successfully" there is - # misleading. Classify via the shared, provider-generic - # residual-CCR signal. - from headroom.ccr.response_handler import ( - RESIDUAL_CCR_SKIPPED_MIXED, - ) + # Continue with original response - residual_status = self.ccr_response_handler.residual_ccr_status( - final_resp_json, "anthropic" + total_latency = (time.time() - start_time) * 1000 + + # Parse response for output token count and cache metrics + output_tokens = 0 + cr_tokens = 0 + cw_tokens = 0 + cw_5m_tokens = 0 + cw_1h_tokens = 0 + uncached_input_tokens = 0 + if resp_json: + usage = resp_json.get("usage", {}) + output_tokens = usage.get("output_tokens", 0) + cr_tokens = usage.get("cache_read_input_tokens", 0) + cw_tokens = usage.get("cache_creation_input_tokens", 0) + cw_5m_tokens, cw_1h_tokens = self._extract_anthropic_cache_ttl_metrics( + usage ) - if residual_status == RESIDUAL_CCR_SKIPPED_MIXED: + uncached_input_tokens = usage.get("input_tokens", 0) + + # Track cache bust: tokens that lost their cache discount due to compression. + # If we had X tokens cached last turn and only Y hit cache this turn, + # then (X - Y) tokens were busted by our modifications. + expected_cached = prefix_tracker._cached_token_count + if expected_cached > 0 and tokens_saved > 0: + bust_tokens = max(0, expected_cached - cr_tokens) + if bust_tokens > 0: logger.info( - f"[{request_id}] CCR: Skipped retrieval — " - "headroom_retrieve returned alongside a client " - "tool for the client to resolve" + f"[{request_id}] CACHE-BUST: " + f"expected_cached={expected_cached:,} actual_read={cr_tokens:,} " + f"tokens_lost={bust_tokens:,} tokens_saved={tokens_saved:,}" ) - else: - logger.info(f"[{request_id}] CCR: Retrieval handled successfully") - except Exception as e: - import traceback + await self.metrics.record_cache_bust(bust_tokens) - logger.error( - f"[{request_id}] CCR: Response handling failed: {e}\n" - f"Traceback: {traceback.format_exc()}" + # Update prefix cache tracker for next turn + next_original_messages = copy.deepcopy(original_client_messages) + next_forwarded_messages = copy.deepcopy(optimized_messages) + assistant_message = self._assistant_message_from_response_json(resp_json) + if assistant_message is not None: + next_original_messages.append(copy.deepcopy(assistant_message)) + next_forwarded_messages.append(copy.deepcopy(assistant_message)) + + # Cache-miss attribution (#1313): when this turn expected a + # prompt-cache hit but got cr_tokens == 0, decide whether the + # cache most likely lapsed (idle > provider TTL → suggest a + # longer TTL) or the cacheable prefix changed (content shifted). + # Classify BEFORE update_from_response, which overwrites the + # last-turn state the classifier reads (idle clock, prefix, + # cached-token count). `optimized_messages` is the prefix we + # forwarded this turn; compare it against last turn's. + # `hasattr` guard: some tests inject a SimpleNamespace stub + # tracker that only implements the freeze API, not the full + # PrefixCacheTracker surface. + if hasattr(prefix_tracker, "classify_cache_miss"): + miss = prefix_tracker.classify_cache_miss( + cache_read_tokens=cr_tokens, + current_forwarded_messages=optimized_messages, ) - raise - - # Memory: Handle memory tool calls in response - if ( - self.memory_handler - and memory_user_id - and resp_json - and response.status_code == 200 - and self.memory_handler.has_memory_tool_calls(resp_json, "anthropic") - ): - logger.info( - f"[{request_id}] Memory: Detected memory tool call, handling..." - ) - - try: - # Execute memory tool calls - tool_results = await self.memory_handler.handle_memory_tool_calls( - resp_json, - memory_user_id, - "anthropic", - request_context=memory_request_ctx, - ) - - if tool_results: - # Create continuation messages - assistant_msg = { - "role": "assistant", - "content": resp_json.get("content", []), - } - user_msg = { - "role": "user", - "content": tool_results, - } - - continuation_messages = optimized_messages + [ - assistant_msg, - user_msg, - ] - - # Make continuation API call - continuation_body = {**body, "messages": continuation_messages} - if tools: - continuation_body["tools"] = tools - - cont_response = await self._retry_request( - "POST", - url, - headers, - continuation_body, - timeout=self._anthropic_buffered_request_timeout(), - ) - - # Update response with continuation - resp_json = cont_response.json() - response = cont_response + if miss.is_miss: logger.info( - f"[{request_id}] Memory: Tool calls handled, continuation complete" + f"[{request_id}] CACHE-MISS-ATTRIBUTION: reason={miss.reason} " + f"idle={miss.idle_seconds:.0f}s ttl={miss.cache_ttl_seconds}s " + f"expected_cached={miss.expected_cached_tokens:,} " + f"prefix_changed={miss.prefix_changed} " + f"ttl_exceeded={miss.ttl_exceeded}" + ) + await self.metrics.record_cache_miss_attribution( + provider_name, miss.reason ) - except Exception as e: - logger.warning(f"[{request_id}] Memory: Tool call handling failed: {e}") - # Continue with original response - - total_latency = (time.time() - start_time) * 1000 - - # Parse response for output token count and cache metrics - output_tokens = 0 - cr_tokens = 0 - cw_tokens = 0 - cw_5m_tokens = 0 - cw_1h_tokens = 0 - uncached_input_tokens = 0 - if resp_json: - usage = resp_json.get("usage", {}) - output_tokens = usage.get("output_tokens", 0) - cr_tokens = usage.get("cache_read_input_tokens", 0) - cw_tokens = usage.get("cache_creation_input_tokens", 0) - cw_5m_tokens, cw_1h_tokens = self._extract_anthropic_cache_ttl_metrics( - usage - ) - uncached_input_tokens = usage.get("input_tokens", 0) - - # Track cache bust: tokens that lost their cache discount due to compression. - # If we had X tokens cached last turn and only Y hit cache this turn, - # then (X - Y) tokens were busted by our modifications. - expected_cached = prefix_tracker._cached_token_count - if expected_cached > 0 and tokens_saved > 0: - bust_tokens = max(0, expected_cached - cr_tokens) - if bust_tokens > 0: - logger.info( - f"[{request_id}] CACHE-BUST: " - f"expected_cached={expected_cached:,} actual_read={cr_tokens:,} " - f"tokens_lost={bust_tokens:,} tokens_saved={tokens_saved:,}" - ) - await self.metrics.record_cache_bust(bust_tokens) - - # Update prefix cache tracker for next turn - next_original_messages = copy.deepcopy(original_client_messages) - next_forwarded_messages = copy.deepcopy(optimized_messages) - assistant_message = self._assistant_message_from_response_json(resp_json) - if assistant_message is not None: - next_original_messages.append(copy.deepcopy(assistant_message)) - next_forwarded_messages.append(copy.deepcopy(assistant_message)) - - # Cache-miss attribution (#1313): when this turn expected a - # prompt-cache hit but got cr_tokens == 0, decide whether the - # cache most likely lapsed (idle > provider TTL → suggest a - # longer TTL) or the cacheable prefix changed (content shifted). - # Classify BEFORE update_from_response, which overwrites the - # last-turn state the classifier reads (idle clock, prefix, - # cached-token count). `optimized_messages` is the prefix we - # forwarded this turn; compare it against last turn's. - # `hasattr` guard: some tests inject a SimpleNamespace stub - # tracker that only implements the freeze API, not the full - # PrefixCacheTracker surface. - if hasattr(prefix_tracker, "classify_cache_miss"): - miss = prefix_tracker.classify_cache_miss( - cache_read_tokens=cr_tokens, - current_forwarded_messages=optimized_messages, - ) - if miss.is_miss: - logger.info( - f"[{request_id}] CACHE-MISS-ATTRIBUTION: reason={miss.reason} " - f"idle={miss.idle_seconds:.0f}s ttl={miss.cache_ttl_seconds}s " - f"expected_cached={miss.expected_cached_tokens:,} " - f"prefix_changed={miss.prefix_changed} " - f"ttl_exceeded={miss.ttl_exceeded}" - ) - await self.metrics.record_cache_miss_attribution( - provider_name, miss.reason - ) - - prefix_tracker.update_from_response( - cache_read_tokens=cr_tokens, - cache_write_tokens=cw_tokens, - messages=next_forwarded_messages, - original_messages=next_original_messages, - ) - - # Cache response under the SAME key it was looked up by: - # cache_lookup_messages is the raw pre-mutation snapshot, not - # the live (compressed/hooked) `messages` (#327). - if self.cache and response.status_code == 200: - await self.cache.set( - cache_lookup_messages, - model, - response.content, - dict(response.headers), - tokens_saved=tokens_saved, - **cache_key_fields, - ) - - # Subscription tracker: update headroom contribution - # counters. Provider-specific OAuth/subscription - # accounting — stays outside the funnel (different - # concern, only fires for Bearer-not-sk-ant tokens). - if _auth_header.startswith("Bearer ") and not _auth_header.startswith( - "Bearer sk-ant-api" - ): - from headroom.subscription.tracker import ( - get_subscription_tracker as _get_sub_tracker, - ) - - _sub_tracker = _get_sub_tracker() - if _sub_tracker is not None: - _sub_tracker.update_contribution( - tokens_submitted=optimized_tokens, - tokens_saved_compression=tokens_saved, - tokens_saved_cache_reads=cr_tokens, - ) - - # The pre-refactor PERF emit (above) read raw usage - # off ``resp_usage`` instead of trusting cr_tokens / - # cw_tokens. Both paths land on identical numbers - # (extraction happens just above the cost_tracker - # call), so the funnel uses the already-computed - # values for consistency. Pre-refactor's - # ``cache_hit`` local was correctly derived from - # cache_read>0; the funnel re-derives via the - # outcome property — same result. - # - # ``attempted_input_tokens`` was MISSING from the - # pre-refactor record_request call here (one of the - # 7-of-18 sites the P0 audit flagged). The funnel - # forces it to a value — using - # ``optimized_tokens + tokens_saved`` as the - # fallback denominator, same as the streaming path - # uses (see _finalize_stream_response). Dashboards - # that were showing 0% active-savings on non- - # streaming Anthropic traffic will now show the - # correct ratio. - await self._record_request_outcome( - RequestOutcome( - request_id=request_id, - provider=provider_name, - model=model, - status_code=response.status_code, - original_tokens=original_tokens, - optimized_tokens=optimized_tokens, - output_tokens=output_tokens, - tokens_saved=tokens_saved, - attempted_input_tokens=optimized_tokens + tokens_saved, + prefix_tracker.update_from_response( cache_read_tokens=cr_tokens, cache_write_tokens=cw_tokens, - cache_write_5m_tokens=cw_5m_tokens, - cache_write_1h_tokens=cw_1h_tokens, - uncached_input_tokens=uncached_input_tokens, - total_latency_ms=total_latency, - overhead_ms=optimization_latency, - pipeline_timing=pipeline_timing, - waste_signals=waste_signals_dict, - transforms_applied=tuple(transforms_applied), - num_messages=len(messages), - tags=tags, - client=client, - turn_id=compute_turn_id( - model, body.get("system"), body.get("messages") - ), - # `original_client_messages` is the deep-copied - # pre-compression snapshot; `body["messages"]` is the - # compressed list sent upstream. Both gated by - # `log_full_messages`. - request_messages=original_client_messages - if self.config.log_full_messages - else None, - compressed_messages=body.get("messages") - if self.config.log_full_messages - else None, + messages=next_forwarded_messages, + original_messages=next_original_messages, ) - ) - # Remove compression headers since httpx already decompressed the response - response_headers = dict(response.headers) - response_headers.pop("content-encoding", None) - response_headers.pop( - "content-length", None - ) # Length changed after decompression - - # Inject Headroom compression metrics (for SaaS metering) - response_headers["x-headroom-tokens-before"] = str(original_tokens) - response_headers["x-headroom-tokens-after"] = str(optimized_tokens) - response_headers["x-headroom-tokens-saved"] = str(tokens_saved) - response_headers["x-headroom-model"] = model - if transforms_applied: - from headroom.proxy.cost import header_safe_transforms - - response_headers["x-headroom-transforms"] = ",".join( - header_safe_transforms(transforms_applied) - ) - if cache_hit: - response_headers["x-headroom-cached"] = "true" - if _compression_failed: - response_headers["x-headroom-compression-failed"] = "true" - - # Enterprise Security: scan response + de-anonymize. - # Gate on a 200 upstream like the sibling CCR/cache/buffered - # blocks below: without this, a non-2xx upstream (rate limit - # 429, overloaded 529, 5xx) whose JSON body is scanned was - # rebuilt as httpx.Response(status_code=200) and returned as - # HTTP 200, so the client's retry/backoff never triggered and - # an error looked like success. - if ( - self.security - and _security_ctx - and resp_json - and response.status_code == 200 - ): - try: - resp_json = self.security.scan_response(resp_json, _security_ctx) - response = httpx.Response( - status_code=200, - content=json.dumps(resp_json).encode(), - headers=response_headers, + # Cache response under the SAME key it was looked up by: + # cache_lookup_messages is the raw pre-mutation snapshot, not + # the live (compressed/hooked) `messages` (#327). + if self.cache and response.status_code == 200: + await self.cache.set( + cache_lookup_messages, + model, + response.content, + dict(response.headers), + tokens_saved=tokens_saved, + **cache_key_fields, ) - if not buffered_stream_ccr: - return Response( - content=response.content, - status_code=response.status_code, + + # Subscription tracker: update headroom contribution + # counters. Provider-specific OAuth/subscription + # accounting — stays outside the funnel (different + # concern, only fires for Bearer-not-sk-ant tokens). + if _auth_header.startswith("Bearer ") and not _auth_header.startswith( + "Bearer sk-ant-api" + ): + from headroom.subscription.tracker import ( + get_subscription_tracker as _get_sub_tracker, + ) + + _sub_tracker = _get_sub_tracker() + if _sub_tracker is not None: + _sub_tracker.update_contribution( + tokens_submitted=optimized_tokens, + tokens_saved_compression=tokens_saved, + tokens_saved_cache_reads=cr_tokens, + ) + + # The pre-refactor PERF emit (above) read raw usage + # off ``resp_usage`` instead of trusting cr_tokens / + # cw_tokens. Both paths land on identical numbers + # (extraction happens just above the cost_tracker + # call), so the funnel uses the already-computed + # values for consistency. Pre-refactor's + # ``cache_hit`` local was correctly derived from + # cache_read>0; the funnel re-derives via the + # outcome property — same result. + # + # ``attempted_input_tokens`` was MISSING from the + # pre-refactor record_request call here (one of the + # 7-of-18 sites the P0 audit flagged). The funnel + # forces it to a value — using + # ``optimized_tokens + tokens_saved`` as the + # fallback denominator, same as the streaming path + # uses (see _finalize_stream_response). Dashboards + # that were showing 0% active-savings on non- + # streaming Anthropic traffic will now show the + # correct ratio. + await self._record_request_outcome( + RequestOutcome( + request_id=request_id, + provider=provider_name, + model=model, + status_code=response.status_code, + original_tokens=original_tokens, + optimized_tokens=optimized_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + attempted_input_tokens=optimized_tokens + tokens_saved, + cache_read_tokens=cr_tokens, + cache_write_tokens=cw_tokens, + cache_write_5m_tokens=cw_5m_tokens, + cache_write_1h_tokens=cw_1h_tokens, + uncached_input_tokens=uncached_input_tokens, + total_latency_ms=total_latency, + overhead_ms=optimization_latency, + pipeline_timing=pipeline_timing, + waste_signals=waste_signals_dict, + transforms_applied=tuple(transforms_applied), + num_messages=len(messages), + tags=tags, + client=client, + turn_id=compute_turn_id( + model, body.get("system"), body.get("messages") + ), + # `original_client_messages` is the deep-copied + # pre-compression snapshot; `body["messages"]` is the + # compressed list sent upstream. Both gated by + # `log_full_messages`. + request_messages=original_client_messages + if self.config.log_full_messages + else None, + compressed_messages=body.get("messages") + if self.config.log_full_messages + else None, + ) + ) + + # Remove compression headers since httpx already decompressed the response + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop( + "content-length", None + ) # Length changed after decompression + + # Inject Headroom compression metrics (for SaaS metering) + response_headers["x-headroom-tokens-before"] = str(original_tokens) + response_headers["x-headroom-tokens-after"] = str(optimized_tokens) + response_headers["x-headroom-tokens-saved"] = str(tokens_saved) + response_headers["x-headroom-model"] = model + if transforms_applied: + from headroom.proxy.cost import header_safe_transforms + + response_headers["x-headroom-transforms"] = ",".join( + header_safe_transforms(transforms_applied) + ) + if cache_hit: + response_headers["x-headroom-cached"] = "true" + if _compression_failed: + response_headers["x-headroom-compression-failed"] = "true" + + # Enterprise Security: scan response + de-anonymize. + # Gate on a 200 upstream like the sibling CCR/cache/buffered + # blocks below: without this, a non-2xx upstream (rate limit + # 429, overloaded 529, 5xx) whose JSON body is scanned was + # rebuilt as httpx.Response(status_code=200) and returned as + # HTTP 200, so the client's retry/backoff never triggered and + # an error looked like success. + if ( + self.security + and _security_ctx + and resp_json + and response.status_code == 200 + ): + try: + resp_json = self.security.scan_response(resp_json, _security_ctx) + response = httpx.Response( + status_code=200, + content=json.dumps(resp_json).encode(), headers=response_headers, ) - except Exception as sec_err: - logger.warning( - f"[{request_id}] Security response scan error: {sec_err}" - ) + if not buffered_stream_ccr: + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + except Exception as sec_err: + logger.warning( + f"[{request_id}] Security response scan error: {sec_err}" + ) - if buffered_stream_ccr and response.status_code == 200 and resp_json: - sse_headers = { - k: v - for k, v in response_headers.items() - if k.lower() - not in ( - "content-encoding", - "content-length", - "transfer-encoding", - "content-type", - ) - } - - def _sse_error_event(message: str) -> bytes: - error_event = { - "type": "error", - "error": {"type": "api_error", "message": message}, + if buffered_stream_ccr and response.status_code == 200 and resp_json: + sse_headers = { + k: v + for k, v in response_headers.items() + if k.lower() + not in ( + "content-encoding", + "content-length", + "transfer-encoding", + "content-type", + ) } - return f"event: error\ndata: {json.dumps(error_event)}\n\n".encode() - # Residual headroom_retrieve is only a hard failure when it - # is NOT an intentional mixed-tool skip. When the model - # emitted headroom_retrieve alongside a client tool (#839), - # the handler deliberately leaves both tool_use blocks in - # place for the client to resolve — a legal turn that the - # non-streaming path returns as 200. Fall through to the - # SSE resynthesis below so the stream:true path matches it - # and preserves both blocks. Use the shared, provider-generic - # residual-CCR signal (not an Anthropic-only branch). - from headroom.ccr.response_handler import RESIDUAL_CCR_ERROR + def _sse_error_event(message: str) -> bytes: + error_event = { + "type": "error", + "error": {"type": "api_error", "message": message}, + } + return f"event: error\ndata: {json.dumps(error_event)}\n\n".encode() - if ( - self.ccr_response_handler - and self.ccr_response_handler.residual_ccr_status( - resp_json, "anthropic" - ) - == RESIDUAL_CCR_ERROR - ): - logger.warning( - f"[{request_id}] CCR: Buffered streaming response still " - "contains an unresolved headroom_retrieve after handling; " - "failing closed" - ) + # Residual headroom_retrieve is only a hard failure when it + # is NOT an intentional mixed-tool skip. When the model + # emitted headroom_retrieve alongside a client tool (#839), + # the handler deliberately leaves both tool_use blocks in + # place for the client to resolve — a legal turn that the + # non-streaming path returns as 200. Fall through to the + # SSE resynthesis below so the stream:true path matches it + # and preserves both blocks. Use the shared, provider-generic + # residual-CCR signal (not an Anthropic-only branch). + from headroom.ccr.response_handler import RESIDUAL_CCR_ERROR - async def _residual_ccr_error_sse(): - yield _sse_error_event( - "Unable to safely complete streamed CCR retrieval." + if ( + self.ccr_response_handler + and self.ccr_response_handler.residual_ccr_status( + resp_json, "anthropic" + ) + == RESIDUAL_CCR_ERROR + ): + logger.warning( + f"[{request_id}] CCR: Buffered streaming response still " + "contains an unresolved headroom_retrieve after handling; " + "failing closed" ) - return StreamingResponse( - _residual_ccr_error_sse(), - media_type="text/event-stream", - headers=sse_headers, - status_code=502, - ) + async def _residual_ccr_error_sse(): + yield _sse_error_event( + "Unable to safely complete streamed CCR retrieval." + ) - try: - sse_events = self._response_to_sse(resp_json, "anthropic") - except ValueError as sse_err: - logger.warning( - f"[{request_id}] CCR: Failed to convert buffered response " - f"to SSE: {sse_err}" - ) - - async def _conversion_error_sse(): - yield _sse_error_event( - "Unable to safely convert buffered response to SSE." + return StreamingResponse( + _residual_ccr_error_sse(), + media_type="text/event-stream", + headers=sse_headers, + status_code=502, ) + try: + sse_events = self._response_to_sse(resp_json, "anthropic") + except ValueError as sse_err: + logger.warning( + f"[{request_id}] CCR: Failed to convert buffered response " + f"to SSE: {sse_err}" + ) + + async def _conversion_error_sse(): + yield _sse_error_event( + "Unable to safely convert buffered response to SSE." + ) + + return StreamingResponse( + _conversion_error_sse(), + media_type="text/event-stream", + headers=sse_headers, + status_code=502, + ) + + async def _buffered_ccr_sse(): + for event in sse_events: + yield event + return StreamingResponse( - _conversion_error_sse(), + _buffered_ccr_sse(), media_type="text/event-stream", headers=sse_headers, - status_code=502, ) - async def _buffered_ccr_sse(): - for event in sse_events: - yield event - - return StreamingResponse( - _buffered_ccr_sse(), - media_type="text/event-stream", - headers=sse_headers, + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, ) - return Response( - content=response.content, - status_code=response.status_code, - headers=response_headers, - ) + if buffered_stream_ccr: + operation = asyncio.create_task(_buffered_ccr_operation()) + record_failed = self.metrics.record_failed + + class _BufferedCCRResponse(Response): + async def __call__(self, scope, receive, send): # noqa: ANN001 + await asyncio.sleep(0) + loop = asyncio.get_running_loop() + keepalive_deadline = loop.time() + 1.0 + started = False + try: + while True: + timeout = ( + 0.25 + if started + else max(0.0, keepalive_deadline - loop.time()) + ) + done, _ = await asyncio.wait({operation}, timeout=timeout) + if done: + try: + result = operation.result() + except Exception as e: + await record_failed(provider=provider_name) + logger.error( + f"[{request_id}] Request failed: {type(e).__name__}: {e}" + ) + if not started: + await send( + { + "type": "http.response.start", + "status": 502, + "headers": [ + (b"content-type", b"application/json") + ], + } + ) + await send( + { + "type": "http.response.body", + "body": json.dumps( + { + "type": "error", + "error": { + "type": "api_error", + "message": "An error occurred while processing your request. Please try again.", + }, + } + ).encode(), + "more_body": False, + } + ) + return + await send( + { + "type": "http.response.body", + "body": b'event: error\ndata: {"type":"error","error":{"type":"api_error","message":"An error occurred while processing the request."}}\n\n', + "more_body": False, + } + ) + return + + if not started: + await result(scope, receive, send) + return + + body_iterator = getattr(result, "body_iterator", None) + if body_iterator is not None: + async for chunk in body_iterator: + await send( + { + "type": "http.response.body", + "body": chunk, + "more_body": True, + } + ) + await send( + { + "type": "http.response.body", + "body": b"", + "more_body": False, + } + ) + return + + await send( + { + "type": "http.response.body", + "body": b'event: error\ndata: {"type":"error","error":{"type":"api_error","message":"An error occurred while processing the request."}}\n\n', + "more_body": False, + } + ) + return + + if not started: + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [ + (b"content-type", b"text/event-stream") + ], + } + ) + started = True + await send( + { + "type": "http.response.body", + "body": b'event: ping\ndata: {"type":"ping"}\n\n', + "more_body": True, + } + ) + except asyncio.CancelledError: + raise + finally: + if not operation.done(): + operation.cancel() + try: + await operation + except asyncio.CancelledError: + pass + except Exception: + pass + + return _BufferedCCRResponse(media_type="text/event-stream") + return await _buffered_ccr_operation() except HTTPException: # FastAPI HTTPException carries its own status code, headers, # and client-facing message (e.g. 429 with Retry-After, 413 for diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 8baa8610f..3637d7d63 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -4911,379 +4911,506 @@ class OpenAIHandlerMixin: waste_signals=waste_signals_dict, ) else: - headers = await apply_copilot_api_auth(headers, url=url) - response = await self._retry_request( - "POST", - url, - headers, - body, - original_body_bytes=original_body_bytes, - body_mutated=body_mutation_tracker.mutated, - mutation_reasons=body_mutation_tracker.reasons, - request_id=request_id, - forwarder_name="openai_responses", - path_for_log=url, - ) - _response_body_for_debug: Any = None - _response_raw_for_debug: str | None = None - try: - _response_body_for_debug = response.json() - except Exception: - try: - _response_raw_for_debug = response.text[:200_000] - except Exception: - _response_raw_for_debug = None - capture_codex_wire_debug( - "http_upstream_response", - request_id=request_id, - transport="http", - direction="upstream_to_headroom", - method="POST", - url=url, - headers=dict(response.headers), - body=_response_body_for_debug, - raw_text=_response_raw_for_debug, - status_code=response.status_code, - metadata={"stream": stream, "auth_mode": auth_mode.value}, - ) - total_latency = (time.time() - start_time) * 1000 - total_input_tokens = original_tokens # fallback - output_tokens = 0 - cache_read_tokens = 0 - try: - resp_json = response.json() - usage = resp_json.get("usage", {}) - - def _usage_int(value: Any, default: int = 0) -> int: - try: - return max(int(value), 0) - except (TypeError, ValueError): - return default - - total_input_tokens = _usage_int( - usage.get("input_tokens"), - original_tokens, - ) - output_tokens = _usage_int(usage.get("output_tokens")) - details = usage.get("input_tokens_details") - if isinstance(details, dict): - cache_read_tokens = _usage_int(details.get("cached_tokens")) - except (KeyError, TypeError, AttributeError) as e: - logger.debug( - f"[{request_id}] Failed to extract cached tokens from OpenAI passthrough response: {e}" - ) - - # CCR Response Handling: intercept headroom_retrieve tool - # calls server-side so a Responses API function_call the - # downstream caller can't resolve (e.g. Strands, or a - # buffered-stream request) never reaches the client. Mirrors - # the chat-completions backend-path block (handle_openai_chat - # ~2775-2848), adapted for the Responses API's flat - # function_call / output[] shape instead of Messages API - # tool_calls. Runs before memory tool handling below so a - # retrieve call never gets treated as an unresolved tool_call - # by the memory-tool branch. - if ( - _ccr_response_handler - and resp_json - and response.status_code == 200 - and _ccr_response_handler.has_ccr_tool_calls(resp_json, "openai_responses") - ): - logger.info( - f"[{request_id}] CCR: Detected retrieval tool call (responses), handling..." - ) - - async def api_call_fn( - items: list[dict[str, Any]], - tls: list[dict[str, Any]] | None, - ) -> dict[str, Any]: - continuation_body = {**body, "input": items} - if tls is not None: - continuation_body["tools"] = tls - # Fresh stateless continuation: resend the full - # item history rather than chaining through - # previous_response_id, matching how - # CCRResponseHandler accumulates `current_messages` - # for every other provider. `body["stream"]` is - # left as-is: for a buffered_stream_ccr request it - # was already forced False above, and continuations - # must stay non-streaming so this handler (not - # `_stream_response`) can parse the JSON reply. - continuation_body.pop("previous_response_id", None) - continuation_body["stream"] = False - - continuation_headers = { - k: v - for k, v in headers.items() - if k.lower() - not in ( - "content-encoding", - "transfer-encoding", - "accept-encoding", - "content-length", - ) - } - logger.info( - f"[{request_id}] CCR: Issuing Responses continuation " - f"({len(items)} input items)" - ) - cont_response = await self._retry_request( - "POST", - url, - continuation_headers, - continuation_body, - request_id=request_id, - forwarder_name="openai_responses_ccr_continuation", - path_for_log=url, - ) - return cont_response.json() - - try: - final_resp_json = await _ccr_response_handler.handle_response( - resp_json, - _responses_input_to_items(body.get("input")), - body.get("tools"), - api_call_fn, - provider="openai_responses", - ) - resp_json = final_resp_json - # Remove encoding headers since content is now - # uncompressed JSON we synthesized. - ccr_response_headers = { - k: v - for k, v in response.headers.items() - if k.lower() not in ("content-encoding", "content-length") - } - response = httpx.Response( - status_code=200, - content=json.dumps(final_resp_json).encode(), - headers=ccr_response_headers, - ) - logger.info( - f"[{request_id}] CCR: Retrieval handled successfully (responses)" - ) - except Exception as e: - logger.error( - f"[{request_id}] CCR: Response handling failed (responses): {e}" - ) - # NO SILENT FALLBACK: re-raise so the client sees a - # clear failure instead of an unresolved tool_call - # it can't act on. Matches the OpenAI backend-path - # block in handle_openai_chat; see - # feedback_no_silent_fallbacks. - raise - - # Memory: handle memory tool calls in Responses API response - if ( - self.memory_handler - and memory_user_id - and responses_memory_tools_allowed - and resp_json - and response.status_code == 200 - and self.memory_handler.has_memory_tool_calls(resp_json, "openai") - ): - try: - # Extract function_call items from output - from headroom.proxy.memory_handler import MEMORY_TOOL_NAMES - - output_items = resp_json.get("output", []) - memory_fc_items = [ - item - for item in output_items - if isinstance(item, dict) - and item.get("type") == "function_call" - and item.get("name") in MEMORY_TOOL_NAMES - ] - - # Execute memory tool calls - tool_outputs: list[dict[str, Any]] = [] - for fc in memory_fc_items: - call_id = fc.get("call_id", fc.get("id", "")) - name = fc.get("name", "") - args_str = fc.get("arguments", "{}") - try: - args = json.loads(args_str) - except json.JSONDecodeError: - args = {} - - await self.memory_handler._ensure_initialized() - if self.memory_handler._backend: - result = await self.memory_handler._execute_memory_tool( - name, args, memory_user_id, "openai" - ) - else: - result = json.dumps({"error": "Memory backend not initialized"}) - - tool_outputs.append( - { - "type": "function_call_output", - "call_id": call_id, - "output": result, - } - ) - - if tool_outputs: - # Make continuation request with tool results - response_id = resp_json.get("id") - continuation_body = { - "model": model, - "input": tool_outputs, - } - if response_id: - continuation_body["previous_response_id"] = response_id - existing_tools = body.get("tools") - if existing_tools: - continuation_body["tools"] = existing_tools - - cont_response = await self._retry_request( - "POST", url, headers, continuation_body - ) - resp_json = cont_response.json() - response = cont_response - logger.info( - f"[{request_id}] Memory: Handled {len(tool_outputs)} " - f"tool call(s) with continuation for user {memory_user_id} (responses)" - ) - except Exception as e: - logger.warning( - f"[{request_id}] Memory tool handling failed (responses): {e}" - ) - - if self.cost_tracker: - cache_write_tokens = _infer_openai_cache_write_tokens( - total_input_tokens, - cache_read_tokens, - ) - uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens) - # (record_tokens clamps negative savings to 0 universally.) - self.cost_tracker.record_tokens( - model, - tokens_saved, - total_input_tokens, - cache_read_tokens=cache_read_tokens, - cache_write_tokens=cache_write_tokens, - uncached_tokens=uncached_input_tokens, - ) - else: - cache_write_tokens = _infer_openai_cache_write_tokens( - total_input_tokens, - cache_read_tokens, - ) - uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens) - - effective_optimized_tokens = ( - total_input_tokens if total_input_tokens > 0 else optimized_tokens - ) - effective_original_tokens = max( - original_tokens, - effective_optimized_tokens + tokens_saved, - ) - - _resp_log_tags = { - **(tags or {}), - "auth_mode": auth_mode.value if auth_mode else "payg", - "endpoint": "responses_http", - } - - # OpenAI Responses HTTP (non-WS, non-streaming). Codex - # uses this path when configured for HTTP transport. - # Pre-refactor `cache_hit` was hardcoded False on - # RequestLog even when cache_read>0 — funnel derives - # it correctly. - from headroom.proxy.helpers import compute_turn_id - - await self._record_request_outcome( - RequestOutcome( + async def _buffered_ccr_operation(): + nonlocal headers + headers = await apply_copilot_api_auth(headers, url=url) + response = await self._retry_request( + "POST", + url, + headers, + body, + original_body_bytes=original_body_bytes, + body_mutated=body_mutation_tracker.mutated, + mutation_reasons=body_mutation_tracker.reasons, request_id=request_id, - provider="openai", - model=model, - status_code=response.status_code, - original_tokens=effective_original_tokens, - optimized_tokens=effective_optimized_tokens, - output_tokens=output_tokens, - tokens_saved=tokens_saved, - attempted_input_tokens=attempted_input_tokens, - cache_read_tokens=cache_read_tokens, - cache_write_tokens=cache_write_tokens, - uncached_input_tokens=uncached_input_tokens, - total_latency_ms=total_latency, - overhead_ms=optimization_latency, - transforms_applied=tuple(transforms_applied), - waste_signals=waste_signals_dict, - num_messages=len(messages) if isinstance(messages, list) else 0, - tags=_resp_log_tags, - turn_id=compute_turn_id(model, body.get("instructions"), messages), - request_messages=messages - if getattr(self.config, "log_full_messages", False) - else None, - client=client, + forwarder_name="openai_responses", + path_for_log=url, ) - ) + _response_body_for_debug: Any = None + _response_raw_for_debug: str | None = None + try: + _response_body_for_debug = response.json() + except Exception: + try: + _response_raw_for_debug = response.text[:200_000] + except Exception: + _response_raw_for_debug = None + capture_codex_wire_debug( + "http_upstream_response", + request_id=request_id, + transport="http", + direction="upstream_to_headroom", + method="POST", + url=url, + headers=dict(response.headers), + body=_response_body_for_debug, + raw_text=_response_raw_for_debug, + status_code=response.status_code, + metadata={"stream": stream, "auth_mode": auth_mode.value}, + ) + total_latency = (time.time() - start_time) * 1000 - logger.info(f"[{request_id}] /v1/responses {model}: {total_input_tokens:,} tokens") + total_input_tokens = original_tokens # fallback + output_tokens = 0 + cache_read_tokens = 0 + try: + resp_json = response.json() + usage = resp_json.get("usage", {}) - # Capture Codex rate-limit window data from response headers - from headroom.subscription.codex_rate_limits import ( - get_codex_rate_limit_state, - ) + def _usage_int(value: Any, default: int = 0) -> int: + try: + return max(int(value), 0) + except (TypeError, ValueError): + return default - get_codex_rate_limit_state().update_from_headers(dict(response.headers)) - - # Remove compression headers - response_headers = _sanitize_forwarded_response_headers(response.headers) - - if buffered_stream_ccr and response.status_code == 200 and resp_json: - sse_headers = { - k: v - for k, v in response_headers.items() - if k.lower() not in ("content-length", "content-type") - } - if _ccr_response_handler and _ccr_response_handler.has_ccr_tool_calls( - resp_json, "openai_responses" - ): - # Handling above didn't fully resolve the retrieve - # call (e.g. max rounds hit, or it was mixed with a - # non-CCR tool call). Fail closed rather than stream - # a response the client can't act on — matches the - # Anthropic buffered path's residual-CCR guard. - logger.warning( - f"[{request_id}] CCR: Buffered streaming Responses " - "reply still contains headroom_retrieve after " - "handling; failing closed" + total_input_tokens = _usage_int( + usage.get("input_tokens"), + original_tokens, + ) + output_tokens = _usage_int(usage.get("output_tokens")) + details = usage.get("input_tokens_details") + if isinstance(details, dict): + cache_read_tokens = _usage_int(details.get("cached_tokens")) + except (KeyError, TypeError, AttributeError) as e: + logger.debug( + f"[{request_id}] Failed to extract cached tokens from OpenAI passthrough response: {e}" ) - async def _residual_ccr_error_sse(): - error_event = { - "type": "error", - "error": { - "message": "Unable to safely complete streamed CCR retrieval.", - }, + # CCR Response Handling: intercept headroom_retrieve tool + # calls server-side so a Responses API function_call the + # downstream caller can't resolve (e.g. Strands, or a + # buffered-stream request) never reaches the client. Mirrors + # the chat-completions backend-path block (handle_openai_chat + # ~2775-2848), adapted for the Responses API's flat + # function_call / output[] shape instead of Messages API + # tool_calls. Runs before memory tool handling below so a + # retrieve call never gets treated as an unresolved tool_call + # by the memory-tool branch. + if ( + _ccr_response_handler + and resp_json + and response.status_code == 200 + and _ccr_response_handler.has_ccr_tool_calls(resp_json, "openai_responses") + ): + logger.info( + f"[{request_id}] CCR: Detected retrieval tool call (responses), handling..." + ) + + async def api_call_fn( + items: list[dict[str, Any]], + tls: list[dict[str, Any]] | None, + ) -> dict[str, Any]: + continuation_body = {**body, "input": items} + if tls is not None: + continuation_body["tools"] = tls + # Fresh stateless continuation: resend the full + # item history rather than chaining through + # previous_response_id, matching how + # CCRResponseHandler accumulates `current_messages` + # for every other provider. `body["stream"]` is + # left as-is: for a buffered_stream_ccr request it + # was already forced False above, and continuations + # must stay non-streaming so this handler (not + # `_stream_response`) can parse the JSON reply. + continuation_body.pop("previous_response_id", None) + continuation_body["stream"] = False + + continuation_headers = { + k: v + for k, v in headers.items() + if k.lower() + not in ( + "content-encoding", + "transfer-encoding", + "accept-encoding", + "content-length", + ) } - yield f"event: error\ndata: {json.dumps(error_event)}\n\n".encode() + logger.info( + f"[{request_id}] CCR: Issuing Responses continuation " + f"({len(items)} input items)" + ) + cont_response = await self._retry_request( + "POST", + url, + continuation_headers, + continuation_body, + request_id=request_id, + forwarder_name="openai_responses_ccr_continuation", + path_for_log=url, + ) + return cont_response.json() + + try: + final_resp_json = await _ccr_response_handler.handle_response( + resp_json, + _responses_input_to_items(body.get("input")), + body.get("tools"), + api_call_fn, + provider="openai_responses", + ) + resp_json = final_resp_json + # Remove encoding headers since content is now + # uncompressed JSON we synthesized. + ccr_response_headers = { + k: v + for k, v in response.headers.items() + if k.lower() not in ("content-encoding", "content-length") + } + response = httpx.Response( + status_code=200, + content=json.dumps(final_resp_json).encode(), + headers=ccr_response_headers, + ) + logger.info( + f"[{request_id}] CCR: Retrieval handled successfully (responses)" + ) + except Exception as e: + logger.error( + f"[{request_id}] CCR: Response handling failed (responses): {e}" + ) + # NO SILENT FALLBACK: re-raise so the client sees a + # clear failure instead of an unresolved tool_call + # it can't act on. Matches the OpenAI backend-path + # block in handle_openai_chat; see + # feedback_no_silent_fallbacks. + raise + + # Memory: handle memory tool calls in Responses API response + if ( + self.memory_handler + and memory_user_id + and responses_memory_tools_allowed + and resp_json + and response.status_code == 200 + and self.memory_handler.has_memory_tool_calls(resp_json, "openai") + ): + try: + # Extract function_call items from output + from headroom.proxy.memory_handler import MEMORY_TOOL_NAMES + + output_items = resp_json.get("output", []) + memory_fc_items = [ + item + for item in output_items + if isinstance(item, dict) + and item.get("type") == "function_call" + and item.get("name") in MEMORY_TOOL_NAMES + ] + + # Execute memory tool calls + tool_outputs: list[dict[str, Any]] = [] + for fc in memory_fc_items: + call_id = fc.get("call_id", fc.get("id", "")) + name = fc.get("name", "") + args_str = fc.get("arguments", "{}") + try: + args = json.loads(args_str) + except json.JSONDecodeError: + args = {} + + await self.memory_handler._ensure_initialized() + if self.memory_handler._backend: + result = await self.memory_handler._execute_memory_tool( + name, args, memory_user_id, "openai" + ) + else: + result = json.dumps({"error": "Memory backend not initialized"}) + + tool_outputs.append( + { + "type": "function_call_output", + "call_id": call_id, + "output": result, + } + ) + + if tool_outputs: + # Make continuation request with tool results + response_id = resp_json.get("id") + continuation_body = { + "model": model, + "input": tool_outputs, + } + if response_id: + continuation_body["previous_response_id"] = response_id + existing_tools = body.get("tools") + if existing_tools: + continuation_body["tools"] = existing_tools + + cont_response = await self._retry_request( + "POST", url, headers, continuation_body + ) + resp_json = cont_response.json() + response = cont_response + logger.info( + f"[{request_id}] Memory: Handled {len(tool_outputs)} " + f"tool call(s) with continuation for user {memory_user_id} (responses)" + ) + except Exception as e: + logger.warning( + f"[{request_id}] Memory tool handling failed (responses): {e}" + ) + + if self.cost_tracker: + cache_write_tokens = _infer_openai_cache_write_tokens( + total_input_tokens, + cache_read_tokens, + ) + uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens) + # (record_tokens clamps negative savings to 0 universally.) + self.cost_tracker.record_tokens( + model, + tokens_saved, + total_input_tokens, + cache_read_tokens=cache_read_tokens, + cache_write_tokens=cache_write_tokens, + uncached_tokens=uncached_input_tokens, + ) + else: + cache_write_tokens = _infer_openai_cache_write_tokens( + total_input_tokens, + cache_read_tokens, + ) + uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens) + + effective_optimized_tokens = ( + total_input_tokens if total_input_tokens > 0 else optimized_tokens + ) + effective_original_tokens = max( + original_tokens, + effective_optimized_tokens + tokens_saved, + ) + + _resp_log_tags = { + **(tags or {}), + "auth_mode": auth_mode.value if auth_mode else "payg", + "endpoint": "responses_http", + } + + # OpenAI Responses HTTP (non-WS, non-streaming). Codex + # uses this path when configured for HTTP transport. + # Pre-refactor `cache_hit` was hardcoded False on + # RequestLog even when cache_read>0 — funnel derives + # it correctly. + from headroom.proxy.helpers import compute_turn_id + + await self._record_request_outcome( + RequestOutcome( + request_id=request_id, + provider="openai", + model=model, + status_code=response.status_code, + original_tokens=effective_original_tokens, + optimized_tokens=effective_optimized_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + attempted_input_tokens=attempted_input_tokens, + cache_read_tokens=cache_read_tokens, + cache_write_tokens=cache_write_tokens, + uncached_input_tokens=uncached_input_tokens, + total_latency_ms=total_latency, + overhead_ms=optimization_latency, + transforms_applied=tuple(transforms_applied), + waste_signals=waste_signals_dict, + num_messages=len(messages) if isinstance(messages, list) else 0, + tags=_resp_log_tags, + turn_id=compute_turn_id(model, body.get("instructions"), messages), + request_messages=messages + if getattr(self.config, "log_full_messages", False) + else None, + client=client, + ) + ) + + logger.info( + f"[{request_id}] /v1/responses {model}: {total_input_tokens:,} tokens" + ) + + # Capture Codex rate-limit window data from response headers + from headroom.subscription.codex_rate_limits import ( + get_codex_rate_limit_state, + ) + + get_codex_rate_limit_state().update_from_headers(dict(response.headers)) + + # Remove compression headers + response_headers = _sanitize_forwarded_response_headers(response.headers) + + if buffered_stream_ccr and response.status_code == 200 and resp_json: + sse_headers = { + k: v + for k, v in response_headers.items() + if k.lower() not in ("content-length", "content-type") + } + if _ccr_response_handler and _ccr_response_handler.has_ccr_tool_calls( + resp_json, "openai_responses" + ): + # Handling above didn't fully resolve the retrieve + # call (e.g. max rounds hit, or it was mixed with a + # non-CCR tool call). Fail closed rather than stream + # a response the client can't act on — matches the + # Anthropic buffered path's residual-CCR guard. + logger.warning( + f"[{request_id}] CCR: Buffered streaming Responses " + "reply still contains headroom_retrieve after " + "handling; failing closed" + ) + + async def _residual_ccr_error_sse(): + error_event = { + "type": "error", + "error": { + "message": "Unable to safely complete streamed CCR retrieval.", + }, + } + yield f"event: error\ndata: {json.dumps(error_event)}\n\n".encode() + + return StreamingResponse( + _residual_ccr_error_sse(), + media_type="text/event-stream", + headers=sse_headers, + status_code=502, + ) + + async def _buffered_ccr_sse(): + for event in _openai_responses_to_sse(resp_json): + yield event return StreamingResponse( - _residual_ccr_error_sse(), + _buffered_ccr_sse(), media_type="text/event-stream", headers=sse_headers, - status_code=502, ) - async def _buffered_ccr_sse(): - for event in _openai_responses_to_sse(resp_json): - yield event - - return StreamingResponse( - _buffered_ccr_sse(), - media_type="text/event-stream", - headers=sse_headers, + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, ) - return Response( - content=response.content, - status_code=response.status_code, - headers=response_headers, - ) + if buffered_stream_ccr: + operation = asyncio.create_task(_buffered_ccr_operation()) + record_failed = self.metrics.record_failed + + class _BufferedCCRResponse(Response): + async def __call__(self, scope, receive, send): # noqa: ANN001 + await asyncio.sleep(0) + loop = asyncio.get_running_loop() + keepalive_deadline = loop.time() + 1.0 + started = False + try: + while True: + timeout = ( + 0.25 if started else max(0.0, keepalive_deadline - loop.time()) + ) + done, _ = await asyncio.wait({operation}, timeout=timeout) + if done: + try: + result = operation.result() + except Exception as e: + await record_failed(provider="openai") + logger.error( + f"[{request_id}] OpenAI responses request failed: {type(e).__name__}: {e}" + ) + if not started: + await send( + { + "type": "http.response.start", + "status": 502, + "headers": [ + (b"content-type", b"application/json") + ], + } + ) + await send( + { + "type": "http.response.body", + "body": json.dumps( + { + "error": { + "message": "An error occurred while processing your request. Please try again.", + "type": "server_error", + "code": "proxy_error", + } + } + ).encode(), + "more_body": False, + } + ) + return + await send( + { + "type": "http.response.body", + "body": b'event: error\ndata: {"type":"error","error":{"message":"An error occurred while processing the request."}}\n\n', + "more_body": False, + } + ) + return + + if not started: + await result(scope, receive, send) + return + + body_iterator = getattr(result, "body_iterator", None) + if body_iterator is not None: + async for chunk in body_iterator: + await send( + { + "type": "http.response.body", + "body": chunk, + "more_body": True, + } + ) + await send( + { + "type": "http.response.body", + "body": b"", + "more_body": False, + } + ) + return + + await send( + { + "type": "http.response.body", + "body": b'event: error\ndata: {"type":"error","error":{"message":"An error occurred while processing the request."}}\n\n', + "more_body": False, + } + ) + return + + if not started: + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"text/event-stream")], + } + ) + started = True + await send( + { + "type": "http.response.body", + "body": b'event: ping\ndata: {"type":"ping"}\n\n', + "more_body": True, + } + ) + except asyncio.CancelledError: + raise + finally: + if not operation.done(): + operation.cancel() + try: + await operation + except asyncio.CancelledError: + pass + except Exception: + pass + + return _BufferedCCRResponse(media_type="text/event-stream") + return await _buffered_ccr_operation() except Exception as e: await self.metrics.record_failed(provider="openai") logger.error(f"[{request_id}] OpenAI responses request failed: {type(e).__name__}: {e}") diff --git a/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py b/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py index 139a6fe09..3fcd3daf4 100644 --- a/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py +++ b/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py @@ -2,7 +2,9 @@ from __future__ import annotations +import asyncio import json +import logging from unittest.mock import AsyncMock, patch import pytest @@ -12,6 +14,7 @@ httpx = pytest.importorskip("httpx") from fastapi.responses import StreamingResponse # noqa: E402 from fastapi.testclient import TestClient # noqa: E402 +from starlette.requests import Request # noqa: E402 from headroom.cache.compression_store import get_compression_store # noqa: E402 from headroom.ccr.tool_injection import create_ccr_tool_definition # noqa: E402 @@ -49,6 +52,10 @@ def _message_response(content: list[dict], *, stop_reason: str = "end_turn") -> } +def _is_client_visible_sse(body: bytes) -> bool: + return b"event:" in body or b"data:" in body + + class _ContinuationClient: def __init__(self, response_json: dict) -> None: self.response_json = response_json @@ -299,9 +306,8 @@ def test_unresolved_ccr_only_streams_through_as_200() -> None: """CCR-only turn that never resolves: the model keeps re-emitting headroom_retrieve so the continuation exhausts its retrieval rounds with a residual marker and no accompanying client tool. Per #2089 the streaming - path no longer hard-502s here — it streams the residual headroom_retrieve - back as a 200 SSE so the client (which owns the tool) can resolve or retry - it, matching the non-streaming path. It must NOT 502.""" + path streams the residual headroom_retrieve back as a 200 SSE so the client + can resolve or retry it, matching the non-streaming path.""" config = _make_config() persistent_ccr = _message_response( [ @@ -343,8 +349,266 @@ def test_unresolved_ccr_only_streams_through_as_200() -> None: }, ) - # Fails closed no longer: residual CCR is handed back to the client as 200 SSE. assert resp.status_code == 200, resp.text assert "text/event-stream" in resp.headers["content-type"] assert "headroom_retrieve" in resp.text - assert "Unable to safely complete streamed CCR retrieval" not in resp.text + + +@pytest.mark.asyncio +async def test_buffered_ccr_emits_keepalive_before_delayed_upstream() -> None: + config = _make_config() + final_response = _message_response([{"type": "text", "text": "done"}]) + started = asyncio.Event() + release = asyncio.Event() + body = { + "model": "claude-sonnet-4-6", + "max_tokens": 64, + "stream": True, + "tools": [create_ccr_tool_definition("anthropic")], + "messages": [{"role": "user", "content": "wait"}], + } + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + + scope = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/messages", + "raw_path": b"/v1/messages", + "query_string": b"", + "headers": [(b"x-api-key", b"test-key"), (b"anthropic-version", b"2023-06-01")], + "server": ("testserver", 80), + "client": ("testclient", 123), + "root_path": "", + } + + with patch("headroom.proxy.server.AnyLLMBackend"): + app = create_app(config) + with TestClient(app): + proxy = app.state.proxy + + async def delayed_retry(*args, **kwargs): # noqa: ANN002, ANN003 + started.set() + await release.wait() + return httpx.Response(200, json=final_response) + + proxy._retry_request = delayed_retry + task = asyncio.create_task(proxy.handle_anthropic_messages(Request(scope, receive))) + await started.wait() + response = await asyncio.wait_for(asyncio.shield(task), 1) + events: list[dict] = [] + first_visible_body = asyncio.Event() + + async def send(message): # noqa: ANN001 + events.append(message) + if message["type"] == "http.response.body" and _is_client_visible_sse( + message["body"] + ): + first_visible_body.set() + + response_task = asyncio.create_task(response(scope, receive, send)) + await asyncio.wait_for(first_visible_body.wait(), 2) + assert not release.is_set() + release.set() + await response_task + + bodies = [event["body"] for event in events if event["type"] == "http.response.body"] + assert bodies[0] == b'event: ping\ndata: {"type":"ping"}\n\n' + assert b"done" in b"".join(bodies) + + +@pytest.mark.asyncio +async def test_buffered_ccr_preserves_early_failure_status_and_headers() -> None: + config = _make_config() + body = { + "model": "claude-sonnet-4-6", + "max_tokens": 64, + "stream": True, + "tools": [create_ccr_tool_definition("anthropic")], + "messages": [{"role": "user", "content": "fail early"}], + } + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + + scope = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/messages", + "raw_path": b"/v1/messages", + "query_string": b"", + "headers": [(b"x-api-key", b"test-key"), (b"anthropic-version", b"2023-06-01")], + "server": ("testserver", 80), + "client": ("testclient", 123), + "root_path": "", + } + + with patch("headroom.proxy.server.AnyLLMBackend"): + app = create_app(config) + with TestClient(app): + proxy = app.state.proxy + + async def early_failure(*args, **kwargs): # noqa: ANN002, ANN003 + await asyncio.sleep(0.05) + return httpx.Response( + 429, + headers={"retry-after": "7"}, + json={"error": {"message": "slow down"}}, + ) + + proxy._retry_request = early_failure + response = await proxy.handle_anthropic_messages(Request(scope, receive)) + events: list[dict] = [] + + async def send(message): # noqa: ANN001 + events.append(message) + + await response(scope, receive, send) + + start = next(event for event in events if event["type"] == "http.response.start") + headers = dict(start["headers"]) + assert start["status"] == 429 + assert headers[b"retry-after"] == b"7" + assert b": headroom-keepalive\n\n" not in b"".join( + event["body"] for event in events if event["type"] == "http.response.body" + ) + + +@pytest.mark.asyncio +async def test_buffered_ccr_late_failure_emits_sanitized_error_event() -> None: + config = _make_config() + started = asyncio.Event() + release = asyncio.Event() + body = { + "model": "claude-sonnet-4-6", + "max_tokens": 64, + "stream": True, + "tools": [create_ccr_tool_definition("anthropic")], + "messages": [{"role": "user", "content": "wait"}], + } + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + + scope = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/messages", + "raw_path": b"/v1/messages", + "query_string": b"", + "headers": [(b"x-api-key", b"test-key"), (b"anthropic-version", b"2023-06-01")], + "server": ("testserver", 80), + "client": ("testclient", 123), + "root_path": "", + } + + with patch("headroom.proxy.server.AnyLLMBackend"): + app = create_app(config) + with TestClient(app): + proxy = app.state.proxy + proxy_logger = logging.getLogger("headroom.proxy") + error_records: list[logging.LogRecord] = [] + log_handler = logging.Handler() + log_handler.setLevel(logging.ERROR) + log_handler.emit = error_records.append + proxy_logger.addHandler(log_handler) + + async def delayed_failure(*args, **kwargs): # noqa: ANN002, ANN003 + started.set() + await release.wait() + raise RuntimeError("boom") + + with patch.object( + proxy.metrics, "record_failed", new_callable=AsyncMock + ) as record_failed: + proxy._retry_request = delayed_failure + task = asyncio.create_task(proxy.handle_anthropic_messages(Request(scope, receive))) + await started.wait() + response = await asyncio.wait_for(asyncio.shield(task), 1) + events: list[dict] = [] + first_body = asyncio.Event() + + async def send(message): # noqa: ANN001 + events.append(message) + if message["type"] == "http.response.body" and message["body"]: + first_body.set() + + response_task = asyncio.create_task(response(scope, receive, send)) + await asyncio.wait_for(first_body.wait(), 2) + release.set() + await response_task + record_failed.assert_awaited_once_with(provider="anthropic") + proxy_logger.removeHandler(log_handler) + + bodies = [event["body"] for event in events if event["type"] == "http.response.body"] + assert bodies[0] == b'event: ping\ndata: {"type":"ping"}\n\n' + assert b"An error occurred while processing the request." in bodies[-1] + assert b"boom" not in bodies[-1] + assert events[-1]["more_body"] is False + assert any( + record.levelno == logging.ERROR and "RuntimeError: boom" in record.getMessage() + for record in error_records + ) + + +@pytest.mark.asyncio +async def test_buffered_ccr_pre_keepalive_exception_returns_json_error() -> None: + config = _make_config() + body = { + "model": "claude-sonnet-4-6", + "max_tokens": 64, + "stream": True, + "tools": [create_ccr_tool_definition("anthropic")], + "messages": [{"role": "user", "content": "fail before keepalive"}], + } + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + + scope = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/messages", + "raw_path": b"/v1/messages", + "query_string": b"", + "headers": [(b"x-api-key", b"test-key"), (b"anthropic-version", b"2023-06-01")], + "server": ("testserver", 80), + "client": ("testclient", 123), + "root_path": "", + } + + with patch("headroom.proxy.server.AnyLLMBackend"): + app = create_app(config) + with TestClient(app): + proxy = app.state.proxy + + async def early_exception(*args, **kwargs): # noqa: ANN002, ANN003 + raise RuntimeError("boom") + + proxy._retry_request = early_exception + response = await proxy.handle_anthropic_messages(Request(scope, receive)) + events: list[dict] = [] + + async def send(message): # noqa: ANN001 + events.append(message) + + await response(scope, receive, send) + + start = next(event for event in events if event["type"] == "http.response.start") + bodies = [event["body"] for event in events if event["type"] == "http.response.body"] + assert start["status"] == 502 + assert dict(start["headers"])[b"content-type"] == b"application/json" + payload = json.loads(bodies[-1].decode()) + assert ( + payload["error"]["message"] + == "An error occurred while processing your request. Please try again." + ) diff --git a/tests/test_proxy/test_openai_responses_ccr.py b/tests/test_proxy/test_openai_responses_ccr.py index fe2602c14..9bae90cd2 100644 --- a/tests/test_proxy/test_openai_responses_ccr.py +++ b/tests/test_proxy/test_openai_responses_ccr.py @@ -10,16 +10,19 @@ tests/test_proxy/test_openai_backend_path.py for that precedent. from __future__ import annotations +import asyncio import json +import logging import pytest fastapi = pytest.importorskip("fastapi") httpx = pytest.importorskip("httpx") -from unittest.mock import AsyncMock, MagicMock # noqa: E402 +from unittest.mock import AsyncMock, MagicMock, patch # noqa: E402 from fastapi.testclient import TestClient # noqa: E402 +from starlette.requests import Request # noqa: E402 from headroom.cache.compression_store import reset_compression_store # noqa: E402 from headroom.ccr.tool_injection import CCR_TOOL_NAME # noqa: E402 @@ -260,3 +263,247 @@ def test_streaming_request_without_retrieve_tool_uses_normal_stream_path(): assert resp.status_code == 200, resp.text assert stream_called["value"] is True + + +@pytest.mark.asyncio +async def test_buffered_responses_ccr_emits_keepalive_before_delayed_upstream(): + app = _make_app() + body = { + "model": "gpt-5-codex", + "input": "please wait", + "tools": [_RETRIEVE_TOOL], + "stream": True, + } + started = asyncio.Event() + release = asyncio.Event() + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + + scope = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/responses", + "raw_path": b"/v1/responses", + "query_string": b"", + "headers": [(b"authorization", b"Bearer sk-test")], + "server": ("testserver", 80), + "client": ("testclient", 123), + "root_path": "", + } + + with TestClient(app): + server = app.state.proxy + + async def delayed_retry(*args, **kwargs): # noqa: ANN002, ANN003 + started.set() + await release.wait() + return _final_response("https://api.openai.com/v1/responses") + + server._retry_request = delayed_retry + task = asyncio.create_task(server.handle_openai_responses(Request(scope, receive))) + await started.wait() + response = await asyncio.wait_for(asyncio.shield(task), 1) + events: list[dict] = [] + first_body = asyncio.Event() + + async def send(message): # noqa: ANN001 + events.append(message) + if message["type"] == "http.response.body" and message["body"]: + first_body.set() + + response_task = asyncio.create_task(response(scope, receive, send)) + await asyncio.wait_for(first_body.wait(), 2) + assert not release.is_set() + release.set() + await response_task + + bodies = [event["body"] for event in events if event["type"] == "http.response.body"] + assert bodies[0] == b'event: ping\ndata: {"type":"ping"}\n\n' + assert b"Resolved!" in b"".join(bodies) + + +@pytest.mark.asyncio +async def test_buffered_responses_ccr_preserves_early_failure_status_and_headers(): + app = _make_app() + body = { + "model": "gpt-5-codex", + "input": "fail early", + "tools": [_RETRIEVE_TOOL], + "stream": True, + } + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + + scope = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/responses", + "raw_path": b"/v1/responses", + "query_string": b"", + "headers": [(b"authorization", b"Bearer sk-test")], + "server": ("testserver", 80), + "client": ("testclient", 123), + "root_path": "", + } + + with TestClient(app): + server = app.state.proxy + + async def early_failure(*args, **kwargs): # noqa: ANN002, ANN003 + await asyncio.sleep(0.05) + return httpx.Response( + 429, + headers={"retry-after": "7"}, + json={"error": {"message": "slow down"}}, + request=httpx.Request("POST", "https://api.openai.com/v1/responses"), + ) + + server._retry_request = early_failure + response = await server.handle_openai_responses(Request(scope, receive)) + events: list[dict] = [] + + async def send(message): # noqa: ANN001 + events.append(message) + + await response(scope, receive, send) + + start = next(event for event in events if event["type"] == "http.response.start") + headers = dict(start["headers"]) + assert start["status"] == 429 + assert headers[b"retry-after"] == b"7" + assert b": headroom-keepalive\n\n" not in b"".join( + event["body"] for event in events if event["type"] == "http.response.body" + ) + + +@pytest.mark.asyncio +async def test_buffered_responses_ccr_late_failure_emits_sanitized_error_event(): + app = _make_app() + body = { + "model": "gpt-5-codex", + "input": "please wait", + "tools": [_RETRIEVE_TOOL], + "stream": True, + } + started = asyncio.Event() + release = asyncio.Event() + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + + scope = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/responses", + "raw_path": b"/v1/responses", + "query_string": b"", + "headers": [(b"authorization", b"Bearer sk-test")], + "server": ("testserver", 80), + "client": ("testclient", 123), + "root_path": "", + } + + with TestClient(app): + server = app.state.proxy + proxy_logger = logging.getLogger("headroom.proxy") + error_records: list[logging.LogRecord] = [] + log_handler = logging.Handler() + log_handler.setLevel(logging.ERROR) + log_handler.emit = error_records.append + proxy_logger.addHandler(log_handler) + + async def delayed_failure(*args, **kwargs): # noqa: ANN002, ANN003 + started.set() + await release.wait() + raise RuntimeError("boom") + + with patch.object(server.metrics, "record_failed", new_callable=AsyncMock) as record_failed: + server._retry_request = delayed_failure + task = asyncio.create_task(server.handle_openai_responses(Request(scope, receive))) + await started.wait() + response = await asyncio.wait_for(asyncio.shield(task), 1) + events: list[dict] = [] + first_body = asyncio.Event() + + async def send(message): # noqa: ANN001 + events.append(message) + if message["type"] == "http.response.body" and message["body"]: + first_body.set() + + response_task = asyncio.create_task(response(scope, receive, send)) + await asyncio.wait_for(first_body.wait(), 2) + release.set() + await response_task + record_failed.assert_awaited_once_with(provider="openai") + proxy_logger.removeHandler(log_handler) + + bodies = [event["body"] for event in events if event["type"] == "http.response.body"] + assert bodies[0] == b'event: ping\ndata: {"type":"ping"}\n\n' + assert b"An error occurred while processing the request." in bodies[-1] + assert b"boom" not in bodies[-1] + assert events[-1]["more_body"] is False + assert any( + record.levelno == logging.ERROR and "RuntimeError: boom" in record.getMessage() + for record in error_records + ) + + +@pytest.mark.asyncio +async def test_buffered_responses_ccr_pre_keepalive_exception_returns_json_error(): + app = _make_app() + body = { + "model": "gpt-5-codex", + "input": "fail before keepalive", + "tools": [_RETRIEVE_TOOL], + "stream": True, + } + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + + scope = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/responses", + "raw_path": b"/v1/responses", + "query_string": b"", + "headers": [(b"authorization", b"Bearer sk-test")], + "server": ("testserver", 80), + "client": ("testclient", 123), + "root_path": "", + } + + with TestClient(app): + server = app.state.proxy + + async def early_exception(*args, **kwargs): # noqa: ANN002, ANN003 + raise RuntimeError("boom") + + server._retry_request = early_exception + response = await server.handle_openai_responses(Request(scope, receive)) + events: list[dict] = [] + + async def send(message): # noqa: ANN001 + events.append(message) + + await response(scope, receive, send) + + start = next(event for event in events if event["type"] == "http.response.start") + bodies = [event["body"] for event in events if event["type"] == "http.response.body"] + assert start["status"] == 502 + assert dict(start["headers"])[b"content-type"] == b"application/json" + payload = json.loads(bodies[-1].decode()) + assert ( + payload["error"]["message"] + == "An error occurred while processing your request. Please try again." + )