diff --git a/docs/content/docs/configuration.mdx b/docs/content/docs/configuration.mdx index 8b6751793..a99a35a4a 100644 --- a/docs/content/docs/configuration.mdx +++ b/docs/content/docs/configuration.mdx @@ -19,6 +19,11 @@ headroom wrap codex --prepare-only Supported values are `rtk` and `lean-ctx`; unset defaults to `rtk`. +The proxy reads RTK lifetime savings with global scope by default so a shared +daemon reports savings across the operator's projects. Set +`HEADROOM_RTK_GAIN_SCOPE=project` to query `rtk gain --project` from the +proxy process working directory. + ## Modes | Mode | Behavior | Use Case | diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index baa0e868e..a3bbc91f6 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -592,6 +592,10 @@ def append_text_to_latest_user_input_item( _CONTEXT_TOOL_ENV = "HEADROOM_CONTEXT_TOOL" _CONTEXT_TOOL_RTK = "rtk" _CONTEXT_TOOL_LEAN_CTX = "lean-ctx" +_RTK_GAIN_SCOPE_ENV = "HEADROOM_RTK_GAIN_SCOPE" +_RTK_GAIN_SCOPE_GLOBAL = "global" +_RTK_GAIN_SCOPE_PROJECT = "project" +_RTK_GAIN_SCOPES = {_RTK_GAIN_SCOPE_GLOBAL, _RTK_GAIN_SCOPE_PROJECT} RTK_STATS_CACHE_TTL_SECONDS = float(os.environ.get("HEADROOM_CONTEXT_TOOL_STATS_TTL_SECONDS", "60")) CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS = RTK_STATS_CACHE_TTL_SECONDS @@ -988,6 +992,36 @@ def _context_tool_label(tool: str) -> str: return "RTK" +def _context_tool_default_scope(tool: str) -> str: + if tool == _CONTEXT_TOOL_LEAN_CTX: + return "local" + return _RTK_GAIN_SCOPE_GLOBAL + + +def _rtk_gain_scope() -> str: + raw = os.environ.get(_RTK_GAIN_SCOPE_ENV, "").strip().lower() + if not raw: + return _RTK_GAIN_SCOPE_GLOBAL + if raw in _RTK_GAIN_SCOPES: + return raw + + logger.warning( + "event=rtk_gain_scope_invalid env=%s value=%r default=%s", + _RTK_GAIN_SCOPE_ENV, + raw, + _RTK_GAIN_SCOPE_GLOBAL, + ) + return _RTK_GAIN_SCOPE_GLOBAL + + +def _rtk_gain_command(rtk_path: Any, scope: str) -> list[str]: + command = [str(rtk_path), "gain"] + if scope == _RTK_GAIN_SCOPE_PROJECT: + command.append("--project") + command.extend(["--format", "json"]) + return command + + def _coerce_int(value: Any, default: int = 0) -> int: try: return int(value or 0) @@ -1013,6 +1047,7 @@ def _context_tool_summary_payload( *, tool: str, installed: bool, + scope: str | None = None, summary: dict[str, Any] | None = None, ) -> dict[str, Any]: """Normalize RTK/lean-ctx lifetime gain output into one schema. @@ -1085,7 +1120,7 @@ def _context_tool_summary_payload( "tool": tool, "label": _context_tool_label(tool), "installed": installed, - "scope": "project" if tool == _CONTEXT_TOOL_RTK else "local", + "scope": scope or _context_tool_default_scope(tool), "total_commands": _coerce_int( _first_value( summary, @@ -1110,30 +1145,37 @@ def _context_tool_summary_payload( } +def _context_tool_zero_payload( + *, + tool: str, + installed: bool, + scope: str | None = None, +) -> dict[str, Any]: + return _context_tool_summary_payload( + tool=tool, + installed=installed, + scope=scope, + summary={}, + ) + + def _read_rtk_lifetime_stats() -> dict[str, Any] | None: - """Read rtk's current project-level lifetime stats.""" + """Read rtk's lifetime stats using the configured gain scope.""" from headroom.rtk import get_rtk_path + scope = _rtk_gain_scope() rtk_path = get_rtk_path() if not rtk_path: - return { - "tool": _CONTEXT_TOOL_RTK, - "label": _context_tool_label(_CONTEXT_TOOL_RTK), - "installed": False, - "scope": "project", - "total_commands": 0, - "input_tokens": 0, - "output_tokens": 0, - "tokens_saved": 0, - "avg_savings_pct": 0.0, - "lifetime_avg_savings_pct": 0.0, - "total_time_ms": 0, - } + return _context_tool_zero_payload( + tool=_CONTEXT_TOOL_RTK, + installed=False, + scope=scope, + ) try: result = subprocess.run( - [str(rtk_path), "gain", "--project", "--format", "json"], + _rtk_gain_command(rtk_path, scope), capture_output=True, text=True, timeout=5, @@ -1144,6 +1186,7 @@ def _read_rtk_lifetime_stats() -> dict[str, Any] | None: payload = _context_tool_summary_payload( tool=_CONTEXT_TOOL_RTK, installed=True, + scope=scope, summary=summary if isinstance(summary, dict) else {}, ) else: @@ -1157,19 +1200,11 @@ def _read_rtk_lifetime_stats() -> dict[str, Any] | None: result.returncode, stderr_excerpt, ) - return { - "tool": _CONTEXT_TOOL_RTK, - "label": _context_tool_label(_CONTEXT_TOOL_RTK), - "installed": True, - "scope": "project", - "total_commands": 0, - "input_tokens": 0, - "output_tokens": 0, - "tokens_saved": 0, - "avg_savings_pct": 0.0, - "lifetime_avg_savings_pct": 0.0, - "total_time_ms": 0, - } + return _context_tool_zero_payload( + tool=_CONTEXT_TOOL_RTK, + installed=True, + scope=scope, + ) except Exception as exc: # PR-G2 remediation (H2): log the exception path too. Reason is the # exception class name (without payload — RTK exceptions can carry @@ -1179,19 +1214,11 @@ def _read_rtk_lifetime_stats() -> dict[str, Any] | None: type(exc).__name__, exc, ) - return { - "tool": _CONTEXT_TOOL_RTK, - "label": _context_tool_label(_CONTEXT_TOOL_RTK), - "installed": True, - "scope": "project", - "total_commands": 0, - "input_tokens": 0, - "output_tokens": 0, - "tokens_saved": 0, - "avg_savings_pct": 0.0, - "lifetime_avg_savings_pct": 0.0, - "total_time_ms": 0, - } + return _context_tool_zero_payload( + tool=_CONTEXT_TOOL_RTK, + installed=True, + scope=scope, + ) return payload @@ -1203,33 +1230,9 @@ def _read_lean_ctx_lifetime_stats() -> dict[str, Any] | None: lean_ctx_path = get_lean_ctx_path() if not lean_ctx_path: - return { - "tool": _CONTEXT_TOOL_LEAN_CTX, - "label": _context_tool_label(_CONTEXT_TOOL_LEAN_CTX), - "installed": False, - "scope": "local", - "total_commands": 0, - "input_tokens": 0, - "output_tokens": 0, - "tokens_saved": 0, - "avg_savings_pct": 0.0, - "lifetime_avg_savings_pct": 0.0, - "total_time_ms": 0, - } + return _context_tool_zero_payload(tool=_CONTEXT_TOOL_LEAN_CTX, installed=False) - base_payload = { - "tool": _CONTEXT_TOOL_LEAN_CTX, - "label": _context_tool_label(_CONTEXT_TOOL_LEAN_CTX), - "installed": True, - "scope": "local", - "total_commands": 0, - "input_tokens": 0, - "output_tokens": 0, - "tokens_saved": 0, - "avg_savings_pct": 0.0, - "lifetime_avg_savings_pct": 0.0, - "total_time_ms": 0, - } + base_payload = _context_tool_zero_payload(tool=_CONTEXT_TOOL_LEAN_CTX, installed=True) try: result = subprocess.run( diff --git a/tests/test_proxy_dashboard_stats_cache.py b/tests/test_proxy_dashboard_stats_cache.py index e4c0fc467..f7f2fa554 100644 --- a/tests/test_proxy_dashboard_stats_cache.py +++ b/tests/test_proxy_dashboard_stats_cache.py @@ -5,6 +5,7 @@ import shutil import subprocess from pathlib import Path from types import SimpleNamespace +from unittest.mock import MagicMock import pytest @@ -31,6 +32,7 @@ class _ToinStub: @pytest.fixture(autouse=True) def _reset_rtk_stats_cache(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + monkeypatch.delenv("HEADROOM_RTK_GAIN_SCOPE", raising=False) monkeypatch.setenv("HEADROOM_REQUIRE_RUST_CORE", "false") proxy_helpers._rtk_stats_cache.update( {"expires_at": 0.0, "has_value": False, "tool": None, "value": None} @@ -72,8 +74,9 @@ def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch }, ] - def _fake_run(*args, **kwargs): + def _fake_run(args, **kwargs): calls["run"] += 1 + assert args == ["/usr/bin/rtk", "gain", "--format", "json"] summary = totals[min(calls["run"] - 1, len(totals) - 1)] return SimpleNamespace( returncode=0, @@ -91,6 +94,7 @@ def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch assert first["tool"] == "rtk" assert first["label"] == "RTK" assert first["installed"] is True + assert first["scope"] == "global" assert first["total_commands"] == 0 assert first["input_tokens"] == 0 assert first["output_tokens"] == 0 @@ -143,6 +147,64 @@ def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch assert calls["run"] == 2 +def test_get_rtk_stats_can_read_project_scoped_gain(monkeypatch: pytest.MonkeyPatch) -> None: + calls = {"run": 0} + + def _fake_run(args, **kwargs): + calls["run"] += 1 + assert args == ["/usr/bin/rtk", "gain", "--project", "--format", "json"] + return SimpleNamespace( + returncode=0, + stdout=json.dumps( + { + "summary": { + "total_commands": 1, + "total_input": 100, + "total_output": 75, + "total_saved": 25, + } + } + ), + ) + + monkeypatch.setenv("HEADROOM_RTK_GAIN_SCOPE", "project") + monkeypatch.setattr(shutil, "which", lambda name: "/usr/bin/rtk") + monkeypatch.setattr(subprocess, "run", _fake_run) + + payload = proxy_helpers._read_rtk_lifetime_stats() + + assert payload is not None + assert payload["scope"] == "project" + assert payload["total_commands"] == 1 + assert payload["tokens_saved"] == 25 + assert calls["run"] == 1 + + +def test_get_rtk_stats_invalid_scope_defaults_to_global( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = {"run": 0} + + def _fake_run(args, **kwargs): + calls["run"] += 1 + assert args == ["/usr/bin/rtk", "gain", "--format", "json"] + return SimpleNamespace(returncode=0, stdout=json.dumps({"summary": {}})) + + mock_warning = MagicMock() + monkeypatch.setenv("HEADROOM_RTK_GAIN_SCOPE", "workspace") + monkeypatch.setattr(proxy_helpers.logger, "warning", mock_warning) + monkeypatch.setattr(shutil, "which", lambda name: "/usr/bin/rtk") + monkeypatch.setattr(subprocess, "run", _fake_run) + + payload = proxy_helpers._read_rtk_lifetime_stats() + + assert payload is not None + assert payload["scope"] == "global" + assert calls["run"] == 1 + warning_calls = " ".join(str(call) for call in mock_warning.call_args_list) + assert "event=rtk_gain_scope_invalid" in warning_calls + + def test_get_context_tool_stats_reads_lean_ctx_gain(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "lean-ctx") now = {"value": 100.0} diff --git a/tests/test_subscription_tracker_rtk_wired.py b/tests/test_subscription_tracker_rtk_wired.py index 6ed116730..d60730986 100644 --- a/tests/test_subscription_tracker_rtk_wired.py +++ b/tests/test_subscription_tracker_rtk_wired.py @@ -580,6 +580,7 @@ def test_rtk_subprocess_failure_logs_structured_warning( payload = _helpers._read_rtk_lifetime_stats() assert payload is not None + assert payload["scope"] == "global" assert payload["tokens_saved"] == 0 # Concatenate all warning call args so the failure message shows what