diff --git a/headroom/cli/learn.py b/headroom/cli/learn.py index 6ea29f2ad..4a469e45f 100644 --- a/headroom/cli/learn.py +++ b/headroom/cli/learn.py @@ -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]] = [] @@ -290,7 +295,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 diff --git a/headroom/learn/analyzer.py b/headroom/learn/analyzer.py index 5fdb44a0e..71acce4ea 100644 --- a/headroom/learn/analyzer.py +++ b/headroom/learn/analyzer.py @@ -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", "--skip-git-repo-check"]), ] @@ -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) @@ -575,7 +598,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 @@ -583,7 +608,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) @@ -593,6 +619,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. @@ -614,7 +642,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( @@ -665,7 +695,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. @@ -677,6 +712,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: @@ -732,6 +771,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 @@ -780,11 +820,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) @@ -833,7 +884,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 @@ -841,7 +914,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 diff --git a/tests/test_cli_learn.py b/tests/test_cli_learn.py index cdb92e123..cad23c017 100644 --- a/tests/test_cli_learn.py +++ b/tests/test_cli_learn.py @@ -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( diff --git a/tests/test_learn/test_analyzer.py b/tests/test_learn/test_analyzer.py index 69308f4e3..fb5573c3d 100644 --- a/tests/test_learn/test_analyzer.py +++ b/tests/test_learn/test_analyzer.py @@ -723,7 +723,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 = [ @@ -1143,14 +1203,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) # =============================================================================