From b1efd6ef1bdf58ea259051791a03d466881d24d4 Mon Sep 17 00:00:00 2001 From: Alex Schapiro Date: Mon, 17 Aug 2026 14:12:38 +0000 Subject: [PATCH] Fix telemetry deltas for resumed runs --- strix/report/state.py | 20 ++++++++ strix/telemetry/posthog.py | 4 +- strix/telemetry/scarf.py | 4 +- tests/test_telemetry_resume.py | 88 ++++++++++++++++++++++++++++++++++ 4 files changed, 112 insertions(+), 4 deletions(-) create mode 100644 tests/test_telemetry_resume.py diff --git a/strix/report/state.py b/strix/report/state.py index ce7662b5..1eff2850 100644 --- a/strix/report/state.py +++ b/strix/report/state.py @@ -39,6 +39,13 @@ def _strix_version() -> str | None: return None +def _number(value: Any) -> int | float: + try: + return float(value or 0) + except (TypeError, ValueError): + return 0 + + def _parse_repo_full_name(uri: str) -> str | None: """Extract ``owner/repo`` from a git URL or slug, else None.""" text = uri.strip().removesuffix(".git") @@ -115,6 +122,7 @@ class ReportState: self.run_name = run_name self.run_id = run_name or f"run-{uuid4().hex[:8]}" self.start_time = datetime.now(UTC).isoformat() + self.process_start_time = self.start_time self.end_time: str | None = None self.vulnerability_reports: list[dict[str, Any]] = [] @@ -123,6 +131,7 @@ class ReportState: self.scan_results: dict[str, Any] | None = None self.scan_config: dict[str, Any] | None = None self._llm_usage = LLMUsageLedger() + self._telemetry_llm_usage_baseline: dict[str, Any] = {} auth_mode = codex.auth_mode(load_settings().llm.model) self._llm_usage.zero_cost = auth_mode == "subscription" self.run_record: dict[str, Any] = { @@ -188,6 +197,7 @@ class ReportState: self.scan_results = scan_results self.final_scan_result = self._format_final_scan_result(scan_results) self._hydrate_llm_usage(data.get("llm_usage")) + self._telemetry_llm_usage_baseline = self._build_llm_usage_record() logger.info("report state hydrated run.json from %s", run_dir) json_path = run_dir / "vulnerabilities.json" @@ -331,6 +341,16 @@ class ReportState: def get_total_llm_usage(self) -> dict[str, Any]: return dict(self.run_record.get("llm_usage") or self._build_llm_usage_record()) + def get_process_llm_usage(self) -> dict[str, int | float]: + """Return LLM usage accumulated since this process started.""" + usage = self._llm_usage.to_record() + return { + key: max( + 0, _number(usage.get(key)) - _number(self._telemetry_llm_usage_baseline.get(key)) + ) + for key in ("requests", "input_tokens", "output_tokens", "total_tokens", "cost") + } + def get_total_llm_cost(self) -> float: """Live accumulated LLM cost, independent of the persisted run-record snapshot.""" return self._llm_usage.total_cost diff --git a/strix/telemetry/posthog.py b/strix/telemetry/posthog.py index 083e2c95..fa89fd0b 100644 --- a/strix/telemetry/posthog.py +++ b/strix/telemetry/posthog.py @@ -107,7 +107,7 @@ def end(report_state: "ReportState", exit_reason: str = "completed") -> None: duration = 0.0 try: - start = datetime.fromisoformat(report_state.start_time.replace("Z", "+00:00")) + start = datetime.fromisoformat(report_state.process_start_time.replace("Z", "+00:00")) end_iso = report_state.end_time or datetime.now(start.tzinfo).isoformat() duration = (datetime.fromisoformat(end_iso.replace("Z", "+00:00")) - start).total_seconds() except (ValueError, TypeError, AttributeError): @@ -115,7 +115,7 @@ def end(report_state: "ReportState", exit_reason: str = "completed") -> None: llm_props: dict[str, int | float] = {} try: - usage = report_state.get_total_llm_usage() + usage = report_state.get_process_llm_usage() if isinstance(usage, dict): llm_props = { "llm_requests": int(usage.get("requests") or 0), diff --git a/strix/telemetry/scarf.py b/strix/telemetry/scarf.py index 161e9980..495c0f65 100644 --- a/strix/telemetry/scarf.py +++ b/strix/telemetry/scarf.py @@ -116,7 +116,7 @@ def end(report_state: ReportState, exit_reason: str = "completed") -> None: duration = 0.0 try: - scan_start = datetime.fromisoformat(report_state.start_time.replace("Z", "+00:00")) + scan_start = datetime.fromisoformat(report_state.process_start_time.replace("Z", "+00:00")) end_iso = report_state.end_time or datetime.now(scan_start.tzinfo).isoformat() duration = ( datetime.fromisoformat(end_iso.replace("Z", "+00:00")) - scan_start @@ -126,7 +126,7 @@ def end(report_state: ReportState, exit_reason: str = "completed") -> None: llm_props: dict[str, int | float] = {} try: - usage = report_state.get_total_llm_usage() + usage = report_state.get_process_llm_usage() if isinstance(usage, dict): llm_props = { "llm_requests": int(usage.get("requests") or 0), diff --git a/tests/test_telemetry_resume.py b/tests/test_telemetry_resume.py new file mode 100644 index 00000000..60804bc7 --- /dev/null +++ b/tests/test_telemetry_resume.py @@ -0,0 +1,88 @@ +"""Regression tests for telemetry emitted by resumed runs.""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from typing import Any + +import pytest +from agents.usage import Usage + +from strix.report.state import ReportState +from strix.telemetry import posthog, scarf + + +def _usage(requests: int, input_tokens: int, output_tokens: int, total_tokens: int) -> Usage: + return Usage( + requests=requests, + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + ) + + +def _capture(sent: list[dict[str, Any]], props: dict[str, Any]) -> bool: + sent.append(props) + return True + + +@pytest.mark.parametrize("telemetry", [posthog, scarf]) +def test_scan_ended_reports_resumed_usage_delta( + telemetry: Any, + tmp_path: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + initial = ReportState(run_name="resumed") + initial.record_sdk_usage( + agent_id="agent", + usage=_usage(10, 1000, 200, 1200), + model="unknown", + ) + initial.record_observed_llm_cost(1.25) + initial.save_run_data() + + resumed = ReportState(run_name="resumed") + resumed.hydrate_from_run_dir() + resumed.record_sdk_usage( + agent_id="agent", + usage=_usage(3, 300, 50, 350), + model="unknown", + ) + resumed.record_observed_llm_cost(0.75) + resumed.end_time = (datetime.now(UTC) + timedelta(seconds=2)).isoformat() + + sent: list[dict[str, Any]] = [] + monkeypatch.setattr(telemetry, "_send", lambda _event, props: _capture(sent, props)) + telemetry.end(resumed) + + assert sent[0]["llm_requests"] == 3 + assert sent[0]["llm_input_tokens"] == 300 + assert sent[0]["llm_output_tokens"] == 50 + assert sent[0]["llm_tokens"] == 350 + assert sent[0]["llm_cost"] == pytest.approx(0.75) + assert sent[0]["duration_seconds"] <= 2 + + +@pytest.mark.parametrize("telemetry", [posthog, scarf]) +def test_scan_ended_reports_all_fresh_run_usage( + telemetry: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + state = ReportState() + state.record_sdk_usage( + agent_id="agent", + usage=_usage(3, 300, 50, 350), + model="unknown", + ) + state.record_observed_llm_cost(0.75) + + sent: list[dict[str, Any]] = [] + monkeypatch.setattr(telemetry, "_send", lambda _event, props: _capture(sent, props)) + telemetry.end(state) + + assert sent[0]["llm_requests"] == 3 + assert sent[0]["llm_input_tokens"] == 300 + assert sent[0]["llm_output_tokens"] == 50 + assert sent[0]["llm_tokens"] == 350 + assert sent[0]["llm_cost"] == pytest.approx(0.75)