From 033c758299d2cdae584d6a8bd9d345dd84d96e22 Mon Sep 17 00:00:00 2001 From: shrey Date: Fri, 30 Oct 2026 15:48:31 -0700 Subject: [PATCH] feat(proxy)!: use Bedrock request metadata as the session id on pass-through routes (#49193) --- litellm/proxy/litellm_pre_call_utils.py | 25 +++ .../unit/proxy/test_litellm_pre_call_utils.py | 193 ++++++++++++++++++ 2 files changed, 218 insertions(+) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index a9cffcc6566..edaacd53d22 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -112,6 +112,7 @@ _EXPLICIT_SESSION_HEADERS: Final = frozenset({"x-litellm-trace-id", "x-litellm-s # ``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") +_BEDROCK_SESSION_ID_METADATA_KEYS: Final = ("sessionId", "threadId", "conversationId") # 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,}$") @@ -1055,6 +1056,22 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None: return session_id +def _get_bedrock_session_id_from_request_metadata( + normalized_headers: Mapping[str, str], request_body: Mapping[str, object] +) -> str | None: + """Read the conversation id a Bedrock SDK client tags its calls with: Converse's ``requestMetadata`` + body field or InvokeModel's ``X-Amzn-Bedrock-Request-Metadata`` JSON header. Only the Bedrock + pass-through route nests the client's body under ``data``.""" + if request_body.get("custom_llm_provider") != "bedrock" or not isinstance(request_body.get("data"), Mapping): + return None + header: Final = normalized_headers.get("x-amzn-bedrock-request-metadata") + body_tags: Final = _request_fields(_request_fields(request_body.get("data")).get("requestMetadata")) + header_tags: Final = _request_fields(safe_json_loads(header) if header else None) + tag_sources: Final = itertools.product((body_tags, header_tags), _BEDROCK_SESSION_ID_METADATA_KEYS) + values: Final = (tags.get(key) for tags, key in tag_sources) + return next((value for value in values if isinstance(value, str) and value), None) + + def _is_llm_inference_route(request: Request) -> bool: route: Final = get_request_route(request) return RouteChecks.is_llm_api_route(route=route) and not RouteChecks.check_route_access( @@ -1967,6 +1984,14 @@ class LiteLLMProxyRequestSetup: body_metadata["user_id"] = session_id verbose_proxy_logger.debug("Extracted session_id from Anthropic metadata.user_id") + if "litellm_session_id" not in data: + bedrock_session_id: Final = _get_bedrock_session_id_from_request_metadata(normalized_headers, request_body) + metadata_session_id: Final = _caller_trace_field(request_body, _metadata_variable_name, "session_id") + if bedrock_session_id and metadata_session_id in (None, bedrock_session_id): + metadata_from_headers["session_id"] = bedrock_session_id + data["litellm_session_id"] = bedrock_session_id # rebind-ok: data is an out-param + verbose_proxy_logger.debug("Extracted session_id from Bedrock request metadata") + # Last-resort fallback: the W3C standards for trace/session propagation # (https://www.w3.org/TR/trace-context/, https://www.w3.org/TR/baggage/). # Lower priority than everything above - only fires when neither the diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 3194124ad1c..91c87c63067 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -4323,6 +4323,186 @@ def test_add_litellm_metadata_from_request_headers_traceparent_and_baggage_toget assert data["litellm_session_id"] == "baggage-session-42" +def _bedrock_passthrough_data( + client_body: dict[str, object], custom_llm_provider: str = "bedrock" +) -> dict[str, object]: + return { + "method": "POST", + "endpoint": "model/us.anthropic.claude-haiku-4-5-20251001-v1:0/converse", + "custom_llm_provider": custom_llm_provider, + "data": client_body, + "litellm_metadata": {}, + } + + +@pytest.mark.parametrize( + ("headers", "client_body"), + [ + pytest.param({}, {"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}}, id="converse_body"), + pytest.param( + {"X-Amzn-Bedrock-Request-Metadata": '{"threadId": "thread-abc-123", "team": "search"}'}, + {"messages": []}, + id="invoke_header", + ), + ], +) +def test_bedrock_request_metadata_thread_id_becomes_session_id( + headers: dict[str, str], client_body: dict[str, object] +): + sent_body: Final = copy.deepcopy(client_body) + data: Final = _bedrock_passthrough_data(client_body) + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert data["litellm_session_id"] == "thread-abc-123" + assert data["litellm_metadata"] == {"session_id": "thread-abc-123"} + assert "litellm_trace_id" not in data, "a conversation tag must not merge the turns into one trace" + assert data["data"] == sent_body, "the tags must still reach Bedrock unchanged" + + +@pytest.mark.parametrize( + ("tags", "expected"), + [ + ({"threadId": "thread-abc-123", "sessionId": "session-xyz-789"}, "session-xyz-789"), + ({"conversationId": "conversation-42", "threadId": "thread-abc-123"}, "thread-abc-123"), + ({"conversationId": "conversation-42", "user": "shrey-1234"}, "conversation-42"), + ({"sessionId": "", "threadId": "thread-abc-123"}, "thread-abc-123"), + ({"sessionId": "conv-42"}, "conv-42"), + ({"sessionId": "user:123 / thread.7"}, "user:123 / thread.7"), + ], +) +def test_bedrock_request_metadata_session_key_priority(tags: dict[str, str], expected: str): + data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": tags}) + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={}, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert data["litellm_session_id"] == expected + + +@pytest.mark.parametrize( + ("headers", "data"), + [ + pytest.param( + {"x-amzn-bedrock-request-metadata": '{"threadId": "thread-abc-123"'}, + _bedrock_passthrough_data({"messages": []}), + id="malformed_header_json", + ), + pytest.param( + {"x-amzn-bedrock-request-metadata": '["thread-abc-123"]'}, + _bedrock_passthrough_data({"messages": []}), + id="header_json_not_an_object", + ), + pytest.param( + {}, _bedrock_passthrough_data({"requestMetadata": {"threadId": 1234567890}}), id="non_string_value" + ), + pytest.param({}, _bedrock_passthrough_data({"requestMetadata": {"threadId": ""}}), id="empty"), + pytest.param({}, _bedrock_passthrough_data({"requestMetadata": {"user": "shrey-1234"}}), id="no_session_key"), + pytest.param( + {"x-amzn-bedrock-request-metadata": '{"threadId": "thread-abc-123"}'}, + {"model": "claude", "custom_llm_provider": "bedrock", "litellm_metadata": {}}, + id="not_a_passthrough_body", + ), + pytest.param( + {}, + _bedrock_passthrough_data( + {"requestMetadata": {"threadId": "thread-abc-123"}}, custom_llm_provider="vertex_ai" + ), + id="not_bedrock", + ), + ], +) +def test_bedrock_request_metadata_without_usable_session_id_sets_no_session( + headers: dict[str, str], data: dict[str, object] +): + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert "litellm_session_id" not in data + + +@pytest.mark.parametrize( + ("body_tags", "header_tags", "expected"), + [ + pytest.param({"team": "search"}, {"threadId": "thread-abc-123"}, "thread-abc-123", id="body_without_session_key"), + pytest.param({"sessionId": "body-session"}, {"sessionId": "header-session"}, "body-session", id="body_wins"), + ], +) +def test_bedrock_request_metadata_reads_body_tags_then_header_tags( + body_tags: dict[str, str], header_tags: dict[str, str], expected: str +): + data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": body_tags}) + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={"x-amzn-bedrock-request-metadata": json.dumps(header_tags)}, + data=data, + _metadata_variable_name="litellm_metadata", + ) + + assert data["litellm_session_id"] == expected + + +@pytest.mark.parametrize("metadata_bucket", ["litellm_metadata", "metadata"]) +def test_bedrock_request_metadata_never_replaces_a_caller_metadata_session(metadata_bucket: str): + data: Final = { + **_bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-B"}}), + metadata_bucket: {"session_id": "conversation-A"}, + } + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={}, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert "litellm_session_id" not in data + assert data[metadata_bucket]["session_id"] == "conversation-A" + + +def test_bedrock_request_metadata_stamps_a_tag_equal_to_the_caller_metadata_session(): + data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}}) + data["litellm_metadata"] = {"session_id": "thread-abc-123"} + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={}, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert data["litellm_session_id"] == "thread-abc-123" + + +def test_bedrock_request_metadata_header_leaves_non_passthrough_caller_session_alone(): + data: Final = {"model": "claude", "custom_llm_provider": "bedrock", "metadata": {"session_id": "caller-session-1"}} + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={"x-amzn-bedrock-request-metadata": '{"threadId": "thread-abc-123"}'}, + data=data, + _metadata_variable_name="metadata", + ) + + assert data["metadata"] == {"session_id": "caller-session-1"} + assert "litellm_session_id" not in data + + +@pytest.mark.parametrize( + ("headers", "expected"), + [ + pytest.param({"x-litellm-session-id": "explicit-session-1"}, "explicit-session-1", id="explicit_header_wins"), + pytest.param({"baggage": "session.id=baggage-session-42"}, "thread-abc-123", id="beats_baggage"), + ], +) +def test_bedrock_request_metadata_session_precedence(headers: dict[str, str], expected: str): + data: Final = _bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}}) + + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="litellm_metadata" + ) + + assert data["litellm_session_id"] == expected + assert data["litellm_metadata"]["session_id"] == expected + + def test_add_litellm_metadata_from_request_headers_explicit_trace_id_beats_traceparent(): """x-litellm-trace-id must win over a traceparent header carrying a different trace-id - explicit litellm headers are always highest priority.""" @@ -9598,6 +9778,19 @@ async def test_missing_session_id_reject_accepts_body_metadata_session_id(): assert updated["metadata"]["session_id"] == "body-session-1" +@pytest.mark.asyncio +async def test_missing_session_id_reject_accepts_bedrock_request_metadata_session_id(): + updated: Final = await add_litellm_data_to_request( + data=_bedrock_passthrough_data({"messages": [], "requestMetadata": {"threadId": "thread-abc-123"}}), + request=_request_for("/bedrock/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/converse"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert _spend_log_session_id(updated, "litellm_metadata") == "thread-abc-123" + + @pytest.mark.asyncio async def test_missing_session_id_reject_returns_400_without_session_id(): with pytest.raises(ProxyException) as exc_info: