From d9aae590e344625a86cdccea2b853f8fda709fa9 Mon Sep 17 00:00:00 2001 From: "Onur C. Cakmak" Date: Wed, 28 Oct 2026 13:57:51 -0400 Subject: [PATCH] fix(chatgpt): derive the upstream session id from prompt_cache_key (#42014) * fix(chatgpt): derive Responses session id from prompt_cache_key ChatGPT partitions the Codex prompt cache by the Responses session-id header (the Codex CLI sends a stable per-conversation key derived from prompt_cache_key). The provider only consulted explicit session ids and otherwise stamped a per-request uuid4, so no turn ever reused the prompt cache -- 0% cache hits on repeat turns, an order-of-magnitude quota/cost multiplier on subscription accounts. Resolution order after this change: explicit session config (session_id / litellm_session_id / metadata.session_id) > prompt_cache_key (sanitized with _safe_header_value) > litellm-internal ids > uuid4. Two guards, mirroring get_fireworks_session_id: - proxy-generated session ids (missing_session_id: "generate") are per-request and are skipped, checked in both "metadata" and "litellm_metadata" (LITELLM_METADATA_ROUTES, responses included) - internal per-request ids (litellm_trace_id / litellm_call_id) lose to a stable conversation key Supersedes the body-key forwarding of #37630: a random session header overrides the body key on the provider side, so the header must carry the caller's key. Related: #29978, #29993. Verified against a real ChatGPT Codex account: stable session id -> 97-100% cached on repeat turns; random per-request -> 0%. * style: ruff format * fix(chatgpt): address review - hash unsafe cache keys, move tests, drop comments - hash prompt_cache_keys containing header-unsafe chars instead of mapping them through _safe_header_value, whose replacement char could collapse distinct keys onto one cache shard and serialize unrelated requests (greptile review) - strip the explanatory comment blocks per AGENTS.md comment policy - move the tests into the existing mapped responses transformation test file per AGENTS.md test-placement rule * fix(chatgpt): keep session typing and test diff scoped * fix(chatgpt): hash surrogate cache keys, type session params, flatten tests * test(chatgpt): record upstream requests with a mock instead of appending to a list * fix(chatgpt): keep proxy-mapped session headers ahead of body session_id and ignore non-string cache keys * fix(chatgpt): hash cache keys with leading or trailing whitespace so the session header stays valid * test(chatgpt): type the session id test parameters * fix(chatgpt): keep trace ids copied into litellm_session_id from outranking prompt_cache_key * fix(chatgpt): mark session ids copied from trace_id so explicit ones keep precedence * fix(chatgpt): keep the trace-id session marker out of provider params The marker now uses the _litellm_ internal prefix, so bridges that pass litellm_params back into completion() drop it instead of sending it as a provider param --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 1 + .../litellm_core_utils/get_litellm_params.py | 6 +- .../litellm_core_utils/provider_affinity.py | 55 +++++++++----- litellm/llms/chatgpt/common_utils.py | 39 ++++++---- .../test_get_litellm_params.py | 9 +++ .../test_chatgpt_responses_transformation.py | 52 +++++++++++++ .../llms/chatgpt/test_chatgpt_common_utils.py | 75 +++++++++++++++++++ 7 files changed, 202 insertions(+), 35 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 7c6d6deba66..2a20e812e68 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1737,6 +1737,7 @@ PASS_THROUGH_HEADER_PREFIX: Final = "x-pass-" TYPESAFE_DANGER_ZONE_HEADER_PREFIX: Final = "typesafe-danger-zone-" INTERNAL_KWARG_PREFIX: Final = "_litellm_" CONTROL_OPTIONS_KEY: Final = f"{INTERNAL_KWARG_PREFIX}control" +SESSION_ID_FROM_TRACE_ID_PARAM_KEY: Final = f"{INTERNAL_KWARG_PREFIX}session_id_from_trace_id" AZURE_SPEECH_CUSTOM_LLM_PROVIDER: Final = "azure_speech" AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX: Final = "/azure_speech" diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index ca8e2af18e6..ba964d200ec 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -6,7 +6,7 @@ from typing import Final, cast from pydantic import TypeAdapter, ValidationError -from litellm.constants import CONTROL_OPTIONS_KEY +from litellm.constants import CONTROL_OPTIONS_KEY, SESSION_ID_FROM_TRACE_ID_PARAM_KEY from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.types.litellm_params import MAX_CONTROL_INT_DIGITS, ControlOptions @@ -212,6 +212,9 @@ def get_litellm_params( # Derive litellm_session_id / litellm_trace_id from metadata when not provided (call chaining) _meta: Final = resolved_metadata or {} + session_id_from_trace_id: Final = ( + litellm_session_id is None and not _meta.get("session_id") and bool(_meta.get("trace_id")) + ) if litellm_session_id is None: litellm_session_id = _meta.get("session_id") or _meta.get("trace_id") if litellm_trace_id is None: @@ -254,6 +257,7 @@ def get_litellm_params( or (get_base_model_from_litellm_call_metadata(metadata=metadata) if metadata else None), "litellm_trace_id": litellm_trace_id, "litellm_session_id": litellm_session_id, + **({SESSION_ID_FROM_TRACE_ID_PARAM_KEY: True} if session_id_from_trace_id else {}), "hf_model_name": hf_model_name, "custom_prompt_dict": custom_prompt_dict, "litellm_metadata": litellm_metadata, diff --git a/litellm/litellm_core_utils/provider_affinity.py b/litellm/litellm_core_utils/provider_affinity.py index e022839c5c7..c831e92c09d 100644 --- a/litellm/litellm_core_utils/provider_affinity.py +++ b/litellm/litellm_core_utils/provider_affinity.py @@ -2,7 +2,7 @@ import re from collections.abc import Mapping from typing import Final -from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import SESSION_ID_FROM_TRACE_ID_PARAM_KEY, SESSION_ID_GENERATED_METADATA_KEY _HTTP_HEADER_NAME_PATTERN: Final = re.compile(r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$") _FORBIDDEN_AFFINITY_HEADERS: Final = frozenset( @@ -53,36 +53,51 @@ def _get_provider_affinity_header_name(litellm_params: object | None, default_he return validate_provider_affinity_header_name(header) -def get_stable_session_id(litellm_params: object | None) -> str | None: - if litellm_params is None: - return None - - direct_session_id: Final = _get_value(litellm_params, "session_id") - if direct_session_id: - return str(direct_session_id) - - metadata_values: Final[tuple[object, ...]] = tuple( +def _metadata_values(litellm_params: object) -> tuple[object, ...]: + return tuple( value for key in ("metadata", "litellm_metadata") if (value := _get_value(litellm_params, key)) is not None ) + + +def _non_generated_litellm_session_id(litellm_params: object) -> str | None: has_generated_session_id: Final = any( - isinstance(metadata, Mapping) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY) - for metadata in metadata_values + _get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) for metadata in _metadata_values(litellm_params) ) - litellm_session_id: Final = _get_value(litellm_params, "litellm_session_id") - if litellm_session_id and not has_generated_session_id: - return str(litellm_session_id) + return str(litellm_session_id) if litellm_session_id and not has_generated_session_id else None - for metadata in metadata_values: - if ( - isinstance(metadata, Mapping) - and not metadata.get(SESSION_ID_GENERATED_METADATA_KEY) - and (value := metadata.get("session_id")) + +def _metadata_session_id(litellm_params: object) -> str | None: + for metadata in _metadata_values(litellm_params): + if not _get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) and ( + value := _get_value(metadata, "session_id") ): return str(value) return None +def get_caller_litellm_session_id(litellm_params: object | None) -> str | None: + if litellm_params is None or _get_value(litellm_params, SESSION_ID_FROM_TRACE_ID_PARAM_KEY): + return None + return _non_generated_litellm_session_id(litellm_params) + + +def get_requested_session_id(litellm_params: object | None) -> str | None: + if litellm_params is None: + return None + direct_session_id: Final = _get_value(litellm_params, "session_id") + return str(direct_session_id) if direct_session_id else _metadata_session_id(litellm_params) + + +def get_stable_session_id(litellm_params: object | None) -> str | None: + if litellm_params is None: + return None + direct_session_id: Final = _get_value(litellm_params, "session_id") + if direct_session_id: + return str(direct_session_id) + return _non_generated_litellm_session_id(litellm_params) or _metadata_session_id(litellm_params) + + def add_provider_affinity_header( headers: Mapping[str, object], litellm_params: object | None, default_header: str | None = None ) -> dict[str, object]: # mutable-ok: downstream handlers add auth and signing headers diff --git a/litellm/llms/chatgpt/common_utils.py b/litellm/llms/chatgpt/common_utils.py index ec7fbfa92da..cd2ea6b9ccc 100644 --- a/litellm/llms/chatgpt/common_utils.py +++ b/litellm/llms/chatgpt/common_utils.py @@ -3,6 +3,7 @@ Constants and helpers for ChatGPT subscription OAuth. """ import base64 +import hashlib import json import os import platform @@ -17,6 +18,7 @@ import httpx from pydantic import JsonValue, TypeAdapter from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils +from litellm.litellm_core_utils.provider_affinity import get_caller_litellm_session_id, get_requested_session_id from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import ProviderSpecificHeader @@ -395,22 +397,31 @@ def _normalize_litellm_params(litellm_params: object) -> Mapping[str, object]: return MappingProxyType({}) +def _prompt_cache_session_id(prompt_cache_key: object) -> str | None: + if not isinstance(prompt_cache_key, str) or not prompt_cache_key: + return None + return ( + prompt_cache_key + if _safe_header_value(prompt_cache_key.strip()) == prompt_cache_key + else hashlib.sha256(prompt_cache_key.encode("utf-8", "surrogatepass")).hexdigest() + ) + + def get_chatgpt_session_id(litellm_params: object) -> str | None: params: Final = _normalize_litellm_params(litellm_params) - for key in ("litellm_session_id", "session_id"): - value = params.get(key) - if value: - return str(value) - metadata: Final = _normalize_litellm_params(params.get("metadata")) - if metadata: - value = metadata.get("session_id") - if value: - return str(value) - for key in ("litellm_trace_id", "litellm_call_id"): - value = params.get(key) - if value: - return str(value) - return None + return ( + get_caller_litellm_session_id(params) + or get_requested_session_id(params) + or _prompt_cache_session_id(params.get("prompt_cache_key")) + or next( + ( + str(value) + for key in ("litellm_session_id", "litellm_trace_id", "litellm_call_id") + if (value := params.get(key)) + ), + None, + ) + ) def ensure_chatgpt_session_id(litellm_params: object) -> str: diff --git a/tests/unit/litellm_core_utils/test_get_litellm_params.py b/tests/unit/litellm_core_utils/test_get_litellm_params.py index 89173956dc6..738ab774b2f 100644 --- a/tests/unit/litellm_core_utils/test_get_litellm_params.py +++ b/tests/unit/litellm_core_utils/test_get_litellm_params.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.get_litellm_params import ( stored_control_options, ) from litellm.types.litellm_params import ControlOptions +from litellm.utils import get_non_default_completion_params def _funnel_kwargs_completion_forwards(monkeypatch, wif_kwargs: dict[str, object]) -> dict[str, object]: @@ -411,6 +412,14 @@ class TestMetadataFallsBackToLitellmMetadata: result["metadata"].pop("trace_id") assert litellm_metadata == {"trace_id": "trace-1"} + def test_session_id_derived_from_trace_id_adds_no_provider_params(self): + from_trace_id: Final = get_litellm_params(metadata={"trace_id": "trace-1"}) + from_session_id: Final = get_litellm_params(metadata={"session_id": "session-1"}) + added_by_trace_id: Final = from_trace_id.keys() - from_session_id.keys() + + assert added_by_trace_id + assert not added_by_trace_id & get_non_default_completion_params(from_trace_id).keys() + @pytest.mark.parametrize( "value, expected", diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index 8985d0b7fd8..a6f7d613f28 100644 --- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -16,6 +16,7 @@ import pytest import litellm from litellm.llms.chatgpt.authenticator import Authenticator +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig from litellm.llms.openai.common_utils import OpenAIError from litellm.main import responses_api_bridge_check @@ -584,3 +585,54 @@ class TestChatGPTResponsesMultiAccount: assert headers_b["ChatGPT-Account-Id"] == "acct-b" assert headers_default["Authorization"] == "Bearer token-default" assert headers_default["ChatGPT-Account-Id"] == "acct-default" + + +@pytest.mark.parametrize("call", ["responses", "completion"]) +def test_prompt_cache_key_keeps_one_session_header_across_requests( + call: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + (tmp_path / "auth.json").write_text( + json.dumps({"access_token": "token", "account_id": "acct", "expires_at": 4102444800}) + ) + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + completed: Final = { + "type": "response.completed", + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.5", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "OK", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + } + upstream: Final = MagicMock( + side_effect=lambda _request: httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=f"event: response.completed\ndata: {json.dumps(completed)}\n\n".encode(), + ) + ) + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(upstream))) + for _ in range(2): + if call == "responses": + litellm.responses(model="chatgpt/gpt-5.5", input="hi", prompt_cache_key="conv-123", client=client) + else: + litellm.completion( + model="chatgpt/gpt-5.5", + messages=[{"role": "user", "content": "hi"}], + prompt_cache_key="conv-123", + client=client, + ) + + assert [call.args[0].headers["session_id"] for call in upstream.call_args_list] == ["conv-123", "conv-123"] diff --git a/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py b/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py index 6b280d97089..375c140187e 100644 --- a/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py +++ b/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py @@ -1,11 +1,17 @@ import base64 import json +from collections.abc import Mapping +from typing import Final +import h11 import pytest +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.llms.chatgpt.common_utils import ( extract_chatgpt_account_id, get_chatgpt_client_credential, + get_chatgpt_session_id, is_chatgpt_oauth_key, select_chatgpt_client_credential_headers, ) @@ -88,3 +94,72 @@ def test_bearer_scheme_is_case_insensitive(scheme): def test_other_authentication_schemes_are_not_chatgpt_logins(): assert not is_chatgpt_oauth_key(f"Basic {CHATGPT_JWT}") + + +class TestGetChatGPTSessionId: + @pytest.mark.parametrize( + "explicit", + [ + {"litellm_session_id": "chosen"}, + {"session_id": "chosen"}, + {"metadata": {"session_id": "chosen"}}, + get_litellm_params(litellm_session_id="chosen", metadata={"trace_id": "chosen"}), + ], + ) + def test_explicit_session_id_wins_over_prompt_cache_key(self, explicit: Mapping[str, object]) -> None: + assert get_chatgpt_session_id({**explicit, "prompt_cache_key": "conv-123"}) == "chosen" + + @pytest.mark.parametrize( + "request_scoped", + [ + {"litellm_trace_id": "trace-1"}, + {"litellm_call_id": "call-1"}, + get_litellm_params(metadata={"trace_id": "trace-1"}), + get_litellm_params(litellm_metadata={"trace_id": "trace-1"}), + {"litellm_session_id": "gen-1", "metadata": {"session_id": "gen-1", SESSION_ID_GENERATED_METADATA_KEY: True}}, + { + "litellm_session_id": "gen-1", + "litellm_metadata": {"session_id": "gen-1", SESSION_ID_GENERATED_METADATA_KEY: True}, + }, + ], + ) + def test_prompt_cache_key_wins_over_request_scoped_ids(self, request_scoped: Mapping[str, object]) -> None: + assert get_chatgpt_session_id({**request_scoped, "prompt_cache_key": "conv-123"}) == "conv-123" + + @pytest.mark.parametrize( + "mapped_header", + [ + { + "litellm_session_id": "header-session", + "metadata": {"trace_id": "header-session", "session_id": "header-session"}, + }, + get_litellm_params(metadata={"trace_id": "header-session", "session_id": "header-session"}), + ], + ) + def test_proxy_mapped_session_header_wins_over_body_session_id(self, mapped_header: Mapping[str, object]) -> None: + assert ( + get_chatgpt_session_id({**mapped_header, "session_id": "body-session", "prompt_cache_key": "conv-123"}) + == "header-session" + ) + + @pytest.mark.parametrize("prompt_cache_key", [{"tenant": "a"}, 12345, True, ["x", "y"]]) + def test_non_string_prompt_cache_key_is_ignored(self, prompt_cache_key: object) -> None: + assert get_chatgpt_session_id({"prompt_cache_key": prompt_cache_key, "litellm_call_id": "call-1"}) == "call-1" + + @pytest.mark.parametrize( + "first, second", + [ + ("a\nb", "a\tb"), + (json.loads('"conv-\\ud800"'), json.loads('"conv-\\udc00"')), + ("conv-123 ", " conv-123"), + ], + ) + def test_header_unsafe_cache_keys_map_to_distinct_stable_header_values(self, first: str, second: str) -> None: + first_id: Final = get_chatgpt_session_id({"prompt_cache_key": first}) + second_id: Final = get_chatgpt_session_id({"prompt_cache_key": second}) + + assert first_id is not None and second_id is not None + assert first_id == get_chatgpt_session_id({"prompt_cache_key": first}) + assert first_id != second_id + for value in (first_id, second_id): + h11.Request(method="POST", target="/", headers=[("host", "chatgpt.com"), ("session_id", value)])