Source code for aidputils.agents.tools.mcp.mcp_client

# Lightweight MCP HTTP client wrapper extracted from mcp_service
# to centralize transport/session-specific logic.
from dataclasses import dataclass, field
from typing import Dict, Optional
from contextlib import asynccontextmanager

from mcp.client.streamable_http import streamablehttp_client
from mcp.client.session import ClientSession
import mcp.types as types
from datetime import timedelta
import logging
import asyncio
from typing import Tuple, Callable, Awaitable
from aidputils.agents.toolkit.metrics_util import MetricUtility
from aidputils.agents.tools.mcp.mcp_auth import build_headers_from_auth

# Gracefully handle ClosedResourceError without hard dependency if anyio is unavailable
try:
    from anyio import ClosedResourceError as AnyioClosedResourceError
except Exception:
    class AnyioClosedResourceError(Exception):
        pass

Key = Tuple[str, str]

# Single registry holding both session and its per-key lock
_ENTRIES: Dict[Key, "CacheEntry"] = {}
_ENTRIES_GUARD: asyncio.Lock = asyncio.Lock()

logger = logging.getLogger(__name__)

_MCP_TOOL_METRICS = MetricUtility("MCP_TOOL")
_MCP_TOOL_SUCCESS_COUNTER = _MCP_TOOL_METRICS.create_counter("MCP_TOOL_SUCCESS_COUNTER")
_MCP_TOOL_FAILURE_COUNTER = _MCP_TOOL_METRICS.create_counter("MCP_TOOL_FAILURE_COUNTER")



@dataclass
class CacheEntry:
    session: Optional[ClientSession] = None
    lock: asyncio.Lock = field(default_factory=asyncio.Lock)




class MCPSessionCache:



    def get(self, server_name: str, session_id: str) -> Optional[ClientSession]:
        """
        Return a cached ClientSession for (server_name, session_id) if present.
        """
        entry = _ENTRIES.get((server_name, session_id))
        return entry.session if entry is not None else None




    def set(self, server_name: str, session_id: str, session: ClientSession) -> None:
        """
        Store a ClientSession for (server_name, session_id).
        """
        key = (server_name, session_id)
        entry = _ENTRIES.get(key)
        if entry is None:
            entry = CacheEntry()
            _ENTRIES[key] = entry
        entry.session = session




    def delete(self, server_name: str, session_id: str) -> None:
        """
        Remove a cached ClientSession for (server_name, session_id).
        """
        _ENTRIES.pop((server_name, session_id), None)


    async def _get_entry_for_key(self, key: Key) -> "CacheEntry":
        """
        Get or create a CacheEntry (with its lock) for a given key in a thread-safe manner.
        """
        async with _ENTRIES_GUARD:
            entry = _ENTRIES.get(key)
            if entry is None:
                entry = CacheEntry()
                _ENTRIES[key] = entry
            return entry



    async def get_or_create(
        self,
        server_name: str,
        session_id: str,
        init_session: Callable[[str], Awaitable[ClientSession]],
    ) -> ClientSession:
        """
        Return a session for (server_name, session_id) if cached, otherwise create one by
        awaiting init_session(session_id), store it, and return it. Uses a per-key lock
        to avoid duplicate concurrent initialization.
        """
        key = (server_name, session_id)
        entry = await self._get_entry_for_key(key)
        if entry.session is not None:
            return entry.session

        async with entry.lock:
            # Re-check within lock
            if entry.session is not None:
                return entry.session

            # Create new session via provided initializer
            session = await init_session(session_id)
            entry.session = session
            return session




    async def close(self, server_name: str, session_id: str) -> None:
        """
        Best-effort closure of a single session; does not enforce underlying transport close semantics.
        """
        key = (server_name, session_id)
        entry = _ENTRIES.pop(key, None)
        if entry is None:
            return
        session = entry.session
        close = getattr(session, "close", None) if session is not None else None
        if callable(close):
            try:
                await close()
            except Exception:
                # Ignore errors on best-effort close
                pass




    async def close_all(self) -> None:
        """
        Best-effort closure of all cached sessions.
        """
        keys = list(_ENTRIES.keys())
        for server_name, session_id in keys:
            await self.close(server_name, session_id)





@dataclass
class MCPHTTPClient:

    def __init__(self,
                 server_name,
                 url,
                 headers,
                 auth: Optional[dict] = None,
                 transport: str = "streamable_http",
                 request_timeout_seconds: float = 30.0,
                 sse_read_timeout_seconds: float = 30.0):
        self.url = url
        self.server_name = server_name
        self.headers = headers
        self.auth = auth or {}
        self.per_session_metadata = {}
        self.transport = transport
        self.request_timeout_seconds = request_timeout_seconds
        self.sse_read_timeout_seconds = sse_read_timeout_seconds
        self.mcp_session_cache = MCPSessionCache()



    def get_all_headers(self) -> Dict[str, str]:
        """
        Return effective headers by augmenting base headers with auth-derived headers.
        Authorization is computed via mcp_auth.build_headers_from_auth using self.auth.
        """
        base: Dict[str, str] = dict(self.headers or {})
        try:
            auth_headers = build_headers_from_auth(getattr(self, "auth", None))
            if isinstance(auth_headers, dict):
                base.update(auth_headers)
        except Exception:
            # Do not fail header construction on auth errors here
            logger.error("Failed to build auth headers", exc_info=True)
            pass
        return base





    async def init_session(self, session_id: str) -> ClientSession:
        """
        Initialize and park a ClientSession for this server and session_id.
        Persist it into the global mcp_session_cache and return the session once ready.
        """
        ready_event = asyncio.Event()
        stop_event = asyncio.Event()
        entry = {"ready_event": ready_event, "stop_event": stop_event, "task": None}
        task = asyncio.create_task(self._run(session_id))
        entry["task"] = task
        self.per_session_metadata[session_id] = entry
        await ready_event.wait()
        if task.done():
            task.result()
        # Return the cached session
        sess =self.mcp_session_cache.get(self.server_name, session_id)
        if sess is None:
            raise RuntimeError("Failed to initialize MCP session")
        return sess


    async def _run(self, session_id):
        try:
            async with streamablehttp_client(
                self.url,
                headers=self.get_all_headers()) as (read_stream, write_stream, _):
                async with ClientSession(read_stream, write_stream) as session_local:
                    try:
                        try:
                            await session_local.initialize()
                        except Exception:
                            # Some client versions initialize lazily; ignore if not supported
                            pass
                        # Persist into cache and signal readiness
                        self.mcp_session_cache.set(self.server_name, session_id, session_local)
                        entry = self.per_session_metadata[session_id]
                        entry.get("ready_event").set()
                        # Park until asked to stop
                        await entry.get("stop_event").wait()
                        logger.info("Stop signal received, exiting scope.")
                    except Exception as e:
                        logger.error("Exception on server '%s' for session_id '%s': %s",
                                     self.server_name, session_id, e, exc_info=True)
        except Exception as e:
            entry = self.per_session_metadata[session_id]
            logger.error("MCPHTTPClient persistent scope crashed for server '%s', session_id '%s': %s", self.server_name, session_id, e, exc_info=True)
            if not entry.get("ready_event").is_set():
                entry.get("ready_event").set()
            raise
        finally:
            self.mcp_session_cache.delete(self.server_name, session_id)



    async def stop_session(self, session_id: str) -> None:
        """
        Stop a parked persistent session for the given session_id, if running,
        and remove it from the cache.
        """
        info = self.per_session_metadata.get(session_id)
        if not info:
            await self.mcp_session_cache.close(self.server_name, session_id)
            return
        task = info.get("task")
        stop_event = info.get("stop_event")
        if task and stop_event and not task.done():
            stop_event.set()
            try:
                await task
            except Exception:
                # Already logged by the scope
                pass
        self.per_session_metadata.pop(session_id, None)
        await self.mcp_session_cache.close(self.server_name, session_id)




    @asynccontextmanager
    async def session(self, _server_name: str):
        """
        Ephemeral session context manager for single-call operations.
        """
        async with streamablehttp_client(
            self.url,
            headers=self.get_all_headers(),
            timeout=timedelta(seconds=float(self.request_timeout_seconds)),
            sse_read_timeout=timedelta(seconds=float(self.sse_read_timeout_seconds)),
            terminate_on_close=True,
        ) as (read_stream, write_stream, _):
            async with ClientSession(read_stream, write_stream) as session:
                try:
                    await session.initialize()
                    yield session
                finally:
                    pass




    async def get_tools(self, server_name: str, use_cached: bool = False, session_id: Optional[str] = None):
        """
        List tools with transparent pagination.
        - If use_cached is False (default): use an on-demand ephemeral session.
        - If use_cached is True: reuse/create a persistent session associated with session_id from the cache.
        Accumulates all pages by following nextCursor when present.
        Returns a list of tool objects.
        """
        async def _accumulate_tools(session: ClientSession):
            all_tools = []
            cursor = None
            while True:
                # Request a page
                if cursor:
                    # Prefer new-style params; fall back to deprecated positional cursor if needed
                    try:
                        result = await session.list_tools(params=types.PaginatedRequestParams(cursor=cursor))
                    except TypeError:
                        result = await session.list_tools(cursor=cursor)
                else:
                    result = await session.list_tools()

                # Collect tools from this page
                tools = getattr(result, "tools", None) or []
                all_tools.extend(tools)

                # Determine if another page exists
                next_cursor = getattr(result, "nextCursor", None) or getattr(result, "next_cursor", None)

                if not next_cursor:
                    break

                cursor = next_cursor

            return all_tools

        if not use_cached:
            async with self.session(server_name) as session:
                return await _accumulate_tools(session)
        else:
            if not session_id:
                raise ValueError("session_id is required when use_cached=True for get_tools")
            session = await self.get_or_create_session(session_id)
            return await _accumulate_tools(session)




    async def get_resources(self, server_name: str, uris=None, use_cached: bool = False, session_id: Optional[str] = None):
        """
        List resources.
        - If use_cached is False (default): use an on-demand ephemeral session.
        - If use_cached is True: reuse/create a persistent session associated with session_id from the cache.
        Returns a list of resource dicts when possible, otherwise raw objects.
        """
        if not use_cached:
            async with self.session(server_name) as session:
                result = await session.list_resources()
        else:
            if not session_id:
                raise ValueError("session_id is required when use_cached=True for get_resources")
            session = await self.get_or_create_session(session_id)
            result = await session.list_resources()
        try:
            return [r.model_dump(mode="python") for r in result.resources]
        except Exception:
            return result.resources




    async def get_prompt(self, server_name: str, prompt: str, use_cached: bool = False, session_id: Optional[str] = None):
        """
        Get a prompt by name.
        - If use_cached is False (default): use an on-demand ephemeral session.
        - If use_cached is True: reuse/create a persistent session associated with session_id from the cache.
        Returns a dict when possible, otherwise the raw result object.
        """
        if not use_cached:
            async with self.session(server_name) as session:
                result = await session.get_prompt(prompt)
        else:
            if not session_id:
                raise ValueError("session_id is required when use_cached=True for get_prompt")
            session = await self.get_or_create_session(session_id)
            result = await session.get_prompt(prompt)
        try:
            return result.model_dump(mode="python")
        except Exception:
            return result




    async def ping(self, server_name: str) -> bool:
        """
        Send a lightweight ping request to validate connectivity/auth.
        Returns True if completed within the configured timeout.

        Note: Uses explicit request_read_timeout_seconds to avoid indefinite waits
        if a server streams or stalls.
        """
        async with self.session(server_name) as session:
            await session.send_request(
                types.ClientRequest(types.PingRequest()),
                types.EmptyResult,
                request_read_timeout_seconds=timedelta(seconds=float(self.request_timeout_seconds)),
            )
            return True




    async def test_connection(self, server_name: str) -> bool:
        """
        Validate basic MCP connectivity/auth by establishing a session.

        Some MCP servers do not implement ping, but a successful session
        initialization is still enough to verify that the endpoint is reachable
        and authentication is working.
        """
        async with self.session(server_name):
            return True





    async def call_tool(self, server_name: str, tool_name: str, arguments: dict | None = None, timeout_secs: float | None = None, use_cached: bool = False, session_id: Optional[str] = None):
        """
        Invoke a tool exposed by the MCP server.
        - If use_cached is False (default): use an on-demand ephemeral session.
        - If use_cached is True: reuse/create a persistent session associated with session_id from the cache.
        Returns the raw CallToolResult (caller is responsible for serialization).
        """
        from datetime import timedelta
        # Enforce a sensible default read timeout to avoid hangs when using cached sessions
        read_timeout = timedelta(seconds=float(timeout_secs)) if timeout_secs is not None else timedelta(seconds=float(self.request_timeout_seconds))

        if not use_cached:
            try:
                async with self.session(server_name) as session:
                    result = await session.call_tool(name=tool_name, arguments=arguments, read_timeout_seconds=read_timeout)
                    _MCP_TOOL_METRICS.increment_counter(_MCP_TOOL_SUCCESS_COUNTER)
                    return result
            except Exception as e:
                _MCP_TOOL_METRICS.increment_counter(_MCP_TOOL_FAILURE_COUNTER)
                logger.error(
                    "MCP call_tool failed (ephemeral session) for server '%s', tool '%s': %s",
                    server_name,
                    tool_name,
                    e,
                    exc_info=True,
                )
                return {"error": str(e)}

        if not session_id:
            raise ValueError("session_id is required when use_cached=True for call_tool.")
        session = await self.get_or_create_session(session_id)
        try:
            tool_output = await session.call_tool(name=tool_name, arguments=arguments, read_timeout_seconds=read_timeout)
            _MCP_TOOL_METRICS.increment_counter(_MCP_TOOL_SUCCESS_COUNTER)
            logger.info(f"Call tool returned {tool_output}")
            return tool_output
        except Exception as e:
            msg = str(e).lower()
            # If the cached session's transport was closed between calls, reset and retry once
            if isinstance(e, AnyioClosedResourceError) or "closed" in msg or "closed resource" in msg:
                try:
                    await self.stop_session(session_id)
                except Exception:
                    pass
                session = await self.get_or_create_session(session_id)
                try:
                    tool_output = await session.call_tool(name=tool_name, arguments=arguments, read_timeout_seconds=read_timeout)
                    _MCP_TOOL_METRICS.increment_counter(_MCP_TOOL_SUCCESS_COUNTER)
                    logger.info(f"Call tool returned {tool_output}")
                    return tool_output
                except Exception as retry_ex:
                    _MCP_TOOL_METRICS.increment_counter(_MCP_TOOL_FAILURE_COUNTER)
                    logger.error(
                        "MCP call_tool failed after cached-session reset for server '%s', session_id '%s', tool '%s': %s",
                        server_name,
                        session_id,
                        tool_name,
                        retry_ex,
                        exc_info=True,
                    )
                    return {"error": str(retry_ex)}

            _MCP_TOOL_METRICS.increment_counter(_MCP_TOOL_FAILURE_COUNTER)
            logger.error(
                "MCP call_tool failed (cached session) for server '%s', session_id '%s', tool '%s': %s",
                server_name,
                session_id,
                tool_name,
                e,
                exc_info=True,
            )
            return {"error": str(e)}




    async def get_or_create_session(self, session_id: str) -> ClientSession:
        """
        Return an existing session for session_id or lazily initialize a new one via init_session.
        """
        return await self.mcp_session_cache.get_or_create(self.server_name, session_id, self.init_session)