import asyncio import base64 import concurrent.futures import hashlib import json import logging import time from typing import Any, Callable, Dict, List, Optional from urllib.parse import parse_qs, urlparse from fastmcp import Client from fastmcp.client.auth import BearerAuth from fastmcp.client.transports import ( SSETransport, StdioTransport, StreamableHttpTransport, ) from mcp.client.auth import OAuthClientProvider, TokenStorage from mcp.shared.auth import ( AuthorizationCodeResult, OAuthClientInformationFull, OAuthClientMetadata, OAuthToken, ) from pydantic import AnyHttpUrl, ValidationError from redis import Redis from docsgpt.agents.tools.base import Tool from docsgpt.api.user.tasks import mcp_oauth_task from docsgpt.cache import get_redis_instance from docsgpt.core.settings import settings from docsgpt.core.url_validation import SSRFError, validate_url from docsgpt.events.keys import stream_key from docsgpt.security.encryption import decrypt_credentials logger = logging.getLogger(__name__) _mcp_clients_cache = {} def forget_cached_clients(*identities: str) -> None: """Drop cached OAuth clients signed in as any of ``identities``. OAuth cache keys name the connection or user whose tokens the client holds (see ``MCPTool._generate_cache_key``). Args: identities: Connection ids and user ids. """ markers = tuple(f"#oauth:{identity}:" for identity in identities if identity) for key in [k for k in list(_mcp_clients_cache) if any(marker in k for marker in markers)]: _mcp_clients_cache.pop(key, None) # A token expiry long past: a stored token of unknown age is renewed before use. _EXPIRED = 1.0 _ANNOTATION_HINTS = ("readOnlyHint", "destructiveHint", "idempotentHint", "openWorldHint") def _secret_fingerprint(secret: str) -> str: """A short digest of a whole secret, for telling cached clients apart. A prefix of the secret itself is not enough: every fine-grained GitHub token starts with ``github_pat_``, so two users' tokens would share (and reuse) one cached client carrying the first user's token. """ return hashlib.sha256(secret.encode("utf-8")).hexdigest()[:16] def _annotation_hints(annotations: Any) -> Dict[str, bool]: """The boolean MCP tool annotation hints, from a model or a dict.""" if annotations is None: return {} if hasattr(annotations, "model_dump"): annotations = annotations.model_dump() if not isinstance(annotations, dict): return {} return {k: annotations[k] for k in _ANNOTATION_HINTS if isinstance(annotations.get(k), bool)} class MCPTool(Tool): """ MCP Tool Connect to remote Model Context Protocol (MCP) servers to access dynamic tools and resources. """ connection_id: Optional[str] = None def __init__(self, config: Dict[str, Any], user_id: Optional[str] = None): """ Initialize the MCP Tool with configuration. Args: config: Dictionary containing MCP server configuration: - server_url: URL of the remote MCP server - transport_type: Transport type (auto, sse, http, stdio) - auth_type: Type of authentication (bearer, oauth, api_key, basic, none) - encrypted_credentials: Encrypted credentials (if available) - timeout: Request timeout in seconds (default: 30) - headers: Custom headers for requests - command: Command for STDIO transport - args: Arguments for STDIO transport - oauth_scopes: OAuth scopes for oauth auth type - oauth_client_name: OAuth client name for oauth auth type - query_mode: If True, use non-interactive OAuth (fail-fast on 401) user_id: User ID for decrypting credentials (required if encrypted_credentials exist) """ self.config = config self.user_id = user_id raw_url = config.get("server_url", "") self.server_url = self._validate_server_url(raw_url) if raw_url else "" self.transport_type = config.get("transport_type", "auto") self.auth_type = config.get("auth_type", "none") self.timeout = config.get("timeout", 30) self.custom_headers = config.get("headers", {}) self.auth_credentials = {} if config.get("encrypted_credentials") and user_id: self.auth_credentials = decrypt_credentials( config["encrypted_credentials"], user_id ) else: self.auth_credentials = config.get("auth_credentials", {}) self.oauth_scopes = config.get("oauth_scopes", []) self.oauth_task_id = config.get("oauth_task_id", None) self.oauth_client_name = config.get("oauth_client_name", "DocsGPT-MCP") # The connection whose OAuth tokens this tool uses. Set by the tool # executor for connection-backed tools, so a shared tool in ``owner`` # mode signs in with its owner's account rather than the invoker's. self.connection_id = config.get("connection_id") self.redirect_uri = self._resolve_redirect_uri(config.get("redirect_uri")) # Pulled out of ``config`` (rather than left in ``self.config``) # because it is a callable supplied by the OAuth worker — not # something the rest of the tool plumbing should marshal or # serialize. ``DocsGPTOAuth`` invokes it from ``redirect_handler`` # so the SSE envelope can carry ``authorization_url``. self.oauth_redirect_publish = config.pop("oauth_redirect_publish", None) self.available_tools = [] self._cache_key = self._generate_cache_key() self._client = None self.query_mode = config.get("query_mode", False) if self.server_url and self.auth_type != "oauth": self._setup_client() @staticmethod def _validate_server_url(server_url: str) -> str: """Validate server_url to prevent SSRF to internal networks. Raises: ValueError: If the URL points to a private/internal address. """ try: return validate_url(server_url) except SSRFError as exc: raise ValueError(f"Invalid MCP server URL: {exc}") from exc def _resolve_redirect_uri(self, configured_redirect_uri: Optional[str]) -> str: # The operator's setting wins over the page's own origin: a page opened # on a plain-HTTP address would otherwise register a callback that # servers such as Linear refuse. explicit = settings.MCP_OAUTH_REDIRECT_URI if explicit: return explicit.rstrip("/") if configured_redirect_uri: return configured_redirect_uri.rstrip("/") connector_base = settings.CONNECTOR_REDIRECT_BASE_URI if connector_base: parsed = urlparse(connector_base) if parsed.scheme and parsed.netloc: return f"{parsed.scheme}://{parsed.netloc}/api/mcp_server/callback" return f"{settings.API_URL.rstrip('/')}/api/mcp_server/callback" def _generate_cache_key(self) -> str: """Generate a unique cache key for this MCP server configuration.""" auth_key = "" if self.auth_type == "oauth": scopes_str = ",".join(self.oauth_scopes) if self.oauth_scopes else "none" # A connection-backed tool shares a client only with calls that # use the same connection's tokens. oauth_identity = self.connection_id or self.user_id or self.oauth_task_id or "anonymous" auth_key = ( f"oauth:{oauth_identity}:{self.oauth_client_name}:{scopes_str}:{self.redirect_uri}" ) elif self.auth_type in ["bearer"]: token = self.auth_credentials.get( "bearer_token", "" ) or self.auth_credentials.get("access_token", "") auth_key = f"bearer:{_secret_fingerprint(token)}" if token else "bearer:none" elif self.auth_type == "api_key": api_key = self.auth_credentials.get("api_key", "") auth_key = f"apikey:{_secret_fingerprint(api_key)}" if api_key else "apikey:none" elif self.auth_type == "basic": username = self.auth_credentials.get("username", "") auth_key = f"basic:{username}" else: auth_key = "none" return f"{self.server_url}#{self.transport_type}#{auth_key}" def _setup_client(self): global _mcp_clients_cache if self._cache_key in _mcp_clients_cache: cached_data = _mcp_clients_cache[self._cache_key] if time.time() - cached_data["created_at"] < 300: self._client = cached_data["client"] return else: del _mcp_clients_cache[self._cache_key] transport = self._create_transport() auth = None if self.auth_type == "oauth": redis_client = get_redis_instance() if self.query_mode: auth = NonInteractiveOAuth( mcp_url=self.server_url, scopes=self.oauth_scopes, redis_client=redis_client, redirect_uri=self.redirect_uri, user_id=self.user_id, connection_id=self.connection_id, ) else: auth = DocsGPTOAuth( mcp_url=self.server_url, scopes=self.oauth_scopes, redis_client=redis_client, redirect_uri=self.redirect_uri, task_id=self.oauth_task_id, user_id=self.user_id, redirect_publish=self.oauth_redirect_publish, connection_id=self.connection_id, ) elif self.auth_type == "bearer": token = self.auth_credentials.get( "bearer_token", "" ) or self.auth_credentials.get("access_token", "") if token: auth = BearerAuth(token) self._client = Client(transport, auth=auth) _mcp_clients_cache[self._cache_key] = { "client": self._client, "created_at": time.time(), } def _create_transport(self): """Create appropriate transport based on configuration.""" headers = {"Content-Type": "application/json", "User-Agent": "DocsGPT-MCP/1.0"} headers.update(self.custom_headers) if self.auth_type == "api_key": api_key = self.auth_credentials.get("api_key", "") header_name = self.auth_credentials.get("api_key_header", "X-API-Key") if api_key: headers[header_name] = api_key elif self.auth_type == "basic": username = self.auth_credentials.get("username", "") password = self.auth_credentials.get("password", "") if username and password: credentials = base64.b64encode( f"{username}:{password}".encode() ).decode() headers["Authorization"] = f"Basic {credentials}" if self.transport_type == "auto": if "sse" in self.server_url.lower() or self.server_url.endswith("/sse"): transport_type = "sse" else: transport_type = "http" else: transport_type = self.transport_type if transport_type == "stdio": raise ValueError("STDIO transport is disabled") if transport_type == "sse": headers.update({"Accept": "text/event-stream", "Cache-Control": "no-cache"}) return SSETransport(url=self.server_url, headers=headers) elif transport_type == "http": return StreamableHttpTransport(url=self.server_url, headers=headers) elif transport_type == "stdio": command = self.config.get("command", "python") args = self.config.get("args", []) env = self.auth_credentials if self.auth_credentials else None return StdioTransport(command=command, args=args, env=env) else: return StreamableHttpTransport(url=self.server_url, headers=headers) def _format_tools(self, tools_response) -> List[Dict]: """Format tools response to match expected format.""" if hasattr(tools_response, "tools"): tools = tools_response.tools elif isinstance(tools_response, list): tools = tools_response else: tools = [] tools_dict = [] for tool in tools: if hasattr(tool, "name"): tool_dict = { "name": tool.name, "description": tool.description, } if hasattr(tool, "inputSchema"): tool_dict["inputSchema"] = tool.inputSchema annotations = _annotation_hints(getattr(tool, "annotations", None)) if annotations: tool_dict["annotations"] = annotations tools_dict.append(tool_dict) elif isinstance(tool, dict): tools_dict.append(tool) else: if hasattr(tool, "model_dump"): tools_dict.append(tool.model_dump()) else: tools_dict.append({"name": str(tool), "description": ""}) return tools_dict async def _execute_with_client(self, operation: str, *args, **kwargs): """Execute operation with FastMCP client.""" if not self._client: raise Exception("FastMCP client not initialized") async with self._client: if operation == "ping": return await self._client.ping() elif operation == "list_tools": tools_response = await self._client.list_tools() self.available_tools = self._format_tools(tools_response) return self.available_tools elif operation == "call_tool": tool_name = args[0] tool_args = kwargs return await self._client.call_tool(tool_name, tool_args) elif operation == "list_resources": return await self._client.list_resources() elif operation == "list_prompts": return await self._client.list_prompts() else: raise Exception(f"Unknown operation: {operation}") _ERROR_MAP = [ (concurrent.futures.TimeoutError, lambda op, t, _: f"Timed out after {t}s"), (ConnectionRefusedError, lambda *_: "Connection refused"), ] _ERROR_PATTERNS = { ("403", "Forbidden"): "Access denied (403 Forbidden)", ("401", "Unauthorized"): "Authentication failed (401 Unauthorized)", ("ECONNREFUSED",): "Connection refused", ("SSL", "certificate"): "SSL/TLS error", } def _run_async_operation(self, operation: str, *args, **kwargs): try: try: asyncio.get_running_loop() with concurrent.futures.ThreadPoolExecutor() as executor: future = executor.submit( self._run_in_new_loop, operation, *args, **kwargs ) return future.result(timeout=self.timeout) except RuntimeError: return self._run_in_new_loop(operation, *args, **kwargs) except Exception as e: raise self._map_error(operation, e) from e raise self._map_error(operation, e) from e def _run_in_new_loop(self, operation, *args, **kwargs): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: return loop.run_until_complete( self._execute_with_client(operation, *args, **kwargs) ) finally: loop.close() def _map_error(self, operation: str, exc: Exception) -> Exception: for exc_type, msg_fn in self._ERROR_MAP: if isinstance(exc, exc_type): return Exception(msg_fn(operation, self.timeout, exc)) error_msg = str(exc) for patterns, friendly in self._ERROR_PATTERNS.items(): if any(p.lower() in error_msg.lower() for p in patterns): return Exception(friendly) logger.error("MCP %s failed: %s", operation, exc) return exc def discover_tools(self) -> List[Dict]: """ Discover available tools from the MCP server using FastMCP. Returns: List of tool definitions from the server """ if not self.server_url: return [] if not self._client: self._setup_client() try: tools = self._run_async_operation("list_tools") self.available_tools = tools return self.available_tools except Exception as e: raise Exception(f"Failed to discover tools from MCP server: {str(e)}") def execute_action(self, action_name: str, **kwargs) -> Any: if not self.server_url: raise Exception("No MCP server configured") if not self._client: self._setup_client() cleaned_kwargs = {} for key, value in kwargs.items(): if value == "" or value is None: continue cleaned_kwargs[key] = value try: result = self._run_async_operation( "call_tool", action_name, **cleaned_kwargs ) return self._format_result(result) except Exception as e: error_msg = str(e) lower_msg = error_msg.lower() is_auth_error = ( "401" in error_msg or "unauthorized" in lower_msg or "session expired" in lower_msg or "re-authorize" in lower_msg ) if is_auth_error: if self.auth_type == "oauth": raise Exception( f"Action '{action_name}' failed: OAuth session expired. " "Please re-authorize this MCP server in tool settings." ) from e global _mcp_clients_cache _mcp_clients_cache.pop(self._cache_key, None) self._client = None self._setup_client() try: result = self._run_async_operation( "call_tool", action_name, **cleaned_kwargs ) return self._format_result(result) except Exception as retry_e: raise Exception( f"Action '{action_name}' failed after re-auth attempt: {retry_e}. " "Your credentials may have expired — please re-authorize in tool settings." ) from retry_e raise Exception( f"Failed to execute action '{action_name}': {error_msg}" ) from e def _format_result(self, result) -> Dict: """Format FastMCP result to match expected format.""" if hasattr(result, "content"): content_list = [] for content_item in result.content: if hasattr(content_item, "text"): content_list.append({"type": "text", "text": content_item.text}) elif hasattr(content_item, "data"): content_list.append({"type": "data", "data": content_item.data}) else: content_list.append( {"type": "unknown", "content": str(content_item)} ) return { "content": content_list, "isError": getattr(result, "isError", False), } else: return result def test_connection(self) -> Dict: if not self.server_url: return { "success": False, "message": "No server URL configured", "tools_count": 0, } try: parsed = urlparse(self.server_url) if parsed.scheme not in ("http", "https"): return { "success": False, "message": f"Invalid URL scheme '{parsed.scheme}' — use http:// or https://", "tools_count": 0, } except Exception: return { "success": False, "message": "Invalid URL format", "tools_count": 0, } if not self._client: try: self._setup_client() except Exception as e: return { "success": False, "message": f"Client init failed: {str(e)}", "tools_count": 0, } try: if self.auth_type == "oauth": return self._test_oauth_connection() else: return self._test_regular_connection() except Exception as e: return { "success": False, "message": f"Connection failed: {str(e)}", "tools_count": 0, } def _test_regular_connection(self) -> Dict: ping_ok = False ping_error = None try: self._run_async_operation("ping") ping_ok = True except Exception as e: ping_error = str(e) try: tools = self.discover_tools() except Exception as e: return { "success": False, "message": f"Connection failed: {ping_error or str(e)}", "tools_count": 0, } if not tools and not ping_ok: return { "success": False, "message": f"Connection failed: {ping_error or 'No tools found'}", "tools_count": 0, } return { "success": True, "message": f"Connected — found {len(tools)} tool{'s' if len(tools) != 1 else ''}.", "tools_count": len(tools), "tools": [ { "name": tool.get("name", "unknown"), "description": tool.get("description", ""), } for tool in tools ], } def _test_oauth_connection(self) -> Dict: storage = DBTokenStorage( server_url=self.server_url, user_id=self.user_id, connection_id=self.connection_id, ) loop = asyncio.new_event_loop() try: tokens = loop.run_until_complete(storage.get_tokens()) finally: loop.close() if tokens and tokens.access_token: self.query_mode = True _mcp_clients_cache.pop(self._cache_key, None) self._client = None self._setup_client() try: tools = self.discover_tools() return { "success": True, "message": f"Connected — found {len(tools)} tool{'s' if len(tools) != 1 else ''}.", "tools_count": len(tools), "tools": [ { "name": t.get("name", "unknown"), "description": t.get("description", ""), } for t in tools ], } except Exception as e: logger.warning("OAuth token validation failed: %s", e) _mcp_clients_cache.pop(self._cache_key, None) self._client = None return self._start_oauth_task() def _start_oauth_task(self) -> Dict: task_config = self.config.copy() task_config.pop("query_mode", None) result = mcp_oauth_task.delay(task_config, self.user_id) return { "success": False, "requires_oauth": True, "task_id": result.id, "message": "OAuth authorization required.", "tools_count": 0, } def get_actions_metadata(self) -> List[Dict]: """ Get metadata for all available actions. Returns: List of action metadata dictionaries """ actions = [] for tool in self.available_tools: input_schema = ( tool.get("inputSchema") or tool.get("input_schema") or tool.get("schema") or tool.get("parameters") ) parameters_schema = { "type": "object", "properties": {}, "required": [], } if input_schema: if isinstance(input_schema, dict): if "properties" in input_schema: parameters_schema = { "type": input_schema.get("type", "object"), "properties": input_schema.get("properties", {}), "required": input_schema.get("required", []), } for key in ["additionalProperties", "description"]: if key in input_schema: parameters_schema[key] = input_schema[key] else: parameters_schema["properties"] = input_schema action = { "name": tool.get("name", ""), "description": tool.get("description", ""), "parameters": parameters_schema, } # ``readOnlyHint`` / ``destructiveHint`` decide whether the action # is a read (always allowed) or a write (needs approval). annotations = _annotation_hints(tool.get("annotations")) if annotations: action["annotations"] = annotations actions.append(action) return actions def get_config_requirements(self) -> Dict: return { "server_url": { "type": "string", "label": "Server URL", "description": "URL of the remote MCP server", "required": True, "secret": False, "order": 1, }, "auth_type": { "type": "string", "label": "Authentication Type", "description": "Authentication method for the MCP server", "enum": ["none", "bearer", "oauth", "api_key", "basic"], "default": "none", "required": True, "secret": False, "order": 2, }, "api_key": { "type": "string", "label": "API Key", "description": "API key for authentication", "required": False, "secret": True, "order": 3, "depends_on": {"auth_type": "api_key"}, }, "api_key_header": { "type": "string", "label": "API Key Header", "description": "Header name for API key (default: X-API-Key)", "default": "X-API-Key", "required": False, "secret": False, "order": 4, "depends_on": {"auth_type": "api_key"}, }, "bearer_token": { "type": "string", "label": "Bearer Token", "description": "Bearer token for authentication", "required": False, "secret": True, "order": 3, "depends_on": {"auth_type": "bearer"}, }, "username": { "type": "string", "label": "Username", "description": "Username for basic authentication", "required": False, "secret": False, "order": 3, "depends_on": {"auth_type": "basic"}, }, "password": { "type": "string", "label": "Password", "description": "Password for basic authentication", "required": False, "secret": True, "order": 4, "depends_on": {"auth_type": "basic"}, }, "oauth_scopes": { "type": "string", "label": "OAuth Scopes", "description": "Comma-separated OAuth scopes to request", "required": False, "secret": False, "order": 3, "depends_on": {"auth_type": "oauth"}, }, "timeout": { "type": "number", "label": "Timeout (seconds)", "description": "Request timeout in seconds (1-300)", "default": 30, "required": False, "secret": False, "order": 10, }, } class MCPReauthorizationRequired(Exception): """An MCP sign-in expired and could not be renewed: its owner must sign in again.""" class DocsGPTOAuth(OAuthClientProvider): """ Custom OAuth handler for DocsGPT that uses frontend redirect instead of browser. """ def __init__( self, mcp_url: str, redirect_uri: str, redis_client: Redis | None = None, redis_prefix: str = "mcp_oauth:", task_id: str = None, scopes: str | list[str] | None = None, client_name: str = "DocsGPT-MCP", user_id=None, additional_client_metadata: dict[str, Any] | None = None, skip_redirect_validation: bool = False, redirect_publish=None, connection_id: Optional[str] = None, ): self.redirect_uri = redirect_uri self.redis_client = redis_client self.redis_prefix = redis_prefix self.task_id = task_id self.user_id = user_id # Worker-supplied callback. Invoked from ``redirect_handler`` # once the authorization URL is known so the SSE envelope can # carry it. ``None`` for any non-worker entrypoint. self.redirect_publish = redirect_publish parsed_url = urlparse(mcp_url) self.server_base_url = f"{parsed_url.scheme}://{parsed_url.netloc}" if isinstance(scopes, list): scopes = " ".join(scopes) client_metadata = OAuthClientMetadata( client_name=client_name, redirect_uris=[AnyHttpUrl(redirect_uri)], grant_types=["authorization_code", "refresh_token"], response_types=["code"], scope=scopes, **(additional_client_metadata or {}), ) storage = DBTokenStorage( server_url=self.server_base_url, user_id=self.user_id, expected_redirect_uri=None if skip_redirect_validation else redirect_uri, connection_id=connection_id, ) # The SDK checks the server's protected-resource metadata against a # resource derived from this URL, so it gets the full MCP endpoint: # Linear and Sentry publish ``https://host/mcp``, and an origin-only # URL fails that check before sign-in. Tokens stay keyed by origin. super().__init__( server_url=mcp_url.rstrip("/") or self.server_base_url, client_metadata=client_metadata, storage=storage, redirect_handler=self.redirect_handler, callback_handler=self.callback_handler, ) self.context.prepare_token_auth = self._one_client_authentication(self.context.prepare_token_auth) self.auth_url = None self.extracted_state = None @staticmethod def _one_client_authentication(prepare: Callable) -> Callable: """Wrap the SDK's token-request auth so the client authenticates one way. With ``client_secret_basic`` the SDK sets the Basic header but leaves ``client_id`` in the body, which servers such as Linear reject as a second method; RFC 6749 names the client in the body only when it does not authenticate otherwise. Args: prepare: The SDK context's ``prepare_token_auth``. Returns: The same function, minus ``client_id`` in the body under Basic auth. """ def prepare_token_auth( data: dict[str, str], headers: dict[str, str] | None = None ) -> tuple[dict[str, str], dict[str, str]]: data, headers = prepare(data, headers) if headers.get("Authorization", "").startswith("Basic "): data = {k: v for k, v in data.items() if k != "client_id"} return data, headers return prepare_token_auth async def _initialize(self) -> None: """Load the stored sign-in along with when its access token expires. The SDK learns a token's expiry only from a token response in this process, so a token read back from the connection gets the expiry saved with it, and an expired one is renewed with its refresh token before it is sent. One saved before expiries were kept is of unknown age and is renewed once. """ await super()._initialize() tokens = self.context.current_tokens if tokens is None: return expires_at = getattr(self.context.storage, "expires_at", None) if expires_at: self.context.token_expiry_time = float(expires_at) elif tokens.refresh_token and tokens.expires_in: self.context.token_expiry_time = _EXPIRED def _process_auth_url(self, authorization_url: str) -> tuple[str, str]: """Process authorization URL to extract state""" try: parsed_url = urlparse(authorization_url) query_params = parse_qs(parsed_url.query) state_params = query_params.get("state", []) if state_params: state = state_params[0] else: raise ValueError("No state in auth URL") return authorization_url, state except Exception as e: raise Exception(f"Failed to process auth URL: {e}") async def redirect_handler(self, authorization_url: str) -> None: """Store auth URL and state in Redis for frontend to use.""" auth_url, state = self._process_auth_url(authorization_url) logger.info("Processed auth_url: %s, state: %s", auth_url, state) self.auth_url = auth_url self.extracted_state = state if self.redis_client and self.extracted_state: key = f"{self.redis_prefix}auth_url:{self.extracted_state}" self.redis_client.setex(key, 600, auth_url) logger.info("Stored auth_url in Redis: %s", key) if self.redirect_publish is not None: # Best-effort: a publish failure must not abort the OAuth # handshake — the user can still authorize via the popup # opened from the legacy polling fallback if the SSE # envelope is lost. try: self.redirect_publish(auth_url) except Exception: logger.warning( "redirect_publish callback raised for task_id=%s", self.task_id, exc_info=True, ) async def callback_handler(self) -> AuthorizationCodeResult: """Wait for auth code from Redis using the state value. Returns: The code, the state it came back with, and the RFC 9207 issuer when the authorization server sent one. """ if not self.redis_client or not self.extracted_state: raise Exception("Redis client or state not configured for OAuth") poll_interval = 1 max_wait_time = 300 code_key = f"{self.redis_prefix}code:{self.extracted_state}" start_time = time.time() while time.time() - start_time < max_wait_time: code_data = self.redis_client.get(code_key) if code_data: code = code_data.decode() returned_state = self.extracted_state iss_key = f"{self.redis_prefix}iss:{self.extracted_state}" iss_data = self.redis_client.get(iss_key) self.redis_client.delete(code_key) self.redis_client.delete(iss_key) self.redis_client.delete( f"{self.redis_prefix}auth_url:{self.extracted_state}" ) self.redis_client.delete( f"{self.redis_prefix}state:{self.extracted_state}" ) return AuthorizationCodeResult( code=code, state=returned_state, iss=iss_data.decode() if iss_data else None, ) error_key = f"{self.redis_prefix}error:{self.extracted_state}" error_data = self.redis_client.get(error_key) if error_data: error_msg = error_data.decode() self.redis_client.delete(error_key) self.redis_client.delete( f"{self.redis_prefix}auth_url:{self.extracted_state}" ) self.redis_client.delete( f"{self.redis_prefix}state:{self.extracted_state}" ) raise Exception(f"OAuth error: {error_msg}") await asyncio.sleep(poll_interval) self.redis_client.delete(f"{self.redis_prefix}auth_url:{self.extracted_state}") self.redis_client.delete(f"{self.redis_prefix}state:{self.extracted_state}") raise Exception("OAuth timeout: no code received within 5 minutes") class NonInteractiveOAuth(DocsGPTOAuth): """OAuth provider that fails fast on 401 instead of starting interactive auth. Used during query execution to prevent the streaming response from blocking while waiting for user authorization that will never come. """ def __init__(self, **kwargs): kwargs.setdefault("task_id", None) kwargs["skip_redirect_validation"] = True super().__init__(**kwargs) async def redirect_handler(self, authorization_url: str) -> None: raise MCPReauthorizationRequired( "OAuth session expired — please re-authorize this MCP server in tool settings." ) async def callback_handler(self) -> AuthorizationCodeResult: raise MCPReauthorizationRequired( "OAuth session expired — please re-authorize this MCP server in tool settings." ) class DBTokenStorage(TokenStorage): """MCP OAuth tokens and client registration, kept encrypted on the connection. Reads and writes go through ``docsgpt.connectors.service``, which stores them in the connection's owner-bound ``encrypted_credentials``. A tool that runs with a specific connection (``owner`` mode on a shared tool) passes its ``connection_id``; otherwise the invoking user's connection for the server's base URL is used. """ def __init__( self, server_url: str, user_id: str, expected_redirect_uri: Optional[str] = None, connection_id: Optional[str] = None, ): self.server_url = server_url self.user_id = user_id self.expected_redirect_uri = expected_redirect_uri self.connection_id = connection_id # When the stored access token expires (epoch seconds), once read. self.expires_at: Optional[float] = None @staticmethod def get_base_url(url: str) -> str: parsed = urlparse(url) return f"{parsed.scheme}://{parsed.netloc}" def _pg_provider(self) -> str: return f"mcp:{self.get_base_url(self.server_url)}" def _fetch_session_data(self) -> dict: """The decrypted ``tokens`` / ``client_info`` for this MCP server.""" from docsgpt.connectors import service return service.read_mcp_secrets( self.user_id, self.get_base_url(self.server_url), self.connection_id, ) async def get_tokens(self) -> OAuthToken | None: data = await asyncio.to_thread(self._fetch_session_data) if not data or "tokens" not in data: return None self.expires_at = data.get("tokens_expires_at") try: return OAuthToken.model_validate(data["tokens"]) except ValidationError as e: logger.error("Could not load tokens: %s", e) return None def _merge(self, patch: dict) -> None: """Merge ``patch`` into the connection's secrets; ``None`` drops a key.""" from docsgpt.connectors import service status = None if patch.get("tokens"): status = service.STATUS_CONNECTED service.update_mcp_secrets( self.user_id, self.get_base_url(self.server_url), patch, connection_id=self.connection_id, status=status, ) def _delete(self) -> None: from docsgpt.storage.db.repositories.connector_sessions import ( ConnectorSessionsRepository, ) from docsgpt.storage.db.session import db_session with db_session() as conn: ConnectorSessionsRepository(conn).delete( self.user_id, self._pg_provider(), ) async def set_tokens(self, tokens: OAuthToken) -> None: base_url = self.get_base_url(self.server_url) token_dump = tokens.model_dump() # Kept beside the token: a later process cannot tell from # ``expires_in`` alone whether it has expired. self.expires_at = time.time() + tokens.expires_in if tokens.expires_in else None await asyncio.to_thread(self._merge, {"tokens": token_dump, "tokens_expires_at": self.expires_at}) logger.info("Saved tokens for %s", base_url) async def get_client_info(self) -> OAuthClientInformationFull | None: data = await asyncio.to_thread(self._fetch_session_data) base_url = self.get_base_url(self.server_url) if not data or "client_info" not in data: logger.debug("No client_info in DB for %s", base_url) return None try: client_info = OAuthClientInformationFull.model_validate(data["client_info"]) if self.expected_redirect_uri: stored_uris = [ str(uri).rstrip("/") for uri in client_info.redirect_uris ] expected_uri = self.expected_redirect_uri.rstrip("/") if expected_uri not in stored_uris: logger.warning( "Redirect URI mismatch for %s: expected=%s stored=%s — clearing.", base_url, expected_uri, stored_uris, ) # Drop ``tokens`` and ``client_info`` from the JSONB # blob via merge_session_data's ``None``-drops-key # semantics — preserves the row + any other keys. await asyncio.to_thread( self._merge, {"tokens": None, "client_info": None}, ) return None return client_info except ValidationError as e: logger.error("Could not load client info: %s", e) return None def _serialize_client_info(self, info: dict) -> dict: if "redirect_uris" in info and isinstance(info["redirect_uris"], list): info["redirect_uris"] = [str(u) for u in info["redirect_uris"]] return info async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: serialized_info = self._serialize_client_info(client_info.model_dump()) base_url = self.get_base_url(self.server_url) await asyncio.to_thread( self._merge, {"client_info": serialized_info}, ) logger.info("Saved client info for %s", base_url) async def clear(self) -> None: await asyncio.to_thread(self._delete) logger.info("Cleared OAuth cache for %s", self.get_base_url(self.server_url)) @classmethod async def clear_all(cls, db_client=None) -> None: """Delete every MCP-tagged connector session row. ``db_client`` retained for call-site compatibility but unused — storage is Postgres-only now. """ from sqlalchemy import text from docsgpt.storage.db.session import db_session def _delete_all() -> None: with db_session() as conn: conn.execute( text( "DELETE FROM connector_sessions " "WHERE provider LIKE 'mcp:%'" ) ) await asyncio.to_thread(_delete_all) logger.info("Cleared all OAuth client cache data.") class MCPOAuthManager: """Manager for handling MCP OAuth callbacks.""" def __init__(self, redis_client: Redis | None, redis_prefix: str = "mcp_oauth:"): self.redis_client = redis_client self.redis_prefix = redis_prefix def handle_oauth_callback( self, state: str, code: str, error: Optional[str] = None, iss: Optional[str] = None ) -> bool: """ Handle OAuth callback from provider. Args: state: The state parameter from OAuth callback code: The authorization code from OAuth callback error: Error message if OAuth failed iss: The RFC 9207 issuer from the callback, if the server sent one Returns: True if successful, False otherwise """ try: if not self.redis_client or not state: raise Exception("Redis client or state not provided") if error: error_key = f"{self.redis_prefix}error:{state}" self.redis_client.setex(error_key, 300, error) raise Exception(f"OAuth error received: {error}") # The issuer goes first: the waiting sign-in reads it once the code lands. if iss: self.redis_client.setex(f"{self.redis_prefix}iss:{state}", 300, iss) code_key = f"{self.redis_prefix}code:{state}" self.redis_client.setex(code_key, 300, code) state_key = f"{self.redis_prefix}state:{state}" self.redis_client.setex(state_key, 300, "completed") return True except Exception as e: logger.error("Error handling OAuth callback: %s", e) return False def get_oauth_status(self, task_id: str, user_id: str) -> Dict[str, Any]: """Return the latest OAuth status for ``task_id`` from the user's SSE journal. Mirrors the legacy polling contract: ``status`` derived from the ``mcp.oauth.*`` event-type suffix, with payload fields surfaced (e.g. ``tools``/``tools_count`` on ``completed``). """ if not task_id: return {"status": "not_started", "message": "OAuth flow not started"} if not user_id: return {"status": "not_found", "message": "User not provided"} if self.redis_client is None: return {"status": "not_found", "message": "Redis unavailable"} try: # OAuth flows are short-lived but a concurrent source # ingest can flood the user channel between the OAuth # popup completing and the user clicking Save, pushing the # completion envelope outside the read window. Bound the # scan by the configured stream cap so we cover the full # journal — XADD MAXLEN keeps that bounded too. scan_count = max(settings.EVENTS_STREAM_MAXLEN, 200) entries = self.redis_client.xrevrange( stream_key(user_id), count=scan_count ) except Exception: logger.exception( "xrevrange failed for oauth status: user_id=%s task_id=%s", user_id, task_id, ) return {"status": "not_found", "message": "Status unavailable"} for _entry_id, fields in entries: if not isinstance(fields, dict): continue # decode_responses=False ⇒ bytes keys; the string-key fallback # covers a future flip of that default without a forced refactor. event_raw = fields.get(b"event") if event_raw is None: event_raw = fields.get("event") if event_raw is None: continue if isinstance(event_raw, bytes): try: event_raw = event_raw.decode("utf-8") except Exception: continue try: envelope = json.loads(event_raw) except Exception: continue if not isinstance(envelope, dict): continue event_type = envelope.get("type", "") if not isinstance(event_type, str) or not event_type.startswith( "mcp.oauth." ): continue scope = envelope.get("scope") or {} if scope.get("kind") != "mcp_oauth" or scope.get("id") != task_id: continue payload = envelope.get("payload") or {} return { "status": event_type[len("mcp.oauth."):], "task_id": task_id, **payload, } return {"status": "not_found", "message": "Status not found"}