From ba87b6239ce38a2ae6c698adf320b91f7b62704a Mon Sep 17 00:00:00 2001
From: Jack Venberg <jack.venberg@rover.com>
Date: Mon, 23 Mar 2026 13:27:53 -0700
Subject: [PATCH] fix(websearch_interception): pass metadata to search calls
 for cost attribution

The websearch interception handler executes search API calls via
litellm.asearch() but does not forward the request metadata (API key,
user, team, etc.) from the original request. This means the proxy's
spend tracking callbacks cannot attribute search costs to the correct
API key or user budget.

This change:
- Adds an optional `metadata` parameter to `_execute_search()`
- Extracts metadata from kwargs in both `_execute_agentic_loop()` and
  `_execute_chat_completion_agentic_loop()`
- Forwards it to `litellm.asearch()` so that the `@client` decorator's
  logging pipeline can attribute costs correctly

Search provider costs (e.g. google_pse/search at $0.005/query) are
already defined in model_prices_and_context_window_backup.json and
calculated by search/cost_calculator.py. This fix closes the gap so
those costs actually reach the proxy's spend tracking system.
---
 .../websearch_interception/handler.py         | 40 ++++++++++++++++---
 1 file changed, 35 insertions(+), 5 deletions(-)

diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py
index 2e5a8734085..2752d7a135c 100644
--- a/litellm/integrations/websearch_interception/handler.py
+++ b/litellm/integrations/websearch_interception/handler.py
@@ -668,6 +668,11 @@ class WebSearchInterceptionLogger(CustomLogger):
     ) -> Any:
         """Execute litellm.search() and make follow-up request"""
 
+        # Extract metadata for cost attribution
+        _metadata = kwargs.get("metadata") or kwargs.get("litellm_params", {}).get(
+            "metadata"
+        )
+
         # Extract search queries from tool_use blocks
         search_tasks = []
         for tool_call in tool_calls:
@@ -676,7 +681,9 @@ class WebSearchInterceptionLogger(CustomLogger):
                 verbose_logger.debug(
                     f"WebSearchInterception: Queuing search for query='{query}'"
                 )
-                search_tasks.append(self._execute_search(query))
+                search_tasks.append(
+                    self._execute_search(query, metadata=_metadata)
+                )
             else:
                 verbose_logger.debug(
                     f"WebSearchInterception: Tool call {tool_call['id']} has no query"
@@ -792,8 +799,18 @@ class WebSearchInterceptionLogger(CustomLogger):
             )
             raise
 
-    async def _execute_search(self, query: str) -> str:
-        """Execute a single web search using router's search tools"""
+    async def _execute_search(
+        self, query: str, metadata: Optional[Dict] = None
+    ) -> str:
+        """Execute a single web search using router's search tools.
+
+        Args:
+            query: The search query string.
+            metadata: Optional metadata dict from the original request. When
+                provided, it is forwarded to ``litellm.asearch`` so that proxy
+                callbacks (e.g. spend tracking) can attribute the search cost
+                to the correct API key / user / team.
+        """
         try:
             # Import router from proxy_server
             try:
@@ -851,7 +868,13 @@ class WebSearchInterceptionLogger(CustomLogger):
             verbose_logger.debug(
                 f"WebSearchInterception: Executing search for '{query}' using provider '{search_provider}'"
             )
-            result = await litellm.asearch(query=query, search_provider=search_provider)
+            search_kwargs: Dict[str, Any] = {
+                "query": query,
+                "search_provider": search_provider,
+            }
+            if metadata:
+                search_kwargs["metadata"] = metadata
+            result = await litellm.asearch(**search_kwargs)
 
             # Format using transformation function
             search_result_text = WebSearchTransformation.format_search_response(result)
@@ -879,6 +902,11 @@ class WebSearchInterceptionLogger(CustomLogger):
     ) -> Any:
         """Execute litellm.search() and make follow-up chat completion request"""
 
+        # Extract metadata for cost attribution
+        _metadata = kwargs.get("metadata") or kwargs.get("litellm_params", {}).get(
+            "metadata"
+        )
+
         # Extract search queries from tool_calls
         search_tasks = []
         for tool_call in tool_calls:
@@ -897,7 +925,9 @@ class WebSearchInterceptionLogger(CustomLogger):
                 verbose_logger.debug(
                     f"WebSearchInterception: Queuing search for query='{query}'"
                 )
-                search_tasks.append(self._execute_search(query))
+                search_tasks.append(
+                    self._execute_search(query, metadata=_metadata)
+                )
             else:
                 verbose_logger.debug(
                     f"WebSearchInterception: Tool call {tool_call.get('id')} has no query"