mortal/strix
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Record usage per SDK LLM response

Browse Source
This commit is contained in:
0xallam
2026-04-26 15:54:45 -07:00
parent 4791feb08e
commit ef50c2dfa6
4 changed files with 107 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
strix/core/execution.py
+15 -33
View File
@@ -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
+54
View File
@@ -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
+5 -15
View File
@@ -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)
strix/interface/utils.py
+33 -32
View File
@@ -56,15 +56,6 @@ def format_token_count(count: float | None) -> str:
return str(value)
def format_cost(cost: float | None) -> str:
value = float(cost or 0.0)
if value < 0.0001:
return f"${value:.6f}"
if value < 1:
return f"${value:.4f}"
return f"${value:.2f}"
def format_vulnerability_report(report: dict[str, Any]) -> Text: # noqa: PLR0915
"""Format a vulnerability report for CLI display with all rich fields."""
field_style = "bold #4ade80"
@@ -289,33 +280,41 @@ def _build_llm_usage_stats(
stats_text: Text,
report_state: Any,
*,
include_requests: bool = False,
live: bool = False,
) -> None:
usage = _llm_usage(report_state)
if not usage or _int_stat(usage, "total_tokens") <= 0:
if not usage or _int_stat(usage, "requests") <= 0:
stats_text.append("\n")
stats_text.append("Cost ", style="dim")
stats_text.append("$0.0000 ", style="#fbbf24")
stats_text.append("· ", style="dim white")
stats_text.append("Tokens ", style="dim")
stats_text.append("0", style="white")
return
input_tokens = _int_stat(usage, "input_tokens")
output_tokens = _int_stat(usage, "output_tokens")
cached_tokens = _detail_value(usage, "input_tokens_details", "cached_tokens")
reasoning_tokens = _detail_value(usage, "output_tokens_details", "reasoning_tokens")
cost = _float_stat(usage, "cost")
stats_text.append("LLM Usage ", style="bold cyan")
if include_requests:
stats_text.append(f"{_int_stat(usage, 'requests')} req", style="white")
stats_text.append(" | ", style="dim white")
stats_text.append(f"In: {format_token_count(input_tokens)}", style="white")
if cached_tokens > 0:
stats_text.append(f" ({format_token_count(cached_tokens)} cached)", style="dim white")
stats_text.append(" | ", style="dim white")
stats_text.append(f"Out: {format_token_count(output_tokens)}", style="white")
if reasoning_tokens > 0:
stats_text.append(f" ({format_token_count(reasoning_tokens)} reasoning)", style="dim white")
if cost > 0:
stats_text.append(" | ", style="dim white")
stats_text.append(f"Cost: {format_cost(cost)}", style="white")
stats_text.append("\n")
stats_text.append("Input Tokens ", style="dim")
stats_text.append(format_token_count(input_tokens), style="white")
if live or cached_tokens > 0:
stats_text.append(" · ", style="dim white")
stats_text.append("Cached Tokens ", style="dim")
stats_text.append(format_token_count(cached_tokens), style="white")
separator = "\n" if live else " · "
stats_text.append(separator, style="dim white")
stats_text.append("Output Tokens ", style="dim")
stats_text.append(format_token_count(output_tokens), style="white")
if live or cost > 0:
stats_text.append(" · ", style="dim white")
stats_text.append("Cost ", style="dim")
stats_text.append(f"${cost:.4f}", style="#fbbf24")
def build_final_stats_text(report_state: Any) -> Text:
@@ -325,7 +324,7 @@ def build_final_stats_text(report_state: Any) -> Text:
return stats_text
_build_vulnerability_stats(stats_text, report_state)
_build_llm_usage_stats(stats_text, report_state, include_requests=True)
_build_llm_usage_stats(stats_text, report_state)
return stats_text
@@ -368,7 +367,7 @@ def build_live_stats_text(report_state: Any) -> Text:
stats_text.append("\n")
_build_llm_usage_stats(stats_text, report_state)
_build_llm_usage_stats(stats_text, report_state, live=True)
return stats_text
@@ -384,12 +383,14 @@ def build_tui_stats_text(report_state: Any) -> Text:
usage = _llm_usage(report_state)
if usage and _int_stat(usage, "total_tokens") > 0:
stats_text.append("\n")
stats_text.append("Tokens: ", style="bold white")
stats_text.append(format_token_count(_int_stat(usage, "total_tokens")), style="white")
stats_text.append(
f"{format_token_count(_int_stat(usage, 'total_tokens'))} tokens",
style="white",
)
cost = _float_stat(usage, "cost")
if cost > 0:
stats_text.append(" Cost: ", style="bold white")
stats_text.append(format_cost(cost), style="white")
stats_text.append(" · ", style="white")
stats_text.append(f"${cost:.2f}", style="white")
caido_url = getattr(report_state, "caido_url", None)
if caido_url: