mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(learn): run project discovery off the event loop (#2731)
## Description
`TrafficLearner.flush_to_file` is a coroutine, but it called
`plugin.discover_projects()` inline. That function walks the filesystem
to decode escaped project directory names — in
`learn/plugins/claude.py`, `_greedy_path_decode` recurses through
`iterdir()` at every level and tries each tokenization of each child,
backtracking on a miss — so on a large home tree it runs for minutes.
Doing that on the event loop freezes uvicorn for the whole window. The
port keeps accepting TCP, but `/readyz` never answers, so a supervisor
health-checking the proxy kills a process that is merely busy.
Field thread dumps show exactly that:
```
Current thread (most recent call first):
File "python3.12/pathlib.py", line 1056 in iterdir
File "headroom/learn/plugins/claude.py", line 454 in _greedy_path_decode
File "headroom/learn/plugins/claude.py", line 478 in _greedy_path_decode
File "headroom/learn/plugins/claude.py", line 478 in _greedy_path_decode
File "headroom/learn/plugins/claude.py", line 426 in _decode_project_path
File "headroom/learn/plugins/claude.py", line 71 in discover_projects
File "headroom/memory/traffic_learner.py", line 591 in flush_to_file
File "headroom/memory/traffic_learner.py", line 535 in _flush_worker
File "python3.12/asyncio/events.py", line 88 in _run
File "python3.12/asyncio/base_events.py", line 1999 in _run_once
File "python3.12/asyncio/base_events.py", line 645 in run_forever
File "uvicorn/server.py", line 75 in run
File "headroom/proxy/server.py", line 4992 in run_server
```
Accompanying signals from the same incidents: port accepts TCP,
`/readyz` times out, process CPU 2-13s across the window (I/O bound, not
spinning), proxy log silent 66-336s.
## Type of Change
- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which 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/memory/traffic_learner.py`: `flush_to_file` now awaits
`asyncio.to_thread(plugin.discover_projects)` instead of calling it
inline. `asyncio` was already imported. The result is cached per learner
(`_project_roots_cache`), so the steady-state flush path pays nothing
for the thread hop.
- `tests/test_memory/test_traffic_learner.py`: added
`test_discover_projects_does_not_block_the_event_loop`.
## 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
$ uv run --frozen --extra dev pytest tests/test_memory/test_traffic_learner.py -q
.................................................. [100%]
============================= 152 passed in 2.93s ==============================
$ uvx ruff check headroom/memory/traffic_learner.py tests/test_memory/test_traffic_learner.py
All checks passed!
$ uvx ruff format --check headroom/memory/traffic_learner.py tests/test_memory/test_traffic_learner.py
2 files already formatted
$ uv run --frozen --extra dev mypy headroom/memory/traffic_learner.py
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: macOS 15.6 arm64, Python 3.10.18, pytest 9.0.3, branch
off `main` @ `01df2452`.
- Exact command / steps: reverted only the one-line source change in the
working tree (`await asyncio.to_thread(plugin.discover_projects)` back
to `plugin.discover_projects()`), left the new test in place, ran `uv
run --frozen --extra dev pytest
tests/test_memory/test_traffic_learner.py -k does_not_block -q`, then
restored the line and re-ran the full file.
- Observed result: without the change the test fails — `flush_to_file`
runs to completion synchronously the moment the task is created, so the
loop never regains control while `discover_projects` is parked on a
`threading.Event`. With the change the loop stays responsive and the
flush completes once discovery returns. Full file: 152 passed.
- Not tested: no live proxy run against a multi-minute real home tree;
the blocking behaviour is reproduced deterministically in the test
instead. The thread dump above is captured field evidence, not a run in
this environment.
Failing output with the fix reverted:
```text
$ uv run --frozen --extra dev pytest tests/test_memory/test_traffic_learner.py -k does_not_block -q
tests/test_memory/test_traffic_learner.py:1254: in test_discover_projects_does_not_block_the_event_loop
assert not flush.done()
E AssertionError: assert not True
E + where True = <built-in method done of _asyncio.Task object at 0x10882dff0>()
E + where <built-in method done of _asyncio.Task object at 0x10882dff0> = <Task finished name='Task-1' coro=<TrafficLearner.flush_to_file() done ...>>.done
========================= 1 failed, 151 deselected in 5.51s =========================
```
## 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
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md`
## Additional Notes
- Documentation: N/A — no user-facing behaviour or interface change.
- Bounding `_greedy_path_decode`'s backtracking is the real cost fix and
belongs in its own change. This one only stops a slow walk from taking
the server's liveness with it.
This commit is contained in:
parent
9cfb00838a
commit
a70e5ff78d
2 changed files with 49 additions and 2 deletions
|
|
@ -629,10 +629,15 @@ class TrafficLearner:
|
|||
if not patterns:
|
||||
return
|
||||
|
||||
# Bucket patterns by project.
|
||||
# Bucket patterns by project. discover_projects() walks the filesystem
|
||||
# to decode escaped project directory names, which on a large home tree
|
||||
# takes minutes; running it inline blocked the event loop, so uvicorn
|
||||
# could not answer /readyz and supervisors killed a proxy that was
|
||||
# merely busy. It is called once per learner (cached below), so the
|
||||
# thread hop costs nothing on the steady-state path.
|
||||
if self._project_roots_cache is None:
|
||||
try:
|
||||
self._project_roots_cache = plugin.discover_projects()
|
||||
self._project_roots_cache = await asyncio.to_thread(plugin.discover_projects)
|
||||
except Exception as e:
|
||||
logger.warning("discover_projects failed: %s", e)
|
||||
self._project_roots_cache = []
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ a real memory backend.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import pytest
|
||||
|
|
@ -1214,6 +1216,46 @@ class TestFlushToFile:
|
|||
await learner.flush_to_file()
|
||||
assert writer.calls == [] # no roots → short-circuits before writer
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discover_projects_does_not_block_the_event_loop(self, tmp_path, monkeypatch):
|
||||
"""A slow discover_projects must not stall other loop work.
|
||||
|
||||
discover_projects walks the filesystem; on a large home tree it takes
|
||||
minutes. Called inline it froze the loop, so uvicorn stopped answering
|
||||
/readyz and supervisors killed a proxy that was only busy.
|
||||
"""
|
||||
writer = _FakeWriter()
|
||||
project_path = tmp_path.resolve()
|
||||
plugin = _FakePlugin(roots=[_make_project(str(project_path))], writer=writer)
|
||||
|
||||
release = threading.Event()
|
||||
|
||||
def slow_discover():
|
||||
release.wait(timeout=5)
|
||||
return [_make_project(str(project_path))]
|
||||
|
||||
plugin.discover_projects = slow_discover # type: ignore[method-assign]
|
||||
_install_plugin_registry(monkeypatch, plugin)
|
||||
|
||||
learner = TrafficLearner(backend=None, agent_type="claude", min_evidence=1)
|
||||
learner._pattern_counts["h"] = (
|
||||
ExtractedPattern(
|
||||
category=PatternCategory.ENVIRONMENT,
|
||||
content=f"Working test command: cd {project_path} && pytest",
|
||||
importance=0.5,
|
||||
evidence_count=2,
|
||||
),
|
||||
2,
|
||||
)
|
||||
|
||||
flush = asyncio.create_task(learner.flush_to_file())
|
||||
# The loop stays responsive while discover_projects is stuck.
|
||||
await asyncio.wait_for(asyncio.sleep(0), timeout=1)
|
||||
assert not flush.done()
|
||||
release.set()
|
||||
await asyncio.wait_for(flush, timeout=5)
|
||||
assert writer.calls, "flush should still complete once discovery returns"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unanchored_patterns_dropped(self, tmp_path, monkeypatch):
|
||||
"""Patterns with no path anchoring are dropped before writer is called."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue