feat(tracing): emit token usage from litellm streaming and auto_send spans

Streaming calls now default stream_options.include_usage=True, the usage-only final chunk is collected, and both auto_send variants attach completion usage to span output. concat_completion_chunks no longer drops choices when a chunk has none.

Author: levilentz

src/agentex/lib/core/services/adk/providers/litellm.py

@@ -24,6 +24,20 @@
 logger = logging.make_logger(__name__)
 
 
+def _stream_kwargs_with_usage(llm_config: LLMConfig) -> dict:
+    """Completion kwargs with usage reporting enabled on the final stream chunk.
+
+    litellm only reports usage for streaming calls when
+    ``stream_options={"include_usage": True}`` is set; default it on so usage
+    reaches the span. Callers can still opt out by explicitly passing
+    ``stream_options={"include_usage": False}``.
+    """
+    kwargs = llm_config.model_dump()
+    stream_options = kwargs.get("stream_options") or {}
+    kwargs["stream_options"] = {"include_usage": True, **stream_options}
+    return kwargs
+
+
 class LiteLLMService:
     def __init__(
         self,
@@ -122,7 +136,11 @@ async def chat_completion_auto_send(
 
             if span:
                 if streaming_context.task_message:
-                    span.output = streaming_context.task_message.model_dump()
+                    output = streaming_context.task_message.model_dump()
+                    # Per-call usage for billing; deduped against any turn aggregate
+                    if completion.usage is not None:
+                        output["usage"] = completion.usage.model_dump()
+                    span.output = output
             return streaming_context.task_message if streaming_context.task_message else None
 
     async def chat_completion_stream(
@@ -151,20 +169,23 @@ async def chat_completion_stream(
         if self.llm_gateway is None:
             raise ValueError("LLM Gateway is not set")
 
+        completion_kwargs = _stream_kwargs_with_usage(llm_config)
         trace = self.tracer.trace(trace_id)
         async with trace.span(
             parent_id=parent_span_id,
             name="chat_completion_stream",
-            input=llm_config.model_dump(),
+            input=completion_kwargs,
         ) as span:
             # Direct streaming outside temporal - yield each chunk as it comes
             chunks: list[Completion] = []
             async for chunk in self.llm_gateway.acompletion_stream(
-                **llm_config.model_dump()
+                **completion_kwargs
             ):
                 chunks.append(chunk)
                 yield chunk
             if span:
+                # The usage-bearing final chunk survives concat, so the dumped
+                # completion carries usage for billing
                 span.output = concat_completion_chunks(chunks).model_dump()
 
     async def chat_completion_stream_auto_send(
@@ -192,11 +213,12 @@ async def chat_completion_stream_auto_send(
         if not llm_config.stream:
             llm_config.stream = True
 
+        completion_kwargs = _stream_kwargs_with_usage(llm_config)
         trace = self.tracer.trace(trace_id)
         async with trace.span(
             parent_id=parent_span_id,
             name="chat_completion_stream_auto_send",
-            input=llm_config.model_dump(),
+            input=completion_kwargs,
         ) as span:
             # Use streaming context manager
             async with self.streaming_service.streaming_task_message_context(
@@ -210,9 +232,12 @@ async def chat_completion_stream_auto_send(
                 # Get the streaming response
                 chunks = []
                 async for response in self.llm_gateway.acompletion_stream(
-                    **llm_config.model_dump()
+                    **completion_kwargs
                 ):
                     heartbeat_if_in_workflow("chat completion streaming")
+                    # Store every chunk for final message assembly, including
+                    # the usage-only final chunk, which has no choices
+                    chunks.append(response)
                     if (
                         response.choices
                         and len(response.choices) > 0
@@ -230,9 +255,6 @@ async def chat_completion_stream_auto_send(
                             )
                             heartbeat_if_in_workflow("content chunk streamed")
 
-                        # Store the chunk for final message assembly
-                        chunks.append(response)
-
                 # Update the final message content
                 complete_message = concat_completion_chunks(chunks)
                 if (
@@ -257,6 +279,10 @@ async def chat_completion_stream_auto_send(
 
             if span:
                 if streaming_context.task_message:
-                    span.output = streaming_context.task_message.model_dump()
+                    output = streaming_context.task_message.model_dump()
+                    # Per-call usage for billing; deduped against any turn aggregate
+                    if complete_message.usage is not None:
+                        output["usage"] = complete_message.usage.model_dump()
+                    span.output = output
 
         return streaming_context.task_message if streaming_context.task_message else None

src/agentex/lib/utils/completions.py

@@ -3,6 +3,7 @@
 from copy import deepcopy
 from typing import Any
 from functools import reduce, singledispatch
+from itertools import zip_longest
 
 from agentex.lib.types.llm_messages import (
     Delta,
@@ -21,7 +22,13 @@ def _concat_chunks(_a: None, b: Any):
 
 @_concat_chunks.register
 def _(a: Completion, b: Completion) -> Completion:
-    a.choices = [_concat_chunks(*c) for c in zip(a.choices, b.choices, strict=False)]
+    # Chunks can have unequal choices: with stream_options.include_usage the
+    # final chunk carries usage but no choices. Keep the unpaired side instead
+    # of truncating to the shorter list.
+    a.choices = [
+        x if y is None else y if x is None else _concat_chunks(x, y)
+        for x, y in zip_longest(a.choices, b.choices)
+    ]
     a.usage = _concat_chunks(a.usage, b.usage)
 
     return a

tests/lib/adk/providers/test_litellm_usage.py

@@ -0,0 +1,213 @@
+"""Tests that LiteLLMService puts LLM token usage on spans for billing.
+
+Covers the paths that previously dropped usage: both auto_send variants (span
+output was only the TaskMessage dump) and streaming (litellm omits usage unless
+``stream_options.include_usage`` is set).
+"""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from contextlib import asynccontextmanager
+from unittest.mock import AsyncMock, MagicMock
+
+from agentex.types.span import Span
+from agentex.types.task_message import TaskMessage
+from agentex.types.task_message_content import TextContent
+from agentex.lib.types.llm_messages import (
+    Delta,
+    Usage,
+    Choice,
+    LLMConfig,
+    Completion,
+    AssistantMessage,
+)
+from agentex.lib.core.services.adk.providers.litellm import (
+    LiteLLMService,
+    _stream_kwargs_with_usage,
+)
+
+
+class FakeTrace:
+    def __init__(self) -> None:
+        self.spans: list[Span] = []
+
+    @asynccontextmanager
+    async def span(self, name, parent_id=None, input=None, data=None, task_id=None):
+        span = Span(
+            id=f"span-{len(self.spans)}",
+            name=name,
+            start_time=datetime.now(UTC),
+            trace_id="trace-1",
+            parent_id=parent_id,
+            input=input,
+        )
+        self.spans.append(span)
+        yield span
+
+
+class FakeTracer:
+    def __init__(self) -> None:
+        self.trace_obj = FakeTrace()
+
+    def trace(self, trace_id):
+        return self.trace_obj
+
+
+def _make_streaming_service():
+    streaming_context = MagicMock()
+    streaming_context.task_message = TaskMessage(
+        id="msg-1",
+        task_id="task-1",
+        content=TextContent(author="agent", content="", format="markdown"),
+    )
+    streaming_context.stream_update = AsyncMock()
+
+    @asynccontextmanager
+    async def fake_context(**kwargs):
+        yield streaming_context
+
+    streaming_service = MagicMock()
+    streaming_service.streaming_task_message_context = fake_context
+    return streaming_service, streaming_context
+
+
+def _make_service(llm_gateway) -> tuple[LiteLLMService, FakeTracer]:
+    streaming_service, _ = _make_streaming_service()
+    tracer = FakeTracer()
+    service = LiteLLMService(
+        agentex_client=MagicMock(),
+        streaming_service=streaming_service,
+        tracer=tracer,
+        llm_gateway=llm_gateway,
+    )
+    return service, tracer
+
+
+def _stream_gateway(chunks, captured_kwargs):
+    gateway = MagicMock()
+
+    def acompletion_stream(**kwargs):
+        captured_kwargs.update(kwargs)
+
+        async def stream():
+            for chunk in chunks:
+                yield chunk
+
+        return stream()
+
+    gateway.acompletion_stream = acompletion_stream
+    return gateway
+
+
+def _delta_chunk(content: str, role: str | None = None) -> Completion:
+    return Completion(choices=[Choice(index=0, delta=Delta(content=content, role=role))])
+
+
+def _usage_only_chunk() -> Completion:
+    return Completion(
+        choices=[],
+        usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
+    )
+
+
+class TestStreamKwargsWithUsage:
+    def test_defaults_include_usage_on(self):
+        config = LLMConfig(model="gpt-4o", messages=[], stream=True)
+        assert _stream_kwargs_with_usage(config)["stream_options"] == {"include_usage": True}
+
+    def test_caller_opt_out_preserved(self):
+        config = LLMConfig(model="gpt-4o", messages=[], stream=True, stream_options={"include_usage": False})
+        assert _stream_kwargs_with_usage(config)["stream_options"] == {"include_usage": False}
+
+    def test_merges_with_other_stream_options(self):
+        config = LLMConfig(model="gpt-4o", messages=[], stream=True, stream_options={"other": 1})
+        assert _stream_kwargs_with_usage(config)["stream_options"] == {"include_usage": True, "other": 1}
+
+
+class TestChatCompletionAutoSend:
+    async def test_span_output_carries_usage(self):
+        completion = Completion(
+            choices=[Choice(index=0, message=AssistantMessage(content="Hello!"), finish_reason="stop")],
+            usage=Usage(prompt_tokens=7, completion_tokens=3, total_tokens=10),
+        )
+        gateway = MagicMock()
+        gateway.acompletion = AsyncMock(return_value=completion)
+        service, tracer = _make_service(gateway)
+
+        await service.chat_completion_auto_send(
+            task_id="task-1",
+            llm_config=LLMConfig(model="gpt-4o", messages=[], stream=False),
+            trace_id="trace-1",
+        )
+
+        span = tracer.trace_obj.spans[0]
+        assert span.output["usage"] == {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}
+
+    async def test_span_output_omits_usage_when_absent(self):
+        completion = Completion(
+            choices=[Choice(index=0, message=AssistantMessage(content="Hello!"), finish_reason="stop")],
+        )
+        gateway = MagicMock()
+        gateway.acompletion = AsyncMock(return_value=completion)
+        service, tracer = _make_service(gateway)
+
+        await service.chat_completion_auto_send(
+            task_id="task-1",
+            llm_config=LLMConfig(model="gpt-4o", messages=[], stream=False),
+            trace_id="trace-1",
+        )
+
+        assert "usage" not in tracer.trace_obj.spans[0].output
+
+
+class TestChatCompletionStream:
+    async def test_stream_requests_usage_and_span_output_carries_it(self):
+        captured_kwargs: dict = {}
+        chunks = [_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!"), _usage_only_chunk()]
+        service, tracer = _make_service(_stream_gateway(chunks, captured_kwargs))
+
+        results = []
+        async for chunk in service.chat_completion_stream(
+            llm_config=LLMConfig(model="gpt-4o", messages=[], stream=True),
+            trace_id="trace-1",
+        ):
+            results.append(chunk)
+
+        assert len(results) == 3
+        assert captured_kwargs["stream_options"] == {"include_usage": True}
+        span = tracer.trace_obj.spans[0]
+        assert span.output["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+        assert span.output["choices"][0]["message"]["content"] == "Hello!"
+
+
+class TestChatCompletionStreamAutoSend:
+    async def test_usage_only_final_chunk_reaches_span_output(self):
+        captured_kwargs: dict = {}
+        chunks = [_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!"), _usage_only_chunk()]
+        service, tracer = _make_service(_stream_gateway(chunks, captured_kwargs))
+
+        await service.chat_completion_stream_auto_send(
+            task_id="task-1",
+            llm_config=LLMConfig(model="gpt-4o", messages=[], stream=True),
+            trace_id="trace-1",
+        )
+
+        assert captured_kwargs["stream_options"] == {"include_usage": True}
+        span = tracer.trace_obj.spans[0]
+        assert span.output["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+        # TaskMessage dump is still the base of the span output
+        assert span.output["id"] == "msg-1"
+
+    async def test_stream_without_usage_chunk_omits_usage(self):
+        captured_kwargs: dict = {}
+        chunks = [_delta_chunk("Hi", role="assistant")]
+        service, tracer = _make_service(_stream_gateway(chunks, captured_kwargs))
+
+        await service.chat_completion_stream_auto_send(
+            task_id="task-1",
+            llm_config=LLMConfig(model="gpt-4o", messages=[], stream=True),
+            trace_id="trace-1",
+        )
+
+        assert "usage" not in tracer.trace_obj.spans[0].output

tests/lib/utils/__init__.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from agentex.lib.utils.completions import concat_completion_chunks
+from agentex.lib.types.llm_messages import Delta, Usage, Choice, Completion
+
+
+def _delta_chunk(content: str, role: str | None = None) -> Completion:
+    return Completion(choices=[Choice(index=0, delta=Delta(content=content, role=role))])
+
+
+def _usage_only_chunk(prompt: int, completion: int) -> Completion:
+    # stream_options.include_usage: litellm/OpenAI send a final chunk that has
+    # usage but an empty choices list
+    return Completion(
+        choices=[],
+        usage=Usage(prompt_tokens=prompt, completion_tokens=completion, total_tokens=prompt + completion),
+    )
+
+
+class TestConcatCompletionChunks:
+    def test_concatenates_delta_content(self):
+        result = concat_completion_chunks([_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!")])
+
+        assert result.choices[0].message.content == "Hello!"
+
+    def test_trailing_usage_only_chunk_keeps_choices_and_usage(self):
+        result = concat_completion_chunks(
+            [_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!"), _usage_only_chunk(10, 5)]
+        )
+
+        assert result.choices[0].message.content == "Hello!"
+        assert result.usage is not None
+        assert result.usage.prompt_tokens == 10
+        assert result.usage.completion_tokens == 5
+        assert result.usage.total_tokens == 15
+
+    def test_usage_summed_across_chunks(self):
+        chunk_a = _delta_chunk("a", role="assistant")
+        chunk_a.usage = Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)
+        chunk_b = _delta_chunk("b")
+        chunk_b.usage = Usage(prompt_tokens=4, completion_tokens=5, total_tokens=9)
+
+        result = concat_completion_chunks([chunk_a, chunk_b])
+
+        assert result.usage.prompt_tokens == 5
+        assert result.usage.completion_tokens == 7
+        assert result.usage.total_tokens == 12
+
+    def test_no_usage_chunks_leave_usage_none(self):
+        result = concat_completion_chunks([_delta_chunk("x", role="assistant")])
+
+        assert result.usage is None