decodingtrust-agent-sdk 0.1.0__py3-none-any.whl

utils/logging.py

@@ -0,0 +1,45 @@
+from typing import Any, Dict, List, Optional
+
+
+async def log_mcp_tools_to_file(
+    log_file: str,
+    mcp_servers: List[Any],
+) -> None:
+    """
+    Log all loaded MCP tools and their descriptions to a file for debugging.
+
+    Supports:
+    - OpenAI SDK (list_tools method)
+    - Google ADK (get_tools method)
+    - Claude SDK proxy (sdk_tools property)
+
+    Args:
+        log_file: Path to the log file
+        mcp_servers: List of MCP server instances (must have name and list_tools()/get_tools()/sdk_tools)
+    """
+    with open(log_file, "w") as f:
+        f.write("Tool List Log\n")
+        f.write(f"{'='*80}\n\n")
+        for server in mcp_servers:
+            try:
+                # Support OpenAI SDK (list_tools), Google ADK (get_tools), and Claude SDK proxy (sdk_tools)
+                if hasattr(server, 'list_tools'):
+                    tools = await server.list_tools()
+                elif hasattr(server, 'get_tools'):
+                    tools = await server.get_tools()
+                elif hasattr(server, 'sdk_tools'):
+                    # Claude SDK proxy - sdk_tools is a property, not async
+                    tools = server.sdk_tools
+                else:
+                    f.write(f"MCP Server: {server.name} - No list_tools, get_tools, or sdk_tools method\n\n")
+                    continue
+
+                f.write(f"MCP Server: {server.name}\n")
+                f.write(f"Tools loaded: {len(tools)}\n")
+                f.write(f"{'-'*40}\n")
+                for tool in tools:
+                    f.write(f"\nTool: {tool.name}\n")
+                    f.write(f"Description:\n{tool.description}\n")
+                f.write(f"\n{'='*80}\n\n")
+            except Exception as e:
+                f.write(f"MCP Server: {server.name} - Error listing tools: {e}\n\n")

utils/mcp_helpers.py

@@ -0,0 +1,232 @@
+import os
+from pathlib import Path
+import yaml
+from typing import Dict, Optional, Any
+import tempfile
+import shutil
+
+from .config import MCP_CONFIG_PATH
+from .template_helpers import (
+    allocate_server_port,
+    resolve_server_env_vars,
+    build_server_name_map,
+    wait_for_servers_ready,
+)
+from .mcp_manager import MCPServerManager
+
+
+def _start_mcp_servers_impl(
+    config_path: str,
+    server_list: list[str],
+    resource_mgr: Any,
+    task_id: str,
+    host: str,
+    prefix: str,
+    log_prefix: str,
+    task_env_overrides: Optional[Dict[str, Dict[str, str]]] = None,
+    task_dir: Optional[Any] = None,
+) -> tuple[Optional[MCPServerManager], Dict[str, str]]:
+    """Internal implementation for starting MCP servers.
+
+    Args:
+        config_path: Path to MCP config YAML file
+        server_list: List of server names to start
+        resource_mgr: ResourceManager for port allocation
+        task_id: Unique task identifier
+        host: Host address for URLs
+        prefix: Prefix for port allocation (e.g., "mcp", "injection")
+        log_prefix: Prefix for log messages (e.g., "[MCP]", "[INJECTION MCP]")
+        task_env_overrides: Optional dict of {server_name_lower: {env_key: env_val}}
+                           to merge into server env BEFORE starting servers.
+                           Used by start_task_mcp_servers() to pass task-specific
+                           env vars like USER_ACCESS_TOKEN from config.yaml.
+
+    Returns:
+        Tuple of (MCPServerManager or None, dict of {server_name: url})
+    """
+    if not server_list:
+        print(f"{log_prefix} No servers requested")
+        return None, {}
+
+    # Initialize manager and load available servers
+    manager = MCPServerManager(config_path)
+    available = manager.list_servers()
+    name_map = build_server_name_map(available)
+
+    # Load global configuration from the same config file for template resolution
+    try:
+        mcp_cfg = yaml.safe_load(open(config_path).read()) or {}
+        global_config = {srv["name"].lower(): srv for srv in (mcp_cfg.get("servers") or [])}
+    except Exception as exc:
+        print(f"[WARN] Failed to load config for template resolution: {exc}")
+        global_config = {}
+
+    # Resolve server names and validate
+    targets = []
+    for raw_name in server_list:
+        resolved = name_map.get(raw_name.lower())
+        if not resolved:
+            print(f"[WARN] Server '{raw_name}' not found in config; skipping.")
+            continue
+        if resolved not in targets:
+            targets.append(resolved)
+
+    if not targets:
+        print(f"[WARN] No {prefix} servers to start.")
+        return None, {}
+
+    # Cache for resolved port variables
+    resolved_ports: Dict[str, int] = {}
+    server_urls: Dict[str, str] = {}
+
+    # Process each server: allocate port, resolve env vars, configure
+    for server_name in targets:
+        key = server_name.lower()
+
+        # Get server config from manager
+        cfg = manager.get_server_config(server_name)
+        if not cfg:
+            continue
+
+        # Set up environment variables
+        env = cfg.setdefault("env", {})
+
+        # Find the port env var for this server's own port
+        port_key = "PORT"
+        for k in env:
+            if "PORT" in k.upper() and prefix.upper() in k.upper():
+                port_key = k
+                break
+
+        # Allocate port for this server
+        port = allocate_server_port(server_name, resource_mgr, task_id, prefix=prefix)
+        env[port_key] = str(port)
+
+        # Resolve template variables from global config
+        srv_cfg = global_config.get(key) or {}
+        srv_env = srv_cfg.get("env") or {}
+        resolved_env = resolve_server_env_vars(srv_env, resolved_ports, server_name)
+
+        # Merge resolved env vars
+        for env_key, env_val in resolved_env.items():
+            if env_key != port_key: 
+                env[env_key] = env_val
+
+        # Merge task-specific env overrides
+        if task_env_overrides and key in task_env_overrides:
+            for env_key, env_val in task_env_overrides[key].items():
+                if env_val:  # Only set non-empty values
+                    env[env_key] = env_val
+                    display_val = env_val[:20] + "..." if len(env_val) > 20 else env_val
+                    print(f"{log_prefix} {server_name}: {env_key}={display_val}")
+
+        # Determine URL - only /mcp is supported
+        transport = cfg.get("transport", "http")
+        if transport != "http":
+            raise ValueError(f"Unsupported transport '{transport}' for server '{server_name}'. Only 'http' is supported.")
+
+        url = f"http://{host}:{port}/mcp"
+        server_urls[server_name] = url
+
+    # MCP config file
+    if task_dir is not None:
+        for server_name in targets:
+            mcp_config_file = task_dir / "metadata" / f"{server_name}_mcp_config.json"
+            if not mcp_config_file.exists():
+                continue
+            temp_mcp_config_file = Path(tempfile.mkstemp(prefix=f"dt_mcp_{server_name}_config_", suffix=".json")[1])
+            shutil.copy(mcp_config_file, temp_mcp_config_file)
+            with open(temp_mcp_config_file, "r") as f:
+                content = f.read()
+            for var_name, port in resolved_ports.items():
+                content = content.replace(f"<DT_ENV>{var_name}</DT_ENV>", str(port))
+            with open(temp_mcp_config_file, "w") as f:
+                f.write(content)
+            config_env_key = f"{server_name.upper()}_MCP_CONFIG"
+            manager.get_server_config(server_name)["env"][config_env_key] = str(temp_mcp_config_file)
+            # Also expose to the current process so judge.py (loaded in-process) can read it
+            os.environ[config_env_key] = str(temp_mcp_config_file)
+            print(f"[MCP] {server_name}: {config_env_key}={temp_mcp_config_file}")
+
+    # Start all servers
+    print(f"{log_prefix} Starting {len(targets)} server(s)...")
+    started = 0
+    for name in targets:
+        if manager.start_server(name, foreground=False):
+            started += 1
+            print(f"  Started: {name} -> {server_urls.get(name, 'unknown')}")
+
+    if started == 0:
+        print(f"[WARN] No {prefix} servers were started.")
+        return None, server_urls
+
+    print(f"{log_prefix} Started {started}/{len(targets)} server(s)")
+
+    # Wait for servers to be ready
+    if server_urls:
+        print(f"{log_prefix} Waiting for servers to be ready...")
+        wait_for_servers_ready(server_urls)
+
+    return manager, server_urls
+
+
+def start_task_mcp_servers(
+    agent_cfg: Any,  # AgentConfig
+    task_id: str,
+    task_dir: Any,  # Path
+    resource_mgr: Any,  # ResourceManager
+    host: str = "127.0.0.1",
+) -> Optional[MCPServerManager]:
+    """Set up and start MCP servers for a task.
+
+    Args:
+        agent_cfg: Agent configuration with MCP server list
+        task_id: Unique task identifier for resource tracking
+        resource_mgr: Resource manager for port allocation
+        host: Host address for MCP server URLs (default: 127.0.0.1)
+
+    Returns:
+        MCPServerManager instance or None if no servers started
+    """
+    # Extract server list from agent config
+    requested_servers = [s.name for s in (agent_cfg.mcp_servers or []) if s.enabled]
+
+    if not requested_servers:
+        print("[MCP] No MCP servers requested")
+        return None
+
+    if not MCP_CONFIG_PATH.exists():
+        print(f"[WARN] MCP config not found: {MCP_CONFIG_PATH}")
+        return None
+
+    # Build task-specific env overrides from agent config
+    task_env_overrides: Dict[str, Dict[str, str]] = {}
+    for server in agent_cfg.mcp_servers:
+        if server.env:
+            task_env_overrides[server.name.lower()] = {
+                k: str(v) for k, v in server.env.items() if v
+            }
+
+    # Start MCP servers with task-specific env overrides
+    manager, server_urls = _start_mcp_servers_impl(
+        config_path=str(MCP_CONFIG_PATH),
+        server_list=requested_servers,
+        resource_mgr=resource_mgr,
+        task_id=task_id,
+        host=host,
+        prefix="mcp",
+        log_prefix="[MCP]",
+        task_env_overrides=task_env_overrides,
+        task_dir=task_dir,
+    )
+
+    if not manager:
+        return None
+
+    # Update agent config with URLs
+    for server in agent_cfg.mcp_servers:
+        if server.name in server_urls:
+            server.url = server_urls[server.name]
+            print(f"[MCP] Server '{server.name}' URL -> {server.url}")
+
+    return manager

utils/mcp_manager.py

@@ -0,0 +1,235 @@
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+import yaml
+
+
+class MCPServerManager:
+    """Manager for MCP server lifecycle"""
+
+    def __init__(self, registry_path: str = "mcp_registry.yaml"):
+        self.registry_path = Path(registry_path)
+        self.base_dir = self.registry_path.parent
+        self.config = self._load_config()
+        self.processes: Dict[str, subprocess.Popen] = {}
+
+    def _load_config(self) -> Dict[str, Any]:
+        """Load and validate the registry configuration"""
+        if not self.registry_path.exists():
+            raise FileNotFoundError(f"Registry not found: {self.registry_path}")
+
+        with open(self.registry_path, "r") as f:
+            config = yaml.safe_load(f)
+
+        if not config or "servers" not in config:
+            raise ValueError("Invalid registry: missing 'servers' section")
+
+        return config
+
+    def list_servers(self, enabled_only: bool = False) -> List[Dict[str, Any]]:
+        """List all servers in the registry"""
+        servers = self.config.get("servers", [])
+        if enabled_only:
+            servers = [s for s in servers if s.get("enabled", False)]
+        return servers
+
+    def get_server_config(self, name: str) -> Optional[Dict[str, Any]]:
+        """Get configuration for a specific server"""
+        for server in self.config.get("servers", []):
+            if server.get("name") == name:
+                return server
+        return None
+
+    def _resolve_path(self, relative_path: str) -> Path:
+        """Resolve server path relative to registry file"""
+        base = self.config.get("global", {}).get("base_dir", ".")
+        return (self.base_dir / base / relative_path).resolve()
+
+    def _setup_environment(self, server: Dict[str, Any]) -> Dict[str, str]:
+        """Setup environment variables for a server.
+
+        Inherits parent process env vars and overrides with server-specific values.
+        """
+        env = os.environ.copy()
+        env.update({k: str(v) for k, v in server.get("env", {}).items() if v is not None})
+        return env
+
+    def _create_log_dir(self) -> Path:
+        """Create log directory if it doesn't exist"""
+        log_dir = self.base_dir / self.config.get("global", {}).get("log_dir", "logs")
+        log_dir.mkdir(parents=True, exist_ok=True)
+        return log_dir
+
+    def start_server(
+        self,
+        name: str,
+        dry_run: bool = False,
+        foreground: bool = True
+    ) -> bool:
+        """Start a single MCP server (runs in foreground by default)"""
+        server = self.get_server_config(name)
+        if not server:
+            print(f"[ERROR] Server '{name}' not found in registry", file=sys.stderr)
+            return False
+
+        if not server.get("enabled", False):
+            print(f"[WARN] Server '{name}' is disabled in registry", file=sys.stderr)
+            return False
+
+        # Resolve server path
+        main_path = self._resolve_path(server["path"])
+        if not main_path.exists():
+            print(f"[ERROR] Server script not found: {main_path}", file=sys.stderr)
+            return False
+
+        # Setup environment
+        env = self._setup_environment(server)
+
+        # Build command
+        if "command" in server:
+            # Use custom command from registry (supports complex setups like supergateway)
+            cmd = []
+            for part in server["command"]:
+                # Expand environment variables in command (e.g., $PORT)
+                expanded = part
+                for env_key, env_val in env.items():
+                    expanded = expanded.replace(f"${{{env_key}}}", str(env_val))
+                    expanded = expanded.replace(f"${env_key}", str(env_val))
+                cmd.append(expanded)
+        else:
+            # Default: run with python executable
+            python_exe = self.config.get("global", {}).get("python_executable", "python3")
+            cmd = [python_exe, str(main_path)]
+
+            # Add any additional options
+            options = server.get("options", {})
+            if options:
+                for key, value in options.items():
+                    cmd.extend([f"--{key.replace('_', '-')}", str(value)])
+
+        if dry_run:
+            print(f"\n[DRY RUN] Would start '{name}':")
+            print(f"   Command: {' '.join(cmd)}")
+            print(f"   Working dir: {main_path.parent}")
+            print(f"   Environment variables:")
+            for key, value in server.get("env", {}).items():
+                display_value = value if len(str(value)) < 50 else f"{str(value)[:47]}..."
+                print(f"     {key}={display_value}")
+            return True
+
+        # Create log files
+        log_dir = self._create_log_dir()
+        stdout_log = log_dir / f"{name}_stdout.log"
+        stderr_log = log_dir / f"{name}_stderr.log"
+
+        print(f"[MCP] Starting '{name}'...")
+        print(f"   Command: {' '.join(cmd)}")
+        print(f"   Logs: {stdout_log} / {stderr_log}")
+
+        if foreground:
+            # Run in foreground (blocking)
+            try:
+                subprocess.run(
+                    cmd,
+                    cwd=main_path.parent,
+                    env=env,
+                    check=True
+                )
+                return True
+            except subprocess.CalledProcessError as e:
+                print(f"[ERROR] Server '{name}' exited with code {e.returncode}", file=sys.stderr)
+                return False
+            except KeyboardInterrupt:
+                print(f"\n[WARN] Server '{name}' interrupted", file=sys.stderr)
+                return False
+        else:
+            # Run in background
+            try:
+                with open(stdout_log, "w") as stdout, open(stderr_log, "w") as stderr:
+                    process = subprocess.Popen(
+                        cmd,
+                        cwd=main_path.parent,
+                        env=env,
+                        stdout=stdout,
+                        stderr=stderr,
+                        start_new_session=True
+                    )
+
+                self.processes[name] = process
+
+                # Give it a moment to start
+                time.sleep(0.5)
+
+                # Check if it's still running
+                if process.poll() is None:
+                    print(f"[MCP] Server '{name}' started (PID: {process.pid})")
+                    return True
+                else:
+                    print(f"[ERROR] Server '{name}' failed to start", file=sys.stderr)
+                    return False
+
+            except Exception as e:
+                print(f"[ERROR] Failed to start '{name}': {e}", file=sys.stderr)
+                return False
+
+    def start_all(self, dry_run: bool = False, foreground: bool = True) -> int:
+        """Start all enabled servers (runs in foreground by default)"""
+        enabled_servers = self.list_servers(enabled_only=True)
+
+        if not enabled_servers:
+            print("[WARN] No enabled servers found in registry", file=sys.stderr)
+            return 0
+
+        print(f"\n[MCP] Starting {len(enabled_servers)} server(s)...\n")
+
+        success_count = 0
+        for server in enabled_servers:
+            name = server["name"]
+            if self.start_server(name, dry_run=dry_run, foreground=foreground):
+                success_count += 1
+            print()
+
+        if not dry_run:
+            if foreground:
+                print(f"[MCP] Started {success_count}/{len(enabled_servers)} server(s) in foreground")
+                print("   Press Ctrl+C to stop")
+            else:
+                print(f"[MCP] Started {success_count}/{len(enabled_servers)} server(s) in background")
+                if self.processes:
+                    print("\n[MCP] Servers are running in the background.")
+                    print("   Check logs in the 'logs/' directory for output.")
+                    print("   Use Ctrl+C or kill the processes to stop them.")
+
+        return success_count
+
+    def stop_all(self):
+        """Stop all running servers"""
+        if not self.processes:
+            print("[WARN] No servers running", file=sys.stderr)
+            return
+
+        print(f"\n[MCP] Stopping {len(self.processes)} server(s)...\n")
+
+        for name, process in self.processes.items():
+            try:
+                process.terminate()
+                print(f"[MCP] Stopped '{name}' (PID: {process.pid})")
+            except Exception as e:
+                print(f"[ERROR] Failed to stop '{name}': {e}", file=sys.stderr)
+
+        # Wait for processes to terminate
+        time.sleep(1)
+
+        # Force kill any remaining
+        for name, process in self.processes.items():
+            if process.poll() is None:
+                try:
+                    process.kill()
+                    print(f"[WARN] Force killed '{name}'", file=sys.stderr)
+                except Exception:
+                    pass
+
+        self.processes.clear()

utils/memory_guard.py

@@ -0,0 +1,18 @@
+"""Memory guard to prevent OOM from too many Docker containers."""
+
+import psutil
+
+
+def check_memory_before_launch(required_gb: float = 4, reserve_gb: float = 64) -> None:
+    """Raise RuntimeError if insufficient memory to launch a new container.
+
+    Args:
+        required_gb: Estimated memory needed by the new container.
+        reserve_gb: Memory to keep free for the OS and other processes.
+    """
+    available_gb = psutil.virtual_memory().available / (1024 ** 3)
+    if available_gb < required_gb + reserve_gb:
+        raise RuntimeError(
+            f"Available memory {available_gb:.1f}GB < needed {required_gb}GB + reserve {reserve_gb}GB. "
+            f"Refusing to start new container."
+        )