mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
3 commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
Abhay Singh
|
29d8a5e563
|
fix(learn/gemini): stop double-counting session tokens (#2230)
## Description
The Gemini `learn` scanner inflates every session's token totals by
double-counting.
In `_parse_messages` the per-message usage accumulation is:
```python
usage = msg.get("usageMetadata", msg.get("usage", {}))
if isinstance(usage, dict):
total_input_tokens += usage.get("promptTokenCount", 0)
total_input_tokens += usage.get("cachedContentTokenCount", 0)
total_output_tokens += usage.get("candidatesTokenCount", 0)
total_output_tokens += (
usage.get("totalTokenCount", 0) - usage.get("promptTokenCount", 0)
if usage.get("totalTokenCount")
else 0
)
```
Both additions on each side double-count, per Gemini's `usageMetadata`
semantics:
- `cachedContentTokenCount` is the cached **subset** of
`promptTokenCount`, not tokens on top of it. Adding both counts the
cached input twice.
- `totalTokenCount == promptTokenCount + candidatesTokenCount`, so
`totalTokenCount - promptTokenCount` is just `candidatesTokenCount`
again. Adding it on top of `candidatesTokenCount` counts the output
twice.
For a turn with 1000 prompt tokens (300 cached) and 500 output tokens
(`totalTokenCount` 1500), the scanner records input 1300 and output 1000
instead of 1000 / 500 — so both totals are materially inflated for any
Gemini session that carries usage metadata.
## Fix
Count the prompt as input and the candidates as output, once each:
```python
total_input_tokens += usage.get("promptTokenCount", 0)
total_output_tokens += usage.get("candidatesTokenCount", 0)
```
Closes #
## 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)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
- `headroom/learn/plugins/gemini.py`: drop the `cachedContentTokenCount`
and `totalTokenCount - promptTokenCount` additions in `_parse_messages`.
- `tests/test_learn/test_gemini_scanner.py`: new test asserting the
input/output totals equal `promptTokenCount` / `candidatesTokenCount`.
- `CHANGELOG.md`: Bug Fixes entry.
## Testing
- [ ] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ uvx ruff@0.15.17 check headroom/learn/plugins/gemini.py tests/test_learn/test_gemini_scanner.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/learn/plugins/gemini.py
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx
mypy@1.20.2`. A full `pytest` OOM-kills this box (ML stack import), so I
reproduced the arithmetic with a dependency-free script and left the
full pytest to CI.
- Exact command / steps: fed a usage dict of `promptTokenCount=1000,
cachedContentTokenCount=300, candidatesTokenCount=500,
totalTokenCount=1500` through the OLD accumulation and the NEW one.
- Observed result: OLD → input 1300, output 1000 (cached and candidates
both counted twice); NEW → input 1000, output 500 (the true figures).
- Not tested: a full `learn` run over a real Gemini history; full local
`pytest` deferred to CI (OOM).
## 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
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] 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
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
The "unit tests pass locally" box is unchecked because the full suite
imports the ML stack, which I can't run here. The new test uses the
existing `GeminiScanner` harness in
`tests/test_learn/test_gemini_scanner.py` so it runs under the normal CI
pytest job; behaviour is additionally verified by the standalone proof
above.
---------
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
|
||
|
Abhay Singh
|
7e83b8da3c
|
fix(learn/gemini): detect the project path for JSONL sessions (#2229)
## Description
The Gemini `learn` plugin can't detect the project path for JSONL
sessions, so it writes its insights to the wrong project.
`discover_projects` globs both `session-*.json` and `session-*.jsonl`
and calls `_detect_project_path`, which reads the file with a single
whole-file `json.load`:
```python
def _detect_project_path(self, session_path: Path) -> Path | None:
try:
with open(session_path, encoding="utf-8", errors="replace") as f:
data = json.load(f)
except (OSError, json.JSONDecodeError):
return None
...
```
A `.jsonl` session is one JSON object per line, so `json.load` on the
whole file raises `json.JSONDecodeError` ("Extra data") on the second
line. The method swallows that and returns `None`, and the caller falls
back to `Path.cwd()`:
```python
project_path = self._detect_project_path(session_files[0])
...
ProjectInfo(
name=project_path.name if project_path else project_dir.name,
project_path=project_path or Path.cwd(), # wrong project
context_file=gemini_md, # None: GEMINI.md never found
...
)
```
So for the JSONL format (Gemini CLI's newer session format — the one
that carries `type: "session_metadata"` records), detection never works:
the learned tool/verbosity insights are attributed to the current
working directory instead of the real project, and the project's
`GEMINI.md` is never located. The sibling `_scan_jsonl_session` already
reads this format line-by-line, and the Claude plugin recovers the
project path from session `cwd` the same way.
## Fix
Route `.jsonl` sessions through a line-by-line reader and share the
field extraction (`projectPath` / `project_path` / `cwd` /
`workingDirectory`) between both formats:
```python
if session_path.suffix == ".jsonl":
return self._detect_project_path_jsonl(session_path)
```
`_detect_project_path_jsonl` parses each line (skipping blanks and
unparseable lines, exactly like `_scan_jsonl_session`) and returns the
first record that yields an existing path. The JSON path is unchanged.
Closes #
## 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)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
- `headroom/learn/plugins/gemini.py`: dispatch `.jsonl` sessions to a
new line-by-line `_detect_project_path_jsonl`; factor the field
extraction into `_project_path_from_entry` shared by both paths.
- `tests/test_learn/test_gemini_scanner.py`: new test asserting a JSONL
session's `cwd` is recovered.
- `CHANGELOG.md`: Bug Fixes entry.
## Testing
- [ ] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ uvx ruff@0.15.17 check headroom/learn/plugins/gemini.py tests/test_learn/test_gemini_scanner.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/learn/plugins/gemini.py
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx
mypy@1.20.2`. A full `pytest` OOM-kills this box (ML stack import), so I
reproduced the behavior with a dependency-free script mirroring both
detection paths and left the full pytest to CI.
- Exact command / steps: wrote a `.jsonl` session whose first record is
`{"type":"session_metadata","cwd":"<project>"}`, then ran the OLD
whole-file `json.load` reader and the NEW line-by-line reader; also
checked a single-object `.json` session still resolves under both.
- Observed result: OLD returns `None` for the JSONL file (the caller
would fall back to cwd); NEW returns the project path; the `.json` case
resolves identically under both.
- Not tested: a full `learn` run over a real Gemini history; full local
`pytest` deferred to CI (OOM).
## 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
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] 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
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
The "unit tests pass locally" box is unchecked because the full suite
imports the ML stack, which I can't run here. The new test uses the
existing `GeminiScanner` harness in
`tests/test_learn/test_gemini_scanner.py` so it runs under the normal CI
pytest job; behaviour is additionally verified by the standalone proof
above.
---------
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
|
||
|
chopratejas
|
446872ed53 |
Plugin architecture for headroom learn + live traffic flush
Refactor headroom learn into a plugin architecture where each coding agent (Claude Code, Codex, Gemini CLI) is a self-contained plugin with scanner, writer, and detection logic. External plugins can register via the headroom.learn_plugin entry point. - Add LearnPlugin ABC (base.py) and plugin registry (registry.py) - Move scanners from monolithic scanner.py into plugins/ directory - Extract shared error classification and tool name map (_shared.py) - Add GeminiScanner for Google Gemini CLI session parsing - CLI uses dynamic agent detection via registry (no hardcoded choices) - All existing imports preserved via backwards-compat re-exports - Wire agent_type through wrap → proxy → TrafficLearner - Flush learned patterns to correct .md file at proxy shutdown - Fix shutdown queue drain bug (patterns were lost on exit) - 97 tests pass (84 existing + 13 new registry/plugin tests) |