agentpool 2.1.9__py3-none-any.whl

agentpool/mcp_server/tool_bridge.py

@@ -0,0 +1,548 @@
+"""MCP server bridge for exposing ToolManager tools to ACP agents.
+
+This module provides a bridge that exposes a ToolManager's tools as an MCP server
+using HTTP transport. This allows ACP agents (external agents like Claude Code,
+Gemini CLI, etc.) to use our internal toolsets like SubagentTools,
+AgentManagementTools, etc.
+
+The bridge runs in-process on the same event loop, providing direct access to
+the pool and avoiding IPC serialization overhead.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import asynccontextmanager, suppress
+from dataclasses import dataclass, field, replace
+import inspect
+from typing import TYPE_CHECKING, Any, Literal, Self, get_args, get_origin
+from uuid import uuid4
+
+import anyio
+from fastmcp import FastMCP
+from fastmcp.tools import Tool as FastMCPTool
+from llmling_models.models.helpers import infer_model
+from pydantic import BaseModel, HttpUrl
+
+from agentpool.agents import Agent, ClaudeCodeAgent
+from agentpool.agents.acp_agent.acp_agent import ACPAgent
+from agentpool.log import get_logger
+from agentpool.utils.signatures import filter_schema_params, get_params_matching_predicate
+
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Callable
+
+    from claude_agent_sdk.types import McpServerConfig
+    from fastmcp import Context
+    from fastmcp.tools.tool import ToolResult
+    from pydantic_ai import RunContext
+    from uvicorn import Server
+
+    from acp.schema.mcp import HttpMcpServer, SseMcpServer
+    from agentpool.agents import AgentContext
+    from agentpool.agents.base_agent import BaseAgent
+    from agentpool.tools.base import Tool
+
+
+logger = get_logger(__name__)
+
+
+def _is_agent_context_type(annotation: Any) -> bool:
+    """Check if annotation is AgentContext."""
+    if annotation is None or annotation is inspect.Parameter.empty:
+        return False
+
+    # Handle string annotations (forward references)
+    if isinstance(annotation, str):
+        base_name = annotation.split("[")[0].strip()
+        return base_name == "AgentContext"
+
+    # Check direct class match by name
+    if isinstance(annotation, type) and annotation.__name__ == "AgentContext":
+        return True
+
+    # Check generic origin (e.g., AgentContext[SomeDeps])
+    origin = get_origin(annotation)
+    if origin is not None:
+        if isinstance(origin, type) and origin.__name__ == "AgentContext":
+            return True
+        # Handle Union types (e.g., AgentContext | None)
+        if origin is type(None) or str(origin) in ("typing.Union", "types.UnionType"):
+            return any(_is_agent_context_type(arg) for arg in get_args(annotation))
+
+    return False
+
+
+def _get_context_param_names(fn: Callable[..., Any]) -> set[str]:
+    """Get parameter names that are AgentContext (to be auto-injected)."""
+    return get_params_matching_predicate(fn, lambda p: _is_agent_context_type(p.annotation))
+
+
+def _is_run_context_type(annotation: Any) -> bool:
+    """Check if annotation is pydantic-ai RunContext."""
+    if annotation is None or annotation is inspect.Parameter.empty:
+        return False
+
+    # Handle string annotations (forward references)
+    if isinstance(annotation, str):
+        base_name = annotation.split("[")[0].strip()
+        return base_name == "RunContext"
+
+    # Check direct class match by name
+    if isinstance(annotation, type) and annotation.__name__ == "RunContext":
+        return True
+
+    # Check generic origin (e.g., RunContext[SomeDeps])
+    origin = get_origin(annotation)
+    if origin is not None:
+        if isinstance(origin, type) and origin.__name__ == "RunContext":
+            return True
+        # Handle Union types (e.g., RunContext | None)
+        if origin is type(None) or str(origin) in ("typing.Union", "types.UnionType"):
+            return any(_is_run_context_type(arg) for arg in get_args(annotation))
+
+    return False
+
+
+def _get_run_context_param_names(fn: Callable[..., Any]) -> set[str]:
+    """Get parameter names that are RunContext (to be auto-injected)."""
+    return get_params_matching_predicate(fn, lambda p: _is_run_context_type(p.annotation))
+
+
+def _create_stub_run_context(ctx: AgentContext[Any]) -> RunContext[Any]:
+    """Create a stub RunContext from AgentContext for MCP bridge tool invocations.
+
+    This provides a best-effort RunContext for tools that require it when
+    invoked via MCP bridge. Not all fields can be populated accurately since
+    we're outside of pydantic-ai's normal execution flow.
+
+    Args:
+        ctx: The AgentContext available in the bridge
+
+    Returns:
+        A RunContext with available information populated
+    """
+    from pydantic_ai import RunContext
+    from pydantic_ai.models.test import TestModel
+    from pydantic_ai.usage import RunUsage
+
+    match ctx.agent:
+        case Agent():
+            model = ctx.agent._model or TestModel()
+        case ACPAgent() | ClaudeCodeAgent():
+            try:
+                model = infer_model(ctx.agent.model_name or "test")
+            except Exception:  # noqa: BLE001
+                model = TestModel()
+        case _:
+            model = TestModel()
+    # Create a minimal usage object
+    return RunContext(
+        deps=ctx.data,
+        model=model,
+        usage=RunUsage(),
+        prompt=None,
+        messages=[],
+        tool_name=ctx.tool_name,
+        tool_call_id=ctx.tool_call_id,
+    )
+
+
+def _convert_to_tool_result(result: Any) -> ToolResult:
+    """Convert a tool's return value to a FastMCP ToolResult.
+
+    Handles different result types appropriately:
+    - ToolResult: Pass through unchanged
+    - dict: Use as structured_content (enables programmatic access by clients)
+    - Pydantic models: Serialize to dict for structured_content
+    - Other types: Pass to ToolResult(content=...) which handles conversion internally
+    """
+    from fastmcp.tools.tool import ToolResult
+
+    # Already a ToolResult - pass through
+    if isinstance(result, ToolResult):
+        return result
+    # Dict - use as structured_content (FastMCP auto-populates content as JSON)
+    if isinstance(result, dict):
+        return ToolResult(structured_content=result)
+    # Pydantic model - serialize to dict for structured_content
+    if isinstance(result, BaseModel):
+        return ToolResult(structured_content=result.model_dump(mode="json"))
+    # All other types (str, list, ContentBlock, Image, None, primitives, etc.)
+    # ToolResult's internal _convert_to_content handles these correctly
+    return ToolResult(content=result if result is not None else "")
+
+
+def _extract_tool_call_id(context: Context | None) -> str:
+    """Extract Claude's original tool_call_id from request metadata.
+
+    Claude Code passes the tool_use_id via the _meta field as
+    'claudecode/toolUseId'. This allows us to maintain consistent
+    tool_call_ids across ToolCallStartEvent and ToolCallCompleteEvent.
+
+    Falls back to generating a UUID if not available.
+    """
+    if context is not None:
+        try:
+            request_ctx = context.request_context
+            if request_ctx and request_ctx.meta:
+                # Access extra fields on the Meta object (extra="allow" in pydantic)
+                meta_dict = request_ctx.meta.model_dump()
+                claude_tool_id = meta_dict.get("claudecode/toolUseId")
+                if isinstance(claude_tool_id, str):
+                    logger.debug("Extracted Claude tool_call_id", tool_call_id=claude_tool_id)
+                    return claude_tool_id
+        except (AttributeError, LookupError) as e:
+            logger.warning("Error extracting tool_call_id from MCP context", error=str(e))
+    # Generate fallback UUID if no tool_call_id found in meta
+    return str(uuid4())
+
+
+@dataclass
+class BridgeConfig:
+    """Configuration for the ToolManager MCP bridge."""
+
+    host: str = "127.0.0.1"
+    """Host to bind the HTTP server to."""
+
+    port: int = 0
+    """Port to bind to (0 = auto-select available port)."""
+
+    transport: Literal["sse", "streamable-http"] = "sse"
+    """Transport protocol: 'sse' or 'streamable-http'."""
+
+    server_name: str = "agentpool-toolmanager"
+    """Name for the MCP server."""
+
+
+@dataclass
+class ToolManagerBridge:
+    """Exposes a node's tools as an MCP server for ACP agents.
+
+    This bridge allows external ACP agents to access our internal toolsets
+    (SubagentTools, AgentManagementTools, etc.) via HTTP MCP transport.
+
+    The node's existing context is used for tool invocations, providing
+    pool access and proper configuration without reconstruction.
+
+    Example:
+        ```python
+        async with AgentPool() as pool:
+            agent = pool.agents["my_agent"]
+            bridge = ToolManagerBridge(node=agent, config=BridgeConfig(port=8765))
+            async with bridge:
+                # Bridge is running, get MCP config for ACP agent
+                mcp_config = bridge.get_mcp_server_config()
+                # Pass to ACP agent...
+        ```
+    """
+
+    node: BaseAgent[Any, Any]
+    """The node whose tools to expose."""
+
+    config: BridgeConfig = field(default_factory=BridgeConfig)
+    """Bridge configuration."""
+
+    _mcp: FastMCP | None = field(default=None, init=False, repr=False)
+    """FastMCP server instance."""
+
+    _server: Server | None = field(default=None, init=False, repr=False)
+    """Uvicorn server instance."""
+
+    _server_task: asyncio.Task[None] | None = field(default=None, init=False, repr=False)
+    """Background task running the server."""
+
+    _actual_port: int | None = field(default=None, init=False, repr=False)
+    """Actual port the server is bound to."""
+
+    async def __aenter__(self) -> Self:
+        """Start the MCP server."""
+        await self.start()
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        """Stop the MCP server."""
+        await self.stop()
+
+    async def start(self) -> None:
+        """Start the HTTP MCP server in the background."""
+        self._mcp = FastMCP(name=self.config.server_name)
+        await self._register_tools()
+        await self._start_server()
+
+    async def stop(self) -> None:
+        """Stop the HTTP MCP server."""
+        if self._server:
+            self._server.should_exit = True
+            if self._server_task:
+                try:
+                    await asyncio.wait_for(self._server_task, timeout=5.0)
+                except TimeoutError:
+                    self._server_task.cancel()
+                    with suppress(asyncio.CancelledError):
+                        await self._server_task
+            self._server = None
+            self._server_task = None
+        self._mcp = None
+        self._actual_port = None
+        logger.info("ToolManagerBridge stopped")
+
+    @property
+    def port(self) -> int:
+        """Get the actual port the server is running on."""
+        if self._actual_port is None:
+            msg = "Server not started"
+            raise RuntimeError(msg)
+        return self._actual_port
+
+    @property
+    def url(self) -> str:
+        """Get the server URL."""
+        path = "/sse" if self.config.transport == "sse" else "/mcp"
+        return f"http://{self.config.host}:{self.port}{path}"
+
+    def get_mcp_server_config(self) -> HttpMcpServer | SseMcpServer:
+        """Get ACP-compatible MCP server configuration.
+
+        Returns config suitable for passing to ACP agent's NewSessionRequest.
+        """
+        from acp.schema import HttpMcpServer, SseMcpServer
+
+        url = HttpUrl(self.url)
+        if self.config.transport == "sse":
+            return SseMcpServer(name=self.config.server_name, url=url, headers=[])
+        return HttpMcpServer(name=self.config.server_name, url=url, headers=[])
+
+    def get_claude_mcp_server_config(self) -> dict[str, McpServerConfig]:
+        """Get Claude Agent SDK-compatible MCP server configuration.
+
+        Returns a dict suitable for passing to ClaudeAgentOptions.mcp_servers.
+        Uses HTTP transport to preserve MCP _meta field which contains
+        claudecode/toolUseId needed for proper tool_call_id correlation.
+
+        Note: SDK transport (passing FastMCP instance directly) would be faster
+        but the Claude Agent SDK's _handle_sdk_mcp_request drops the _meta field,
+        breaking tool_call_id propagation for progress events.
+
+        Returns:
+            Dict mapping server name to McpHttpServerConfig
+
+        Raises:
+            RuntimeError: If bridge not started (no server running)
+        """
+        if self._actual_port is None:
+            msg = "Bridge not started - call start() first"
+            raise RuntimeError(msg)
+
+        # Use HTTP transport to preserve _meta field with claudecode/toolUseId
+        # SDK transport drops _meta in Claude Agent SDK's query.py
+        url = f"http://{self.config.host}:{self.port}/mcp"
+        return {self.config.server_name: {"type": "http", "url": url}}
+
+    async def _register_tools(self) -> None:
+        """Register all node tools with the FastMCP server."""
+        if not self._mcp:
+            return
+
+        tools = await self.node.tools.get_tools(state="enabled")
+        for tool in tools:
+            self._register_single_tool(tool)
+        logger.info("Registered tools with MCP bridge", tools=[t.name for t in tools])
+
+    def _register_single_tool(self, tool: Tool) -> None:
+        """Register a single tool with the FastMCP server."""
+        if not self._mcp:
+            return
+        # Create a custom FastMCP Tool that wraps our tool
+        bridge_tool = _BridgeTool(tool=tool, bridge=self)
+        self._mcp.add_tool(bridge_tool)
+
+    async def invoke_tool_with_context(
+        self,
+        tool: Tool,
+        ctx: AgentContext[Any],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        """Invoke a tool with proper context injection.
+
+        Handles tools that expect AgentContext, RunContext, or neither.
+        """
+        fn = tool.callable
+
+        # Inject AgentContext parameters
+        context_param_names = _get_context_param_names(fn)
+        for param_name in context_param_names:
+            if param_name not in kwargs:
+                kwargs[param_name] = ctx
+
+        # Inject RunContext parameters (as stub since we're outside pydantic-ai)
+        run_context_param_names = _get_run_context_param_names(fn)
+        if run_context_param_names:
+            stub_run_ctx = _create_stub_run_context(ctx)
+            for param_name in run_context_param_names:
+                if param_name not in kwargs:
+                    kwargs[param_name] = stub_run_ctx
+
+        # Execute the tool
+        result = fn(**kwargs)
+        if inspect.isawaitable(result):
+            result = await result
+        return result
+
+    async def _start_server(self) -> None:
+        """Start the uvicorn server in the background."""
+        import socket
+
+        import uvicorn
+
+        if not self._mcp:
+            msg = "MCP server not initialized"
+            raise RuntimeError(msg)
+
+        # Determine actual port (auto-select if 0)
+        port = self.config.port
+        if port == 0:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                s.bind((self.config.host, 0))
+                port = s.getsockname()[1]
+        self._actual_port = port
+        # Create the ASGI app
+        app = self._mcp.http_app(transport=self.config.transport)
+        # Configure uvicorn
+        cfg = uvicorn.Config(app=app, host=self.config.host, port=port, log_level="warning")
+        self._server = uvicorn.Server(cfg)
+        # Start server in background task
+        name = f"mcp-bridge-{self.config.server_name}"
+        self._server_task = asyncio.create_task(self._server.serve(), name=name)
+        await anyio.sleep(0.1)  # Wait briefly for server to start
+        msg = "ToolManagerBridge started"
+        logger.info(msg, url=self.url, transport=self.config.transport)
+
+
+class _BridgeTool(FastMCPTool):
+    """Custom FastMCP Tool that wraps a agentpool Tool.
+
+    This allows us to use our own schema and invoke tools with AgentContext.
+    """
+
+    def __init__(self, tool: Tool, bridge: ToolManagerBridge) -> None:
+        # Get input schema from our tool
+        schema = tool.schema["function"]
+        input_schema = schema.get("parameters", {"type": "object", "properties": {}})
+        # Filter out context parameters - they're auto-injected by the bridge
+        context_params = _get_context_param_names(tool.callable)
+        run_context_params = _get_run_context_param_names(tool.callable)
+        all_context_params = context_params | run_context_params
+        filtered_schema = filter_schema_params(input_schema, all_context_params)
+        desc = tool.description or "No description"
+        super().__init__(name=tool.name, description=desc, parameters=filtered_schema)
+        # Set these AFTER super().__init__() to avoid being overwritten
+        self._tool = tool
+        self._bridge = bridge
+
+    async def run(self, arguments: dict[str, Any]) -> ToolResult:
+        """Execute the wrapped tool with context bridging."""
+        from fastmcp.server.dependencies import get_context
+
+        # Get FastMCP context from context variable (not passed as parameter)
+        try:
+            mcp_context: Context | None = get_context()
+        except LookupError:
+            mcp_context = None
+
+        # Try to get Claude's original tool_call_id from request metadata
+        tool_call_id = _extract_tool_call_id(mcp_context)
+        # Create context with tool-specific metadata from node's context.
+        ctx = replace(
+            self._bridge.node.get_context(),
+            tool_name=self._tool.name,
+            tool_call_id=tool_call_id,
+            tool_input=arguments,
+        )
+
+        # Invoke with context - copy arguments since invoke_tool_with_context
+        # modifies kwargs in-place to inject context parameters
+        result = await self._bridge.invoke_tool_with_context(self._tool, ctx, arguments.copy())
+        return _convert_to_tool_result(result)
+
+
+@asynccontextmanager
+async def create_tool_bridge(
+    node: BaseAgent[Any, Any],
+    *,
+    host: str = "127.0.0.1",
+    port: int = 0,
+    transport: Literal["sse", "streamable-http"] = "sse",
+) -> AsyncIterator[ToolManagerBridge]:
+    """Create and start a ToolManagerBridge as a context manager.
+
+    Args:
+        node: The node whose tools to expose
+        host: Host to bind to
+        port: Port to bind to (0 = auto-select)
+        transport: Transport protocol ('sse' or 'streamable-http')
+
+    Yields:
+        Running ToolManagerBridge instance
+    """
+    config = BridgeConfig(host=host, port=port, transport=transport)
+    bridge = ToolManagerBridge(node=node, config=config)
+    async with bridge:
+        yield bridge
+
+
+class ToolBridgeRegistry:
+    """Registry for managing multiple tool bridges.
+
+    Useful when multiple ACP agents need access to different toolsets.
+    """
+
+    def __init__(self) -> None:
+        self._bridges: dict[str, ToolManagerBridge] = {}
+        self._port_counter = 18000  # Start port range for auto-allocation
+
+    async def create_bridge(self, name: str, node: BaseAgent[Any, Any]) -> ToolManagerBridge:
+        """Create and register a new bridge.
+
+        Args:
+            name: Unique name for this bridge
+            node: The node whose tools to expose
+
+        Returns:
+            Started ToolManagerBridge
+        """
+        if name in self._bridges:
+            msg = f"Bridge {name!r} already exists"
+            raise ValueError(msg)
+
+        config = BridgeConfig(port=self._port_counter, server_name=f"agentpool-{name}")
+        self._port_counter += 1
+
+        bridge = ToolManagerBridge(node=node, config=config)
+        await bridge.start()
+        self._bridges[name] = bridge
+        return bridge
+
+    async def get_bridge(self, name: str) -> ToolManagerBridge:
+        """Get a bridge by name."""
+        if name not in self._bridges:
+            msg = f"Bridge {name!r} not found"
+            raise KeyError(msg)
+        return self._bridges[name]
+
+    async def remove_bridge(self, name: str) -> None:
+        """Stop and remove a bridge."""
+        if name in self._bridges:
+            await self._bridges[name].stop()
+            del self._bridges[name]
+
+    async def close_all(self) -> None:
+        """Stop all bridges."""
+        for bridge in list(self._bridges.values()):
+            await bridge.stop()
+        self._bridges.clear()
+
+    def get_all_mcp_configs(self) -> list[HttpMcpServer | SseMcpServer]:
+        """Get MCP server configs for all active bridges."""
+        return [bridge.get_mcp_server_config() for bridge in self._bridges.values()]

agentpool/messaging/__init__.py

@@ -0,0 +1,58 @@
+"""Core messsaging classes for AgentPool."""
+
+from agentpool.messaging.messages import ChatMessage, TokenCost, AgentResponse, TeamResponse
+from agentpool.messaging.message_container import ChatMessageList
+from agentpool.messaging.event_manager import EventManager
+from agentpool.messaging.messagenode import MessageNode
+from agentpool.messaging.message_history import MessageHistory
+from agentpool.messaging.compaction import (
+    CompactionPipeline,
+    CompactionPipelineConfig,
+    CompactionStep,
+    FilterBinaryContent,
+    FilterEmptyMessages,
+    FilterRetryPrompts,
+    FilterThinking,
+    FilterToolCalls,
+    KeepFirstAndLast,
+    KeepFirstMessages,
+    KeepLastMessages,
+    Summarize,
+    TokenBudget,
+    TruncateTextParts,
+    TruncateToolOutputs,
+    WhenMessageCountExceeds,
+    balanced_context,
+    minimal_context,
+    summarizing_context,
+)
+
+__all__ = [
+    "AgentResponse",
+    "ChatMessage",
+    "ChatMessageList",
+    "CompactionPipeline",
+    "CompactionPipelineConfig",
+    "CompactionStep",
+    "EventManager",
+    "FilterBinaryContent",
+    "FilterEmptyMessages",
+    "FilterRetryPrompts",
+    "FilterThinking",
+    "FilterToolCalls",
+    "KeepFirstAndLast",
+    "KeepFirstMessages",
+    "KeepLastMessages",
+    "MessageHistory",
+    "MessageNode",
+    "Summarize",
+    "TeamResponse",
+    "TokenBudget",
+    "TokenCost",
+    "TruncateTextParts",
+    "TruncateToolOutputs",
+    "WhenMessageCountExceeds",
+    "balanced_context",
+    "minimal_context",
+    "summarizing_context",
+]