diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..26ce978 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,80 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +permissions: + contents: read + +jobs: + lint: + name: ruff + mypy + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install package with dev dependencies + run: | + uv venv + uv pip install -e ".[dev]" + + - name: ruff check + run: uv run ruff check src tests + + - name: mypy + run: uv run mypy src/mcp_reverse_proxy + + test: + name: pytest (Python ${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install package with dev dependencies + run: | + uv venv + uv pip install -e ".[dev]" + + - name: Run tests + run: uv run pytest tests/ -q + + build: + name: Build wheel + sdist + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + + - name: Build distributions + run: uv build + + - name: Upload distributions + uses: actions/upload-artifact@v4 + with: + name: dist + path: dist/ + retention-days: 7 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..6bb5841 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,45 @@ +name: Release + +on: + push: + tags: + - "v*" + +permissions: + contents: read + +jobs: + build: + name: Build distributions + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + + - name: Build distributions + run: uv build + + - name: Upload distributions + uses: actions/upload-artifact@v4 + with: + name: dist + path: dist/ + + publish: + name: Publish to PyPI (trusted publishing) + needs: build + runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write # OIDC token for PyPI trusted publishing (no API token) + steps: + - name: Download distributions + uses: actions/download-artifact@v4 + with: + name: dist + path: dist/ + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/README.md b/README.md index 93d953b..e08aa50 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,7 @@ A gateway-side integration package, `mcp-reverse-proxy-contextforge`, is planned ## Documentation - [docs/README.md](docs/README.md): architecture, transport abstraction, and health-check strategy -- [docs/USER_README.md](docs/USER_README.md): user guide with installation, TLS configuration, and troubleshooting +- [USER_README.md](USER_README.md): user guide with installation, TLS configuration, and troubleshooting ## Development diff --git a/USER_README.md b/USER_README.md new file mode 100644 index 0000000..57a07b1 --- /dev/null +++ b/USER_README.md @@ -0,0 +1,676 @@ +# MCP Reverse Proxy + +Bridge MCP servers to remote gateways with multi-transport support. + +## Overview + +The MCP Reverse Proxy is a standalone component that enables you to connect local or remote MCP servers to a ContextForge gateway (or any compatible MCP gateway). It supports multiple transport protocols for both MCP server connections and gateway connections. + +### Key Features + +- **Multi-Transport Support**: Connect to MCP servers via stdio, HTTP/2 streaming, or SSE +- **WebSocket Gateway**: Persistent connection to remote gateways +- **Automatic Reconnection**: Built-in retry logic with exponential backoff +- **Health Monitoring**: Active health checks with automatic recovery +- **Flexible Configuration**: CLI arguments, environment variables, or config files +- **Production Ready**: Comprehensive error handling and logging + +## Installation + +### From PyPI (when published) + +```bash +pip install mcp-reverse-proxy +``` + +### From Source + +```bash +# Clone the repository +git clone https://github.com/your-org/contextforge.git +cd contextforge/mcp_reverse_proxy + +# Clean any old installations +rm -rf .venv mcp_reverse_proxy.egg-info __pycache__ + +# Create fresh virtual environment +python3 -m venv .venv +source .venv/bin/activate # On Windows: .venv\Scripts\activate + +# Install the package (REQUIRED to get the mcp-reverse-proxy command) +pip install -e . +``` + +**Important Notes**: +- The package uses a **src/ layout** - all Python modules are in `src/mcp_reverse_proxy/` +- You **must** run `pip install -e .` to install the package and create the `mcp-reverse-proxy` command +- If you get `ModuleNotFoundError`, clean old installations: `rm -rf .venv mcp_reverse_proxy.egg-info src/mcp_reverse_proxy.egg-info` then reinstall +- The `.venv` directory name (with dot prefix) avoids Python import conflicts +- Using a virtual environment is strongly recommended to avoid conflicts with system packages +- Python 3.11+ is required (Python 3.14 is supported) + +**Alternative: Run without installing** +```bash +# If you don't want to install, you can run directly as a module +python -m mcp_reverse_proxy.cli --config config.json +``` + +**Troubleshooting** + +If you see `ModuleNotFoundError: No module named 'mcp_reverse_proxy'`: + +```bash +# 1. Clean old installations +cd mcp_reverse_proxy +rm -rf .venv mcp_reverse_proxy.egg-info __pycache__ transports/__pycache__ + +# 2. Recreate venv and reinstall +python3 -m venv .venv +source .venv/bin/activate +pip install -e . + +# 3. Verify installation +which mcp-reverse-proxy # Should show path in .venv/bin/ +mcp-reverse-proxy --help # Should show help message +``` + +If you see `[SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed: self-signed certificate`: + +```bash +# Option 1: Provide CA certificate file path (recommended for production) +# The --cert parameter is used for BOTH MCP server HTTPS/WSS and gateway WSS connections +mcp-reverse-proxy \ + --local-streamable-http https://mcp-server.example.com/mcp \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token YOUR_TOKEN \ + --cert /path/to/ca-certificate.pem + +# Option 2: Provide certificate inline in config file (JSON or YAML) +# The cert field accepts the full certificate chain as a string +{ + "gateway": "wss://gateway.example.com/reverse-proxy", + "token": "YOUR_TOKEN", + "cert": "-----BEGIN CERTIFICATE-----\n...\n-----END CERTIFICATE-----\n" +} + +# Option 3: For development/testing only - disable SSL verification +# WARNING: Only use in trusted development environments +export PYTHONHTTPSVERIFY=0 +mcp-reverse-proxy --local-streamable-http http://localhost:9020/mcp --gateway wss://gateway.example.com/reverse-proxy --token YOUR_TOKEN +``` + +**Certificate Format Notes**: +- The `--cert` option accepts either a file path OR inline certificate data +- When using inline certificates in config files, include the complete certificate chain +- Certificates must be in PEM format with proper `\n` line breaks +- For certificate chains, concatenate all certificates (root CA, intermediate CAs, server cert) +- The code uses Python's `ssl.create_default_context(cadata=cert)` which validates the chain + +**Troubleshooting Certificate Issues**: + +If you see `[SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed: self-signed certificate in certificate chain`: + +**The package now automatically loads system CAs as fallback** when you provide a custom certificate. This should resolve most certificate chain issues. Simply ensure your config has the `cert` field with the full certificate chain. + +1. **Verify Your Setup (Recommended First Step)**: + ```bash + # Your config.json should have the cert field with full chain + # The package will now use both your custom cert AND system CAs + mcp-reverse-proxy --config config.json + ``` + +2. **If Still Failing - Disable SSL Verification (Development Only)**: + ```bash + # This won't work with the current code - see option 3 instead + export PYTHONHTTPSVERIFY=0 + mcp-reverse-proxy --config config.json + ``` + +2. **Proper Fix - Use System CA Bundle**: + ```bash + # Extract cert from config to file + python3 mcp_reverse_proxy/fix_ssl_cert.py config.json + + # This creates ca-bundle.pem - now use it + mcp-reverse-proxy --cert ca-bundle.pem --config config.json + ``` + +3. **Alternative - Load System Certificates**: + The issue is that Python's SSL context doesn't trust the self-signed root CA. You can: + - Add the root CA to your system's trust store (macOS: Keychain Access, Linux: `/etc/ssl/certs/`) + - Or use `--cert` with the full certificate chain from your config + +**Common Issues**: +- Certificate chain incomplete (missing root or intermediate CAs) +- Line breaks not properly escaped as `\n` in JSON +- Certificate expired or not yet valid +- Hostname mismatch between cert and gateway URL + +**Verify Certificate**: +```bash +# Extract and test certificate +python3 mcp_reverse_proxy/fix_ssl_cert.py config.json + +# Verify with OpenSSL +openssl verify -CAfile ca-bundle.pem ca-bundle.pem +``` + +### Development Installation + +```bash +# Install with development dependencies +pip install -e ".[dev]" +``` + +## HTTPS/SSL Support + +The reverse proxy **fully supports HTTPS** for both MCP server connections and gateway connections. + +### Supported Protocols + +- **MCP Server Connections**: `http://`, `https://` (for SSE and Streamable HTTP transports) +- **Gateway Connections**: `ws://`, `wss://` (WebSocket) + +### Certificate Configuration + +You have three options for configuring SSL certificates: + +#### Option 1: Single Certificate (Simple) +Use `--cert` for both MCP server and gateway connections when they share the same CA: + +```bash +mcp-reverse-proxy \ + --local-streamable-http https://mcp-server.example.com/mcp \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token YOUR_TOKEN \ + --cert /path/to/shared-ca.pem +``` + +#### Option 2: Separate Certificates (Recommended for Different CAs) +Use `--mcp-cert` and `--gateway-cert` when MCP server and gateway use different CAs: + +```bash +mcp-reverse-proxy \ + --local-streamable-http https://mcp-server.example.com/mcp \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token YOUR_TOKEN \ + --mcp-cert /path/to/mcp-ca.pem \ + --gateway-cert /path/to/gateway-ca.pem +``` + +#### Option 3: Mixed Configuration +Use `--cert` as fallback with specific override: + +```bash +# Use shared cert for gateway, but specific cert for MCP server +mcp-reverse-proxy \ + --local-sse https://mcp-server.example.com/sse \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token YOUR_TOKEN \ + --cert /path/to/gateway-ca.pem \ + --mcp-cert /path/to/mcp-ca.pem +``` + +### Certificate Precedence + +- `--mcp-cert` takes precedence over `--cert` for MCP server connections +- `--gateway-cert` takes precedence over `--cert` for gateway connections +- If only `--cert` is provided, it's used for both connections (backward compatible) + +### Self-Signed Certificates + +When no certificate is provided for a connection, SSL verification is **disabled** (insecure, for development only): +- Hostname verification: disabled +- Certificate verification: disabled + +**⚠️ Warning**: Only use this in trusted development environments. Always provide CA certificates for production deployments. + +### Certificate Bundle Format + +For multiple CAs in a single file, concatenate certificates in PEM format: + +```bash +cat root-ca.pem intermediate-ca.pem > ca-bundle.pem +mcp-reverse-proxy --cert ca-bundle.pem ... +``` + +## Quick Start + +### Basic Usage - Stdio Transport + +Connect a local MCP server (running via stdio) to a remote gateway: + +```bash +mcp-reverse-proxy \ + --local-stdio "uvx mcp-server-git" \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token YOUR_AUTH_TOKEN +``` + +### HTTP/2 Streaming Transport + +Connect to an MCP server running on HTTP: + +```bash +mcp-reverse-proxy \ + --local-streamable-http http://localhost:8000/mcp \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token YOUR_AUTH_TOKEN +``` + +Connect to an MCP server running on HTTPS: + +```bash +mcp-reverse-proxy \ + --local-streamable-http https://mcp-server.example.com/mcp \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token YOUR_AUTH_TOKEN \ + --cert /path/to/ca-certificate.pem +``` + +### SSE Transport + +Connect to an MCP server using Server-Sent Events (HTTP): + +```bash +mcp-reverse-proxy \ + --local-sse http://localhost:9020/sse \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token YOUR_AUTH_TOKEN +``` + +Connect to an MCP server using Server-Sent Events (HTTPS): + +```bash +mcp-reverse-proxy \ + --local-sse https://mcp-server.example.com/sse \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token YOUR_AUTH_TOKEN \ + --cert /path/to/ca-certificate.pem +``` + +## Configuration + +### Environment Variables + +Set these environment variables to avoid passing them on the command line: + +```bash +export REVERSE_PROXY_GATEWAY="wss://gateway.example.com/reverse-proxy" +export REVERSE_PROXY_TOKEN="your-auth-token" +``` + +Then run: + +```bash +mcp-reverse-proxy --local-stdio "uvx mcp-server-git" +``` + +### Configuration File + +Create a YAML or JSON configuration file: + +**config.yaml:** +```yaml +local-stdio: "uvx mcp-server-git" +gateway: "wss://gateway.example.com/reverse-proxy" +token: "your-auth-token" +server-name: "My Git Server" +server-description: "Git repository access" +keepalive: 30 +reconnect-delay: 2.0 +max-retries: 0 +log-level: "INFO" +``` + +**config.json:** +```json +{ + "local-stdio": "uvx mcp-server-git", + "gateway": "wss://gateway.example.com/reverse-proxy", + "token": "your-auth-token", + "server-name": "My Git Server", + "server-description": "Git repository access", + "keepalive": 30, + "reconnect-delay": 2.0, + "max-retries": 0, + "log-level": "INFO" +} +``` + +Use the config file: + +```bash +# With YAML config +mcp-reverse-proxy --config config.yaml + +# With JSON config +mcp-reverse-proxy --config config.json +``` + +**Tip**: Configuration files are useful for: +- Managing multiple proxy instances with different settings +- Version controlling your proxy configurations +- Avoiding long command lines with many arguments + +## Command Line Options + +### MCP Server Transport + +Choose **one** of these options: + +- `--local-stdio COMMAND` - Run MCP server as subprocess (e.g., `"uvx mcp-server-git"`) +- `--local-streamable-http URL` - Connect to HTTP/2 streaming endpoint (e.g., `http://localhost:8000/mcp`) +- `--local-sse URL` - Connect to SSE endpoint (e.g., `http://localhost:9020/sse`) + +### Gateway Connection + +- `--gateway URL` - Gateway WebSocket URL (or use `REVERSE_PROXY_GATEWAY` env var) +- `--token TOKEN` - Bearer token for authentication (or use `REVERSE_PROXY_TOKEN` env var) +- `--server-id ID` - Session identifier (auto-generated if not provided) +- `--server-name NAME` - Server name for registration +- `--server-description DESC` - Server description for registration +- `--cert PATH` - CA certificate for SSL verification (used for both MCP and gateway if specific certs not provided) +- `--mcp-cert PATH` - CA certificate specifically for MCP server HTTPS connections (overrides --cert) +- `--gateway-cert PATH` - CA certificate specifically for gateway WSS connections (overrides --cert) + +### Connection Options + +- `--reconnect-delay SECONDS` - Initial reconnection delay (default: 1.0) +- `--max-retries COUNT` - Maximum reconnection attempts, 0=infinite (default: 0) +- `--keepalive SECONDS` - Heartbeat interval (default: 2) +- `--mcp-health-check-timeout SECONDS` - MCP health check timeout (default: 5.0) +- `--mcp-health-check-retry-interval SECONDS` - Retry interval during MCP outage (default: 10.0) + +### Logging + +- `--log-level LEVEL` - Log level: DEBUG, INFO, WARNING, ERROR, CRITICAL (default: INFO) +- `--verbose` - Enable verbose logging (same as `--log-level DEBUG`) + +Set `LOG_FORMAT=json` environment variable for JSON-formatted logs. + +### Configuration File + +- `--config PATH` - Load configuration from YAML or JSON file + +## Health Monitoring + +The reverse proxy implements a two-layer health monitoring system: + +### Layer 1: MCP Server Health Checks + +Before sending heartbeats to the gateway, the proxy actively checks if the MCP server is healthy by sending a `tools/list` request: + +- ✅ **MCP responds** → Proxy sends heartbeat to gateway +- ❌ **MCP timeout** → Proxy skips heartbeat (gateway detects missing heartbeat) +- 🔄 **During outage** → Proxy continues probing with shorter retry interval +- ✅ **MCP recovers** → Proxy automatically reconnects to gateway + +### Layer 2: Gateway Health Monitoring + +The gateway tracks reverse proxy sessions through heartbeat monitoring: + +- Expects heartbeats every 30 seconds (configurable) +- Marks session as stale after 90 seconds without heartbeat (configurable) +- After 3 consecutive failures, marks server as unreachable (configurable) +- Automatically recovers when heartbeats resume + +## Examples + +### Example 1: Local Git Server + +```bash +mcp-reverse-proxy \ + --local-stdio "uvx mcp-server-git" \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token $TOKEN \ + --server-name "Git Server" \ + --server-description "Access to local git repositories" +``` + +### Example 2: Remote HTTPS Server with Shared Certificate + +```bash +# Single certificate for both MCP server and gateway +mcp-reverse-proxy \ + --local-streamable-http https://mcp-server.internal:8000/mcp \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token $TOKEN \ + --cert /path/to/shared-ca.pem \ + --keepalive 30 \ + --mcp-health-check-timeout 10.0 \ + --mcp-health-check-retry-interval 15.0 +``` + +### Example 3: SSE Server with Separate Certificates + +```bash +# Different certificates for MCP server and gateway +mcp-reverse-proxy \ + --local-sse https://mcp-server.internal:9020/sse \ + --gateway wss://gateway.example.com/reverse-proxy \ + --token $TOKEN \ + --mcp-cert /path/to/mcp-ca.pem \ + --gateway-cert /path/to/gateway-ca.pem +``` + +### Example 4: Using Configuration File + +```bash +# Create config file +cat > reverse-proxy-config.yaml < MCP server responds ✓ │ +│ │ +│ 2. Client sends heartbeat to gateway │ +│ └─> Gateway receives heartbeat ✓ │ +│ │ +│ 3. Gateway marks session as healthy │ +│ └─> Database: reachable = true │ +│ │ +└─────────────────────────────────────────────────────────────────┘ + +┌─────────────────────────────────────────────────────────────────┐ +│ MCP Server Failure │ +├─────────────────────────────────────────────────────────────────┤ +│ │ +│ 1. Client checks MCP server health (tools/list) │ +│ └─> MCP server timeout ✗ │ +│ │ +│ 2. Client SKIPS heartbeat to gateway │ +│ └─> Gateway detects missing heartbeat │ +│ │ +│ 3. Gateway increments failure counter │ +│ └─> After 3 failures: reachable = false │ +│ │ +│ 4. Client continues probing MCP server (10s interval) │ +│ └─> Waiting for recovery... │ +│ │ +└─────────────────────────────────────────────────────────────────┘ + +┌─────────────────────────────────────────────────────────────────┐ +│ Automatic Recovery │ +├─────────────────────────────────────────────────────────────────┤ +│ │ +│ 1. MCP server comes back online │ +│ └─> Client detects via health check ✓ │ +│ │ +│ 2. Client reconnects to gateway │ +│ └─> Sends registration + heartbeat │ +│ │ +│ 3. Gateway marks session as healthy │ +│ └─> Database: reachable = true │ +│ │ +│ 4. Normal operation resumes │ +│ │ +└─────────────────────────────────────────────────────────────────┘ +``` + +### Benefits + +✅ **Fail-Fast Detection**: Gateway quickly detects MCP server outages via missing heartbeats +✅ **Self-Healing**: Automatic recovery when MCP server comes back online +✅ **No False Positives**: Only sends heartbeats when MCP server is actually healthy +✅ **Configurable**: All timeouts and intervals can be tuned for your environment +✅ **Production-Ready**: Comprehensive error handling and logging + +### Configuration Reference + +**Client-Side (Reverse Proxy):** +- `--keepalive`: Heartbeat interval in seconds (default: 30) +- `--mcp-health-check-timeout`: MCP health check timeout in seconds (default: 5.0) +- `--mcp-health-check-retry-interval`: Retry interval during MCP outage in seconds (default: 10.0) + +**Server-Side (Gateway):** +- `MCPGATEWAY_REVERSE_PROXY_HEARTBEAT_TIMEOUT`: Max time without heartbeat before marking stale (default: 90s) +- `MCPGATEWAY_REVERSE_PROXY_HEALTH_CHECK_INTERVAL`: How often to check session health (default: 60s) +- `MCPGATEWAY_REVERSE_PROXY_FAILURE_THRESHOLD`: Consecutive failures before marking unreachable (default: 3) + +## Key Features + +### SOLID Principles + +- **Single Responsibility**: Each transport handles one protocol +- **Open/Closed**: New transports can be added without modifying existing code +- **Liskov Substitution**: All transports are interchangeable via interfaces +- **Interface Segregation**: Separate interfaces for MCP server vs gateway +- **Dependency Inversion**: Client depends on abstractions, not implementations + +### Extensibility + +Adding a new MCP server transport: + +1. Create adapter class inheriting from `McpServerTransport` +2. Implement required methods: `start()`, `stop()`, `send()`, `add_message_handler()` +3. Add to transport factory in `cli.py` + +### Configuration + +- **Environment variables**: `REVERSE_PROXY_GATEWAY`, `REVERSE_PROXY_TOKEN` +- **CLI arguments**: See `--help` for full list +- **Automatic reconnection**: Configurable backoff and retry limits + +## Transport Details + +### Stdio Adapter + +- Spawns subprocess with stdin/stdout pipes +- Line-based JSON-RPC message protocol +- Automatic process cleanup on shutdown + +### Streamable HTTP Adapter + +- HTTP/2 streaming for bidirectional communication +- Inline responses for request/response pattern +- POST for sending messages to server +- Supports both HTTP and HTTPS +- Optional SSL certificate verification + +### SSE Adapter + +- Server-Sent Events (SSE) for server-to-client streaming +- HTTP POST for client-to-server messages +- Automatic endpoint discovery from SSE stream +- Session management via headers +- Supports both HTTP and HTTPS +- Optional SSL certificate verification + +### WebSocket Adapter + +- Persistent WebSocket connection to gateway +- Automatic ping/pong keepalive +- SSL/TLS support with optional certificate verification +- Graceful reconnection on disconnect + +## Message Flow + +``` +MCP Server → McpServerTransport → ReverseProxyClient → GatewayTransport → Gateway + ↓ + Message Routing + - Registration + - Heartbeat + - Request/Response + - Notifications +``` + +## Error Handling + +- Transport-level errors trigger reconnection +- Configurable retry limits and backoff delays +- Graceful shutdown on SIGINT/SIGTERM +- Automatic cleanup of resources + +## Future Enhancements + +- Additional gateway transports (HTTP, gRPC) +- Metrics and observability +- Connection pooling +- Load balancing across multiple MCP servers \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..820a9f8 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,129 @@ +# ---------------------------------------------------------------- +# 💡 Build system (PEP 517) +# ---------------------------------------------------------------- +[build-system] +requires = ["setuptools>=78.1.1", "wheel>=0.46.2"] +build-backend = "setuptools.build_meta" + +# ---------------------------------------------------------------- +# 📦 Core project metadata (PEP 621) +# ---------------------------------------------------------------- +[project] +name = "mcp-reverse-proxy" +version = "1.0.0" +description = "MCP Reverse Proxy - Bridge MCP servers to remote gateways with multi-transport support" +keywords = [ + "MCP", + "reverse-proxy", + "model-context-protocol", + "gateway", + "stdio", + "websocket", + "sse", + "http", + "streaming" +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Framework :: AsyncIO", + "Topic :: Software Development :: Libraries :: Application Frameworks", + "Topic :: Internet :: WWW/HTTP" +] +readme = "USER_README.md" +requires-python = ">=3.11" + +# SPDX licence expression + explicit licence file (PEP 639) +license = "Apache-2.0" + +# Maintainers +maintainers = [ + {name = "Mihai Criveti", email = "redacted@ibm.com"} +] + +# ---------------------------------------------------------------- +# Runtime dependencies +# ---------------------------------------------------------------- +dependencies = [ + "httpx>=0.28.1", + "httpx[http2]>=0.28.1", + "orjson>=3.11.9", + "python-json-logger>=3.2.1", + "websockets>=14.1", + "pyyaml>=6.0.2" +] + +# ---------------------------------------------------------------- +# Optional dependencies +# ---------------------------------------------------------------- +[project.optional-dependencies] +dev = [ + "pytest>=8.3.4", + "pytest-asyncio>=0.25.2", + "pytest-cov>=6.0.0", + "types-PyYAML>=6.0.12", + "mypy>=1.14.1", + "ruff>=0.9.1", + "black>=25.1.0" +] + +# ---------------------------------------------------------------- +# Console scripts +# ---------------------------------------------------------------- +[project.scripts] +mcp-reverse-proxy = "mcp_reverse_proxy.cli:run" + +# ---------------------------------------------------------------- +# Package discovery +# ---------------------------------------------------------------- +[tool.setuptools.packages.find] +where = ["src"] +include = ["mcp_reverse_proxy*"] + +# ---------------------------------------------------------------- +# Black formatter configuration +# ---------------------------------------------------------------- +[tool.black] +line-length = 120 +target-version = ["py311", "py312", "py313"] + +# ---------------------------------------------------------------- +# Ruff linter configuration +# ---------------------------------------------------------------- +[tool.ruff] +line-length = 120 +target-version = "py311" + +[tool.ruff.lint] +select = ["E", "F", "W", "I", "N", "UP", "B", "A", "C4", "DTZ", "T10", "ISC", "ICN", "PIE", "PT", "RSE", "RET", "SIM", "TID", "ARG", "PLE", "PLW", "RUF"] +ignore = ["E501"] + +# ---------------------------------------------------------------- +# MyPy type checker configuration +# ---------------------------------------------------------------- +[tool.mypy] +python_version = "3.11" +warn_return_any = true +warn_unused_configs = true +disallow_untyped_defs = false +disallow_incomplete_defs = false +check_untyped_defs = true +no_implicit_optional = true +warn_redundant_casts = true +warn_unused_ignores = true +warn_no_return = true + +# ---------------------------------------------------------------- +# Pytest configuration +# ---------------------------------------------------------------- +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] +python_files = ["test_*.py"] +python_classes = ["Test*"] +python_functions = ["test_*"] +addopts = "-v --cov=mcp_reverse_proxy --cov-report=term-missing" \ No newline at end of file diff --git a/src/mcp_reverse_proxy/__init__.py b/src/mcp_reverse_proxy/__init__.py new file mode 100644 index 0000000..0d9a8d6 --- /dev/null +++ b/src/mcp_reverse_proxy/__init__.py @@ -0,0 +1,22 @@ +"""Location: ./mcp_reverse_proxy/__init__.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +MCP Reverse Proxy - Bridge MCP servers to remote gateways. +""" + +from mcp_reverse_proxy.base import ( + ConnectionState, + GatewayTransport, + McpServerTransport, + MessageType, +) + +__all__ = [ + "ConnectionState", + "GatewayTransport", + "McpServerTransport", + "MessageType", +] + diff --git a/src/mcp_reverse_proxy/__main__.py b/src/mcp_reverse_proxy/__main__.py new file mode 100644 index 0000000..fe34a7c --- /dev/null +++ b/src/mcp_reverse_proxy/__main__.py @@ -0,0 +1,15 @@ +"""Location: ./mcp_reverse_proxy/__main__.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Entry point for running reverse proxy as a module. +Allows: python -m mcp_reverse_proxy +""" + +# First-Party +from mcp_reverse_proxy.cli import run + +if __name__ == "__main__": + run() + diff --git a/src/mcp_reverse_proxy/base.py b/src/mcp_reverse_proxy/base.py new file mode 100644 index 0000000..3e123e3 --- /dev/null +++ b/src/mcp_reverse_proxy/base.py @@ -0,0 +1,140 @@ +"""Location: ./mcp_reverse_proxy/base.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Base classes and interfaces for reverse proxy transports. +This module defines abstract base classes for local and remote transports, +enabling extensible support for multiple MCP transport protocols. +""" + +# Future +from __future__ import annotations + +# Standard +from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable +from enum import Enum + + +class ConnectionState(Enum): + """Connection state enumeration.""" + + DISCONNECTED = "disconnected" + CONNECTING = "connecting" + CONNECTED = "connected" + RECONNECTING = "reconnecting" + SHUTTING_DOWN = "shutting_down" + + +class MessageType(Enum): + """Control message types for the reverse proxy protocol.""" + + # Control messages + REGISTER = "register" + UNREGISTER = "unregister" + HEARTBEAT = "heartbeat" + ERROR = "error" + + # MCP messages + REQUEST = "request" + RESPONSE = "response" + NOTIFICATION = "notification" + + +class McpServerTransport(ABC): + """Abstract base class for MCP server transports. + + Handles communication with MCP servers via various protocols + (stdio, HTTP, SSE, etc.). The server can be local or remote. + """ + + @abstractmethod + async def start(self) -> None: + """Start the MCP server transport. + + Raises: + RuntimeError: If transport fails to start. + """ + + @abstractmethod + async def stop(self) -> None: + """Stop the MCP server transport gracefully.""" + + @abstractmethod + async def send(self, message: str) -> None: + """Send a message to the MCP server. + + Args: + message: JSON-RPC message to send. + + Raises: + RuntimeError: If transport is not running. + """ + + @abstractmethod + def add_message_handler(self, handler: Callable[[str], Awaitable[None]]) -> None: + """Add a handler for messages from the MCP server. + + Args: + handler: Async function to handle messages. + """ + + def set_authentication(self, auth_headers: dict[str, str], auth_type: str | None = None) -> None: # noqa: B027 - optional no-op default; HTTP transports override + """Set authentication headers for subsequent requests to the MCP server. + + Args: + auth_headers: Dictionary of HTTP headers to use for authentication. + auth_type: Type of authentication (basic, bearer, authheaders, etc.) + + Note: + This is optional and only used by HTTP-based transports. + Stdio-based transports can ignore this as they don't use HTTP headers. + """ + # Default implementation does nothing (for stdio) + + +class GatewayTransport(ABC): + """Abstract base class for gateway transports. + + Handles communication with the remote gateway (currently WebSocket only). + """ + + @abstractmethod + async def connect(self) -> None: + """Establish connection to gateway. + + Raises: + Exception: If connection fails. + """ + + @abstractmethod + async def disconnect(self) -> None: + """Close connection to gateway.""" + + @abstractmethod + async def send(self, message: str | bytes) -> None: + """Send a message to the gateway. + + Args: + message: Message to send (str or bytes). + + Raises: + RuntimeError: If not connected. + """ + + @abstractmethod + def add_message_handler(self, handler: Callable[[str], Awaitable[None]]) -> None: + """Add a handler for messages from the gateway. + + Args: + handler: Async function to handle messages. + """ + + @abstractmethod + async def is_connected(self) -> bool: + """Check if transport is connected. + + Returns: + True if connected to gateway. + """ diff --git a/src/mcp_reverse_proxy/cert_utils.py b/src/mcp_reverse_proxy/cert_utils.py new file mode 100644 index 0000000..267e204 --- /dev/null +++ b/src/mcp_reverse_proxy/cert_utils.py @@ -0,0 +1,45 @@ +"""Location: ./mcp_reverse_proxy/cert_utils.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Certificate utilities for MCP reverse proxy transports. +""" + +# Standard +import os + + +def load_cert_data(cert: str) -> str: + """Load certificate data from file path or return as-is if already PEM content. + + This handles certificates from config.json which can be either file paths or PEM content. + CLI arguments are pre-loaded by the CLI layer. + + Args: + cert: Either a file path to a certificate or PEM-encoded certificate content. + + Returns: + PEM-encoded certificate content as string. + + Raises: + FileNotFoundError: If cert is a file path but file doesn't exist. + ValueError: If cert content is invalid. + """ + # First check if it's already PEM content (most reliable indicator) + if '-----BEGIN CERTIFICATE-----' in cert: + return cert + + # Not PEM content, treat as file path + cert_path = os.path.expanduser(cert) + if not os.path.isfile(cert_path): + raise FileNotFoundError(f"Certificate file not found: {cert_path}") + + with open(cert_path, encoding='utf-8') as f: + cert_data = f.read() + + # Validate it looks like PEM content + if '-----BEGIN CERTIFICATE-----' not in cert_data: + raise ValueError(f"Certificate file does not contain valid PEM data: {cert_path}") + + return cert_data diff --git a/src/mcp_reverse_proxy/cli.py b/src/mcp_reverse_proxy/cli.py new file mode 100644 index 0000000..9b2fb87 --- /dev/null +++ b/src/mcp_reverse_proxy/cli.py @@ -0,0 +1,496 @@ +"""Location: ./mcp_reverse_proxy/cli.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +CLI entry point and transport factory for reverse proxy. +""" + +# Future +from __future__ import annotations + +# Standard +import argparse +import asyncio +import logging +import os +import signal +import sys +from contextlib import suppress + +# Third-Party +import orjson + +try: + # Third-Party + import yaml +except ImportError: + yaml = None # type: ignore[assignment] + +# First-Party +from mcp_reverse_proxy.base import GatewayTransport, McpServerTransport +from mcp_reverse_proxy.client import ( + DEFAULT_MCP_HEALTH_CHECK_RETRY_INTERVAL, + DEFAULT_MCP_HEALTH_CHECK_TIMEOUT, + ReverseProxyClient, + StdioSubprocessTerminatedError, +) +from mcp_reverse_proxy.logging_config import LoggingService +from mcp_reverse_proxy.transports.sse_adapter import SseAdapter +from mcp_reverse_proxy.transports.stdio_adapter import StdioAdapter +from mcp_reverse_proxy.transports.streamablehttp_adapter import ( + StreamableHttpAdapter, +) +from mcp_reverse_proxy.transports.websocket_adapter import WebSocketAdapter + +# Initialize logging +logging_service = LoggingService() +LOGGER = logging_service.get_logger("mcp_reverse_proxy.cli") + +# Environment variable names +ENV_GATEWAY = "REVERSE_PROXY_GATEWAY" +ENV_TOKEN = "REVERSE_PROXY_TOKEN" # nosec B105 + +# Defaults +DEFAULT_RECONNECT_DELAY = 1.0 +DEFAULT_MAX_RETRIES = 0 +# CRITICAL: Must be less than gateway's MCPGATEWAY_REVERSE_PROXY_HEARTBEAT_TIMEOUT +# For gateway with 5s timeout, use 2s. For default 90s timeout, can use 30s. +DEFAULT_KEEPALIVE_INTERVAL = 2 + + +def _load_cert_from_cli_arg(cert_path: str) -> str: + """Load certificate content from a file path provided via CLI argument. + + CLI arguments (--cert, --mcp-cert, --gateway-cert) are always treated as file paths. + + Args: + cert_path: Path to certificate file. + + Returns: + PEM-encoded certificate content as string. + + Raises: + FileNotFoundError: If certificate file doesn't exist. + ValueError: If certificate file doesn't contain valid PEM data. + """ + cert_path_expanded = os.path.expanduser(cert_path) + if not os.path.isfile(cert_path_expanded): + raise FileNotFoundError(f"Certificate file not found: {cert_path_expanded}") + + LOGGER.info(f"Loading certificate from CLI argument: {cert_path_expanded}") + with open(cert_path_expanded, encoding='utf-8') as f: + cert_data = f.read() + + # Validate it looks like PEM content + if '-----BEGIN CERTIFICATE-----' not in cert_data: + raise ValueError(f"Certificate file does not contain valid PEM data: {cert_path_expanded}") + + return cert_data + + +def create_mcp_transport( + local_stdio: str | None = None, + local_streamable_http: str | None = None, + local_sse: str | None = None, + cert: str | None = None, + mcp_cert: str | None = None, + cert_from_cli: bool = False, +) -> McpServerTransport: + """Create MCP server transport based on configuration. + + Args: + local_stdio: Stdio command for MCP server. + local_streamable_http: Streamable HTTP URL for MCP server (http(s)://.../mcp). + local_sse: SSE URL for MCP server (http(s)://.../sse). + cert: Optional CA certificate (file path or PEM content). + mcp_cert: Optional CA certificate specifically for MCP server connections (file path or PEM content). + cert_from_cli: If True, treat cert/mcp_cert as file paths from CLI args and load content. + + Returns: + Configured MCP server transport. + + Raises: + ValueError: If no transport is specified or multiple are specified. + FileNotFoundError: If cert file path is invalid. + """ + # Use mcp_cert if provided, otherwise fall back to cert for backward compatibility + mcp_certificate = mcp_cert or cert + + # If certificate came from CLI argument, load the file content + if mcp_certificate and cert_from_cli: + try: + mcp_certificate = _load_cert_from_cli_arg(mcp_certificate) + except (FileNotFoundError, ValueError) as e: + LOGGER.error(f"Failed to load MCP certificate: {e}") + raise + transports = [local_stdio, local_streamable_http, local_sse] + specified = [t for t in transports if t is not None] + + if len(specified) == 0: + raise ValueError( + "Must specify one MCP server transport (--local-stdio, --local-streamable-http, or --local-sse)") + if len(specified) > 1: + raise ValueError("Can only specify one MCP server transport") + + if local_stdio: + LOGGER.info(f"Using stdio transport: {local_stdio}") + return StdioAdapter(local_stdio) + if local_streamable_http: + LOGGER.info(f"Using Streamable HTTP transport: {local_streamable_http}") + if mcp_certificate: + LOGGER.info("Using MCP-specific certificate for server connection") + return StreamableHttpAdapter(local_streamable_http, cert=mcp_certificate) + if local_sse: + LOGGER.info(f"Using SSE transport: {local_sse}") + if mcp_certificate: + LOGGER.info("Using MCP-specific certificate for server connection") + return SseAdapter(local_sse, cert=mcp_certificate) + + raise ValueError("No valid transport specified") + + +def create_gateway_transport( + gateway_url: str, + session_id: str, + token: str | None = None, + cert: str | None = None, + gateway_cert: str | None = None, + cert_from_cli: bool = False, +) -> GatewayTransport: + """Create gateway transport (currently WebSocket only). + + Args: + gateway_url: Gateway URL. + session_id: Session identifier. + token: Optional bearer token. + cert: Optional CA certificate (file path or PEM content). + gateway_cert: Optional CA certificate specifically for gateway connections (file path or PEM content). + cert_from_cli: If True, treat cert/gateway_cert as file paths from CLI args and load content. + + Returns: + Configured gateway transport. + + Raises: + FileNotFoundError: If cert file path is invalid. + """ + # Use gateway_cert if provided, otherwise fall back to cert for backward compatibility + gateway_certificate = gateway_cert or cert + + # If certificate came from CLI argument, load the file content + if gateway_certificate and cert_from_cli: + try: + gateway_certificate = _load_cert_from_cli_arg(gateway_certificate) + except (FileNotFoundError, ValueError) as e: + LOGGER.error(f"Failed to load gateway certificate: {e}") + raise + LOGGER.info(f"Using WebSocket gateway transport: {gateway_url}") + if gateway_certificate: + LOGGER.info("Using gateway-specific certificate for WebSocket connection") + return WebSocketAdapter(gateway_url, session_id, token=token, cert=gateway_certificate) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + prog="mcp-reverse-proxy", + description="Bridge MCP servers to remote gateways", + ) + + # MCP server transport options + mcp_group = parser.add_argument_group("MCP Server Transport") + mcp_group.add_argument( + "--local-stdio", + help="MCP server command to run via stdio", + ) + mcp_group.add_argument( + "--local-streamable-http", + help="MCP server Streamable HTTP URL (e.g., https://server.com/mcp)", + ) + mcp_group.add_argument( + "--local-sse", + help="MCP server SSE URL (e.g., https://server.com/sse)", + ) + + # Gateway options + gateway_group = parser.add_argument_group("Gateway Connection") + gateway_group.add_argument( + "--gateway", + help=f"Gateway URL (or use {ENV_GATEWAY} env var)", + ) + gateway_group.add_argument( + "--token", + help=f"Bearer token for authentication (or use {ENV_TOKEN} env var)", + ) + gateway_group.add_argument( + "--server-id", + help="Session identifier (auto-generated if not provided)", + ) + gateway_group.add_argument( + "--server-name", + help="Server name for registration", + ) + gateway_group.add_argument( + "--server-description", + help="Server description for registration", + ) + gateway_group.add_argument( + "--cert", + help="CA certificate for SSL verification (used for both MCP and gateway if specific certs not provided)", + ) + gateway_group.add_argument( + "--mcp-cert", + help="CA certificate specifically for MCP server HTTPS connections (overrides --cert for MCP)", + ) + gateway_group.add_argument( + "--gateway-cert", + help="CA certificate specifically for gateway WSS connections (overrides --cert for gateway)", + ) + + # Connection options + conn_group = parser.add_argument_group("Connection Options") + conn_group.add_argument( + "--reconnect-delay", + type=float, + default=DEFAULT_RECONNECT_DELAY, + help=f"Initial reconnection delay in seconds (default: {DEFAULT_RECONNECT_DELAY})", + ) + conn_group.add_argument( + "--max-retries", + type=int, + default=DEFAULT_MAX_RETRIES, + help=f"Maximum reconnection attempts, 0=infinite (default: {DEFAULT_MAX_RETRIES})", + ) + conn_group.add_argument( + "--keepalive", + type=int, + default=DEFAULT_KEEPALIVE_INTERVAL, + help=f"Keepalive interval in seconds (default: {DEFAULT_KEEPALIVE_INTERVAL})", + ) + conn_group.add_argument( + "--mcp-health-check-timeout", + type=float, + default=DEFAULT_MCP_HEALTH_CHECK_TIMEOUT, + help=f"Timeout for MCP health check calls in seconds (default: {DEFAULT_MCP_HEALTH_CHECK_TIMEOUT})", + ) + conn_group.add_argument( + "--mcp-health-check-retry-interval", + type=float, + default=DEFAULT_MCP_HEALTH_CHECK_RETRY_INTERVAL, + help=f"Interval between MCP health check retries when server is down (default: {DEFAULT_MCP_HEALTH_CHECK_RETRY_INTERVAL})", + ) + + # Configuration file + parser.add_argument( + "--config", + help="Configuration file (YAML or JSON)", + ) + + # Logging + parser.add_argument( + "--log-level", + choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], + default="INFO", + help="Log level (default: INFO)", + ) + parser.add_argument( + "--verbose", + action="store_true", + help="Enable verbose logging (same as --log-level DEBUG)", + ) + + args = parser.parse_args(argv) + + # Load configuration file if provided + if args.config: + try: + config_path = os.path.expanduser(args.config) + with open(config_path, encoding="utf-8") as f: + # Determine format by file extension + if config_path.endswith((".yaml", ".yml")): + if not yaml: + parser.error("PyYAML package required for YAML configuration file support") + config = yaml.safe_load(f) + else: + config = orjson.loads(f.read()) + + # Validate config is a dict + if not isinstance(config, dict): + parser.error("Configuration file must contain a JSON/YAML object at the top level") + + # Merge configuration (command line takes precedence) + for key, value in config.items(): + key_normalized = key.replace("-", "_") + if not hasattr(args, key_normalized) or getattr(args, key_normalized) is None: + setattr(args, key_normalized, value) + except FileNotFoundError: + parser.error(f"Configuration file not found: {config_path}") + except Exception as e: + parser.error(f"Error loading configuration file: {e}") + + # Handle verbose flag + if args.verbose: + args.log_level = "DEBUG" + + # Get gateway from environment if not provided + if not args.gateway: + args.gateway = os.getenv(ENV_GATEWAY) + if not args.gateway: + parser.error(f"--gateway or {ENV_GATEWAY} environment variable required") + + # Get token from environment if not provided + if not args.token: + args.token = os.getenv(ENV_TOKEN) + + # Generate session ID if not provided + if not args.server_id: + # Standard + import uuid + + args.server_id = str(uuid.uuid4()) + + return args + + +async def main(argv: list[str] | None = None) -> None: + """Main entry point for reverse proxy.""" + args = parse_args(argv) + + # Configure logging using LoggingService to respect JSON format settings + # This ensures the reverse proxy CLI uses the same logging format as the main gateway + root_logger = logging.getLogger() + root_logger.handlers.clear() + + # Set log level from args + log_level = getattr(logging, args.log_level) + root_logger.setLevel(log_level) + + # Use JSON formatter if LOG_FORMAT=json, otherwise use text formatter + log_format = os.getenv("LOG_FORMAT", "text").lower() + if log_format == "json": + # First-Party + from mcp_reverse_proxy.logging_config import json_formatter + + console_handler = logging.StreamHandler(sys.stderr) + console_handler.setFormatter(json_formatter) + else: + # First-Party + from mcp_reverse_proxy.logging_config import text_formatter + + console_handler = logging.StreamHandler(sys.stderr) + console_handler.setFormatter(text_formatter) + + console_handler.setLevel(log_level) + root_logger.addHandler(console_handler) + + # Track which specific certificates came from CLI vs config file + # We need to check if each cert was explicitly provided as a CLI argument + # Parse the original command line to see what was actually passed + cli_args = sys.argv[1:] + + # Check if specific cert arguments were provided on command line + mcp_cert_from_cli = '--mcp-cert' in cli_args + gateway_cert_from_cli = '--gateway-cert' in cli_args + cert_from_cli = '--cert' in cli_args and not args.config + + # Determine whether the MCP cert came from the command line + mcp_cert_is_cli = mcp_cert_from_cli or (cert_from_cli and not getattr(args, "mcp_cert", None)) + + # Determine whether the gateway cert came from the command line + gateway_cert_is_cli = gateway_cert_from_cli or (cert_from_cli and not getattr(args, "gateway_cert", None)) + + # Create transports + mcp_transport = create_mcp_transport( + local_stdio=args.local_stdio, + local_streamable_http=args.local_streamable_http, + local_sse=args.local_sse, + cert=args.cert, + mcp_cert=getattr(args, "mcp_cert", None), + cert_from_cli=mcp_cert_is_cli, + ) + + gateway_transport = create_gateway_transport( + gateway_url=args.gateway, + session_id=args.server_id, + token=args.token, + cert=args.cert, + gateway_cert=getattr(args, "gateway_cert", None), + cert_from_cli=gateway_cert_is_cli, + ) + + # Create client + client = ReverseProxyClient( + mcp_transport=mcp_transport, + gateway_transport=gateway_transport, + session_id=args.server_id, + server_name=args.server_name, + server_description=args.server_description, + reconnect_delay=args.reconnect_delay, + max_retries=args.max_retries, + keepalive_interval=args.keepalive, + mcp_health_check_timeout=args.mcp_health_check_timeout, + mcp_health_check_retry_interval=args.mcp_health_check_retry_interval, + ) + + # Handle shutdown signals + shutdown_event = asyncio.Event() + + def signal_handler(*_args: object) -> None: + """Handle shutdown signals gracefully.""" + LOGGER.info("Shutdown signal received") + shutdown_event.set() + + # Register signal handlers + loop = asyncio.get_running_loop() + for sig in (signal.SIGINT, signal.SIGTERM): + with suppress(NotImplementedError): + loop.add_signal_handler(sig, signal_handler) + + # Run client with reconnection + client_task = asyncio.create_task(client.run_with_reconnect()) + + try: + # Wait for either shutdown signal or client task to complete + done, _pending = await asyncio.wait([asyncio.create_task(shutdown_event.wait()), client_task], + return_when=asyncio.FIRST_COMPLETED) + + # Check if client_task completed with an exception + if client_task in done: + try: + client_task.result() # This will re-raise any exception + except StdioSubprocessTerminatedError as e: + LOGGER.error(f"Client task failed with StdioSubprocessTerminatedError: {e}") + raise # Re-raise to trigger clean exit + except Exception as e: + LOGGER.error(f"Client task failed with unexpected exception: {e}", exc_info=True) + raise + finally: + await client.disconnect() + client_task.cancel() + with suppress(asyncio.CancelledError): + await client_task + + +def run() -> None: + """Console script entry point.""" + try: + asyncio.run(main()) + except KeyboardInterrupt: + LOGGER.info("Shutdown complete") + sys.exit(0) + except Exception as e: + # Check if this is a StdioSubprocessTerminatedError exception + # First-Party + from mcp_reverse_proxy.client import StdioSubprocessTerminatedError + + if isinstance(e, StdioSubprocessTerminatedError): + LOGGER.error(f"Stdio subprocess terminated: {e}") + LOGGER.info("Exiting so process supervisor can restart with fresh subprocess") + sys.exit(1) + else: + LOGGER.error(f"Error: {e}", exc_info=True) + sys.exit(1) + + +if __name__ == "__main__": + run() diff --git a/src/mcp_reverse_proxy/client.py b/src/mcp_reverse_proxy/client.py new file mode 100644 index 0000000..1e2b0de --- /dev/null +++ b/src/mcp_reverse_proxy/client.py @@ -0,0 +1,750 @@ +"""Location: ./mcp_reverse_proxy/client.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Refactored reverse proxy client using transport abstractions. +""" + +# Future +from __future__ import annotations + +# Standard +import asyncio +from contextlib import suppress +from typing import Any + +# Third-Party +import orjson + +# First-Party +from mcp_reverse_proxy.base import ( + ConnectionState, + GatewayTransport, + McpServerTransport, + MessageType, +) +from mcp_reverse_proxy.logging_config import LoggingService +from mcp_reverse_proxy.transports.streamablehttp_adapter import SessionExpiredError + +# Initialize logging +logging_service = LoggingService() +LOGGER = logging_service.get_logger("mcp_reverse_proxy.client") + +# Default configuration +DEFAULT_RECONNECT_DELAY = 1.0 +DEFAULT_MAX_RETRIES = 0 +# Client sends heartbeats at this interval (MUST be less than gateway's HEARTBEAT_TIMEOUT) +# CRITICAL: Gateway with MCPGATEWAY_REVERSE_PROXY_HEARTBEAT_TIMEOUT=5 requires client < 5s +# Default 2s provides safety margin for network latency and processing delays +# For production with default gateway settings (90s timeout), this can be increased to 30s +DEFAULT_KEEPALIVE_INTERVAL = 2 +DEFAULT_MCP_HEALTH_CHECK_TIMEOUT = 5.0 +DEFAULT_MCP_HEALTH_CHECK_RETRY_INTERVAL = 10.0 + + +class StdioSubprocessTerminatedError(Exception): + """Exception raised when stdio subprocess terminates and cannot be recovered.""" + + +class ReverseProxyClient: + """Reverse proxy client using transport abstractions. + + Bridges MCP servers to remote gateways using pluggable transports. + """ + + def __init__( + self, + mcp_transport: McpServerTransport, + gateway_transport: GatewayTransport, + session_id: str, + server_name: str | None = None, + server_description: str | None = None, + reconnect_delay: float = DEFAULT_RECONNECT_DELAY, + max_retries: int = DEFAULT_MAX_RETRIES, + keepalive_interval: float = DEFAULT_KEEPALIVE_INTERVAL, + mcp_health_check_timeout: float = DEFAULT_MCP_HEALTH_CHECK_TIMEOUT, + mcp_health_check_retry_interval: float = DEFAULT_MCP_HEALTH_CHECK_RETRY_INTERVAL, + ): + """Initialize reverse proxy client. + + Args: + mcp_transport: Transport for MCP server communication. + gateway_transport: Transport for gateway communication. + session_id: Session identifier. + server_name: Optional server name. + server_description: Optional server description. + reconnect_delay: Initial reconnection delay in seconds. + max_retries: Maximum reconnection attempts (0 = infinite). + keepalive_interval: Heartbeat interval in seconds. + mcp_health_check_timeout: Timeout for MCP health check calls in seconds. + mcp_health_check_retry_interval: Interval between MCP health check retries when server is down. + """ + self.mcp_transport = mcp_transport + self.gateway_transport = gateway_transport + self.session_id = session_id + self.reconnect_delay = reconnect_delay + self.max_retries = max_retries + self.keepalive_interval = keepalive_interval + self.mcp_health_check_timeout = mcp_health_check_timeout + self.mcp_health_check_retry_interval = mcp_health_check_retry_interval + + self.server_name = server_name or f"reverse-proxy-{session_id[:8]}" + self.description = server_description or "Reverse proxied MCP server" + + self.state = ConnectionState.DISCONNECTED + self.retry_count = 0 + self._mcp_server_healthy = True + self._consecutive_mcp_failures = 0 + self._registration_successful = False + + # Session expiration recovery: When an MCP server restarts, its session IDs become invalid. + # If a request (e.g., tool call) fails with 404, we need to: + # 1. Re-register with the gateway to trigger a new initialize sequence + # 2. Save the failed request so we can retry it after the new session is established + # Without this, the original request would be lost and the gateway would timeout waiting for a response. + self._pending_reregistration_request: dict[str, Any] | None = None + + self._keepalive_task: asyncio.Task[None] | None = None + self._pending_requests: dict[Any, asyncio.Future[Any]] = {} + + # Register message handlers + self.mcp_transport.add_message_handler(self._handle_mcp_message) + self.gateway_transport.add_message_handler(self._handle_gateway_message) + + async def connect(self) -> None: + """Establish connection to gateway and MCP server.""" + if self.state != ConnectionState.DISCONNECTED: + return + + self.state = ConnectionState.CONNECTING + LOGGER.info("Establishing reverse proxy connection...") + + try: + # Start MCP server transport + LOGGER.info("Starting MCP server transport...") + await self.mcp_transport.start() + + # Check MCP server health before connecting to gateway + LOGGER.info("[CONNECT] Checking MCP server health before connecting to gateway...") + mcp_healthy = await self._check_mcp_server_health() + + if not mcp_healthy: + LOGGER.warning("[CONNECT] MCP server is not reachable, aborting gateway connection | Will retry connection later") + self.state = ConnectionState.DISCONNECTED + raise RuntimeError("MCP server is not reachable") + + LOGGER.info("[CONNECT] MCP server is healthy, proceeding with gateway connection") + + # Connect to gateway + LOGGER.info("Connecting to gateway...") + await self.gateway_transport.connect() + + self.state = ConnectionState.CONNECTED + self.retry_count = 0 + + # Register with gateway + LOGGER.info("Registering with gateway...") + await self._register() + + # Start keepalive (only if not already running) + if self._keepalive_task is None or self._keepalive_task.done(): + LOGGER.info("Starting new keepalive task") + self._keepalive_task = asyncio.create_task(self._keepalive_loop()) + else: + LOGGER.warning("Keepalive task already running, not starting a new one") + + LOGGER.info("Reverse proxy connected successfully") + + except Exception as e: + LOGGER.error(f"Connection failed: {e}") + self.state = ConnectionState.DISCONNECTED + raise + + async def disconnect(self) -> None: + """Disconnect from gateway and stop MCP server.""" + if self.state == ConnectionState.SHUTTING_DOWN: + return + + self.state = ConnectionState.SHUTTING_DOWN + LOGGER.info("Disconnecting reverse proxy...") + + if self._keepalive_task: + self._keepalive_task.cancel() + with suppress(asyncio.CancelledError): + await self._keepalive_task + + # Send unregister message + if await self.gateway_transport.is_connected(): + try: + unregister = { + "type": MessageType.UNREGISTER.value, + "sessionId": self.session_id, + } + await self.gateway_transport.send(orjson.dumps(unregister).decode()) + except Exception: + pass # nosec B110 + + await self.gateway_transport.disconnect() + await self.mcp_transport.stop() + + self.state = ConnectionState.DISCONNECTED + LOGGER.info("Reverse proxy disconnected") + + async def run_with_reconnect(self) -> None: + """Run the reverse proxy with automatic reconnection.""" + while True: + try: + if self.state == ConnectionState.SHUTTING_DOWN: + break + + await self.connect() + + # Wait for disconnection by monitoring both gateway and MCP connections + while self.state == ConnectionState.CONNECTED: + # Check if keepalive task has failed with an exception + if self._keepalive_task and self._keepalive_task.done(): + try: + # This will re-raise any exception from the task + self._keepalive_task.result() + except StdioSubprocessTerminatedError: + LOGGER.error("[RUN_WITH_RECONNECT] Keepalive task failed with StdioSubprocessTerminatedError, re-raising") + raise + except Exception as e: + LOGGER.error(f"[RUN_WITH_RECONNECT] Keepalive task failed: {e}") + self.state = ConnectionState.DISCONNECTED + break + + # Check if gateway is still connected + if not await self.gateway_transport.is_connected(): + LOGGER.warning("Gateway connection lost, triggering reconnection") + self.state = ConnectionState.DISCONNECTED + break + + # Check if MCP transport is still connected (for SSE/HTTP transports) + # This detects when the MCP server restarts and the SSE stream disconnects + # The keepalive loop will handle stopping heartbeats when MCP is unhealthy + # which causes the gateway to mark this gateway as unreachable + mcp_connected = getattr(self.mcp_transport, "_connected", True) + if not mcp_connected: + LOGGER.info("[RUN_WITH_RECONNECT] MCP transport disconnected (server likely restarted)") + LOGGER.info("[RUN_WITH_RECONNECT] Keepalive loop will stop sending heartbeats, gateway will mark as unreachable") + LOGGER.info("[RUN_WITH_RECONNECT] Monitoring for MCP server recovery...") + # Mark MCP server as unhealthy so recovery logic triggers re-registration + if self._mcp_server_healthy: + LOGGER.info("[RUN_WITH_RECONNECT] Marking MCP server as unhealthy to trigger re-registration on recovery") + self._mcp_server_healthy = False + self._consecutive_mcp_failures += 1 + # Don't break or disconnect - let the keepalive loop handle health checks + # It will stop sending heartbeats when MCP is unhealthy + + await asyncio.sleep(1) + + if self.state == ConnectionState.SHUTTING_DOWN: + break + + except StdioSubprocessTerminatedError as e: + # Re-raise to trigger proxy shutdown + LOGGER.error(f"[RUN_WITH_RECONNECT] Caught StdioSubprocessTerminatedError, re-raising: {e}") + LOGGER.error("[RUN_WITH_RECONNECT] About to raise - this should exit the function") + raise + except Exception as e: + LOGGER.error(f"Connection error: {e}") + LOGGER.error("[RUN_WITH_RECONNECT] After logging connection error, continuing to retry logic") + + # Check retry limit + self.retry_count += 1 + if self.max_retries > 0 and self.retry_count >= self.max_retries: + LOGGER.error(f"Max retries ({self.max_retries}) exceeded") + break + + # Calculate backoff delay + delay = min(self.reconnect_delay * (2**self.retry_count), 60) + LOGGER.info(f"Reconnecting in {delay}s (attempt {self.retry_count})") + + self.state = ConnectionState.RECONNECTING + await asyncio.sleep(delay) + + # Before reconnecting to gateway, verify MCP server is reachable + LOGGER.info("[RECONNECT] Checking MCP server health before reconnecting to gateway...") + mcp_healthy = await self._check_mcp_server_health() + + if not mcp_healthy: + LOGGER.warning(f"[RECONNECT] MCP server still unreachable, delaying gateway reconnection | Will retry in {self.mcp_health_check_retry_interval}s") + # Wait before checking again + await asyncio.sleep(self.mcp_health_check_retry_interval) + # Don't increment retry count for MCP health check failures + self.retry_count -= 1 + continue + + LOGGER.info("[RECONNECT] MCP server is healthy, proceeding with gateway reconnection") + self.state = ConnectionState.DISCONNECTED + + async def _register(self) -> None: + """Register MCP server with gateway.""" + register_msg = { + "type": MessageType.REGISTER.value, + "sessionId": self.session_id, + "server": { + "name": self.server_name, + "description": self.description, + "protocol": "mcp", + }, + } + LOGGER.info(f"Sending registration: session={self.session_id}, name={self.server_name}") + await self.gateway_transport.send(orjson.dumps(register_msg).decode()) + LOGGER.debug(f"Registration message sent: {register_msg}") + + async def _handle_mcp_message(self, message: str) -> None: + """Handle message from MCP server.""" + try: + LOGGER.debug(f"Handling MCP message: {message[:200]}...") + data = orjson.loads(message) + + result = data.get("result") + LOGGER.info(f"MCP response result: {result}, type: {type(result)}") + request_id = data.get("id") + + # Check if this is a health check response (don't forward to gateway) + is_health_check = request_id and str(request_id).startswith("health_check_") + + if request_id and request_id in self._pending_requests: + LOGGER.info(f"Request ID {request_id} found in pending requests") + future = self._pending_requests.pop(request_id) + LOGGER.info(f"Future retrieved: {future}") + if not future.done(): + LOGGER.info("Setting result on future") + # For health checks, just resolve the future (don't forward to gateway) + # For normal requests, create envelope and resolve future + if is_health_check: + future.set_result(data) + else: + envelope = { + "type": MessageType.RESPONSE.value, + "sessionId": self.session_id, + "payload": data, + } + future.set_result(envelope) + elif not is_health_check: + # Forward to gateway (but not health check responses) + envelope = { + "type": (MessageType.RESPONSE.value if "id" in data else MessageType.NOTIFICATION.value), + "sessionId": self.session_id, + "payload": data, + } + LOGGER.debug(f"Forwarding MCP message to gateway: {envelope}") + await self.gateway_transport.send(orjson.dumps(envelope).decode()) + + except Exception as e: + LOGGER.error(f"Error handling MCP message: {e}") + + async def _handle_gateway_message(self, message: str) -> None: + """Handle message from gateway.""" + try: + LOGGER.debug(f"Handling gateway message: {message[:200]}...") + data = orjson.loads(message) + msg_type = data.get("type") + + if msg_type == MessageType.REQUEST.value: + LOGGER.info("=" * 80) + LOGGER.info("[REVERSE_PROXY_CLIENT] Received REQUEST from gateway") + payload = data.get("payload", {}) + authentication = data.get("authentication") + auth_type = data.get("authType") + + LOGGER.info(f"[REVERSE_PROXY_CLIENT] Message keys: {list(data.keys())}") + LOGGER.info(f"[REVERSE_PROXY_CLIENT] Authentication present: {authentication is not None}") + + if authentication: + LOGGER.info(f"[REVERSE_PROXY_CLIENT] ✓ Gateway provided authentication (type: {auth_type})") + LOGGER.info(f"[REVERSE_PROXY_CLIENT] ✓ Auth headers: {list(authentication.keys())}") + # Store authentication for this request, passing auth_type for proper formatting + self.mcp_transport.set_authentication(authentication, auth_type) + else: + LOGGER.warning("[REVERSE_PROXY_CLIENT] ✗ NO authentication in gateway message") + + LOGGER.info(f"[REVERSE_PROXY_CLIENT] Gateway request payload: {payload}") + LOGGER.info("=" * 80) + + try: + await self.mcp_transport.send(orjson.dumps(payload).decode()) + except (SessionExpiredError, RuntimeError) as e: + # Session Expiration / Disconnection Recovery Flow: + # When an MCP server restarts, it loses all session state. This can manifest as: + # - SessionExpiredError: Streamable-HTTP gets 404 with stale session ID + # - RuntimeError("Not connected"): SSE transport is disconnected + # - RuntimeError("Failed to send message"): All connection attempts failed + # - RuntimeError("Subprocess not running"): Stdio process terminated + # - RuntimeError("Subprocess terminated..."): Stdio process crashed + # - RuntimeError("Subprocess stdin closed..."): Broken pipe to stdio process + # + # Recovery steps: + # 1. Save the failed request (tool call, etc.) so we don't lose it + # 2. Trigger re-registration with the gateway + # 3. Gateway will send a new initialize request to establish a fresh session + # 4. After successful re-registration, retry the saved request with the new session + # + # If the MCP server is completely stopped (not restarting), we'll send an error + # response after a timeout to prevent the gateway from hanging indefinitely. + + # Check if this is a connection-related RuntimeError + is_connection_error = isinstance(e, RuntimeError) and ( + "Not connected" in str(e) + or "Failed to send message" in str(e) + or "All connection attempts failed" in str(e) + or "Subprocess not running" in str(e) + or "Subprocess terminated" in str(e) + or "Subprocess stdin closed" in str(e) + ) + + if isinstance(e, RuntimeError) and not is_connection_error: + LOGGER.warning(f"[REVERSE_PROXY_CLIENT] Re-raising non-connection RuntimeError: {e}") + raise + + LOGGER.warning(f"[REVERSE_PROXY_CLIENT] MCP transport unavailable: {e}") + + # Check if MCP server is healthy - if not, send immediate error response + mcp_healthy = await self._check_mcp_server_health() + if not mcp_healthy: + LOGGER.warning("[REVERSE_PROXY_CLIENT] MCP server is down, sending error response to gateway") + await self._send_error_response(payload, f"MCP server is unavailable: {e}") + return + + LOGGER.info("[REVERSE_PROXY_CLIENT] MCP server appears healthy, storing request for retry after re-registration...") + + # Store the full message data for retry after re-registration completes + self._pending_reregistration_request = { + "payload": payload, + "authentication": authentication, + "authType": auth_type, + } + + LOGGER.info("[REVERSE_PROXY_CLIENT] Triggering re-registration with gateway...") + self._registration_successful = False + await self._register() + LOGGER.info("[REVERSE_PROXY_CLIENT] Re-registration triggered, returning from handler") + return + + elif msg_type == MessageType.HEARTBEAT.value: + # Gateway heartbeat is just an acknowledgment, no pong needed + LOGGER.debug("Received HEARTBEAT acknowledgment from gateway") + + elif msg_type == "register_ack": + # Gateway acknowledged receipt of registration request + LOGGER.info(f"Gateway registration acknowledged: {data.get('status', 'unknown')}") + + elif msg_type == "register_complete": + # Gateway completed registration (success or error) + status = data.get("status", "unknown") + if status == "success": + self._registration_successful = True + LOGGER.info(f"Gateway registration completed successfully for session {data.get('sessionId')}") + + # Session Expiration Recovery - Part 2: Retry the saved request + # After successful re-registration, the gateway has sent an initialize request and + # a new session has been established with the MCP server. Now we can safely retry + # the request (e.g., tool call) that originally failed with 404. + # + # This ensures the gateway receives a response and doesn't timeout waiting for one. + if self._pending_reregistration_request: + LOGGER.info("[REVERSE_PROXY_CLIENT] Retrying stored request after re-registration...") + pending = self._pending_reregistration_request + self._pending_reregistration_request = None + + # Restore authentication headers for the retry + if pending.get("authentication"): + self.mcp_transport.set_authentication(pending["authentication"], pending.get("authType")) + + # Retry the original request with the new session + try: + await self.mcp_transport.send(orjson.dumps(pending["payload"]).decode()) + LOGGER.info("[REVERSE_PROXY_CLIENT] Successfully retried request after re-registration") + except Exception as retry_error: + LOGGER.error(f"[REVERSE_PROXY_CLIENT] Failed to retry request after re-registration: {retry_error}") + # Send error response back to gateway so it doesn't hang waiting + await self._send_error_response(pending["payload"], f"MCP server unavailable after re-registration: {retry_error}") + else: + self._registration_successful = False + error_msg = data.get("message", "Unknown error") + LOGGER.error(f"Gateway registration failed for session {data.get('sessionId')}: {error_msg}") + LOGGER.error("Disconnecting due to registration failure...") + # Schedule disconnect to avoid blocking message handler + self._disconnect_task = asyncio.create_task(self.disconnect()) + + elif msg_type == MessageType.ERROR.value: + LOGGER.error(f"Gateway error: {data.get('message', 'Unknown')}") + + else: + LOGGER.warning(f"Unknown message type from gateway: {msg_type}") + + except Exception as e: + LOGGER.error(f"Error handling gateway message: {e}") + + async def _check_mcp_server_health(self) -> bool: + """Check MCP server health using transport-specific connectivity checks. + + For HTTP-based transports: performs a simple HTTP connectivity check without authentication. + For stdio transports: checks if the process is still running. + + This avoids authentication failures when the reverse proxy doesn't have credentials + (credentials are stored in ContextForge and forwarded per-request). + + Returns: + True if MCP server is reachable, False otherwise. + """ + try: + # Check if MCP transport is connected + is_connected = getattr(self.mcp_transport, "_connected", True) + LOGGER.info(f"[MCP_HEALTH] Starting health check | is_connected={is_connected}") + + # If not connected, try to connect to see if server is back online + if not is_connected: + LOGGER.info("[MCP_HEALTH] MCP transport not connected, attempting to connect to check if server is available") + try: + await self.mcp_transport.stop() + await self.mcp_transport.start() + # Give SSE stream a moment to establish + await asyncio.sleep(0.5) + # Re-check connection status after start attempt + is_connected = getattr(self.mcp_transport, "_connected", False) + LOGGER.info(f"[MCP_HEALTH] After connection attempt: is_connected={is_connected}") + if not is_connected: + LOGGER.info("[MCP_HEALTH] Failed to connect - server is down") + return False + except Exception as e: + LOGGER.warning(f"[MCP_HEALTH] Failed to connect to MCP server: {e}") + return False + + # Get transport type for health check logic + transport_type = type(self.mcp_transport).__name__ + + # Check if MCP transport is ready (has message endpoint for SSE/HTTP transports) + has_endpoint_attr = hasattr(self.mcp_transport, "_message_endpoint") + endpoint_value = getattr(self.mcp_transport, "_message_endpoint", None) if has_endpoint_attr else None + LOGGER.info(f"[MCP_HEALTH] Endpoint check | has_attr={has_endpoint_attr}, value={endpoint_value}") + + if has_endpoint_attr and not endpoint_value: + LOGGER.info("[MCP_HEALTH] MCP transport not ready yet (waiting for endpoint)") + return False + + # For HTTP-based transports, check if session is initialized + has_session_id = hasattr(self.mcp_transport, "_session_id") + session_id_value = getattr(self.mcp_transport, "_session_id", None) if has_session_id else None + + # Perform transport-specific health check + + if transport_type == "StdioAdapter": + # For stdio: check if process is still running + process = getattr(self.mcp_transport, "process", None) + if process and process.returncode is None: + LOGGER.info("[MCP_HEALTH] Stdio process is running ✓") + return True + # Stdio subprocess has terminated - no recovery possible + returncode = process.returncode if process else "N/A" + LOGGER.error(f"[MCP_HEALTH] Stdio subprocess terminated (returncode={returncode}) | Session {self.session_id[:8]}... | No automatic recovery possible for stdio transport") + LOGGER.error("[MCP_HEALTH] Raising exception to trigger proxy shutdown | Process supervisor will restart with fresh subprocess") + + # Raise exception to trigger clean shutdown + LOGGER.error("[MCP_HEALTH] About to raise StdioSubprocessTerminatedError exception") + raise StdioSubprocessTerminatedError(f"Stdio subprocess terminated with returncode={returncode}") + + if transport_type in ("SseAdapter", "StreamableHttpAdapter"): + # For SSE transports: verify the SSE stream is actually established and stable + # The _connected flag is set immediately in start(), but the SSE stream connects + # asynchronously. We need to check if the receive task is running and healthy. + if transport_type == "SseAdapter": + receive_task = getattr(self.mcp_transport, "_receive_task", None) + + # Check if receive task exists and is running (not done/failed) + if receive_task and not receive_task.done(): + # SSE stream is actively running and connected + # Healthy once endpoint is available, even without session ID + # The gateway will send initialize to establish the session + if is_connected and endpoint_value: + if session_id_value: + LOGGER.info("[MCP_HEALTH] SSE stream active with session ID - healthy ✓") + else: + LOGGER.info("[MCP_HEALTH] SSE stream active with endpoint (no session yet - gateway will initialize) ✓") + return True + LOGGER.info("[MCP_HEALTH] SSE stream active but no endpoint yet - waiting for endpoint event") + return False + # No receive task or it's done (failed/cancelled) + LOGGER.warning("[MCP_HEALTH] SSE receive task not running - server unreachable") + return False + + # For HTTP-based transports (including StreamableHttpAdapter): perform HTTP connectivity check + # to verify the server is actually reachable, not just configured + client = getattr(self.mcp_transport, "_client", None) + server_url = getattr(self.mcp_transport, "server_url", None) + + if not client or not server_url: + LOGGER.warning("[MCP_HEALTH] HTTP transport not properly initialized") + return False + + try: + # Simple HEAD or GET request to check if server is reachable + # Use a short timeout for health checks + LOGGER.info(f"[MCP_HEALTH] Checking HTTP connectivity to {server_url}") + response = await client.head(server_url, timeout=self.mcp_health_check_timeout) + + # Accept any response (including 401/403 auth errors) as "healthy" + # because it means the server is reachable + if response.status_code < 500: + LOGGER.info(f"[MCP_HEALTH] HTTP server is reachable (status: {response.status_code}) ✓") + return True + LOGGER.warning(f"[MCP_HEALTH] HTTP server returned error: {response.status_code}") + return False + + except Exception as e: + LOGGER.warning(f"[MCP_HEALTH] HTTP connectivity check failed: {e}") + return False + else: + # Unknown transport type - assume healthy if connected + LOGGER.info(f"[MCP_HEALTH] Unknown transport type {transport_type}, assuming healthy if connected") + return is_connected + + except StdioSubprocessTerminatedError as e: + # Re-raise this exception to trigger proxy shutdown + LOGGER.error(f"[MCP_HEALTH] Caught StdioSubprocessTerminatedError in health check, re-raising: {e}") + raise + except Exception as e: + LOGGER.warning(f"[MCP_HEALTH] MCP server health check failed: {e}") + return False + + async def _keepalive_loop(self) -> None: + """Send periodic heartbeat messages, conditional on MCP server health. + + This implements the transport-aware heartbeat strategy: + 1. Check MCP server health using transport-specific checks: + - HTTP transports: HTTP HEAD request to verify server connectivity + - Stdio transports: verify subprocess is still running (process.returncode is None) + 2. Only send heartbeat to gateway if MCP server is healthy + 3. If MCP server is unhealthy, skip heartbeat (gateway will detect timeout) + 4. Continue checking MCP server health and reconnect when it recovers + """ + LOGGER.info( + f"[HEARTBEAT_LOOP] Session {self.session_id[:8]}... | " + f"Starting keepalive loop | Interval: {self.keepalive_interval}s | " + f"MCP health check timeout: {self.mcp_health_check_timeout}s | " + f"MCP retry interval: {self.mcp_health_check_retry_interval}s" + ) + + heartbeat_count = 0 + + while self.state == ConnectionState.CONNECTED: + await asyncio.sleep(self.keepalive_interval) + + # Check MCP server health before sending heartbeat + LOGGER.debug(f"[HEARTBEAT_LOOP] Session {self.session_id[:8]}... | Checking MCP server health (heartbeat #{heartbeat_count + 1})") + mcp_healthy = await self._check_mcp_server_health() + + if mcp_healthy: + # MCP server is healthy - send heartbeat to gateway + if not self._mcp_server_healthy: + # MCP server recovered - always trigger re-registration to establish new session + LOGGER.info( + f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | " + f"MCP server recovered after {self._consecutive_mcp_failures} failures | " + f"Triggering re-registration to establish new session" + ) + self._mcp_server_healthy = True + self._consecutive_mcp_failures = 0 + + # Always re-register when MCP recovers to trigger new initialization + if not await self.gateway_transport.is_connected(): + LOGGER.info(f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | Gateway disconnected, reconnecting before re-registration") + try: + await self.gateway_transport.connect() + except Exception as e: + LOGGER.error(f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | Failed to reconnect to gateway: {e}") + continue + + # Re-register to trigger new MCP initialization sequence + try: + LOGGER.info(f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | Sending re-registration to gateway") + self._registration_successful = False + await self._register() + LOGGER.info(f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | Re-registration sent, gateway will initialize MCP server") + except Exception as e: + LOGGER.error(f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | Failed to re-register with gateway: {e}") + continue + + # Send heartbeat + heartbeat = { + "type": MessageType.HEARTBEAT.value, + "sessionId": self.session_id, + } + + try: + heartbeat_count += 1 + LOGGER.info(f"[HEARTBEAT_SENT] Session {self.session_id[:8]}... | Sending heartbeat #{heartbeat_count} to gateway | MCP server: healthy | Consecutive failures: 0") + await self.gateway_transport.send(orjson.dumps(heartbeat).decode()) + except Exception as e: + LOGGER.warning(f"[HEARTBEAT_ERROR] Session {self.session_id[:8]}... | Failed to send heartbeat to gateway: {e}") + break + + else: + # MCP server is unhealthy - skip heartbeat + self._consecutive_mcp_failures += 1 + + if self._mcp_server_healthy: + # First failure detected + LOGGER.warning( + f"[HEARTBEAT_SKIPPED] Session {self.session_id[:8]}... | " + f"MCP server unhealthy - skipping heartbeat #{heartbeat_count + 1} | " + f"Gateway will detect timeout and mark unreachable | " + f"Consecutive failures: {self._consecutive_mcp_failures}" + ) + self._mcp_server_healthy = False + else: + # Ongoing failure + LOGGER.info( + f"[HEARTBEAT_SKIPPED] Session {self.session_id[:8]}... | " + f"MCP server still unhealthy - skipping heartbeat #{heartbeat_count + 1} | " + f"Consecutive failures: {self._consecutive_mcp_failures}" + ) + + # Continue checking MCP server with shorter interval during outage + # This allows faster recovery detection + # Note: We already slept for keepalive_interval at the start of the loop, + # so we don't need to sleep again here. The next iteration will sleep + # for keepalive_interval before checking health again. + LOGGER.debug(f"[HEARTBEAT_RETRY] Session {self.session_id[:8]}... | Will retry MCP health check in {self.keepalive_interval}s (next loop iteration)") + + LOGGER.info(f"[HEARTBEAT_LOOP] Session {self.session_id[:8]}... | Keepalive loop ended | Total heartbeats sent: {heartbeat_count} | Final state: {self.state.value}") + + async def _send_error_response(self, request_payload: dict[str, Any], error_message: str) -> None: + """Send an error response back to the gateway for a failed request. + + Args: + request_payload: The original request payload that failed + error_message: Description of the error + """ + try: + request_id = request_payload.get("id") + if not request_id: + LOGGER.warning("[SEND_ERROR] Cannot send error response - no request ID in payload") + return + + # Create JSON-RPC error response + error_response = { + "jsonrpc": "2.0", + "id": request_id, + "error": { + "code": -32603, # Internal error + "message": error_message, + }, + } + + # Wrap in gateway envelope + envelope = { + "type": MessageType.RESPONSE.value, + "sessionId": self.session_id, + "payload": error_response, + } + + LOGGER.info(f"[SEND_ERROR] Sending error response to gateway for request {request_id}") + await self.gateway_transport.send(orjson.dumps(envelope).decode()) + + except Exception as e: + LOGGER.error(f"[SEND_ERROR] Failed to send error response to gateway: {e}") diff --git a/src/mcp_reverse_proxy/logging_config.py b/src/mcp_reverse_proxy/logging_config.py new file mode 100644 index 0000000..4b20be1 --- /dev/null +++ b/src/mcp_reverse_proxy/logging_config.py @@ -0,0 +1,80 @@ +"""Location: ./mcp_reverse_proxy/logging_config.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Standalone logging configuration for reverse proxy. +This module provides a simplified, self-contained logging setup with no parent-package dependencies. +""" + +# Standard +import logging +import os +import socket + +# Third-Party +from pythonjsonlogger import jsonlogger + +# Standard log format +LOG_FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s" +LOG_DATE_FORMAT = "%Y-%m-%dT%H:%M:%S" + +# Cache static values +_CACHED_HOSTNAME: str = socket.gethostname() +_CACHED_PID: int = os.getpid() + + +class CorrelationIdJsonFormatter(jsonlogger.JsonFormatter): + """JSON formatter with hostname and PID.""" + + def add_fields(self, log_record: dict, record: logging.LogRecord, message_dict: dict) -> None: + """Add custom fields to the log record. + + Args: + log_record: The log record dictionary to modify. + record: The original LogRecord object. + message_dict: Additional message fields. + """ + super().add_fields(log_record, record, message_dict) + + # Add standard fields + log_record["timestamp"] = self.formatTime(record, self.datefmt) + log_record["level"] = record.levelname + log_record["logger"] = record.name + log_record["hostname"] = _CACHED_HOSTNAME + log_record["pid"] = _CACHED_PID + + # Add exception info if present + if record.exc_info: + log_record["exception"] = self.formatException(record.exc_info) + + +# Create formatters +text_formatter = logging.Formatter(LOG_FORMAT, datefmt=LOG_DATE_FORMAT) +json_formatter = CorrelationIdJsonFormatter( + "%(timestamp)s %(level)s %(logger)s %(message)s", + datefmt=LOG_DATE_FORMAT +) + + +class LoggingService: + """Simplified logging service for reverse proxy.""" + + def __init__(self) -> None: + """Initialize logging service.""" + self._loggers: dict[str, logging.Logger] = {} + + def get_logger(self, name: str) -> logging.Logger: + """Get or create a logger with the given name. + + Args: + name: Logger name. + + Returns: + Configured logger instance. + """ + if name not in self._loggers: + logger = logging.getLogger(name) + self._loggers[name] = logger + return self._loggers[name] + diff --git a/src/mcp_reverse_proxy/transports/__init__.py b/src/mcp_reverse_proxy/transports/__init__.py new file mode 100644 index 0000000..9d8f9b8 --- /dev/null +++ b/src/mcp_reverse_proxy/transports/__init__.py @@ -0,0 +1,14 @@ +"""Location: ./mcp_reverse_proxy/transports/__init__.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Transport implementations for reverse proxy. +""" + +from mcp_reverse_proxy.transports.sse_adapter import SseAdapter +from mcp_reverse_proxy.transports.stdio_adapter import StdioAdapter +from mcp_reverse_proxy.transports.streamablehttp_adapter import StreamableHttpAdapter +from mcp_reverse_proxy.transports.websocket_adapter import WebSocketAdapter + +__all__ = ["SseAdapter", "StdioAdapter", "StreamableHttpAdapter", "WebSocketAdapter"] diff --git a/src/mcp_reverse_proxy/transports/sse_adapter.py b/src/mcp_reverse_proxy/transports/sse_adapter.py new file mode 100644 index 0000000..16ae9e9 --- /dev/null +++ b/src/mcp_reverse_proxy/transports/sse_adapter.py @@ -0,0 +1,439 @@ +"""Location: ./mcp_reverse_proxy/transports/sse_adapter.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +SSE transport adapter for MCP servers. +Implements Server-Sent Events (SSE) based communication with MCP servers. +This adapter connects to SSE endpoints and handles bidirectional communication +via SSE for server-to-client streaming and HTTP POST for client-to-server messages. +""" + +# Future +from __future__ import annotations + +# Standard +import asyncio +import ssl +from collections.abc import Awaitable, Callable +from contextlib import suppress + +# Third-Party +import httpx +import orjson + +# First-Party +from mcp_reverse_proxy.base import McpServerTransport +from mcp_reverse_proxy.cert_utils import load_cert_data +from mcp_reverse_proxy.logging_config import LoggingService + +# Initialize logging +logging_service = LoggingService() +LOGGER = logging_service.get_logger("mcp_reverse_proxy.transports.sse_adapter") + + +class SseAdapter(McpServerTransport): + """Transport adapter for SSE-based MCP servers. + + Communicates with MCP servers via Server-Sent Events (SSE) for server-to-client + streaming and HTTP POST for client-to-server messages. The server can be local + or remote - this adapter connects via HTTP/SSE. + + SSE Protocol Flow: + 1. Connect to SSE endpoint to receive server messages + 2. Extract message endpoint URL from initial SSE event + 3. Send client messages via HTTP POST to message endpoint + 4. Receive responses via SSE stream + + Authentication: + - Supports Basic, Bearer, and custom header authentication + - Auth headers are included in both SSE connection and POST requests + - Session management via mcp-session-id header + """ + + def __init__( + self, + server_url: str, + cert: str | None = None, + timeout: float = 90.0, + ): + """Initialize SSE adapter. + + Args: + server_url: MCP server SSE endpoint URL (can be local or remote). + cert: Optional CA certificate for SSL verification. + timeout: Request timeout in seconds. + """ + self.server_url = server_url.rstrip("/") + self.cert = cert + self.timeout = timeout + + self._client: httpx.AsyncClient | None = None + self._connected = False + self._message_handlers: list[Callable[[str], Awaitable[None]]] = [] + self._receive_task: asyncio.Task[None] | None = None + self._message_endpoint: str | None = None + self._session_id: str | None = None + self._protocol_version: str | None = None + self._auth_headers: dict[str, str] = {} + self._shutdown_event = asyncio.Event() + + async def start(self) -> None: + """Start SSE connection to MCP server. + + Establishes HTTP client, connects to SSE endpoint, and starts + receiving messages from the server. + + Raises: + RuntimeError: If connection fails or adapter is already connected. + """ + if self._connected: + return + + LOGGER.info(f"Connecting to MCP server via SSE: {self.server_url}") + + # Configure SSL context only for HTTPS + is_https = self.server_url.startswith("https://") + ssl_context = None + + if is_https: + if self.cert is not None: + # Load certificate data (from file or use as-is if already PEM content) + try: + cert_data = load_cert_data(self.cert) + LOGGER.info("Certificate loaded successfully for HTTPS connection") + except (FileNotFoundError, ValueError) as e: + LOGGER.error(f"Failed to load certificate: {e}") + raise RuntimeError(f"Certificate error: {e}") from e + + # Create SSL context with ONLY the custom CA (not system CAs) + # This is critical for self-signed certificates to work properly + # Using create_default_context() would load system CAs which reject self-signed certs + ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ssl_context.check_hostname = True + ssl_context.verify_mode = ssl.CERT_REQUIRED + # Load ONLY our custom CA, not system CAs + ssl_context.load_verify_locations(cadata=cert_data) + LOGGER.info("SSL context configured with custom CA bundle (self-signed CA support enabled)") + else: + ssl_context = ssl.create_default_context() + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + + # Create HTTP client with HTTP/2 support for better SSE performance + self._client = httpx.AsyncClient( + timeout=httpx.Timeout(self.timeout), + http2=True, + # ssl_context is always set when is_https (both branches above) + verify=ssl_context if (is_https and ssl_context is not None) else False, + ) + + self._connected = True + self._shutdown_event.clear() + + # Start receiving SSE messages + self._receive_task = asyncio.create_task(self._receive_sse_stream()) + + LOGGER.info("SSE connection to MCP server established") + + async def stop(self) -> None: + """Stop SSE connection gracefully. + + Cancels receive task, closes HTTP client, and cleans up resources. + Clears session state to ensure clean reconnection. + """ + if not self._connected: + return + + LOGGER.info("Disconnecting from MCP server via SSE") + self._connected = False + self._shutdown_event.set() + + if self._receive_task: + self._receive_task.cancel() + with suppress(asyncio.CancelledError): + await self._receive_task + + if self._client: + await self._client.aclose() + self._client = None + + # Clear session state to prevent using stale session ID on reconnection + self._session_id = None + self._message_endpoint = None + self._protocol_version = None + + LOGGER.info("SSE connection closed and session state cleared") + + async def send(self, message: str) -> None: + """Send a message to the MCP server via HTTP POST. + + Messages are sent to the message endpoint URL received from the + initial SSE connection. Responses are received via the SSE stream. + + Args: + message: JSON-RPC message to send. + + Raises: + RuntimeError: If not connected or message endpoint not available. + httpx.HTTPError: If HTTP request fails. + """ + if not self._connected or not self._client: + raise RuntimeError("Not connected to MCP server") + + if not self._message_endpoint: + raise RuntimeError("Message endpoint not yet available from SSE stream") + + LOGGER.debug(f"→ SSE POST: {message[:200]}...") + + headers = { + "Content-Type": "application/json", + "Accept": "application/json", + } + + # Add authentication headers from gateway if available + if self._auth_headers: + LOGGER.info(f"Using authentication headers from gateway: {list(self._auth_headers.keys())}") + headers.update(self._auth_headers) + + # Add session headers if available (required after initialization) + if self._session_id: + headers["mcp-session-id"] = self._session_id + if self._protocol_version: + headers["mcp-protocol-version"] = self._protocol_version + + try: + response = await self._client.post( + self._message_endpoint, + content=message, + headers=headers, + ) + response.raise_for_status() + + # Extract session ID from response headers (first request) + session_id = response.headers.get("mcp-session-id") + if session_id and not self._session_id: + self._session_id = session_id + LOGGER.info(f"Received session ID: {self._session_id}") + + LOGGER.debug(f"SSE POST successful: status={response.status_code}") + + except httpx.HTTPStatusError as e: + # If we get a 404, the session is invalid (server restarted) + # Clear session state to force reconnection + if e.response.status_code == 404: + LOGGER.warning("SSE POST returned 404 - session invalid, clearing state to force reconnection") + self._session_id = None + self._message_endpoint = None + self._protocol_version = None + LOGGER.error(f"SSE POST error: {e}") + raise RuntimeError(f"Failed to send message: {e}") from e + except httpx.HTTPError as e: + LOGGER.error(f"SSE POST error: {e}") + raise RuntimeError(f"Failed to send message: {e}") from e + + def add_message_handler(self, handler: Callable[[str], Awaitable[None]]) -> None: + """Add a handler for messages from the MCP server. + + Handlers are called when messages are received via the SSE stream. + + Args: + handler: Async function to handle messages. + """ + self._message_handlers.append(handler) + + def set_authentication(self, auth_headers: dict[str, str], auth_type: str | None = None) -> None: + """Set authentication headers for subsequent requests to the MCP server. + + Converts various authentication types to standard HTTP headers. + Headers are used for both SSE connection and POST requests. + + Args: + auth_headers: Dictionary of HTTP headers to use for authentication. + auth_type: Type of authentication (basic, bearer, authheaders, etc.) + """ + # Convert basic auth credentials to standard Authorization header + if auth_type == "basic" and "username" in auth_headers and "password" in auth_headers: + # Standard + import base64 + + username = auth_headers["username"] + password = auth_headers["password"] + credentials = f"{username}:{password}" + encoded = base64.b64encode(credentials.encode()).decode() + self._auth_headers = {"Authorization": f"Basic {encoded}"} + LOGGER.info("Converted basic auth credentials to Authorization header") + elif auth_type == "bearer" and "token" in auth_headers: + # Convert bearer token to Authorization header + self._auth_headers = {"Authorization": f"Bearer {auth_headers['token']}"} + LOGGER.info("Converted bearer token to Authorization header") + else: + # For other auth types (authheaders, custom), use headers as-is + self._auth_headers = auth_headers + LOGGER.info(f"Authentication headers set ({auth_type or 'custom'}): {list(auth_headers.keys())}") + + async def _receive_sse_stream(self) -> None: + """Receive and process SSE events from the MCP server. + + Connects to the SSE endpoint and processes incoming events: + - endpoint: Extracts message endpoint URL for POST requests + - message: Forwards JSON-RPC messages to handlers + - keepalive: Maintains connection (no action needed) + - error: Logs error events + + The stream runs until disconnection or cancellation. + """ + if not self._client: + LOGGER.error("Cannot start SSE stream: client not initialized") + return + + LOGGER.info(f"Starting SSE stream from {self.server_url}") + + headers = { + "Accept": "text/event-stream", + "Cache-Control": "no-cache", + } + + # Add authentication headers for SSE connection + if self._auth_headers: + LOGGER.info(f"Using authentication headers for SSE connection: {list(self._auth_headers.keys())}") + headers.update(self._auth_headers) + + try: + async with self._client.stream("GET", self.server_url, headers=headers) as response: + response.raise_for_status() + LOGGER.info(f"SSE stream connected: status={response.status_code}") + + # Process SSE events + event_type = None + data_lines = [] + + async for line in response.aiter_lines(): + # Check for shutdown + if self._shutdown_event.is_set(): + LOGGER.info("Shutdown requested, closing SSE stream") + break + + stripped_line = line.strip() + + if not stripped_line: + # Empty line marks end of event + if event_type and data_lines: + await self._process_sse_event(event_type, "\n".join(data_lines)) + event_type = None + data_lines = [] + continue + + if stripped_line.startswith("event:"): + event_type = stripped_line[6:].strip() + elif stripped_line.startswith("data:"): + data_lines.append(stripped_line[5:].strip()) + elif line.startswith("retry:"): + # Retry timeout - informational only + pass + elif line.startswith(":"): + # Comment - ignore + pass + + except httpx.HTTPError as e: + if self._connected: + LOGGER.error(f"SSE stream error: {e}") + # Notify handlers of connection loss + for handler in self._message_handlers: + try: + error_msg = orjson.dumps( + { + "jsonrpc": "2.0", + "error": { + "code": -32000, + "message": f"SSE connection lost: {e}", + }, + } + ).decode() + await handler(error_msg) + except Exception as handler_error: + LOGGER.error(f"Error notifying handler of connection loss: {handler_error}") + except asyncio.CancelledError: + LOGGER.info("SSE stream cancelled") + raise + except Exception as e: + LOGGER.error(f"Unexpected error in SSE stream: {e}", exc_info=True) + finally: + # Mark as disconnected when stream ends + self._connected = False + LOGGER.info("SSE stream ended, marked as disconnected") + + async def _process_sse_event(self, event_type: str, data: str) -> None: + """Process a single SSE event. + + Args: + event_type: Type of SSE event (endpoint, message, keepalive, error). + data: Event data payload. + """ + LOGGER.debug(f"← SSE event: type={event_type}, data={data[:200]}...") + + if event_type == "endpoint": + # Extract message endpoint URL + # If it's a relative URL, construct full URL from server_url + if data.startswith("/"): + # Parse base URL to get scheme and host + # Standard + from urllib.parse import urlparse + + parsed = urlparse(self.server_url) + self._message_endpoint = f"{parsed.scheme}://{parsed.netloc}{data}" + else: + self._message_endpoint = data + LOGGER.info(f"Received message endpoint: {self._message_endpoint}") + + # Extract session ID from endpoint URL query parameter if present + # FastMCP SSE servers include session_id in the endpoint URL + # Standard + from urllib.parse import parse_qs, urlparse + + parsed_endpoint = urlparse(self._message_endpoint) + query_params = parse_qs(parsed_endpoint.query) + if query_params.get("session_id"): + self._session_id = query_params["session_id"][0] + LOGGER.info(f"Extracted session ID from endpoint URL: {self._session_id}") + + elif event_type == "message": + # Forward JSON-RPC message to handlers + try: + # Extract protocol version from initialize response + if not self._protocol_version: + try: + msg_data = orjson.loads(data) + if msg_data.get("result", {}).get("protocolVersion"): + self._protocol_version = msg_data["result"]["protocolVersion"] + LOGGER.info(f"Negotiated protocol version: {self._protocol_version}") + except Exception: + pass # Not an initialize response or parsing failed + + # Notify all message handlers + for idx, handler in enumerate(self._message_handlers): + try: + LOGGER.debug(f"Calling handler {idx + 1}/{len(self._message_handlers)}") + await handler(data) + except Exception as handler_error: + LOGGER.error(f"Handler {idx + 1} failed: {handler_error}", exc_info=True) + + except Exception as e: + LOGGER.error(f"Error processing SSE message: {e}", exc_info=True) + + elif event_type == "keepalive": + # Keepalive event - no action needed + LOGGER.debug("Received keepalive event") + + elif event_type == "error": + # Error event from server + LOGGER.error(f"SSE error event: {data}") + # Forward error to handlers + for handler in self._message_handlers: + try: + await handler(data) + except Exception as handler_error: + LOGGER.error(f"Error forwarding SSE error to handler: {handler_error}") + + else: + LOGGER.warning(f"Unknown SSE event type: {event_type}") diff --git a/src/mcp_reverse_proxy/transports/stdio_adapter.py b/src/mcp_reverse_proxy/transports/stdio_adapter.py new file mode 100644 index 0000000..259c1a3 --- /dev/null +++ b/src/mcp_reverse_proxy/transports/stdio_adapter.py @@ -0,0 +1,151 @@ +"""Location: ./mcp_reverse_proxy/transports/stdio_adapter.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Stdio transport adapter for local MCP servers. +Wraps subprocess communication via stdin/stdout. +""" + +# Future +from __future__ import annotations + +# Standard +import asyncio +import shlex +import sys +from collections.abc import Awaitable, Callable +from contextlib import suppress + +# First-Party +from mcp_reverse_proxy.base import McpServerTransport +from mcp_reverse_proxy.logging_config import LoggingService + +# Initialize logging +logging_service = LoggingService() +LOGGER = logging_service.get_logger("mcp_reverse_proxy.transports.stdio_adapter") + + +class StdioAdapter(McpServerTransport): + """Transport adapter for stdio-based MCP servers. + + Manages subprocess lifecycle and stdio stream communication. + """ + + def __init__(self, command: str): + """Initialize stdio adapter. + + Args: + command: The command to run as a subprocess. + """ + self.command = command + self.process: asyncio.subprocess.Process | None = None + self._stdout_reader_task: asyncio.Task[None] | None = None + self._message_handlers: list[Callable[[str], Awaitable[None]]] = [] + + async def start(self) -> None: + """Start the stdio subprocess.""" + LOGGER.info(f"Starting local MCP server: {self.command}") + + try: + self.process = await asyncio.create_subprocess_exec( + *shlex.split(self.command), + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=sys.stderr, + ) + except FileNotFoundError as e: + raise RuntimeError(f"Command not found: {self.command}") from e + except Exception as e: + raise RuntimeError(f"Failed to start subprocess '{self.command}': {e}") from e + + if not self.process.stdin or not self.process.stdout: + raise RuntimeError(f"Failed to create subprocess with stdio: {self.command}") + + self._stdout_reader_task = asyncio.create_task(self._read_stdout()) + LOGGER.info(f"Local MCP server started (PID: {self.process.pid})") + + # Give the process a moment to initialize and check if it crashes immediately + # Use a longer delay to catch processes that fail during startup + await asyncio.sleep(0.5) + if self.process.returncode is not None: + raise RuntimeError(f"Subprocess terminated immediately after start (exit code: {self.process.returncode}). Command: {self.command}") + + async def stop(self) -> None: + """Stop the stdio subprocess gracefully.""" + if not self.process: + return + + LOGGER.info(f"Stopping local MCP server (PID: {self.process.pid})") + + if self._stdout_reader_task: + self._stdout_reader_task.cancel() + with suppress(asyncio.CancelledError): + await self._stdout_reader_task + + # Check if process is already terminated before trying to terminate it + if self.process.returncode is None: + try: + self.process.terminate() + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(self.process.wait(), timeout=5) + + if self.process.returncode is None: + LOGGER.warning("Force killing subprocess") + self.process.kill() + await self.process.wait() + except ProcessLookupError: + # Process already terminated, this is fine + LOGGER.debug("Process already terminated") + else: + LOGGER.debug(f"Process already terminated with exit code {self.process.returncode}") + + async def send(self, message: str) -> None: + """Send a message to the subprocess stdin.""" + if not self.process or not self.process.stdin: + raise RuntimeError("Subprocess not running") + + # Check if process has terminated + if self.process.returncode is not None: + raise RuntimeError(f"Subprocess terminated with exit code {self.process.returncode}") + + LOGGER.debug(f"→ stdio: {message[:200]}...") + try: + self.process.stdin.write((message + "\n").encode()) + await self.process.stdin.drain() + except (BrokenPipeError, ConnectionResetError) as e: + # Process terminated while we were trying to write + returncode = self.process.returncode if self.process else "unknown" + raise RuntimeError(f"Subprocess stdin closed (process exit code: {returncode})") from e + + def add_message_handler(self, handler: Callable[[str], Awaitable[None]]) -> None: + """Add a handler for messages from stdout.""" + self._message_handlers.append(handler) + + async def _read_stdout(self) -> None: + """Read messages from subprocess stdout.""" + if not self.process or not self.process.stdout: + return + + try: + while True: + line = await self.process.stdout.readline() + if not line: + break + + message = line.decode().strip() + if not message: + continue + + LOGGER.debug(f"← stdio: {message[:200]}...") + + for handler in self._message_handlers: + try: + await handler(message) + except Exception as e: + LOGGER.error(f"Handler error: {e}") + + except asyncio.CancelledError: # pylint: disable=try-except-raise + raise + except Exception as e: + LOGGER.error(f"Error reading stdout: {e}") diff --git a/src/mcp_reverse_proxy/transports/streamablehttp_adapter.py b/src/mcp_reverse_proxy/transports/streamablehttp_adapter.py new file mode 100644 index 0000000..7b0abca --- /dev/null +++ b/src/mcp_reverse_proxy/transports/streamablehttp_adapter.py @@ -0,0 +1,404 @@ +"""Location: ./mcp_reverse_proxy/transports/streamablehttp_adapter.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Streamable HTTP transport adapter for MCP servers. +Implements HTTP-based communication with MCP servers using httpx. +This is an alternative to stdio for servers that expose HTTP endpoints. +The MCP server can be local or remote - the proxy connects via HTTP. +""" + +# Future +from __future__ import annotations + +# Standard +import asyncio +import ssl +from collections.abc import Awaitable, Callable +from contextlib import suppress + +# Third-Party +import httpx + +# First-Party +from mcp_reverse_proxy.base import McpServerTransport +from mcp_reverse_proxy.cert_utils import load_cert_data +from mcp_reverse_proxy.logging_config import LoggingService + + +class SessionExpiredError(Exception): + """Raised when MCP server session has expired and re-registration is needed.""" + + +# Initialize logging +logging_service = LoggingService() +LOGGER = logging_service.get_logger("mcp_reverse_proxy.transports.streamablehttp_adapter") + + +class StreamableHttpAdapter(McpServerTransport): + """Transport adapter for Streamable HTTP MCP servers. + + Communicates with MCP servers via HTTP/2 streaming instead of stdio. + The server can be local or remote - this adapter connects via HTTP. + """ + + def __init__( + self, + server_url: str, + cert: str | None = None, + timeout: float = 90.0, + ): + """Initialize Streamable HTTP adapter. + + Args: + server_url: MCP server HTTP URL (can be local or remote). + cert: Optional CA certificate for SSL verification. + timeout: Request timeout in seconds. + """ + self.server_url = server_url.rstrip("/") + self.cert = cert + self.timeout = timeout + + self._client: httpx.AsyncClient | None = None + self._connected = False + self._message_handlers: list[Callable[[str], Awaitable[None]]] = [] + self._receive_task: asyncio.Task[None] | None = None + # Streamable HTTP uses the main endpoint for communication + self._endpoint_url = self.server_url + # Message endpoint (for consistency with SSE adapter health checks) + self._message_endpoint: str | None = None + # Session management (MCP protocol requirement) + self._session_id: str | None = None + self._protocol_version: str | None = None + # Authentication headers from gateway + self._auth_headers: dict[str, str] = {} + + async def start(self) -> None: + """Start HTTP client connection to MCP server.""" + if self._connected: + return + + LOGGER.info(f"Connecting to MCP server via HTTP: {self.server_url}") + + # Configure SSL context only for HTTPS + is_https = self.server_url.startswith("https://") + ssl_context = None + + if is_https: + if self.cert is not None: + # Load certificate data (from file or use as-is if already PEM content) + try: + cert_data = load_cert_data(self.cert) + LOGGER.info("Certificate loaded successfully for HTTPS connection") + except (FileNotFoundError, ValueError) as e: + LOGGER.error(f"Failed to load certificate: {e}") + raise RuntimeError(f"Certificate error: {e}") from e + + # Create SSL context with ONLY the custom CA (not system CAs) + # This is critical for self-signed certificates to work properly + # Using create_default_context() would load system CAs which reject self-signed certs + ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ssl_context.check_hostname = True + ssl_context.verify_mode = ssl.CERT_REQUIRED + # Load ONLY our custom CA, not system CAs + ssl_context.load_verify_locations(cadata=cert_data) + LOGGER.info("SSL context configured with custom CA bundle (self-signed CA support enabled)") + else: + ssl_context = ssl.create_default_context() + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + + # Create HTTP client with HTTP/2 support + self._client = httpx.AsyncClient( + timeout=httpx.Timeout(self.timeout), + http2=True, + # ssl_context is always set when is_https (both branches above) + verify=ssl_context if (is_https and ssl_context is not None) else False, + ) + + self._connected = True + + # Set message endpoint (streamable HTTP knows endpoint immediately) + self._message_endpoint = self._endpoint_url + + # Start receiving messages via SSE streaming + self._receive_task = asyncio.create_task(self._receive_stream()) + + LOGGER.info("HTTP connection to MCP server established") + + async def stop(self) -> None: + """Stop HTTP client connection. + + Clears session state to ensure clean reconnection. + """ + if not self._connected: + return + + LOGGER.info("Disconnecting from MCP server") + self._connected = False + + if self._receive_task: + self._receive_task.cancel() + with suppress(asyncio.CancelledError): + await self._receive_task + + if self._client: + await self._client.aclose() + self._client = None + + # Clear session state to prevent using stale session ID on reconnection + self._session_id = None + self._message_endpoint = None + self._protocol_version = None + + LOGGER.info("HTTP connection closed and session state cleared") + + async def send(self, message: str) -> None: + """Send a message to the MCP server via HTTP POST and handle inline response. + + Args: + message: JSON-RPC message to send to the MCP server. + + Raises: + RuntimeError: If not connected to MCP server or HTTP request fails. + SessionExpiredError: If MCP server session has expired (404) and re-registration is needed. + """ + if not self._connected or not self._client: + raise RuntimeError("Not connected to MCP server") + + LOGGER.debug(f"→ HTTP: {message[:200]}...") + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + + # Add authentication headers from gateway if available + if self._auth_headers: + LOGGER.info(f"Using authentication headers from gateway: {list(self._auth_headers.keys())}") + headers.update(self._auth_headers) + + # IMPORTANT: Only add session headers AFTER we have received them from the server + # The first request (initialize) should NOT include session headers - let the server create the session + # Subsequent requests must include BOTH session ID and protocol version + if self._session_id and self._protocol_version: + headers["mcp-session-id"] = self._session_id + headers["mcp-protocol-version"] = self._protocol_version + LOGGER.debug(f"Including session headers: session_id={self._session_id}, protocol={self._protocol_version}") + else: + # No session - check if this is an initialize request + # Standard + import json + + try: + msg_data = json.loads(message) + is_initialize = msg_data.get("method") == "initialize" + except Exception: + is_initialize = False + + if not is_initialize: + # Non-initialize request without a session - this should not happen + # The gateway should have re-initialized before sending other requests + LOGGER.error("Attempted to send non-initialize request without a valid session") + raise RuntimeError("No valid session - gateway must send initialize request first") + + LOGGER.info("Initialize request - no session headers, server will create session and return headers") + + try: + response = await self._client.post( + self._endpoint_url, + content=message, + headers=headers, + ) + response.raise_for_status() + + # Extract session ID from response headers (first request) + session_id = response.headers.get("mcp-session-id") + if session_id and not self._session_id: + self._session_id = session_id + LOGGER.info(f"Received session ID: {self._session_id}") + LOGGER.info( + f"HTTP POST successful: status={response.status_code}, content_length={len(response.content) if response.content else 0}") + + # Streamable HTTP returns responses inline - forward to handlers + if response.content: + response_text = response.text + LOGGER.info(f"← HTTP response received: {response_text[:200]}... (total length: {len(response_text)})") + LOGGER.info(f"Number of message handlers: {len(self._message_handlers)}") + + # Parse SSE format if present (streamable HTTP may return SSE-formatted responses) + # SSE format: "event: message\ndata: {json}\n\n" + json_message = response_text + if response_text.startswith(("event:", "data:")): + # Extract JSON from SSE format + lines = response_text.strip().split("\n") + for line in lines: + if line.startswith("data:"): + json_message = line[5:].strip() # Remove "data:" prefix + LOGGER.info(f"Extracted JSON from SSE format: {json_message[:200]}...") + break + + # Extract protocol version from initialize response + if not self._protocol_version: + try: + # Standard + import json + + msg_data = json.loads(json_message) + if msg_data.get("result", {}).get("protocolVersion"): + self._protocol_version = msg_data["result"]["protocolVersion"] + LOGGER.info(f"Negotiated protocol version: {self._protocol_version}") + except Exception: + pass # Not an initialize response or parsing failed + + # Notify all message handlers of the response + for idx, handler in enumerate(self._message_handlers): + try: + LOGGER.info(f"Calling handler {idx + 1}/{len(self._message_handlers)}") + await handler(json_message) + LOGGER.info(f"Handler {idx + 1} completed successfully") + except Exception as handler_error: + LOGGER.error(f"Handler {idx + 1} failed: {handler_error}", exc_info=True) + else: + LOGGER.warning("HTTP response has no content - this may indicate a problem with the MCP server") + + except httpx.HTTPStatusError as e: + # If we get a 404, the session is invalid (server restarted or session expired) + # Only retry if this is an initialize request - other requests need gateway to re-initialize + if e.response.status_code == 404 and (self._session_id or self._protocol_version): + # Standard + import json + + try: + msg_data = json.loads(message) + is_initialize = msg_data.get("method") == "initialize" + except Exception: + is_initialize = False + + if is_initialize: + # Initialize requests can be retried without session headers + LOGGER.warning("Initialize request returned 404 - retrying without session headers") + self._session_id = None + self._protocol_version = None + + try: + LOGGER.info("Retrying initialize request without session headers") + retry_headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + if self._auth_headers: + retry_headers.update(self._auth_headers) + + response = await self._client.post( + self._endpoint_url, + content=message, + headers=retry_headers, + ) + response.raise_for_status() + + # Extract session ID from response headers + session_id = response.headers.get("mcp-session-id") + if session_id: + self._session_id = session_id + LOGGER.info(f"Received new session ID after retry: {self._session_id}") + + LOGGER.info(f"Initialize retry successful: status={response.status_code}") + + # Process the response + if response.content: + response_text = response.text + LOGGER.info(f"← HTTP response received: {response_text[:200]}...") + + # Parse SSE format if present + json_message = response_text + if response_text.startswith(("event:", "data:")): + lines = response_text.strip().split("\n") + for line in lines: + if line.startswith("data:"): + json_message = line[5:].strip() + break + + # Extract protocol version from initialize response + if not self._protocol_version: + try: + msg_data = json.loads(json_message) + if msg_data.get("result", {}).get("protocolVersion"): + self._protocol_version = msg_data["result"]["protocolVersion"] + LOGGER.info(f"Negotiated protocol version: {self._protocol_version}") + except Exception: + pass + + # Notify handlers + for handler in self._message_handlers: + try: + await handler(json_message) + except Exception as handler_error: + LOGGER.error(f"Handler failed: {handler_error}", exc_info=True) + + return # Success, exit the method + + except Exception as retry_error: + LOGGER.error(f"Initialize retry after 404 failed: {retry_error}") + raise RuntimeError(f"Failed to send initialize after retry: {retry_error}") from retry_error + else: + # For non-initialize requests, clear session state and raise special exception to trigger re-registration + LOGGER.warning( + "Non-initialize request returned 404 - clearing session state and triggering re-registration") + self._session_id = None + self._protocol_version = None + raise SessionExpiredError("MCP server session expired (404), re-registration required") from e + + LOGGER.error(f"HTTP send error: {e}") + raise RuntimeError(f"Failed to send message: {e}") from e + except httpx.HTTPError as e: + LOGGER.error(f"HTTP send error: {e}") + raise RuntimeError(f"Failed to send message: {e}") from e + + def add_message_handler(self, handler: Callable[[str], Awaitable[None]]) -> None: + """Add a handler for messages from the MCP server.""" + self._message_handlers.append(handler) + + def set_authentication(self, auth_headers: dict[str, str], auth_type: str | None = None) -> None: + """Set authentication headers for subsequent requests to the MCP server. + + Args: + auth_headers: Dictionary of HTTP headers to use for authentication. + auth_type: Type of authentication (basic, bearer, authheaders, etc.) + """ + # Convert basic auth credentials to standard Authorization header + if auth_type == "basic" and "username" in auth_headers and "password" in auth_headers: + # Standard + import base64 + + username = auth_headers["username"] + password = auth_headers["password"] + credentials = f"{username}:{password}" + encoded = base64.b64encode(credentials.encode()).decode() + self._auth_headers = {"Authorization": f"Basic {encoded}"} + LOGGER.info("Converted basic auth credentials to Authorization header") + elif auth_type == "bearer" and "token" in auth_headers: + # Convert bearer token to Authorization header + self._auth_headers = {"Authorization": f"Bearer {auth_headers['token']}"} + LOGGER.info("Converted bearer token to Authorization header") + else: + # For other auth types (authheaders, custom), use headers as-is + self._auth_headers = auth_headers + LOGGER.info(f"Authentication headers set ({auth_type or 'custom'}): {list(auth_headers.keys())}") + + async def _receive_stream(self) -> None: + """Monitor connection for streamable HTTP. + + Streamable HTTP protocol handles bidirectional communication through + the main endpoint with proper Accept headers. Responses come back + inline with POST requests, not via a separate SSE stream. + """ + LOGGER.debug("Streamable HTTP uses inline responses, monitoring connection") + + # Keep the task alive to maintain connection state + try: + while self._connected: + await asyncio.sleep(1) + except asyncio.CancelledError: + LOGGER.debug("Connection monitoring cancelled") + raise diff --git a/src/mcp_reverse_proxy/transports/websocket_adapter.py b/src/mcp_reverse_proxy/transports/websocket_adapter.py new file mode 100644 index 0000000..cac55ac --- /dev/null +++ b/src/mcp_reverse_proxy/transports/websocket_adapter.py @@ -0,0 +1,219 @@ +"""Location: ./mcp_reverse_proxy/transports/websocket_adapter.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +WebSocket transport adapter for gateway connections. +Implements WebSocket-based communication with the remote gateway. +""" + +# Future +from __future__ import annotations + +# Standard +import asyncio +import ssl +from collections.abc import Awaitable, Callable +from contextlib import suppress +from typing import Any, cast +from urllib.parse import urljoin + +try: + # Third-Party + import websockets +except ImportError: + websockets = None # type: ignore[assignment] + +# First-Party +from mcp_reverse_proxy.base import GatewayTransport +from mcp_reverse_proxy.cert_utils import load_cert_data +from mcp_reverse_proxy.logging_config import LoggingService + +# Initialize logging +logging_service = LoggingService() +LOGGER = logging_service.get_logger("mcp_reverse_proxy.transports.websocket_adapter") + +# Type alias for websocket client protocol +WSClientProtocol = Any + + +class WebSocketAdapter(GatewayTransport): + """Transport adapter for WebSocket gateway connections. + + Handles WebSocket communication with the remote gateway. + """ + + def __init__( + self, + gateway_url: str, + session_id: str, + token: str | None = None, + cert: str | None = None, + ): + """Initialize WebSocket adapter. + + Args: + gateway_url: Remote gateway base URL. + session_id: Session identifier for this connection. + token: Optional bearer token for authentication. + cert: Optional CA certificate for SSL verification. + """ + self.gateway_url = gateway_url + self.session_id = session_id + self.token = token + self.cert = cert + + self._connection: WSClientProtocol | None = None + self._connected = False + self._message_handlers: list[Callable[[str], Awaitable[None]]] = [] + self._receive_task: asyncio.Task[None] | None = None + + async def connect(self) -> None: + """Establish WebSocket connection to gateway.""" + if not websockets: + raise ImportError("websockets package required for WebSocket support") + + # Build WebSocket URL + ws_url = self.gateway_url.replace("http://", "ws://").replace("https://", "wss://") + if not ws_url.startswith(("ws://", "wss://")): + ws_url = f"wss://{ws_url}" + + # Add reverse proxy endpoint + if "/reverse-proxy" not in ws_url: + ws_url = urljoin(ws_url, "/reverse-proxy/ws") + + # Build headers + headers = {} + if self.token: + headers["Authorization"] = f"Bearer {self.token}" + headers["X-Session-ID"] = self.session_id + + LOGGER.info(f"Connecting to WebSocket: {ws_url}") + + # Configure SSL context only for wss:// URLs + is_secure = ws_url.startswith("wss://") + ssl_context = None + + if is_secure: + if self.cert is not None: + # Load certificate data (from file or use as-is if already PEM content) + try: + cert_data = load_cert_data(self.cert) + LOGGER.info("Certificate loaded successfully for WebSocket connection") + except (FileNotFoundError, ValueError) as e: + LOGGER.error(f"Failed to load certificate: {e}") + raise RuntimeError(f"Certificate error: {e}") from e + + # Create SSL context with ONLY the custom CA (not system CAs) + # This is critical for self-signed certificates to work properly + # Using create_default_context() would load system CAs which reject self-signed certs + ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ssl_context.check_hostname = True + ssl_context.verify_mode = ssl.CERT_REQUIRED + # Load ONLY our custom CA, not system CAs + ssl_context.load_verify_locations(cadata=cert_data) + LOGGER.info("SSL context configured with custom CA bundle (self-signed CA support enabled)") + else: + # No cert provided - disable verification (insecure, for development only) + ssl_context = ssl.create_default_context() + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + + # Connect + self._connection = await websockets.connect( + ws_url, + additional_headers=headers, + ping_interval=20, + ping_timeout=10, + ssl=ssl_context if is_secure else None, + ) + + # Mark as connected + self._connected = True + + # Start receiving messages + self._receive_task = asyncio.create_task(self._receive_messages()) + + LOGGER.info("WebSocket connection established") + + async def disconnect(self) -> None: + """Close WebSocket connection.""" + if not self._connected: + return + + LOGGER.info("Disconnecting WebSocket") + self._connected = False + + if self._receive_task: + self._receive_task.cancel() + with suppress(asyncio.CancelledError): + await self._receive_task + + if self._connection: + await cast(Any, self._connection).close() + self._connection = None + + async def send(self, message: str | bytes) -> None: + """Send a message to the gateway via WebSocket.""" + if not self._connected or not self._connection: + raise RuntimeError("Not connected to gateway") + + # Ensure message is string for WebSocket text frames + if isinstance(message, bytes): + message = message.decode("utf-8") + + await cast(Any, self._connection).send(message) + LOGGER.debug(f"→ WS: {message[:200]}...") + + def add_message_handler(self, handler: Callable[[str], Awaitable[None]]) -> None: + """Add a handler for messages from the gateway.""" + self._message_handlers.append(handler) + + async def is_connected(self) -> bool: + """Check if WebSocket is connected. + + Uses _connected flag for consistency with SSE/StreamableHTTP adapters. + This prevents race conditions where the receive loop has ended but + the connection object hasn't been cleared yet. + """ + return self._connected + + async def _receive_messages(self) -> None: + """Receive messages from WebSocket connection.""" + if not self._connection: + return + + try: + conn = cast(Any, self._connection) + async for message in conn: + # Ensure message is string + text = message.decode("utf-8") if isinstance(message, bytes) else message + + LOGGER.debug(f"← WS: {text[:200]}...") + + for handler in self._message_handlers: + try: + await handler(text) + except Exception as e: + LOGGER.error(f"Handler error: {e}") + + except Exception as e: + # Check for ConnectionClosed exception + closed_exc = None + if websockets is not None: + ex_mod = getattr(websockets, "exceptions", None) + if ex_mod is not None: + closed_exc = getattr(ex_mod, "ConnectionClosed", None) + + if closed_exc and isinstance(e, closed_exc): + LOGGER.warning("WebSocket connection closed by remote") + else: + LOGGER.error(f"WebSocket receive error: {e}") + except asyncio.CancelledError: + LOGGER.debug("WebSocket receive cancelled") + raise + finally: + # Mark connection as closed when receive loop exits + LOGGER.info("WebSocket receive loop ended, marking connection as closed") + self._connected = False + self._connection = None diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..d5acf82 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,38 @@ +# MCP Reverse Proxy Tests + +This directory contains unit tests for the `mcp_reverse_proxy` package. + +## Test Organization + +The tests were moved from `tests/unit/mcpgateway/test_mcp_reverse_proxy_*` to this location because `mcp_reverse_proxy` is a standalone package, not a submodule of `mcpgateway`. + +## Test Files + +- `test_mcp_reverse_proxy_base.py` - Tests for base transport classes +- `test_mcp_reverse_proxy_cli.py` - Tests for CLI module +- `test_mcp_reverse_proxy_client.py` - Tests for reverse proxy client +- `test_mcp_reverse_proxy_sse_adapter.py` - Tests for SSE transport adapter +- `test_mcp_reverse_proxy_stdio_adapter.py` - Tests for stdio transport adapter +- `test_mcp_reverse_proxy_streamablehttp_adapter.py` - Tests for streamable HTTP transport adapter +- `test_mcp_reverse_proxy_websocket_adapter.py` - Tests for WebSocket transport adapter + +## Running Tests + +From the repository root: +```bash +pytest tests/ +``` + +## Import Changes + +All imports have been updated from: +```python +from mcpgateway.mcp_reverse_proxy.* import ... +``` + +To: +```python +from mcp_reverse_proxy.* import ... +``` + +This reflects the correct package structure where `mcp_reverse_proxy` is a standalone package. \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..c53526a --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,8 @@ +"""Location: ./mcp_reverse_proxy/tests/__init__.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Test package for mcp_reverse_proxy. +""" + diff --git a/tests/test_mcp_reverse_proxy_base.py b/tests/test_mcp_reverse_proxy_base.py new file mode 100644 index 0000000..ae56bde --- /dev/null +++ b/tests/test_mcp_reverse_proxy_base.py @@ -0,0 +1,151 @@ +"""Location: ./mcp_reverse_proxy/tests/test_mcp_reverse_proxy_base.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Unit tests for reverse proxy transport base classes. +""" + +# Future +from __future__ import annotations + +# Standard +from unittest.mock import AsyncMock + +# Third-Party +import pytest + +# First-Party +from mcp_reverse_proxy.base import ConnectionState, GatewayTransport, McpServerTransport, MessageType + + +class DummyMcpTransport(McpServerTransport): + """Concrete MCP transport for exercising base class behavior.""" + + def __init__(self) -> None: + self.handlers = [] + self.auth_calls = [] + + async def start(self) -> None: + """Start the transport.""" + + async def stop(self) -> None: + """Stop the transport.""" + + async def send(self, message: str) -> None: + """Send a message.""" + self.last_message = message + + def add_message_handler(self, handler) -> None: + """Register a message handler.""" + self.handlers.append(handler) + + def set_authentication(self, auth_headers: dict[str, str], auth_type: str | None = None) -> None: + """Delegate to the base implementation and record the call.""" + super().set_authentication(auth_headers, auth_type) + self.auth_calls.append((auth_headers, auth_type)) + + +class DummyGatewayTransport(GatewayTransport): + """Concrete gateway transport for exercising abstract contract behavior.""" + + def __init__(self) -> None: + self.handlers = [] + self.sent_messages = [] + self.connected = False + + async def connect(self) -> None: + """Open the connection.""" + self.connected = True + + async def disconnect(self) -> None: + """Close the connection.""" + self.connected = False + + async def send(self, message: str | bytes) -> None: + """Record a sent message.""" + self.sent_messages.append(message) + + def add_message_handler(self, handler) -> None: + """Register a message handler.""" + self.handlers.append(handler) + + async def is_connected(self) -> bool: + """Return current connection status.""" + return self.connected + + +def test_connection_state_enum_values_are_stable() -> None: + """Connection state enum values should match protocol expectations.""" + assert ConnectionState.DISCONNECTED.value == "disconnected" + assert ConnectionState.CONNECTING.value == "connecting" + assert ConnectionState.CONNECTED.value == "connected" + assert ConnectionState.RECONNECTING.value == "reconnecting" + assert ConnectionState.SHUTTING_DOWN.value == "shutting_down" + + +def test_message_type_enum_values_are_stable() -> None: + """Message type enum values should match the reverse proxy protocol.""" + assert MessageType.REGISTER.value == "register" + assert MessageType.UNREGISTER.value == "unregister" + assert MessageType.HEARTBEAT.value == "heartbeat" + assert MessageType.ERROR.value == "error" + assert MessageType.REQUEST.value == "request" + assert MessageType.RESPONSE.value == "response" + assert MessageType.NOTIFICATION.value == "notification" + + +def test_mcp_server_transport_cannot_be_instantiated_directly() -> None: + """Abstract MCP transport should reject direct instantiation.""" + with pytest.raises(TypeError): + McpServerTransport() + + +def test_gateway_transport_cannot_be_instantiated_directly() -> None: + """Abstract gateway transport should reject direct instantiation.""" + with pytest.raises(TypeError): + GatewayTransport() + + +@pytest.mark.asyncio +async def test_mcp_server_transport_default_set_authentication_is_noop() -> None: + """Default MCP auth hook should accept headers without mutating behavior.""" + transport = DummyMcpTransport() + + transport.set_authentication({"Authorization": "Bearer token"}, "bearer") + + assert transport.auth_calls == [({"Authorization": "Bearer token"}, "bearer")] + + +@pytest.mark.asyncio +async def test_dummy_mcp_transport_implements_base_contract() -> None: + """Concrete MCP transport should satisfy the abstract interface.""" + transport = DummyMcpTransport() + handler = AsyncMock() + + await transport.start() + transport.add_message_handler(handler) + await transport.send('{"jsonrpc":"2.0"}') + await transport.stop() + + assert transport.handlers == [handler] + assert transport.last_message == '{"jsonrpc":"2.0"}' + + +@pytest.mark.asyncio +async def test_dummy_gateway_transport_implements_base_contract() -> None: + """Concrete gateway transport should satisfy the abstract interface.""" + transport = DummyGatewayTransport() + handler = AsyncMock() + + await transport.connect() + transport.add_message_handler(handler) + await transport.send("message") + assert await transport.is_connected() is True + + await transport.disconnect() + + assert transport.handlers == [handler] + assert transport.sent_messages == ["message"] + assert await transport.is_connected() is False + diff --git a/tests/test_mcp_reverse_proxy_cli.py b/tests/test_mcp_reverse_proxy_cli.py new file mode 100644 index 0000000..1e33634 --- /dev/null +++ b/tests/test_mcp_reverse_proxy_cli.py @@ -0,0 +1,316 @@ +"""Location: ./mcp_reverse_proxy/tests/test_mcp_reverse_proxy_cli.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Unit tests for the reverse proxy multi-transport CLI module. +""" + +# Future +from __future__ import annotations + +# Standard +from pathlib import Path +from unittest.mock import AsyncMock + +# Third-Party +import pytest + +# First-Party +import mcp_reverse_proxy.cli as cli + + +def test_create_mcp_transport_requires_exactly_one_transport() -> None: + """No local transport should raise a validation error.""" + with pytest.raises(ValueError, match="Must specify one MCP server transport"): + cli.create_mcp_transport() + + +def test_create_mcp_transport_rejects_multiple_transports() -> None: + """Multiple local transports should raise a validation error.""" + with pytest.raises(ValueError, match="Can only specify one MCP server transport"): + cli.create_mcp_transport(local_stdio="cmd", local_sse="https://example.com/sse") + + +def test_create_mcp_transport_stdio(monkeypatch) -> None: + """Stdio transport should instantiate the stdio adapter.""" + captured = {} + + def _fake_stdio_adapter(command: str): + captured["command"] = command + return "stdio-transport" + + monkeypatch.setattr(cli, "StdioAdapter", _fake_stdio_adapter) + + result = cli.create_mcp_transport(local_stdio="python -m server") + + assert result == "stdio-transport" + assert captured["command"] == "python -m server" + + +def test_create_mcp_transport_streamable_http(monkeypatch) -> None: + """Streamable HTTP transport should pass URL and cert to adapter.""" + captured = {} + + def _fake_streamable_http_adapter(url: str, cert: str | None = None): + captured["url"] = url + captured["cert"] = cert + return "streamable-http-transport" + + monkeypatch.setattr(cli, "StreamableHttpAdapter", _fake_streamable_http_adapter) + + result = cli.create_mcp_transport(local_streamable_http="https://example.com/mcp", cert="/tmp/ca.pem") + + assert result == "streamable-http-transport" + assert captured == {"url": "https://example.com/mcp", "cert": "/tmp/ca.pem"} + + +def test_create_mcp_transport_sse(monkeypatch) -> None: + """SSE transport should pass URL and cert to adapter.""" + captured = {} + + def _fake_sse_adapter(url: str, cert: str | None = None): + captured["url"] = url + captured["cert"] = cert + return "sse-transport" + + monkeypatch.setattr(cli, "SseAdapter", _fake_sse_adapter) + + result = cli.create_mcp_transport(local_sse="https://example.com/sse", cert="/tmp/ca.pem") + + assert result == "sse-transport" + assert captured == {"url": "https://example.com/sse", "cert": "/tmp/ca.pem"} + + +def test_create_gateway_transport_uses_websocket_adapter(monkeypatch) -> None: + """Gateway transport factory should build the WebSocket adapter.""" + captured = {} + + def _fake_websocket_adapter(gateway_url: str, session_id: str, token: str | None = None, cert: str | None = None): + captured["gateway_url"] = gateway_url + captured["session_id"] = session_id + captured["token"] = token + captured["cert"] = cert + return "gateway-transport" + + monkeypatch.setattr(cli, "WebSocketAdapter", _fake_websocket_adapter) + + result = cli.create_gateway_transport("wss://gateway.example/ws", "session-123", token="secret", cert="/tmp/ca.pem") + + assert result == "gateway-transport" + assert captured == { + "gateway_url": "wss://gateway.example/ws", + "session_id": "session-123", + "token": "secret", + "cert": "/tmp/ca.pem", + } + + +def test_parse_args_reads_gateway_and_token_from_environment(monkeypatch) -> None: + """Environment variables should supply missing gateway and token values.""" + monkeypatch.setenv(cli.ENV_GATEWAY, "wss://gateway.example/ws") + monkeypatch.setenv(cli.ENV_TOKEN, "env-token") + monkeypatch.setattr("uuid.uuid4", lambda: "generated-uuid") + + args = cli.parse_args(["--local-stdio", "python -m server"]) + + assert args.gateway == "wss://gateway.example/ws" + assert args.token == "env-token" + assert args.server_id == "generated-uuid" + + +def test_parse_args_verbose_sets_debug_log_level(monkeypatch) -> None: + """Verbose flag should override the configured log level.""" + monkeypatch.setenv(cli.ENV_GATEWAY, "wss://gateway.example/ws") + monkeypatch.delenv(cli.ENV_TOKEN, raising=False) + monkeypatch.setattr("uuid.uuid4", lambda: "generated-uuid") + + args = cli.parse_args(["--local-stdio", "python -m server", "--log-level", "ERROR", "--verbose"]) + + assert args.log_level == "DEBUG" + + +def test_parse_args_json_config_merges_missing_values_only(tmp_path: Path, monkeypatch) -> None: + """JSON config should fill missing args but not override explicit CLI values.""" + config_file = tmp_path / "config.json" + config_file.write_text( + """ +{ + "gateway": "wss://config.example/ws", + "token": "config-token", + "server-name": "config-server", + "reconnect-delay": 5.5 +} +""".strip(), + encoding="utf-8", + ) + monkeypatch.delenv(cli.ENV_GATEWAY, raising=False) + monkeypatch.delenv(cli.ENV_TOKEN, raising=False) + monkeypatch.setattr("uuid.uuid4", lambda: "generated-uuid") + + args = cli.parse_args( + [ + "--config", + str(config_file), + "--local-stdio", + "python -m server", + "--gateway", + "wss://cli.example/ws", + ] + ) + + assert args.gateway == "wss://cli.example/ws" + assert args.token == "config-token" + assert args.server_name == "config-server" + assert args.reconnect_delay == cli.DEFAULT_RECONNECT_DELAY + assert args.server_id == "generated-uuid" + + +def test_parse_args_yaml_config_requires_yaml_support(tmp_path: Path, monkeypatch) -> None: + """YAML config should fail clearly when PyYAML is unavailable.""" + config_file = tmp_path / "config.yaml" + config_file.write_text("gateway: wss://config.example/ws\n", encoding="utf-8") + monkeypatch.setattr(cli, "yaml", None) + + with pytest.raises(SystemExit): + cli.parse_args(["--config", str(config_file), "--local-stdio", "python -m server"]) + + +def test_parse_args_rejects_non_mapping_config(tmp_path: Path, monkeypatch) -> None: + """Top-level config must be an object/mapping.""" + config_file = tmp_path / "config.json" + config_file.write_text('["not", "an", "object"]', encoding="utf-8") + monkeypatch.delenv(cli.ENV_GATEWAY, raising=False) + monkeypatch.delenv(cli.ENV_TOKEN, raising=False) + + with pytest.raises(SystemExit): + cli.parse_args(["--config", str(config_file), "--local-stdio", "python -m server"]) + + +def test_parse_args_requires_gateway_when_not_in_args_or_environment(monkeypatch) -> None: + """Gateway is mandatory unless provided by CLI or environment.""" + monkeypatch.delenv(cli.ENV_GATEWAY, raising=False) + monkeypatch.delenv(cli.ENV_TOKEN, raising=False) + + with pytest.raises(SystemExit): + cli.parse_args(["--local-stdio", "python -m server"]) + + +@pytest.mark.asyncio +async def test_main_creates_client_and_disconnects_on_shutdown(monkeypatch) -> None: + """Main should build transports/client and disconnect during cleanup.""" + args = cli.argparse.Namespace( + local_stdio="python -m server", + local_streamable_http=None, + local_sse=None, + cert="/tmp/ca.pem", + gateway="wss://gateway.example/ws", + server_id="server-123", + token="token-123", + server_name="Server Name", + server_description="Server Description", + reconnect_delay=2.5, + max_retries=3, + keepalive=7, + mcp_health_check_timeout=11.0, + mcp_health_check_retry_interval=1.5, + log_level="INFO", + ) + monkeypatch.setattr(cli, "parse_args", lambda _argv=None: args) + + created: dict[str, object] = {} + disconnect_mock = AsyncMock() + + def _fake_create_mcp_transport(**kwargs): + created["mcp_transport_kwargs"] = kwargs + return "mcp-transport" + + def _fake_create_gateway_transport(**kwargs): + created["gateway_transport_kwargs"] = kwargs + return "gateway-transport" + + class _FakeClient: + def __init__(self, **kwargs): + created["client_kwargs"] = kwargs + self.disconnect = disconnect_mock + + async def run_with_reconnect(self): + return None + + monkeypatch.setattr(cli, "create_mcp_transport", _fake_create_mcp_transport) + monkeypatch.setattr(cli, "create_gateway_transport", _fake_create_gateway_transport) + monkeypatch.setattr(cli, "ReverseProxyClient", _FakeClient) + + await cli.main([]) + + assert created["mcp_transport_kwargs"] == { + "local_stdio": "python -m server", + "local_streamable_http": None, + "local_sse": None, + "cert": "/tmp/ca.pem", + "mcp_cert": None, + "cert_from_cli": False, + } + assert created["gateway_transport_kwargs"] == { + "gateway_url": "wss://gateway.example/ws", + "session_id": "server-123", + "token": "token-123", + "cert": "/tmp/ca.pem", + "gateway_cert": None, + "cert_from_cli": False, + } + assert created["client_kwargs"] == { + "mcp_transport": "mcp-transport", + "gateway_transport": "gateway-transport", + "session_id": "server-123", + "server_name": "Server Name", + "server_description": "Server Description", + "reconnect_delay": 2.5, + "max_retries": 3, + "keepalive_interval": 7, + "mcp_health_check_timeout": 11.0, + "mcp_health_check_retry_interval": 1.5, + } + disconnect_mock.assert_awaited_once() + + +def test_run_exits_zero_on_keyboard_interrupt(monkeypatch) -> None: + """KeyboardInterrupt should produce a clean zero exit code.""" + + def _raise_keyboard_interrupt(_coro) -> None: + raise KeyboardInterrupt + + monkeypatch.setattr(cli.asyncio, "run", _raise_keyboard_interrupt) + + with pytest.raises(SystemExit) as exc_info: + cli.run() + + assert exc_info.value.code == 0 + + +def test_run_exits_one_on_stdio_subprocess_terminated(monkeypatch) -> None: + """Stdio subprocess termination should exit non-zero for supervisor restart.""" + + def _raise_stdio_terminated(_coro) -> None: + raise cli.StdioSubprocessTerminatedError("subprocess died") + + monkeypatch.setattr(cli.asyncio, "run", _raise_stdio_terminated) + + with pytest.raises(SystemExit) as exc_info: + cli.run() + + assert exc_info.value.code == 1 + + +def test_run_exits_one_on_unexpected_exception(monkeypatch) -> None: + """Unexpected exceptions should exit with status 1.""" + + def _raise_runtime_error(_coro) -> None: + raise RuntimeError("boom") + + monkeypatch.setattr(cli.asyncio, "run", _raise_runtime_error) + + with pytest.raises(SystemExit) as exc_info: + cli.run() + + assert exc_info.value.code == 1 diff --git a/tests/test_mcp_reverse_proxy_client.py b/tests/test_mcp_reverse_proxy_client.py new file mode 100644 index 0000000..410238a --- /dev/null +++ b/tests/test_mcp_reverse_proxy_client.py @@ -0,0 +1,848 @@ +"""Location: ./mcp_reverse_proxy/tests/test_mcp_reverse_proxy_client.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Unit tests for the reverse proxy multi-transport client module. +""" + +# Future +from __future__ import annotations + +# Standard +import asyncio +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock + +# Third-Party +import pytest + +import mcp_reverse_proxy.client as client_mod + +# First-Party +from mcp_reverse_proxy.base import ConnectionState, MessageType +from mcp_reverse_proxy.client import ReverseProxyClient, SessionExpiredError, StdioSubprocessTerminatedError + + +class FakeMcpTransport: + """Simple MCP transport fake for unit tests.""" + + def __init__(self) -> None: + self.handlers: list[Any] = [] + self.start = AsyncMock() + self.stop = AsyncMock() + self.send = AsyncMock() + self.set_authentication = AsyncMock() + self._connected = True + self._message_endpoint = "/messages" + self._session_id = "mcp-session" + self.process = SimpleNamespace(returncode=None) + + def add_message_handler(self, handler) -> None: + self.handlers.append(handler) + + +class FakeGatewayTransport: + """Simple gateway transport fake for unit tests.""" + + def __init__(self) -> None: + self.handlers: list[Any] = [] + self.connect = AsyncMock() + self.disconnect = AsyncMock() + self.send = AsyncMock() + self.is_connected = AsyncMock(return_value=True) + + def add_message_handler(self, handler) -> None: + self.handlers.append(handler) + + +@pytest.fixture +def transports(): + """Create paired fake transports.""" + return FakeMcpTransport(), FakeGatewayTransport() + + +@pytest.fixture +def proxy_client(transports) -> ReverseProxyClient: + """Create a reverse proxy client with fake transports.""" + mcp_transport, gateway_transport = transports + return ReverseProxyClient( + mcp_transport=mcp_transport, + gateway_transport=gateway_transport, + session_id="session-12345678", + server_name="Test Server", + server_description="Test Description", + reconnect_delay=0.01, + keepalive_interval=0.01, + ) + + +def test_init_registers_message_handlers_and_defaults(transports) -> None: + """Client initialization should wire handlers and defaults.""" + mcp_transport, gateway_transport = transports + + client = ReverseProxyClient( + mcp_transport=mcp_transport, + gateway_transport=gateway_transport, + session_id="abcdef123456", + ) + + assert client.server_name == "reverse-proxy-abcdef12" + assert client.description == "Reverse proxied MCP server" + assert client.state == ConnectionState.DISCONNECTED + assert len(mcp_transport.handlers) == 1 + assert len(gateway_transport.handlers) == 1 + + +@pytest.mark.asyncio +async def test_connect_starts_transports_registers_and_creates_keepalive_task(proxy_client, monkeypatch) -> None: + """Connect should start transports, register, and launch keepalive.""" + register_mock = AsyncMock() + keepalive_mock = AsyncMock() + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=True)) + monkeypatch.setattr(proxy_client, "_register", register_mock) + monkeypatch.setattr(proxy_client, "_keepalive_loop", keepalive_mock) + + await proxy_client.connect() + + assert proxy_client.state == ConnectionState.CONNECTED + assert proxy_client.retry_count == 0 + proxy_client.mcp_transport.start.assert_awaited_once() + proxy_client.gateway_transport.connect.assert_awaited_once() + register_mock.assert_awaited_once() + assert proxy_client._keepalive_task is not None + + proxy_client._keepalive_task.cancel() + with pytest.raises(asyncio.CancelledError): + await proxy_client._keepalive_task + + +@pytest.mark.asyncio +async def test_connect_raises_when_mcp_health_check_fails(proxy_client, monkeypatch) -> None: + """Connect should abort before gateway connection when MCP is unhealthy.""" + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=False)) + + with pytest.raises(RuntimeError, match="MCP server is not reachable"): + await proxy_client.connect() + + assert proxy_client.state == ConnectionState.DISCONNECTED + proxy_client.mcp_transport.start.assert_awaited_once() + proxy_client.gateway_transport.connect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_disconnect_sends_unregister_and_stops_transports(proxy_client) -> None: + """Disconnect should cancel keepalive, unregister, and stop both transports.""" + proxy_client.state = ConnectionState.CONNECTED + proxy_client._keepalive_task = asyncio.create_task(asyncio.sleep(10)) + proxy_client.gateway_transport.is_connected.return_value = True + + await proxy_client.disconnect() + + assert proxy_client.state == ConnectionState.DISCONNECTED + proxy_client.gateway_transport.send.assert_awaited_once() + unregister_payload = proxy_client.gateway_transport.send.await_args.args[0] + assert MessageType.UNREGISTER.value in unregister_payload + proxy_client.gateway_transport.disconnect.assert_awaited_once() + proxy_client.mcp_transport.stop.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_disconnect_returns_immediately_when_already_shutting_down(proxy_client) -> None: + """Disconnect should be a no-op when shutdown is already in progress.""" + proxy_client.state = ConnectionState.SHUTTING_DOWN + + await proxy_client.disconnect() + + proxy_client.gateway_transport.disconnect.assert_not_awaited() + proxy_client.mcp_transport.stop.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_register_sends_registration_envelope(proxy_client) -> None: + """Register should send the correct gateway registration message.""" + await proxy_client._register() + + proxy_client.gateway_transport.send.assert_awaited_once() + register_payload = proxy_client.gateway_transport.send.await_args.args[0] + assert MessageType.REGISTER.value in register_payload + assert proxy_client.session_id in register_payload + assert proxy_client.server_name in register_payload + + +@pytest.mark.asyncio +async def test_handle_mcp_message_resolves_pending_health_check_future(proxy_client) -> None: + """Health check responses should resolve pending future without gateway forwarding.""" + future: asyncio.Future[Any] = asyncio.Future() + proxy_client._pending_requests["health_check_123"] = future + + await proxy_client._handle_mcp_message('{"jsonrpc":"2.0","id":"health_check_123","result":{"ok":true}}') + + assert future.done() + assert future.result()["result"] == {"ok": True} + proxy_client.gateway_transport.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_handle_mcp_message_resolves_pending_request_with_response_envelope(proxy_client) -> None: + """Pending non-health responses should resolve with a gateway envelope.""" + future: asyncio.Future[Any] = asyncio.Future() + proxy_client._pending_requests[7] = future + + await proxy_client._handle_mcp_message('{"jsonrpc":"2.0","id":7,"result":{"value":1}}') + + assert future.done() + result = future.result() + assert result["type"] == MessageType.RESPONSE.value + assert result["sessionId"] == proxy_client.session_id + assert result["payload"]["result"] == {"value": 1} + + +@pytest.mark.asyncio +async def test_handle_mcp_message_forwards_notification_to_gateway(proxy_client) -> None: + """Non-pending notifications should be forwarded to the gateway.""" + await proxy_client._handle_mcp_message('{"jsonrpc":"2.0","method":"notifications/test","params":{"x":1}}') + + proxy_client.gateway_transport.send.assert_awaited_once() + forwarded = proxy_client.gateway_transport.send.await_args.args[0] + assert MessageType.NOTIFICATION.value in forwarded + assert proxy_client.session_id in forwarded + + +@pytest.mark.asyncio +async def test_handle_gateway_request_sets_auth_and_forwards_to_mcp(proxy_client) -> None: + """Gateway request should set auth and forward the payload to MCP.""" + message = """ + { + "type": "request", + "payload": {"jsonrpc": "2.0", "id": 1, "method": "tools/call"}, + "authentication": {"Authorization": "Bearer token"}, + "authType": "bearer" + } + """.strip() + + await proxy_client._handle_gateway_message(message) + + proxy_client.mcp_transport.set_authentication.assert_called_once_with({"Authorization": "Bearer token"}, "bearer") + proxy_client.mcp_transport.send.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_handle_gateway_request_sends_error_when_transport_unavailable_and_health_fails(proxy_client, monkeypatch) -> None: + """Connection-related MCP errors should return an error when health check fails.""" + proxy_client.mcp_transport.send.side_effect = RuntimeError("Not connected") + send_error_mock = AsyncMock() + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=False)) + monkeypatch.setattr(proxy_client, "_send_error_response", send_error_mock) + + await proxy_client._handle_gateway_message('{"type":"request","payload":{"jsonrpc":"2.0","id":5,"method":"x"}}') + + send_error_mock.assert_awaited_once() + assert proxy_client._pending_reregistration_request is None + + +@pytest.mark.asyncio +async def test_handle_gateway_request_stores_pending_and_reregisters_when_health_recovers(proxy_client, monkeypatch) -> None: + """Recoverable transport errors should save the request and trigger re-registration.""" + proxy_client.mcp_transport.send.side_effect = SessionExpiredError("expired") + register_mock = AsyncMock() + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=True)) + monkeypatch.setattr(proxy_client, "_register", register_mock) + + await proxy_client._handle_gateway_message( + '{"type":"request","payload":{"jsonrpc":"2.0","id":9,"method":"x"},"authentication":{"Authorization":"Bearer t"},"authType":"bearer"}' + ) + + assert proxy_client._pending_reregistration_request == { + "payload": {"jsonrpc": "2.0", "id": 9, "method": "x"}, + "authentication": {"Authorization": "Bearer t"}, + "authType": "bearer", + } + assert proxy_client._registration_successful is False + register_mock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_handle_gateway_request_reraises_non_connection_runtime_error(proxy_client, monkeypatch) -> None: + """Non-connection runtime errors should be re-raised then logged by the outer handler.""" + proxy_client.mcp_transport.send.side_effect = RuntimeError("different failure") + check_health_mock = AsyncMock() + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", check_health_mock) + + await proxy_client._handle_gateway_message('{"type":"request","payload":{"jsonrpc":"2.0","id":11,"method":"x"}}') + + check_health_mock.assert_not_awaited() + assert proxy_client._pending_reregistration_request is None + + +@pytest.mark.asyncio +async def test_handle_gateway_register_complete_retries_pending_request(proxy_client) -> None: + """Successful register_complete should retry a stored request and restore auth.""" + proxy_client._pending_reregistration_request = { + "payload": {"jsonrpc": "2.0", "id": 21, "method": "retry"}, + "authentication": {"Authorization": "Bearer token"}, + "authType": "bearer", + } + + await proxy_client._handle_gateway_message('{"type":"register_complete","status":"success","sessionId":"session-12345678"}') + + assert proxy_client._registration_successful is True + assert proxy_client._pending_reregistration_request is None + proxy_client.mcp_transport.set_authentication.assert_called_once_with({"Authorization": "Bearer token"}, "bearer") + proxy_client.mcp_transport.send.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_handle_gateway_register_complete_sends_error_if_retry_fails(proxy_client, monkeypatch) -> None: + """Retry failure after successful registration should send an error response.""" + proxy_client._pending_reregistration_request = { + "payload": {"jsonrpc": "2.0", "id": 22, "method": "retry"}, + "authentication": None, + "authType": None, + } + proxy_client.mcp_transport.send.side_effect = RuntimeError("retry failed") + send_error_mock = AsyncMock() + monkeypatch.setattr(proxy_client, "_send_error_response", send_error_mock) + + await proxy_client._handle_gateway_message('{"type":"register_complete","status":"success","sessionId":"session-12345678"}') + + send_error_mock.assert_awaited_once() + assert proxy_client._pending_reregistration_request is None + + +@pytest.mark.asyncio +async def test_handle_gateway_register_complete_failure_schedules_disconnect(proxy_client, monkeypatch) -> None: + """Failed registration completion should schedule disconnect.""" + created_tasks: list[asyncio.Task[Any]] = [] + + original_create_task = asyncio.create_task + + def _capture_task(coro): + task = original_create_task(coro) + created_tasks.append(task) + return task + + monkeypatch.setattr(asyncio, "create_task", _capture_task) + disconnect_mock = AsyncMock() + monkeypatch.setattr(proxy_client, "disconnect", disconnect_mock) + + await proxy_client._handle_gateway_message('{"type":"register_complete","status":"error","message":"bad","sessionId":"session-12345678"}') + + assert proxy_client._registration_successful is False + assert len(created_tasks) == 1 + await created_tasks[0] + disconnect_mock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_for_stdio_running_process(proxy_client) -> None: + """Healthy stdio process should report healthy.""" + proxy_client.mcp_transport.process = SimpleNamespace(returncode=None) + proxy_client.mcp_transport.__class__.__name__ = "StdioAdapter" + + assert await proxy_client._check_mcp_server_health() is True + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_for_stdio_terminated_process_raises(proxy_client) -> None: + """Terminated stdio subprocess should raise shutdown exception.""" + proxy_client.mcp_transport.process = SimpleNamespace(returncode=17) + proxy_client.mcp_transport.__class__.__name__ = "StdioAdapter" + + with pytest.raises(StdioSubprocessTerminatedError, match="returncode=17"): + await proxy_client._check_mcp_server_health() + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_for_sse_requires_receive_task(proxy_client) -> None: + """SSE transport without a running receive task should be unhealthy.""" + proxy_client.mcp_transport.__class__.__name__ = "SseAdapter" + proxy_client.mcp_transport._receive_task = None + + assert await proxy_client._check_mcp_server_health() is False + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_for_http_transport_uses_head_request(proxy_client) -> None: + """HTTP transport should use HEAD request reachability check.""" + proxy_client.mcp_transport.__class__.__name__ = "StreamableHttpAdapter" + proxy_client.mcp_transport._client = SimpleNamespace(head=AsyncMock(return_value=SimpleNamespace(status_code=401))) + proxy_client.mcp_transport.server_url = "https://server.example/mcp" + + assert await proxy_client._check_mcp_server_health() is True + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_attempts_reconnect_when_transport_disconnected(proxy_client, monkeypatch) -> None: + """Disconnected transport should attempt stop/start recovery before failing.""" + proxy_client.mcp_transport._connected = False + proxy_client.mcp_transport._message_endpoint = "/messages" + start_mock = AsyncMock(side_effect=lambda: setattr(proxy_client.mcp_transport, "_connected", False)) + stop_mock = AsyncMock() + proxy_client.mcp_transport.start = start_mock + proxy_client.mcp_transport.stop = stop_mock + sleep_mock = AsyncMock() + monkeypatch.setattr(client_mod.asyncio, "sleep", sleep_mock) + + assert await proxy_client._check_mcp_server_health() is False + stop_mock.assert_awaited_once() + start_mock.assert_awaited_once() + sleep_mock.assert_awaited_once_with(0.5) + + +@pytest.mark.asyncio +async def test_send_error_response_requires_request_id(proxy_client) -> None: + """Missing request ID should skip sending an error envelope.""" + await proxy_client._send_error_response({"jsonrpc": "2.0"}, "broken") + + proxy_client.gateway_transport.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_send_error_response_sends_gateway_envelope(proxy_client) -> None: + """Error response should be wrapped and sent to the gateway.""" + await proxy_client._send_error_response({"jsonrpc": "2.0", "id": 33}, "broken") + + proxy_client.gateway_transport.send.assert_awaited_once() + payload = proxy_client.gateway_transport.send.await_args.args[0] + assert MessageType.RESPONSE.value in payload + assert '"id":33' in payload or '"id": 33' in payload + assert "broken" in payload + + + +@pytest.mark.asyncio +async def test_run_with_reconnect_breaks_immediately_when_shutting_down(proxy_client) -> None: + """Reconnect loop should exit immediately during shutdown.""" + proxy_client.state = ConnectionState.SHUTTING_DOWN + + await proxy_client.run_with_reconnect() + + proxy_client.mcp_transport.start.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_with_reconnect_breaks_after_connect_if_shutdown_state_set(proxy_client, monkeypatch) -> None: + """Reconnect loop should stop once connect transitions to shutting down.""" + connect_mock = AsyncMock(side_effect=lambda: setattr(proxy_client, "state", ConnectionState.SHUTTING_DOWN)) + monkeypatch.setattr(proxy_client, "connect", connect_mock) + + await proxy_client.run_with_reconnect() + + connect_mock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_run_with_reconnect_reraises_stdio_exception_from_connect(proxy_client, monkeypatch) -> None: + """Stdio termination during connect should be re-raised.""" + monkeypatch.setattr(proxy_client, "connect", AsyncMock(side_effect=StdioSubprocessTerminatedError("dead"))) + + with pytest.raises(StdioSubprocessTerminatedError): + await proxy_client.run_with_reconnect() + + +@pytest.mark.asyncio +async def test_run_with_reconnect_disconnects_when_keepalive_task_fails(proxy_client, monkeypatch) -> None: + """Reconnect loop should disconnect when keepalive task fails.""" + first_connect = True + + async def _fake_connect() -> None: + nonlocal first_connect + proxy_client.state = ConnectionState.CONNECTED + if first_connect: + # Create a failed task directly by creating a done task with an exception + task = asyncio.create_task(asyncio.sleep(0)) + await task # Let it complete + # Now create a new task and set its exception + proxy_client._keepalive_task = asyncio.Future() + proxy_client._keepalive_task.set_exception(RuntimeError("keepalive failed")) + first_connect = False + else: + proxy_client.state = ConnectionState.SHUTTING_DOWN + + monkeypatch.setattr(proxy_client, "connect", _fake_connect) + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=True)) + monkeypatch.setattr(client_mod.asyncio, "sleep", AsyncMock()) + + await proxy_client.run_with_reconnect() + + assert proxy_client.retry_count == 1 + assert proxy_client.state == ConnectionState.SHUTTING_DOWN + + +@pytest.mark.asyncio +async def test_run_with_reconnect_breaks_when_gateway_disconnects(proxy_client, monkeypatch) -> None: + """Gateway disconnect should trigger reconnect handling.""" + call_count = 0 + + async def _fake_connect() -> None: + nonlocal call_count + call_count += 1 + if call_count == 1: + proxy_client.state = ConnectionState.CONNECTED + proxy_client._keepalive_task = None + proxy_client.gateway_transport.is_connected.return_value = False + else: + proxy_client.state = ConnectionState.SHUTTING_DOWN + + monkeypatch.setattr(proxy_client, "connect", _fake_connect) + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=True)) + monkeypatch.setattr(client_mod.asyncio, "sleep", AsyncMock()) + + await proxy_client.run_with_reconnect() + + assert proxy_client.retry_count == 1 + assert call_count == 2 + + +@pytest.mark.asyncio +async def test_run_with_reconnect_marks_mcp_unhealthy_when_transport_disconnects(proxy_client, monkeypatch) -> None: + """MCP transport disconnection should mark health false and continue loop.""" + call_count = 0 + + async def _fake_connect() -> None: + nonlocal call_count + call_count += 1 + proxy_client._keepalive_task = None + if call_count == 1: + proxy_client.state = ConnectionState.CONNECTED + proxy_client.gateway_transport.is_connected.return_value = True + proxy_client.mcp_transport._connected = False + else: + proxy_client.state = ConnectionState.SHUTTING_DOWN + + async def _fake_sleep(_delay: float) -> None: + if proxy_client.state == ConnectionState.CONNECTED: + proxy_client.state = ConnectionState.DISCONNECTED + + monkeypatch.setattr(proxy_client, "connect", _fake_connect) + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=True)) + monkeypatch.setattr(client_mod.asyncio, "sleep", _fake_sleep) + + await proxy_client.run_with_reconnect() + + assert proxy_client._mcp_server_healthy is False + assert proxy_client._consecutive_mcp_failures == 1 + + +@pytest.mark.asyncio +async def test_run_with_reconnect_honors_max_retries(proxy_client, monkeypatch) -> None: + """Reconnect loop should stop when max retries is exceeded.""" + proxy_client.max_retries = 1 + monkeypatch.setattr(proxy_client, "connect", AsyncMock(side_effect=RuntimeError("boom"))) + monkeypatch.setattr(client_mod.asyncio, "sleep", AsyncMock()) + + await proxy_client.run_with_reconnect() + + assert proxy_client.retry_count == 1 + + +@pytest.mark.asyncio +async def test_run_with_reconnect_decrements_retry_when_mcp_still_unhealthy(proxy_client, monkeypatch) -> None: + """Health-check failures during reconnect should not consume a retry.""" + call_count = 0 + + async def _fake_connect() -> None: + nonlocal call_count + call_count += 1 + if call_count == 1: + raise RuntimeError("boom") + proxy_client.state = ConnectionState.SHUTTING_DOWN + + health_mock = AsyncMock(side_effect=[False, True]) + sleep_mock = AsyncMock() + monkeypatch.setattr(proxy_client, "connect", _fake_connect) + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", health_mock) + monkeypatch.setattr(client_mod.asyncio, "sleep", sleep_mock) + + await proxy_client.run_with_reconnect() + + assert proxy_client.retry_count == 0 + + +@pytest.mark.asyncio +async def test_handle_mcp_message_invalid_json_is_swallowed(proxy_client) -> None: + """Malformed MCP messages should be logged and swallowed.""" + await proxy_client._handle_mcp_message("{invalid-json") + + proxy_client.gateway_transport.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_handle_gateway_heartbeat_ack_does_nothing(proxy_client) -> None: + """Heartbeat acknowledgments should be ignored.""" + await proxy_client._handle_gateway_message('{"type":"heartbeat"}') + + proxy_client.mcp_transport.send.assert_not_awaited() + proxy_client.gateway_transport.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_handle_gateway_register_ack_does_nothing(proxy_client) -> None: + """Register acknowledgments should not trigger side effects.""" + await proxy_client._handle_gateway_message('{"type":"register_ack","status":"ok"}') + + proxy_client.mcp_transport.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_handle_gateway_error_message_does_nothing(proxy_client) -> None: + """Gateway error messages should be logged without raising.""" + await proxy_client._handle_gateway_message('{"type":"error","message":"bad"}') + + proxy_client.mcp_transport.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_handle_gateway_unknown_message_type_is_ignored(proxy_client) -> None: + """Unknown gateway messages should be ignored safely.""" + await proxy_client._handle_gateway_message('{"type":"mystery"}') + + proxy_client.mcp_transport.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_handle_gateway_invalid_json_is_swallowed(proxy_client) -> None: + """Malformed gateway messages should be logged and swallowed.""" + await proxy_client._handle_gateway_message("{broken") + + proxy_client.mcp_transport.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_returns_false_when_endpoint_not_ready(proxy_client) -> None: + """Missing message endpoint should make connected HTTP/SSE transport unhealthy.""" + proxy_client.mcp_transport.__class__.__name__ = "StreamableHttpAdapter" + proxy_client.mcp_transport._message_endpoint = None + + assert await proxy_client._check_mcp_server_health() is False + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_for_sse_running_receive_task_without_endpoint_returns_false(proxy_client) -> None: + """SSE receive task without endpoint should not yet be healthy.""" + proxy_client.mcp_transport.__class__.__name__ = "SseAdapter" + proxy_client.mcp_transport._message_endpoint = None + proxy_client.mcp_transport._receive_task = asyncio.create_task(asyncio.sleep(10)) + + try: + assert await proxy_client._check_mcp_server_health() is False + finally: + proxy_client.mcp_transport._receive_task.cancel() + with pytest.raises(asyncio.CancelledError): + await proxy_client.mcp_transport._receive_task + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_for_sse_running_receive_task_with_endpoint_is_true(proxy_client) -> None: + """SSE receive task with endpoint should be healthy.""" + proxy_client.mcp_transport.__class__.__name__ = "SseAdapter" + proxy_client.mcp_transport._message_endpoint = "/messages" + proxy_client.mcp_transport._session_id = None + proxy_client.mcp_transport._receive_task = asyncio.create_task(asyncio.sleep(10)) + + try: + assert await proxy_client._check_mcp_server_health() is True + finally: + proxy_client.mcp_transport._receive_task.cancel() + with pytest.raises(asyncio.CancelledError): + await proxy_client.mcp_transport._receive_task + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_for_http_uninitialized_transport_returns_false(proxy_client) -> None: + """HTTP transport without client or URL should be unhealthy.""" + proxy_client.mcp_transport.__class__.__name__ = "StreamableHttpAdapter" + proxy_client.mcp_transport._client = None + proxy_client.mcp_transport.server_url = None + + assert await proxy_client._check_mcp_server_health() is False + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_for_http_5xx_returns_false(proxy_client) -> None: + """HTTP 5xx responses should be considered unhealthy.""" + proxy_client.mcp_transport.__class__.__name__ = "StreamableHttpAdapter" + proxy_client.mcp_transport._client = SimpleNamespace(head=AsyncMock(return_value=SimpleNamespace(status_code=503))) + proxy_client.mcp_transport.server_url = "https://server.example/mcp" + + assert await proxy_client._check_mcp_server_health() is False + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_for_http_exception_returns_false(proxy_client) -> None: + """HTTP HEAD exceptions should be considered unhealthy.""" + proxy_client.mcp_transport.__class__.__name__ = "StreamableHttpAdapter" + proxy_client.mcp_transport._client = SimpleNamespace(head=AsyncMock(side_effect=RuntimeError("bad"))) + proxy_client.mcp_transport.server_url = "https://server.example/mcp" + + assert await proxy_client._check_mcp_server_health() is False + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_unknown_transport_returns_connected_state(proxy_client) -> None: + """Unknown transport types should fall back to connected state.""" + proxy_client.mcp_transport.__class__.__name__ = "MysteryAdapter" + proxy_client.mcp_transport._connected = True + + assert await proxy_client._check_mcp_server_health() is True + + +@pytest.mark.asyncio +async def test_check_mcp_server_health_returns_false_on_unexpected_exception(proxy_client, monkeypatch) -> None: + """Unexpected health-check exceptions should return false.""" + monkeypatch.delattr(proxy_client, "mcp_transport", raising=False) + + assert await proxy_client._check_mcp_server_health() is False + + +@pytest.mark.asyncio +async def test_keepalive_loop_sends_heartbeat_when_mcp_is_healthy(proxy_client, monkeypatch) -> None: + """Healthy MCP server should produce a heartbeat.""" + proxy_client.state = ConnectionState.CONNECTED + proxy_client.keepalive_interval = 0 + check_health_mock = AsyncMock(return_value=True) + + async def _fake_sleep(_delay: float) -> None: + proxy_client.state = ConnectionState.SHUTTING_DOWN + + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", check_health_mock) + monkeypatch.setattr(client_mod.asyncio, "sleep", _fake_sleep) + + await proxy_client._keepalive_loop() + + proxy_client.gateway_transport.send.assert_awaited_once() + payload = proxy_client.gateway_transport.send.await_args.args[0] + assert MessageType.HEARTBEAT.value in payload + + +@pytest.mark.asyncio +async def test_keepalive_loop_recovers_and_reregisters_when_mcp_returns(proxy_client, monkeypatch) -> None: + """Recovered MCP server should reconnect gateway if needed and re-register.""" + proxy_client.state = ConnectionState.CONNECTED + proxy_client.keepalive_interval = 0 + proxy_client._mcp_server_healthy = False + proxy_client._consecutive_mcp_failures = 2 + proxy_client.gateway_transport.is_connected.return_value = False + check_health_mock = AsyncMock(return_value=True) + register_mock = AsyncMock() + + async def _fake_sleep(_delay: float) -> None: + proxy_client.state = ConnectionState.SHUTTING_DOWN + + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", check_health_mock) + monkeypatch.setattr(proxy_client, "_register", register_mock) + monkeypatch.setattr(client_mod.asyncio, "sleep", _fake_sleep) + + await proxy_client._keepalive_loop() + + proxy_client.gateway_transport.connect.assert_awaited_once() + register_mock.assert_awaited_once() + assert proxy_client._mcp_server_healthy is True + assert proxy_client._consecutive_mcp_failures == 0 + proxy_client.gateway_transport.send.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_keepalive_loop_continues_when_gateway_reconnect_fails(proxy_client, monkeypatch) -> None: + """Gateway reconnect failure during recovery should skip heartbeat send.""" + proxy_client.state = ConnectionState.CONNECTED + proxy_client.keepalive_interval = 0 + proxy_client._mcp_server_healthy = False + proxy_client.gateway_transport.is_connected.return_value = False + proxy_client.gateway_transport.connect.side_effect = RuntimeError("cannot reconnect") + check_health_mock = AsyncMock(return_value=True) + + async def _fake_sleep(_delay: float) -> None: + proxy_client.state = ConnectionState.SHUTTING_DOWN + + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", check_health_mock) + monkeypatch.setattr(client_mod.asyncio, "sleep", _fake_sleep) + + await proxy_client._keepalive_loop() + + proxy_client.gateway_transport.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_keepalive_loop_continues_when_reregister_fails(proxy_client, monkeypatch) -> None: + """Re-registration failure during recovery should skip heartbeat send.""" + proxy_client.state = ConnectionState.CONNECTED + proxy_client.keepalive_interval = 0 + proxy_client._mcp_server_healthy = False + proxy_client.gateway_transport.is_connected.return_value = True + check_health_mock = AsyncMock(return_value=True) + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", check_health_mock) + monkeypatch.setattr(proxy_client, "_register", AsyncMock(side_effect=RuntimeError("register fail"))) + + async def _fake_sleep(_delay: float) -> None: + proxy_client.state = ConnectionState.SHUTTING_DOWN + + monkeypatch.setattr(client_mod.asyncio, "sleep", _fake_sleep) + + await proxy_client._keepalive_loop() + + proxy_client.gateway_transport.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_keepalive_loop_breaks_when_heartbeat_send_fails(proxy_client, monkeypatch) -> None: + """Heartbeat send failures should break the loop.""" + proxy_client.state = ConnectionState.CONNECTED + proxy_client.keepalive_interval = 0 + proxy_client.gateway_transport.send.side_effect = RuntimeError("send fail") + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=True)) + + async def _fake_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(client_mod.asyncio, "sleep", _fake_sleep) + + await proxy_client._keepalive_loop() + + +@pytest.mark.asyncio +async def test_keepalive_loop_marks_first_unhealthy_failure(proxy_client, monkeypatch) -> None: + """First unhealthy check should flip the health flag and increment failures.""" + proxy_client.state = ConnectionState.CONNECTED + proxy_client.keepalive_interval = 0 + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=False)) + + async def _fake_sleep(_delay: float) -> None: + proxy_client.state = ConnectionState.SHUTTING_DOWN + + monkeypatch.setattr(client_mod.asyncio, "sleep", _fake_sleep) + + await proxy_client._keepalive_loop() + + assert proxy_client._mcp_server_healthy is False + assert proxy_client._consecutive_mcp_failures == 1 + proxy_client.gateway_transport.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_keepalive_loop_tracks_ongoing_unhealthy_failures(proxy_client, monkeypatch) -> None: + """Repeated unhealthy checks should continue incrementing failure count.""" + proxy_client.state = ConnectionState.CONNECTED + proxy_client.keepalive_interval = 0 + proxy_client._mcp_server_healthy = False + proxy_client._consecutive_mcp_failures = 1 + monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=False)) + + async def _fake_sleep(_delay: float) -> None: + proxy_client.state = ConnectionState.SHUTTING_DOWN + + monkeypatch.setattr(client_mod.asyncio, "sleep", _fake_sleep) + + await proxy_client._keepalive_loop() + + assert proxy_client._consecutive_mcp_failures == 2 + + +@pytest.mark.asyncio +async def test_send_error_response_swallow_send_failures(proxy_client) -> None: + """Gateway send failures while returning an error should be swallowed.""" + proxy_client.gateway_transport.send.side_effect = RuntimeError("cannot send") + + await proxy_client._send_error_response({"jsonrpc": "2.0", "id": 34}, "broken") diff --git a/tests/test_mcp_reverse_proxy_sse_adapter.py b/tests/test_mcp_reverse_proxy_sse_adapter.py new file mode 100644 index 0000000..96f523b --- /dev/null +++ b/tests/test_mcp_reverse_proxy_sse_adapter.py @@ -0,0 +1,373 @@ +"""Location: ./mcp_reverse_proxy/tests/test_mcp_reverse_proxy_sse_adapter.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Unit tests for the SSE reverse proxy transport adapter. +""" + +# Future +from __future__ import annotations + +# Standard +import asyncio +import ssl +from collections.abc import AsyncIterator +from unittest.mock import AsyncMock, Mock + +# Third-Party +import httpx +import pytest + +# First-Party +from mcp_reverse_proxy.transports.sse_adapter import SseAdapter + + +class FakeSseResponse: + """Minimal async SSE response fake.""" + + def __init__(self, lines: list[str], status_code: int = 200) -> None: + self._lines = lines + self.status_code = status_code + self.headers: dict[str, str] = {} + + def raise_for_status(self) -> None: + """Mimic successful HTTP responses.""" + + async def aiter_lines(self) -> AsyncIterator[str]: + """Yield configured SSE lines.""" + for line in self._lines: + yield line + + +class FakeStreamContext: + """Async context manager wrapper for SSE response fakes.""" + + def __init__(self, response: FakeSseResponse) -> None: + self._response = response + + async def __aenter__(self) -> FakeSseResponse: + """Return the fake response.""" + return self._response + + async def __aexit__(self, exc_type, exc, tb) -> bool: + """Do not suppress exceptions.""" + return False + + +@pytest.mark.asyncio +async def test_start_creates_http_client_and_receive_task_for_http(monkeypatch) -> None: + """Start should initialize the client and spawn the receive loop task.""" + created_clients: list[dict[str, object]] = [] + created_tasks: list[asyncio.Task[None]] = [] + + class FakeClient: + """Async client fake capturing constructor parameters.""" + + def __init__(self, **kwargs) -> None: + created_clients.append(kwargs) + + async def aclose(self) -> None: + """No-op close.""" + + original_create_task = asyncio.create_task + + def capture_task(coro): + task = original_create_task(coro) + created_tasks.append(task) + return task + + monkeypatch.setattr("mcp_reverse_proxy.transports.sse_adapter.httpx.AsyncClient", FakeClient) + monkeypatch.setattr("mcp_reverse_proxy.transports.sse_adapter.asyncio.create_task", capture_task) + + adapter = SseAdapter("http://server.example/sse", timeout=12.5) + + await adapter.start() + + assert adapter._connected is True + assert len(created_clients) == 1 + assert created_clients[0]["http2"] is True + assert created_clients[0]["verify"] is False + assert isinstance(created_clients[0]["timeout"], httpx.Timeout) + assert adapter._receive_task is created_tasks[0] + + adapter._receive_task.cancel() + with pytest.raises(asyncio.CancelledError): + await adapter._receive_task + + +@pytest.mark.asyncio +async def test_start_configures_https_ssl_context_with_custom_cert(monkeypatch) -> None: + """HTTPS start should use an SSL context created from the provided CA data.""" + fake_ssl_context = Mock(spec=ssl.SSLContext) + fake_ssl_context.check_hostname = True + fake_ssl_context.verify_mode = ssl.CERT_REQUIRED + fake_ssl_context.load_verify_locations = Mock() # Mock the load_verify_locations method + async_client_mock = Mock() + receive_task = asyncio.create_task(asyncio.sleep(10)) + create_task_mock = Mock(return_value=receive_task) + + # Mock load_cert_data to return test certificate data + test_cert_data = "MOCK_CERT_DATA" + monkeypatch.setattr("mcp_reverse_proxy.transports.sse_adapter.load_cert_data", lambda _cert: test_cert_data) + + # Mock SSLContext constructor to return our fake context + monkeypatch.setattr("mcp_reverse_proxy.transports.sse_adapter.ssl.SSLContext", lambda _protocol: fake_ssl_context) + monkeypatch.setattr("mcp_reverse_proxy.transports.sse_adapter.httpx.AsyncClient", async_client_mock) + monkeypatch.setattr("mcp_reverse_proxy.transports.sse_adapter.asyncio.create_task", create_task_mock) + + adapter = SseAdapter("https://secure.example/sse", cert="test_cert_input") + + await adapter.start() + + # Verify load_verify_locations was called with the cert data + fake_ssl_context.load_verify_locations.assert_called_once_with(cadata=test_cert_data) + async_client_mock.assert_called_once() + assert async_client_mock.call_args.kwargs["verify"] is fake_ssl_context + assert fake_ssl_context.check_hostname is True + assert fake_ssl_context.verify_mode == ssl.CERT_REQUIRED + + receive_task.cancel() + with pytest.raises(asyncio.CancelledError): + await receive_task + + +@pytest.mark.asyncio +async def test_stop_cancels_task_closes_client_and_clears_session_state() -> None: + """Stop should cancel the receive task, close the client, and clear session state.""" + adapter = SseAdapter("http://server.example/sse") + adapter._connected = True + adapter._session_id = "session-1" + adapter._message_endpoint = "http://server.example/messages" + adapter._protocol_version = "2025-11-05" + + adapter._receive_task = asyncio.create_task(asyncio.sleep(10)) + adapter._client = AsyncMock() + adapter._shutdown_event.clear() + + await adapter.stop() + + assert adapter._connected is False + assert adapter._shutdown_event.is_set() + adapter._client = None + assert adapter._session_id is None + assert adapter._message_endpoint is None + assert adapter._protocol_version is None + + +@pytest.mark.asyncio +async def test_send_posts_payload_with_auth_session_and_protocol_headers() -> None: + """Send should POST JSON using auth, session, and protocol headers.""" + adapter = SseAdapter("http://server.example/sse") + adapter._connected = True + adapter._message_endpoint = "http://server.example/messages" + adapter._session_id = "session-1" + adapter._protocol_version = "2025-11-05" + adapter._auth_headers = {"Authorization": "Bearer token"} + + response = Mock() + response.headers = {"mcp-session-id": "session-2"} + response.status_code = 200 + response.raise_for_status = Mock() + + client = AsyncMock() + client.post = AsyncMock(return_value=response) + adapter._client = client + + await adapter.send('{"jsonrpc":"2.0","id":1}') + + client.post.assert_awaited_once() + assert client.post.await_args.args[0] == "http://server.example/messages" + assert client.post.await_args.kwargs["content"] == '{"jsonrpc":"2.0","id":1}' + assert client.post.await_args.kwargs["headers"] == { + "Content-Type": "application/json", + "Accept": "application/json", + "Authorization": "Bearer token", + "mcp-session-id": "session-1", + "mcp-protocol-version": "2025-11-05", + } + assert adapter._session_id == "session-1" + + +@pytest.mark.asyncio +async def test_send_stores_session_id_from_first_successful_response() -> None: + """First successful POST should capture the session ID from response headers.""" + adapter = SseAdapter("http://server.example/sse") + adapter._connected = True + adapter._message_endpoint = "http://server.example/messages" + + response = Mock() + response.headers = {"mcp-session-id": "new-session"} + response.status_code = 200 + response.raise_for_status = Mock() + + adapter._client = AsyncMock() + adapter._client.post = AsyncMock(return_value=response) + + await adapter.send('{"jsonrpc":"2.0","id":1}') + + assert adapter._session_id == "new-session" + + +@pytest.mark.asyncio +async def test_send_clears_session_state_on_http_404() -> None: + """A 404 from POST should clear session-related state before raising.""" + adapter = SseAdapter("http://server.example/sse") + adapter._connected = True + adapter._client = AsyncMock() + adapter._message_endpoint = "http://server.example/messages" + adapter._session_id = "session-1" + adapter._protocol_version = "2025-11-05" + + request = httpx.Request("POST", adapter._message_endpoint) + response = httpx.Response(404, request=request) + error = httpx.HTTPStatusError("missing", request=request, response=response) + adapter._client.post = AsyncMock(side_effect=error) + + with pytest.raises(RuntimeError, match="Failed to send message"): + await adapter.send('{"jsonrpc":"2.0","id":1}') + + assert adapter._session_id is None + assert adapter._message_endpoint is None + assert adapter._protocol_version is None + + +def test_set_authentication_supports_basic_bearer_and_passthrough_headers() -> None: + """Authentication helper should normalize supported auth styles.""" + adapter = SseAdapter("http://server.example/sse") + + adapter.set_authentication({"username": "alice", "password": "secret"}, "basic") + assert adapter._auth_headers == {"Authorization": "Basic YWxpY2U6c2VjcmV0"} + + adapter.set_authentication({"token": "abc123"}, "bearer") + assert adapter._auth_headers == {"Authorization": "Bearer abc123"} + + adapter.set_authentication({"X-Api-Key": "key"}, "custom") + assert adapter._auth_headers == {"X-Api-Key": "key"} + + +@pytest.mark.asyncio +async def test_receive_sse_stream_processes_events_and_forwards_auth_headers(monkeypatch) -> None: + """SSE receive loop should parse streamed events and pass connection headers.""" + adapter = SseAdapter("http://server.example/sse") + adapter._connected = True + adapter._auth_headers = {"Authorization": "Bearer token"} + + processed_events: list[tuple[str, str]] = [] + + async def fake_process(event_type: str, data: str) -> None: + processed_events.append((event_type, data)) + + stream_mock = Mock( + return_value=FakeStreamContext( + FakeSseResponse( + [ + "event: endpoint", + "data: /messages?session_id=abc", + "", + "retry: 5000", + ": comment", + "event: message", + 'data: {"jsonrpc":"2.0"}', + "", + ] + ) + ) + ) + + adapter._client = Mock() + adapter._client.stream = stream_mock + monkeypatch.setattr(adapter, "_process_sse_event", fake_process) + + await adapter._receive_sse_stream() + + stream_mock.assert_called_once_with( + "GET", + "http://server.example/sse", + headers={ + "Accept": "text/event-stream", + "Cache-Control": "no-cache", + "Authorization": "Bearer token", + }, + ) + assert processed_events == [ + ("endpoint", "/messages?session_id=abc"), + ("message", '{"jsonrpc":"2.0"}'), + ] + assert adapter._connected is False + + +@pytest.mark.asyncio +async def test_receive_sse_stream_notifies_handlers_on_http_error() -> None: + """HTTP stream failures while connected should notify registered handlers.""" + adapter = SseAdapter("http://server.example/sse") + adapter._connected = True + + request = httpx.Request("GET", adapter.server_url) + stream_error = httpx.ConnectError("boom", request=request) + + handler = AsyncMock() + adapter.add_message_handler(handler) + + class RaisingStreamContext: + """Context manager that raises the configured HTTP error.""" + + async def __aenter__(self): + raise stream_error + + async def __aexit__(self, exc_type, exc, tb) -> bool: + return False + + adapter._client = Mock() + adapter._client.stream = Mock(return_value=RaisingStreamContext()) + + await adapter._receive_sse_stream() + + handler.assert_awaited_once() + payload = handler.await_args.args[0] + assert "SSE connection lost: boom" in payload + assert adapter._connected is False + + +@pytest.mark.asyncio +async def test_process_sse_event_endpoint_builds_absolute_url_and_extracts_session() -> None: + """Endpoint events should resolve relative URLs and extract session query parameters.""" + adapter = SseAdapter("https://server.example/sse") + + await adapter._process_sse_event("endpoint", "/messages?session_id=session-123") + + assert adapter._message_endpoint == "https://server.example/messages?session_id=session-123" + assert adapter._session_id == "session-123" + + +@pytest.mark.asyncio +async def test_process_sse_event_message_sets_protocol_and_calls_all_handlers() -> None: + """Message events should negotiate protocol version and notify handlers.""" + adapter = SseAdapter("http://server.example/sse") + handler_one = AsyncMock() + handler_two = AsyncMock() + adapter.add_message_handler(handler_one) + adapter.add_message_handler(handler_two) + + payload = '{"jsonrpc":"2.0","result":{"protocolVersion":"2025-11-05"}}' + + await adapter._process_sse_event("message", payload) + + assert adapter._protocol_version == "2025-11-05" + handler_one.assert_awaited_once_with(payload) + handler_two.assert_awaited_once_with(payload) + + +@pytest.mark.asyncio +async def test_process_sse_event_error_forwards_to_handlers_even_if_one_fails() -> None: + """Error events should continue forwarding after an individual handler failure.""" + adapter = SseAdapter("http://server.example/sse") + failing_handler = AsyncMock(side_effect=RuntimeError("handler failed")) + successful_handler = AsyncMock() + adapter.add_message_handler(failing_handler) + adapter.add_message_handler(successful_handler) + + await adapter._process_sse_event("error", '{"error":"bad"}') + + failing_handler.assert_awaited_once_with('{"error":"bad"}') + successful_handler.assert_awaited_once_with('{"error":"bad"}') diff --git a/tests/test_mcp_reverse_proxy_stdio_adapter.py b/tests/test_mcp_reverse_proxy_stdio_adapter.py new file mode 100644 index 0000000..2320c67 --- /dev/null +++ b/tests/test_mcp_reverse_proxy_stdio_adapter.py @@ -0,0 +1,268 @@ +"""Location: ./mcp_reverse_proxy/tests/test_mcp_reverse_proxy_stdio_adapter.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Unit tests for the stdio reverse proxy transport adapter. +""" + +# Future +from __future__ import annotations + +# Standard +import asyncio +from unittest.mock import AsyncMock, Mock + +# Third-Party +import pytest + +# First-Party +from mcp_reverse_proxy.transports.stdio_adapter import StdioAdapter + + +class FakeStdout: + """Simple stdout fake backed by queued byte lines.""" + + def __init__(self, lines: list[bytes]) -> None: + self._lines = list(lines) + + async def readline(self) -> bytes: + """Return the next configured line or EOF.""" + if self._lines: + return self._lines.pop(0) + return b"" + + +class FakeStdin: + """Simple stdin fake recording written bytes.""" + + def __init__(self) -> None: + self.writes: list[bytes] = [] + self.drain = AsyncMock() + + def write(self, data: bytes) -> None: + """Record write calls.""" + self.writes.append(data) + + +@pytest.mark.asyncio +async def test_start_launches_subprocess_and_starts_stdout_reader(monkeypatch) -> None: + """Start should spawn the subprocess and create the stdout reader task.""" + stdin = FakeStdin() + stdout = FakeStdout([]) + process = Mock() + process.stdin = stdin + process.stdout = stdout + process.stderr = None + process.pid = 1234 + process.returncode = None + + create_subprocess_mock = AsyncMock(return_value=process) + stdout_task = asyncio.create_task(asyncio.sleep(10)) + create_task_mock = Mock(return_value=stdout_task) + sleep_mock = AsyncMock() + + monkeypatch.setattr("mcp_reverse_proxy.transports.stdio_adapter.asyncio.create_subprocess_exec", create_subprocess_mock) + monkeypatch.setattr("mcp_reverse_proxy.transports.stdio_adapter.asyncio.create_task", create_task_mock) + monkeypatch.setattr("mcp_reverse_proxy.transports.stdio_adapter.asyncio.sleep", sleep_mock) + + adapter = StdioAdapter("python -m example") + + await adapter.start() + + create_subprocess_mock.assert_awaited_once() + assert create_subprocess_mock.await_args.args[:3] == ("python", "-m", "example") + assert adapter.process is process + assert adapter._stdout_reader_task is create_task_mock.return_value + sleep_mock.assert_awaited_once_with(0.5) + + stdout_task.cancel() + with pytest.raises(asyncio.CancelledError): + await stdout_task + + +@pytest.mark.asyncio +async def test_start_raises_runtime_error_when_command_missing(monkeypatch) -> None: + """Missing command errors should be wrapped in a RuntimeError.""" + create_subprocess_mock = AsyncMock(side_effect=FileNotFoundError("missing")) + monkeypatch.setattr("mcp_reverse_proxy.transports.stdio_adapter.asyncio.create_subprocess_exec", create_subprocess_mock) + + adapter = StdioAdapter("does-not-exist") + + with pytest.raises(RuntimeError, match="Command not found: does-not-exist"): + await adapter.start() + + +@pytest.mark.asyncio +async def test_start_raises_runtime_error_when_process_exits_immediately(monkeypatch) -> None: + """Processes that die during startup should raise a clear RuntimeError.""" + process = Mock() + process.stdin = FakeStdin() + process.stdout = FakeStdout([]) + process.pid = 1234 + process.returncode = 7 + + monkeypatch.setattr( + "mcp_reverse_proxy.transports.stdio_adapter.asyncio.create_subprocess_exec", + AsyncMock(return_value=process), + ) + stdout_task = asyncio.create_task(asyncio.sleep(10)) + monkeypatch.setattr("mcp_reverse_proxy.transports.stdio_adapter.asyncio.create_task", Mock(return_value=stdout_task)) + monkeypatch.setattr("mcp_reverse_proxy.transports.stdio_adapter.asyncio.sleep", AsyncMock()) + + adapter = StdioAdapter("python -m crash") + + with pytest.raises(RuntimeError, match="Subprocess terminated immediately after start"): + await adapter.start() + + stdout_task.cancel() + with pytest.raises(asyncio.CancelledError): + await stdout_task + + +@pytest.mark.asyncio +async def test_stop_cancels_reader_and_terminates_running_process(monkeypatch) -> None: + """Stop should cancel the reader task and terminate a live process.""" + adapter = StdioAdapter("python -m example") + process = Mock() + process.pid = 4321 + process.returncode = None + process.wait = AsyncMock(side_effect=[None]) + + reader_task = asyncio.create_task(asyncio.sleep(10)) + adapter.process = process + adapter._stdout_reader_task = reader_task + + wait_for_mock = AsyncMock(return_value=None) + monkeypatch.setattr("mcp_reverse_proxy.transports.stdio_adapter.asyncio.wait_for", wait_for_mock) + + await adapter.stop() + + process.terminate.assert_called_once_with() + wait_for_mock.assert_awaited_once() + process.kill.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_stop_force_kills_process_after_timeout(monkeypatch) -> None: + """Stop should kill the process when graceful termination times out.""" + adapter = StdioAdapter("python -m example") + process = Mock() + process.pid = 4321 + process.returncode = None + process.wait = AsyncMock(side_effect=[None, None]) + + adapter.process = process + adapter._stdout_reader_task = None + + async def fake_wait_for(awaitable, timeout: float): # noqa: ARG001 - signature must match asyncio.wait_for (keyword call site) + await awaitable + raise TimeoutError + + monkeypatch.setattr("mcp_reverse_proxy.transports.stdio_adapter.asyncio.wait_for", fake_wait_for) + + await adapter.stop() + + process.terminate.assert_called_once_with() + process.kill.assert_called_once_with() + assert process.wait.await_count == 2 + + +@pytest.mark.asyncio +async def test_send_writes_newline_delimited_message_and_drains() -> None: + """Send should write a newline-delimited message to stdin.""" + adapter = StdioAdapter("python -m example") + stdin = FakeStdin() + process = Mock() + process.stdin = stdin + process.returncode = None + adapter.process = process + + await adapter.send('{"jsonrpc":"2.0","id":1}') + + assert stdin.writes == [b'{"jsonrpc":"2.0","id":1}\n'] + stdin.drain.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_send_raises_when_process_not_running() -> None: + """Send should fail when no subprocess is available.""" + adapter = StdioAdapter("python -m example") + + with pytest.raises(RuntimeError, match="Subprocess not running"): + await adapter.send("{}") + + +@pytest.mark.asyncio +async def test_send_raises_when_process_has_already_terminated() -> None: + """Terminated processes should fail before attempting to write to stdin.""" + adapter = StdioAdapter("python -m example") + stdin = FakeStdin() + process = Mock() + process.stdin = stdin + process.returncode = 9 + adapter.process = process + + with pytest.raises(RuntimeError, match="Subprocess terminated with exit code 9"): + await adapter.send("{}") + + +@pytest.mark.asyncio +async def test_send_wraps_broken_pipe_with_return_code() -> None: + """Broken pipes during write should be surfaced as runtime errors with exit code context.""" + adapter = StdioAdapter("python -m example") + stdin = FakeStdin() + stdin.drain = AsyncMock(side_effect=BrokenPipeError()) + process = Mock() + process.stdin = stdin + process.returncode = None + adapter.process = process + + with pytest.raises(RuntimeError, match="Subprocess stdin closed"): + await adapter.send("{}") + + +@pytest.mark.asyncio +async def test_read_stdout_forwards_non_empty_messages_and_skips_blank_lines() -> None: + """Stdout reader should ignore blank lines and forward decoded messages.""" + adapter = StdioAdapter("python -m example") + handler = AsyncMock() + adapter.add_message_handler(handler) + + process = Mock() + process.stdout = FakeStdout([b'{"jsonrpc":"2.0"}\n', b"\n", b'{"jsonrpc":"2.0","id":1}\n', b""]) + adapter.process = process + + await adapter._read_stdout() + + assert handler.await_count == 2 + handler.assert_any_await('{"jsonrpc":"2.0"}') + handler.assert_any_await('{"jsonrpc":"2.0","id":1}') + + +@pytest.mark.asyncio +async def test_read_stdout_continues_when_handler_raises() -> None: + """A failing handler should not prevent later handlers from running.""" + adapter = StdioAdapter("python -m example") + failing_handler = AsyncMock(side_effect=RuntimeError("boom")) + successful_handler = AsyncMock() + adapter.add_message_handler(failing_handler) + adapter.add_message_handler(successful_handler) + + process = Mock() + process.stdout = FakeStdout([b'{"jsonrpc":"2.0"}\n', b""]) + adapter.process = process + + await adapter._read_stdout() + + failing_handler.assert_awaited_once_with('{"jsonrpc":"2.0"}') + successful_handler.assert_awaited_once_with('{"jsonrpc":"2.0"}') + + +@pytest.mark.asyncio +async def test_read_stdout_returns_immediately_without_stdout() -> None: + """Reader should no-op when process stdout is unavailable.""" + adapter = StdioAdapter("python -m example") + adapter.process = Mock(stdout=None) + + await adapter._read_stdout() diff --git a/tests/test_mcp_reverse_proxy_streamablehttp_adapter.py b/tests/test_mcp_reverse_proxy_streamablehttp_adapter.py new file mode 100644 index 0000000..e833a61 --- /dev/null +++ b/tests/test_mcp_reverse_proxy_streamablehttp_adapter.py @@ -0,0 +1,296 @@ +"""Location: ./mcp_reverse_proxy/tests/test_mcp_reverse_proxy_streamablehttp_adapter.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Unit tests for the streamable HTTP reverse proxy transport adapter. +""" + +# Future +from __future__ import annotations + +# Standard +import asyncio +import ssl +from unittest.mock import AsyncMock, Mock + +# Third-Party +import httpx +import pytest + +# First-Party +from mcp_reverse_proxy.transports.streamablehttp_adapter import SessionExpiredError, StreamableHttpAdapter + + +@pytest.mark.asyncio +async def test_start_creates_http_client_sets_endpoint_and_receive_task(monkeypatch) -> None: + """Start should initialize the client, message endpoint, and receive task.""" + created_tasks: list[asyncio.Task[None]] = [] + + class FakeClient: + """Minimal async client fake.""" + + async def aclose(self) -> None: + """No-op close.""" + + async_client_mock = Mock(return_value=FakeClient()) + original_create_task = asyncio.create_task + + def capture_task(coro): + task = original_create_task(coro) + created_tasks.append(task) + return task + + monkeypatch.setattr("mcp_reverse_proxy.transports.streamablehttp_adapter.httpx.AsyncClient", async_client_mock) + monkeypatch.setattr("mcp_reverse_proxy.transports.streamablehttp_adapter.asyncio.create_task", capture_task) + + adapter = StreamableHttpAdapter("http://server.example/mcp", timeout=12.5) + + await adapter.start() + + assert adapter._connected is True + assert adapter._message_endpoint == "http://server.example/mcp" + assert async_client_mock.call_args.kwargs["http2"] is True + assert async_client_mock.call_args.kwargs["verify"] is False + assert isinstance(async_client_mock.call_args.kwargs["timeout"], httpx.Timeout) + assert adapter._receive_task is created_tasks[0] + + adapter._receive_task.cancel() + with pytest.raises(asyncio.CancelledError): + await adapter._receive_task + + +@pytest.mark.asyncio +async def test_start_configures_https_ssl_context_with_custom_cert(monkeypatch) -> None: + """HTTPS start should use an SSL context created from the provided CA data.""" + fake_ssl_context = Mock(spec=ssl.SSLContext) + fake_ssl_context.check_hostname = True + fake_ssl_context.verify_mode = ssl.CERT_REQUIRED + fake_ssl_context.load_verify_locations = Mock() # Mock the load_verify_locations method + receive_task = asyncio.create_task(asyncio.sleep(10)) + async_client_mock = Mock() + + # Mock load_cert_data to return test certificate data + test_cert_data = "MOCK_CERT_DATA" + monkeypatch.setattr("mcp_reverse_proxy.transports.streamablehttp_adapter.load_cert_data", + lambda _cert: test_cert_data) + + # Mock SSLContext constructor to return our fake context + monkeypatch.setattr("mcp_reverse_proxy.transports.streamablehttp_adapter.ssl.SSLContext", + lambda _protocol: fake_ssl_context) + monkeypatch.setattr("mcp_reverse_proxy.transports.streamablehttp_adapter.httpx.AsyncClient", async_client_mock) + monkeypatch.setattr( + "mcp_reverse_proxy.transports.streamablehttp_adapter.asyncio.create_task", + Mock(return_value=receive_task), + ) + + adapter = StreamableHttpAdapter("https://secure.example/mcp", cert="test_cert_input") + + await adapter.start() + + # Verify load_verify_locations was called with the cert data + fake_ssl_context.load_verify_locations.assert_called_once_with(cadata=test_cert_data) + assert async_client_mock.call_args.kwargs["verify"] is fake_ssl_context + assert fake_ssl_context.check_hostname is True + assert fake_ssl_context.verify_mode == ssl.CERT_REQUIRED + + receive_task.cancel() + with pytest.raises(asyncio.CancelledError): + await receive_task + + +@pytest.mark.asyncio +async def test_stop_cancels_task_closes_client_and_clears_session_state() -> None: + """Stop should cancel the receive task, close the client, and clear session state.""" + adapter = StreamableHttpAdapter("http://server.example/mcp") + adapter._connected = True + adapter._session_id = "session-1" + adapter._message_endpoint = "http://server.example/mcp" + adapter._protocol_version = "2025-11-05" + adapter._receive_task = asyncio.create_task(asyncio.sleep(10)) + adapter._client = AsyncMock() + + await adapter.stop() + + assert adapter._connected is False + assert adapter._session_id is None + assert adapter._message_endpoint is None + assert adapter._protocol_version is None + assert adapter._client is None + + +@pytest.mark.asyncio +async def test_send_initialize_stores_session_protocol_and_notifies_handlers() -> None: + """Initialize response should set session state and forward inline JSON responses.""" + adapter = StreamableHttpAdapter("http://server.example/mcp") + adapter._connected = True + adapter._client = AsyncMock() + handler = AsyncMock() + adapter.add_message_handler(handler) + + response = Mock() + response.headers = {"mcp-session-id": "session-1"} + response.status_code = 200 + response.content = b'{"jsonrpc":"2.0","result":{"protocolVersion":"2025-11-05"}}' + response.text = '{"jsonrpc":"2.0","result":{"protocolVersion":"2025-11-05"}}' + response.raise_for_status = Mock() + adapter._client.post = AsyncMock(return_value=response) + + await adapter.send('{"jsonrpc":"2.0","id":1,"method":"initialize"}') + + assert adapter._session_id == "session-1" + assert adapter._protocol_version == "2025-11-05" + handler.assert_awaited_once_with('{"jsonrpc":"2.0","result":{"protocolVersion":"2025-11-05"}}') + headers = adapter._client.post.await_args.kwargs["headers"] + assert "mcp-session-id" not in headers + assert "mcp-protocol-version" not in headers + + +@pytest.mark.asyncio +async def test_send_with_existing_session_adds_session_headers_and_parses_sse_response() -> None: + """Established sessions should send MCP headers and parse SSE-formatted inline responses.""" + adapter = StreamableHttpAdapter("http://server.example/mcp") + adapter._connected = True + adapter._client = AsyncMock() + adapter._session_id = "session-1" + adapter._protocol_version = "2025-11-05" + adapter._auth_headers = {"Authorization": "Bearer token"} + handler = AsyncMock() + adapter.add_message_handler(handler) + + response = Mock() + response.headers = {} + response.status_code = 200 + response.content = b'event: message\ndata: {"jsonrpc":"2.0","id":2,"result":{"ok":true}}\n\n' + response.text = 'event: message\ndata: {"jsonrpc":"2.0","id":2,"result":{"ok":true}}\n\n' + response.raise_for_status = Mock() + adapter._client.post = AsyncMock(return_value=response) + + await adapter.send('{"jsonrpc":"2.0","id":2,"method":"tools/call"}') + + handler.assert_awaited_once_with('{"jsonrpc":"2.0","id":2,"result":{"ok":true}}') + headers = adapter._client.post.await_args.kwargs["headers"] + assert headers["Authorization"] == "Bearer token" + assert headers["mcp-session-id"] == "session-1" + assert headers["mcp-protocol-version"] == "2025-11-05" + + +@pytest.mark.asyncio +async def test_send_raises_for_non_initialize_request_without_session() -> None: + """Non-initialize messages without a session should fail fast.""" + adapter = StreamableHttpAdapter("http://server.example/mcp") + adapter._connected = True + adapter._client = AsyncMock() + + with pytest.raises(RuntimeError, match="No valid session"): + await adapter.send('{"jsonrpc":"2.0","id":2,"method":"tools/call"}') + + +@pytest.mark.asyncio +async def test_send_raises_runtime_error_when_not_connected() -> None: + """Send should fail if the HTTP client connection has not been started.""" + adapter = StreamableHttpAdapter("http://server.example/mcp") + + with pytest.raises(RuntimeError, match="Not connected to MCP server"): + await adapter.send('{"jsonrpc":"2.0","id":1}') + + +@pytest.mark.asyncio +async def test_send_retries_initialize_after_404_and_updates_session_state() -> None: + """Initialize requests should retry once without session headers after a 404.""" + adapter = StreamableHttpAdapter("http://server.example/mcp") + adapter._connected = True + adapter._client = AsyncMock() + adapter._session_id = "stale-session" + adapter._protocol_version = "old-version" + handler = AsyncMock() + adapter.add_message_handler(handler) + + request = httpx.Request("POST", "http://server.example/mcp") + first_response = httpx.Response(404, request=request) + http_error = httpx.HTTPStatusError("missing", request=request, response=first_response) + + retry_response = Mock() + retry_response.headers = {"mcp-session-id": "new-session"} + retry_response.status_code = 200 + retry_response.content = b'{"jsonrpc":"2.0","result":{"protocolVersion":"2025-11-05"}}' + retry_response.text = '{"jsonrpc":"2.0","result":{"protocolVersion":"2025-11-05"}}' + retry_response.raise_for_status = Mock() + + adapter._client.post = AsyncMock(side_effect=[http_error, retry_response]) + + await adapter.send('{"jsonrpc":"2.0","id":1,"method":"initialize"}') + + assert adapter._session_id == "new-session" + assert adapter._protocol_version == "2025-11-05" + assert adapter._client.post.await_count == 2 + handler.assert_awaited_once_with('{"jsonrpc":"2.0","result":{"protocolVersion":"2025-11-05"}}') + + +@pytest.mark.asyncio +async def test_send_raises_session_expired_for_non_initialize_404() -> None: + """Non-initialize 404 responses should clear session state and raise SessionExpiredError.""" + adapter = StreamableHttpAdapter("http://server.example/mcp") + adapter._connected = True + adapter._client = AsyncMock() + adapter._session_id = "session-1" + adapter._protocol_version = "2025-11-05" + + request = httpx.Request("POST", "http://server.example/mcp") + response = httpx.Response(404, request=request) + http_error = httpx.HTTPStatusError("missing", request=request, response=response) + adapter._client.post = AsyncMock(side_effect=http_error) + + with pytest.raises(SessionExpiredError, match="session expired"): + await adapter.send('{"jsonrpc":"2.0","id":2,"method":"tools/call"}') + + assert adapter._session_id is None + assert adapter._protocol_version is None + + +@pytest.mark.asyncio +async def test_send_wraps_generic_http_error() -> None: + """Non-status HTTP client failures should be wrapped in RuntimeError.""" + adapter = StreamableHttpAdapter("http://server.example/mcp") + adapter._connected = True + adapter._client = AsyncMock() + + request = httpx.Request("POST", "http://server.example/mcp") + adapter._client.post = AsyncMock(side_effect=httpx.ConnectError("boom", request=request)) + + with pytest.raises(RuntimeError, match="Failed to send message"): + await adapter.send('{"jsonrpc":"2.0","id":1,"method":"initialize"}') + + +def test_set_authentication_supports_basic_bearer_and_passthrough_headers() -> None: + """Authentication helper should normalize supported auth styles.""" + adapter = StreamableHttpAdapter("http://server.example/mcp") + + adapter.set_authentication({"username": "alice", "password": "secret"}, "basic") + assert adapter._auth_headers == {"Authorization": "Basic YWxpY2U6c2VjcmV0"} + + adapter.set_authentication({"token": "abc123"}, "bearer") + assert adapter._auth_headers == {"Authorization": "Bearer abc123"} + + adapter.set_authentication({"X-Api-Key": "key"}, "custom") + assert adapter._auth_headers == {"X-Api-Key": "key"} + + +@pytest.mark.asyncio +async def test_receive_stream_sleeps_until_cancelled(monkeypatch) -> None: + """Receive loop should keep sleeping while connected and propagate cancellation.""" + adapter = StreamableHttpAdapter("http://server.example/mcp") + adapter._connected = True + call_count = 0 + + async def fake_sleep(_delay: float) -> None: + nonlocal call_count + call_count += 1 + raise asyncio.CancelledError + + monkeypatch.setattr("mcp_reverse_proxy.transports.streamablehttp_adapter.asyncio.sleep", fake_sleep) + + with pytest.raises(asyncio.CancelledError): + await adapter._receive_stream() + + assert call_count == 1 diff --git a/tests/test_mcp_reverse_proxy_websocket_adapter.py b/tests/test_mcp_reverse_proxy_websocket_adapter.py new file mode 100644 index 0000000..fdbe9f6 --- /dev/null +++ b/tests/test_mcp_reverse_proxy_websocket_adapter.py @@ -0,0 +1,269 @@ +"""Location: ./mcp_reverse_proxy/tests/test_mcp_reverse_proxy_websocket_adapter.py +Copyright 2025 +SPDX-License-Identifier: Apache-2.0 +Authors: Mihai Criveti + +Unit tests for the WebSocket reverse proxy transport adapter. +""" + +# Future +from __future__ import annotations + +# Standard +import asyncio +import ssl +from unittest.mock import AsyncMock, Mock + +# Third-Party +import pytest + +import mcp_reverse_proxy.transports.websocket_adapter as websocket_adapter_mod + +# First-Party +from mcp_reverse_proxy.transports.websocket_adapter import WebSocketAdapter + + +class FakeWebSocketConnection: + """Async iterable WebSocket connection fake.""" + + def __init__(self, messages: list[str | bytes] | None = None) -> None: + self._messages = list(messages or []) + self.send = AsyncMock() + self.close = AsyncMock() + + def __aiter__(self): + """Return self as async iterator.""" + return self + + async def __anext__(self) -> str | bytes: + """Yield queued messages until exhausted.""" + if self._messages: + return self._messages.pop(0) + raise StopAsyncIteration + + +@pytest.mark.asyncio +async def test_connect_builds_ws_url_headers_and_receive_task(monkeypatch) -> None: + """Connect should normalize URL, pass headers, and start the receive task.""" + connection = FakeWebSocketConnection() + connect_mock = AsyncMock(return_value=connection) + created_tasks: list[asyncio.Task[None]] = [] + original_create_task = asyncio.create_task + + def capture_task(coro): + task = original_create_task(coro) + created_tasks.append(task) + return task + + monkeypatch.setattr(websocket_adapter_mod, "websockets", Mock(connect=connect_mock)) + monkeypatch.setattr(websocket_adapter_mod.asyncio, "create_task", capture_task) + + adapter = WebSocketAdapter("http://gateway.example/api", "session-1", token="token-123") + + await adapter.connect() + + connect_mock.assert_awaited_once() + assert connect_mock.await_args.args[0] == "ws://gateway.example/reverse-proxy/ws" + assert connect_mock.await_args.kwargs["additional_headers"] == { + "Authorization": "Bearer token-123", + "X-Session-ID": "session-1", + } + assert connect_mock.await_args.kwargs["ssl"] is None + assert adapter._connected is True + assert adapter._connection is connection + assert adapter._receive_task is created_tasks[0] + + adapter._receive_task.cancel() + with pytest.raises(asyncio.CancelledError): + await adapter._receive_task + + +@pytest.mark.asyncio +async def test_connect_uses_secure_defaults_and_custom_cert(monkeypatch) -> None: + """Secure WebSocket URLs should use an SSL context configured from the CA data.""" + connection = FakeWebSocketConnection() + connect_mock = AsyncMock(return_value=connection) + receive_task = asyncio.create_task(asyncio.sleep(10)) + fake_ssl_context = Mock(spec=ssl.SSLContext) + fake_ssl_context.check_hostname = True + fake_ssl_context.verify_mode = ssl.CERT_REQUIRED + fake_ssl_context.load_verify_locations = Mock() # Mock the load_verify_locations method + + # Mock load_cert_data to return test certificate data + test_cert_data = "MOCK_CERT_DATA" + monkeypatch.setattr(websocket_adapter_mod, "load_cert_data", lambda _cert: test_cert_data) + + monkeypatch.setattr(websocket_adapter_mod, "websockets", Mock(connect=connect_mock)) + # Mock SSLContext constructor to return our fake context + monkeypatch.setattr(websocket_adapter_mod.ssl, "SSLContext", lambda _protocol: fake_ssl_context) + monkeypatch.setattr(websocket_adapter_mod.asyncio, "create_task", Mock(return_value=receive_task)) + + adapter = WebSocketAdapter("gateway.example/base", "session-1", cert="test_cert_input") + + await adapter.connect() + + # Verify load_verify_locations was called with the cert data + fake_ssl_context.load_verify_locations.assert_called_once_with(cadata=test_cert_data) + assert connect_mock.await_args.args[0] == "wss://gateway.example/reverse-proxy/ws" + assert connect_mock.await_args.kwargs["ssl"] is fake_ssl_context + assert fake_ssl_context.check_hostname is True + assert fake_ssl_context.verify_mode == ssl.CERT_REQUIRED + + receive_task.cancel() + with pytest.raises(asyncio.CancelledError): + await receive_task + + +@pytest.mark.asyncio +async def test_connect_requires_websockets_dependency(monkeypatch) -> None: + """Missing websockets dependency should raise ImportError.""" + monkeypatch.setattr(websocket_adapter_mod, "websockets", None) + + adapter = WebSocketAdapter("http://gateway.example", "session-1") + + with pytest.raises(ImportError, match="websockets package required"): + await adapter.connect() + + +@pytest.mark.asyncio +async def test_disconnect_cancels_receive_task_and_closes_connection() -> None: + """Disconnect should cancel receive task, close the connection, and clear state.""" + adapter = WebSocketAdapter("http://gateway.example", "session-1") + adapter._connected = True + adapter._connection = FakeWebSocketConnection() + adapter._receive_task = asyncio.create_task(asyncio.sleep(10)) + + await adapter.disconnect() + + assert adapter._connected is False + assert adapter._connection is None + + +@pytest.mark.asyncio +async def test_send_converts_bytes_and_forwards_to_connection() -> None: + """Send should normalize bytes to text and forward the message.""" + adapter = WebSocketAdapter("http://gateway.example", "session-1") + connection = FakeWebSocketConnection() + adapter._connected = True + adapter._connection = connection + + await adapter.send(b'{"type":"heartbeat"}') + + connection.send.assert_awaited_once_with('{"type":"heartbeat"}') + + +@pytest.mark.asyncio +async def test_send_raises_when_not_connected() -> None: + """Send should fail if the adapter has no active connection.""" + adapter = WebSocketAdapter("http://gateway.example", "session-1") + + with pytest.raises(RuntimeError, match="Not connected to gateway"): + await adapter.send("message") + + +@pytest.mark.asyncio +async def test_is_connected_returns_flag_state() -> None: + """Connection status should reflect the internal connected flag.""" + adapter = WebSocketAdapter("http://gateway.example", "session-1") + adapter._connected = True + assert await adapter.is_connected() is True + + adapter._connected = False + assert await adapter.is_connected() is False + + +@pytest.mark.asyncio +async def test_receive_messages_decodes_bytes_and_notifies_handlers() -> None: + """Receive loop should decode byte frames and forward all messages to handlers.""" + adapter = WebSocketAdapter("http://gateway.example", "session-1") + handler = AsyncMock() + adapter.add_message_handler(handler) + adapter._connected = True + adapter._connection = FakeWebSocketConnection([b'{"type":"one"}', '{"type":"two"}']) + + await adapter._receive_messages() + + handler.assert_any_await('{"type":"one"}') + handler.assert_any_await('{"type":"two"}') + assert handler.await_count == 2 + assert adapter._connected is False + assert adapter._connection is None + + +@pytest.mark.asyncio +async def test_receive_messages_continues_when_handler_raises() -> None: + """One failing handler should not prevent later handlers from receiving a message.""" + adapter = WebSocketAdapter("http://gateway.example", "session-1") + failing_handler = AsyncMock(side_effect=RuntimeError("boom")) + successful_handler = AsyncMock() + adapter.add_message_handler(failing_handler) + adapter.add_message_handler(successful_handler) + adapter._connected = True + adapter._connection = FakeWebSocketConnection(['{"type":"test"}']) + + await adapter._receive_messages() + + failing_handler.assert_awaited_once_with('{"type":"test"}') + successful_handler.assert_awaited_once_with('{"type":"test"}') + + +@pytest.mark.asyncio +async def test_receive_messages_treats_connection_closed_as_clean_shutdown(monkeypatch) -> None: + """ConnectionClosed exceptions should end the loop and clear adapter state.""" + + class FakeConnectionClosedError(Exception): + """Fake connection-closed exception type.""" + + class RaisingConnection: + """Async iterator that raises the configured close exception.""" + + def __aiter__(self): + return self + + async def __anext__(self): + raise FakeConnectionClosedError("closed") + + fake_websockets = Mock() + fake_websockets.exceptions = Mock(ConnectionClosed=FakeConnectionClosedError) + monkeypatch.setattr(websocket_adapter_mod, "websockets", fake_websockets) + + adapter = WebSocketAdapter("http://gateway.example", "session-1") + adapter._connected = True + adapter._connection = RaisingConnection() + + await adapter._receive_messages() + + assert adapter._connected is False + assert adapter._connection is None + + +@pytest.mark.asyncio +async def test_receive_messages_re_raises_cancelled_error() -> None: + """Cancelled receive loops should propagate cancellation after cleanup.""" + + class CancelledConnection: + """Async iterator that raises CancelledError.""" + + def __aiter__(self): + return self + + async def __anext__(self): + raise asyncio.CancelledError + + adapter = WebSocketAdapter("http://gateway.example", "session-1") + adapter._connected = True + adapter._connection = CancelledConnection() + + with pytest.raises(asyncio.CancelledError): + await adapter._receive_messages() + + assert adapter._connected is False + assert adapter._connection is None + + +@pytest.mark.asyncio +async def test_receive_messages_returns_immediately_without_connection() -> None: + """Receive loop should no-op when no connection is present.""" + adapter = WebSocketAdapter("http://gateway.example", "session-1") + + await adapter._receive_messages()