akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy): read RTK gain stats globally by default (#957)

Browse source 
## Description

Closes #900.

The proxy now reads RTK lifetime savings with global scope by default.
This matches shared daemon deployments where the proxy process cwd is
often `$HOME` or a service directory, while RTK savings are accumulated
across the operator's projects.

`HEADROOM_RTK_GAIN_SCOPE=project` keeps the old `rtk gain --project`
behavior for operators who explicitly want the proxy process working
directory as the scope.

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [x] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- Default RTK stats subprocess command to `rtk gain --format json`
- Add `HEADROOM_RTK_GAIN_SCOPE=project` to opt into `rtk gain --project
--format json`
- Keep fallback/synthetic-zero payload `scope` aligned with the queried
scope
- Deduplicate context-tool zero payload construction
- Document the new RTK gain scope environment variable

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
uv run --extra dev python -m pytest tests/test_proxy_dashboard_stats_cache.py tests/test_subscription_tracker_rtk_wired.py -q
============================= test session starts ==============================
platform darwin -- Python 3.14.5, pytest-9.0.3, pluggy-1.6.0
rootdir: /Users/joshuasiu/vibe/temp/headroom
configfile: pyproject.toml
plugins: anyio-4.12.1, langsmith-0.8.0, asyncio-1.3.0, cov-7.0.0
asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 26 items

tests/test_proxy_dashboard_stats_cache.py ..........                     [ 38%]
tests/test_subscription_tracker_rtk_wired.py ................            [100%]

============================== 26 passed in 0.40s ==============================

uv run --extra dev python -m pytest tests/test_proxy_stats_recent_requests.py tests/test_proxy_healthchecks.py -q
============================= test session starts ==============================
platform darwin -- Python 3.14.5, pytest-9.0.3, pluggy-1.6.0
rootdir: /Users/joshuasiu/vibe/temp/headroom
configfile: pyproject.toml
plugins: anyio-4.12.1, langsmith-0.8.0, asyncio-1.3.0, cov-7.0.0
asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 15 items

tests/test_proxy_stats_recent_requests.py ...                            [ 20%]
tests/test_proxy_healthchecks.py ............                            [100%]

============================= 15 passed in 10.41s ==============================

uv run --extra dev ruff check headroom/proxy/helpers.py tests/test_proxy_dashboard_stats_cache.py tests/test_subscription_tracker_rtk_wired.py
All checks passed!

uv run --extra dev ruff format --check headroom/proxy/helpers.py tests/test_proxy_dashboard_stats_cache.py tests/test_subscription_tracker_rtk_wired.py
3 files already formatted

uv run --extra dev mypy headroom
headroom/proxy/server.py:1141: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
headroom/proxy/server.py:1211: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
headroom/proxy/server.py:1215: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
Success: no issues found in 358 source files
```

## Real Behavior Proof

- Environment: macOS, Python 3.14.5 via `uv run --extra dev`
- Exact command / steps: mocked RTK subprocess argv in unit tests
- Observed result: default command is `rtk gain --format json`; project
scope command is `rtk gain --project --format json`; invalid scope logs
`event=rtk_gain_scope_invalid` and falls back to global
- Not tested: full repository pytest suite

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable

## Additional Notes

The unchecked comment and changelog items are not applicable for this
scoped proxy stats fix.
This commit is contained in:
skblue 2026-06-14 12:21:38 +08:00 • committed by GitHub
parent b51cda10d7
commit b70fccbe17
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 140 additions and 69 deletions
Download patch file Download diff file Expand all files Collapse all files

5
docs/content/docs/configuration.mdx
View file

@ -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 |

139
headroom/proxy/helpers.py
View file

@ -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(

64
tests/test_proxy_dashboard_stats_cache.py
View file

@ -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}

1
tests/test_subscription_tracker_rtk_wired.py
View file

@ -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