mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
5e0ce242e9
commit
905a08fe9e
4 changed files with 199 additions and 20 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue