fix(websearch_interception): bill intercepted searches to the calling key

litellm/integrations/websearch_interception/handler.py

@@ -42,7 +42,12 @@
     WebSearchInterceptionConfig,
 )
 from litellm.types.llms.openai import AllMessageValues
-from litellm.types.utils import AgenticLoopParams, CallTypes, LlmProviders
+from litellm.types.utils import (
+    AgenticLoopParams,
+    CallTypes,
+    LlmProviders,
+    StandardLoggingUserAPIKeyMetadata,
+)
 from litellm.utils import ProviderConfigManager
 
 if TYPE_CHECKING:
@@ -1318,6 +1323,7 @@ async def _execute_search(
             search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router)
             search_provider: str | None = None
             search_litellm_params: dict[str, Any] = {}
+            search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool)
             if search_tool is not None:
                 await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs)
                 search_litellm_params = dict(search_tool.get("litellm_params", {}) or {})
@@ -1334,12 +1340,30 @@ async def _execute_search(
             verbose_logger.debug(
                 "WebSearchInterception: Executing search for '%s' using provider '%s'", query, search_provider
             )
+            user_api_key_auth: Final = self._get_user_api_key_auth_from_kwargs(kwargs)
+            search_metadata: Final = (
+                None
+                if user_api_key_auth is None
+                else self._build_search_request_metadata(
+                    user_api_key_auth=user_api_key_auth,
+                    search_tool_name=search_tool_name,
+                )
+            )
             search_kwargs: Final = {
                 key: value
                 for key, value in search_litellm_params.items()
                 if key != "search_provider" and value is not None
             }
-            result: Final = await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs)
+            result: Final = (
+                await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs)
+                if search_metadata is None
+                else await litellm.asearch(
+                    query=query,
+                    search_provider=search_provider,
+                    litellm_metadata=search_metadata,
+                    **search_kwargs,
+                )
+            )
 
             # Format using transformation function
             search_result_text: Final = WebSearchTransformation.format_search_response(result)
@@ -1396,6 +1420,35 @@ async def _authorize_search_tool(
                 team_object=team_object,
             )
 
+    @staticmethod
+    def _build_search_request_metadata(
+        user_api_key_auth: "UserAPIKeyAuth",
+        search_tool_name: str | None,
+    ) -> Mapping[str, object]:
+        """
+        Spend-tracking metadata for the intercepted search, so its provider cost is logged
+        and billed against the key/user/team that made the originating LLM request instead
+        of being dropped by the proxy's spend hook for lack of an owner.
+        """
+        from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
+
+        user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = (
+            LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth)
+        )
+        return {  # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches
+            **user_api_key_metadata,
+            "model_group": search_tool_name,
+            "user_api_key": user_api_key_auth.api_key,
+            "user_api_key_auth": user_api_key_auth,
+        }
+
+    @staticmethod
+    def _selected_search_tool_name(search_tool: Mapping[str, object] | None) -> str | None:
+        if search_tool is None:
+            return None
+        search_tool_name: Final = search_tool.get("search_tool_name")
+        return search_tool_name if isinstance(search_tool_name, str) and search_tool_name else None
+
     @staticmethod
     def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object] | None) -> "UserAPIKeyAuth | None":
         if not kwargs:

tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py

@@ -221,14 +221,97 @@ async def test_execute_search_passes_selected_search_tool_litellm_params(monkeyp
         kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}},
     )
 
-    mock_asearch.assert_awaited_once_with(
-        query="what is litellm",
-        search_provider="tavily",
-        api_key="fake-ui-key",
-        api_base="https://api.tavily.com",
-        timeout=10.0,
-        max_retries=2,
+    forwarded_kwargs = mock_asearch.await_args.kwargs
+    assert forwarded_kwargs["query"] == "what is litellm"
+    assert forwarded_kwargs["search_provider"] == "tavily"
+    assert forwarded_kwargs["api_key"] == "fake-ui-key"
+    assert forwarded_kwargs["api_base"] == "https://api.tavily.com"
+    assert forwarded_kwargs["timeout"] == 10.0
+    assert forwarded_kwargs["max_retries"] == 2
+
+
+@pytest.mark.asyncio
+async def test_execute_search_attributes_spend_to_the_calling_key(monkeypatch):
+    """An intercepted search is billed and logged against the key that made the LLM request.
+
+    Without the forwarded attribution metadata the proxy's spend hook skips the search
+    entirely, so its provider cost never reaches SpendLogs or any budget.
+    """
+    import litellm
+    from litellm.proxy import proxy_server
+    from litellm.proxy.hooks.proxy_track_cost_callback import _should_track_cost_callback
+
+    logger = WebSearchInterceptionLogger(
+        enabled_providers=["bedrock"],
+        search_tool_name="perplexity-sonar-pro",
     )
+    router = MagicMock()
+    router.search_tools = [
+        {
+            "search_tool_name": "perplexity-sonar-pro",
+            "litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"},
+        }
+    ]
+    mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[]))
+    user_api_key_auth = UserAPIKeyAuth(
+        api_key="hashed-sk-1234",
+        key_alias="alice-key",
+        user_id="user-alice",
+        org_id="org-1",
+    )
+
+    monkeypatch.setattr(proxy_server, "llm_router", router)
+    monkeypatch.setattr(litellm, "asearch", mock_asearch)
+
+    await logger._execute_search(
+        "what is litellm",
+        kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}},
+    )
+
+    forwarded_metadata = mock_asearch.await_args.kwargs["litellm_metadata"]
+    assert forwarded_metadata["user_api_key"] == "hashed-sk-1234"
+    assert forwarded_metadata["user_api_key_hash"] == "hashed-sk-1234"
+    assert forwarded_metadata["user_api_key_alias"] == "alice-key"
+    assert forwarded_metadata["user_api_key_user_id"] == "user-alice"
+    assert forwarded_metadata["user_api_key_org_id"] == "org-1"
+    assert forwarded_metadata["model_group"] == "perplexity-sonar-pro"
+    assert (
+        _should_track_cost_callback(
+            user_api_key=forwarded_metadata["user_api_key"],
+            user_id=forwarded_metadata["user_api_key_user_id"],
+            team_id=forwarded_metadata["user_api_key_team_id"],
+            end_user_id=None,
+            call_type="asearch",
+        )
+        is True
+    )
+
+
+@pytest.mark.asyncio
+async def test_execute_search_without_proxy_auth_context_stays_sdk_only(monkeypatch):
+    """SDK callers have no key to attribute the search to, so no proxy metadata is invented."""
+    import litellm
+    from litellm.proxy import proxy_server
+
+    logger = WebSearchInterceptionLogger(
+        enabled_providers=["bedrock"],
+        search_tool_name="perplexity-sonar-pro",
+    )
+    router = MagicMock()
+    router.search_tools = [
+        {
+            "search_tool_name": "perplexity-sonar-pro",
+            "litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"},
+        }
+    ]
+    mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[]))
+
+    monkeypatch.setattr(proxy_server, "llm_router", router)
+    monkeypatch.setattr(litellm, "asearch", mock_asearch)
+
+    await logger._execute_search("what is litellm", kwargs={"litellm_params": {}})
+
+    assert "litellm_metadata" not in mock_asearch.await_args.kwargs
 
 
 @pytest.mark.asyncio