monkeyscode 1.0.0__py3-none-any.whl

monkeyscode/cli_process.py

@@ -0,0 +1,730 @@
+"""
+Drive an installed ``mc`` over stdio (the local agent, not the cloud API).
+
+Spawns ``mc -p --input-format stream-json --output-format stream-json
+--permission-prompt stdio`` and speaks the same NDJSON the CLI prints
+(``schema_version`` 1). Hooks, MCP servers, settings and permission rules
+all apply, because the real CLI runs them.
+
+One-shot::
+
+    from monkeyscode.cli_process import query
+
+    async for event in query("explain src/main.ts"):
+        if isinstance(event, ResultMessage):
+            print(event.result)
+
+Multi-turn with approvals::
+
+    async def can_use_tool(tool, tool_input, ctx):
+        if tool == "create_file" and tool_input.get("path", "").startswith("docs/"):
+            return PermissionAllow()
+        return PermissionDeny(message="only docs/ may be written")
+
+    async with MonkeysCodeClient(CliOptions(can_use_tool=can_use_tool)) as client:
+        await client.send("write docs/intro.md")
+        async for event in client.receive_turn():
+            ...
+        await client.send("now write src/x.ts")
+        result = await client.receive_result()
+
+Standard library only. The TypeScript (``@monkeyscode/sdk/subprocess``) and
+Go (``monkeyscode.Query``) clients implement the same protocol.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+import os
+import shutil
+import tempfile
+from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+SUPPORTED_SCHEMA_VERSION = 1
+
+# ── Events ───────────────────────────────────────────────────────────
+
+
+@dataclass
+class TextBlock:
+    text: str
+    type: Literal["text"] = "text"
+
+
+@dataclass
+class ToolUseBlock:
+    id: str
+    name: str
+    input: dict[str, Any]
+    type: Literal["tool_use"] = "tool_use"
+
+
+@dataclass
+class ToolResultBlock:
+    tool_use_id: str
+    content: str
+    is_error: bool
+    type: Literal["tool_result"] = "tool_result"
+
+
+@dataclass
+class PermissionDenial:
+    tool_name: str
+    tool_input: dict[str, Any]
+    reason: str
+    tool_use_id: str | None = None
+    suggested_rule: str | None = None
+
+
+@dataclass
+class Usage:
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cached_tokens: int = 0
+    thinking_tokens: int = 0
+
+
+@dataclass
+class InitMessage:
+    session_id: str
+    schema_version: int
+    cwd: str
+    model: str
+    permission_mode: str
+    tools: list[str]
+    mcp_servers: list[dict[str, str]]
+    cli_version: str
+    resumed_from: str | None = None
+    forked_from: str | None = None
+    raw: dict[str, Any] = field(default_factory=dict, repr=False)
+    type: Literal["system"] = "system"
+    subtype: Literal["init"] = "init"
+
+
+@dataclass
+class AssistantMessage:
+    session_id: str
+    content: list[TextBlock | ToolUseBlock]
+    raw: dict[str, Any] = field(default_factory=dict, repr=False)
+    type: Literal["assistant"] = "assistant"
+
+    @property
+    def text(self) -> str:
+        return "".join(b.text for b in self.content if isinstance(b, TextBlock))
+
+
+@dataclass
+class UserMessage:
+    """A prompt (``content`` is a str) or tool output (a list of ToolResultBlock)."""
+
+    session_id: str
+    content: str | list[ToolResultBlock]
+    raw: dict[str, Any] = field(default_factory=dict, repr=False)
+    type: Literal["user"] = "user"
+
+
+@dataclass
+class ResultMessage:
+    subtype: str
+    is_error: bool
+    result: str
+    session_id: str
+    num_turns: int
+    duration_ms: int
+    usage: Usage
+    cost_usd: float
+    is_estimated: bool
+    permission_denials: list[PermissionDenial]
+    exit_code: int
+    structured: Any = None
+    error: str | None = None
+    schema_errors: list[str] | None = None
+    raw: dict[str, Any] = field(default_factory=dict, repr=False)
+    type: Literal["result"] = "result"
+
+
+@dataclass
+class SystemMessage:
+    """Informational ``system`` events (retrying, compacted, hook, warning, …)."""
+
+    subtype: str
+    session_id: str
+    data: dict[str, Any]
+    type: Literal["system"] = "system"
+
+
+@dataclass
+class ErrorMessage:
+    error: str
+    session_id: str | None = None
+    type: Literal["error"] = "error"
+
+
+Event = (
+    InitMessage | AssistantMessage | UserMessage | ResultMessage | SystemMessage | ErrorMessage
+)
+
+
+def parse_event(e: Mapping[str, Any]) -> Event:
+    """Typed view of one stream-json line. Unknown fields are kept in ``raw``."""
+    t = e.get("type")
+    sid = str(e.get("session_id", ""))
+    if t == "system" and e.get("subtype") == "init":
+        return InitMessage(
+            session_id=sid,
+            schema_version=int(e.get("schema_version", 0)),
+            cwd=str(e.get("cwd", "")),
+            model=str(e.get("model", "")),
+            permission_mode=str(e.get("permission_mode", "")),
+            tools=list(e.get("tools", [])),
+            mcp_servers=list(e.get("mcp_servers", [])),
+            cli_version=str(e.get("cli_version", "")),
+            resumed_from=e.get("resumed_from"),
+            forked_from=e.get("forked_from"),
+            raw=dict(e),
+        )
+    if t == "assistant":
+        blocks: list[TextBlock | ToolUseBlock] = []
+        for b in e.get("message", {}).get("content", []):
+            if b.get("type") == "text":
+                blocks.append(TextBlock(text=b.get("text", "")))
+            elif b.get("type") == "tool_use":
+                blocks.append(
+                    ToolUseBlock(id=b["id"], name=b["name"], input=dict(b.get("input", {}))),
+                )
+        return AssistantMessage(session_id=sid, content=blocks, raw=dict(e))
+    if t == "user":
+        c = e.get("message", {}).get("content", "")
+        if isinstance(c, list):
+            results = [
+                ToolResultBlock(
+                    tool_use_id=b["tool_use_id"],
+                    content=str(b.get("content", "")),
+                    is_error=bool(b.get("is_error")),
+                )
+                for b in c
+                if b.get("type") == "tool_result"
+            ]
+            return UserMessage(session_id=sid, content=results, raw=dict(e))
+        return UserMessage(session_id=sid, content=str(c), raw=dict(e))
+    if t == "result":
+        u = e.get("usage", {}) or {}
+        return ResultMessage(
+            subtype=str(e.get("subtype", "")),
+            is_error=bool(e.get("is_error")),
+            result=str(e.get("result", "")),
+            session_id=sid,
+            num_turns=int(e.get("num_turns", 0)),
+            duration_ms=int(e.get("duration_ms", 0)),
+            usage=Usage(
+                input_tokens=int(u.get("input_tokens", 0)),
+                output_tokens=int(u.get("output_tokens", 0)),
+                cached_tokens=int(u.get("cached_tokens", 0)),
+                thinking_tokens=int(u.get("thinking_tokens", 0)),
+            ),
+            cost_usd=float(e.get("cost_usd", 0) or 0),
+            is_estimated=bool(e.get("is_estimated")),
+            permission_denials=[
+                PermissionDenial(
+                    tool_name=d.get("tool_name", ""),
+                    tool_input=dict(d.get("tool_input", {})),
+                    reason=d.get("reason", ""),
+                    tool_use_id=d.get("tool_use_id"),
+                    suggested_rule=d.get("suggested_rule"),
+                )
+                for d in e.get("permission_denials", [])
+            ],
+            exit_code=int(e.get("exit_code", 0)),
+            structured=e.get("structured"),
+            error=e.get("error"),
+            schema_errors=e.get("schema_errors"),
+            raw=dict(e),
+        )
+    if t == "error":
+        return ErrorMessage(error=str(e.get("error", "")), session_id=e.get("session_id"))
+    data = {k: v for k, v in e.items() if k not in ("type", "subtype", "session_id")}
+    return SystemMessage(subtype=str(e.get("subtype", t or "")), session_id=sid, data=data)
+
+
+# ── Permissions ──────────────────────────────────────────────────────
+
+
+@dataclass
+class PermissionAllow:
+    updated_input: dict[str, Any] | None = None
+    behavior: Literal["allow"] = "allow"
+
+
+@dataclass
+class PermissionDeny:
+    message: str | None = None
+    behavior: Literal["deny"] = "deny"
+
+
+PermissionResult = PermissionAllow | PermissionDeny
+
+
+@dataclass
+class ToolPermissionContext:
+    description: str
+    suggestions: list[str]
+    tool_use_id: str | None = None
+
+
+CanUseTool = Callable[[str, dict[str, Any], ToolPermissionContext], Awaitable[PermissionResult]]
+
+# ── Options ──────────────────────────────────────────────────────────
+
+
+@dataclass
+class CliOptions:
+    """Options for the CLI subprocess. ``None`` means "use the CLI's default"."""
+
+    #: ``mc`` executable, or argv for a wrapper (``["node", "cli.cjs"]``).
+    #: Default: ``$MC_PATH``, else ``mc`` on PATH.
+    path_to_mc: str | list[str] | None = None
+    cwd: str | None = None
+    model: str | None = None
+    permission_mode: Literal["default", "acceptEdits", "plan", "bypass"] | None = None
+    effort: Literal["low", "medium", "high"] | None = None
+    allowed_tools: list[str] | None = None
+    disallowed_tools: list[str] | None = None
+    add_dirs: list[str] | None = None
+    #: Approve every ordinary ask when no ``can_use_tool`` is given (deny rules still apply).
+    auto_approve: bool = False
+    can_use_tool: CanUseTool | None = None
+    permission_timeout_sec: float | None = None
+    mcp_config: list[str] | None = None
+    strict_mcp_config: bool = False
+    #: Ignore the repository's .monkeyscode settings, hooks and MCP servers (untrusted code).
+    ignore_repo_settings: bool = False
+    #: Extra settings layer: a dict (sent inline as JSON) or a file path.
+    settings: dict[str, Any] | str | None = None
+    append_system_prompt: str | None = None
+    system_prompt_file: str | None = None
+    max_cost_usd: float | None = None
+    max_tokens: int | None = None
+    #: JSON Schema for the final answer: a dict, or a path.
+    json_schema: dict[str, Any] | str | None = None
+    resume: str | None = None
+    continue_session: bool = False
+    fork_session: bool = False
+    no_hooks: bool = False
+    #: Wall-clock limit in ms (default: none).
+    timeout_ms: int | None = None
+    env: dict[str, str | None] | None = None
+    extra_args: list[str] | None = None
+    #: Called with each stderr line of the CLI.
+    stderr: Callable[[str], None] | None = None
+
+
+class McError(Exception):
+    """Base class for subprocess-client errors."""
+
+
+class McSchemaError(McError):
+    def __init__(self, schema_version: int, cli_version: str) -> None:
+        super().__init__(
+            f"mc {cli_version} speaks stream-json schema_version {schema_version}; this client "
+            f"supports {SUPPORTED_SCHEMA_VERSION}. "
+            "Upgrade the monkeyscode package (or install a matching mc).",
+        )
+        self.schema_version = schema_version
+        self.cli_version = cli_version
+
+
+class McProcessError(McError):
+    def __init__(self, message: str, exit_code: int | None, stderr_tail: str = "") -> None:
+        super().__init__(f"{message}\n{stderr_tail}" if stderr_tail else message)
+        self.exit_code = exit_code
+        self.stderr_tail = stderr_tail
+
+
+def build_args(o: CliOptions, write_json: Callable[[Any], str] | None = None) -> list[str]:
+    """CLI argv (no executable). Always an argv list, never a shell string."""
+    a = [
+        "-p",
+        "--input-format",
+        "stream-json",
+        "--output-format",
+        "stream-json",
+        "--permission-prompt",
+        "stdio",
+    ]
+
+    def push(flag: str, v: object) -> None:
+        if v is not None and v != "":
+            a.extend([flag, str(v)])
+
+    push("--workspace", o.cwd)
+    push("--model", o.model)
+    push("--permission-mode", o.permission_mode)
+    push("--effort", o.effort)
+    for flag, values in (
+        ("--allowed-tools", o.allowed_tools),
+        ("--disallowed-tools", o.disallowed_tools),
+        ("--add-dir", o.add_dirs),
+        ("--mcp-config", o.mcp_config),
+    ):
+        if values:
+            a.extend([flag, *values])
+    if o.strict_mcp_config:
+        a.append("--strict-mcp-config")
+    if o.ignore_repo_settings:
+        a.append("--ignore-repo-settings")
+    if o.settings is not None:
+        push("--settings", o.settings if isinstance(o.settings, str) else json.dumps(o.settings))
+    push("--append-system-prompt", o.append_system_prompt)
+    push("--system-prompt-file", o.system_prompt_file)
+    push("--max-cost", o.max_cost_usd)
+    push("--max-tokens", o.max_tokens)
+    if o.json_schema is not None:
+        path = (
+            o.json_schema
+            if isinstance(o.json_schema, str)
+            else (write_json or _write_temp_json)(o.json_schema)
+        )
+        push("--json-schema", path)
+    push("--resume", o.resume)
+    if o.continue_session:
+        a.append("--continue")
+    if o.fork_session:
+        a.append("--fork-session")
+    if o.no_hooks:
+        a.append("--no-hooks")
+    push("--permission-timeout", o.permission_timeout_sec)
+    a.extend(["--timeout", str(o.timeout_ms or 0)])
+    if o.extra_args:
+        a.extend(o.extra_args)
+    return a
+
+
+_temp_dirs: list[str] = []
+
+
+def _write_temp_json(v: Any) -> str:
+    d = tempfile.mkdtemp(prefix="mc-sdk-")
+    _temp_dirs.append(d)
+    p = os.path.join(d, "schema.json")
+    with open(p, "w", encoding="utf-8") as f:
+        json.dump(v, f)
+    return p
+
+
+def _resolve_command(o: CliOptions) -> list[str]:
+    c = o.path_to_mc or os.environ.get("MC_PATH") or "mc"
+    if isinstance(c, list):
+        if not c:
+            raise ValueError("path_to_mc must not be empty")
+        return list(c)
+    found = shutil.which(c) if os.sep not in c else c
+    return [found or c]
+
+
+# ── Client ───────────────────────────────────────────────────────────
+
+_END = object()
+
+
+class MonkeysCodeClient:
+    """
+    A long-lived ``mc`` session. Use as an async context manager.
+
+    ``send()`` queues a user message; ``receive()`` yields every event;
+    ``receive_turn()`` yields until (and including) the next result;
+    ``interrupt()`` stops the current turn; ``close()`` ends stdin and waits.
+    """
+
+    def __init__(self, options: CliOptions | None = None) -> None:
+        self.options = options or CliOptions()
+        self._proc: asyncio.subprocess.Process | None = None
+        self._events: asyncio.Queue[object] = asyncio.Queue()
+        self._tasks: list[asyncio.Task[None]] = []
+        self._stderr_tail = ""
+        self._interrupts = 0
+        self._sent = 0
+        self._results = 0
+        self._ending = False
+        self._stdin_closed = False
+        self._write_lock = asyncio.Lock()
+        self.session_id: str | None = None
+        self.last_result: ResultMessage | None = None
+        self.exit_code: int | None = None
+        self._error: BaseException | None = None
+
+    async def __aenter__(self) -> MonkeysCodeClient:
+        await self.connect()
+        return self
+
+    async def __aexit__(self, *exc: object) -> None:
+        await self.close()
+
+    async def connect(self) -> None:
+        cmd = _resolve_command(self.options)
+        argv = [*cmd, *build_args(self.options)]
+        env = dict(os.environ)
+        for k, v in (self.options.env or {}).items():
+            if v is None:
+                env.pop(k, None)
+            else:
+                env[k] = v
+        try:
+            self._proc = await asyncio.create_subprocess_exec(
+                *argv,
+                cwd=self.options.cwd,
+                env=env,
+                stdin=asyncio.subprocess.PIPE,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                limit=16 * 1024 * 1024,
+            )
+        except (FileNotFoundError, PermissionError) as err:
+            raise McProcessError(
+                f"could not start {cmd[0]}: {err}. Install the CLI "
+                "(npm i -g monkeyscode-cli) or set path_to_mc / MC_PATH.",
+                None,
+            ) from err
+        self._tasks = [
+            asyncio.create_task(self._read_stdout()),
+            asyncio.create_task(self._read_stderr()),
+        ]
+
+    # ── writing ────────────────────────────────────────────────
+
+    async def _write(self, obj: Mapping[str, Any]) -> None:
+        proc = self._proc
+        if proc is None or proc.stdin is None or self._stdin_closed:
+            return
+        async with self._write_lock:
+            try:
+                proc.stdin.write((json.dumps(obj) + "\n").encode())
+                await proc.stdin.drain()
+            except (BrokenPipeError, ConnectionResetError):
+                self._stdin_closed = True
+
+    async def send(self, message: str) -> None:
+        """Queue a user message."""
+        if self._ending:
+            raise McError("send() after end()")
+        self._sent += 1
+        await self._write({"type": "user", "message": {"role": "user", "content": message}})
+
+    async def interrupt(self) -> None:
+        """Stop the current turn (its result has subtype ``error_interrupted``)."""
+        self._interrupts += 1
+        await self._write(
+            {
+                "type": "control_request",
+                "request_id": f"int_{self._interrupts}",
+                "request": {"subtype": "interrupt"},
+            },
+        )
+
+    async def end(self) -> None:
+        """No more messages. stdin closes once every sent message has its result."""
+        self._ending = True
+        await self._maybe_close_stdin()
+
+    async def _maybe_close_stdin(self) -> None:
+        proc = self._proc
+        if (
+            self._ending
+            and not self._stdin_closed
+            and self._results >= self._sent
+            and proc
+            and proc.stdin
+        ):
+            self._stdin_closed = True
+            with contextlib.suppress(Exception):
+                proc.stdin.close()
+
+    # ── reading ────────────────────────────────────────────────
+
+    async def _answer(self, e: Mapping[str, Any]) -> None:
+        req = e.get("request") or {}
+        if req.get("subtype") != "can_use_tool":
+            return
+        suggestions = list(req.get("suggestions") or [])
+        cb = self.options.can_use_tool
+        r: PermissionResult
+        try:
+            if cb is not None:
+                ctx = ToolPermissionContext(
+                    description=str(req.get("description", "")),
+                    suggestions=suggestions,
+                    tool_use_id=req.get("tool_use_id"),
+                )
+                r = await cb(str(req.get("tool_name", "")), dict(req.get("input") or {}), ctx)
+            elif self.options.auto_approve:
+                r = PermissionAllow()
+            else:
+                hint = (
+                    f' To allow it, add allowed_tools=["{suggestions[0]}"].' if suggestions else ""
+                )
+                r = PermissionDeny(
+                    message=(
+                        f"Not approved (headless): requires approval.{hint} "
+                        "Continue without this action if possible."
+                    ),
+                )
+        except Exception as err:
+            r = PermissionDeny(message=f"Permission callback failed: {err}")
+        if isinstance(r, PermissionAllow):
+            resp: dict[str, Any] = {"behavior": "allow"}
+            if r.updated_input is not None:
+                resp["updated_input"] = r.updated_input
+        else:
+            resp = {"behavior": "deny"}
+            if r.message:
+                resp["message"] = r.message
+        await self._write(
+            {"type": "control_response", "request_id": e.get("request_id"), "response": resp},
+        )
+
+    async def _read_stdout(self) -> None:
+        assert self._proc and self._proc.stdout
+        try:
+            while True:
+                line = await self._proc.stdout.readline()
+                if not line:
+                    break
+                text = line.decode("utf-8", "replace").strip()
+                if not text:
+                    continue
+                try:
+                    raw = json.loads(text)
+                except json.JSONDecodeError:
+                    if self.options.stderr:
+                        self.options.stderr(f"[stdout] {text}")
+                    continue
+                t = raw.get("type")
+                if t == "control_request":
+                    # Answer concurrently: a slow callback must not stall the stream.
+                    self._tasks.append(asyncio.create_task(self._answer(raw)))
+                    continue
+                if t == "control_response":
+                    continue
+                ev = parse_event(raw)
+                if isinstance(ev, InitMessage):
+                    if ev.schema_version != SUPPORTED_SCHEMA_VERSION:
+                        self._error = McSchemaError(ev.schema_version, ev.cli_version)
+                        with contextlib.suppress(ProcessLookupError):
+                            self._proc.kill()
+                        break
+                    self.session_id = ev.session_id
+                if isinstance(ev, ResultMessage):
+                    self.last_result = ev
+                    self._results += 1
+                await self._events.put(ev)
+                if isinstance(ev, ResultMessage):
+                    await self._maybe_close_stdin()
+        finally:
+            self.exit_code = await self._proc.wait()
+            for d in _temp_dirs:
+                shutil.rmtree(d, ignore_errors=True)
+            _temp_dirs.clear()
+            if self._error is None and self.session_id is None:
+                self._error = McProcessError(
+                    f"mc exited ({self.exit_code}) before system/init",
+                    self.exit_code,
+                    self._stderr_tail.strip(),
+                )
+            await self._events.put(_END)
+
+    async def _read_stderr(self) -> None:
+        assert self._proc and self._proc.stderr
+        while True:
+            line = await self._proc.stderr.readline()
+            if not line:
+                return
+            s = line.decode("utf-8", "replace").rstrip("\n")
+            self._stderr_tail = (self._stderr_tail + s + "\n")[-4000:]
+            if self.options.stderr:
+                self.options.stderr(s)
+
+    async def receive(self) -> AsyncIterator[Event]:
+        """Every event until the CLI exits."""
+        while True:
+            ev = await self._events.get()
+            if ev is _END:
+                self._events.put_nowait(_END)  # later receive() calls end too
+                if self._error is not None:
+                    raise self._error
+                return
+            yield ev  # type: ignore[misc,unused-ignore]
+
+    async def receive_turn(self) -> AsyncIterator[Event]:
+        """Events up to and including the next result."""
+        async for ev in self.receive():
+            yield ev
+            if isinstance(ev, ResultMessage):
+                return
+
+    async def receive_result(self) -> ResultMessage:
+        """Skip to the next result."""
+        async for ev in self.receive_turn():
+            if isinstance(ev, ResultMessage):
+                return ev
+        raise McProcessError(
+            f"mc exited ({self.exit_code}) without a result",
+            self.exit_code,
+            self._stderr_tail.strip(),
+        )
+
+    async def close(self) -> None:
+        """End input and wait for the CLI to exit (killed after 10 s)."""
+        self._ending = True
+        proc = self._proc
+        if proc is None:
+            return
+        if not self._stdin_closed and proc.stdin:
+            self._stdin_closed = True
+            with contextlib.suppress(Exception):
+                proc.stdin.close()
+        try:
+            await asyncio.wait_for(proc.wait(), timeout=10)
+        except asyncio.TimeoutError:
+            with contextlib.suppress(ProcessLookupError):
+                proc.kill()
+            await proc.wait()
+        for t in self._tasks:
+            if not t.done():
+                with contextlib.suppress(asyncio.TimeoutError, asyncio.CancelledError):
+                    await asyncio.wait_for(t, timeout=2)
+
+
+async def query(
+    prompt: str | AsyncIterable[str],
+    options: CliOptions | None = None,
+) -> AsyncIterator[Event]:
+    """
+    Run a prompt (or a stream of prompts) and yield every event.
+
+    With a string, the session ends after its result. With an async iterable,
+    each item is a turn; the session ends when the iterable is exhausted.
+    """
+    client = MonkeysCodeClient(options)
+    await client.connect()
+    feeder: asyncio.Task[None] | None = None
+    try:
+        if isinstance(prompt, str):
+            await client.send(prompt)
+            await client.end()
+        else:
+
+            async def feed() -> None:
+                async for p in prompt:
+                    await client.send(p)
+                await client.end()
+
+            feeder = asyncio.create_task(feed())
+        async for ev in client.receive():
+            yield ev
+    finally:
+        if feeder and not feeder.done():
+            feeder.cancel()
+        await client.close()