From 76799179d9437dcd133a40ffcf5eecc524063c2f Mon Sep 17 00:00:00 2001 From: h Date: Tue, 28 Jul 2026 02:56:17 +0200 Subject: [PATCH] feat(stream): headless `claude -p` transport behind a transport flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The PTY transport stands in for a protocol that did not exist when it was written: it pastes bracketed text into claude's TUI, guesses when the Ink render loop has settled, re-presses Enter when the paste is swallowed, and tails the session JSONL at 100ms. `claude -p` with stream-json on both pipes is that protocol, so add it as a second transport and let callers pick with `BackendOptions.transport`. Measured on one host, same model and prompt, cold single-reply turn: pty burns 2.3s on TUI readiness before the prompt is even submitted (first event 2.5s, turn 5.2s); stream_json burns none (first event 0.3s, first token 1.9s, turn 2.7s). The gap is what `startup_delay`'s 60s cap exists to survive on slow hardware. Everything the old transport relies on carries over unchanged and was verified against a real CLI: multi-turn over one live process, MCP via --mcp-config, --resume, and `native_jsonl` history seeding. Two things are new rather than equal: `result` is native (so usage is the turn's aggregate and durations are real, not synthesized), and `include_partial_messages` yields token-level `StreamEvent`s — which the JSONL, holding only finished blocks, could never provide. Notably, assistant records carry `stop_reason: null` in this mode even at the end of a turn, so `result` is not just tidier than the PTY path's terminal-stop_reason heuristic, it is the only correct signal. Default stays `pty`; nothing changes for existing callers. Also: extract process-group bookkeeping into `procgroup` so both transports sweep claude's children on shutdown, and fix a stale assertion in test_backend that still expected the `is_error: None` that `_block_to_dict` deliberately stopped emitting. --- README.md | 61 ++- src/claude_code_api/__init__.py | 28 +- src/claude_code_api/backend.py | 86 +++- src/claude_code_api/events.py | 28 +- src/claude_code_api/procgroup.py | 106 +++++ src/claude_code_api/pty.py | 137 ++---- src/claude_code_api/stream.py | 781 +++++++++++++++++++++++++++++++ tests/test_backend.py | 4 +- tests/test_stream.py | 429 +++++++++++++++++ 9 files changed, 1549 insertions(+), 111 deletions(-) create mode 100644 src/claude_code_api/procgroup.py create mode 100644 src/claude_code_api/stream.py create mode 100644 tests/test_stream.py diff --git a/README.md b/README.md index f64fa9b..677cd0d 100644 --- a/README.md +++ b/README.md @@ -2,10 +2,27 @@ [![AI Slop Inside](https://sladge.net/badge.svg)](https://sladge.net) Python wrapper around the `claude` CLI for subscription-mode (no API key) -backends. Drives one long-running interactive `claude` per conversation via -a PTY and reads events from the JSONL session file; the public surface is -Anthropic-Messages-API shaped so a gateway in front of it is a one-liner -serializer away. +backends. Drives one long-running `claude` per conversation and yields +typed events; the public surface is Anthropic-Messages-API shaped so a +gateway in front of it is a one-liner serializer away. + +Two transports, picked with `BackendOptions.transport`: + +- **`stream_json`** — `claude -p` with stream-json on stdin and stdout. + Prompts go into a pipe, events come back as they happen, and a native + `result` record closes each turn. Prefer this. +- **`pty`** (default, for backwards compatibility) — drives the + interactive TUI through a pseudo-tty and tails the session JSONL. It + has to guess when the TUI is ready to accept a paste, re-press Enter + when the paste is swallowed, and poll a file for output. Keep it only + if you need something the TUI alone has (e.g. `--remote-control`). + +Measured on the same host, same model, same prompt — cold turn, single +reply: `pty` spends 2.3s on TUI readiness before the prompt is even +submitted (first event 2.5s, turn 5.2s); `stream_json` spends none +(first event 0.3s, first token 1.9s, turn 2.7s). The gap widens on slow +hardware, which is what the PTY transport's 60s `startup_delay` cap +exists to survive. Not affiliated with Anthropic. You need a working subscription, the `claude` CLI on PATH, and to have run `claude /login` once. @@ -27,7 +44,12 @@ import asyncio from claude_code_api import BackendOptions, ClaudeCodeBackend async def main() -> None: - opts = BackendOptions(cwd="/path/to/project", dangerously_skip_permissions=True) + opts = BackendOptions( + cwd="/path/to/project", + dangerously_skip_permissions=True, + transport="stream_json", + include_partial_messages=True, # token-level deltas, stream_json only + ) async with ClaudeCodeBackend(opts) as backend: async for event in backend.complete( [{"role": "user", "content": "say hi"}] @@ -72,19 +94,38 @@ session orchestration. ## How a turn works +Steps 1–3 are shared by both transports: + 1. The backend looks up a live session by `hash_history(messages[:-1])`. - If one matches, the new user message goes straight into its PTY. + If one matches, the new user message goes straight to it. 2. If nothing matches and `messages[:-1]` is empty, a fresh `claude` is spawned with a brand-new `--session-id`. 3. If `messages[:-1]` is non-empty (a continuation we don't have a live - PTY for — e.g. after restart), the backend writes a hand-crafted + process for — e.g. after restart), the backend writes a hand-crafted JSONL transcript at `~/.claude/projects//.jsonl` and spawns `claude --resume `. That is the `native_jsonl` injection mode; the fallback is `concat_message`, which folds the prior history into one large first prompt. + +Then, on `stream_json`: + +4. The prompt is written to stdin as one JSON line. No readiness wait — + the CLI buffers stdin, so this works even before it has finished + booting. +5. stdout is read line by line and each record normalized into a typed + `Event`. With `include_partial_messages` the raw Anthropic + `message_stream` events also come through as `StreamEvent`, with + `index` rebased across the turn's several API requests. +6. The turn closes on the native `result` record, which supplies the + `ResultMessage` verbatim — aggregate usage, real durations, cost. + +Or, on `pty`: + 4. The PTY's stdout is drained continuously by a background thread; we - never read events from there. The JSONL file is tailed at 100ms - cadence and each new record is normalized into a typed `Event`. + never read events from there. The prompt is pasted in + bracketed-paste framing followed by a separate Enter, and re-sent if + the JSONL shows no sign it registered. The JSONL file is tailed at + 100ms cadence and each new record is normalized into a typed `Event`. 5. The turn closes on the first `assistant` record with `stop_reason ∈ {end_turn, max_tokens, stop_sequence, refusal}`. A `ResultMessage` is synthesized from its `usage` and yielded last. @@ -100,7 +141,7 @@ session orchestration. ```bash uv run pytest # unit tests (fast, no real claude) -RUN_CLAUDE_SMOKE=1 uv run pytest tests/test_pty.py tests/test_turn.py tests/test_backend.py +RUN_CLAUDE_SMOKE=1 uv run pytest tests/test_pty.py tests/test_turn.py tests/test_backend.py tests/test_stream.py ``` The smoke-marked tests spawn a real `claude` process and need a logged-in diff --git a/src/claude_code_api/__init__.py b/src/claude_code_api/__init__.py index 243dd9c..e64957e 100644 --- a/src/claude_code_api/__init__.py +++ b/src/claude_code_api/__init__.py @@ -1,15 +1,27 @@ -"""PTY-based wrapper around the `claude` CLI for subscription-mode backends. +"""Wrapper around the `claude` CLI for subscription-mode backends. `ClaudeCodeBackend` + `BackendOptions` is the surface a gateway -consumes. `TurnManager` and the typed events / errors are re-exported for -callers that want to assemble the lower layers directly (e.g. tests, custom -session orchestration). +consumes. Two transports sit underneath, selected by +`BackendOptions.transport`: + +* ``pty`` (default) — drives the interactive TUI through a pseudo-tty and + reads the conversation back out of the session JSONL + (:mod:`claude_code_api.pty`, :mod:`claude_code_api.turn`); +* ``stream_json`` — drives ``claude -p`` with stream-json on stdin and + stdout (:mod:`claude_code_api.stream`), which removes the TUI-timing + heuristics entirely and can additionally emit token-level deltas. + +Both yield the same typed events. `TurnManager` / `StreamTurnManager` and +the typed events / errors are re-exported for callers that want to +assemble the lower layers directly (e.g. tests, custom session +orchestration). """ from claude_code_api.backend import ( BackendOptions, ClaudeCodeBackend, HistoryInjectionMode, + Transport, synthesize_turn_messages, ) from claude_code_api.errors import ( @@ -27,6 +39,7 @@ from claude_code_api.events import ( ContentBlock, Event, ResultMessage, + StreamEvent, SystemMessage, TextBlock, ThinkingBlock, @@ -43,7 +56,8 @@ from claude_code_api.models import ( is_valid_model, ) from claude_code_api.normalizer import normalize -from claude_code_api.pty import kill_orphaned_processes +from claude_code_api.procgroup import kill_orphaned_processes +from claude_code_api.stream import StreamClaudeProcess, StreamTurnManager from claude_code_api.turn import TurnManager __version__ = "0.1.0" @@ -68,11 +82,15 @@ __all__ = [ "RateLimitError", "ResultMessage", "SessionError", + "StreamClaudeProcess", + "StreamEvent", + "StreamTurnManager", "SystemMessage", "TextBlock", "ThinkingBlock", "ToolResultBlock", "ToolUseBlock", + "Transport", "TurnManager", "UserMessage", "classify_pty_failure", diff --git a/src/claude_code_api/backend.py b/src/claude_code_api/backend.py index d74f1b2..61cfb30 100644 --- a/src/claude_code_api/backend.py +++ b/src/claude_code_api/backend.py @@ -49,11 +49,21 @@ from claude_code_api.injection import ( ) from claude_code_api.paths import resolve_jsonl_path from claude_code_api.pty import PtyClaudeProcess, PtyProcessOptions +from claude_code_api.stream import StreamClaudeProcess, StreamTurnManager from claude_code_api.turn import TurnManager from claude_code_api.watcher import JsonlWatcher HistoryInjectionMode = Literal["native_jsonl", "concat_message"] +Transport = Literal["pty", "stream_json"] +"""Which claude process the backend drives. + +``pty`` spawns the interactive TUI and reads the session JSONL — the +original transport, kept as the default so existing callers are +unaffected. ``stream_json`` spawns ``claude -p`` with stream-json on both +pipes; see :mod:`claude_code_api.stream` for why that is strictly less +machinery. Both honour the same options and yield the same events.""" + ParseErrorCallback = Callable[[MessageParseError, dict[str, Any]], None] # How often the reaper wakes up to look for idle sessions. Far shorter @@ -86,6 +96,24 @@ class BackendOptions: extra_env: Mapping[str, str] = field(default_factory=dict) preserve_provider_env: bool = False + transport: Transport = "pty" + """Which claude process to drive. See :data:`Transport`. + + ``stream_json`` ignores the PTY-only timing knobs below + (``startup_delay`` / ``file_wait_timeout`` / + ``turn_duration_timeout`` / ``wait_for_turn_duration``): the + conditions they wait for — a rendered TUI, a materialized JSONL + file, a turn-end heartbeat — do not exist in headless mode. They are + left on the dataclass rather than rejected so a config can flip + ``transport`` without being rewritten.""" + + include_partial_messages: bool = False + """Emit :class:`~claude_code_api.events.StreamEvent` for token-level + deltas. Requires ``transport="stream_json"``; silently inert on the + PTY transport, which has no finer signal than a finished block. + Whole-block ``AssistantMessage`` events are emitted either way, so a + consumer that ignores ``StreamEvent`` sees no difference.""" + history_injection_mode: HistoryInjectionMode = "native_jsonl" wait_for_turn_duration: bool = False include_meta_user: bool = False @@ -109,11 +137,19 @@ class BackendOptions: @dataclass class _LiveSession: - """One live PTY + watcher + turn manager. Created per conversation.""" + """One live claude process (+ watcher, on PTY) and its turn manager. - pty: PtyClaudeProcess - watcher: JsonlWatcher - tm: TurnManager + Created per conversation. The two transports fill this in + differently — ``stream_json`` has no JSONL to tail, so ``watcher`` is + ``None`` there — but everything the backend and the gateway's admin + surfaces read off ``pty`` (``session_id``, ``pid``, ``created_at``, + ``last_activity_at``, ``captured_output``, output listeners) is + implemented by both process classes. + """ + + pty: PtyClaudeProcess | StreamClaudeProcess + watcher: JsonlWatcher | None + tm: TurnManager | StreamTurnManager @property def session_id(self) -> str: @@ -189,7 +225,7 @@ class ClaudeCodeBackend: return len(self._sessions) @property - def live_sessions(self) -> dict[str, PtyClaudeProcess]: + def live_sessions(self) -> dict[str, PtyClaudeProcess | StreamClaudeProcess]: """Snapshot of live PTY processes keyed by ``session_id``. Returned dict is a copy — caller may iterate freely without @@ -203,7 +239,9 @@ class ClaudeCodeBackend: *during* its turn, not only after it's repooled. ``_active`` wins on key collision, but a session is never in both at once. """ - out = {s.session_id: s.pty for s in self._sessions.values()} + out: dict[str, PtyClaudeProcess | StreamClaudeProcess] = { + s.session_id: s.pty for s in self._sessions.values() + } out.update({sid: s.pty for sid, s in self._active.items()}) return out @@ -467,11 +505,16 @@ class ClaudeCodeBackend: self, *, session_id: str, resume: bool, jsonl_path: Path, start_offset: int ) -> _LiveSession: _log.info( - "_spawn_real_session: session_id=%s resume=%s start_offset=%d", + "_spawn_real_session: session_id=%s transport=%s resume=%s start_offset=%d", session_id, + self._opts.transport, resume, start_offset, ) + if self._opts.transport == "stream_json": + return await self._spawn_stream_session( + session_id=session_id, resume=resume + ) pty_opts = self._build_pty_options(session_id=session_id, resume=resume) pty = PtyClaudeProcess(pty_opts) watcher = JsonlWatcher(jsonl_path, start_offset=start_offset) @@ -493,6 +536,35 @@ class ClaudeCodeBackend: ) return _LiveSession(pty=pty, watcher=watcher, tm=tm) + async def _spawn_stream_session( + self, *, session_id: str, resume: bool + ) -> _LiveSession: + """Spawn the headless counterpart of :meth:`_spawn_real_session`. + + No watcher: stdout *is* the event stream, so there is no file to + tail and no start offset to seek past. History seeding still + happens the same way — the caller has already written the seed + JSONL and set ``resume``, and ``claude -p --resume`` picks it up + exactly as the TUI does. + """ + opts = self._build_pty_options(session_id=session_id, resume=resume) + proc = StreamClaudeProcess( + opts, include_partial_messages=self._opts.include_partial_messages + ) + tm = StreamTurnManager( + proc, + include_meta_user=self._opts.include_meta_user, + include_partial_messages=self._opts.include_partial_messages, + on_parse_error=self._on_parse_error, + ) + await tm.start() + _log.info( + "_spawn_stream_session: session_id=%s STARTED pid=%s", + session_id, + proc.pid, + ) + return _LiveSession(pty=proc, watcher=None, tm=tm) + def _build_pty_options(self, *, session_id: str, resume: bool) -> PtyProcessOptions: mcp_config = self._mcp_config_argument() kwargs: dict[str, Any] = { diff --git a/src/claude_code_api/events.py b/src/claude_code_api/events.py index dfe7422..1df95c8 100644 --- a/src/claude_code_api/events.py +++ b/src/claude_code_api/events.py @@ -92,6 +92,31 @@ class SystemMessage: uuid: str | None = None +@dataclass +class StreamEvent: + """A raw Anthropic ``message_stream`` event, mid-turn. + + Only the headless stream-json transport can produce these, and only + when ``include_partial_messages`` is on — the session JSONL the PTY + transport reads contains finished blocks and nothing finer. Consumers + that want token-level streaming forward :attr:`event` as-is; those + that don't ignore this type and keep using the whole-block + ``AssistantMessage`` records, which are still emitted either way. + + :attr:`event` is the Anthropic event verbatim + (``content_block_start`` / ``content_block_delta`` / + ``content_block_stop``), with one adjustment: ``index`` is rebased + onto the whole turn. A turn spans several API requests — each one + numbers its content blocks from zero — while the caller sees a single + message envelope, so raw indices would collide. + """ + + event: dict[str, Any] + session_id: str | None = None + parent_tool_use_id: str | None = None + uuid: str | None = None + + @dataclass class ResultMessage: """Synthesized turn-completion summary. @@ -112,7 +137,7 @@ class ResultMessage: uuid: str | None = None -Event = UserMessage | AssistantMessage | SystemMessage | ResultMessage +Event = UserMessage | AssistantMessage | SystemMessage | ResultMessage | StreamEvent __all__ = [ @@ -120,6 +145,7 @@ __all__ = [ "ContentBlock", "Event", "ResultMessage", + "StreamEvent", "SystemMessage", "TextBlock", "ThinkingBlock", diff --git a/src/claude_code_api/procgroup.py b/src/claude_code_api/procgroup.py new file mode 100644 index 0000000..30d9955 --- /dev/null +++ b/src/claude_code_api/procgroup.py @@ -0,0 +1,106 @@ +"""Process-group bookkeeping shared by both transports. + +Every `claude` we spawn becomes a session leader in its own process group +— `pty.fork()` calls `setsid()` for the PTY transport, and the stream +transport asks for `start_new_session=True` explicitly. Two consequences +follow, and both are why this module exists: + +* a Ctrl-C in the gateway's terminal goes to the *gateway's* foreground + process group and never reaches claude, so orphans survive unless we + kill them explicitly; +* killing the group (rather than just the pid) also takes down whatever + claude spawned — MCP stdio servers, ripgrep, node workers. + +We keep the pgids of every live process here so an ``atexit`` hook can +sweep anything a disorderly shutdown left behind. Nothing saves us from +SIGKILL on the gateway itself; everything short of that is covered. +""" + +from __future__ import annotations + +import atexit +import contextlib +import logging +import os +import signal +import threading +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterable + +_log = logging.getLogger("claude_code_api.procgroup") + +_LIVE_PGIDS: set[int] = set() +_LIVE_PGIDS_LOCK = threading.Lock() + + +def register_live(pgid: int) -> None: + """Track a process group so :func:`kill_orphaned_processes` can sweep it.""" + with _LIVE_PGIDS_LOCK: + _LIVE_PGIDS.add(pgid) + + +def unregister_live(pgid: int) -> None: + """Stop tracking a process group we have reaped ourselves.""" + with _LIVE_PGIDS_LOCK: + _LIVE_PGIDS.discard(pgid) + + +def pgid_of(pid: int) -> int: + """Process-group id of `pid`. + + Both spawn paths make the child a session leader, so its pgid equals + its pid; we still ask the kernel rather than assume it. + """ + try: + return os.getpgid(pid) + except OSError: + return pid + + +def signal_group(pid: int, sig: int) -> None: + """Signal `pid`'s whole group, falling back to the bare pid. + + Signalling only the pid leaves claude's own children behind — MCP + stdio servers, node workers, long-running Bash tools — which is how + a "terminated" session keeps holding hundreds of MB. + """ + try: + os.killpg(pgid_of(pid), sig) + except (OSError, ProcessLookupError): + with contextlib.suppress(OSError, ProcessLookupError): + os.kill(pid, sig) + + +def kill_orphaned_processes() -> int: + """SIGKILL every process group we spawned and haven't reaped. + + Registered as an ``atexit`` hook, and safe to call directly (e.g. + from a SIGTERM handler). Returns the number of groups signalled. + """ + with _LIVE_PGIDS_LOCK: + pgids = sorted(_LIVE_PGIDS) + _LIVE_PGIDS.clear() + killed = 0 + for pgid in pgids: + try: + os.killpg(pgid, signal.SIGKILL) + except (OSError, ProcessLookupError): + continue + killed += 1 + if killed: + _log.warning("kill_orphaned_processes: SIGKILLed %d orphan group(s)", killed) + return killed + + +atexit.register(kill_orphaned_processes) + + +__all__: Iterable[str] = ( + "kill_orphaned_processes", + "pgid_of", + "register_live", + "signal_group", + "unregister_live", +) diff --git a/src/claude_code_api/pty.py b/src/claude_code_api/pty.py index 010841f..b5e708d 100644 --- a/src/claude_code_api/pty.py +++ b/src/claude_code_api/pty.py @@ -13,7 +13,6 @@ This module knows nothing about turns, JSONL, or event normalization. from __future__ import annotations import asyncio -import atexit import contextlib import errno import logging @@ -31,61 +30,18 @@ from typing import Self from ptyprocess import PtyProcess from claude_code_api.errors import CLINotFoundError +from claude_code_api.procgroup import ( + kill_orphaned_processes, + pgid_of, + register_live, + signal_group, + unregister_live, +) _log = logging.getLogger("claude_code_api.pty") PtyOutputCallback = Callable[[bytes], None] -# ptyprocess spawns via ``pty.fork()``, which calls ``setsid()`` in the -# child — every claude lands in its own session and process group, with -# pgid == pid. Two consequences: -# -# * a Ctrl-C in the gateway's terminal goes to the *gateway's* foreground -# process group and never reaches claude, so orphans survive unless we -# kill them explicitly; -# * killing the group (rather than just the pid) also takes down whatever -# claude spawned — MCP stdio servers, ripgrep, node workers. -# -# We keep the pgids of every live process here so an ``atexit`` hook can -# sweep anything a disorderly shutdown left behind. Nothing saves us from -# SIGKILL on the gateway itself; everything short of that is covered. -_LIVE_PGIDS: set[int] = set() -_LIVE_PGIDS_LOCK = threading.Lock() - - -def _register_live(pgid: int) -> None: - with _LIVE_PGIDS_LOCK: - _LIVE_PGIDS.add(pgid) - - -def _unregister_live(pgid: int) -> None: - with _LIVE_PGIDS_LOCK: - _LIVE_PGIDS.discard(pgid) - - -def kill_orphaned_processes() -> int: - """SIGKILL every process group we spawned and haven't reaped. - - Registered as an ``atexit`` hook, and safe to call directly (e.g. - from a SIGTERM handler). Returns the number of groups signalled. - """ - with _LIVE_PGIDS_LOCK: - pgids = sorted(_LIVE_PGIDS) - _LIVE_PGIDS.clear() - killed = 0 - for pgid in pgids: - try: - os.killpg(pgid, signal.SIGKILL) - except (OSError, ProcessLookupError): - continue - killed += 1 - if killed: - _log.warning("kill_orphaned_processes: SIGKILLed %d orphan group(s)", killed) - return killed - - -atexit.register(kill_orphaned_processes) - _PROVIDER_ENV_VARS: tuple[str, ...] = ( "ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", @@ -173,22 +129,28 @@ class PtyProcessOptions: raise ValueError(msg) -def build_argv(opts: PtyProcessOptions, session_id: str) -> list[str]: - """Materialize CLI argv for `claude` interactive mode. +def build_session_flags(opts: PtyProcessOptions, session_id: str) -> list[str]: + """Session-identity flags: `--resume ` or `--session-id `. - Subscription-mode TUI must NOT pass `--print`, `--output-format`, or - `--input-format` — they either force headless mode or are silently - ignored by interactive claude. - - When `opts.resume_session_id` is set, emit `--resume ` instead of - `--session-id ` — claude rejects the two flags together unless - `--fork-session` is also passed, which would branch the session into a - new JSONL. + When `opts.resume_session_id` is set we emit `--resume` — claude + rejects the two flags together unless `--fork-session` is also + passed, which would branch the session into a new JSONL. """ if opts.resume_session_id is not None: - argv: list[str] = [opts.executable, "--resume", opts.resume_session_id] - else: - argv = [opts.executable, "--session-id", session_id] + return ["--resume", opts.resume_session_id] + return ["--session-id", session_id] + + +def build_common_flags(opts: PtyProcessOptions) -> list[str]: + """Transport-agnostic claude flags — everything but session identity. + + Shared verbatim by the PTY transport (:func:`build_argv`) and the + headless stream-json transport (:func:`claude_code_api.stream.build_stream_argv`) + so the two spawn an otherwise identical claude. `extra_args` is + appended last; a transport that cannot honour a given escape-hatch + flag is responsible for filtering it out. + """ + argv: list[str] = [] if opts.dangerously_skip_permissions: argv.append("--dangerously-skip-permissions") else: @@ -215,6 +177,22 @@ def build_argv(opts: PtyProcessOptions, session_id: str) -> list[str]: return argv +def build_argv(opts: PtyProcessOptions, session_id: str) -> list[str]: + """Materialize CLI argv for `claude` interactive (TUI) mode. + + Subscription-mode TUI must NOT pass `--print`, `--output-format`, or + `--input-format` — they either force headless mode or are silently + ignored by interactive claude. See + :func:`claude_code_api.stream.build_stream_argv` for the headless + counterpart. + """ + return [ + opts.executable, + *build_session_flags(opts, session_id), + *build_common_flags(opts), + ] + + def build_env( opts: PtyProcessOptions, base: Mapping[str, str] | None = None ) -> dict[str, str]: @@ -385,7 +363,7 @@ class PtyClaudeProcess: self._session_id, self._pty.pid, ) - _register_live(self._pgid()) + register_live(self._pgid()) self._drain_stop.clear() self._drain_thread = threading.Thread( target=self._drain_loop, @@ -685,36 +663,19 @@ class PtyClaudeProcess: return await asyncio.to_thread(pty.wait) def _pgid(self) -> int: - """Process-group id of the child. - - ``pty.fork()`` makes the child a session leader, so its pgid - equals its pid; we still ask the kernel rather than assume it. - """ + """Process-group id of the child.""" pty = self._pty if pty is None: msg = "PtyClaudeProcess not started" raise RuntimeError(msg) - try: - return os.getpgid(pty.pid) - except OSError: - return pty.pid + return pgid_of(pty.pid) def _signal_group(self, sig: int) -> None: - """Signal claude *and everything it spawned*. - - Signalling only the pid leaves claude's own children behind — - MCP stdio servers, node workers, long-running Bash tools — which - is how a "terminated" session keeps holding hundreds of MB. - Falls back to a plain pid kill if the group is already gone. - """ + """Signal claude *and everything it spawned*.""" pty = self._pty if pty is None: return - try: - os.killpg(self._pgid(), sig) - except (OSError, ProcessLookupError): - with contextlib.suppress(OSError): - pty.kill(sig) + signal_group(pty.pid, sig) async def terminate(self, *, grace: float = 5.0) -> int | None: """SIGTERM → wait up to `grace` seconds → SIGKILL ladder.""" @@ -744,7 +705,7 @@ class PtyClaudeProcess: if pty is None: return None with contextlib.suppress(RuntimeError): - _unregister_live(self._pgid()) + unregister_live(self._pgid()) exit_status = await asyncio.to_thread(pty.wait) self._drain_stop.set() thread = self._drain_thread @@ -781,6 +742,8 @@ __all__: Iterable[str] = ( "PtyOutputCallback", "PtyProcessOptions", "build_argv", + "build_common_flags", "build_env", + "build_session_flags", "kill_orphaned_processes", ) diff --git a/src/claude_code_api/stream.py b/src/claude_code_api/stream.py new file mode 100644 index 0000000..5f8b944 --- /dev/null +++ b/src/claude_code_api/stream.py @@ -0,0 +1,781 @@ +r"""Headless `claude -p` transport: stream-json in, stream-json out. + +The PTY transport in :mod:`claude_code_api.pty` drives claude's *TUI*: +it pastes bracketed text into a terminal, guesses when the Ink render +loop has settled, re-presses Enter when the paste is swallowed, and +reads the conversation back out of the session JSONL file at a 100ms +poll. Every one of those steps is a heuristic standing in for a +protocol that did not exist when it was written. + +It does now. ``claude -p --input-format stream-json --output-format +stream-json`` is that protocol: + +* **stdin is a pipe, not a keyboard.** One JSON object per line, each a + complete user message. No paste framing, no submit key, no + confirm-and-retry — a write either lands in the pipe or raises. There + is no readiness handshake to wait for either: the CLI buffers stdin, + so the first prompt can be written before the process has finished + booting. +* **stdout is the event stream.** Records arrive as they happen instead + of being tailed out of a file, which removes both the poll latency and + the "did the JSONL file ever appear?" failure mode. +* **the turn boundary is explicit.** A native ``result`` record closes + every turn and carries aggregated usage, cost and timings. Notably, + ``assistant`` records in this mode carry ``stop_reason: null`` even at + the end of a turn — so ``result`` is not merely more convenient than + the PTY path's terminal-``stop_reason`` heuristic, it is the *only* + correct signal here. +* **partial messages are available.** With + ``include_partial_messages=True`` the CLI also emits the raw Anthropic + ``message_stream`` events, giving genuine token-level streaming. The + JSONL the PTY transport reads only ever contains finished blocks, so + this is a capability the old transport cannot have. + +What is deliberately kept identical to the PTY transport: the argv +(:func:`claude_code_api.pty.build_common_flags`), the environment +(:func:`claude_code_api.pty.build_env`), session identity and +``--resume``-based history seeding, and the typed :mod:`events` this +module yields. Swapping transports must not change what the model sees +or what the caller gets back. + +TUI-only escape-hatch flags in ``extra_args`` (``--remote-control``, +``--ide``, ``--worktree``, ``--tmux``) are dropped with a warning rather +than passed through: ``claude -p --remote-control`` would drop out of +print mode entirely and hang the turn. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import logging +import os +import signal +import time +import uuid +from collections.abc import Callable, Iterable, Mapping +from typing import TYPE_CHECKING, Any, Self + +from claude_code_api.errors import ( + AuthError, + BackendError, + CLINotFoundError, + MessageParseError, + ProcessError, + classify_pty_failure, +) +from claude_code_api.events import AssistantMessage, Event, ResultMessage, StreamEvent +from claude_code_api.normalizer import normalize +from claude_code_api.procgroup import ( + pgid_of, + register_live, + signal_group, + unregister_live, +) +from claude_code_api.pty import ( + PtyProcessOptions, + build_common_flags, + build_env, + build_session_flags, +) + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + +_log = logging.getLogger("claude_code_api.stream") + +OutputCallback = Callable[[bytes], None] + +_DEFAULT_OUTPUT_BUFFER_CAP = 1_000_000 + +# ``asyncio.StreamReader`` defaults to a 64 KiB line limit and raises +# ``LimitOverrunError`` past it. A single stream-json record routinely +# blows through that — a Read tool_result carrying a whole file, or a +# base64 image block. 32 MiB is far above any record we have observed +# and still bounded. +_STDOUT_LINE_LIMIT = 32 * 1024 * 1024 + +# `system` subtypes that exist for the CLI's own bookkeeping. The PTY +# transport never saw these (they are stdout-only, not JSONL records); +# surfacing them now would be a behaviour change on transport swap, so +# they are dropped before they reach `normalize`. +_BOOKKEEPING_SUBTYPES: frozenset[str] = frozenset( + { + "init", + "status", + "hook_started", + "hook_response", + "thinking_tokens", + "compact_boundary", + } +) + +# Flags that only make sense for an interactive claude. Passing any of +# them alongside `-p` either errors out or silently drops the process out +# of print mode, which looks like a hung turn from here. +_TUI_ONLY_FLAGS: frozenset[str] = frozenset( + {"--remote-control", "--ide", "--worktree", "-w", "--tmux", "--chrome"} +) + +# `result.subtype` values that mean "claude finished, here is the turn". +# Anything else is a failure the caller needs to see as an exception. +_OK_RESULT_SUBTYPES: frozenset[str] = frozenset({"success"}) + + +def strip_tui_only_flags(extra_args: Iterable[str]) -> tuple[str, ...]: + """Drop interactive-only flags from `extra_args`, loudly. + + Handles both bare (``--ide``) and ``=``-joined (``--worktree=x``) + spellings. A flag that takes a separate value argument would leave + that value orphaned, so we do not attempt to be clever: the flags in + :data:`_TUI_ONLY_FLAGS` are all bare or ``=``-joined in practice. + """ + out: list[str] = [] + for arg in extra_args: + head = arg.split("=", 1)[0] + if head in _TUI_ONLY_FLAGS: + _log.warning( + "stream transport: dropping TUI-only flag %r from extra_args " + "(it would take claude out of print mode)", + arg, + ) + continue + out.append(arg) + return tuple(out) + + +def build_stream_argv( + opts: PtyProcessOptions, + session_id: str, + *, + include_partial_messages: bool = False, +) -> list[str]: + """Materialize CLI argv for headless stream-json mode. + + ``--verbose`` is not optional: the CLI refuses + ``--output-format stream-json`` without it. + """ + argv = [ + opts.executable, + "-p", + "--input-format", + "stream-json", + "--output-format", + "stream-json", + "--verbose", + ] + if include_partial_messages: + argv.append("--include-partial-messages") + argv += build_session_flags(opts, session_id) + sanitized = PtyProcessOptions( + **{ + **{ + f: getattr(opts, f) + for f in opts.__dataclass_fields__ + if f != "extra_args" + }, + "extra_args": strip_tui_only_flags(opts.extra_args), + } + ) + argv += build_common_flags(sanitized) + return argv + + +class StreamClaudeProcess: + """A live headless `claude -p` process speaking stream-json. + + Public lifecycle mirrors :class:`~claude_code_api.pty.PtyClaudeProcess` + closely enough that :class:`~claude_code_api.backend.ClaudeCodeBackend` + and the gateway's admin surfaces can treat the two interchangeably:: + + proc = StreamClaudeProcess(opts) + await proc.start() + await proc.send_user_message("hi") + rec = await proc.next_record() + await proc.terminate() + + Records are parsed off stdout by a reader task into an unbounded + queue; ``next_record()`` pops from it and returns ``None`` once the + process's stdout has closed. + """ + + def __init__( + self, + options: PtyProcessOptions, + *, + include_partial_messages: bool = False, + on_output: OutputCallback | None = None, + output_buffer_cap: int = _DEFAULT_OUTPUT_BUFFER_CAP, + ) -> None: + self._opts = options + self._output_buffer_cap = output_buffer_cap + self._output_listeners: list[OutputCallback] = [] + if on_output is not None: + self._output_listeners.append(on_output) + + if options.resume_session_id is not None: + self._session_id = options.resume_session_id + else: + self._session_id = options.session_id or str(uuid.uuid4()) + self._argv = build_stream_argv( + options, + self._session_id, + include_partial_messages=include_partial_messages, + ) + self._env = build_env(options) + + self._proc: asyncio.subprocess.Process | None = None + self._queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue() + self._reader_task: asyncio.Task[None] | None = None + self._stderr_task: asyncio.Task[None] | None = None + self._output_buffer = bytearray() + self._stderr_buffer = bytearray() + self._eof = False + self._created_at = time.time() + self._last_activity_at = self._created_at + + # ---- introspection (duck-compatible with PtyClaudeProcess) --------- + + @property + def session_id(self) -> str: + return self._session_id + + @property + def created_at(self) -> float: + return self._created_at + + @property + def last_activity_at(self) -> float: + """Unix timestamp of the last prompt written to this process.""" + return self._last_activity_at + + def touch(self) -> None: + self._last_activity_at = time.time() + + @property + def argv(self) -> list[str]: + return list(self._argv) + + @property + def env(self) -> dict[str, str]: + return dict(self._env) + + @property + def cwd(self) -> str: + return os.fspath(self._opts.cwd) + + @property + def pid(self) -> int | None: + return self._proc.pid if self._proc is not None else None + + def is_alive(self) -> bool: + return self._proc is not None and self._proc.returncode is None + + def captured_output(self) -> bytes: + """Rolling buffer of raw stdout bytes (one JSON record per line). + + The PTY transport's equivalent holds ANSI TUI frames; here it + holds the structured event stream, which is what the admin + terminal viewer renders. + """ + return bytes(self._output_buffer) + + def stderr_text(self) -> str: + return self._stderr_buffer.decode("utf-8", errors="replace") + + def add_output_listener(self, listener: OutputCallback) -> None: + """Subscribe to every stdout chunk. Callbacks run on the loop.""" + self._output_listeners.append(listener) + + def remove_output_listener(self, listener: OutputCallback) -> None: + with contextlib.suppress(ValueError): + self._output_listeners.remove(listener) + + # ---- lifecycle ---------------------------------------------------- + + async def start(self) -> None: + if self._proc is not None: + msg = "StreamClaudeProcess.start() called twice" + raise RuntimeError(msg) + + _log.info( + "start: session_id=%s spawning argv=%r cwd=%s", + self._session_id, + self._argv, + self.cwd, + ) + try: + self._proc = await asyncio.create_subprocess_exec( + *self._argv, + cwd=self.cwd, + env=self._env, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + # Own session/process group, so terminate() can sweep + # claude's children (MCP stdio servers, node workers) + # and a Ctrl-C in the parent's terminal doesn't race us. + start_new_session=True, + limit=_STDOUT_LINE_LIMIT, + ) + except FileNotFoundError as exc: + _log.exception( + "start: session_id=%s claude CLI not found: %s", + self._session_id, + self._opts.executable, + ) + raise CLINotFoundError(executable=self._opts.executable) from exc + + _log.info( + "start: session_id=%s spawned pid=%s", self._session_id, self._proc.pid + ) + register_live(pgid_of(self._proc.pid)) + self._reader_task = asyncio.create_task( + self._read_stdout(), name=f"stream-stdout-{self._session_id[:8]}" + ) + self._stderr_task = asyncio.create_task( + self._read_stderr(), name=f"stream-stderr-{self._session_id[:8]}" + ) + + async def _read_stdout(self) -> None: + proc = self._proc + if proc is None or proc.stdout is None: + return + try: + while True: + try: + line = await proc.stdout.readline() + except (ValueError, asyncio.LimitOverrunError): + # A record longer than the (very generous) line limit. + # Dropping one record beats killing the turn, and the + # log line makes it diagnosable. + _log.warning( + "read_stdout: session_id=%s record exceeded %d bytes — dropped", + self._session_id, + _STDOUT_LINE_LIMIT, + ) + continue + if not line: + _log.info("read_stdout: session_id=%s EOF", self._session_id) + return + self._append_output(line) + stripped = line.strip() + if not stripped: + continue + try: + record = json.loads(stripped) + except json.JSONDecodeError: + # claude prints the odd non-JSON line (update + # notices, node warnings). Not fatal. + _log.warning( + "read_stdout: session_id=%s non-JSON line: %r", + self._session_id, + stripped[:200], + ) + continue + await self._queue.put(record) + finally: + self._eof = True + await self._queue.put(None) + + async def _read_stderr(self) -> None: + proc = self._proc + if proc is None or proc.stderr is None: + return + while True: + chunk = await proc.stderr.read(65536) + if not chunk: + return + self._stderr_buffer.extend(chunk) + overflow = len(self._stderr_buffer) - self._output_buffer_cap + if overflow > 0: + del self._stderr_buffer[:overflow] + _log.warning( + "stderr: session_id=%s %s", + self._session_id, + chunk.decode("utf-8", errors="replace").rstrip(), + ) + + def _append_output(self, chunk: bytes) -> None: + self._output_buffer.extend(chunk) + overflow = len(self._output_buffer) - self._output_buffer_cap + if overflow > 0: + del self._output_buffer[:overflow] + for cb in list(self._output_listeners): + with contextlib.suppress(Exception): + cb(chunk) + + # ---- I/O ---------------------------------------------------------- + + async def send_user_message(self, text: str) -> None: + """Write one user message to stdin as a stream-json line.""" + proc = self._proc + if proc is None or proc.stdin is None: + msg = "StreamClaudeProcess not started" + raise RuntimeError(msg) + self.touch() + payload = { + "type": "user", + "message": { + "role": "user", + "content": [{"type": "text", "text": text}], + }, + } + line = json.dumps(payload, ensure_ascii=False) + "\n" + _log.info( + "send_user_message: session_id=%s writing %d chars to stdin", + self._session_id, + len(text), + ) + proc.stdin.write(line.encode("utf-8")) + await proc.stdin.drain() + + async def next_record(self) -> dict[str, Any] | None: + """Pop the next stdout record, or ``None`` once stdout has closed.""" + return await self._queue.get() + + async def write(self, data: str | bytes, *, newline: bool = True) -> int: + """Admin-terminal compatibility shim — this transport is read-only. + + The PTY transport exposes ``write()`` so the admin's xterm.js + bridge can type into a stuck session by hand. Headless claude + has no keyboard: stdin carries the turn protocol, and injecting + arbitrary bytes would desynchronize it. We record the attempt in + the output buffer so the operator sees why nothing happened + rather than typing into a void. + """ + _ = data, newline + notice = ( + b"\r\n[beaver] this session runs the headless stream-json " + b"transport; the terminal is read-only\r\n" + ) + self._append_output(notice) + return 0 + + # ---- shutdown ----------------------------------------------------- + + async def terminate(self, *, grace: float = 5.0) -> int | None: + """SIGTERM → wait up to `grace` seconds → SIGKILL ladder.""" + proc = self._proc + if proc is None: + return None + if proc.returncode is None: + signal_group(proc.pid, signal.SIGTERM) + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(proc.wait(), timeout=grace) + if proc.returncode is None: + signal_group(proc.pid, signal.SIGKILL) + return await self._reap() + + async def kill(self) -> int | None: + proc = self._proc + if proc is None: + return None + if proc.returncode is None: + signal_group(proc.pid, signal.SIGKILL) + return await self._reap() + + async def _reap(self) -> int | None: + proc = self._proc + if proc is None: + return None + unregister_live(pgid_of(proc.pid)) + with contextlib.suppress(Exception): + if proc.stdin is not None and not proc.stdin.is_closing(): + proc.stdin.close() + returncode = await proc.wait() + for task in (self._reader_task, self._stderr_task): + if task is not None and not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + self._reader_task = None + self._stderr_task = None + return returncode + + async def aclose(self) -> int | None: + """Idempotent shutdown — terminate if alive, otherwise reap.""" + if self._proc is None: + return None + if self._proc.returncode is None: + return await self.terminate() + return await self._reap() + + async def __aenter__(self) -> Self: + await self.start() + return self + + async def __aexit__(self, _t: object, _e: object, _tb: object) -> None: + await self.aclose() + + +class StreamTurnManager: + """Drive one turn at a time over a long-lived headless claude. + + The counterpart to :class:`~claude_code_api.turn.TurnManager`, with + the same public shape (``start`` / ``send_user_message`` / ``aclose`` + / ``turn_count``) and a fraction of the machinery: there is no + readiness wait, no submit confirmation, no JSONL poll, and no + synthesized result — the CLI hands us all four for free. + """ + + def __init__( + self, + proc: StreamClaudeProcess, + *, + include_meta_user: bool = False, + include_partial_messages: bool = False, + on_parse_error: Callable[[MessageParseError, dict[str, Any]], None] + | None = None, + owns_proc: bool = True, + ) -> None: + self._proc = proc + self._include_meta_user = include_meta_user + self._include_partial_messages = include_partial_messages + self._on_parse_error = on_parse_error + self._owns_proc = owns_proc + self._started = False + self._turn_count = 0 + self._turn_in_progress = False + + @property + def proc(self) -> StreamClaudeProcess: + return self._proc + + @property + def turn_count(self) -> int: + return self._turn_count + + async def start(self) -> None: + """Spawn the process. Idempotent, and does not wait for readiness. + + Nothing to wait *for*: stdin is a pipe the CLI drains once it is + up, so a prompt written microseconds after spawn is picked up + just the same. This is the single biggest difference from the + PTY transport, whose ``start()`` blocks until claude's TUI has + rendered and gone quiet. + """ + if self._started: + return + await self._proc.start() + self._started = True + _log.info("start: session_id=%s READY", self._proc.session_id) + + async def send_user_message(self, text: str) -> AsyncIterator[Event]: + """Send `text` and stream typed events until the turn's ``result``.""" + if not self._started: + msg = "StreamTurnManager.send_user_message() called before start()" + raise RuntimeError(msg) + if self._turn_in_progress: + msg = "send_user_message() called while a turn is in progress" + raise RuntimeError(msg) + self._turn_in_progress = True + self._turn_count += 1 + sid = self._proc.session_id + _log.info( + "send_user_message: session_id=%s turn=%d text_len=%d preview=%r", + sid, + self._turn_count, + len(text), + text[:80].replace("\n", "\\n"), + ) + try: + try: + await self._proc.send_user_message(text) + except (OSError, RuntimeError) as exc: + msg = "claude process not accepting input" + raise self._failure(msg) from exc + + # Each API request inside a turn restarts its content-block + # indices at 0, but the caller sees one message envelope per + # turn — so partial events get rebased onto a running offset. + block_offset = 0 + max_index = -1 + n_assistant = 0 + n_yielded = 0 + + while True: + record = await self._proc.next_record() + if record is None: + _log.error( + "send_user_message: session_id=%s stdout closed before result " + "(assistant=%d yielded=%d)", + sid, + n_assistant, + n_yielded, + ) + msg = "claude process exited before completing the turn" + raise self._failure(msg) + + rtype = record.get("type") + + if rtype == "result": + result = self._build_result(record) + if result.is_error and n_assistant == 0: + raise self._failure( + f"claude turn failed: {record.get('subtype')}" + + ( + f" ({record['result']})" + if isinstance(record.get("result"), str) + else "" + ) + ) + _log.info( + "send_user_message: session_id=%s RESULT subtype=%s " + "stop=%s turns=%s dur=%sms assistant=%d yielded=%d", + sid, + record.get("subtype"), + result.stop_reason, + record.get("num_turns"), + result.duration_ms, + n_assistant, + n_yielded, + ) + yield result + return + + if rtype == "stream_event": + if not self._include_partial_messages: + continue + event = record.get("event") + if not isinstance(event, dict): + continue + etype = event.get("type") + if etype in ("content_block_start", "content_block_delta", + "content_block_stop"): + index = event.get("index") + if not isinstance(index, int): + continue + max_index = max(max_index, index) + yield StreamEvent( + event={**event, "index": index + block_offset}, + session_id=record.get("session_id"), + parent_tool_use_id=record.get("parent_tool_use_id"), + uuid=record.get("uuid"), + ) + n_yielded += 1 + elif etype == "message_stop": + # Close of one API request inside the turn: the + # next one starts numbering from 0 again. + block_offset += max_index + 1 + max_index = -1 + # message_start / message_delta describe the inner + # request's envelope; the caller owns the outer one. + continue + + if rtype == "rate_limit_event": + _log.info( + "send_user_message: session_id=%s rate_limit %s", + sid, + json.dumps(record.get("rate_limit_info"))[:200], + ) + continue + + if rtype == "system" and record.get("subtype") in _BOOKKEEPING_SUBTYPES: + continue + + try: + event_obj = normalize( + _adapt_envelope(record), + include_meta_user=self._include_meta_user, + ) + except MessageParseError as exc: + _log.warning( + "send_user_message: session_id=%s parse error: %s", sid, exc + ) + if self._on_parse_error is not None: + with contextlib.suppress(Exception): + self._on_parse_error(exc, record) + continue + if event_obj is None: + continue + if isinstance(event_obj, AssistantMessage): + n_assistant += 1 + _log.info( + "send_user_message: session_id=%s yield AssistantMessage " + "blocks=%d [%s]", + sid, + len(event_obj.content), + ",".join(type(b).__name__ for b in event_obj.content), + ) + n_yielded += 1 + yield event_obj + finally: + self._turn_in_progress = False + + def _build_result(self, record: dict[str, Any]) -> ResultMessage: + """Map a native ``result`` record onto :class:`ResultMessage`. + + Unlike the PTY transport, nothing here is synthesized: the usage + is the turn's aggregate (not the last request's), and the + durations are the CLI's own measurements. + """ + subtype = record.get("subtype") or "success" + usage = record.get("usage") + return ResultMessage( + subtype=subtype, + duration_ms=int(record.get("duration_ms") or 0), + num_turns=int(record.get("num_turns") or self._turn_count), + session_id=record.get("session_id") or self._proc.session_id, + is_error=bool(record.get("is_error")) + or subtype not in _OK_RESULT_SUBTYPES, + stop_reason=record.get("stop_reason"), + usage=usage if isinstance(usage, dict) else None, + uuid=record.get("uuid"), + ) + + def _failure(self, message: str) -> BackendError: + """Build the typed exception that fits the process's current state. + + stderr is the honest error channel in headless mode — unlike the + PTY transport, which had to scrape claude's TUI chrome. + """ + captured = self._proc.stderr_text() or self._proc.captured_output().decode( + "utf-8", errors="replace" + ) + cls = classify_pty_failure(captured) or ProcessError + if issubclass(cls, ProcessError): + proc = getattr(self._proc, "_proc", None) + return cls( + message, + exit_code=getattr(proc, "returncode", None), + stderr=captured or None, + ) + if issubclass(cls, AuthError): + return cls() + return cls(message) + + async def aclose(self) -> None: + if self._owns_proc: + await self._proc.aclose() + + async def __aenter__(self) -> Self: + await self.start() + return self + + async def __aexit__(self, _t: object, _e: object, _tb: object) -> None: + await self.aclose() + + +def _adapt_envelope(record: Mapping[str, Any]) -> dict[str, Any]: + """Bridge stdout's envelope naming to the JSONL naming `normalize` knows. + + The ``message`` payload is byte-identical between the two; only the + envelope differs (``session_id`` vs ``sessionId``). Rather than teach + the normalizer two dialects — it is a pure function shared by both + transports and worth keeping that way — we translate here. + """ + out = dict(record) + if "sessionId" not in out and "session_id" in out: + out["sessionId"] = out["session_id"] + return out + + +__all__: Iterable[str] = ( + "OutputCallback", + "StreamClaudeProcess", + "StreamTurnManager", + "build_stream_argv", + "strip_tui_only_flags", +) diff --git a/tests/test_backend.py b/tests/test_backend.py index fed8ee9..4cc1a00 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -836,7 +836,9 @@ def test_synthesize_turn_messages_covers_whole_cycle() -> None: "type": "tool_result", "tool_use_id": "tu_1", "content": "result body", - "is_error": None, + # `is_error` is omitted rather than nulled — the + # Anthropic API rejects a null with "Input should be + # a valid boolean". See `_block_to_dict`. } ], }, diff --git a/tests/test_stream.py b/tests/test_stream.py new file mode 100644 index 0000000..f5a0e62 --- /dev/null +++ b/tests/test_stream.py @@ -0,0 +1,429 @@ +"""Unit tests for the headless stream-json transport. + +No real `claude` here — `StreamTurnManager` talks to a fake process that +replays a canned record list, which is enough to pin the parts that are +genuinely ours: turn-boundary detection, the content-block index rebase +across multi-request turns, bookkeeping filtering, and error mapping. +The live end-to-end path is covered by the smoke test at the bottom. +""" + +from __future__ import annotations + +import os +from typing import Any + +import pytest + +from claude_code_api.errors import AuthError, ProcessError +from claude_code_api.events import ( + AssistantMessage, + ResultMessage, + StreamEvent, + UserMessage, +) +from claude_code_api.pty import PtyProcessOptions +from claude_code_api.stream import ( + StreamTurnManager, + build_stream_argv, + strip_tui_only_flags, +) + + +class FakeStreamProcess: + """Minimal stand-in for `StreamClaudeProcess`. + + Records sent prompts, replays a scripted list of stdout records, and + returns ``None`` (EOF) once exhausted. + """ + + def __init__(self, records: list[dict[str, Any] | None], *, stderr: str = "") -> None: + self._records = list(records) + self.sent: list[str] = [] + self.started = False + self.closed = False + self.session_id = "sess-fake" + self._stderr = stderr + + async def start(self) -> None: + self.started = True + + async def send_user_message(self, text: str) -> None: + self.sent.append(text) + + async def next_record(self) -> dict[str, Any] | None: + if not self._records: + return None + return self._records.pop(0) + + def stderr_text(self) -> str: + return self._stderr + + def captured_output(self) -> bytes: + return b"" + + async def aclose(self) -> None: + self.closed = True + + +def _assistant(*blocks: dict[str, Any], stop: str | None = None) -> dict[str, Any]: + return { + "type": "assistant", + "session_id": "sess-fake", + "uuid": "u1", + "message": { + "role": "assistant", + "model": "claude-test", + "content": list(blocks), + "stop_reason": stop, + "usage": {"input_tokens": 1, "output_tokens": 2}, + }, + } + + +def _result(**over: Any) -> dict[str, Any]: + base: dict[str, Any] = { + "type": "result", + "subtype": "success", + "is_error": False, + "stop_reason": "end_turn", + "num_turns": 1, + "duration_ms": 1234, + "session_id": "sess-fake", + "usage": {"input_tokens": 10, "output_tokens": 20}, + } + base.update(over) + return base + + +def _stream(etype: str, index: int | None = None, **extra: Any) -> dict[str, Any]: + event: dict[str, Any] = {"type": etype, **extra} + if index is not None: + event["index"] = index + return { + "type": "stream_event", + "session_id": "sess-fake", + "parent_tool_use_id": None, + "event": event, + } + + +async def _drain(tm: StreamTurnManager, text: str = "hi") -> list[Any]: + return [ev async for ev in tm.send_user_message(text)] + + +# --- argv ---------------------------------------------------------------- + + +def test_build_stream_argv_is_headless_and_carries_common_flags() -> None: + opts = PtyProcessOptions( + cwd="/vault", + session_id="SID", + model="claude-opus-5", + system_prompt="SP", + effort="high", + dangerously_skip_permissions=True, + ) + argv = build_stream_argv(opts, "SID") + assert argv[:7] == [ + "claude", + "-p", + "--input-format", + "stream-json", + "--output-format", + "stream-json", + "--verbose", + ] + assert "--include-partial-messages" not in argv + assert argv[argv.index("--session-id") + 1] == "SID" + assert argv[argv.index("--model") + 1] == "claude-opus-5" + assert argv[argv.index("--system-prompt") + 1] == "SP" + assert argv[argv.index("--effort") + 1] == "high" + assert "--dangerously-skip-permissions" in argv + + +def test_build_stream_argv_resume_wins_over_session_id() -> None: + opts = PtyProcessOptions(cwd="/vault", resume_session_id="RID") + argv = build_stream_argv(opts, "ignored") + assert "--session-id" not in argv + assert argv[argv.index("--resume") + 1] == "RID" + + +def test_build_stream_argv_include_partial_messages() -> None: + opts = PtyProcessOptions(cwd="/vault", session_id="SID") + argv = build_stream_argv(opts, "SID", include_partial_messages=True) + assert "--include-partial-messages" in argv + + +def test_strip_tui_only_flags_drops_interactive_escape_hatches() -> None: + kept = strip_tui_only_flags( + ("--remote-control", "--keep", "--worktree=wt", "--ide", "--also-keep") + ) + assert kept == ("--keep", "--also-keep") + + +def test_build_stream_argv_strips_tui_only_extra_args() -> None: + opts = PtyProcessOptions( + cwd="/vault", session_id="SID", extra_args=("--remote-control", "--keep") + ) + argv = build_stream_argv(opts, "SID") + assert "--remote-control" not in argv + assert "--keep" in argv + + +# --- turn loop ----------------------------------------------------------- + + +async def test_result_record_closes_the_turn_and_is_not_synthesized() -> None: + proc = FakeStreamProcess( + [_assistant({"type": "text", "text": "hello"}), _result()] + ) + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + events = await _drain(tm) + + assert [type(e).__name__ for e in events] == ["AssistantMessage", "ResultMessage"] + result = events[-1] + assert isinstance(result, ResultMessage) + # Straight off the wire, not fabricated from the last assistant. + assert result.duration_ms == 1234 + assert result.stop_reason == "end_turn" + assert result.usage == {"input_tokens": 10, "output_tokens": 20} + assert result.is_error is False + assert proc.sent == ["hi"] + + +async def test_assistant_stop_reason_null_does_not_end_the_turn() -> None: + """The distinguishing property of this transport. + + Headless claude leaves `stop_reason` null on every assistant record, + so a terminal-stop_reason heuristic would either end the turn at the + first record or never end it. Only `result` counts. + """ + proc = FakeStreamProcess( + [ + _assistant({"type": "tool_use", "id": "t1", "name": "Bash", "input": {}}), + { + "type": "user", + "session_id": "sess-fake", + "message": { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "ok"} + ], + }, + }, + _assistant({"type": "text", "text": "done"}), + _result(num_turns=2), + ] + ) + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + events = await _drain(tm) + + assert [type(e).__name__ for e in events] == [ + "AssistantMessage", + "UserMessage", + "AssistantMessage", + "ResultMessage", + ] + assert isinstance(events[1], UserMessage) + assert isinstance(events[2], AssistantMessage) + assert events[2].content[0].text == "done" # type: ignore[union-attr] + + +async def test_bookkeeping_system_records_are_dropped() -> None: + proc = FakeStreamProcess( + [ + {"type": "system", "subtype": "init", "session_id": "s"}, + {"type": "system", "subtype": "hook_started", "session_id": "s"}, + {"type": "system", "subtype": "status", "session_id": "s"}, + {"type": "rate_limit_event", "rate_limit_info": {"status": "allowed"}}, + _assistant({"type": "text", "text": "x"}), + _result(), + ] + ) + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + events = await _drain(tm) + assert [type(e).__name__ for e in events] == ["AssistantMessage", "ResultMessage"] + + +async def test_partial_messages_are_suppressed_by_default() -> None: + proc = FakeStreamProcess( + [ + _stream("content_block_delta", 0, delta={"type": "text_delta", "text": "a"}), + _assistant({"type": "text", "text": "a"}), + _result(), + ] + ) + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + events = await _drain(tm) + assert not any(isinstance(e, StreamEvent) for e in events) + + +async def test_partial_message_indices_are_rebased_across_requests() -> None: + """A turn spans several API requests; the caller sees one envelope. + + Each request numbers its content blocks from zero, so without a + rebase the second request's block 0 would collide with the first + request's block 0 in the consumer's accumulator. + """ + proc = FakeStreamProcess( + [ + # request 1: two blocks (thinking + tool_use) + _stream("message_start", None), + _stream("content_block_start", 0, content_block={"type": "thinking"}), + _stream("content_block_stop", 0), + _stream("content_block_start", 1, content_block={"type": "tool_use"}), + _stream("content_block_stop", 1), + _stream("message_delta", None), + _stream("message_stop", None), + # request 2: one block, numbered from zero again + _stream("message_start", None), + _stream("content_block_start", 0, content_block={"type": "text"}), + _stream( + "content_block_delta", 0, delta={"type": "text_delta", "text": "hi"} + ), + _stream("content_block_stop", 0), + _stream("message_stop", None), + _result(), + ], + ) + tm = StreamTurnManager(proc, include_partial_messages=True) # type: ignore[arg-type] + await tm.start() + events = await _drain(tm) + + partials = [e for e in events if isinstance(e, StreamEvent)] + assert [(e.event["type"], e.event["index"]) for e in partials] == [ + ("content_block_start", 0), + ("content_block_stop", 0), + ("content_block_start", 1), + ("content_block_stop", 1), + # request 2's block 0 lands at 2, after request 1's two blocks + ("content_block_start", 2), + ("content_block_delta", 2), + ("content_block_stop", 2), + ] + # Inner envelopes never leak — the caller owns the outer one. + assert not any( + e.event["type"].startswith("message_") for e in partials + ) + + +# --- failures ------------------------------------------------------------ + + +async def test_eof_before_result_raises_process_error() -> None: + proc = FakeStreamProcess([_assistant({"type": "text", "text": "partial"})]) + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + with pytest.raises(ProcessError, match="exited before completing the turn"): + await _drain(tm) + + +async def test_error_result_with_no_content_raises() -> None: + proc = FakeStreamProcess( + [_result(subtype="error_during_execution", is_error=True, result="boom")] + ) + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + with pytest.raises(ProcessError, match="error_during_execution"): + await _drain(tm) + + +async def test_error_result_after_content_is_surfaced_not_raised() -> None: + """Partial output beats an exception that would discard it.""" + proc = FakeStreamProcess( + [ + _assistant({"type": "text", "text": "got this far"}), + _result(subtype="error_max_turns", is_error=True), + ] + ) + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + events = await _drain(tm) + assert isinstance(events[-1], ResultMessage) + assert events[-1].is_error is True + assert events[-1].subtype == "error_max_turns" + + +async def test_auth_failure_is_classified_from_stderr() -> None: + proc = FakeStreamProcess([], stderr="API Error: 403 please run /login") + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + with pytest.raises(AuthError): + await _drain(tm) + + +async def test_second_turn_reuses_the_same_process() -> None: + proc = FakeStreamProcess( + [ + _assistant({"type": "text", "text": "one"}), + _result(), + _assistant({"type": "text", "text": "two"}), + _result(num_turns=1), + ] + ) + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + await _drain(tm, "first") + await _drain(tm, "second") + assert proc.sent == ["first", "second"] + assert tm.turn_count == 2 + + +async def test_concurrent_turns_are_rejected() -> None: + proc = FakeStreamProcess([_result()]) + tm = StreamTurnManager(proc) # type: ignore[arg-type] + await tm.start() + gen = tm.send_user_message("a") + await anext(gen) + with pytest.raises(RuntimeError, match="turn is in progress"): + await _drain(tm, "b") + await gen.aclose() + + +# --- smoke (real claude) ------------------------------------------------- + +_SMOKE_ENV = "RUN_CLAUDE_SMOKE" + + +@pytest.mark.live +@pytest.mark.skipif( + not os.environ.get(_SMOKE_ENV), reason=f"set {_SMOKE_ENV}=1 to run against claude" +) +async def test_live_stream_transport_multi_turn(tmp_path: Any) -> None: + """Two turns over one live headless claude, second recalling the first.""" + from claude_code_api.backend import BackendOptions, ClaudeCodeBackend + + opts = BackendOptions( + cwd=str(tmp_path), + model="sonnet", + system_prompt="Answer in one short sentence.", + dangerously_skip_permissions=True, + transport="stream_json", + include_partial_messages=True, + ) + async with ClaudeCodeBackend(opts) as backend: + history: list[dict[str, Any]] = [ + {"role": "user", "content": "Remember the codeword: OKAPI-13. Acknowledge."} + ] + events = [ev async for ev in backend.complete(history)] + assert any(isinstance(e, StreamEvent) for e in events) + assert isinstance(events[-1], ResultMessage) + + from claude_code_api.backend import synthesize_turn_messages + + history = [*history, *synthesize_turn_messages(events)] + history.append({"role": "user", "content": "What was the codeword?"}) + + text = "" + async for ev in backend.complete(history): + if isinstance(ev, AssistantMessage): + text += "".join( + b.text for b in ev.content if hasattr(b, "text") # type: ignore[attr-defined] + ) + assert "OKAPI" in text.upper() + # One process served both turns. + assert backend.live_session_count == 1