fix(codex): offload project resolution

Author: krzkraw

headroom/providers/codex/project_context.py

@@ -2,13 +2,15 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 import logging
 import os
 import sqlite3
 import time
 from collections.abc import Mapping
 from dataclasses import dataclass
+from functools import lru_cache
 from pathlib import Path
 from typing import Any, Literal, cast
 
@@ -26,6 +28,7 @@
 _MAX_METADATA_BYTES = 16 * 1024
 _SQLITE_TIMEOUT_SECONDS = 0.1
 _SQLITE_ATTEMPTS = 2
+_ROLLOUT_CACHE_MAX_ENTRIES = 256
 
 
 @dataclass(frozen=True, slots=True)
@@ -121,6 +124,27 @@ def resolve(
             metadata_reason,
         )
 
+    async def resolve_async(
+        self,
+        *,
+        headers: Mapping[str, str],
+        body: Mapping[str, Any],
+        pinned_cwd: Path | None = None,
+        project_root_override: str | None = None,
+    ) -> CodexResolvedProject:
+        """Resolve optional project context without blocking async model traffic."""
+        try:
+            return await asyncio.to_thread(
+                self.resolve,
+                headers=headers,
+                body=body,
+                pinned_cwd=pinned_cwd,
+                project_root_override=project_root_override,
+            )
+        except Exception:
+            logger.warning("event=codex_project_resolution_failed", exc_info=True)
+            return self._skip("resolver_failed")
+
     @staticmethod
     def _header(headers: Mapping[str, str], name: str) -> str | None:
         lowered = name.lower()
@@ -150,6 +174,10 @@ def _turn_identity(
             metadata = container.get("client_metadata")
             identity = self._identity_from_metadata(metadata)
             if identity is not None:
+                if sum(len(value.encode("utf-8")) for value in identity if value) > (
+                    _MAX_METADATA_BYTES
+                ):
+                    return None, "codex-client-metadata", "metadata_too_large"
                 return identity, "codex-client-metadata", "resolved"
         return None, "unresolved", "metadata_missing"
 
@@ -359,7 +387,30 @@ def _rollout_for_thread(state_path: Path, thread_id: str) -> tuple[Path | None,
         return None, "state_locked"
 
     def _cwd_from_rollout(self, rollout: Path, turn_id: str) -> tuple[Path | None, str]:
-        matches: set[Path] = set()
+        try:
+            metadata = rollout.stat()
+        except OSError:
+            return None, "rollout_stale"
+        fingerprint = (
+            metadata.st_dev,
+            metadata.st_ino,
+            metadata.st_size,
+            metadata.st_mtime_ns,
+            metadata.st_ctime_ns,
+        )
+        raw_cwds, reason = _cached_raw_cwds_from_rollout(rollout, fingerprint, turn_id)
+        if reason != "resolved":
+            return None, reason
+        matches = {cwd for raw_cwd in raw_cwds if (cwd := self._canonical_cwd(raw_cwd))}
+        if len(matches) > 1:
+            return None, "turn_ambiguous"
+        if not matches:
+            return None, "turn_context_missing"
+        return next(iter(matches)), "resolved"
+
+    @staticmethod
+    def _read_raw_cwds_from_rollout(rollout: Path, turn_id: str) -> tuple[tuple[str, ...], str]:
+        matches: set[str] = set()
         truncated = False
         try:
             with rollout.open(encoding="utf-8") as handle:
@@ -378,18 +429,14 @@ def _cwd_from_rollout(self, rollout: Path, turn_id: str) -> tuple[Path | None, s
                         continue
                     cwd = payload.get("cwd")
                     if isinstance(cwd, str):
-                        canonical = self._canonical_cwd(cwd)
-                        if canonical is not None:
-                            matches.add(canonical)
+                        matches.add(cwd)
         except OSError:
-            return None, "rollout_stale"
+            return (), "rollout_stale"
         if truncated:
-            return None, "rollout_truncated"
-        if len(matches) > 1:
-            return None, "turn_ambiguous"
+            return (), "rollout_truncated"
         if matches:
-            return next(iter(matches)), "resolved"
-        return None, "turn_context_missing"
+            return tuple(sorted(matches)), "resolved"
+        return (), "turn_context_missing"
 
     @staticmethod
     def _skip(reason: str) -> CodexResolvedProject:
@@ -402,4 +449,13 @@ def _skip(reason: str) -> CodexResolvedProject:
         )
 
 
+@lru_cache(maxsize=_ROLLOUT_CACHE_MAX_ENTRIES)
+def _cached_raw_cwds_from_rollout(
+    rollout: Path,
+    _fingerprint: tuple[int, int, int, int, int],
+    turn_id: str,
+) -> tuple[tuple[str, ...], str]:
+    return CodexProjectContextResolver._read_raw_cwds_from_rollout(rollout, turn_id)
+
+
 __all__ = ["CodexProjectContextResolver", "CodexResolvedProject"]

headroom/proxy/handlers/openai.py

@@ -5163,7 +5163,7 @@ async def handle_openai_responses(
             or None
         )
         if client == "codex":
-            codex_project = CodexProjectContextResolver().resolve(
+            codex_project = await CodexProjectContextResolver().resolve_async(
                 headers=dict(request.headers),
                 body=body,
                 project_root_override=codex_project_root_override,
@@ -7057,7 +7057,7 @@ async def _prepare_memory_frame(frame_body: dict[str, Any], frame_raw: str) -> s
 
                 resolved_project_root_override = ws_project_root_override
                 if client == "codex":
-                    resolved_project = ws_project_resolver.resolve(
+                    resolved_project = await ws_project_resolver.resolve_async(
                         headers=ws_headers,
                         body=frame_body,
                         pinned_cwd=ws_pinned_project_cwd,

tests/test_codex_project_context.py

@@ -4,6 +4,7 @@
 import json
 import sqlite3
 import sys
+import threading
 from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import patch
@@ -129,6 +130,42 @@ def test_codex_http_projects_are_isolated_and_ws_mismatch_fails_closed(
     assert mismatch.cwd is None
     assert mismatch.reason == "project_mismatch"
 
+    _append_turn(rollout_a, "turn-a", project_b)
+    changed_rollout = resolver.resolve(
+        headers={},
+        body={"client_metadata": {"thread_id": "thread-a", "turn_id": "turn-a"}},
+    )
+    assert changed_rollout.cwd is None
+    assert changed_rollout.reason == "turn_ambiguous"
+
+
+def test_cached_rollout_recanonicalizes_symlink_before_pinned_check(
+    monkeypatch, tmp_path: Path
+) -> None:
+    codex_home = tmp_path / "codex"
+    codex_home.mkdir()
+    project_a = tmp_path / "a"
+    project_b = tmp_path / "b"
+    project_a.mkdir()
+    project_b.mkdir()
+    project_link = tmp_path / "current-project"
+    project_link.symlink_to(project_a, target_is_directory=True)
+    rollout = codex_home / "rollout.jsonl"
+    _seed_rollout(rollout, "turn", project_link)
+    _seed_thread(codex_home, "thread", rollout)
+    monkeypatch.setenv("CODEX_HOME", str(codex_home))
+    resolver = CodexProjectContextResolver()
+    body = {"client_metadata": {"thread_id": "thread", "turn_id": "turn"}}
+
+    first = resolver.resolve(headers={}, body=body)
+    project_link.unlink()
+    project_link.symlink_to(project_b, target_is_directory=True)
+    second = resolver.resolve(headers={}, body=body, pinned_cwd=first.cwd)
+
+    assert first.cwd == project_a.resolve()
+    assert second.cwd is None
+    assert second.reason == "project_mismatch"
+
 
 def test_turn_specific_resume_and_fork_use_exact_context(monkeypatch, tmp_path: Path) -> None:
     codex_home = tmp_path / "codex"
@@ -284,7 +321,7 @@ def test_explicit_overrides_precede_state_and_body_cwd(monkeypatch, tmp_path: Pa
     )
 
     assert project_id.source == "x-headroom-project-id"
-    assert project_id.project_key == "chosen"
+    assert project_id.project_key.startswith("chosen-")
     assert explicit_cwd.cwd == project_a.resolve()
     assert explicit_cwd.source == "x-headroom-cwd"
 
@@ -465,6 +502,49 @@ async def _record_request_outcome(self, _outcome):
     assert unresolved_handler.observed == []
 
 
+@pytest.mark.asyncio
+async def test_http_project_resolution_does_not_block_event_loop(monkeypatch) -> None:
+    resolver_entered = threading.Event()
+    resolver_release = threading.Event()
+
+    def blocked_resolve(self, **kwargs):
+        resolver_entered.set()
+        assert resolver_release.wait(timeout=3)
+        raise RuntimeError("blocked resolver failed")
+
+    monkeypatch.setattr(CodexProjectContextResolver, "resolve", blocked_resolve)
+    monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _DummyTokenizer())
+    request = _build_request(
+        {"model": "gpt-5.4", "input": "hello"},
+        {"Authorization": "Bearer test", "X-Client": "codex"},
+    )
+    handler = _HTTPHandler()
+    progressed_before_release = False
+
+    async def unrelated_work() -> None:
+        nonlocal progressed_before_release
+        while not resolver_entered.is_set():
+            await asyncio.sleep(0)
+        progressed_before_release = not resolver_release.is_set()
+        resolver_release.set()
+
+    release_timer = threading.Timer(2, resolver_release.set)
+    release_timer.start()
+    try:
+        response, _ = await asyncio.gather(
+            handler.handle_openai_responses(request),
+            unrelated_work(),
+        )
+    finally:
+        resolver_release.set()
+        release_timer.cancel()
+        release_timer.join()
+
+    assert resolver_entered.is_set()
+    assert progressed_before_release
+    assert response.status_code == 200
+
+
 @pytest.mark.asyncio
 async def test_ws_mismatch_skips_project_memory_but_forwards_main_traffic(
     monkeypatch, tmp_path: Path
@@ -521,7 +601,8 @@ def compute_memory_tool_definitions(self, _provider):
         [
             json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
             json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
-        ]
+        ],
+        hold_after_events=True,
     )
     websocket = _FakeWebSocket(frames=[first, mismatch])
     handler = _WSHandler()

tests/test_codex_ws_per_frame_memory.py

@@ -147,7 +147,8 @@ async def test_memory_lookup_runs_for_each_issue_artifact_frame_and_preserves_no
         [
             json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
             json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
-        ]
+        ],
+        hold_after_events=True,
     )
     first_turn, later_turn = _issue_2059_turns()
     first_input, later_input = _issue_2059_inputs()
@@ -183,7 +184,8 @@ async def test_memory_lookup_skips_input_bearing_non_create_first_frame():
         [
             json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
             json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
-        ]
+        ],
+        hold_after_events=True,
     )
     _first_input, later_input = _issue_2059_inputs()
     cancel_frame = json.dumps(
@@ -214,7 +216,8 @@ async def test_memory_lookup_skips_bypassed_frames():
         [
             json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
             json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
-        ]
+        ],
+        hold_after_events=True,
     )
     first, later = _issue_2059_turns()
     client_ws = _FakeWebSocket(
@@ -238,7 +241,8 @@ async def test_memory_lookup_keeps_legacy_direct_first_frame():
         [
             json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
             json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
-        ]
+        ],
+        hold_after_events=True,
     )
     first_input, later_input = _issue_2059_inputs()
     first = _direct_turn(first_input)
@@ -265,7 +269,8 @@ async def test_memory_lookup_skips_disabled_memory(monkeypatch):
         [
             json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
             json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
-        ]
+        ],
+        hold_after_events=True,
     )
     first, later = _issue_2059_turns()
     client_ws = _FakeWebSocket(frames=[first, later])