# -*- coding: utf-8 -*-
"""
MCP Session Pool Implementation.
Provides session pooling for MCP ClientSessions to reduce per-request overhead.
Sessions are isolated per user/tenant via identity hashing to prevent session collision.
Performance Impact:
- Without pooling: 20-23ms per tool call (new session each time)
- With pooling: 1-2ms per tool call (10-20x improvement)
Security:
Sessions are isolated by (url, identity_hash, transport_type) to prevent:
- Cross-user session sharing
- Cross-tenant data leakage
- Authentication bypass
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti
"""
# flake8: noqa: DAR101, DAR201, DAR401
# Future
from __future__ import annotations
# Standard
import asyncio
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from enum import Enum
import hashlib
import logging
import os
import re
import socket
import time
from typing import Any, Callable, Dict, Optional, Set, Tuple, TYPE_CHECKING
import uuid
# Third-Party
import anyio
import httpx
from mcp import ClientSession, McpError
from mcp.client.sse import sse_client
from mcp.client.streamable_http import streamablehttp_client
from mcp.shared.session import RequestResponder
import mcp.types as mcp_types
import orjson
# First-Party
from mcpgateway.config import settings
from mcpgateway.utils.url_auth import sanitize_url_for_logging
# JSON-RPC standard error code for method not found
METHOD_NOT_FOUND = -32601
# Shared session-id validation (downstream MCP session IDs used for affinity).
# Intentionally strict: protects Redis key/channel construction and log lines.
_MCP_SESSION_ID_PATTERN = re.compile(r"^[a-zA-Z0-9_-]{1,128}$")
# Worker ID for multi-worker session affinity
# Uses hostname + PID to be unique across Docker containers (each container has PID 1)
# and across gunicorn workers within the same container
WORKER_ID = f"{socket.gethostname()}:{os.getpid()}"
def _get_cleanup_timeout() -> float:
"""Get session cleanup timeout from config (lazy import to avoid circular deps).
This timeout controls how long to wait for session/transport __aexit__ calls
when closing sessions. It prevents CPU spin loops when internal tasks don't
respond to cancellation (anyio's _deliver_cancellation issue).
Returns:
Cleanup timeout in seconds (default: 5.0)
"""
try:
# Lazy import to avoid circular dependency during startup
return settings.mcp_session_pool_cleanup_timeout
except Exception:
return 5.0 # Fallback default
if TYPE_CHECKING:
# Standard
from collections.abc import AsyncIterator # pragma: no cover
logger = logging.getLogger(__name__)
class TransportType(Enum):
"""Supported MCP transport types."""
SSE = "sse"
STREAMABLE_HTTP = "streamablehttp"
@dataclass(eq=False) # eq=False makes instances hashable by object identity
class PooledSession:
"""A pooled MCP session with metadata for lifecycle management.
Note: eq=False is required because we store these in Sets for active session
tracking. This makes instances hashable by their object id (identity).
"""
session: ClientSession
transport_context: Any # The transport context manager (kept open)
url: str
transport_type: TransportType
headers: Dict[str, str] # Original headers (for reconnection)
identity_key: str # Identity hash component for headers
user_identity: str = "anonymous" # for user isolation
gateway_id: str = "" # Gateway ID for notification attribution
created_at: float = field(default_factory=time.time)
last_used: float = field(default_factory=time.time)
use_count: int = 0
_closed: bool = field(default=False, repr=False)
@property
def age_seconds(self) -> float:
"""Return session age in seconds.
Returns:
float: Session age in seconds since creation.
"""
return time.time() - self.created_at
@property
def idle_seconds(self) -> float:
"""Return seconds since last use.
Returns:
float: Seconds since last use of this session.
"""
return time.time() - self.last_used
@property
def is_closed(self) -> bool:
"""Return whether this session has been closed.
Returns:
bool: True if session is closed, False otherwise.
"""
return self._closed
def mark_closed(self) -> None:
"""Mark this session as closed."""
self._closed = True
# Type aliases
# Pool key includes transport type and gateway_id to prevent returning wrong transport for same URL
# and to ensure correct notification attribution when notifications are enabled
PoolKey = Tuple[str, str, str, str, str] # (user_identity_hash, url, identity_hash, transport_type, gateway_id)
# Session affinity mapping key: (mcp_session_id, url, transport_type, gateway_id)
SessionMappingKey = Tuple[str, str, str, str]
HttpxClientFactory = Callable[
[Optional[Dict[str, str]], Optional[httpx.Timeout], Optional[httpx.Auth]],
httpx.AsyncClient,
]
# Type alias for identity extractor callback
# Extracts stable identity from headers (e.g., decode JWT to get user_id)
IdentityExtractor = Callable[[Dict[str, str]], Optional[str]]
# Type alias for message handler factory
# Factory that creates message handlers given URL and optional gateway_id
# The handler receives ServerNotification, ServerRequest responders, or Exceptions
MessageHandlerFactory = Callable[
[str, Optional[str]], # (url, gateway_id)
Callable[
[RequestResponder[mcp_types.ServerRequest, mcp_types.ClientResult] | mcp_types.ServerNotification | Exception],
Any, # Coroutine
],
]
class MCPSessionPool: # pylint: disable=too-many-instance-attributes
"""
Pool of MCP ClientSessions keyed by (user_identity, server URL, identity hash, transport type, gateway_id).
Thread-Safety:
This pool is designed for asyncio concurrency. It uses asyncio.Lock
for synchronization, which is safe for coroutine-based concurrency
but NOT for multi-threaded access.
Session Isolation:
Sessions are isolated per user/tenant to prevent session collision.
The identity hash is derived from authentication headers ensuring
that different users never share MCP sessions.
Transport Isolation:
Sessions are also isolated by transport type (SSE vs STREAMABLE_HTTP).
The same URL with different transports will use separate pools.
Gateway Isolation:
Sessions are isolated by gateway_id for correct notification attribution.
When notifications are enabled, each gateway gets its own pooled sessions
even if they share the same URL and authentication.
Features:
- Session reuse across requests (10-20x latency improvement)
- Per-user/tenant session isolation (prevents session collision)
- Per-transport session isolation (prevents transport mismatch)
- TTL-based expiration with configurable lifetime
- Health checks on acquire for stale sessions
- Configurable pool size per URL+identity+transport
- Circuit breaker for failing endpoints
- Idle pool key eviction to prevent unbounded growth
- Custom identity extractor for rotating tokens (e.g., JWT decode)
- Metrics for monitoring (hits, misses, evictions)
- Graceful shutdown with close_all()
Usage:
pool = MCPSessionPool()
# Use as context manager for lifecycle management
async with pool:
pooled = await pool.acquire(url, headers)
try:
result = await pooled.session.call_tool("my_tool", {})
finally:
await pool.release(pooled)
# With custom identity extractor for JWT tokens:
def extract_user_id(headers: dict) -> str:
token = headers.get("Authorization", "").replace("Bearer ", "")
claims = jwt.decode(token, options={"verify_signature": False})
return claims.get("sub") or claims.get("user_id")
pool = MCPSessionPool(identity_extractor=extract_user_id)
"""
# Headers that contribute to session identity (case-insensitive)
DEFAULT_IDENTITY_HEADERS: frozenset[str] = frozenset(
[
"authorization",
"x-tenant-id",
"x-user-id",
"x-api-key",
"cookie",
"x-mcp-session-id",
]
)
def __init__(
self,
max_sessions_per_key: int = 10,
session_ttl_seconds: float = 300.0,
health_check_interval_seconds: float = 60.0,
acquire_timeout_seconds: float = 30.0,
session_create_timeout_seconds: float = 30.0,
circuit_breaker_threshold: int = 5,
circuit_breaker_reset_seconds: float = 60.0,
identity_headers: Optional[frozenset[str]] = None,
identity_extractor: Optional[IdentityExtractor] = None,
idle_pool_eviction_seconds: float = 600.0,
default_transport_timeout_seconds: float = 30.0,
health_check_methods: Optional[list[str]] = None,
health_check_timeout_seconds: float = 5.0,
message_handler_factory: Optional[MessageHandlerFactory] = None,
):
"""
Initialize the session pool.
Args:
max_sessions_per_key: Maximum pooled sessions per (URL, identity, transport).
session_ttl_seconds: Session TTL in seconds before forced expiration.
health_check_interval_seconds: Seconds of idle time before health check.
acquire_timeout_seconds: Timeout for waiting when pool is exhausted.
session_create_timeout_seconds: Timeout for creating new sessions.
circuit_breaker_threshold: Consecutive failures before circuit opens.
circuit_breaker_reset_seconds: Seconds before circuit breaker resets.
identity_headers: Headers that contribute to identity hash.
identity_extractor: Optional callback to extract stable identity from headers.
Use this when tokens rotate frequently (e.g., short-lived JWTs).
Should return a stable user/tenant ID string.
idle_pool_eviction_seconds: Evict empty pool keys after this many seconds of no use.
default_transport_timeout_seconds: Default timeout for transport connections.
health_check_methods: Ordered list of health check methods to try.
Options: ping, list_tools, list_prompts, list_resources, skip.
Default: ["ping", "skip"] (try ping, skip if unsupported).
health_check_timeout_seconds: Timeout for each health check attempt.
message_handler_factory: Optional factory for creating message handlers.
Called with (url, gateway_id) to create handlers for
each new session. Enables notification handling.
"""
# Configuration
self._max_sessions = max_sessions_per_key
self._session_ttl = session_ttl_seconds
self._health_check_interval = health_check_interval_seconds
self._acquire_timeout = acquire_timeout_seconds
self._session_create_timeout = session_create_timeout_seconds
self._circuit_breaker_threshold = circuit_breaker_threshold
self._circuit_breaker_reset = circuit_breaker_reset_seconds
self._identity_headers = identity_headers or self.DEFAULT_IDENTITY_HEADERS
self._identity_extractor = identity_extractor
self._idle_pool_eviction = idle_pool_eviction_seconds
self._default_transport_timeout = default_transport_timeout_seconds
self._health_check_methods = health_check_methods or ["ping", "skip"]
self._health_check_timeout = health_check_timeout_seconds
self._message_handler_factory = message_handler_factory
# State - protected by _global_lock for creation, per-key locks for access
self._global_lock = asyncio.Lock()
self._pools: Dict[PoolKey, asyncio.Queue[PooledSession]] = {}
self._active: Dict[PoolKey, Set[PooledSession]] = {}
self._locks: Dict[PoolKey, asyncio.Lock] = {}
self._semaphores: Dict[PoolKey, asyncio.Semaphore] = {}
self._pool_last_used: Dict[PoolKey, float] = {} # Track last use time per pool key
# Circuit breaker state
self._failures: Dict[str, int] = {} # url -> consecutive failure count
self._circuit_open_until: Dict[str, float] = {} # url -> timestamp
# Eviction throttling - only run eviction once per interval
self._last_eviction_run: float = 0.0
self._eviction_run_interval: float = 60.0 # Run eviction at most every 60 seconds
# Metrics
self._hits = 0
self._misses = 0
self._evictions = 0
self._health_check_failures = 0
self._circuit_breaker_trips = 0
self._pool_keys_evicted = 0
self._sessions_reaped = 0 # Sessions closed during background eviction
self._anonymous_identity_count = 0 # Count of requests with no identity headers
# Lifecycle
self._closed = False
# Pre-registered session mappings for session affinity
# Mapping from (mcp_session_id, url, transport_type, gateway_id) -> pool_key
# Set by broadcast() before acquire() is called to enable session affinity lookup
self._mcp_session_mapping: Dict[SessionMappingKey, PoolKey] = {}
self._mcp_session_mapping_lock = asyncio.Lock()
# Multi-worker session affinity via Redis pub/sub
# Track pending responses for forwarded RPC requests
self._rpc_listener_task: Optional[asyncio.Task[None]] = None
self._pending_responses: Dict[str, asyncio.Future[Dict[str, Any]]] = {}
# Session affinity metrics
self._session_affinity_local_hits = 0
self._session_affinity_redis_hits = 0
self._session_affinity_misses = 0
self._forwarded_requests = 0
self._forwarded_request_failures = 0
self._forwarded_request_timeouts = 0
async def __aenter__(self) -> "MCPSessionPool":
"""Async context manager entry.
Returns:
MCPSessionPool: This pool instance.
"""
return self
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
"""Async context manager exit - closes all sessions.
Args:
exc_type: Exception type if an exception was raised.
exc_val: Exception value if an exception was raised.
exc_tb: Exception traceback if an exception was raised.
"""
await self.close_all()
def _compute_identity_hash(self, headers: Optional[Dict[str, str]]) -> str:
"""
Compute a hash of identity-relevant headers.
This ensures sessions are isolated per user/tenant. Different users
with different Authorization headers will have different identity hashes
and thus separate session pools.
Identity resolution order:
1. Custom identity_extractor (if configured) - for rotating tokens like JWTs
2. x-mcp-session-id header (if present) - for session affinity, ensures
requests with the same downstream session ID get the same upstream
session even when JWT tokens rotate (different jti values)
3. Configured identity headers - fallback to hashing all identity headers
Args:
headers: Request headers dict.
Returns:
Identity hash string, or "anonymous" if no identity headers present.
"""
if not headers:
self._anonymous_identity_count += 1
logger.debug("Session pool identity collapsed to 'anonymous' (no headers provided). " + "Sessions will be shared. Ensure this is intentional for stateless MCP servers.")
return "anonymous"
# Try custom identity extractor first (for rotating tokens like JWTs)
if self._identity_extractor:
try:
extracted = self._identity_extractor(headers)
if extracted:
return hashlib.sha256(extracted.encode()).hexdigest()
except Exception as e:
logger.debug(f"Identity extractor failed, falling back to header hash: {e}")
# Normalize headers for case-insensitive lookup
headers_lower = {k.lower(): v for k, v in headers.items()}
# Session affinity: prioritize x-mcp-session-id for stable identity
# When present, use ONLY the session ID for identity hash. This ensures
# requests with the same downstream session ID get the same upstream session,
# even when JWT tokens rotate (different jti values per request).
if settings.mcpgateway_session_affinity_enabled:
session_id = headers_lower.get("x-mcp-session-id")
if session_id:
logger.debug(f"Using x-mcp-session-id for session affinity: {session_id[:8]}...")
return hashlib.sha256(session_id.encode()).hexdigest()
# Fallback: extract identity from configured headers
identity_parts = []
for header in sorted(self._identity_headers):
if header in headers_lower:
identity_parts.append(f"{header}:{headers_lower[header]}")
if not identity_parts:
self._anonymous_identity_count += 1
logger.debug(
"Session pool identity collapsed to 'anonymous' (no identity headers found). " + "Expected headers: %s. Sessions will be shared.",
list(self._identity_headers),
)
return "anonymous"
# Create a stable, deterministic hash using JSON serialization
# Prevents delimiter-collision or injection issues present in string joining
serialized_identity = orjson.dumps(identity_parts)
return hashlib.sha256(serialized_identity).hexdigest()
def _make_pool_key(
self,
url: str,
headers: Optional[Dict[str, str]],
transport_type: TransportType,
user_identity: str,
gateway_id: Optional[str] = None,
) -> PoolKey:
"""Create composite pool key from URL, identity, transport type, user identity, and gateway_id.
Including gateway_id ensures correct notification attribution when multiple gateways
share the same URL/auth. Sessions are isolated per gateway for proper event routing.
"""
identity_hash = self._compute_identity_hash(headers)
# Anonymize user identity by hashing it (unless it's commonly "anonymous")
# Use full hash for collision resistance - truncate only for display in logs/metrics
if user_identity == "anonymous":
user_hash = "anonymous"
else:
user_hash = hashlib.sha256(user_identity.encode()).hexdigest()
# Use empty string for None gateway_id to maintain consistent key type
gw_id = gateway_id or ""
return (user_hash, url, identity_hash, transport_type.value, gw_id)
async def _get_or_create_lock(self, pool_key: PoolKey) -> asyncio.Lock:
"""Get or create a lock for the given pool key (thread-safe)."""
async with self._global_lock:
if pool_key not in self._locks:
self._locks[pool_key] = asyncio.Lock()
return self._locks[pool_key]
async def _get_or_create_pool(self, pool_key: PoolKey) -> asyncio.Queue[PooledSession]:
"""Get or create a pool queue for the given key (thread-safe)."""
async with self._global_lock:
if pool_key not in self._pools:
self._pools[pool_key] = asyncio.Queue(maxsize=self._max_sessions)
self._active[pool_key] = set()
self._semaphores[pool_key] = asyncio.Semaphore(self._max_sessions)
return self._pools[pool_key]
def _is_circuit_open(self, url: str) -> bool:
"""Check if circuit breaker is open for a URL."""
if url not in self._circuit_open_until:
return False
if time.time() >= self._circuit_open_until[url]:
# Circuit breaker reset
del self._circuit_open_until[url]
self._failures[url] = 0
logger.info(f"Circuit breaker reset for {sanitize_url_for_logging(url)}")
return False
return True
def _record_failure(self, url: str) -> None:
"""Record a failure and potentially trip circuit breaker."""
self._failures[url] = self._failures.get(url, 0) + 1
if self._failures[url] >= self._circuit_breaker_threshold:
self._circuit_open_until[url] = time.time() + self._circuit_breaker_reset
self._circuit_breaker_trips += 1
logger.warning(f"Circuit breaker opened for {sanitize_url_for_logging(url)} after {self._failures[url]} failures. " f"Will reset in {self._circuit_breaker_reset}s")
def _record_success(self, url: str) -> None:
"""Record a success, resetting failure count."""
self._failures[url] = 0
@staticmethod
def is_valid_mcp_session_id(session_id: str) -> bool:
"""Validate downstream MCP session ID format for affinity.
Used for:
- Redis key construction (ownership + mapping)
- Pub/Sub channel naming
- Avoiding log spam / injection
"""
if not session_id:
return False
return bool(_MCP_SESSION_ID_PATTERN.match(session_id))
def _sanitize_redis_key_component(self, value: str) -> str:
"""Sanitize a value for use in Redis key construction.
Replaces any characters that could cause key collision or injection.
Args:
value: The value to sanitize.
Returns:
Sanitized value safe for Redis key construction.
"""
if not value:
return ""
# Replace problematic characters with underscores
return re.sub(r"[^a-zA-Z0-9_-]", "_", value)
def _session_mapping_redis_key(self, mcp_session_id: str, url: str, transport_type: str, gateway_id: str) -> str:
"""Compute a bounded Redis key for session mapping.
The URL is hashed to keep keys small and avoid special character issues.
"""
sanitized_session_id = self._sanitize_redis_key_component(mcp_session_id)
url_hash = hashlib.sha256(url.encode()).hexdigest()[:16]
return f"mcpgw:session_mapping:{sanitized_session_id}:{url_hash}:{transport_type}:{gateway_id}"
@staticmethod
def _pool_owner_key(mcp_session_id: str) -> str:
"""Return Redis key for session ownership tracking."""
return f"mcpgw:pool_owner:{mcp_session_id}"
async def register_session_mapping(
self,
mcp_session_id: str,
url: str,
gateway_id: str,
transport_type: str,
user_email: Optional[str] = None,
) -> None:
"""Pre-register session mapping for session affinity.
Called from respond() to set up mapping BEFORE acquire() is called.
This ensures acquire() can find the correct pool key for session affinity.
The mapping stores the relationship between an incoming MCP session ID
and the pool key that should be used for upstream connections. This
enables session affinity even when JWT tokens rotate (different jti values
per request).
For multi-worker deployments, the mapping is also stored in Redis with TTL
so that any worker can look it up during acquire().
Args:
mcp_session_id: The downstream MCP session ID from x-mcp-session-id header.
url: The upstream MCP server URL.
gateway_id: The gateway ID.
transport_type: The transport type (sse, streamablehttp).
user_email: The email of the authenticated user (or "system" for unauthenticated).
"""
if not settings.mcpgateway_session_affinity_enabled:
return
# Validate mcp_session_id to prevent Redis key injection
if not self.is_valid_mcp_session_id(mcp_session_id):
logger.warning(f"Invalid mcp_session_id format, skipping session mapping: {mcp_session_id[:20]}...")
return
# Use user email for user_identity, or "anonymous" if not provided
user_identity = user_email or "anonymous"
# Normalize gateway_id to empty string if None for consistent key matching
normalized_gateway_id = gateway_id or ""
mapping_key: SessionMappingKey = (mcp_session_id, url, transport_type, normalized_gateway_id)
# Compute what the pool_key will be for this session
# Use mcp_session_id as the identity basis for affinity
identity_hash = hashlib.sha256(mcp_session_id.encode()).hexdigest()
# Hash user identity for privacy (unless it's "anonymous")
if user_identity == "anonymous":
user_hash = "anonymous"
else:
user_hash = hashlib.sha256(user_identity.encode()).hexdigest()
pool_key: PoolKey = (user_hash, url, identity_hash, transport_type, normalized_gateway_id)
# Store in local memory
async with self._mcp_session_mapping_lock:
self._mcp_session_mapping[mapping_key] = pool_key
logger.debug(f"Session affinity pre-registered (local): {mcp_session_id[:8]}... → {url}, user={user_identity}")
# Store in Redis for multi-worker support AND register ownership atomically
# Registering ownership HERE (during mapping) instead of in acquire() prevents
# a race condition where two workers could both start creating sessions before
# either registers ownership
try:
# First-Party
from mcpgateway.utils.redis_client import get_redis_client # pylint: disable=import-outside-toplevel
redis = await get_redis_client()
if redis:
redis_key = self._session_mapping_redis_key(mcp_session_id, url, transport_type, normalized_gateway_id)
# Store pool_key as JSON for easy deserialization
pool_key_data = {
"user_hash": user_hash,
"url": url,
"identity_hash": identity_hash,
"transport_type": transport_type,
"gateway_id": normalized_gateway_id,
}
await redis.setex(redis_key, settings.mcpgateway_session_affinity_ttl, orjson.dumps(pool_key_data)) # TTL from config
# CRITICAL: Register ownership atomically with mapping.
# This claims ownership BEFORE any session creation attempt, preventing
# the race condition where two workers both start creating sessions
owner_key = self._pool_owner_key(mcp_session_id)
# Atomic claim with TTL (avoids the SETNX/EXPIRE crash window).
was_set = await redis.set(owner_key, WORKER_ID, nx=True, ex=settings.mcpgateway_session_affinity_ttl)
if was_set:
logger.debug(f"Session ownership claimed (SET NX): {mcp_session_id[:8]}... → worker {WORKER_ID}")
else:
# Another worker already claimed ownership
existing_owner = await redis.get(owner_key)
owner_id = existing_owner.decode() if isinstance(existing_owner, bytes) else existing_owner
logger.debug(f"Session ownership already claimed by {owner_id}: {mcp_session_id[:8]}...")
logger.debug(f"Session affinity pre-registered (Redis): {mcp_session_id[:8]}... TTL={settings.mcpgateway_session_affinity_ttl}s")
except Exception as e:
# Redis failure is non-fatal - local mapping still works for same-worker requests
logger.debug(f"Failed to store session mapping in Redis: {e}")
async def acquire(
self,
url: str,
headers: Optional[Dict[str, str]] = None,
transport_type: TransportType = TransportType.STREAMABLE_HTTP,
httpx_client_factory: Optional[HttpxClientFactory] = None,
timeout: Optional[float] = None,
user_identity: Optional[str] = None,
gateway_id: Optional[str] = None,
) -> PooledSession:
"""
Acquire a session for the given URL, identity, and transport type.
Sessions are isolated by identity (derived from auth headers) AND
transport type. Returns an initialized, healthy session ready for tool calls.
Args:
url: The MCP server URL.
headers: Request headers (used for identity hashing and passed to server).
transport_type: The transport type (SSE or STREAMABLE_HTTP).
httpx_client_factory: Optional factory for creating httpx clients
(for custom SSL/timeout configuration).
timeout: Optional timeout in seconds for transport connection.
gateway_id: Optional gateway ID for notification handler context.
Returns:
PooledSession ready for use.
Raises:
asyncio.TimeoutError: If acquire times out waiting for available session.
RuntimeError: If pool is closed or circuit breaker is open.
Exception: If session creation fails.
"""
if self._closed:
raise RuntimeError("Session pool is closed")
if self._is_circuit_open(url):
raise RuntimeError(f"Circuit breaker open for {url}")
# Use default timeout if not provided
effective_timeout = timeout if timeout is not None else self._default_transport_timeout
user_id = user_identity or "anonymous"
pool_key: Optional[PoolKey] = None
# Check pre-registered mapping first (set by respond() for session affinity)
if settings.mcpgateway_session_affinity_enabled and headers:
headers_lower = {k.lower(): v for k, v in headers.items()}
mcp_session_id = headers_lower.get("x-mcp-session-id")
if mcp_session_id and self.is_valid_mcp_session_id(mcp_session_id):
normalized_gateway_id = gateway_id or ""
mapping_key: SessionMappingKey = (mcp_session_id, url, transport_type.value, normalized_gateway_id)
# Check local memory first (fast path - same worker)
async with self._mcp_session_mapping_lock:
pool_key = self._mcp_session_mapping.get(mapping_key)
if pool_key:
self._session_affinity_local_hits += 1
logger.debug(f"Session affinity hit (local): {mcp_session_id[:8]}...")
# If not in local memory, check Redis (multi-worker support)
if pool_key is None:
try:
# First-Party
from mcpgateway.utils.redis_client import get_redis_client # pylint: disable=import-outside-toplevel
redis = await get_redis_client()
if redis:
redis_key = self._session_mapping_redis_key(mcp_session_id, url, transport_type.value, normalized_gateway_id)
pool_key_data = await redis.get(redis_key)
if pool_key_data:
# Deserialize pool_key from JSON
data = orjson.loads(pool_key_data)
pool_key = (
data["user_hash"],
data["url"],
data["identity_hash"],
data["transport_type"],
data["gateway_id"],
)
# Cache in local memory for future requests
async with self._mcp_session_mapping_lock:
self._mcp_session_mapping[mapping_key] = pool_key
self._session_affinity_redis_hits += 1
logger.debug(f"Session affinity hit (Redis): {mcp_session_id[:8]}...")
except Exception as e:
logger.debug(f"Failed to check Redis for session mapping: {e}")
# Fallback to normal pool key computation
if pool_key is None:
self._session_affinity_misses += 1
pool_key = self._make_pool_key(url, headers, transport_type, user_id, gateway_id)
pool = await self._get_or_create_pool(pool_key)
# Update pool key last used time IMMEDIATELY after getting pool
# This prevents race with eviction removing keys between awaits
self._pool_last_used[pool_key] = time.time()
lock = await self._get_or_create_lock(pool_key)
# Guard semaphore access - eviction may have removed it between awaits
# If so, re-create the pool structures
if pool_key not in self._semaphores:
pool = await self._get_or_create_pool(pool_key)
self._pool_last_used[pool_key] = time.time()
semaphore = self._semaphores[pool_key]
# Throttled eviction - only run if enough time has passed (inline, not spawned)
await self._maybe_evict_idle_pool_keys()
# Try to get from pool first (quick path, no lock needed for queue get)
while True:
try:
pooled = pool.get_nowait()
except asyncio.QueueEmpty:
break
# Validate the session outside the lock
if await self._validate_session(pooled):
pooled.last_used = time.time()
pooled.use_count += 1
self._hits += 1
async with lock:
self._active[pool_key].add(pooled)
logger.debug(f"Pool hit for {sanitize_url_for_logging(url)} (identity={pool_key[2][:8]}, transport={transport_type.value})")
return pooled
# Session invalid, close it
await self._close_session(pooled)
self._evictions += 1
semaphore.release() # Free up a slot
# No valid session in pool - try to create one or wait
try:
# Use semaphore with timeout to limit concurrent sessions
acquired = await asyncio.wait_for(semaphore.acquire(), timeout=self._acquire_timeout)
if not acquired:
raise asyncio.TimeoutError("Failed to acquire session slot")
except asyncio.TimeoutError:
raise asyncio.TimeoutError(f"Timeout waiting for available session for {sanitize_url_for_logging(url)}") from None
# Create new session (semaphore acquired)
try:
# Verify we own this session before creating (prevents race condition)
# If another worker already claimed ownership, we should not create a new session
# Note: Ownership is registered atomically in register_session_mapping() using SETNX
if settings.mcpgateway_session_affinity_enabled and headers:
headers_lower = {k.lower(): v for k, v in headers.items()}
mcp_session_id = headers_lower.get("x-mcp-session-id")
if mcp_session_id and self.is_valid_mcp_session_id(mcp_session_id):
owner = await self._get_pool_session_owner(mcp_session_id)
if owner and owner != WORKER_ID:
# Another worker claimed ownership - should have been forwarded
# Release semaphore and raise to trigger forwarding
semaphore.release()
logger.warning(f"Session {mcp_session_id[:8]}... owned by worker {owner}, not us ({WORKER_ID})")
raise RuntimeError(f"Session owned by another worker: {owner}")
pooled = await asyncio.wait_for(
self._create_session(url, headers, transport_type, httpx_client_factory, effective_timeout, gateway_id),
timeout=self._session_create_timeout,
)
# Store identity components for key reconstruction
pooled.identity_key = pool_key[2]
pooled.user_identity = user_id
# Note: Ownership is now registered atomically in register_session_mapping()
# before acquire() is called, so we don't need to register it here
self._misses += 1
self._record_success(url)
async with lock:
self._active[pool_key].add(pooled)
logger.debug(f"Pool miss for {sanitize_url_for_logging(url)} - created new session (transport={transport_type.value})")
return pooled
except BaseException as e:
# Release semaphore on ANY failure (including CancelledError)
semaphore.release()
if not isinstance(e, asyncio.CancelledError):
self._record_failure(url)
logger.warning(f"Failed to create session for {sanitize_url_for_logging(url)}: {e}")
raise
async def release(self, pooled: PooledSession) -> None:
"""
Return a session to the pool for reuse.
Args:
pooled: The session to release.
"""
if pooled.is_closed:
logger.warning("Attempted to release already-closed session")
return
# Pool key includes transport type, user identity, and gateway_id
# Re-compute user hash from stored raw identity (full hash for collision resistance)
user_hash = "anonymous"
if pooled.user_identity != "anonymous":
user_hash = hashlib.sha256(pooled.user_identity.encode()).hexdigest()
pool_key = (user_hash, pooled.url, pooled.identity_key, pooled.transport_type.value, pooled.gateway_id)
lock = await self._get_or_create_lock(pool_key)
pool = await self._get_or_create_pool(pool_key)
async with lock:
# Update last-used FIRST to prevent eviction race:
# If eviction runs between removing from _active and putting back in pool,
# it would see key as idle + inactive and evict it. By updating last-used
# while still holding the lock and before removing from _active, we ensure
# eviction sees recent activity.
self._pool_last_used[pool_key] = time.time()
self._active.get(pool_key, set()).discard(pooled)
# Check if session should be returned to pool
if self._closed or pooled.age_seconds > self._session_ttl:
await self._close_session(pooled)
if pool_key in self._semaphores:
self._semaphores[pool_key].release()
if pooled.age_seconds > self._session_ttl:
self._evictions += 1
return
# Return to pool (pool may have been evicted in edge case, recreate if needed)
if pool_key not in self._pools:
pool = await self._get_or_create_pool(pool_key)
self._pool_last_used[pool_key] = time.time()
try:
pool.put_nowait(pooled)
logger.debug(f"Session returned to pool for {sanitize_url_for_logging(pooled.url)}")
except asyncio.QueueFull:
# Pool full (shouldn't happen with semaphore), close session
await self._close_session(pooled)
if pool_key in self._semaphores:
self._semaphores[pool_key].release()
async def _maybe_evict_idle_pool_keys(self) -> None:
"""
Reap stale sessions and evict idle pool keys.
This method is throttled - it only runs eviction if enough time has
passed since the last run (default: 60 seconds). This prevents:
- Unbounded task spawning on every acquire
- Lock contention under high load
Two-phase cleanup:
1. Close expired/stale sessions parked in idle pools (frees connections)
2. Evict pool keys that are now empty and have no active sessions
This prevents unbounded connection and pool key growth when using
rotating tokens (e.g., short-lived JWTs with unique identifiers).
"""
if self._closed:
return
now = time.time()
# Throttle: only run eviction once per interval
if now - self._last_eviction_run < self._eviction_run_interval:
return
self._last_eviction_run = now
# Collect sessions to close and keys to evict (minimize time holding lock)
sessions_to_close: list[PooledSession] = []
keys_to_evict: list[PoolKey] = []
async with self._global_lock:
for pool_key, last_used in list(self._pool_last_used.items()):
# Skip recently-used pools
if now - last_used < self._idle_pool_eviction:
continue
pool = self._pools.get(pool_key)
active = self._active.get(pool_key, set())
# Skip if there are active sessions (in use)
if active:
continue
if pool:
# Phase 1: Drain and collect expired/stale sessions from idle pools
while not pool.empty():
try:
session = pool.get_nowait()
# Close if expired OR idle too long (defense in depth)
if session.age_seconds > self._session_ttl or session.idle_seconds > self._idle_pool_eviction:
sessions_to_close.append(session)
# Release semaphore slot for this session
if pool_key in self._semaphores:
self._semaphores[pool_key].release()
else:
# Session still valid, put it back
pool.put_nowait(session)
break # Stop draining if we find a valid session
except asyncio.QueueEmpty:
break
# Phase 2: Evict pool key if now empty
if pool.empty():
keys_to_evict.append(pool_key)
# Remove evicted keys from all tracking dicts
for pool_key in keys_to_evict:
self._pools.pop(pool_key, None)
self._active.pop(pool_key, None)
self._locks.pop(pool_key, None)
self._semaphores.pop(pool_key, None)
self._pool_last_used.pop(pool_key, None)
self._pool_keys_evicted += 1
logger.debug(f"Evicted idle pool key: {pool_key[0][:8]}|{pool_key[1]}|{pool_key[2][:8]}")
# Close sessions outside the lock (I/O operations)
for session in sessions_to_close:
await self._close_session(session)
self._sessions_reaped += 1
logger.debug(f"Reaped stale session for {sanitize_url_for_logging(session.url)} (age={session.age_seconds:.1f}s)")
async def _validate_session(self, pooled: PooledSession) -> bool:
"""
Validate a session is still usable.
Checks TTL and performs health check if session is stale.
Args:
pooled: The session to validate.
Returns:
True if session is valid, False otherwise.
"""
if pooled.is_closed:
return False
# Check TTL
if pooled.age_seconds > self._session_ttl:
logger.debug(f"Session expired (age={pooled.age_seconds:.1f}s)")
return False
# Health check if stale
if pooled.idle_seconds > self._health_check_interval:
return await self._run_health_check_chain(pooled)
return True
async def _run_health_check_chain(self, pooled: PooledSession) -> bool:
"""
Run health check methods in configured order until one succeeds.
The health check chain allows configuring which methods to try and in what order.
This supports both modern servers (with ping support) and legacy servers
(that may only support list_tools or no health check at all).
Args:
pooled: The session to health check.
Returns:
True if any health check method succeeds, False if all fail.
"""
for method in self._health_check_methods:
try:
if method == "ping":
await asyncio.wait_for(pooled.session.send_ping(), timeout=self._health_check_timeout)
logger.debug(f"Health check passed: ping (url={sanitize_url_for_logging(pooled.url)})")
return True
if method == "list_tools":
await asyncio.wait_for(pooled.session.list_tools(), timeout=self._health_check_timeout)
logger.debug(f"Health check passed: list_tools (url={sanitize_url_for_logging(pooled.url)})")
return True
if method == "list_prompts":
await asyncio.wait_for(pooled.session.list_prompts(), timeout=self._health_check_timeout)
logger.debug(f"Health check passed: list_prompts (url={sanitize_url_for_logging(pooled.url)})")
return True
if method == "list_resources":
await asyncio.wait_for(pooled.session.list_resources(), timeout=self._health_check_timeout)
logger.debug(f"Health check passed: list_resources (url={sanitize_url_for_logging(pooled.url)})")
return True
if method == "skip":
logger.debug(f"Health check skipped per configuration (url={sanitize_url_for_logging(pooled.url)})")
return True
logger.warning(f"Unknown health check method '{method}', skipping")
continue
except McpError as e:
# METHOD_NOT_FOUND (-32601) means the method isn't supported - try next
if e.error.code == METHOD_NOT_FOUND:
logger.debug(f"Health check method '{method}' not supported by server, trying next")
continue
# Other MCP errors are real failures
logger.debug(f"Health check '{method}' failed with MCP error: {e}")
self._health_check_failures += 1
return False
except asyncio.TimeoutError:
logger.debug(f"Health check '{method}' timed out after {self._health_check_timeout}s, trying next")
continue
except Exception as e:
logger.debug(f"Health check '{method}' failed: {e}")
self._health_check_failures += 1
return False
# All methods failed or were unsupported
logger.warning(f"All health check methods failed or unsupported (methods={self._health_check_methods})")
self._health_check_failures += 1
return False
async def _create_session(
self,
url: str,
headers: Optional[Dict[str, str]],
transport_type: TransportType,
httpx_client_factory: Optional[HttpxClientFactory],
timeout: Optional[float] = None,
gateway_id: Optional[str] = None,
) -> PooledSession:
"""
Create a new initialized MCP session.
Args:
url: Server URL.
headers: Request headers.
transport_type: Transport type to use.
httpx_client_factory: Optional factory for httpx clients.
timeout: Optional timeout in seconds for transport connection.
gateway_id: Optional gateway ID for notification handler context.
Returns:
Initialized PooledSession.
Raises:
RuntimeError: If session creation or initialization fails.
asyncio.CancelledError: If cancelled during creation.
"""
# Merge headers with defaults
merged_headers = {"Accept": "application/json, text/event-stream"}
if headers:
merged_headers.update(headers)
# Strip gateway-internal session affinity headers before sending to upstream
# x-mcp-session-id is our internal representation, mcp-session-id is the MCP protocol header
# Neither should be forwarded to upstream servers
keys_to_remove = [k for k in merged_headers if k.lower() in ("x-mcp-session-id", "mcp-session-id")]
for k in keys_to_remove:
del merged_headers[k]
identity_key = self._compute_identity_hash(headers)
transport_ctx = None
session = None
success = False
try:
# Create transport context
if transport_type == TransportType.SSE:
if httpx_client_factory:
transport_ctx = sse_client(url=url, headers=merged_headers, httpx_client_factory=httpx_client_factory, timeout=timeout)
else:
transport_ctx = sse_client(url=url, headers=merged_headers, timeout=timeout)
# pylint: disable=unnecessary-dunder-call,no-member
streams = await transport_ctx.__aenter__() # Must call directly for manual lifecycle management
read_stream, write_stream = streams[0], streams[1]
else: # STREAMABLE_HTTP
if httpx_client_factory:
transport_ctx = streamablehttp_client(url=url, headers=merged_headers, httpx_client_factory=httpx_client_factory, timeout=timeout)
else:
transport_ctx = streamablehttp_client(url=url, headers=merged_headers, timeout=timeout)
# pylint: disable=unnecessary-dunder-call,no-member
read_stream, write_stream, _ = await transport_ctx.__aenter__() # Must call directly for manual lifecycle management
# Create message handler if factory is configured
message_handler = None
if self._message_handler_factory:
try:
message_handler = self._message_handler_factory(url, gateway_id)
logger.debug(f"Created message handler for session {sanitize_url_for_logging(url)} (gateway={gateway_id})")
except Exception as e:
logger.warning(f"Failed to create message handler for {sanitize_url_for_logging(url)}: {e}")
# Create and initialize session
session = ClientSession(read_stream, write_stream, message_handler=message_handler)
# pylint: disable=unnecessary-dunder-call
await session.__aenter__() # Must call directly for manual lifecycle management
await session.initialize()
logger.info(f"Created new MCP session for {sanitize_url_for_logging(url)} (transport={transport_type.value})")
success = True
return PooledSession(
session=session,
transport_context=transport_ctx,
url=url,
transport_type=transport_type,
headers=merged_headers,
identity_key=identity_key,
gateway_id=gateway_id or "",
)
except asyncio.CancelledError: # pylint: disable=try-except-raise
# Re-raise CancelledError after cleanup (handled in finally)
raise
except Exception as e:
raise RuntimeError(f"Failed to create MCP session for {url}: {e}") from e
finally:
# Clean up on ANY failure (Exception, CancelledError, etc.)
# Only clean up if we didn't succeed
# Use anyio.move_on_after instead of asyncio.wait_for to properly propagate
# cancellation through anyio's cancel scope system (prevents orphaned spinning tasks)
if not success:
cleanup_timeout = _get_cleanup_timeout()
if session is not None:
with anyio.move_on_after(cleanup_timeout):
try:
await session.__aexit__(None, None, None) # pylint: disable=unnecessary-dunder-call
except Exception: # nosec B110 - Best effort cleanup on connection failure
pass
if transport_ctx is not None:
with anyio.move_on_after(cleanup_timeout):
try:
await transport_ctx.__aexit__(None, None, None) # pylint: disable=unnecessary-dunder-call
except Exception: # nosec B110 - Best effort cleanup on connection failure
pass
async def _close_session(self, pooled: PooledSession) -> None:
"""
Close a session and its transport.
Uses timeouts to prevent indefinite blocking if session/transport tasks
don't respond to cancellation. This prevents CPU spin loops in anyio's
_deliver_cancellation which can occur when async iterators or blocking
operations don't properly handle CancelledError.
Args:
pooled: The session to close.
"""
if pooled.is_closed:
return
pooled.mark_closed()
# Use anyio's move_on_after instead of asyncio.wait_for to properly propagate
# cancellation through anyio's cancel scope system. asyncio.wait_for() creates
# orphaned anyio tasks that keep spinning in _deliver_cancellation.
cleanup_timeout = _get_cleanup_timeout()
# Close session with anyio timeout
with anyio.move_on_after(cleanup_timeout) as session_scope:
try:
await pooled.session.__aexit__(None, None, None) # pylint: disable=unnecessary-dunder-call
except Exception as e:
logger.debug(f"Error closing session: {e}")
if session_scope.cancelled_caught:
logger.warning(f"Session cleanup timed out for {sanitize_url_for_logging(pooled.url)} - proceeding anyway")
# Close transport with anyio timeout
with anyio.move_on_after(cleanup_timeout) as transport_scope:
try:
await pooled.transport_context.__aexit__(None, None, None) # pylint: disable=unnecessary-dunder-call
except Exception as e:
logger.debug(f"Error closing transport: {e}")
if transport_scope.cancelled_caught:
logger.warning(f"Transport cleanup timed out for {sanitize_url_for_logging(pooled.url)} - proceeding anyway")
logger.debug(f"Closed session for {sanitize_url_for_logging(pooled.url)} (uses={pooled.use_count})")
# Clean up pool_owner key in Redis for session affinity
if settings.mcpgateway_session_affinity_enabled and pooled.headers:
headers_lower = {k.lower(): v for k, v in pooled.headers.items()}
mcp_session_id = headers_lower.get("x-mcp-session-id")
if mcp_session_id and self.is_valid_mcp_session_id(mcp_session_id):
await self._cleanup_pool_session_owner(mcp_session_id)
async def _cleanup_pool_session_owner(self, mcp_session_id: str) -> None:
"""Clean up pool_owner key in Redis when session is closed.
Only deletes the key if this worker owns it (to prevent removing other workers' ownership).
Args:
mcp_session_id: The MCP session ID from x-mcp-session-id header.
"""
try:
# First-Party
from mcpgateway.utils.redis_client import get_redis_client # pylint: disable=import-outside-toplevel
redis = await get_redis_client()
if redis:
key = self._pool_owner_key(mcp_session_id)
# Only delete if we own it
owner = await redis.get(key)
if owner:
owner_id = owner.decode() if isinstance(owner, bytes) else owner
if owner_id == WORKER_ID:
await redis.delete(key)
logger.debug(f"Cleaned up pool session owner: {mcp_session_id[:8]}...")
except Exception as e:
# Cleanup failure is non-fatal
logger.debug(f"Failed to cleanup pool session owner in Redis: {e}")
async def close_all(self) -> None:
"""
Gracefully close all pooled and active sessions.
Should be called during application shutdown.
"""
self._closed = True
logger.info("Closing all pooled sessions...")
async with self._global_lock:
# Close all pooled sessions
for _pool_key, pool in list(self._pools.items()):
while not pool.empty():
try:
pooled = pool.get_nowait()
await self._close_session(pooled)
except asyncio.QueueEmpty:
break
# Close all active sessions
for _pool_key, active_set in list(self._active.items()):
for pooled in list(active_set):
await self._close_session(pooled)
self._pools.clear()
self._active.clear()
self._locks.clear()
self._semaphores.clear()
# Stop RPC listener if running
if self._rpc_listener_task and not self._rpc_listener_task.done():
self._rpc_listener_task.cancel()
try:
await self._rpc_listener_task
except asyncio.CancelledError:
pass
self._rpc_listener_task = None
logger.info("All sessions closed")
async def register_pool_session_owner(self, mcp_session_id: str) -> None:
"""Register this worker as owner of a pool session in Redis.
This enables multi-worker session affinity by tracking which worker owns
which pool session. When a request with x-mcp-session-id arrives at a
different worker, it can forward the request to the owner worker.
Note: This method is now primarily used for refreshing TTL on existing ownership.
Initial ownership is claimed atomically in register_session_mapping() using SETNX.
Args:
mcp_session_id: The MCP session ID from x-mcp-session-id header.
"""
if not settings.mcpgateway_session_affinity_enabled:
return
if not self.is_valid_mcp_session_id(mcp_session_id):
logger.debug("Invalid mcp_session_id for owner registration, skipping")
return
try:
# First-Party
from mcpgateway.utils.redis_client import get_redis_client # pylint: disable=import-outside-toplevel
redis = await get_redis_client()
if redis:
key = self._pool_owner_key(mcp_session_id)
# Do not steal ownership: only claim if missing, or refresh TTL if we already own.
# Lua keeps this atomic.
script = """
local cur = redis.call('GET', KEYS[1])
if not cur then
redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])
return 1
end
if cur == ARGV[1] then
redis.call('EXPIRE', KEYS[1], ARGV[2])
return 2
end
return 0
"""
ttl = int(settings.mcpgateway_session_affinity_ttl)
outcome = await redis.eval(script, 1, key, WORKER_ID, ttl)
logger.debug(f"Owner registration outcome={outcome} for session {mcp_session_id[:8]}...")
except Exception as e:
# Redis failure is non-fatal - single worker mode still works
logger.debug(f"Failed to register pool session owner in Redis: {e}")
async def _get_pool_session_owner(self, mcp_session_id: str) -> Optional[str]:
"""Get the worker ID that owns a pool session.
Args:
mcp_session_id: The MCP session ID from x-mcp-session-id header.
Returns:
The worker ID that owns this session, or None if not found or Redis unavailable.
"""
if not settings.mcpgateway_session_affinity_enabled:
return None
if not self.is_valid_mcp_session_id(mcp_session_id):
return None
try:
# First-Party
from mcpgateway.utils.redis_client import get_redis_client # pylint: disable=import-outside-toplevel
redis = await get_redis_client()
if redis:
key = self._pool_owner_key(mcp_session_id)
owner = await redis.get(key)
if owner:
decoded = owner.decode() if isinstance(owner, bytes) else owner
return decoded
except Exception as e:
logger.debug(f"Failed to get pool session owner from Redis: {e}")
return None
async def forward_request_to_owner(
self,
mcp_session_id: str,
request_data: Dict[str, Any],
timeout: Optional[float] = None,
) -> Optional[Dict[str, Any]]:
"""Forward RPC request to the worker that owns the pool session.
This method checks Redis to find which worker owns the pool session for
the given mcp_session_id. If owned by another worker, it forwards the
request via Redis pub/sub and waits for the response.
Args:
mcp_session_id: The MCP session ID from x-mcp-session-id header.
request_data: The RPC request data to forward.
timeout: Optional timeout in seconds (default from config).
Returns:
The response from the owner worker, or None if we own the session
(caller should execute locally) or if Redis is unavailable.
Raises:
asyncio.TimeoutError: If the forwarded request times out.
"""
if not settings.mcpgateway_session_affinity_enabled:
return None
if not self.is_valid_mcp_session_id(mcp_session_id):
return None
effective_timeout = timeout if timeout is not None else settings.mcpgateway_pool_rpc_forward_timeout
try:
# First-Party
from mcpgateway.utils.redis_client import get_redis_client # pylint: disable=import-outside-toplevel
redis = await get_redis_client()
if not redis:
return None # Execute locally - no Redis
# Check who owns this session
owner = await redis.get(self._pool_owner_key(mcp_session_id))
method = request_data.get("method", "unknown")
if not owner:
logger.info(f"[AFFINITY] Worker {WORKER_ID} | Session {mcp_session_id[:8]}... | Method: {method} | No owner → execute locally (new session)")
return None # No owner registered - execute locally (new session)
owner_id = owner.decode() if isinstance(owner, bytes) else owner
if owner_id == WORKER_ID:
logger.info(f"[AFFINITY] Worker {WORKER_ID} | Session {mcp_session_id[:8]}... | Method: {method} | We own it → execute locally")
return None # We own it - execute locally
logger.info(f"[AFFINITY] Worker {WORKER_ID} | Session {mcp_session_id[:8]}... | Method: {method} | Owner: {owner_id} → forwarding")
# Forward to owner worker via pub/sub
response_id = str(uuid.uuid4())
response_channel = f"mcpgw:pool_rpc_response:{response_id}"
# Subscribe to response channel
pubsub = redis.pubsub()
await pubsub.subscribe(response_channel)
try:
# Prepare request with response channel
forward_data = {
"type": "rpc_forward",
**request_data,
"response_channel": response_channel,
"mcp_session_id": mcp_session_id,
}
# Publish request to owner's channel
await redis.publish(f"mcpgw:pool_rpc:{owner_id}", orjson.dumps(forward_data))
self._forwarded_requests += 1
logger.info(f"[AFFINITY] Worker {WORKER_ID} | Session {mcp_session_id[:8]}... | Method: {method} | Published to worker {owner_id}")
# Wait for response
async with asyncio.timeout(effective_timeout):
while True:
msg = await pubsub.get_message(ignore_subscribe_messages=True, timeout=0.1)
if msg and msg["type"] == "message":
return orjson.loads(msg["data"])
finally:
await pubsub.unsubscribe(response_channel)
except asyncio.TimeoutError:
self._forwarded_request_timeouts += 1
logger.warning(f"Timeout forwarding request to owner for session {mcp_session_id[:8]}...")
raise
except Exception as e:
self._forwarded_request_failures += 1
logger.debug(f"Error forwarding request to owner: {e}")
return None # Execute locally on error
async def start_rpc_listener(self) -> None:
"""Start listening for forwarded RPC and HTTP requests on this worker's channels.
This method subscribes to Redis pub/sub channels specific to this worker
and processes incoming forwarded requests from other workers:
- mcpgw:pool_rpc:{WORKER_ID} - for SSE transport JSON-RPC forwards
- mcpgw:pool_http:{WORKER_ID} - for Streamable HTTP request forwards
"""
if not settings.mcpgateway_session_affinity_enabled:
return
try:
# First-Party
from mcpgateway.utils.redis_client import get_redis_client # pylint: disable=import-outside-toplevel
redis = await get_redis_client()
if not redis:
logger.debug("Redis not available, RPC listener not started")
return
rpc_channel = f"mcpgw:pool_rpc:{WORKER_ID}"
http_channel = f"mcpgw:pool_http:{WORKER_ID}"
pubsub = redis.pubsub()
await pubsub.subscribe(rpc_channel, http_channel)
logger.info(f"RPC/HTTP listener started for worker {WORKER_ID} on channels: {rpc_channel}, {http_channel}")
while not self._closed:
try:
msg = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0)
if msg and msg["type"] == "message":
request = orjson.loads(msg["data"])
forward_type = request.get("type")
response_channel = request.get("response_channel")
if response_channel:
if forward_type == "rpc_forward":
# Execute forwarded RPC request for SSE transport
response = await self._execute_forwarded_request(request)
await redis.publish(response_channel, orjson.dumps(response))
logger.debug(f"Processed forwarded RPC request, response sent to {response_channel}")
elif forward_type == "http_forward":
# Execute forwarded HTTP request for Streamable HTTP transport
await self._execute_forwarded_http_request(request, redis)
else:
logger.warning(f"Unknown forward type: {forward_type}")
except asyncio.CancelledError:
break
except Exception as e:
logger.warning(f"Error processing forwarded request: {e}")
await pubsub.unsubscribe(rpc_channel, http_channel)
logger.info(f"RPC/HTTP listener stopped for worker {WORKER_ID}")
except Exception as e:
logger.warning(f"RPC/HTTP listener failed: {e}")
async def _execute_forwarded_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
"""Execute a forwarded RPC request locally via internal HTTP call.
This method handles RPC requests that were forwarded from another worker.
Instead of handling specific methods here, we make an internal HTTP call
to the local /rpc endpoint which reuses ALL existing method handling logic.
The x-forwarded-internally header prevents infinite forwarding loops.
Args:
request: The forwarded RPC request containing method, params, headers, req_id, etc.
Returns:
The JSON-RPC response from the local endpoint.
"""
try:
method = request.get("method")
params = request.get("params", {})
headers = request.get("headers", {})
req_id = request.get("req_id", 1)
mcp_session_id = request.get("mcp_session_id", "unknown")
session_short = mcp_session_id[:8] if len(mcp_session_id) >= 8 else mcp_session_id
logger.info(f"[AFFINITY] Worker {WORKER_ID} | Session {session_short}... | Method: {method} | Received forwarded request, executing locally")
# Make internal HTTP call to local /rpc endpoint
# This reuses ALL existing method handling logic without duplication
async with httpx.AsyncClient() as client:
# Build headers for internal request - forward original headers
# but add x-forwarded-internally to prevent infinite loops
internal_headers = dict(headers)
internal_headers["x-forwarded-internally"] = "true"
# Ensure content-type is set
internal_headers["content-type"] = "application/json"
response = await client.post(
f"http://127.0.0.1:{settings.port}/rpc",
json={"jsonrpc": "2.0", "method": method, "params": params, "id": req_id},
headers=internal_headers,
timeout=settings.mcpgateway_pool_rpc_forward_timeout,
)
# Parse response
response_data = response.json()
# Extract result or error from JSON-RPC response
if "error" in response_data:
logger.info(f"[AFFINITY] Worker {WORKER_ID} | Session {session_short}... | Method: {method} | Forwarded execution completed with error")
return {"error": response_data["error"]}
logger.info(f"[AFFINITY] Worker {WORKER_ID} | Session {session_short}... | Method: {method} | Forwarded execution completed successfully")
return {"result": response_data.get("result", {})}
except httpx.TimeoutException:
logger.warning(f"Timeout executing forwarded request: {request.get('method')}")
return {"error": {"code": -32603, "message": "Internal request timeout"}}
except Exception as e:
logger.warning(f"Error executing forwarded request: {e}")
return {"error": {"code": -32603, "message": str(e)}}
async def _execute_forwarded_http_request(self, request: Dict[str, Any], redis: Any) -> None:
"""Execute a forwarded HTTP request locally and return response via Redis.
This method handles full HTTP requests forwarded from other workers for
Streamable HTTP transport session affinity. It reconstructs the HTTP request,
makes an internal call to the appropriate endpoint, and publishes the response
back through Redis.
Args:
request: Serialized HTTP request data from Redis Pub/Sub containing:
- type: "http_forward"
- response_channel: Redis channel to publish response to
- mcp_session_id: Session identifier
- method: HTTP method (GET, POST, DELETE)
- path: Request path (e.g., /mcp)
- query_string: Query parameters
- headers: Request headers dict
- body: Hex-encoded request body
redis: Redis client for publishing response
"""
response_channel = None
try:
response_channel = request.get("response_channel")
method = request.get("method")
path = request.get("path")
query_string = request.get("query_string", "")
headers = request.get("headers", {})
body_hex = request.get("body", "")
mcp_session_id = request.get("mcp_session_id")
# Decode hex body back to bytes
body = bytes.fromhex(body_hex) if body_hex else b""
session_short = mcp_session_id[:8] if mcp_session_id and len(mcp_session_id) >= 8 else "unknown"
logger.debug(f"[HTTP_AFFINITY] Worker {WORKER_ID} | Session {session_short}... | Received forwarded HTTP request: {method} {path}")
# Add internal forwarding headers to prevent loops
internal_headers = dict(headers)
internal_headers["x-forwarded-internally"] = "true"
internal_headers["x-original-worker"] = request.get("original_worker", "unknown")
# Make internal HTTP request to local endpoint
url = f"http://127.0.0.1:{settings.port}{path}"
if query_string:
url = f"{url}?{query_string}"
async with httpx.AsyncClient() as client:
response = await client.request(
method=method,
url=url,
headers=internal_headers,
content=body,
timeout=settings.mcpgateway_pool_rpc_forward_timeout,
)
logger.debug(f"[HTTP_AFFINITY] Worker {WORKER_ID} | Session {session_short}... | Executed locally: {response.status_code}")
# Serialize response for Redis transport
response_data = {
"status": response.status_code,
"headers": dict(response.headers),
"body": response.content.hex(), # Hex encode binary response
}
# Publish response back to requesting worker
if redis and response_channel:
await redis.publish(response_channel, orjson.dumps(response_data))
logger.debug(f"[HTTP_AFFINITY] Published HTTP response to Redis channel: {response_channel}")
except Exception as e:
logger.error(f"Error executing forwarded HTTP request: {e}")
# Try to send error response if possible
if redis and response_channel:
error_response = {
"status": 500,
"headers": {"content-type": "application/json"},
"body": orjson.dumps({"error": "Internal forwarding error"}).hex(),
}
try:
await redis.publish(response_channel, orjson.dumps(error_response))
except Exception as publish_error:
logger.debug(f"Failed to publish error response via Redis: {publish_error}")
async def get_streamable_http_session_owner(self, mcp_session_id: str) -> Optional[str]:
"""Get the worker ID that owns a Streamable HTTP session.
This is a public wrapper around _get_pool_session_owner for use by
streamablehttp_transport to check session ownership before handling requests.
Args:
mcp_session_id: The MCP session ID from mcp-session-id header.
Returns:
Worker ID if found, None otherwise.
"""
return await self._get_pool_session_owner(mcp_session_id)
async def forward_streamable_http_to_owner(
self,
owner_worker_id: str,
mcp_session_id: str,
method: str,
path: str,
headers: Dict[str, str],
body: bytes,
query_string: str = "",
) -> Optional[Dict[str, Any]]:
"""Forward a Streamable HTTP request to the worker that owns the session via Redis Pub/Sub.
This method forwards the entire HTTP request to another worker using Redis
Pub/Sub channels, similar to forward_request_to_owner() for SSE transport.
This ensures session affinity works correctly in single-host multi-worker
deployments where hostname-based routing fails.
Args:
owner_worker_id: The worker ID that owns the session.
mcp_session_id: The MCP session ID.
method: HTTP method (GET, POST, DELETE).
path: Request path (e.g., /mcp).
headers: Request headers.
body: Request body bytes.
query_string: Query string if any.
Returns:
Dict with 'status', 'headers', and 'body' from the owner worker's response,
or None if forwarding fails.
"""
if not settings.mcpgateway_session_affinity_enabled:
return None
if not self.is_valid_mcp_session_id(mcp_session_id):
return None
session_short = mcp_session_id[:8] if len(mcp_session_id) >= 8 else mcp_session_id
logger.debug(f"[HTTP_AFFINITY] Worker {WORKER_ID} | Session {session_short}... | {method} {path} | Forwarding to worker {owner_worker_id}")
try:
# First-Party
from mcpgateway.utils.redis_client import get_redis_client # pylint: disable=import-outside-toplevel
redis = await get_redis_client()
if not redis:
logger.warning("Redis unavailable for HTTP forwarding, executing locally")
return None # Fall back to local execution
# Generate unique response channel for this request
response_uuid = uuid.uuid4().hex
response_channel = f"mcpgw:pool_http_response:{response_uuid}"
# Serialize HTTP request for Redis transport
forward_data = {
"type": "http_forward",
"response_channel": response_channel,
"mcp_session_id": mcp_session_id,
"method": method,
"path": path,
"query_string": query_string,
"headers": headers,
"body": body.hex() if body else "", # Hex encode binary body
"original_worker": WORKER_ID,
"timestamp": time.time(),
}
# Subscribe to response channel BEFORE publishing request (prevent race)
pubsub = redis.pubsub()
await pubsub.subscribe(response_channel)
try:
# Publish forwarded request to owner worker's HTTP channel
owner_channel = f"mcpgw:pool_http:{owner_worker_id}"
await redis.publish(owner_channel, orjson.dumps(forward_data))
logger.debug(f"[HTTP_AFFINITY] Published HTTP request to Redis channel: {owner_channel}")
# Wait for response with timeout
timeout = settings.mcpgateway_pool_rpc_forward_timeout
async with asyncio.timeout(timeout):
while True:
msg = await pubsub.get_message(ignore_subscribe_messages=True, timeout=0.1)
if msg and msg["type"] == "message":
response_data = orjson.loads(msg["data"])
logger.debug(f"[HTTP_AFFINITY] Received HTTP response via Redis: status={response_data.get('status')}")
# Decode hex body back to bytes
body_hex = response_data.get("body", "")
response_data["body"] = bytes.fromhex(body_hex) if body_hex else b""
self._forwarded_requests += 1
return response_data
finally:
await pubsub.unsubscribe(response_channel)
except asyncio.TimeoutError:
self._forwarded_request_timeouts += 1
logger.warning(f"Timeout forwarding HTTP request to owner {owner_worker_id}")
return None
except Exception as e:
self._forwarded_request_failures += 1
logger.warning(f"Error forwarding HTTP request via Redis: {e}")
return None
def get_metrics(self) -> Dict[str, Any]:
"""
Return pool metrics for monitoring.
Returns:
Dict with hits, misses, evictions, hit_rate, and per-pool stats.
"""
total_requests = self._hits + self._misses
total_affinity_requests = self._session_affinity_local_hits + self._session_affinity_redis_hits + self._session_affinity_misses
return {
"hits": self._hits,
"misses": self._misses,
"evictions": self._evictions,
"health_check_failures": self._health_check_failures,
"circuit_breaker_trips": self._circuit_breaker_trips,
"pool_keys_evicted": self._pool_keys_evicted,
"sessions_reaped": self._sessions_reaped,
"anonymous_identity_count": self._anonymous_identity_count,
"hit_rate": self._hits / total_requests if total_requests > 0 else 0.0,
"pool_key_count": len(self._pools),
# Session affinity metrics
"session_affinity": {
"local_hits": self._session_affinity_local_hits,
"redis_hits": self._session_affinity_redis_hits,
"misses": self._session_affinity_misses,
"hit_rate": (self._session_affinity_local_hits + self._session_affinity_redis_hits) / total_affinity_requests if total_affinity_requests > 0 else 0.0,
"forwarded_requests": self._forwarded_requests,
"forwarded_failures": self._forwarded_request_failures,
"forwarded_timeouts": self._forwarded_request_timeouts,
},
"pools": {
f"{url}|{identity[:8]}|{transport}|{user}|{gw_id[:8] if gw_id else 'none'}": {
"available": pool.qsize(),
"active": len(self._active.get((user, url, identity, transport, gw_id), set())),
"max": self._max_sessions,
}
for (user, url, identity, transport, gw_id), pool in self._pools.items()
},
"circuit_breakers": {
url: {
"failures": self._failures.get(url, 0),
"open_until": self._circuit_open_until.get(url),
}
for url in set(self._failures.keys()) | set(self._circuit_open_until.keys())
},
}
@asynccontextmanager
async def session(
self,
url: str,
headers: Optional[Dict[str, str]] = None,
transport_type: TransportType = TransportType.STREAMABLE_HTTP,
httpx_client_factory: Optional[HttpxClientFactory] = None,
timeout: Optional[float] = None,
user_identity: Optional[str] = None,
gateway_id: Optional[str] = None,
) -> "AsyncIterator[PooledSession]":
"""
Context manager for acquiring and releasing a session.
Usage:
async with pool.session(url, headers) as pooled:
result = await pooled.session.call_tool("my_tool", {})
Args:
url: The MCP server URL.
headers: Request headers.
transport_type: Transport type to use.
httpx_client_factory: Optional factory for httpx clients.
timeout: Optional timeout in seconds for transport connection.
user_identity: Optional user identity for strict isolation.
gateway_id: Optional gateway ID for notification handler context.
Yields:
PooledSession ready for use.
"""
pooled = await self.acquire(url, headers, transport_type, httpx_client_factory, timeout, user_identity, gateway_id)
try:
yield pooled
finally:
await self.release(pooled)
# Global pool instance - initialized by FastAPI lifespan
_mcp_session_pool: Optional[MCPSessionPool] = None
def get_mcp_session_pool() -> MCPSessionPool:
"""Get the global MCP session pool instance.
Returns:
The global MCPSessionPool instance.
Raises:
RuntimeError: If pool has not been initialized.
"""
if _mcp_session_pool is None:
raise RuntimeError("MCP session pool not initialized. Call init_mcp_session_pool() first.")
return _mcp_session_pool
def init_mcp_session_pool(
max_sessions_per_key: int = 10,
session_ttl_seconds: float = 300.0,
health_check_interval_seconds: float = 60.0,
acquire_timeout_seconds: float = 30.0,
session_create_timeout_seconds: float = 30.0,
circuit_breaker_threshold: int = 5,
circuit_breaker_reset_seconds: float = 60.0,
identity_headers: Optional[frozenset[str]] = None,
identity_extractor: Optional[IdentityExtractor] = None,
idle_pool_eviction_seconds: float = 600.0,
default_transport_timeout_seconds: float = 30.0,
health_check_methods: Optional[list[str]] = None,
health_check_timeout_seconds: float = 5.0,
message_handler_factory: Optional[MessageHandlerFactory] = None,
enable_notifications: bool = True,
notification_debounce_seconds: float = 5.0,
) -> MCPSessionPool:
"""Initialize the global MCP session pool.
Args:
See MCPSessionPool.__init__ for argument descriptions.
enable_notifications: Enable automatic notification service for list_changed events.
notification_debounce_seconds: Debounce interval for notification-triggered refreshes.
Returns:
The initialized MCPSessionPool instance.
"""
global _mcp_session_pool # pylint: disable=global-statement
# Auto-create notification service if enabled and no custom handler provided
effective_handler_factory = message_handler_factory
if enable_notifications and message_handler_factory is None:
# First-Party
from mcpgateway.services.notification_service import ( # pylint: disable=import-outside-toplevel
init_notification_service,
)
# Initialize notification service (will be started during acquire with gateway context)
notification_svc = init_notification_service(debounce_seconds=notification_debounce_seconds)
# Create default handler factory that uses notification service
def default_handler_factory(url: str, gateway_id: Optional[str]):
"""Create a message handler for MCP session notifications.
Args:
url: The MCP server URL for the session.
gateway_id: Optional gateway ID for attribution, falls back to URL if not provided.
Returns:
A message handler that forwards notifications to the notification service.
"""
return notification_svc.create_message_handler(gateway_id or url, url)
effective_handler_factory = default_handler_factory
logger.info("MCP notification service created (debounce=%ss)", notification_debounce_seconds)
_mcp_session_pool = MCPSessionPool(
max_sessions_per_key=max_sessions_per_key,
session_ttl_seconds=session_ttl_seconds,
health_check_interval_seconds=health_check_interval_seconds,
acquire_timeout_seconds=acquire_timeout_seconds,
session_create_timeout_seconds=session_create_timeout_seconds,
circuit_breaker_threshold=circuit_breaker_threshold,
circuit_breaker_reset_seconds=circuit_breaker_reset_seconds,
identity_headers=identity_headers,
identity_extractor=identity_extractor,
idle_pool_eviction_seconds=idle_pool_eviction_seconds,
default_transport_timeout_seconds=default_transport_timeout_seconds,
health_check_methods=health_check_methods,
health_check_timeout_seconds=health_check_timeout_seconds,
message_handler_factory=effective_handler_factory,
)
logger.info("MCP session pool initialized")
return _mcp_session_pool
async def close_mcp_session_pool() -> None:
"""Close the global MCP session pool and notification service."""
global _mcp_session_pool # pylint: disable=global-statement
if _mcp_session_pool is not None:
await _mcp_session_pool.close_all()
_mcp_session_pool = None
logger.info("MCP session pool closed")
# Close notification service if it was initialized
try:
# First-Party
from mcpgateway.services.notification_service import ( # pylint: disable=import-outside-toplevel
close_notification_service,
)
await close_notification_service()
except (ImportError, RuntimeError):
pass # Notification service not initialized
async def start_pool_notification_service(gateway_service: Any = None) -> None:
"""Start the notification service background worker.
Call this after gateway_service is initialized to enable event-driven refresh.
Args:
gateway_service: Optional GatewayService instance for triggering refreshes.
"""
try:
# First-Party
from mcpgateway.services.notification_service import ( # pylint: disable=import-outside-toplevel
get_notification_service,
)
notification_svc = get_notification_service()
await notification_svc.initialize(gateway_service)
logger.info("MCP notification service started")
except RuntimeError:
logger.debug("Notification service not configured, skipping start")
def register_gateway_capabilities_for_notifications(gateway_id: str, capabilities: Dict[str, Any]) -> None:
"""Register gateway capabilities for notification handling.
Call this after gateway initialization to enable list_changed notifications.
Args:
gateway_id: The gateway ID.
capabilities: Server capabilities from initialization response.
"""
try:
# First-Party
from mcpgateway.services.notification_service import ( # pylint: disable=import-outside-toplevel
get_notification_service,
)
notification_svc = get_notification_service()
notification_svc.register_gateway_capabilities(gateway_id, capabilities)
except RuntimeError:
pass # Notification service not initialized
def unregister_gateway_from_notifications(gateway_id: str) -> None:
"""Unregister a gateway from notification handling.
Call this when a gateway is deleted.
Args:
gateway_id: The gateway ID to unregister.
"""
try:
# First-Party
from mcpgateway.services.notification_service import ( # pylint: disable=import-outside-toplevel
get_notification_service,
)
notification_svc = get_notification_service()
notification_svc.unregister_gateway(gateway_id)
except RuntimeError:
pass # Notification service not initialized
