From 4791feb08e7b645541ffd71f1735bd6b6ee1c9e4 Mon Sep 17 00:00:00 2001 From: 0xallam Date: Sun, 26 Apr 2026 15:38:41 -0700 Subject: [PATCH] Track SDK LLM usage --- strix/core/execution.py | 33 ++++++ strix/core/inputs.py | 3 +- strix/core/runner.py | 15 +++ strix/interface/utils.py | 96 +++++++++++++++ strix/report/dedupe.py | 14 ++- strix/report/state.py | 37 ++++++ strix/report/usage.py | 234 +++++++++++++++++++++++++++++++++++++ strix/telemetry/posthog.py | 15 +++ 8 files changed, 444 insertions(+), 3 deletions(-) create mode 100644 strix/report/usage.py diff --git a/strix/core/execution.py b/strix/core/execution.py index b5222c2..bd28cce 100644 --- a/strix/core/execution.py +++ b/strix/core/execution.py @@ -30,6 +30,7 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) StreamEventSink = Callable[[str, Any], None] +UsageSink = Callable[[str, str | None, Any], None] async def run_agent_loop( @@ -45,6 +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, ) -> RunResultBase | None: await coordinator.attach_runtime( agent_id, @@ -66,6 +68,7 @@ async def run_agent_loop( session=session, interactive=interactive, event_sink=event_sink, + usage_sink=usage_sink, ) else: result = await _run_noninteractive_until_lifecycle( @@ -78,6 +81,7 @@ async def run_agent_loop( max_turns=max_turns, session=session, event_sink=event_sink, + usage_sink=usage_sink, ) if not interactive: @@ -101,6 +105,7 @@ async def run_agent_loop( session=session, interactive=interactive, event_sink=event_sink, + usage_sink=usage_sink, ) @@ -119,6 +124,7 @@ async def spawn_child_agent( skills: list[str], parent_history: list[Any], event_sink: StreamEventSink | None = None, + usage_sink: UsageSink | None = None, ) -> dict[str, Any]: parent_id = parent_ctx.get("agent_id") if not isinstance(parent_id, str): @@ -155,6 +161,7 @@ async def spawn_child_agent( parent_history=parent_history, ), event_sink=event_sink, + usage_sink=usage_sink, ) return { @@ -178,6 +185,7 @@ async def respawn_subagents( parent_ctx: dict[str, Any], root_id: str, event_sink: StreamEventSink | None = None, + usage_sink: UsageSink | None = None, ) -> None: """Re-spawn subagent runners from a restored coordinator snapshot.""" async with coordinator._lock: @@ -234,6 +242,7 @@ async def respawn_subagents( initial_input=[], start_parked=start_parked, event_sink=event_sink, + usage_sink=usage_sink, ) logger.info( "respawned %s (%s) parent=%s task_len=%d", @@ -259,6 +268,7 @@ async def _run_noninteractive_until_lifecycle( max_turns: int, session: Session | None, event_sink: StreamEventSink | None, + usage_sink: UsageSink | None, ) -> RunResultBase | None: """Non-chat mode keeps running until finish_scan / agent_finish settles status.""" result: RunResultBase | None = None @@ -278,6 +288,7 @@ async def _run_noninteractive_until_lifecycle( session=session, interactive=False, event_sink=event_sink, + usage_sink=usage_sink, ) status = await _agent_status(coordinator, agent_id) @@ -322,6 +333,7 @@ async def _run_cycle( session: Session | None, interactive: bool, event_sink: StreamEventSink | None, + usage_sink: UsageSink | None, ) -> RunResultBase | None: try: await coordinator.mark_running(agent_id) @@ -344,6 +356,8 @@ 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: @@ -463,6 +477,7 @@ async def _start_child_runner( initial_input: Any, start_parked: bool = False, event_sink: StreamEventSink | None = None, + usage_sink: UsageSink | None = None, ) -> None: session = open_agent_session(child_id, agents_db_path) sessions_to_close.append(session) @@ -486,7 +501,25 @@ async def _start_child_runner( session=session, start_parked=start_parked, event_sink=event_sink, + usage_sink=usage_sink, ), 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/inputs.py b/strix/core/inputs.py index 87bf007..3d65803 100644 --- a/strix/core/inputs.py +++ b/strix/core/inputs.py @@ -12,7 +12,7 @@ from strix.config.models import DEFAULT_MODEL_RETRY # Default max_turns budget passed to the SDK runner. -DEFAULT_MAX_TURNS = 300 +DEFAULT_MAX_TURNS = 500 def build_root_task(scan_config: dict[str, Any]) -> str: @@ -112,6 +112,7 @@ def make_model_settings( parallel_tool_calls=False, tool_choice="required", retry=DEFAULT_MODEL_RETRY, + include_usage=True, ) if reasoning_effort is not None: model_settings = model_settings.resolve( diff --git a/strix/core/runner.py b/strix/core/runner.py index 6240df9..04fbd65 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -36,6 +36,7 @@ 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,6 +168,17 @@ async def run_strix_scan( 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, + ) + scope_context = build_scope_context(scan_config) root_agent = build_strix_agent( @@ -205,6 +217,7 @@ async def run_strix_scan( max_turns=max_turns, interactive=interactive, event_sink=event_sink, + usage_sink=usage_sink, **kwargs, ) @@ -236,6 +249,7 @@ async def run_strix_scan( parent_ctx=context, root_id=root_id, event_sink=event_sink, + usage_sink=usage_sink, ) initial_input: Any = [] if is_resume else root_task @@ -276,6 +290,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, ) except BaseException: logger.exception("Strix scan %s failed", scan_id) diff --git a/strix/interface/utils.py b/strix/interface/utils.py index fa208a9..44137bb 100644 --- a/strix/interface/utils.py +++ b/strix/interface/utils.py @@ -47,6 +47,24 @@ def get_cvss_color(cvss_score: float) -> str: return "#6b7280" +def format_token_count(count: float | None) -> str: + value = int(count or 0) + if value >= 1_000_000: + return f"{value / 1_000_000:.1f}M" + if value >= 1_000: + return f"{value / 1_000:.1f}K" + 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" @@ -235,6 +253,71 @@ def _build_vulnerability_stats(stats_text: Text, report_state: Any) -> None: stats_text.append("\n") +def _llm_usage(report_state: Any) -> dict[str, Any]: + if hasattr(report_state, "get_total_llm_usage"): + usage = report_state.get_total_llm_usage() + return usage if isinstance(usage, dict) else {} + usage = getattr(report_state, "run_record", {}).get("llm_usage") + return usage if isinstance(usage, dict) else {} + + +def _int_stat(usage: dict[str, Any], key: str) -> int: + try: + return max(0, int(usage.get(key) or 0)) + except (TypeError, ValueError): + return 0 + + +def _float_stat(usage: dict[str, Any], key: str) -> float: + try: + value = float(usage.get(key) or 0.0) + except (TypeError, ValueError): + return 0.0 + return value if value > 0 else 0.0 + + +def _detail_value(usage: dict[str, Any], detail_key: str, value_key: str) -> int: + details = usage.get(detail_key) + if isinstance(details, list): + details = details[0] if details and isinstance(details[0], dict) else {} + if not isinstance(details, dict): + return 0 + return _int_stat(details, value_key) + + +def _build_llm_usage_stats( + stats_text: Text, + report_state: Any, + *, + include_requests: bool = False, +) -> None: + usage = _llm_usage(report_state) + if not usage or _int_stat(usage, "total_tokens") <= 0: + 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") + + def build_final_stats_text(report_state: Any) -> Text: """Build final stats from Strix-owned scan artifacts.""" stats_text = Text() @@ -242,6 +325,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) return stats_text @@ -284,6 +368,8 @@ def build_live_stats_text(report_state: Any) -> Text: stats_text.append("\n") + _build_llm_usage_stats(stats_text, report_state) + return stats_text @@ -295,6 +381,16 @@ def build_tui_stats_text(report_state: Any) -> Text: model = load_settings().llm.model or "unknown" stats_text.append(str(model), style="white") + 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") + cost = _float_stat(usage, "cost") + if cost > 0: + stats_text.append(" Cost: ", style="bold white") + stats_text.append(format_cost(cost), style="white") + caido_url = getattr(report_state, "caido_url", None) if caido_url: stats_text.append("\n") diff --git a/strix/report/dedupe.py b/strix/report/dedupe.py index 0d5703b..af468a9 100644 --- a/strix/report/dedupe.py +++ b/strix/report/dedupe.py @@ -17,6 +17,7 @@ from strix.config.models import ( configure_sdk_model_defaults, normalize_model_name, ) +from strix.report.state import get_global_report_state if TYPE_CHECKING: @@ -187,11 +188,12 @@ async def check_duplicate( ) configure_sdk_model_defaults(settings) - model = MultiProvider().get_model(normalize_model_name(model_name)) + resolved_model = normalize_model_name(model_name) + model = MultiProvider().get_model(resolved_model) response = await model.get_response( system_instructions=DEDUPE_SYSTEM_PROMPT, input=user_msg, - model_settings=ModelSettings(retry=DEFAULT_MODEL_RETRY), + model_settings=ModelSettings(retry=DEFAULT_MODEL_RETRY, include_usage=True), tools=[], output_schema=None, handoffs=[], @@ -200,6 +202,14 @@ async def check_duplicate( conversation_id=None, prompt=None, ) + report_state = get_global_report_state() + if report_state is not None: + report_state.record_sdk_usage( + agent_id="dedupe", + agent_name="dedupe", + model=resolved_model, + usage=response.usage, + ) content = _extract_text(response) if not content: return { diff --git a/strix/report/state.py b/strix/report/state.py index 8b84cd1..c9541d0 100644 --- a/strix/report/state.py +++ b/strix/report/state.py @@ -6,7 +6,10 @@ from pathlib import Path from typing import Any, Optional from uuid import uuid4 +from agents.usage import Usage + from strix.core.paths import run_dir_for +from strix.report.usage import LLMUsageLedger from strix.report.writer import ( read_run_record, write_executive_report, @@ -51,6 +54,7 @@ class ReportState: self.scan_results: dict[str, Any] | None = None self.scan_config: dict[str, Any] | None = None + self._llm_usage = LLMUsageLedger() self.run_record: dict[str, Any] = { "run_id": self.run_id, "run_name": self.run_name, @@ -58,6 +62,7 @@ class ReportState: "end_time": None, "status": "running", "targets_info": [], + "llm_usage": self._build_llm_usage_record(), } self._run_dir: Path | None = None self._saved_vuln_ids: set[str] = set() @@ -104,6 +109,7 @@ class ReportState: if isinstance(scan_results, dict): self.scan_results = scan_results self.final_scan_result = self._format_final_scan_result(scan_results) + self._hydrate_llm_usage(data.get("llm_usage")) logger.info("report state hydrated run.json from %s", run_dir) json_path = run_dir / "vulnerabilities.json" @@ -206,6 +212,26 @@ class ReportState: def get_existing_vulnerabilities(self) -> list[dict[str, Any]]: return list(self.vulnerability_reports) + def record_sdk_usage( + self, + *, + agent_id: str, + usage: Usage | None, + agent_name: str | None = None, + model: str | None = None, + ) -> None: + """Record SDK-native token usage for one completed model run/cycle.""" + if self._llm_usage.record( + agent_id=agent_id, + agent_name=agent_name, + model=model, + usage=usage, + ): + self.save_run_data() + + def get_total_llm_usage(self) -> dict[str, Any]: + return dict(self.run_record.get("llm_usage") or self._build_llm_usage_record()) + def update_scan_final_fields( self, executive_summary: str, @@ -264,6 +290,7 @@ class ReportState: self.run_record["end_time"] = self.end_time self.run_record["status"] = status + self._sync_llm_usage_record() self._save_artifacts() def cleanup(self, status: str = "stopped") -> None: @@ -304,3 +331,13 @@ class ReportState: logger.info("Essential scan data saved to: %s", run_dir) except (OSError, RuntimeError): logger.exception("Failed to save scan data") + + def _sync_llm_usage_record(self) -> None: + self.run_record["llm_usage"] = self._build_llm_usage_record() + + def _build_llm_usage_record(self) -> dict[str, Any]: + return self._llm_usage.to_record() + + def _hydrate_llm_usage(self, raw_usage: Any) -> None: + self._llm_usage.hydrate(raw_usage) + self._sync_llm_usage_record() diff --git a/strix/report/usage.py b/strix/report/usage.py new file mode 100644 index 0000000..69f6f28 --- /dev/null +++ b/strix/report/usage.py @@ -0,0 +1,234 @@ +"""SDK-native LLM usage aggregation for scan reports.""" + +from __future__ import annotations + +import logging +from typing import Any + +from agents.usage import Usage, deserialize_usage, serialize_usage + + +logger = logging.getLogger(__name__) + + +class LLMUsageLedger: + """Aggregate SDK ``Usage`` objects and attach best-effort cost estimates.""" + + def __init__(self) -> None: + self._total_usage = Usage() + self._agent_usage: dict[str, Usage] = {} + self._agent_metadata: dict[str, dict[str, str]] = {} + self._total_cost = 0.0 + self._agent_cost: dict[str, float] = {} + + def record( + self, + *, + agent_id: str, + usage: Usage | None, + agent_name: str | None = None, + model: str | None = None, + ) -> bool: + if usage is None or not _usage_has_activity(usage): + return False + + normalized_agent_id = str(agent_id or "unknown") + estimated_cost = _estimate_litellm_cost(usage, model) + + self._total_usage.add(usage) + self._agent_usage.setdefault(normalized_agent_id, Usage()).add(usage) + + metadata = self._agent_metadata.setdefault(normalized_agent_id, {}) + if agent_name: + metadata["agent_name"] = agent_name + if model: + metadata["model"] = model + + if estimated_cost is not None: + self._total_cost += estimated_cost + self._agent_cost[normalized_agent_id] = ( + self._agent_cost.get(normalized_agent_id, 0.0) + estimated_cost + ) + + return True + + def to_record(self) -> dict[str, Any]: + record = serialize_usage(self._total_usage) + record["cost"] = _round_cost(self._total_cost) + record["cost_source"] = "litellm_estimate" + record["agents"] = [] + + for agent_id in sorted(self._agent_usage): + usage = self._agent_usage[agent_id] + metadata = self._agent_metadata.get(agent_id, {}) + agent_record = serialize_usage(usage) + agent_record.update( + { + "agent_id": agent_id, + "agent_name": metadata.get("agent_name") or agent_id, + "model": metadata.get("model"), + "cost": _round_cost(self._agent_cost.get(agent_id, 0.0)), + "cost_source": "litellm_estimate", + } + ) + record["agents"].append(agent_record) + + return record + + def hydrate(self, raw_usage: Any) -> None: + self._total_usage = Usage() + self._agent_usage.clear() + self._agent_metadata.clear() + self._total_cost = 0.0 + self._agent_cost.clear() + + if not isinstance(raw_usage, dict): + return + + try: + self._total_usage = deserialize_usage(raw_usage) + except Exception: + logger.exception("Failed to hydrate aggregate llm_usage from run.json") + self._total_usage = Usage() + + self._total_cost = _float_or_zero(raw_usage.get("cost")) + agents = raw_usage.get("agents") or [] + if not isinstance(agents, list): + return + + for raw_agent in agents: + if not isinstance(raw_agent, dict): + continue + agent_id = str(raw_agent.get("agent_id") or "").strip() + if not agent_id: + continue + try: + self._agent_usage[agent_id] = deserialize_usage(raw_agent) + except Exception: + logger.exception("Failed to hydrate llm_usage for agent %s", agent_id) + self._agent_usage[agent_id] = Usage() + + metadata: dict[str, str] = {} + agent_name = raw_agent.get("agent_name") + model = raw_agent.get("model") + if isinstance(agent_name, str) and agent_name: + metadata["agent_name"] = agent_name + if isinstance(model, str) and model: + metadata["model"] = model + self._agent_metadata[agent_id] = metadata + self._agent_cost[agent_id] = _float_or_zero(raw_agent.get("cost")) + + +def _usage_has_activity(usage: Usage) -> bool: + return bool( + usage.requests + or usage.input_tokens + or usage.output_tokens + or usage.total_tokens + or usage.request_usage_entries + ) + + +def _estimate_litellm_cost(usage: Usage, model: str | None) -> float | None: + litellm_model = _litellm_model_name(model) + if not litellm_model: + return None + + entries = list(usage.request_usage_entries) + if not entries: + return _estimate_litellm_entry_cost(usage, litellm_model) + + total = 0.0 + estimated_any = False + for entry in entries: + cost = _estimate_litellm_entry_cost(entry, litellm_model) + if cost is None: + continue + total += cost + estimated_any = True + + return total if estimated_any else None + + +def _estimate_litellm_entry_cost(entry: Any, model: str) -> float | None: + prompt_tokens = _int_or_zero(getattr(entry, "input_tokens", 0)) + completion_tokens = _int_or_zero(getattr(entry, "output_tokens", 0)) + total_tokens = _int_or_zero(getattr(entry, "total_tokens", 0)) + if total_tokens <= 0: + total_tokens = prompt_tokens + completion_tokens + if total_tokens <= 0: + return None + + usage_payload: dict[str, Any] = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": total_tokens, + } + prompt_details = _details_to_dict(getattr(entry, "input_tokens_details", None)) + completion_details = _details_to_dict(getattr(entry, "output_tokens_details", None)) + if prompt_details: + usage_payload["prompt_tokens_details"] = prompt_details + if completion_details: + usage_payload["completion_tokens_details"] = completion_details + + try: + from litellm import completion_cost + + cost = completion_cost( + completion_response={ + "model": model.split("/", 1)[-1], + "usage": usage_payload, + }, + model=model, + ) + except Exception: # noqa: BLE001 - LiteLLM raises plain Exception for unknown model prices. + logger.debug("LiteLLM cost estimate unavailable for model %s", model, exc_info=True) + return None + + return cost if isinstance(cost, int | float) and cost >= 0 else None + + +def _litellm_model_name(model: str | None) -> str | None: + if not model: + return None + normalized = model.strip() + for prefix in ("litellm/", "any-llm/", "openai/"): + if normalized.startswith(prefix): + normalized = normalized.removeprefix(prefix) + break + return normalized or None + + +def _details_to_dict(details: Any) -> dict[str, Any]: + if details is None: + return {} + if isinstance(details, list): + for item in details: + result = _details_to_dict(item) + if result: + return result + return {} + if hasattr(details, "model_dump"): + return _details_to_dict(details.model_dump()) + if not isinstance(details, dict): + return {} + return {str(k): v for k, v in details.items() if v is not None} + + +def _int_or_zero(value: Any) -> int: + try: + return max(0, int(value or 0)) + except (TypeError, ValueError): + return 0 + + +def _float_or_zero(value: Any) -> float: + try: + result = float(value or 0.0) + except (TypeError, ValueError): + return 0.0 + return result if result >= 0 else 0.0 + + +def _round_cost(cost: float) -> float: + return round(max(0.0, cost), 10) diff --git a/strix/telemetry/posthog.py b/strix/telemetry/posthog.py index e8f4481..3a2a669 100644 --- a/strix/telemetry/posthog.py +++ b/strix/telemetry/posthog.py @@ -131,6 +131,20 @@ def end(report_state: "ReportState", exit_reason: str = "completed") -> None: except (ValueError, TypeError, AttributeError): pass + llm_props: dict[str, int | float] = {} + try: + usage = report_state.get_total_llm_usage() + if isinstance(usage, dict): + llm_props = { + "llm_requests": int(usage.get("requests") or 0), + "llm_input_tokens": int(usage.get("input_tokens") or 0), + "llm_output_tokens": int(usage.get("output_tokens") or 0), + "llm_tokens": int(usage.get("total_tokens") or 0), + "llm_cost": float(usage.get("cost") or 0.0), + } + except (TypeError, ValueError, AttributeError): + pass + _send( "scan_ended", { @@ -139,6 +153,7 @@ def end(report_state: "ReportState", exit_reason: str = "completed") -> None: "duration_seconds": round(duration), "vulnerabilities_total": len(report_state.vulnerability_reports), **{f"vulnerabilities_{k}": v for k, v in vulnerabilities_counts.items()}, + **llm_props, }, )