jupyter-ai-acp-client 0.0.1__py3-none-any.whl

jupyter_ai_acp_client/__init__.py

@@ -0,0 +1,25 @@
+try:
+    from ._version import __version__
+except ImportError:
+    # Fallback when using the package in dev mode without installing
+    # in editable mode with pip. It is highly recommended to install
+    # the package from a stable release or in editable mode: https://pip.pypa.io/en/stable/topics/local-project-installs/#editable-installs
+    import warnings
+    warnings.warn("Importing 'jupyter_ai_acp_client' outside a proper installation.")
+    __version__ = "dev"
+
+from .extension_app import JaiAcpClientExtension
+
+
+def _jupyter_labextension_paths():
+    return [{
+        "src": "labextension",
+        "dest": "@jupyter-ai/acp-client"
+    }]
+
+
+def _jupyter_server_extension_points():
+    return [{
+        "module": "jupyter_ai_acp_client",
+        "app": JaiAcpClientExtension
+    }]

jupyter_ai_acp_client/_version.py

@@ -0,0 +1,4 @@
+# This file is auto-generated by Hatchling. As such, do not:
+#   - modify
+#   - track in version control e.g. be sure to add to .gitignore
+__version__ = VERSION = '0.0.1'

jupyter_ai_acp_client/acp_personas/claude.py

@@ -0,0 +1,30 @@
+import shutil
+from jupyter_ai_persona_manager import PersonaRequirementsUnmet
+if shutil.which("claude-code-acp") is None:
+    raise PersonaRequirementsUnmet(
+        "This persona requires the Claude Code ACP adapter to be installed."
+        " Install it via `npm install -g @zed-industries/claude-code-acp`"
+        " then restart."
+    )
+
+import os
+from ..base_acp_persona import BaseAcpPersona
+from jupyter_ai_persona_manager import PersonaDefaults
+
+class ClaudeAcpPersona(BaseAcpPersona):
+    def __init__(self, *args, **kwargs):
+        executable = ["claude-code-acp"]
+        super().__init__(*args, executable=executable, **kwargs)
+    
+    @property
+    def defaults(self) -> PersonaDefaults:
+        avatar_path = str(os.path.abspath(
+            os.path.join(os.path.dirname(__file__), "..", "static", "claude.svg")
+        ))
+
+        return PersonaDefaults(
+            name="Claude-ACP",
+            description="Claude Code as an ACP agent persona.",
+            avatar_path=avatar_path,
+            system_prompt="unused"
+        )

jupyter_ai_acp_client/acp_personas/test.py

@@ -0,0 +1,26 @@
+import os
+import sys
+from ..base_acp_persona import BaseAcpPersona
+from jupyter_ai_persona_manager import PersonaDefaults
+
+class TestAcpPersona(BaseAcpPersona):
+    def __init__(self, *args, **kwargs):
+        # Get absolute path to agent.py
+        agent_path = os.path.abspath(
+            os.path.join(os.path.dirname(__file__), "..", "..", "examples", "agent.py")
+        )
+        executable = [sys.executable, agent_path]
+        super().__init__(*args, executable=executable, **kwargs)
+    
+    @property
+    def defaults(self) -> PersonaDefaults:
+        avatar_path = str(os.path.abspath(
+            os.path.join(os.path.dirname(__file__), "..", "static", "test.svg")
+        ))
+
+        return PersonaDefaults(
+            name="Test-ACP",
+            description="A test ACP persona",
+            avatar_path=avatar_path,
+            system_prompt="unused"
+        )

jupyter_ai_acp_client/base_acp_persona.py

@@ -0,0 +1,163 @@
+from jupyter_ai_persona_manager import BasePersona
+from jupyterlab_chat.models import Message
+import asyncio
+import sys
+from asyncio.subprocess import Process
+from typing import Awaitable, ClassVar
+from acp import NewSessionResponse
+from acp.schema import AvailableCommand
+
+from .default_acp_client import JaiAcpClient
+
+
+class BaseAcpPersona(BasePersona):
+    _subprocess_future: ClassVar[Awaitable[Process] | None] = None
+    """
+    The task that yields the agent subprocess once complete. This is a class
+    attribute because multiple instances of the same ACP persona may share an
+    ACP agent subprocess.
+    
+    Developers should always use `self.get_agent_subprocess()`.
+    """
+
+    _client_future: ClassVar[Awaitable[JaiAcpClient] | None] = None
+    """
+    The future that yields the ACP Client once complete. This is a class
+    attribute because multiple instances of the same ACP persona may share an
+    ACP client as well. ACP agent subprocesses and clients map 1-to-1.
+
+    Developers should always use `self.get_client()`.
+    """
+
+    _client_session_future: Awaitable[NewSessionResponse]
+    """
+    The future that yields the ACP client session info. Each instance of an ACP
+    persona has a unique session ID, i.e. each chat reserves a unique session.
+
+    Developers should always call `self.get_session()` or `self.get_session_id()`.
+    """
+
+    _acp_slash_commands: list[AvailableCommand]
+
+    def __init__(self, *args, executable: list[str], **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self._executable = executable
+
+        # Ensure each subclass has its own subprocess and client by checking if the
+        # class variable is defined directly on this class (not inherited)
+        if '_subprocess_future' not in self.__class__.__dict__ or self.__class__._subprocess_future is None:
+            self.__class__._subprocess_future = self.event_loop.create_task(
+                self._init_agent_subprocess()
+            )
+        if '_client_future' not in self.__class__.__dict__ or self.__class__._client_future is None:
+            self.__class__._client_future = self.event_loop.create_task(
+                self._init_client()
+            )
+
+        self._client_session_future = self.event_loop.create_task(
+            self._init_client_session()
+        )
+        self._acp_slash_commands = []
+
+    async def _init_agent_subprocess(self) -> Process:
+        process = await asyncio.create_subprocess_exec(
+            *self._executable,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=sys.stderr,
+        )
+        self.log.info(f"Spawned ACP agent subprocess for '{self.__class__.__name__}'.")
+        return process
+
+    async def _init_client(self) -> JaiAcpClient:
+        agent_subprocess = await self.get_agent_subprocess()
+        client = JaiAcpClient(agent_subprocess=agent_subprocess, event_loop=self.event_loop)
+        self.log.info(f"Initialized ACP client for '{self.__class__.__name__}'.")
+        return client
+    
+    async def _init_client_session(self) -> NewSessionResponse:
+        client = await self.get_client()
+        session = await client.create_session(persona=self)
+        self.log.info(
+            f"Initialized new ACP client session for '{self.__class__.__name__}'"
+            f" with ID '{session.session_id}'."
+        )
+        return session
+
+    async def get_agent_subprocess(self) -> asyncio.subprocess.Process:
+        """
+        Safely returns the ACP agent subprocess for this persona.
+        """
+        return await self.__class__._subprocess_future
+    
+    async def get_client(self) -> JaiAcpClient:
+        """
+        Safely returns the ACP client for this persona.
+        """
+        return await self.__class__._client_future
+    
+    async def get_session(self) -> NewSessionResponse:
+        """
+        Safely returns the ACP client session for this chat.
+        """
+        return await self._client_session_future
+    
+    async def get_session_id(self) -> str:
+        """
+        Safely returns the ACP client ID assigned to this chat.
+        """
+        session = await self._client_session_future
+        return session.session_id
+    
+    async def process_message(self, message: Message) -> None:
+        """
+        A default implementation for the `BasePersona.process_message()` method
+        for ACP agents.
+
+        This method may be overriden by child classes.
+        """
+        client = await self.get_client()
+        session_id = await self.get_session_id()
+
+        # TODO: add attachments!
+        prompt = message.body.replace("@" + self.as_user().mention_name, "").strip()
+        await client.prompt_and_reply(
+            session_id=session_id,
+            prompt=prompt,
+        )
+    
+    @property
+    def acp_slash_commands(self) -> list[AvailableCommand]:
+        """
+        Returns the list of slash commands advertised by the ACP agent in the
+        current session.
+
+        This initializes to an empty list, and should be updated **only** by the
+        ACP client upon receiving a `session/update` request containing an
+        `AvailableCommandsUpdate` payload from the ACP agent.
+        """
+        return self._acp_slash_commands
+    
+    @acp_slash_commands.setter
+    def acp_slash_commands(self, commands: list[AvailableCommand]):
+        self.log.info(
+            f"Setting {len(commands)} slash commands for '{self.name}' in room '{self.parent.room_id}'."
+        )
+        self._acp_slash_commands = commands
+
+    def shutdown(self):
+        # TODO: allow shutdown() to be async
+        self.event_loop.create_task(self._shutdown())
+
+    async def _shutdown(self):
+        self.log.info(f"Closing ACP agent and client for '{self.__class__.__name__}'.")
+        client = await self.get_client()
+        conn = await client.get_connection()
+        await conn.close()
+        subprocess = await self.get_agent_subprocess()
+        try:
+            subprocess.kill()
+        except ProcessLookupError:
+            pass
+        self.log.info(f"Completed closed ACP agent and client for '{self.__class__.__name__}'.")

jupyter_ai_acp_client/default_acp_client.py

@@ -0,0 +1,368 @@
+import asyncio
+import contextlib
+import logging
+import os
+import sys
+from pathlib import Path
+from typing import Any, AsyncIterator
+
+from acp import (
+    PROTOCOL_VERSION,
+    Client,
+    RequestError,
+    connect_to_agent,
+    text_block,
+)
+from acp.core import ClientSideConnection
+from acp.schema import (
+    AgentMessageChunk,
+    AgentPlanUpdate,
+    AgentThoughtChunk,
+    AudioContentBlock,
+    AvailableCommandsUpdate,
+    ClientCapabilities,
+    CreateTerminalResponse,
+    CurrentModeUpdate,
+    EmbeddedResourceContentBlock,
+    EnvVariable,
+    FileSystemCapability,
+    ImageContentBlock,
+    Implementation,
+    KillTerminalCommandResponse,
+    NewSessionResponse,
+    PermissionOption,
+    PromptResponse,
+    ReadTextFileResponse,
+    ReleaseTerminalResponse,
+    RequestPermissionResponse,
+    ResourceContentBlock,
+    TerminalOutputResponse,
+    TextContentBlock,
+    ToolCall,
+    ToolCallProgress,
+    ToolCallStart,
+    UserMessageChunk,
+    WaitForTerminalExitResponse,
+    WriteTextFileResponse,
+    AllowedOutcome
+)
+from jupyter_ai_persona_manager import BasePersona
+from typing import Awaitable
+from asyncio.subprocess import Process
+
+from .terminal_manager import TerminalManager
+
+async def queue_to_iterator(queue: asyncio.Queue[str], sentinel: str = "__end__") -> AsyncIterator[str]:
+    """Convert an asyncio queue to an async iterator."""
+    while True:
+        item = await queue.get()
+        if item == sentinel:
+            break
+        yield item
+
+class JaiAcpClient(Client):
+    """
+    The default ACP client. The client should be stored as a class attribute on each
+    ACP persona, such that each ACP agent subprocess is communicated through
+    exactly one ACP client (an instance of this class).
+    """
+
+    agent_subprocess: Process
+    _connection_future: Awaitable[ClientSideConnection]
+    event_loop: asyncio.AbstractEventLoop
+    _personas_by_session: dict[str, BasePersona]
+    _queues_by_session: dict[str, asyncio.Queue[str]]
+    _terminal_manager: TerminalManager
+
+    def __init__(self, *args, agent_subprocess: Awaitable[Process], event_loop: asyncio.AbstractEventLoop, **kwargs):
+        """
+        :param agent_subprocess: The ACP agent subprocess
+        (`asyncio.subprocess.Process`) assigned to this client.
+
+        :param event_loop: The `asyncio` event loop running this process.
+        """
+        self.agent_subprocess = agent_subprocess
+        # Each client instance needs its own connection to its own subprocess
+        self._connection_future = event_loop.create_task(
+            self._init_connection()
+        )
+        self.event_loop = event_loop
+        # Each client instance maintains its own session mappings
+        self._personas_by_session = {}
+        self._queues_by_session = {}
+        self._terminal_manager = TerminalManager(event_loop)
+        super().__init__(*args, **kwargs)
+    
+
+    async def _init_connection(self) -> ClientSideConnection:
+        proc = self.agent_subprocess
+        conn = connect_to_agent(self, proc.stdin, proc.stdout)
+        await conn.initialize(
+            protocol_version=PROTOCOL_VERSION,
+            client_capabilities=ClientCapabilities(
+                fs=FileSystemCapability(read_text_file=True, write_text_file=True),
+                terminal=True,
+            ),
+            client_info=Implementation(name="Jupyter AI", title="Jupyter AI ACP Client", version="0.1.0"),
+        )
+        return conn
+    
+    async def get_connection(self) -> ClientSideConnection:
+        return await self._connection_future
+
+    async def create_session(self, persona: BasePersona) -> NewSessionResponse:
+        """
+        Create an ACP agent session through this client scoped to a
+        `BasePersona` instance.
+        """
+        conn = await self.get_connection()
+        # TODO: change this to Jupyter preferred dir
+        session = await conn.new_session(mcp_servers=[], cwd=os.getcwd())
+        self._personas_by_session[session.session_id] = persona
+        return session
+    
+    async def prompt_and_reply(self, session_id: str, prompt: str, attachments: list[dict] = []) -> PromptResponse:
+        """
+        A helper method that sends a prompt with an optional list of attachments
+        to the assigned ACP server. This method writes back to the chat by
+        calling methods on the persona corresponding to this session ID.
+        """
+        assert session_id in self._personas_by_session
+        conn = await self.get_connection()
+
+        # ensure an asyncio Queue exists for this session
+        # the `session_update()` method will push chunks to this queue
+        queue = self._queues_by_session.get(session_id, None)
+        if queue is None:
+            queue: asyncio.Queue[str] = asyncio.Queue()
+            self._queues_by_session[session_id] = queue
+
+        # create async iterator that yields until the response is complete
+        aiter = queue_to_iterator(queue)
+
+        # create background task to stream message back to client using the
+        # dedicated persona method
+        persona = self._personas_by_session[session_id]
+        self.event_loop.create_task(
+            persona.stream_message(aiter)
+        )
+
+        # call the model and await
+        # TODO: add attachments!
+        response = await conn.prompt(
+            prompt=[
+                TextContentBlock(text=prompt, type="text"),
+            ],
+            session_id=session_id
+        )
+
+        # push sentinel value to queue to close the async iterator
+        queue.put_nowait("__end__")
+
+        return response
+
+    async def session_update(
+        self,
+        session_id: str,
+        update: UserMessageChunk
+        | AgentMessageChunk
+        | AgentThoughtChunk
+        | ToolCallStart
+        | ToolCallProgress
+        | AgentPlanUpdate
+        | AvailableCommandsUpdate
+        | CurrentModeUpdate,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Handles `session/update` requests from the ACP agent. There must be an
+        `asyncio.Queue` corresponding to this session ID - this should be set by
+        the `prompt_and_reply()` method.
+        """
+
+        if isinstance(update, AvailableCommandsUpdate):
+            if not update.available_commands:
+                return
+            persona = self._personas_by_session.get(session_id)
+            if persona and hasattr(persona, 'acp_slash_commands'):
+                persona.acp_slash_commands = update.available_commands
+            return
+
+        if not isinstance(update, AgentMessageChunk):
+            return
+
+        if session_id not in self._queues_by_session:
+            logging.error(f"No queue found for session_id: {session_id}")
+            return
+
+        content = update.content
+        text: str
+        if isinstance(content, TextContentBlock):
+            text = content.text
+        elif isinstance(content, ImageContentBlock):
+            text = "<image>"
+        elif isinstance(content, AudioContentBlock):
+            text = "<audio>"
+        elif isinstance(content, ResourceContentBlock):
+            text = content.uri or "<resource>"
+        elif isinstance(content, EmbeddedResourceContentBlock):
+            text = "<resource>"
+        else:
+            text = "<content>"
+
+        queue = self._queues_by_session[session_id]
+        queue.put_nowait(text)
+
+    async def request_permission(
+        self, options: list[PermissionOption], session_id: str, tool_call: ToolCall, **kwargs: Any
+    ) -> RequestPermissionResponse:
+        """
+        Handles `session/request_permission` requests from the ACP agent.
+
+        TODO: This currently always gives the agent permission. We will need to
+        add some tool call approval UI and handle permission requests properly.
+        """
+        option_id = ""
+        for o in options:
+            if "allow" in o.option_id.lower():
+                option_id = o.option_id
+                break
+
+        return RequestPermissionResponse(
+            outcome=AllowedOutcome(option_id=option_id, outcome='selected')
+        )
+
+    async def write_text_file(
+        self, content: str, path: str, session_id: str, **kwargs: Any
+    ) -> WriteTextFileResponse | None:
+        # Validate path parameter
+        if not path or not path.strip():
+            raise RequestError.invalid_params({"path": "path cannot be empty"})
+
+        file_path = Path(path)
+
+        # Check if path is a directory
+        if file_path.is_dir():
+            raise RequestError.invalid_params({"path": "path cannot be a directory"})
+
+        try:
+            file_path.parent.mkdir(parents=True, exist_ok=True)
+            await asyncio.to_thread(file_path.write_text, content, encoding="utf-8")
+        except PermissionError as e:
+            raise RequestError.internal_error({"path": path, "error": f"Permission denied: {e}"})
+        except OSError as e:
+            raise RequestError.internal_error({"path": path, "error": str(e)})
+
+        return WriteTextFileResponse()
+
+    async def read_text_file(
+        self, path: str, session_id: str, limit: int | None = None, line: int | None = None, **kwargs: Any
+    ) -> ReadTextFileResponse:
+        # Validate path parameter
+        if not path or not path.strip():
+            raise RequestError.invalid_params({"path": "path cannot be empty"})
+
+        # Validate line parameter (must be >= 1 if provided)
+        if line is not None and line < 1:
+            raise RequestError.invalid_params({"line": "line must be >= 1 (1-indexed)"})
+
+        # Validate limit parameter (must be >= 1 if provided)
+        if limit is not None and limit < 1:
+            raise RequestError.invalid_params({"limit": "limit must be >= 1"})
+
+        file_path = Path(path)
+
+        # Check if file exists
+        if not file_path.exists():
+            raise RequestError.resource_not_found(path)
+
+        # Check if path is a directory
+        if file_path.is_dir():
+            raise RequestError.invalid_params({"path": "path cannot be a directory"})
+
+        try:
+            text = await asyncio.to_thread(file_path.read_text, encoding="utf-8")
+        except PermissionError as e:
+            raise RequestError.internal_error({"path": path, "error": f"Permission denied: {e}"})
+        except OSError as e:
+            raise RequestError.internal_error({"path": path, "error": str(e)})
+
+        lines = text.splitlines(keepends=True)
+
+        # line is 1-indexed; default to line 1 if not specified
+        start_index = (line - 1) if line is not None else 0
+
+        if limit is not None:
+            lines = lines[start_index : start_index + limit]
+        else:
+            lines = lines[start_index:]
+
+        content = "".join(lines)
+        return ReadTextFileResponse(content=content)
+
+    ##############################
+    # Terminal methods
+    ##############################
+
+    async def create_terminal(
+        self,
+        command: str,
+        session_id: str,
+        args: list[str] | None = None,
+        cwd: str | None = None,
+        env: list[EnvVariable] | None = None,
+        output_byte_limit: int | None = None,
+        **kwargs: Any,
+    ) -> CreateTerminalResponse:
+        return await self._terminal_manager.create_terminal(
+            command=command,
+            session_id=session_id,
+            args=args,
+            cwd=cwd,
+            env=env,
+            output_byte_limit=output_byte_limit,
+            **kwargs,
+        )
+
+    async def terminal_output(
+        self, session_id: str, terminal_id: str, **kwargs: Any
+    ) -> TerminalOutputResponse:
+        return await self._terminal_manager.terminal_output(
+            session_id=session_id,
+            terminal_id=terminal_id,
+            **kwargs,
+        )
+
+    async def release_terminal(
+        self, session_id: str, terminal_id: str, **kwargs: Any
+    ) -> ReleaseTerminalResponse | None:
+        return await self._terminal_manager.release_terminal(
+            session_id=session_id,
+            terminal_id=terminal_id,
+            **kwargs,
+        )
+
+    async def wait_for_terminal_exit(
+        self, session_id: str, terminal_id: str, **kwargs: Any
+    ) -> WaitForTerminalExitResponse:
+        return await self._terminal_manager.wait_for_terminal_exit(
+            session_id=session_id,
+            terminal_id=terminal_id,
+            **kwargs,
+        )
+
+    async def kill_terminal(
+        self, session_id: str, terminal_id: str, **kwargs: Any
+    ) -> KillTerminalCommandResponse | None:
+        return await self._terminal_manager.kill_terminal(
+            session_id=session_id,
+            terminal_id=terminal_id,
+            **kwargs,
+        )
+
+    async def ext_method(self, method: str, params: dict) -> dict:
+        raise RequestError.method_not_found(method)
+
+    async def ext_notification(self, method: str, params: dict) -> None:
+        raise RequestError.method_not_found(method)
+

jupyter_ai_acp_client/extension_app.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from jupyter_server.extension.application import ExtensionApp
+from .routes import AcpSlashCommandsHandler
+
+
+class JaiAcpClientExtension(ExtensionApp):
+    """
+    Jupyter AI ACP client extension.
+    """
+
+    name = "jupyter_ai_acp_client"
+    handlers = [
+        (r"ai/acp/slash_commands/?([^/]*)?", AcpSlashCommandsHandler),
+    ]
+
+    def initialize_settings(self):
+        """Initialize router settings and event listeners."""
+        # # Ensure 'jupyter-ai' dictionary is in `self.settings`, which gets
+        # # copied to `self.serverapp.web_app.settings` after this method returns
+        # if 'jupyter-ai' not in self.settings:
+        #     self.settings['jupyter-ai'] = {}
+        
+        # self.settings['jupyter-ai']['acp-client']
+        return
+
+    async def stop_extension(self):
+        """Clean up router when extension stops."""
+        return