From 83266b0feadbc9262f9417f078a5810fdb43e779 Mon Sep 17 00:00:00 2001 From: Asadullo Ganiev <62354884+solvemproblr@users.noreply.github.com> Date: Sun, 9 Aug 2026 03:25:07 +0500 Subject: [PATCH 1/3] fix(claude): stream CLI steps in real time --- hud/agents/claude/sdk/agent.py | 370 +++++++++++++++++----- hud/agents/tests/test_claude_sdk_agent.py | 99 +++++- hud/capabilities/ssh.py | 24 ++ hud/capabilities/tests/test_ssh.py | 15 + 4 files changed, 420 insertions(+), 88 deletions(-) diff --git a/hud/agents/claude/sdk/agent.py b/hud/agents/claude/sdk/agent.py index 6b1900784..10408779f 100644 --- a/hud/agents/claude/sdk/agent.py +++ b/hud/agents/claude/sdk/agent.py @@ -10,16 +10,22 @@ from __future__ import annotations +import asyncio +import contextlib import json import logging import shlex from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast +import asyncssh +import mcp.types as mcp_types + from hud.agents.base import Agent -from hud.agents.types import AgentStep, ClaudeSDKConfig, Usage +from hud.agents.types import AgentStep, ClaudeSDKConfig, ToolStep, Usage from hud.settings import settings -from hud.types import Step +from hud.types import MCPToolCall, MCPToolResult, Step +from hud.utils.time import now_iso if TYPE_CHECKING: from hud.capabilities import RFBClient, SSHClient @@ -34,6 +40,7 @@ "{ curl -fsSL https://claude.ai/install.sh | bash -s -- 2>/dev/null; " 'export PATH="$HOME/.local/bin:$PATH"; }' ) +_PROCESS_CLOSE_TIMEOUT_S = 5.0 @dataclass(slots=True) @@ -49,6 +56,257 @@ class RemoteInvocation: script_body: str | None = None +@dataclass(slots=True) +class _PendingToolCall: + call: MCPToolCall + started_at: str + + +class _ClaudeStreamParser: + """Translate Claude CLI stream messages into canonical HUD steps.""" + + def __init__(self, run: Run, *, model: str, started_at: str) -> None: + self._run = run + self._model = model + self._agent_started_at = started_at + self._pending_calls: dict[str, _PendingToolCall] = {} + self._messages: list[dict[str, Any]] = [] + self._last_agent_content = "" + self._saw_result = False + self._error_recorded = False + + @property + def message_count(self) -> int: + return len(self._messages) + + def feed_line(self, line: str) -> None: + line = line.strip() + if not line: + return + try: + raw = json.loads(line) + except json.JSONDecodeError: + logger.warning("Ignoring non-JSON Claude stream output") + return + if not isinstance(raw, dict): + logger.warning("Ignoring non-object Claude stream message") + return + + message = cast("dict[str, Any]", raw) + self._messages.append(message) + received_at = now_iso() + match message.get("type"): + case "system" if message.get("subtype") == "init": + self._agent_started_at = received_at + case "assistant": + self._record_assistant(message, received_at) + case "user": + self._record_tool_results(message, received_at) + case "result": + self._record_result(message, received_at) + + def finish(self, *, exit_status: int, stderr: str) -> None: + trace = self._run.trace + trace.extra["messages"] = self._messages + trace.extra["exit_status"] = exit_status + if stderr: + trace.extra["stderr"] = stderr + if not trace.content and self._last_agent_content: + trace.content = self._last_agent_content + + if exit_status != 0: + trace.status = "error" + self._record_error(stderr or f"claude CLI exited with status {exit_status}") + elif not self._saw_result: + trace.status = "error" + self._record_error("claude CLI exited without a result message") + elif self._pending_calls: + trace.status = "error" + missing = ", ".join(sorted(self._pending_calls)) + self._record_error(f"claude CLI exited without results for tool calls: {missing}") + + def _record_assistant(self, event: dict[str, Any], received_at: str) -> None: + message = event.get("message") + if not isinstance(message, dict): + return + + text_parts: list[str] = [] + thinking_parts: list[str] = [] + tool_calls: list[MCPToolCall] = [] + content = message.get("content") + if isinstance(content, list): + for raw_block in content: + if not isinstance(raw_block, dict): + continue + block = cast("dict[str, Any]", raw_block) + match block.get("type"): + case "text": + if isinstance(block.get("text"), str): + text_parts.append(block["text"]) + case "thinking": + if isinstance(block.get("thinking"), str): + thinking_parts.append(block["thinking"]) + case "tool_use": + call = _tool_call(block) + if call is not None: + tool_calls.append(call) + + text = "".join(text_parts) + if text: + self._last_agent_content = text + model = message.get("model") + stop_reason = message.get("stop_reason") + step = AgentStep( + content=text, + reasoning="\n".join(thinking_parts) if thinking_parts else None, + tool_calls=tool_calls, + done=not tool_calls, + finish_reason=stop_reason if isinstance(stop_reason, str) else None, + model=model if isinstance(model, str) else self._model, + usage=_usage(message.get("usage")), + started_at=self._agent_started_at, + ended_at=received_at, + extra=_event_metadata(event, message), + ) + self._run.record(step) + for call in tool_calls: + self._pending_calls[call.id] = _PendingToolCall(call=call, started_at=received_at) + + def _record_tool_results(self, event: dict[str, Any], received_at: str) -> None: + message = event.get("message") + if not isinstance(message, dict): + return + content = message.get("content") + if not isinstance(content, list): + return + + saw_result = False + for raw_block in content: + if not isinstance(raw_block, dict) or raw_block.get("type") != "tool_result": + continue + block = cast("dict[str, Any]", raw_block) + call_id = block.get("tool_use_id") + if not isinstance(call_id, str): + continue + pending = self._pending_calls.pop(call_id, None) + if pending is None: + logger.warning("Claude returned a result for unknown tool call %s", call_id) + continue + saw_result = True + self._run.record( + ToolStep( + call=pending.call, + result=MCPToolResult( + call_id=call_id, + content=_tool_result_content(block.get("content")), + isError=block.get("is_error") is True, + ), + started_at=pending.started_at, + ended_at=received_at, + extra=_event_metadata(event, message), + ) + ) + if saw_result: + self._agent_started_at = received_at + + def _record_result(self, event: dict[str, Any], received_at: str) -> None: + self._saw_result = True + trace = self._run.trace + result = event.get("result") + trace.content = result if isinstance(result, str) else self._last_agent_content + is_error = event.get("is_error") is True + trace.status = "error" if is_error else "completed" + for key in ( + "subtype", + "session_id", + "duration_ms", + "duration_api_ms", + "stop_reason", + "num_turns", + "total_cost_usd", + ): + value = event.get(key) + if value is not None: + trace.extra[key] = value + if is_error: + self._record_error(trace.content or "claude CLI reported an error", received_at) + + def _record_error(self, error: str, at: str | None = None) -> None: + if self._error_recorded: + return + timestamp = at or now_iso() + self._run.record( + Step(source="system", error=error, started_at=timestamp, ended_at=timestamp) + ) + self._error_recorded = True + + +def _tool_call(block: dict[str, Any]) -> MCPToolCall | None: + call_id = block.get("id") + name = block.get("name") + if not isinstance(call_id, str) or not isinstance(name, str): + logger.warning("Ignoring malformed Claude tool call") + return None + raw_arguments = block.get("input") + if isinstance(raw_arguments, dict): + arguments: dict[str, Any] | str = cast("dict[str, Any]", raw_arguments) + elif isinstance(raw_arguments, str): + arguments = raw_arguments + else: + arguments = json.dumps(raw_arguments, ensure_ascii=False) + return MCPToolCall(id=call_id, name=name, arguments=arguments) + + +def _tool_result_content(value: Any) -> list[mcp_types.ContentBlock]: + values = value if isinstance(value, list) else [value] + content: list[mcp_types.ContentBlock] = [] + for item in values: + if isinstance(item, str): + content.append(mcp_types.TextContent(type="text", text=item)) + elif ( + isinstance(item, dict) + and item.get("type") == "text" + and isinstance(item.get("text"), str) + ): + content.append(mcp_types.TextContent(type="text", text=item["text"])) + elif item is not None: + content.append( + mcp_types.TextContent( + type="text", + text=json.dumps(item, ensure_ascii=False, separators=(",", ":")), + ) + ) + return content + + +def _usage(value: Any) -> Usage | None: + if not isinstance(value, dict): + return None + usage = cast("dict[str, Any]", value) + normalized = Usage( + prompt_tokens=_integer(usage.get("input_tokens")), + completion_tokens=_integer(usage.get("output_tokens")), + cached_tokens=_integer(usage.get("cache_read_input_tokens")), + ) + return normalized if any(v is not None for v in normalized.model_dump().values()) else None + + +def _integer(value: Any) -> int | None: + return value if isinstance(value, int) and not isinstance(value, bool) else None + + +def _event_metadata(event: dict[str, Any], message: dict[str, Any]) -> dict[str, Any]: + metadata: dict[str, Any] = {} + for key in ("session_id", "uuid", "parent_tool_use_id"): + value = event.get(key) + if value is not None: + metadata[key] = value + message_id = message.get("id") + if message_id is not None: + metadata["message_id"] = message_id + return metadata + + def build_remote_invocation(shell: str, run_cmd: str) -> RemoteInvocation: """Build the remote exec command for ``run_cmd`` under the given login shell. @@ -152,20 +410,39 @@ async def _exec( logger.info("SSH exec claude CLI (%d chars)", len(full_cmd)) logger.info("Full command: %s", full_cmd) - completed = await self._ssh.run(full_cmd, check=False) - stdout = completed.stdout if isinstance(completed.stdout, str) else "" - stderr = completed.stderr if isinstance(completed.stderr, str) else "" - - logger.info("exit=%s stdout=%d stderr=%d", completed.exit_status, len(stdout), len(stderr)) - - if completed.exit_status != 0 and not stdout.strip(): - error = stderr or f"claude CLI exited with status {completed.exit_status}" - run.trace.status = "error" - run.trace.extra.update({"exit_status": completed.exit_status, "stderr": stderr}) - run.record(Step(source="system", error=error)) - return - - self._parse_stream_json(run, stdout, stderr) + parser = _ClaudeStreamParser(run, model=self.config.model, started_at=now_iso()) + process = await self._ssh.create_process(full_cmd) + stderr_task = asyncio.create_task(process.stderr.read()) + try: + while line := await process.stdout.readline(): + parser.feed_line(line if isinstance(line, str) else line.decode(errors="replace")) + await process.wait_closed() + stderr_output = await stderr_task + except BaseException: + process.close() + stderr_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await stderr_task + with contextlib.suppress(OSError, TimeoutError, asyncssh.Error): + async with asyncio.timeout(_PROCESS_CLOSE_TIMEOUT_S): + await process.wait_closed() + raise + + stderr = ( + stderr_output + if isinstance(stderr_output, str) + else stderr_output.decode(errors="replace") + ) + exit_status = process.exit_status + if exit_status is None: + raise RuntimeError("claude CLI process closed without an exit status") + logger.info( + "exit=%s messages=%d stderr=%d", + exit_status, + parser.message_count, + len(stderr), + ) + parser.finish(exit_status=exit_status, stderr=stderr) def _build_env_vars(self) -> dict[str, str]: env: dict[str, str] = {} @@ -258,66 +535,5 @@ def _build_cli_command( env_prefix = " ".join(f"{k}={shlex.quote(v)}" for k, v in env_vars.items()) return f'export PATH="$HOME/.local/bin:$PATH"; {env_prefix} {cli_cmd}' - def _parse_stream_json(self, run: Run, stdout: str, stderr: str) -> None: - messages: list[dict[str, Any]] = [] - content_parts: list[str] = [] - is_error = False - info: dict[str, Any] = {} - cost_usd: float | None = None - num_turns: int | None = None - - for line in stdout.splitlines(): - line = line.strip() - if not line: - continue - try: - msg = json.loads(line) - except json.JSONDecodeError: - continue - - messages.append(msg) - msg_type = msg.get("type") - - if msg_type == "assistant" and isinstance(msg.get("message"), dict): - for raw_block in msg["message"].get("content", []): - if not isinstance(raw_block, dict): - continue - block = cast("dict[str, Any]", raw_block) - if block.get("type") == "text" and block.get("text"): - content_parts.append(str(block["text"])) - - elif msg_type == "result": - is_error = msg.get("is_error", False) - result_text = msg.get("result") - if result_text: - content_parts.append(result_text) - info["session_id"] = msg.get("session_id") - info["duration_ms"] = msg.get("duration_ms") - info["stop_reason"] = msg.get("stop_reason") - num_turns = msg.get("num_turns") - cost_usd = msg.get("total_cost_usd") - - content = "\n".join(content_parts) - trace = run.trace - trace.status = "error" if is_error else "completed" - trace.content = content - # Raw CLI stream kept locally; a claude-native serializer can take over - # per-turn fidelity later (the CLI session is its own span vocabulary). - trace.extra["messages"] = messages - if stderr: - trace.extra["stderr"] = stderr - - # The CLI run collapses to one coarse agent step with aggregate usage. - run.record( - AgentStep( - content=content, - done=True, - model=self.config.model, - usage=Usage(cost_usd=cost_usd, llm_call_count=num_turns), - error=content if is_error else None, - extra={k: v for k, v in info.items() if v is not None}, - ), - ) - __all__ = ["ClaudeSDKAgent", "ClaudeSDKConfig", "RemoteInvocation", "build_remote_invocation"] diff --git a/hud/agents/tests/test_claude_sdk_agent.py b/hud/agents/tests/test_claude_sdk_agent.py index 0b9e19617..3ca6f8bbf 100644 --- a/hud/agents/tests/test_claude_sdk_agent.py +++ b/hud/agents/tests/test_claude_sdk_agent.py @@ -9,6 +9,7 @@ from __future__ import annotations +import asyncio import base64 import re from types import SimpleNamespace @@ -19,7 +20,7 @@ from hud.agents.claude.sdk import computer_mcp from hud.agents.claude.sdk.agent import ClaudeSDKAgent, build_remote_invocation -from hud.agents.types import ClaudeSDKConfig +from hud.agents.types import AgentStep, ClaudeSDKConfig, ToolStep from hud.capabilities import Capability, RFBClient, SSHClient from hud.capabilities.rfb import WebPScreenshotEncoding @@ -48,10 +49,50 @@ def test_posix_shell_runs_inline_with_install_check() -> None: # ─── _exec end-to-end over a fake SSH workspace ──────────────────────── +class _FakeProcess: + def __init__( + self, + stdout: str, + *, + stderr: str = "", + exit_status: int = 0, + pause_after: int | None = None, + ) -> None: + self.stdout = self + self.stderr = self + self._lines = stdout.splitlines(keepends=True) + self._stderr = stderr + self._pause_after = pause_after + self._index = 0 + self.exit_status = exit_status + self.blocked = asyncio.Event() + self.release = asyncio.Event() + + async def readline(self) -> str: + if self._pause_after == self._index: + self.blocked.set() + await self.release.wait() + self._pause_after = None + if self._index == len(self._lines): + return "" + line = self._lines[self._index] + self._index += 1 + return line + + async def read(self) -> str: + return self._stderr + + def close(self) -> None: + pass + + async def wait_closed(self) -> None: + pass + + class _FakeConn: - def __init__(self, sink: dict[str, bytes], result: Any) -> None: + def __init__(self, sink: dict[str, bytes], process: _FakeProcess) -> None: self._sink = sink - self._result = result + self._process = process self.ran: list[str] = [] self.write_commands: list[str] = [] @@ -83,8 +124,12 @@ async def run( else: self._sink[name] = b"" return SimpleNamespace(stdout="", stderr="", exit_status=0) + raise AssertionError(f"unexpected buffered command: {cmd}") + + async def create_process(self, cmd: str, **kwargs: Any) -> _FakeProcess: + assert kwargs == {} self.ran.append(cmd) - return self._result + return self._process def _fake_run() -> Any: @@ -94,9 +139,12 @@ def _fake_run() -> Any: _STREAM_JSON = ( - '{"type":"assistant","message":{"content":[{"type":"text","text":"working"}]}}\n' - '{"type":"result","is_error":false,"result":"done","session_id":"s",' - '"duration_ms":5,"num_turns":2,"total_cost_usd":0.01}\n' + '{"type":"assistant","message":{"content":[{"type":"text","text":"editing"},' + '{"type":"tool_use","id":"tool-1","name":"Write","input":{}}]}}\n' + '{"type":"user","message":{"content":[{"type":"tool_result",' + '"tool_use_id":"tool-1","content":"wrote a.txt","is_error":false}]}}\n' + '{"type":"assistant","message":{"content":[{"type":"text","text":"finished"}]}}\n' + '{"type":"result","is_error":false,"result":"finished"}\n' ) @@ -115,7 +163,7 @@ def _agent_with_conn(shell: str, conn: _FakeConn) -> ClaudeSDKAgent: async def test_exec_on_windows_writes_batch_and_execs_via_cmd() -> None: sink: dict[str, bytes] = {} - conn = _FakeConn(sink, SimpleNamespace(stdout=_STREAM_JSON, stderr="", exit_status=0)) + conn = _FakeConn(sink, _FakeProcess(_STREAM_JSON)) agent = _agent_with_conn("cmd", conn) run = _fake_run() @@ -126,12 +174,12 @@ async def test_exec_on_windows_writes_batch_and_execs_via_cmd() -> None: assert sink[".hud_run.bat"].startswith(b"@echo off\r\n") assert sink[".hud_prompt.txt"] == b"build it" assert run.trace.status == "completed" - assert "done" in run.trace.content + assert run.trace.content == "finished" async def test_exec_on_bash_runs_inline_without_batch() -> None: sink: dict[str, bytes] = {} - conn = _FakeConn(sink, SimpleNamespace(stdout=_STREAM_JSON, stderr="", exit_status=0)) + conn = _FakeConn(sink, _FakeProcess(_STREAM_JSON)) agent = _agent_with_conn("bash", conn) run = _fake_run() @@ -147,7 +195,7 @@ async def test_exec_on_bash_runs_inline_without_batch() -> None: async def test_exec_nonzero_exit_with_no_stdout_records_system_error() -> None: sink: dict[str, bytes] = {} - conn = _FakeConn(sink, SimpleNamespace(stdout="", stderr="boom", exit_status=1)) + conn = _FakeConn(sink, _FakeProcess("", stderr="boom", exit_status=1)) agent = _agent_with_conn("cmd", conn) run = _fake_run() @@ -158,6 +206,35 @@ async def test_exec_nonzero_exit_with_no_stdout_records_system_error() -> None: assert run.steps[0].error == "boom" +async def test_exec_records_claude_turn_before_process_exit() -> None: + sink: dict[str, bytes] = {} + process = _FakeProcess(_STREAM_JSON, pause_after=1) + conn = _FakeConn(sink, process) + agent = _agent_with_conn("bash", conn) + run = _fake_run() + + execution = asyncio.create_task(agent._exec(run, prompt="edit it", max_steps=5)) + await process.blocked.wait() + + assert not execution.done() + assert len(run.steps) == 1 + first = run.steps[0] + assert isinstance(first, AgentStep) + assert first.content == "editing" + assert first.tool_calls[0].id == "tool-1" + + process.release.set() + await execution + + assert [type(step) for step in run.steps] == [AgentStep, ToolStep, AgentStep] + tool = cast("ToolStep", run.steps[1]) + assert tool.started_at == first.ended_at + final = cast("AgentStep", run.steps[2]) + assert final.started_at == tool.ended_at + assert run.trace.status == "completed" + assert run.trace.content == "finished" + + @pytest.mark.parametrize( ("transport", "claude_type"), [("streamable-http", "http"), ("sse", "sse")], diff --git a/hud/capabilities/ssh.py b/hud/capabilities/ssh.py index 37e634191..29692cdc9 100644 --- a/hud/capabilities/ssh.py +++ b/hud/capabilities/ssh.py @@ -101,6 +101,30 @@ async def run(self, *args: object, **kwargs: Any) -> asyncssh.SSHCompletedProces raise SSHConnectionError("SSH connection lost during operation") from exc raise + async def create_process( + self, + *args: object, + **kwargs: Any, + ) -> asyncssh.SSHClientProcess[Any]: + """Start a streaming command, reconnecting first when the transport is closed. + + The caller owns the returned process and must close it when interrupted. + Commands are never replayed after the process has been created. + """ + conn: asyncssh.SSHClientConnection | None = None + try: + conn = await self._connection() + assert conn is not None + return await conn.create_process(*args, **kwargs) + except asyncssh.ChannelOpenError as exc: + raise SSHConnectionError("SSH server rejected the session") from exc + except asyncssh.ConnectionLost as exc: + raise SSHConnectionError("SSH connection lost while opening session") from exc + except (OSError, asyncssh.Error) as exc: + if conn is not None and _is_closed(conn): + raise SSHConnectionError("SSH connection lost while opening session") from exc + raise + async def _connection(self) -> asyncssh.SSHClientConnection: if self._closing: raise SSHConnectionError("SSH client is closed") diff --git a/hud/capabilities/tests/test_ssh.py b/hud/capabilities/tests/test_ssh.py index 7549065ba..5f6a26e12 100644 --- a/hud/capabilities/tests/test_ssh.py +++ b/hud/capabilities/tests/test_ssh.py @@ -151,6 +151,21 @@ async def test_run_does_not_replay_a_command_lost_in_flight( assert replacement.commands == ["next-command"] +async def test_create_process_reconnects_before_opening_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + dropped = _Connection(closed=True) + replacement = _Connection() + client = _client(dropped) + reconnect = AsyncMock(return_value=replacement) + monkeypatch.setattr(client, "_connect", reconnect) + + process = await client.create_process("stream-command") + + assert process is replacement.process + reconnect.assert_awaited_once_with(client.capability) + + async def test_run_classifies_rejected_session_as_connection_error() -> None: connection = _Connection( open_error=asyncssh.ChannelOpenError(asyncssh.OPEN_RESOURCE_SHORTAGE, "busy") From 5a9441f87eb0ad7dce18a71e1e74af8ee8939f7a Mon Sep 17 00:00:00 2001 From: Asadullo Ganiev <62354884+solvemproblr@users.noreply.github.com> Date: Sun, 9 Aug 2026 03:54:03 +0500 Subject: [PATCH 2/3] fix(claude): link CLI inference to trace --- hud/agents/claude/sdk/agent.py | 3 +++ hud/agents/tests/test_claude_sdk_agent.py | 15 +++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/hud/agents/claude/sdk/agent.py b/hud/agents/claude/sdk/agent.py index 10408779f..d2dfce167 100644 --- a/hud/agents/claude/sdk/agent.py +++ b/hud/agents/claude/sdk/agent.py @@ -24,6 +24,7 @@ from hud.agents.base import Agent from hud.agents.types import AgentStep, ClaudeSDKConfig, ToolStep, Usage from hud.settings import settings +from hud.telemetry.context import get_current_trace_id from hud.types import MCPToolCall, MCPToolResult, Step from hud.utils.time import now_iso @@ -450,6 +451,8 @@ def _build_env_vars(self) -> dict[str, str]: if settings.api_key: env["ANTHROPIC_BASE_URL"] = settings.hud_gateway_url env["ANTHROPIC_API_KEY"] = settings.api_key + if trace_id := get_current_trace_id(): + env["ANTHROPIC_CUSTOM_HEADERS"] = f"Trace-Id: {trace_id}" elif settings.anthropic_api_key: env["ANTHROPIC_API_KEY"] = settings.anthropic_api_key diff --git a/hud/agents/tests/test_claude_sdk_agent.py b/hud/agents/tests/test_claude_sdk_agent.py index 3ca6f8bbf..337181fff 100644 --- a/hud/agents/tests/test_claude_sdk_agent.py +++ b/hud/agents/tests/test_claude_sdk_agent.py @@ -23,6 +23,8 @@ from hud.agents.types import AgentStep, ClaudeSDKConfig, ToolStep from hud.capabilities import Capability, RFBClient, SSHClient from hud.capabilities.rfb import WebPScreenshotEncoding +from hud.settings import settings +from hud.telemetry.context import set_trace_context # ─── build_remote_invocation (pure) ─────────────────────────────────── @@ -193,6 +195,19 @@ async def test_exec_on_bash_runs_inline_without_batch() -> None: assert run.trace.status == "completed" +async def test_exec_forwards_current_trace_id_to_hud_gateway( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(settings, "api_key", "hud-key") + conn = _FakeConn({}, _FakeProcess(_STREAM_JSON)) + agent = _agent_with_conn("bash", conn) + + with set_trace_context("trace-123"): + await agent._exec(_fake_run(), prompt="build it", max_steps=5) + + assert "ANTHROPIC_CUSTOM_HEADERS='Trace-Id: trace-123'" in conn.ran[0] + + async def test_exec_nonzero_exit_with_no_stdout_records_system_error() -> None: sink: dict[str, bytes] = {} conn = _FakeConn(sink, _FakeProcess("", stderr="boom", exit_status=1)) From bd59c732c980a9e5e4db94f8f5b9e10e0032d15d Mon Sep 17 00:00:00 2001 From: Asadullo Ganiev <62354884+solvemproblr@users.noreply.github.com> Date: Sun, 9 Aug 2026 03:59:16 +0500 Subject: [PATCH 3/3] refactor(claude): reuse message step parser --- hud/agents/claude/agent.py | 34 ++++---- hud/agents/claude/sdk/agent.py | 95 ++++------------------- hud/agents/tests/test_claude_sdk_agent.py | 12 ++- 3 files changed, 46 insertions(+), 95 deletions(-) diff --git a/hud/agents/claude/agent.py b/hud/agents/claude/agent.py index 4294aa270..7aad09954 100644 --- a/hud/agents/claude/agent.py +++ b/hud/agents/claude/agent.py @@ -257,26 +257,36 @@ async def get_response( if response is None: raise ValueError("Claude response missing after retries") - result = AgentStep(content="", done=True) - result.model = response.model - result.usage = Usage( - prompt_tokens=response.usage.input_tokens, - completion_tokens=response.usage.output_tokens, - cached_tokens=response.usage.cache_read_input_tokens, + return self._message_to_agent_step(response, citations_enabled=citations_enabled) + + @classmethod + def _message_to_agent_step( + cls, + response: BetaMessage, + *, + citations_enabled: bool = False, + ) -> AgentStep: + result = AgentStep( + content="", + done=True, + model=response.model, + usage=Usage( + prompt_tokens=response.usage.input_tokens, + completion_tokens=response.usage.output_tokens, + cached_tokens=response.usage.cache_read_input_tokens, + ), ) text_parts: list[str] = [] thinking_parts: list[str] = [] - citations: list[Citation] = [] for block in response.content: match block.type: case "tool_use": - arguments = dict(block.input) if block.input else {} result.tool_calls.append( MCPToolCall( id=block.id, name=block.name, - arguments=arguments, + arguments=dict(block.input) if block.input else {}, _meta=mcp_types.RequestParams.Meta.model_validate( {"citations_enabled": citations_enabled}, ), @@ -284,9 +294,8 @@ async def get_response( ) result.done = False case "text": - text_block = block - text_parts.append(text_block.text) - citations.extend(self._citation(c) for c in (text_block.citations or [])) + text_parts.append(block.text) + result.citations.extend(cls._citation(c) for c in (block.citations or [])) case "thinking": if block.thinking: thinking_parts.append(block.thinking) @@ -294,7 +303,6 @@ async def get_response( pass result.content = "".join(text_parts) - result.citations = citations if thinking_parts: result.reasoning = "\n".join(thinking_parts) result.finish_reason = response.stop_reason diff --git a/hud/agents/claude/sdk/agent.py b/hud/agents/claude/sdk/agent.py index d2dfce167..093d0c7ee 100644 --- a/hud/agents/claude/sdk/agent.py +++ b/hud/agents/claude/sdk/agent.py @@ -20,9 +20,11 @@ import asyncssh import mcp.types as mcp_types +from anthropic.types.beta import BetaMessage from hud.agents.base import Agent -from hud.agents.types import AgentStep, ClaudeSDKConfig, ToolStep, Usage +from hud.agents.claude.agent import ClaudeAgent +from hud.agents.types import ClaudeSDKConfig, ToolStep from hud.settings import settings from hud.telemetry.context import get_current_trace_id from hud.types import MCPToolCall, MCPToolResult, Step @@ -66,9 +68,8 @@ class _PendingToolCall: class _ClaudeStreamParser: """Translate Claude CLI stream messages into canonical HUD steps.""" - def __init__(self, run: Run, *, model: str, started_at: str) -> None: + def __init__(self, run: Run, *, started_at: str) -> None: self._run = run - self._model = model self._agent_started_at = started_at self._pending_calls: dict[str, _PendingToolCall] = {} self._messages: list[dict[str, Any]] = [] @@ -127,50 +128,18 @@ def finish(self, *, exit_status: int, stderr: str) -> None: self._record_error(f"claude CLI exited without results for tool calls: {missing}") def _record_assistant(self, event: dict[str, Any], received_at: str) -> None: - message = event.get("message") - if not isinstance(message, dict): - return - - text_parts: list[str] = [] - thinking_parts: list[str] = [] - tool_calls: list[MCPToolCall] = [] - content = message.get("content") - if isinstance(content, list): - for raw_block in content: - if not isinstance(raw_block, dict): - continue - block = cast("dict[str, Any]", raw_block) - match block.get("type"): - case "text": - if isinstance(block.get("text"), str): - text_parts.append(block["text"]) - case "thinking": - if isinstance(block.get("thinking"), str): - thinking_parts.append(block["thinking"]) - case "tool_use": - call = _tool_call(block) - if call is not None: - tool_calls.append(call) - - text = "".join(text_parts) - if text: - self._last_agent_content = text - model = message.get("model") - stop_reason = message.get("stop_reason") - step = AgentStep( - content=text, - reasoning="\n".join(thinking_parts) if thinking_parts else None, - tool_calls=tool_calls, - done=not tool_calls, - finish_reason=stop_reason if isinstance(stop_reason, str) else None, - model=model if isinstance(model, str) else self._model, - usage=_usage(message.get("usage")), - started_at=self._agent_started_at, - ended_at=received_at, - extra=_event_metadata(event, message), - ) + raw_message = event.get("message") + if not isinstance(raw_message, dict): + raise ValueError("Claude assistant event is missing its message payload") + message = BetaMessage.model_validate(raw_message) + step = ClaudeAgent._message_to_agent_step(message) + step.started_at = self._agent_started_at + step.ended_at = received_at + step.extra = _event_metadata(event, raw_message) + if step.content: + self._last_agent_content = step.content self._run.record(step) - for call in tool_calls: + for call in step.tool_calls: self._pending_calls[call.id] = _PendingToolCall(call=call, started_at=received_at) def _record_tool_results(self, event: dict[str, Any], received_at: str) -> None: @@ -242,22 +211,6 @@ def _record_error(self, error: str, at: str | None = None) -> None: self._error_recorded = True -def _tool_call(block: dict[str, Any]) -> MCPToolCall | None: - call_id = block.get("id") - name = block.get("name") - if not isinstance(call_id, str) or not isinstance(name, str): - logger.warning("Ignoring malformed Claude tool call") - return None - raw_arguments = block.get("input") - if isinstance(raw_arguments, dict): - arguments: dict[str, Any] | str = cast("dict[str, Any]", raw_arguments) - elif isinstance(raw_arguments, str): - arguments = raw_arguments - else: - arguments = json.dumps(raw_arguments, ensure_ascii=False) - return MCPToolCall(id=call_id, name=name, arguments=arguments) - - def _tool_result_content(value: Any) -> list[mcp_types.ContentBlock]: values = value if isinstance(value, list) else [value] content: list[mcp_types.ContentBlock] = [] @@ -280,22 +233,6 @@ def _tool_result_content(value: Any) -> list[mcp_types.ContentBlock]: return content -def _usage(value: Any) -> Usage | None: - if not isinstance(value, dict): - return None - usage = cast("dict[str, Any]", value) - normalized = Usage( - prompt_tokens=_integer(usage.get("input_tokens")), - completion_tokens=_integer(usage.get("output_tokens")), - cached_tokens=_integer(usage.get("cache_read_input_tokens")), - ) - return normalized if any(v is not None for v in normalized.model_dump().values()) else None - - -def _integer(value: Any) -> int | None: - return value if isinstance(value, int) and not isinstance(value, bool) else None - - def _event_metadata(event: dict[str, Any], message: dict[str, Any]) -> dict[str, Any]: metadata: dict[str, Any] = {} for key in ("session_id", "uuid", "parent_tool_use_id"): @@ -411,7 +348,7 @@ async def _exec( logger.info("SSH exec claude CLI (%d chars)", len(full_cmd)) logger.info("Full command: %s", full_cmd) - parser = _ClaudeStreamParser(run, model=self.config.model, started_at=now_iso()) + parser = _ClaudeStreamParser(run, started_at=now_iso()) process = await self._ssh.create_process(full_cmd) stderr_task = asyncio.create_task(process.stderr.read()) try: diff --git a/hud/agents/tests/test_claude_sdk_agent.py b/hud/agents/tests/test_claude_sdk_agent.py index 337181fff..00283f048 100644 --- a/hud/agents/tests/test_claude_sdk_agent.py +++ b/hud/agents/tests/test_claude_sdk_agent.py @@ -141,11 +141,17 @@ def _fake_run() -> Any: _STREAM_JSON = ( - '{"type":"assistant","message":{"content":[{"type":"text","text":"editing"},' - '{"type":"tool_use","id":"tool-1","name":"Write","input":{}}]}}\n' + '{"type":"assistant","message":{"id":"msg-1","type":"message",' + '"role":"assistant","model":"claude-test","content":[{"type":"text",' + '"text":"editing"},{"type":"tool_use","id":"tool-1","name":"Write","input":{}}],' + '"stop_reason":"tool_use","stop_sequence":null,"usage":{"input_tokens":11,' + '"output_tokens":7,"cache_read_input_tokens":3}}}\n' '{"type":"user","message":{"content":[{"type":"tool_result",' '"tool_use_id":"tool-1","content":"wrote a.txt","is_error":false}]}}\n' - '{"type":"assistant","message":{"content":[{"type":"text","text":"finished"}]}}\n' + '{"type":"assistant","message":{"id":"msg-2","type":"message",' + '"role":"assistant","model":"claude-test","content":[{"type":"text",' + '"text":"finished"}],"stop_reason":"end_turn","stop_sequence":null,' + '"usage":{"input_tokens":11,"output_tokens":7,"cache_read_input_tokens":3}}}\n' '{"type":"result","is_error":false,"result":"finished"}\n' )