agentpool 2.1.9__py3-none-any.whl → 2.5.0__py3-none-any.whl

agentpool_server/acp_server/session.py

@@ -7,54 +7,26 @@ between agents and ACP clients through the JSON-RPC protocol.
 from __future__ import annotations
 
 import asyncio
-from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 import re
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
+import anyio
 from exxec.acp_provider import ACPExecutionEnvironment
 import logfire
-from pydantic_ai import (
-    FinalResultEvent,
-    FunctionToolCallEvent,
-    FunctionToolResultEvent,
-    PartDeltaEvent,
-    PartStartEvent,
-    RetryPromptPart,
-    TextPart,
-    TextPartDelta,
-    ThinkingPart,
-    ThinkingPartDelta,
-    ToolCallPartDelta,
-    ToolReturnPart,
-    UsageLimitExceeded,
-    UserPromptPart,
-)
+from pydantic_ai import UsageLimitExceeded, UserPromptPart
 from slashed import Command, CommandStore
+from tokonomics.model_discovery.model_info import ModelInfo
 
 from acp import RequestPermissionRequest
 from acp.acp_requests import ACPRequests
 from acp.filesystem import ACPFileSystem
 from acp.notifications import ACPNotifications
-from acp.schema import (
-    AvailableCommand,
-    ClientCapabilities,
-    ContentToolCallContent,
-    PlanEntry,
-    TerminalToolCallContent,
-    ToolCallLocation,
-)
-from acp.tool_call_state import ToolCallState
-from acp.utils import generate_tool_title, infer_tool_kind, to_acp_content_blocks
+from acp.schema import AvailableCommand, ClientCapabilities, SessionNotification
 from agentpool import Agent, AgentContext  # noqa: TC001
 from agentpool.agents import SlashedAgent
 from agentpool.agents.acp_agent import ACPAgent
-from agentpool.agents.events import (
-    PlanUpdateEvent,
-    StreamCompleteEvent,
-    ToolCallProgressEvent,
-    ToolCallStartEvent,
-)
+from agentpool.agents.modes import ModeInfo
 from agentpool.log import get_logger
 from agentpool_commands import get_commands
 from agentpool_commands.base import NodeCommand
@@ -62,20 +34,20 @@ from agentpool_server.acp_server.converters import (
     convert_acp_mcp_server_to_config,
     from_acp_content,
 )
+from agentpool_server.acp_server.event_converter import ACPEventConverter
 from agentpool_server.acp_server.input_provider import ACPInputProvider
 
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
 
-    from pydantic_ai import (
-        SystemPromptPart,
-        UserContent,
-    )
+    from pydantic_ai import SystemPromptPart, UserContent
     from slashed import CommandContext
 
     from acp import Client, RequestPermissionResponse
     from acp.schema import (
+        AvailableCommandsUpdate,
+        ConfigOptionUpdate,
         ContentBlock,
         Implementation,
         McpServer,
@@ -84,7 +56,6 @@ if TYPE_CHECKING:
     from agentpool import AgentPool
     from agentpool.agents import AGUIAgent
     from agentpool.agents.claude_code_agent import ClaudeCodeAgent
-    from agentpool.agents.events import RichAgentStreamEvent
     from agentpool.prompts.manager import PromptManager
     from agentpool.prompts.prompts import MCPClientPrompt
     from agentpool_server.acp_server.acp_agent import AgentPoolACPAgent
@@ -129,13 +100,27 @@ def _is_slash_command(text: str) -> bool:
 
 def split_commands(
     contents: Sequence[UserContent],
+    command_store: CommandStore,
 ) -> tuple[list[str], list[UserContent]]:
+    """Split content into local slash commands and pass-through content.
+
+    Only commands that exist in the local command_store are extracted.
+    Remote commands (from nested ACP agents) stay in non_command_content
+    so they flow through to the agent and reach the nested server.
+    """
     commands: list[str] = []
     non_command_content: list[UserContent] = []
     for item in contents:
-        if isinstance(item, str) and _is_slash_command(item):
+        # Check if this is a LOCAL command we handle
+        if (
+            isinstance(item, str)
+            and _is_slash_command(item)
+            and (match := SLASH_PATTERN.match(item.strip()))
+            and command_store.get_command(match.group(1))
+        ):
             commands.append(item.strip())
         else:
+            # Not a local command - pass through (may be remote command or regular text)
             non_command_content.append(item)
     return commands, non_command_content
 
@@ -223,6 +208,12 @@ class ACPSession:
     manager: ACPSessionManager | None = None
     """Session manager for managing sessions. Used for session management commands."""
 
+    subagent_display_mode: Literal["inline", "tool_box"] = "tool_box"
+    """How to display subagent output:
+    - 'inline': Subagent output flows into main message stream
+    - 'tool_box': Subagent output contained in the tool call's progress box (default)
+    """
+
     def __post_init__(self) -> None:
         """Initialize session state and set up providers."""
         from agentpool_server.acp_server.commands import get_commands as get_acp_commands
@@ -231,8 +222,7 @@ class ACPSession:
         self.log = logger.bind(session_id=self.session_id)
         self._task_lock = asyncio.Lock()
         self._cancelled = False
-        self._current_tool_inputs: dict[str, dict[str, Any]] = {}
-        self._tool_call_states: dict[str, ToolCallState] = {}
+        self._current_converter: ACPEventConverter | None = None
         self.fs = ACPFileSystem(self.client, session_id=self.session_id)
         cmds = [
             *get_commands(
@@ -243,10 +233,10 @@ class ACPSession:
             ),
             *get_acp_commands(),
         ]
-        self.command_store = CommandStore(enable_system_commands=True, commands=cmds)
+        self.command_store = CommandStore(commands=cmds)
         self.command_store._initialize_sync()
         self._update_callbacks: list[Callable[[], None]] = []
-
+        self._remote_commands: list[AvailableCommand] = []  # Commands from nested ACP agents
         self.staged_content = StagedContent()
         # Inject Zed-specific instructions if client is Zed
         if self.client_info and self.client_info.name and "zed" in self.client_info.name.lower():
@@ -294,44 +284,47 @@ class ACPSession:
                         return response
 
                 agent.acp_permission_callback = permission_callback
-        self.log.info("Created ACP session", current_agent=self.current_agent_name)
 
-    async def initialize(self) -> None:
-        """Initialize async resources. Must be called after construction."""
-        await self.acp_env.__aenter__()
+            # Subscribe to state change signal for all agents
+            agent.state_updated.connect(self._on_state_updated)
+        self.log.info("Created ACP session", current_agent=self.current_agent_name)
 
-    def _get_or_create_tool_state(
-        self,
-        tool_call_id: str,
-        tool_name: str,
-        tool_input: dict[str, Any],
-    ) -> ToolCallState:
-        """Get existing tool call state or create a new one.
+    async def _on_state_updated(
+        self, state: ModeInfo | ModelInfo | AvailableCommandsUpdate | ConfigOptionUpdate
+    ) -> None:
+        """Handle state update signal from agent - forward to ACP client."""
+        from acp.schema import (
+            AvailableCommandsUpdate as ACPAvailableCommandsUpdate,
+            ConfigOptionUpdate as ACPConfigOptionUpdate,
+            CurrentModelUpdate,
+            CurrentModeUpdate,
+            SessionNotification,
+        )
 
-        Args:
-            tool_call_id: Unique identifier for the tool call
-            tool_name: Name of the tool being called
-            tool_input: Input parameters for the tool
+        update: CurrentModeUpdate | CurrentModelUpdate | ACPConfigOptionUpdate
+        match state:
+            case ModeInfo(id=mode_id):
+                update = CurrentModeUpdate(current_mode_id=mode_id)
+                self.log.debug("Forwarding mode change to client", mode_id=mode_id)
+            case ModelInfo(id=model_id):
+                update = CurrentModelUpdate(current_model_id=model_id)
+                self.log.debug("Forwarding model change to client", model_id=model_id)
+            case ACPAvailableCommandsUpdate(available_commands=cmds):
+                # Store remote commands and send merged list
+                self._remote_commands = list(cmds)
+                await self.send_available_commands_update()
+                self.log.debug("Merged and sent commands update to client")
+                return
+            case ACPConfigOptionUpdate():
+                update = state
+                self.log.debug("Forwarding config option update to client")
 
-        Returns:
-            ToolCallState instance (existing or newly created)
-        """
-        if tool_call_id not in self._tool_call_states:
-            state = ToolCallState(
-                notifications=self.notifications,
-                tool_call_id=tool_call_id,
-                tool_name=tool_name,
-                title=generate_tool_title(tool_name, tool_input),
-                kind=infer_tool_kind(tool_name),
-                raw_input=tool_input,
-            )
-            self._tool_call_states[tool_call_id] = state
-        return self._tool_call_states[tool_call_id]
+        notification = SessionNotification(session_id=self.session_id, update=update)
+        await self.client.session_update(notification)  # pyright: ignore[reportArgumentType]
 
-    def _cleanup_tool_state(self, tool_call_id: str) -> None:
-        """Remove tool call state after completion."""
-        self._tool_call_states.pop(tool_call_id, None)
-        self._current_tool_inputs.pop(tool_call_id, None)
+    async def initialize(self) -> None:
+        """Initialize async resources. Must be called after construction."""
+        await self.acp_env.__aenter__()
 
     async def initialize_mcp_servers(self) -> None:
         """Initialize MCP servers if any are configured."""
@@ -339,26 +332,32 @@ class ACPSession:
             return
         self.log.info("Initializing MCP servers", server_count=len(self.mcp_servers))
         cfgs = [convert_acp_mcp_server_to_config(s) for s in self.mcp_servers]
-        # Define accessible roots for MCP servers
-        # root = Path(self.cwd).resolve().as_uri() if self.cwd else None
-        for _cfg in cfgs:
+        # Add each MCP server to the current agent's MCP manager dynamically
+        for cfg in cfgs:
             try:
-                # Server will be initialized when MCP manager enters context
-                self.log.info("Added MCP servers", server_count=len(cfgs))
-                await self._register_mcp_prompts_as_commands()
+                await self.agent.mcp.setup_server(cfg)
+                self.log.info(
+                    "Added MCP server to agent", server_name=cfg.name, agent=self.current_agent_name
+                )
             except Exception:
-                self.log.exception("Failed to initialize MCP manager")
+                self.log.exception("Failed to setup MCP server", server_name=cfg.name)
                 # Don't fail session creation, just log the error
+        # Register MCP prompts as commands after all servers are added
+        try:
+            await self._register_mcp_prompts_as_commands()
+        except Exception:
+            self.log.exception("Failed to register MCP prompts as commands")
 
     async def init_project_context(self) -> None:
-        """Load AGENTS.md file and inject project context into all agents.
+        """Load AGENTS.md/CLAUDE.md file and stage as initial context.
 
-        TODO: Consider moving this to __aenter__
+        The project context is staged as user message content rather than system prompts,
+        which ensures it's available for all agent types and avoids timing issues with
+        agent initialization.
         """
         if info := await self.requests.read_agent_rules(self.cwd):
-            for agent in self.agent_pool.agents.values():
-                prompt = f"## Project Information\n\n{info}"
-                agent.sys_prompts.prompts.append(prompt)
+            # Stage as user message to be prepended to first prompt
+            self.staged_content.add_text(f"## Project Information\n\n{info}")
 
     async def init_client_skills(self) -> None:
         """Discover and load skills from client-side .claude/skills directory.
@@ -377,7 +376,14 @@ class ACPSession:
             self.log.exception("Failed to discover client-side skills", error=e)
 
     @property
-    def agent(self) -> Agent[ACPSession, str] | ACPAgent | AGUIAgent | ClaudeCodeAgent:
+    def agent(
+        self,
+    ) -> (
+        Agent[ACPSession, str]
+        | ACPAgent[ACPSession]
+        | AGUIAgent[ACPSession]
+        | ClaudeCodeAgent[ACPSession]
+    ):
         """Get the currently active agent."""
         if self.current_agent_name in self.agent_pool.agents:
             return self.agent_pool.get_agent(self.current_agent_name, deps_type=ACPSession)
@@ -421,6 +427,10 @@ class ACPSession:
         This actively interrupts the running agent by calling its interrupt() method,
         which handles protocol-specific cancellation (e.g., sending CancelNotification
         for ACP agents, calling SDK interrupt for ClaudeCodeAgent, etc.).
+
+        Note:
+            Tool call cleanup is handled in process_prompt() to avoid race conditions
+            with the converter state being modified from multiple async contexts.
         """
         self._cancelled = True
         self.log.info("Session cancelled, interrupting agent")
@@ -450,7 +460,7 @@ class ACPSession:
         if not contents:
             self.log.warning("Empty prompt received")
             return "refusal"
-        commands, non_command_content = split_commands(contents)
+        commands, non_command_content = split_commands(contents, self.command_store)
         async with self._task_lock:
             if commands:  # Process commands if found
                 for command in commands:
@@ -470,19 +480,62 @@ class ACPSession:
                 has_staged=staged is not None,
             )
             event_count = 0
-            self._current_tool_inputs.clear()  # Reset tool inputs for new stream
+            # Create a new event converter for this prompt
+            converter = ACPEventConverter(subagent_display_mode=self.subagent_display_mode)
+            self._current_converter = converter  # Track for cancellation
 
             try:  # Use the session's persistent input provider
                 async for event in self.agent.run_stream(
-                    *all_content, input_provider=self.input_provider
+                    *all_content,
+                    input_provider=self.input_provider,
+                    deps=self,
+                    conversation_id=self.session_id,  # Tie agent conversation to ACP session
                 ):
                     if self._cancelled:
+                        self.log.info("Cancelled during event loop, cleaning up tool calls")
+                        # Send cancellation notifications for any pending tool calls
+                        # This happens in the same async context as the converter
+                        async for cancel_update in converter.cancel_pending_tools():
+                            notification = SessionNotification(
+                                session_id=self.session_id, update=cancel_update
+                            )
+                            await self.client.session_update(notification)  # pyright: ignore[reportArgumentType]
+                        # CRITICAL: Allow time for client to process tool completion notifications
+                        # before sending PromptResponse. Without this delay, the client may receive
+                        # and process the PromptResponse before the tool notifications, causing UI
+                        # state desync where subsequent prompts appear stuck/unresponsive.
+                        # This is needed because even though send() awaits the write, the client
+                        # may process messages asynchronously or out of order.
+                        await anyio.sleep(0.05)
+                        self._current_converter = None
                         return "cancelled"
 
                     event_count += 1
-                    await self.handle_event(event)
+                    async for update in converter.convert(event):
+                        notification = SessionNotification(
+                            session_id=self.session_id, update=update
+                        )
+                        await self.client.session_update(notification)  # pyright: ignore[reportArgumentType]
+                    # Yield control to allow notifications to be sent immediately
+                    await anyio.sleep(0.01)
                 self.log.info("Streaming finished", events_processed=event_count)
 
+            except asyncio.CancelledError:
+                # Task was cancelled (e.g., via interrupt()) - return proper stop reason
+                # This is critical: CancelledError doesn't inherit from Exception,
+                # so we must catch it explicitly to send the PromptResponse
+                self.log.info("Stream cancelled via CancelledError, cleaning up tool calls")
+                # Send cancellation notifications for any pending tool calls
+                async for cancel_update in converter.cancel_pending_tools():
+                    notification = SessionNotification(
+                        session_id=self.session_id, update=cancel_update
+                    )
+                    await self.client.session_update(notification)  # pyright: ignore[reportArgumentType]
+                # CRITICAL: Allow time for client to process tool completion notifications
+                # before sending PromptResponse. See comment in cancellation branch above.
+                await anyio.sleep(0.05)
+                self._current_converter = None
+                return "cancelled"
             except UsageLimitExceeded as e:
                 self.log.info("Usage limit exceeded", error=str(e))
                 error_msg = str(e)  # Determine which limit was hit based on error
@@ -495,6 +548,7 @@ class ACPSession:
                     return "refusal"
                 return "max_tokens"  # Default to max_tokens for other usage limits
             except Exception as e:
+                self._current_converter = None  # Clear converter reference
                 self.log.exception("Error during streaming")
                 # Send error notification asynchronously to avoid blocking response
                 self.acp_agent.tasks.create_task(
@@ -503,228 +557,21 @@ class ACPSession:
                 )
                 return "end_turn"
             else:
+                # Title generation is now handled automatically by log_conversation
+                self._current_converter = None  # Clear converter reference
                 return "end_turn"
 
     async def _send_error_notification(self, message: str) -> None:
         """Send error notification, with exception handling."""
+        if self._cancelled:
+            return
         try:
             await self.notifications.send_agent_text(message)
         except Exception:
             self.log.exception("Failed to send error notification")
 
-    async def handle_event(self, event: RichAgentStreamEvent[Any]) -> None:  # noqa: PLR0915
-        match event:
-            case (
-                PartStartEvent(part=TextPart(content=delta))
-                | PartDeltaEvent(delta=TextPartDelta(content_delta=delta))
-            ):
-                await self.notifications.send_agent_text(delta)
-
-            case (
-                PartStartEvent(part=ThinkingPart(content=delta))
-                | PartDeltaEvent(delta=ThinkingPartDelta(content_delta=delta))
-            ):
-                await self.notifications.send_agent_thought(delta or "\n")
-
-            case PartStartEvent(part=part):
-                self.log.debug("Received unhandled PartStartEvent", part=part)
-
-            # Tool call streaming delta - create/update state and start notification
-            case PartDeltaEvent(delta=ToolCallPartDelta() as delta):
-                if delta_part := delta.as_part():
-                    tool_call_id = delta_part.tool_call_id
-                    try:
-                        tool_input = delta_part.args_as_dict()
-                    except ValueError:
-                        # Args still streaming, not valid JSON yet - skip this delta
-                        pass
-                    else:
-                        self._current_tool_inputs[tool_call_id] = tool_input
-                        # Create state and send initial notification
-                        state = self._get_or_create_tool_state(
-                            tool_call_id=tool_call_id,
-                            tool_name=delta_part.tool_name,
-                            tool_input=tool_input,
-                        )
-                        await state.start()
-
-            # Tool call started - create/update state and start notification
-            case FunctionToolCallEvent(part=part):
-                tool_call_id = part.tool_call_id
-                try:
-                    tool_input = part.args_as_dict()
-                except ValueError as e:
-                    # Args might be malformed - use empty dict and log
-                    self.log.warning(
-                        "Failed to parse tool args", tool_name=part.tool_name, error=str(e)
-                    )
-                    tool_input = {}
-                self._current_tool_inputs[tool_call_id] = tool_input
-                # Create state and send initial notification
-                state = self._get_or_create_tool_state(
-                    tool_call_id=tool_call_id,
-                    tool_name=part.tool_name,
-                    tool_input=tool_input,
-                )
-                await state.start()
-
-            # Tool completed successfully - update state and finalize
-            case FunctionToolResultEvent(
-                result=ToolReturnPart(content=content, tool_name=tool_name) as result,
-                tool_call_id=tool_call_id,
-            ):
-                if isinstance(content, AsyncGenerator):
-                    full_content = ""
-                    async for chunk in content:
-                        full_content += str(chunk)
-                        # Stream progress through state
-                        if tool_state := self._tool_call_states.get(tool_call_id):
-                            await tool_state.update(status="in_progress", raw_output=chunk)
-
-                    # Replace the AsyncGenerator with the full content to prevent errors
-                    result.content = full_content
-                    final_output = full_content
-                else:
-                    final_output = result.content
-
-                # Complete tool call through state (preserves accumulated content/locations)
-                if complete_state := self._tool_call_states.get(tool_call_id):
-                    # Only add return value as content if no content was emitted during execution
-                    if complete_state.has_content:
-                        # Content already provided via progress events - just set raw_output
-                        await complete_state.complete(raw_output=final_output)
-                    else:
-                        # No content yet - convert return value for display
-                        converted_blocks = to_acp_content_blocks(final_output)
-                        content_items = [
-                            ContentToolCallContent(content=block) for block in converted_blocks
-                        ]
-                        await complete_state.complete(
-                            raw_output=final_output,
-                            content=content_items,
-                        )
-                self._cleanup_tool_state(tool_call_id)
-
-            # Tool failed with retry - update state with error
-            case FunctionToolResultEvent(
-                result=RetryPromptPart(tool_name=tool_name) as result,
-                tool_call_id=tool_call_id,
-            ):
-                error_message = result.model_response()
-                if fail_state := self._tool_call_states.get(tool_call_id):
-                    await fail_state.fail(error=error_message)
-                self._cleanup_tool_state(tool_call_id)
-
-            # Tool emits its own start event - update state with better title/content
-            case ToolCallStartEvent(
-                tool_call_id=tool_call_id,
-                tool_name=tool_name,
-                title=title,
-                kind=kind,
-                locations=loc_items,
-                raw_input=raw_input,
-            ):
-                self.log.debug(
-                    "Tool call start event", tool_name=tool_name, tool_call_id=tool_call_id
-                )
-                # Get or create state (may already exist from FunctionToolCallEvent)
-                state = self._get_or_create_tool_state(
-                    tool_call_id=tool_call_id,
-                    tool_name=tool_name,
-                    tool_input=raw_input or {},
-                )
-                # Convert LocationContentItem objects to ACP format
-                acp_locations = [
-                    ToolCallLocation(path=loc.path, line=loc.line) for loc in loc_items
-                ]
-                # Update state with tool-provided details (better title, content, locations)
-                await state.update(title=title, kind=kind, locations=acp_locations or None)
-
-            # Tool progress event - update state with title and content
-            case ToolCallProgressEvent(
-                tool_call_id=tool_call_id,
-                title=title,
-                status=status,
-                items=items,
-            ) if tool_call_id and tool_call_id in self._tool_call_states:
-                progress_state = self._tool_call_states[tool_call_id]
-                self.log.debug("Progress event", tool_call_id=tool_call_id, title=title)
-
-                # Convert items to ACP content
-                from agentpool.agents.events import (
-                    DiffContentItem,
-                    FileContentItem,
-                    LocationContentItem,
-                    TerminalContentItem,
-                    TextContentItem,
-                )
-                from agentpool_server.acp_server.syntax_detection import (
-                    format_zed_code_block,
-                )
-
-                acp_content: list[Any] = []
-                location_paths: list[str] = []
-
-                for item in items:
-                    match item:
-                        case TerminalContentItem(terminal_id=tid):
-                            acp_content.append(TerminalToolCallContent(terminal_id=tid))
-                        case TextContentItem(text=text):
-                            acp_content.append(ContentToolCallContent.text(text=text))
-                        case FileContentItem(
-                            content=file_content,
-                            path=file_path,
-                            start_line=start_line,
-                            end_line=end_line,
-                        ):
-                            # Format as Zed-compatible code block with clickable path
-                            formatted = format_zed_code_block(
-                                file_content, file_path, start_line, end_line
-                            )
-                            acp_content.append(ContentToolCallContent.text(text=formatted))
-                            # Also add path to locations for "follow along" feature
-                            location_paths.append(file_path)
-                        case DiffContentItem(path=diff_path, old_text=old, new_text=new):
-                            # Send diff via direct notification
-                            await self.notifications.file_edit_progress(
-                                tool_call_id=tool_call_id,
-                                path=diff_path,
-                                old_text=old or "",
-                                new_text=new,
-                                status=status,
-                                changed_lines=[],
-                            )
-                        case LocationContentItem(path=loc_path):
-                            location_paths.append(loc_path)
-
-                await progress_state.update(
-                    title=title,
-                    status="in_progress",
-                    content=acp_content if acp_content else None,
-                    locations=location_paths if location_paths else None,
-                )
-
-            case FinalResultEvent():
-                self.log.debug("Final result received")
-
-            case StreamCompleteEvent():
-                pass
-
-            case PlanUpdateEvent(entries=entries, tool_call_id=tool_call_id):
-                acp_entries = [
-                    PlanEntry(content=e.content, priority=e.priority, status=e.status)
-                    for e in entries
-                ]
-                await self.notifications.update_plan(acp_entries)
-
-            case _:
-                self.log.debug("Unhandled event", event_type=type(event).__name__)
-
     async def close(self) -> None:
         """Close the session and cleanup resources."""
-        self._current_tool_inputs.clear()
-        self._tool_call_states.clear()
-
         try:
             await self.acp_env.__aexit__(None, None, None)
         except Exception:
@@ -743,9 +590,14 @@ class ACPSession:
             self.log.exception("Error closing session")
 
     async def send_available_commands_update(self) -> None:
-        """Send current available commands to client."""
+        """Send current available commands to client.
+
+        Merges local commands from command_store with any remote commands
+        from nested ACP agents.
+        """
         try:
-            commands = self.get_acp_commands()
+            commands = self.get_acp_commands()  # Local commands
+            commands.extend(self._remote_commands)  # Merge remote commands
             await self.notifications.update_commands(commands)
         except Exception:
             self.log.exception("Failed to send available commands update")
@@ -802,11 +654,10 @@ class ACPSession:
         Returns:
             List of ACP AvailableCommand objects compatible with current node
         """
-        all_commands = self.command_store.list_commands()
         current_node = self.agent
         # Filter commands by node compatibility
         compatible_commands = []
-        for cmd in all_commands:
+        for cmd in self.command_store.list_commands():
             cmd_cls = cmd if isinstance(cmd, type) else type(cmd)
             # Check if command supports current node type
             if issubclass(cmd_cls, NodeCommand) and not cmd_cls.supports_node(current_node):  # type: ignore[union-attr]
@@ -860,11 +711,7 @@ class ACPSession:
             )
 
     def register_update_callback(self, callback: Callable[[], None]) -> None:
-        """Register callback for command updates.
-
-        Args:
-            callback: Function to call when commands are updated
-        """
+        """Register callback for command updates."""
         self._update_callbacks.append(callback)
 
     def create_mcp_command(self, prompt: MCPClientPrompt) -> Command:
@@ -885,16 +732,13 @@ class ACPSession:
         ) -> None:
             """Execute the MCP prompt with parsed arguments."""
             # Map parsed args to prompt parameters
-
-            result = {}
             # Map positional args to prompt parameter names
-            for i, arg_value in enumerate(args):
-                if i < len(prompt.arguments):
-                    param_name = prompt.arguments[i]["name"]
-                    result[param_name] = arg_value
-            result.update(kwargs)
-            try:
-                # Get prompt components
+            result = {
+                prompt.arguments[i]["name"]: arg_value
+                for i, arg_value in enumerate(args)
+                if i < len(prompt.arguments)
+            } | kwargs
+            try:  # Get prompt components
                 components = await prompt.get_components(result or None)
                 self.staged_content.add(components)
                 # Send confirmation
@@ -905,15 +749,13 @@ class ACPSession:
                 logger.exception("MCP prompt execution failed", prompt=prompt.name)
                 await ctx.print(f"❌ Prompt error: {e}")
 
-        usage_hint = (
-            " ".join(f"<{arg['name']}>" for arg in prompt.arguments) if prompt.arguments else None
-        )
-        return Command(
-            execute_func=execute_prompt,
+        usage = " ".join(f"<{i['name']}>" for i in args) if (args := prompt.arguments) else None
+        return Command.from_raw(
+            execute_prompt,
             name=prompt.name,
             description=prompt.description or f"MCP prompt: {prompt.name}",
             category="mcp",
-            usage=usage_hint,
+            usage=usage,
         )
 
     def create_prompt_hub_command(
@@ -960,8 +802,8 @@ class ACPSession:
         # Create command name - prefix with provider if not builtin
         command_name = f"{provider}_{name}" if provider != "builtin" else name
 
-        return Command(
-            execute_func=execute_prompt,
+        return Command.from_raw(
+            execute_prompt,
             name=command_name,
             description=f"Prompt hub: {provider}:{name}",
             category="prompts",