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

agentpool_server/acp_server/acp_agent.py

@@ -16,31 +16,27 @@ from acp.schema import (
     NewSessionResponse,
     PromptResponse,
     ResumeSessionResponse,
+    SessionConfigOption,
+    SessionConfigSelectOption,
     SessionInfo,
+    SessionMode,
     SessionModelState,
     SessionModeState,
+    SetSessionConfigOptionResponse,
     SetSessionModelRequest,
     SetSessionModelResponse,
     SetSessionModeRequest,
     SetSessionModeResponse,
 )
-from agentpool import Agent
 from agentpool.log import get_logger
 from agentpool.utils.tasks import TaskManager
-from agentpool_server.acp_server.converters import (
-    # agent_to_mode,  # TODO: Re-enable when supporting agent switching via modes
-    get_confirmation_modes,
-    mode_id_to_confirmation_mode,
-)
 from agentpool_server.acp_server.session_manager import ACPSessionManager
 
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
-
     from pydantic_ai import ModelRequest, ModelResponse
 
-    from acp import AgentSideConnection, Client
+    from acp import Client
     from acp.schema import (
         AuthenticateRequest,
         CancelNotification,
@@ -53,6 +49,7 @@ if TYPE_CHECKING:
         NewSessionRequest,
         PromptRequest,
         ResumeSessionRequest,
+        SetSessionConfigOptionRequest,
         SetSessionModelRequest,
         SetSessionModeRequest,
     )
@@ -63,62 +60,155 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
-# Claude Code model definitions - simple IDs that the SDK understands
-CLAUDE_CODE_MODELS: list[ACPModelInfo] = [
-    ACPModelInfo(
-        model_id="opus",
-        name="Claude Opus",
-        description="Claude Opus - most capable model",
-    ),
-    ACPModelInfo(
-        model_id="sonnet",
-        name="Claude Sonnet",
-        description="Claude Sonnet - balanced performance and speed",
-    ),
-    ACPModelInfo(
-        model_id="haiku",
-        name="Claude Haiku",
-        description="Claude Haiku - fast and cost-effective",
-    ),
-]
-
-
-def create_claude_code_model_state(current_model: str | None = None) -> SessionModelState:
-    """Create SessionModelState for Claude Code agents.
+async def get_session_model_state(
+    agent: Any, current_model: str | None = None
+) -> SessionModelState | None:
+    """Get SessionModelState from an agent using its get_available_models() method.
+
+    Converts tokonomics ModelInfo to ACP ModelInfo format.
 
     Args:
-        current_model: Currently active model ID (e.g., "sonnet")
+        agent: Any agent with get_available_models() method
+        current_model: Currently active model ID (defaults to first available)
 
     Returns:
-        SessionModelState with Claude Code models
+        SessionModelState with all available models, None if no models available
     """
-    model_ids = [m.model_id for m in CLAUDE_CODE_MODELS]
-    current = current_model if current_model in model_ids else "sonnet"
-    return SessionModelState(available_models=CLAUDE_CODE_MODELS, current_model_id=current)
+    from agentpool.agents.base_agent import BaseAgent
 
+    if not isinstance(agent, BaseAgent):
+        return None
 
-def create_session_model_state(
-    available_models: Sequence[ACPModelInfo], current_model: str | None = None
-) -> SessionModelState | None:
-    """Create a SessionModelState from available ACP models.
+    try:
+        toko_models = await agent.get_available_models()
+    except Exception:
+        logger.exception("Failed to get available models from agent")
+        return None
+
+    if not toko_models:
+        return None
+
+    # Convert tokonomics ModelInfo to ACP ModelInfo
+    acp_models: list[ACPModelInfo] = []
+    for toko in toko_models:
+        # Use id_override if set (e.g., "opus" for Claude Code), otherwise use id
+        model_id = toko.id_override if toko.id_override else toko.id
+        info = ACPModelInfo(model_id=model_id, name=toko.name, description=toko.description or "")
+        acp_models.append(info)
+    if not acp_models:
+        return None
+
+    # Ensure current model is in the list
+    all_ids = [model.model_id for model in acp_models]
+    if current_model and current_model not in all_ids:
+        # Add current model to the list so the UI shows it
+        desc = "Currently configured model"
+        model_info = ACPModelInfo(model_id=current_model, name=current_model, description=desc)
+        acp_models.insert(0, model_info)
+        current_model_id = current_model
+    else:
+        current_model_id = current_model if current_model in all_ids else all_ids[0]
+
+    return SessionModelState(available_models=acp_models, current_model_id=current_model_id)
+
+
+async def get_session_mode_state(agent: Any) -> SessionModeState | None:
+    """Get SessionModeState from an agent using its get_modes() method.
+
+    Converts agentpool ModeCategory to ACP SessionModeState format.
+    Uses the first category that looks like permissions (not model).
 
     Args:
-        available_models: List of ACP ModelInfo objects (already converted from tokonomics)
-        current_model: The currently active model (defaults to first available)
+        agent: Any agent with get_modes() method
 
     Returns:
-        SessionModelState with all available models, None if no models provided
+        SessionModeState from agent's modes, None if no modes available
     """
-    if not available_models:
+    from agentpool.agents.base_agent import BaseAgent
+
+    if not isinstance(agent, BaseAgent):
+        return None
+
+    try:
+        mode_categories = await agent.get_modes()
+    except Exception:
+        logger.exception("Failed to get modes from agent")
+        return None
+
+    if not mode_categories:
+        return None
+
+    # Find the permissions category (not model)
+    category = next((c for c in mode_categories if c.id != "model"), None)
+    if not category:
         return None
-    # Use first model as current if not specified or not found
-    all_ids = [model.model_id for model in available_models]
-    current_model_id = current_model if current_model in all_ids else all_ids[0]
-    return SessionModelState(
-        available_models=list(available_models), current_model_id=current_model_id
+
+    # Convert ModeInfo to ACP SessionMode
+    acp_modes = [
+        SessionMode(
+            id=mode.id,
+            name=mode.name,
+            description=mode.description,
+        )
+        for mode in category.available_modes
+    ]
+
+    return SessionModeState(
+        available_modes=acp_modes,
+        current_mode_id=category.current_mode_id,
     )
 
 
+async def get_session_config_options(agent: Any) -> list[SessionConfigOption]:
+    """Get SessionConfigOptions from an agent using its get_modes() method.
+
+    Converts all agentpool ModeCategories to ACP SessionConfigOption format.
+
+    Args:
+        agent: Any agent with get_modes() method
+
+    Returns:
+        List of SessionConfigOption from agent's mode categories
+    """
+    from agentpool.agents.base_agent import BaseAgent
+
+    if not isinstance(agent, BaseAgent):
+        return []
+
+    try:
+        mode_categories = await agent.get_modes()
+    except Exception:
+        logger.exception("Failed to get modes from agent")
+        return []
+
+    if not mode_categories:
+        return []
+
+    # Convert each ModeCategory to a SessionConfigOption
+    config_options: list[SessionConfigOption] = []
+    for category in mode_categories:
+        options = [
+            SessionConfigSelectOption(
+                value=mode.id,
+                name=mode.name,
+                description=mode.description,
+            )
+            for mode in category.available_modes
+        ]
+        config_options.append(
+            SessionConfigOption(
+                id=category.id,
+                name=category.name,
+                description=None,
+                category=category.category,
+                current_value=category.current_mode_id,
+                options=options,
+            )
+        )
+
+    return config_options
+
+
 @dataclass
 class AgentPoolACPAgent(ACPAgent):
     """Implementation of ACP Agent protocol interface for AgentPool.
@@ -129,7 +219,7 @@ class AgentPoolACPAgent(ACPAgent):
 
     PROTOCOL_VERSION: ClassVar = 1
 
-    connection: AgentSideConnection
+    client: Client
     """ACP connection for client communication."""
 
     agent_pool: AgentPool[Any]
@@ -137,9 +227,6 @@ class AgentPoolACPAgent(ACPAgent):
 
     _: KW_ONLY
 
-    available_models: Sequence[ACPModelInfo] = field(default_factory=list)
-    """List of available ACP ModelInfo objects (converted from tokonomics at startup)."""
-
     file_access: bool = True
     """Whether agent can access filesystem."""
 
@@ -160,7 +247,6 @@ class AgentPoolACPAgent(ACPAgent):
 
     def __post_init__(self) -> None:
         """Initialize derived attributes and setup after field assignment."""
-        self.client: Client = self.connection
         self.client_capabilities: ClientCapabilities | None = None
         self.client_info: Implementation | None = None
         self.session_manager = ACPSessionManager(pool=self.agent_pool)
@@ -176,17 +262,12 @@ class AgentPoolACPAgent(ACPAgent):
 
     async def initialize(self, params: InitializeRequest) -> InitializeResponse:
         """Initialize the agent and negotiate capabilities."""
-        logger.info("Initializing ACP agent implementation")
         version = min(params.protocol_version, self.PROTOCOL_VERSION)
         self.client_capabilities = params.client_capabilities
         self.client_info = params.client_info
-        logger.info(
-            "Client capabilities",
-            capabilities=self.client_capabilities,
-            client_info=self.client_info,
-        )
+        logger.info("Client info", request=params.model_dump_json())
         self._initialized = True
-        response = InitializeResponse.create(
+        return InitializeResponse.create(
             protocol_version=version,
             name="agentpool",
             title="AgentPool",
@@ -199,10 +280,8 @@ class AgentPoolACPAgent(ACPAgent):
             embedded_context_prompts=True,
             image_prompts=True,
         )
-        logger.info("ACP agent initialized successfully", response=response)
-        return response
 
-    async def new_session(self, params: NewSessionRequest) -> NewSessionResponse:  # noqa: PLR0915
+    async def new_session(self, params: NewSessionRequest) -> NewSessionResponse:
         """Create a new session."""
         if not self._initialized:
             raise RuntimeError("Agent not initialized")
@@ -230,48 +309,26 @@ class AgentPoolACPAgent(ACPAgent):
                 client_info=self.client_info,
             )
 
-            # Get mode information - pass through from ACPAgent or use our confirmation modes
+            # Get mode and model information from the agent
             from agentpool.agents.acp_agent import ACPAgent as ACPAgentClient
 
-            if session := self.session_manager.get_session(session_id):
-                if isinstance(session.agent, ACPAgentClient):
-                    # Pass through nested agent's modes (e.g., Claude Code's modes)
-                    if session.agent._state and session.agent._state.modes:
-                        state = session.agent._state.modes
-                        modes = state.available_modes
-                    else:
-                        # Fallback to our confirmation modes if nested agent has none
-                        modes = get_confirmation_modes()
-                        state = SessionModeState(current_mode_id="default", available_modes=modes)
-                else:
-                    # Native Agent - use our tool confirmation modes
-                    modes = get_confirmation_modes()
-                    state = SessionModeState(current_mode_id="default", available_modes=modes)
-            else:
-                modes = get_confirmation_modes()
-                state = SessionModeState(current_mode_id="default", available_modes=modes)
-            # TODO: Re-enable agent switching via modes when needed:
-            # modes = [agent_to_mode(agent) for agent in self.agent_pool.all_agents.values()]
-            # state = SessionModeState(current_mode_id=default_name, available_modes=modes)
+            state: SessionModeState | None = None
+            models: SessionModelState | None = None
+            config_options: list[SessionConfigOption] = []
 
-            # Get model information from the default agent
             if session := self.session_manager.get_session(session_id):
-                from agentpool.agents.claude_code_agent import ClaudeCodeAgent
-
-                if isinstance(session.agent, ClaudeCodeAgent):
-                    # Claude Code uses simple model IDs (sonnet, opus, haiku)
-                    models = create_claude_code_model_state(session.agent.model_name)
-                elif isinstance(session.agent, ACPAgentClient):
-                    # Nested ACP agent - pass through its model state directly
-                    if session.agent._state and session.agent._state.models:
+                if isinstance(session.agent, ACPAgentClient):
+                    # Nested ACP agent - pass through its state directly
+                    if session.agent._state:
                         models = session.agent._state.models
-                    else:
-                        models = None
+                        state = session.agent._state.modes
+                    # Also get config_options from nested agent
+                    config_options = await get_session_config_options(session.agent)
                 else:
-                    current_model = session.agent.model_name
-                    models = create_session_model_state(self.available_models, current_model)
-            else:
-                models = None
+                    # Use unified helpers for all other agents
+                    models = await get_session_model_state(session.agent, session.agent.model_name)
+                    state = await get_session_mode_state(session.agent)
+                    config_options = await get_session_config_options(session.agent)
         except Exception:
             logger.exception("Failed to create new session")
             raise
@@ -285,8 +342,14 @@ class AgentPoolACPAgent(ACPAgent):
                 if self.load_skills:
                     coro_4 = session.init_client_skills()
                     self.tasks.create_task(coro_4, name=f"init_client_skills_{session_id}")
-            logger.info("Created session", session_id=session_id, agent_count=len(modes))
-            return NewSessionResponse(session_id=session_id, modes=state, models=models)
+            logger.info("Created session", session_id=session_id)
+
+            return NewSessionResponse(
+                session_id=session_id,
+                modes=state,
+                models=models,
+                config_options=config_options if config_options else None,
+            )
 
     async def load_session(self, params: LoadSessionRequest) -> LoadSessionResponse:
         """Load an existing session from storage.
@@ -298,13 +361,14 @@ class AgentPoolACPAgent(ACPAgent):
         4. Replay conversation history via ACP notifications
         5. Return session state (modes, models)
         """
+        from agentpool.agents.acp_agent import ACPAgent as ACPAgentClient
+
         if not self._initialized:
             raise RuntimeError("Agent not initialized")
 
         try:
             # First check if session is already active
             session = self.session_manager.get_session(params.session_id)
-
             if not session:
                 # Try to resume from storage
                 msg = "Attempting to resume session from storage"
@@ -331,46 +395,30 @@ class AgentPoolACPAgent(ACPAgent):
                 session.mcp_servers = params.mcp_servers
                 await session.initialize_mcp_servers()
 
-            # Build response with current session state
-            # Get mode information - pass through from ACPAgent or use our confirmation modes
-            from agentpool.agents.acp_agent import ACPAgent as ACPAgentClient
+            mode_state: SessionModeState | None = None
+            models: SessionModelState | None = None
+            config_opts: list[SessionConfigOption] = []
 
             if isinstance(session.agent, ACPAgentClient):
-                # Pass through nested agent's modes (e.g., Claude Code's modes)
-                if session.agent._state and session.agent._state.modes:
+                # Nested ACP agent - pass through its state directly
+                if session.agent._state:
                     mode_state = session.agent._state.modes
-                else:
-                    # Fallback to our confirmation modes if nested agent has none
-                    modes = get_confirmation_modes()
-                    mode_state = SessionModeState(current_mode_id="default", available_modes=modes)
-            else:
-                # Native Agent - use our tool confirmation modes
-                modes = get_confirmation_modes()
-                mode_state = SessionModeState(current_mode_id="default", available_modes=modes)
-            # TODO: Re-enable agent switching via modes when needed:
-            # modes = [agent_to_mode(agent) for agent in self.agent_pool.all_agents.values()]
-            # mode_state = SessionModeState(
-            #     current_mode_id=session.current_agent_name,
-            #     available_modes=modes,
-            # )
-
-            # Get model information based on agent type
-            from agentpool.agents.claude_code_agent import ClaudeCodeAgent
-
-            models: SessionModelState | None
-            if session.agent and isinstance(session.agent, ClaudeCodeAgent):
-                # Claude Code uses simple model IDs (sonnet, opus, haiku)
-                models = create_claude_code_model_state(session.agent.model_name)
+                    models = session.agent._state.models
+                config_opts = await get_session_config_options(session.agent)
+            elif session.agent:
+                # Use unified helpers for all other agents
+                mode_state = await get_session_mode_state(session.agent)
+                models = await get_session_model_state(session.agent, session.agent.model_name)
+                config_opts = await get_session_config_options(session.agent)
             else:
-                current_model = session.agent.model_name if session.agent else None
-                models = create_session_model_state(self.available_models, current_model)
+                models = None
             # Schedule post-load initialization tasks
             self.tasks.create_task(session.send_available_commands_update())
             self.tasks.create_task(session.init_project_context())
             # Replay conversation history via ACP notifications
             self.tasks.create_task(self._replay_conversation_history(session))
             logger.info("Session loaded successfully", agent=session.current_agent_name)
-            return LoadSessionResponse(models=models, modes=mode_state)
+            return LoadSessionResponse(models=models, modes=mode_state, config_options=config_opts)
 
         except Exception:
             logger.exception("Failed to load session", session_id=params.session_id)
@@ -629,14 +677,8 @@ class AgentPoolACPAgent(ACPAgent):
     ) -> SetSessionModeResponse | None:
         """Set the session mode (change tool confirmation level).
 
-        Maps ACP mode IDs to ToolConfirmationMode and calls set_tool_confirmation_mode
-        on the session's agent. Each agent type handles the mode change appropriately:
-        - Agent/AGUIAgent: Updates local confirmation mode
-        - ACPAgent: Updates local mode AND forwards to remote ACP server
-
-        Mode mappings:
-        - "default": per_tool (confirm tools marked as requiring it)
-        - "acceptEdits": never (auto-approve all tool calls)
+        Calls set_mode directly on the agent with the mode_id, allowing the agent
+        to handle mode-specific logic (e.g., acceptEdits auto-allowing edit tools).
         """
         from agentpool.agents.acp_agent import ACPAgent as ACPAgentClient
 
@@ -646,15 +688,8 @@ class AgentPoolACPAgent(ACPAgent):
                 logger.warning("Session not found for mode switch", session_id=params.session_id)
                 return None
 
-            # Map mode_id to confirmation mode
-            confirmation_mode = mode_id_to_confirmation_mode(params.mode_id)
-            if not confirmation_mode:
-                logger.error("Invalid mode_id", mode_id=params.mode_id)
-                return None
-
-            # All agent types support set_tool_confirmation_mode
-            # ACPAgent handles forwarding to remote server internally
-            await session.agent.set_tool_confirmation_mode(confirmation_mode)
+            # Call set_mode directly - agent handles mode-specific logic
+            await session.agent.set_mode(params.mode_id)
 
             # Update stored mode state for ACPAgent
             if (
@@ -665,9 +700,8 @@ class AgentPoolACPAgent(ACPAgent):
                 session.agent._state.modes.current_mode_id = params.mode_id
 
             logger.info(
-                "Set tool confirmation mode",
+                "Set mode",
                 mode_id=params.mode_id,
-                confirmation_mode=confirmation_mode,
                 session_id=params.session_id,
                 agent_type=type(session.agent).__name__,
             )
@@ -677,19 +711,15 @@ class AgentPoolACPAgent(ACPAgent):
             logger.exception("Failed to set session mode", session_id=params.session_id)
             return None
 
-    async def set_session_model(  # noqa: PLR0911
+    async def set_session_model(
         self, params: SetSessionModelRequest
     ) -> SetSessionModelResponse | None:
         """Set the session model.
 
         Changes the model for the active agent in the session.
-        Validates that the requested model is in the available models list:
-        - For ClaudeCodeAgent: validates against Claude Code models (sonnet, opus, haiku)
-        - For native Agent: validates against server's available_models
-        - For ACPAgent: validates against the nested agent's model list
+        Validates that the requested model is available via agent.get_available_models().
         """
         from agentpool.agents.acp_agent import ACPAgent as ACPAgentClient
-        from agentpool.agents.claude_code_agent import ClaudeCodeAgent
 
         try:
             session = self.session_manager.get_session(params.session_id)
@@ -698,23 +728,8 @@ class AgentPoolACPAgent(ACPAgent):
                 logger.warning(msg, session_id=params.session_id)
                 return None
 
-            # Validate model based on agent type
-            if isinstance(session.agent, ClaudeCodeAgent):
-                # For ClaudeCodeAgent, validate against our hardcoded models
-                available_ids = [m.model_id for m in CLAUDE_CODE_MODELS]
-                if params.model_id not in available_ids:
-                    logger.warning(
-                        "Model not in Claude Code available models",
-                        model_id=params.model_id,
-                        available=available_ids,
-                    )
-                    return None
-                await session.agent.set_model(params.model_id)
-                logger.info("Set model", model_id=params.model_id, session_id=params.session_id)
-                return SetSessionModelResponse()
-
+            # Handle ACPAgent specially - pass through to nested agent's model state
             if isinstance(session.agent, ACPAgentClient):
-                # For ACPAgent, validate against nested agent's model list
                 if session.agent._state and session.agent._state.models:
                     available_ids = [
                         m.model_id for m in session.agent._state.models.available_models
@@ -727,28 +742,66 @@ class AgentPoolACPAgent(ACPAgent):
                         )
                         return None
                 # TODO: Use ACP protocol once set_session_model stabilizes
-                # For now, we can't actually change the model on ACPAgent
                 logger.warning(
                     "Model change for ACPAgent not yet supported (ACP protocol UNSTABLE)",
                     model_id=params.model_id,
                 )
                 return None
 
-            if isinstance(session.agent, Agent):
-                # For native Agent, validate against server's available models
-                available_ids = [m.model_id for m in self.available_models]
+            # Get available models from agent and validate
+            toko_models = await session.agent.get_available_models()
+            if toko_models:
+                # Build list of valid model IDs (using id_override if set)
+                available_ids = [m.id_override if m.id_override else m.id for m in toko_models]
                 if params.model_id not in available_ids:
-                    msg = "Model not in available models"
-                    logger.warning(msg, model_id=params.model_id, available=available_ids)
+                    logger.warning(
+                        "Model not in available models",
+                        model_id=params.model_id,
+                        available=available_ids,
+                    )
                     return None
-                session.agent.set_model(params.model_id)
 
+            # Set the model on the agent (all agents now have async set_model)
+            await session.agent.set_model(params.model_id)
             logger.info("Set model", model_id=params.model_id, session_id=params.session_id)
             return SetSessionModelResponse()
         except Exception:
             logger.exception("Failed to set session model", session_id=params.session_id)
             return None
 
+    async def set_session_config_option(
+        self, params: SetSessionConfigOptionRequest
+    ) -> SetSessionConfigOptionResponse | None:
+        """Set a session config option.
+
+        Forwards the config option change to the agent's set_mode method
+        and returns the updated config options.
+        """
+        try:
+            session = self.session_manager.get_session(params.session_id)
+            if not session or not session.agent:
+                msg = "Session not found for config option change"
+                logger.warning(msg, session_id=params.session_id)
+                return None
+
+            logger.info(
+                "Set config option",
+                config_id=params.config_id,
+                value=params.value,
+                session_id=params.session_id,
+            )
+
+            # Forward to agent's set_mode method
+            # config_id maps to category_id, value maps to mode_id
+            await session.agent.set_mode(params.value, category_id=params.config_id)
+
+            # Return updated config options
+            config_options = await get_session_config_options(session.agent)
+            return SetSessionConfigOptionResponse(config_options=config_options)
+        except Exception:
+            logger.exception("Failed to set session config option", session_id=params.session_id)
+            return None
+
     async def swap_pool(self, config_path: str, agent: str | None = None) -> list[str]:
         """Swap the agent pool with a new one from configuration.