diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 0854a6495..f73cf3c6b 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -627,6 +627,39 @@ def _responses_request_allows_memory_tool_continuation(payload: dict[str, Any]) 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. + return not is_chatgpt_auth + + +def _ensure_chatgpt_responses_store_false( + payload: dict[str, Any], + *, + is_chatgpt_auth: bool, +) -> bool: + """Return True when ChatGPT auth requires and receives a store rewrite.""" + + if is_chatgpt_auth and payload.get("store") is not False: + payload["store"] = False + return True + return False + + def _responses_input_item_text_bytes(item: Any) -> int: if not isinstance(item, dict): return _json_byte_len(item) @@ -4145,6 +4178,12 @@ class OpenAIHandlerMixin: stripped_count=_pre_strip_count_resp, request_id=request_id, ) + headers, is_chatgpt_auth = _resolve_codex_routing_headers(headers) + if is_chatgpt_auth: + client = "codex" + if _ensure_chatgpt_responses_store_false(body, is_chatgpt_auth=is_chatgpt_auth): + logger.info(f"[{request_id}] Responses: forced store=false for ChatGPT auth") + responses_memory_tools_allowed = _allow_responses_memory_tools(is_chatgpt_auth) # PR-A6 (P5-50, preps P0-6): session-sticky `OpenAI-Beta` merge # for /v1/responses. Compute a session_id off the same store the @@ -4365,7 +4404,7 @@ class OpenAIHandlerMixin: memory_tool_defs_chat = ( self.memory_handler.compute_memory_tool_definitions("openai") - if self.memory_handler.config.inject_tools + if self.memory_handler.config.inject_tools and responses_memory_tools_allowed else [] ) memory_tool_defs_responses: list[dict[str, Any]] = [] @@ -4383,7 +4422,10 @@ class OpenAIHandlerMixin: else: memory_tool_defs_responses.append(t) - if _responses_request_allows_memory_tool_continuation(body): + if ( + responses_memory_tools_allowed + and _responses_request_allows_memory_tool_continuation(body) + ): resp_tools = body.get("tools") or [] resp_tools, mem_tools_injected = _apply_sticky_mem_tools_resp( provider="openai", @@ -4399,6 +4441,14 @@ class OpenAIHandlerMixin: 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", @@ -4420,10 +4470,6 @@ class OpenAIHandlerMixin: f"(backend '{self.anthropic_backend.name}' not used for Responses API)" ) - headers, is_chatgpt_auth = _resolve_codex_routing_headers(headers) - if is_chatgpt_auth: - client = "codex" - # Route to correct endpoint based on auth mode. # ChatGPT session auth (codex login) uses chatgpt.com, not api.openai.com. if is_chatgpt_auth: @@ -4818,6 +4864,7 @@ class OpenAIHandlerMixin: 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") @@ -5196,6 +5243,8 @@ class OpenAIHandlerMixin: for key, value in upstream_headers.items() if key.lower() != _CODEX_RESPONSES_LITE_HEADER } + ws_memory_tools_allowed = _allow_responses_memory_tools(is_chatgpt_auth) + _lower_headers = {k.lower(): v for k, v in upstream_headers.items()} # Build upstream WebSocket URL based on auth mode if is_chatgpt_auth: @@ -5584,6 +5633,16 @@ class OpenAIHandlerMixin: except json.JSONDecodeError: # Not JSON — pass through as-is pass + if isinstance(body, dict) and body: + ws_response_body_for_store = ( + body["response"] if isinstance(body.get("response"), dict) else body + ) + if _ensure_chatgpt_responses_store_false( + ws_response_body_for_store, + is_chatgpt_auth=is_chatgpt_auth, + ): + first_msg_raw = json.dumps(body) + logger.info(f"[{request_id}] WS Responses: forced store=false for ChatGPT auth") ws_input_tokens_total = 0 ws_output_tokens_total = 0 ws_cache_read_tokens_total = 0 @@ -5854,7 +5913,7 @@ class OpenAIHandlerMixin: ws_mem_defs_chat = ( self.memory_handler.compute_memory_tool_definitions("openai") - if self.memory_handler.config.inject_tools + if self.memory_handler.config.inject_tools and ws_memory_tools_allowed else [] ) ws_mem_defs_responses: list[dict[str, Any]] = [] @@ -5875,11 +5934,13 @@ class OpenAIHandlerMixin: ws_tools = ws_response_body.get("tools") or [] ws_tools, mem_injected = _apply_sticky_mem_tools_ws( provider="openai", - session_id=session_id, + session_id=session_id if ws_memory_tools_allowed else None, request_id=request_id, existing_tools=ws_tools, memory_tools_to_inject=ws_mem_defs_responses, - inject_this_turn=bool(self.memory_handler.config.inject_tools), + inject_this_turn=bool( + self.memory_handler.config.inject_tools and ws_memory_tools_allowed + ), ) if mem_injected: ws_response_body["tools"] = ws_tools @@ -6222,7 +6283,7 @@ class OpenAIHandlerMixin: when its `type` is `response.create`. Other event types (response.cancel, session.update, etc.) pass through unchanged. Errors are - warned and the original frame is returned — + warned and the safest frame available is returned — fail loud in logs, fail safe on the wire. Updates outer-scope ``tokens_saved``, ``transforms_applied``, and @@ -6232,20 +6293,6 @@ class OpenAIHandlerMixin: """ nonlocal tokens_saved, transforms_applied, attempted_input_tokens_total nonlocal ws_frames_compressed - if _ws_bypass: - _log_ws_passthrough( - "bypass_header", - frame_index=frame_index, - raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), - ) - return raw_msg, False, "bypass_header" - if not self.config.optimize: - _log_ws_passthrough( - "optimize_disabled", - frame_index=frame_index, - raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), - ) - return raw_msg, False, "optimize_disabled" _preflight_started = time.perf_counter() try: parsed_frame = json.loads(raw_msg) @@ -6281,6 +6328,44 @@ class OpenAIHandlerMixin: frame_type="response.create", ) return raw_msg, False, "invalid_inner_payload" + store_forced = _ensure_chatgpt_responses_store_false( + inner_payload, + is_chatgpt_auth=is_chatgpt_auth, + ) + raw_after_store = raw_msg + if store_forced: + raw_after_store = json.dumps(parsed_frame) + logger.info( + "[%s] WS Responses: forced store=false for ChatGPT auth frame=%d", + request_id, + frame_index, + ) + if _ws_bypass: + _log_ws_passthrough( + "bypass_header", + frame_index=frame_index, + raw_bytes=len(raw_after_store.encode("utf-8", errors="replace")), + frame_type="response.create", + model=str(inner_payload.get("model") or "unknown"), + ) + return ( + raw_after_store, + store_forced, + "chatgpt_store_false" if store_forced else "bypass_header", + ) + if not self.config.optimize: + _log_ws_passthrough( + "optimize_disabled", + frame_index=frame_index, + raw_bytes=len(raw_after_store.encode("utf-8", errors="replace")), + frame_type="response.create", + model=str(inner_payload.get("model") or "unknown"), + ) + return ( + raw_after_store, + store_forced, + "chatgpt_store_false" if store_forced else "optimize_disabled", + ) frame_compression_elapsed_ms = 0.0 try: model_for_frame = inner_payload.get("model") or "" @@ -6396,11 +6481,10 @@ class OpenAIHandlerMixin: _log_ws_passthrough( "compression_exception", frame_index=frame_index, - raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), + raw_bytes=len(raw_after_store.encode("utf-8", errors="replace")), frame_type="response.create", model=str(inner_payload.get("model") or "unknown"), ) - return raw_msg, False, "compression_exception" # Record transform labels even when the frame bytes are # unchanged: control-arm output-shaper labels # (output_shaper:control:*) must reach the outcome @@ -6408,25 +6492,44 @@ class OpenAIHandlerMixin: for t in frame_transforms: if t not in transforms_applied: transforms_applied.append(t) + return ( + raw_after_store, + store_forced, + "chatgpt_store_false" if store_forced else "compression_exception", + ) if not modified: reason = frame_reason or "no_compression" _log_ws_passthrough( reason, frame_index=frame_index, - raw_bytes=bytes_before, + raw_bytes=len(raw_after_store.encode("utf-8", errors="replace")), frame_type="response.create", model=str(inner_payload.get("model") or "unknown"), ) - return raw_msg, False, reason + return ( + raw_after_store, + store_forced, + "chatgpt_store_false" if store_forced else reason, + ) if not isinstance(new_inner, dict): _log_ws_passthrough( "compressed_payload_not_dict", frame_index=frame_index, - raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), + raw_bytes=len(raw_after_store.encode("utf-8", errors="replace")), frame_type="response.create", model=str(inner_payload.get("model") or "unknown"), ) - return raw_msg, False, "compressed_payload_not_dict" + return ( + raw_after_store, + store_forced, + "chatgpt_store_false" + if store_forced + else "compressed_payload_not_dict", + ) + _ensure_chatgpt_responses_store_false( + new_inner, + is_chatgpt_auth=is_chatgpt_auth, + ) if wrapped_frame: _rewrite_started = time.perf_counter() parsed_frame["response"] = new_inner @@ -6637,7 +6740,9 @@ class OpenAIHandlerMixin: nonlocal ws_upstream_frames_total, ws_last_upstream_frame_type nonlocal ws_ttfb_ms - memory_enabled = bool(self.memory_handler and memory_user_id) + memory_enabled = bool( + self.memory_handler and memory_user_id and ws_memory_tools_allowed + ) # Per-response state (reset after each response.completed) event_buffer: list[str] = [] @@ -7464,7 +7569,8 @@ class OpenAIHandlerMixin: Codex work immediately instead of exhausting its WS retry budget. """ # Route to correct endpoint based on auth mode - if has_chatgpt_account_header(upstream_headers): + is_chatgpt_fallback = has_chatgpt_account_header(upstream_headers) + if is_chatgpt_fallback: http_url = codex_responses_http_url() else: http_url = build_copilot_upstream_url(self.OPENAI_API_URL, "/v1/responses") @@ -7497,6 +7603,12 @@ class OpenAIHandlerMixin: # Ensure streaming is enabled so we get SSE events http_body["stream"] = True + mutation_reasons = ["ws_http_fallback_resynthesized"] + if _ensure_chatgpt_responses_store_false( + http_body, + is_chatgpt_auth=is_chatgpt_fallback, + ): + mutation_reasons.append("chatgpt_store_false") # Build HTTP headers from the upstream headers (already stripped of WS # hop-by-hop headers by the caller). @@ -7521,7 +7633,7 @@ class OpenAIHandlerMixin: path=http_url, body_bytes_count=len(outbound_bytes), body_mutated=True, - mutation_reasons=["ws_http_fallback_resynthesized"], + mutation_reasons=mutation_reasons, request_id=request_id, source=outbound_source, ) diff --git a/tests/test_openai_codex_routing.py b/tests/test_openai_codex_routing.py index 9dc261548..47839983d 100644 --- a/tests/test_openai_codex_routing.py +++ b/tests/test_openai_codex_routing.py @@ -177,6 +177,7 @@ class _DummyOpenAIHandler(OpenAIHandlerMixin): self.anthropic_backend = None self.cost_tracker = None self.memory_handler = None + self.traffic_learner = None # PR-A6 wires session-sticky `OpenAI-Beta` merging into the # responses HTTP handler — it reads `compute_session_id` to key # the SessionBetaTracker. The routing tests don't exercise the @@ -238,6 +239,33 @@ class _DummyOpenAIHandler(OpenAIHandlerMixin): ) +class _MemoryToolsOnlyHandler: + def __init__(self) -> None: + self.config = SimpleNamespace( + inject_context=False, + inject_tools=True, + project_root_override="", + ) + self.compute_calls = 0 + + def compute_memory_tool_definitions(self, provider: str) -> list[dict]: + self.compute_calls += 1 + assert provider == "openai" + return [ + { + "type": "function", + "function": { + "name": "memory_search", + "description": "Search memory.", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def has_memory_tool_calls(self, response: dict, provider: str) -> bool: + return False + + def _build_request(body: dict, headers: dict[str, str]) -> Request: payload = json.dumps(body).encode("utf-8") @@ -286,6 +314,7 @@ def test_handle_openai_responses_routes_chatgpt_auth_to_backend_api(monkeypatch) assert url == "https://chatgpt.com/backend-api/codex/responses" assert headers["ChatGPT-Account-ID"] == "acct-from-jwt" assert body["input"] == "hello" + assert body["store"] is False assert response.status_code == 200 @@ -322,6 +351,38 @@ def test_handle_openai_responses_strips_codex_lite_header_upstream(monkeypatch): assert lowered.get("x-openai-debug") == "keep-me" +def test_handle_openai_responses_chatgpt_auth_skips_memory_tools(monkeypatch): + token = _jwt( + { + "https://api.openai.com/auth": { + "chatgpt_account_id": "acct-from-jwt", + } + } + ) + request = _build_request( + {"model": "gpt-5.4", "input": "hello", "store": True}, + {"Authorization": f"Bearer {token}", "x-headroom-user-id": "user-1"}, + ) + handler = _DummyOpenAIHandler() + memory_handler = _MemoryToolsOnlyHandler() + handler.memory_handler = memory_handler + handler.session_tracker_store = SimpleNamespace( + compute_session_id=lambda *a, **k: "sess-chatgpt-no-memory-tools", + ) + + monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _DummyTokenizer()) + + response = anyio.run(handler.handle_openai_responses, request) + + assert response.status_code == 200 + assert handler.captured_request is not None + _, url, _, body = handler.captured_request + assert url == "https://chatgpt.com/backend-api/codex/responses" + assert body["store"] is False + assert "tools" not in body + assert memory_handler.compute_calls == 0 + + def test_handle_openai_responses_chatgpt_codex_timeout_fails_open(monkeypatch): token = _jwt( { @@ -351,6 +412,32 @@ def test_handle_openai_responses_chatgpt_codex_timeout_fails_open(monkeypatch): assert method == "POST" assert url == "https://chatgpt.com/backend-api/codex/responses" assert body["input"] == "large context" + assert body["store"] is False + + +def test_handle_openai_responses_api_auth_store_false_skips_memory_tools(monkeypatch): + request = _build_request( + {"model": "gpt-4o-mini", "input": "hello", "store": False}, + {"Authorization": "Bearer sk-test", "x-headroom-user-id": "user-1"}, + ) + handler = _DummyOpenAIHandler() + memory_handler = _MemoryToolsOnlyHandler() + handler.memory_handler = memory_handler + handler.session_tracker_store = SimpleNamespace( + compute_session_id=lambda *a, **k: "sess-api-memory-tools", + ) + + monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _DummyTokenizer()) + + response = anyio.run(handler.handle_openai_responses, request) + + assert response.status_code == 200 + assert handler.captured_request is not None + _, url, _, body = handler.captured_request + assert url == "https://api.openai.com/v1/responses" + assert body["store"] is False + assert "tools" not in body + assert memory_handler.compute_calls == 1 def test_handle_openai_responses_routes_api_key_auth_direct_to_openai(monkeypatch): diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index d08c9fb38..40d54680a 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -1067,8 +1067,10 @@ async def test_ws_connect_failure_falls_back_to_http(): # Fallback ran with the first frame. assert len(fallback_calls) == 1 _body, _first_raw = fallback_calls[0] - assert _first_raw == first - assert _body == json.loads(first) + expected = json.loads(first) + expected["response"]["store"] = False + assert json.loads(_first_raw) == expected + assert _body == expected # Clean teardown. assert handler.ws_sessions.active_count() == 0 diff --git a/tests/test_openai_codex_ws_timings.py b/tests/test_openai_codex_ws_timings.py index 242eea566..d38458795 100644 --- a/tests/test_openai_codex_ws_timings.py +++ b/tests/test_openai_codex_ws_timings.py @@ -46,11 +46,36 @@ class _DummyOpenAIHandler(OpenAIHandlerMixin): self.anthropic_backend = None self.cost_tracker = None self.memory_handler = None + self.traffic_learner = None async def _next_request_id(self) -> str: return "req-ws-test" +class _MemoryToolsOnlyHandler: + def __init__(self) -> None: + self.config = SimpleNamespace( + inject_context=False, + inject_tools=True, + project_root_override="", + ) + self.compute_calls = 0 + + def compute_memory_tool_definitions(self, provider: str) -> list[dict]: + self.compute_calls += 1 + assert provider == "openai" + return [ + { + "type": "function", + "function": { + "name": "memory_search", + "description": "Search memory.", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + class _FakeWebSocket: """Minimal async WebSocket stub that delivers a scripted frame list.""" @@ -226,6 +251,44 @@ def test_codex_ws_happy_path_emits_all_stage_timings(stage_log_capture): assert "total_session" in emitted +def test_codex_ws_chatgpt_auth_skips_memory_tools(stage_log_capture): + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "resp_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "resp_1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + + first_frame = json.dumps( + { + "type": "response.create", + "response": {"model": "gpt-5.4", "input": "hello", "store": True}, + } + ) + client_ws = _FakeWebSocket( + frames=[first_frame], + headers={ + "authorization": "Bearer chatgpt-session-token", + "chatgpt-account-id": "acct_123", + "x-headroom-user-id": "user-1", + }, + ) + handler = _DummyOpenAIHandler() + memory_handler = _MemoryToolsOnlyHandler() + handler.memory_handler = memory_handler + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + anyio.run(handler.handle_openai_responses_ws, client_ws) + + assert len(upstream.sent) == 1 + sent = json.loads(upstream.sent[0]) + response_body = sent["response"] + assert response_body["store"] is False + assert "tools" not in response_body + assert "## Memory" not in response_body.get("instructions", "") + assert memory_handler.compute_calls == 0 + + def test_codex_ws_upstream_connect_failure_still_logs_timings(stage_log_capture): """A session that never connects upstream still logs a timing line with ``upstream_first_event`` absent (null).""" diff --git a/tests/test_ws_http_fallback.py b/tests/test_ws_http_fallback.py index 69da0d654..459246fb3 100644 --- a/tests/test_ws_http_fallback.py +++ b/tests/test_ws_http_fallback.py @@ -266,6 +266,30 @@ class TestWsHttpFallback: assert "chatgpt.com" in captured_url["url"] assert "api.openai.com" not in captured_url["url"] + def test_fallback_chatgpt_auth_forces_store_false(self): + """ChatGPT Responses backend requires explicit store=false.""" + handler = _make_handler() + ws = FakeWebSocket() + captured_kwargs: dict = {} + + class CapturingClient: + def stream(self, method, url, **kwargs): + captured_kwargs.update(kwargs) + return FakeStreamResponse(200, ["data: [DONE]\n\n"]) + + handler.http_client = CapturingClient() + + body = {"model": "gpt-5.4", "input": "test", "store": True} + headers = { + "Authorization": "Bearer chatgpt-session-token", + "ChatGPT-Account-ID": "acct_abc123", + } + asyncio.run(handler._ws_http_fallback(ws, body, json.dumps(body), headers, "req_store")) + + posted = json.loads(captured_kwargs["content"]) + assert posted["store"] is False + assert posted["stream"] is True + def test_fallback_routes_api_key_to_openai(self): """API key auth should route to api.openai.com.""" handler = _make_handler()