Merge pull request #23686 from BerriAI/litellm_oss_staging_03_14_2026

Litellm oss staging 03 14 2026

litellm/main.py

@@ -7528,8 +7528,15 @@ def stream_chunk_builder(  # noqa: PLR0915
         ]
 
         if len(annotation_chunks) > 0:
-            annotations = annotation_chunks[0]["choices"][0]["delta"]["annotations"]
-            response["choices"][0]["message"]["annotations"] = annotations
+            # Merge annotations from ALL chunks — providers may spread
+            # them across multiple streaming chunks or send them only in
+            # the final chunk.
+            all_annotations: list = []
+            for ac in annotation_chunks:
+                all_annotations.extend(
+                    ac["choices"][0]["delta"]["annotations"]
+                )
+            response["choices"][0]["message"]["annotations"] = all_annotations
 
         audio_chunks = [
             chunk

litellm/proxy/auth/auth_utils.py

@@ -662,11 +662,12 @@ def _has_user_setup_sso():
     return sso_setup
 
 
-def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]:
+def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[list]:
     """Return the header_name mapped to CUSTOMER role, if any (dict-based)."""
     if not user_id_mapping:
         return None
     items = user_id_mapping if isinstance(user_id_mapping, list) else [user_id_mapping]
+    customer_headers_mappings = []
     for item in items:
         if not isinstance(item, dict):
             continue
@@ -675,7 +676,11 @@ def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]:
         if role is None or not header_name:
             continue
         if str(role).lower() == str(LitellmUserRoles.CUSTOMER).lower():
-            return header_name
+            customer_headers_mappings.append(header_name.lower())
+
+    if customer_headers_mappings:
+        return customer_headers_mappings
+    
     return None
 
 
@@ -724,7 +729,7 @@ def get_end_user_id_from_request_body(
     # User query: "system not respecting user_header_name property"
     # This implies the key in general_settings is 'user_header_name'.
     if request_headers is not None:
-        custom_header_name_to_check: Optional[str] = None
+        custom_header_name_to_check: Optional[Union[list, str]] = None
 
         # Prefer user mappings (new behavior)
         user_id_mapping = general_settings.get("user_header_mappings", None)
@@ -741,13 +746,21 @@ def get_end_user_id_from_request_body(
                 custom_header_name_to_check = value
 
         # If we have a header name to check, try to read it from request headers
-        if isinstance(custom_header_name_to_check, str):
+        if isinstance(custom_header_name_to_check, list):
+            headers_lower = {k.lower(): v for k, v in request_headers.items()}
+            for expected_header in custom_header_name_to_check:
+                header_value = headers_lower.get(expected_header)
+                if header_value is not None:
+                    user_id_str = str(header_value)
+                    if user_id_str.strip():
+                        return user_id_str
+                    
+        elif isinstance(custom_header_name_to_check, str):
             for header_name, header_value in request_headers.items():
                 if header_name.lower() == custom_header_name_to_check.lower():
-                    user_id_from_header = header_value
                     user_id_str = (
-                        str(user_id_from_header)
-                        if user_id_from_header is not None
+                        str(header_value)
+                        if header_value is not None
                         else ""
                     )
                     if user_id_str.strip():

tests/local_testing/test_auth_utils.py

@@ -268,7 +268,7 @@ def test_get_customer_user_header_from_mapping_returns_customer_header():
         {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"},
     ]
     result = get_customer_user_header_from_mapping(mappings)
-    assert result == "X-OpenWebUI-User-Email"
+    assert result == ["x-openwebui-user-email"]
 
 
 def test_get_customer_user_header_from_mapping_no_customer_returns_none():

tests/test_litellm/proxy/auth/test_auth_utils.py

@@ -209,3 +209,109 @@ def test_get_model_from_request_supports_google_model_names_with_slashes():
 def test_get_model_from_request_vertex_passthrough_still_works():
     route = "/vertex_ai/v1/projects/p/locations/l/publishers/google/models/gemini-1.5-pro:generateContent"
     assert get_model_from_request(request_data={}, route=route) == "gemini-1.5-pro"
+
+
+def test_get_customer_user_header_returns_none_when_no_customer_role():
+    from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping
+
+    mappings = [
+        {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}
+    ]
+    result = get_customer_user_header_from_mapping(mappings)
+    assert result is None
+
+
+def test_get_customer_user_header_returns_none_for_single_non_customer_mapping():
+    from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping
+
+    mapping = {"header_name": "X-Only-Internal", "litellm_user_role": "internal_user"}
+    result = get_customer_user_header_from_mapping(mapping)
+    assert result is None
+
+def test_get_customer_user_header_from_mapping_returns_customer_header():
+    from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping
+
+    mappings = [
+        {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"},
+        {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"},
+    ]
+    result = get_customer_user_header_from_mapping(mappings)
+    assert result == ["x-openwebui-user-email"]
+
+
+def test_get_customer_user_header_returns_customers_header_in_config_order_when_multiple_exist():
+    from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping
+
+    mappings = [
+        {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"},
+        {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"},
+        {"header_name": "X-User-Id", "litellm_user_role": "customer"},
+    ]
+    result = get_customer_user_header_from_mapping(mappings)
+    assert result == ['x-openwebui-user-email', 'x-user-id']
+    
+
+def test_get_end_user_id_returns_id_from_user_header_mappings():
+    from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
+
+    mappings = [
+        {"header_name": "x-openwebui-user-id", "litellm_user_role": "internal_user"},
+        {"header_name": "x-openwebui-user-email", "litellm_user_role": "customer"},
+    ]
+    general_settings = {"user_header_mappings": mappings}
+    headers = {"x-openwebui-user-email": "1234"}
+
+    with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \
+         patch("litellm.proxy.proxy_server.general_settings", general_settings):
+        result = get_end_user_id_from_request_body(request_body={}, request_headers=headers)
+
+    assert result == "1234"
+
+
+def test_get_end_user_id_returns_first_customer_header_when_multiple_mappings_exist():
+    from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
+
+    mappings = [
+        {"header_name": "x-openwebui-user-id", "litellm_user_role": "internal_user"},
+        {"header_name": "x-user-id", "litellm_user_role": "customer"},
+        {"header_name": "x-openwebui-user-email", "litellm_user_role": "customer"},
+    ]
+    general_settings = {"user_header_mappings": mappings}
+    headers = {
+        "x-user-id": "user-456",
+        "x-openwebui-user-email": "user@example.com",
+    }
+
+    with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \
+         patch("litellm.proxy.proxy_server.general_settings", general_settings):
+        result = get_end_user_id_from_request_body(request_body={}, request_headers=headers)
+
+    assert result == "user-456"
+
+
+def test_get_end_user_id_returns_none_when_no_customer_role_in_mappings():
+    from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
+
+    mappings = [
+        {"header_name": "x-openwebui-user-id", "litellm_user_role": "internal_user"},
+    ]
+    general_settings = {"user_header_mappings": mappings}
+    headers = {"x-openwebui-user-id": "user-789"}
+
+    with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \
+         patch("litellm.proxy.proxy_server.general_settings", general_settings):
+        result = get_end_user_id_from_request_body(request_body={}, request_headers=headers)
+
+    assert result is None
+
+def test_get_end_user_id_falls_back_to_deprecated_user_header_name():
+    from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
+
+    general_settings = {"user_header_name": "x-custom-user-id"}
+    headers = {"x-custom-user-id": "user-legacy"}
+
+    with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \
+         patch("litellm.proxy.proxy_server.general_settings", general_settings):
+        result = get_end_user_id_from_request_body(request_body={}, request_headers=headers)
+
+    assert result == "user-legacy"

tests/test_litellm/test_stream_chunk_builder_annotations.py

@@ -0,0 +1,191 @@
+"""
+Tests for stream_chunk_builder annotation merging.
+
+Previously, stream_chunk_builder only took annotations from the FIRST
+annotation chunk, losing any annotations that arrived in later chunks.
+This fix merges annotations from ALL chunks.
+"""
+
+from litellm import stream_chunk_builder
+from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+
+def test_stream_chunk_builder_merges_annotations_from_multiple_chunks():
+    """
+    stream_chunk_builder must merge annotations from ALL streaming chunks,
+    not just take them from the first annotation chunk.
+
+    Providers may spread annotations across multiple chunks (e.g. Gemini
+    sends grounding metadata in the final chunk, while intermediate chunks
+    may carry different annotations).
+    """
+    annotation_a = {
+        "type": "url_citation",
+        "url_citation": {
+            "url": "https://example.com/a",
+            "title": "Source A",
+            "start_index": 0,
+            "end_index": 10,
+        },
+    }
+    annotation_b = {
+        "type": "url_citation",
+        "url_citation": {
+            "url": "https://example.com/b",
+            "title": "Source B",
+            "start_index": 20,
+            "end_index": 30,
+        },
+    }
+
+    chunks = [
+        ModelResponseStream(
+            id="chatcmpl-test",
+            created=1700000000,
+            model="test-model",
+            object="chat.completion.chunk",
+            choices=[
+                StreamingChoices(
+                    finish_reason=None,
+                    index=0,
+                    delta=Delta(
+                        content="Part one. ",
+                        role="assistant",
+                        annotations=[annotation_a],
+                    ),
+                )
+            ],
+        ),
+        ModelResponseStream(
+            id="chatcmpl-test",
+            created=1700000000,
+            model="test-model",
+            object="chat.completion.chunk",
+            choices=[
+                StreamingChoices(
+                    finish_reason=None,
+                    index=0,
+                    delta=Delta(content="Part two."),
+                )
+            ],
+        ),
+        ModelResponseStream(
+            id="chatcmpl-test",
+            created=1700000000,
+            model="test-model",
+            object="chat.completion.chunk",
+            choices=[
+                StreamingChoices(
+                    finish_reason="stop",
+                    index=0,
+                    delta=Delta(
+                        content=None,
+                        annotations=[annotation_b],
+                    ),
+                )
+            ],
+        ),
+    ]
+
+    response = stream_chunk_builder(chunks=chunks)
+    assert response is not None
+
+    message = response["choices"][0]["message"]
+    assert message.annotations is not None
+    assert len(message.annotations) == 2
+    assert message.annotations[0] == annotation_a
+    assert message.annotations[1] == annotation_b
+
+
+def test_stream_chunk_builder_single_annotation_chunk_still_works():
+    """
+    When annotations come from a single chunk (most common case),
+    stream_chunk_builder must still work correctly (no regression).
+    """
+    annotation = {
+        "type": "url_citation",
+        "url_citation": {
+            "url": "https://example.com/only",
+            "title": "Only Source",
+            "start_index": 0,
+            "end_index": 5,
+        },
+    }
+
+    chunks = [
+        ModelResponseStream(
+            id="chatcmpl-test",
+            created=1700000000,
+            model="test-model",
+            object="chat.completion.chunk",
+            choices=[
+                StreamingChoices(
+                    finish_reason=None,
+                    index=0,
+                    delta=Delta(content="Hello", role="assistant"),
+                )
+            ],
+        ),
+        ModelResponseStream(
+            id="chatcmpl-test",
+            created=1700000000,
+            model="test-model",
+            object="chat.completion.chunk",
+            choices=[
+                StreamingChoices(
+                    finish_reason="stop",
+                    index=0,
+                    delta=Delta(content=None, annotations=[annotation]),
+                )
+            ],
+        ),
+    ]
+
+    response = stream_chunk_builder(chunks=chunks)
+    assert response is not None
+
+    message = response["choices"][0]["message"]
+    assert message.annotations is not None
+    assert len(message.annotations) == 1
+    assert message.annotations[0] == annotation
+
+
+def test_stream_chunk_builder_no_annotations():
+    """
+    When no chunks contain annotations, the message should not have
+    an annotations key (no regression).
+    """
+    chunks = [
+        ModelResponseStream(
+            id="chatcmpl-test",
+            created=1700000000,
+            model="test-model",
+            object="chat.completion.chunk",
+            choices=[
+                StreamingChoices(
+                    finish_reason=None,
+                    index=0,
+                    delta=Delta(content="Hello", role="assistant"),
+                )
+            ],
+        ),
+        ModelResponseStream(
+            id="chatcmpl-test",
+            created=1700000000,
+            model="test-model",
+            object="chat.completion.chunk",
+            choices=[
+                StreamingChoices(
+                    finish_reason="stop",
+                    index=0,
+                    delta=Delta(content=None),
+                )
+            ],
+        ),
+    ]
+
+    response = stream_chunk_builder(chunks=chunks)
+    assert response is not None
+
+    message = response["choices"][0]["message"]
+    assert not hasattr(message, "annotations") or message.annotations is None