alita-sdk 0.3.351__py3-none-any.whl → 0.3.499__py3-none-any.whl

alita_sdk/runtime/utils/mcp_oauth.py

@@ -0,0 +1,244 @@
+import json
+import logging
+import re
+from typing import Any, Dict, Optional
+from urllib.parse import urlparse
+
+import requests
+from langchain_core.tools import ToolException
+
+logger = logging.getLogger(__name__)
+
+
+class McpAuthorizationRequired(ToolException):
+    """Raised when an MCP server requires OAuth authorization before use."""
+
+    def __init__(
+        self,
+        message: str,
+        server_url: str,
+        resource_metadata_url: Optional[str] = None,
+        www_authenticate: Optional[str] = None,
+        resource_metadata: Optional[Dict[str, Any]] = None,
+        status: Optional[int] = None,
+        tool_name: Optional[str] = None,
+    ):
+        super().__init__(message)
+        self.server_url = server_url
+        self.resource_metadata_url = resource_metadata_url
+        self.www_authenticate = www_authenticate
+        self.resource_metadata = resource_metadata
+        self.status = status
+        self.tool_name = tool_name
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "message": str(self),
+            "server_url": self.server_url,
+            "resource_metadata_url": self.resource_metadata_url,
+            "www_authenticate": self.www_authenticate,
+            "resource_metadata": self.resource_metadata,
+            "status": self.status,
+            "tool_name": self.tool_name,
+        }
+
+
+def extract_resource_metadata_url(www_authenticate: Optional[str], server_url: Optional[str] = None) -> Optional[str]:
+    """
+    Pull the resource_metadata URL from a WWW-Authenticate header if present.
+    If not found and server_url is provided, try to construct resource metadata URLs.
+    """
+    if not www_authenticate and not server_url:
+        return None
+
+    # RFC9728 returns `resource_metadata="<url>"` inside the header value
+    if www_authenticate:
+        match = re.search(r'resource_metadata\s*=\s*\"?([^\", ]+)\"?', www_authenticate)
+        if match:
+            return match.group(1)
+    
+    # For servers that don't provide resource_metadata in WWW-Authenticate,
+    # we'll return None and rely on inferring authorization servers from the realm
+    # or using well-known OAuth discovery endpoints directly
+    return None
+
+
+def fetch_oauth_authorization_server_metadata(base_url: str, timeout: int = 10) -> Optional[Dict[str, Any]]:
+    """
+    Fetch OAuth authorization server metadata from well-known endpoints.
+    Tries both oauth-authorization-server and openid-configuration discovery endpoints.
+    """
+    discovery_endpoints = [
+        f"{base_url}/.well-known/oauth-authorization-server",
+        f"{base_url}/.well-known/openid-configuration",
+    ]
+    
+    for endpoint in discovery_endpoints:
+        try:
+            resp = requests.get(endpoint, timeout=timeout)
+            if resp.status_code == 200:
+                return resp.json()
+        except Exception as exc:
+            logger.debug(f"Failed to fetch OAuth metadata from {endpoint}: {exc}")
+            continue
+    
+    return None
+
+
+def infer_authorization_servers_from_realm(www_authenticate: Optional[str], server_url: str) -> Optional[list]:
+    """
+    Infer authorization server URLs from WWW-Authenticate realm or server URL.
+    This is used when the server doesn't provide resource_metadata endpoint.
+    """
+    if not www_authenticate and not server_url:
+        return None
+    
+    authorization_servers = []
+    
+    # Try to extract realm from WWW-Authenticate header
+    realm = None
+    if www_authenticate:
+        realm_match = re.search(r'realm\s*=\s*\"([^\"]+)\"', www_authenticate)
+        if realm_match:
+            realm = realm_match.group(1)
+    
+    # Parse the server URL to get base domain
+    parsed = urlparse(server_url)
+    base_url = f"{parsed.scheme}://{parsed.netloc}"
+    
+    # Return the base authorization server URL (not the discovery endpoint)
+    # The client will append .well-known paths when fetching metadata
+    authorization_servers.append(base_url)
+    
+    return authorization_servers if authorization_servers else None
+
+
+def fetch_resource_metadata(resource_metadata_url: str, timeout: int = 10) -> Optional[Dict[str, Any]]:
+    """Fetch and parse the protected resource metadata document."""
+    try:
+        resp = requests.get(resource_metadata_url, timeout=timeout)
+        resp.raise_for_status()
+        return resp.json()
+    except Exception as exc:  # broad catch – we want to surface auth requirement even if this fails
+        logger.warning("Failed to fetch resource metadata from %s: %s", resource_metadata_url, exc)
+        return None
+
+
+async def fetch_resource_metadata_async(resource_metadata_url: str, session=None, timeout: int = 10) -> Optional[Dict[str, Any]]:
+    """Async variant for fetching protected resource metadata."""
+    try:
+        import aiohttp
+
+        client_timeout = aiohttp.ClientTimeout(total=timeout)
+        if session:
+            async with session.get(resource_metadata_url, timeout=client_timeout) as resp:
+                text = await resp.text()
+        else:
+            async with aiohttp.ClientSession(timeout=client_timeout) as local_session:
+                async with local_session.get(resource_metadata_url) as resp:
+                    text = await resp.text()
+
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError:
+            logger.warning("Resource metadata at %s is not valid JSON: %s", resource_metadata_url, text[:200])
+            return None
+    except Exception as exc:
+        logger.warning("Failed to fetch resource metadata from %s: %s", resource_metadata_url, exc)
+        return None
+
+
+def canonical_resource(server_url: str) -> str:
+    """Produce a canonical resource identifier for the MCP server."""
+    parsed = urlparse(server_url)
+    # Normalize scheme/host casing per RFC guidance
+    normalized = parsed._replace(
+        scheme=parsed.scheme.lower(),
+        netloc=parsed.netloc.lower(),
+    )
+    resource = normalized.geturl()
+
+    # Prefer form without trailing slash unless path is meaningful
+    if resource.endswith("/") and parsed.path in ("", "/"):
+        resource = resource[:-1]
+    return resource
+
+
+def exchange_oauth_token(
+    token_endpoint: str,
+    code: str,
+    redirect_uri: str,
+    client_id: str,
+    client_secret: Optional[str] = None,
+    code_verifier: Optional[str] = None,
+    scope: Optional[str] = None,
+    timeout: int = 30,
+) -> Dict[str, Any]:
+    """
+    Exchange an OAuth authorization code for access tokens.
+    
+    This function performs the OAuth token exchange on the server side,
+    avoiding CORS issues that would occur if done from a browser.
+    
+    Args:
+        token_endpoint: OAuth token endpoint URL
+        code: Authorization code from OAuth provider
+        redirect_uri: Redirect URI used in authorization request
+        client_id: OAuth client ID
+        client_secret: OAuth client secret (optional for public clients)
+        code_verifier: PKCE code verifier (optional)
+        scope: OAuth scope (optional)
+        timeout: Request timeout in seconds
+        
+    Returns:
+        Token response from OAuth provider containing access_token, etc.
+        
+    Raises:
+        requests.RequestException: If the HTTP request fails
+        ValueError: If the token exchange fails
+    """
+    # Build the token request body
+    token_body = {
+        "grant_type": "authorization_code",
+        "code": code,
+        "redirect_uri": redirect_uri,
+        "client_id": client_id,
+    }
+    
+    if client_secret:
+        token_body["client_secret"] = client_secret
+    if code_verifier:
+        token_body["code_verifier"] = code_verifier
+    if scope:
+        token_body["scope"] = scope
+
+    logger.info(f"MCP OAuth: exchanging code at {token_endpoint}")
+    
+    # Make the token exchange request
+    response = requests.post(
+        token_endpoint,
+        data=token_body,
+        headers={
+            "Content-Type": "application/x-www-form-urlencoded",
+            "Accept": "application/json",
+        },
+        timeout=timeout
+    )
+    
+    # Try to parse as JSON
+    try:
+        token_data = response.json()
+    except Exception:
+        # Some providers return URL-encoded response
+        from urllib.parse import parse_qs
+        token_data = {k: v[0] if len(v) == 1 else v 
+                     for k, v in parse_qs(response.text).items()}
+    
+    if response.ok:
+        logger.info("MCP OAuth: token exchange successful")
+        return token_data
+    else:
+        error_msg = token_data.get("error_description") or token_data.get("error") or response.text
+        logger.error(f"MCP OAuth: token exchange failed - {response.status_code}: {error_msg}")
+        raise ValueError(f"Token exchange failed: {error_msg}")
+

alita_sdk/runtime/utils/mcp_sse_client.py

@@ -0,0 +1,405 @@
+"""
+MCP SSE (Server-Sent Events) Client
+Handles persistent SSE connections for MCP servers like Atlassian
+"""
+import asyncio
+import json
+import logging
+from typing import Dict, Any, Optional, AsyncIterator
+import aiohttp
+
+logger = logging.getLogger(__name__)
+
+
+class McpSseClient:
+    """
+    Client for MCP servers using SSE (Server-Sent Events) transport.
+    
+    For Atlassian-style SSE (dual-connection model):
+    - GET request opens persistent SSE stream for receiving events
+    - POST requests send commands (return 202 Accepted immediately)
+    - Responses come via the GET stream
+    
+    This client handles:
+    - Opening persistent SSE connection via GET
+    - Sending JSON-RPC requests via POST
+    - Reading SSE event streams
+    - Matching responses to requests by ID
+    """
+    
+    def __init__(self, url: str, session_id: str, headers: Optional[Dict[str, str]] = None, timeout: int = 300):
+        """
+        Initialize SSE client.
+        
+        Args:
+            url: Base URL of the MCP SSE server
+            session_id: Client-generated UUID for session
+            headers: Additional headers (e.g., Authorization)
+            timeout: Request timeout in seconds
+        """
+        self.url = url
+        self.session_id = session_id
+        self.headers = headers or {}
+        self.timeout = timeout
+        self.url_with_session = f"{url}?sessionId={session_id}"
+        self._stream_task = None
+        self._pending_requests = {}  # request_id -> asyncio.Future
+        self._stream_session = None
+        self._stream_response = None
+        self._endpoint_ready = asyncio.Event()  # Signal when endpoint is received
+        
+        logger.info(f"[MCP SSE Client] Initialized for {url} with session {session_id}")
+    
+    async def _ensure_stream_connected(self):
+        """Ensure the GET stream is connected and reading events."""
+        if self._stream_task is None or self._stream_task.done():
+            logger.info(f"[MCP SSE Client] Opening persistent SSE stream...")
+            self._stream_session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=None))
+            
+            headers = {
+                "Accept": "text/event-stream",
+                **self.headers
+            }
+            
+            self._stream_response = await self._stream_session.get(self.url_with_session, headers=headers)
+            
+            logger.info(f"[MCP SSE Client] Stream opened: status={self._stream_response.status}")
+            
+            # Handle 401 Unauthorized - need OAuth
+            if self._stream_response.status == 401:
+                from ..utils.mcp_oauth import (
+                    McpAuthorizationRequired,
+                    canonical_resource,
+                    extract_resource_metadata_url,
+                    fetch_resource_metadata_async,
+                    infer_authorization_servers_from_realm,
+                    fetch_oauth_authorization_server_metadata
+                )
+                
+                auth_header = self._stream_response.headers.get('WWW-Authenticate', '')
+                resource_metadata_url = extract_resource_metadata_url(auth_header, self.url)
+                
+                metadata = None
+                if resource_metadata_url:
+                    metadata = await fetch_resource_metadata_async(
+                        resource_metadata_url,
+                        session=self._stream_session,
+                        timeout=30
+                    )
+                
+                # Infer authorization servers if not in metadata
+                if not metadata or not metadata.get('authorization_servers'):
+                    inferred_servers = infer_authorization_servers_from_realm(auth_header, self.url)
+                    if inferred_servers:
+                        if not metadata:
+                            metadata = {}
+                        metadata['authorization_servers'] = inferred_servers
+                        logger.info(f"[MCP SSE Client] Inferred authorization servers: {inferred_servers}")
+                        
+                        # Fetch OAuth metadata
+                        auth_server_metadata = fetch_oauth_authorization_server_metadata(inferred_servers[0], timeout=30)
+                        if auth_server_metadata:
+                            metadata['oauth_authorization_server'] = auth_server_metadata
+                            logger.info(f"[MCP SSE Client] Fetched OAuth metadata")
+                
+                raise McpAuthorizationRequired(
+                    message=f"MCP server {self.url} requires OAuth authorization",
+                    server_url=canonical_resource(self.url),
+                    resource_metadata_url=resource_metadata_url,
+                    www_authenticate=auth_header,
+                    resource_metadata=metadata,
+                    status=self._stream_response.status,
+                    tool_name=self.url,
+                )
+            
+            if self._stream_response.status != 200:
+                error_text = await self._stream_response.text()
+                raise Exception(f"Failed to open SSE stream: HTTP {self._stream_response.status}: {error_text}")
+            
+            # Start background task to read stream
+            self._stream_task = asyncio.create_task(self._read_stream())
+    
+    async def _read_stream(self):
+        """Background task that continuously reads the SSE stream."""
+        logger.info(f"[MCP SSE Client] Starting stream reader...")
+        
+        try:
+            buffer = ""
+            current_event = {}
+            
+            async for chunk in self._stream_response.content.iter_chunked(1024):
+                chunk_str = chunk.decode('utf-8')
+                buffer += chunk_str
+                
+                # Process complete lines
+                while '\n' in buffer:
+                    line, buffer = buffer.split('\n', 1)
+                    line_str = line.strip()
+                    
+                    # Empty line indicates end of event
+                    if not line_str:
+                        if current_event and 'data' in current_event:
+                            self._process_event(current_event)
+                            current_event = {}
+                        continue
+                    
+                    # Parse SSE fields
+                    if line_str.startswith('event:'):
+                        current_event['event'] = line_str[6:].strip()
+                    elif line_str.startswith('data:'):
+                        data_str = line_str[5:].strip()
+                        current_event['data'] = data_str
+                    elif line_str.startswith('id:'):
+                        current_event['id'] = line_str[3:].strip()
+                    
+        except Exception as e:
+            logger.error(f"[MCP SSE Client] Stream reader error: {e}")
+            # Fail all pending requests
+            for future in self._pending_requests.values():
+                if not future.done():
+                    future.set_exception(e)
+        finally:
+            logger.info(f"[MCP SSE Client] Stream reader stopped")
+    
+    def _process_event(self, event: Dict[str, str]):
+        """Process a complete SSE event."""
+        event_type = event.get('event', 'message')
+        data_str = event.get('data', '')
+        
+        # Handle 'endpoint' event - server provides the actual session URL to use
+        if event_type == 'endpoint':
+            # Extract session ID from endpoint URL
+            # Format: /v1/sse?sessionId=<uuid>
+            if 'sessionId=' in data_str:
+                new_session_id = data_str.split('sessionId=')[1].split('&')[0]
+                logger.info(f"[MCP SSE Client] Server provided session ID: {new_session_id}")
+                self.session_id = new_session_id
+                self.url_with_session = f"{self.url}?sessionId={new_session_id}"
+                self._endpoint_ready.set()  # Signal that we can now send requests
+            return
+        
+        # Skip other non-message events
+        if event_type != 'message' and not data_str.startswith('{'):
+            return
+        
+        if not data_str:
+            return
+        
+        try:
+            data = json.loads(data_str)
+            request_id = data.get('id')
+            
+            logger.debug(f"[MCP SSE Client] Received response for request {request_id}")
+            
+            # Resolve pending request
+            if request_id and request_id in self._pending_requests:
+                future = self._pending_requests.pop(request_id)
+                if not future.done():
+                    future.set_result(data)
+                
+        except json.JSONDecodeError as e:
+            logger.warning(f"[MCP SSE Client] Failed to parse SSE data: {e}, data: {repr(data_str)[:200]}")
+                            
+        except Exception as e:
+            logger.error(f"[MCP SSE Client] Stream reader error: {e}")
+            # Fail all pending requests
+            for future in self._pending_requests.values():
+                if not future.done():
+                    future.set_exception(e)
+        finally:
+            logger.info(f"[MCP SSE Client] Stream reader stopped")
+    
+    async def send_request(self, method: str, params: Optional[Dict[str, Any]] = None, request_id: Optional[str] = None) -> Dict[str, Any]:
+        """
+        Send a JSON-RPC request and wait for response via SSE stream.
+        
+        Uses dual-connection model:
+        1. GET stream is kept open to receive responses
+        2. POST request sends the command (returns 202 immediately)
+        3. Response comes via the GET stream
+        
+        Args:
+            method: JSON-RPC method name (e.g., "tools/list", "tools/call")
+            params: Method parameters
+            request_id: Optional request ID (auto-generated if not provided)
+            
+        Returns:
+            Parsed JSON-RPC response
+            
+        Raises:
+            Exception: If request fails or times out
+        """
+        import time
+        if request_id is None:
+            request_id = f"{method.replace('/', '_')}_{int(time.time() * 1000)}"
+        
+        request = {
+            "jsonrpc": "2.0",
+            "id": request_id,
+            "method": method,
+            "params": params or {}
+        }
+        
+        logger.debug(f"[MCP SSE Client] Sending request: {method} (id={request_id})")
+        
+        # Ensure stream is connected
+        await self._ensure_stream_connected()
+        
+        # Wait for endpoint event (server provides the actual session ID to use)
+        await asyncio.wait_for(self._endpoint_ready.wait(), timeout=10)
+        
+        # Create future for this request
+        future = asyncio.Future()
+        self._pending_requests[request_id] = future
+        
+        # Send POST request
+        headers = {
+            "Content-Type": "application/json",
+            **self.headers
+        }
+        
+        timeout = aiohttp.ClientTimeout(total=30)
+        
+        try:
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+                async with session.post(self.url_with_session, json=request, headers=headers) as response:
+                    if response.status == 404:
+                        error_text = await response.text()
+                        raise Exception(f"HTTP 404: {error_text}")
+                    
+                    # 202 is expected - response will come via stream
+                    if response.status not in [200, 202]:
+                        error_text = await response.text()
+                        raise Exception(f"HTTP {response.status}: {error_text}")
+            
+            # Wait for response from stream (with timeout)
+            result = await asyncio.wait_for(future, timeout=self.timeout)
+            
+            # Check for JSON-RPC error
+            if 'error' in result:
+                error = result['error']
+                raise Exception(f"MCP Error: {error.get('message', str(error))}")
+            
+            return result
+                    
+        except asyncio.TimeoutError:
+            self._pending_requests.pop(request_id, None)
+            logger.error(f"[MCP SSE Client] Request timeout after {self.timeout}s")
+            raise Exception(f"SSE request timeout after {self.timeout}s")
+        except Exception as e:
+            self._pending_requests.pop(request_id, None)
+            logger.error(f"[MCP SSE Client] Request failed: {e}")
+            raise
+    
+    async def close(self):
+        """Close the persistent SSE stream."""
+        logger.info(f"[MCP SSE Client] Closing connection...")
+        
+        # Cancel background stream reader task
+        if self._stream_task and not self._stream_task.done():
+            self._stream_task.cancel()
+            try:
+                await self._stream_task
+            except (asyncio.CancelledError, Exception) as e:
+                logger.debug(f"[MCP SSE Client] Stream task cleanup: {e}")
+        
+        # Close response stream
+        if self._stream_response and not self._stream_response.closed:
+            try:
+                self._stream_response.close()
+            except Exception as e:
+                logger.debug(f"[MCP SSE Client] Response close error: {e}")
+        
+        # Close session
+        if self._stream_session and not self._stream_session.closed:
+            try:
+                await self._stream_session.close()
+                # Give aiohttp time to cleanup
+                await asyncio.sleep(0.1)
+            except Exception as e:
+                logger.debug(f"[MCP SSE Client] Session close error: {e}")
+        
+        logger.info(f"[MCP SSE Client] Connection closed")
+    
+    async def __aenter__(self):
+        """Async context manager entry."""
+        return self
+    
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.close()
+    
+    async def initialize(self) -> Dict[str, Any]:
+        """
+        Send initialize request to establish MCP protocol session.
+        
+        Returns:
+            Server capabilities and info
+        """
+        response = await self.send_request(
+            method="initialize",
+            params={
+                "protocolVersion": "2024-11-05",
+                "capabilities": {
+                    "roots": {"listChanged": True},
+                    "sampling": {}
+                },
+                "clientInfo": {
+                    "name": "Alita MCP Client",
+                    "version": "1.0.0"
+                }
+            }
+        )
+        
+        logger.info(f"[MCP SSE Client] MCP session initialized")
+        return response.get('result', {})
+    
+    async def list_tools(self) -> list:
+        """
+        Discover available tools from the MCP server.
+        
+        Returns:
+            List of tool definitions
+        """
+        response = await self.send_request(method="tools/list")
+        result = response.get('result', {})
+        tools = result.get('tools', [])
+        
+        logger.info(f"[MCP SSE Client] Discovered {len(tools)} tools")
+        return tools
+    
+    async def list_prompts(self) -> list:
+        """
+        Discover available prompts from the MCP server.
+        
+        Returns:
+            List of prompt definitions
+        """
+        response = await self.send_request(method="prompts/list")
+        result = response.get('result', {})
+        prompts = result.get('prompts', [])
+        
+        logger.debug(f"[MCP SSE Client] Discovered {len(prompts)} prompts")
+        return prompts
+    
+    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """
+        Execute a tool on the MCP server.
+        
+        Args:
+            tool_name: Name of the tool to call
+            arguments: Tool arguments
+            
+        Returns:
+            Tool execution result
+        """
+        response = await self.send_request(
+            method="tools/call",
+            params={
+                "name": tool_name,
+                "arguments": arguments
+            }
+        )
+        
+        result = response.get('result', {})
+        return result