@ngocsangairvds/vsaf 3.2.16 → 4.0.0

package/tools/vds-scripts/memory_orchestrator/src/vds_memory_orchestrator/service/mcp/official_mcp_client.py

@@ -1,1574 +0,0 @@
-"""
-Official MCP Client Implementation using the Official MCP Python SDK
-
-This module provides a proper MCP client implementation that follows the MCP
-specification using the official Python SDK instead of custom HTTP workarounds.
-"""
-
-import json
-import logging
-import os
-import time
-from collections.abc import Callable
-from contextlib import asynccontextmanager, suppress
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any
-
-import yaml
-
-# Initialize logger for import error handling
-
-logger = logging.getLogger(__name__)
-
-try:
-    from mcp import ClientSession
-    from mcp.client.streamable_http import streamablehttp_client
-    from mcp.types import CallToolResult, ReadResourceResult, Tool
-except ImportError as e:
-    logger.exception(f"MCP SDK not properly installed: {e}")
-    raise
-
-from ..core.ai_memory_exceptions import (
-    ConfigurationError,
-    ConnectionError,
-    ToolExecutionError,
-)
-from ..mcp.env_config_parser import load_mcp_servers_from_env
-from ..models.ai_memory_models import SessionConsolidationRequest
-
-
-@dataclass
-class MCPServerCapability:
-    """Capability metadata describing an MCP tool or resource."""
-
-    name: str
-    description: str | None = None
-    category: str | None = None
-    input_schema: dict[str, Any] = field(default_factory=dict)
-    output_schema: dict[str, Any] = field(default_factory=dict)
-
-
-@dataclass
-class MCPServerMetadata:
-    """Metadata describing a configured MCP server."""
-
-    name: str
-    version: str | None = None
-    description: str | None = None
-    author: str | None = None
-    homepage: str | None = None
-    tags: list[str] = field(default_factory=list)
-    categories: list[str] = field(default_factory=list)
-    capabilities: list[MCPServerCapability] = field(default_factory=list)
-    supported_domains: list[str] = field(default_factory=list)
-    raw: dict[str, Any] = field(default_factory=dict)
-
-
-@dataclass
-class MCPConnectionConfig:
-    """Configuration for MCP server connection."""
-
-    name: str
-    transport: str
-    url: str | None = None
-    api_key: str | None = None
-    timeout: float = 30.0
-    max_connections: int = 5
-    command: str | None = None
-    args: list[str] = field(default_factory=list)
-    env: dict[str, str] = field(default_factory=dict)
-    headers: dict[str, str] = field(default_factory=dict)
-    authentication: dict[str, Any] = field(default_factory=dict)
-    config_source: str = "env"
-    extra: dict[str, Any] = field(default_factory=dict)
-
-
-@dataclass
-class MCPConnection:
-    """Maintains MCP connection state including streams and session."""
-
-    client_session: ClientSession
-    read_stream: Any
-    write_stream: Any
-    session_id_callback: Callable[..., Any] | None = None
-    config: MCPConnectionConfig | None = None
-    metadata: MCPServerMetadata | None = None
-    # Store the stdio context manager to keep it alive
-    _stdio_context: Any | None = None
-    # Store the HTTP connection client to keep it alive
-    _connection_client: Any | None = None
-
-
-class OfficialMCPClient:
-    """
-    Proper MCP Client implementation using the official MCP Python SDK.
-
-    This client follows MCP specification correctly by using:
-    - ClientSession for session management
-    - Proper transport initialization (stdio, streamablehttp)
-    - Official MCP protocol communication
-    - Correct error handling and timeouts
-    """
-
-    def __init__(self):
-        """Initialize the official MCP client."""
-        self.connections: dict[str, MCPConnection] = {}
-        self.connection_configs: dict[str, MCPConnectionConfig] = {}
-        self.server_metadata: dict[str, MCPServerMetadata] = {}
-        self._capability_index: dict[str, list[str]] = {}
-        self._tool_cache: dict[str, dict[str, Any]] = {}
-        self._tool_cache_ttl: float = 300.0  # seconds
-        self._config_file: Path | None = None
-        self.initialized = False
-        self._stats = {
-            "connection_attempts": 0,
-            "successful_connections": 0,
-            "failed_connections": 0,
-            "tool_calls": 0,
-            "successful_tool_calls": 0,
-        }
-        self._server_stats: dict[str, dict[str, Any]] = {}
-
-    async def initialize(self, config_file: str | None = None, force_reload: bool = False):
-        """Initialize MCP client with configuration from env and optional YAML file."""
-
-        if self.initialized:
-            requested_path = Path(config_file) if config_file else self._config_file
-            if not force_reload and requested_path == self._config_file:
-                return
-            await self.close()
-
-        if config_file:
-            self._config_file = Path(config_file)
-
-        try:
-            logger.info(
-                "Initializing Official MCP Client%s",
-                f" using {self._config_file}" if self._config_file else "",
-            )
-            self._load_configuration_sources()
-
-            if not self.connection_configs:
-                logger.warning(
-                    "No MCP servers configured — proceeding with official client in idle state"
-                )
-
-            await self._initialize_connections()
-            self.initialized = True
-            logger.info(
-                "Official MCP Client initialized with %d configured servers (%d active connections)",
-                len(self.connection_configs),
-                len(self.connections),
-            )
-
-        except Exception as exc:
-            logger.exception("Failed to initialize Official MCP Client: %s", exc)
-            raise
-
-    def _load_configuration_sources(self):
-        """Load server configurations from environment variables and optional YAML file."""
-
-        self.connection_configs.clear()
-        self.server_metadata.clear()
-        self._capability_index.clear()
-        self._tool_cache.clear()
-
-        for metadata, config in self._load_env_configuration():
-            self._register_configuration(metadata, config)
-
-        if self._config_file:
-            for metadata, config in self._load_file_configuration(self._config_file):
-                self._register_configuration(metadata, config)
-
-    def _load_env_configuration(
-        self,
-    ) -> list[tuple[MCPServerMetadata, MCPConnectionConfig]]:
-        """Parse environment-driven MCP server configuration."""
-
-        entries: list[tuple[MCPServerMetadata, MCPConnectionConfig]] = []
-
-        for server in load_mcp_servers_from_env():
-            capabilities = [MCPServerCapability(name=name) for name in server.capabilities]
-
-            metadata = MCPServerMetadata(
-                name=server.name,
-                description=f"Environment-configured MCP server: {server.name}",
-                capabilities=capabilities,
-                raw={"source": "env"},
-            )
-
-            connection_config = MCPConnectionConfig(
-                name=server.name,
-                transport=server.transport,
-                url=str(server.url) if server.url else None,
-                api_key=server.api_key,
-                timeout=server.timeout,
-                max_connections=server.max_connections,
-                command=server.command,
-                args=server.args,
-                env=server.env,
-                config_source="env",
-            )
-
-            entries.append((metadata, connection_config))
-
-        return entries
-
-    def _load_file_configuration(
-        self, config_path: Path
-    ) -> list[tuple[MCPServerMetadata, MCPConnectionConfig]]:
-        """Parse YAML configuration for MCP servers."""
-
-        entries: list[tuple[MCPServerMetadata, MCPConnectionConfig]] = []
-
-        if not config_path.exists():
-            logger.warning("MCP configuration file %s not found", config_path)
-            return entries
-
-        try:
-            data = yaml.safe_load(config_path.read_text()) or {}
-        except Exception as exc:
-            logger.exception("Failed to parse MCP configuration file %s: %s", config_path, exc)
-            raise
-
-        for server_data in data.get("servers", []):
-            metadata_dict = self._interpolate_env(server_data.get("metadata", {}))
-            config_dict = self._interpolate_env(server_data.get("config", {}))
-
-            if not config_dict:
-                logger.debug(
-                    "Skipping MCP server with empty config block: %s",
-                    metadata_dict.get("name", "unknown"),
-                )
-                continue
-
-            metadata = self._build_metadata(metadata_dict)
-            metadata.raw.setdefault("source", str(config_path))
-            connection_config = self._build_connection_config(config_dict, source="file")
-            entries.append((metadata, connection_config))
-
-        return entries
-
-    def _build_metadata(self, metadata_dict: dict[str, Any]) -> MCPServerMetadata:
-        raw_copy = dict(metadata_dict) if metadata_dict else {}
-        capabilities_data = raw_copy.pop("capabilities", [])
-        capabilities = self._normalise_capabilities(capabilities_data)
-
-        filtered_fields = {
-            key: raw_copy.get(key)
-            for key in (
-                "name",
-                "version",
-                "description",
-                "author",
-                "homepage",
-                "tags",
-                "categories",
-                "supported_domains",
-            )
-            if raw_copy.get(key) is not None
-        }
-
-        return MCPServerMetadata(
-            capabilities=capabilities,
-            raw=metadata_dict or {},
-            **filtered_fields,
-        )
-
-    def _build_connection_config(
-        self, config_dict: dict[str, Any], source: str
-    ) -> MCPConnectionConfig:
-        headers = config_dict.pop("headers", {}) or {}
-        authentication = config_dict.pop("authentication", {}) or {}
-        command = config_dict.pop("command", None)
-        args = config_dict.pop("args", []) or []
-        env = config_dict.pop("env", {}) or {}
-        extra = {
-            k: v
-            for k, v in config_dict.items()
-            if k not in {"name", "transport", "url", "api_key", "timeout", "max_connections"}
-        }
-
-        name = config_dict.get("name")
-        if not name:
-            raise ConfigurationError("MCP configuration entry missing 'name'")
-
-        return MCPConnectionConfig(
-            name=name,
-            transport=config_dict.get("transport", "http"),
-            url=config_dict.get("url"),
-            api_key=config_dict.get("api_key"),
-            timeout=float(config_dict.get("timeout", 30.0)),
-            max_connections=int(config_dict.get("max_connections", 5)),
-            command=command,
-            args=args if isinstance(args, list) else [args],
-            env=env,
-            headers=headers,
-            authentication=authentication,
-            config_source=source,
-            extra=extra,
-        )
-
-    def _normalise_capabilities(self, capabilities_data: Any) -> list[MCPServerCapability]:
-        capabilities: list[MCPServerCapability] = []
-
-        if isinstance(capabilities_data, list):
-            for capability in capabilities_data:
-                if isinstance(capability, str):
-                    capabilities.append(MCPServerCapability(name=capability))
-                elif isinstance(capability, dict):
-                    capabilities.append(
-                        MCPServerCapability(
-                            name=capability.get("name"),
-                            description=capability.get("description"),
-                            category=capability.get("category"),
-                            input_schema=capability.get("input_schema", {}),
-                            output_schema=capability.get("output_schema", {}),
-                        )
-                    )
-        return capabilities
-
-    def _interpolate_env(self, value: Any) -> Any:
-        if isinstance(value, dict):
-            return {key: self._interpolate_env(val) for key, val in value.items()}
-        if isinstance(value, list):
-            return [self._interpolate_env(item) for item in value]
-        if isinstance(value, str):
-            return os.path.expandvars(value)
-        return value
-
-    def _register_configuration(self, metadata: MCPServerMetadata, config: MCPConnectionConfig):
-        if not config.name:
-            raise ConfigurationError("MCP server configuration missing required 'name'")
-
-        if config.transport == "http" and (
-            not config.url or not str(config.url).startswith(("http://", "https://"))
-        ):
-            logger.warning(
-                "Skipping MCP server '%s': invalid or missing HTTP URL (%r)",
-                metadata.name,
-                config.url,
-            )
-            return
-
-        if config.transport == "stdio" and not config.command:
-            logger.warning(
-                "Skipping MCP server '%s': stdio transport requires 'command' configuration",
-                metadata.name,
-            )
-            return
-
-        self.connection_configs[config.name] = config
-        self.server_metadata[config.name] = metadata
-        self._index_capabilities(config.name, metadata.capabilities)
-        self._server_stats.setdefault(
-            config.name,
-            {
-                "tool_calls": 0,
-                "successful_tool_calls": 0,
-                "failed_tool_calls": 0,
-                "avg_execution_time": 0.0,
-            },
-        )
-
-    def _index_capabilities(self, server_name: str, capabilities: list[MCPServerCapability]):
-        for capability in capabilities:
-            if not capability.name:
-                continue
-            self._capability_index.setdefault(capability.name, [])
-            if server_name not in self._capability_index[capability.name]:
-                self._capability_index[capability.name].append(server_name)
-
-    async def _initialize_connections(self):
-        """Initialize MCP connections for all configured servers."""
-        self.connections.clear()
-
-        for server_name, config in self.connection_configs.items():
-            try:
-                await self._create_connection(server_name, config)
-            except ConnectionError:
-                self._stats["failed_connections"] += 1
-            except Exception as exc:  # pragma: no cover - defensive guard
-                self._stats["failed_connections"] += 1
-                logger.exception(
-                    "Unexpected error creating MCP connection to %s: %s",
-                    server_name,
-                    exc,
-                )
-
-    async def _create_connection(self, server_name: str, config: MCPConnectionConfig):
-        """
-        Create a proper MCP connection using the official SDK.
-
-        Args:
-            server_name: Name of the MCP server
-            config: Connection configuration
-        """
-        try:
-            self._stats["connection_attempts"] += 1
-
-            if config.transport == "http":
-                # Use streamablehttp_client for HTTP transport
-                session = await self._create_http_connection(config)
-            elif config.transport == "stdio":
-                # Use stdio_client for local process transport
-                session = await self._create_stdio_connection(config)
-            else:
-                raise ConfigurationError(f"Unsupported transport type: {config.transport}")
-
-            metadata = self.server_metadata.get(server_name)
-            session.metadata = metadata
-            session.config = config
-            self.connections[server_name] = session
-            self._stats["successful_connections"] += 1
-
-            logger.info(
-                f"Successfully created MCP connection to {server_name} via {config.transport}"
-            )
-
-        except Exception as e:
-            logger.exception(f"Failed to create MCP connection to {server_name}: {e}")
-            raise
-
-    def _serialize_message(message) -> str:
-        """
-        Serialize a SessionMessage or JSONRPCMessage to JSON-RPC string format.
-
-        This method handles the different message types that can come from
-        the MCP ClientSession and converts them to proper JSON-RPC strings
-        that can be sent to MCP servers via stdin.
-        """
-        try:
-            # Try the simplest approach first - direct model_dump if available
-            if hasattr(message, "model_dump"):
-                try:
-                    data = message.model_dump(exclude_none=True, by_alias=True)
-                    return json.dumps(data)
-                except Exception:
-                    pass  # Continue to other methods
-
-            # Try with root attribute for SessionMessage
-            if hasattr(message, "root") and hasattr(message.root, "model_dump"):
-                try:
-                    data = message.root.model_dump(exclude_none=True, by_alias=True)
-                    return json.dumps(data)
-                except Exception:
-                    pass
-
-            # Simple field extraction for JSONRPCMessage
-            obj = message
-            if hasattr(message, "root"):
-                obj = message.root
-
-            # Build JSON-RPC message manually with common fields
-            result = {}
-
-            # Try to get common JSON-RPC fields
-            for field in ["jsonrpc", "method", "params", "result", "error", "id"]:
-                if hasattr(obj, field):
-                    value = getattr(obj, field)
-                    if value is not None:
-                        result[field] = value
-
-            # Default jsonrpc version if not present
-            if "jsonrpc" not in result:
-                result["jsonrpc"] = "2.0"
-
-            return json.dumps(result)
-
-        except Exception as e:
-            logger.exception(f"Failed to serialize message {type(message)}: {e}")
-            # Return a valid JSON-RPC error message as fallback
-            return json.dumps(
-                {
-                    "jsonrpc": "2.0",
-                    "error": {
-                        "code": -32603,
-                        "message": f"Serialization error: {str(e)}",
-                    },
-                }
-            )
-
-    async def _create_http_connection(self, config: MCPConnectionConfig) -> MCPConnection:
-        """Create HTTP connection using streamablehttp_client."""
-        try:
-            if not config.url:
-                raise ConfigurationError(f"HTTP transport for {config.name} requires a URL")
-
-            headers = self._prepare_http_headers(config)
-
-            connection_client = streamablehttp_client(
-                url=config.url, headers=headers, timeout=config.timeout
-            )
-
-            # Get the streams from the connection client
-            (
-                read_stream,
-                write_stream,
-                session_id_callback,
-            ) = await connection_client.__aenter__()
-
-            # Create ClientSession with the streams
-            client_session = ClientSession(read_stream, write_stream)
-
-            # Initialize the session
-            await client_session.initialize()
-
-            # Store complete connection state
-            connection = MCPConnection(
-                client_session=client_session,
-                read_stream=read_stream,
-                write_stream=write_stream,
-                session_id_callback=session_id_callback,
-                config=config,
-                metadata=self.server_metadata.get(config.name),
-            )
-
-            # Store the connection client for cleanup
-            connection._connection_client = connection_client
-
-            client_session._transport = "streamable_http"
-            client_session._get_session_id = session_id_callback
-            client_session._url = config.url
-            client_session._headers = headers
-            client_session._config_timeout = config.timeout
-            client_session._config_name = config.name
-
-            return connection
-
-        except Exception as e:
-            logger.exception(f"HTTP connection creation failed for {config.name}: {e}")
-            raise
-
-    async def _create_filtered_read_stream(self, original_read_stream, server_name: str):
-        """Create a filtered read stream that skips non-JSON startup messages."""
-        from mcp.types import JSONRPCMessage
-
-        class FilteredReadStream:
-            def __init__(self, original_stream, server_name):
-                self.original_stream = original_stream
-                self.server_name = server_name
-                self.buffer = []
-                self.filtered_startup = False
-                self.logger = logging.getLogger(f"{__name__}.{server_name}")
-
-            async def read(self):
-                """Read and filter messages from the stream."""
-                while True:
-                    # Read a message from the original stream
-                    message = await self.original_stream.read()
-
-                    # If we haven't filtered startup messages yet, check if this is JSON
-                    if not self.filtered_startup:
-                        try:
-                            # Try to parse as JSON to validate it's a proper JSON-RPC message
-                            if isinstance(message, str):
-                                JSONRPCMessage.model_validate_json(message)
-                                # If we get here, it's valid JSON
-                                self.filtered_startup = True
-                                self.logger.debug(
-                                    f"Received first valid JSON-RPC message from {self.server_name}"
-                                )
-                                return message
-                            else:
-                                # Binary data, return as-is
-                                self.filtered_startup = True
-                                return message
-                        except Exception:
-                            # This is not valid JSON, likely a startup message
-                            if isinstance(message, str) and message.strip():
-                                self.logger.info(
-                                    f"Filtering startup message from {self.server_name}: {message.strip()}"
-                                )
-                            # Skip this message and continue reading
-                            continue
-                    else:
-                        # Already filtered startup messages, return everything as-is
-                        return message
-
-            async def aclose(self):
-                """Close the underlying stream."""
-                if hasattr(self.original_stream, "aclose"):
-                    await self.original_stream.aclose()
-
-        return FilteredReadStream(original_read_stream, server_name)
-
-    async def _create_stdio_connection(self, config: MCPConnectionConfig) -> MCPConnection:
-        """Create stdio connection for local processes with startup message filtering."""
-        try:
-            import asyncio
-            import os
-            import subprocess
-
-            from mcp.client.session import ClientSession
-
-            # For stdio, we need a command and args
-            if not config.command:
-                raise ConfigurationError("stdio transport requires command configuration")
-
-            logger.info(
-                f"Starting filtered stdio server: {config.command} {' '.join(config.args or [])}"
-            )
-
-            # Prepare environment
-            env = os.environ.copy()
-            if config.env:
-                env.update(config.env)
-
-            # Create custom context manager for filtered stdio client
-            class FilteredStdioContext:
-                """Async context manager for stdio client with startup message filtering."""
-
-                def __init__(self, server_params, config_name):
-                    self.server_params = server_params
-                    self.config_name = config_name
-                    self.process = None
-                    self.task_group = None
-
-                async def __aenter__(self):
-                    """Start the process and create filtered streams."""
-                    try:
-                        import anyio
-                        from mcp.shared.session import SessionMessage
-
-                        # Start the process manually to have full control over stdout
-                        self.process = await asyncio.create_subprocess_exec(
-                            self.server_params["command"],
-                            *(self.server_params.get("args", [])),
-                            stdin=subprocess.PIPE,
-                            stdout=subprocess.PIPE,
-                            stderr=subprocess.PIPE,
-                            env=self.server_params.get("env", env),
-                        )
-
-                        # Create memory streams for ClientSession communication
-                        read_stream_writer, read_stream_reader = anyio.create_memory_object_stream(
-                            0
-                        )
-                        write_stream_writer, write_stream_reader = (
-                            anyio.create_memory_object_stream(0)
-                        )
-
-                        # stdout reader - filters startup messages and parses JSON-RPC
-                        async def stdout_reader():
-                            """Read from stdout, filter startup messages, and send valid JSON-RPC to ClientSession."""
-                            try:
-                                import json
-
-                                from mcp.types import JSONRPCMessage
-
-                                filtered_startup = False
-                                logger = logging.getLogger(f"{__name__}.{self.config_name}")
-
-                                while True:
-                                    line = await self.process.stdout.readline()
-                                    if not line:  # EOF
-                                        break
-
-                                    line_str = line.decode("utf-8").rstrip()
-
-                                    if not filtered_startup:
-                                        # Try to validate as JSON-RPC message
-                                        try:
-                                            # Try parsing as JSON to verify it's a valid message
-                                            parsed = json.loads(line_str)
-                                            # Check if it has JSON-RPC structure
-                                            if isinstance(parsed, dict) and "jsonrpc" in parsed:
-                                                filtered_startup = True
-                                                logger.debug(
-                                                    f"Received first valid JSON-RPC message from {self.config_name}"
-                                                )
-                                                # Parse as JSONRPCMessage and send to ClientSession
-                                                jsonrpc_message = (
-                                                    JSONRPCMessage.model_validate_json(line_str)
-                                                )
-                                                await read_stream_writer.send(
-                                                    SessionMessage(jsonrpc_message)
-                                                )
-                                            else:
-                                                # Valid JSON but not JSON-RPC, still filter it
-                                                if line_str.strip():
-                                                    logger.info(
-                                                        f"Filtering non-JSON-RPC message from {self.config_name}: {line_str}"
-                                                    )
-                                                continue
-                                        except json.JSONDecodeError:
-                                            # Not valid JSON, filter it as startup message
-                                            if line_str.strip():
-                                                logger.info(
-                                                    f"Filtering startup message from {self.config_name}: {line_str}"
-                                                )
-                                            continue
-                                        except Exception as e:
-                                            logger.exception(
-                                                f"Error parsing JSON-RPC message from {self.config_name}: {e}"
-                                            )
-                                            await read_stream_writer.send(e)
-                                    else:
-                                        # Already filtered startup, parse all remaining lines as JSON-RPC
-                                        try:
-                                            jsonrpc_message = JSONRPCMessage.model_validate_json(
-                                                line_str
-                                            )
-                                            await read_stream_writer.send(
-                                                SessionMessage(jsonrpc_message)
-                                            )
-                                        except Exception as e:
-                                            logger.exception(
-                                                f"Error parsing JSON-RPC message from {self.config_name}: {e}"
-                                            )
-                                            await read_stream_writer.send(e)
-                            except Exception as e:
-                                logger.exception(f"stdout_reader error for {self.config_name}: {e}")
-                                await read_stream_writer.send(e)
-                            finally:
-                                await read_stream_writer.aclose()
-
-                        # stdin writer - writes messages from ClientSession to process stdin
-                        async def stdin_writer():
-                            """Write messages from ClientSession to process stdin."""
-
-                            def _serialize_message(message) -> str:
-                                """
-                                Serialize a SessionMessage or JSONRPCMessage to JSON-RPC string format.
-
-                                This method handles the different message types that can come from
-                                the MCP ClientSession and converts them to proper JSON-RPC strings
-                                that can be sent to MCP servers via stdin.
-                                """
-                                try:
-                                    # Handle different message types from the MCP ClientSession
-                                    if hasattr(message, "root"):
-                                        # SessionMessage.root should be the JSONRPCMessage
-                                        jsonrpc_message = message.root
-                                    elif hasattr(message, "jsonrpc"):
-                                        # Alternative access path
-                                        jsonrpc_message = message.jsonrpc
-                                    else:
-                                        # message itself might be the JSONRPCMessage
-                                        jsonrpc_message = message
-
-                                    logger.debug(
-                                        f"Attempting to serialize message: {type(jsonrpc_message)} - {jsonrpc_message}"
-                                    )
-
-                                    # Check if this is a JSONRPCMessage with built-in serialization
-                                    if hasattr(jsonrpc_message, "model_dump_json"):
-                                        # Pydantic v2 model - use its JSON serialization with by_alias=True
-                                        result = jsonrpc_message.model_dump_json(by_alias=True)
-                                        logger.debug(f"Serialized using model_dump_json: {result}")
-                                        return result
-                                    elif isinstance(jsonrpc_message, dict):
-                                        # Already a dictionary
-                                        result = json.dumps(jsonrpc_message)
-                                        logger.debug(f"Serialized as dict: {result}")
-                                        return result
-                                    else:
-                                        # Last resort - try to extract what we can
-                                        if hasattr(jsonrpc_message, "__dict__"):
-                                            data = jsonrpc_message.__dict__
-                                            # Remove any internal/private attributes
-                                            filtered_data = {
-                                                k: v
-                                                for k, v in data.items()
-                                                if not k.startswith("_")
-                                            }
-                                            result = json.dumps(filtered_data)
-                                            logger.debug(f"Serialized via __dict__: {result}")
-                                            return result
-                                        else:
-                                            # Ultimate fallback - create a basic JSON-RPC message
-                                            result = json.dumps(
-                                                {
-                                                    "jsonrpc": "2.0",
-                                                    "id": getattr(jsonrpc_message, "id", 1),
-                                                    "method": getattr(
-                                                        jsonrpc_message,
-                                                        "method",
-                                                        "unknown",
-                                                    ),
-                                                    "params": getattr(
-                                                        jsonrpc_message, "params", {}
-                                                    ),
-                                                }
-                                            )
-                                            logger.debug(f"Serialized via fallback: {result}")
-                                            return result
-
-                                except Exception as e:
-                                    logger.exception(
-                                        f"Failed to serialize message {type(message)}: {e}"
-                                    )
-                                    logger.debug(f"Message details: {message}")
-                                    # Ultimate fallback - ensure we always return valid JSON-RPC
-                                    fallback = json.dumps(
-                                        {
-                                            "jsonrpc": "2.0",
-                                            "id": 1,
-                                            "error": {
-                                                "code": -32603,
-                                                "message": f"Internal serialization error: {str(e)}",
-                                            },
-                                        }
-                                    )
-                                    logger.debug(f"Using fallback serialization: {fallback}")
-                                    return fallback
-
-                            try:
-                                async for message in write_stream_reader:
-                                    # Convert SessionMessage to JSON-RPC format for the MCP server
-                                    # The MCP SDK expects to send raw JSON-RPC strings
-                                    json_data = _serialize_message(message)
-                                    logger.debug(f"Sending to {self.config_name}: {json_data}")
-
-                                    self.process.stdin.write(json_data.encode("utf-8"))
-                                    await self.process.stdin.drain()
-                            except Exception as e:
-                                logger.exception(f"stdin_writer error for {self.config_name}: {e}")
-                            finally:
-                                with suppress(BaseException):
-                                    self.process.stdin.close()
-
-                        # Create task group for reader and writer
-                        self.task_group = anyio.create_task_group()
-
-                        # Start reader and writer tasks
-                        await self.task_group.__aenter__()
-                        self.task_group.start_soon(stdout_reader)
-                        self.task_group.start_soon(stdin_writer)
-
-                        # Return the streams for ClientSession
-                        return read_stream_reader, write_stream_writer
-
-                    except Exception as e:
-                        logger.exception(
-                            f"Failed to start filtered stdio client for {self.config_name}: {e}"
-                        )
-                        raise
-
-                async def __aexit__(self, exc_type, exc_val, exc_tb):
-                    """Cleanup process and task group."""
-                    try:
-                        # Cleanup task group
-                        if self.task_group:
-                            await self.task_group.__aexit__(exc_type, exc_val, exc_tb)
-
-                        # Cleanup process
-                        if self.process and self.process.returncode is None:
-                            try:
-                                self.process.terminate()
-                                await asyncio.wait_for(self.process.wait(), timeout=5.0)
-                            except TimeoutError:
-                                self.process.kill()
-                                await self.process.wait()
-                    except Exception as e:
-                        logger.exception(f"Error during cleanup for {self.config_name}: {e}")
-
-            # Use the filtered stdio client
-            server_params = {
-                "command": config.command,
-                "args": config.args or [],
-                "env": env,
-            }
-
-            async with FilteredStdioContext(server_params, config.name) as (
-                read_stream,
-                write_stream,
-            ):
-                # Create ClientSession with filtered streams
-                client_session = ClientSession(read_stream, write_stream)
-
-                # Initialize the session
-                await client_session.initialize()
-                logger.info(
-                    f"Successfully initialized {config.name} with startup message filtering"
-                )
-
-                # Store complete connection state
-                connection = MCPConnection(
-                    client_session=client_session,
-                    read_stream=read_stream,
-                    write_stream=write_stream,
-                    session_id_callback=None,
-                    config=config,
-                    metadata=self.server_metadata.get(config.name),
-                )
-
-                logger.info(f"Successfully created stdio connection for {config.name}")
-                return connection
-
-        except Exception as e:
-            logger.exception(f"stdio connection creation failed for {config.name}: {e}")
-            raise
-
-    def _prepare_http_headers(self, config: MCPConnectionConfig) -> dict[str, str]:
-        headers = dict(config.headers or {})
-
-        headers.setdefault("User-Agent", "VDS-AI-Memory/2.0.0")
-        headers.setdefault("Accept", "application/json, text/event-stream")
-
-        if config.api_key and "Authorization" not in headers:
-            headers["Authorization"] = f"Bearer {config.api_key}"
-
-        if config.authentication:
-            self._apply_authentication(headers, config.authentication, config.api_key)
-
-        return headers
-
-    def _apply_authentication(
-        self,
-        headers: dict[str, str],
-        authentication: dict[str, Any],
-        fallback_api_key: str | None,
-    ):
-        auth = self._interpolate_env(authentication)
-        auth_type = (auth.get("type") or "").lower()
-
-        if auth_type in {"api_key", "api_key_in_header"}:
-            header_name = auth.get("header_name", "Authorization")
-            header_format = auth.get("header_format", "Bearer ${api_key}")
-            api_key = auth.get("api_key") or fallback_api_key
-            if api_key:
-                headers[header_name] = header_format.replace("${api_key}", api_key)
-        elif auth_type == "bearer_token":
-            token = auth.get("token") or fallback_api_key
-            if token:
-                headers["Authorization"] = f"Bearer {token}"
-        elif auth_type == "multi_api_key":
-            for key_name, key_value in (auth.get("api_keys") or {}).items():
-                resolved_value = self._interpolate_env(key_value)
-                if resolved_value:
-                    headers[key_name] = resolved_value
-        elif auth_type == "custom_headers":
-            for key_name, key_value in (auth.get("headers") or {}).items():
-                resolved_value = self._interpolate_env(key_value)
-                if resolved_value:
-                    headers[key_name] = resolved_value
-
-    async def execute_tool(
-        self,
-        tool_name: str,
-        arguments: dict[str, Any],
-        server_name: str | None = None,
-        timeout: float = 30.0,
-    ) -> CallToolResult:
-        """
-        Execute a tool using proper MCP protocol.
-
-        Args:
-            tool_name: Name of the tool to execute
-            arguments: Tool arguments
-            server_name: Specific server to use (optional)
-            timeout: Execution timeout
-
-        Returns:
-            CallToolResult from MCP protocol
-        """
-        await self._ensure_initialized()
-
-        self._stats["tool_calls"] += 1
-        start_time = time.perf_counter()
-        evaluated_server = server_name
-
-        try:
-            # Check if this is an internal VDS Memory service tool
-            if self._is_internal_memory_tool(tool_name):
-                logger.info(f"Executing internal memory tool: {tool_name}")
-                result = await self._execute_internal_memory_tool(
-                    tool_name, arguments, timeout=timeout
-                )
-                self._stats["successful_tool_calls"] += 1
-                elapsed = time.perf_counter() - start_time
-                self._record_server_success("vds-memory-service", elapsed)
-                result.server_name = "vds-memory-service"
-                return result
-
-            # Find appropriate server for the tool
-            target_server = server_name or await self._find_server_for_tool(tool_name)
-            evaluated_server = target_server
-
-            if not target_server:
-                raise ToolExecutionError(f"No server available for tool: {tool_name}")
-
-            connection = self.connections.get(target_server)
-            if not connection:
-                raise ConnectionError(f"No connection to server: {target_server}")
-
-            logger.info(f"Executing tool {tool_name} on server {target_server}")
-
-            # Execute tool using proper MCP protocol
-            result = await connection.client_session.call_tool(
-                tool_name, arguments, timeout=timeout
-            )
-
-            self._stats["successful_tool_calls"] += 1
-            elapsed = time.perf_counter() - start_time
-            self._record_server_success(target_server, elapsed)
-            result.server_name = target_server
-
-            logger.info(f"Tool {tool_name} executed successfully on {target_server}")
-            return result
-
-        except TimeoutError:
-            logger.exception(f"Tool execution timeout for {tool_name}")
-            self._record_server_failure(
-                evaluated_server or "unknown", time.perf_counter() - start_time
-            )
-            raise TimeoutError
-        except Exception as e:
-            logger.exception(f"Tool execution failed for {tool_name}: {e}")
-            self._record_server_failure(
-                evaluated_server or "unknown", time.perf_counter() - start_time
-            )
-            raise
-
-    async def list_tools(self, server_name: str | None = None) -> list[Tool]:
-        """
-        List available tools from MCP servers.
-
-        Args:
-            server_name: Specific server to query (optional)
-
-        Returns:
-            List of available tools
-        """
-        await self._ensure_initialized()
-
-        all_tools = []
-
-        servers_to_query = [server_name] if server_name else list(self.connection_configs.keys())
-
-        for server in servers_to_query:
-            try:
-                tools = await self._get_tools_for_server(server)
-                all_tools.extend(tools)
-            except Exception as exc:
-                logger.warning("Failed to list tools from %s: %s", server, exc)
-
-        return all_tools
-
-    async def _get_tools_for_server(
-        self, server_name: str, force_refresh: bool = False
-    ) -> list[Tool]:
-        if server_name not in self.connection_configs:
-            raise ConfigurationError(f"Server {server_name} is not configured")
-
-        cache_entry = self._tool_cache.get(server_name)
-        if cache_entry and not force_refresh and time.time() - cache_entry["timestamp"] < self._tool_cache_ttl:
-                return cache_entry["tools"]
-
-        connection = self.connections.get(server_name)
-        if not connection:
-            config = self.connection_configs.get(server_name)
-            if not config:
-                raise ConnectionError(f"No connection config for server: {server_name}")
-            await self._create_connection(server_name, config)
-            connection = self.connections.get(server_name)
-
-        if not connection:
-            raise ConnectionError(f"Unable to establish connection to server: {server_name}")
-
-        tools_result = await connection.client_session.list_tools()
-        tools: list[Tool] = []
-        for tool in tools_result.tools:
-            tool_dict = tool.model_dump()
-            tool_dict["server_name"] = server_name
-            tools.append(Tool(**tool_dict))
-
-        self._tool_cache[server_name] = {
-            "tools": tools,
-            "timestamp": time.time(),
-        }
-
-        return tools
-
-    async def get_resource(self, uri: str, server_name: str | None = None) -> ReadResourceResult:
-        """
-        Get a resource using proper MCP protocol.
-
-        Args:
-            uri: Resource URI
-            server_name: Specific server to query (optional)
-
-        Returns:
-            ReadResourceResult from MCP protocol
-        """
-        await self._ensure_initialized()
-
-        target_server = server_name or await self._find_server_for_resource(uri)
-
-        if not target_server:
-            raise ToolExecutionError(f"No server available for resource: {uri}")
-
-        connection = self.connections.get(target_server)
-        if not connection:
-            await self._create_connection(target_server, self.connection_configs[target_server])
-            connection = self.connections.get(target_server)
-        if not connection:
-            raise ConnectionError(f"No connection to server: {target_server}")
-
-        try:
-            # Get resource using proper MCP protocol
-            result = await connection.client_session.read_resource(uri)
-            result.server_name = target_server
-            return result
-
-        except Exception as e:
-            logger.exception(f"Resource access failed for {uri}: {e}")
-            raise
-
-    async def search_resources(
-        self,
-        query: str,
-        resource_types: list[str] | None = None,
-        limit: int = 50,
-    ) -> list[dict[str, Any]]:
-        """Perform a simple resource search across all configured servers."""
-
-        await self._ensure_initialized()
-
-        results: list[dict[str, Any]] = []
-        normalized_query = (query or "").lower()
-        type_filters = [rt.lower() for rt in (resource_types or [])]
-
-        for server_name in self.connection_configs:
-            if len(results) >= limit:
-                break
-
-            connection = self.connections.get(server_name)
-            if not connection:
-                await self._create_connection(server_name, self.connection_configs[server_name])
-                connection = self.connections.get(server_name)
-
-            if not connection or not hasattr(connection.client_session, "list_resources"):
-                continue
-
-            try:
-                resources_result = await connection.client_session.list_resources()
-            except Exception as exc:  # pragma: no cover - diagnostic only
-                logger.debug("Failed to list resources on %s: %s", server_name, exc)
-                continue
-
-            for resource in getattr(resources_result, "resources", []):
-                resource_dict = (
-                    resource.model_dump() if hasattr(resource, "model_dump") else resource
-                )
-                searchable = " ".join(
-                    str(resource_dict.get(field, "")) for field in ("uri", "name", "description")
-                ).lower()
-
-                if normalized_query and normalized_query not in searchable:
-                    continue
-
-                resource_type = (
-                    resource_dict.get("mimeType") or resource_dict.get("type") or ""
-                ).lower()
-                if type_filters and resource_type not in type_filters:
-                    continue
-
-                resource_dict["server_name"] = server_name
-                results.append(resource_dict)
-
-                if len(results) >= limit:
-                    break
-
-        return results
-
-    async def _find_server_for_tool(self, tool_name: str) -> str | None:
-        """Find a server that provides the specified tool."""
-        candidates = self._capability_index.get(tool_name, [])
-        for candidate in candidates:
-            if candidate in self.connection_configs:
-                return candidate
-
-        for server_name in self.connection_configs:
-            try:
-                tools = await self._get_tools_for_server(server_name)
-                if any(tool.name == tool_name for tool in tools):
-                    self._capability_index.setdefault(tool_name, []).append(server_name)
-                    return server_name
-            except Exception as exc:
-                logger.debug("Failed to resolve tool %s on %s: %s", tool_name, server_name, exc)
-        return None
-
-    async def _find_server_for_resource(self, uri: str) -> str | None:
-        """Find a server that provides the specified resource."""
-        for server_name, connection in self.connections.items():
-            try:
-                # Try to read the resource
-                await connection.client_session.read_resource(uri)
-                return server_name
-            except Exception:
-                continue
-        return None
-
-    async def get_server_capabilities(self, server_name: str) -> dict[str, Any]:
-        """Get capabilities from a specific server."""
-        await self._ensure_initialized()
-
-        if server_name not in self.connection_configs:
-            raise ConnectionError(f"No server named {server_name} is configured")
-
-        connection = self.connections.get(server_name)
-        if not connection:
-            await self._create_connection(server_name, self.connection_configs[server_name])
-            connection = self.connections.get(server_name)
-
-        if not connection:
-            raise ConnectionError(f"Unable to establish connection to server: {server_name}")
-
-        try:
-            server_info = None
-            if hasattr(connection.client_session, "get_server_info"):
-                server_info = connection.client_session.get_server_info()
-
-            tools = await self._get_tools_for_server(server_name)
-
-            metadata = self.server_metadata.get(server_name)
-            capabilities = [cap.__dict__ for cap in (metadata.capabilities if metadata else [])]
-
-            return {
-                "server_info": server_info,
-                "tools": [tool.model_dump() for tool in tools],
-                "tool_count": len(tools),
-                "metadata": metadata.raw if metadata else {},
-                "capabilities": capabilities,
-            }
-
-        except Exception as e:
-            logger.exception(f"Failed to get capabilities from {server_name}: {e}")
-            return {"error": str(e)}
-
-    def list_servers(self, _status: str | None = None) -> list[str]:
-        """List configured MCP servers including internal services."""
-        servers = list(self.connection_configs.keys())
-
-        # Include internal vds-memory-service
-        if "vds-memory-service" not in servers:
-            servers.append("vds-memory-service")
-
-        return servers
-
-    async def find_servers_for_tool(self, tool_name: str) -> list[str]:
-        """Return all servers that provide the requested tool."""
-        await self._ensure_initialized()
-
-        # Check if this is an internal VDS Memory service tool
-        if self._is_internal_memory_tool(tool_name):
-            return ["vds-memory-service"]
-
-        servers = set(self._capability_index.get(tool_name, []))
-        if servers:
-            return list(servers)
-
-        for server_name in self.connection_configs:
-            try:
-                tools = await self._get_tools_for_server(server_name)
-                if any(tool.name == tool_name for tool in tools):
-                    servers.add(server_name)
-            except Exception:
-                continue
-
-        if servers:
-            self._capability_index[tool_name] = list(servers)
-
-        return list(servers)
-
-    async def register_server(
-        self,
-        metadata: MCPServerMetadata,
-        config: MCPConnectionConfig,
-        connect: bool = True,
-    ):
-        """Register a new MCP server at runtime."""
-        self._register_configuration(metadata, config)
-        self._tool_cache.pop(config.name, None)
-        if self.initialized and connect:
-            await self._create_connection(config.name, config)
-
-    async def unregister_server(self, server_name: str):
-        """Remove a server and close any active connections."""
-        await self._close_connection(server_name)
-        self.connection_configs.pop(server_name, None)
-        self.server_metadata.pop(server_name, None)
-        self._tool_cache.pop(server_name, None)
-        self._server_stats.pop(server_name, None)
-        for capability, servers in list(self._capability_index.items()):
-            if server_name in servers:
-                servers.remove(server_name)
-            if not servers:
-                self._capability_index.pop(capability, None)
-
-    async def reload(self, config_file: str | None = None):
-        """Reload configuration from scratch."""
-        await self.initialize(
-            config_file or (str(self._config_file) if self._config_file else None),
-            force_reload=True,
-        )
-
-    async def get_statistics(self) -> dict[str, Any]:
-        """Get MCP client statistics."""
-
-        # Include internal vds-memory-service in available servers
-        available_servers = list(self.connection_configs.keys())
-        internal_capabilities = 0
-
-        # Check if vds-memory-service should be listed as an internal server
-        if "vds-memory-service" not in available_servers:
-            available_servers.append("vds-memory-service")
-            internal_capabilities = (
-                4  # search_memory, create_memory, get_session_context, consolidate_memories
-            )
-
-        # Ensure server_stats includes vds-memory-service
-        server_stats = dict(self._server_stats)
-        if "vds-memory-service" not in server_stats:
-            server_stats["vds-memory-service"] = {
-                "tool_calls": 0,
-                "successful_tool_calls": 0,
-                "failed_tool_calls": 0,
-                "avg_execution_time": 0.0,
-            }
-
-        return {
-            "initialized": self.initialized,
-            "connection_stats": {
-                "total_attempts": self._stats["connection_attempts"],
-                "successful": self._stats["successful_connections"],
-                "failed": self._stats["failed_connections"],
-                "active_connections": len(self.connections),
-            },
-            "execution_stats": {
-                "total_tool_calls": self._stats["tool_calls"],
-                "successful_tool_calls": self._stats["successful_tool_calls"],
-                "success_rate": (
-                    self._stats["successful_tool_calls"] / max(self._stats["tool_calls"], 1) * 100
-                ),
-            },
-            "server_count": len(self.connection_configs),
-            "internal_capabilities": internal_capabilities,
-            "available_servers": available_servers,
-            "server_stats": server_stats,
-            "vds_capabilities": [],  # Will be populated with external server tools
-        }
-
-    def get_performance_metrics(self) -> dict[str, Any]:
-        """Return lightweight performance metrics for monitoring APIs."""
-
-        total_tool_calls = self._stats.get("tool_calls", 0)
-        successful_tool_calls = self._stats.get("successful_tool_calls", 0)
-        failed_tool_calls = max(total_tool_calls - successful_tool_calls, 0)
-
-        per_server: list[dict[str, Any]] = []
-        for server_name, stats in self._server_stats.items():
-            server_total = stats.get("tool_calls", 0)
-            server_success = stats.get("successful_tool_calls", 0)
-            per_server.append(
-                {
-                    "server": server_name,
-                    "tool_calls": server_total,
-                    "successful_tool_calls": server_success,
-                    "failed_tool_calls": stats.get("failed_tool_calls", 0),
-                    "avg_execution_time": stats.get("avg_execution_time", 0.0),
-                    "success_rate": (server_success / max(server_total, 1)) * 100,
-                    "last_success_at": stats.get("last_success_at"),
-                    "last_failure_at": stats.get("last_failure_at"),
-                }
-            )
-
-        return {
-            "initialized": self.initialized,
-            "tool_calls": total_tool_calls,
-            "successful_tool_calls": successful_tool_calls,
-            "failed_tool_calls": failed_tool_calls,
-            "success_rate": (successful_tool_calls / max(total_tool_calls, 1)) * 100,
-            "connection_attempts": self._stats.get("connection_attempts", 0),
-            "successful_connections": self._stats.get("successful_connections", 0),
-            "failed_connections": self._stats.get("failed_connections", 0),
-            "per_server": per_server,
-        }
-
-    async def close(self):
-        """Close all MCP connections."""
-        for server_name in list(self.connections.keys()):
-            await self._close_connection(server_name)
-
-        self.initialized = False
-        self._tool_cache.clear()
-        logger.info("All MCP connections closed")
-
-    async def _close_connection(self, server_name: str):
-        connection = self.connections.pop(server_name, None)
-        if not connection:
-            return
-
-        try:
-            await connection.client_session.close()
-        except Exception as exc:
-            logger.warning("Error closing client session for %s: %s", server_name, exc)
-
-        # Handle stdio context cleanup - properly exit the context manager
-        if connection._stdio_context:
-            try:
-                await connection._stdio_context.__aexit__(None, None, None)  # type: ignore[attr-defined]
-            except Exception as cleanup_error:
-                logger.warning(
-                    "Error cleaning up stdio context for %s: %s",
-                    server_name,
-                    cleanup_error,
-                )
-
-        # Handle HTTP connection cleanup
-        if connection._connection_client:
-            try:
-                await connection._connection_client.__aexit__(None, None, None)  # type: ignore[attr-defined]
-            except Exception as cleanup_error:
-                logger.warning(
-                    "Error cleaning up connection client for %s: %s",
-                    server_name,
-                    cleanup_error,
-                )
-
-        self._tool_cache.pop(server_name, None)
-        logger.info("Closed connection to %s", server_name)
-
-    def _record_server_success(self, server_name: str, elapsed: float):
-        stats = self._server_stats.setdefault(
-            server_name,
-            {
-                "tool_calls": 0,
-                "successful_tool_calls": 0,
-                "failed_tool_calls": 0,
-                "avg_execution_time": 0.0,
-            },
-        )
-
-        stats["tool_calls"] += 1
-        stats["successful_tool_calls"] += 1
-        success_count = stats["successful_tool_calls"]
-        current_avg = stats.get("avg_execution_time", 0.0)
-        stats["avg_execution_time"] = ((current_avg * (success_count - 1)) + elapsed) / max(
-            success_count, 1
-        )
-        stats["last_success_at"] = time.time()
-
-    def _record_server_failure(self, server_name: str, _elapsed: float):
-        stats = self._server_stats.setdefault(
-            server_name,
-            {
-                "tool_calls": 0,
-                "successful_tool_calls": 0,
-                "failed_tool_calls": 0,
-                "avg_execution_time": 0.0,
-            },
-        )
-
-        stats["tool_calls"] += 1
-        stats["failed_tool_calls"] += 1
-        stats["last_failure_at"] = time.time()
-
-    async def _ensure_initialized(self):
-        if not self.initialized:
-            config_path = str(self._config_file) if self._config_file else None
-            await self.initialize(config_path)
-
-    def _is_internal_memory_tool(self, tool_name: str) -> bool:
-        """Check if this is an internal VDS Memory service tool."""
-        internal_tools = {
-            "search_memory",
-            "create_memory",
-            "get_session_context",
-            "consolidate_memories",
-        }
-        return tool_name in internal_tools
-
-    async def _execute_internal_memory_tool(
-        self, tool_name: str, arguments: dict[str, Any], _timeout: float = 30.0
-    ) -> CallToolResult:
-        """Execute internal VDS Memory service tool."""
-        from ..models.ai_memory_models import (
-            MemoryCreateRequest,
-            MemorySearchRequest,
-            SessionContextRequest,
-        )
-        from ..services.ai_memory_service import get_ai_memory_service
-
-        try:
-            # Get the AI memory service instance
-            ai_memory_service = await get_ai_memory_service()
-
-            # Map tool names to service methods with proper request objects
-            if tool_name == "search_memory":
-                request = MemorySearchRequest(
-                    query=arguments.get("query", ""),
-                    user_id=arguments.get("user_id", "default"),
-                    limit=arguments.get("limit", 10),
-                    session_id=arguments.get("session_id"),
-                    offset=arguments.get("offset", 0),
-                )
-                result = await ai_memory_service.search_memories(request=request)
-            elif tool_name == "create_memory":
-                request = MemoryCreateRequest(
-                    memory=arguments.get("memory", ""),
-                    user_id=arguments.get("user_id", "default"),
-                    metadata=arguments.get("metadata", {}),
-                    session_id=arguments.get("session_id"),
-                    category=arguments.get("category"),
-                    importance=arguments.get("importance"),
-                )
-                result = await ai_memory_service.create_memory(request=request)
-            elif tool_name == "get_session_context":
-                request = SessionContextRequest(
-                    session_id=arguments.get("session_id"),
-                    user_id=arguments.get("user_id", "default"),
-                )
-                result = await ai_memory_service.get_session_context(request=request)
-            elif tool_name == "consolidate_memories":
-                request = SessionConsolidationRequest(
-                    session_id=arguments.get("session_id"),
-                    user_id=arguments.get("user_id", "default"),
-                )
-                result = await ai_memory_service.consolidate_session(request=request)
-            else:
-                raise ToolExecutionError(f"Unknown internal tool: {tool_name}")
-
-            # Convert result to CallToolResult format
-            return CallToolResult(
-                content=[
-                    {
-                        "type": "text",
-                        "text": json.dumps(
-                            result.model_dump() if hasattr(result, "model_dump") else result,
-                            default=str,
-                        ),
-                    }
-                ],
-                isError=False,
-            )
-
-        except Exception as e:
-            logger.exception(f"Internal tool execution failed for {tool_name}: {e}")
-            return CallToolResult(
-                content=[{"type": "text", "text": json.dumps({"error": str(e)})}],
-                isError=True,
-            )
-
-
-# Global official MCP client instance
-official_mcp_client = OfficialMCPClient()
-
-
-@asynccontextmanager
-async def get_mcp_client():
-    """Context manager for getting MCP client with proper initialization."""
-    try:
-        if not official_mcp_client.initialized:
-            await official_mcp_client.initialize()
-        yield official_mcp_client
-    except Exception as e:
-        logger.exception(f"MCP client error: {e}")
-        raise
-    finally:
-        # Don't close here - keep connections alive for reuse
-        pass