fix(learn): echo live progress during claude-cli analysis

`headroom learn` streams claude-cli's stream-json output internally to
drive an idle-timeout watchdog, but never surfaced any of that
liveness on stdout. Wrapper UIs that tail the CLI's output (like the
Headroom desktop app) saw nothing print between "Analyzing with
claude-cli..." and the final result, often for minutes, making the
scan look hung.

This adds an optional on_progress callback threaded through
SessionAnalyzer.analyze -> _call_llm -> _call_cli_llm ->
_call_claude_cli_streaming, defaulting to None everywhere so no
existing caller is affected. The streaming loop now maps each
non-terminal stream-json event to a short phrase and invokes the
callback, throttled to once per 3 seconds so a burst of partial-message
events doesn't spam the console. headroom/cli/learn.py supplies the
concrete callback, reusing the exact "  Analyzing with ..." prefix
with the detail appended, so prefix-matching wrapper UIs keep working
unchanged. The callback is invoked defensively: an exception from a
broken UI pipe is logged and swallowed rather than aborting an
otherwise-successful analysis.

Also adds --include-partial-messages to the claude-cli invocation, as
suggested in the issue, since without it claude only emits about three
events total for the whole run, too sparse for the throttle to produce
useful incremental updates.

Fixes #3105

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Parideboy 2026-08-21 03:51:19 +02:00
parent 5e0ce242e9
commit 905a08fe9e
4 changed files with 199 additions and 20 deletions

View file

@ -208,6 +208,11 @@ def learn(
analyzer = SessionAnalyzer(model=resolved_model)
def _on_progress(detail: str) -> None:
# Reuses the exact " Analyzing with ..." prefix so wrapper UIs that
# whitelist known stage-line prefixes keep parsing without changes.
click.echo(f" Analyzing with {resolved_model}... ({detail})")
# Determine which agents to scan
agent_configs: list[tuple[str, LearnPlugin]] = []
@ -289,7 +294,7 @@ def learn(
continue
click.echo(f" Analyzing with {resolved_model}...")
result_data = analyzer.analyze(proj, sessions)
result_data = analyzer.analyze(proj, sessions, on_progress=_on_progress)
total_projects += 1
total_failures += result_data.total_failures

View file

@ -53,8 +53,22 @@ _MAX_DIGEST_TOKENS = 80_000 # Budget for the digest (leave room for prompt + ou
# Each entry: (binary_name, model_identifier, command_prefix). The claude-cli
# command uses stream-json output so the analyzer can detect progress and
# enforce an idle (rather than wall-clock-only) timeout — see _call_cli_llm.
# --include-partial-messages adds incremental "stream_event" ticks during the
# assistant's response; without it claude only emits ~3 events total (system
# init, one assistant message, result) — too sparse to report live progress.
_CLI_BACKENDS: list[tuple[str, str, list[str]]] = [
("claude", "claude-cli", ["claude", "-p", "--output-format", "stream-json", "--verbose"]),
(
"claude",
"claude-cli",
[
"claude",
"-p",
"--output-format",
"stream-json",
"--verbose",
"--include-partial-messages",
],
),
("gemini", "gemini-cli", ["gemini", "-p"]),
("codex", "codex-cli", ["codex", "exec"]),
]
@ -71,6 +85,10 @@ _CLI_TIMEOUT = 300
# this long. Lets us catch genuine hangs quickly while letting long-but-active
# analyses run to completion. Override with HEADROOM_LEARN_CLI_IDLE_TIMEOUT_SECS.
_CLI_IDLE_TIMEOUT = 60
# Minimum time between progress-echo callbacks during claude-cli streaming —
# the CLI can emit several partial-message events per second, but a wrapper
# UI heartbeat only needs one update every few seconds, not a play-by-play.
_PROGRESS_THROTTLE_SECS = 3.0
def _resolve_windows_cli_shim(cmd: list[str]) -> list[str] | None:
@ -166,7 +184,12 @@ class SessionAnalyzer:
def __init__(self, model: str | None = None):
self.model = model
def analyze(self, project: ProjectInfo, sessions: list[SessionData]) -> AnalysisResult:
def analyze(
self,
project: ProjectInfo,
sessions: list[SessionData],
on_progress: typing.Callable[[str], None] | None = None,
) -> AnalysisResult:
"""Analyze sessions and produce recommendations via LLM."""
all_calls = [tc for s in sessions for tc in s.tool_calls]
failed_calls = [tc for tc in all_calls if tc.is_error]
@ -195,7 +218,7 @@ class SessionAnalyzer:
# Call LLM for analysis
try:
raw = _call_llm(digest, model)
raw = _call_llm(digest, model, on_progress=on_progress)
result.recommendations = _parse_llm_response(raw)
# Weight loop guardrails above one-off rules using MEASURED waste.
apply_loop_weighting(result.recommendations, loops)
@ -573,7 +596,9 @@ def _failure_detail(
return "\n".join(parts) if parts else "(no output captured)"
def _call_cli_llm(digest: str, model: str) -> dict:
def _call_cli_llm(
digest: str, model: str, on_progress: typing.Callable[[str], None] | None = None
) -> dict:
"""Call a locally installed CLI tool as the LLM backend.
Enables keyless usage for subscription-based CLI tools that handle
@ -581,7 +606,8 @@ def _call_cli_llm(digest: str, model: str) -> dict:
OS ``ARG_MAX`` limits and argument-injection risks.
CLI invocations:
claude-cli claude -p --output-format stream-json --verbose (idle-timeout)
claude-cli claude -p --output-format stream-json --verbose
--include-partial-messages (idle-timeout)
gemini-cli gemini -p (wall-clock timeout)
codex-cli codex exec (wall-clock timeout)
@ -591,6 +617,8 @@ def _call_cli_llm(digest: str, model: str) -> dict:
Args:
digest: Token-efficient session digest to analyze.
model: CLI model identifier (e.g. ``claude-cli``).
on_progress: Optional callback invoked with a short progress phrase
while claude-cli streams (throttled). Ignored for other backends.
Returns:
Parsed JSON recommendations from the CLI tool.
@ -612,7 +640,9 @@ def _call_cli_llm(digest: str, model: str) -> dict:
if model == "claude-cli":
idle_cap = _resolve_timeout_secs("HEADROOM_LEARN_CLI_IDLE_TIMEOUT_SECS", _CLI_IDLE_TIMEOUT)
return _call_claude_cli_streaming(cmd, prompt, hard_cap=hard_cap, idle_cap=idle_cap)
return _call_claude_cli_streaming(
cmd, prompt, hard_cap=hard_cap, idle_cap=idle_cap, on_progress=on_progress
)
try:
result = run(
@ -663,7 +693,12 @@ def _call_cli_llm(digest: str, model: str) -> dict:
def _call_claude_cli_streaming(
cmd: list[str], prompt: str, *, hard_cap: int, idle_cap: int
cmd: list[str],
prompt: str,
*,
hard_cap: int,
idle_cap: int,
on_progress: typing.Callable[[str], None] | None = None,
) -> dict:
"""Run claude-cli with stream-json output and an idle-timeout watchdog.
@ -675,6 +710,10 @@ def _call_claude_cli_streaming(
Threads (rather than ``select``) drain stdout/stderr so the watchdog works
on Windows too, where ``select`` does not support pipe handles.
If *on_progress* is given, it is called (throttled to at most once per
``_PROGRESS_THROTTLE_SECS``) with a short phrase for non-terminal events,
so a caller can surface liveness during the otherwise-silent analysis.
"""
def _popen(cmd: list[str]) -> subprocess.Popen:
@ -730,6 +769,7 @@ def _call_claude_cli_streaming(
start = time.monotonic()
last_activity = start
last_progress_at = 0.0 # 0 so the first eligible event fires immediately
stdout_lines: list[str] = []
stderr_lines: list[str] = []
final_result: str | None = None
@ -778,11 +818,22 @@ def _call_claude_cli_streaming(
if tag == "stdout":
stdout_lines.append(line)
event = _parse_stream_event(line)
if event is not None and event.get("type") == "result":
# Last result event wins if multiple are emitted.
result_text = event.get("result")
if isinstance(result_text, str):
final_result = result_text
if event is not None:
if event.get("type") == "result":
# Last result event wins if multiple are emitted.
result_text = event.get("result")
if isinstance(result_text, str):
final_result = result_text
elif on_progress is not None:
now = time.monotonic()
detail = _progress_detail(event, now - start)
if detail is not None and now - last_progress_at >= _PROGRESS_THROTTLE_SECS:
last_progress_at = now
try:
on_progress(detail)
except Exception as exc: # pragma: no cover — defensive, a UI
# callback must never abort a successful analysis.
logger.debug("on_progress callback failed: %s", exc)
else:
stderr_lines.append(line)
@ -831,7 +882,29 @@ def _parse_stream_event(line: str) -> dict | None:
return parsed if isinstance(parsed, dict) else None
def _call_llm(digest: str, model: str) -> dict:
def _progress_detail(event: dict, elapsed: float) -> str | None:
"""Map a claude-cli stream-json event to a short progress phrase.
Returns None for event types with nothing progress-worthy to report the
terminal "result" event is handled by the caller before reaching here, and
any other unrecognized type (including future ones) stays silent rather
than guessed at.
"""
event_type = event.get("type")
if event_type == "system":
# subtype "init" is the session start; anything else (e.g. an
# API-retry notice) still deserves a heartbeat, just not "started".
return "session started" if event.get("subtype") == "init" else f"retrying, {elapsed:.0f}s"
if event_type in ("assistant", "stream_event"):
return f"assistant responding, {elapsed:.0f}s"
if event_type == "user":
return f"tool running, {elapsed:.0f}s"
return None
def _call_llm(
digest: str, model: str, on_progress: typing.Callable[[str], None] | None = None
) -> dict:
"""Call LLM with the session digest and return parsed JSON.
Uses LiteLLM for provider-agnostic access. The model string determines
@ -839,7 +912,7 @@ def _call_llm(digest: str, model: str) -> dict:
For CLI-based models (ending in "-cli"), delegates to ``_call_cli_llm``.
"""
if model in _CLI_MODEL_IDS:
return _call_cli_llm(digest, model)
return _call_cli_llm(digest, model, on_progress=on_progress)
import litellm

View file

@ -65,7 +65,7 @@ class FakeAnalyzer:
self.model = model
self.calls: list[tuple[object, list[object]]] = []
def analyze(self, project, sessions): # noqa: ANN001, ANN201
def analyze(self, project, sessions, on_progress=None): # noqa: ANN001, ANN201
self.calls.append((project, sessions))
return SimpleNamespace(
total_sessions=len(sessions),
@ -176,6 +176,47 @@ def test_learn_project_lookup_and_apply_flow(
assert plugin.writer.calls[0][2] is False
class ProgressEchoingAnalyzer(FakeAnalyzer):
def analyze(self, project, sessions, on_progress=None): # noqa: ANN001, ANN201
self.calls.append((project, sessions))
if on_progress is not None:
on_progress("session started")
on_progress("assistant responding, 5s")
return SimpleNamespace(
total_sessions=len(sessions),
total_calls=3,
total_failures=1,
failure_rate=1 / 3,
recommendations=[SimpleNamespace(section="Rules")],
)
def test_learn_analyzing_line_gets_progress_detail_appended(
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
) -> None:
project_path = tmp_path / "project-a"
project_path.mkdir()
matched = SimpleNamespace(name="project-a", project_path=project_path)
plugin = FakePlugin("codex", "Codex", [matched])
analyzer = ProgressEchoingAnalyzer()
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
monkeypatch.setattr("headroom.learn.registry.get_plugin", lambda name: plugin)
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", lambda model=None: analyzer)
result = runner.invoke(
main,
["learn", "--agent", "codex", "--project", str(project_path)],
catch_exceptions=False,
)
assert result.exit_code == 0, result.output
assert " Analyzing with gpt-4o... (session started)" in result.output
assert " Analyzing with gpt-4o... (assistant responding, 5s)" in result.output
# Final result reporting still appears unmodified after the progress lines.
assert "Recommendations: 1" in result.output
def test_verbosity_all_apply_aggregates_baselines_across_projects(
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
) -> None:
@ -337,7 +378,7 @@ def test_learn_handles_empty_sessions_and_no_pattern_outputs(
return [SimpleNamespace(events=["event"], tool_calls=[], failure_count=0)]
class BranchingAnalyzer(FakeAnalyzer):
def analyze(self, project, sessions): # noqa: ANN001, ANN201
def analyze(self, project, sessions, on_progress=None): # noqa: ANN001, ANN201
self.calls.append((project, sessions))
if project is no_failures:
return SimpleNamespace(

View file

@ -722,7 +722,67 @@ class TestCallCliLlm:
result = _call_cli_llm("test digest", "claude-cli")
assert result == {"context_file_rules": [], "memory_file_rules": []}
cmd = popen.call_args[0][0]
assert cmd == ["claude", "-p", "--output-format", "stream-json", "--verbose"]
assert cmd == [
"claude",
"-p",
"--output-format",
"stream-json",
"--verbose",
"--include-partial-messages",
]
def test_claude_cli_progress_callback_is_throttled(self):
stdout = [
_stream_event("system", subtype="init"),
*[_stream_event("assistant", message={"content": "..."}) for _ in range(5)],
_result_event('{"context_file_rules": [], "memory_file_rules": []}'),
]
progress: list[str] = []
with patch(
"headroom.learn.analyzer.subprocess.Popen", _fake_claude_popen(stdout_lines=stdout)
):
result = _call_cli_llm("test digest", "claude-cli", on_progress=progress.append)
assert result == {"context_file_rules": [], "memory_file_rules": []}
# All 6 progress-worthy events arrive well within one 3s throttle
# window, so only the first ("session started") should be echoed.
assert progress == ["session started"]
def test_claude_cli_progress_callback_reports_spaced_events(self, monkeypatch):
monkeypatch.setattr("headroom.learn.analyzer._PROGRESS_THROTTLE_SECS", 0.01)
stdout = [
_stream_event("system", subtype="init"),
_stream_event("assistant", message={"content": "thinking..."}),
_stream_event("stream_event", event={"type": "content_block_delta"}),
_stream_event("user", message={"content": []}),
_result_event('{"context_file_rules": [], "memory_file_rules": []}'),
]
progress: list[str] = []
with patch(
"headroom.learn.analyzer.subprocess.Popen",
_fake_claude_popen(stdout_lines=stdout, stdout_delay=0.03),
):
result = _call_cli_llm("test digest", "claude-cli", on_progress=progress.append)
assert result == {"context_file_rules": [], "memory_file_rules": []}
assert progress[0] == "session started"
assert progress[1].startswith("assistant responding, ")
assert progress[2].startswith("assistant responding, ") # stream_event maps like assistant
assert progress[3].startswith("tool running, ")
assert len(progress) == 4 # the terminal "result" event is never echoed as progress
def test_claude_cli_progress_callback_exception_does_not_abort_analysis(self):
stdout = [
_stream_event("system", subtype="init"),
_result_event('{"context_file_rules": [], "memory_file_rules": []}'),
]
def _boom(_detail: str) -> None:
raise RuntimeError("wrapper UI pipe closed")
with patch(
"headroom.learn.analyzer.subprocess.Popen", _fake_claude_popen(stdout_lines=stdout)
):
result = _call_cli_llm("test digest", "claude-cli", on_progress=_boom)
assert result == {"context_file_rules": [], "memory_file_rules": []}
def test_claude_cli_parses_fenced_result(self):
stdout = [
@ -1142,14 +1202,14 @@ class TestCallLlmRouting:
def test_routes_cli_model_to_cli_backend(self, mock_cli: MagicMock):
mock_cli.return_value = {"context_file_rules": [], "memory_file_rules": []}
result = _call_llm("test digest", "claude-cli")
mock_cli.assert_called_once_with("test digest", "claude-cli")
mock_cli.assert_called_once_with("test digest", "claude-cli", on_progress=None)
assert result == {"context_file_rules": [], "memory_file_rules": []}
@patch("headroom.learn.analyzer._call_cli_llm")
def test_routes_codex_cli(self, mock_cli: MagicMock):
mock_cli.return_value = {}
_call_llm("digest", "codex-cli")
mock_cli.assert_called_once_with("digest", "codex-cli")
mock_cli.assert_called_once_with("digest", "codex-cli", on_progress=None)
# =============================================================================