Merge 8e6a54cfb3 into 2e427c6247

litellm/proxy/litellm_pre_call_utils.py

@@ -854,11 +854,13 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None:
     if not isinstance(metadata, dict):
         return None
 
-    user_id: Final = metadata.get("user_id")
-    if isinstance(user_id, dict):
-        session_id = user_id.get("session_id")
-        if isinstance(session_id, str) and _ANTHROPIC_SESSION_ID_VALUE_RE.fullmatch(session_id):
-            return session_id
+    metadata_map: Final = cast(Mapping[object, object], metadata)  # cast-ok: dict checked above
+    user_id: Final = metadata_map.get("user_id")
+    identity: Final = safe_json_loads(user_id) if isinstance(user_id, str) else user_id
+    if isinstance(identity, dict):
+        identity_session_id: Final = identity.get("session_id")
+        if isinstance(identity_session_id, str) and _ANTHROPIC_SESSION_ID_VALUE_RE.fullmatch(identity_session_id):
+            return identity_session_id
         return None
     if not isinstance(user_id, str):
         return None

tests/unit/proxy/test_litellm_pre_call_utils.py

@@ -3446,6 +3446,55 @@ def test_add_litellm_metadata_from_anthropic_user_id_sets_session_id():
     assert "litellm_trace_id" not in data
 
 
+@pytest.mark.parametrize("metadata_key", ("metadata", "litellm_metadata"))
+@pytest.mark.parametrize("session_id", ("session-one", "session-two"))
+def test_anthropic_json_user_id_extracts_session_without_rewriting_payload(metadata_key: str, session_id: str) -> None:
+    user_id: Final = json.dumps({"device_id": "device", "account_uuid": "", "session_id": session_id})
+    data: Final = {"metadata": {"user_id": user_id}, "litellm_metadata": {}}
+    LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
+        headers={}, data=data, _metadata_variable_name=metadata_key
+    )
+    assert data[metadata_key]["session_id"] == session_id
+    assert data["litellm_session_id"] == session_id
+    assert data["metadata"]["user_id"] == user_id
+
+
+@pytest.mark.parametrize("header", ("x-litellm-session-id", "x-claude-code-session-id"))
+def test_session_header_beats_anthropic_json_user_id(header: str) -> None:
+    user_id: Final = json.dumps({"session_id": "body-session-id"})
+    data: Final = {"metadata": {"user_id": user_id}, "litellm_metadata": {}}
+    LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
+        headers={header: "header-session-id"}, data=data, _metadata_variable_name="litellm_metadata"
+    )
+    assert data["litellm_metadata"]["session_id"] == "header-session-id"
+    assert data["litellm_session_id"] == "header-session-id"
+    assert data["metadata"]["user_id"] == user_id
+
+
+@pytest.mark.parametrize(
+    "user_id",
+    (
+        '{"session_id":',
+        "[]",
+        "null",
+        "123",
+        '"session-123"',
+        "{}",
+        '{"session_id": null}',
+        '{"session_id": 123}',
+        '{"session_id": ""}',
+        '{"session_id": "invalid session!", "parent_session_id": "parent-session"}',
+        '{"device_id": "device", "parent_session_id": "parent-session"}',
+    ),
+)
+def test_anthropic_json_user_id_ignores_invalid_session(user_id: str) -> None:
+    data: Final = {"metadata": {"user_id": user_id}}
+    LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
+        headers={}, data=data, _metadata_variable_name="metadata"
+    )
+    assert data == {"metadata": {"user_id": user_id}}
+
+
 def test_add_litellm_metadata_from_anthropic_user_id_dict_sets_session_id():
     data = {
         "metadata": {