diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 1541b8acfdc..903974363e8 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -64,6 +64,15 @@ _TRANSPORT_ONLY_CREDENTIAL_KEYS: Final = frozenset({"provider_specific_header", # Excludes the two explicit litellm headers which are handled with higher priority. _GENERIC_SESSION_ID_HEADER_RE: Final = re.compile(r"^x-.+-session-id$", re.IGNORECASE) _EXPLICIT_SESSION_HEADERS: Final = frozenset({"x-litellm-trace-id", "x-litellm-session-id"}) +# Codex carries its conversation uuid in unprefixed headers, so the +# x--session-id convention above never matches it. Current builds send +# ``session-id``/``thread-id``; builds before the codex-api split sent +# ``session_id``/``conversation_id``. Ordered session before thread. +_CODEX_SESSION_ID_HEADERS: Final = ("session-id", "session_id", "thread-id", "conversation_id") +# Matches every first-party Codex originator: codex-tui, codex_cli_rs, codex_exec, +# codex_vscode, "Codex ...". A separator is required so an unrelated "codexfoo" client +# does not read as Codex. +_CODEX_CLIENT_PREFIX_RE: Final = re.compile(r"^codex[-_ /]", re.IGNORECASE) # Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores # (covers UUIDs and most common session-id formats). _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") @@ -583,6 +592,35 @@ def _extract_generic_session_id_from_headers( return None +def _extract_codex_session_id_from_headers( + normalized: Mapping[str, str], +) -> str | None: + """ + Read Codex's conversation uuid off one of ``_CODEX_SESSION_ID_HEADERS``. + + Codex sends no request metadata the Anthropic path could parse and no + ``x-``-prefixed session header, so without this every turn of a Codex session + falls through to a freshly generated per-call trace id and lands as its own + row in the logs instead of grouping. + + Unprefixed names like ``session-id`` are generic enough that another client + could send one meaning something unrelated, and colliding values across + callers would merge their traces, so this only applies to callers that + identify as Codex. + """ + user_agent: Final = normalized.get("user-agent") + if not isinstance(user_agent, str) or not is_codex_user_agent(user_agent): + return None + return next( + ( + value + for value in (normalized.get(header) for header in _CODEX_SESSION_ID_HEADERS) + if isinstance(value, str) and _SESSION_ID_VALUE_RE.match(value) + ), + None, + ) + + def get_chain_id_from_headers(headers: dict[str, str] | None) -> str | None: """ Extract chain id for call chaining from request headers. @@ -592,6 +630,7 @@ def get_chain_id_from_headers(headers: dict[str, str] | None) -> str | None: 2. ``x-litellm-session-id`` (explicit) 3. Any ``x--session-id`` header whose value looks like a session id (alphanumeric / UUID, at least 8 chars). E.g. ``x-claude-code-session-id``. + 4. Codex's unprefixed ``session-id`` / ``thread-id``, for Codex callers only. Header keys are matched case-insensitively so this works with raw header dicts from any transport. @@ -606,6 +645,7 @@ def get_chain_id_from_headers(headers: dict[str, str] | None) -> str | None: normalized.get("x-litellm-trace-id") or normalized.get("x-litellm-session-id") or _extract_generic_session_id_from_headers(normalized) + or _extract_codex_session_id_from_headers(normalized) ) @@ -640,10 +680,13 @@ def is_claude_code_user_agent(user_agent: str) -> bool: def is_codex_user_agent(user_agent: str) -> bool: - """Codex identifies itself as ``codex_cli_rs/ ...`` (TUI), - ``codex_exec/ ...`` (exec mode), or ``codex_vscode/ ...`` - (IDE extension); all share the ``codex_`` prefix.""" - return user_agent.startswith("codex_") + """Codex builds its user agent as ``/ ...`` and ships + several first-party originators: ``codex-tui``, ``codex_cli_rs``, + ``codex_exec`` (exec mode), ``codex_vscode`` (IDE extension) and ``Codex ...`` + (see ``is_first_party_originator`` in codex-rs). They agree only on the + ``codex`` stem, and the TUI sends a bare ``codex-tui`` with no version at all, + so match the stem plus a separator rather than any one spelling.""" + return bool(_CODEX_CLIENT_PREFIX_RE.match(user_agent)) def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_config: ProxyConfig) -> bool: diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index d14c89342e9..264376495ec 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -3167,6 +3167,129 @@ def test_get_chain_id_from_headers_generic_vendor_session_id(): ) +CODEX_USER_AGENT = "codex_cli_rs/0.62.0 (Mac OS 25.5.0; arm64) Apple_Terminal" +CODEX_SESSION_UUID = "0199f0c2-8b41-7c3e-9a52-6d1f4b8e2a77" + + +@pytest.mark.parametrize( + "user_agent", + [ + "codex-tui", + "codex-tui/0.149.0 (Mac OS 26.5.1; arm64) ghostty/1.3.1 (codex-tui; 0.149.0)", + "codex_cli_rs/0.62.0 (Mac OS 25.5.0; arm64) Apple_Terminal", + "codex_exec/0.62.0 (Linux 6.1; x86_64) unknown", + "codex_vscode/0.62.0 (Mac OS 26.5.1; arm64) vscode/1.99.0", + "Codex CLI/1.0", + ], +) +def test_is_codex_user_agent_accepts_every_first_party_originator(user_agent: str): + """Codex ships several originators sharing only the `codex` stem, and the TUI + sends a bare `codex-tui` with no version, so matching one spelling misses real clients.""" + from litellm.proxy.litellm_pre_call_utils import is_codex_user_agent + + assert is_codex_user_agent(user_agent) is True + + +@pytest.mark.parametrize( + "user_agent", + ["codexify/1.0", "mycodex-tui/1.0", "curl/8.7.1", "claude-cli/2.1.0 (external, cli)", ""], +) +def test_is_codex_user_agent_rejects_non_codex_clients(user_agent: str): + from litellm.proxy.litellm_pre_call_utils import is_codex_user_agent + + assert is_codex_user_agent(user_agent) is False + + +def test_get_chain_id_from_headers_codex_tui_user_agent(): + """The real Codex TUI user agent must group turns, not just the codex_cli_rs spelling.""" + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + ua = "codex-tui/0.149.0 (Mac OS 26.5.1; arm64) ghostty/1.3.1 (codex-tui; 0.149.0)" + assert get_chain_id_from_headers({"user-agent": ua, "session-id": CODEX_SESSION_UUID}) == CODEX_SESSION_UUID + assert ( + get_chain_id_from_headers({"user-agent": "codex-tui", "session-id": CODEX_SESSION_UUID}) == CODEX_SESSION_UUID + ) + + +@pytest.mark.parametrize( + "header", + ["session-id", "session_id", "thread-id", "conversation_id", "Session-Id"], +) +def test_get_chain_id_from_headers_codex_unprefixed_session_id(header: str): + """Codex sends its conversation uuid unprefixed, so the x--session-id regex misses it.""" + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + assert get_chain_id_from_headers({"user-agent": CODEX_USER_AGENT, header: CODEX_SESSION_UUID}) == CODEX_SESSION_UUID + + +@pytest.mark.parametrize( + "user_agent", + ["curl/8.7.1", "claude-cli/2.1.0 (external, cli)", "OpenAI/Python 1.0.0"], +) +def test_get_chain_id_from_headers_unprefixed_session_id_requires_codex(user_agent: str): + """An unprefixed session-id from a non-Codex caller must not group traces. + + The name is generic enough that two unrelated callers could collide on a value + and have their sessions merged, so the bare-header path is Codex-only. + """ + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + assert get_chain_id_from_headers({"user-agent": user_agent, "session-id": CODEX_SESSION_UUID}) is None + assert get_chain_id_from_headers({"session-id": CODEX_SESSION_UUID}) is None + + +def test_get_chain_id_from_headers_codex_prefers_session_over_thread(): + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + assert ( + get_chain_id_from_headers( + { + "user-agent": CODEX_USER_AGENT, + "thread-id": "e96634a3-fa28-4083-b354-55542e2dca01", + "session-id": CODEX_SESSION_UUID, + } + ) + == CODEX_SESSION_UUID + ) + + +def test_get_chain_id_from_headers_codex_ignores_implausible_value(): + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + assert get_chain_id_from_headers({"user-agent": CODEX_USER_AGENT, "session-id": "short"}) is None + assert get_chain_id_from_headers({"user-agent": CODEX_USER_AGENT, "session-id": "has spaces!!"}) is None + + +def test_get_chain_id_from_headers_explicit_beats_codex_header(): + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + assert ( + get_chain_id_from_headers( + { + "user-agent": CODEX_USER_AGENT, + "x-litellm-trace-id": "explicit-id-value", + "session-id": CODEX_SESSION_UUID, + } + ) + == "explicit-id-value" + ) + + +def test_add_litellm_metadata_groups_codex_turns_into_one_session(): + """Every turn of a Codex session must log under one session id, not a fresh per-call trace id.""" + headers = {"user-agent": CODEX_USER_AGENT, "session-id": CODEX_SESSION_UUID} + turns = [{"litellm_metadata": {}}, {"litellm_metadata": {}}] + for turn in turns: + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=turn, _metadata_variable_name="litellm_metadata" + ) + + for turn in turns: + assert turn["litellm_session_id"] == CODEX_SESSION_UUID + assert turn["litellm_trace_id"] == CODEX_SESSION_UUID + assert turn["litellm_metadata"]["session_id"] == CODEX_SESSION_UUID + + def test_trace_id_from_traceparent_valid(): from litellm.proxy.litellm_pre_call_utils import _trace_id_from_traceparent