Enkrypt AI Secure MCP Gateway

"""MCP Gateway client module.""" import hashlib import json import sys import threading import time from datetime import datetime import redis as external_cache_server from mcp import ClientSession, StdioServerParameters # https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/client/stdio/__init__.py from mcp.client.stdio import stdio_client from secure_mcp_gateway.services.oauth.integration import ( inject_oauth_into_args, inject_oauth_into_env, prepare_oauth_for_server, ) from secure_mcp_gateway.utils import get_common_config, logger from secure_mcp_gateway.version import __version__ # logger.info(f"Initializing Enkrypt Secure MCP Gateway Client Module v{__version__}") common_config = get_common_config() ENKRYPT_LOG_LEVEL = common_config.get("enkrypt_log_level", "INFO").lower() IS_DEBUG_LOG_LEVEL = ENKRYPT_LOG_LEVEL == "debug" # --- Cache Configuration --- ENKRYPT_MCP_USE_EXTERNAL_CACHE = common_config.get( "enkrypt_mcp_use_external_cache", False ) ENKRYPT_CACHE_HOST = common_config.get("enkrypt_cache_host", "localhost") ENKRYPT_CACHE_PORT = int(common_config.get("enkrypt_cache_port", "6379")) ENKRYPT_CACHE_DB = int(common_config.get("enkrypt_cache_db", "0")) ENKRYPT_CACHE_PASSWORD = common_config.get("enkrypt_cache_password", None) # Cache expiration times (in hours) ENKRYPT_TOOL_CACHE_EXPIRATION = int( common_config.get("enkrypt_tool_cache_expiration", 4) ) # 4 hours ENKRYPT_GATEWAY_CACHE_EXPIRATION = int( common_config.get("enkrypt_gateway_cache_expiration", 24) ) # 24 hours (1 day) local_cache = {} local_cache_lock = threading.Lock() # Add local registries for gateway config and servers local_key_map = {} local_server_registry = {} local_gateway_config_registry = set() # --- Cache connection --- def initialize_cache(): """ Initializes and tests the connection to the Redis cache server. This function creates a Redis client instance with the configured connection parameters and verifies the connection is working. If external cache is disabled, it returns None. Returns: Redis: A configured Redis client instance if external cache is enabled None: If external cache is disabled or connection fails Raises: ConnectionError: If unable to connect to the Redis server when external cache is enabled """ # Initialize Cache client cache_client = external_cache_server.Redis( host=ENKRYPT_CACHE_HOST, port=ENKRYPT_CACHE_PORT, db=ENKRYPT_CACHE_DB, password=ENKRYPT_CACHE_PASSWORD, decode_responses=True, # Automatically decode responses to strings ) # Test Cache connection try: cache_client.ping() logger.info( f"[external_cache] Successfully connected to External Cache at {ENKRYPT_CACHE_HOST}:{ENKRYPT_CACHE_PORT}" ) except external_cache_server.ConnectionError as e: logger.error(f"[external_cache] Failed to connect to External Cache: {e}") logger.error( "[external_cache] Exiting as External Cache is required for this gateway" ) sys.exit(1) # Exit if External Cache is unavailable return cache_client # --- Cache key patterns with hashing --- def hash_key(key): """ Creates an MD5 hash of the given key for secure cache storage. Args: key (str): The key to be hashed Returns: str: MD5 hash of the input key """ return hashlib.md5(key.encode()).hexdigest() def get_server_hashed_key(id, server_name): """ Generates a hashed cache key for server tools. Args: id (str): ID includes project_id, user_id, and mcp_config_id context. server_name (str): Name of the server Returns: str: Hashed cache key for the server tools """ raw_key = f"{id}-{server_name}-tools" return hash_key(raw_key) def get_gateway_config_hashed_key(id): """ Generates a hashed cache key for gateway configuration. Args: id (str): The ID of the Gateway or User Returns: str: Hashed cache key for the gateway configuration """ raw_key = f"{id}-mcp-config" return hash_key(raw_key) def get_hashed_key(key): """ Generates a hashed cache key for key to gateway/user ID mapping. Args: key (str): The Gateway/API key Returns: str: Hashed cache key for the API key mapping """ raw_key = f"gateway_key-{key}" return hash_key(raw_key) # --- Registry cache keys for gateway/user-server and gateway/user tracking --- def get_gateway_servers_registry_hashed_key(id): """ Generates a hashed cache key for the servers registry. Args: id (str): The ID of the Gateway or User Returns: str: Hashed cache key for the gateway/user's servers registry """ return hash_key(f"registry:{id}:servers") def get_gateway_registry_hashed_key(): """ Generates a hashed cache key for the global gateway/user registry. Returns: str: Hashed cache key for the gateway/user registry """ return hash_key("registry:gateway/user") # --- Tool forwarding function --- # This discovers tools and also invokes a specific tool if tool_name is provided async def get_server_metadata_only(server_name, gateway_config=None): """ Gets only server metadata (description, name, version) without discovering tools. This function is used for config servers that already have tools defined, but we need to get their dynamic description for validation. Args: server_name (str): Name of the server gateway_config (dict): Gateway/user's configuration containing server details Returns: dict: Server metadata including description, name, and version """ if not gateway_config: logger.error("[get_server_metadata_only] Error: No gateway_config provided") raise ValueError("No gateway configuration provided") mcp_config = gateway_config.get("mcp_config", []) server_entry = next( (s for s in mcp_config if s.get("server_name") == server_name), None ) if not server_entry: raise ValueError(f"No config found for server: {server_name}") config = server_entry["config"] command = config["command"] command_args = config["args"] env = config.get("env", None) logger.info( f"[get_server_metadata_only] Getting metadata for server: {server_name}" ) # Prepare OAuth for this server if configured # Extract project_id and mcp_config_id from gateway_config project_id = gateway_config.get("project_id") mcp_config_id = gateway_config.get("mcp_config_id") oauth_data, oauth_error = await prepare_oauth_for_server( server_name=server_name, server_entry=server_entry, config_id=mcp_config_id, project_id=project_id, ) if oauth_error: logger.error( f"[get_server_metadata_only] OAuth preparation failed for {server_name}: {oauth_error}" ) # Continue without OAuth - let the server handle authentication failure elif oauth_data: logger.info( f"[get_server_metadata_only] OAuth configured for {server_name}, injecting credentials" ) # Inject OAuth environment variables env = inject_oauth_into_env(env, oauth_data) # Inject OAuth header arguments for remote servers command_args = inject_oauth_into_args(command_args, oauth_data) if IS_DEBUG_LOG_LEVEL: logger.debug(f"[get_server_metadata_only] Command: {command}") logger.debug(f"[get_server_metadata_only] Command args: {command_args}") # Mask sensitive environment variables from secure_mcp_gateway.utils import mask_sensitive_data masked_env = mask_sensitive_data(env or {}) if env else None logger.debug(f"[get_server_metadata_only] Env: {masked_env}") async with stdio_client( StdioServerParameters(command=command, args=command_args, env=env) ) as (read, write): async with ClientSession(read, write) as session: # Initialize and capture server metadata ONLY init_result = await session.initialize() # Extract server description from initialization response server_info = getattr(init_result, "serverInfo", {}) if hasattr(server_info, "description"): server_description = server_info.description else: server_description = getattr(server_info, "description", "") if hasattr(server_info, "name"): server_name_from_server = server_info.name else: server_name_from_server = getattr(server_info, "name", "unknown") if hasattr(server_info, "version"): server_version = server_info.version else: server_version = getattr(server_info, "version", "unknown") logger.info( f"[get_server_metadata_only] Connected successfully to {server_name}" ) logger.info("[get_server_metadata_only] 🔍 Dynamic Server Metadata:") logger.info( f"[get_server_metadata_only] 📝 Description: '{server_description}'" ) logger.info( f"[get_server_metadata_only] 🏷️ Name: '{server_name_from_server}'" ) logger.info(f"[get_server_metadata_only] 📦 Version: '{server_version}'") # Return only metadata, NO tool discovery return { "server_metadata": { "description": server_description, "name": server_name_from_server, "version": server_version, } } async def forward_tool_call(server_name, tool_name, args=None, gateway_config=None): """ Forwards tool calls to the appropriate MCP server. This function handles both tool discovery (when tool_name is None) and tool invocation. It uses the gateway/user's configuration to determine the correct server and connection details. Args: server_name (str): Name of the server to call tool_name (str): Name of the tool to call (None for discovery) args (dict, optional): Arguments to pass to the tool gateway_config (dict): Gateway/user's configuration containing server details Returns: dict/ListToolsResult: Tool discovery results or tool call response Raises: ValueError: If gateway_config is missing or server not found """ if not gateway_config: logger.error("[forward_tool_call] Error: No gateway_config provided") raise ValueError("No gateway configuration provided") mcp_config = gateway_config.get("mcp_config", []) server_entry = next( (s for s in mcp_config if s.get("server_name") == server_name), None ) if not server_entry: raise ValueError(f"No config found for server: {server_name}") config = server_entry["config"] command = config["command"] command_args = config["args"] env = config.get("env", None) logger.info( f"[forward_tool_call] Starting tool call for server: {server_name} and tool: {tool_name}" ) # Prepare OAuth for this server if configured # Extract project_id and mcp_config_id from gateway_config project_id = gateway_config.get("project_id") mcp_config_id = gateway_config.get("mcp_config_id") oauth_data, oauth_error = await prepare_oauth_for_server( server_name=server_name, server_entry=server_entry, config_id=mcp_config_id, project_id=project_id, ) if oauth_error: logger.error( f"[forward_tool_call] OAuth preparation failed for {server_name}: {oauth_error}" ) # Continue without OAuth - let the server handle authentication failure elif oauth_data: logger.info( f"[forward_tool_call] OAuth configured for {server_name}, injecting credentials" ) # Inject OAuth environment variables env = inject_oauth_into_env(env, oauth_data) # Inject OAuth header arguments for remote servers command_args = inject_oauth_into_args(command_args, oauth_data) if IS_DEBUG_LOG_LEVEL: logger.debug(f"[forward_tool_call] Command: {command}") logger.debug(f"[forward_tool_call] Command args: {command_args}") # Mask sensitive environment variables from secure_mcp_gateway.utils import mask_sensitive_data masked_env = mask_sensitive_data(env or {}) if env else None logger.debug(f"[forward_tool_call] Env: {masked_env}") async with stdio_client( StdioServerParameters(command=command, args=command_args, env=env) ) as (read, write): async with ClientSession(read, write) as session: # Initialize and capture server metadata init_result = await session.initialize() # Extract server description from initialization response # InitializeResult is a Pydantic model, so we access attributes directly server_info = getattr(init_result, "serverInfo", {}) if hasattr(server_info, "description"): server_description = server_info.description else: server_description = getattr(server_info, "description", "") if hasattr(server_info, "name"): server_name_from_server = server_info.name else: server_name_from_server = getattr(server_info, "name", "unknown") if hasattr(server_info, "version"): server_version = server_info.version else: server_version = getattr(server_info, "version", "unknown") # Store the dynamic server metadata for later use session._server_description = server_description session._server_name = server_name_from_server session._server_version = server_version session._server_info = server_info logger.info(f"[forward_tool_call] Connected successfully to {server_name}") logger.info("[forward_tool_call] 🔍 Dynamic Server Metadata:") logger.info(f"[forward_tool_call] 📝 Description: '{server_description}'") logger.info(f"[forward_tool_call] 🏷️ Name: '{server_name_from_server}'") logger.info(f"[forward_tool_call] 📦 Version: '{server_version}'") # If tool_name is None, this is a discovery request if tool_name is None: logger.info( "[forward_tool_call] Starting tool discovery as tool_name is None" ) # Request tool listing from the MCP server tools_result = await session.list_tools() try: # Safely print the tools result to avoid async context issues tools_summary = f"[forward_tool_call] Discovered {len(getattr(tools_result or {}, 'tools', []))} tools for {server_name}" logger.info(tools_summary) if IS_DEBUG_LOG_LEVEL and tools_result is not None: logger.debug( f"[forward_tool_call] Tool details for {server_name}: {tools_result}" ) except Exception as e: logger.error( f"[forward_tool_call] Tools discovered for {server_name}, but encountered error when printing details: {e}" ) # Return tools result with dynamic server metadata return { "tools": tools_result, "server_metadata": { "description": server_description, "name": server_name_from_server, "version": server_version, "server_info": server_info, }, } # Normal tool call logger.info(f"[forward_tool_call] Calling specific tool: {tool_name}") return await session.call_tool(tool_name, arguments=args) # --- Cache management functions --- def set_local_cache(key, value, expires_in_seconds): """ Stores a value in the local in-memory cache with expiration. Args: key (str): The cache key value: The value to cache expires_in_seconds (int): Time in seconds until the cache entry expires """ expires_at = time.time() + expires_in_seconds with local_cache_lock: local_cache[key] = (value, expires_at) def get_local_cache(key): """ Retrieves a value from the local in-memory cache. Args: key (str): The cache key Returns: tuple: (value, expiration_time) if found and not expired, None otherwise """ with local_cache_lock: item = local_cache.get(key) if not item: return None value, expires_at = item if time.time() > expires_at: del local_cache[key] return None return value, expires_at def get_cached_tools(cache_client, id, server_name): """ Retrieves cached tools for a specific gateway/user and server. Args: cache_client: The cache client instance id (str): ID of the Gateway or User server_name (str): Name of the server Returns: tuple: (tools_data, expiration_time) if found and not expired, None otherwise """ key = get_server_hashed_key(id, server_name) if not ENKRYPT_MCP_USE_EXTERNAL_CACHE: return get_local_cache(key) if cache_client is None: return None cached_data = cache_client.get(key) if not cached_data: return None try: tool_data = json.loads(cached_data) if IS_DEBUG_LOG_LEVEL: logger.debug( f"[external_cache] Using cached tools for id '{id}', server '{server_name}' with key hash: {key}" ) return tool_data except json.JSONDecodeError: logger.error( f"[external_cache] Error deserializing tools cache for hash key: {key}" ) return None def cache_tools(cache_client, id, server_name, tools): """ Caches tools for a specific gateway/user and server. Args: cache_client: The cache client instance id (str): ID of the Gateway or User server_name (str): Name of the server tools: The tools data to cache """ expires_in_seconds = int(ENKRYPT_TOOL_CACHE_EXPIRATION * 3600) key = get_server_hashed_key(id, server_name) if not ENKRYPT_MCP_USE_EXTERNAL_CACHE: set_local_cache(key, tools, expires_in_seconds) # Also set the server name in the local_server_registry as set of server in id if id not in local_server_registry: local_server_registry[id] = set() local_server_registry[id].add(server_name) return if cache_client is None: return raw_key = f"{id}-{server_name}-tools" # For logging only # Convert ListToolsResult to a serializable format if needed if hasattr(tools, "__class__") and tools.__class__.__name__ == "ListToolsResult": # Extract data from ListToolsResult into a serializable dictionary serializable_tools = { "tools": [ { "name": tool.name, "description": tool.description, "inputSchema": tool.inputSchema, } for tool in getattr(tools, "tools", []) ] } # Add any other attributes that might be useful serialized_data = json.dumps(serializable_tools) else: # Try to serialize directly try: serialized_data = json.dumps(tools) except TypeError as e: logger.error(f"[external_cache] Warning: Cannot serialize tools - {e}") logger.error("[external_cache] Using simplified tool representation") # Fall back to a simplified format if serialization fails # This handles cases where tools has complex objects inside if ( isinstance(tools, dict) and "tools" in tools and isinstance(tools["tools"], list) ): # Handle case where tools is a dict with a tools list simplified = { "tools": [ { "name": t.get("name") if isinstance(t, dict) else getattr(t, "name", str(t)), "description": t.get("description") if isinstance(t, dict) else getattr(t, "description", ""), "inputSchema": t.get("inputSchema") if isinstance(t, dict) else getattr(t, "inputSchema", {}), } for t in tools["tools"] ] } elif hasattr(tools, "__iter__") and not isinstance(tools, (str, bytes)): # Handle case where tools is an iterable (list, etc.) simplified = { "tools": [ { "name": t.get("name") if isinstance(t, dict) else getattr(t, "name", str(t)), "description": t.get("description") if isinstance(t, dict) else getattr(t, "description", ""), "inputSchema": t.get("inputSchema") if isinstance(t, dict) else getattr(t, "inputSchema", {}), } for t in tools ] } else: # Fall back to an empty dict if we can't make sense of the structure simplified = {"tools": []} serialized_data = json.dumps(simplified) # Store in External Cache with expiration cache_client.setex(key, expires_in_seconds, serialized_data) if IS_DEBUG_LOG_LEVEL: expiration_time = datetime.fromtimestamp( time.time() + expires_in_seconds ).strftime("%Y-%m-%d %H:%M:%S") logger.debug( f"[external_cache] Cached tools for gateway/user '{id}', server '{server_name}' with key '{raw_key}' (hash: {key}) until {expiration_time}" ) # Maintain a registry of servers for this gateway/user registry_key = get_gateway_servers_registry_hashed_key(id) cache_client.sadd(registry_key, server_name) cache_client.expire(registry_key, expires_in_seconds) # Register the gateway/user in the gateway/user registry if not already there gateway_registry = get_gateway_registry_hashed_key() cache_client.sadd(gateway_registry, id) def get_cached_gateway_config(cache_client, id): """ Retrieves cached gateway configuration. Args: cache_client: The cache client instance id (str): ID of the Gateway or User Returns: tuple: (config_data, expiration_time) if found and not expired, None otherwise """ config_key = get_gateway_config_hashed_key(id) if not ENKRYPT_MCP_USE_EXTERNAL_CACHE: return get_local_cache(config_key) if cache_client is None: return None cached_data = cache_client.get(config_key) if not cached_data: return None try: config_data = json.loads(cached_data) if IS_DEBUG_LOG_LEVEL: logger.debug( f"[external_cache] Using cached config for id '{id}' with key hash: {config_key}" ) return config_data except json.JSONDecodeError: logger.error( f"[external_cache] Error deserializing config cache for hash key: {config_key}" ) return None def cache_gateway_config(cache_client, id, config): """ Caches gateway configuration. Args: cache_client: The cache client instance id (str): ID of the Gateway or User config (dict): The gateway configuration to cache """ expires_in_seconds = int(ENKRYPT_GATEWAY_CACHE_EXPIRATION * 3600) config_key = get_gateway_config_hashed_key(id) if not ENKRYPT_MCP_USE_EXTERNAL_CACHE: set_local_cache(config_key, config, expires_in_seconds) local_gateway_config_registry.add(id) return if cache_client is None: return serialized_data = json.dumps(config) cache_client.setex(config_key, expires_in_seconds, serialized_data) if IS_DEBUG_LOG_LEVEL: expiration_time = datetime.fromtimestamp( time.time() + expires_in_seconds ).strftime("%Y-%m-%d %H:%M:%S") logger.debug( f"[external_cache] Cached gateway config for '{id}' with key '{id}-mcp-config' (hash: {config_key}) until {expiration_time}" ) gateway_registry = get_gateway_registry_hashed_key() cache_client.sadd(gateway_registry, id) def cache_key_to_id(cache_client, gateway_key, id): """ Caches the mapping between a key and gateway/user ID. Args: cache_client: The cache client instance gateway_key (str): The key for gateway/user id (str): ID of the Gateway or User """ expires_in_seconds = int(ENKRYPT_GATEWAY_CACHE_EXPIRATION * 3600) key = get_hashed_key(gateway_key) if not ENKRYPT_MCP_USE_EXTERNAL_CACHE: local_key_map[key] = id return if cache_client is None: return cache_client.setex(key, expires_in_seconds, id) if IS_DEBUG_LOG_LEVEL: logger.debug( f"[external_cache] Cached key mapping with key 'gateway_key-****{gateway_key[-4:]}' (hash: {key})" ) def get_id_from_key(cache_client, gateway_key): """ Retrieves the gateway/user ID associated with a key. Args: cache_client: The cache client instance gateway_key (str): The key for gateway/user Returns: str: The associated gateway/user ID if found, None otherwise """ key = get_hashed_key(gateway_key) if not ENKRYPT_MCP_USE_EXTERNAL_CACHE: return local_key_map.get(key) if cache_client is None: return None id = cache_client.get(key) if id: if IS_DEBUG_LOG_LEVEL: logger.debug(f"[external_cache] Found id for key with hash: {key}") return id def clear_cache_for_servers(cache_client, id, server_name=None): """ Clears tool cache for specific or all servers for a gateway/user. Args: cache_client: The cache client instance id (str): ID of the Gateway or User server_name (str, optional): Name of the server to clear cache for Returns: int: Number of cache entries cleared """ if IS_DEBUG_LOG_LEVEL: logger.debug( f"[clear_cache_for_servers] Clearing cache for servers for gateway/user: {id} with current local_server_registry: {local_server_registry}" ) count = 0 # Local cache clear if server_name: if IS_DEBUG_LOG_LEVEL: logger.info( f"[clear_cache_for_servers] Clearing cache for server: {server_name}" ) key = get_server_hashed_key(id, server_name) if key in local_cache: del local_cache[key] count += 1 if id in local_server_registry: local_server_registry[id].discard(server_name) else: if IS_DEBUG_LOG_LEVEL: logger.debug( "[clear_cache_for_servers] Clearing cache for all servers for gateway/user" ) # Clear all servers for a gateway/user if id in local_server_registry: if IS_DEBUG_LOG_LEVEL: logger.debug( f"[clear_cache_for_servers] Clearing cache for all servers for gateway/user found in local_server_registry: {id}" ) for server_name in list(local_server_registry[id]): if IS_DEBUG_LOG_LEVEL: logger.debug( f"[clear_cache_for_servers] Clearing cache for server: {server_name}" ) key = get_server_hashed_key(id, server_name) if key in local_cache: if IS_DEBUG_LOG_LEVEL: logger.debug( f"[clear_cache_for_servers] Clearing cache for server: {server_name} found in local_cache" ) del local_cache[key] count += 1 local_server_registry[id].clear() else: if IS_DEBUG_LOG_LEVEL: logger.debug( f"[clear_cache_for_servers] Clearing cache for all servers for gateway/user not found in local_server_registry: {id}" ) if cache_client is None: return count # External cache clear count = 0 # Resetting as it is external cache if server_name: key = get_server_hashed_key(id, server_name) if cache_client.exists(key): cache_client.delete(key) registry_key = get_gateway_servers_registry_hashed_key(id) cache_client.srem(registry_key, server_name) count += 1 return count else: registry_key = get_gateway_servers_registry_hashed_key(id) servers = ( cache_client.smembers(registry_key) if cache_client.exists(registry_key) else [] ) for server_name in servers: key = get_server_hashed_key(id, server_name) if cache_client.exists(key): cache_client.delete(key) count += 1 cache_client.delete(registry_key) return count def clear_gateway_config_cache(cache_client, id, gateway_key): """ Clears all cache entries for a gateway/user, including config, tools, and key mapping. Args: cache_client: The cache client instance id (str): ID of the Gateway or User gateway_key (str): The gateway/user's key Returns: bool: True if any cache entries were cleared """ if IS_DEBUG_LOG_LEVEL: logger.debug( "[clear_gateway_config_cache] Clearing all cache entries for gateway/user" ) # 1. Clear all tool caches for the gateway/user registry_key = get_gateway_servers_registry_hashed_key(id) if IS_DEBUG_LOG_LEVEL: logger.debug( f"[clear_gateway_config_cache] Clearing all tool caches for gateway/user: {id} with registry key: {registry_key}" ) servers = ( cache_client.smembers(registry_key) if cache_client and cache_client.exists(registry_key) else [] ) if IS_DEBUG_LOG_LEVEL: logger.debug( f"[clear_gateway_config_cache] Clearing all tool caches for gateway/user: {id} with servers: {servers}" ) for server_name in servers: tool_key = get_server_hashed_key(id, server_name) if IS_DEBUG_LOG_LEVEL: logger.debug( f"[clear_gateway_config_cache] Clearing tool cache for server: {server_name} with tool_key: {tool_key}" ) if cache_client and cache_client.exists(tool_key): cache_client.delete(tool_key) if cache_client and cache_client.exists(registry_key): cache_client.delete(registry_key) # 2. Clear gateway config cache (local and external cache) config_key = get_gateway_config_hashed_key(id) if config_key in local_cache: del local_cache[config_key] local_gateway_config_registry.discard(id) if cache_client and cache_client.exists(config_key): cache_client.delete(config_key) # 3. Remove key mapping if gateway_key is provided if gateway_key: gateway_key_hash = get_hashed_key(gateway_key) if cache_client and cache_client.exists(gateway_key_hash): cache_client.delete(gateway_key_hash) if gateway_key_hash in local_key_map: del local_key_map[gateway_key_hash] # 4. Remove gateway/user from gateway/user registry (local and external cache) if cache_client: gateway_registry = get_gateway_registry_hashed_key() cache_client.srem(gateway_registry, id) local_server_registry.pop(id, None) return True def get_cache_statistics(cache_client): """ Retrieves statistics about the current cache state. Args: cache_client: The cache client instance Returns: dict: Cache statistics including: - total_gateways: Number of gateway/users in cache - total_tool_caches: Number of tool caches - total_config_caches: Number of config caches - cache_type: Type of cache being used """ if not ENKRYPT_MCP_USE_EXTERNAL_CACHE: total_gateways = len(local_gateway_config_registry) total_tool_caches = sum(len(s) for s in local_server_registry.values()) total_config_caches = len(local_gateway_config_registry) return { "total_gateways": total_gateways, "total_tool_caches": total_tool_caches, "total_config_caches": total_config_caches, "cache_type": "local", } if cache_client is None: return { "total_gateways": 0, "total_tool_caches": 0, "total_config_caches": 0, "cache_type": "none", } gateway_registry = get_gateway_registry_hashed_key() total_gateways = cache_client.scard(gateway_registry) total_tool_caches = 0 total_config_caches = 0 gateways = cache_client.smembers(gateway_registry) for id in gateways: config_key = get_gateway_config_hashed_key(id) if cache_client.exists(config_key): total_config_caches += 1 servers_registry = get_gateway_servers_registry_hashed_key(id) if cache_client.exists(servers_registry): server_count = cache_client.scard(servers_registry) total_tool_caches += server_count return { "total_gateways": total_gateways, "total_tool_caches": total_tool_caches, "total_config_caches": total_config_caches, "cache_type": "external_cache", }