diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 16af50b54..79045fa84 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -4195,6 +4195,25 @@ class AnthropicHandlerMixin: headers=relay_headers, ) + if buffered_stream_ccr and response.status_code == 200 and not resp_json: + logger.warning( + f"[{request_id}] CCR: rejecting malformed buffered 200 reply " + f"(content-type={response.headers.get('content-type')!r}, " + f"body_bytes={len(response.content)})" + ) + return Response( + content=json.dumps( + { + "error": { + "type": "upstream_protocol_error", + "message": "Upstream returned an invalid buffered response.", + } + } + ), + status_code=502, + media_type="application/json", + ) + if buffered_stream_ccr and response.status_code == 200 and resp_json: sse_headers = { k: v @@ -4293,122 +4312,50 @@ class AnthropicHandlerMixin: 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 + # Send nothing until the buffered operation resolves. + # The previous keepalive preamble committed + # `200 text/event-stream` after 1s, i.e. before the + # outcome was known: any upstream reply that then + # failed to become SSE (non-200, unparseable body) + # reached the client as a 200 whose body carried no + # `message_start`, which Claude Code reports as "API + # returned an empty or malformed response (HTTP 200) — + # check for a proxy or gateway intercepting the + # request". The real status was lost with it, so + # client-side 429/5xx backoff never fired. Clients + # budget minutes for a turn (Claude Code sends + # `x-stainless-timeout: 600`), so waiting is free. 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( + result = await operation + except Exception as e: + await record_failed(provider=provider_name) + logger.error( + f"[{request_id}] Request failed: {type(e).__name__}: {e}" + ) + 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": "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, + "type": "error", + "error": { + "type": "api_error", + "message": "An error occurred while processing your request. Please try again.", + }, } - ) - 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 + ).encode(), + "more_body": False, + } + ) + return + await result(scope, receive, send) return _BufferedCCRResponse(media_type="text/event-stream") return await _buffered_ccr_operation() diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 827b7f0f9..c011b4984 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -763,36 +763,9 @@ def _compact_openai_responses_tools( return compact_tools(payload) -def _responses_request_allows_memory_tool_continuation(payload: dict[str, Any]) -> bool: - """Return whether Responses memory tools may rely on stored continuations. - - Headroom memory tools use ``previous_response_id`` continuations after a - tool call. Those continuations require the originating response to be - stored. When a client explicitly sends ``store=false``, preserve that - contract and skip the Responses memory-tool injection path instead of - mutating the request. - """ - - return payload.get("store") is not False - - -def _ensure_responses_store_for_memory_tools( - payload: dict[str, Any], - *, - memory_tools_injected: bool, -) -> bool: - """Return True when memory-tool injection requires and receives store=true.""" - - if memory_tools_injected and payload.get("store") is not True: - payload["store"] = True - return True - return False - - def _allow_responses_memory_tools(is_chatgpt_auth: bool) -> bool: - # ChatGPT Codex rejects Responses payloads unless store=false. The - # transparent memory-tool continuation flow needs stored responses, so keep - # it on the regular API path only. + # Preserve the ChatGPT Codex route's existing store policy and memory-tool + # exclusion while API Responses memory continuations stay stateless. return not is_chatgpt_auth @@ -1047,6 +1020,23 @@ def _responses_input_to_items(input_data: Any) -> list[dict[str, Any]]: return [] +def _responses_stateless_output_items(output_items: Any) -> list[dict[str, Any]]: + """Return response items that can be replayed without provider state.""" + if not isinstance(output_items, list): + return [] + return [ + item + for item in output_items + if isinstance(item, dict) + and not (item.get("type") == "reasoning" and not item.get("encrypted_content")) + ] + + +def _responses_stateless_input_items(input_data: Any) -> list[dict[str, Any]]: + """Normalize input and remove reasoning items without encrypted content.""" + return _responses_stateless_output_items(_responses_input_to_items(input_data)) + + def _dedup_responses_output_items( items: list[dict[str, Any]], output_types: frozenset[str], @@ -5345,10 +5335,7 @@ class OpenAIHandlerMixin: else: memory_tool_defs_responses.append(t) - if ( - responses_memory_tools_allowed - and _responses_request_allows_memory_tool_continuation(body) - ): + if responses_memory_tools_allowed: resp_tools = body.get("tools") or [] resp_tools, mem_tools_injected = _apply_sticky_mem_tools_resp( provider="openai", @@ -5360,23 +5347,16 @@ class OpenAIHandlerMixin: ) if mem_tools_injected: body["tools"] = resp_tools + include = body.get("include") + if isinstance(include, list): + if "reasoning.encrypted_content" not in include: + body["include"] = [*include, "reasoning.encrypted_content"] + elif include is None: + body["include"] = ["reasoning.encrypted_content"] body_mutation_tracker.mark_mutated("responses_memory_tools") logger.info( f"[{request_id}] Memory: Injected memory tools (openai/responses)" ) - if _ensure_responses_store_for_memory_tools( - body, - memory_tools_injected=True, - ): - body_mutation_tracker.mark_mutated("responses_memory_store") - logger.info( - f"[{request_id}] Memory: forced store=true for Responses memory tool continuation" - ) - elif self.memory_handler.config.inject_tools: - logger.info( - "[%s] Memory: skipped Responses memory tools because client set store=false", - request_id, - ) except Exception as e: logger.warning(f"[{request_id}] Memory injection failed (responses): {e}") elif self.memory_handler and memory_user_id and _bypass: @@ -5808,6 +5788,7 @@ class OpenAIHandlerMixin: total_input_tokens = original_tokens # fallback output_tokens = 0 cache_read_tokens = 0 + resp_json = None try: resp_json = response.json() usage = resp_json.get("usage", {}) @@ -5826,7 +5807,13 @@ class OpenAIHandlerMixin: 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: + except ( + json.JSONDecodeError, + ValueError, + KeyError, + TypeError, + AttributeError, + ) as e: logger.debug( f"[{request_id}] Failed to extract cached tokens from OpenAI passthrough response: {e}" ) @@ -5996,17 +5983,17 @@ class OpenAIHandlerMixin: ) if tool_outputs: - # Make continuation request with tool results - response_id = resp_json.get("id") + # Make a stateless continuation with the complete + # item history instead of relying on retained state. continuation_body = { - "model": model, - "input": tool_outputs, + **body, + "input": [ + *_responses_stateless_input_items(body.get("input")), + *_responses_stateless_output_items(output_items), + *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 + continuation_body.pop("previous_response_id", None) cont_response = await self._retry_request( "POST", url, headers, continuation_body @@ -6144,6 +6131,25 @@ class OpenAIHandlerMixin: headers=sse_headers, ) + if buffered_stream_ccr and response.status_code == 200 and not resp_json: + logger.warning( + f"[{request_id}] CCR: rejecting malformed buffered Responses 200 " + f"reply (content-type={response.headers.get('content-type')!r}, " + f"body_bytes={len(response.content)})" + ) + return Response( + content=json.dumps( + { + "error": { + "type": "upstream_protocol_error", + "message": "Upstream returned an invalid buffered response.", + } + } + ), + status_code=502, + media_type="application/json", + ) + # Inline marker resolution, non-streaming only. Runs # outside the has_ccr_tool_calls gate above on purpose: # the #2509 case has no retrieve tool call at all. @@ -6171,118 +6177,44 @@ class OpenAIHandlerMixin: 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 + # Send nothing until the buffered operation resolves — + # see the AnthropicHandler twin for the full rationale. + # Committing `200 text/event-stream` on a keepalive timer, + # before the outcome is known, turns every non-200 or + # unparseable upstream reply into a 200 with no usable + # body and discards the status the client needs to back + # off on. 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( + result = await operation + except Exception as e: + await record_failed(provider="openai") + logger.error( + f"[{request_id}] OpenAI responses request failed: {type(e).__name__}: {e}" + ) + 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": "http.response.body", - "body": b'event: error\ndata: {"type":"error","error":{"message":"An error occurred while processing the request."}}\n\n', - "more_body": False, + "error": { + "message": "An error occurred while processing your request. Please try again.", + "type": "server_error", + "code": "proxy_error", + } } - ) - 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 + ).encode(), + "more_body": False, + } + ) + return + await result(scope, receive, send) return _BufferedCCRResponse(media_type="text/event-stream") return await _buffered_ccr_operation() @@ -6827,6 +6759,8 @@ class OpenAIHandlerMixin: return json.dumps(flattened, ensure_ascii=False) body: dict[str, Any] = {} + current_response_input: list[dict[str, Any]] = [] + current_response_template: dict[str, Any] = {} tokens_saved = 0 # Session-scoped accumulator for tokens we *attempted* to # compress (extracted units + schema). Drives the active- @@ -7025,7 +6959,8 @@ class OpenAIHandlerMixin: try: # Unwrap response.create envelope to access the response body ws_response_body = frame_body.get("response", frame_body) - + if not isinstance(ws_response_body, dict): + return frame_raw # Per-project memory routing (GH #462). For WS, # ``ws_response_body`` carries ``instructions`` — # that's the system-prompt-equivalent we feed to the @@ -7171,6 +7106,15 @@ class OpenAIHandlerMixin: ) if mem_injected: ws_response_body["tools"] = ws_tools + include = ws_response_body.get("include") + if isinstance(include, list): + if "reasoning.encrypted_content" not in include: + ws_response_body["include"] = [ + *include, + "reasoning.encrypted_content", + ] + elif include is None: + ws_response_body["include"] = ["reasoning.encrypted_content"] # Add memory instruction so the model uses # memory tools as persistent cross-session knowledge. @@ -7210,6 +7154,12 @@ class OpenAIHandlerMixin: body.get("type") == "response.create" or ("type" not in body and "input" in body) ): first_msg_raw = await _prepare_memory_frame(body, first_msg_raw) + first_response_body = body.get("response", body) + current_response_input = _responses_input_to_items( + first_response_body.get("input") + if isinstance(first_response_body, dict) + else None + ) # Hot-fix follow-up to PR #406 — inline Rust compression on the # WS first frame before forwarding upstream. PR #406 enabled @@ -7469,6 +7419,19 @@ class OpenAIHandlerMixin: ) first_msg_raw = _normalize_ws_response_create_for_upstream(first_msg_raw) + try: + final_first_body = json.loads(first_msg_raw) + except json.JSONDecodeError: + final_first_body = None + if isinstance(final_first_body, dict): + final_first_response = final_first_body.get("response", final_first_body) + if isinstance(final_first_response, dict): + current_response_template = dict(final_first_response) + current_response_template.pop("input", None) + current_response_template.pop("previous_response_id", None) + current_response_input = _responses_input_to_items( + final_first_response.get("input") + ) _first_upstream_body: Any = None try: _first_upstream_body = json.loads(first_msg_raw) @@ -7807,6 +7770,8 @@ class OpenAIHandlerMixin: nonlocal ws_client_frames_total, ws_cancel_frames nonlocal ws_frames_compressed nonlocal ws_last_client_frame_type, ws_client_disconnect_seen + nonlocal current_response_input + nonlocal current_response_template client_frame_index = 1 try: while True: @@ -7859,6 +7824,14 @@ class OpenAIHandlerMixin: and _inbound_frame_body.get("type") == "response.create" ): ws_response_create_frames += 1 + inbound_response = _inbound_frame_body.get( + "response", _inbound_frame_body + ) + current_response_input = _responses_input_to_items( + inbound_response.get("input") + if isinstance(inbound_response, dict) + else None + ) msg = await _prepare_memory_frame(_inbound_frame_body, msg) ( msg, @@ -7903,6 +7876,27 @@ class OpenAIHandlerMixin: ) msg = _normalize_ws_response_create_for_upstream(msg) + if ws_last_client_frame_type == "response.create": + try: + outbound_body = json.loads(msg) + outbound_body = ( + outbound_body.get("response", outbound_body) + if isinstance(outbound_body, dict) + else {} + ) + if isinstance(outbound_body, dict): + current_response_template = dict(outbound_body) + current_response_template.pop("input", None) + current_response_template.pop( + "previous_response_id", None + ) + candidate_input = _responses_input_to_items( + outbound_body.get("input") + ) + if candidate_input: + current_response_input = candidate_input + except (json.JSONDecodeError, AttributeError, TypeError): + current_response_input = [] _outbound_frame_body: Any = None try: _outbound_frame_body = json.loads(msg) @@ -7959,10 +7953,11 @@ class OpenAIHandlerMixin: """Relay upstream→client with transparent memory tool handling. Uses a buffer-then-decide approach: - 1. Buffer events until first output item arrives - 2. If first output is a memory tool → suppress entire response, - execute tools silently, send continuation upstream - 3. If first output is non-memory → flush buffer, stream normally + 1. Buffer events until the first non-reasoning output item arrives + 2. If a memory tool arrives, suppress the response and execute it + silently before sending a stateless continuation upstream + 3. Flush ordinary output as soon as the response is known not to + begin with a reasoning-only prelude 4. Continuation response events are relayed to Codex seamlessly This prevents orphaned response.created events from confusing Codex. @@ -7987,25 +7982,22 @@ class OpenAIHandlerMixin: nonlocal ws_recorded_overhead_ms_total, ws_recorded_ttfb_ms nonlocal ws_upstream_frames_total, ws_last_upstream_frame_type nonlocal ws_ttfb_ms - - memory_enabled = bool( - self.memory_handler and memory_user_id and ws_memory_tools_allowed - ) + nonlocal current_response_input # Per-response state (reset after each response.completed) event_buffer: list[str] = [] decided = False suppress_response = False pending_fcs: list[dict[str, Any]] = [] - resp_id: str | None = None + response_output_items: list[dict[str, Any]] = [] def _reset() -> None: - nonlocal decided, suppress_response, resp_id + nonlocal decided, suppress_response event_buffer.clear() decided = False suppress_response = False pending_fcs.clear() - resp_id = None + response_output_items.clear() response_started_ms: float | None = None @@ -8233,6 +8225,11 @@ class OpenAIHandlerMixin: ws_cache_write_tokens_total += usage_cache_write_tokens ws_uncached_input_tokens_total += usage_uncached_tokens + memory_enabled = bool( + self.memory_handler + and memory_user_id + and ws_memory_tools_allowed + ) if not memory_enabled: if event_type == "response.completed": response_completed_seen = True @@ -8240,17 +8237,25 @@ class OpenAIHandlerMixin: await websocket.send_text(msg_str) continue - # --- Phase 1: Buffer until first output item --- - if not decided: - event_buffer.append(msg_str) - - if event_type == "response.output_item.added": - item = event.get("item", {}) + if event_type == "response.output_item.done": + item = event.get("item", {}) + if isinstance(item, dict): + response_output_items.append(item) if ( item.get("type") == "function_call" and item.get("name") in MEMORY_TOOL_NAMES ): - # Memory tool first → suppress entire response + pending_fcs.append(item) + + if not decided: + event_buffer.append(msg_str) + if event_type == "response.output_item.added": + item = event.get("item", {}) + if ( + isinstance(item, dict) + and item.get("type") == "function_call" + and item.get("name") in MEMORY_TOOL_NAMES + ): suppress_response = True decided = True event_buffer.clear() @@ -8258,106 +8263,108 @@ class OpenAIHandlerMixin: f"[{request_id}] WS Memory: Detected " f"{item.get('name')} — suppressing response" ) - else: - # Non-memory first → flush buffer, pass through + elif not ( + isinstance(item, dict) + and item.get("type") == "reasoning" + ): decided = True for buf in event_buffer: await websocket.send_text(buf) event_buffer.clear() - + continue elif event_type == "response.completed": - # No output items at all — flush decided = True - for buf in event_buffer: - await websocket.send_text(buf) - event_buffer.clear() - await _record_ws_response_metrics() - _reset() - response_completed_seen = True - - continue - - # --- Phase 2a: Suppress mode (memory response) --- - if suppress_response: - if event_type == "response.output_item.done": - item = event.get("item", {}) - if ( - item.get("type") == "function_call" - and item.get("name") in MEMORY_TOOL_NAMES - ): - pending_fcs.append(item) - - elif event_type == "response.completed": response_completed_seen = True await _record_ws_response_metrics() - resp = event.get("response", {}) - resp_id = resp.get("id") - - if pending_fcs: - logger.info( - f"[{request_id}] WS Memory: Executing " - f"{len(pending_fcs)} tool(s) transparently" - ) - - # Execute memory tool calls - tool_outputs: list[dict[str, Any]] = [] - for fc in pending_fcs: - call_id = fc.get("call_id", fc.get("id", "")) - fc_name = fc.get("name", "") - args_str = fc.get("arguments") or "{}" - try: - fc_args = json.loads(args_str) - except (json.JSONDecodeError, TypeError): - fc_args = {} - - await self.memory_handler._ensure_initialized() - if self.memory_handler._backend: - result = await self.memory_handler._execute_memory_tool( - fc_name, - fc_args, - memory_user_id, - "openai", - ) - else: - result = json.dumps( - {"error": "backend not ready"} - ) - - tool_outputs.append( - { - "type": "function_call_output", - "call_id": call_id, - "output": result, - } - ) - logger.info( - f"[{request_id}] WS Memory: Executed " - f"{fc_name} for user {memory_user_id}" - ) - - # Send continuation upstream - cont: dict[str, Any] = { - "type": "response.create", - "response": {"input": tool_outputs}, - } - if resp_id: - cont["response"]["previous_response_id"] = resp_id - await upstream.send( - _normalize_ws_response_create_for_upstream( - json.dumps(cont) - ) - ) - logger.info( - f"[{request_id}] WS Memory: Sent continuation " - f"with {len(tool_outputs)} result(s)" - ) - + for buf in event_buffer: + await websocket.send_text(buf) + _reset() + continue + if not decided: + continue + + if not suppress_response: + await websocket.send_text(msg_str) + if event_type == "response.completed": + response_completed_seen = True + await _record_ws_response_metrics() _reset() - # All events suppressed in this mode continue - # --- Phase 2b: Pass-through mode --- - await websocket.send_text(msg_str) + if event_type != "response.completed": + continue + + response_completed_seen = True + await _record_ws_response_metrics() + + logger.info( + f"[{request_id}] WS Memory: Executing " + f"{len(pending_fcs)} tool(s) transparently" + ) + + # Execute memory tool calls. + tool_outputs: list[dict[str, Any]] = [] + for fc in pending_fcs: + call_id = fc.get("call_id", fc.get("id", "")) + fc_name = fc.get("name", "") + args_str = fc.get("arguments") or "{}" + try: + fc_args = json.loads(args_str) + except (json.JSONDecodeError, TypeError): + fc_args = {} + + await self.memory_handler._ensure_initialized() + if self.memory_handler._backend: + result = await self.memory_handler._execute_memory_tool( + fc_name, + fc_args, + memory_user_id, + "openai", + ) + else: + result = json.dumps({"error": "backend not ready"}) + + tool_outputs.append( + { + "type": "function_call_output", + "call_id": call_id, + "output": result, + } + ) + logger.info( + f"[{request_id}] WS Memory: Executed " + f"{fc_name} for user {memory_user_id}" + ) + + stateless_input = _responses_stateless_input_items( + current_response_input + ) + stateless_output = _responses_stateless_output_items( + response_output_items + ) + current_response_input = [ + *stateless_input, + *stateless_output, + *tool_outputs, + ] + continuation_response = { + **current_response_template, + "input": current_response_input, + } + continuation_response.pop("previous_response_id", None) + cont: dict[str, Any] = { + "type": "response.create", + "response": continuation_response, + } + continuation_raw = _strip_codex_lite_metadata(json.dumps(cont)) + await upstream.send( + _normalize_ws_response_create_for_upstream(continuation_raw) + ) + logger.info( + f"[{request_id}] WS Memory: Sent continuation " + f"with {len(tool_outputs)} result(s)" + ) + _reset() except asyncio.CancelledError: raise diff --git a/tests/test_openai_codex_routing.py b/tests/test_openai_codex_routing.py index 1666184f8..471d0d8ca 100644 --- a/tests/test_openai_codex_routing.py +++ b/tests/test_openai_codex_routing.py @@ -2,10 +2,12 @@ import asyncio import base64 import json import sys +from copy import deepcopy from types import SimpleNamespace from unittest.mock import MagicMock, patch import anyio +import httpx import pytest from fastapi import Request @@ -14,6 +16,7 @@ from headroom.proxy.handlers.openai import ( _is_allowed_websocket_origin, _openai_responses_unit_cache_key, _resolve_codex_routing_headers, + _responses_stateless_output_items, ) @@ -135,6 +138,30 @@ def test_openai_responses_unit_cache_key_includes_target_ratio() -> None: assert aggressive_key != balanced_key +def test_responses_stateless_output_items_drop_unencrypted_reasoning() -> None: + assert _responses_stateless_output_items(None) == [] + assert _responses_stateless_output_items( + [ + {"type": "reasoning", "id": "rs-unusable", "summary": []}, + { + "type": "reasoning", + "id": "rs-reusable", + "summary": [], + "encrypted_content": "encrypted", + }, + {"type": "function_call", "call_id": "call-1"}, + ] + ) == [ + { + "type": "reasoning", + "id": "rs-reusable", + "summary": [], + "encrypted_content": "encrypted", + }, + {"type": "function_call", "call_id": "call-1"}, + ] + + class _DummyMetrics: async def record_request(self, **kwargs): # noqa: ANN003 return None @@ -274,6 +301,106 @@ class _MemoryToolsOnlyHandler: return False +class _MemoryContinuationHandler(_MemoryToolsOnlyHandler): + async def _ensure_initialized(self) -> None: + self._backend = True + + async def _execute_memory_tool( + self, + name: str, + args: dict, + user_id: str, + provider: str, + ) -> str: + assert (name, args, user_id, provider) == ( + "memory_search", + {}, + "user-1", + "openai", + ) + return '{"memories": []}' + + def has_memory_tool_calls(self, response: dict, provider: str) -> bool: + assert provider == "openai" + return any( + item.get("name") == "memory_search" + for item in response.get("output", []) + if isinstance(item, dict) + ) + + +class _ZdrResponsesHandler(_DummyOpenAIHandler): + def __init__(self) -> None: + super().__init__() + self.memory_handler = _MemoryContinuationHandler() + self.requests: list[dict] = [] + + async def _retry_request(self, method: str, url: str, headers: dict, body: dict, **kwargs): + assert (method, url) == ("POST", "https://api.openai.com/v1/responses") + self.requests.append(deepcopy(body)) + request = httpx.Request(method, url) + if len(self.requests) == 1: + return httpx.Response( + 200, + json={ + "id": "resp-initial", + "output": [ + { + "type": "reasoning", + "id": "reasoning-1", + "summary": [], + "encrypted_content": "encrypted-1", + }, + { + "type": "function_call", + "id": "fc-1", + "call_id": "call-1", + "name": "memory_search", + "arguments": "{}", + }, + ], + "usage": {"input_tokens": 3, "output_tokens": 2}, + }, + request=request, + ) + if any( + isinstance(item, dict) + and item.get("type") == "reasoning" + and not item.get("encrypted_content") + for item in body.get("input", []) + if isinstance(body.get("input"), list) + ): + return httpx.Response( + 400, + json={ + "error": {"message": "Reasoning item is not reusable without encrypted_content"} + }, + request=request, + ) + if "previous_response_id" in body: + return httpx.Response( + 400, + json={ + "error": { + "message": "Unknown parameter: 'previous_response_id'.", + "type": "invalid_request_error", + "param": "previous_response_id", + "code": "unsupported_parameter", + } + }, + request=request, + ) + return httpx.Response( + 200, + json={ + "id": "resp-final", + "output": [{"type": "message", "id": "message-1"}], + "usage": {"input_tokens": 8, "output_tokens": 3}, + }, + request=request, + ) + + def _build_request(body: dict, headers: dict[str, str]) -> Request: payload = json.dumps(body).encode("utf-8") @@ -423,7 +550,7 @@ def test_handle_openai_responses_chatgpt_codex_timeout_fails_open(monkeypatch): assert body["store"] is False -def test_handle_openai_responses_api_auth_store_false_skips_memory_tools(monkeypatch): +def test_handle_openai_responses_api_auth_store_false_injects_stateless_memory_tools(monkeypatch): request = _build_request( {"model": "gpt-4o-mini", "input": "hello", "store": False}, {"Authorization": "Bearer sk-test", "x-headroom-user-id": "user-1"}, @@ -444,10 +571,113 @@ def test_handle_openai_responses_api_auth_store_false_skips_memory_tools(monkeyp _, url, _, body = handler.captured_request assert url == "https://api.openai.com/v1/responses" assert body["store"] is False - assert "tools" not in body + assert [tool["name"] for tool in body["tools"]] == ["memory_search"] + assert body["include"] == ["reasoning.encrypted_content"] assert memory_handler.compute_calls == 1 +@pytest.mark.parametrize("store", [pytest.param(None, id="omitted"), True, False]) +@pytest.mark.parametrize( + "include", + [ + pytest.param(None, id="omitted"), + pytest.param(["response.output_text.done"], id="missing-marker"), + pytest.param( + ["response.output_text.done", "reasoning.encrypted_content"], + id="existing-marker", + ), + pytest.param("not-a-list", id="non-list"), + ], +) +def test_openai_responses_memory_continuation_is_zdr_safe(store, include, monkeypatch): + body = { + "model": "gpt-5.4", + "previous_response_id": "resp-inherited", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}], + }, + { + "type": "reasoning", + "id": "prior-reasoning", + "summary": [], + "encrypted_content": "prior-encrypted", + }, + { + "type": "reasoning", + "id": "prior-unencrypted-reasoning", + "summary": [], + }, + ], + } + if include is not None: + body["include"] = include + if store is not None: + body["store"] = store + request = _build_request( + body, + {"Authorization": "Bearer sk-test", "x-headroom-user-id": "user-1"}, + ) + handler = _ZdrResponsesHandler() + + monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _DummyTokenizer()) + + response = anyio.run(handler.handle_openai_responses, request) + + assert response.status_code == 200 + assert len(handler.requests) == 2 + first_body, continuation_body = handler.requests + expected_include = ( + ["reasoning.encrypted_content"] + if include is None + else ( + include + if not isinstance(include, list) + else ( + include + if "reasoning.encrypted_content" in include + else [*include, "reasoning.encrypted_content"] + ) + ) + ) + assert first_body["include"] == expected_include + assert continuation_body["include"] == expected_include + assert first_body["previous_response_id"] == "resp-inherited" + assert ("store" in first_body) is (store is not None) + if store is not None: + assert first_body["store"] is store + else: + assert "store" not in first_body + assert continuation_body["input"] == [ + body["input"][0], + body["input"][1], + { + "type": "reasoning", + "id": "reasoning-1", + "summary": [], + "encrypted_content": "encrypted-1", + }, + { + "type": "function_call", + "id": "fc-1", + "call_id": "call-1", + "name": "memory_search", + "arguments": "{}", + }, + { + "type": "function_call_output", + "call_id": "call-1", + "output": '{"memories": []}', + }, + ] + assert "previous_response_id" not in continuation_body + assert ("store" in continuation_body) is (store is not None) + if store is not None: + assert continuation_body["store"] is store + + def test_handle_openai_responses_routes_api_key_auth_direct_to_openai(monkeypatch): request = _build_request( {"model": "gpt-4o-mini", "input": "hello"}, diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index dd60f5ed7..259c58838 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -70,6 +70,47 @@ class _DummyMetrics: self.codex_ws_frames.append(dict(kwargs)) +class _MemoryWsHandler: + def __init__(self) -> None: + self.config = SimpleNamespace( + inject_context=False, + inject_tools=True, + project_root_override="", + ) + self._backend = False + + def compute_memory_tool_definitions(self, provider: str) -> list[dict]: + assert provider == "openai" + return [ + { + "type": "function", + "function": { + "name": "memory_search", + "description": "Search memory.", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + async def _ensure_initialized(self) -> None: + self._backend = True + + async def _execute_memory_tool( + self, + name: str, + args: dict, + user_id: str, + provider: str, + ) -> str: + assert (name, args, user_id, provider) == ( + "memory_search", + {}, + user_id, + "openai", + ) + return '{"memories": []}' + + class _DummyOpenAIHandler(OpenAIHandlerMixin): OPENAI_API_URL = "https://api.openai.com" @@ -1739,3 +1780,466 @@ async def test_ws_recognized_client_with_real_path_is_not_restamped(): # the caller already self-identifies via its User-Agent. assert "x-client" not in {k.lower() for k in client_ws.headers} assert handler.ws_sessions.active_count() == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("store", [True, False]) +@pytest.mark.parametrize( + "include", + [ + pytest.param(None, id="omitted"), + pytest.param(["response.output_text.done"], id="missing-marker"), + pytest.param( + ["response.output_text.done", "reasoning.encrypted_content"], + id="existing-marker", + ), + pytest.param("not-a-list", id="non-list"), + ], +) +async def test_ws_memory_continuation_replays_history_without_previous_response_id(include, store): + function_call = { + "type": "function_call", + "id": "fc-1", + "call_id": "call-1", + "name": "memory_search", + "arguments": "{}", + } + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r-1"}}), + json.dumps({"type": "response.output_item.added", "item": function_call}), + json.dumps({"type": "response.output_item.done", "item": function_call}), + json.dumps({"type": "response.completed", "response": {"id": "r-1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + client_ws = _FakeWebSocket( + frames=[ + json.dumps( + { + "type": "response.create", + "response": { + "model": "gpt-5.4", + "input": "remember this", + "store": store, + }, + } + ) + ], + hold_after_initial=True, + ) + if include is not None: + client_ws._frames[0] = json.dumps( + { + "type": "response.create", + "response": { + "model": "gpt-5.4", + "input": "remember this", + "store": store, + "include": include, + }, + } + ) + client_ws.headers["x-headroom-user-id"] = "user-1" + handler = _DummyOpenAIHandler() + handler.memory_handler = _MemoryWsHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert len(upstream.sent) >= 2 + expected_include = ( + ["reasoning.encrypted_content"] + if include is None + else ( + include + if not isinstance(include, list) + else ( + include + if "reasoning.encrypted_content" in include + else [*include, "reasoning.encrypted_content"] + ) + ) + ) + assert json.loads(upstream.sent[0])["response"]["include"] == expected_include + continuation = json.loads(upstream.sent[1]) + assert "previous_response_id" not in continuation["response"] + assert continuation["response"]["model"] == "gpt-5.4" + assert continuation["response"]["store"] is store + assert continuation["response"]["include"] == expected_include + assert continuation["response"]["tools"] + assert continuation["response"]["instructions"] + assert continuation["response"]["input"] == [ + {"role": "user", "content": "remember this"}, + function_call, + { + "type": "function_call_output", + "call_id": "call-1", + "output": '{"memories": []}', + }, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "initial_frame", + [ + pytest.param("not-json", id="initial-non-json"), + pytest.param( + json.dumps({"type": "response.create", "response": []}), + id="initial-non-mapping-response", + ), + ], +) +async def test_ws_memory_frame_shape_guards_fail_open(initial_frame): + later_frames = [ + json.dumps( + { + "type": "response.create", + "response": {"model": "gpt-5.4", "input": []}, + } + ), + json.dumps({"type": "response.create", "response": "invalid"}), + ] + frames = [initial_frame, *later_frames] + upstream = _FakeUpstream([], hold_after_events=True) + fake_ws_mod = _make_fake_websockets_module(upstream) + client_ws = _FakeWebSocket(frames=frames, hold_after_initial=True) + client_ws.headers["x-headroom-user-id"] = "user-1" + handler = _DummyOpenAIHandler() + handler.memory_handler = _MemoryWsHandler() + + async def _trigger_disconnect() -> None: + await asyncio.sleep(0.05) + client_ws.trigger_disconnect() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + trigger_task = asyncio.create_task(_trigger_disconnect()) + try: + await asyncio.wait_for(handler.handle_openai_responses_ws(client_ws), timeout=2.0) + finally: + trigger_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await trigger_task + + assert upstream.sent[0] == initial_frame + forwarded_valid = json.loads(upstream.sent[1])["response"] + assert forwarded_valid["model"] == "gpt-5.4" + assert forwarded_valid["input"] == [] + assert forwarded_valid["tools"] + assert forwarded_valid["include"] == ["reasoning.encrypted_content"] + assert upstream.sent[2] == later_frames[1] + + +@pytest.mark.asyncio +async def test_ws_memory_enabled_non_memory_response_streams_completion(): + message_item = { + "type": "message", + "id": "message-1", + "role": "assistant", + "content": [{"type": "output_text", "text": "hello"}], + } + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r-1"}}), + json.dumps({"type": "response.output_item.added", "item": message_item}), + json.dumps({"type": "response.output_item.done", "item": message_item}), + json.dumps({"type": "response.completed", "response": {"id": "r-1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + client_ws = _FakeWebSocket(frames=[_first_frame()]) + client_ws.headers["x-headroom-user-id"] = "user-1" + handler = _DummyOpenAIHandler() + handler.memory_handler = _MemoryWsHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + forwarded_initial = json.loads(upstream.sent[0])["response"] + assert forwarded_initial["model"] == "gpt-5.4" + assert forwarded_initial["input"] == "hi" + assert forwarded_initial["tools"] + assert forwarded_initial["include"] == ["reasoning.encrypted_content"] + assert client_ws.sent_text == upstream_events + assert len(upstream.sent) == 1 + + +@pytest.mark.asyncio +async def test_ws_late_memory_call_after_streamed_message_passes_through(): + message_item = { + "type": "message", + "id": "message-1", + "role": "assistant", + "content": [{"type": "output_text", "text": "searching"}], + } + function_call = { + "type": "function_call", + "id": "fc-1", + "call_id": "call-1", + "name": "memory_search", + "arguments": "{}", + } + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r-1"}}), + json.dumps({"type": "response.output_item.added", "item": message_item}), + json.dumps({"type": "response.output_item.done", "item": message_item}), + json.dumps({"type": "response.output_item.added", "item": function_call}), + json.dumps({"type": "response.output_item.done", "item": function_call}), + json.dumps({"type": "response.completed", "response": {"id": "r-1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + client_ws = _FakeWebSocket(frames=[_first_frame()]) + client_ws.headers["x-headroom-user-id"] = "user-1" + handler = _DummyOpenAIHandler() + handler.memory_handler = _MemoryWsHandler() + executed: list[tuple[str, dict, str, str]] = [] + + async def _execute_memory_tool(name, args, user_id, provider): + executed.append((name, args, user_id, provider)) + return '{"memories": []}' + + handler.memory_handler._execute_memory_tool = _execute_memory_tool + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert client_ws.sent_text == upstream_events + assert len(upstream.sent) == 1 + assert executed == [] + + +@pytest.mark.asyncio +async def test_ws_memory_continuation_handles_invalid_item_arguments_and_unavailable_backend(): + function_call = { + "type": "function_call", + "id": "fc-1", + "call_id": "call-1", + "name": "memory_search", + "arguments": "{malformed", + } + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r-1"}}), + json.dumps({"type": "response.output_item.done", "item": "invalid"}), + json.dumps({"type": "response.output_item.added", "item": function_call}), + json.dumps({"type": "response.output_item.done", "item": function_call}), + json.dumps({"type": "response.completed", "response": {"id": "r-1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + client_ws = _FakeWebSocket(frames=[_first_frame()]) + client_ws.headers["x-headroom-user-id"] = "user-1" + handler = _DummyOpenAIHandler() + handler.memory_handler = _MemoryWsHandler() + + async def _leave_backend_unavailable(): + return None + + handler.memory_handler._ensure_initialized = _leave_backend_unavailable + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert len(upstream.sent) == 2 + continuation = json.loads(upstream.sent[1])["response"] + assert continuation["input"] == [ + {"role": "user", "content": "hi"}, + function_call, + { + "type": "function_call_output", + "call_id": "call-1", + "output": '{"error": "backend not ready"}', + }, + ] + assert continuation["input"][-1] == { + "type": "function_call_output", + "call_id": "call-1", + "output": '{"error": "backend not ready"}', + } + + +@pytest.mark.asyncio +async def test_ws_memory_continuation_normalizes_malformed_arguments(): + function_call = { + "type": "function_call", + "id": "fc-1", + "call_id": "call-1", + "name": "memory_search", + "arguments": "{malformed", + } + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r-1"}}), + json.dumps({"type": "response.output_item.done", "item": "invalid"}), + json.dumps({"type": "response.output_item.added", "item": function_call}), + json.dumps({"type": "response.output_item.done", "item": function_call}), + json.dumps({"type": "response.completed", "response": {"id": "r-1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + client_ws = _FakeWebSocket(frames=[_first_frame()]) + client_ws.headers["x-headroom-user-id"] = "user-1" + handler = _DummyOpenAIHandler() + handler.memory_handler = _MemoryWsHandler() + executed: list[tuple[str, dict, str, str]] = [] + + async def _execute_memory_tool(name, args, user_id, provider): + executed.append((name, args, user_id, provider)) + return '{"memories": []}' + + handler.memory_handler._execute_memory_tool = _execute_memory_tool + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert executed == [("memory_search", {}, "user-1", "openai")] + + +@pytest.mark.asyncio +async def test_ws_memory_tools_preserve_explicit_store_false_while_injecting(): + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r-1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r-1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + client_ws = _FakeWebSocket( + frames=[ + json.dumps( + { + "type": "response.create", + "response": { + "model": "gpt-5.4", + "input": "use stateless memory", + "store": False, + }, + } + ) + ], + hold_after_initial=True, + ) + client_ws.headers["x-headroom-user-id"] = "user-1" + handler = _DummyOpenAIHandler() + handler.memory_handler = _MemoryWsHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert len(upstream.sent) == 1 + initial = json.loads(upstream.sent[0])["response"] + assert initial["store"] is False + assert [tool["name"] for tool in initial["tools"]] == ["memory_search"] + assert initial["include"] == ["reasoning.encrypted_content"] + + +@pytest.mark.asyncio +async def test_ws_memory_continuation_continues_pre_stream_and_passes_late_call(): + function_call_one = { + "type": "function_call", + "id": "fc-1", + "call_id": "call-1", + "name": "memory_search", + "arguments": "{}", + } + function_call_two = { + "type": "function_call", + "id": "fc-2", + "call_id": "call-2", + "name": "memory_search", + "arguments": "{}", + } + reasoning_without_encryption = { + "type": "reasoning", + "id": "reasoning-1", + "summary": [], + } + reasoning_with_encryption = { + "type": "reasoning", + "id": "reasoning-2", + "summary": [], + "encrypted_content": "encrypted-2", + } + message_item = { + "type": "message", + "id": "message-2", + "role": "assistant", + "content": [{"type": "output_text", "text": "searching"}], + } + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r-1"}}), + json.dumps({"type": "response.output_item.added", "item": reasoning_without_encryption}), + json.dumps({"type": "response.output_item.done", "item": reasoning_without_encryption}), + json.dumps({"type": "response.output_item.added", "item": function_call_one}), + json.dumps({"type": "response.output_item.done", "item": function_call_one}), + json.dumps({"type": "response.completed", "response": {"id": "r-1"}}), + json.dumps({"type": "response.created", "response": {"id": "r-2"}}), + json.dumps({"type": "response.output_item.added", "item": reasoning_with_encryption}), + json.dumps({"type": "response.output_item.done", "item": reasoning_with_encryption}), + json.dumps({"type": "response.output_item.added", "item": message_item}), + json.dumps({"type": "response.output_item.done", "item": message_item}), + json.dumps({"type": "response.output_item.added", "item": function_call_two}), + json.dumps({"type": "response.output_item.done", "item": function_call_two}), + json.dumps({"type": "response.completed", "response": {"id": "r-2"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + client_ws = _FakeWebSocket( + frames=[ + json.dumps( + { + "type": "response.create", + "response": { + "model": "gpt-5.4", + "input": "remember this", + "client_metadata": { + "ws_request_header_x_openai_internal_codex_responses_lite": "true", + "keep": "yes", + }, + }, + } + ) + ], + hold_after_initial=True, + ) + client_ws.headers["x-headroom-user-id"] = "user-1" + handler = _DummyOpenAIHandler() + handler.memory_handler = _MemoryWsHandler() + executed: list[tuple[str, dict, str, str]] = [] + + async def _execute_memory_tool(name, args, user_id, provider): + executed.append((name, args, user_id, provider)) + return '{"memories": []}' + + handler.memory_handler._execute_memory_tool = _execute_memory_tool + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert len(upstream.sent) == 2 + first_continuation = json.loads(upstream.sent[1])["response"]["input"] + assert reasoning_without_encryption not in first_continuation + assert function_call_one in first_continuation + assert { + "type": "function_call_output", + "call_id": "call-1", + "output": '{"memories": []}', + } in first_continuation + assert json.loads(upstream.sent[1])["response"]["client_metadata"] == {"keep": "yes"} + + second_response = [json.loads(frame) for frame in client_ws.sent_text] + assert [event["type"] for event in second_response] == [ + "response.created", + "response.output_item.added", + "response.output_item.done", + "response.output_item.added", + "response.output_item.done", + "response.output_item.added", + "response.output_item.done", + "response.completed", + ] + assert second_response[0]["response"]["id"] == "r-2" + assert second_response[2]["item"] == reasoning_with_encryption + assert second_response[3]["item"] == message_item + assert second_response[4]["item"] == message_item + assert second_response[5]["item"] == function_call_two + assert second_response[6]["item"] == function_call_two + assert second_response[7]["response"]["id"] == "r-2" + assert executed == [("memory_search", {}, "user-1", "openai")] diff --git a/tests/test_openai_responses_context_compaction.py b/tests/test_openai_responses_context_compaction.py index e29b81dac..74dc99cba 100644 --- a/tests/test_openai_responses_context_compaction.py +++ b/tests/test_openai_responses_context_compaction.py @@ -7,7 +7,6 @@ from headroom.proxy.handlers.openai import ( OpenAIHandlerMixin, _compact_openai_responses_tools, _openai_responses_context_budget, - _responses_request_allows_memory_tool_continuation, ) from headroom.transforms.content_router import ( CompressionStrategy, @@ -438,30 +437,6 @@ def test_content_router_retries_kompress_when_structured_strategy_noops(monkeypa assert strategy_chain == ["smart_crusher", "kompress"] -def test_responses_memory_tools_skip_explicit_store_false() -> None: - """Regression: explicit store=false must block Responses memory-tool injection.""" - - payload = {"model": "gpt-5.5", "input": "remember this", "store": False} - - assert _responses_request_allows_memory_tool_continuation(payload) is False - assert payload["store"] is False - - -def test_responses_memory_tools_allow_default_and_stored_requests() -> None: - no_memory_payload = {"model": "gpt-5.5", "input": "plain", "store": False} - already_stored_payload = {"model": "gpt-5.5", "input": "plain", "store": True} - default_store_payload = {"model": "gpt-5.5", "input": "plain"} - - assert _responses_request_allows_memory_tool_continuation(no_memory_payload) is False - assert no_memory_payload["store"] is False - - assert _responses_request_allows_memory_tool_continuation(already_stored_payload) is True - assert already_stored_payload["store"] is True - - assert _responses_request_allows_memory_tool_continuation(default_store_payload) is True - assert "store" not in default_store_payload - - def test_responses_turn_hook_message_fold_is_applied_and_counted() -> None: """On the Responses path a turn hook may fold the `input` items (in place), not just tools. The fold must be written back to the outbound payload AND its diff --git a/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py b/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py index 3fcd3daf4..478eaceba 100644 --- a/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py +++ b/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py @@ -52,10 +52,6 @@ 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 @@ -355,7 +351,14 @@ def test_unresolved_ccr_only_streams_through_as_200() -> None: @pytest.mark.asyncio -async def test_buffered_ccr_emits_keepalive_before_delayed_upstream() -> None: +async def test_buffered_ccr_withholds_output_until_delayed_upstream_resolves() -> None: + """Nothing is sent — no status, no body — until the buffered result exists. + + The response used to commit ``200 text/event-stream`` on a 1s keepalive + timer, which made every later failure unreportable: the client saw a 200 + with no ``message_start`` and the real status was gone. See + ``test_buffered_ccr_preserves_late_failure_status_and_headers``. + """ config = _make_config() final_response = _message_response([{"type": "text", "text": "done"}]) started = asyncio.Event() @@ -367,8 +370,17 @@ async def test_buffered_ccr_emits_keepalive_before_delayed_upstream() -> None: "tools": [create_ccr_tool_definition("anthropic")], "messages": [{"role": "user", "content": "wait"}], } + request_delivered = False async def receive(): + # Mirror a real ASGI server: the body arrives once, then the channel + # stays open because the client is still connected. Returning instantly + # on every call spins `StreamingResponse.listen_for_disconnect`, which + # never yields, so the response body would never be scheduled. + nonlocal request_delivered + if request_delivered: + await asyncio.Event().wait() + request_delivered = True return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} scope = { @@ -400,24 +412,23 @@ async def test_buffered_ccr_emits_keepalive_before_delayed_upstream() -> None: 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() + # Longer than the deleted 1.0s keepalive deadline: an unresolved + # upstream must still have produced no ASGI message at all. + await asyncio.sleep(1.1) + assert events == [] 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) + start = next(event for event in events if event["type"] == "http.response.start") + assert start["status"] == 200 + bodies = b"".join(event["body"] for event in events if event["type"] == "http.response.body") + assert b"event: ping" not in bodies + assert b"done" in bodies @pytest.mark.asyncio @@ -480,7 +491,115 @@ async def test_buffered_ccr_preserves_early_failure_status_and_headers() -> None @pytest.mark.asyncio -async def test_buffered_ccr_late_failure_emits_sanitized_error_event() -> None: +async def test_buffered_ccr_preserves_late_failure_status_and_headers() -> None: + """The reported failure: a non-200 landing after the old keepalive deadline. + + Headroom had already committed ``200 text/event-stream`` by then, so the 429 + reached Claude Code as a 200 whose body carried no ``message_start`` — shown + as "API returned an empty or malformed response (HTTP 200) — check for a + proxy or gateway intercepting the request" — and ``retry-after`` was dropped, + so the client never backed off. + """ + 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": "fail late"}], + } + + 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_failure(*args, **kwargs): # noqa: ANN002, ANN003 + started.set() + await release.wait() + return httpx.Response( + 429, + headers={"retry-after": "7"}, + json={"error": {"message": "slow down"}}, + ) + + 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] = [] + + async def send(message): # noqa: ANN001 + events.append(message) + + response_task = asyncio.create_task(response(scope, receive, send)) + # Past the deleted 1.0s keepalive deadline before the upstream fails. + await asyncio.sleep(1.1) + assert events == [] + release.set() + await response_task + + start = next(event for event in events if event["type"] == "http.response.start") + assert start["status"] == 429 + assert dict(start["headers"])[b"retry-after"] == b"7" + assert b"slow down" in b"".join( + event["body"] for event in events if event["type"] == "http.response.body" + ) + + +def test_buffered_ccr_rejects_malformed_success_as_502() -> None: + """A non-SSE, non-JSON 200 is an upstream protocol error, not success.""" + config = _make_config() + with patch("headroom.proxy.server.AnyLLMBackend"): + app = create_app(config) + with TestClient(app) as client: + proxy = app.state.proxy + proxy._retry_request = AsyncMock( + return_value=httpx.Response( + 200, + content=b"gateway timeout", + headers={"content-type": "text/html"}, + ) + ) + response = client.post( + "/v1/messages", + headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"}, + json={ + "model": "claude-sonnet-4-6", + "max_tokens": 64, + "stream": True, + "tools": [create_ccr_tool_definition("anthropic")], + "messages": [{"role": "user", "content": "fail safely"}], + }, + ) + + assert response.status_code == 502 + assert response.json()["error"]["type"] == "upstream_protocol_error" + assert b"gateway timeout" not in response.content + + +@pytest.mark.asyncio +async def test_buffered_ccr_late_failure_returns_sanitized_json_error() -> None: + """A slow crash gets the same 502 the fast one does, not a downgraded 200.""" config = _make_config() started = asyncio.Event() release = asyncio.Event() @@ -533,24 +652,23 @@ async def test_buffered_ccr_late_failure_emits_sanitized_error_event() -> None: 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) + await asyncio.sleep(0) + assert events == [] 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] + start = next(event for event in events if event["type"] == "http.response.start") + assert start["status"] == 502 + bodies = b"".join(event["body"] for event in events if event["type"] == "http.response.body") + assert b"An error occurred while processing your request." in bodies + assert b"boom" not in bodies assert events[-1]["more_body"] is False assert any( record.levelno == logging.ERROR and "RuntimeError: boom" in record.getMessage() diff --git a/tests/test_proxy/test_openai_responses_ccr.py b/tests/test_proxy/test_openai_responses_ccr.py index 9bae90cd2..7e63496b0 100644 --- a/tests/test_proxy/test_openai_responses_ccr.py +++ b/tests/test_proxy/test_openai_responses_ccr.py @@ -266,7 +266,8 @@ def test_streaming_request_without_retrieve_tool_uses_normal_stream_path(): @pytest.mark.asyncio -async def test_buffered_responses_ccr_emits_keepalive_before_delayed_upstream(): +async def test_buffered_responses_ccr_withholds_output_until_upstream_resolves(): + """Nothing is sent — no status, no body — until the buffered result exists.""" app = _make_app() body = { "model": "gpt-5-codex", @@ -276,8 +277,17 @@ async def test_buffered_responses_ccr_emits_keepalive_before_delayed_upstream(): } started = asyncio.Event() release = asyncio.Event() + request_delivered = False async def receive(): + # Mirror a real ASGI server: the body arrives once, then the channel + # stays open because the client is still connected. Returning instantly + # on every call spins `StreamingResponse.listen_for_disconnect`, which + # never yields, so the response body would never be scheduled. + nonlocal request_delivered + if request_delivered: + await asyncio.Event().wait() + request_delivered = True return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} scope = { @@ -307,22 +317,23 @@ async def test_buffered_responses_ccr_emits_keepalive_before_delayed_upstream(): 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() + # Longer than the deleted 1.0s keepalive deadline: an unresolved upstream + # must still have produced no ASGI message at all. + await asyncio.sleep(1.1) + assert events == [] 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) + start = next(event for event in events if event["type"] == "http.response.start") + assert start["status"] == 200 + bodies = b"".join(event["body"] for event in events if event["type"] == "http.response.body") + assert b"event: ping" not in bodies + assert b"Resolved!" in bodies @pytest.mark.asyncio @@ -382,6 +393,35 @@ async def test_buffered_responses_ccr_preserves_early_failure_status_and_headers ) +def test_buffered_responses_ccr_rejects_malformed_success_as_502(): + """A malformed upstream 200 must not become a successful streamed turn.""" + app = _make_app() + with TestClient(app) as client: + server = app.state.proxy + server._retry_request = AsyncMock( + return_value=httpx.Response( + 200, + content=b"gateway timeout", + headers={"content-type": "text/html"}, + request=httpx.Request("POST", "https://api.openai.com/v1/responses"), + ) + ) + response = client.post( + "/v1/responses", + headers={"authorization": "Bearer sk-test"}, + json={ + "model": "gpt-5-codex", + "input": "fail safely", + "stream": True, + "tools": [_RETRIEVE_TOOL], + }, + ) + + assert response.status_code == 502 + assert response.json()["error"]["type"] == "upstream_protocol_error", response.text + assert b"gateway timeout" not in response.content + + @pytest.mark.asyncio async def test_buffered_responses_ccr_late_failure_emits_sanitized_error_event(): app = _make_app() @@ -431,24 +471,23 @@ async def test_buffered_responses_ccr_late_failure_emits_sanitized_error_event() 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) + await asyncio.sleep(0) + assert events == [] 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] + start = next(event for event in events if event["type"] == "http.response.start") + assert start["status"] == 502 + bodies = b"".join(event["body"] for event in events if event["type"] == "http.response.body") + assert b"An error occurred while processing your request." in bodies + assert b"boom" not in bodies assert events[-1]["more_body"] is False assert any( record.levelno == logging.ERROR and "RuntimeError: boom" in record.getMessage()