diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 1f40e8e2f..ef839890e 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -5331,39 +5331,37 @@ class OpenAIHandlerMixin: ), ) - # --- Memory: inject context, tools, and instructions --- - # Gated on MemoryDecision — uniform bypass-respect across - # all five sites. WS sets memory_user_id only on the inject - # path (matches pre-PR behaviour); MemoryDecision is the - # canonical gate. memory_user_id: str | None = None memory_request_ctx = None - if self.memory_handler and body: - _ws_memory_user_id_candidate = ws_headers.get( - "x-headroom-user-id", - os.environ.get("USER", os.environ.get("USERNAME", "default")), - ) - else: - _ws_memory_user_id_candidate = None from headroom.proxy.helpers import get_memory_injection_mode from headroom.proxy.memory_decision import MemoryDecision from headroom.proxy.memory_query import MemoryQuery - ws_memory_decision = MemoryDecision.decide( - headers=ws_headers, - memory_handler=self.memory_handler if body else None, - memory_user_id=_ws_memory_user_id_candidate, - mode_name=get_memory_injection_mode(), - ) - # ws_tags was extracted at handler entry (L3028); applying - # the memory skip reason here so per-turn RequestOutcomes - # carry it for dashboard slicing. - ws_memory_decision.apply_to_tags(ws_tags) - if ws_memory_decision.inject: - memory_user_id = _ws_memory_user_id_candidate + async def _prepare_memory_frame(frame_body: dict[str, Any], frame_raw: str) -> str: + nonlocal memory_user_id, memory_request_ctx + + memory_user_id_candidate = ( + ws_headers.get( + "x-headroom-user-id", + os.environ.get("USER", os.environ.get("USERNAME", "default")), + ) + if self.memory_handler + else None + ) + memory_decision = MemoryDecision.decide( + headers=ws_headers, + memory_handler=self.memory_handler, + memory_user_id=memory_user_id_candidate, + mode_name=get_memory_injection_mode(), + ) + memory_decision.apply_to_tags(ws_tags) + if not memory_decision.inject: + return frame_raw + + memory_user_id = memory_user_id_candidate try: # Unwrap response.create envelope to access the response body - ws_response_body = body.get("response", body) + ws_response_body = frame_body.get("response", frame_body) # Per-project memory routing (GH #462). For WS, # ``ws_response_body`` carries ``instructions`` — @@ -5525,19 +5523,20 @@ class OpenAIHandlerMixin: ) # Write back into envelope if it was wrapped - if "response" in body and isinstance(body["response"], dict): - body["response"] = ws_response_body + if "response" in frame_body and isinstance(frame_body["response"], dict): + frame_body["response"] = ws_response_body else: - body = ws_response_body + frame_body = ws_response_body - first_msg_raw = json.dumps(body) + return json.dumps(frame_body) except Exception as e: logger.warning(f"[{request_id}] WS Memory injection failed: {e}") - elif self.memory_handler and body and _ws_bypass: - logger.info( - "[%s] WS memory passthrough reason=bypass_header", - request_id, - ) + return frame_raw + + if isinstance(body, dict) and ( + 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) # Hot-fix follow-up to PR #406 — inline Rust compression on the # WS first frame before forwarding upstream. PR #406 enabled @@ -6128,6 +6127,7 @@ class OpenAIHandlerMixin: and _inbound_frame_body.get("type") == "response.create" ): ws_response_create_frames += 1 + msg = await _prepare_memory_frame(_inbound_frame_body, msg) ( msg, _frame_modified, diff --git a/tests/fixtures/issues/headroom_issue_2059.json b/tests/fixtures/issues/headroom_issue_2059.json new file mode 100644 index 000000000..7db50dfbf --- /dev/null +++ b/tests/fixtures/issues/headroom_issue_2059.json @@ -0,0 +1 @@ +{"author":"superdiaodiao","body":"## Description\n\nFor a long-lived Codex WebSocket session, Headroom performs memory context lookup and tool injection only for the first client frame. Later `response.create` frames are compressed and output-shaped, but they do not run the memory decision/query/injection path again.\n\nThis means a user prompt sent on a later turn cannot retrieve newly relevant memories automatically, even though memory is enabled and the same WebSocket remains open.\n\n## To Reproduce\n\n1. Start Headroom with OpenAI/Codex memory enabled, including context injection.\n2. Connect Codex through `/v1/responses` WebSocket.\n3. Send a first `response.create` frame that establishes the session.\n4. Save a durable memory, or choose an existing memory relevant only to a later prompt.\n5. Send a second `response.create` frame on the same WebSocket whose user input should match that memory.\n6. Observe that the second frame is forwarded without a new memory lookup or context injection.\n\nA static source check shows the same lifecycle gap: the memory block in `handle_openai_responses_websocket` runs before the relay loop, while `_client_to_upstream` applies `_maybe_compress_response_create_frame` and output shaping to subsequent frames but does not rerun memory lookup/injection.\n\n## Expected Behavior\n\nEvery new `response.create` turn should independently:\n\n- resolve project/user memory scope;\n- build a query from that turn's current user input;\n- search and inject relevant memory;\n- preserve session-sticky memory tool definitions without duplicating them.\n\n## Actual Behavior\n\nOnly the first frame gets the memory pipeline. Later turns on the same WebSocket miss automatic retrieval.\n\n## Code Sample\n\nConceptual frame sequence:\n\n```json\n{\"type\":\"response.create\",\"response\":{\"input\":\"initial turn\"}}\n{\"type\":\"response.create\",\"response\":{\"input\":\"later turn requiring stored preference\"}}\n```\n\nThe second frame reaches the compression/shaping path but not the memory lookup path.\n\n## Error Output\n\nNo exception is emitted. This is a silent behavior gap; logs show no memory lookup/injection event for later `response.create` frames.\n\n## Environment\n\n- **Headroom version**: 0.31.0 and current main at `868b88bc6400c98f11134dbbe3cb03d1ecff7e1d`\n- **Python version**: 3.13\n- **OS**: macOS\n- **LLM Provider**: OpenAI Responses / Codex subscription WebSocket\n\n## Additional Context\n\nThis appears independent of project DB routing fixes such as GH #462/#1147. The correct DB may be selected, but the lookup is never invoked for the later turn.\n\nA regression test could open one WS connection, send two `response.create` frames with distinct inputs, and assert that memory query construction/injection runs once per frame while tool definitions remain deduplicated.\n","labels":["bug"],"number":2059,"state":"OPEN","title":"[BUG] Codex WebSocket memory lookup only runs on the first response.create frame","updatedAt":"2026-07-12T13:32:05Z","url":"https://github.com/headroomlabs-ai/headroom/issues/2059"} diff --git a/tests/test_codex_ws_per_frame_memory.py b/tests/test_codex_ws_per_frame_memory.py new file mode 100644 index 000000000..278b00ec3 --- /dev/null +++ b/tests/test_codex_ws_per_frame_memory.py @@ -0,0 +1,411 @@ +from __future__ import annotations + +import asyncio +import json +import re +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from tests.test_openai_codex_ws_lifecycle import ( + _DummyOpenAIHandler, + _FakeUpstream, + _FakeWebSocket, + _make_fake_websockets_module, +) + + +class _MemoryHandler: + def __init__(self) -> None: + self.config = SimpleNamespace( + inject_context=True, + inject_tools=True, + project_root_override="", + ) + self.queries: list[str] = [] + + async def search_and_format_context(self, _user_id, messages, **_kwargs): + current_turn = messages[-1]["content"] if messages else "" + self.queries.append(current_turn) + return f"current memory: {current_turn}" + + def compute_memory_tool_definitions(self, _provider): + return [ + { + "type": "function", + "function": { + "name": "memory_search", + "description": "search", + "parameters": {"type": "object"}, + }, + }, + { + "type": "function", + "function": { + "name": "memory_save", + "description": "save", + "parameters": {"type": "object"}, + }, + }, + ] + + +def _expected_memory_response_tools() -> list[dict[str, object]]: + expected: list[dict[str, object]] = [] + for tool in _MemoryHandler().compute_memory_tool_definitions("openai"): + function = tool["function"] + expected.append( + { + "type": "function", + "name": function["name"], + "description": function["description"], + "parameters": function["parameters"], + } + ) + return expected + + +def _turn(text: str) -> str: + return json.dumps({"type": "response.create", "response": {"input": text}}) + + +def _direct_turn(text: str) -> str: + return json.dumps({"input": text}) + + +def _issue_2059_artifact_path() -> Path: + return Path(__file__).resolve().parent / "fixtures" / "issues" / "headroom_issue_2059.json" + + +def _issue_2059_turns() -> tuple[str, str]: + issue_path = _issue_2059_artifact_path() + issue = json.loads(issue_path.read_text(encoding="utf-8")) + match = re.search(r"```json\s*(.*?)```", issue["body"], re.DOTALL) + assert match is not None, "issue 2059 artifact must contain a JSON code sample" + frames = [line.strip() for line in match.group(1).splitlines() if line.strip()] + assert len(frames) == 2, "issue 2059 artifact must contain exactly two frames" + return frames[0], frames[1] + + +def _issue_2059_inputs() -> tuple[str, str]: + first, later = _issue_2059_turns() + return ( + json.loads(first)["response"]["input"], + json.loads(later)["response"]["input"], + ) + + +def _list_turn(text: str, *, instructions: str) -> str: + return json.dumps( + { + "type": "response.create", + "response": { + "instructions": instructions, + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": text}], + } + ], + }, + } + ) + + +class _FlakyMemoryHandler(_MemoryHandler): + def __init__(self, *, fail_on: set[str]) -> None: + super().__init__() + self.fail_on = set(fail_on) + + async def search_and_format_context(self, _user_id, messages, **_kwargs): + current_turn = messages[-1]["content"] if messages else "" + self.queries.append(current_turn) + if current_turn in self.fail_on: + raise RuntimeError(f"memory failed for {current_turn}") + return f"current memory: {current_turn}" + + +class _ToolFailingMemoryHandler(_MemoryHandler): + def __init__(self) -> None: + super().__init__() + self._fail_next_tools = True + + def compute_memory_tool_definitions(self, _provider): + if self._fail_next_tools: + self._fail_next_tools = False + raise RuntimeError("memory tool preparation failed") + return super().compute_memory_tool_definitions(_provider) + + +@pytest.mark.asyncio +async def test_memory_lookup_runs_for_each_issue_artifact_frame_and_preserves_non_create_frames(): + upstream = _FakeUpstream( + [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + ) + first_turn, later_turn = _issue_2059_turns() + first_input, later_input = _issue_2059_inputs() + client_frames = [ + first_turn, + json.dumps({"type": "response.cancel"}), + later_turn, + ] + client_ws = _FakeWebSocket(frames=client_frames) + handler = _DummyOpenAIHandler() + memory = _MemoryHandler() + handler.memory_handler = memory + + with patch.dict(sys.modules, {"websockets": _make_fake_websockets_module(upstream)}): + await handler.handle_openai_responses_ws(client_ws) + + assert memory.queries == [first_input, later_input] + assert upstream.sent[1] == client_frames[1] + forwarded_turns = [ + json.loads(frame) for frame in upstream.sent if "response" in json.loads(frame) + ] + assert f"current memory: {first_input}" in forwarded_turns[0]["response"]["input"] + assert f"current memory: {later_input}" in forwarded_turns[1]["response"]["input"] + expected_tools = _expected_memory_response_tools() + for frame in forwarded_turns: + assert frame["response"]["tools"] == expected_tools + assert forwarded_turns[0]["response"]["tools"] == forwarded_turns[1]["response"]["tools"] + + +@pytest.mark.asyncio +async def test_memory_lookup_skips_input_bearing_non_create_first_frame(): + upstream = _FakeUpstream( + [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + ) + _first_input, later_input = _issue_2059_inputs() + cancel_frame = json.dumps( + { + "type": "response.cancel", + "response_id": "r_1", + "input": "must not query", + } + ) + later_turn = _issue_2059_turns()[1] + client_ws = _FakeWebSocket(frames=[cancel_frame, later_turn]) + handler = _DummyOpenAIHandler() + memory = _MemoryHandler() + handler.memory_handler = memory + + with patch.dict(sys.modules, {"websockets": _make_fake_websockets_module(upstream)}): + await handler.handle_openai_responses_ws(client_ws) + + assert memory.queries == [later_input] + assert upstream.sent[0] == cancel_frame + forwarded_later = json.loads(upstream.sent[1]) + assert f"current memory: {later_input}" in forwarded_later["response"]["input"] + + +@pytest.mark.asyncio +async def test_memory_lookup_skips_bypassed_frames(): + upstream = _FakeUpstream( + [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + ) + first, later = _issue_2059_turns() + client_ws = _FakeWebSocket( + frames=[first, later], + headers={"authorization": "Bearer test", "x-headroom-bypass": "true"}, + ) + handler = _DummyOpenAIHandler() + memory = _MemoryHandler() + handler.memory_handler = memory + + with patch.dict(sys.modules, {"websockets": _make_fake_websockets_module(upstream)}): + await handler.handle_openai_responses_ws(client_ws) + + assert memory.queries == [] + assert upstream.sent == [first, later] + + +@pytest.mark.asyncio +async def test_memory_lookup_keeps_legacy_direct_first_frame(): + upstream = _FakeUpstream( + [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + ) + first_input, later_input = _issue_2059_inputs() + first = _direct_turn(first_input) + later = _issue_2059_turns()[1] + client_ws = _FakeWebSocket(frames=[first, later]) + handler = _DummyOpenAIHandler() + memory = _MemoryHandler() + handler.memory_handler = memory + + with patch.dict(sys.modules, {"websockets": _make_fake_websockets_module(upstream)}): + await handler.handle_openai_responses_ws(client_ws) + + assert memory.queries == [first_input, later_input] + forwarded_first = json.loads(upstream.sent[0]) + forwarded_later = json.loads(upstream.sent[1]) + assert f"current memory: {first_input}" in forwarded_first["input"] + assert f"current memory: {later_input}" in forwarded_later["response"]["input"] + + +@pytest.mark.asyncio +async def test_memory_lookup_skips_disabled_memory(monkeypatch): + monkeypatch.setenv("HEADROOM_MEMORY_INJECTION_MODE", "disabled") + upstream = _FakeUpstream( + [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + ) + first, later = _issue_2059_turns() + client_ws = _FakeWebSocket(frames=[first, later]) + handler = _DummyOpenAIHandler() + memory = _MemoryHandler() + handler.memory_handler = memory + + with patch.dict(sys.modules, {"websockets": _make_fake_websockets_module(upstream)}): + await handler.handle_openai_responses_ws(client_ws) + + assert memory.queries == [] + assert upstream.sent == [first, later] + + +@pytest.mark.asyncio +async def test_memory_lookup_fails_open_and_recovers_on_later_frame(): + first, later = _issue_2059_turns() + first_input, later_input = _issue_2059_inputs() + upstream = _FakeUpstream([], hold_after_events=True) + client_ws = _FakeWebSocket(frames=[first, later], hold_after_initial=True) + handler = _DummyOpenAIHandler() + memory = _FlakyMemoryHandler(fail_on={first_input}) + handler.memory_handler = memory + + async def _trigger() -> None: + await asyncio.sleep(0.05) + client_ws.trigger_disconnect() + + with patch.dict(sys.modules, {"websockets": _make_fake_websockets_module(upstream)}): + trigger_task = asyncio.create_task(_trigger()) + try: + await handler.handle_openai_responses_ws(client_ws) + finally: + trigger_task.cancel() + try: + await trigger_task + except asyncio.CancelledError: + pass + + assert memory.queries == [first_input, later_input] + assert upstream.sent[0] == first + assert f"current memory: {later_input}" in json.loads(upstream.sent[1])["response"]["input"] + + +@pytest.mark.asyncio +async def test_memory_lookup_fails_open_when_tool_preparation_raises(): + first, later = _issue_2059_turns() + first_input, later_input = _issue_2059_inputs() + upstream = _FakeUpstream([], hold_after_events=True) + client_ws = _FakeWebSocket(frames=[first, later], hold_after_initial=True) + handler = _DummyOpenAIHandler() + memory = _ToolFailingMemoryHandler() + handler.memory_handler = memory + + async def _trigger() -> None: + await asyncio.sleep(0.05) + client_ws.trigger_disconnect() + + with patch.dict(sys.modules, {"websockets": _make_fake_websockets_module(upstream)}): + trigger_task = asyncio.create_task(_trigger()) + try: + await handler.handle_openai_responses_ws(client_ws) + finally: + trigger_task.cancel() + try: + await trigger_task + except asyncio.CancelledError: + pass + + assert memory.queries == [first_input, later_input] + assert upstream.sent[0] == first + assert f"current memory: {later_input}" in json.loads(upstream.sent[1])["response"]["input"] + + +@pytest.mark.asyncio +async def test_memory_lookup_preserves_list_shaped_later_frame_input(): + first, _later = _issue_2059_turns() + list_frame = _list_turn( + "later turn with list payload", + instructions="list payload instructions", + ) + expected_input = json.loads(list_frame)["response"]["input"] + upstream = _FakeUpstream([], hold_after_events=True) + client_ws = _FakeWebSocket(frames=[first, list_frame], hold_after_initial=True) + handler = _DummyOpenAIHandler() + memory = _MemoryHandler() + handler.memory_handler = memory + + async def _trigger() -> None: + await asyncio.sleep(0.05) + client_ws.trigger_disconnect() + + with patch.dict(sys.modules, {"websockets": _make_fake_websockets_module(upstream)}): + trigger_task = asyncio.create_task(_trigger()) + try: + await handler.handle_openai_responses_ws(client_ws) + finally: + trigger_task.cancel() + try: + await trigger_task + except asyncio.CancelledError: + pass + + forwarded_later = json.loads(upstream.sent[1]) + assert forwarded_later["response"]["input"] == expected_input + assert memory.queries[-1] == "list payload instructions" + + +@pytest.mark.asyncio +async def test_later_frame_compression_receives_memory_prepared_input(): + first, later = _issue_2059_turns() + _first_input, later_input = _issue_2059_inputs() + upstream = _FakeUpstream([], hold_after_events=True) + client_ws = _FakeWebSocket(frames=[first, later], hold_after_initial=True) + handler = _DummyOpenAIHandler() + handler.config.optimize = True + memory = _MemoryHandler() + handler.memory_handler = memory + seen_inputs: list[object] = [] + + def _capture_compress(payload, *, model, request_id, timing=None): + seen_inputs.append(payload["input"]) + return payload, False, 0, [], "test_noop", 10, 10, 0 + + async def _trigger() -> None: + await asyncio.sleep(0.05) + client_ws.trigger_disconnect() + + handler._compress_openai_responses_payload = _capture_compress # type: ignore[method-assign] + + with patch.dict(sys.modules, {"websockets": _make_fake_websockets_module(upstream)}): + trigger_task = asyncio.create_task(_trigger()) + try: + await handler.handle_openai_responses_ws(client_ws) + finally: + trigger_task.cancel() + try: + await trigger_task + except asyncio.CancelledError: + pass + + assert len(seen_inputs) == 2 + assert f"current memory: {later_input}" in str(seen_inputs[1])