code-puppy-core-plugins 0.0.2__py3-none-any.whl

code_puppy_core_plugins/acp/register_callbacks.py

@@ -0,0 +1,110 @@
+"""Plugin registration: the ``--acp`` entry point.
+
+This is the only place the ACP plugin touches Code Puppy's CLI. It:
+  1. contributes a ``--acp`` flag via the ``register_cli_args`` hook, and
+  2. in ``handle_cli_args``, when ``--acp`` is present, boots the ACP server
+     over stdio and returns the short-circuit sentinel so the normal
+     interactive TUI never starts.
+
+Per the contributing guide, all of this rides on existing hooks — no edits to
+``code_puppy/command_line`` or the core CLI runner.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, Dict, Optional
+
+from code_puppy.callbacks import register_callback
+
+logger = logging.getLogger(__name__)
+
+
+def _register_cli_args(parser: Any) -> None:
+    """Add the ``--acp`` flag to Code Puppy's argument parser."""
+    parser.add_argument(
+        "--acp",
+        action="store_true",
+        help="Run as a native agent over the Agent Client Protocol (ACP), speaking JSON-RPC over stdio.",
+    )
+
+
+def _handle_cli_args(args: Any) -> Optional[Dict[str, Any]]:
+    """Boot the ACP server when ``--acp`` is set; otherwise stand down.
+
+    Returns the ``{"handled": True, "exit_code": ...}`` sentinel so the CLI
+    runner exits after the ACP session ends instead of falling through to the
+    interactive TUI. Any other invocation returns ``None`` (not ours).
+
+    This hook is invoked *synchronously from within* Code Puppy's already-
+    running ``asyncio.run(main())`` loop, so we cannot start a nested loop
+    here. Instead we run the server on its own event loop in a dedicated
+    thread and block until it finishes — which keeps ``main()`` parked so the
+    TUI never starts.
+    """
+    if not getattr(args, "acp", False):
+        return None
+
+    import threading
+
+    box: Dict[str, int] = {"exit_code": 0}
+
+    def _run() -> None:
+        try:
+            box["exit_code"] = asyncio.run(_serve())
+        except Exception:  # noqa: BLE001 - never leak a traceback onto stdout
+            logger.exception("ACP server crashed")
+            box["exit_code"] = 1
+
+    thread = threading.Thread(target=_run, name="acp-server")
+    thread.start()
+    thread.join()
+    return {"handled": True, "exit_code": box["exit_code"]}
+
+
+async def _serve() -> int:
+    """Bind stdio via the ACP SDK and run one connection to completion.
+
+    Before serving we protect stdout — the JSON-RPC channel — by pointing
+    Code Puppy's streaming console and the root logger at stderr. In ACP mode
+    the interactive path that normally calls ``set_streaming_console`` is
+    short-circuited, so without this the streaming console would fall back to a
+    plain ``Console()`` on stdout and corrupt the protocol.
+
+    ``acp.run_agent`` owns the wire (stdio binding, JSON-RPC framing, param
+    parsing); we hand it a ``CodePuppyAgent`` and tear down the approval + I/O
+    backends afterwards so the process leaves no global seams set.
+    """
+    import sys
+
+    from acp import run_agent
+    from rich.console import Console
+
+    from code_puppy.agents.event_stream_handler import set_streaming_console
+    from code_puppy_core_plugins.acp import io_delegation, permissions, state
+    from code_puppy_core_plugins.acp.agent import CodePuppyAgent
+
+    # Route all incidental console output to stderr; stdout is JSON-RPC only.
+    set_streaming_console(Console(file=sys.stderr))
+    logging.basicConfig(stream=sys.stderr, level=logging.WARNING)
+
+    agent = CodePuppyAgent()
+    try:
+        await run_agent(agent)
+    finally:
+        agent.shutdown()
+        permissions.uninstall()
+        io_delegation.uninstall()
+        state.set_connection(None, None)
+    return 0
+
+
+def register() -> None:
+    """Register the ACP CLI hooks."""
+    register_callback("register_cli_args", _register_cli_args)
+    register_callback("handle_cli_args", _handle_cli_args)
+    logger.debug("ACP plugin callbacks registered")
+
+
+register()

code_puppy_core_plugins/acp/replay.py

@@ -0,0 +1,114 @@
+"""Replay a rehydrated session's history back to the client on load/resume.
+
+The ACP ``session/load`` (and ``session/resume``) contract is not just "rebuild
+the agent's memory" -- the client rebuilds its *own* thread UI from the
+``session/update`` notifications the agent streams while handling the request.
+An agent that only rehydrates its internal ``_message_history`` (and streams
+nothing) leaves the client with an empty thread: the conversation looks gone,
+and clients that treat an empty replayed thread as dead will discard it
+outright.
+
+So on load/resume we walk the rehydrated pydantic-ai message history and emit
+one ``session/update`` per visible turn -- user prompts, assistant text,
+assistant thinking, and past tool calls (as already-``completed`` entries) --
+in order, before the load response returns. System prompts and tool-return
+plumbing are intentionally skipped: they are model-plumbing, not conversation
+the user needs to see re-rendered.
+
+Everything here is best-effort: a replay hiccup must never fail the load, so
+per-update send failures are logged and swallowed.
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from typing import Any, Iterable, List
+
+from acp.helpers import (
+    start_tool_call,
+    update_agent_message_text,
+    update_agent_thought_text,
+    update_user_message_text,
+)
+
+from code_puppy_core_plugins.acp import state
+
+logger = logging.getLogger(__name__)
+
+
+async def replay_history(session_id: str, history: Iterable[Any]) -> None:
+    """Stream ``history`` back to the client as ordered ``session/update``s.
+
+    No-op when there is no connection or nothing to replay.
+    """
+    connection = state.get_connection()
+    if connection is None:
+        return
+    for message in history or []:
+        for update in _updates_for_message(message):
+            try:
+                await connection.session_update(session_id, update)
+            except Exception:  # noqa: BLE001
+                logger.debug("ACP: history replay update failed", exc_info=True)
+
+
+def _updates_for_message(message: Any) -> List[Any]:
+    """Map one pydantic-ai message to zero or more ACP updates.
+
+    Discriminates on each part's ``part_kind`` (stable across pydantic-ai
+    versions) rather than importing concrete part classes, so a version bump
+    that renames a class doesn't silently drop history.
+    """
+    updates: List[Any] = []
+    for part in getattr(message, "parts", None) or []:
+        kind = getattr(part, "part_kind", None)
+        if kind == "user-prompt":
+            text = _text_of(getattr(part, "content", None))
+            if text:
+                updates.append(update_user_message_text(text))
+        elif kind == "text":
+            text = _text_of(getattr(part, "content", None))
+            if text:
+                updates.append(update_agent_message_text(text))
+        elif kind == "thinking":
+            text = _text_of(getattr(part, "content", None))
+            if text:
+                updates.append(update_agent_thought_text(text))
+        elif kind == "tool-call":
+            updates.append(_tool_call_update(part))
+        # system-prompt / tool-return: model plumbing, not replayed.
+    return updates
+
+
+def _tool_call_update(part: Any) -> Any:
+    """Render a past tool call as an already-``completed`` tool-call entry.
+
+    Replayed tool calls are historical, so they open and close in one update
+    (status ``completed``). A fresh id is minted -- the original call id isn't
+    needed since nothing will update this entry again.
+    """
+    tool_name = str(getattr(part, "tool_name", "") or "tool")
+    return start_tool_call(
+        f"tc_{uuid.uuid4().hex[:12]}",
+        tool_name,
+        kind="other",
+        status="completed",
+    )
+
+
+def _text_of(content: Any) -> str:
+    """Flatten a part's content to plain text.
+
+    Handles the string case and pydantic-ai's multimodal list form (where a
+    user prompt is a list of str + binary/image blocks): string items are
+    kept, non-text blocks are dropped, since replay is a text reconstruction.
+    """
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content.strip()
+    if isinstance(content, (list, tuple)):
+        parts = [item for item in content if isinstance(item, str)]
+        return "\n".join(p for p in parts if p).strip()
+    return str(content).strip()

code_puppy_core_plugins/acp/session.py

@@ -0,0 +1,329 @@
+"""``ACPSession`` — one client thread bound to one Code Puppy agent.
+
+Each ACP session id maps to a single ``BaseAgent`` instance whose
+``_message_history`` persists across ``session/prompt`` calls, giving native
+multi-turn conversations. The session also carries the ``cwd`` the client
+opened the thread in, so tools resolve paths relative to the right project.
+
+Cancellation (``session/cancel``) is real: the prompt runs as an
+``asyncio.Task`` we hold a handle to, so a cancel notification cancels that
+task. pydantic-ai surfaces the ``CancelledError``, which we translate into the
+``cancelled`` stop reason for the in-flight ``session/prompt`` reply.
+
+Each turn also reports token usage (mapped from the pydantic-ai run result)
+back on the ``PromptResponse`` so the client can show cost/usage.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from acp.schema import Usage
+
+from code_puppy_core_plugins.acp import commands, content, persistence, state
+
+if TYPE_CHECKING:
+    from code_puppy.agents.base_agent import BaseAgent
+
+logger = logging.getLogger(__name__)
+
+# ACP stop reasons a ``session/prompt`` may resolve to.
+STOP_END_TURN = "end_turn"
+STOP_CANCELLED = "cancelled"
+STOP_REFUSAL = "refusal"
+
+
+@dataclass
+class PromptResult:
+    """Outcome of one turn: an ACP stop reason plus optional token usage."""
+
+    stop_reason: str
+    usage: Optional[Usage] = None
+
+
+class ACPSession:
+    """State + run lifecycle for a single client agent thread."""
+
+    def __init__(
+        self,
+        session_id: str,
+        agent: "BaseAgent",
+        cwd: Optional[str] = None,
+        additional_directories: Optional[List[str]] = None,
+        mcp_specs: Optional[List[Any]] = None,
+    ) -> None:
+        self.session_id = session_id
+        self.agent = agent
+        self.cwd = cwd
+        self.additional_directories = list(additional_directories or [])
+        # Raw client-injected ACP MCP specs, retained so they can be re-attached
+        # if the session's agent is rebuilt (e.g. on ``session/set_mode``).
+        self.mcp_specs = list(mcp_specs or [])
+        # Set while a ``session/prompt`` is in flight so ``cancel`` has a task
+        # to cancel. ``None`` means idle.
+        self._task: Optional["asyncio.Task[Any]"] = None
+
+    async def prompt(self, blocks: List[Any]) -> PromptResult:
+        """Run the agent on a user turn and return the stop reason + usage.
+
+        Flow:
+          1. Flatten ACP content blocks into the user text.
+          2. Point the messaging session context + run context at this session
+             so streamed deltas and tool calls are tagged/correlated.
+          3. Run ``agent.run_with_mcp(text)`` as a cancellable task — history
+             stays on the agent instance, so the next turn continues.
+          4. If nothing streamed (streaming disabled/empty), send the final
+             result text as one ``agent_message_chunk`` so the client sees it.
+        """
+        from code_puppy.messaging import set_session_context
+
+        parsed = content.parse_prompt(blocks)
+        text = parsed.text
+        if not text.strip() and not parsed.attachments and not parsed.link_attachments:
+            return PromptResult(STOP_END_TURN)
+        # A leading /slash command runs via Code Puppy's command handler (not the
+        # model); if it expands into a prompt (string result), run the model on it,
+        # exactly as cli_runner does with a string command result.
+        if (
+            not parsed.attachments
+            and not parsed.link_attachments
+            and commands.is_command(text)
+        ):
+            expanded = await commands.run(self.session_id, text)
+            if not expanded:
+                return PromptResult(STOP_END_TURN)
+            text = expanded
+        # Resolve relative paths against this session's own cwd -- no process-global
+        # os.chdir (would corrupt the SDK's I/O + any concurrent session).
+        from code_puppy.tools.common import (
+            reset_working_directory,
+            set_working_directory,
+        )
+        from code_puppy.tools.subagent_context import (
+            reset_conversation_root_id,
+            set_conversation_root_id,
+        )
+
+        cwd_token = set_working_directory(self.cwd) if self.cwd else None
+        # ContextVar, not set_session_context: conversation_scope reads it and needs
+        # real per-task isolation across concurrent ACP sessions, which the shared
+        # mutable attribute can't provide.
+        root_id_token = set_conversation_root_id(self.session_id)
+
+        set_session_context(self.session_id)
+        state.begin_run(self.session_id)
+        result: Any = None
+        error: Optional[BaseException] = None
+        try:
+            self._task = asyncio.ensure_future(
+                self.agent.run_with_mcp(
+                    text,
+                    attachments=parsed.attachments or None,
+                    link_attachments=parsed.link_attachments or None,
+                )
+            )
+            result = await self._task
+            self._absorb_history(result)
+            stop_reason = STOP_END_TURN
+        except asyncio.CancelledError:
+            # Two ways here: our own cancel cancelled the inner task (report
+            # ``cancelled``), or THIS coroutine was cancelled (SDK teardown) while the
+            # task lives on -> cancel it so it can't outlive us, and propagate.
+            task = self._task
+            if task is not None and not task.cancelled():
+                task.cancel()
+                raise
+            stop_reason = STOP_CANCELLED
+        except Exception as exc:  # noqa: BLE001
+            logger.exception("ACP: agent run failed")
+            error = exc
+            stop_reason = STOP_REFUSAL
+        finally:
+            streamed = state.streamed_text()
+            self._task = None
+            state.end_run()
+            set_session_context(None)
+            reset_conversation_root_id(root_id_token)
+            if cwd_token is not None:
+                reset_working_directory(cwd_token)
+
+        # Persist off the event loop so pickling a large history can't stall other
+        # sessions' streaming. Best-effort; never fails the turn.
+        try:
+            await asyncio.to_thread(
+                persistence.save,
+                self.session_id,
+                self.agent,
+                self.cwd,
+                self.additional_directories,
+            )
+        except Exception:  # noqa: BLE001
+            logger.debug("ACP: async persist failed", exc_info=True)
+
+        if stop_reason == STOP_END_TURN and not streamed and result is not None:
+            await self._send_final_result(result)
+        elif stop_reason == STOP_REFUSAL:
+            await self._send_error_notice(error)
+        return PromptResult(stop_reason, _to_acp_usage(result))
+
+    def _absorb_history(self, result: Any) -> None:
+        """Fold a completed run's full message list back into the agent.
+
+        ``run_with_mcp`` (the shared runtime) does NOT write the turn's
+        request+response back onto ``agent._message_history`` on the normal
+        path -- the caller must, exactly as ``cli_runner`` does after each
+        interactive turn. Without this the agent forgets every turn the moment
+        it ends: the next prompt runs with empty history, and persistence saves
+        a history containing only the user's prompt (no assistant reply). So we
+        replace the agent's history with ``result.all_messages()`` here, which
+        is what makes ACP multi-turn memory *and* load/resume replay real.
+        """
+        if result is None:
+            return
+        try:
+            messages = result.all_messages()
+        except Exception:  # noqa: BLE001
+            return
+        if messages:
+            self.agent.set_message_history(list(messages))
+
+    def cancel(self) -> None:
+        """Cancel the in-flight run, if any. No-op when idle.
+
+        Besides cancelling the asyncio task, force-kill any local shell
+        processes the run spawned so they don't orphan. (Shells delegated to
+        the client run client-side and are unaffected.)
+
+        Only *requests* cancellation -- delivery happens at the task's next
+        await point, whenever that is, not immediately. Callers that need
+        the run to have actually finished unwinding (its ``finally`` block
+        included) before they act should use :meth:`cancel_and_wait` instead.
+        """
+        task = self._task
+        if task is not None and not task.done():
+            task.cancel()
+        try:
+            from code_puppy.tools.command_runner import (
+                kill_all_running_shell_processes,
+            )
+
+            kill_all_running_shell_processes()
+        except Exception:  # noqa: BLE001
+            logger.debug("ACP: shell kill on cancel failed", exc_info=True)
+
+    async def cancel_and_wait(self, timeout: float = 2.0) -> None:
+        """Cancel the in-flight run and wait for it to actually unwind.
+
+        ``cancel()`` alone only *requests* cancellation; asyncio only
+        delivers it at the task's next await point, with no guarantee of
+        when that is. A caller that needs the run's own cleanup to have
+        genuinely finished first -- notably ``close_session`` purging this
+        session's warn-dedup bucket, which would otherwise race a
+        still-in-flight ``_invoke_agent_impl`` synchronously writing into
+        that same bucket a moment later -- must await this instead of the
+        bare ``cancel()``.
+
+        Bounded by ``timeout`` so a run stuck ignoring cancellation (blocked
+        in genuinely non-cancellable I/O, say) can't hang session close
+        forever; on timeout we give up waiting and return anyway -- the task
+        is still cancelled, just not confirmed finished.
+        """
+        task = self._task
+        self.cancel()
+        if task is None:
+            return
+        try:
+            await asyncio.wait_for(task, timeout=timeout)
+        except (asyncio.CancelledError, asyncio.TimeoutError):
+            pass
+        except Exception:  # noqa: BLE001
+            # prompt()'s own except block already logs/handles run failures;
+            # a close shouldn't raise because the run it interrupted failed.
+            logger.debug("ACP: cancel_and_wait observed run failure", exc_info=True)
+
+    async def _send_error_notice(self, error: Optional[BaseException]) -> None:
+        """Tell the client a turn failed, so it isn't a silent empty response.
+
+        A ``refusal`` stop reason alone renders as nothing in most clients; a
+        short message chunk makes the failure visible without leaking a stack
+        trace (the full traceback is logged to stderr).
+        """
+        from acp.helpers import update_agent_message_text
+
+        connection = state.get_connection()
+        if connection is None:
+            return
+        detail = f": {error}" if error else "."
+        try:
+            await connection.session_update(
+                self.session_id,
+                update_agent_message_text(f"\u26a0\ufe0f The agent run failed{detail}"),
+            )
+        except Exception:  # noqa: BLE001
+            logger.debug("ACP: failed to send error notice", exc_info=True)
+
+    async def _send_final_result(self, result: Any) -> None:
+        """Emit the run's final text when nothing streamed during the turn."""
+        from acp.helpers import update_agent_message_text
+
+        final_text = self._extract_result_text(result)
+        if not final_text:
+            return
+        connection = state.get_connection()
+        if connection is None:
+            return
+        try:
+            await connection.session_update(
+                self.session_id, update_agent_message_text(final_text)
+            )
+        except Exception:  # noqa: BLE001
+            logger.exception("ACP: failed to send final result")
+
+    @staticmethod
+    def _extract_result_text(result: Any) -> str:
+        """Pull display text out of a pydantic-ai run result.
+
+        Mirrors ``_runtime``'s output extraction: prefer ``output``, then the
+        legacy ``data`` attribute, then ``str()``.
+        """
+        for attr in ("output", "data"):
+            value = getattr(result, attr, None)
+            if value:
+                return str(value)
+        return str(result) if result else ""
+
+
+def _to_acp_usage(result: Any) -> Optional[Usage]:
+    """Map a pydantic-ai run result's usage onto ACP's ``Usage`` model.
+
+    Field names have drifted across pydantic-ai versions, so each ACP field is
+    resolved from a list of candidate attributes. Returns ``None`` when no
+    usage is available so we simply omit it from the response.
+    """
+    if result is None:
+        return None
+    try:
+        usage = result.usage()
+    except Exception:  # noqa: BLE001
+        return None
+    if usage is None:
+        return None
+
+    def pick(*names: str) -> Optional[int]:
+        for name in names:
+            value = getattr(usage, name, None)
+            if value is not None:
+                return int(value)
+        return None
+
+    return Usage(
+        input_tokens=pick("input_tokens", "request_tokens", "prompt_tokens"),
+        output_tokens=pick("output_tokens", "response_tokens", "completion_tokens"),
+        total_tokens=pick("total_tokens"),
+        cached_read_tokens=pick("cache_read_tokens", "cached_read_tokens"),
+        cached_write_tokens=pick("cache_write_tokens", "cached_write_tokens"),
+        thought_tokens=pick("thinking_tokens", "thought_tokens", "reasoning_tokens"),
+    )