Record usage per SDK LLM response

strix/core/execution.py

@@ -21,6 +21,7 @@ if TYPE_CHECKING:
     from pathlib import Path
 
     from agents.items import TResponseInputItem
+    from agents.lifecycle import RunHooks
     from agents.memory import Session, SQLiteSession
     from agents.result import RunResultBase
 
@@ -30,7 +31,6 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 StreamEventSink = Callable[[str, Any], None]
-UsageSink = Callable[[str, str | None, Any], None]
 
 
 async def run_agent_loop(
@@ -46,7 +46,7 @@ async def run_agent_loop(
     session: Session | None = None,
     start_parked: bool = False,
     event_sink: StreamEventSink | None = None,
-    usage_sink: UsageSink | None = None,
+    hooks: RunHooks[dict[str, Any]] | None = None,
 ) -> RunResultBase | None:
     await coordinator.attach_runtime(
         agent_id,
@@ -68,7 +68,7 @@ async def run_agent_loop(
                 session=session,
                 interactive=interactive,
                 event_sink=event_sink,
-                usage_sink=usage_sink,
+                hooks=hooks,
             )
         else:
             result = await _run_noninteractive_until_lifecycle(
@@ -81,7 +81,7 @@ async def run_agent_loop(
                 max_turns=max_turns,
                 session=session,
                 event_sink=event_sink,
-                usage_sink=usage_sink,
+                hooks=hooks,
             )
 
     if not interactive:
@@ -105,7 +105,7 @@ async def run_agent_loop(
             session=session,
             interactive=interactive,
             event_sink=event_sink,
-            usage_sink=usage_sink,
+            hooks=hooks,
         )
 
 
@@ -124,7 +124,7 @@ async def spawn_child_agent(
     skills: list[str],
     parent_history: list[Any],
     event_sink: StreamEventSink | None = None,
-    usage_sink: UsageSink | None = None,
+    hooks: RunHooks[dict[str, Any]] | None = None,
 ) -> dict[str, Any]:
     parent_id = parent_ctx.get("agent_id")
     if not isinstance(parent_id, str):
@@ -161,7 +161,7 @@ async def spawn_child_agent(
             parent_history=parent_history,
         ),
         event_sink=event_sink,
-        usage_sink=usage_sink,
+        hooks=hooks,
     )
 
     return {
@@ -185,7 +185,7 @@ async def respawn_subagents(
     parent_ctx: dict[str, Any],
     root_id: str,
     event_sink: StreamEventSink | None = None,
-    usage_sink: UsageSink | None = None,
+    hooks: RunHooks[dict[str, Any]] | None = None,
 ) -> None:
     """Re-spawn subagent runners from a restored coordinator snapshot."""
     async with coordinator._lock:
@@ -242,7 +242,7 @@ async def respawn_subagents(
                 initial_input=[],
                 start_parked=start_parked,
                 event_sink=event_sink,
-                usage_sink=usage_sink,
+                hooks=hooks,
             )
             logger.info(
                 "respawned %s (%s) parent=%s task_len=%d",
@@ -268,7 +268,7 @@ async def _run_noninteractive_until_lifecycle(
     max_turns: int,
     session: Session | None,
     event_sink: StreamEventSink | None,
-    usage_sink: UsageSink | None,
+    hooks: RunHooks[dict[str, Any]] | None,
 ) -> RunResultBase | None:
     """Non-chat mode keeps running until finish_scan / agent_finish settles status."""
     result: RunResultBase | None = None
@@ -288,7 +288,7 @@ async def _run_noninteractive_until_lifecycle(
             session=session,
             interactive=False,
             event_sink=event_sink,
-            usage_sink=usage_sink,
+            hooks=hooks,
         )
 
         status = await _agent_status(coordinator, agent_id)
@@ -333,7 +333,7 @@ async def _run_cycle(
     session: Session | None,
     interactive: bool,
     event_sink: StreamEventSink | None,
-    usage_sink: UsageSink | None,
+    hooks: RunHooks[dict[str, Any]] | None,
 ) -> RunResultBase | None:
     try:
         await coordinator.mark_running(agent_id)
@@ -344,6 +344,7 @@ async def _run_cycle(
             context=context,
             max_turns=max_turns,
             session=session,
+            hooks=hooks,
         )
         await coordinator.attach_stream(agent_id, stream)
         try:
@@ -356,8 +357,6 @@ async def _run_cycle(
             if stream.run_loop_exception is not None:
                 raise stream.run_loop_exception
         finally:
-            if usage_sink is not None:
-                _emit_stream_usage(agent, agent_id, stream, usage_sink)
             await coordinator.detach_stream(agent_id, stream)
     except Exception as exc:
         if not interactive:
@@ -477,7 +476,7 @@ async def _start_child_runner(
     initial_input: Any,
     start_parked: bool = False,
     event_sink: StreamEventSink | None = None,
-    usage_sink: UsageSink | None = None,
+    hooks: RunHooks[dict[str, Any]] | None = None,
 ) -> None:
     session = open_agent_session(child_id, agents_db_path)
     sessions_to_close.append(session)
@@ -501,25 +500,8 @@ async def _start_child_runner(
             session=session,
             start_parked=start_parked,
             event_sink=event_sink,
-            usage_sink=usage_sink,
+            hooks=hooks,
         ),
         name=f"agent-{name}-{child_id}",
     )
     await coordinator.attach_runtime(child_id, task=task_handle)
-
-
-def _emit_stream_usage(
-    agent: Any,
-    agent_id: str,
-    stream: RunResultBase,
-    usage_sink: UsageSink,
-) -> None:
-    context_wrapper = getattr(stream, "context_wrapper", None)
-    usage = getattr(context_wrapper, "usage", None)
-    if usage is None:
-        return
-    agent_name = getattr(agent, "name", None)
-    try:
-        usage_sink(agent_id, agent_name if isinstance(agent_name, str) else None, usage)
-    except Exception:
-        logger.exception("usage sink failed for %s", agent_id)

strix/core/hooks.py

@@ -0,0 +1,54 @@
+"""SDK run hooks used by Strix orchestration."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from agents.lifecycle import RunHooks
+
+from strix.report.state import get_global_report_state
+
+
+if TYPE_CHECKING:
+    from agents import RunContextWrapper
+    from agents.agent import Agent
+    from agents.items import ModelResponse
+
+
+logger = logging.getLogger(__name__)
+
+
+class ReportUsageHooks(RunHooks[dict[str, Any]]):
+    """Persist SDK-native usage after every model response."""
+
+    def __init__(self, *, model: str) -> None:
+        self._model = model
+
+    async def on_llm_end(
+        self,
+        context: RunContextWrapper[dict[str, Any]],
+        agent: Agent[dict[str, Any]],
+        response: ModelResponse,
+    ) -> None:
+        report_state = get_global_report_state()
+        if report_state is None:
+            return
+
+        ctx = context.context if isinstance(context.context, dict) else {}
+        agent_name = getattr(agent, "name", None)
+        if not isinstance(agent_name, str):
+            agent_name = None
+        agent_id = ctx.get("agent_id")
+        if not isinstance(agent_id, str) or not agent_id:
+            agent_id = agent_name or "unknown"
+
+        try:
+            report_state.record_sdk_usage(
+                agent_id=agent_id,
+                agent_name=agent_name,
+                model=self._model,
+                usage=response.usage,
+            )
+        except Exception:
+            logger.exception("failed to record SDK usage for agent %s", agent_id)

strix/core/runner.py

@@ -28,6 +28,7 @@ from strix.core.execution import (
 from strix.core.execution import (
     spawn_child_agent as start_child_agent,
 )
+from strix.core.hooks import ReportUsageHooks
 from strix.core.inputs import (
     DEFAULT_MAX_TURNS,
     build_root_task,
@@ -36,7 +37,6 @@ from strix.core.inputs import (
 )
 from strix.core.paths import run_dir_for, runtime_state_dir
 from strix.core.sessions import open_agent_session
-from strix.report.state import get_global_report_state
 from strix.runtime import session_manager
 from strix.telemetry.logging import set_scan_id, setup_scan_logging
 
@@ -167,17 +167,7 @@ async def run_strix_scan(
             sandbox=SandboxRunConfig(client=bundle["client"], session=bundle["session"]),
             trace_include_sensitive_data=False,
         )
-
-        def usage_sink(agent_id: str, agent_name: str | None, usage: Any) -> None:
-            report_state = get_global_report_state()
-            if report_state is None:
-                return
-            report_state.record_sdk_usage(
-                agent_id=agent_id,
-                agent_name=agent_name,
-                model=resolved_model,
-                usage=usage,
-            )
+        hooks = ReportUsageHooks(model=resolved_model)
 
         scope_context = build_scope_context(scan_config)
 
@@ -217,7 +207,7 @@ async def run_strix_scan(
                 max_turns=max_turns,
                 interactive=interactive,
                 event_sink=event_sink,
-                usage_sink=usage_sink,
+                hooks=hooks,
                 **kwargs,
             )
 
@@ -249,7 +239,7 @@ async def run_strix_scan(
                 parent_ctx=context,
                 root_id=root_id,
                 event_sink=event_sink,
-                usage_sink=usage_sink,
+                hooks=hooks,
             )
 
         initial_input: Any = [] if is_resume else root_task
@@ -290,7 +280,7 @@ async def run_strix_scan(
             session=root_session,
             start_parked=bool(interactive and is_resume and root_status != "running"),
             event_sink=event_sink,
-            usage_sink=usage_sink,
+            hooks=hooks,
         )
     except BaseException:
         logger.exception("Strix scan %s failed", scan_id)