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

agentpool/agents/agui_agent/agui_agent.py

@@ -34,8 +34,9 @@ from agentpool.agents.events import RunStartedEvent, StreamCompleteEvent
 from agentpool.log import get_logger
 from agentpool.messaging import ChatMessage
 from agentpool.messaging.processing import prepare_prompts
-from agentpool.talk.stats import MessageStats
 from agentpool.tools import ToolManager
+from agentpool.utils.streams import FileTracker
+from agentpool.utils.token_breakdown import calculate_usage_from_parts
 
 
 if TYPE_CHECKING:
@@ -45,10 +46,13 @@ if TYPE_CHECKING:
 
     from ag_ui.core import Message, ToolMessage
     from evented.configs import EventConfig
+    from slashed import BaseCommand
+    from tokonomics.model_discovery.model_info import ModelInfo
 
     from agentpool.agents.base_agent import ToolConfirmationMode
     from agentpool.agents.context import AgentContext
     from agentpool.agents.events import RichAgentStreamEvent
+    from agentpool.agents.modes import ModeCategory, ModeInfo
     from agentpool.common_types import (
         BuiltinEventHandlerType,
         IndividualEventHandler,
@@ -132,6 +136,7 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         event_configs: Sequence[EventConfig] | None = None,
         event_handlers: Sequence[IndividualEventHandler | BuiltinEventHandlerType] | None = None,
         tool_confirmation_mode: ToolConfirmationMode = "per_tool",
+        commands: Sequence[BaseCommand] | None = None,
     ) -> None:
         """Initialize AG-UI agent client.
 
@@ -153,6 +158,7 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
             event_configs: Event trigger configurations
             event_handlers: Sequence of event handlers to register
             tool_confirmation_mode: Tool confirmation mode
+            commands: Slash commands
         """
         super().__init__(
             name=name,
@@ -164,6 +170,7 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
             event_configs=event_configs,
             tool_confirmation_mode=tool_confirmation_mode,
             event_handlers=event_handlers,
+            commands=commands,
         )
 
         # AG-UI specific configuration
@@ -379,11 +386,11 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         )
 
         from agentpool.agents.agui_agent.agui_converters import (
-            ToolCallAccumulator,
             agui_to_native_event,
             to_agui_input_content,
             to_agui_tool,
         )
+        from agentpool.agents.tool_call_accumulator import ToolCallAccumulator
 
         if not self._client or not self._thread_id:
             msg = "Agent not initialized - use async context manager"
@@ -394,7 +401,11 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         self._current_stream_task = asyncio.current_task()
 
         conversation = message_history if message_history is not None else self.conversation
-        user_msg, processed_prompts, _original_message = await prepare_prompts(*prompts)
+        # Get parent_id from last message in history for tree structure
+        last_msg_id = conversation.get_last_message_id()
+        user_msg, processed_prompts, _original_message = await prepare_prompts(
+            *prompts, parent_id=last_msg_id
+        )
         self._run_id = str(uuid4())  # New run ID for each run
         self._chunk_transformer.reset()  # Reset chunk transformer
         # Track messages in pydantic-ai format: ModelRequest -> ModelResponse -> ModelRequest...
@@ -411,10 +422,9 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
             run_id=self._run_id or str(uuid4()),
             agent_name=self.name,
         )
-        for handler in self.event_handler._wrapped_handlers:
-            await handler(None, run_started)
-        yield run_started
 
+        await self.event_handler(None, run_started)
+        yield run_started
         # Get pending parts from conversation and convert them
         pending_parts = conversation.get_pending_parts()
         pending_content = to_agui_input_content(pending_parts)
@@ -432,6 +442,8 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         tool_accumulator = ToolCallAccumulator()
         pending_tool_results: list[ToolMessage] = []
         self.log.debug("Sending prompt to AG-UI agent", tool_names=[t.name for t in agui_tools])
+        # Track files modified during this run
+        file_tracker = FileTracker()
         # Loop to handle tool calls - agent may request multiple rounds
         try:
             while True:
@@ -509,25 +521,24 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
 
                                 # Convert to native event and distribute to handlers
                                 if native_event := agui_to_native_event(event):
+                                    # Track file modifications
+                                    file_tracker.process_event(native_event)
                                     # Check for queued custom events first
                                     while not self._event_queue.empty():
                                         try:
                                             custom_event = self._event_queue.get_nowait()
-                                            for handler in self.event_handler._wrapped_handlers:
-                                                await handler(None, custom_event)
+                                            await self.event_handler(None, custom_event)
                                             yield custom_event
                                         except asyncio.QueueEmpty:
                                             break
                                     # Distribute to handlers
-                                    for handler in self.event_handler._wrapped_handlers:
-                                        await handler(None, native_event)
+                                    await self.event_handler(None, native_event)
                                     yield native_event
 
                         # Flush any pending chunk events at end of stream
                         for event in self._chunk_transformer.flush():
                             if native_event := agui_to_native_event(event):
-                                for handler in self.event_handler._wrapped_handlers:
-                                    await handler(None, native_event)
+                                await self.event_handler(None, native_event)
                                 yield native_event
 
                 except httpx.HTTPError:
@@ -592,12 +603,13 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
                 name=self.name,
                 message_id=message_id or str(uuid4()),
                 conversation_id=self.conversation_id,
+                parent_id=user_msg.message_id,
                 messages=model_messages,
                 finish_reason="stop",
+                metadata=file_tracker.get_metadata(),
             )
             complete_event = StreamCompleteEvent(message=final_message)
-            for handler in self.event_handler._wrapped_handlers:
-                await handler(None, complete_event)
+            await self.event_handler(None, complete_event)
             yield complete_event
             self._current_stream_task = None
             return
@@ -610,24 +622,36 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         while not self._event_queue.empty():
             try:
                 queued_event = self._event_queue.get_nowait()
-                for handler in self.event_handler._wrapped_handlers:
-                    await handler(None, queued_event)
+                await self.event_handler(None, queued_event)
                 yield queued_event
             except asyncio.QueueEmpty:
                 break
 
         text_content = "".join(text_chunks)
+
+        # Calculate approximate token usage from what we can observe
+        input_parts = [*processed_prompts, *pending_parts]
+        usage, cost_info = await calculate_usage_from_parts(
+            input_parts=input_parts,
+            response_parts=current_response_parts,
+            text_content=text_content,
+            model_name=self.model_name,
+        )
+
         final_message = ChatMessage[str](
             content=text_content,
             role="assistant",
             name=self.name,
             message_id=message_id or str(uuid4()),
             conversation_id=self.conversation_id,
+            parent_id=user_msg.message_id,
             messages=model_messages,
+            metadata=file_tracker.get_metadata(),
+            usage=usage,
+            cost_info=cost_info,
         )
         complete_event = StreamCompleteEvent(message=final_message)
-        for handler in self.event_handler._wrapped_handlers:
-            await handler(None, complete_event)
+        await self.event_handler(None, complete_event)
         yield complete_event
         # Record to conversation history
         conversation.add_chat_messages([user_msg, final_message])
@@ -657,9 +681,54 @@ class AGUIAgent[TDeps = None](BaseAgent[TDeps, str]):
         """Get model name (AG-UI doesn't expose this)."""
         return None
 
-    async def get_stats(self) -> MessageStats:
-        """Get message statistics for this node."""
-        return MessageStats()
+    async def set_model(self, model: str) -> None:
+        """Set model (no-op for AG-UI as model is controlled by remote server)."""
+        # AG-UI agents don't support model selection - the model is
+        # determined by the remote server configuration
+
+    async def get_available_models(self) -> list[ModelInfo] | None:
+        """Get available models for AG-UI agent.
+
+        AG-UI doesn't expose model information, so returns a placeholder model
+        indicating the model is determined by the remote server.
+
+        Returns:
+            List with a single placeholder ModelInfo
+        """
+        from tokonomics.model_discovery.model_info import ModelInfo
+
+        return [
+            ModelInfo(
+                id="server-determined",
+                name="Determined by server",
+                description="The model is determined by the remote AG-UI server",
+            )
+        ]
+
+    def get_modes(self) -> list[ModeCategory]:
+        """Get available modes for AG-UI agent.
+
+        AG-UI doesn't expose mode information, so returns an empty list.
+
+        Returns:
+            Empty list - no modes supported
+        """
+        return []
+
+    async def set_mode(self, mode: ModeInfo | str, category_id: str | None = None) -> None:
+        """Set a mode for AG-UI agent.
+
+        AG-UI doesn't support modes, so this always raises an error.
+
+        Args:
+            mode: The mode to set (not supported)
+            category_id: Optional category ID (not supported)
+
+        Raises:
+            ValueError: Always - AG-UI doesn't support modes
+        """
+        msg = "AG-UI agent does not support modes"
+        raise ValueError(msg)
 
 
 if __name__ == "__main__":

agentpool/agents/agui_agent/agui_converters.py

@@ -12,7 +12,6 @@ from __future__ import annotations
 import base64
 from typing import TYPE_CHECKING, Any
 
-import anyenv
 from pydantic_ai import (
     AudioUrl,
     BinaryContent,
@@ -291,133 +290,3 @@ def to_agui_tool(tool: Tool) -> AGUITool:
         description=func_schema.get("description", ""),
         parameters=func_schema.get("parameters", {"type": "object", "properties": {}}),
     )
-
-
-def _repair_partial_json(buffer: str) -> str:
-    """Attempt to repair truncated JSON for preview purposes.
-
-    Handles common truncation cases:
-    - Unclosed strings
-    - Missing closing braces/brackets
-    - Trailing commas
-
-    Args:
-        buffer: Potentially incomplete JSON string
-
-    Returns:
-        Repaired JSON string (may still be invalid in edge cases)
-    """
-    if not buffer:
-        return "{}"
-
-    result = buffer.rstrip()
-
-    # Check if we're in the middle of a string by counting unescaped quotes
-    in_string = False
-    i = 0
-    while i < len(result):
-        char = result[i]
-        if char == "\\" and i + 1 < len(result):
-            i += 2  # Skip escaped character
-            continue
-        if char == '"':
-            in_string = not in_string
-        i += 1
-
-    # Close unclosed string
-    if in_string:
-        result += '"'
-
-    # Remove trailing comma (invalid JSON)
-    result = result.rstrip()
-    if result.endswith(","):
-        result = result[:-1]
-
-    # Balance braces and brackets
-    open_braces = result.count("{") - result.count("}")
-    open_brackets = result.count("[") - result.count("]")
-
-    result += "]" * max(0, open_brackets)
-    result += "}" * max(0, open_braces)
-
-    return result
-
-
-class ToolCallAccumulator:
-    """Accumulates streamed tool call arguments.
-
-    AG-UI streams tool call arguments as deltas, this class accumulates them
-    and provides the complete arguments when the tool call ends.
-    """
-
-    def __init__(self) -> None:
-        self._calls: dict[str, dict[str, Any]] = {}
-
-    def start(self, tool_call_id: str, tool_name: str) -> None:
-        """Start tracking a new tool call."""
-        self._calls[tool_call_id] = {"name": tool_name, "args_buffer": ""}
-
-    def add_args(self, tool_call_id: str, delta: str) -> None:
-        """Add argument delta to a tool call."""
-        if tool_call_id in self._calls:
-            self._calls[tool_call_id]["args_buffer"] += delta
-
-    def complete(self, tool_call_id: str) -> tuple[str, dict[str, Any]] | None:
-        """Complete a tool call and return (tool_name, parsed_args).
-
-        Returns:
-            Tuple of (tool_name, args_dict) or None if call not found
-        """
-        if tool_call_id not in self._calls:
-            return None
-
-        call_data = self._calls.pop(tool_call_id)
-        args_str = call_data["args_buffer"]
-        try:
-            args = anyenv.load_json(args_str) if args_str else {}
-        except anyenv.JsonLoadError:
-            args = {"raw": args_str}
-        return call_data["name"], args
-
-    def get_pending(self, tool_call_id: str) -> tuple[str, str] | None:
-        """Get pending call data (tool_name, args_buffer) without completing."""
-        if tool_call_id not in self._calls:
-            return None
-        data = self._calls[tool_call_id]
-        return data["name"], data["args_buffer"]
-
-    def get_partial_args(self, tool_call_id: str) -> dict[str, Any]:
-        """Get best-effort parsed args from incomplete JSON stream.
-
-        Uses heuristics to complete truncated JSON for preview purposes.
-        Handles unclosed strings, missing braces/brackets, and trailing commas.
-
-        Args:
-            tool_call_id: Tool call ID
-
-        Returns:
-            Partially parsed arguments or empty dict
-        """
-        if tool_call_id not in self._calls:
-            return {}
-
-        buffer = self._calls[tool_call_id]["args_buffer"]
-        if not buffer:
-            return {}
-
-        # Try direct parse first
-        try:
-            return anyenv.load_json(buffer)
-        except anyenv.JsonLoadError:
-            pass
-
-        # Try to repair truncated JSON
-        try:
-            repaired = _repair_partial_json(buffer)
-            return anyenv.load_json(repaired)
-        except anyenv.JsonLoadError:
-            return {}
-
-    def clear(self) -> None:
-        """Clear all pending tool calls."""
-        self._calls.clear()

agentpool/agents/base_agent.py

@@ -7,6 +7,7 @@ import asyncio
 from typing import TYPE_CHECKING, Any, Literal
 
 from anyenv import MultiEventHandler
+from anyenv.signals import BoundSignal
 from exxec import LocalExecutionEnvironment
 
 from agentpool.agents.events import resolve_event_handlers
@@ -20,17 +21,30 @@ if TYPE_CHECKING:
 
     from evented.configs import EventConfig
     from exxec import ExecutionEnvironment
+    from slashed import BaseCommand, CommandStore
+    from tokonomics.model_discovery.model_info import ModelInfo
 
+    from acp.schema import AvailableCommandsUpdate, ConfigOptionUpdate
     from agentpool.agents.context import AgentContext
     from agentpool.agents.events import RichAgentStreamEvent
-    from agentpool.common_types import BuiltinEventHandlerType, IndividualEventHandler
+    from agentpool.agents.modes import ModeCategory, ModeInfo
+    from agentpool.common_types import (
+        BuiltinEventHandlerType,
+        IndividualEventHandler,
+        MCPServerStatus,
+    )
     from agentpool.delegation import AgentPool
+    from agentpool.talk.stats import MessageStats
     from agentpool.ui.base import InputProvider
     from agentpool_config.mcp_server import MCPServerConfig
 
+    # Union type for state updates emitted via state_updated signal
+    type StateUpdate = ModeInfo | ModelInfo | AvailableCommandsUpdate | ConfigOptionUpdate
+
 
 logger = get_logger(__name__)
 
+
 ToolConfirmationMode = Literal["always", "never", "per_tool"]
 
 
@@ -64,6 +78,7 @@ class BaseAgent[TDeps = None, TResult = str](MessageNode[TDeps, TResult]):
         output_type: type[TResult] = str,  # type: ignore[assignment]
         tool_confirmation_mode: ToolConfirmationMode = "per_tool",
         event_handlers: Sequence[IndividualEventHandler | BuiltinEventHandlerType] | None = None,
+        commands: Sequence[BaseCommand] | None = None,
     ) -> None:
         """Initialize base agent with shared infrastructure.
 
@@ -80,6 +95,7 @@ class BaseAgent[TDeps = None, TResult = str](MessageNode[TDeps, TResult]):
             output_type: Output type for this agent
             tool_confirmation_mode: How tool execution confirmation is handled
             event_handlers: Event handlers for this agent
+            commands: Slash commands to register with this agent
         """
         super().__init__(
             name=name,
@@ -108,6 +124,31 @@ class BaseAgent[TDeps = None, TResult = str](MessageNode[TDeps, TResult]):
         self._cancelled = False
         self._current_stream_task: asyncio.Task[Any] | None = None
 
+        # State change signal - emitted when mode/model/commands change
+        # Uses union type for different state update kinds
+        self.state_updated: BoundSignal[StateUpdate] = BoundSignal()
+
+        # Command store for slash commands
+        from slashed import CommandStore
+
+        from agentpool_commands import get_commands
+
+        self._command_store: CommandStore = CommandStore()
+
+        # Register default agent commands
+        for command in get_commands():
+            self._command_store.register_command(command)
+
+        # Register additional provided commands
+        if commands:
+            for command in commands:
+                self._command_store.register_command(command)
+
+    @property
+    def command_store(self) -> CommandStore:
+        """Get the command store for slash commands."""
+        return self._command_store
+
     @abstractmethod
     def get_context(self, data: Any = None) -> AgentContext[Any]:
         """Create a new context for this agent.
@@ -126,6 +167,15 @@ class BaseAgent[TDeps = None, TResult = str](MessageNode[TDeps, TResult]):
         """Get the model name used by this agent."""
         ...
 
+    @abstractmethod
+    async def set_model(self, model: str) -> None:
+        """Set the model for this agent.
+
+        Args:
+            model: New model identifier to use
+        """
+        ...
+
     @abstractmethod
     def run_stream(
         self,
@@ -175,3 +225,115 @@ class BaseAgent[TDeps = None, TResult = str](MessageNode[TDeps, TResult]):
         if self._current_stream_task and not self._current_stream_task.done():
             self._current_stream_task.cancel()
             logger.info("Interrupted agent stream", agent=self.name)
+
+    async def get_stats(self) -> MessageStats:
+        """Get message statistics."""
+        from agentpool.talk.stats import MessageStats
+
+        return MessageStats(messages=list(self.conversation.chat_messages))
+
+    def get_mcp_server_info(self) -> dict[str, MCPServerStatus]:
+        """Get information about configured MCP servers.
+
+        Returns a dict mapping server names to their status info. Used by
+        the OpenCode /mcp endpoint to display MCP servers in the UI.
+
+        The default implementation checks external_providers on the tool manager.
+        Subclasses may override to provide agent-specific MCP server info
+        (e.g., ClaudeCodeAgent has its own MCP server handling).
+
+        Returns:
+            Dict mapping server name to MCPServerStatus
+        """
+        from agentpool.common_types import MCPServerStatus
+        from agentpool.mcp_server.manager import MCPManager
+        from agentpool.resource_providers import AggregatingResourceProvider
+        from agentpool.resource_providers.mcp_provider import MCPResourceProvider
+
+        def add_status(provider: MCPResourceProvider, result: dict[str, MCPServerStatus]) -> None:
+            status_dict = provider.get_status()
+            status_type = status_dict.get("status", "disabled")
+            if status_type == "connected":
+                result[provider.name] = MCPServerStatus(
+                    name=provider.name, status="connected", server_type="stdio"
+                )
+            elif status_type == "failed":
+                error = status_dict.get("error", "Unknown error")
+                result[provider.name] = MCPServerStatus(
+                    name=provider.name, status="error", error=error
+                )
+            else:
+                result[provider.name] = MCPServerStatus(name=provider.name, status="disconnected")
+
+        result: dict[str, MCPServerStatus] = {}
+        try:
+            for provider in self.tools.external_providers:
+                if isinstance(provider, MCPResourceProvider):
+                    add_status(provider, result)
+                elif isinstance(provider, AggregatingResourceProvider):
+                    for nested in provider.providers:
+                        if isinstance(nested, MCPResourceProvider):
+                            add_status(nested, result)
+                elif isinstance(provider, MCPManager):
+                    for mcp_provider in provider.get_mcp_providers():
+                        add_status(mcp_provider, result)
+        except Exception:  # noqa: BLE001
+            pass
+
+        return result
+
+    @abstractmethod
+    async def get_available_models(self) -> list[ModelInfo] | None:
+        """Get available models for this agent.
+
+        Returns a list of models that can be used with this agent, or None
+        if model discovery is not supported for this agent type.
+
+        Uses tokonomics.ModelInfo which includes pricing, capabilities,
+        and limits. Can be converted to protocol-specific formats (OpenCode, ACP).
+
+        Returns:
+            List of tokonomics ModelInfo, or None if not supported
+        """
+        ...
+
+    @abstractmethod
+    def get_modes(self) -> list[ModeCategory]:
+        """Get available mode categories for this agent.
+
+        Returns a list of mode categories that can be switched. Each category
+        represents a group of mutually exclusive modes (e.g., permissions,
+        behavior presets).
+
+        Different agent types expose different modes:
+        - Native Agent: Tool confirmation modes (default, acceptEdits)
+        - ClaudeCodeAgent: Claude Code SDK modes (plan, code, etc.)
+        - ACPAgent: Passthrough from remote server
+        - AGUIAgent: Empty list (no modes)
+
+        Returns:
+            List of ModeCategory, empty list if no modes supported
+        """
+        ...
+
+    @abstractmethod
+    async def set_mode(self, mode: ModeInfo | str, category_id: str | None = None) -> None:
+        """Set a mode within a category.
+
+        Each agent type handles mode switching according to its own semantics:
+        - Native Agent: Maps to tool confirmation mode
+        - ClaudeCodeAgent: Maps to SDK permission mode
+        - ACPAgent: Forwards to remote server
+        - AGUIAgent: No-op (no modes supported)
+
+        Args:
+            mode: The mode to activate - either a ModeInfo object or mode ID string.
+                  If ModeInfo, category_id is extracted from it (unless overridden).
+            category_id: Optional category ID. If None and mode is a string,
+                         uses the first category. If None and mode is ModeInfo,
+                         uses the mode's category_id.
+
+        Raises:
+            ValueError: If mode_id or category_id is invalid
+        """
+        ...