coding-agent-cli 0.1.0__py3-none-any.whl

coding_agent_cli/_process.py

@@ -0,0 +1,338 @@
+"""Asyncio subprocess transport for agent CLIs."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import time
+from collections import deque
+from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass
+
+from ._platform import kill_process_tree, prepare_command, process_group_kwargs, track_process_tree
+
+log = logging.getLogger("coding_agent_cli.process")
+
+TAIL_BYTES = 64 * 1024
+LINE_LIMIT_BYTES = 1024 * 1024
+EXIT_GRACE_S = 2.0
+POLL_INTERVAL_S = 0.1
+CLEANUP_S = 10.0
+
+LineHandler = Callable[[str], Awaitable[None]]
+SpawnHandler = Callable[[int], Awaitable[None]]
+
+
+@dataclass(slots=True)
+class ProcessOutcome:
+    """Exit status, newest output tail, and the failure that ended the run, if any."""
+
+    returncode: int | None = None
+    tail: str = ""
+    error: str = ""
+    error_class: str = ""
+    spawned: bool = False
+
+
+class OutputTail:
+    """Newest lines within ``TAIL_BYTES``. The newest line is always kept, even when larger."""
+
+    def __init__(self) -> None:
+        self._lines: deque[tuple[str, int]] = deque()
+        self._size = 0
+
+    def append(self, line: str) -> None:
+        size = len(line.encode("utf-8")) + 1
+        self._lines.append((line, size))
+        self._size += size
+        while self._size > TAIL_BYTES and len(self._lines) > 1:
+            self._size -= self._lines.popleft()[1]
+
+    def text(self) -> str:
+        return "\n".join(line for line, _ in self._lines)
+
+
+async def _write_stdin(stdin: asyncio.StreamWriter, payload: bytes) -> None:
+    stdin.write(payload)
+    await stdin.drain()
+    stdin.close()
+
+
+async def _discard_until_eof(stdout: asyncio.StreamReader) -> None:
+    """Drain the killed group's leftover output so asyncio can close the pipe transport."""
+    try:
+        while await asyncio.wait_for(stdout.read(64 * 1024), EXIT_GRACE_S):
+            pass
+    except (TimeoutError, OSError):
+        pass
+
+
+def _stdin_failure(writer: asyncio.Future[None] | None) -> BaseException | None:
+    if writer is None or not writer.done() or writer.cancelled():
+        return None
+    return writer.exception()
+
+
+def _stdin_outcome(failure: BaseException) -> tuple[str, str]:
+    return f"stdin write failed: {type(failure).__name__}: {failure}", "stdin_failed"
+
+
+def _expired(deadline: float | None, now: float) -> bool:
+    return deadline is not None and now >= deadline
+
+
+def _stop_reason(
+    now: float,
+    deadline: float | None,
+    cancel: asyncio.Event | None,
+    writer: asyncio.Future[None] | None,
+    outcome: ProcessOutcome,
+) -> str:
+    """Return the first stop reason in check order: cancel, deadline, stdin failure. Empty means keep going."""
+    if cancel is not None and cancel.is_set():
+        return "cancelled"
+    if _expired(deadline, now):
+        return "timeout"
+    failure = _stdin_failure(writer)
+    if failure is None:
+        return ""
+    outcome.error, outcome.error_class = _stdin_outcome(failure)
+    return "failed"
+
+
+def _note_exit(process: asyncio.subprocess.Process, exited_at: float | None, now: float) -> float | None:
+    """Return when the leader was first seen exited, or ``exited_at`` unchanged while it still runs."""
+    if process.returncode is not None and exited_at is None:
+        return now
+    return exited_at
+
+
+async def _read_one(
+    stdout: asyncio.StreamReader,
+    tail: OutputTail,
+    on_line: LineHandler,
+    outcome: ProcessOutcome,
+) -> str:
+    """Read and handle one line. Return "" to keep looping, "eof" at end of output, or "failed"."""
+    try:
+        raw = await asyncio.wait_for(stdout.readline(), POLL_INTERVAL_S)
+    except TimeoutError:
+        return ""
+    except ValueError:
+        outcome.error, outcome.error_class = "output line exceeds 1 MiB", "execution_error"
+        return "failed"
+    if not raw:
+        return "eof"
+    line = raw.decode("utf-8", errors="replace").rstrip("\r\n")
+    tail.append(line)
+    try:
+        await on_line(line)
+    except Exception as exc:  # noqa: BLE001 - a parser or callback failure fails the attempt
+        outcome.error, outcome.error_class = f"{type(exc).__name__}: {exc}", "execution_error"
+        return "failed"
+    return ""
+
+
+async def _pump(
+    process: asyncio.subprocess.Process,
+    writer: asyncio.Future[None] | None,
+    deadline: float | None,
+    cancel: asyncio.Event | None,
+    on_line: LineHandler,
+    tail: OutputTail,
+    outcome: ProcessOutcome,
+) -> str:
+    """Poll the child until a stop reason, leader exit plus grace, or EOF. Return the stop reason or ""."""
+    assert process.stdout is not None
+    exited_at: float | None = None
+    eof = False
+    while True:
+        now = time.monotonic()
+        stop = _stop_reason(now, deadline, cancel, writer, outcome)
+        if stop:
+            return stop
+        exited_at = _note_exit(process, exited_at, now)
+        if exited_at is not None and (eof or now - exited_at >= EXIT_GRACE_S):
+            if writer is not None:
+                return await _finish_stdin(writer, outcome)
+            return ""
+        if eof:
+            await asyncio.sleep(POLL_INTERVAL_S)
+            continue
+        step = await _read_one(process.stdout, tail, on_line, outcome)
+        if step == "failed":
+            return step
+        eof = step == "eof"
+
+
+async def _watch(deadline: float | None, cancel: asyncio.Event | None) -> str:
+    """Return "cancelled" or "timeout" when one fires, independent of the line handler."""
+    event = cancel if cancel is not None else asyncio.Event()
+    timeout = None if deadline is None else max(0.0, deadline - time.monotonic())
+    try:
+        async with asyncio.timeout(timeout):
+            await event.wait()
+    except TimeoutError:
+        return "timeout"
+    return "cancelled"
+
+
+async def _finish_stdin(writer: asyncio.Future[None], outcome: ProcessOutcome) -> str:
+    """Give the stdin writer ``EXIT_GRACE_S`` to finish. Return "failed" when it failed, else ""."""
+    await asyncio.wait({writer}, timeout=EXIT_GRACE_S)
+    failure = _stdin_failure(writer)
+    if failure is None:
+        return ""
+    outcome.error, outcome.error_class = _stdin_outcome(failure)
+    return "failed"
+
+
+def _map_stop(outcome: ProcessOutcome, stop: str, name: str) -> None:
+    """Set the final error from the stop reason, or from a nonzero exit when nothing stopped the run."""
+    if stop == "cancelled":
+        outcome.error, outcome.error_class = "command cancelled", "cancelled"
+    elif stop == "timeout":
+        outcome.error, outcome.error_class = f"{name} subprocess timed out", "timeout"
+    elif not stop and outcome.returncode:
+        outcome.error = f"command exited with status {outcome.returncode}"
+        outcome.error_class = "execution_error"
+
+
+async def _spawn(
+    command: list[str],
+    payload: bytes | None,
+    env: Mapping[str, str] | None,
+    cwd: str | None,
+) -> asyncio.subprocess.Process:
+    return await asyncio.create_subprocess_exec(
+        *command,
+        stdin=asyncio.subprocess.DEVNULL if payload is None else asyncio.subprocess.PIPE,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.STDOUT,
+        cwd=cwd,
+        env=None if env is None else {**os.environ, **env},
+        limit=LINE_LIMIT_BYTES,
+        **process_group_kwargs(),
+    )
+
+
+async def _supervise(
+    process: asyncio.subprocess.Process,
+    writer: asyncio.Future[None] | None,
+    deadline: float | None,
+    cancel: asyncio.Event | None,
+    on_line: LineHandler,
+    tail: OutputTail,
+    outcome: ProcessOutcome,
+) -> str:
+    """Run the pump while an independent watcher enforces cancel and the deadline."""
+    pump = asyncio.ensure_future(_pump(process, writer, deadline, cancel, on_line, tail, outcome))
+    watcher = asyncio.ensure_future(_watch(deadline, cancel))
+    try:
+        await asyncio.wait({pump, watcher}, return_when=asyncio.FIRST_COMPLETED)
+        if pump.done():
+            return pump.result()
+        return watcher.result()
+    finally:
+        for task in (pump, watcher):
+            if not task.done():
+                task.cancel()
+        await asyncio.gather(pump, watcher, return_exceptions=True)
+
+
+async def _cleanup(process: asyncio.subprocess.Process, writer: asyncio.Future[None] | None) -> None:
+    """Kill the tree, reap the leader, and drain output, all within ``CLEANUP_S``."""
+    assert process.stdout is not None
+    if writer is not None and not writer.done():
+        writer.cancel()
+    try:
+        async with asyncio.timeout(CLEANUP_S):
+            await kill_process_tree(process.pid)
+            if process.returncode is None:
+                try:
+                    process.kill()
+                except ProcessLookupError:
+                    pass
+            await process.wait()
+            if process.stdin is not None:
+                process.stdin.close()
+            await _discard_until_eof(process.stdout)
+    except TimeoutError:
+        log.warning("process cleanup exceeded %.1fs pid=%d", CLEANUP_S, process.pid)
+
+
+async def _shielded(cleanup: asyncio.Future[None]) -> None:
+    """Await ``cleanup`` to completion even when the caller is cancelled, then re-raise the cancellation."""
+    cancelled = False
+    while not cleanup.done():
+        try:
+            await asyncio.shield(cleanup)
+        except asyncio.CancelledError:
+            if cleanup.done():
+                break
+            cancelled = True
+    if cancelled:
+        raise asyncio.CancelledError
+
+
+async def run_process(
+    argv: list[str],
+    stdin: str | None,
+    env: Mapping[str, str] | None,
+    cwd: str | None,
+    deadline: float | None,
+    cancel: asyncio.Event | None,
+    on_line: LineHandler,
+    *,
+    on_spawn: SpawnHandler | None = None,
+) -> ProcessOutcome:
+    """Run ``argv`` in a new process group and await ``on_line`` for each output line.
+
+    stderr is merged into stdout. ``stdin`` is UTF-8 encoded before spawn, written, then
+    closed. ``env`` overlays ``os.environ``. A ``cancel`` already set prevents the spawn.
+    ``deadline`` is a ``time.monotonic`` value. Cancel and the deadline are enforced even
+    while ``on_line`` stalls. The process tree is killed on every exit path, and a cancelled
+    caller task gets ``CancelledError`` after cleanup. Exceptions from ``on_spawn`` propagate
+    after cleanup. A spawn ``OSError`` is returned as ``missing_cli`` or ``spawn_failed``.
+    """
+    if cancel is not None and cancel.is_set():
+        return ProcessOutcome(error="command cancelled", error_class="cancelled")
+    try:
+        payload = None if stdin is None else stdin.encode("utf-8")
+    except UnicodeEncodeError as exc:
+        return ProcessOutcome(error=f"prompt is not valid UTF-8: {exc}", error_class="invalid_prompt")
+
+    name = os.path.basename(argv[0])
+    started_at = time.monotonic()
+    try:
+        process = await _spawn(prepare_command(list(argv)), payload, env, cwd)
+    except FileNotFoundError as exc:
+        return ProcessOutcome(error=f"{name} could not be started: {exc.strerror or exc}", error_class="missing_cli")
+    except OSError as exc:
+        return ProcessOutcome(error=f"{name} could not be started: {exc.strerror or exc}", error_class="spawn_failed")
+    track_process_tree(process.pid)
+    log.debug("spawned %s pid=%d", name, process.pid)
+    writer = None
+    if payload is not None:
+        assert process.stdin is not None
+        writer = asyncio.ensure_future(_write_stdin(process.stdin, payload))
+    tail = OutputTail()
+    outcome = ProcessOutcome(spawned=True)
+    stop = ""
+    try:
+        if on_spawn is not None:
+            await on_spawn(process.pid)
+        stop = await _supervise(process, writer, deadline, cancel, on_line, tail, outcome)
+    finally:
+        await _shielded(asyncio.ensure_future(_cleanup(process, writer)))
+        outcome.returncode = process.returncode
+        outcome.tail = tail.text()
+
+    _map_stop(outcome, stop, name)
+    if outcome.error_class:
+        log.info("%s subprocess ended class=%s after %.2fs", name, outcome.error_class, time.monotonic() - started_at)
+    return outcome
+
+
+__all__ = ["LineHandler", "OutputTail", "ProcessOutcome", "SpawnHandler", "run_process"]

coding_agent_cli/api.py

@@ -0,0 +1,190 @@
+"""Public entry points: ``run``, ``stream``, ``run_sync``, and ``check``."""
+
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+import logging
+from collections.abc import Awaitable, Callable
+from types import TracebackType
+
+from .events import Event, RunCompleted
+from .registry import DriverRegistry, SpawnHandler
+from .types import Capabilities, Request, Result
+
+log = logging.getLogger("coding_agent_cli")
+
+STREAM_QUEUE_SIZE = 256
+
+_default_registry: DriverRegistry | None = None
+
+
+def _registry(registry: DriverRegistry | None) -> DriverRegistry:
+    global _default_registry
+    if registry is not None:
+        return registry
+    if _default_registry is None:
+        _default_registry = DriverRegistry()
+    return _default_registry
+
+
+async def run(
+    req: Request,
+    *,
+    on_event: Callable[[Event], Awaitable[None]] | None = None,
+    cancel: asyncio.Event | None = None,
+    on_spawn: SpawnHandler | None = None,
+    registry: DriverRegistry | None = None,
+) -> Result:
+    """Run ``req`` to completion and return its ``Result``.
+
+    ``on_event`` receives each event with a strictly increasing ``sequence``, ending with
+    exactly one ``RunCompleted`` after the process tree is gone. An exception from
+    ``on_event`` or ``on_spawn`` is raised after cleanup. Setting ``cancel`` ends the run
+    with a ``cancelled`` result; from another thread use ``loop.call_soon_threadsafe(cancel.set)``.
+    Raises ``KeyError`` for an unknown driver.
+    """
+    driver = _registry(registry).get(req.driver)
+    sequence = 0
+    failures: list[BaseException] = []
+
+    async def emit(event: Event) -> None:
+        nonlocal sequence
+        if isinstance(event, RunCompleted):
+            return
+        sequence += 1
+        if on_event is None:
+            return
+        try:
+            await on_event(dataclasses.replace(event, sequence=sequence))
+        except Exception as exc:
+            failures.append(exc)
+            raise
+
+    log.debug("run start driver=%s", req.driver)
+    result = await driver.run(req, on_event=emit, cancel=cancel or asyncio.Event(), on_spawn=on_spawn)
+    if failures:
+        raise failures[0]
+    log.debug("run done driver=%s outcome=%s", req.driver, result.outcome)
+    if on_event is not None:
+        await on_event(RunCompleted(result=result, sequence=sequence + 1))
+    return result
+
+
+class EventStream:
+    """Async iterator over one run's events, ending after ``RunCompleted``.
+
+    Use it as ``async with stream(req) as events: async for event in events: ...``.
+    Leaving the block or calling ``aclose()`` cancels the run and kills the process tree.
+    A bare ``async for`` without ``async with`` is safe only when it runs to the end.
+    """
+
+    def __init__(
+        self,
+        req: Request,
+        *,
+        cancel: asyncio.Event | None,
+        on_spawn: SpawnHandler | None,
+        registry: DriverRegistry | None,
+    ) -> None:
+        self._req = req
+        self._cancel = cancel
+        self._on_spawn = on_spawn
+        self._registry = registry
+        self._queue: asyncio.Queue[Event] = asyncio.Queue(STREAM_QUEUE_SIZE)
+        self._task: asyncio.Task[Result] | None = None
+        self._finished = False
+
+    def _start(self) -> asyncio.Task[Result]:
+        if self._task is None:
+            self._task = asyncio.ensure_future(
+                run(
+                    self._req,
+                    on_event=self._queue.put,
+                    cancel=self._cancel,
+                    on_spawn=self._on_spawn,
+                    registry=self._registry,
+                )
+            )
+        return self._task
+
+    def __aiter__(self) -> EventStream:
+        return self
+
+    async def __anext__(self) -> Event:
+        if self._finished:
+            raise StopAsyncIteration
+        task = self._start()
+        while True:
+            if not self._queue.empty():
+                return self._queue.get_nowait()
+            if task.done():
+                self._finished = True
+                if not task.cancelled() and task.exception() is not None:
+                    raise task.exception()  # type: ignore[misc]
+                raise StopAsyncIteration
+            getter = asyncio.ensure_future(self._queue.get())
+            try:
+                await asyncio.wait({getter, task}, return_when=asyncio.FIRST_COMPLETED)
+            finally:
+                if not getter.done():
+                    getter.cancel()
+            if getter.done() and not getter.cancelled():
+                return getter.result()
+
+    async def aclose(self) -> None:
+        """Cancel the run if it is still going and wait for its cleanup."""
+        self._finished = True
+        task = self._task
+        if task is None:
+            return
+        if not task.done():
+            task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            current = asyncio.current_task()
+            if current is not None and current.cancelling():
+                raise
+        except Exception:
+            pass
+
+    async def __aenter__(self) -> EventStream:
+        self._start()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        await self.aclose()
+
+
+def stream(
+    req: Request,
+    *,
+    cancel: asyncio.Event | None = None,
+    on_spawn: SpawnHandler | None = None,
+    registry: DriverRegistry | None = None,
+) -> EventStream:
+    """Return an ``EventStream`` for ``req``. The run starts on ``async with`` or the first ``__anext__``."""
+    return EventStream(req, cancel=cancel, on_spawn=on_spawn, registry=registry)
+
+
+def run_sync(req: Request, *, registry: DriverRegistry | None = None) -> Result:
+    """Run ``req`` with ``asyncio.run``. Raises ``RuntimeError`` inside a running event loop."""
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return asyncio.run(run(req, registry=registry))
+    raise RuntimeError("run_sync() cannot be called from a running event loop; await run() instead")
+
+
+async def check(req: Request, *, registry: DriverRegistry | None = None, refresh: bool = False) -> Capabilities:
+    """Report whether the driver's CLI is available and supports structured streaming."""
+    return await _registry(registry).get(req.driver).check(req, refresh=refresh)
+
+
+__all__ = ["EventStream", "check", "run", "run_sync", "stream"]

coding_agent_cli/events.py

@@ -0,0 +1,109 @@
+"""Typed run events. Consumers must ignore kinds they do not know."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, ClassVar, Literal
+
+from .types import Result, Usage
+
+SCHEMA_VERSION = 1
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class SessionIdentified:
+    kind: ClassVar[str] = "session.identified"
+    session_id: str
+    sequence: int = 0
+    entity_id: str | None = None
+    schema_version: int = SCHEMA_VERSION
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class MessageStarted:
+    kind: ClassVar[str] = "assistant.message.started"
+    sequence: int = 0
+    entity_id: str | None = None
+    schema_version: int = SCHEMA_VERSION
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class TextDelta:
+    kind: ClassVar[str] = "assistant.text.delta"
+    text: str
+    sequence: int = 0
+    entity_id: str | None = None
+    schema_version: int = SCHEMA_VERSION
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class MessageCompleted:
+    kind: ClassVar[str] = "assistant.message.completed"
+    sequence: int = 0
+    entity_id: str | None = None
+    schema_version: int = SCHEMA_VERSION
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ToolStarted:
+    kind: ClassVar[str] = "tool.started"
+    tool: str
+    activity: Literal["command", "tool_call"]
+    input: Any = None
+    sequence: int = 0
+    entity_id: str | None = None
+    schema_version: int = SCHEMA_VERSION
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ToolCompleted:
+    kind: ClassVar[str] = "tool.completed"
+    status: Literal["completed", "error", "unknown"]
+    sequence: int = 0
+    entity_id: str | None = None
+    schema_version: int = SCHEMA_VERSION
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class UsageUpdated:
+    """A full replacement snapshot of usage so far, not an increment."""
+
+    kind: ClassVar[str] = "usage.updated"
+    usage: Usage
+    sequence: int = 0
+    entity_id: str | None = None
+    schema_version: int = SCHEMA_VERSION
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class RunCompleted:
+    kind: ClassVar[str] = "run.completed"
+    result: Result
+    sequence: int = 0
+    entity_id: str | None = None
+    schema_version: int = SCHEMA_VERSION
+
+
+type Event = (
+    SessionIdentified
+    | MessageStarted
+    | TextDelta
+    | MessageCompleted
+    | ToolStarted
+    | ToolCompleted
+    | UsageUpdated
+    | RunCompleted
+)
+
+__all__ = [
+    "SCHEMA_VERSION",
+    "Event",
+    "MessageCompleted",
+    "MessageStarted",
+    "RunCompleted",
+    "SessionIdentified",
+    "TextDelta",
+    "ToolCompleted",
+    "ToolStarted",
+    "UsageUpdated",
+]

coding_agent_cli/registry.py

@@ -0,0 +1,67 @@
+"""Driver protocol and the instance-scoped driver registry."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from typing import Protocol, runtime_checkable
+
+from .events import Event
+from .types import Capabilities, Request, Result
+
+EventHandler = Callable[[Event], Awaitable[None]]
+SpawnHandler = Callable[[int], Awaitable[None]]
+
+
+@runtime_checkable
+class Driver(Protocol):
+    """Runs one request through one CLI.
+
+    ``run`` awaits ``on_event`` for each event it produces, leaving ``sequence`` to the API,
+    and never emits ``RunCompleted``. It kills the process tree before returning.
+    """
+
+    async def check(self, req: Request, *, refresh: bool = False) -> Capabilities: ...
+
+    async def run(
+        self,
+        req: Request,
+        *,
+        on_event: EventHandler,
+        cancel: asyncio.Event,
+        on_spawn: SpawnHandler | None,
+    ) -> Result: ...
+
+
+def _builtin_drivers() -> dict[str, Driver]:
+    # Imported here: the drivers import this module for the handler aliases.
+    from ._drivers.agy import AgyDriver
+    from ._drivers.claude import ClaudeDriver
+    from ._drivers.codex import CodexDriver
+    from ._drivers.pi import PiDriver
+
+    return {"claude": ClaudeDriver(), "pi": PiDriver(), "codex": CodexDriver(), "agy": AgyDriver()}
+
+
+class DriverRegistry:
+    """Maps driver names to drivers. Each instance is independent."""
+
+    def __init__(self, *, builtins: bool = True) -> None:
+        self._drivers: dict[str, Driver] = _builtin_drivers() if builtins else {}
+
+    def register(self, name: str, driver: Driver, *, replace: bool = False) -> None:
+        if name in self._drivers and not replace:
+            raise ValueError(f"driver {name!r} is already registered")
+        self._drivers[name] = driver
+
+    def get(self, name: str) -> Driver:
+        try:
+            return self._drivers[name]
+        except KeyError:
+            raise KeyError(f"unknown driver {name!r}") from None
+
+    def names(self) -> tuple[str, ...]:
+        return tuple(sorted(self._drivers))
+
+
+__all__ = ["Driver", "DriverRegistry", "EventHandler", "SpawnHandler"]