diff --git a/strix/core/execution.py b/strix/core/execution.py index bd28cce..788fcfb 100644 --- a/strix/core/execution.py +++ b/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) diff --git a/strix/core/hooks.py b/strix/core/hooks.py new file mode 100644 index 0000000..f7f888f --- /dev/null +++ b/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) diff --git a/strix/core/runner.py b/strix/core/runner.py index 04fbd65..e0f9be3 100644 --- a/strix/core/runner.py +++ b/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) diff --git a/strix/interface/utils.py b/strix/interface/utils.py index 44137bb..76b7d3e 100644 --- a/strix/interface/utils.py +++ b/strix/interface/utils.py @@ -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: