ags-cli 0.1.0__py3-none-any.whl

harness/adapters/claude_code.py

@@ -0,0 +1,522 @@
+"""ClaudeCodeAdapter — drives the official ``claude`` CLI (verified against
+v2.1.187).
+
+Invocation (no token hijacking — the CLI uses its own supported auth; an optional
+``ANTHROPIC_API_KEY`` is injected from a ``secret_ref``):
+
+    claude --print --output-format stream-json --verbose \
+           --model <model> --permission-mode <mode> --add-dir <cwd> \
+           --session-id <run_id> --append-system-prompt <system+memory> "<message>"
+
+``--output-format stream-json`` emits NDJSON; we map each line onto the
+``NormalizedEvent`` union. The terminal ``{"type":"result", ...}`` object (schema
+confirmed empirically) carries the final text, ``session_id``, ``total_cost_usd``
+and token ``usage``.
+
+Warm-pool path (HARNESS_WARM_CLI=1):
+  When enabled, same-session turns are routed through ``WarmCliPool`` so the
+  ``claude`` process is kept alive between turns and the CLI's KV cache is reused,
+  reducing turn-2+ latency.
+
+  Respawn-resumes-on-disk (verified empirically, claude 2.1.190): a FRESH
+    ``claude --print --verbose --input-format stream-json --output-format
+    stream-json --resume <session_id>`` process RESUMES the on-disk conversation
+    created by a previous (now-dead) process — the resumed process correctly
+    recalls turn-1 content. So a warm-proc respawn after an idle-reap does NOT
+    need engine-side transcript replay: the warm argv includes ``--resume
+    <native_session_id>`` whenever a native id is known. The first-ever spawn of a
+    session has no native id → fresh conversation; the engine captures the
+    ``session_id`` from the init event as before; any respawn gets the native id
+    from the engine and resumes the on-disk conversation.
+
+  Fail-closed mid-turn: a warm turn that dies AFTER we have already yielded parsed
+    events (a partial answer reached the bus) does NOT fall back to a fresh spawn
+    (that would double-emit the answer to the ledger/bus). Instead we emit an error
+    + failed status, kill the broken proc (next turn respawns and ``--resume``s),
+    and let the engine persist a failed run. Only a death BEFORE any event was
+    yielded falls back to the one-shot spawn path (the answer never reached the bus).
+"""
+
+from __future__ import annotations
+
+import contextlib
+import itertools
+import json
+import shutil
+import tempfile
+from collections.abc import AsyncIterator, Iterable
+from pathlib import Path
+
+from harness_sdk.base import Capabilities, Cost, ProviderSession, SessionContext, Usage
+from harness_sdk.cli_base import CliSubprocessAdapter, ParserState
+from harness.adapters.probe import probe_cli, supports_flag
+from harness_sdk.events import (
+    AdapterMetaEvent,
+    CostEvent,
+    ErrorEvent,
+    MessageEvent,
+    NormalizedEvent,
+    StatusEvent,
+    ToolCallEvent,
+    ToolResultEvent,
+)
+
+from harness.models.enums import AdapterKind
+from harness.observability.logging import get_logger
+
+_log = get_logger("claude-code-adapter")
+
+# Public Anthropic list price for Opus-class models ($/token). Override per
+# provider via ``params.price_in`` / ``params.price_out`` for other models.
+_DEFAULT_PRICE_IN = 15.0 / 1_000_000
+_DEFAULT_PRICE_OUT = 75.0 / 1_000_000
+
+
+class ClaudeCodeAdapter(CliSubprocessAdapter):
+    kind = AdapterKind.claude_code
+    default_use_pty = False  # claude streams fine on a plain pipe
+
+    def __init__(self, config) -> None:  # noqa: ANN001 - AdapterConfig
+        super().__init__(config)
+        # Per-key dirs holding generated .mcp.json (with resolved secrets); keyed by
+        # run_id for one-shot runs, by pool key for warm procs. One-shot dirs are
+        # removed in _reap; warm dirs via the pool's on_close callback.
+        self._mcp_dirs: dict[str, str] = {}
+        # Active warm turns: run_id → pool key, so cancel_run can reach the warm proc.
+        self._warm_runs: dict[str, str] = {}
+        self._probe: dict | None = None
+
+    async def ensure_probe(self) -> dict:
+        """Run ``probe_cli`` once, cache the result on the instance, and return it."""
+        if self._probe is not None:
+            return self._probe
+        try:
+            result = await probe_cli(self.binary)
+        except Exception:
+            return {"ok": False, "flags": set(), "version": None}
+        self._probe = result
+        return self._probe
+
+    @property
+    def binary(self) -> str:
+        return self.config.params.get("bin", "claude")
+
+    def env_overrides(self) -> dict[str, str]:
+        env: dict[str, str] = {}
+        if self.config.api_key:
+            env["ANTHROPIC_API_KEY"] = self.config.api_key
+        # Opt-in managed-skill home (off by default; relocating the config dir means the
+        # CLI uses the harness-supplied API key rather than the user's own login).
+        from harness.skills import provision_config_home
+        if (home := provision_config_home(self.config.name)):
+            env[self.config.params.get("config_dir_env", "CLAUDE_CONFIG_DIR")] = home
+        return env
+
+    def health_argv(self) -> list[str]:
+        return [self.binary, "--version"]
+
+    async def list_models(self) -> list[str]:
+        """The ``claude`` CLI has no machine-readable model list — no subcommand
+        enumerates it and ``--model`` only documents the alias forms — so we
+        surface a known set instead. It is a static snapshot: it does NOT track
+        the models your CLI login is actually entitled to, and it goes stale as
+        new models ship. Override it with ``params.models`` in config.yaml
+        (that list is used verbatim), which is also how you expose account- or
+        CLI-specific variants such as ``claude-fable-5[1m]``.
+
+        The CLI also accepts short aliases (``opus``/``sonnet``/``haiku``/
+        ``fable``) for ``--model``, which always resolve to the latest model in
+        that tier."""
+        configured = self.config.params.get("models")
+        if configured:
+            return list(configured)
+        return [
+            "claude-fable-5",
+            "claude-opus-4-8",
+            "claude-opus-4-7",
+            "claude-sonnet-5",
+            "claude-sonnet-4-6",
+            "claude-haiku-4-5",
+        ]
+
+    @staticmethod
+    def _image_attachments(ctx: SessionContext | None) -> list[dict]:
+        return list((ctx.extra or {}).get("attachments_images") or []) if ctx else []
+
+    def _user_content(self, message: str, ctx: SessionContext | None) -> list[dict]:
+        """Content-block array for a user turn: the text plus any image
+        attachments as base64 blocks. A vanished/unreadable file is skipped with
+        a warning — an attachment must never sink the run."""
+        import base64
+        import logging
+
+        content: list[dict] = [{"type": "text", "text": message}]
+        for att in self._image_attachments(ctx):
+            try:
+                data = Path(att["path"]).read_bytes()
+            except OSError as exc:
+                logging.getLogger(__name__).warning(
+                    "attachment image unreadable, skipping: %s (%s)", att.get("path"), exc
+                )
+                continue
+            content.append({
+                "type": "image",
+                "source": {
+                    "type": "base64",
+                    "media_type": att.get("media_type", "image/png"),
+                    "data": base64.b64encode(data).decode(),
+                },
+            })
+        return content
+
+    def stdin_payload(self, message: str, ctx: SessionContext | None = None) -> bytes | None:
+        # Feed the prompt via stdin so it can't be swallowed by the variadic
+        # ``--add-dir`` flag (verified failure mode with v2.1.187).
+        if self._image_attachments(ctx):
+            # Image blocks need the stream-json input format (see build_argv).
+            line = json.dumps({
+                "type": "user",
+                "message": {"role": "user", "content": self._user_content(message, ctx)},
+            })
+            return (line + "\n").encode()
+        return message.encode()
+
+    def _permission_mode(self, ctx: SessionContext) -> str:
+        """Resolve the effective ``--permission-mode`` for this context. Safe by
+        default (plan = no edits); write mode raises it to acceptEdits. The effective
+        mode is the per-run setting on the context (workspace mode), set by the engine.
+        """
+        p = self.config.params
+        mode = p.get("permission_mode", "plan")
+        if not ctx.read_only:
+            mode = p.get("write_permission_mode", "acceptEdits")
+        return mode
+
+    def build_argv(self, *, run_id: str, ctx: SessionContext, message: str) -> list[str]:
+        p = self.config.params
+        mode = self._permission_mode(ctx)
+        nsid = ctx.extra.get("native_session_id")
+        if nsid:
+            session_args = ["--resume", str(nsid)]
+        else:
+            session_args = ["--session-id", run_id]
+        # Image attachments ride as content blocks, which require stream-json input.
+        input_format = "stream-json" if self._image_attachments(ctx) else "text"
+        argv = [
+            self.binary, "--print",
+            "--output-format", "stream-json", "--verbose",
+            "--input-format", input_format,
+            "--permission-mode", mode,
+            *session_args,
+        ]
+        if self.config.model:
+            argv += ["--model", self.config.model]
+        if (cwd := self.working_dir(ctx)):
+            argv += ["--add-dir", cwd]
+        if p.get("dangerously_skip_permissions"):
+            argv.append("--dangerously-skip-permissions")
+        argv += self._mcp_argv(run_id, ctx)
+        system = "\n\n".join(x for x in (ctx.system_prompt, ctx.memory_context) if x)
+        if system:
+            argv += ["--append-system-prompt", system]
+        # Prompt is supplied via stdin (see stdin_payload) — not a positional arg.
+        return argv
+
+    def _mcp_argv(self, key: str, ctx: SessionContext) -> list[str]:
+        """Claude is its own MCP client: write the harness-resolved servers to a
+        ``.mcp.json`` and point the CLI at it. ``--strict-mcp-config`` makes the CLI
+        use ONLY these servers (ignoring any project/user ``.mcp.json``), and
+        ``--allowedTools`` pre-authorizes their tools so they run under plan/acceptEdits
+        without an interactive prompt.
+
+        ``key`` scopes the generated dir: a ``run_id`` for one-shot runs, a pool key
+        for warm procs. If a dir already exists for this key (e.g. a warm proc being
+        respawned into a new argv), it is removed first so resolved-secret temp dirs
+        don't leak."""
+        from harness.mcp import config_gen
+
+        servers = ctx.extra.get("mcp_servers") or []
+        if not servers:
+            return []
+        if self.remote_host(ctx):
+            # The generated .mcp.json is a local temp file the remote CLI can't
+            # read — MCP injection is a local-only feature for now.
+            _log.warning("mcp_skipped_remote", host=self.remote_host(ctx),
+                         servers=len(servers))
+            return []
+
+        # Probe gating for MCP config flag
+        explicit_flag = self.config.params.get("mcp_config_flag")
+        flag = explicit_flag or "--mcp-config"
+        if not explicit_flag and self._probe is not None:
+            if not supports_flag(self._probe, flag):
+                _log.warning("Skipping MCP servers: flag %r not found in %r help output", flag, self.binary)
+                return []
+
+        # Remove any prior dir for this key before creating a new one (leak guard).
+        if (old := self._mcp_dirs.pop(key, None)):
+            shutil.rmtree(old, ignore_errors=True)
+        d = tempfile.mkdtemp(prefix=f"ags-mcp-{key[:8]}-")
+        self._mcp_dirs[key] = d
+        path = config_gen.write_json(config_gen.claude_mcp_config(servers), d, "mcp.json")
+        argv = [flag, path, "--strict-mcp-config"]
+
+        # Check if --allowedTools is supported
+        if (patterns := config_gen.allowed_tool_patterns(servers)):
+            tools_flag = "--allowedTools"
+            if not explicit_flag and self._probe is not None and not supports_flag(self._probe, tools_flag):
+                _log.warning("Skipping allowed tools pre-authorization: flag %r not supported", tools_flag)
+            else:
+                argv += [tools_flag, " ".join(patterns)]
+        return argv
+
+    def _drop_mcp_dir(self, key: str) -> None:
+        """Remove the generated MCP dir for ``key`` (it may hold resolved secrets)."""
+        if (d := self._mcp_dirs.pop(key, None)):
+            shutil.rmtree(d, ignore_errors=True)
+
+    def build_warm_argv(self, *, ctx: SessionContext) -> list[str]:
+        """Build the argv for the warm-pool process: like ``build_argv`` but uses
+        ``--input-format stream-json`` (persistent stdin protocol).
+
+        When a ``native_session_id`` is present in ``ctx.extra`` (any turn after the
+        first-ever spawn of the session), the warm argv includes ``--resume
+        <native_session_id>`` so a freshly respawned process (after an idle-reap)
+        resumes the on-disk conversation created by the previous (dead) process —
+        verified empirically. The first-ever spawn has no native id → fresh
+        conversation; the engine captures the id from the init event as before. We
+        never pass ``--session-id`` on the warm path (a warm process tracks one
+        conversation for its lifetime; on respawn ``--resume`` restores it).
+        """
+        p = self.config.params
+        mode = self._permission_mode(ctx)
+        argv = [
+            self.binary, "--print",
+            "--output-format", "stream-json", "--verbose",
+            "--input-format", "stream-json",
+            "--permission-mode", mode,
+        ]
+        nsid = ctx.extra.get("native_session_id")
+        if nsid:
+            argv += ["--resume", str(nsid)]
+        if self.config.model:
+            argv += ["--model", self.config.model]
+        if (cwd := self.working_dir(ctx)):
+            argv += ["--add-dir", cwd]
+        if p.get("dangerously_skip_permissions"):
+            argv.append("--dangerously-skip-permissions")
+        # MCP: scope the generated .mcp.json to the pool key (session:mode) so it is
+        # created once per warm proc and torn down by the pool's on_close callback.
+        argv += self._mcp_argv(f"{ctx.session_id}:{mode}", ctx)
+        system = "\n\n".join(x for x in (ctx.system_prompt, ctx.memory_context) if x)
+        if system:
+            argv += ["--append-system-prompt", system]
+        return argv
+
+    async def stream_events(
+        self,
+        session: ProviderSession,
+        message: str,
+        *,
+        tools=None,
+        stream: bool = True,
+    ) -> AsyncIterator[NormalizedEvent]:
+        """Override to optionally route through the warm pool.
+
+        When ``settings.warm_cli`` is True the turn is routed through ``WarmCliPool``.
+        A warm-turn death BEFORE any parsed event has been yielded falls back to the
+        one-shot spawn path (the answer never reached the bus). A death AFTER events
+        were yielded is fail-closed: emit an error + failed status and kill the broken
+        proc (so the next turn respawns and ``--resume``s the on-disk conversation)
+        rather than double-emitting a second answer to the ledger/bus.
+        """
+        from harness.settings import get_settings
+
+        settings = get_settings()
+        # Remote (ssh) sessions always take the one-shot path: the warm pool spawns
+        # local processes, and Claude's session state lives on the remote host anyway
+        # (--resume works there across one-shot invocations).
+        remote = bool((session.meta.get("extra") or {}).get("ssh_host"))
+        if not settings.warm_cli or remote:
+            async for ev in super().stream_events(session, message, tools=tools, stream=stream):
+                yield ev
+            return
+
+        run_id = session.provider_session_id
+        extra = session.meta.get("extra") or {}
+        session_id = session.meta.get("session_id", run_id)
+
+        # Build context for warm argv.
+        ctx = SessionContext(
+            run_id=run_id,
+            session_id=session_id,
+            workspace_slug=session.meta.get("workspace_slug", "default"),
+            objective=message,
+            system_prompt=session.meta.get("system_prompt"),
+            memory_context=session.meta.get("memory_context", ""),
+            workspace_dir=session.meta.get("workspace_dir"),
+            read_only=session.meta.get("read_only", True),
+            extra=extra,
+        )
+        # Pool key includes the resolved permission mode so a plan→write switch spawns
+        # a separate warm proc (the old one is idle-reaped); context is restored on the
+        # fresh proc via --resume.
+        mode = self._permission_mode(ctx)
+        pool_key = f"{session_id}:{mode}"
+        argv = self.build_warm_argv(ctx=ctx)
+        env = self.env_overrides()
+
+        from harness.adapters.warm_pool import get_warm_pool
+
+        pool = get_warm_pool()
+        yielded_any = False
+        self._warm_runs[run_id] = pool_key
+        try:
+            proc = await pool.acquire(
+                pool_key, argv, env,
+                on_close=lambda: self._drop_mcp_dir(pool_key),
+            )
+            user_line = json.dumps({
+                "type": "user",
+                "message": {"role": "user", "content": self._user_content(message, ctx)},
+            })
+            seq = itertools.count()
+            state: ParserState = {"text": []}
+
+            def _tag(ev: NormalizedEvent) -> NormalizedEvent:
+                ev.seq = next(seq)
+                return ev
+
+            yield _tag(StatusEvent(seq=0, status="running"))
+            async for line in pool.send_turn(proc, user_line):
+                for ev in self.parse_chunk(line, state):
+                    yielded_any = True
+                    yield _tag(ev)
+            for ev in self.finalize(state):
+                yielded_any = True
+                yield _tag(ev)
+            yield _tag(StatusEvent(seq=0, status="succeeded"))
+            return
+        except Exception as exc:  # noqa: BLE001 — fail soft (pre-yield) / fail closed (mid-turn)
+            if yielded_any:
+                # A partial answer already reached the bus/ledger. Falling back to a
+                # fresh one-shot spawn would double-emit the answer — instead fail
+                # closed: kill the broken proc (next turn respawns + --resume) and let
+                # the engine persist a failed run. No duplicate answer ever reaches
+                # the ledger/bus.
+                _log.warning("warm_turn_failed_mid_stream", error=str(exc),
+                             session_id=session_id, key=pool_key)
+                with contextlib.suppress(Exception):
+                    await pool.kill(pool_key)
+                yield ErrorEvent(
+                    seq=0, message=f"warm turn failed mid-stream: {exc}", retriable=True
+                )
+                yield StatusEvent(seq=0, status="failed")
+                return
+            _log.warning("warm_pool_error_fallback", error=str(exc), session_id=session_id)
+            # Nothing yielded yet → safe to fall through to the normal spawn path.
+        finally:
+            self._warm_runs.pop(run_id, None)
+
+        async for ev in super().stream_events(session, message, tools=tools, stream=stream):
+            yield ev
+
+    async def cancel_run(self, run_id: str) -> None:
+        # If this run is an in-flight warm turn, kill its warm proc so the CLI stops
+        # generating; the next turn respawns and --resumes the on-disk conversation.
+        key = self._warm_runs.pop(run_id, None)
+        if key is not None:
+            from harness.adapters.warm_pool import get_warm_pool
+            with contextlib.suppress(Exception):
+                await get_warm_pool().kill(key)
+        # Also run the base cancellation for the one-shot spawn path.
+        await super().cancel_run(run_id)
+
+    async def _reap(self, proc, run_id: str) -> None:  # noqa: ANN001
+        # Remove the run's generated MCP config (it may hold resolved secrets) once the
+        # subprocess is gone, then run the base reap (process-group teardown).
+        self._drop_mcp_dir(run_id)
+        await super()._reap(proc, run_id)
+
+    def parse_chunk(self, line: str, state: ParserState) -> Iterable[NormalizedEvent]:
+        line = line.strip()
+        if not line:
+            return
+        try:
+            obj = json.loads(line)
+        except json.JSONDecodeError:
+            return  # non-JSON noise (shouldn't happen with stream-json)
+        t = obj.get("type")
+
+        if t == "system":
+            yield AdapterMetaEvent(seq=0, data={"subtype": obj.get("subtype"),
+                                                "session_id": obj.get("session_id")})
+        elif t == "assistant":
+            for block in obj.get("message", {}).get("content", []):
+                bt = block.get("type")
+                if bt == "text" and block.get("text"):
+                    state["had_message"] = True
+                    yield MessageEvent(seq=0, role="assistant", text=block["text"])
+                elif bt == "tool_use":
+                    yield ToolCallEvent(seq=0, name=block.get("name", ""),
+                                        arguments=block.get("input", {}) or {},
+                                        call_id=block.get("id"))
+        elif t == "user":
+            for block in obj.get("message", {}).get("content", []):
+                if block.get("type") == "tool_result":
+                    yield ToolResultEvent(
+                        seq=0, call_id=block.get("tool_use_id"),
+                        output=_stringify(block.get("content")),
+                        is_error=bool(block.get("is_error")),
+                    )
+        elif t == "result":
+            is_error = obj.get("is_error")
+            result_content = obj.get("result") or obj.get("content") or obj.get("text")
+            if is_error:
+                yield ErrorEvent(seq=0, message=str(result_content or "claude error"),
+                                 retriable=obj.get("api_error_status") in (429, 500, 503))
+            elif not state.get("had_message") and result_content:
+                state["had_message"] = True
+                yield MessageEvent(seq=0, role="assistant", text=str(result_content))
+            usage = obj.get("usage", {}) or {}
+            prompt_tokens = usage.get("input_tokens") or usage.get("prompt_tokens") or 0
+            completion_tokens = usage.get("output_tokens") or usage.get("completion_tokens") or 0
+            cost_usd = obj.get("total_cost_usd")
+            if cost_usd is None:
+                cost_usd = obj.get("cost_usd") or obj.get("cost")
+            yield CostEvent(
+                seq=0,
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=prompt_tokens + completion_tokens,
+                usd=cost_usd,
+            )
+
+    def list_capabilities(self) -> Capabilities:
+        return Capabilities(
+            streaming=True, tool_calling=True,
+            max_context=int(self.config.params.get("max_context", 200_000)),
+            reasoning=True, cost_known=True,
+            notes="Official Claude Code CLI (stream-json).",
+        )
+
+    def estimate_cost(self, usage: Usage) -> Cost | None:
+        p = self.config.params
+        price_in = float(p.get("price_in", _DEFAULT_PRICE_IN))
+        price_out = float(p.get("price_out", _DEFAULT_PRICE_OUT))
+        return Cost(
+            prompt_tokens=usage.prompt_tokens,
+            completion_tokens=usage.completion_tokens,
+            total_tokens=usage.prompt_tokens + usage.completion_tokens,
+            usd=usage.prompt_tokens * price_in + usage.completion_tokens * price_out,
+        )
+
+
+def _stringify(content: object) -> str:
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        return "\n".join(
+            b.get("text", "") if isinstance(b, dict) else str(b) for b in content
+        )
+    return str(content) if content is not None else ""