From 62b14cd8dd740240f3c849555b273d55e05f7a55 Mon Sep 17 00:00:00 2001
From: kcncb <128760047+kcncb@users.noreply.github.com>
Date: Fri, 30 Oct 2026 11:38:42 +0800
Subject: [PATCH] fix(cost): detect file citations in provider_specific_fields
 again

After the provider-specific annotation split, file citations no longer live on
`message.annotations` -- that field is typed for `url_citation` only, and the
split moved everything else under `provider_specific_fields["annotations"]`.
The file-search detector still read only `message.annotations`, so a streamed
response with a `file_citation` was no longer recognised as a file search and
was priced at 0.0 instead of the tool cost.

The detector now walks both places, sharing one helper with the rest of the
check, so a citation is found wherever the split put it. A `url_citation` is
still not a file search call.
---
 .../llm_cost_calc/tool_call_cost_tracking.py  | 24 ++++++++---
 .../test_tool_call_cost_tracking.py           | 42 +++++++++++++++++++
 2 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py
index 045869e4547..3dbd862ae81 100644
--- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py
+++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py
@@ -613,13 +613,27 @@ class StandardBuiltInToolCostTracking:
                 message: Message | None = getattr(choice, "message", None)
                 if message is None:
                     continue
-                annotations: list[ChatCompletionAnnotation] | None = getattr(message, "annotations", None)
-                if annotations:
-                    for annotation in annotations:
-                        if annotation.get("type", None) == annotation_type:
-                            return True
+                for annotation in StandardBuiltInToolCostTracking._message_annotations(message):
+                    if annotation.get("type", None) == annotation_type:
+                        return True
         return False
 
+    @staticmethod
+    def _message_annotations(message: Message) -> list[ChatCompletionAnnotation]:
+        """Every annotation on the message, wherever the split put it.
+
+        `Message.annotations` is typed for `url_citation` only, so file citations travel
+        under `provider_specific_fields["annotations"]` and were invisible to the cost
+        detector, which billed streamed file search at 0.0.
+        """
+        annotations: list[ChatCompletionAnnotation] = list(getattr(message, "annotations", None) or [])
+        provider_specific_fields = getattr(message, "provider_specific_fields", None)
+        if isinstance(provider_specific_fields, Mapping):
+            nested = cast("Mapping[str, object]", provider_specific_fields).get("annotations")
+            if isinstance(nested, list):
+                annotations.extend(cast("list[ChatCompletionAnnotation]", nested))
+        return annotations
+
     @staticmethod
     def response_includes_output_type(
         response_object: ResponsesAPIResponse,
diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
index e2452a928a4..ad91db692f4 100644
--- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
+++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
@@ -1060,3 +1060,45 @@ def test_azure_ai_non_claude_web_search_keeps_generic_fallback() -> None:
 def test_built_in_tool_lookups_ignore_malformed_request_values(kwargs: dict[str, object]) -> None:
     assert StandardBuiltInToolCostTracking.get_web_search_options(kwargs) is None
     assert StandardBuiltInToolCostTracking.get_file_search_tool_call(kwargs) is None
+
+
+def _response_with_annotations(**message_kwargs) -> ModelResponse:
+    from litellm.types.utils import Choices, Message
+
+    message = Message(content="", role="assistant", **message_kwargs)
+    return ModelResponse(choices=[Choices(message=message)], model="gpt-4o")
+
+
+def test_file_citation_under_provider_specific_fields_is_a_file_search_call():
+    """A streamed file citation must still be billed.
+
+    `Message.annotations` is typed for `url_citation` only, so file citations sit in
+    `provider_specific_fields["annotations"]` after the provider-specific split. The
+    detector only read the former, so streamed file search was priced at 0.
+    """
+    response = _response_with_annotations(
+        provider_specific_fields={
+            "annotations": [{"type": "file_citation", "file_id": "file-abc123", "filename": "doc.pdf", "index": 0}]
+        }
+    )
+
+    assert StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response) is True
+
+
+def test_url_citation_is_not_a_file_search_call():
+    """A url citation is a web search result, not a file search one."""
+    response = _response_with_annotations(
+        provider_specific_fields={
+            "annotations": [
+                {
+                    "type": "url_citation",
+                    "url": "https://example.com",
+                    "title": "example",
+                    "start_index": 0,
+                    "end_index": 1,
+                }
+            ]
+        }
+    )
+
+    assert StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response) is False