From 8ede419dccf6742aa0e0c4fe3e7faf11c471ff9a Mon Sep 17 00:00:00 2001 From: alex s <46074070+bearsyankees@users.noreply.github.com> Date: Mon, 17 Aug 2026 16:55:27 -0400 Subject: [PATCH] handle resume tokens gracefully (#1097) * Fix telemetry deltas for resumed runs * Fix resumed telemetry duration --- strix/report/state.py | 29 +++++++++++ strix/telemetry/posthog.py | 11 +---- strix/telemetry/scarf.py | 13 +---- tests/test_telemetry_resume.py | 89 ++++++++++++++++++++++++++++++++++ 4 files changed, 122 insertions(+), 20 deletions(-) create mode 100644 tests/test_telemetry_resume.py diff --git a/strix/report/state.py b/strix/report/state.py index ce7662b5..6f111178 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,25 @@ 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_process_duration_seconds(self) -> float: + """Return this process's elapsed wall time for telemetry.""" + try: + start = datetime.fromisoformat(self.process_start_time.replace("Z", "+00:00")) + duration = (datetime.now(start.tzinfo) - start).total_seconds() + return max(0.0, duration) + except (ValueError, TypeError, AttributeError): + return 0.0 + 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..756f163c 100644 --- a/strix/telemetry/posthog.py +++ b/strix/telemetry/posthog.py @@ -1,5 +1,4 @@ import logging -from datetime import datetime from typing import TYPE_CHECKING, Any import requests @@ -105,17 +104,11 @@ def end(report_state: "ReportState", exit_reason: str = "completed") -> None: if sev in vulnerabilities_counts: vulnerabilities_counts[sev] += 1 - duration = 0.0 - try: - start = datetime.fromisoformat(report_state.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): - pass + duration = report_state.get_process_duration_seconds() 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..22767424 100644 --- a/strix/telemetry/scarf.py +++ b/strix/telemetry/scarf.py @@ -2,7 +2,6 @@ from __future__ import annotations import logging import urllib.parse -from datetime import datetime from typing import TYPE_CHECKING, Any import requests @@ -114,19 +113,11 @@ def end(report_state: ReportState, exit_reason: str = "completed") -> None: if sev in vulnerabilities_counts: vulnerabilities_counts[sev] += 1 - duration = 0.0 - try: - scan_start = datetime.fromisoformat(report_state.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 - ).total_seconds() - except (ValueError, TypeError, AttributeError): - pass + duration = report_state.get_process_duration_seconds() 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..9a0f92fd --- /dev/null +++ b/tests/test_telemetry_resume.py @@ -0,0 +1,89 @@ +"""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.end_time = (datetime.now(UTC) - timedelta(hours=1)).isoformat() + initial.run_record["end_time"] = initial.end_time + 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) + + 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 0 <= 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)