[Feat] Agent Gateway - Add token counting non streaming + streaming (#17779)

* fix calculate_a2a_cost

* add cost_per_query

* add test_asend_message_uses_cost_per_query

* fix: _initialize_slack_alerting_jobs

* feat: add token tracking for agents invoke

* add A2ARequestUtils

* add _set_usage_on_logging_obj

* test_asend_message_token_tracking

* add _handle_a2a_response_logging

* test_asend_message_streaming_token_tracking

* add A2AStreamingIterator

litellm/a2a_protocol/main.py

@@ -7,7 +7,9 @@ Provides standalone functions with @client decorator for LiteLLM logging integra
 import asyncio
 from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union
 
+import litellm
 from litellm._logging import verbose_logger
+from litellm.a2a_protocol.utils import A2ARequestUtils
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
@@ -38,6 +40,29 @@ except ImportError:
     pass
 
 
+def _set_usage_on_logging_obj(
+    kwargs: Dict[str, Any],
+    prompt_tokens: int,
+    completion_tokens: int,
+) -> None:
+    """
+    Set usage on litellm_logging_obj for standard logging payload.
+
+    Args:
+        kwargs: The kwargs dict containing litellm_logging_obj
+        prompt_tokens: Number of input tokens
+        completion_tokens: Number of output tokens
+    """
+    litellm_logging_obj = kwargs.get("litellm_logging_obj")
+    if litellm_logging_obj is not None:
+        usage = litellm.Usage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_tokens + completion_tokens,
+        )
+        litellm_logging_obj.model_call_details["usage"] = usage
+
+
 def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str:
     """
     Extract agent info and set model/custom_llm_provider for cost tracking.
@@ -123,6 +148,20 @@ async def asend_message(
     # Wrap in LiteLLM response type for _hidden_params support
     response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response)
 
+    # Calculate token usage from request and response
+    response_dict = a2a_response.model_dump(mode="json", exclude_none=True)
+    prompt_tokens, completion_tokens, _ = A2ARequestUtils.calculate_usage_from_request_response(
+        request=request,
+        response_dict=response_dict,
+    )
+
+    # Set usage on logging obj for standard logging payload
+    _set_usage_on_logging_obj(
+        kwargs=kwargs,
+        prompt_tokens=prompt_tokens,
+        completion_tokens=completion_tokens,
+    )
+
     return response
 
 

litellm/a2a_protocol/streaming_iterator.py

@@ -0,0 +1,165 @@
+"""
+A2A Streaming Iterator with token tracking and logging support.
+"""
+
+import asyncio
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.a2a_protocol.utils import A2ARequestUtils
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.thread_pool_executor import executor
+
+if TYPE_CHECKING:
+    from a2a.types import SendStreamingMessageRequest, SendStreamingMessageResponse
+
+
+class A2AStreamingIterator:
+    """
+    Async iterator for A2A streaming responses with token tracking.
+
+    Collects chunks, extracts text, and logs usage on completion.
+    """
+
+    def __init__(
+        self,
+        stream: AsyncIterator["SendStreamingMessageResponse"],
+        request: "SendStreamingMessageRequest",
+        logging_obj: LiteLLMLoggingObj,
+        agent_name: str = "unknown",
+    ):
+        self.stream = stream
+        self.request = request
+        self.logging_obj = logging_obj
+        self.agent_name = agent_name
+        self.start_time = datetime.now()
+
+        # Collect chunks for token counting
+        self.chunks: List[Any] = []
+        self.collected_text_parts: List[str] = []
+        self.final_chunk: Optional[Any] = None
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> "SendStreamingMessageResponse":
+        try:
+            chunk = await self.stream.__anext__()
+
+            # Store chunk
+            self.chunks.append(chunk)
+
+            # Extract text from chunk for token counting
+            self._collect_text_from_chunk(chunk)
+
+            # Check if this is the final chunk (completed status)
+            if self._is_completed_chunk(chunk):
+                self.final_chunk = chunk
+
+            return chunk
+
+        except StopAsyncIteration:
+            # Stream ended - handle logging
+            if self.final_chunk is None and self.chunks:
+                self.final_chunk = self.chunks[-1]
+            await self._handle_stream_complete()
+            raise
+
+    def _collect_text_from_chunk(self, chunk: Any) -> None:
+        """Extract text from a streaming chunk and add to collected parts."""
+        try:
+            chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
+            text = A2ARequestUtils.extract_text_from_response(chunk_dict)
+            if text:
+                self.collected_text_parts.append(text)
+        except Exception:
+            verbose_logger.debug("Failed to extract text from A2A streaming chunk")
+
+    def _is_completed_chunk(self, chunk: Any) -> bool:
+        """Check if chunk indicates stream completion."""
+        try:
+            chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
+            result = chunk_dict.get("result", {})
+            if isinstance(result, dict):
+                status = result.get("status", {})
+                if isinstance(status, dict):
+                    return status.get("state") == "completed"
+        except Exception:
+            pass
+        return False
+
+    async def _handle_stream_complete(self) -> None:
+        """Handle logging and token counting when stream completes."""
+        try:
+            end_time = datetime.now()
+
+            # Calculate tokens from collected text
+            input_message = A2ARequestUtils.get_input_message_from_request(self.request)
+            input_text = A2ARequestUtils.extract_text_from_message(input_message)
+            prompt_tokens = A2ARequestUtils.count_tokens(input_text)
+
+            # Use the last (most complete) text from chunks
+            output_text = self.collected_text_parts[-1] if self.collected_text_parts else ""
+            completion_tokens = A2ARequestUtils.count_tokens(output_text)
+
+            total_tokens = prompt_tokens + completion_tokens
+
+            # Create usage object
+            usage = litellm.Usage(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=total_tokens,
+            )
+
+            # Set usage on logging obj
+            self.logging_obj.model_call_details["usage"] = usage
+
+            # Build result for logging
+            result = self._build_logging_result(usage)
+
+            # Call success handlers
+            asyncio.create_task(
+                self.logging_obj.async_success_handler(
+                    result=result,
+                    start_time=self.start_time,
+                    end_time=end_time,
+                    cache_hit=None,
+                )
+            )
+
+            executor.submit(
+                self.logging_obj.success_handler,
+                result=result,
+                cache_hit=None,
+                start_time=self.start_time,
+                end_time=end_time,
+            )
+
+            verbose_logger.info(
+                f"A2A streaming completed: prompt_tokens={prompt_tokens}, "
+                f"completion_tokens={completion_tokens}, total_tokens={total_tokens}"
+            )
+
+        except Exception as e:
+            verbose_logger.debug(f"Error in A2A streaming completion handler: {e}")
+
+    def _build_logging_result(self, usage: litellm.Usage) -> Dict[str, Any]:
+        """Build a result dict for logging."""
+        result: Dict[str, Any] = {
+            "id": getattr(self.request, "id", "unknown"),
+            "jsonrpc": "2.0",
+            "usage": usage.model_dump() if hasattr(usage, "model_dump") else dict(usage),
+        }
+
+        # Add final chunk result if available
+        if self.final_chunk:
+            try:
+                chunk_dict = self.final_chunk.model_dump(mode="json", exclude_none=True)
+                result["result"] = chunk_dict.get("result", {})
+            except Exception:
+                pass
+
+        return result
+

litellm/a2a_protocol/utils.py

@@ -0,0 +1,138 @@
+"""
+Utility functions for A2A protocol.
+"""
+
+from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union
+
+import litellm
+from litellm._logging import verbose_logger
+
+if TYPE_CHECKING:
+    from a2a.types import SendMessageRequest, SendStreamingMessageRequest
+
+
+class A2ARequestUtils:
+    """Utility class for A2A request/response processing."""
+
+    @staticmethod
+    def extract_text_from_message(message: Any) -> str:
+        """
+        Extract text content from A2A message parts.
+
+        Args:
+            message: A2A message dict or object with 'parts' containing text parts
+
+        Returns:
+            Concatenated text from all text parts
+        """
+        if message is None:
+            return ""
+
+        # Handle both dict and object access
+        if isinstance(message, dict):
+            parts = message.get("parts", [])
+        else:
+            parts = getattr(message, "parts", []) or []
+
+        text_parts: List[str] = []
+        for part in parts:
+            if isinstance(part, dict):
+                if part.get("kind") == "text":
+                    text_parts.append(part.get("text", ""))
+            else:
+                if getattr(part, "kind", None) == "text":
+                    text_parts.append(getattr(part, "text", ""))
+
+        return " ".join(text_parts)
+
+    @staticmethod
+    def extract_text_from_response(response_dict: Dict[str, Any]) -> str:
+        """
+        Extract text content from A2A response result.
+
+        Args:
+            response_dict: A2A response dict with 'result' containing message
+
+        Returns:
+            Text from response message parts
+        """
+        result = response_dict.get("result", {})
+        if not isinstance(result, dict):
+            return ""
+
+        message = result.get("message", {})
+        return A2ARequestUtils.extract_text_from_message(message)
+
+    @staticmethod
+    def get_input_message_from_request(
+        request: "Union[SendMessageRequest, SendStreamingMessageRequest]",
+    ) -> Any:
+        """
+        Extract the input message from an A2A request.
+
+        Args:
+            request: The A2A SendMessageRequest or SendStreamingMessageRequest
+
+        Returns:
+            The message object/dict or None
+        """
+        params = getattr(request, "params", None)
+        if params is None:
+            return None
+        return getattr(params, "message", None)
+
+    @staticmethod
+    def count_tokens(text: str) -> int:
+        """
+        Count tokens in text using litellm.token_counter.
+
+        Args:
+            text: Text to count tokens for
+
+        Returns:
+            Token count, or 0 if counting fails
+        """
+        if not text:
+            return 0
+        try:
+            return litellm.token_counter(text=text)
+        except Exception:
+            verbose_logger.debug("Failed to count tokens")
+            return 0
+
+    @staticmethod
+    def calculate_usage_from_request_response(
+        request: "Union[SendMessageRequest, SendStreamingMessageRequest]",
+        response_dict: Dict[str, Any],
+    ) -> Tuple[int, int, int]:
+        """
+        Calculate token usage from A2A request and response.
+
+        Args:
+            request: The A2A SendMessageRequest or SendStreamingMessageRequest
+            response_dict: The A2A response as a dict
+
+        Returns:
+            Tuple of (prompt_tokens, completion_tokens, total_tokens)
+        """
+        # Count input tokens
+        input_message = A2ARequestUtils.get_input_message_from_request(request)
+        input_text = A2ARequestUtils.extract_text_from_message(input_message)
+        prompt_tokens = A2ARequestUtils.count_tokens(input_text)
+
+        # Count output tokens
+        output_text = A2ARequestUtils.extract_text_from_response(response_dict)
+        completion_tokens = A2ARequestUtils.count_tokens(output_text)
+
+        total_tokens = prompt_tokens + completion_tokens
+
+        return prompt_tokens, completion_tokens, total_tokens
+
+
+# Backwards compatibility aliases
+def extract_text_from_a2a_message(message: Any) -> str:
+    return A2ARequestUtils.extract_text_from_message(message)
+
+
+def extract_text_from_a2a_response(response_dict: Dict[str, Any]) -> str:
+    return A2ARequestUtils.extract_text_from_response(response_dict)

litellm/litellm_core_utils/litellm_logging.py

@@ -1651,6 +1651,11 @@ class Logging(LiteLLMLoggingBaseClass):
                 result = self._handle_non_streaming_google_genai_generate_content_response_logging(
                     result=result
                 )
+            elif (
+                self.call_type == CallTypes.asend_message.value
+                or self.call_type == CallTypes.send_message.value
+            ):
+                result = self._handle_a2a_response_logging(result=result)
 
             logging_result = self.normalize_logging_result(result=result)
 
@@ -3243,6 +3248,29 @@ class Logging(LiteLLMLoggingBaseClass):
         )
         return result
 
+    def _handle_a2a_response_logging(self, result: Any) -> Any:
+        """
+        Handles logging for A2A (Agent-to-Agent) responses.
+
+        Adds usage from model_call_details to the result if available.
+        Uses Pydantic's model_copy to avoid modifying the original response.
+
+        Args:
+            result: The LiteLLMSendMessageResponse from the A2A call
+
+        Returns:
+            The response object with usage added if available
+        """
+        # Get usage from model_call_details (set by asend_message)
+        usage = self.model_call_details.get("usage")
+        if usage is None:
+            return result
+
+        # Deep copy result and add usage
+        result_copy = result.model_copy(deep=True)
+        result_copy.usage = usage.model_dump() if hasattr(usage, "model_dump") else dict(usage)
+        return result_copy
+
 
 def _get_masked_values(
     sensitive_object: dict,
@@ -3806,7 +3834,10 @@ def _init_custom_logger_compatible_class(  # noqa: PLR0915
             from litellm.integrations.opentelemetry import (
                 OpenTelemetryConfig,
             )
-            from litellm.integrations.weave.weave_otel import WeaveOtelLogger, get_weave_otel_config
+            from litellm.integrations.weave.weave_otel import (
+                WeaveOtelLogger,
+                get_weave_otel_config,
+            )
 
             weave_otel_config = get_weave_otel_config()
 

litellm/types/agents.py

@@ -221,6 +221,9 @@ class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase):
     result: Optional[Dict[str, Any]] = None
     error: Optional[Dict[str, Any]] = None
 
+    # LiteLLM usage tracking
+    usage: Optional[Dict[str, Any]] = None
+
     model_config = {"extra": "allow"}
 
     # LiteLLM private attributes for logging/cost tracking