gobby 0.2.5__py3-none-any.whl

gobby/mcp_proxy/transports/base.py

@@ -0,0 +1,95 @@
+"""Base transport connection abstract class."""
+
+import asyncio
+from collections.abc import Callable, Coroutine
+from datetime import UTC, datetime
+from typing import Any
+
+from mcp import ClientSession
+
+from gobby.mcp_proxy.models import ConnectionState, MCPServerConfig
+
+
+class BaseTransportConnection:
+    """
+    Base class for MCP transport connections.
+
+    All transport implementations must provide:
+    - connect() -> ClientSession
+    - disconnect()
+    - is_connected property
+    - state property
+    """
+
+    def __init__(
+        self,
+        config: MCPServerConfig,
+        auth_token: str | None = None,
+        token_refresh_callback: Callable[[], Coroutine[Any, Any, str]] | None = None,
+    ):
+        """
+        Initialize transport connection.
+
+        Args:
+            config: Server configuration
+            auth_token: Optional auth token
+            token_refresh_callback: Optional callback for token refresh
+        """
+        self.config = config
+        self._auth_token = auth_token
+        self._token_refresh_callback = token_refresh_callback
+        self._session: Any | None = None  # ClientSession
+        self._transport_context: Any | None = None  # Transport-specific context manager
+        self._state = ConnectionState.DISCONNECTED
+        self._last_health_check: datetime | None = None
+        self._consecutive_failures = 0
+
+    async def connect(self) -> Any:
+        """Connect and return ClientSession. Must be implemented by subclasses."""
+        raise NotImplementedError
+
+    async def disconnect(self) -> None:
+        """Disconnect from server. Must be implemented by subclasses."""
+        raise NotImplementedError
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if connection is active."""
+        return self._state == ConnectionState.CONNECTED and self._session is not None
+
+    @property
+    def state(self) -> ConnectionState:
+        """Get current connection state."""
+        return self._state
+
+    @property
+    def session(self) -> ClientSession | None:
+        """Get the current client session, if connected."""
+        return self._session
+
+    def set_auth_token(self, token: str) -> None:
+        """Update authentication token."""
+        self._auth_token = token
+
+    async def health_check(self, timeout: float = 5.0) -> bool:
+        """
+        Check connection health.
+
+        Args:
+            timeout: Health check timeout in seconds
+
+        Returns:
+            True if healthy, False otherwise
+        """
+        if not self.is_connected or not self._session:
+            return False
+
+        try:
+            # Use asyncio.wait_for for timeout
+            await asyncio.wait_for(self._session.list_tools(), timeout)
+            self._last_health_check = datetime.now(UTC)
+            self._consecutive_failures = 0
+            return True
+        except (TimeoutError, Exception):
+            self._consecutive_failures += 1
+            return False

gobby/mcp_proxy/transports/factory.py

@@ -0,0 +1,44 @@
+"""Transport connection factory."""
+
+from collections.abc import Callable, Coroutine
+from typing import Any
+
+from gobby.mcp_proxy.models import MCPServerConfig
+from gobby.mcp_proxy.transports.base import BaseTransportConnection
+from gobby.mcp_proxy.transports.http import HTTPTransportConnection
+from gobby.mcp_proxy.transports.stdio import StdioTransportConnection
+from gobby.mcp_proxy.transports.websocket import WebSocketTransportConnection
+
+
+def create_transport_connection(
+    config: MCPServerConfig,
+    auth_token: str | None = None,
+    token_refresh_callback: Callable[[], Coroutine[Any, Any, str]] | None = None,
+) -> BaseTransportConnection:
+    """
+    Factory function to create appropriate transport connection.
+
+    Args:
+        config: Server configuration
+        auth_token: Optional auth token
+        token_refresh_callback: Optional token refresh callback
+
+    Returns:
+        Transport-specific connection instance
+
+    Raises:
+        ValueError: If transport type is unsupported
+    """
+    transport_map: dict[str, type[BaseTransportConnection]] = {
+        "http": HTTPTransportConnection,
+        "stdio": StdioTransportConnection,
+        "websocket": WebSocketTransportConnection,
+    }
+
+    transport_class = transport_map.get(config.transport)
+    if not transport_class:
+        raise ValueError(
+            f"Unsupported transport: {config.transport}. Supported: {list(transport_map.keys())}"
+        )
+
+    return transport_class(config, auth_token, token_refresh_callback)

gobby/mcp_proxy/transports/http.py

@@ -0,0 +1,139 @@
+"""HTTP transport connection."""
+
+import asyncio
+import logging
+from typing import Any
+
+from mcp import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+
+from gobby.mcp_proxy.models import ConnectionState, MCPError
+from gobby.mcp_proxy.transports.base import BaseTransportConnection
+
+logger = logging.getLogger("gobby.mcp.client")
+
+
+class HTTPTransportConnection(BaseTransportConnection):
+    """HTTP/Streamable HTTP transport connection using MCP SDK.
+
+    Uses a dedicated background task to own the streamablehttp_client lifecycle,
+    ensuring that context entry and exit happen in the same task (required by anyio).
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._owner_task: asyncio.Task[None] | None = None
+        self._disconnect_event: asyncio.Event | None = None
+        self._session_ready: asyncio.Event | None = None
+        self._connection_error: Exception | None = None
+        self._session_context: ClientSession | None = None
+
+    async def connect(self) -> Any:
+        """Connect via HTTP transport using a dedicated owner task."""
+        if self._state == ConnectionState.CONNECTED and self._session is not None:
+            return self._session
+
+        # Clean up old connection if reconnecting
+        if self._owner_task is not None:
+            await self.disconnect()
+
+        self._state = ConnectionState.CONNECTING
+        self._connection_error = None
+
+        # Create synchronization events
+        self._disconnect_event = asyncio.Event()
+        self._session_ready = asyncio.Event()
+
+        # Start owner task that manages the connection lifecycle
+        self._owner_task = asyncio.create_task(
+            self._run_connection(), name=f"http-conn-{self.config.name}"
+        )
+
+        # Wait for connection to be ready or fail
+        timeout = self.config.connect_timeout
+        try:
+            await asyncio.wait_for(self._session_ready.wait(), timeout=timeout)
+        except TimeoutError as e:
+            self._disconnect_event.set()
+            await self._cleanup_owner_task()
+            self._state = ConnectionState.FAILED
+            raise MCPError(f"Connection timeout for {self.config.name} after {timeout}s") from e
+
+        if self._connection_error is not None:
+            error = self._connection_error
+            self._connection_error = None
+            await self._cleanup_owner_task()
+            self._state = ConnectionState.FAILED
+            raise error
+
+        return self._session
+
+    async def _run_connection(self) -> None:
+        """Background task that owns the streamablehttp_client lifecycle."""
+        if self._disconnect_event is None or self._session_ready is None:
+            raise RuntimeError("Connection events not initialized")
+
+        try:
+            # URL is required for HTTP transport
+            if not self.config.url:
+                raise ValueError("URL is required for HTTP transport")
+
+            async with streamablehttp_client(
+                self.config.url,
+                headers=self.config.headers,
+            ) as (read_stream, write_stream, _):
+                self._session_context = ClientSession(read_stream, write_stream)
+                async with self._session_context as session:
+                    self._session = session
+                    await self._session.initialize()
+
+                    self._state = ConnectionState.CONNECTED
+                    self._consecutive_failures = 0
+                    logger.debug(f"Connected to HTTP MCP server: {self.config.name}")
+
+                    # Signal that connection is ready
+                    self._session_ready.set()
+
+                    # Wait until disconnect is requested
+                    await self._disconnect_event.wait()
+
+                    logger.debug(f"Disconnect requested for {self.config.name}")
+
+        except Exception as e:
+            error_msg = str(e) if str(e) else f"{type(e).__name__}: Connection closed or timed out"
+            logger.error(f"Failed to connect to HTTP server '{self.config.name}': {error_msg}")
+
+            if isinstance(e, MCPError):
+                self._connection_error = e
+            else:
+                self._connection_error = MCPError(f"HTTP connection failed: {error_msg}")
+
+            self._session_ready.set()  # Unblock waiter with error
+
+        finally:
+            self._session = None
+            self._session_context = None
+            self._state = ConnectionState.DISCONNECTED
+
+    async def _cleanup_owner_task(self) -> None:
+        """Clean up the owner task."""
+        if self._owner_task is not None:
+            if not self._owner_task.done():
+                self._owner_task.cancel()
+                try:
+                    await asyncio.wait_for(self._owner_task, timeout=2.0)
+                except asyncio.CancelledError:
+                    logger.debug(f"Owner task cancelled for {self.config.name}")
+                except TimeoutError:
+                    logger.warning(f"Owner task cleanup timed out for {self.config.name}")
+            self._owner_task = None
+        self._disconnect_event = None
+        self._session_ready = None
+
+    async def disconnect(self) -> None:
+        """Disconnect from HTTP server by signaling the owner task."""
+        if self._disconnect_event is not None:
+            self._disconnect_event.set()
+
+        await self._cleanup_owner_task()
+        self._state = ConnectionState.DISCONNECTED

gobby/mcp_proxy/transports/stdio.py

@@ -0,0 +1,213 @@
+"""Stdio transport connection."""
+
+import asyncio
+import logging
+import os
+import re
+from collections.abc import Callable, Coroutine
+from typing import TYPE_CHECKING, Any
+
+from mcp import ClientSession
+from mcp.client.stdio import StdioServerParameters, stdio_client
+
+from gobby.mcp_proxy.models import ConnectionState, MCPError
+from gobby.mcp_proxy.transports.base import BaseTransportConnection
+
+# Pattern for ${VAR} or ${VAR:-default} environment variable expansion
+ENV_VAR_PATTERN = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-([^}]*))?\}")
+
+
+def _expand_env_var(value: str) -> str:
+    """Expand ${VAR} and ${VAR:-default} patterns in a string.
+
+    Args:
+        value: String that may contain ${VAR} patterns
+
+    Returns:
+        String with environment variables expanded
+    """
+
+    def replace_match(match: re.Match[str]) -> str:
+        var_name = match.group(1)
+        default_value = match.group(2)  # None if no default specified
+
+        env_value = os.environ.get(var_name)
+
+        if env_value is not None and env_value != "":
+            return env_value
+        elif default_value is not None:
+            return default_value
+        else:
+            # Leave unchanged if no value and no default
+            return match.group(0)
+
+    return ENV_VAR_PATTERN.sub(replace_match, value)
+
+
+def _expand_env_dict(env: dict[str, str] | None) -> dict[str, str] | None:
+    """Expand environment variables in env dict values.
+
+    Args:
+        env: Dictionary of environment variables (may contain ${VAR} patterns)
+
+    Returns:
+        Dictionary with expanded values, or None if input is None
+    """
+    if env is None:
+        return None
+
+    return {key: _expand_env_var(value) for key, value in env.items()}
+
+
+def _expand_args(args: list[str] | None) -> list[str] | None:
+    """Expand environment variables in command args.
+
+    Args:
+        args: List of command arguments (may contain ${VAR} patterns)
+
+    Returns:
+        List with expanded values, or None if input is None
+    """
+    if args is None:
+        return None
+
+    return [_expand_env_var(arg) for arg in args]
+
+
+if TYPE_CHECKING:
+    from gobby.mcp_proxy.models import MCPServerConfig
+
+logger = logging.getLogger("gobby.mcp.client")
+
+
+class StdioTransportConnection(BaseTransportConnection):
+    """Stdio transport connection using MCP SDK."""
+
+    def __init__(
+        self,
+        config: "MCPServerConfig",
+        auth_token: str | None = None,
+        token_refresh_callback: Callable[[], Coroutine[Any, Any, str]] | None = None,
+    ) -> None:
+        """Initialize stdio transport connection."""
+        super().__init__(config, auth_token, token_refresh_callback)
+        self._session_context: ClientSession | None = None
+        # Explicitly initialize transport context (inherited from base class, but
+        # ensures the attribute exists with proper type annotation for this transport)
+        self._transport_context: Any | None = None
+
+    async def connect(self) -> Any:
+        """Connect via stdio transport."""
+        if self._state == ConnectionState.CONNECTED:
+            return self._session
+
+        self._state = ConnectionState.CONNECTING
+
+        # Track what was entered for cleanup
+        transport_entered = False
+        session_entered = False
+
+        try:
+            # Create stdio server parameters
+            if self.config.command is None:
+                raise RuntimeError("Command is required for stdio transport")
+
+            # Expand ${VAR} patterns in args and env values
+            expanded_args = _expand_args(self.config.args) or []
+            expanded_env = _expand_env_dict(self.config.env)
+
+            params = StdioServerParameters(
+                command=self.config.command,
+                args=expanded_args,
+                env=expanded_env,
+            )
+
+            # Create stdio client context
+            self._transport_context = stdio_client(params)
+
+            # Enter the transport context to get streams
+            read_stream, write_stream = await self._transport_context.__aenter__()
+            transport_entered = True
+
+            # Save the context manager itself so we can call __aexit__ on it later
+            self._session_context = ClientSession(read_stream, write_stream)
+            self._session = await self._session_context.__aenter__()
+            session_entered = True
+
+            await self._session.initialize()
+
+            self._state = ConnectionState.CONNECTED
+            self._consecutive_failures = 0
+            logger.debug(f"Connected to stdio MCP server: {self.config.name}")
+
+            return self._session
+
+        except Exception as e:
+            # Handle exceptions with empty str() (EndOfStream, ClosedResourceError, CancelledError)
+            error_msg = str(e) if str(e) else f"{type(e).__name__}: Connection closed or timed out"
+            logger.error(f"Failed to connect to stdio server '{self.config.name}': {error_msg}")
+
+            # Cleanup in reverse order - session first, then transport
+            # Cleanup in reverse order - session first, then transport
+            session_ctx = self._session_context
+            if session_entered and session_ctx is not None:
+                try:
+                    await session_ctx.__aexit__(None, None, None)
+                except Exception as cleanup_error:
+                    logger.warning(
+                        f"Error during session cleanup for {self.config.name}: {cleanup_error}"
+                    )
+
+            transport_ctx = self._transport_context
+            if transport_entered and transport_ctx is not None:
+                try:
+                    await transport_ctx.__aexit__(None, None, None)
+                except Exception as cleanup_error:
+                    logger.warning(
+                        f"Error during transport cleanup for {self.config.name}: {cleanup_error}"
+                    )
+
+            # Reset state before raising
+            self._session = None
+            self._session_context = None
+            self._transport_context = None
+            self._state = ConnectionState.FAILED
+
+            # Re-raise wrapped in MCPError (don't double-wrap)
+            if isinstance(e, MCPError):
+                raise
+            raise MCPError(f"Stdio connection failed: {error_msg}") from e
+
+    async def disconnect(self) -> None:
+        """Disconnect from stdio server."""
+        # Exit session context manager (not the session object itself)
+        session_ctx = self._session_context
+        if session_ctx is not None:
+            try:
+                await asyncio.wait_for(session_ctx.__aexit__(None, None, None), timeout=2.0)
+            except TimeoutError:
+                logger.warning(f"Session close timed out for {self.config.name}")
+            except RuntimeError as e:
+                # Expected when exiting cancel scope from different task
+                if "cancel scope" not in str(e):
+                    logger.warning(f"Error closing session for {self.config.name}: {e}")
+            except Exception as e:
+                logger.warning(f"Error closing session for {self.config.name}: {e}")
+            self._session_context = None
+            self._session = None
+
+        transport_ctx = self._transport_context
+        if transport_ctx is not None:
+            try:
+                await asyncio.wait_for(transport_ctx.__aexit__(None, None, None), timeout=2.0)
+            except TimeoutError:
+                logger.warning(f"Transport close timed out for {self.config.name}")
+            except RuntimeError as e:
+                # Expected when exiting cancel scope from different task
+                if "cancel scope" not in str(e):
+                    logger.warning(f"Error closing transport for {self.config.name}: {e}")
+            except Exception as e:
+                logger.warning(f"Error closing transport for {self.config.name}: {e}")
+            self._transport_context = None
+
+        self._state = ConnectionState.DISCONNECTED

gobby/mcp_proxy/transports/websocket.py

@@ -0,0 +1,136 @@
+"""WebSocket transport connection."""
+
+import asyncio
+import logging
+from collections.abc import Callable, Coroutine
+from typing import TYPE_CHECKING, Any
+
+from mcp import ClientSession
+from mcp.client.websocket import websocket_client
+
+from gobby.mcp_proxy.models import ConnectionState, MCPError
+from gobby.mcp_proxy.transports.base import BaseTransportConnection
+
+if TYPE_CHECKING:
+    from gobby.config.mcp import MCPServerConfig
+
+logger = logging.getLogger("gobby.mcp.client")
+
+
+class WebSocketTransportConnection(BaseTransportConnection):
+    """WebSocket transport connection using MCP SDK."""
+
+    def __init__(
+        self,
+        config: "MCPServerConfig",
+        auth_token: str | None = None,
+        token_refresh_callback: Callable[[], Coroutine[Any, Any, str]] | None = None,
+    ) -> None:
+        """Initialize WebSocket transport connection."""
+        super().__init__(config, auth_token, token_refresh_callback)
+        self._session_context: ClientSession | None = None
+
+    async def connect(self) -> Any:
+        """Connect via WebSocket transport."""
+        if self._state == ConnectionState.CONNECTED:
+            return self._session
+
+        self._state = ConnectionState.CONNECTING
+
+        # Track what was entered for cleanup
+        transport_entered = False
+        session_entered = False
+
+        try:
+            # URL is required for WebSocket transport
+            if self.config.url is None:
+                raise RuntimeError("URL is required for WebSocket transport")
+
+            # Create WebSocket client context
+            self._transport_context = websocket_client(self.config.url)
+
+            # Enter the transport context to get streams
+            read_stream, write_stream = await self._transport_context.__aenter__()
+            transport_entered = True
+
+            # Save the context manager itself so we can call __aexit__ on it later
+            self._session_context = ClientSession(read_stream, write_stream)
+            self._session = await self._session_context.__aenter__()
+            session_entered = True
+
+            await self._session.initialize()
+
+            self._state = ConnectionState.CONNECTED
+            self._consecutive_failures = 0
+            logger.debug(f"Connected to WebSocket MCP server: {self.config.name}")
+
+            return self._session
+
+        except Exception as e:
+            # Handle exceptions with empty str() (EndOfStream, ClosedResourceError, CancelledError)
+            error_msg = str(e) if str(e) else f"{type(e).__name__}: Connection closed or timed out"
+            logger.error(f"Failed to connect to WebSocket server '{self.config.name}': {error_msg}")
+
+            # Cleanup in reverse order - session first, then transport
+            if session_entered and self._session_context is not None:
+                try:
+                    await self._session_context.__aexit__(None, None, None)
+                except Exception as cleanup_error:
+                    logger.warning(
+                        f"Error during session cleanup for {self.config.name}: {cleanup_error}"
+                    )
+
+            if transport_entered and self._transport_context is not None:
+                try:
+                    await self._transport_context.__aexit__(None, None, None)
+                except Exception as cleanup_error:
+                    logger.warning(
+                        f"Error during transport cleanup for {self.config.name}: {cleanup_error}"
+                    )
+
+            # Reset state before raising
+            self._session = None
+            self._session_context = None
+            self._transport_context = None
+            self._state = ConnectionState.FAILED
+
+            # Re-raise wrapped in MCPError (don't double-wrap)
+            if isinstance(e, MCPError):
+                raise
+            raise MCPError(f"WebSocket connection failed: {error_msg}") from e
+
+    async def disconnect(self) -> None:
+        """Disconnect from WebSocket server."""
+        # Exit session context manager (not the session object itself)
+        if self._session_context is not None:
+            try:
+                await asyncio.wait_for(
+                    self._session_context.__aexit__(None, None, None), timeout=2.0
+                )
+            except TimeoutError:
+                logger.warning(f"Session close timed out for {self.config.name}")
+            except RuntimeError as e:
+                # Expected when exiting cancel scope from different task
+                if "cancel scope" not in str(e):
+                    logger.warning(f"Error closing session for {self.config.name}: {e}")
+            except Exception as e:
+                logger.warning(f"Error closing session for {self.config.name}: {e}")
+            self._session_context = None
+            self._session = None
+
+        if self._transport_context is not None:
+            try:
+                await asyncio.wait_for(
+                    self._transport_context.__aexit__(None, None, None), timeout=2.0
+                )
+            except TimeoutError:
+                logger.warning(f"Transport close timed out for {self.config.name}")
+            except RuntimeError as e:
+                # Expected when exiting cancel scope from different task
+                if "cancel scope" not in str(e):
+                    logger.warning(f"Error closing transport for {self.config.name}: {e}")
+            except Exception as e:
+                logger.warning(f"Error closing transport for {self.config.name}: {e}")
+            self._transport_context = None
+
+        self._state = ConnectionState.DISCONNECTED