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

agentpool_server/acp_server/acp_agent.py

@@ -17,6 +17,7 @@ from acp.schema import (
     PromptResponse,
     ResumeSessionResponse,
     SessionInfo,
+    SessionMode,
     SessionModelState,
     SessionModeState,
     SetSessionModelRequest,
@@ -24,20 +25,16 @@ from acp.schema import (
     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
@@ -63,59 +60,111 @@ 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
+        acp_models.append(
+            ACPModelInfo(
+                model_id=model_id,
+                name=toko.name,
+                description=toko.description or "",
+            )
+        )
+
+    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
+        acp_models.insert(
+            0,
+            ACPModelInfo(
+                model_id=current_model,
+                name=current_model,
+                description="Currently configured model",
+            ),
+        )
+        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)
+
+
+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.
+    Currently uses the first mode category (ACP only supports one dropdown for now).
 
     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 = agent.get_modes()
+    except Exception:
+        logger.exception("Failed to get modes from agent")
+        return None
+
+    if not mode_categories:
         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
+
+    # Use first category for now (ACP currently only supports single mode dropdown)
+    category = mode_categories[0]
+
+    # 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,
     )
 
 
@@ -137,9 +186,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."""
 
@@ -176,17 +222,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 +240,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 +269,22 @@ 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
 
-            # 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
                 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 = get_session_mode_state(session.agent)
         except Exception:
             logger.exception("Failed to create new session")
             raise
@@ -285,7 +298,7 @@ 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))
+            logger.info("Created session", session_id=session_id)
             return NewSessionResponse(session_id=session_id, modes=state, models=models)
 
     async def load_session(self, params: LoadSessionRequest) -> LoadSessionResponse:
@@ -332,38 +345,22 @@ class AgentPoolACPAgent(ACPAgent):
                 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
+
             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)
+                    models = session.agent._state.models
+            elif session.agent:
+                # Use unified helpers for all other agents
+                mode_state = get_session_mode_state(session.agent)
+                models = await get_session_model_state(session.agent, session.agent.model_name)
             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)
-            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())
@@ -677,19 +674,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 +691,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,21 +705,27 @@ 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()