From aa2036f22942eed3ecd142169289a256a2e7acb7 Mon Sep 17 00:00:00 2001 From: Ian Date: Tue, 20 Oct 2026 11:08:44 -0400 Subject: [PATCH] fix(llm): keep request_log.install(); record tool call timings; both only with STRIX_RUN_RECORD_TIMINGS --- docs/advanced/configuration.mdx | 4 ++++ strix/config/models.py | 4 +++- strix/config/settings.py | 1 + strix/core/hooks.py | 39 ++++++++++++++++++++++++++++++++- strix/core/runner.py | 1 + strix/report/state.py | 17 ++++++++++++++ tests/test_cost_tracking.py | 20 +++++++++++++++++ 7 files changed, 84 insertions(+), 2 deletions(-) diff --git a/docs/advanced/configuration.mdx b/docs/advanced/configuration.mdx index 29b003ec..f14b110f 100644 --- a/docs/advanced/configuration.mdx +++ b/docs/advanced/configuration.mdx @@ -141,6 +141,10 @@ affecting the agents that do the actual testing. Telemetry toggle. Set to `0`, `false`, `no`, or `off` to disable telemetry (PostHog, Scarf, OTEL). + + Record every LLM request's and tool call's timing in `run.json`, under `llm_requests` and `tool_calls`. Adds about 0.5 MB to a typical scan's `run.json`. + + OTLP/Traceloop base URL for remote OpenTelemetry export. If unset, Strix keeps traces local only. diff --git a/strix/config/models.py b/strix/config/models.py index 8a11695f..d9127ee8 100644 --- a/strix/config/models.py +++ b/strix/config/models.py @@ -707,7 +707,9 @@ def configure_sdk_model_defaults(settings: Settings) -> None: set_tracing_disabled(True) from strix.report.state import record_llm_request - request_log.register_sink(record_llm_request) + request_log.install() + if settings.telemetry.run_record_timings: + request_log.register_sink(record_llm_request) if codex.subscription_model(llm.model): return _configure_litellm_compatibility() diff --git a/strix/config/settings.py b/strix/config/settings.py index b965281a..22e6c596 100644 --- a/strix/config/settings.py +++ b/strix/config/settings.py @@ -163,6 +163,7 @@ class TelemetrySettings(BaseSettings): model_config = _BASE_CONFIG enabled: bool = Field(default=True, alias="STRIX_TELEMETRY") + run_record_timings: bool = Field(default=False, alias="STRIX_RUN_RECORD_TIMINGS") WebSearchProvider = Literal["auto", "perplexity", "exa"] diff --git a/strix/core/hooks.py b/strix/core/hooks.py index 13fc98f9..daddda1b 100644 --- a/strix/core/hooks.py +++ b/strix/core/hooks.py @@ -4,18 +4,21 @@ from __future__ import annotations import logging import math +import time +from datetime import UTC, datetime from typing import TYPE_CHECKING, Any from agents.lifecycle import RunHooks from strix.core.agents import BudgetPolicy, coordinator_from_context -from strix.report.state import get_global_report_state +from strix.report.state import get_global_report_state, record_tool_call if TYPE_CHECKING: from agents import RunContextWrapper from agents.agent import Agent from agents.items import ModelResponse, TResponseInputItem + from agents.tool import Tool logger = logging.getLogger(__name__) @@ -163,6 +166,7 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]): max_turns: int | None = None, interactive: bool = False, budget_policy: BudgetPolicy = "stop", + record_tool_timings: bool = False, ) -> None: _validate_budget(max_budget_usd) if max_turns is not None and max_turns <= 0: @@ -175,6 +179,8 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]): self._max_turns = max_turns self._interactive = interactive self._budget_policy: BudgetPolicy = budget_policy + self._record_tool_timings = record_tool_timings + self._tool_starts: dict[str, tuple[datetime, float]] = {} @property def max_budget_usd(self) -> float | None: @@ -315,6 +321,37 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]): ) input_items.append({"role": "user", "content": content}) + async def on_tool_start( + self, + context: RunContextWrapper[dict[str, Any]], + agent: Agent[dict[str, Any]], # noqa: ARG002 + tool: Tool, # noqa: ARG002 + ) -> None: + call_id = getattr(context, "tool_call_id", None) + if self._record_tool_timings and isinstance(call_id, str): + self._tool_starts[call_id] = (datetime.now(UTC), time.monotonic()) + + async def on_tool_end( + self, + context: RunContextWrapper[dict[str, Any]], + agent: Agent[dict[str, Any]], # noqa: ARG002 + tool: Tool, + result: object, # noqa: ARG002 + ) -> None: + call_id = getattr(context, "tool_call_id", None) + start = self._tool_starts.pop(call_id, None) if isinstance(call_id, str) else None + if start is None: + return + started_at, started = start + ctx = context.context if isinstance(context.context, dict) else {} + agent_id = ctx.get("agent_id") + record_tool_call( + started_at=started_at, + agent_id=agent_id if isinstance(agent_id, str) else None, + tool=tool.name, + duration_ms=round((time.monotonic() - started) * 1000), + ) + async def on_llm_end( self, context: RunContextWrapper[dict[str, Any]], diff --git a/strix/core/runner.py b/strix/core/runner.py index 7c9c1d98..ae6dc08b 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -419,6 +419,7 @@ async def run_strix_scan( max_turns=max_turns, interactive=interactive, budget_policy=budget_policy, + record_tool_timings=settings.telemetry.run_record_timings, ) coordinator.set_budget_limit_setter(hooks.set_max_budget_usd) if interactive and budget_policy != "pause": diff --git a/strix/report/state.py b/strix/report/state.py index 5daf2849..104e4af4 100644 --- a/strix/report/state.py +++ b/strix/report/state.py @@ -1045,6 +1045,23 @@ def record_llm_request(event: "LlmRequestEvent") -> None: ) +def record_tool_call( + *, started_at: datetime, agent_id: str | None, tool: str, duration_ms: int +) -> None: + """Each tool call's timing, kept in run.json's ``tool_calls``.""" + report_state = get_global_report_state() + if report_state is None: + return + report_state.run_record.setdefault("tool_calls", []).append( + { + "started_at": started_at.isoformat(), + "agent_id": agent_id, + "tool": tool, + "duration_ms": duration_ms, + } + ) + + def record_openrouter_provider(provider: Any, usage: Any) -> None: """Tally which upstream provider served a stream, from its final usage chunk. diff --git a/tests/test_cost_tracking.py b/tests/test_cost_tracking.py index 66eb040d..4f12bc0e 100644 --- a/tests/test_cost_tracking.py +++ b/tests/test_cost_tracking.py @@ -19,6 +19,7 @@ from strix.config.models import ( _configure_litellm_compatibility, _install_openrouter_stream_cost_capture, ) +from strix.core.hooks import ReportUsageHooks from strix.llm import request_log from strix.report.state import ( ReportState, @@ -508,3 +509,22 @@ def test_request_timings_land_in_run_record() -> None: "retry_attempt": 0, } ] + + +async def test_tool_call_timings_land_in_run_record() -> None: + hooks = ReportUsageHooks(model="m", record_tool_timings=True) + context = SimpleNamespace(tool_call_id="call-1", context={"agent_id": "agent-1"}) + tool = SimpleNamespace(name="terminal_execute") + state = ReportState.__new__(ReportState) + state.run_record = {} + try: + set_global_report_state(state) + await hooks.on_tool_start(context, MagicMock(), tool) # type: ignore[arg-type] + await hooks.on_tool_end(context, MagicMock(), tool, "ok") # type: ignore[arg-type] + finally: + set_global_report_state(None) + + [entry] = state.run_record["tool_calls"] + assert entry["agent_id"] == "agent-1" + assert entry["tool"] == "terminal_execute" + assert entry["duration_ms"] >= 0