mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(memory): bound the TrafficLearner pending-pattern accumulator (memory leak) (#2579)
## Description
`TrafficLearner` (the memory/learning subsystem that accumulates
patterns from proxy traffic) has an unbounded in-memory accumulator.
`_pattern_counts` maps `content_hash -> (pattern, count)`. A pattern is
added on first sighting, its count is bumped on each re-sighting, and it
is **removed only when it reaches `min_evidence`** (default 5), at which
point it is promoted and its hash moves to `_saved_hashes`:
```python
if h in self._pattern_counts:
existing, count = self._pattern_counts[h]
count += 1
self._pattern_counts[h] = (existing, count)
else:
self._pattern_counts[h] = (pattern, 1)
return # first sighting — wait for more evidence
...
if count >= self._min_evidence:
del self._pattern_counts[h] # only removal path
self._saved_hashes.add(h)
if len(self._saved_hashes) > self._dedup_window: # sibling IS trimmed
self._saved_hashes.pop()
```
A pattern seen **once but never corroborated** — the common case for
one-off traffic (a unique error string, an ad-hoc shell command, a
distinct file path) — never reaches `min_evidence`, so it is **never
removed**. Over a long-lived proxy processing varied traffic,
`_pattern_counts` grows without bound and RSS climbs. The sibling
`_saved_hashes` is explicitly trimmed to `dedup_window` ("prevent
unbounded growth"); `_pattern_counts` was missed.
Reproduced directly: feeding 500 distinct one-off patterns leaves 500
entries in `_pattern_counts` (one per pattern, forever).
## Fix
Make `_pattern_counts` an LRU-ordered `OrderedDict` capped at a new
`max_pending_patterns` (default 2048):
- On each corroboration, `move_to_end(h)` so an actively-accumulating
pattern stays "fresh" and is never evicted before it can be promoted.
- On a first sighting when the accumulator is full, evict the
least-recently-corroborated pending entry (`popitem(last=False)`).
Evicting a stale one-off is safe: if it recurs it simply restarts
accumulation (delayed promotion at worst) — the same tradeoff
`_saved_hashes` already makes. Promotion at `min_evidence` is unchanged,
and the cap (2048) is generous enough that any pattern receiving repeat
sightings within a normal window reaches `min_evidence=5` long before
eviction.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature
- [ ] Breaking change
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring
## Changes Made
- `headroom/memory/traffic_learner.py`: `_pattern_counts` becomes a
capped LRU `OrderedDict`; add `max_pending_patterns` (default 2048);
`move_to_end` on corroboration and evict-oldest on overflow.
- `tests/test_memory/test_traffic_learner.py`: a regression that 500
one-off patterns keep the accumulator at its cap, and one that a
corroborated pattern still promotes into `_saved_hashes` (both sync via
`asyncio.run` so they run without the pytest-asyncio plugin).
## Testing
- [x] 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
$ python -m pytest tests/test_memory/test_traffic_learner.py -q
35 failed, 109 passed
# the 35 failures are pre-existing @pytest.mark.asyncio tests that need
# pytest-asyncio (not configured in this environment); they fail identically
# on clean main (35 failed, 107 passed) and pass in CI. My two new tests are
# synchronous and pass; they add +2 passing with no new failures.
# with the fix reverted, test_pending_accumulator_is_bounded fails
# (the accumulator holds all 500 one-off patterns)
$ uvx ruff@0.15.17 check headroom/memory/traffic_learner.py tests/test_memory/test_traffic_learner.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/memory/traffic_learner.py
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, project venv (`uv sync --extra
proxy`), `uvx ruff@0.15.17` / `uvx mypy@1.20.2`, pytest in the venv.
- Exact command / steps: built a `TrafficLearner(backend=None,
min_evidence=5, max_pending_patterns=8)` and drove `_accumulate` with
500 distinct one-off `ExtractedPattern`s; separately corroborated one
pattern to `min_evidence`; then reverted the source and re-ran.
- Observed result: with the fix `len(_pattern_counts)` stays at the cap
(8) after 500 one-offs, the corroborated pattern is removed from pending
and present in `_saved_hashes`, and an actively-bumped pattern survives
LRU eviction; with the fix reverted the accumulator holds all 500
one-off entries (the unbounded leak). Ran against the actual module.
- Not tested: a live multi-day proxy run measuring RSS (the leak is
inferred from the removed unbounded-growth path; the accumulator bound
is verified directly).
## 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
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable
This commit is contained in:
parent
b97c7c6e99
commit
1f5fefffd3
2 changed files with 67 additions and 2 deletions
|
|
@ -26,6 +26,7 @@ import os
|
|||
import re
|
||||
import sqlite3
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
|
|
@ -450,6 +451,7 @@ class TrafficLearner:
|
|||
max_history: int = 20,
|
||||
dedup_window: int = 100,
|
||||
min_evidence: int = 5,
|
||||
max_pending_patterns: int = 2048,
|
||||
) -> None:
|
||||
"""Initialize the traffic learner.
|
||||
|
||||
|
|
@ -468,12 +470,19 @@ class TrafficLearner:
|
|||
self.agent_type = agent_type
|
||||
self._max_history = max_history
|
||||
self._min_evidence = min_evidence
|
||||
self._max_pending_patterns = max_pending_patterns
|
||||
|
||||
# Recent tool call history for error→recovery matching
|
||||
self._tool_history: list[dict[str, Any]] = []
|
||||
|
||||
# Pattern accumulator: hash → (pattern, count)
|
||||
self._pattern_counts: dict[str, tuple[ExtractedPattern, int]] = {}
|
||||
# Pattern accumulator: hash → (pattern, count). LRU-ordered and capped:
|
||||
# a pattern that is seen once but never reaches ``min_evidence`` would
|
||||
# otherwise linger here forever, so this dict grew unbounded over a
|
||||
# long-lived proxy's traffic (the sibling ``_saved_hashes`` is trimmed
|
||||
# to ``dedup_window`` for the same reason; this one was missed). Evicting
|
||||
# the least-recently-corroborated pending pattern is safe: if it recurs
|
||||
# it simply restarts accumulation.
|
||||
self._pattern_counts: OrderedDict[str, tuple[ExtractedPattern, int]] = OrderedDict()
|
||||
|
||||
# Dedup: hashes of patterns already saved to DB
|
||||
self._saved_hashes: set[str] = set()
|
||||
|
|
@ -1250,7 +1259,13 @@ class TrafficLearner:
|
|||
existing, count = self._pattern_counts[h]
|
||||
count += 1
|
||||
self._pattern_counts[h] = (existing, count)
|
||||
# Mark as most-recently-corroborated so it survives LRU eviction.
|
||||
self._pattern_counts.move_to_end(h)
|
||||
else:
|
||||
# Bound the pending accumulator so one-off patterns can't grow it
|
||||
# without limit; drop the least-recently-corroborated pending entry.
|
||||
if len(self._pattern_counts) >= self._max_pending_patterns:
|
||||
self._pattern_counts.popitem(last=False)
|
||||
self._pattern_counts[h] = (pattern, 1)
|
||||
return # First sighting — wait for more evidence
|
||||
|
||||
|
|
|
|||
|
|
@ -362,6 +362,56 @@ class TestTrafficLearner:
|
|||
stats = learner.get_stats()
|
||||
assert stats["patterns_extracted"] >= 3
|
||||
|
||||
def test_pending_accumulator_is_bounded(self):
|
||||
"""One-off patterns that never reach ``min_evidence`` must not grow the
|
||||
pending ``_pattern_counts`` accumulator without bound — the sibling
|
||||
``_saved_hashes`` is already trimmed to ``dedup_window`` and this one was
|
||||
missed, so a long-lived proxy leaked memory across varied traffic. It is
|
||||
now LRU-capped at ``max_pending_patterns``.
|
||||
|
||||
Sync test (drives the async accumulate via ``asyncio.run``) so it runs
|
||||
without the pytest-asyncio plugin.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
learner = TrafficLearner(backend=None, min_evidence=5, max_pending_patterns=8)
|
||||
|
||||
async def feed_one_offs() -> None:
|
||||
for i in range(500):
|
||||
await learner._accumulate(
|
||||
ExtractedPattern(
|
||||
category=PatternCategory.PREFERENCE,
|
||||
content=f"one-off pattern number {i}",
|
||||
importance=0.5,
|
||||
)
|
||||
)
|
||||
|
||||
asyncio.run(feed_one_offs())
|
||||
assert len(learner._pattern_counts) <= 8 # capped, not 500
|
||||
|
||||
def test_pending_accumulator_lru_still_promotes_corroborated_pattern(self):
|
||||
"""Capping the accumulator must not break promotion: a pattern
|
||||
corroborated to ``min_evidence`` without interruption is still removed
|
||||
from pending and recorded in ``_saved_hashes``."""
|
||||
import asyncio
|
||||
|
||||
learner = TrafficLearner(backend=None, min_evidence=3, max_pending_patterns=100)
|
||||
pattern = ExtractedPattern(
|
||||
category=PatternCategory.PREFERENCE,
|
||||
content="corroborated preference",
|
||||
importance=0.5,
|
||||
)
|
||||
|
||||
async def corroborate() -> None:
|
||||
await learner._accumulate(pattern) # count 1
|
||||
await learner._accumulate(pattern) # count 2
|
||||
assert pattern.content_hash in learner._pattern_counts
|
||||
await learner._accumulate(pattern) # count 3 == min_evidence -> promote
|
||||
|
||||
asyncio.run(corroborate())
|
||||
assert pattern.content_hash not in learner._pattern_counts # removed on promotion
|
||||
assert pattern.content_hash in learner._saved_hashes
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dedup(self, learner: TrafficLearner):
|
||||
"""Test that identical patterns are deduplicated."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue