rossum-agent 1.0.0rc0__py3-none-any.whl

rossum_agent/tools/spawn_mcp.py

@@ -0,0 +1,227 @@
+"""Tools for spawning, calling, and closing MCP connections to different Rossum environments.
+
+This module provides tools to manage secondary MCP connections to different Rossum
+environments at runtime, enabling cross-environment operations like deployments.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import threading
+from concurrent.futures import TimeoutError as FuturesTimeoutError
+from dataclasses import dataclass
+
+from anthropic import beta_tool
+from fastmcp import Client
+
+from rossum_agent.rossum_mcp_integration import MCPConnection, create_mcp_transport
+from rossum_agent.tools.core import get_mcp_event_loop
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SpawnedConnection:
+    """Record for a spawned MCP connection to a different environment."""
+
+    connection: MCPConnection
+    client: Client
+    api_base_url: str
+
+
+# Secondary MCP connections spawned at runtime for different environments
+_spawned_connections: dict[str, SpawnedConnection] = {}
+_spawned_connections_lock = threading.Lock()
+
+
+def get_spawned_connections() -> dict[str, SpawnedConnection]:
+    """Get the spawned connections dict (for internal use only)."""
+    return _spawned_connections
+
+
+def get_spawned_connections_lock() -> threading.Lock:
+    """Get the spawned connections lock (for internal use only)."""
+    return _spawned_connections_lock
+
+
+def clear_spawned_connections() -> None:
+    """Clear all spawned connections. Called when MCP connection is reset."""
+    with _spawned_connections_lock:
+        _spawned_connections.clear()
+
+
+async def _spawn_connection_async(
+    connection_id: str, api_token: str, api_base_url: str, mcp_mode: str = "read-write"
+) -> SpawnedConnection:
+    """Spawn a new MCP connection asynchronously.
+
+    Raises:
+        ValueError: If connection_id already exists.
+    """
+    with _spawned_connections_lock:
+        if connection_id in _spawned_connections:
+            raise ValueError(f"Connection '{connection_id}' already exists")
+
+    transport = create_mcp_transport(api_token, api_base_url, mcp_mode)  # type: ignore[arg-type]
+    client = Client(transport)
+
+    await client.__aenter__()
+    connection = MCPConnection(client=client)
+
+    record = SpawnedConnection(connection=connection, client=client, api_base_url=api_base_url)
+
+    with _spawned_connections_lock:
+        _spawned_connections[connection_id] = record
+
+    return record
+
+
+async def _close_spawned_connection_async(connection_id: str) -> None:
+    """Close a spawned MCP connection."""
+    with _spawned_connections_lock:
+        record = _spawned_connections.pop(connection_id, None)
+
+    if record is not None:
+        await record.client.__aexit__(None, None, None)
+
+
+def cleanup_all_spawned_connections() -> None:
+    """Cleanup all spawned connections. Call this when the agent session ends.
+
+    Should only be called once at session teardown, after no more tool calls are expected.
+    """
+    if (mcp_event_loop := get_mcp_event_loop()) is None:
+        return
+
+    with _spawned_connections_lock:
+        conn_ids = list(_spawned_connections.keys())
+
+    for conn_id in conn_ids:
+        try:
+            future = asyncio.run_coroutine_threadsafe(_close_spawned_connection_async(conn_id), mcp_event_loop)
+            future.result(timeout=10)
+        except FuturesTimeoutError:
+            future.cancel()
+            logger.warning(f"Timeout cleaning up connection {conn_id}")
+        except Exception as e:
+            logger.warning(f"Failed to cleanup connection {conn_id}: {e}")
+
+
+@beta_tool
+def spawn_mcp_connection(connection_id: str, api_token: str, api_base_url: str, mcp_mode: str = "read-write") -> str:
+    """Spawn a new MCP connection to a different Rossum environment.
+
+    Use this when you need to make changes to a different Rossum environment than the one the agent was initialized with.
+    For example, when deploying changes from a source environment to a target environment.
+
+    Args:
+        connection_id: A unique identifier for this connection (e.g., 'target', 'sandbox')
+
+    Returns:
+        Success message with available tools, or error message if failed.
+    """
+    if (mcp_event_loop := get_mcp_event_loop()) is None:
+        return "Error: MCP event loop not set. Agent not properly initialized."
+
+    if not connection_id or not connection_id.strip():
+        return "Error: connection_id must be non-empty."
+
+    if not api_base_url or not api_base_url.startswith("https://"):
+        return "Error: api_base_url must start with https://"
+
+    try:
+        future = asyncio.run_coroutine_threadsafe(
+            _spawn_connection_async(connection_id, api_token, api_base_url, mcp_mode),
+            mcp_event_loop,
+        )
+        record = future.result(timeout=30)
+
+        tools_future = asyncio.run_coroutine_threadsafe(record.connection.get_tools(), mcp_event_loop)
+        tools = tools_future.result(timeout=30)
+        tool_names = [t.name for t in tools]
+
+        return f"Successfully spawned MCP connection '{connection_id}' to {api_base_url}. Available tools: {', '.join(tool_names[:10])}{'...' if len(tool_names) > 10 else ''}"
+    except ValueError as e:
+        return f"Error: {e}"
+    except FuturesTimeoutError:
+        future.cancel()
+        return "Error: Timed out while spawning MCP connection."
+    except RuntimeError as e:
+        logger.exception("Error scheduling MCP call")
+        return f"Error: Failed to schedule MCP call: {e}"
+    except Exception as e:
+        logger.error(f"Failed to spawn connection: {e}")
+        return f"Error spawning connection: {e}"
+
+
+@beta_tool
+def call_on_connection(connection_id: str, tool_name: str, arguments: str | dict) -> str:
+    """Call a tool on a spawned MCP connection.
+
+    Use this to execute MCP tools on a connection that was previously spawned with spawn_mcp_connection.
+
+    Args:
+        connection_id: The identifier of the spawned connection.
+        tool_name: The name of the MCP tool to call.
+        arguments: Arguments to pass to the tool (JSON string or dict).
+
+    Returns:
+        The result of the tool call as a JSON string, or error message.
+    """
+    if (mcp_event_loop := get_mcp_event_loop()) is None:
+        return "Error: MCP event loop not set."
+
+    with _spawned_connections_lock:
+        if connection_id not in _spawned_connections:
+            available = list(_spawned_connections.keys())
+            return f"Error: Connection '{connection_id}' not found. Available: {available}"
+        record = _spawned_connections[connection_id]
+
+    logger.debug(f"call_on_connection: Using connection '{connection_id}' - API URL: {record.api_base_url}")
+
+    if isinstance(arguments, dict):
+        args = arguments
+    elif arguments:
+        try:
+            args = json.loads(arguments)
+        except json.JSONDecodeError as e:
+            return f"Error parsing arguments JSON: {e}"
+    else:
+        args = {}
+
+    try:
+        future = asyncio.run_coroutine_threadsafe(record.connection.call_tool(tool_name, args), mcp_event_loop)
+        result = future.result(timeout=60)
+
+        if isinstance(result, (dict, list)):
+            return f"[{tool_name}] {json.dumps(result, indent=2, default=str)}"
+        return f"[{tool_name}] {result}" if result is not None else f"[{tool_name}] Tool executed successfully"
+    except FuturesTimeoutError:
+        future.cancel()
+        return f"Error: Timed out calling {tool_name} after 60 seconds."
+    except Exception as e:
+        logger.error(f"Error calling tool on connection: {e}")
+        return f"Error calling {tool_name}: {e}"
+
+
+@beta_tool
+def close_connection(connection_id: str) -> str:
+    """Close a spawned MCP connection."""
+    if (mcp_event_loop := get_mcp_event_loop()) is None:
+        return "Error: MCP event loop not set."
+
+    with _spawned_connections_lock:
+        if connection_id not in _spawned_connections:
+            return f"Connection '{connection_id}' not found."
+
+    try:
+        future = asyncio.run_coroutine_threadsafe(_close_spawned_connection_async(connection_id), mcp_event_loop)
+        future.result(timeout=10)
+        return f"Successfully closed connection '{connection_id}'."
+    except FuturesTimeoutError:
+        future.cancel()
+        return f"Error: Timed out closing connection '{connection_id}'."
+    except Exception as e:
+        return f"Error closing connection: {e}"

rossum_agent/tools/subagents/__init__.py

@@ -0,0 +1,31 @@
+"""Sub-agents for the Rossum Agent.
+
+Opus-powered sub-agents for complex iterative tasks:
+- Hook debugging with sandboxed execution
+- Knowledge base search with AI analysis
+- Schema patching with programmatic bulk updates
+"""
+
+from __future__ import annotations
+
+from rossum_agent.bedrock_client import OPUS_MODEL_ID
+from rossum_agent.tools.subagents.base import SubAgent, SubAgentConfig, SubAgentResult
+from rossum_agent.tools.subagents.hook_debug import HookDebugSubAgent, debug_hook, evaluate_python_hook
+from rossum_agent.tools.subagents.knowledge_base import WebSearchError, search_knowledge_base
+from rossum_agent.tools.subagents.mcp_helpers import call_mcp_tool
+from rossum_agent.tools.subagents.schema_patching import SchemaPatchingSubAgent, patch_schema_with_subagent
+
+__all__ = [
+    "OPUS_MODEL_ID",
+    "HookDebugSubAgent",
+    "SchemaPatchingSubAgent",
+    "SubAgent",
+    "SubAgentConfig",
+    "SubAgentResult",
+    "WebSearchError",
+    "call_mcp_tool",
+    "debug_hook",
+    "evaluate_python_hook",
+    "patch_schema_with_subagent",
+    "search_knowledge_base",
+]

rossum_agent/tools/subagents/base.py

@@ -0,0 +1,303 @@
+"""Shared base module for sub-agents.
+
+Provides common infrastructure for sub-agents that use iterative LLM calls with tool use:
+- Unified iteration loop with token tracking
+- Context saving for debugging
+- Consistent logging patterns
+- Progress and token usage reporting
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+
+from rossum_agent.bedrock_client import create_bedrock_client, get_model_id
+from rossum_agent.tools.core import (
+    SubAgentProgress,
+    SubAgentTokenUsage,
+    get_output_dir,
+    report_progress,
+    report_token_usage,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SubAgentConfig:
+    """Configuration for a sub-agent iteration loop."""
+
+    tool_name: str
+    system_prompt: str
+    tools: list[dict[str, Any]]
+    max_iterations: int = 15
+    max_tokens: int = 16384
+
+
+@dataclass
+class SubAgentResult:
+    """Result from a sub-agent execution."""
+
+    analysis: str
+    input_tokens: int
+    output_tokens: int
+    iterations_used: int
+
+
+def save_iteration_context(
+    tool_name: str,
+    iteration: int,
+    max_iterations: int,
+    messages: list[dict[str, Any]],
+    system_prompt: str,
+    tools: list[dict[str, Any]],
+    max_tokens: int,
+) -> None:
+    """Save agent input context to file for debugging.
+
+    Args:
+        tool_name: Name of the sub-agent tool (e.g., "debug_hook", "patch_schema").
+        iteration: Current iteration number (1-indexed).
+        max_iterations: Maximum number of iterations.
+        messages: Current conversation messages.
+        system_prompt: System prompt used.
+        tools: Tool definitions.
+        max_tokens: Max tokens setting.
+    """
+    try:
+        output_dir = get_output_dir()
+        context_file = output_dir / f"{tool_name}_context_iter_{iteration}.json"
+        context_data = {
+            "iteration": iteration,
+            "max_iterations": max_iterations,
+            "model": get_model_id(),
+            "max_tokens": max_tokens,
+            "system_prompt": system_prompt,
+            "messages": messages,
+            "tools": tools,
+        }
+        context_file.write_text(json.dumps(context_data, indent=2, default=str))
+        logger.info(f"{tool_name} sub-agent: saved context to {context_file}")
+    except Exception as e:
+        logger.warning(f"Failed to save {tool_name} context: {e}")
+
+
+class SubAgent(ABC):
+    """Base class for sub-agents with iterative tool use.
+
+    Provides a unified iteration loop with:
+    - Token tracking and reporting
+    - Progress reporting
+    - Context saving for debugging
+    - Consistent logging
+    """
+
+    def __init__(self, config: SubAgentConfig) -> None:
+        """Initialize the sub-agent.
+
+        Args:
+            config: Configuration for the sub-agent.
+        """
+        self.config = config
+        self._client = None
+
+    @property
+    def client(self):
+        """Lazily create the Bedrock client."""
+        if self._client is None:
+            client_start = time.perf_counter()
+            self._client = create_bedrock_client()
+            elapsed_ms = (time.perf_counter() - client_start) * 1000
+            logger.info(f"{self.config.tool_name}: Bedrock client created in {elapsed_ms:.1f}ms")
+        return self._client
+
+    @abstractmethod
+    def execute_tool(self, tool_name: str, tool_input: dict[str, Any]) -> str:
+        """Execute a tool call from the LLM.
+
+        Args:
+            tool_name: Name of the tool to execute.
+            tool_input: Input arguments for the tool.
+
+        Returns:
+            Tool result as a string.
+        """
+
+    @abstractmethod
+    def process_response_block(self, block: Any, iteration: int, max_iterations: int) -> dict[str, Any] | None:
+        """Process a response block for special handling (e.g., web search).
+
+        Args:
+            block: Response content block.
+            iteration: Current iteration number (1-indexed).
+            max_iterations: Maximum iterations.
+
+        Returns:
+            Tool result dict if the block was processed, None otherwise.
+        """
+
+    def run(self, initial_message: str) -> SubAgentResult:
+        """Run the sub-agent iteration loop."""
+        messages: list[dict[str, Any]] = [{"role": "user", "content": initial_message}]
+        total_input_tokens = 0
+        total_output_tokens = 0
+        current_iteration = 0
+
+        response = None
+        try:
+            for iteration in range(self.config.max_iterations):
+                current_iteration = iteration + 1
+                iter_start = time.perf_counter()
+
+                logger.info(
+                    f"{self.config.tool_name} sub-agent: iteration {current_iteration}/{self.config.max_iterations}"
+                )
+
+                report_progress(
+                    SubAgentProgress(
+                        tool_name=self.config.tool_name,
+                        iteration=current_iteration,
+                        max_iterations=self.config.max_iterations,
+                        status="thinking",
+                    )
+                )
+
+                save_iteration_context(
+                    tool_name=self.config.tool_name,
+                    iteration=current_iteration,
+                    max_iterations=self.config.max_iterations,
+                    messages=messages,
+                    system_prompt=self.config.system_prompt,
+                    tools=self.config.tools,
+                    max_tokens=self.config.max_tokens,
+                )
+
+                llm_start = time.perf_counter()
+                response = self.client.messages.create(
+                    model=get_model_id(),
+                    max_tokens=self.config.max_tokens,
+                    system=self.config.system_prompt,
+                    messages=messages,
+                    tools=self.config.tools,
+                )
+                llm_elapsed_ms = (time.perf_counter() - llm_start) * 1000
+
+                input_tokens = response.usage.input_tokens
+                output_tokens = response.usage.output_tokens
+                total_input_tokens += input_tokens
+                total_output_tokens += output_tokens
+
+                logger.info(
+                    f"{self.config.tool_name} [iter {current_iteration}]: "
+                    f"LLM {llm_elapsed_ms:.1f}ms, tokens in={input_tokens} out={output_tokens}"
+                )
+
+                report_token_usage(
+                    SubAgentTokenUsage(
+                        tool_name=self.config.tool_name,
+                        input_tokens=input_tokens,
+                        output_tokens=output_tokens,
+                        iteration=current_iteration,
+                    )
+                )
+
+                has_tool_use = any(hasattr(block, "type") and block.type == "tool_use" for block in response.content)
+
+                if response.stop_reason == "end_of_turn" or not has_tool_use:
+                    iter_elapsed_ms = (time.perf_counter() - iter_start) * 1000
+                    logger.info(
+                        f"{self.config.tool_name}: completed after {current_iteration} iterations "
+                        f"in {iter_elapsed_ms:.1f}ms (stop_reason={response.stop_reason}, has_tool_use={has_tool_use})"
+                    )
+                    report_progress(
+                        SubAgentProgress(
+                            tool_name=self.config.tool_name,
+                            iteration=current_iteration,
+                            max_iterations=self.config.max_iterations,
+                            status="completed",
+                        )
+                    )
+                    text_parts = [block.text for block in response.content if hasattr(block, "text")]
+                    return SubAgentResult(
+                        analysis="\n".join(text_parts) if text_parts else "No analysis provided",
+                        input_tokens=total_input_tokens,
+                        output_tokens=total_output_tokens,
+                        iterations_used=current_iteration,
+                    )
+
+                messages.append({"role": "assistant", "content": response.content})
+
+                tool_results: list[dict[str, Any]] = []
+                iteration_tool_calls: list[str] = []
+
+                for block in response.content:
+                    special_result = self.process_response_block(block, current_iteration, self.config.max_iterations)
+                    if special_result:
+                        tool_results.append(special_result)
+
+                for block in response.content:
+                    if hasattr(block, "type") and block.type == "tool_use":
+                        tool_name = block.name
+                        tool_input = block.input
+                        iteration_tool_calls.append(tool_name)
+
+                        logger.info(f"{self.config.tool_name} [iter {current_iteration}]: calling tool '{tool_name}'")
+
+                        report_progress(
+                            SubAgentProgress(
+                                tool_name=self.config.tool_name,
+                                iteration=current_iteration,
+                                max_iterations=self.config.max_iterations,
+                                current_tool=tool_name,
+                                tool_calls=iteration_tool_calls.copy(),
+                                status="running_tool",
+                            )
+                        )
+
+                        try:
+                            tool_start = time.perf_counter()
+                            result = self.execute_tool(tool_name, tool_input)
+                            tool_elapsed_ms = (time.perf_counter() - tool_start) * 1000
+                            logger.info(
+                                f"{self.config.tool_name}: tool '{tool_name}' executed in {tool_elapsed_ms:.1f}ms"
+                            )
+                            tool_results.append({"type": "tool_result", "tool_use_id": block.id, "content": result})
+                        except Exception as e:
+                            logger.warning(f"Tool {tool_name} failed: {e}")
+                            tool_results.append(
+                                {
+                                    "type": "tool_result",
+                                    "tool_use_id": block.id,
+                                    "content": f"Error: {e}",
+                                    "is_error": True,
+                                }
+                            )
+
+                if tool_results:
+                    messages.append({"role": "user", "content": tool_results})
+
+            logger.warning(f"{self.config.tool_name}: max iterations ({self.config.max_iterations}) reached")
+            text_parts = [block.text for block in response.content if hasattr(block, "text")] if response else []
+            return SubAgentResult(
+                analysis="\n".join(text_parts) if text_parts else "Max iterations reached",
+                input_tokens=total_input_tokens,
+                output_tokens=total_output_tokens,
+                iterations_used=self.config.max_iterations,
+            )
+
+        except Exception as e:
+            logger.exception(f"Error in {self.config.tool_name} sub-agent")
+            return SubAgentResult(
+                analysis=f"Error calling Opus sub-agent: {e}",
+                input_tokens=total_input_tokens,
+                output_tokens=total_output_tokens,
+                iterations_used=current_iteration,
+            )