core/agent: emit UsageEvent from stream and non-stream rounds

Request stream_options.include_usage; yield usage after each model
completion when the provider reports it.

src/plyngent/agent/events.py

@@ -6,6 +6,8 @@ from plyngent.lmproto.openai_compatible.model import (  # noqa: TC001
     ToolChatMessage,
 )
 
+from .usage import TokenUsage  # noqa: TC001
+
 
 class TextDeltaEvent(Struct, tag_field="type", tag="text_delta"):
     content: str
@@ -38,6 +40,12 @@ class CancelledEvent(Struct, tag_field="type", tag="cancelled"):
     reason: str = ""
 
 
+class UsageEvent(Struct, tag_field="type", tag="usage"):
+    """Token usage for one model completion (one tool-loop round)."""
+
+    usage: TokenUsage
+
+
 type AgentEvent = (
     TextDeltaEvent
     | AssistantMessageEvent
@@ -46,4 +54,5 @@ type AgentEvent = (
     | MaxRoundsEvent
     | ErrorEvent
     | CancelledEvent
+    | UsageEvent
 )

src/plyngent/agent/loop.py

@@ -4,6 +4,7 @@ import asyncio
 import inspect
 from typing import TYPE_CHECKING, cast
 
+import msgspec
 from msgspec import UNSET
 
 from plyngent.lmproto.openai_compatible.client import merge_stream_tool_calls
@@ -12,6 +13,7 @@ from plyngent.lmproto.openai_compatible.model import (
     AssistantChatMessage,
     AssistantFunctionToolCall,
     ChatCompletionsParam,
+    StreamOptions,
     StreamToolCallDelta,
     ToolChatMessage,
 )
@@ -31,7 +33,9 @@ from .events import (
     TextDeltaEvent,
     ToolCallEvent,
     ToolResultEvent,
+    UsageEvent,
 )
+from .usage import token_usage_from_api
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
@@ -110,19 +114,21 @@ async def _execute_tool_calls(
         yield ToolResultEvent(message=tool_msg)
 
 
-async def _non_stream_assistant(
+async def _non_stream_round(
     client: ChatClient,
     param: ChatCompletionsParam,
-) -> tuple[AssistantChatMessage, list[TextDeltaEvent]]:
+) -> AsyncIterator[AgentEvent]:
     response = await client.chat_completions(param, stream=False)
     if not response.choices:
         msg = "chat completion response contained no choices"
         raise RuntimeError(msg)
     assistant = response.choices[0].message
-    text_events: list[TextDeltaEvent] = []
     if isinstance(assistant.content, str) and assistant.content:
-        text_events.append(TextDeltaEvent(content=assistant.content))
-    return assistant, text_events
+        yield TextDeltaEvent(content=assistant.content)
+    yield AssistantMessageEvent(message=assistant)
+    usage = token_usage_from_api(response.usage)
+    if usage is not None:
+        yield UsageEvent(usage=usage)
 
 
 async def _stream_round(
@@ -133,12 +139,21 @@ async def _stream_round(
 
     Pattern: ``stream = await client.chat_completions(..., stream=True)`` then
     ``async for chunk in stream``. Tool-call deltas are merged after the stream.
+    Requests ``stream_options.include_usage`` so providers may send a final usage chunk.
     """
-    stream = await client.chat_completions(param, stream=True)
+    stream_param = msgspec.structs.replace(
+        param,
+        stream_options=StreamOptions(include_usage=True),
+    )
+    stream = await client.chat_completions(stream_param, stream=True)
     content_parts: list[str] = []
     tool_deltas: list[StreamToolCallDelta] = []
+    last_usage = None
 
     async for chunk in stream:
+        usage = token_usage_from_api(chunk.usage)
+        if usage is not None:
+            last_usage = usage
         if not chunk.choices:
             continue
         choice = chunk.choices[0]
@@ -161,6 +176,8 @@ async def _stream_round(
         tool_calls=tool_calls,
     )
     yield AssistantMessageEvent(message=assistant)
+    if last_usage is not None:
+        yield UsageEvent(usage=last_usage)
 
 
 async def _assistant_round(
@@ -177,11 +194,10 @@ async def _assistant_round(
                 messages.append(event.message)
             yield event
         return
-    assistant, text_events = await _non_stream_assistant(client, param)
-    for event in text_events:
+    async for event in _non_stream_round(client, param):
+        if isinstance(event, AssistantMessageEvent):
+            messages.append(event.message)
         yield event
-    messages.append(assistant)
-    yield AssistantMessageEvent(message=assistant)
 
 
 def _last_assistant(messages: list[AnyChatMessage], pre_len: int) -> AssistantChatMessage:

tests/test_agent/test_loop.py

@@ -13,6 +13,7 @@ from plyngent.agent import (
     ToolCallEvent,
     ToolRegistry,
     ToolResultEvent,
+    UsageEvent,
     run_chat_loop,
     tool,
 )
@@ -146,7 +147,11 @@ class ScriptedClient:
             yield chunk
 
 
-def _response(message: AssistantChatMessage) -> ChatCompletionResponse:
+def _response(
+    message: AssistantChatMessage,
+    *,
+    usage: dict[str, int] | None = None,
+) -> ChatCompletionResponse:
     return ChatCompletionResponse(
         id="1",
         object="chat.completion",
@@ -161,7 +166,7 @@ def _response(message: AssistantChatMessage) -> ChatCompletionResponse:
             )
         ],
         system_fingerprint="",
-        usage={},
+        usage=usage if usage is not None else {},
     )
 
 
@@ -180,6 +185,46 @@ async def test_run_chat_loop_text_only() -> None:
     assert len(client.calls) == 1
 
 
+async def test_non_stream_emits_usage() -> None:
+    client = ScriptedClient(
+        [
+            _response(
+                AssistantChatMessage(content="hi"),
+                usage={"prompt_tokens": 9, "completion_tokens": 2, "total_tokens": 11},
+            ),
+        ]
+    )
+    messages: list[AnyChatMessage] = [UserChatMessage(content="x")]
+    events = [e async for e in run_chat_loop(client, messages, model="m", stream=False)]
+    usages = [e for e in events if isinstance(e, UsageEvent)]
+    assert len(usages) == 1
+    assert usages[0].usage.prompt_tokens == 9
+    assert usages[0].usage.completion_tokens == 2
+    assert usages[0].usage.total_tokens == 11
+
+
+async def test_chat_agent_accumulates_session_usage() -> None:
+    client = ScriptedClient(
+        [
+            _response(
+                AssistantChatMessage(content="a"),
+                usage={"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
+            ),
+            _response(
+                AssistantChatMessage(content="b"),
+                usage={"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10},
+            ),
+        ]
+    )
+    agent = ChatAgent(client, model="m", stream=False)
+    _ = [e async for e in agent.run("one")]
+    assert agent.last_turn_usage.total_tokens == 6
+    assert agent.session_usage.total_tokens == 6
+    _ = [e async for e in agent.run("two")]
+    assert agent.last_turn_usage.total_tokens == 10
+    assert agent.session_usage.total_tokens == 16
+
+
 async def test_stream_yields_deltas_incrementally() -> None:
     """Text deltas are yielded as chunks arrive, not only after the full stream."""